Update app.py

app.py

@@ -114,28 +114,27 @@ def predict(input_image):
     
     # Filter results by score and class
     segments_info = []
-    scores = result['scores'].cpu().numpy()
-    labels = result['labels'].cpu().numpy()
-    masks = result['masks'].cpu().numpy()
-    
-    # Get bounding boxes from masks
-    for i in range(len(scores)):
-        score = scores[i]
-        label_id = labels[i]
+    segmentation = result["segmentation"].cpu().numpy()
+    segments_info_raw = result["segments_info"]
+
+    for segment in segments_info_raw:
+        score = segment["score"]
+        label_id = segment["label_id"]
+        segment_id = segment["id"]
         label_name = model.config.id2label[label_id]
 
-        # Filter out low-confidence scores and unwanted classes
         if score > 0.9 and label_name in TARGET_CLASSES:
-            mask = masks[i]
-            
+            # Create a binary mask for this segment
+            mask = (segmentation == segment_id)
+
             # Calculate bounding box from mask
             pos = np.where(mask)
-            if pos[0].size > 0 and pos[1].size > 0: # Ensure mask is not empty
+            if pos[0].size > 0 and pos[1].size > 0:
                 xmin = np.min(pos[1])
                 xmax = np.max(pos[1])
                 ymin = np.min(pos[0])
                 ymax = np.max(pos[0])
-                
+
                 segments_info.append({
                     "score": score,
                     "label": label_name,
@@ -158,9 +157,21 @@ def predict(input_image):
 # Load some example images
 # Note: You must upload these images to your Hugging Face Space repository in a folder named 'examples'
 example_paths = [
-    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/cats-vs-dogs.png",
-    "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/instance-segmentation-input.jpg",
-    "https://placehold.co/800x600/FFF/333?text=A+busy+street+scene", # Corrected URL (removed space)
+    "examples/street1.jpeg",
+    "examples/street2.jpeg",
+    "examples/street3.jpeg",
+    "examples/street4.jpeg",
+    "examples/street5.jpeg",
+    "examples/street6.jpeg",
+    "examples/street7.jpeg",
+    "examples/street8.jpeg",
+    "examples/street9.jpeg",
+    "examples/street10.jpeg",
+    "examples/catsanddogs1.jpeg",
+    "examples/catsanddogs2.jpeg",
+    "examples/catsanddogs3.jpeg",
+    "examples/catsanddogs4.jpeg",
+    "examples/catsanddogs5.jpeg",
 ]