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

# Directory and file configurations
MODELS_DIR = "models"
MODELS_INFO_FILE = "models_info.json"
TEMP_DIR = "temp"
OUTPUT_DIR = "outputs"

# New files for storing ratings, detections, and recommended datasets
RATINGS_FILE = "ratings.json"
DETECTIONS_FILE = "detections.json"
RECOMMENDED_DATASETS_FILE = "recommended_datasets.json"

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()  
        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")  
        download_url = model_info['download_url']

        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  

        try:
            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, ratings_info):
    """
    Retrieve formatted model information for display, including average rating.
    Args:
        model_info (dict): The model's information dictionary.
        ratings_info (dict): The ratings information for the model.
    Returns:
        str: A formatted string containing model details and average rating.
    """
    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])

    # Calculate average rating
    total_rating = ratings_info.get('total', 0)
    count_rating = ratings_info.get('count', 0)
    average_rating = (total_rating / count_rating) if count_rating > 0 else "No ratings yet"

    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}\n\n"
        f"**Average Rating:** {average_rating} ⭐"
    )
    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:
        os.makedirs(TEMP_DIR, exist_ok=True)
        os.makedirs(OUTPUT_DIR, exist_ok=True)

        input_image_path = os.path.join(TEMP_DIR, f"{model_name}_input_image.jpg")
        image.save(input_image_path)

        results = model(input_image_path, save=True, save_txt=False, conf=confidence)

        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):
            output_image_path = results[0].save()[0]

        final_output_path = os.path.join(OUTPUT_DIR, f"{model_name}_output_image.jpg")
        shutil.copy(output_image_path, final_output_path)

        output_image = Image.open(final_output_path)

        # Calculate number of detections
        detections = len(results[0].boxes)
        return f"✅ Prediction completed successfully. **Detections:** {detections}", output_image, final_output_path
    except Exception as e:
        return f"❌ Error during prediction: {str(e)}", None, None

def load_or_initialize_json(file_path, default_data):
    """
    Load JSON data from a file or initialize it with default data if the file doesn't exist.
    Args:
        file_path (str): Path to the JSON file.
        default_data (dict or list): Default data to initialize if file doesn't exist.
    Returns:
        dict or list: The loaded or initialized data.
    """
    if os.path.isfile(file_path):
        with open(file_path, 'r') as f:
            return json.load(f)
    else:
        with open(file_path, 'w') as f:
            json.dump(default_data, f, indent=4)
        return default_data

def save_json(file_path, data):
    """
    Save data to a JSON file.
    Args:
        file_path (str): Path to the JSON file.
        data (dict or list): Data to save.
    """
    with open(file_path, 'w') as f:
        json.dump(data, f, indent=4)

def is_valid_roboflow_url(url):
    """
    Validate if the provided URL is a Roboflow URL.
    Args:
        url (str): The URL to validate.
    Returns:
        bool: True if valid, False otherwise.
    """
    return url.startswith("https://roboflow.com/") or url.startswith("http://roboflow.com/")

def get_top_model(detections_per_model, models):
    """
    Determine the top model based on the highest number of detections.
    Args:
        detections_per_model (dict): Dictionary with model names as keys and detection counts as values.
        models (dict): Dictionary of loaded models.
    Returns:
        str: The display name of the top model or a message if no detections exist.
    """
    if not detections_per_model:
        return "No detections yet."
    top_model_name = max(detections_per_model, key=detections_per_model.get)
    top_model_display = models[top_model_name]['display_name']
    top_detections = detections_per_model[top_model_name]
    return f"**Top Model:** {top_model_display} with **{top_detections}** detections."

