import gradio as gr
import numpy as np
from PIL import Image
import torch
import cv2
from torchvision import transforms
from transformers import AutoModel, AutoProcessor  # Import Hugging Face model

# Load pre-trained U2-Net from Hugging Face
model_name = "netradrishti/u2net-saliency"
model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
processor = AutoProcessor.from_pretrained(model_name)
model.eval()

def get_dress_mask(image_np):
    """Segment the dress using the Hugging Face U2-Net model."""
    
    # Convert NumPy image to PIL image
    image = Image.fromarray(image_np).convert("RGB")

    # Preprocess image using the Hugging Face processor
    inputs = processor(images=image, return_tensors="pt")

    with torch.no_grad():
        output = model(**inputs).logits  # Get model output

    # Convert the predicted mask into a binary mask
    mask = output[0, 0].squeeze().cpu().numpy()
    mask = (mask > 0.5).astype(np.uint8)  # Convert to binary mask

    return mask

def detect_dress_and_warp(hoodie_img, design_img, warp_strength, scale):
    if hoodie_img is None or design_img is None:
        return None
    
    hoodie = Image.open(hoodie_img).convert("RGBA")
    design = Image.open(design_img).convert("RGBA")
    hoodie_np = np.array(hoodie)
    
    dress_mask = get_dress_mask(hoodie_np)
    if dress_mask is None:
        return hoodie  # No mask, return original
    
    # Find bounding box of the mask
    y, x = np.where(dress_mask > 0)
    if len(x) == 0 or len(y) == 0:
        return hoodie  # Empty mask, return original

    x_min, x_max, y_min, y_max = x.min(), x.max(), y.min(), y.max()
    center_x, center_y = (x_min + x_max) // 2, (y_min + y_max) // 2
    w, h = x_max - x_min, y_max - y_min
    
    # Scale the design to match the bounding box
    new_w, new_h = int(w * scale), int(h * scale)
    design = design.resize((new_w, new_h))
    design_np = np.array(design)
    
    # Warp the design using TPS
    from tps_warp import apply_tps_warp
    warped_design = apply_tps_warp(design_np, dress_mask, warp_strength)
    design = Image.fromarray(warped_design)
    
    # Compute position to paste the design
    paste_x = center_x - new_w // 2
    paste_y = center_y - new_h // 2
    
    # Create a temporary layer and paste the warped design onto it
    temp_layer = Image.new("RGBA", hoodie.size, (0, 0, 0, 0))
    temp_layer.paste(design, (paste_x, paste_y), design)
    
    # Combine the hoodie with the design layer
    final_image = Image.alpha_composite(hoodie, temp_layer)
    return final_image

def generate_mockup(hoodie, design, warp_strength=5.0, scale=1.0):
    return detect_dress_and_warp(hoodie, design, warp_strength, scale)

demo = gr.Interface(
    fn=generate_mockup,
    inputs=[
        gr.Image(type="filepath", label="Upload Hoodie Mockup"),
        gr.Image(type="filepath", label="Upload Your Design"),
        gr.Slider(0.0, 20.0, value=5.0, label="Warp Strength"),
        gr.Slider(0.5, 2.0, value=1.0, label="Design Scale")
    ],
    outputs=gr.Image(type="pil", label="Generated Mockup"),
    title="Rookus Hoodie Mockup Generator",
    description="Upload a hoodie mockup and your design to generate a realistic preview with automatic dress detection and warping using TPS transformation."
)

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