Update app.py

app.py

@@ -94,69 +94,87 @@ def detect_layout(img, conf_threshold, iou_threshold):
     classes = [mapping[i] for i in classes]
     return bboxes , classes 
     
+from PIL import Image, ImageDraw, ImageFont
 
 def draw_bboxes_on_image(image_path, bboxes, classes, reading_order):
-    # Define a color map for each class name
+    # Define a color palette for classes
     class_colors = {
-        'CheckBox': 'orange',
-        'List': 'blue',
-        'P': 'green',
-        'abandon': 'purple',
-        'figure': 'cyan',
-        'gridless_table': 'yellow',
-        'handwritten_signature': 'magenta',
-        'qr_code': 'red',
-        'table': 'brown',
-        'title': 'pink'
+        'CheckBox': '#FFA500',  # Orange
+        'List': '#4682B4',      # Steel Blue
+        'P': '#32CD32',         # Lime Green
+        'abandon': '#8A2BE2',   # Blue Violet
+        'figure': '#00CED1',    # Dark Turquoise
+        'gridless_table': '#FFD700',  # Gold
+        'handwritten_signature': '#FF69B4',  # Hot Pink
+        'qr_code': '#FF4500',   # Orange Red
+        'table': '#8B4513',     # Saddle Brown
+        'title': '#FF1493'      # Deep Pink
     }
 
     # Open the image using PIL
-    image = Image.open(image_path)
-    
-    # Prepare to draw on the image
-    draw = ImageDraw.Draw(image)
+    image = Image.open(image_path).convert("RGBA")
+    overlay = Image.new("RGBA", image.size, (255, 255, 255, 0))  # Transparent overlay
+    draw = ImageDraw.Draw(overlay)
     
-    # Try loading a default font, if it fails, use a basic font
+    # Try loading a modern font, or fall back to default
     try:
         font = ImageFont.truetype("arial.ttf", 20)
-        title_font = ImageFont.truetype("arial.ttf", 30)  # Larger font for titles
+        title_font = ImageFont.truetype("arial.ttf", 30)
     except IOError:
-        font = ImageFont.load_default(size = 30)
-        title_font = font  # Use the same font for title if custom font fails
+        font = ImageFont.load_default()
+        title_font = font
 
     # Loop through the bounding boxes and corresponding labels
     for i in range(len(bboxes)):
         x1, y1, x2, y2 = bboxes[i]
         class_name = classes[i]
         order = reading_order[i]
-        
-        # Get the color for the class
-        color = class_colors[class_name]
-        
-        # If it's a title, make the bounding box thicker and text larger
+        color = class_colors.get(class_name, "#FFFFFF")  # Default white if class is unknown
+
+        # Determine box and label styles
         if class_name == 'title':
-            box_thickness = 4  # Thicker box for title
-            label_font = title_font  # Larger font for title
+            box_thickness = 6
+            label_font = title_font
         else:
-            box_thickness = 2  # Default box thickness
-            label_font = font  # Default font for other classes
-        
-        # Draw the rectangle with the class color and box thickness
-        draw.rectangle([x1, y1, x2, y2], outline=color, width=box_thickness)
-        
-        # Label the box with the class and order
+            box_thickness = 3
+            label_font = font
+
+        # Draw a rounded rectangle for the bounding box
+        draw.rounded_rectangle(
+            [(x1, y1), (x2, y2)],
+            radius=10,  # Rounded corners
+            outline=color,
+            width=box_thickness
+        )
+
+        # Create the label with the class and order
         label = f"{class_name}-{order}"
+        text_width, text_height = draw.textsize(label, font=label_font)
+
+        # Add padding to the text background
+        padding = 5
+        label_bg_coords = [
+            x1, y1 - text_height - 2 * padding,
+            x1 + text_width + 2 * padding, y1
+        ]
         
-        # Calculate text size using textbbox() to get the bounding box of the text
-        bbox = draw.textbbox((x1, y1 - 20), label, font=label_font)
-        label_width = bbox[2] - bbox[0]
-        label_height = bbox[3] - bbox[1]
+        # Draw a semi-transparent rectangle for the label background
+        draw.rectangle(label_bg_coords, fill=(0, 0, 0, 128))  # Semi-transparent black
         
-        # Draw the text above the box
-        draw.text((x1, y1 - label_height), label, fill="black", font=label_font)
-    
-    # Return the modified image as a PIL image object
-    return image
+        # Draw the label text
+        draw.text(
+            (x1 + padding, y1 - text_height - padding),
+            label,
+            fill="white",
+            font=label_font
+        )
+
+    # Merge the overlay with the original image
+    image_with_overlay = Image.alpha_composite(image, overlay)
+
+    # Convert back to RGB mode for saving/display
+    return image_with_overlay.convert("RGB")
+
 
 
 def scale_and_normalize_boxes(bboxes, old_width = 1024, old_height= 1024, new_width=640, new_height=640, normalize_width=1000, normalize_height=1000):