123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143 |
- import os
- import io
- import json
- import numpy as np
- import cv2
- import gradio as gr
- import modules.scripts as scripts
- from modules import script_callbacks
- from modules.shared import opts
- from modules.paths import models_path
- from basicsr.utils.download_util import load_file_from_url
- from scripts.openpose.body import Body
- from PIL import Image
- body_estimation = None
- presets_file = os.path.join(scripts.basedir(), "presets.json")
- presets = {}
- try:
- with open(presets_file) as file:
- presets = json.load(file)
- except FileNotFoundError:
- pass
- 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")
- # delete = gr.Button(value="Delete")
- with gr.Row():
- reset_btn = gr.Button(value="Reset")
- json_input = gr.UploadButton(label="Load from JSON", file_types=[".json"], elem_id="openpose_json_button")
- png_input = gr.UploadButton(label="Detect from Image", file_types=["image"], type="bytes", elem_id="openpose_detect_button")
- bg_input = gr.UploadButton(label="Add Background Image", file_types=["image"], elem_id="openpose_bg_button")
- with gr.Row():
- preset_list = gr.Dropdown(label="Presets", choices=sorted(presets.keys()), interactive=True)
- preset_load = gr.Button(value="Load Preset")
- preset_save = gr.Button(value="Save Preset")
- with gr.Column():
- # gradioooooo...
- 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="jsonbox", visible=False)
- with gr.Row():
- json_output = gr.Button(value="Save JSON")
- png_output = gr.Button(value="Save PNG")
- send_t2t = gr.Button(value="Send to txt2img")
- send_i2i = gr.Button(value="Send to img2img")
- control_net_max_models_num = getattr(opts, 'control_net_max_models_num', 0)
- select_target_index = gr.Dropdown([str(i) for i in range(control_net_max_models_num)], label="Send to", value="0", interactive=True, visible=(control_net_max_models_num > 1))
- def estimate(file):
- global body_estimation
- if body_estimation is None:
- model_path = os.path.join(models_path, "openpose", "body_pose_model.pth")
- if not os.path.isfile(model_path):
- 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=os.path.join(models_path, "openpose"))
- body_estimation = Body(model_path)
-
- stream = io.BytesIO(file)
- img = Image.open(stream)
- candidate, subset = body_estimation(pil2cv(img))
- result = {
- "candidate": candidate2li(candidate),
- "subset": subset2li(subset),
- }
-
- return str(result).replace("'", '"')
- def savePreset(name, data):
- if name:
- presets[name] = json.loads(data)
- with open(presets_file, "w") as file:
- json.dump(presets, file)
- return gr.update(choices=sorted(presets.keys()), value=name), json.dumps(data)
- return gr.update(), gr.update()
- dummy_component = gr.Label(visible=False)
- preset = gr.Text(visible=False)
- 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.upload(None, [bg_input], [width, height], _js="addBackground")
- png_input.upload(estimate, png_input, jsonbox)
- png_input.upload(None, png_input, [width, height], _js="addBackground")
- add.click(None, [], None, _js="addPose")
- send_t2t.click(None, select_target_index, None, _js="(i) => {sendImage('txt2img', i)}")
- send_i2i.click(None, select_target_index, None, _js="(i) => {sendImage('img2img', i)}")
- reset_btn.click(None, [], None, _js="resetCanvas")
- json_input.upload(None, json_input, [width, height], _js="loadJSON")
- json_output.click(None, None, None, _js="saveJSON")
- preset_save.click(savePreset, [dummy_component, dummy_component], [preset_list, preset], _js="savePreset")
- preset_load.click(None, preset, [width, height], _js="loadPreset")
- preset_list.change(lambda selected: json.dumps(presets[selected]), preset_list, preset)
- return [(openpose_editor, "OpenPose Editor", "openpose_editor")]
- script_callbacks.on_ui_tabs(on_ui_tabs)
|