Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from ultralytics import YOLO
 from PIL import Image, ImageDraw
+import numpy as np
 
 # Load YOLO model
 YOLO_MODEL_PATH = "best.pt"
@@ -47,7 +48,7 @@ def detect_lines(image, resize=False, target_size=(640, 640)):
         resize: Boolean, whether to resize image before detection
         target_size: Tuple (width, height) for resizing
     Returns:
-        Image with bounding boxes drawn
+        Image with bounding boxes drawn, Number of detected boxes
     """
     image = Image.fromarray(image)  # Convert NumPy array to PIL Image
     original_image = image.copy()  # Keep a copy of the original image
@@ -71,11 +72,12 @@ def detect_lines(image, resize=False, target_size=(640, 640)):
         draw.rectangle([x1, y1, x2, y2], outline="blue", width=2)
         draw.text((x1, y1 - 10), f"Line {idx}", fill="blue")
 
-    return image_with_boxes
+    return image_with_boxes, len(merged_boxes)
 
 # Define Gradio interface with two options: Original & Resized detection
 with gr.Blocks() as iface:
-    gr.Markdown("# Text Line Detection with Merging (YOLOv8)")
+    gr.Markdown("# Text Line Detection")
+    gr.Markdown("## Input your custom image for text line detection")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -85,15 +87,17 @@ with gr.Blocks() as iface:
         with gr.Column(scale=1):
             gr.Markdown("### YOLO on Original Image")
             output_original = gr.Image(type="pil", label="Detected Lines (Original Size)")
+            count_original = gr.Textbox(label="Number of Detected Lines (Original Size)")
 
         with gr.Column(scale=1):
             gr.Markdown("### YOLO on Resized Image (640x640)")
             output_resized = gr.Image(type="pil", label="Detected Lines (Resized to 640x640)")
+            count_resized = gr.Textbox(label="Number of Detected Lines (Resized)")
 
     image_input.upload(
-        lambda img: (detect_lines(img, resize=False), detect_lines(img, resize=True, target_size=(640, 640))),
+        lambda img: (*detect_lines(img, resize=False), *detect_lines(img, resize=True, target_size=(640, 640))),
         inputs=image_input,
-        outputs=[output_original, output_resized]
+        outputs=[output_original, count_original, output_resized, count_resized]
     )
 
 # Launch Gradio interface