Create app.py

app.py

@@ -0,0 +1,107 @@
+!pip install gradio
+import gradio as gr
+from transformers import AutoImageProcessor, SiglipForImageClassification
+from torch.optim import AdamW
+from PIL import Image
+import torch
+from torch.utils.data import Dataset, DataLoader
+import os
+
+# Load model and processor
+model_name = "prithivMLmods/deepfake-detector-model-v1"
+processor = AutoImageProcessor.from_pretrained(model_name)
+model = SiglipForImageClassification.from_pretrained(model_name)
+model.train()
+
+# Device setup
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Labels mapping
+id2label = {0: "FAKE", 1: "REAL"}
+label2id = {"FAKE": 0, "REAL": 1}
+
+# Optimizer for fine-tuning
+optimizer = AdamW(model.parameters(), lr=5e-6)
+
+# Dataset class for single example fine-tuning
+class SingleImageDataset(Dataset):
+    def __init__(self, image, label):
+        self.image = image
+        self.label = label
+
+    def __len__(self):
+        return 1
+
+    def __getitem__(self, idx):
+        inputs = processor(images=self.image, return_tensors="pt")
+        inputs = {k: v.squeeze(0) for k,v in inputs.items()}
+        inputs['labels'] = torch.tensor(self.label)
+        return inputs
+
+def fine_tune(image, correct_label):
+    dataset = SingleImageDataset(image, correct_label)
+    dataloader = DataLoader(dataset, batch_size=1)
+
+    model.train()
+    for epoch in range(1):  # just 1 epoch for fast feedback
+        for batch in dataloader:
+            batch = {k: v.to(device) for k,v in batch.items()}
+            outputs = model(**batch)
+            loss = outputs.loss
+            optimizer.zero_grad()
+            loss.backward()
+            optimizer.step()
+    # Save the updated model locally
+    save_path = "./fine_tuned_model"
+    os.makedirs(save_path, exist_ok=True)
+    model.save_pretrained(save_path)
+    processor.save_pretrained(save_path)
+    return
+
+def predict(image):
+    model.eval()
+    inputs = processor(images=image, return_tensors="pt").to(device)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        pred_class = logits.argmax(-1).item()
+    return id2label[pred_class]
+
+def inference(image, feedback, correct_label_text):
+    if image is None:
+        return "Please upload an image.", None
+
+    prediction = predict(image)
+    message = f"Prediction: {prediction}"
+
+    if feedback == "Wrong":
+        if correct_label_text.upper() in label2id:
+            correct_label = label2id[correct_label_text.upper()]
+            fine_tune(image, correct_label)
+            message += f" | Model fine-tuned with correct label: {correct_label_text.upper()}"
+        else:
+            message += " | Please enter a valid correct label (REAL or FAKE)."
+
+    return message, image
+
+# Gradio UI setup
+title = "Deepfake Detector with Interactive Feedback and Fine-tuning"
+
+iface = gr.Interface(
+    fn=inference,
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Radio(["Correct", "Wrong"], label="Is the prediction correct?", value="Correct"),
+        gr.Textbox(label="If Wrong, enter correct label (REAL or FAKE)", lines=1, placeholder="REAL or FAKE")
+    ],
+    outputs=[
+        gr.Textbox(label="Output"),
+        gr.Image(type="pil", label="Uploaded Image")
+    ],
+    title=title,
+    live=False,
+    allow_flagging="never"
+)
+
+iface.launch()