import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image
from torchvision import transforms
from transformers import SegformerForSemanticSegmentation, AutoImageProcessor

# Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Load model and processor once
processor = AutoImageProcessor.from_pretrained("nvidia/segformer-b2-finetuned-ade-512-512")
model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b2-finetuned-ade-512-512").to(device)

def process(room_img, tile_img):
    room_img = room_img.convert("RGB")
    tile_img = tile_img.convert("RGB")
    room_np = np.array(room_img)

    # Segmentation
    inputs = processor(images=room_img, return_tensors="pt").to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    segmentation = outputs.logits.argmax(dim=1).squeeze().cpu().numpy()
    segmentation_resized = cv2.resize(segmentation.astype(np.uint8), (room_np.shape[1], room_np.shape[0]), interpolation=cv2.INTER_NEAREST)

    # Mask for floor (ADE20K class index 3)
    floor_class_index = 3
    mask_bin = (segmentation_resized == floor_class_index).astype(np.uint8)

    # Largest contour
    contours, _ = cv2.findContours(mask_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        return room_img, Image.fromarray(mask_bin*255), tile_img, room_img
    contour = max(contours, key=cv2.contourArea)
    if len(contour) < 4:
        return room_img, Image.fromarray(mask_bin*255), tile_img, room_img

    x, y, w, h = cv2.boundingRect(contour)
    src_pts = np.array([[x, y + h], [x + w, y + h], [x + w, y], [x, y]], dtype=np.float32)
    offset = h * 0.5
    dst_pts = np.array([[x - offset, y + h], [x + w + offset, y + h], [x + w, y], [x, y]], dtype=np.float32)
    H = cv2.getPerspectiveTransform(src_pts, dst_pts)

    # Resize tile
    target_tile_width = room_np.shape[1] // 10
    tile_aspect_ratio = tile_img.height / tile_img.width
    target_tile_height = int(target_tile_width * tile_aspect_ratio)
    tile_resized = tile_img.resize((target_tile_width, target_tile_height), Image.LANCZOS)
    tile_np = np.array(tile_resized)

    # Tile texture
    tile_h, tile_w = tile_np.shape[:2]
    room_h, room_w = room_np.shape[:2]
    reps_y = room_h // tile_h + 2
    reps_x = room_w // tile_w + 2
    tiled_texture = np.tile(tile_np, (reps_y, reps_x, 1))[:room_h, :room_w]
    warped_texture = cv2.warpPerspective(tiled_texture, H, (room_w, room_h))

    # Blend
    room_float = room_np.astype(np.float32) / 255.0
    texture_float = warped_texture.astype(np.float32) / 255.0
    room_gray = cv2.cvtColor(room_float, cv2.COLOR_RGB2GRAY)
    lighting = np.stack([room_gray]*3, axis=-1)
    lighting = np.clip(lighting * 1.2, 0, 1)
    lit_texture = np.clip(texture_float * lighting, 0, 1)
    mask_3ch = np.stack([mask_bin]*3, axis=-1)
    blended = np.where(mask_3ch == 1, lit_texture, room_float)
    blended_img = (blended * 255).astype(np.uint8)

    return Image.fromarray(room_np), Image.fromarray(mask_bin * 255), Image.fromarray(warped_texture), Image.fromarray(blended_img)


demo = gr.Interface(
    fn=process,
    inputs=[gr.Image(label="Room Image", type="pil"), gr.Image(label="Tile Image", type="pil")],
    outputs=[
        gr.Image(label="Original Room"),
        gr.Image(label="Floor Mask"),
        gr.Image(label="Warped Texture"),
        gr.Image(label="Final Overlay")
    ],
    examples=[
        ["https://www.thespruce.com/thmb/GtlHim5EsWERYoVi62TnWpu6JTA=/5472x3648/filters:fill(auto,1)/GettyImages-9261821821-5c69c1b7c9e77c0001675a49.jpg", "https://renovlange.de/images/marbles/arabescato.jpg"],
        ["https://st.hzcdn.com/simgs/57717d160282e776_9-0651/home-design.jpg", "https://renovlange.de/images/marbles/grigio-cenere.jpg"]
    ],
    title="Room Floor Tiler",
    description="Upload a room image and a tile texture. The floor is automatically detected and overlaid with your selected tile using SegFormer and perspective warping."
)

if __name__ == "__main__":
    demo.launch()