DETR-DEMO / main.py
cvarandas's picture
first commit - demo
53a413a
Raw
History Blame Contribute Delete
1.1 kB
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")