add zip downloading

app.py

@@ -1,8 +1,12 @@
-import gradio as gr
+
 import os
 import shutil
 import cv2
+import zipfile
+import uuid
+
 import numpy as np
+import gradio as gr
 
 from naming import im2c
 from collections import Counter
@@ -31,45 +35,64 @@ def get_top_names(img):
     filtered_counts = Counter(name_idx_img[name_idx_img <= 10])
     sorted_counts = sorted(filtered_counts.items(), key=lambda x: x[1], reverse=True)
     top_3_values = [num for num, count in sorted_counts[:3]]
+    top_3_counts = [count/(dim[0]*dim[1])  for num, count in sorted_counts[:3]]
     top_3_colors = [COLOR_NAME[i] for i in top_3_values]
-    # print("Top 3 colors:", top_3_colors)
-    return top_3_values, top_3_colors
+    # print("Top 3 colors:", top_3_counts)
+    return top_3_values, top_3_counts, top_3_colors
 
 
 def classify_and_log(images):
-    output_folder = "classified_results"
-    os.makedirs(output_folder, exist_ok=True)
+    # output_folder = "classified_results"
+    # os.makedirs(output_folder, exist_ok=True)
+
+     # create a temporary directory
+    session_id = str(uuid.uuid4())
+    output_dir = f"temp_{session_id}"
+    os.makedirs(output_dir, exist_ok=True)
 
-    category_folders = {i: os.path.join(output_folder, COLOR_NAME[i]) for i in range(11)}
+
+    category_folders = {i: os.path.join(output_dir, COLOR_NAME[i]) for i in range(11)}
     for folder in category_folders.values():
         os.makedirs(folder, exist_ok=True)
 
-    log_file = os.path.join(output_folder, "top3colors.txt")
+    log_file = os.path.join(output_dir, "top3colors.txt")
     
     results = {i: [] for i in range(11)}
 
+
     with open(log_file, "w") as log:
         for id_img, img in enumerate(images):
             filename = os.path.basename(img.name)
             img_array = cv2.imread(img).astype(np.float32)
             img_array = cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB)
 
-            cat_id, category = get_top_names(img_array)
-            # print(category_folders[cat_id[0]], filename)
-            target_path = os.path.join(category_folders[cat_id[0]], filename)
+            cat_id, cat_counts, category = get_top_names(img_array)
 
-            cv2.imwrite(target_path, cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB))
+            for i in range(3):
+                if cat_counts[i] > 0.15:
+                    target_path = os.path.join(category_folders[cat_id[i]], filename)
+                    cv2.imwrite(target_path, cv2.cvtColor(img_array, cv2.COLOR_BGR2RGB))
 
-            print(f"Image:{filename} -> Top 3 colors:{category}\n")
+            # print(f"Image:{filename} -> Top 3 colors:{category}\n")
 
-            log.write(f"{filename} -> 1 {category[0]}, 2 {category[1]}, 3 {category[2]}\n")
+            log.write(f"{filename} -> 1 {category[0]} {100*cat_counts[0]:.2f}%, 2 {category[1]} {100*cat_counts[1]:.2f}%, 3 {category[2]} {100*cat_counts[2]:.2f}%\n")
 
             results[cat_id[0]].append(target_path)
 
-    # output_images = [results[i] for i in range(11)]
-    # print(output_images)
 
-    return log_file
+    # compile all images into a zip file
+    zip_path = f"{output_dir}.zip"
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as zipf:
+        for root, _, files in os.walk(output_dir):
+            for file in files:
+                file_path = os.path.join(root, file)
+                arcname = os.path.relpath(file_path, start=output_dir)
+                zipf.write(file_path, arcname)
+
+    # optional: clean up the output directory
+    shutil.rmtree(output_dir)
+
+    return zip_path
 
 
 def swap_to_gallery(images):
@@ -85,6 +108,12 @@ def remove_back_to_files():
 
 with gr.Blocks() as demo:
     gr.Markdown("## Image color categorization")
+    gr.Markdown("We categorize images into different classes based on the frequency of different colors appearing in each image.")
+    gr.Markdown("The 11 color catergories include: black, brown, blue, gray, green, orange, pink, purple, red, white and yellow.")
+    gr.Markdown("The classification is based on the color naming model from paper _Van De Weijer, Joost, et al. Learning color names for real-world applications. IEEE Transactions on Image Processing 18.7 (2009): 1512-1523._")
+    gr.Markdown("The output results are in a zip file with all the images in the correspoding folders.")
+    gr.Markdown("Note that one image can be classified into multiple categories (top 3 categories and frequency > 15%), and the top 3 categories are listed in the log file.")
+
 
     with gr.Row():
         with gr.Column():