Create app.py

app.py

@@ -0,0 +1,101 @@
+import gradio as gr
+from ultralytics import YOLO
+import os
+import shutil
+
+# 1. Load the fine-tuned YOLO model
+model_path = "/content/runs/detect/train/weights/best.pt"
+model = YOLO(model_path)
+
+# Define the directory to save inference results
+output_dir = "inference_results"
+os.makedirs(output_dir, exist_ok=True)
+
+# 2. Define an inference function
+def predict_image_or_video(input_file, conf_threshold):
+    if input_file is None:
+        return "No input file provided."
+
+    print(f"Processing: {input_file}")
+    print(f"Confidence threshold: {conf_threshold}")
+
+    # Clear previous results if any
+    for f in os.listdir(output_dir):
+        os.remove(os.path.join(output_dir, f))
+
+    results = model.predict(source=input_file, conf=conf_threshold, save=True, project=output_dir, name="run", exist_ok=True)
+    
+    # The results are saved in a subdirectory within output_dir/run/
+    # We need to find the actual path to the saved file.
+    # Ultralytics saves to runs/detect/runX (where X is an incrementing number)
+    # We need to find the latest run folder.
+
+    # Get the latest run folder within the output_dir
+    run_folders = [d for d in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, d))]
+    run_folders.sort(key=lambda x: os.path.getmtime(os.path.join(output_dir, x)), reverse=True)
+
+    if not run_folders:
+        return "No detection results saved."
+
+    latest_run_path = os.path.join(output_dir, run_folders[0])
+    
+    # Now find the actual image or video file inside this run_folder
+    # For images, it saves directly. For videos, it creates a new video file.
+    detected_files = [f for f in os.listdir(latest_run_path) if f.endswith(('.jpg', '.jpeg', '.png', '.mp4', '.avi', '.mov'))]
+    
+    if not detected_files:
+        return "No detected image/video found in results."
+        
+    # Assuming only one file is processed at a time
+    result_path = os.path.join(latest_run_path, detected_files[0])
+    
+    print(f"Results saved to: {result_path}")
+    return result_path
+
+# 3. Create a Gradio interface
+# Determine if the input is an image or video based on file extension for the output
+def get_output_component(input_file):
+    if input_file and (input_file.endswith(('.mp4', '.avi', '.mov'))):
+        return gr.Video(label="Detection Results")
+    else:
+        return gr.Image(label="Detection Results")
+
+
+# The Gradio interface setup
+# We will use two separate interfaces, one for image and one for video, and use gr.Tab to switch between them.
+# Alternatively, a single interface with conditional logic inside predict_image_or_video can work, but Gradio is a bit tricky with multiple output types based on input.
+# For simplicity and direct instruction fulfillment, let's make two functions and use gr.Tab.
+
+
+def predict_image(image_file, conf_threshold):
+    if image_file is None:
+        return None
+    return predict_image_or_video(image_file, conf_threshold)
+
+def predict_video(video_file, conf_threshold):
+    if video_file is None:
+        return None
+    return predict_image_or_video(video_file, conf_threshold)
+
+
+with gr.Blocks() as demo:
+    gr.Markdown("# YOLOv8 Signature Detection")
+    gr.Markdown("Upload an image or video to perform signature detection using a fine-tuned YOLOv8n model.")
+
+    with gr.Tab("Image Detection"):
+        image_input = gr.Image(type="filepath", label="Upload Image")
+        image_conf_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold")
+        image_output = gr.Image(label="Detection Results")
+        image_button = gr.Button("Detect Signature in Image")
+        image_button.click(predict_image, inputs=[image_input, image_conf_slider], outputs=image_output)
+
+    with gr.Tab("Video Detection"):
+        video_input = gr.Video(type="filepath", label="Upload Video")
+        video_conf_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.25, step=0.05, label="Confidence Threshold")
+        video_output = gr.Video(label="Detection Results")
+        video_button = gr.Button("Detect Signature in Video")
+        video_button.click(predict_video, inputs=[video_input, video_conf_slider], outputs=video_output)
+
+
+# 4. Launch the Gradio interface
+demo.launch(share=True)