import cv2
import numpy as np
import torch
import gradio as gr
from segment_anything import sam_model_registry, SamPredictor
from PIL import Image
from sklearn.cluster import KMeans

# Load SAM model
sam_checkpoint = "sam_vit_h_4b8939.pth"
model_type = "vit_h"
device = "cuda" if torch.cuda.is_available() else "cpu"
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)

def extract_dominant_color(img):
    pixels = img.reshape(-1, 3)
    kmeans = KMeans(n_clusters=1, n_init=10)
    kmeans.fit(pixels)
    return kmeans.cluster_centers_[0]

def create_circular_swatch(color, size=200):
    swatch = np.zeros((size, size, 3), dtype=np.uint8)
    color_tuple = tuple(map(int, color))
    cv2.circle(swatch, (size//2, size//2), size//2, color_tuple, -1)
    return swatch

def find_clean_area(img, mask, min_area_size=100*100):
    # Apply mask to image
    masked_img = cv2.bitwise_and(img, img, mask=mask)
    
    gray = cv2.cvtColor(masked_img, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, 30, 100)
    
    kernel = np.ones((5,5), np.uint8)
    dilated_edges = cv2.dilate(edges, kernel, iterations=1)
    
    clean_mask = (dilated_edges == 0).astype(np.uint8) & mask
    
    contours, _ = cv2.findContours(clean_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    valid_contours = [c for c in contours if cv2.contourArea(c) >= min_area_size]
    
    if not valid_contours:
        return None
    
    largest_contour = max(valid_contours, key=cv2.contourArea)
    x, y, w, h = cv2.boundingRect(largest_contour)
    
    return (x, y, w, h)

def process_image(image, input_point):
    predictor.set_image(image)
    input_label = np.array([1])
    masks, _, _ = predictor.predict(
        point_coords=input_point,
        point_labels=input_label,
        multimask_output=False,
    )
    mask = masks[0]

    # Extract dominant color and create swatch
    masked_image = cv2.bitwise_and(image, image, mask=mask.astype(np.uint8))
    dominant_color = extract_dominant_color(masked_image)
    swatch = create_circular_swatch(dominant_color)

    # Find clean area and crop
    clean_area = find_clean_area(image, mask.astype(np.uint8))
    if clean_area is not None:
        x, y, w, h = clean_area
        detailed_fabric = image[y:y+h, x:x+w]
        detailed_fabric = cv2.resize(detailed_fabric, (400, 600))
    else:
        detailed_fabric = np.zeros((400, 600, 3), dtype=np.uint8)

    # Visualize mask on image
    visualization = image.copy()
    visualization[mask] = visualization[mask] * 0.5 + np.array([0, 0, 255]) * 0.5

    return visualization, swatch, detailed_fabric

def gradio_interface(input_image, click_x, click_y):
    input_point = np.array([[click_x, click_y]])
    visualization, swatch, detailed_fabric = process_image(input_image, input_point)
    return visualization, swatch, detailed_fabric

# Create Gradio interface
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Image(type="numpy"),
        gr.Slider(0, 1000, label="Click X"),
        gr.Slider(0, 1000, label="Click Y")
    ],
    outputs=[
        gr.Image(type="numpy", label="Segmentation"),
        gr.Image(type="numpy", label="Color Swatch"),
        gr.Image(type="numpy", label="Detailed Fabric")
    ],
    title="Fabric Analyzer",
    description="Upload an image or choose from examples. Click on the garment to analyze fabric and color.",
    examples=[
        ["blue_shirt.png", 400, 500],
        ["polo.png", 400, 500],
        ["dress.jpg", 400, 500]
    ]
)

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