from pathlib import Path
from PIL import Image
import torch
from transformers import AutoImageProcessor, ViTForImageClassification
from fastapi import FastAPI, File, UploadFile
from io import BytesIO
import shutil

model_path = "./best_model"

app = FastAPI()

processor = AutoImageProcessor.from_pretrained(
    model_path,
    local_files_only=True,
    use_fast=True
)

model = ViTForImageClassification.from_pretrained(
    model_path,
    local_files_only=True,
    id2label={"0": "real", "1": "fake"},
    label2id={"real": 0, "fake": 1}
)

def predict_image(image: Image.Image):
    inputs = processor(image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        pred_id = torch.argmax(outputs.logits, dim=1).item()
        pred_label = model.config.id2label[str(pred_id)]
    return pred_label

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