app-WorkFine(Debug).py

import gradio as gr
from ultralytics import YOLO
from PIL import Image, ImageDraw, ImageFont
import random

# Load YOLO model (ensure best.pt exists in the working directory)
YOLO_MODEL_PATH = "my_best_355epoch.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]

# 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 class names."""
    image = Image.fromarray(image)
    original_image = image.copy()

    # Run YOLO text detection
    results = model.predict(image, conf=0.7, 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)
    
    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 Classes 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 and Labels")
            output_annotated = gr.Image(type="pil", label="Detected Text Lines")

    gr.Markdown("### 🔢 Detection Results")
    output_objects = gr.Textbox(label="Total Objects Detected", lines=1)
    output_classes = gr.Textbox(label="Total Classes 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)


# 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)