Update utils functions

src/utils.py

@@ -1,30 +1,7 @@
-from PIL import ImageDraw
-import numpy as np
+# Necessary imports
 import re
-
-
-# Use a color map for bounding boxes
-colormap = [
-    "#0000FF",
-    "#FFA500",
-    "#008000",
-    "#800080",
-    "#A52A2A",
-    "#FFC0CB",
-    "#808080",
-    "#808000",
-    "#00FFFF",
-    "#FF0000",
-    "#00FF00",
-    "#4B0082",
-    "#4B0082",
-    "#EE82EE",
-    "#00FFFF",
-    "#FF00FF",
-    "#FF7F50",
-    "#FFD700",
-    "#87CEEB",
-]
+import supervision as sv
+from PIL import Image
 
 
 # Text cleaning function
@@ -62,111 +39,40 @@ def clean_text(text):
     return cleaned_text
 
 
-# Convert hex color to RGBA with the given alpha
-def hex_to_rgba(hex_color, alpha):
-    """
-    Convert a hexadecimal color code to RGBA format.
-
-    Args:
-        hex_color (str): The hexadecimal color code (e.g., "#FF0000").
-        alpha (int): The alpha value for the RGBA color (0-255).
-
-    Returns:
-        tuple: A tuple representing the RGBA color values (red, green, blue, alpha).
-    """
-    hex_color = hex_color.lstrip("#")
-    r, g, b = int(hex_color[0:2], 16), int(hex_color[2:4], 16), int(hex_color[4:6], 16)
-    return (r, g, b, alpha)
-
-
 # Draw OCR bounding boxes with enhanced visual elements
-def draw_ocr_bboxes(image, prediction):
+def draw_ocr_bboxes(image: Image, detections: sv.Detections) -> Image:
     """
-    Draw bounding boxes with enhanced visual elements on the given image based on the OCR prediction.
+    Draws bounding boxes and labels on the input image based on the OCR detections.
 
     Args:
-        image (PIL.Image.Image): The input image on which the bounding boxes will be drawn.
-        prediction (dict): The OCR prediction containing 'quad_boxes' and 'labels'.
+        image (PIL.Image): The input image on which to draw the bounding boxes and labels.
+        detections (sv.Detections): The OCR detections containing the bounding box coordinates and labels.
 
     Returns:
-        PIL.Image.Image: The image with the bounding boxes drawn.
+        PIL.Image: The annotated image with bounding boxes and labels.
+
     """
 
-    # Create a drawing object for the image with RGBA mode
-    draw = ImageDraw.Draw(image, "RGBA")
-
-    # Extract bounding boxes and labels from the prediction
-    bboxes, labels = prediction["quad_boxes"], prediction["labels"]
-
-    for i, (box, label) in enumerate(zip(bboxes, labels)):
-        # Select color for the bounding box and label
-        color = colormap[i % len(colormap)]
-        new_box = (np.array(box)).tolist()
-
-        # Define the outline width and corner radius for the bounding box
-        box_outline_width = 3
-        corner_radius = 10
-
-        # Draw rounded corners for the bounding box
-        for j in range(4):
-            start_x, start_y = new_box[j * 2], new_box[j * 2 + 1]
-            end_x, end_y = new_box[(j * 2 + 2) % 8], new_box[(j * 2 + 3) % 8]
-
-            # Draw the arcs for the rounded corners
-            draw.arc(
-                [
-                    (start_x - corner_radius, start_y - corner_radius),
-                    (start_x + corner_radius, start_y + corner_radius),
-                ],
-                90 + j * 90,
-                180 + j * 90,
-                fill=color,
-                width=box_outline_width,
-            )
-            draw.arc(
-                [
-                    (end_x - corner_radius, end_y - corner_radius),
-                    (end_x + corner_radius, end_y + corner_radius),
-                ],
-                j * 90,
-                90 + j * 90,
-                fill=color,
-                width=box_outline_width,
-            )
-
-            # Draw the lines connecting the arcs
-            if j in [0, 1, 2]:
-                draw.line(
-                    [
-                        (start_x + corner_radius if j != 1 else start_x, start_y),
-                        (end_x - corner_radius if j != 1 else end_x, end_y),
-                    ],
-                    fill=color,
-                    width=box_outline_width,
-                )
-            else:
-                draw.line(
-                    [
-                        (start_x, start_y + corner_radius),
-                        (end_x, end_y - corner_radius),
-                    ],
-                    fill=color,
-                    width=box_outline_width,
-                )
-
-        # Calculate the position for the text label
-        text_x, text_y = min(new_box[0::2]), min(new_box[1::2]) - 20
-        text_w, text_h = draw.textsize(label)
-        rgba_color = hex_to_rgba(color, 200)  # Semi-transparent background for text
-
-        # Draw the background rectangle for the text
-        draw.rectangle(
-            [text_x, text_y, text_x + text_w + 10, text_y + text_h + 10],
-            fill=rgba_color,
-        )
-
-        # Draw the text label
-        draw.text((text_x + 5, text_y + 5), label, fill=(0, 0, 0, 255))
-
-    # Return the image with the OCR boxes drawn
-    return image
+    # Copy the input image to avoid modifying the original image
+    annotated_image = image.copy()
+
+    # Calculate the optimal line thickness and text scale based on the image resolution
+    thickness = sv.calculate_optimal_line_thickness(resolution_wh=image.size)
+    text_scale = sv.calculate_optimal_text_scale(resolution_wh=image.size)
+
+    # Initialize the bounding box and label annotators
+    bounding_box_annotator = sv.BoundingBoxAnnotator(
+        color_lookup=sv.ColorLookup.INDEX, thickness=thickness
+    )
+    label_annotator = sv.LabelAnnotator(
+        color_lookup=sv.ColorLookup.INDEX,
+        text_scale=text_scale,
+        text_thickness=thickness,
+    )
+
+    # Annotate the image with bounding boxes and labels
+    annotated_image = bounding_box_annotator.annotate(annotated_image, detections)
+    annotated_image = label_annotator.annotate(annotated_image, detections)
+
+    # Return the annotated image
+    return annotated_image