Update app.py

app.py

@@ -9,6 +9,7 @@ import supervision as sv
 from segment_anything import sam_model_registry, SamPredictor
 from werkzeug.utils import secure_filename
 
+import torchvision.ops
 import argparse
 from io import BytesIO
 from PIL import Image
@@ -89,9 +90,10 @@ def run_grounding(input_image, grounding_caption, box_threshold, text_threshold)
     image_with_box = Image.fromarray(cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB))
 
     # run sam
-    boxes_np = boxes.cpu().numpy()
-    boxes.mask = segment(sam_predictor, image_tensor, boxes_np)
-    sam_segment = mask_annotator.annotate(scene=np.asarray(image_pil),detections=boxes)
+    areas = torchvision.ops.boxes.box_area(boxes)
+    sorted_indices = torch.flip(torch.argsort(areas), [0])
+    sorted_boxes = boxes[sorted_indices]
+    sam_segment = mask_annotator.annotate(scene=np.asarray(image_pil), detections=sorted_boxes)
     sam_output = Image.fromarray(cv2.cvtColor(sam_segment, cv2.COLOR_BGR2RGB))
     
     return image_with_box, sam_output