import gradio as gr import os from pathlib import Path import sys import torch from PIL import Image, ImageOps , ImageDraw import numpy as np from utils_ootd import get_mask_location from cloths_db import cloths_map, modeL_db PROJECT_ROOT = Path(__file__).absolute().parents[1].absolute() sys.path.insert(0, str(PROJECT_ROOT)) from preprocess.openpose.run_openpose import OpenPose from preprocess.humanparsing.run_parsing import Parsing from ootd.inference_ootd_hd import OOTDiffusionHD from ootd.inference_ootd_dc import OOTDiffusionDC from preprocess.openpose.annotator.openpose.util import draw_bodypose # Set default dtype to float64 # torch.set_default_dtype(torch.float16) openpose_model = OpenPose(0) parsing_model_hd = Parsing(0) ootd_model_hd = OOTDiffusionHD(0) parsing_model_dc = Parsing(0) ootd_model_dc = OOTDiffusionDC(0) category_dict = ['upperbody', 'lowerbody', 'dress'] category_dict_utils = ['upper_body', 'lower_body', 'dresses'] example_path = os.path.join(os.path.dirname(__file__), 'examples') garment_path = os.path.join(os.path.dirname(__file__), 'examples','garment') model_hd = os.path.join(example_path, 'model/model_1.png') garment_hd = os.path.join(example_path, 'garment/03244_00.jpg') model_dc = os.path.join(example_path, 'model/model_8.png') garment_dc = os.path.join(example_path, 'garment/048554_1.jpg') openpose_model.preprocessor.body_estimation.model.to('cuda') #model dc ootd_model_dc.pipe.to('cuda') ootd_model_dc.image_encoder.to('cuda') ootd_model_dc.text_encoder.to('cuda') #model hd # ootd_model_hd.pipe.to('cuda') # ootd_model_hd.image_encoder.to('cuda') # ootd_model_hd.text_encoder.to('cuda') def convert_to_image(image_array): if isinstance(image_array, np.ndarray): # Normalize the data to the range [0, 255] image_array = 255 * (image_array - np.min(image_array)) / (np.max(image_array) - np.min(image_array)) # Convert to uint8 image_array = image_array.astype(np.uint8) return Image.fromarray(image_array) else: # Convert to NumPy array first if necessary image_array = np.array(image_array) # Normalize and convert to uint8 image_array = 255 * (image_array - np.min(image_array)) / (np.max(image_array) - np.min(image_array)) image_array = image_array.astype(np.uint8) return Image.fromarray(image_array) # import spaces # @spaces.GPU def process_hd(vton_img, garm_img, n_samples, n_steps, image_scale, seed): model_type = 'hd' category = 0 # 0:upperbody; 1:lowerbody; 2:dress with torch.no_grad(): # openpose_model_hd.preprocessor.body_estimation.model.to('cuda') # ootd_model_hd.pipe.to('cuda') # ootd_model_hd.image_encoder.to('cuda') # ootd_model_hd.text_encoder.to('cuda') # garm_img = Image.open(garm_img).resize((768, 1024)) # vton_img = Image.open(vton_img).resize((768, 1024)) # keypoints = openpose_model(vton_img.resize((384, 512))) garm_img = Image.open(garm_img).resize((768, 1024)) vton_img = Image.open(vton_img).resize((768, 1024)) keypoints ,candidate , subset = openpose_model(vton_img.resize((384, 512))) print(len(keypoints["pose_keypoints_2d"])) print(keypoints["pose_keypoints_2d"]) model_parse, _ = parsing_model_hd(vton_img.resize((384, 512))) mask, mask_gray = get_mask_location(model_type, category_dict_utils[category], model_parse, keypoints) mask = mask.resize((768, 1024), Image.NEAREST) mask_gray = mask_gray.resize((768, 1024), Image.NEAREST) masked_vton_img = Image.composite(mask_gray, vton_img, mask) images = ootd_model_hd( model_type=model_type, category=category_dict[category], image_garm=garm_img, image_vton=masked_vton_img, mask=mask, image_ori=vton_img, num_samples=n_samples, num_steps=n_steps, image_scale=2.0, seed=42, ) return images def create_bw_mask(size): width, height = size mask = Image.new('L', (width, height)) draw = ImageDraw.Draw(mask) draw.rectangle([0, 0, width, height // 2], fill=255) # top half white draw.rectangle([0, height // 2, width, height], fill=0) # bottom half black return mask def create_mask(vton_img, garm_img, category): model_type = 'dc' if category == 'Upper-body': category = 0 elif category == 'Lower-body': category = 1 else: category =2 with torch.no_grad(): # openpose_model_dc.preprocessor.body_estimation.model.to('cuda') # ootd_model_dc.pipe.to('cuda') # ootd_model_dc.image_encoder.to('cuda') # ootd_model_dc.text_encoder.to('cuda') garm_img = Image.open(garm_img).resize((768, 1024)) vton_img = Image.open(vton_img).resize((768, 1024)) keypoints = openpose_model(vton_img.resize((384, 512))) print(len(keypoints["pose_keypoints_2d"])) print(keypoints["pose_keypoints_2d"]) # person_image = np.asarray(vton_img) # print(len(person_image)) # person_image = np.asarray(Image.open(vton_img).resize((768, 1024))) # output = draw_bodypose(canvas=person_image,candidate=candidate, subset=subset ) # output_image = Image.fromarray(output) # output_image.save('keypose.png') left_point = keypoints["pose_keypoints_2d"][2] right_point = keypoints["pose_keypoints_2d"][5] neck_point = keypoints["pose_keypoints_2d"][1] hip_point = keypoints["pose_keypoints_2d"][8] print(f'left shoulder - {left_point}') print(f'right shoulder - {right_point}') # #find disctance using Euclidian distance shoulder_width_pixels = round(np.sqrt( np.power((right_point[0]-left_point[0]),2) + np.power((right_point[1]-left_point[1]),2)),2) height_pixels = round(np.sqrt( np.power((neck_point[0]-hip_point[0]),2) + np.power((neck_point[1]-hip_point[1]),2)),2) *2 # # Assuming an average human height average_height_cm = 172.72 *1.5 # Conversion factor from pixels to cm conversion_factor = average_height_cm / height_pixels # Convert shoulder width to real-world units shoulder_width_cm = shoulder_width_pixels * conversion_factor print(f'Shoulder width (in pixels): {shoulder_width_pixels}') print(f'Estimated height (in pixels): {height_pixels}') print(f'Conversion factor (pixels to cm): {conversion_factor}') print(f'Shoulder width (in cm): {shoulder_width_cm}') print(f'Shoulder width (in INCH): {round(shoulder_width_cm/2.54,1)}') model_parse, face_mask = parsing_model_dc(vton_img.resize((384, 512))) model_parse_image = convert_to_image(model_parse) face_mask_image = convert_to_image(face_mask) # Save the images model_parse_image.save('model_parse_image.png') face_mask_image.save('face_mask_image.png') mask, mask_gray = get_mask_location(model_type, category_dict_utils[category], model_parse, keypoints) # up_mask, up_mask_gray = get_mask_location(model_type, category_dict_utils[0], model_parse, keypoints) # lo_mask, lo_mask_gray = get_mask_location(model_type, category_dict_utils[1], model_parse, keypoints) # mask = Image.composite(up_mask,lo_mask,up_mask) # mask_gray = Image.composite(up_mask_gray, lo_mask_gray,up_mask) mask = mask.resize((768, 1024), Image.NEAREST) mask_gray = mask_gray.resize((768, 1024), Image.NEAREST) # if modeL_db[vton_img] == 0: # Create a black-and-white mask bw_mask = create_bw_mask((768, 1024)) #crete empty black image with mode L temp_img = Image.new("L", (768, 1024), 0) mask = Image.composite(mask, temp_img, bw_mask) # print(mask) # Save the resized masks mask.save("mask_resized.png") mask_gray.save("mask_gray_resized.png") return [mask, mask_gray], mask, mask_gray # @spaces.GPU def process_dc(vton_img, garm_img, category, mask,mask_gray): model_type = 'dc' if category == 'Upper-body': category = 0 elif category == 'Lower-body': category = 1 else: category =2 # Extract the composite images from the edit data edited_mask = mask['composite'] edited_mask_gray = mask_gray['composite'] # print(edited_mask) garm_img = Image.open(garm_img).resize((768, 1024)) vton_img = Image.open(vton_img).resize((768, 1024)) # print(f'vton_img is {vton_img}') with torch.no_grad(): # Ensure both masks are in 'L' mode (grayscale) if edited_mask.mode != 'L': edited_mask = edited_mask.convert('L') if edited_mask_gray.mode != 'L': edited_mask_gray = edited_mask_gray.convert('L') # Ensure all images and masks are the same size edited_mask = edited_mask.resize((768, 1024), Image.NEAREST) edited_mask_gray = edited_mask_gray.resize((768, 1024), Image.NEAREST) print(f'mask: {edited_mask}') print(f'vton_img: {vton_img}') masked_vton_img = Image.composite(edited_mask_gray, vton_img, edited_mask) masked_vton_img.save("masked_vton_img.png") print(f'category is {category}') images = ootd_model_dc( model_type=model_type, category=category_dict[category], image_garm=garm_img, image_vton=masked_vton_img, mask=edited_mask, image_ori=vton_img, num_samples=1, num_steps=10, image_scale= 2.0, seed=-1, ) # return None return images # is_upper = False block = gr.Blocks().queue() with block: mask_state = gr.State() mask_gray_state = gr.State() with gr.Row(): gr.Markdown("# ") with gr.Row(): gr.Markdown("## Virtual Trial Room") # with gr.Row(): # gr.Markdown("") with gr.Row(): with gr.Column(): vton_img_dc = gr.Image(label="Model", sources='upload', type="filepath", height=384, value=model_dc) # Hidden component to store is_upper value # is_upper = gr.State(value=True) # #set is_upper variable to True when user selects examples from gr.examples upper/lower body # def check_image_type(image_path): # if image_path: # filename = os.path.basename(image_path) # image_type = modeL_db.get(filename, "no-dress") # Default to "no-dress" if not found # return image_type == "no-dress" # return False # Default to True if no image # vton_img_dc.change(fn=check_image_type, inputs=vton_img_dc) example = gr.Examples( label="Select for Upper/Lower Body", inputs=vton_img_dc, examples_per_page=7, examples=[ os.path.join(example_path, 'model/model_8.png'), os.path.join(example_path, 'model/049447_0.jpg'), os.path.join(example_path, 'model/049713_0.jpg'), os.path.join(example_path, 'model/051482_0.jpg'), os.path.join(example_path, 'model/051918_0.jpg'), os.path.join(example_path, 'model/051962_0.jpg'), os.path.join(example_path, 'model/049205_0.jpg'), os.path.join(example_path, 'model/05997_00.jpg'), ], ) example = gr.Examples( label="Select for Full Body Dress", inputs=vton_img_dc, examples_per_page=7, examples=[ os.path.join(example_path, 'model/model_9.png'), os.path.join(example_path, 'model/052767_0.jpg'), os.path.join(example_path, 'model/052472_0.jpg'), os.path.join(example_path, 'model/053514_0.jpg'), os.path.join(example_path, 'model/053228_0.jpg'), os.path.join(example_path, 'model/06802_00.jpg'), os.path.join(example_path, 'model/053700_0.jpg'), ], ) with gr.Column(): garm_img_dc = gr.Image(label="Garment", sources='upload', type="filepath", height=384, value=garment_dc) category_dc = gr.Dropdown(label="Garment category (important option!!!)", choices=["Upper-body", "Lower-body", "Dress"], value="Upper-body") def update_category(image_path): if image_path: filename = os.path.basename(image_path) return cloths_map.get(filename, None) # Return None if not found return None # Return None if no image garm_img_dc.change(fn=update_category, inputs=garm_img_dc, outputs=category_dc) example = gr.Examples( label="Examples (upper-body)", inputs=garm_img_dc, examples_per_page=7, examples=[ os.path.join(garment_path,'01260_00.jpg'), os.path.join(garment_path,'01430_00.jpg'), os.path.join(garment_path,'02783_00.jpg'), os.path.join(garment_path,'03751_00.jpg'), os.path.join(garment_path,'06429_00.jpg'), os.path.join(garment_path,'06802_00.jpg'), os.path.join(garment_path,'07429_00.jpg'), os.path.join(garment_path,'08348_00.jpg'), os.path.join(garment_path,'09933_00.jpg'), os.path.join(garment_path,'11028_00.jpg'), os.path.join(garment_path,'11351_00.jpg'), os.path.join(garment_path,'11791_00.jpg'), os.path.join(garment_path, '048554_1.jpg'), os.path.join(garment_path, '049920_1.jpg'), os.path.join(garment_path, '049965_1.jpg'), os.path.join(garment_path, '049949_1.jpg'), os.path.join(garment_path, '050181_1.jpg'), os.path.join(garment_path, '049805_1.jpg'), os.path.join(garment_path, '050105_1.jpg'), os.path.join(garment_path, 'male_tshirt1.png'), ]) example = gr.Examples( label="Examples (lower-body)", inputs=garm_img_dc, examples_per_page=7, examples=[ os.path.join(garment_path, '051827_1.jpg'), os.path.join(garment_path, '051946_1.jpg'), os.path.join(garment_path, '051473_1.jpg'), os.path.join(garment_path, '051515_1.jpg'), os.path.join(garment_path, '051517_1.jpg'), os.path.join(garment_path, '051988_1.jpg'), os.path.join(garment_path, '051412_1.jpg'), ]) example = gr.Examples( label="Examples (dress)", inputs=garm_img_dc, examples_per_page=7, examples=[ os.path.join(garment_path, '053290_1.jpg'), os.path.join(garment_path, '053744_1.jpg'), os.path.join(garment_path, '053742_1.jpg'), os.path.join(garment_path, '053786_1.jpg'), os.path.join(garment_path, '053790_1.jpg'), os.path.join(garment_path, '053319_1.jpg'), os.path.join(garment_path, '052234_1.jpg'), ]) with gr.Column(): mask_gallery = gr.Gallery(label="Created Masks") result_gallery_dc = gr.Gallery(label='Output', show_label=False, elem_id="gallery", preview=True, scale=1) with gr.Row(): # Add ImageEditor for mask editing mask_editor = gr.ImageEditor(label="Edit Mask", type="pil") # Add ImageEditor for mask_gray editing mask_gray_editor = gr.ImageEditor(label="Edit Mask Gray", type="pil") with gr.Column(): create_mask_button = gr.Button(value="Create Mask") run_button_dc = gr.Button(value="Run") # n_samples_dc = gr.Slider(label="Images", minimum=1, maximum=4, value=1, step=1) # n_steps_dc = gr.Slider(label="Steps", minimum=20, maximum=40, value=20, step=1) # scale_dc = gr.Slider(label="Scale", minimum=1.0, maximum=12.0, value=5.0, step=0.1) # image_scale_dc = gr.Slider(label="Guidance scale", minimum=1.0, maximum=5.0, value=2.0, step=0.1) # seed_dc = gr.Slider(label="Seed", minimum=-1, maximum=2147483647, step=1, value=-1) # ips_dc = [vton_img_dc, garm_img_dc, category_dc] ips_dc = [vton_img_dc, garm_img_dc ,category_dc] # create_mask_button.click( # fn=create_mask, # inputs=ips_dc, # outputs=[mask_gallery, mask_state, mask_gray_state] # ) create_mask_button.click( fn=create_mask, inputs=ips_dc, outputs=[mask_gallery, mask_editor, mask_gray_editor] ) # run_button_dc.click(fn=process_dc, inputs=ips_dc, outputs=[result_gallery_dc]) # run_button_dc.click( # fn=process_dc, # inputs=ips_dc + [mask_state, mask_gray_state], # outputs=[result_gallery_dc]) run_button_dc.click( fn=process_dc, inputs=[vton_img_dc, garm_img_dc, category_dc, mask_editor, mask_gray_editor], outputs=[result_gallery_dc] ) block.launch(server_name="0.0.0.0", server_port=7860 )