app.py

import gradio as gr
import cv2
import torch
import time
import numpy as np
from ultralytics import YOLO
import os

# Optimize CPU usage
torch.set_num_threads(8)
MODEL_DIR = "models"

stop_processing = False  # Global flag to stop processing

def get_model_options():
    models = {}
    for root, dirs, files in os.walk(MODEL_DIR):
        for file in files:
            if file.endswith(".pt"):
                model_name = os.path.basename(os.path.dirname(root))
                models[model_name] = os.path.join(root, file)
    return models

model_options = get_model_options()

def annotate_frame(frame, results):
    for box in results[0].boxes:
        xyxy = box.xyxy[0].numpy()
        class_id = int(box.cls[0].item())
        label = results[0].names[class_id]

        start_point = (int(xyxy[0]), int(xyxy[1]))
        end_point = (int(xyxy[2]), int(xyxy[3]))
        color = (0, 255, 0)
        thickness = 2
        cv2.rectangle(frame, start_point, end_point, color, thickness)

        font = cv2.FONT_HERSHEY_SIMPLEX
        font_scale = 0.5
        font_thickness = 1
        label_position = (int(xyxy[0]), int(xyxy[1] - 10))
        cv2.putText(frame, label, label_position, font, font_scale, color, font_thickness)
    return frame

def process_image(model_name, image, confidence_threshold, iou_threshold):
    model_path = model_options[model_name]
    model = YOLO(model_path).to('cpu')

    frame = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    with torch.inference_mode():
        results = model(frame, conf=confidence_threshold, iou=iou_threshold)
    annotated_frame = annotate_frame(frame, results)
    annotated_frame = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
    return annotated_frame, "N/A"

def run_inference(mode, model_name, image, video, confidence_threshold, iou_threshold):
    global stop_processing
    stop_processing = False  # Reset stop flag at the start

    if mode == "Image":
        if image is None:
            yield None, None, "Please upload an image."
            return
        annotated_img, fps = process_image(model_name, image, confidence_threshold, iou_threshold)
        yield annotated_img, None, fps
    else:
        if video is None:
            yield None, None, "Please upload a video."
            return

        model_path = model_options[model_name]
        model = YOLO(model_path).to('cpu')
        cap = cv2.VideoCapture(video)

        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if frame_count <= 0:
            frame_count = 1

        output_frames = []
        fps_list = []
        processed_count = 0

        while not stop_processing:
            ret, frame = cap.read()
            if not ret:
                break

            start_time = time.time()
            with torch.inference_mode():
                results = model(frame, conf=confidence_threshold, iou=iou_threshold)

            annotated_frame = annotate_frame(frame, results)
            output_frames.append(annotated_frame)

            fps_val = 1 / (time.time() - start_time)
            fps_list.append(fps_val)

            processed_count += 1
            progress_fraction = processed_count / frame_count

            # Yield progress every few frames
            if processed_count % 5 == 0:
                yield None, None, f"Processing... {progress_fraction * 100:.2f}%"

            if stop_processing:
                yield None, None, "Processing canceled."
                return

        cap.release()

        if len(output_frames) > 0 and not stop_processing:
            avg_fps = sum(fps_list) / len(fps_list) if fps_list else 0
            height, width, _ = output_frames[0].shape
            output_video_path = "output.mp4"
            out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), 30, (width, height))
            for frame in output_frames:
                out.write(frame)
            out.release()

            yield None, output_video_path, f"Average FPS: {avg_fps:.2f}"
        elif not stop_processing:
            yield None, None, "No frames processed."

def cancel_processing():
    global stop_processing
    stop_processing = True
    return "Cancel signal sent."

def start_app():
    model_names = list(model_options.keys())

    with gr.Blocks() as app:
        # **Instructional Message Added Here**
        gr.Markdown("""
        ### Welcome to the YOLO Inference App!
        **How to Use:**
        1. **Select Mode:**
           - Choose between **Image** or **Video** processing.
        2. **Select Model:**
           - Pick a pre-trained YOLO model from the dropdown menu.
        3. **Upload Your File:**
           - For **Image** mode, upload an image (e.g., `pothole.jpg`).
           - For **Video** mode, upload a video (e.g., `potholeall.mp4` or `electric bus fire.mp4`).
        4. **Adjust Thresholds:**
           - **Confidence Threshold:** Determines the minimum confidence for detections.
           - **IoU Threshold:** Determines the Intersection over Union threshold for non-maximum suppression.
        5. **Start Processing:**
           - Click on **Start Processing** to begin inference.
           - You can cancel the processing at any time by clicking **Cancel Processing**.
        **Example Files:**
        - **Image:** `pothole.jpg`
        - **Videos:** `potholeall.mp4`, `electric bus fire.mp4`
        """)

        gr.Markdown("## YOLO Inference (Image or Video) with Progress & Cancel")
        
        with gr.Row():
            mode = gr.Radio(["Image", "Video"], value="Image", label="Mode")
            model_selector = gr.Dropdown(choices=model_names, label="Select Model", value=model_names[0])
        
        image_input = gr.Image(label="Upload Image", visible=True)
        video_input = gr.Video(label="Upload Video", visible=False)

        confidence_slider = gr.Slider(0.1, 1.0, value=0.3, step=0.1, label="Confidence Threshold")
        iou_slider = gr.Slider(0.1, 1.0, value=0.001, step=0.001, label="IoU Threshold")

        annotated_image_output = gr.Image(label="Annotated Image", visible=True)
        annotated_video_output = gr.Video(label="Output Video", visible=False)
        fps_output = gr.Textbox(label="Status / Average FPS", interactive=False)

        start_button = gr.Button("Start Processing")
        cancel_button = gr.Button("Cancel Processing", variant="stop")

        # Updated example files with 'examples/' path and renamed video file
        examples = gr.Examples(
            examples=[
                ["examples/pothole.jpg", None, 0.3, 0.001],                  # Example for image
                [None, "examples/potholeall.mp4", 0.3, 0.001],              # Renamed video example
                [None, "examples/electric bus fire.mp4", 0.5, 0.001]        # Updated confidence threshold for new video example
            ],
            inputs=[image_input, video_input, confidence_slider, iou_slider]
        )

        def update_visibility(selected_mode):
            if selected_mode == "Image":
                return (
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(visible=True),
                    gr.update(visible=False)
                )
            else:
                return (
                    gr.update(visible=False),
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(visible=True)
                )

        mode.change(
            update_visibility, 
            inputs=mode, 
            outputs=[image_input, video_input, annotated_image_output, annotated_video_output]
        )

        start_button.click(
            fn=run_inference,
            inputs=[mode, model_selector, image_input, video_input, confidence_slider, iou_slider],
            outputs=[annotated_image_output, annotated_video_output, fps_output],
            queue=True
        )

        cancel_button.click(
            fn=cancel_processing,
            inputs=[],
            outputs=[fps_output],
            queue=False
        )

    return app

if __name__ == "__main__":
    app = start_app()
    app.launch()