Update app.py

app.py

@@ -51,32 +51,30 @@ def infer(
 def extract_bounding_boxes(result):
     """
     Extract bounding boxes and labels from the model result.
-    Each bounding box is represented by two locXXXX tags and a label.
+    
+    Extracts two x,y coordinate pairs from <loc> tags and associates them with the corresponding labels.
     
     Example return: [((x1, y1, x2, y2), "Label")]
     """
-    bounding_boxes = []
+    # Regular expression to capture the <loc> tags and their associated labels
+    loc_pattern = re.compile(r"<loc(\d{4})><loc(\d{4})><loc(\d{4})><loc(\d{4})>\s*([a-zA-Z\-]+)")
     
-    # Regular expression to find <locXXXX> tags and labels
-    pattern = re.compile(r'<loc(\d{4})><loc(\d{4})>\s*(\S.+?)\s*(?=<loc|\Z)')
+    # Find all matches of bounding box coordinates and labels in the result string
+    matches = loc_pattern.findall(result)
     
-    matches = pattern.findall(result)
+    bounding_boxes = []
     
     for match in matches:
-        # Extract x1, y1 from the first loc tag
-        x1, y1 = int(match[0][:2]), int(match[0][2:])
+        x1, y1, x2, y2, label = match
         
-        # Extract x2, y2 from the second loc tag
-        x2, y2 = int(match[1][:2]), int(match[1][2:])
+        # Convert coordinates from string to integer
+        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
         
-        # Get the label
-        label = match[2].strip()
-        
-        # Append the bounding box with the label
+        # Append the bounding box and label as a tuple
         bounding_boxes.append(((x1, y1, x2, y2), label))
     
     return bounding_boxes
-
+    
 ######## Demo
 
 INTRO_TEXT = """## Curacel Auto Damage demo\n\n