Update app.py

app.py

@@ -1,122 +1,96 @@
 import gradio as gr
 import numpy as np
-import torch
 import cv2
 from PIL import Image
-from torchvision import transforms
-from cloth_segmentation.networks.u2net import U2NET  # Import U²-Net
+import torch
+import torch.nn as nn
+import torchvision.transforms as T
+from skimage import color
+from sklearn.cluster import KMeans
+from cloth_segmentation.networks.u2net import U2NET
 
-# Load U²-Net model
-model_path = "cloth_segmentation/networks/u2net.pth"
 model = U2NET(3, 1)
-state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}  # Remove 'module.' prefix
-model.load_state_dict(state_dict)
+model.load_state_dict(torch.load("u2net.pth", map_location=torch.device('cpu')))
 model.eval()
 
-def segment_dress(image_np):
-    """Segment the dress using U²-Net & refine with Lab color space."""
-    
-    # Convert to Lab space
-    img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
-    L, A, B = cv2.split(img_lab)
-
-    # Use K-means clustering to detect dominant dress region
-    pixel_values = img_lab.reshape((-1, 3)).astype(np.float32)
-    k = 3  # Three clusters: background, skin, dress
-    _, labels, centers = cv2.kmeans(pixel_values, k, None, (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0), 10, cv2.KMEANS_RANDOM_CENTERS)
-    labels = labels.reshape(image_np.shape[:2])
-    
-    # Assume dress is the largest non-background cluster
-    unique_labels, counts = np.unique(labels, return_counts=True)
-    dress_label = unique_labels[np.argmax(counts[1:]) + 1]  # Avoid background
-    
-    # Create dress mask
-    mask = (labels == dress_label).astype(np.uint8) * 255
+# Preprocessing
+transform = T.Compose([
+    T.Resize((320, 320)),
+    T.ToTensor(),
+    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+])
 
-    # Use U²-Net prediction to refine segmentation
-    transform_pipeline = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Resize((320, 320))
-    ])
-    
-    image = Image.fromarray(image_np).convert("RGB")
-    input_tensor = transform_pipeline(image).unsqueeze(0)
+# Segmentation mask
+@torch.no_grad()
+def get_dress_mask(image_pil):
+    img = transform(image_pil).unsqueeze(0)
+    pred = model(img)[0]
+    pred = pred.squeeze().cpu().numpy()
+    mask = (pred > 0.5).astype(np.uint8)
+    mask = cv2.resize(mask, image_pil.size[::-1])
+    return mask
 
-    with torch.no_grad():
-        output = model(input_tensor)[0][0].squeeze().cpu().numpy()
+# Color parsing (extract target color from prompt)
+def extract_target_color(prompt):
+    # Basic keyword matching (can be replaced with NLP-based color detection)
+    import re
+    colors = ["red", "blue", "green", "yellow", "pink", "black", "white", "sky blue", "purple"]
+    for c in colors:
+        if re.search(c, prompt.lower()):
+            return c
+    return "red"  # default fallback
 
-    u2net_mask = (output > 0.5).astype(np.uint8) * 255
-    u2net_mask = cv2.resize(u2net_mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST)
-    
-    # Combine K-means and U²-Net masks
-    refined_mask = cv2.bitwise_and(mask, u2net_mask)
+# Recoloring function
+def recolor_dress(image_pil, prompt):
+    image_np = np.array(image_pil.convert("RGB")) / 255.0
+    lab = color.rgb2lab(image_np)
 
-    return refined_mask
+    mask = get_dress_mask(image_pil)
 
-def detect_design(image_np, dress_mask):
-    """Detect the design part of the dress and separate it from fabric."""
-    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
-    edges = cv2.Canny(gray, 50, 150)
-    
-    # Expand detected edges to mask the design area
-    kernel = np.ones((5, 5), np.uint8)
-    design_mask = cv2.dilate(edges, kernel, iterations=2)
-    
-    # Keep only the design within the dress area
-    design_mask = cv2.bitwise_and(design_mask, dress_mask)
-    return design_mask
+    # Get mean a, b values in masked region
+    a_mean = lab[:, :, 1][mask == 1].mean()
+    b_mean = lab[:, :, 2][mask == 1].mean()
 
