Update app.py

app.py

@@ -2,38 +2,25 @@ import gradio as gr
 from ultralytics import YOLO
 from PIL import Image, ImageDraw
 
-# Load YOLO model
 YOLO_MODEL_PATH = "best.pt"
-model = YOLO(YOLO_MODEL_PATH, task='detect').to("cpu")  # Force CPU usage
+model = YOLO(YOLO_MODEL_PATH, task='detect').to("cpu")
 
 def merge_boxes_into_lines(boxes, y_threshold=10):
-    """
-    Merge bounding boxes that are on the same row but not merge different row lines.
-    Args:
-        boxes: List of bounding boxes [x1, y1, x2, y2]
-        y_threshold: Max difference in y1 position to be considered the same row
-    Returns:
-        List of merged line bounding boxes
-    """
     if len(boxes) == 0:
         return []
 
-    # 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 only if y position is very close (same row)
         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
+            current_line[0] = min(current_line[0], x1)
+            current_line[2] = max(current_line[2], x2)
+            current_line[3] = max(current_line[3], y2)
         else:
-            # Store previous line and start a new one
             merged_lines.append(current_line)
             current_line = list(boxes[i])
 
@@ -41,25 +28,15 @@ def merge_boxes_into_lines(boxes, y_threshold=10):
     return merged_lines
 
 def detect_and_crop_lines(image):
-    """
-    Detects text lines using YOLO, merges them, and crops each line.
-    Args:
-        image: Input image (PIL format)
-    Returns:
-        Annotated image with bounding boxes, List of cropped images
-    """
-    image = Image.fromarray(image)  # Convert NumPy array to PIL Image
-    original_image = image.copy()  # Keep a copy of the original image
-
-    # Run YOLO detection on the original image
+    image = Image.fromarray(image)
+    original_image = image.copy()
+
     results = model.predict(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
+    detected_boxes = [list(map(int, box)) for box in detected_boxes]
 
-    # Merge bounding boxes based on row position
     merged_boxes = merge_boxes_into_lines(detected_boxes)
 
-    # Draw bounding boxes
     draw = ImageDraw.Draw(original_image)
     cropped_lines = []
 
@@ -67,13 +44,11 @@ def detect_and_crop_lines(image):
         draw.rectangle([x1, y1, x2, y2], outline="blue", width=2)
         draw.text((x1, y1 - 10), f"Line {idx}", fill="blue")
 
-        # Crop the detected text line
         cropped_line = image.crop((x1, y1, x2, y2))
         cropped_lines.append(cropped_line)
 
     return original_image, cropped_lines
 
-# Define Gradio interface
 with gr.Blocks() as iface:
     gr.Markdown("# Text Line Detection")
     gr.Markdown("## Input your custom image for text line detection")
@@ -87,15 +62,22 @@ with gr.Blocks() as iface:
             gr.Markdown("### Annotated Image with Detected Lines")
             output_annotated = gr.Image(type="pil", label="Detected Text Lines")
 
-    gr.Markdown("### Cropped Text Lines (Each Line Detected Separately)")
+    gr.Markdown("### Cropped Text Lines (Displayed Row by Row)")
+    
+    cropped_output_rows = []
+    for i in range(20):
+        with gr.Row():
+            cropped_output_rows.append(gr.Image(type="pil", label=f"Line {i+1}"))
 
-    cropped_gallery = gr.Gallery(label="Cropped Lines Gallery", columns=3, preview=True)
+    def process_and_display(image):
+        annotated_img, cropped_imgs = detect_and_crop_lines(image)
+        cropped_imgs += [None] * (20 - len(cropped_imgs))
+        return [annotated_img] + cropped_imgs[:20]
 
     image_input.upload(
-        lambda img: detect_and_crop_lines(img),
+        process_and_display,
         inputs=image_input,
-        outputs=[output_annotated, cropped_gallery]
+        outputs=[output_annotated] + cropped_output_rows
     )
 
-# Launch Gradio interface
 iface.launch()