allfahad's picture
Update app.py
8a50b81 verified
Raw
History Blame Contribute Delete
1.05 kB
from fastapi import FastAPI, File, UploadFile
from transformers import pipeline
from PIL import Image
import io
app = FastAPI()
# Load zero-shot image classification model
classifier = pipeline("zero-shot-image-classification",
model="google/siglip2-base-patch16-224",
device=-1) # -1 = CPU
@app.get("/")
def read_root():
return {"message": "Fish Detector API is running. Use POST /detect"}
@app.post("/detect")
async def detect(file: UploadFile = File(...)):
# Read the uploaded image
contents = await file.read()
image = Image.open(io.BytesIO(contents)).convert("RGB")
# Candidate labels
candidate_labels = [
"a clear photo of a real fish",
"a photo with absolutely no fish in it"
]
# Perform zero-shot classification
predictions = classifier(image, candidate_labels=candidate_labels)
# Check if the top prediction is "fish"
if predictions[0]['label'] == candidate_labels[0]:
return "fish"
else:
return "not a fish"