Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from ultralytics import YOLO
-from PIL import Image, ImageDraw
+from PIL import Image, ImageDraw, ImageFont
 import random
 
 # Load YOLO model (ensure best.pt exists in the working directory)
@@ -16,8 +16,11 @@ def get_class_color(class_id):
         CLASS_COLORS[class_id] = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
     return CLASS_COLORS[class_id]
 
+# Class Names (Modify based on your dataset)
+CLASS_NAMES = {0: "Text Line", 1: "Heading", 2: "Signature"}  # Example labels
+
 def detect_text_lines(image):
-    """Detects text lines and draws bounding boxes with different colors for each class."""
+    """Detects text lines and draws bounding boxes with class names."""
     image = Image.fromarray(image)
     original_image = image.copy()
 
@@ -29,15 +32,33 @@ def detect_text_lines(image):
 
     # Draw bounding boxes on the image
     draw = ImageDraw.Draw(original_image)
+    
+    try:
+        font = ImageFont.truetype("arial.ttf", 18)  # Load a font (ensure arial.ttf is available)
+    except:
+        font = ImageFont.load_default()  # Fallback in case font is missing
+
     for idx, (x1, y1, x2, y2) in enumerate(detected_boxes):
         class_id = int(class_ids[idx])
         color = get_class_color(class_id)
+        class_name = CLASS_NAMES.get(class_id, f"Class {class_id}")
+
+        # Draw bounding box
         draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
 
+        # Draw label with background
+        text_size = draw.textbbox((0, 0), class_name, font=font)
+        text_width = text_size[2] - text_size[0]
+        text_height = text_size[3] - text_size[1]
+
+        # Draw filled rectangle behind text for better visibility
+        draw.rectangle([x1, y1 - text_height - 4, x1 + text_width + 6, y1], fill=color)
+        draw.text((x1 + 3, y1 - text_height - 2), class_name, fill="white", font=font)
+
     total_objects = len(detected_boxes)
     total_classes = len(set(class_ids))
 
-    return original_image, f"Total Objects Detected: {total_objects}", f"Total Class Detected: {total_classes}"
+    return original_image, f"Total Objects Detected: {total_objects}", f"Total Classes Detected: {total_classes}"
 
 # Gradio UI
 with gr.Blocks() as iface:
@@ -50,13 +71,12 @@ with gr.Blocks() as iface:
             image_input = gr.Image(type="numpy", label="Upload an image")
 
         with gr.Column(scale=1):
-            gr.Markdown("### 🖼 Annotated Image with Bounding Boxes")
+            gr.Markdown("### 🖼 Annotated Image with Bounding Boxes and Labels")
             output_annotated = gr.Image(type="pil", label="Detected Text Lines")
 
-    gr.Markdown("### 🔢 Total Objects Detected")
+    gr.Markdown("### 🔢 Detection Results")
     output_objects = gr.Textbox(label="Total Objects Detected", lines=1)
-    gr.Markdown("### 🔢 Total Class Detected")
-    output_classes = gr.Textbox(label="Total Class Detected", lines=1)
+    output_classes = gr.Textbox(label="Total Classes Detected", lines=1)
 
     image_input.upload(
         detect_text_lines,
@@ -66,4 +86,75 @@ with gr.Blocks() as iface:
 
 # 🚀 Ensure the app runs properly in Hugging Face Spaces
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch(server_name="0.0.0.0", server_port=7860)
+
+
+# import gradio as gr
+# from ultralytics import YOLO
+# from PIL import Image, ImageDraw
+# import random
+
+# # Load YOLO model (ensure best.pt exists in the working directory)
+# YOLO_MODEL_PATH = "Yolov12s-trained.pt"
+# model = YOLO(YOLO_MODEL_PATH, task='detect').to("cpu")
+
+# # Define a set of colors for different classes
+# CLASS_COLORS = {}
+
+# def get_class_color(class_id):
+#     """Assign a random color to each class."""
+#     if class_id not in CLASS_COLORS:
+#         CLASS_COLORS[class_id] = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+#     return CLASS_COLORS[class_id]
+
+# def detect_text_lines(image):
+#     """Detects text lines and draws bounding boxes with different colors for each class."""
+#     image = Image.fromarray(image)
+#     original_image = image.copy()
+
+#     # Run YOLO text detection
+#     results = model.predict(image, conf=0.4, device="cpu")
+#     detected_boxes = results[0].boxes.xyxy.tolist()
+#     class_ids = results[0].boxes.cls.tolist()
+#     detected_boxes = [list(map(int, box)) for box in detected_boxes]
+
+#     # Draw bounding boxes on the image
+#     draw = ImageDraw.Draw(original_image)
+#     for idx, (x1, y1, x2, y2) in enumerate(detected_boxes):
+#         class_id = int(class_ids[idx])
+#         color = get_class_color(class_id)
+#         draw.rectangle([x1, y1, x2, y2], outline=color, width=2)
+
+#     total_objects = len(detected_boxes)
+#     total_classes = len(set(class_ids))
+
+#     return original_image, f"Total Objects Detected: {total_objects}", f"Total Class Detected: {total_classes}"
+
+# # Gradio UI
+# with gr.Blocks() as iface:
+#     gr.Markdown("# 📜 Text Line Detection with YOLO")
+#     gr.Markdown("## 📷 Upload an image to detect text lines")
+
+#     with gr.Row():
+#         with gr.Column(scale=1):
+#             gr.Markdown("### 📤 Upload Image")
+#             image_input = gr.Image(type="numpy", label="Upload an image")
+
+#         with gr.Column(scale=1):
+#             gr.Markdown("### 🖼 Annotated Image with Bounding Boxes")
+#             output_annotated = gr.Image(type="pil", label="Detected Text Lines")
+
+#     gr.Markdown("### 🔢 Total Objects Detected")
+#     output_objects = gr.Textbox(label="Total Objects Detected", lines=1)
+#     gr.Markdown("### 🔢 Total Class Detected")
+#     output_classes = gr.Textbox(label="Total Class Detected", lines=1)
+
+#     image_input.upload(
+#         detect_text_lines,
+#         inputs=image_input,
+#         outputs=[output_annotated, output_objects, output_classes]
+#     )
+
+# # 🚀 Ensure the app runs properly in Hugging Face Spaces
+# if __name__ == "__main__":
+#     iface.launch(server_name="0.0.0.0", server_port=7860)