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=False)