main.py

from fastapi import FastAPI, status, HTTPException, File
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from ClassifyRetina import classify
from io import BytesIO
import numpy as np
import cv2

app = FastAPI()

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

@app.get("/")
def root():
    return "Welcome to DSMATICS!"

@app.post("/image", status_code=status.HTTP_201_CREATED)
def create_image(Image: bytes = File(...)):
    
    if Image:
        try:
            nparr = np.frombuffer(Image, dtype=np.uint8)
            img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
            count, labels, image = classify(img)
            _, im_png = cv2.imencode(".png", image)
            return StreamingResponse(BytesIO(im_png.tobytes()), media_type="image/png", headers={"Human-Count": str(count), "Labels":str(labels)})
        except Exception as e:
            raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error in processing image...")
    else:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND, detail="Image not found...")