| 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" | |
| async def predict(file: UploadFile = File(...)): | |
| image = Image.open(io.BytesIO(await file.read())).convert("RGB") | |
| result = predict_image(image) | |
| return {"prediction": result} |