Fix logic

main.py

@@ -88,14 +88,16 @@ def predict(image_input):
     mask_images = result.pred_instances.masks.cpu().numpy()
     scores = result.pred_instances.scores.cpu().numpy()
 
-    # If there are no tables, return an empty list.
-    if len(mask_images) == 0:
-        return []
+    bbox_list = []
 
-    # Get mask with highest score
-    mask_image = mask_images[scores.argmax()]
+    # Filter out the masks with a score less than 0.5.
+    filtered_mask_images = mask_images[scores > 0.5]
 
-    return get_bbox(mask_image.astype(np.uint8))
+    # Get the bounding boxes of the tables.
+    for mask_image in filtered_mask_images:
+        bbox_list.extend(get_bbox(mask_image.astype(np.uint8)))
+
+    return bbox_list
 
 
 def run():