Update app.py

app.py

@@ -7,6 +7,7 @@ import shutil
 from ultralytics import YOLO
 import requests
 
+# Constants
 MODELS_DIR = "models"
 MODELS_INFO_FILE = "models_info.json"
 TEMP_DIR = "temp"
@@ -15,17 +16,17 @@ 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()  
+        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)
@@ -39,35 +40,36 @@ 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")  
+        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  
-
+                continue  # Skip loading this model
+        
         try:
-
+            # Load the YOLO model
             model = YOLO(model_path)
             models[model_name] = {
                 'display_name': display_name,
@@ -77,16 +79,16 @@ def load_models(models_dir=MODELS_DIR, info_file=MODELS_INFO_FILE):
             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.
     """
@@ -94,11 +96,11 @@ def get_model_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"
@@ -118,13 +120,13 @@ def get_model_info(model_info):
 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,26 +135,32 @@ def predict_image(model_name, image, confidence, models):
     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
         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
@@ -160,13 +168,13 @@ def predict_image(model_name, image, confidence, models):
 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.
     """
@@ -175,35 +183,42 @@ def predict_video(model_name, video, confidence, models):
     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)
-
-        results = model(input_video_path, save=True, save_txt=False, conf=confidence)
-
+        
+        # Perform prediction with user-specified confidence and specify output format
+        # Here, we set save_format to 'avi' to ensure compatibility
+        results = model(input_video_path, save=True, save_txt=False, conf=confidence, save_format='avi')
+        
+        # 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)
+        output_video_path = os.path.join(latest_run, f"{model_name}_input_video.avi")
         if not os.path.isfile(output_video_path):
-
+            # Alternative method to get the output path
             output_video_path = results[0].save()[0]
-
-        final_output_path = os.path.join(OUTPUT_DIR, f"{model_name}_output_video.mp4")
+        
+        # Copy the output video to OUTPUT_DIR
+        final_output_path = os.path.join(OUTPUT_DIR, f"{model_name}_output_video.avi")
         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(
@@ -211,7 +226,8 @@ def main():
             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],
@@ -219,9 +235,11 @@ def main():
                 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."
@@ -230,13 +248,14 @@ def main():
                 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,
@@ -246,33 +265,37 @@ def main():
                 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(
@@ -282,27 +305,30 @@ def main():
                     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()
+    main()