Update app.py

app.py

@@ -3,15 +3,10 @@ import gradio as gr
 from PIL import Image
 import os
 import zipfile
-import io
-import shutil
 
-# Clear the torch hub cache to prevent any previous corrupted downloads
-shutil.rmtree(torch.hub._get_torch_home(), ignore_errors=True)
-
-# Load the YOLOv5 model from the uploaded best.pt file
+# Load the YOLOv5 model (ensure the path is correct)
 model = torch.hub.load(
-    'ultralytics/yolov5', 'custom', path='best-3.pt', trust_repo=True
+    'ultralytics/yolov5', 'custom', path='best.pt'
 )
 
 # Global variables to store images and labels for carousel functionality
@@ -21,10 +16,7 @@ processed_image_paths = []
 
 # Temporary folder for saving processed files
 TEMP_DIR = "temp_processed"
-
-# Ensure the temp directory exists
-if not os.path.exists(TEMP_DIR):
-    os.makedirs(TEMP_DIR)
+os.makedirs(TEMP_DIR, exist_ok=True)
 
 # Function to process all uploaded images
 def process_images(files):
@@ -79,10 +71,16 @@ def process_images(files):
     else:
         return None, "No images found.", 0
 
-# Function to create and return a ZIP file for download
+# Function to create and return the path to the ZIP file for download
 def create_zip():
-    zip_buffer = io.BytesIO()
-    with zipfile.ZipFile(zip_buffer, 'w') as z:
+    zip_filename = "processed_images_annotations.zip"
+    zip_path = os.path.join(TEMP_DIR, zip_filename)
+
+    # Remove existing ZIP file if it exists
+    if os.path.exists(zip_path):
+        os.remove(zip_path)
+
+    with zipfile.ZipFile(zip_path, 'w') as z:
         # Add images and labels to the ZIP file
         for image_path in processed_image_paths:
             z.write(image_path, os.path.basename(image_path))
@@ -95,9 +93,7 @@ def create_zip():
             label_path = os.path.join(TEMP_DIR, label_filename)
             z.write(label_path, label_filename)
 
-    zip_buffer.seek(0)  # Go to the start of the buffer
-    # Return the bytes of the ZIP file
-    return zip_buffer.getvalue()
+    return zip_path  # Return the file path as a string
 
 # Function to navigate through images
 def next_image(index):
@@ -130,11 +126,9 @@ with gr.Blocks() as interface:
     # Hidden state to store current index
     current_index = gr.State(0)
 
-    # Button to prepare the ZIP file for download
-    prepare_download_button = gr.Button("Prepare Download")
-
-    # Download button
-    download_button = gr.DownloadButton(label="Download ZIP", visible=False)
+    # Button to download all processed images and annotations as a ZIP file
+    download_button = gr.Button("Prepare and Download All")
+    download_file = gr.File()
 
     # Define functionality when files are uploaded
     file_input.change(
@@ -155,15 +149,15 @@ with gr.Blocks() as interface:
         outputs=[image_display, label_display, current_index]
     )
 
-    # Define functionality for the prepare download button
+    # Define functionality for the download button to zip the files and allow download
     def prepare_download():
-        data = create_zip()
-        return gr.update(value={"name": "processed_images_annotations.zip", "data": data}, visible=True)
+        zip_path = create_zip()
+        return zip_path
 
-    prepare_download_button.click(
+    download_button.click(
         prepare_download,
-        outputs=download_button
+        outputs=download_file
     )
 
 # Launch the interface
-interface.launch()
+interface.launch(share=True)