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"