Update app.py

app.py

@@ -176,8 +176,7 @@ class ImageDataset(Dataset):
         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)
@@ -221,24 +220,75 @@ 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 tempfile.TemporaryDirectory() as temp_dir:
+        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'])
 
-    with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
-        zip_ref.extractall(temp_dir)
+    return temp_file.name, df
+
+@spaces.GPU(duration=120)  # Reduced GPU duration for less wait time...
+def process_images_light(images, threshold):
+    dataset = ImageDataset(images, transform)
+
+    dataloader = DataLoader(dataset, batch_size=16, num_workers=0, pin_memory=True, drop_last=False)
     
-    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'])
+    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 process_zip_light(zip_file, threshold):
+    if zip_file is None:
+        return None, None
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        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_light(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:
@@ -286,6 +336,21 @@ with gr.Blocks(css=".output-class { display: none; }") as demo:
                 inputs=[zip_input, multi_threshold_slider],
                 outputs=[zip_output, dataframe_output]
             )
+        with gr.TabItem("Multiple Images (Light)"):
+            with gr.Row():
+                with gr.Column():
+                    zip_input_light = gr.File(label="Upload ZIP file", file_types=['.zip'])
+                    multi_threshold_slider_light = gr.Slider(minimum=0.00, maximum=1.00, step=0.01, value=0.20, label="Threshold")
+                    process_button_light = gr.Button("Process Images (Light)")
+                with gr.Column():
+                    zip_output_light = gr.File(label="Download Tagged Text Files (ZIP)")
+                    dataframe_output_light = gr.Dataframe(label="Image Tags Summary")
+            
+            process_button_light.click(
+                fn=process_zip_light,
+                inputs=[zip_input_light, multi_threshold_slider_light],
+                outputs=[zip_output_light, dataframe_output_light]
+            )
 
 if __name__ == "__main__":
     demo.queue().launch()