File size: 714 Bytes
03767fe
 
 
 
 
 
 
 
5bf4ffc
 
03767fe
 
 
 
 
 
5bf4ffc
 
 
 
99cf18a
03767fe
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from transformers import pipeline
from PIL import Image
import io

app = FastAPI()

# Initialize SmolVLM for image-to-text tasks
pipe = pipeline("image-to-text", model="HuggingFaceTB/SmolVLM-Instruct", device=-1)

@app.post("/predict")
async def predict_gender(file: UploadFile = File(...)):
    image_bytes = await file.read()
    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
    
    # Prompt the model to identify gender
    prompt = "Is the person on this ID male or female?"
    result = pipe(image, prompt)
    answer = result[0]['generated_text'].strip()
    
    return JSONResponse({"gender": answer})