add download

app.py

@@ -1,12 +1,14 @@
 # Author: Huzheng Yang
 # %%
 import copy
+from datetime import datetime
 import pickle
 from functools import partial
 from io import BytesIO
 import json
 import os
 import uuid
+import zipfile
 
 from einops import rearrange
 from matplotlib import pyplot as plt
@@ -1280,15 +1282,15 @@ def make_input_video_section():
     return input_gallery, submit_button, clear_images_button, max_frames_number
 
 
-def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_random=False):
+def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_random=False, allow_download=False):
     gr.Markdown('### Input Images')
     input_gallery = gr.Gallery(value=None, label="Input images", show_label=True, elem_id="input_images", columns=[cols], rows=[rows], object_fit="contain", height=height, type="pil", show_share_button=False)
     
     submit_button = gr.Button("🔴 RUN", elem_id="submit_button", variant='primary')
     with gr.Row():
-        clear_images_button = gr.Button("🗑️Clear", elem_id='clear_button', variant='stop')
+        clear_images_button = gr.Button("🗑️ Clear", elem_id='clear_button', variant='stop')
         clear_images_button.click(fn=lambda: gr.update(value=None), outputs=[input_gallery])
-        upload_button = gr.UploadButton(elem_id="upload_button", label="⬆️Upload", variant='secondary', file_types=["image"], file_count="multiple")
+        upload_button = gr.UploadButton(elem_id="upload_button", label="⬆️ Upload", variant='secondary', file_types=["image"], file_count="multiple")
         
         def convert_to_pil_and_append(images, new_images):
             if images is None:
