Update app.py

app.py

@@ -34,7 +34,7 @@ def infer(
     result = processor.decode(generated_ids[0], skip_special_tokens=True)
 
     # Placeholder to extract bounding box info from the result (you should replace this with actual bounding box extraction)
-    bounding_boxes = extract_bounding_boxes(result)
+    bounding_boxes = extract_bounding_boxes(result, image)
     
     # Draw bounding boxes on the image
     annotated_image = image.copy()
@@ -48,17 +48,26 @@ def infer(
     
     return result, annotated_image
 
-def extract_bounding_boxes(result):
+
+def extract_bounding_boxes(result, image):
     """
     Extract bounding boxes and labels from the model result.
     
-    Extracts two x,y coordinate pairs from <loc> tags and associates them with the corresponding labels.
+    Coordinates are scaled by dividing by 1024 and then multiplying by the image dimensions.
     
-    Example return: [((x1, y1, x2, y2), "Label")]
+    Args:
+        result (str): The model's output string containing bounding box data.
+        image (PIL.Image.Image): The image to use for scaling the bounding boxes.
+        
+    Returns:
+        List[Tuple[Tuple[int, int, int, int], str]]: A list of bounding boxes and labels.
     """
     # 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\-]+)")
     
+    # Get image dimensions
+    width, height = image.size
+    
     # Find all matches of bounding box coordinates and labels in the result string
     matches = loc_pattern.findall(result)
     
@@ -70,7 +79,13 @@ def extract_bounding_boxes(result):
         # Convert coordinates from string to integer
         x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
         
-        # Append the bounding box and label as a tuple
+        # Scale coordinates
+        x1 = int((x1 / 1024) * width)
+        y1 = int((y1 / 1024) * height)
+        x2 = int((x2 / 1024) * width)
+        y2 = int((y2 / 1024) * height)
+        
+        # Append the scaled bounding box and label as a tuple
         bounding_boxes.append(((x1, y1, x2, y2), label))
     
     return bounding_boxes