Skip to content

Commit

Permalink
Add healthcheck and perform resizing on server-side (#80)
Browse files Browse the repository at this point in the history
  • Loading branch information
bgoelTT authored Jan 28, 2025
1 parent 0bbdde0 commit 70320a5
Showing 1 changed file with 15 additions and 0 deletions.
15 changes: 15 additions & 0 deletions tt-metal-yolov4/server/fast_api_yolov4.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import os
import logging
from fastapi import FastAPI, File, HTTPException, Request, status, UploadFile
from fastapi.responses import JSONResponse
from functools import wraps
from io import BytesIO
import jwt
Expand Down Expand Up @@ -63,6 +64,8 @@ def load_class_names(namesfile):
class_names = load_class_names(file_path)

global model
global ready
ready = False
if ("WH_ARCH_YAML" in os.environ) and os.environ[
"WH_ARCH_YAML"
] == "wormhole_b0_80_arch_eth_dispatch.yaml":
Expand Down Expand Up @@ -92,6 +95,17 @@ def load_class_names(namesfile):
ttnn.enable_program_cache(device)
model = Yolov4Trace2CQ()
model.initialize_yolov4_trace_2cqs_inference(device)
ready = True


@app.get("/health")
async def health_check():
if not ready:
raise HTTPException(
status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
detail="Server is not ready yet",
)
return JSONResponse(content={"message": "OK\n"}, status_code=status.HTTP_200_OK)


@app.on_event("shutdown")
Expand Down Expand Up @@ -259,6 +273,7 @@ async def objdetection_v2(request: Request, file: UploadFile = File(...)):
contents = await file.read()
# Load and convert the image to RGB
image = Image.open(BytesIO(contents)).convert("RGB")
image = image.resize((320, 320)) # Resize to target dimensions
image = np.array(image)
if isinstance(image, np.ndarray) and len(image.shape) == 3: # cv2 image
image = torch.from_numpy(image).float().div(255.0).unsqueeze(0)
Expand Down

0 comments on commit 70320a5

Please sign in to comment.