Update app.py

app.py

@@ -26,21 +26,23 @@ def segment_dress(image):
 def warp_design(design, mask, warp_scale):
     """Warp the design using TPS and scale control."""
     h, w = mask.shape[:2]
+    
+    # Resize design to match mask dimensions
     design_resized = cv2.resize(design, (w, h))
 
-    # Normalize mask, convert to uint8
-    scaled_mask = (mask * 255).astype(np.uint8)
-
-    # Ensure single-channel mask
-    if len(scaled_mask.shape) == 3 and scaled_mask.shape[2] == 3:
-        scaled_mask = cv2.cvtColor(scaled_mask, cv2.COLOR_BGR2GRAY)
+    # Ensure mask is single-channel grayscale
+    if len(mask.shape) == 3:
+        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+    
+    # Ensure mask is uint8
+    if mask.dtype != np.uint8:
+        mask = (mask * 255).astype(np.uint8)
 
-    # Resize the mask if needed
-    if scaled_mask.shape != (h, w):
-        scaled_mask = cv2.resize(scaled_mask, (w, h), interpolation=cv2.INTER_NEAREST)
+    # Convert design_resized to 3-channel if it's grayscale
+    if len(design_resized.shape) == 2:
+        design_resized = cv2.cvtColor(design_resized, cv2.COLOR_GRAY2BGR)
 
-    print(f"Design Resized Shape: {design_resized.shape}, Mask Shape: {scaled_mask.shape}")
-    return cv2.bitwise_and(design_resized, design_resized, mask=scaled_mask)
+    return cv2.bitwise_and(design_resized, design_resized, mask=mask)
 
 def blend_images(base, overlay, mask):
     """Blends the design onto the dress using seamless cloning."""