File size: 1,419 Bytes
dc39c99 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 |
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image
import io
from pathlib import Path
from model import YOLOModel
import shutil
yolo = YOLOModel()
UPLOAD_FOLDER = Path("./uploads")
UPLOAD_FOLDER.mkdir(exist_ok=True)
app = FastAPI()
@app.post("/upload")
async def upload_image(image: UploadFile = File(...)):
# print(f'\n\t\tUPLOADED!!!!')
try:
file_path = UPLOAD_FOLDER / image.filename
with file_path.open("wb") as buffer:
shutil.copyfileobj(image.file, buffer)
# print(f'Starting to pass into model, {file_path}')
# Perform YOLO inference
predictions = yolo.predict(str(file_path))
print(f'\n\n\n{predictions}\n\n\ \n\t\t\t\tare predictions')
# Clean up uploaded file
file_path.unlink() # Remove file after processing
return JSONResponse(content={"items": predictions})
except Exception as e:
return JSONResponse(content={"error": str(e)}, status_code=500)
# code to accept the localhost to get images from
app.add_middleware(
CORSMiddleware,
allow_origins=["http://192.168.56.1:3000", "http://192.168.56.1:3001"],
allow_methods=["*"],
allow_headers=["*"],
)
if __name__ == "__main__":
import uvicorn
uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)
|