Update app.py

app.py

@@ -14,36 +14,29 @@ state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}  # Rem
 model.load_state_dict(state_dict)
 model.eval()
 
-def remove_background(image_np):
-    """Removes background using U²-Net and replaces it with white."""
-    
-    transform_pipeline = transforms.Compose([
-        transforms.ToTensor(),
-        transforms.Resize((320, 320))
-    ])
-    
-    image = Image.fromarray(image_np).convert("RGB")
-    input_tensor = transform_pipeline(image).unsqueeze(0)
+def detect_design(image_np):
+    """Detects if a design exists on the dress using edge detection & clustering."""
+    gray = cv2.cvtColor(image_np, cv2.COLOR_RGB2GRAY)
+    edges = cv2.Canny(gray, 50, 150)
 
-    with torch.no_grad():
-        output = model(input_tensor)[0][0].squeeze().cpu().numpy()
-
-    mask = (output > 0.5).astype(np.uint8) * 255
-    mask = cv2.resize(mask, (image_np.shape[1], image_np.shape[0]), interpolation=cv2.INTER_NEAREST)
+    # Dilation to highlight patterns
+    kernel = np.ones((3, 3), np.uint8)
+    edges = cv2.dilate(edges, kernel, iterations=1)
 
-    white_bg = np.ones_like(image_np) * 255  # White background
-    segmented_image = np.where(mask[..., None] > 128, image_np, white_bg)
+    # Count edge density
+    design_ratio = np.sum(edges > 0) / (image_np.shape[0] * image_np.shape[1])
 
-    return segmented_image, mask
+    return design_ratio > 0.02, edges  # If edge density is high, assume a design is present
 
 def segment_dress(image_np):
-    """Segments the dress using K-means and refines with U²-Net."""
+    """Segment the dress using U²-Net & refine with Lab color space."""
     
-    # Convert to Lab color space
+    # Convert to Lab space
     img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
-    pixel_values = img_lab.reshape((-1, 3)).astype(np.float32)
+    L, A, B = cv2.split(img_lab)
 
-    # K-means clustering to detect dress region
+    # 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])
@@ -55,7 +48,7 @@ def segment_dress(image_np):
     # Create dress mask
     mask = (labels == dress_label).astype(np.uint8) * 255
 
-    # Refine with U²-Net prediction
+    # Use U²-Net prediction to refine segmentation
     transform_pipeline = transforms.Compose([
         transforms.ToTensor(),
         transforms.Resize((320, 320))
@@ -70,7 +63,7 @@ def segment_dress(image_np):
     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 masks
+    # Combine K-means and U²-Net masks
     refined_mask = cv2.bitwise_and(mask, u2net_mask)
 
     # Morphological operations for smoothness
@@ -80,11 +73,15 @@ def segment_dress(image_np):
     
     return refined_mask
 
-def recolor_dress(image_np, mask, target_color):
-    """Change dress color while keeping the white background intact."""
+def recolor_dress(image_np, mask, target_color, edges):
+    """Change dress color while preserving texture, shadows, and designs."""
     
     img_lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
     target_color_lab = cv2.cvtColor(np.uint8([[target_color]]), cv2.COLOR_BGR2LAB)[0][0]
+    
+    # Exclude design from recoloring
+    design_mask = (edges > 0).astype(np.uint8) * 255
+    mask = cv2.bitwise_and(mask, cv2.bitwise_not(design_mask))
 
     # Preserve lightness (L) and change only chromatic channels (A & B)
     blend_factor = 0.7
@@ -92,27 +89,26 @@ def recolor_dress(image_np, mask, target_color):
     img_lab[..., 2] = np.where(mask > 128, 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
 
-def process_image(image_path, color):
-    """Remove background, segment dress, and recolor while keeping background white."""
+def change_dress_color(image_path, color):
+    """Change the dress color naturally while keeping designs intact."""
     if image_path is None:
         return None
 
     img = Image.open(image_path).convert("RGB")
     img_np = np.array(img)
 
-    # Remove background
-    img_segmented, _ = remove_background(img_np)
+    # Detect if a design is present
+    design_present, edges = detect_design(img_np)
 
     # Get dress segmentation mask
     mask = segment_dress(img_np)
     
     if mask is None:
-        return Image.fromarray(img_segmented)  # No dress detected, return only background removal
-
-    # Convert selected color to BGR
+        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),
@@ -120,21 +116,26 @@ def process_image(image_path, color):
     }
     new_color_bgr = np.array(color_map.get(color, (0, 0, 255)), dtype=np.uint8)  # Default to Red
     
-    # Apply recoloring
-    img_recolored = recolor_dress(img_segmented, mask, new_color_bgr)
+    # Apply recoloring logic
+    if design_present:
+        print("Design detected! Coloring only non-design areas.")
+        img_recolored = recolor_dress(img_np, mask, new_color_bgr, edges)
+    else:
+        print("No design detected. Coloring entire dress.")
+        img_recolored = recolor_dress(img_np, mask, new_color_bgr, np.zeros_like(mask))  # No design mask
 
     return Image.fromarray(img_recolored)
 
 # Gradio Interface
 demo = gr.Interface(
-    fn=process_image,
+    fn=change_dress_color,
     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")
     ],
-    outputs=gr.Image(type="pil", label="Final Dress Image"),
-    title="Dress Color Changer with Background Removal",
-    description="Upload an image of a dress, remove its background, and recolor it naturally while keeping the background white."
+    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 designs."
 )
 
 if __name__ == "__main__":