Upload utils_mask.py

utils_mask.py

@@ -0,0 +1,498 @@
+import numpy as np
+import cv2
+from PIL import Image, ImageDraw, ImageOps
+import torch
+import matplotlib.pyplot as plt
+
+import apply_net
+from detectron2.data.detection_utils import convert_PIL_to_numpy,_apply_exif_orientation
+
+label_map = {
+    "background": 0,
+    "hat": 1,
+    "hair": 2,
+    "sunglasses": 3,
+    "upper_clothes": 4,
+    "skirt": 5,
+    "pants": 6,
+    "dress": 7,
+    "belt": 8,
+    "left_shoe": 9,
+    "right_shoe": 10,
+    "head": 11,
+    "left_leg": 12,
+    "right_leg": 13,
+    "left_arm": 14,
+    "right_arm": 15,
+    "bag": 16,
+    "scarf": 17,
+}
+
+dense_map = {
+    "background" : [0],
+    "torso" : [1,2],
+    "right_hand" : [3],
+    "left_hand" : [4],
+    "left_foot" : [5],
+    "right_foot" : [6],
+    "upper_leg_right" : [7,9],
+    "upper_leg_left" : [8,10],
+    "lower_leg_right" : [11,13],
+    "lower_leg_left" : [12,14],
+    "upper_arm_left" : [15,17],
+    "upper_arm_right" : [16,18],
+    "lower_arm_left" : [19,21],
+    "lower_arm_right" : [20,22],
+    "head" : [23,24]
+}
+
+def extend_arm_mask(wrist, elbow, scale):
+  wrist = elbow + scale * (wrist - elbow)
+  return wrist
+
+def hole_fill(img):
+    img = np.pad(img[1:-1, 1:-1], pad_width = 1, mode = 'constant', constant_values=0)
+    img_copy = img.copy()
+    mask = np.zeros((img.shape[0] + 2, img.shape[1] + 2), dtype=np.uint8)
+
+    cv2.floodFill(img, mask, (0, 0), 255)
+    img_inverse = cv2.bitwise_not(img)
+    dst = cv2.bitwise_or(img_copy, img_inverse)
+    return dst
+
+def refine_mask(mask):
+    contours, hierarchy = cv2.findContours(mask.astype(np.uint8),
+                                           cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
+    area = []
+    for j in range(len(contours)):
+        a_d = cv2.contourArea(contours[j], True)
+        area.append(abs(a_d))
+    refine_mask = np.zeros_like(mask).astype(np.uint8)
+    if len(area) != 0:
+        i = area.index(max(area))
+        cv2.drawContours(refine_mask, contours, i, color=255, thickness=-1)
+
+    return refine_mask
+
+def get_mask_location_new(category, model_parse: Image.Image, keypoint: dict, width=384,height=512, dense_pose = None):
+
+    if category != 'lower_body_shoes' and category != 'lower_body_boots' and category != 'full_body' and category != 'dresses' and category != 'upper_clothes' and category != 'lower_body_pants' and category != 'lower_body_skirts':
+        raise ValueError("Category not found")
+    
+    
+    #mask for lower_body_shoes, lower_body_boots
+    if category == 'lower_body_shoes':
+        dense_mask = np.zeros((height, width))
+        dense_mask += (dense_pose == 5).astype(np.float32) + \
+                        (dense_pose == 6).astype(np.float32)
+        
+        dense_mask = cv2.dilate(dense_mask, np.ones((5, 5), np.uint16), iterations=5)
+
+        mask = Image.fromarray(dense_mask.astype(np.uint8) * 255)
+        mask_gray = Image.fromarray(dense_mask.astype(np.uint8) * 127)
+
+        return mask, mask_gray, dense_mask
+
+
+    if category == 'lower_body_boots':
+        dense_mask = np.zeros((height, width))
+
+        dense_mask += (dense_pose == 5).astype(np.float32) + \
+                        (dense_pose == 6).astype(np.float32) + \
+                        (dense_pose == 11).astype(np.float32) + \
+                        (dense_pose == 12).astype(np.float32) + \
+                        (dense_pose == 13).astype(np.float32) + \
+                        (dense_pose == 14).astype(np.float32)
+
+        dense_mask = cv2.dilate(dense_mask, np.ones((5, 5), np.uint16), iterations=5)
+
+        mask = Image.fromarray(dense_mask.astype(np.uint8) * 255)
+        mask_gray = Image.fromarray(dense_mask.astype(np.uint8) * 127)
+
+        return mask, mask_gray, dense_mask
+    
+    #mask others category
+    im_parse = model_parse.resize((width, height), Image.NEAREST)
+    parse_array = np.array(im_parse)
+
+    arm_width = 40
+
+    parse_head = (parse_array == 1).astype(np.float32) + \
+                 (parse_array == 3).astype(np.float32) + \
+                 (parse_array == 11).astype(np.float32)
+
+    parser_mask_fixed = (parse_array == label_map["left_shoe"]).astype(np.float32) + \
+                        (parse_array == label_map["right_shoe"]).astype(np.float32) + \
+                        (parse_array == label_map["hat"]).astype(np.float32) + \
+                        (parse_array == label_map["sunglasses"]).astype(np.float32) + \
+                        (parse_array == label_map["bag"]).astype(np.float32)
+
+    parser_mask_changeable = (parse_array == label_map["background"]).astype(np.float32)
+
+    arms_left = (parse_array == 14).astype(np.float32)
+    arms_right = (parse_array == 15).astype(np.float32)
+    arms = arms_left + arms_right
+
+    if category == 'dresses' or category == 'full_body': # upper_clothes + lower_body_skirts
+        parse_mask = (parse_array == 7).astype(np.float32) + \
+                     (parse_array == 4).astype(np.float32) + \
+                     (parse_array == 5).astype(np.float32) + \
+                     (parse_array == 6).astype(np.float32)
+
+        parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+
+    elif category == 'upper_clothes' : # -> upper_clothes
+        parse_mask = (parse_array == 4).astype(np.float32) 
+        parser_mask_fixed_lower_cloth = (parse_array == label_map["skirt"]).astype(np.float32) + \
+                                        (parse_array == label_map["pants"]).astype(np.float32)
+        # parser_mask_fixed += parser_mask_fixed_lower_cloth
+        parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+    elif category == 'lower_body_pants' or category == 'lower_body_skirts': # -> remove
+        parse_mask = (parse_array == 6).astype(np.float32) + \
+                     (parse_array == 12).astype(np.float32) + \
+                     (parse_array == 13).astype(np.float32) + \
+                     (parse_array == 5).astype(np.float32)
+        parser_mask_fixed += (parse_array == label_map["upper_clothes"]).astype(np.float32) + \
+                             (parse_array == 14).astype(np.float32) + \
+                             (parse_array == 15).astype(np.float32)
+        parser_mask_changeable += np.logical_and(parse_array, np.logical_not(parser_mask_fixed))
+    else:
+        raise NotImplementedError
+
+    # Load pose points
+    pose_data = keypoint["pose_keypoints_2d"]
+    pose_data = np.array(pose_data)
+    pose_data = pose_data.reshape((-1, 2))
+
+    im_arms_left = Image.new('L', (width, height))
+    im_arms_right = Image.new('L', (width, height))
+    arms_draw_left = ImageDraw.Draw(im_arms_left)
+    arms_draw_right = ImageDraw.Draw(im_arms_right)
+    if category == 'dresses' or category == 'upper_clothes' or category == 'full_body':
+        shoulder_right = np.multiply(tuple(pose_data[2][:2]), height / 512.0)
+        shoulder_left = np.multiply(tuple(pose_data[5][:2]), height / 512.0)
+        elbow_right = np.multiply(tuple(pose_data[3][:2]), height / 512.0)
+        elbow_left = np.multiply(tuple(pose_data[6][:2]), height / 512.0)
+        wrist_right = np.multiply(tuple(pose_data[4][:2]), height / 512.0)
+        wrist_left = np.multiply(tuple(pose_data[7][:2]), height / 512.0)
+        ARM_LINE_WIDTH = int(arm_width / 512 * height)
+        size_left = [shoulder_left[0] - ARM_LINE_WIDTH // 2, shoulder_left[1] - ARM_LINE_WIDTH // 2, shoulder_left[0] + ARM_LINE_WIDTH // 2, shoulder_left[1] + ARM_LINE_WIDTH // 2]
+        size_right = [shoulder_right[0] - ARM_LINE_WIDTH // 2, shoulder_right[1] - ARM_LINE_WIDTH // 2, shoulder_right[0] + ARM_LINE_WIDTH // 2,
+                      shoulder_right[1] + ARM_LINE_WIDTH // 2]
+        
+    
+        if wrist_right[0] <= 1. and wrist_right[1] <= 1.:
+            im_arms_right = arms_right
+        else:
+            wrist_right = extend_arm_mask(wrist_right, elbow_right, 1.2)
+            arms_draw_right.line(np.concatenate((shoulder_right, elbow_right, wrist_right)).astype(np.uint16).tolist(), 'white', ARM_LINE_WIDTH, 'curve')
+            arms_draw_right.arc(size_right, 0, 360, 'white', ARM_LINE_WIDTH // 2)
+
+        if wrist_left[0] <= 1. and wrist_left[1] <= 1.:
+            im_arms_left = arms_left
+        else:
+            wrist_left = extend_arm_mask(wrist_left, elbow_left, 1.2)
+            arms_draw_left.line (np.concatenate((wrist_left, elbow_left, shoulder_left)).astype(np.uint16).tolist(), 'white', ARM_LINE_WIDTH, 'curve')
+            arms_draw_left.arc(size_left, 0, 360, 'white', ARM_LINE_WIDTH // 2)
+
+        hands_left = np.logical_and(np.logical_not(im_arms_left), arms_left)
+        hands_right = np.logical_and(np.logical_not(im_arms_right), arms_right)
+        parser_mask_fixed += hands_left + hands_right
+
+    parser_mask_fixed = np.logical_or(parser_mask_fixed, parse_head)
+    parse_mask = cv2.dilate(parse_mask, np.ones((5, 5), np.uint16), iterations=5)
+    if category == 'dresses' or category == 'upper_clothes' or category == 'full_body':
+        neck_mask = (parse_array == 18).astype(np.float32)
+        neck_mask = cv2.dilate(neck_mask, np.ones((5, 5), np.uint16), iterations=1)
+        neck_mask = np.logical_and(neck_mask, np.logical_not(parse_head))
+        parse_mask = np.logical_or(parse_mask, neck_mask)
+        arm_mask = cv2.dilate(np.logical_or(im_arms_left, im_arms_right).astype('float32'), np.ones((5, 5), np.uint16), iterations=4)
+        parse_mask += np.logical_or(parse_mask, arm_mask)
+
+    # parse_mask_img = Image.fromarray(parse_mask.astype(np.uint8) * 255)
+    # parse_mask_img.save("mask_their_pre.png")
+
+    # parser_mask_changeable_img = Image.fromarray(parse_mask.astype(np.uint8) * 255)
+    # parser_mask_changeable_img.save("mask_change.png")
+
+    parse_mask = np.logical_and(parser_mask_changeable, np.logical_not(parse_mask))
+
+    #convert parse_mask to iamge and save
+    # parse_mask_img = Image.fromarray(parse_mask.astype(np.uint8) * 255)
+    # parse_mask_img.save("mask_their.png")
+
+    #my code
+
+    #get pose points
+    hip_right = np.multiply(tuple(pose_data[8][:2]), height / 512.0)
+    hip_left = np.multiply(tuple(pose_data[11][:2]), height / 512.0)
+    knee_right = np.multiply(tuple(pose_data[9][:2]), height / 512.0)
+    knee_left = np.multiply(tuple(pose_data[12][:2]), height / 512.0)
+    ankle_right = np.multiply(tuple(pose_data[10][:2]), height / 512.0)
+    ankle_left = np.multiply(tuple(pose_data[13][:2]), height / 512.0)
+
+    #for upper clothes
+    mid_point_left = hip_left + (knee_left - hip_left) / 5
+    mid_point_right = hip_right + (knee_right - hip_right) / 5
+
+    extra_mask = Image.new('L', (width, height))
+    extra_draw = ImageDraw.Draw(extra_mask)
+
+
+    #mask for dresses category
+    if category == 'dresses' or category == 'lower_body_skirts' or category == 'lower_body_pants':
+        
+        #draw line from 6 points
+        if ankle_left[0] != 0 and ankle_right[0] != 0 and ankle_left[1] != 0 and ankle_right[1] != 0:
+            extra_draw.line(np.concatenate((ankle_right, ankle_left)).astype(np.uint16).tolist(), 'white', 1, 'curve')
+            extra_draw.line(np.concatenate((hip_right, knee_right, ankle_right)).astype(np.uint16).tolist(), 'white', arm_width+20, 'curve')
+            extra_draw.line(np.concatenate((hip_left, knee_left, ankle_left)).astype(np.uint16).tolist(), 'white', arm_width+20, 'curve')
+            extra_draw.line(np.concatenate((hip_right, hip_left)).astype(np.uint16).tolist(), 'white', 1, 'curve')
+            extra_draw.line(np.concatenate((knee_right, knee_left)).astype(np.uint16).tolist(), 'white', 1, 'curve')
+
+        elif knee_left[0] != 0 and knee_right[0] != 0 and knee_left[1] != 0 and knee_right[1] != 0:
+            extra_draw.line(np.concatenate((hip_right, knee_right)).astype(np.uint16).tolist(), 'white', 1, 'curve')
+            extra_draw.line(np.concatenate((hip_left, knee_left)).astype(np.uint16).tolist(), 'white', arm_width, 'curve')
+            extra_draw.line(np.concatenate((hip_right, hip_left)).astype(np.uint16).tolist(), 'white', arm_width, 'curve')
+        else:
+            pass 
+
+        if category == 'lower_body_pants':
+            extra_mask = hole_fill(np.array(extra_mask))
+            extra_mask = cv2.dilate(np.array(extra_mask), np.ones((5, 5), np.uint16), iterations=int((knee_right[1] - hip_right[1])/10))
+            
+            dense = (dense_pose == 1).astype(np.float32) +\
+                    (dense_pose == 2).astype(np.float32) +\
+                    (dense_pose == 7).astype(np.float32) +\
+                    (dense_pose == 8).astype(np.float32) +\
+                    (dense_pose == 9).astype(np.float32) +\
+                    (dense_pose == 10).astype(np.float32)
+            extra_mask = np.logical_and(extra_mask, dense)
+            extra_mask = cv2.dilate((extra_mask * 255).astype(np.uint8), np.ones((5, 5), np.uint16), iterations=5)
+            extra_mask = Image.fromarray((extra_mask * 255).astype(np.uint8), 'L')   
+
+    #mask for upper_clothes
+    if category == "upper_clothes":
+        if knee_left[0] != 0 and knee_right[0] != 0 and knee_left[1] != 0 and knee_right[1] != 0:
+
+            extra_draw.line(np.concatenate((hip_right, hip_left)).astype(np.uint16).tolist(), 'white', 1, 'curve')
+            extra_draw.line(np.concatenate((mid_point_right, mid_point_left)).astype(np.uint16).tolist(), 'white', 1, 'curve')
+            extra_draw.line(np.concatenate((hip_right, mid_point_right)).astype(np.uint16).tolist(), 'white', 40, 'curve')
+            extra_draw.line(np.concatenate((hip_left, mid_point_left)).astype(np.uint16).tolist(), 'white', 40, 'curve')
+        else:
+            pass
+        extra_mask = cv2.dilate(np.array(extra_mask), np.ones((5, 5), np.uint16), iterations=4)
+
+    extra_mask = Image.fromarray(hole_fill(np.array(extra_mask))) 
+
+    extra_mask = ImageOps.invert(extra_mask)
+    extra_mask.save("mask_mine.png")
+
+    if category == 'lower_body_pants':
+        parse_mask = np.logical_or(parse_mask, parser_mask_fixed)
+        parse_mask = np.logical_and(parse_mask, extra_mask)
+    else:
+        parse_mask = np.logical_and(parse_mask, extra_mask)
+        parse_mask = np.logical_or(parse_mask, parser_mask_fixed)
+
+
+    parse_mask_img = Image.fromarray(parse_mask.astype(np.uint8) * 255)
+    parse_mask_img.save("mask_all.png")
+
+    inpaint_mask = 1 - parse_mask
+    
+    #densepose
+    if dense_pose is not None:
+
+        dense_mask = np.zeros((height, width))
+        dense_fixed = np.zeros((height, width))
+
+        dense_foot = (dense_pose == 5).astype(np.float32) + \
+                        (dense_pose == 6).astype(np.float32)
+
+        dense_hand = (dense_pose == 3).astype(np.float32) + \
+                        (dense_pose == 4).astype(np.float32)
+
+        dense_fixed =  dense_foot + dense_hand
+
+        #resolving users' upper clothes in hand
+        up_clothes = (parse_array == 4).astype(np.float32)
+        low_clothes = (parse_array == 6).astype(np.float32) + \
+                        (parse_array == 5).astype(np.float32) +\
+                        (parse_array == 7).astype(np.float32) 
+        up_clothes = cv2.dilate(up_clothes, np.ones((5, 5), np.uint16), iterations=3)
+        low_clothes = cv2.dilate(low_clothes, np.ones((5, 5), np.uint16), iterations=3)
+
+        dense_fixed = np.logical_and(dense_fixed, np.logical_not(up_clothes))
+        dense_fixed = np.logical_and(dense_fixed, np.logical_not(low_clothes))
+        dense_fixed = (dense_fixed).astype(np.float32)
+
+    #masking for upper_clothes and lower_body
+        if category == 'upper_clothes' or category == 'full_body' or category == 'dresses':
+            dense_mask += (dense_pose == 1).astype(np.float32) + \
+                    (dense_pose == 2).astype(np.float32) + \
+                    (dense_pose == 15).astype(np.float32) + \
+                    (dense_pose == 16).astype(np.float32) + \
+                    (dense_pose == 17).astype(np.float32) + \
+                    (dense_pose == 18).astype(np.float32) + \
+                    (dense_pose == 19).astype(np.float32) + \
+                    (dense_pose == 20).astype(np.float32) + \
+                    (dense_pose == 21).astype(np.float32) + \
+                    (dense_pose == 22).astype(np.float32)
+        if category == 'lower_body_pants' or category == 'lower_body_skirts' or category == 'full_body' or category == 'dresses':
+            dense_mask += (dense_pose == 7).astype(np.float32) + \
+                    (dense_pose == 8).astype(np.float32) + \
+                    (dense_pose == 9).astype(np.float32) + \
+                    (dense_pose == 10).astype(np.float32) + \
+                    (dense_pose == 11).astype(np.float32) + \
+                    (dense_pose == 12).astype(np.float32) + \
+                    (dense_pose == 13).astype(np.float32) + \
+                    (dense_pose == 14).astype(np.float32)
+            
+            # if category == 'lower_body_pants' or category == 'lower_body_skirts':
+            #     dense_fixed += (dense_pose == 15).astype(np.float32) + \
+            #             (dense_pose == 16).astype(np.float32) + \
+            #             (dense_pose == 17).astype(np.float32) + \
+            #             (dense_pose == 18).astype(np.float32) + \
+            #             (dense_pose == 19).astype(np.float32) + \
+            #             (dense_pose == 20).astype(np.float32) + \
+            #             (dense_pose == 21).astype(np.float32) + \
+            #             (dense_pose == 22).astype(np.float32)
+            #     dense_fixed = cv2.dilate(dense_fixed, np.ones((5, 5), np.uint16), iterations=1)
+
+            
+            if category == 'lower_body_skirts' or category == 'dresses':
+                #masking giữa 2 chân
+                extra_mask = ImageOps.invert(extra_mask)
+                extra_mask = np.array(extra_mask)
+                extra_mask = cv2.dilate(extra_mask, np.ones((5, 5), np.uint16), iterations=9) 
+                dense_mask = np.logical_or(dense_mask, extra_mask)
+                dense_mask = dense_mask.astype(np.float32)
+
+            if category == "lower_body_pants" :
+                extra_dense_mask = cv2.dilate(dense_mask, np.ones((5, 5), np.uint16), iterations=5)
+                backgroud_mask = (dense_pose == 0).astype(np.float32)
+                extra_dense_mask = np.logical_and(extra_dense_mask, np.logical_not(backgroud_mask))
+                
+                dense_mask = np.logical_or(dense_mask, extra_dense_mask)
+                dense_mask = dense_mask.astype(np.float32)
+        
+        #grow the mask
+        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * 10 + 1, 2 * 10 + 1))
+        dense_mask = cv2.dilate(dense_mask, kernel, iterations=1)
+
+        dense_mask_img = Image.fromarray(dense_mask.astype(np.uint8) * 255)
+        dense_mask_img.save("mask_new.png")
+
+        #refine for upper_clothes
+        if category == 'upper_clothes':
+            mid_y = max(mid_point_left[1], mid_point_right[1])
+            y_grid = np.arange(dense_mask.shape[0]).reshape(-1, 1)
+            lower_half_mask = y_grid > mid_y
+            lower_half_mask = np.tile(lower_half_mask, (1, dense_mask.shape[1]))
+            dense_mask[lower_half_mask] = 0
+
+    inpaint_mask = np.logical_or(inpaint_mask, dense_mask)
+        
+    img = np.where(inpaint_mask, 255, 0)
+    dst = hole_fill(img.astype(np.uint8))
+
+    # inpaint_mask = dst / 255 * 1
+    # inpaint_mask_img = Image.fromarray(inpaint_mask.astype(np.uint8) * 255)
+    # inpaint_mask_img.save("mask_inpaint_before.png")
+
+    dst = refine_mask(dst)
+    inpaint_mask = dst / 255 * 1
+
+    inpaint_mask_img = Image.fromarray(inpaint_mask.astype(np.uint8) * 255)
+    inpaint_mask_img.save("mask_inpaint.png")
+        #refine for upper_clothes
+    
+    #keep hand, foot, head
+    inpaint_mask = np.logical_and(inpaint_mask, np.logical_not(dense_fixed))
+
+    mask = Image.fromarray(inpaint_mask.astype(np.uint8) * 255)
+    mask_gray = Image.fromarray(inpaint_mask.astype(np.uint8) * 127)
+
+    return mask, mask_gray, inpaint_mask
+
+def merge_mask_image(image, mask):
+    mask = mask.convert("L")
+    white_image = Image.new("RGB", image.size, "white")
+    inverted_mask = Image.eval(mask, lambda x: 255 - x)
+    combined_image = Image.composite(image, white_image, inverted_mask)
+    
+    return combined_image
+
+#get bbox from densepose
+def get_bbox_from_densepose(image, densepose_array, padding=0):
+    body_pixels = np.column_stack(np.where(densepose_array > 0))
+
+    if body_pixels.size == 0:
+        return None  # No body pixels found
+
+    min_y, min_x = body_pixels.min(axis=0)
+    max_y, max_x = body_pixels.max(axis=0)
+
+    min_x = max(0, min_x - padding)
+    min_y = max(0, min_y - padding)
+    max_x = min(densepose_array.shape[1], max_x + padding)
+    max_y = min(densepose_array.shape[0], max_y + padding)
+
+    bbox = (min_x, min_y, max_x, max_y)
+
+    mask = np.zeros_like(image)
+    min_x, min_y, max_x, max_y = bbox
+    mask[min_y:max_y, min_x:max_x, :] = 255
+    masked_image = np.where(mask == 255, image, 0)
+    masked_image = Image.fromarray(masked_image)
+    
+    return masked_image
+    
+
+#testing
+import matplotlib.pyplot as plt
+from preprocess.openpose.run_openpose import OpenPose
+from preprocess.humanparsing.run_parsing import Parsing
+
+# from humanparsing.run_parsing import Parsing
+
+if  __name__ == '__main__':
+    device = "cuda" if torch.cuda.is_available() else "cpu" 
+    openpose_model = OpenPose(0)
+    openpose_model.preprocessor.body_estimation.model.to(device)
+    
+    model_image = Image.open('../model1.jpg').copy()
+    model_image = model_image.resize((768, 1024))
+
+    human_img_arg = _apply_exif_orientation(model_image.resize((384,512)))
+    human_img_arg = convert_PIL_to_numpy(human_img_arg, format="BGR")
+
+    args = apply_net.create_argument_parser().parse_args(('show', './configs/densepose_rcnn_R_50_FPN_s1x.yaml', './ckpt/densepose/model_final_162be9.pkl', 'dp_segm', '-v', '--opts', 'MODEL.DEVICE', 'cuda'))
+    dense_pose = args.func(args,human_img_arg)   
+
+    Image.fromarray(dense_pose[0][:,:,::-1]).resize((768,1024)).save("densepose.png")
+
+    dense_pose = dense_pose[1]
+
+    bbox_image = get_bbox_from_densepose(model_image.resize((384,512)), dense_pose, 15)
+    bbox_image.save("zzz.png")
+
+    #get keypoints
+    keypoints = openpose_model(bbox_image)
+
+    parsing_model = Parsing(0)
+    model_parse, _ = parsing_model(model_image.resize((384,512)))
+    model_parse.save("model_parse.png")
+
+    cate = ['upper_clothes', 'lower_body_pants', 'lower_body_skirts', 'dresses', 'full_body', 'lower_body_shoes', 'lower_body_boots']
+    # cate = ['lower_body_pants']
+    for category in cate:
+        mask, mask_gray, mask_arr = get_mask_location_new(category, model_parse, keypoints, width=384, height=512, dense_pose = dense_pose)
+        mask.resize((768, 1024)).save(f"mask_{category}.png")
+
+        model_image = model_image.resize((384, 512))
+        # print("kkkkkkkkk")
+        # mask = Image.open("mask_fixed.png")
+        model_image_end = merge_mask_image(model_image, mask)
+        model_image_end.save(f"model_image_{category}.png")