Update app.py

app.py

@@ -99,83 +99,67 @@ def detect_layout(img, conf_threshold, iou_threshold):
 from PIL import Image, ImageDraw, ImageFont
 
 def draw_bboxes_on_image(image_path, bboxes, classes, reading_order):
-    # Define a color palette for classes
+    # Define a color map for each class name
     class_colors = {
-        '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
+        'CheckBox': 'orange',
+        'List': 'blue',
+        'P': 'green',
+        'abandon': 'purple',
+        'figure': 'cyan',
+        'gridless_table': 'yellow',
+        'handwritten_signature': 'magenta',
+        'qr_code': 'red',
+        'table': 'brown',
+        'title': 'pink'
     }
 
     # Open the image using PIL
-    image = image_path.convert("RGBA")
-    overlay = Image.new("RGBA", image.size, (255, 255, 255, 0))  # Transparent overlay
-    draw = ImageDraw.Draw(overlay)
+    image = image_path
     
-    # Try loading a modern font, or fall back to default
+    # Prepare to draw on the image
+    draw = ImageDraw.Draw(image)
+    
+    # Try loading a default font, if it fails, use a basic font
     try:
         font = ImageFont.truetype("arial.ttf", 20)
-        title_font = ImageFont.truetype("arial.ttf", 30)
+        title_font = ImageFont.truetype("arial.ttf", 30)  # Larger font for titles
     except IOError:
-        font = ImageFont.load_default()
-        title_font = font
+        font = ImageFont.load_default(size = 30)
+        title_font = font  # Use the same font for title if custom font fails
 
     # 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]
-        color = class_colors.get(class_name, "#FFFFFF")  # Default white if class is unknown
-
-        # Determine box and label styles
+        
+        # Get the color for the class
+        color = class_colors[class_name]
+        
+        # If it's a title, make the bounding box thicker and text larger
         if class_name == 'title':
-            box_thickness = 6
-            label_font = title_font
+            box_thickness = 4  # Thicker box for title
+            label_font = title_font  # Larger font for title
         else:
-            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
+            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
         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
-        ]
         
-        # Draw a semi-transparent rectangle for the label background
-        draw.rectangle(label_bg_coords, fill=(0, 0, 0, 128))  # Semi-transparent black
+        # 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 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")
+        # 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