Update app.py

app.py

@@ -7,19 +7,23 @@ 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.
     """
@@ -39,11 +43,9 @@ 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.
     """
@@ -67,7 +69,6 @@ def load_models(models_dir=MODELS_DIR, info_file=MODELS_INFO_FILE):
                 continue  
 
         try:
-
             model = YOLO(model_path)
             models[model_name] = {
                 'display_name': display_name,
@@ -80,15 +81,14 @@ def load_models(models_dir=MODELS_DIR, info_file=MODELS_INFO_FILE):
 
     return models
 
-def get_model_info(model_info):
+def get_model_info(model_info, ratings_info):
     """
-    Retrieve formatted model information for display.
-
+    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.
+        str: A formatted string containing model details and average rating.
     """
     info = model_info
     class_ids = info.get('class_ids', {})
@@ -99,6 +99,11 @@ def get_model_info(model_info):
     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"
@@ -111,20 +116,19 @@ def get_model_info(model_info):
         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}"
+        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.
     """
@@ -133,7 +137,6 @@ def predict_image(model_name, image, confidence, models):
     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)
 
@@ -145,7 +148,6 @@ def predict_image(model_name, image, confidence, models):
         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")
@@ -153,17 +155,86 @@ def predict_image(model_name, image, confidence, models):
 
         output_image = Image.open(final_output_path)
 
-        return "✅ Prediction completed successfully.", output_image, 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 main():
+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(
@@ -172,6 +243,15 @@ def main():
             """
         )
 
+        # 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],
@@ -189,7 +269,8 @@ def main():
             if not model_name:
                 return "Model information not available."
             model_entry = models[model_name]['info']
-            return get_model_info(model_entry)
+            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,
@@ -212,7 +293,6 @@ def main():
                 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.**")
@@ -223,7 +303,32 @@ def main():
                 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)
+                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,
@@ -231,6 +336,107 @@ def main():
                 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(
             """
             ---