Update app.py

app.py

@@ -10,6 +10,7 @@ import io
 import uuid
 import time
 import random
+import zipfile
 
 DATA_DIR = "/data"
 IMAGES_DIR = os.path.join(DATA_DIR, "images")
@@ -92,6 +93,29 @@ class DatasetBuilder:
         with open(dataset_file, 'w') as f:
             json.dump(self.dataset, f)
 
+    def create_downloadable_dataset(self):
+        if not self.dataset:
+            return None, "Dataset is empty. Add some images first."
+
+        try:
+            # Create a temporary ZIP file
+            zip_filename = f"{self.dataset_name}.zip"
+            zip_path = os.path.join(DATA_DIR, zip_filename)
+            
+            with zipfile.ZipFile(zip_path, 'w') as zipf:
+                # Add the dataset JSON file
+                dataset_file = self.get_dataset_file()
+                zipf.write(dataset_file, os.path.basename(dataset_file))
+                
+                # Add all images
+                for item in self.dataset:
+                    image_path = os.path.join(IMAGES_DIR, item['image'])
+                    zipf.write(image_path, os.path.join("images", item['image']))
+
+            return zip_path, f"Dataset '{self.dataset_name}' ready for download."
+        except Exception as e:
+            return None, f"Error creating downloadable dataset: {str(e)}"
+
     def add_image(self, url, cookies=None):
         try:
             response = make_request(url, cookies)
@@ -201,7 +225,12 @@ def view_dataset(dataset_name):
 
 def upload_huggingface_dataset(dataset_name, privacy):
     builder = DatasetBuilder(dataset_name)
-    return builder.upload_to_huggingface(private=privacy)    
+    return builder.upload_to_huggingface(private=privacy) 
+
+def download_dataset(dataset_name):
+    builder = DatasetBuilder(dataset_name)
+    zip_path, message = builder.create_downloadable_dataset()
+    return zip_path, message
 
 # Create Gradio interface
 with gr.Blocks(theme="huggingface") as iface:
@@ -236,5 +265,16 @@ with gr.Blocks(theme="huggingface") as iface:
     hf_upload_result = gr.Textbox(label="Upload Result")
     upload_hf_button.click(upload_huggingface_dataset, inputs=[dataset_name_input, privacy_radio], outputs=hf_upload_result)
 
+    gr.Markdown("## Download Dataset")
+    download_button = gr.Button("Download Dataset")
+    download_output = gr.File(label="Download")
+    download_message = gr.Textbox(label="Download Status")
+    
+    download_button.click(
+        download_dataset,
+        inputs=[dataset_name_input],
+        outputs=[download_output, download_message]
+    )
+
 # Launch the interface
 iface.launch()