Update app.py

app.py

@@ -1,6 +1,11 @@
 import json
+import os
+import zipfile
+from io import BytesIO
+from tempfile import NamedTemporaryFile
 
 import gradio as gr
+import pandas as pd
 from PIL import Image
 import safetensors.torch
 import spaces
@@ -10,6 +15,8 @@ import torch
 from torchvision.transforms import transforms
 from torchvision.transforms import InterpolationMode
 import torchvision.transforms.functional as TF
+from torch.utils.data import Dataset, DataLoader
+
 
 torch.set_grad_enabled(False)
 
@@ -132,12 +139,11 @@ for idx, tag in enumerate(allowed_tags):
 
 sorted_tag_score = {}
 
-@spaces.GPU(duration=5)
+@spaces.GPU(duration=9)
 def run_classifier(image, threshold):
     global sorted_tag_score
     img = image.convert('RGB')
-    tensor = transform(img).unsqueeze(0)
-    tensor = tensor.to(device)
+    tensor = transform(img).unsqueeze(0).to(device)
     with torch.no_grad():
         logits = model(tensor)
         probabilities = torch.nn.functional.sigmoid(logits[0])
@@ -156,7 +162,84 @@ def create_tags(threshold):
     filtered_tag_score = {key: value for key, value in sorted_tag_score.items() if value > threshold}
     text_no_impl = ", ".join(filtered_tag_score.keys())
     return text_no_impl, filtered_tag_score
+
+    
+class ImageDataset(Dataset):
+    def __init__(self, image_files, transform):
+        self.image_files = image_files
+        self.transform = transform
+    
+    def __len__(self):
+        return len(self.image_files)
     
+    def __getitem__(self, idx):
+        img_path = self.image_files[idx]
+        img = Image.open(img_path).convert('RGB')
+        return self.transform(img), os.path.basename(img_path)
+
+
+@spaces.GPU(duration=299)
+def process_images(images, threshold):
+    dataset = ImageDataset(images, transform)
+
+    dataloader = DataLoader(dataset, batch_size=64, num_workers=0, pin_memory=True, drop_last=False)
+    
+    all_results = []
+
+    with torch.no_grad():
+        for batch, filenames in dataloader:
+            
+            batch = batch.to(device) 
+            with torch.no_grad():
+                logits = model(batch)
+                probabilities = torch.nn.functional.sigmoid(logits)
+            
+                for i, prob in enumerate(probabilities):
+                    indices = torch.where(prob > threshold)[0]
+                    values = prob[indices]
+                    
+                    temp = []
+                    tag_score = dict()
+                    for j in range(indices.size(0)):
+                        temp.append([allowed_tags[indices[j]], values[j].item()])
+                        tag_score[allowed_tags[indices[j]]] = values[j].item()
+                    
+                    tags = ", ".join([t[0] for t in temp])
+                    all_results.append((filenames[i], tags, tag_score))
+
+    return all_results
+
+def is_valid_image(file_path):
+    try:
+        with Image.open(file_path) as img:
+            img.verify()
+        return True
+    except:
+        return False
+
+def process_zip(zip_file, threshold):
+    if zip_file is None:
+        return None, None
+
+    temp_dir = "temp_images"
+    os.makedirs(temp_dir, exist_ok=True)
+
+    with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
+        zip_ref.extractall(temp_dir)
+    
+    all_files = [os.path.join(temp_dir, f) for f in os.listdir(temp_dir)]
+    image_files = [f for f in all_files if is_valid_image(f)]
+    results = process_images(image_files, threshold)
+
+    temp_file = NamedTemporaryFile(delete=False, suffix=".zip")
+    with zipfile.ZipFile(temp_file, "w") as zip_ref:
+        for image_name, text_no_impl, _ in results:
+            with zip_ref.open(''.join(image_name.split('.')[:-1]) + ".txt", 'w') as file:
+                file.write(text_no_impl.encode())
+    temp_file.seek(0)
+    df = pd.DataFrame([(os.path.basename(f), t) for f, t, _ in results], columns=['Image', 'Tags'])
+    
+    return temp_file.name, df
 
 with gr.Blocks(css=".output-class { display: none; }") as demo:
     gr.Markdown("""
@@ -165,25 +248,44 @@ with gr.Blocks(css=".output-class { display: none; }") as demo:
 
     This tagger is the result of joint efforts between members of the RedRocket team.  Special thanks to Minotoro at frosting.ai for providing the compute power for this project.
     """)
-    with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(label="Source", sources=['upload'], type='pil', height=512, show_label=False)
-            threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
-        with gr.Column():
-            tag_string = gr.Textbox(label="Tag String")
-            label_box = gr.Label(label="Tag Predictions", num_top_classes=250, show_label=False)
-
-    image_input.upload(
-        fn=run_classifier,
-        inputs=[image_input, threshold_slider],
-        outputs=[tag_string, label_box]
-    )
-
-    threshold_slider.input(
-        fn=create_tags,
-        inputs=[threshold_slider],
-        outputs=[tag_string, label_box]
-    )
+    
+    with gr.Tabs():
+        with gr.TabItem("Single Image"):
+            with gr.Row():
+                with gr.Column():
+                    image_input = gr.Image(label="Source", sources=['upload'], type='pil', height=512, show_label=False)
+                    threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
+                with gr.Column():
+                    tag_string = gr.Textbox(label="Tag String")
+                    label_box = gr.Label(label="Tag Predictions", num_top_classes=250, show_label=False)
+
+            image_input.upload(
+                fn=run_classifier,
+                inputs=[image_input, threshold_slider],
+                outputs=[tag_string, label_box]
+            )
+
+            threshold_slider.input(
+                fn=create_tags,
+                inputs=[threshold_slider],
+                outputs=[tag_string, label_box]
+            )
+        
+        with gr.TabItem("Multiple Images"):
+            with gr.Row():
+                with gr.Column():
+                    zip_input = gr.File(label="Upload ZIP file", file_types=['.zip'])
+                    multi_threshold_slider = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
+                    process_button = gr.Button("Process Images")
+                with gr.Column():
+                    zip_output = gr.File(label="Download Tagged Text Files (ZIP)")
+                    dataframe_output = gr.Dataframe(label="Image Tags Summary")
+            
+            process_button.click(
+                fn=process_zip,
+                inputs=[zip_input, multi_threshold_slider],
+                outputs=[zip_output, dataframe_output]
+            )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.queue().launch()