Update app.py

app.py

@@ -3,6 +3,8 @@ from transformers import ViTForImageClassification, ViTImageProcessor
 from PIL import Image
 import torch
 import os
+import shutil
+import zipfile
 
 # Load model and processor
 model_name = "prithivMLmods/Deepfake-QualityAssess-88M"
@@ -10,11 +12,14 @@ model = ViTForImageClassification.from_pretrained(model_name)
 processor = ViTImageProcessor.from_pretrained(model_name)
 
 # Create output directories
-os.makedirs("Issue In Deepfake", exist_ok=True)
-os.makedirs("High Quality Deepfake", exist_ok=True)
+issue_dir = "Issue In Deepfake"
+high_quality_dir = "High Quality Deepfake"
+zip_filename = "classified_images.zip"
+os.makedirs(issue_dir, exist_ok=True)
+os.makedirs(high_quality_dir, exist_ok=True)
 
 def deepfake_detection(images):
-    """Predicts deepfake probability scores for multiple images and saves them in respective folders."""
+    """Predicts deepfake probability scores for multiple images, saves them in respective folders, and returns a zip file."""
     results = {}
     
     for image in images:
@@ -36,20 +41,28 @@ def deepfake_detection(images):
         
         # Save image in respective folder
         filename = f"{max_label}_{max_score:.3f}.png"
-        save_path = os.path.join("Issue In Deepfake" if "issue" in max_label.lower() else "High Quality Deepfake", filename)
+        save_dir = issue_dir if "issue" in max_label.lower() else high_quality_dir
+        save_path = os.path.join(save_dir, filename)
         image.save(save_path)
         
         results[image_name] = {"predictions": predictions, "saved_as": filename}
     
-    return results
+    # Create a zip file
+    with zipfile.ZipFile(zip_filename, "w") as zipf:
+        for folder in [issue_dir, high_quality_dir]:
+            for root, _, files in os.walk(folder):
+                for file in files:
+                    zipf.write(os.path.join(root, file), os.path.relpath(os.path.join(root, file)))
+    
+    return results, zip_filename
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=deepfake_detection,
     inputs=gr.File(file_types=["image"], file_count="multiple", label="Upload Images"),
-    outputs=gr.JSON(label="Prediction Scores"),
+    outputs=[gr.JSON(label="Prediction Scores"), gr.File(label="Download Classified Images (ZIP)")],
     title="Deepfake Quality Detection",
-    description="Upload multiple images to check their deepfake probability scores. Images will be classified and stored accordingly."
+    description="Upload multiple images to check their deepfake probability scores. Images will be classified, stored, and available for download as a zip file."
 )
 
 # Launch the app