| import os |
| import numpy as np |
| import cv2 |
|
|
| import gradio as gr |
|
|
| import modules.scripts as scripts |
| from modules import script_callbacks |
| from modules import extensions |
|
|
| from basicsr.utils.download_util import load_file_from_url |
|
|
| from scripts.openpose.body import Body |
|
|
| body_estimation = None |
|
|
| def pil2cv(in_image): |
| out_image = np.array(in_image, dtype=np.uint8) |
|
|
| if out_image.shape[2] == 3: |
| out_image = cv2.cvtColor(out_image, cv2.COLOR_RGB2BGR) |
| return out_image |
|
|
| def candidate2li(li): |
| res = [] |
| for x, y, *_ in li: |
| res.append([x, y]) |
| return res |
|
|
| def subset2li(li): |
| res = [] |
| for r in li: |
| for c in r: |
| res.append(c) |
| return res |
|
|
| class Script(scripts.Script): |
| def __init__(self) -> None: |
| super().__init__() |
|
|
| def title(self): |
| return "OpenPose Editor" |
|
|
| def show(self, is_img2img): |
| return scripts.AlwaysVisible |
|
|
| def ui(self, is_img2img): |
| return () |
|
|
| def on_ui_tabs(): |
| with gr.Blocks(analytics_enabled=False) as openpose_editor: |
| with gr.Row(): |
| with gr.Column(): |
| width = gr.Slider(label="width", minimum=64, maximum=2048, value=512, step=64, interactive=True) |
| height = gr.Slider(label="height", minimum=64, maximum=2048, value=512, step=64, interactive=True) |
| with gr.Row(): |
| add = gr.Button(value="Add", variant="primary") |
| |
| with gr.Row(): |
| reset_btn = gr.Button(value="Reset") |
| json_input = gr.Button(value="Load from JSON") |
| png_input = gr.Button(value="Detect from image") |
| png_input_area = gr.Image(label="Detect from image", elem_id="openpose_editor_input") |
| bg_input = gr.Button(value="Add Background image") |
|
|
| with gr.Column(): |
| |
| canvas = gr.HTML('<canvas id="openpose_editor_canvas" width="512" height="512" style="margin: 0.25rem; border-radius: 0.25rem; border: 0.5px solid"></canvas>') |
| jsonbox = gr.Text(label="json", elem_id="hide_json") |
| with gr.Row(): |
| json_output = gr.Button(value="Save JSON") |
| png_output = gr.Button(value="Save PNG") |
| send_output = gr.Button(value="Send to ControlNet") |
|
|
| def estimate(img): |
| global body_estimation |
|
|
| modeldir = os.path.join(extensions.extensions_dir, "sd-webui-controlnet", "annotator", "openpose") |
| if body_estimation is None: |
| if not os.path.isfile(os.path.join(modeldir, "body_pose_model.pth")): |
| body_model_path = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/body_pose_model.pth" |
| load_file_from_url(body_model_path, model_dir=modeldir) |
| body_estimation = Body(os.path.join(modeldir, "body_pose_model.pth")) |
| |
| candidate, subset = body_estimation(pil2cv(img)) |
|
|
| result = { |
| "candidate": candidate2li(candidate), |
| "subset": subset2li(subset) |
| } |
| |
| return result |
|
|
|
|
| width.change(None, [width, height], None, _js="(w, h) => {resizeCanvas(w, h)}") |
| height.change(None, [width, height], None, _js="(w, h) => {resizeCanvas(w, h)}") |
| png_output.click(None, [], None, _js="savePNG") |
| bg_input.click(None, [], None, _js="addBackground") |
| png_input.click(None, [], None, _js="detectImage") |
| add.click(None, [], None, _js="addPose") |
| png_input_area.change(estimate, [png_input_area], [jsonbox]) |
| send_output.click(None, [], None, _js="sendImage") |
| reset_btn.click(None, [], None, _js="resetCanvas") |
| json_input.click(None, None, [width, height], _js="loadJSON") |
| json_output.click(None, None, None, _js="saveJSON") |
|
|
| return [(openpose_editor, "OpenPose Editor", "openpose_editor")] |
|
|
| script_callbacks.on_ui_tabs(on_ui_tabs) |