-def recolor_dress(image_np, mask, design_mask, target_color):
-    """Change dress color while preserving texture and design."""
-    img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
-    target_color_lab = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2LAB)[0][0]
-    
-    # Preserve lightness (L) and change only chromatic channels (A & B)
-    blend_factor = 0.7
-    img_lab[..., 1] = np.where((mask > 128) & (design_mask == 0), img_lab[..., 1] * (1 - blend_factor) + target_color_lab[1] * blend_factor, img_lab[..., 1])
-    img_lab[..., 2] = np.where((mask > 128) & (design_mask == 0), img_lab[..., 2] * (1 - blend_factor) + target_color_lab[2] * blend_factor, img_lab[..., 2])
-
-    img_recolored = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
-    return img_recolored
+    # Target a, b (from a small predefined palette)
+    target_color_map = {
+        "red": [60, 40],
+        "blue": [20, -60],
+        "green": [-60, 60],
+        "yellow": [10, 70],
+        "pink": [50, 10],
+        "purple": [40, -40],
+        "black": [0, 0],
+        "white": [0, 0],
+        "sky blue": [0, -50],
+    }
+    target = extract_target_color(prompt)
+    target_a, target_b = target_color_map.get(target, [60, 40])
 
-def change_dress_color(image_path, color):
-    """Change the dress color naturally while keeping textures and design."""
-    if image_path is None:
-        return None
+    # Apply color shift only to dress region
+    lab_new = lab.copy()
+    delta_a = target_a - a_mean
+    delta_b = target_b - b_mean
+    lab_new[:, :, 1][mask == 1] += delta_a
+    lab_new[:, :, 2][mask == 1] += delta_b
 
-    img = Image.open(image_path).convert("RGB")
-    img_np = np.array(img)
-    dress_mask = segment_dress(img_np)
-    design_mask = detect_design(img_np, dress_mask)
-    
-    if dress_mask is None:
-        return img  # No dress detected
-    
-    # Convert the selected color to BGR
-    color_map = {
-        "Red": (0, 0, 255), "Blue": (255, 0, 0), "Green": (0, 255, 0), "Yellow": (0, 255, 255),
-        "Purple": (128, 0, 128), "Orange": (0, 165, 255), "Cyan": (255, 255, 0), "Magenta": (255, 0, 255),
-        "White": (255, 255, 255), "Black": (0, 0, 0)
-    }
-    new_color_bgr = np.array(color_map.get(color, (0, 0, 255)), dtype=np.uint8)  # Default to Red
-    
-    # Recolor the dress naturally
-    img_recolored = recolor_dress(img_np, dress_mask, design_mask, new_color_bgr)
+    rgb_new = color.lab2rgb(lab_new)
+    rgb_new = (rgb_new * 255).astype(np.uint8)
+    return Image.fromarray(rgb_new)
 
-    return Image.fromarray(img_recolored)
+# Gradio UI
+def interface_fn(image, prompt):
+    return recolor_dress(image, prompt)
 
-# Gradio Interface
-demo = gr.Interface(
-    fn=change_dress_color,
+interface = gr.Interface(
+    fn=interface_fn,
     inputs=[
-        gr.Image(type="filepath", label="Upload Dress Image"),
-        gr.Radio(["Red", "Blue", "Green", "Yellow", "Purple", "Orange", "Cyan", "Magenta", "White", "Black"], label="Choose New Dress Color")
+        gr.Image(label="Upload Image", type="pil"),
+        gr.Textbox(label="Describe Dress Color Change", placeholder="e.g. Make the dress sky blue")
     ],
-    outputs=gr.Image(type="pil", label="Color Changed Dress"),
-    title="Dress Color Changer",
-    description="Upload an image of a dress and select a new color to change its appearance naturally while preserving the design."
+    outputs=gr.Image(label="Recolored Dress Output"),
+    title="Natural Dress Recoloring with Design Preservation (No GPU)",
+    description="Upload a fashion photo and use prompts like “Make the dress blue” or “Change to red silk”. This CPU-based app changes the color of dresses while preserving textures, embroidery, and shadows.",
+    allow_flagging="never"
 )
 
-if __name__ == "__main__":
-    demo.launch()
+interface.launch()