Spaces:
Sleeping
Sleeping
File size: 1,639 Bytes
5714fc1 17f85a1 0eeb953 17f85a1 5714fc1 0eeb953 |
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 63 64 |
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)
# 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)
|