File size: 2,661 Bytes
d33766a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
from typing import Dict, List, Any
from io import BytesIO
from PIL import Image
import torch
import base64
import numpy as np
import cv2
import albumentations as A
from albumentations.pytorch import ToTensorV2
from safetensors.torch import load_file
# Import your model definition
from models import DeepfakeDetector
class EndpointHandler:
def __init__(self, path="."):
# Load model definition
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.device = device
self.model = DeepfakeDetector(pretrained=False) # Architecture only
# Load weights
try:
# Try loading safetensors
state_dict = load_file(f"{path}/best_model.safetensors")
self.model.load_state_dict(state_dict, strict=False)
except Exception as e:
print(f"Error loading weights: {e}")
# Fallback path if necessary
state_dict = load_file("best_model.safetensors")
self.model.load_state_dict(state_dict, strict=False)
self.model.to(device)
self.model.eval()
# Define transform
self.transform = A.Compose([
A.Resize(224, 224),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2(),
])
def __call__(self, data: Any) -> List[Dict[str, Any]]:
inputs = data.pop("inputs", data)
# Decode image
image = None
if isinstance(inputs, Image.Image):
image = inputs
elif isinstance(inputs, str):
# Try base64
try:
if "base64," in inputs:
inputs = inputs.split("base64,")[1]
image_bytes = base64.b64decode(inputs)
image = Image.open(BytesIO(image_bytes))
except:
# Url?
pass
elif isinstance(inputs, bytes):
image = Image.open(BytesIO(inputs))
if image is None:
return [{"error": "Invalid input format"}]
image = image.convert("RGB")
image_np = np.array(image)
# Augmentations expect numpy array
augmented = self.transform(image=image_np)
image_tensor = augmented['image'].unsqueeze(0).to(self.device)
# Inference
with torch.no_grad():
output = self.model(image_tensor)
prob = torch.sigmoid(output).item()
label = "FAKE" if prob > 0.5 else "REAL"
score = prob if prob > 0.5 else 1 - prob
return [{"label": label, "score": score}]
|