app

Dockerfile

@@ -6,6 +6,8 @@ RUN pip install ultralytics
 RUN pip install opencv-python==4.6.0.66
 RUN pip install Pillow==10.3.0
 RUN pip install uvicorn fastapi
+RUN apt update && apt install fonts-dejavu -y
+RUN wget https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8l-pose.pt -O /app/yolov8l-pose.pt
 
 COPY . .
 

main.py

@@ -10,8 +10,54 @@ import utils
 from drawing import draw_keypoints
 
 
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print('Using device:', device)
+
+model_pose = YOLO('yolov8l-pose.pt')
+model_pose.to(device)
+
 app = FastAPI()
 
+
 @app.get("/")
-def read_root():
-    return {"Hello": "World!"}
+async def health():
+    return JSONResponse(content={"status": "ok"})
+
+
+@app.post("/predict-image")
+async def predict_image(file: UploadFile = File(...)):
+    contents = await file.read()
+    input_image = Image.open(io.BytesIO(contents)).convert("RGB")
+    input_image = ImageOps.exif_transpose(input_image)
+
+    # predict
+    result = model_pose(input_image)[0]
+    keypoints = utils.get_keypoints(result)
+
+    # draw keypoints
+    output_image = draw_keypoints(input_image, keypoints).convert("RGB")
+
+    # calculate angles
+    lea, rea = utils.get_eye_angles(keypoints)
+    lba, rba = utils.get_elbow_angles(keypoints)
+    angles  = {'left_eye_angle': lea, 'right_eye_angle': rea, 'left_elbow_angle': lba, 'right_elbow_angle': rba}
+
+    # encode to base64
+    img_buffer = io.BytesIO()
+    output_image.save(img_buffer, format="JPEG")
+    img_buffer.seek(0)
+    img_base64 = base64.b64encode(img_buffer.getvalue()).decode("utf-8")
+
+    # prepare json response
+    json_data = {
+        "keypoints": keypoints,
+        "angles": angles,
+        "output_image": img_base64
+    }
+
+    return JSONResponse(content=json_data)
+
+
+# if __name__ == "__main__":
+#     import uvicorn
+#     uvicorn.run(app, host="0.0.0.0", port=7860)