Spaces:
Sleeping
Sleeping
File size: 1,635 Bytes
4ca2511 ad29a66 6fda8a4 ad29a66 6fda8a4 ad29a66 6fda8a4 ff2b03c 6fda8a4 |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
import io
import base64
from ultralytics import YOLO
import torch
from PIL import Image, ImageOps
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("/")
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)
if not keypoints:
return JSONResponse(content={"error": "No person detected"})
# 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)
|