Create app.py

app.py

@@ -0,0 +1,224 @@
+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()