changes

app.py

@@ -2,7 +2,6 @@ import streamlit as st
 import numpy as np
 import cv2
 from PIL import Image
-from torchvision import transforms, models
 from segment_anything import sam_model_registry, SamPredictor
 from io import BytesIO
 import base64
@@ -27,21 +26,34 @@ def preprocess_image(image):
 
 # Segment garment using SAM
 def segment_garment(image, predictor):
+    # Convert image to NumPy and set it in SAM
     image_np = np.array(image.convert("RGB"))
     predictor.set_image(image_np)
 
+    # Define multiple input points (center, upper body, lower body)
     height, width, _ = image_np.shape
-    input_point = np.array([[width // 2, height // 2]])
-    input_label = np.array([1])
-
-    masks, _, _ = predictor.predict(point_coords=input_point, point_labels=input_label)
+    input_points = np.array([
+        [width // 2, height // 3],         # upper torso
+        [width // 2, height // 2],         # center
+        [width // 2, 2 * height // 3]      # lower torso
+    ])
+    input_labels = np.array([1, 1, 1])     # all positive prompts
+
+    # Generate mask using multiple points
+    masks, scores, logits = predictor.predict(
+        point_coords=input_points,
+        point_labels=input_labels,
+        multimask_output=False
+    )
+
+    # Post-process mask
     mask = masks[0]
     mask_resized = cv2.resize(mask.astype(np.uint8) * 255, (width, height), interpolation=cv2.INTER_NEAREST)
-    mask_resized = np.stack([mask_resized] * 3, axis=-1)
+    mask_rgb = np.stack([mask_resized] * 3, axis=-1)
 
-    segmented = np.where(mask_resized > 0, image_np, 0)
+    # Apply mask to original image
+    segmented = np.where(mask_rgb > 0, image_np, 0)
     return Image.fromarray(segmented)
-
 # AI garment analysis
 def analyze_garment(image, style_pref=None, feedback=None, generate_variations=False):
     image_bytes = BytesIO()