import torch
import timm
from PIL import Image
from torchvision import transforms
from fastapi import FastAPI, File, UploadFile
import io

app = FastAPI()

# load model
model = timm.create_model("efficientnet_b7", pretrained=False, num_classes=2)
model.load_state_dict(torch.load("best_model.pth", map_location="cpu")["model_state_dict"])
model.eval()

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485,0.456,0.406],[0.229,0.224,0.225])
])

def predict_image(image):
    image = transform(image).unsqueeze(0)
    with torch.no_grad():
        output = model(image)
        pred = torch.argmax(output, dim=1).item()
    return "Real" if pred == 1 else "Fake"

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
    result = predict_image(image)
    return {"prediction": result}