Update app.py

app.py

@@ -97,43 +97,57 @@ def change_dress_color(img, color):
     if img is None:
         return None
 
-    img_np = np.array(img)
+    # Convert color name to BGR using a safer method
+    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)
+    }
+    
+    # Safely get color with fallback to red
+    new_color_bgr = color_map.get(color, (0, 0, 255))
 
+    img_np = np.array(img)
+    
     # Get dress segmentation mask
     dress_mask = segment_dress(img_np)
-
-    if dress_mask is None:
-        return img  # No dress detected
+    
+    if dress_mask is None or np.sum(dress_mask) == 0:
+        return img  # Return original if no mask found
 
     # Further refine mask with GrabCut
     dress_mask = apply_grabcut(img_np, dress_mask)
-
-    # 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)
-
+    
     # Apply recoloring with blending
     img_recolored = recolor_dress(img_np, dress_mask, new_color_bgr)
 
     return Image.fromarray(img_recolored)
 
-# Gradio Interface
-demo = gr.Interface(
-    fn=change_dress_color,
-    inputs=[
-        gr.Image(type = "pil", label="Upload Dress Image"),
-        gr.Dropdown(choices=["Red", "Blue", "Green", "Yellow", "Purple", "Orange", "Cyan", "Magenta", "White", "Black"], 
-            label="Choose New Dress Color",
-            value="Red") 
-    ],
-    outputs=gr.Image(type="pil", label="Color Changed Dress"),
-    title="AI-Powered Dress Color Changer",
-    description="Upload an image of a dress and select a new color. The AI will change the dress color naturally while keeping the fabric texture."
-)
+# Create Gradio Blocks interface instead of simple Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# AI-Powered Dress Color Changer")
+    gr.Markdown("Upload an image of a dress and select a new color. The AI will change the dress color naturally while keeping the fabric texture.")
+    
+    with gr.Row():
+        with gr.Column():
+            input_image = gr.Image(type="pil", label="Upload Dress Image")
+            color_choice = gr.Dropdown(
+                choices=["Red", "Blue", "Green", "Yellow", "Purple", 
+                        "Orange", "Cyan", "Magenta", "White", "Black"],
+                value="Red",
+                label="Choose New Dress Color"
+            )
+            submit_btn = gr.Button("Change Color")
+        
+        with gr.Column():
+            output_image = gr.Image(type="pil", label="Result")
+
+    submit_btn.click(
+        fn=change_dress_color,
+        inputs=[input_image, color_choice],
+        outputs=output_image
+    )
 
 if __name__ == "__main__":
     demo.launch()