File size: 1,517 Bytes
3dabe4a |
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 |
from modules import script_callbacks
import gradio as gr
from PIL import Image
import numpy as np
from imgutils.detect import detect_halfbody, detect_heads
def on_ui_tab_called():
with gr.Blocks() as transparent_interface:
with gr.Row():
with gr.Tabs():
with gr.TabItem("CropHead"):
image_upload_input = gr.Image(label="Upload Image", source="upload",type="pil")
button = gr.Button(label="Convert")
image_upload_output = gr.Image(label="Output Image",type="numpy")
def convert_image(image:Image.Image):
# first convert to RGB
if image.mode == "RGBA":
# convert transparent pixels to white
white_image = Image.new("RGB", image.size, (255, 255, 255))
white_image.paste(image, mask=image.split()[3])
image = white_image
else:
image = image.convert("RGB")
result = detect_heads(image)
new_image = image.crop(result[0][0])
return new_image # return the new image
button.click(convert_image, inputs=[image_upload_input], outputs=[image_upload_output])
return (transparent_interface, "CropHead", "script_crophead_interface"),
script_callbacks.on_ui_tabs(on_ui_tab_called)
|