app.py

import gradio as gr
import json
import os
from pathlib import Path
from PIL import Image
import shutil
from ultralytics import YOLO
import requests

# Constants
MODELS_DIR = "models"
MODELS_INFO_FILE = "models_info.json"
TEMP_DIR = "temp"
OUTPUT_DIR = "outputs"

def download_file(url, dest_path):
    """
    Download a file from a URL to the destination path.
    
    Args:
        url (str): The URL to download from.
        dest_path (str): The local path to save the file.
    
    Returns:
        bool: True if download succeeded, False otherwise.
    """
    try:
        response = requests.get(url, stream=True)
        response.raise_for_status()  # Raise an error on bad status
        with open(dest_path, 'wb') as f:
            for chunk in response.iter_content(chunk_size=8192):
                f.write(chunk)
        print(f"Downloaded {url} to {dest_path}.")
        return True
    except Exception as e:
        print(f"Failed to download {url}. Error: {e}")
        return False

def load_models(models_dir=MODELS_DIR, info_file=MODELS_INFO_FILE):
    """
    Load YOLO models and their information from the specified directory and JSON file.
    Downloads models if they are not already present.
    
    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']
        display_name = model_info.get('display_name', model_name)
        model_dir = os.path.join(models_dir, model_name)
        os.makedirs(model_dir, exist_ok=True)
        model_path = os.path.join(model_dir, f"{model_name}.pt")  # e.g., models/human/human.pt
        download_url = model_info['download_url']
        
        # Check if the model file exists
        if not os.path.isfile(model_path):
            print(f"Model '{display_name}' not found locally. Downloading from {download_url}...")
            success = download_file(download_url, model_path)
            if not success:
                print(f"Skipping model '{display_name}' due to download failure.")
                continue  # Skip loading this model
        
        try:
            # Load the YOLO model
            model = YOLO(model_path)
            models[model_name] = {
                'display_name': display_name,
                'model': model,
                'info': model_info
            }
            print(f"Loaded model '{display_name}' from '{model_path}'.")
        except Exception as e:
            print(f"Error loading model '{display_name}': {e}")
    
    return models

def get_model_info(model_info):
    """
    Retrieve formatted model information for display.
    
    Args:
        model_info (dict): The model's information dictionary.
    
    Returns:
        str: A formatted string containing model details.
    """
    info = model_info
    class_ids = info.get('class_ids', {})
    class_image_counts = info.get('class_image_counts', {})
    datasets_used = info.get('datasets_used', [])
    
    class_ids_formatted = "\n".join([f"{cid}: {cname}" for cid, cname in class_ids.items()])
    class_image_counts_formatted = "\n".join([f"{cname}: {count}" for cname, count in class_image_counts.items()])
    datasets_used_formatted = "\n".join([f"- {dataset}" for dataset in datasets_used])
    
    info_text = (
        f"**{info.get('display_name', 'Model Name')}**\n\n"
        f"**Architecture:** {info.get('architecture', 'N/A')}\n\n"
        f"**Training Epochs:** {info.get('training_epochs', 'N/A')}\n\n"
        f"**Batch Size:** {info.get('batch_size', 'N/A')}\n\n"
        f"**Optimizer:** {info.get('optimizer', 'N/A')}\n\n"
        f"**Learning Rate:** {info.get('learning_rate', 'N/A')}\n\n"
        f"**Data Augmentation Level:** {info.get('data_augmentation_level', 'N/A')}\n\n"
        f"**[email protected]:** {info.get('mAP_score', 'N/A')}\n\n"
        f"**Number of Images Trained On:** {info.get('num_images', 'N/A')}\n\n"
        f"**Class IDs:**\n{class_ids_formatted}\n\n"
        f"**Datasets Used:**\n{datasets_used_formatted}\n\n"
        f"**Class Image Counts:**\n{class_image_counts_formatted}"
    )
    return info_text

def predict_image(model_name, image, confidence, 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.
        confidence (float): The confidence threshold for detections.
        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_entry = models.get(model_name, {})
    model = model_entry.get('model', None)
    if not model:
        return "Error: Model not found.", None, None
    try:
        # Ensure temporary and output directories exist
        os.makedirs(TEMP_DIR, exist_ok=True)
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        
        # Save the uploaded image to a temporary path
        input_image_path = os.path.join(TEMP_DIR, f"{model_name}_input_image.jpg")
        image.save(input_image_path)
        
        # Perform prediction with user-specified confidence
        results = model(input_image_path, save=True, save_txt=False, conf=confidence)
        
        # Determine the output path
        # Ultralytics YOLO saves the results in 'runs/detect/predict' by default
        # We'll move the result to OUTPUT_DIR with a unique name
        latest_run = sorted(Path("runs/detect").glob("predict*"), key=os.path.getmtime)[-1]
        output_image_path = os.path.join(latest_run, Path(input_image_path).name)
        if not os.path.isfile(output_image_path):
            # Alternative method to get the output path
            output_image_path = results[0].save()[0]
        
        # Copy the output image to OUTPUT_DIR
        final_output_path = os.path.join(OUTPUT_DIR, f"{model_name}_output_image.jpg")
        shutil.copy(output_image_path, final_output_path)
        
        # Open the output image
        output_image = Image.open(final_output_path)
        
        return "✅ Prediction completed successfully.", output_image, final_output_path
    except Exception as e:
        return f"❌ Error during prediction: {str(e)}", None, None

def predict_video(model_name, video, confidence, 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.
        confidence (float): The confidence threshold for detections.
        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_entry = models.get(model_name, {})
    model = model_entry.get('model', None)
    if not model:
        return "Error: Model not found.", None, None
    try:
        # Ensure temporary and output directories exist
        os.makedirs(TEMP_DIR, exist_ok=True)
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        
        # Save the uploaded video to a temporary path
        input_video_path = os.path.join(TEMP_DIR, f"{model_name}_input_video.mp4")
        shutil.copy(video, input_video_path)
        
        # Perform prediction with user-specified confidence
        results = model(input_video_path, save=True, save_txt=False, conf=confidence)
        
        # Determine the output path
        # Ultralytics YOLO saves the results in 'runs/detect/predict' by default
        latest_run = sorted(Path("runs/detect").glob("predict*"), key=os.path.getmtime)[-1]
        output_video_path = os.path.join(latest_run, Path(input_video_path).name)
        if not os.path.isfile(output_video_path):
            # Alternative method to get the output path
            output_video_path = results[0].save()[0]
        
        # Copy the output video to OUTPUT_DIR
        final_output_path = os.path.join(OUTPUT_DIR, f"{model_name}_output_video.mp4")
        shutil.copy(output_video_path, final_output_path)
        
        return "✅ Prediction completed successfully.", final_output_path, final_output_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()
    if not models:
        print("No models loaded. Please check your models_info.json and model URLs.")
        return
    
    # Initialize Gradio Blocks interface
    with gr.Blocks() as demo:
        gr.Markdown("# 🧪 YOLOv11 Model Tester")
        gr.Markdown(
            """
            Upload images or videos to test different YOLOv11 models. Select a model from the dropdown to see its details.
            """
        )
        
        # Model selection and info
        with gr.Row():
            model_dropdown = gr.Dropdown(
                choices=[models[m]['display_name'] for m in models],
                label="Select Model",
                value=None
            )
            model_info = gr.Markdown("**Model Information will appear here.**")
        
        # Mapping from display_name to model_name
        display_to_name = {models[m]['display_name']: m for m in models}
        
        # Update model_info when a model is selected
        def update_model_info(selected_display_name):
            if not selected_display_name:
                return "Please select a model."
            model_name = display_to_name.get(selected_display_name)
            if not model_name:
                return "Model information not available."
            model_entry = models[model_name]['info']
            return get_model_info(model_entry)
        
        model_dropdown.change(
            fn=update_model_info,
            inputs=model_dropdown,
            outputs=model_info
        )
        
        # Confidence Threshold Slider
        with gr.Row():
            confidence_slider = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                step=0.01,
                value=0.25,
                label="Confidence Threshold",
                info="Adjust the minimum confidence required for detections to be displayed."
            )
        
        # 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"
                        # Removed 'tool' parameter
                    )
                    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(selected_display_name, image, confidence):
                    if not selected_display_name:
                        return "❌ Please select a model.", None, None
                    model_name = display_to_name.get(selected_display_name)
                    return predict_image(model_name, image, confidence, models)
                
                # Connect the predict button
                image_predict_btn.click(
                    fn=process_image,
                    inputs=[model_dropdown, image_input, confidence_slider],
                    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(selected_display_name, video, confidence):
                    if not selected_display_name:
                        return "❌ Please select a model.", None, None
                    model_name = display_to_name.get(selected_display_name)
                    return predict_video(model_name, video, confidence, models)
                
                # Connect the predict button
                video_predict_btn.click(
                    fn=process_video,
                    inputs=[model_dropdown, video_input, confidence_slider],
                    outputs=[video_status, video_output, video_download_btn]
                )
        
        gr.Markdown(
            """
            ---
            **Note:** Models are downloaded from GitHub upon first use. Ensure that you have a stable internet connection and sufficient storage space.
            """
        )
    
    # Launch the Gradio app
    demo.launch()

if __name__ == "__main__":
    main()