File size: 2,360 Bytes
bcf59c3 615de65 bcf59c3 615de65 bcf59c3 615de65 ace3ff2 bcf59c3 615de65 bcf59c3 |
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 |
from dotenv import load_dotenv
load_dotenv()
import glob
import os
from PIL import Image
import gradio as gr
import preprocess
from huggingface_hub import login
def extract_clothes(input_img, cls):
print(input_img, type(input_img), cls)
input_dir = "input_image"
output_dir = "output_image"
os.makedirs(input_dir, exist_ok=True)
os.makedirs(output_dir, exist_ok=True)
for f in glob.glob(input_dir + "/*.*"):
os.remove(f)
for f in glob.glob(output_dir + "/*.*"):
os.remove(f)
for f in glob.glob("cloth-mask/*.*"):
os.remove(f)
input_img.save(os.path.join(input_dir, "img.jpg"))
preprocess.extract_garment(inputs_dir=input_dir, outputs_dir=output_dir, cls=cls)
return Image.open(glob.glob(output_dir + "/*.*")[0])
css = """
#col-container {
margin: 0 auto;
max-width: 720px;
}
"""
with gr.Blocks(css=css) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(f"""
# Clothes Extraction using U2Net
Pull out clothes like tops, bottoms, and dresses from a photo. This implementation is based on the [U2Net](https://github.com/xuebinqin/U-2-Net) model.
""")
with gr.Row():
with gr.Column():
input_image = gr.Image(label="Input Image", type='pil', height="400px", show_label=True)
dropdown = gr.Dropdown(["upper", "lower", "dress"], value="upper", label="Extract clothes",
info="Select the clothes type you wish to extract!")
output_image = gr.Image(label="Extracted clothes", type='pil', height="400px", show_label=True,
show_download_button=True, format="png")
with gr.Row():
submit_button = gr.Button("Submit", variant='primary', scale=1)
reset_button = gr.ClearButton(value="Reset", scale=1)
gr.on(
triggers=[submit_button.click],
fn=extract_clothes,
inputs=[input_image, dropdown],
outputs=[output_image]
)
reset_button.click(
fn=lambda: (None, "upper", None),
inputs=[],
outputs=[input_image, dropdown, output_image],
concurrency_limit=1,
)
if __name__ == '__main__':
# login to hugging face
login(os.environ.get("HF_TOKEN"))
demo.launch(show_api=True)
|