def main():
    # Load models
    models = load_models()
    if not models:
        print("No models loaded. Please check your models_info.json and model URLs.")
        return

    # Load or initialize ratings
    ratings_data = load_or_initialize_json(RATINGS_FILE, {})
    # Initialize ratings for each model if not present
    for model_name in models:
        if model_name not in ratings_data:
            ratings_data[model_name] = {"total": 0, "count": 0}
    save_json(RATINGS_FILE, ratings_data)

    # Load or initialize detections counter
    detections_data = load_or_initialize_json(DETECTIONS_FILE, {"total_detections": 0, "detections_per_model": {}})

    # Load or initialize recommended datasets
    recommended_datasets = load_or_initialize_json(RECOMMENDED_DATASETS_FILE, [])

    with gr.Blocks() as demo:
        gr.Markdown("# 🧪 YOLOv11 Model Tester")
        gr.Markdown(
            """
            Upload images to test different YOLOv11 models. Select a model from the dropdown to see its details.
            """
        )

        # Display total detections counter and top model
        with gr.Row():
            detections_counter = gr.Markdown(
                f"**Total Detections Across All Users:** {detections_data.get('total_detections', 0)}"
            )
            top_model_display = gr.Markdown(
                get_top_model(detections_data.get('detections_per_model', {}), models)
            )

        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.**")

        display_to_name = {models[m]['display_name']: m for m in models}

        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']
            ratings_info = ratings_data.get(model_name, {"total": 0, "count": 0})
            return get_model_info(model_entry, ratings_info)

        model_dropdown.change(
            fn=update_model_info,
            inputs=model_dropdown,
            outputs=model_info
        )

        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."
            )

        with gr.Tab("🖼️ Image"):
            with gr.Column():
                image_input = gr.Image(
                    type='pil',
                    label="Upload Image for Prediction"
                )
                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")

            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)
                status, output_img, output_path = predict_image(model_name, image, confidence, models)
                
                # Extract number of detections from the status message
                detections = 0
                if "Detections:" in status:
                    try:
                        detections = int(status.split("Detections:")[1].strip())
                    except:
                        pass

                # Update detections counter
                try:
                    detections_data['total_detections'] += detections
                    if model_name in detections_data['detections_per_model']:
                        detections_data['detections_per_model'][model_name] += detections
                    else:
                        detections_data['detections_per_model'][model_name] = detections
                    save_json(DETECTIONS_FILE, detections_data)
                except Exception as e:
                    print(f"Error updating detections counter: {e}")

                # Update detections and top model display
                detections_counter.value = f"**Total Detections Across All Users:** {detections_data.get('total_detections', 0)}"
                top_model_display.value = get_top_model(detections_data.get('detections_per_model', {}), models)

                return status, output_img, output_path

            image_predict_btn.click(
                fn=process_image,
                inputs=[model_dropdown, image_input, confidence_slider],
                outputs=[image_status, image_output, image_download_btn]
            )

        with gr.Tab("⭐ Rate Model"):
            with gr.Column():
                selected_model = gr.Dropdown(
                    choices=[models[m]['display_name'] for m in models],
                    label="Select Model to Rate",
                    value=None
                )
                rating = gr.Slider(
                    minimum=1,
                    maximum=5,
                    step=1,
                    label="Rate the Model (1-5 Stars)",
                    info="Select a star rating between 1 and 5."
                )
                submit_rating_btn = gr.Button("Submit Rating")
                rating_status = gr.Markdown("**Your rating will be submitted here.**")

            def submit_rating(selected_display_name, user_rating):
                if not selected_display_name:
                    return "❌ Please select a model to rate."
                if not user_rating:
                    return "❌ Please provide a rating."
                model_name = display_to_name.get(selected_display_name)
                if not model_name:
                    return "❌ Invalid model selected."

                # Update ratings data
                ratings_info = ratings_data.get(model_name, {"total": 0, "count": 0})
                ratings_info['total'] += user_rating
                ratings_info['count'] += 1
                ratings_data[model_name] = ratings_info
                save_json(RATINGS_FILE, ratings_data)

                # Update model info display if the rated model is currently selected
                if model_dropdown.value == selected_display_name:
                    updated_info = get_model_info(models[model_name]['info'], ratings_info)
                    model_info.value = updated_info

                average = (ratings_info['total'] / ratings_info['count'])
                return f"✅ Thank you for rating! Current Average Rating: {average:.2f} ⭐"

            submit_rating_btn.click(
                fn=submit_rating,
                inputs=[selected_model, rating],
                outputs=rating_status
            )

        with gr.Tab("💡 Recommend Dataset"):
            with gr.Column():
                dataset_name = gr.Textbox(
                    label="Dataset Name",
                    placeholder="Enter the name of the dataset"
                )
                dataset_url = gr.Textbox(
                    label="Dataset URL",
                    placeholder="Enter the Roboflow dataset URL"
                )
                recommend_btn = gr.Button("Recommend Dataset")
                recommend_status = gr.Markdown("**Your recommendation status will appear here.**")

            def recommend_dataset(name, url):
                if not name or not url:
                    return "❌ Please provide both the dataset name and URL."

                if not is_valid_roboflow_url(url):
                    return "❌ Invalid URL. Please provide a valid Roboflow dataset URL."

                # Check for duplicates
                for dataset in recommended_datasets:
                    if dataset['name'].lower() == name.lower() or dataset['url'] == url:
                        return "❌ This dataset has already been recommended."

                # Add to recommended datasets
                recommended_datasets.append({"name": name, "url": url})
                save_json(RECOMMENDED_DATASETS_FILE, recommended_datasets)

                return f"✅ Thank you for recommending the dataset **{name}**!"

            recommend_btn.click(
                fn=recommend_dataset,
                inputs=[dataset_name, dataset_url],
                outputs=recommend_status
            )

        with gr.Tab("📄 Recommended Datasets"):
            with gr.Column():
                recommended_display = gr.Markdown("### Recommended Roboflow Datasets\n")

                def display_recommended_datasets():
                    if not recommended_datasets:
                        return "No datasets have been recommended yet."
                    dataset_md = "\n".join([f"- [{dataset['name']}]({dataset['url']})" for dataset in recommended_datasets])
                    return dataset_md

                # Display the recommended datasets
                recommended_display.value = display_recommended_datasets()

        with gr.Tab("🏆 Top Model"):
            with gr.Column():
                top_model_md = gr.Markdown(get_top_model(detections_data.get('detections_per_model', {}), models))

        gr.Markdown(
            """
            ---
            **Note:** Models are downloaded from GitHub upon first use. Ensure that you have a stable internet connection and sufficient storage space.
            """
        )

    demo.launch()

if __name__ == "__main__":
    main()