PoseMaker / app.py
jonigata's picture
add body estimation
ec80d0f
raw
history blame
3.66 kB
import gradio as gr
import numpy as np
import cv2
from fastapi import FastAPI, Request, Response
from src.body import Body
body_estimation = Body('model/body_pose_model.pth')
def pil2cv(image):
''' PIL型 -> OpenCV型 '''
new_image = np.array(image, dtype=np.uint8)
if new_image.ndim == 2: # モノクロ
pass
elif new_image.shape[2] == 3: # カラー
new_image = cv2.cvtColor(new_image, cv2.COLOR_RGB2BGR)
elif new_image.shape[2] == 4: # 透過
new_image = cv2.cvtColor(new_image, cv2.COLOR_RGBA2BGRA)
return new_image
with open("static/poseEditor.js", "r") as f:
file_contents = f.read()
app = FastAPI()
@app.middleware("http")
async def some_fastapi_middleware(request: Request, call_next):
path = request.scope['path'] # get the request route
response = await call_next(request)
if path == "/":
response_body = ""
async for chunk in response.body_iterator:
response_body += chunk.decode()
some_javascript = f"""
<script type="text/javascript" defer>
{file_contents}
</script>
"""
response_body = response_body.replace("</body>", some_javascript + "</body>")
del response.headers["content-length"]
return Response(
content=response_body,
status_code=response.status_code,
headers=dict(response.headers),
media_type=response.media_type
)
return response
# make cndidate to json
def candidate_to_json_string(arr):
a = [f'[{x:.2f}, {y:.2f}]' for x, y, *_ in arr]
return '[' + ', '.join(a) + ']'
# make subset to json
def subset_to_json_string(arr):
arr_str = ','.join(['[' + ','.join([f'{num:.2f}' for num in row]) + ']' for row in arr])
return '[' + arr_str + ']'
def estimate_body(source):
print("estimate_body")
if source == None:
return None
candidate, subset = body_estimation(pil2cv(source))
print(candidate_to_json_string(candidate))
print(subset_to_json_string(subset))
return "{ \"candidate\": " + candidate_to_json_string(candidate) + ", \"subset\": " + subset_to_json_string(subset) + " }"
def image_changed(image):
if (image == None):
return None
json = estimate_body(image)
return json, image.width, image.height
html_text = f"""
<canvas id="canvas" width="512" height="512"></canvas>
<script type="text/javascript" defer>{file_contents}</script>
"""
with gr.Blocks() as demo:
with gr.Row():
with gr.Column(scale=1):
source = gr.Image(type="pil")
width = gr.Slider(label="Width", mininmum=512, maximum=1024, step=64, value=512, key="Width", interactive=True)
height = gr.Slider(label="Height", mininmum=512, maximum=1024, step=64, value=512, key="Height", interactive=True)
startBtn = gr.Button(value="Start edit")
json = gr.JSON(label="Body")
with gr.Column(scale=2):
gr.HTML("<ul><li>ctrl + drag to scale</li><li>alt + drag to translate</li><li>shift + drag to rotate(move right first, then up or down)</li></ul>")
html = gr.HTML(html_text)
saveBtn = gr.Button(value="Save")
source.change(
fn = image_changed,
inputs = [source],
outputs = [json, width, height])
startBtn.click(
fn = None,
inputs = [json, width, height],
outputs = [],
_js="(json, w, h) => { initializePose(json,w,h); return []; }")
saveBtn.click(
fn = None,
inputs = [], outputs = [],
_js="() => { savePose(); }")
gr.mount_gradio_app(app, demo, path="/")