deepfake-recognition / inference.py
SivaMallikarjun's picture
Create inference.py
ac2f4c2 verified
raw
history blame contribute delete
632 Bytes
import torch
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
from PIL import Image
# Load model and feature extractor
model = AutoModelForImageClassification.from_pretrained("your-username/deepfake-recognition")
feature_extractor = AutoFeatureExtractor.from_pretrained("your-username/deepfake-recognition")
# Load an image
image = Image.open("sample_image.jpg")
inputs = feature_extractor(images=image, return_tensors="pt")
# Predict
outputs = model(**inputs)
predicted_class = torch.argmax(outputs.logits, dim=1).item()
print(f"Predicted Class: {'Deepfake' if predicted_class == 1 else 'Real'}")