Update app.py

app.py

@@ -16,26 +16,37 @@ def segment_dress(image):
     """Segments the dress from an input image using SAM."""
     input_points = [[[image.size[0] // 2, image.size[1] // 2]]]
     inputs = processor(image, input_points=input_points, return_tensors="pt").to(DEVICE)
+    
     with torch.no_grad():
         outputs = model(**inputs)
+
     masks = processor.image_processor.post_process_masks(
-        outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu()
+        outputs.pred_masks.cpu(), 
+        inputs["original_sizes"].cpu(), 
+        inputs["reshaped_input_sizes"].cpu()
     )
-    return masks[0][0].numpy() if masks else None
+
+    if masks:
+        mask = masks[0][0].numpy()
+        # Convert boolean mask to uint8 (0-255)
+        mask = (mask * 255).astype(np.uint8)
+        return mask
+    return None
 
 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))
 
-    # Convert boolean mask to uint8 (0-255)
-    mask = (mask * 255).astype(np.uint8)
-
     # 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)
 
     # Convert design_resized to 3-channel if it's grayscale
     if len(design_resized.shape) == 2:
@@ -61,7 +72,6 @@ def blend_images(base, overlay, mask):
     if mask.dtype == np.bool_:
         mask = mask.astype(np.uint8) * 255
 
-    # Ensure mask is single-channel grayscale
     # Ensure mask is single-channel grayscale
     if len(mask.shape) == 3 and mask.shape[2] == 3:
         mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
@@ -71,7 +81,7 @@ def blend_images(base, overlay, mask):
         mask = (mask * 255).astype(np.uint8)
 
     # Compute center of the mask for seamless cloning
-    center = tuple(np.array(base.shape[:2]) // 2)
+    center = (base.shape[1] // 2, base.shape[0] // 2)
 
     return cv2.seamlessClone(overlay, base, mask, center, cv2.NORMAL_CLONE)
 
@@ -79,13 +89,14 @@ def apply_design(image_path, design_path, warp_scale):
     """Pipeline to segment, warp, and blend design onto dress."""
     image = Image.open(image_path).convert("RGB")
     design = cv2.imread(design_path)
+
     mask = segment_dress(image)
-    
     if mask is None:
         return "Segmentation Failed!"
     
     warped_design = warp_design(design, mask, warp_scale)
     blended = blend_images(np.array(image), warped_design, mask)
+    
     return Image.fromarray(blended)
 
 def main(image, design, warp_scale):