File size: 1,011 Bytes
17bbba1 | 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 | from transformers import pipeline
from PIL import Image
import base64
import io
from typing import Dict, List, Any
class EndpointHandler():
def __init__(self, path=""): # Load model
self.classifier = pipeline("image-classification", model="kpeng-05/fake-face-classification-lite")
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, float]]:
"""
Args:
data (Dict[str, Any]):
Dictionary containing "inputs" key with image url as value
Return:
A list of dictionaries containing:
- "label": A string representing the class (e.g., "real" or "fake")
- "score": A float between 0 and 1 for confidence
"""
# Extract URL
data_url = data.get("inputs", data)
encoded_data = data_url.split(',')[1]
decoded_data = base64.b64decode(encoded_data)
image = Image.open(io.BytesIO(decoded_data))
results = self.classifier(image)
return results |