from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse, HTMLResponse
from transformers import pipeline
from PIL import Image
import io, os, traceback
# Ensure Hugging Face cache directories are writable
os.environ["HF_HOME"] = "/app/cache"
os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
os.environ["HF_HUB_CACHE"] = "/app/cache/hub"
app = FastAPI()
# Load BLIP (stable image captioning model)
pipe = pipeline(
"image-to-text",
model="Salesforce/blip-image-captioning-base",
device=-1 # CPU (works on free tier)
)
@app.get("/")
def home():
return {
"message": "API is running. Use POST /predict with an image, or visit /upload to test in browser."
}
@app.get("/upload", response_class=HTMLResponse)
def upload_form():
return """
Upload an ID Image
"""
@app.post("/predict")
async def predict_gender(file: UploadFile = File(...)):
try:
# Read uploaded image
image_bytes = await file.read()
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# Run BLIP captioning
result = pipe(image, max_new_tokens=32)
caption = result[0]["generated_text"].strip()
# Simple heuristic for gender detection
gender = "unknown"
lower_caption = caption.lower()
if "male" in lower_caption or "man" in lower_caption:
gender = "male"
elif "female" in lower_caption or "woman" in lower_caption:
gender = "female"
return JSONResponse({
"caption": caption,
"gender": gender
})
except Exception as e:
traceback.print_exc()
return JSONResponse({"error": str(e)}, status_code=500)