Update app.py

app.py

@@ -94,17 +94,14 @@ def process_and_display(gender, num_images):
     # Fetch image URLs
     image_urls = fetch_image_urls(gender, num_images)
     if not image_urls:
-        return "Failed to fetch image URLs.", None, None, None, None
+        return "Failed to fetch image URLs.", None, None, None
 
     # Download images
     downloaded_count, image_paths = download_images(image_urls)
     if downloaded_count == 0:
-        return "No images were successfully downloaded.", None, None, None, None
+        return "No images were successfully downloaded.", None, None, None
 
-    # Prepare gallery data (image paths for thumbnails)
-    gallery_data = [(path, f"Image {i+1}") for i, path in enumerate(image_paths)]
-
-    return f"Successfully downloaded {downloaded_count}/{num_images} images. Select images to include in ZIP below.", None, None, gallery_data, image_paths
+    return f"Successfully downloaded {downloaded_count}/{num_images} images. Select images to include in ZIP below.", None, image_paths, image_paths
 
 def process_zip_submission(image_paths, *checkbox_states):
     """Create a ZIP file based on the selected images."""
@@ -143,34 +140,42 @@ with gr.Blocks(title="Image Downloader") as demo:
     zip_output = gr.File(label="Download ZIP", visible=False)
     
     gr.Markdown("### Image Gallery (Click Thumbnails to View Full Size)")
-    gallery_output = gr.Gallery(label="Gallery", show_label=True, elem_id="gallery", columns=[5], height="auto")
     image_paths_state = gr.State()  # Store image paths for later use
     
-    # Checkbox components for toggling inclusion in ZIP
-    checkbox_outputs = [gr.Checkbox(label=f"Include Image {i+1} in ZIP", value=True, visible=False) for i in range(25)]  # Max 25 images
+    # Create a grid layout for images and checkboxes
+    image_outputs = []
+    checkbox_outputs = []
+    with gr.Row():
+        for i in range(25):  # Max 25 images
+            with gr.Column():
+                image_outputs.append(gr.Image(label=f"Image {i+1}", visible=False, width=100, height=100))
+                checkbox_outputs.append(gr.Checkbox(label=f"Include in ZIP", value=True, visible=False))
     
     gr.Markdown("### Submit Selections")
     submit_button = gr.Button("Create ZIP of Selected Images")
     
     def on_download(gender, num_images):
-        status, zip_path, image_path, gallery_data, image_paths = process_and_display(gender, num_images)
+        status, zip_path, image_paths, image_paths_for_display = process_and_display(gender, num_images)
         if image_paths:
-            # Show checkboxes for the downloaded images
-            checkbox_updates = {checkbox_outputs[i]: gr.Checkbox(value=True, visible=True, label=f"Include Image {i+1} in ZIP") for i in range(len(image_paths))}
-            # Hide remaining checkboxes
-            checkbox_updates.update({checkbox_outputs[i]: gr.Checkbox(value=False, visible=False) for i in range(len(image_paths), 25)})
-            return {
+            # Update image and checkbox components for the downloaded images
+            updates = {
                 status_output: status,
                 zip_output: gr.File(value=None, visible=False),  # Hide ZIP initially
-                gallery_output: gallery_data,
-                image_paths_state: image_paths,
-                **checkbox_updates
+                image_paths_state: image_paths
             }
+            for i in range(25):
+                if i < len(image_paths):
+                    updates[image_outputs[i]] = gr.Image(value=image_paths[i], visible=True, label=f"Image {i+1}", width=100, height=100)
+                    updates[checkbox_outputs[i]] = gr.Checkbox(value=True, visible=True, label=f"Include in ZIP")
+                else:
+                    updates[image_outputs[i]] = gr.Image(value=None, visible=False)
+                    updates[checkbox_outputs[i]] = gr.Checkbox(value=False, visible=False)
+            return updates
         return {
             status_output: status,
             zip_output: gr.File(visible=False),
-            gallery_output: None,
             image_paths_state: None,
+            **{image_outputs[i]: gr.Image(value=None, visible=False) for i in range(25)},
             **{checkbox_outputs[i]: gr.Checkbox(value=False, visible=False) for i in range(25)}
         }
     
@@ -184,7 +189,7 @@ with gr.Blocks(title="Image Downloader") as demo:
     download_button.click(
         fn=on_download,
         inputs=[gender_input, num_images_input],
-        outputs=[status_output, zip_output, gallery_output, image_paths_state] + checkbox_outputs
+        outputs=[status_output, zip_output, image_paths_state] + image_outputs + checkbox_outputs
     )
     
     submit_button.click(