Spaces:
Running
Running
| #This code is largely based on: https://huggingface.co/spaces/diffusers/controlnet-3d-pose | |
| import gradio as gr | |
| from PIL import Image, ImageFilter, ImageOps | |
| from io import BytesIO | |
| import numpy as np | |
| import base64 | |
| import cv2 | |
| import os | |
| import uuid | |
| uid=uuid.uuid4() | |
| canvas_html = '<div style="height:50%"><pose-maker/></div>' | |
| low_threshold = 100 | |
| high_threshold = 200 | |
| load_js = """ | |
| async () => { | |
| const url = "https://huggingface.co/datasets/mishig/gradio-components/raw/main/mannequinAll.js" | |
| fetch(url) | |
| .then(res => res.text()) | |
| .then(text => { | |
| const script = document.createElement('script'); | |
| script.type = "module" | |
| script.src = URL.createObjectURL(new Blob([text], { type: 'application/javascript' })); | |
| document.head.appendChild(script); | |
| }); | |
| } | |
| """ | |
| get_js_image = """ | |
| async (canvas, prompt) => { | |
| const poseMakerEl = document.querySelector("pose-maker"); | |
| const imgBase64 = poseMakerEl.captureScreenshotDepthMap(); | |
| return [imgBase64, prompt] | |
| } | |
| """ | |
| js_change_rotation_axis = """ | |
| async (axis) => { | |
| const poseMakerEl = document.querySelector("pose-maker"); | |
| poseMakerEl.changeRotationAxis(axis); | |
| } | |
| """ | |
| js_pose_template = """ | |
| async (pose) => { | |
| const poseMakerEl = document.querySelector("pose-maker"); | |
| poseMakerEl.setPose(pose); | |
| } | |
| """ | |
| def get_canny_filter(image): | |
| if not isinstance(image, np.ndarray): | |
| image = np.array(image) | |
| image = cv2.Canny(image, low_threshold, high_threshold) | |
| image = image[:, :, None] | |
| image = np.concatenate([image, image, image], axis=2) | |
| canny_image = Image.fromarray(image) | |
| return canny_image | |
| def generate_images(canvas,output_type): | |
| base64_img = canvas | |
| image_data = base64.b64decode(base64_img.split(',')[1]) | |
| left =128 | |
| top = 0 | |
| right = 384 | |
| bottom = 512 | |
| input_img = Image.open(BytesIO(image_data)).convert( | |
| 'RGB').resize((512,512)).crop((left, top, right, bottom)) | |
| #input_img = input_img.filter(ImageFilter.GaussianBlur(radius=2)) | |
| input_img = ImageOps.invert(get_canny_filter(input_img)) if output_type == "Canny" else input_img | |
| input_img = ImageOps.invert(input_img) if output_type == "Depth" else input_img | |
| #input_img = Image.open(BytesIO(base64_img)).resize((512,512)).crop((left, top, right, bottom)) if output_type == "Base" else input_img | |
| #input_img = ImageOps.invert(input_img) | |
| input_img.save(f"{uid}-pose.png") | |
| out = os.path.abspath(f"{uid}-pose.png") | |
| out_url = f'https://omnibus-model-mover.hf.space/file={out}' | |
| return f"{uid}-pose.png",f"{uid}-pose.png",out_url | |
| def placeholder_fn(axis): | |
| pass | |
| with gr.Blocks() as b: | |
| with gr.Row(): | |
| canvas = gr.HTML(canvas_html, elem_id="canvas_html", visible=True) | |
| with gr.Column(): | |
| rotation_axis = gr.Radio(choices=["x", "y", "z"], value="x", label="Joint rotation axis") | |
| pose_template = gr.Radio(choices=["regular", "ballet", "handstand", "split", "kick", "chilling"], value="regular", label="Pose template") | |
| output_type = gr.Radio(choices=["Depth", "Canny"], value="Depth", label="Image Output") | |
| run_button = gr.Button("Generate") | |
| out_text = gr.Textbox(label="Image URL") | |
| out_file=gr.File() | |
| out_im = gr.Image(type='filepath',show_share_button=True) | |
| rotation_axis.change(fn=placeholder_fn, | |
| inputs=[rotation_axis], | |
| outputs=[], | |
| queue=False, | |
| _js=js_change_rotation_axis) | |
| pose_template.change(fn=placeholder_fn, | |
| inputs=[pose_template], | |
| outputs=[], | |
| queue=False, | |
| _js=js_pose_template) | |
| run_button.click(fn=generate_images, | |
| inputs=[canvas,output_type], | |
| outputs=[out_im,out_file,out_text], | |
| _js=get_js_image) | |
| b.load(None,None,None,_js=load_js) | |
| b.launch() |