utils_mask.py

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")