Update app.py

app.py

@@ -1,63 +1,54 @@
 import gradio as gr
-import torch
 from ultralytics import YOLO
 from PIL import Image, ImageDraw
 
-# Load YOLO model for line detection
+# Load YOLO model (trained on 640x640)
 YOLO_MODEL_PATH = "best.pt"
 model = YOLO(YOLO_MODEL_PATH, task='detect').to("cpu")  # Force CPU usage
 
-def merge_boxes_into_lines(boxes, y_threshold=20):
-    """
-    Merge bounding boxes that are close together in the y-axis (same line).
-    Args:
-        boxes: List of bounding boxes [x1, y1, x2, y2]
-        y_threshold: Max distance between words to consider as the same line
-    Returns:
-        List of merged line bounding boxes without overlap
-    """
-    if len(boxes) == 0:
-        return []
+def resize_and_pad(image, target_size=(640, 640)):
+    """Resize image while keeping aspect ratio and padding to fit target size."""
+    original_size = image.size  # (width, height)
+    image.thumbnail((target_size[0], target_size[1]), Image.ANTIALIAS)  # Resize while keeping aspect ratio
 
-    # Sort boxes by y1 (top position)
-    boxes = sorted(boxes, key=lambda b: b[1])
-
-    merged_lines = []
-    current_line = list(boxes[0])
-
-    for i in range(1, len(boxes)):
-        x1, y1, x2, y2 = boxes[i]
-
-        # Merge boxes that are close in the y-axis
-        if abs(y1 - current_line[1]) < y_threshold:
-            current_line[0] = min(current_line[0], x1)  # Expand left boundary
-            current_line[2] = max(current_line[2], x2)  # Expand right boundary
-            current_line[3] = max(current_line[3], y2)  # Expand bottom boundary
-        else:
-            merged_lines.append(current_line)
-            current_line = list(boxes[i])
+    # Create a new white background image
+    new_image = Image.new("RGB", target_size, (255, 255, 255))
+    
+    # Paste the resized image in the center
+    paste_x = (target_size[0] - image.size[0]) // 2
+    paste_y = (target_size[1] - image.size[1]) // 2
+    new_image.paste(image, (paste_x, paste_y))
 
-    merged_lines.append(current_line)
-    return merged_lines
+    return new_image, original_size, paste_x, paste_y
 
 def detect_lines(image):
     """Runs YOLOv8 detection on the input image to detect lines."""
     image = Image.fromarray(image)  # Convert NumPy array to PIL Image
     
-    # Run YOLO detection with confidence threshold 0.2
-    results = model.predict(image, conf=0.3, iou=0.5, device="cpu")
+    # Resize & pad A4 images
+    resized_image, original_size, pad_x, pad_y = resize_and_pad(image)
+
+    # Run YOLO detection
+    results = model.predict(resized_image, conf=0.3, iou=0.5, device="cpu")
     
     detected_boxes = results[0].boxes.xyxy.tolist()
     detected_boxes = [list(map(int, box)) for box in detected_boxes]  # Convert to integer
 
-    # Merge detected bounding boxes into line-level boxes
-    merged_boxes = merge_boxes_into_lines(detected_boxes)
+    # Scale boxes back to original image size
+    width_ratio = original_size[0] / 640
+    height_ratio = original_size[1] / 640
+
+    scaled_boxes = [
+        [int((x1 - pad_x) * width_ratio), int((y1 - pad_y) * height_ratio), 
+         int((x2 - pad_x) * width_ratio), int((y2 - pad_y) * height_ratio)]
+        for x1, y1, x2, y2 in detected_boxes
+    ]
 
-    # Draw bounding boxes
+    # Draw bounding boxes on the original image
     image_with_boxes = image.copy()
     draw = ImageDraw.Draw(image_with_boxes)
 
-    for idx, (x1, y1, x2, y2) in enumerate(merged_boxes):
+    for idx, (x1, y1, x2, y2) in enumerate(scaled_boxes):
         draw.rectangle([x1, y1, x2, y2], outline="blue", width=2)
         draw.text((x1, y1 - 10), f"Line {idx}", fill="blue")