update

app.py

@@ -13,6 +13,64 @@ from utils import (
     overlay_image_top_left,
     annotate_image_with_phrase_and_label
 )
+import numpy as np
+
+
+def compute_iou(boxA, boxB):
+    xA = max(boxA[0], boxB[0])
+    yA = max(boxA[1], boxB[1])
+    xB = min(boxA[2], boxB[2])
+    yB = min(boxA[3], boxB[3])
+
+    interArea = max(0, xB - xA) * max(0, yB - yA)
+    if interArea == 0:
+        return 0.0
+
+    boxAArea = (boxA[2] - boxA[0]) * (boxA[3] - boxA[1])
+    boxBArea = (boxB[2] - boxB[0]) * (boxB[3] - boxB[1])
+
+    iou = interArea / float(boxAArea + boxBArea - interArea)
+    return iou
+
+def extract_voted_bounding_box(image_path, description, num_attempts=5, iou_threshold=0.5):
+    bboxes = []
+
+    # Extract multiple bounding boxes
+    for _ in range(num_attempts):
+        response = extract_bounding_box(image_path, description)
+        match = re.search(r'\((\d+),\s*(\d+),\s*(\d+),\s*(\d+)\)', response)
+        if match:
+            bbox = tuple(map(int, match.groups()))
+            bboxes.append(bbox)
+
+    if not bboxes:
+        return None
+
+    # Group boxes by IoU
+    groups = []
+    used = set()
+
+    for i in range(len(bboxes)):
+        if i in used:
+            continue
+        group = [bboxes[i]]
+        used.add(i)
+        for j in range(i+1, len(bboxes)):
+            if j in used:
+                continue
+            if compute_iou(bboxes[i], bboxes[j]) > iou_threshold:
+                group.append(bboxes[j])
+                used.add(j)
+        if len(group) >= 3:
+            groups.append(group)
+
+    if not groups:
+        return bboxes[0]  # Fallback to the first bounding box if no groups found
+
+    # Use the group with the most members
+    best_group = max(groups, key=len)
+    voted_bbox = np.mean(np.array(best_group), axis=0).astype(int)
+    return tuple(voted_bbox)
 
 # Set your API key
 os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY") if os.getenv("GOOGLE_API_KEY") else st.secrets["GOOGLE_API_KEY"]
@@ -42,9 +100,7 @@ if st.button("Generate Thumbnail"):
         extract_frame(video_path, timestamp, image_path)
 
         # BOUNDING BOX
-        response = extract_bounding_box(image_path, description)
-        bbox_match = re.search(r'\((\d+),\s*(\d+),\s*(\d+),\s*(\d+)\)', response)
-        bbox = tuple(map(int, bbox_match.groups())) if bbox_match else None
+        bbox = extract_voted_bounding_box(image_path, description)
 
         center_bbox_in_circle(image_path, bbox, image_path)
         overlay_image_bottom_right(image_path, "assets/corner.png", image_path)