Update app.py

app.py

@@ -6,15 +6,11 @@ from PIL import Image
 import shutil
 from ultralytics import YOLO
 import requests
-import zipfile
-import uuid
 
-# Constants
 MODELS_DIR = "models"
 MODELS_INFO_FILE = "models_info.json"
 TEMP_DIR = "temp"
 OUTPUT_DIR = "outputs"
-ZIP_DIR = "zips"
 
 def download_file(url, dest_path):
     """
@@ -29,7 +25,7 @@ def download_file(url, dest_path):
     """
     try:
         response = requests.get(url, stream=True)
-        response.raise_for_status()  # Raise an error on bad status
+        response.raise_for_status()  
         with open(dest_path, 'wb') as f:
             for chunk in response.iter_content(chunk_size=8192):
                 f.write(chunk)
@@ -60,19 +56,18 @@ def load_models(models_dir=MODELS_DIR, info_file=MODELS_INFO_FILE):
         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
+        model_path = os.path.join(model_dir, f"{model_name}.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
+                continue  
 
         try:
-            # Load the YOLO model
+
             model = YOLO(model_path)
             models[model_name] = {
                 'display_name': display_name,
@@ -120,108 +115,103 @@ def get_model_info(model_info):
     )
     return info_text
 
-def zip_processed_images(processed_image_paths, model_name):
+def predict_image(model_name, image, confidence, models):
     """
-    Create a ZIP file containing all processed images.
+    Perform prediction on an uploaded image using the selected YOLO model.
 
     Args:
-        processed_image_paths (list): List of file paths to processed images.
-        model_name (str): Name of the model used for processing.
+        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:
-        str: Path to the created ZIP file.
+        tuple: A status message, the processed image, and the path to the output image.
     """
-    os.makedirs(ZIP_DIR, exist_ok=True)
-    zip_filename = f"{model_name}_processed_images_{uuid.uuid4().hex}.zip"
-    zip_path = os.path.join(ZIP_DIR, zip_filename)
+    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):
 
-    with zipfile.ZipFile(zip_path, 'w') as zipf:
-        for img_path in processed_image_paths:
-            arcname = os.path.basename(img_path)
-            zipf.write(img_path, arcname)
+            output_image_path = results[0].save()[0]
 
-    print(f"Created ZIP file at {zip_path}.")
-    return zip_path
+        final_output_path = os.path.join(OUTPUT_DIR, f"{model_name}_output_image.jpg")
+        shutil.copy(output_image_path, final_output_path)
 
-def predict_image(model_name, images, confidence, models):
+        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 uploaded images using the selected YOLO model.
+    Perform prediction on an uploaded video using the selected YOLO model.
 
     Args:
         model_name (str): The name of the selected model.
-        images (list): List of uploaded PIL.Image.Image objects.
+        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, list of processed images, and a ZIP file for download.
+        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)
 
-        processed_image_paths = []
-        processed_images = []
-
-        for idx, image in enumerate(images):
-            # Generate unique filenames to avoid conflicts
-            unique_id = uuid.uuid4().hex
-            input_image_path = os.path.join(TEMP_DIR, f"{model_name}_input_image_{unique_id}.jpg")
-            output_image_path = os.path.join(OUTPUT_DIR, f"{model_name}_output_image_{unique_id}.jpg")
-
-            # Save the uploaded image to a temporary path
-            image.save(input_image_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_image_path, save=True, save_txt=False, conf=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]
-            detected_image_path = os.path.join(latest_run, Path(input_image_path).name)
+        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):
 
-            if not os.path.isfile(detected_image_path):
-                # Alternative method to get the output path
-                detected_image_path = results[0].save()[0]
+            output_video_path = results[0].save()[0]
 
-            # Copy the output image to OUTPUT_DIR with a unique name
-            shutil.copy(detected_image_path, output_image_path)
-            processed_image_paths.append(output_image_path)
-
-            # Open the processed image for display
-            processed_image = Image.open(output_image_path)
-            processed_images.append(processed_image)
-
-        # Create a ZIP file containing all processed images
-        zip_path = zip_processed_images(processed_image_paths, model_name)
-
-        return "✅ Prediction completed successfully.", processed_images, zip_path
+        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 one or multiple images to test different YOLOv11 models. Select a model from the dropdown to see its details.
+            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],
@@ -230,10 +220,8 @@ def main():
             )
             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."
@@ -249,7 +237,6 @@ def main():
             outputs=model_info
         )
 
-        # Confidence Threshold Slider
         with gr.Row():
             confidence_slider = gr.Slider(
                 minimum=0.0,
@@ -260,33 +247,53 @@ def main():
                 info="Adjust the minimum confidence required for detections to be displayed."
             )
 
-        # Image Prediction Tab (now supporting multiple images)
-        with gr.Tab("🖼️ Image"):
-            with gr.Column():
-                image_input = gr.Images(
-                    label="Upload Images for Prediction",
-                    type='pil'
+        with gr.Tabs():
+
+            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)
+                    return predict_image(model_name, image, confidence, models)
+
+                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("🎥 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")
+
+                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)
+
+                video_predict_btn.click(
+                    fn=process_video,
+                    inputs=[model_dropdown, video_input, confidence_slider],
+                    outputs=[video_status, video_output, video_download_btn]
                 )
-                image_predict_btn = gr.Button("🔍 Predict on Images")
-                image_status = gr.Markdown("**Status will appear here.**")
-                image_gallery = gr.Gallery(label="Predicted Images").style(grid=[2], height="auto")
-                image_download_btn = gr.File(label="⬇️ Download All Processed Images (ZIP)")
-
-            # Define the image prediction function
-            def process_image(selected_display_name, images, confidence):
-                if not selected_display_name:
-                    return "❌ Please select a model.", None, None
-                if not images:
-                    return "❌ Please upload at least one image.", None, None
-                model_name = display_to_name.get(selected_display_name)
-                return predict_image(model_name, images, confidence, models)
-
-            # Connect the predict button
-            image_predict_btn.click(
-                fn=process_image,
-                inputs=[model_dropdown, image_input, confidence_slider],
-                outputs=[image_status, image_gallery, image_download_btn]
-            )
 
         gr.Markdown(
             """
@@ -295,8 +302,7 @@ def main():
             """
         )
 
-    # Launch the Gradio app
     demo.launch()
 
 if __name__ == "__main__":
-    main()
+    main()