from fastapi import FastAPI, UploadFile, File
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from typing import Annotated
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline
from PIL import Image

# Use a pipeline as a high-level helper
# https://huggingface.co/docs/transformers/main_classes/pipelines#transformers.ObjectDetectionPipeline

# FastAPI tutorial
# https://fastapi.tiangolo.com/tutorial/

app = FastAPI()

# Allow CORS for localhost
origins = [
    #"http://localhost:7860",
    "*"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

detector = pipeline(model="facebook/detr-resnet-50")

@app.post("/detect-image")
async def detect_image(image: Annotated[UploadFile, File()]):
    return detector(Image.open(image.file))

app.mount("/", StaticFiles(directory="static", html=True), name="static")

@app.get("/")
def index() -> FileResponse:
    return FileResponse(path="/app/static/index.html", media_type="text/html")