@@ -1304,7 +1306,9 @@ def make_input_images_section(rows=1, cols=3, height="auto", advanced=False, is_
             gr.Info(f"Total images: {len(images)}")
             return images
         upload_button.upload(convert_to_pil_and_append, inputs=[input_gallery, upload_button], outputs=[input_gallery])
-        
+    
+    if allow_download:
+        create_file_button, download_button = add_download_button(input_gallery, "input_images")
     
     gr.Markdown('### Load Datasets')
     load_images_button = gr.Button("🔴 Load Images", elem_id="load-images-button", variant='primary')
@@ -1596,7 +1600,114 @@ def add_output_images_buttons(output_gallery):
         flip_button = gr.Button("🔃 Flip", elem_id="flip_button", variant='secondary')
         flip_button.click(flip_rgb_gallery, inputs=[output_gallery], outputs=[output_gallery])
     return rotate_button, flip_button
+
+def add_download_button(gallery, filename_prefix="output"):
+    
+    def make_3x5_plot(images):
+        plot_list = []
+        
+        # Split the list of images into chunks of 15
+        chunks = [images[i:i + 15] for i in range(0, len(images), 15)]
+        
+        for chunk in chunks:
+            fig, axs = plt.subplots(3, 4, figsize=(12, 9))
+            for ax in axs.flatten():
+                ax.axis("off")
+            for ax, img in zip(axs.flatten(), chunk):
+                img = img.convert("RGB")
+                ax.imshow(img)
+            
+            plt.tight_layout(h_pad=0.5, w_pad=0.3)
+
+            # Generate a unique filename
+            filename = uuid.uuid4()
+            tmp_path = f"/tmp/{filename}.png"
+            
+            # Save the plot to the temporary file
+            plt.savefig(tmp_path, bbox_inches='tight', dpi=144)
+            
+            # Open the saved image
+            img = Image.open(tmp_path)
+            img = img.convert("RGB")
+            img = copy.deepcopy(img)
+            
+            # Remove the temporary file
+            os.remove(tmp_path)
+
+            plot_list.append(img)
+            plt.close()
+        
+        return plot_list
+    
+    def delete_file_after_delay(file_path, delay):
+        def delete_file():
+            if os.path.exists(file_path):
+                os.remove(file_path)
+        
+        timer = threading.Timer(delay, delete_file)
+        timer.start()
+    
+    def create_zip_file(images, filename_prefix=filename_prefix):
+        if images is None or len(images) == 0:
+            gr.Warning("No images selected.")
+            return None
+        gr.Info("Creating zip file for download...")
+        images = [image[0] for image in images]
+        if isinstance(images[0], str):
+            images = [Image.open(image) for image in images]
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        
+        zip_filename = f"/tmp/gallery_download/{filename_prefix}_{timestamp}.zip"
+        os.makedirs(os.path.dirname(zip_filename), exist_ok=True)
+        
+        plots = make_3x5_plot(images)
+        
+        
+        
+        with zipfile.ZipFile(zip_filename, 'w') as zipf:
+            # Create a temporary directory to store images and plots
+            temp_dir = f"/tmp/gallery_download/images/{uuid.uuid4()}"
+            os.makedirs(temp_dir)
+            
+            try:
+                # Save images to the temporary directory
+                for i, img in enumerate(images):
+                    img = img.convert("RGB")
+                    img_path = os.path.join(temp_dir, f"single_{i:04d}.jpg")
+                    img.save(img_path)
+                    zipf.write(img_path, f"single_{i:04d}.jpg")
+                
+                # Save plots to the temporary directory
+                for i, plot in enumerate(plots):
+                    plot = plot.convert("RGB")
+                    plot_path = os.path.join(temp_dir, f"grid_{i:04d}.jpg")
+                    plot.save(plot_path)
+                    zipf.write(plot_path, f"grid_{i:04d}.jpg")
+            finally:
+                # Clean up the temporary directory
+                for file in os.listdir(temp_dir):
+                    os.remove(os.path.join(temp_dir, file))
+                os.rmdir(temp_dir)
+        
+        # Schedule the deletion of the zip file after 24 hours (86400 seconds)
+        delete_file_after_delay(zip_filename, 86400)
+        gr.Info(f"File is ready for download: {os.path.basename(zip_filename)}")
+        return gr.update(value=zip_filename, interactive=True)
     
+    with gr.Row():
+        create_file_button = gr.Button("📦 Pack", elem_id="create_file_button", variant='secondary')
+        download_button = gr.DownloadButton(label="📥 Download", value=None, variant='secondary', elem_id="download_button", interactive=False)
+        
+        create_file_button.click(create_zip_file, inputs=[gallery], outputs=[download_button])
+        def warn_on_click(filename):
+            if filename is None:
+                gr.Warning("No file to download, please `📦 Pack` first.")
+            interactive = filename is not None
+            return gr.update(interactive=interactive)
+        download_button.click(warn_on_click, inputs=[download_button], outputs=[download_button])
+    
+    return create_file_button, download_button
+
 
 def make_output_images_section():
     gr.Markdown('### Output Images')
@@ -1752,14 +1863,17 @@ with demo:
 
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
-                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section()
+                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section(allow_download=True)
                 num_images_slider.value = 30
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information", autofocus=False, autoscroll=False, lines=20)
                 
             with gr.Column(scale=5, min_width=200):
                 output_gallery = make_output_images_section()
+                add_download_button(output_gallery, "ncut_embed")
                 norm_gallery = gr.Gallery(value=[], label="Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
+                add_download_button(norm_gallery, "eig_norm")
                 cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=600, show_share_button=True, preview=True, interactive=False)
+                add_download_button(cluster_gallery, "clusters")
                 [
                     model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
                     affinity_focal_gamma_slider, num_sample_ncut_slider, ncut_knn_slider, ncut_indirect_connection, ncut_make_orthogonal, 
@@ -1921,27 +2035,36 @@ with demo:
 
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
-                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section()
+                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section(allow_download=True)
                 num_images_slider.value = 100
                 logging_text = gr.Textbox("Logging information", label="Logging", elem_id="logging", type="text", placeholder="Logging information", lines=20)
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #1)')
                 l1_gallery = gr.Gallery(format='png', value=[], label="Recursion #1", show_label=True, elem_id="ncut_l1", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
                 add_output_images_buttons(l1_gallery)
+                add_download_button(l1_gallery, "ncut_embed_recur1")
                 l1_norm_gallery = gr.Gallery(value=[], label="Recursion #1 Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
+                add_download_button(l1_norm_gallery, "eig_norm_recur1")
                 l1_cluster_gallery = gr.Gallery(value=[], label="Recursion #1 Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=500, show_share_button=True, preview=True, interactive=False)
+                add_download_button(l1_cluster_gallery, "clusters_recur1")
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #2)')
                 l2_gallery = gr.Gallery(format='png', value=[], label="Recursion #2", show_label=True, elem_id="ncut_l2", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
                 add_output_images_buttons(l2_gallery)
+                add_download_button(l2_gallery, "ncut_embed_recur2")
                 l2_norm_gallery = gr.Gallery(value=[], label="Recursion #2 Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
+                add_download_button(l2_norm_gallery, "eig_norm_recur2")
                 l2_cluster_gallery = gr.Gallery(value=[], label="Recursion #2 Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=500, show_share_button=True, preview=True, interactive=False)
+                add_download_button(l2_cluster_gallery, "clusters_recur2")
             with gr.Column(scale=5, min_width=200):
                 gr.Markdown('### Output (Recursion #3)')
                 l3_gallery = gr.Gallery(format='png', value=[], label="Recursion #3", show_label=True, elem_id="ncut_l3", columns=[3], rows=[5], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
                 add_output_images_buttons(l3_gallery)
+                add_download_button(l3_gallery, "ncut_embed_recur3")
                 l3_norm_gallery = gr.Gallery(value=[], label="Recursion #3 Eigenvector Magnitude", show_label=True, elem_id="eig_norm", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
+                add_download_button(l3_norm_gallery, "eig_norm_recur3")
                 l3_cluster_gallery = gr.Gallery(value=[], label="Recursion #3 Clusters", show_label=True, elem_id="clusters", columns=[2], rows=[4], object_fit="contain", height=500, show_share_button=True, preview=True, interactive=False)
+                add_download_button(l3_cluster_gallery, "clusters_recur3")
         
         with gr.Row():        
             with gr.Column(scale=5, min_width=200):
@@ -2165,16 +2288,19 @@ with demo:
         gr.Markdown('### Output (Recursion #1)')
         l1_gallery = gr.Gallery(format='png', value=[], label="Recursion #1", show_label=True, elem_id="ncut_l1", columns=[100], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False, preview=True)
         add_output_images_buttons(l1_gallery)
+        add_download_button(l1_gallery, "modelaligned_recur1")
         gr.Markdown('### Output (Recursion #2)')
         l2_gallery = gr.Gallery(format='png', value=[], label="Recursion #2", show_label=True, elem_id="ncut_l2", columns=[100], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False, preview=True)
         add_output_images_buttons(l2_gallery)
+        add_download_button(l2_gallery, "modelaligned_recur2")
         gr.Markdown('### Output (Recursion #3)')
         l3_gallery = gr.Gallery(format='png', value=[], label="Recursion #3", show_label=True, elem_id="ncut_l3", columns=[100], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False, preview=True)
         add_output_images_buttons(l3_gallery)  
+        add_download_button(l3_gallery, "modelaligned_recur3")
     
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
-                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section()
+                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section(allow_download=True)
                 num_images_slider.value = 100
                 
                 
@@ -2311,8 +2437,11 @@ with demo:
                 output_gallery = gr.Gallery(format='png', value=[], label="NCUT Embedding", show_label=True, elem_id=f"ncut{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto", show_fullscreen_button=True, interactive=False)
                 submit_button = gr.Button("🔴 RUN", elem_id=f"submit_button{i_model}", variant='primary')
                 add_output_images_buttons(output_gallery)
+                add_download_button(output_gallery, f"ncut_embed")
                 norm_gallery = gr.Gallery(value=[], label="Eigenvector Magnitude", show_label=True, elem_id=f"eig_norm{i_model}", columns=[3], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=False, interactive=False)
+                add_download_button(norm_gallery, f"eig_norm")
                 cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=True, elem_id=f"clusters{i_model}", columns=[2], rows=[4], object_fit="contain", height=500, show_share_button=True, preview=True, interactive=False)
+                add_download_button(cluster_gallery, f"clusters")
                 [
                     model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
                     affinity_focal_gamma_slider, num_sample_ncut_slider, ncut_knn_slider, ncut_indirect_connection, ncut_make_orthogonal, 
@@ -2342,7 +2471,7 @@ with demo:
 
         with gr.Row():
             with gr.Column(scale=5, min_width=200):
-                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section()
+                input_gallery, submit_button, clear_images_button, dataset_dropdown, num_images_slider, random_seed_slider, load_images_button = make_input_images_section(allow_download=True)
                 submit_button.visible = False