File size: 5,466 Bytes
34097e9 |
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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
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)
|