Create app.py

app.py

@@ -0,0 +1,78 @@
+import gradio as gr
+from transformers import AutoImageProcessor, SiglipForImageClassification
+from PIL import Image
+import torch
+import os
+import zipfile
+
+# Load model and processor
+model_name = "prithivMLmods/Watermark-Detection-SigLIP2"
+model = SiglipForImageClassification.from_pretrained(model_name)
+processor = AutoImageProcessor.from_pretrained(model_name)
+
+# Label mapping
+id2label = {
+    "0": "No Watermark",
+    "1": "Watermark"
+}
+
+# Output folders
+watermark_dir = "Watermarked"
+no_watermark_dir = "No_Watermark"
+zip_filename = "watermark_classified_images.zip"
+os.makedirs(watermark_dir, exist_ok=True)
+os.makedirs(no_watermark_dir, exist_ok=True)
+
+def classify_and_save_watermarks(images):
+    results = {}
+
+    for image_file in images:
+        image_name = os.path.basename(image_file.name)
+        image = Image.open(image_file).convert("RGB")
+        inputs = processor(images=image, return_tensors="pt")
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+            logits = outputs.logits
+            probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
+
+        prediction = {id2label[str(i)]: round(probs[i], 3) for i in range(len(probs))}
+        max_label = max(prediction, key=prediction.get)
+        max_score = prediction[max_label]
+
+        # Save image to appropriate folder
+        filename = f"{max_label.replace(' ', '_')}_{max_score:.3f}.png"
+        save_dir = watermark_dir if "Watermark" in max_label else no_watermark_dir
+        save_path = os.path.join(save_dir, filename)
+        image.save(save_path)
+
+        results[image_name] = {
+            "predictions": prediction,
+            "saved_as": filename
+        }
+
+    # Create zip of both folders
+    with zipfile.ZipFile(zip_filename, "w") as zipf:
+        for folder in [watermark_dir, no_watermark_dir]:
+            for root, _, files in os.walk(folder):
+                for file in files:
+                    file_path = os.path.join(root, file)
+                    arcname = os.path.relpath(file_path, start=os.path.dirname(folder))
+                    zipf.write(file_path, arcname)
+
+    return results, zip_filename
+
+# Gradio interface
+iface = gr.Interface(
+    fn=classify_and_save_watermarks,
+    inputs=gr.File(file_types=["image"], file_count="multiple", label="Upload Images"),
+    outputs=[
+        gr.JSON(label="Watermark Predictions"),
+        gr.File(label="Download Classified Images (ZIP)")
+    ],
+    title="Watermark Detection and Classification",
+    description="Upload multiple images to detect watermarks. Images will be saved in 'Watermarked' or 'No Watermark' folders and available for download as a ZIP."
+)
+
+if __name__ == "__main__":
+    iface.launch(ssr_mode=True)