| import gradio as gr | |
| from transformers import ViTForImageClassification, ViTImageProcessor | |
| from PIL import Image | |
| import torch | |
| MODEL_NAME = "prithivMLmods/Deep-Fake-Detector-v2-Model" | |
| processor = ViTImageProcessor.from_pretrained(MODEL_NAME) | |
| model = ViTForImageClassification.from_pretrained(MODEL_NAME) | |
| def classify_image(image): | |
| image = Image.fromarray(image).convert("RGB") | |
| inputs = processor(images=image, return_tensors="pt") | |
| with torch.no_grad(): | |
| outputs = model(**inputs) | |
| probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0] | |
| labels = model.config.id2label | |
| return { labels[i]: float(probs[i]) for i in range(len(probs)) } | |
| demo = gr.Interface( | |
| fn=classify_image, | |
| inputs=gr.Image(type="numpy"), | |
| outputs=gr.Label(num_top_classes=2, label="Prediction (Real vs Deepfake)"), | |
| title="Deepfake Detector (ViT)", | |
| description="Upload an image — model classifies it as Real or Deepfake." | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() | |