Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import requests
 import shutil
 import gradio as gr
-from concurrent.futures import ThreadPoolExecutor
 from zipfile import ZipFile
 import logging
 
@@ -24,6 +23,8 @@ def get_image_ids(batch_id: str) -> list[str]:
             image_ids.append(image_id)
         else:
             logging.warning(f"Unexpected id format: {item['id']}")
+    if not image_ids:
+        raise ValueError("No images found in the manifest.")
     logging.info(f"Found {len(image_ids)} images in batch {batch_id}")
     return image_ids
 
@@ -42,6 +43,7 @@ def download_image(url: str, dest: str) -> None:
         logging.info(f"Successfully downloaded image: {dest}")
     else:
         logging.error(f"Failed to download image from {url}. Status code: {response.status_code}")
+        raise Exception(f"Failed to download image from {url}. Status code: {response.status_code}")
     del response
 
 def download_image_by_image_id(image_id: str):
@@ -56,28 +58,22 @@ def download_image_by_image_id(image_id: str):
     dest = os.path.join(batch_id, image_id + ".jpg")
     download_image(url, dest)
 
-def download_batch_images(batch_id: str, workers: int = 2, progress=None):
+def download_batch_images(batch_id: str, progress=None):
     logging.info(f"Starting download for batch {batch_id}")
-    image_ids = get_image_ids(batch_id)
-    if not image_ids:
-        raise ValueError("No images found to download.")
-    total_images = len(image_ids)
 
-    if progress:
+    if progress is not None:
         progress(0, desc=f"Starting download for {batch_id}...")
 
-    def track_download(image_id, idx):
+    image_ids = get_image_ids(batch_id)
+    total_images = len(image_ids)
+
+    for idx, image_id in enumerate(image_ids):
         download_image_by_image_id(image_id)
         logging.info(f"Downloaded image {image_id}")
-        if progress:
-            # Update progress after each image
+        if progress is not None:
             current_progress = (idx + 1) / total_images
             progress(current_progress, desc=f"Downloading {image_id}...")
 
-    with ThreadPoolExecutor(max_workers=workers) as executor:
-        for idx, image_id in enumerate(image_ids):
-            executor.submit(track_download, image_id, idx)
-
     logging.info(f"Zipping downloaded images for batch {batch_id}")
     zip_filename = f"{batch_id}.zip"
     with ZipFile(zip_filename, 'w') as zipf:
@@ -88,20 +84,19 @@ def download_batch_images(batch_id: str, workers: int = 2, progress=None):
             else:
                 logging.warning(f"Image {img_path} does not exist and will not be zipped.")
 
-    if progress:
+    if progress is not None:
         progress(1, desc=f"Completed {batch_id}")
 
     logging.info(f"Completed download and zip for batch {batch_id}")
     return zip_filename
 
 def gradio_interface(batch_id_input, progress=gr.Progress()):
-    batch_id = batch_id_input.strip()
     try:
-        zip_file = download_batch_images(batch_id, progress=progress)
-        return zip_file, ""  # No error message
+        zip_file = download_batch_images(batch_id_input, progress=progress)
+        return zip_file 
     except Exception as e:
         logging.error(f"Error processing batch: {e}")
-        return None, f"Error: {str(e)}"  # Return error message
+        raise gr.Error(f"Error: {str(e)}")
 
 with gr.Blocks() as app:
     gr.Markdown("# Batch Image Downloader")
@@ -112,12 +107,11 @@ with gr.Blocks() as app:
             download_button = gr.Button("Download Images")
         with gr.Column():
             output_file = gr.File(label="Download Zip File")
-            error_message = gr.Textbox(label="Error Message", interactive=False)
 
     download_button.click(
         gradio_interface,
         inputs=[batch_id_input],
-        outputs=[output_file, error_message]
+        outputs=[output_file]
     )
 
 app.queue()