| 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} |
