Create main.py

main.py

@@ -0,0 +1,217 @@
+import gradio as gr
+import json
+import os
+from pathlib import Path
+from PIL import Image
+import shutil
+from ultralytics import YOLO
+
+def load_models(models_dir='models', info_file='models_info.json'):
+    """
+    Load YOLO models and their information from the specified directory and JSON file.
+
+    Args:
+        models_dir (str): Path to the models directory.
+        info_file (str): Path to the JSON file containing model info.
+
+    Returns:
+        dict: A dictionary of models and their associated information.
+    """
+    with open(info_file, 'r') as f:
+        models_info = json.load(f)
+    
+    models = {}
+    for model_info in models_info:
+        model_name = model_info['model_name']
+        model_path = os.path.join(models_dir, model_name, 'best.pt')  # Assuming 'best.pt' as the weight file
+        if os.path.isfile(model_path):
+            try:
+                # Load the YOLO model
+                model = YOLO(model_path)
+                models[model_name] = {
+                    'model': model,
+                    'mAP': model_info.get('mAP_score', 'N/A'),
+                    'num_images': model_info.get('num_images', 'N/A')
+                }
+                print(f"Loaded model '{model_name}' from '{model_path}'.")
+            except Exception as e:
+                print(f"Error loading model '{model_name}': {e}")
+        else:
+            print(f"Model weight file for '{model_name}' not found at '{model_path}'. Skipping.")
+    return models
+
+def get_model_info(model_name, models):
+    """
+    Retrieve model information for the selected model.
+
+    Args:
+        model_name (str): The name of the model.
+        models (dict): The dictionary containing models and their info.
+
+    Returns:
+        str: A formatted string containing model information.
+    """
+    model_info = models.get(model_name, {})
+    if not model_info:
+        return "Model information not available."
+    info_text = (
+        f"**Model Name:** {model_name}\n\n"
+        f"**mAP Score:** {model_info.get('mAP', 'N/A')}\n\n"
+        f"**Number of Images Trained On:** {model_info.get('num_images', 'N/A')}"
+    )
+    return info_text
+
+def predict_image(model_name, image, models):
+    """
+    Perform prediction on an uploaded image using the selected YOLO model.
+
+    Args:
+        model_name (str): The name of the selected model.
+        image (PIL.Image.Image): The uploaded image.
+        models (dict): The dictionary containing models and their info.
+
+    Returns:
+        tuple: A status message, the processed image, and the path to the output image.
+    """
+    model = models.get(model_name, {}).get('model', None)
+    if not model:
+        return "Error: Model not found.", None, None
+    try:
+        # Save the uploaded image to a temporary path
+        input_image_path = f"temp/{model_name}_input_image.jpg"
+        os.makedirs(os.path.dirname(input_image_path), exist_ok=True)
+        image.save(input_image_path)
+
+        # Perform prediction
+        results = model(input_image_path, save=True, save_txt=False, conf=0.25)
+        # Ultralytics saves the result images in 'runs/detect/predict'
+        output_image_path = results[0].save()[0]  # Get the path to the saved image
+
+        # Open the output image
+        output_image = Image.open(output_image_path)
+
+        return "Prediction completed successfully.", output_image, output_image_path
+    except Exception as e:
+        return f"Error during prediction: {str(e)}", None, None
+
+def predict_video(model_name, video, models):
+    """
+    Perform prediction on an uploaded video using the selected YOLO model.
+
+    Args:
+        model_name (str): The name of the selected model.
+        video (str): Path to the uploaded video file.
+        models (dict): The dictionary containing models and their info.
+
+    Returns:
+        tuple: A status message, the processed video, and the path to the output video.
+    """
+    model = models.get(model_name, {}).get('model', None)
+    if not model:
+        return "Error: Model not found.", None, None
+    try:
+        # Ensure the video is saved in a temporary location
+        input_video_path = video.name
+        if not os.path.isfile(input_video_path):
+            # If the video is a temp file provided by Gradio
+            shutil.copy(video.name, input_video_path)
+        
+        # Perform prediction
+        results = model(input_video_path, save=True, save_txt=False, conf=0.25)
+        # Ultralytics saves the result videos in 'runs/detect/predict'
+        output_video_path = results[0].save()[0]  # Get the path to the saved video
+
+        return "Prediction completed successfully.", output_video_path, output_video_path
+    except Exception as e:
+        return f"Error during prediction: {str(e)}", None, None
+
+def main():
+    # Load the models and their information
+    models = load_models()
+
+    # Initialize Gradio Blocks interface
+    with gr.Blocks() as demo:
+        gr.Markdown("# 🧪 YOLO Model Tester")
+
+        gr.Markdown(
+            """
+            Upload images or videos to test different YOLO models. Select a model from the dropdown to see its details.
+            """
+        )
+
+        # Model selection and info
+        with gr.Row():
+            model_dropdown = gr.Dropdown(
+                choices=list(models.keys()),
+                label="Select Model",
+                value=None
+            )
+            model_info = gr.Markdown("**Model Information will appear here.**")
+
+        # Update model_info when a model is selected
+        model_dropdown.change(
+            fn=lambda model_name: get_model_info(model_name, models) if model_name else "Please select a model.",
+            inputs=model_dropdown,
+            outputs=model_info
+        )
+
+        # Tabs for different input types
+        with gr.Tabs():
+            # Image Prediction Tab
+            with gr.Tab("🖼️ Image"):
+                with gr.Column():
+                    image_input = gr.Image(
+                        type='pil',
+                        label="Upload Image for Prediction",
+                        tool="editor"
+                    )
+                    image_predict_btn = gr.Button("🔍 Predict on Image")
+                    image_status = gr.Markdown("**Status will appear here.**")
+                    image_output = gr.Image(label="Predicted Image")
+                    image_download_btn = gr.File(label="⬇️ Download Predicted Image")
+
+                # Define the image prediction function
+                def process_image(model_name, image):
+                    return predict_image(model_name, image, models)
+
+                # Connect the predict button
+                image_predict_btn.click(
+                    fn=process_image,
+                    inputs=[model_dropdown, image_input],
+                    outputs=[image_status, image_output, image_download_btn]
+                )
+
+            # Video Prediction Tab
+            with gr.Tab("🎥 Video"):
+                with gr.Column():
+                    video_input = gr.Video(
+                        label="Upload Video for Prediction"
+                    )
+                    video_predict_btn = gr.Button("🔍 Predict on Video")
+                    video_status = gr.Markdown("**Status will appear here.**")
+                    video_output = gr.Video(label="Predicted Video")
+                    video_download_btn = gr.File(label="⬇️ Download Predicted Video")
+
+                # Define the video prediction function
+                def process_video(model_name, video):
+                    return predict_video(model_name, video, models)
+
+                # Connect the predict button
+                video_predict_btn.click(
+                    fn=process_video,
+                    inputs=[model_dropdown, video_input],
+                    outputs=[video_status, video_output, video_download_btn]
+                )
+
+        gr.Markdown(
+            """
+            ---
+            **Note:** Ensure that the YOLO models are correctly placed in the `models/` directory and that `models_info.json` is properly configured.
+            """
+        )
+
+    # Launch the Gradio app
+    demo.launch()
+
+if __name__ == "__main__":
+    main()