preprocess/humanparsing/parsing_api.py

import pdb
from pathlib import Path
import sys

PROJECT_ROOT = Path(__file__).absolute().parents[0].absolute()
sys.path.insert(0, str(PROJECT_ROOT))
import os
import torch
import numpy as np
import cv2
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from datasets.simple_extractor_dataset import SimpleFolderDataset
from utils.transforms import transform_logits
from tqdm import tqdm
from PIL import Image


def colorize_parsing(parsing_result):
    label_map = {
        0: "background", 1: "hat", 2: "hair", 3: "sunglasses", 4: "upper_clothes",
        5: "skirt", 6: "pants", 7: "dress", 8: "belt", 9: "left_shoe",
        10: "right_shoe", 11: "head", 12: "left_leg", 13: "right_leg",
        14: "left_arm", 15: "right_arm", 16: "bag", 17: "scarf"
    }
    
    # Define colors for each part (RGB)
    color_map = {
        0: (0, 0, 0),       # Background
        1: (128, 0, 0),     # Hat
        2: (255, 0, 0),     # Hair
        3: (0, 255, 0),     # Sunglasses
        4: (0, 0, 255),     # Upper-clothes
        5: (255, 255, 0),   # Skirt
        6: (255, 0, 255),   # Pants
        7: (0, 255, 255),   # Dress
        8: (128, 128, 0),   # Belt
        9: (0, 128, 128),   # Left-shoe
        10: (128, 0, 128),  # Right-shoe
        11: (128, 128, 128),# Head
        12: (64, 0, 0),     # Left-leg
        13: (192, 0, 0),    # Right-leg
        14: (64, 128, 0),   # Left-arm
        15: (192, 128, 0),  # Right-arm
        16: (64, 0, 128),   # Bag
        17: (192, 0, 128),  # Scarf
    }
    
    height, width = parsing_result.shape
    colored_parsing = np.zeros((height, width, 3), dtype=np.uint8)
    
    for label, color in color_map.items():
        colored_parsing[parsing_result == label] = color
    
    return colored_parsing

def add_numbers_to_image(colored_parsing, parsing_result):
    label_map = {
        0: "background", 1: "hat", 2: "hair", 3: "sunglasses", 4: "upper_clothes",
        5: "skirt", 6: "pants", 7: "dress", 8: "belt", 9: "left_shoe",
        10: "right_shoe", 11: "head", 12: "left_leg", 13: "right_leg",
        14: "left_arm", 15: "right_arm", 16: "bag", 17: "scarf"
    }
    
    height, width = parsing_result.shape
    numbered_image = colored_parsing.copy()
    
    for label in range(18):  # 0 to 17
        mask = (parsing_result == label)
        if np.any(mask):
            y, x = np.where(mask)
            center_y, center_x = int(np.mean(y)), int(np.mean(x))
            
            cv2.putText(numbered_image, str(label), (center_x, center_y),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
    
    return numbered_image

def get_palette(num_cls):
    """ Returns the color map for visualizing the segmentation mask.
    Args:
        num_cls: Number of classes
    Returns:
        The color map
    """
    n = num_cls
    palette = [0] * (n * 3)
    for j in range(0, n):
        lab = j
        palette[j * 3 + 0] = 0
        palette[j * 3 + 1] = 0
        palette[j * 3 + 2] = 0
        i = 0
        while lab:
            palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
            palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
            palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
            i += 1
            lab >>= 3
    return palette


def delete_irregular(logits_result):
    parsing_result = np.argmax(logits_result, axis=2)
    upper_cloth = np.where(parsing_result == 4, 255, 0)
    contours, hierarchy = cv2.findContours(upper_cloth.astype(np.uint8),
                                           cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
    area = []
    for i in range(len(contours)):
        a = cv2.contourArea(contours[i], True)
        area.append(abs(a))
    if len(area) != 0:
        top = area.index(max(area))
        M = cv2.moments(contours[top])
        cY = int(M["m01"] / M["m00"])

    dresses = np.where(parsing_result == 7, 255, 0)
    contours_dress, hierarchy_dress = cv2.findContours(dresses.astype(np.uint8),
                                                       cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
    area_dress = []
    for j in range(len(contours_dress)):
        a_d = cv2.contourArea(contours_dress[j], True)
        area_dress.append(abs(a_d))
    if len(area_dress) != 0:
        top_dress = area_dress.index(max(area_dress))
        M_dress = cv2.moments(contours_dress[top_dress])
        cY_dress = int(M_dress["m01"] / M_dress["m00"])
    wear_type = "dresses"
    if len(area) != 0:
        if len(area_dress) != 0 and cY_dress > cY:
            irregular_list = np.array([4, 5, 6])
            logits_result[:, :, irregular_list] = -1
        else:
            irregular_list = np.array([5, 6, 7, 8, 9, 10, 12, 13])
            logits_result[:cY, :, irregular_list] = -1
            wear_type = "cloth_pant"
        parsing_result = np.argmax(logits_result, axis=2)
    # pad border
    parsing_result = np.pad(parsing_result, pad_width=1, mode='constant', constant_values=0)
    return parsing_result, wear_type



def hole_fill(img):
    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)
        # keep large area in skin case
        for j in range(len(area)):
          if j != i and area[i] > 2000:
             cv2.drawContours(refine_mask, contours, j, color=255, thickness=-1)
    return refine_mask

def refine_hole(parsing_result_filled, parsing_result, arm_mask):
    filled_hole = cv2.bitwise_and(np.where(parsing_result_filled == 4, 255, 0),
                                  np.where(parsing_result != 4, 255, 0)) - arm_mask * 255
    contours, hierarchy = cv2.findContours(filled_hole, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_L1)
    refine_hole_mask = np.zeros_like(parsing_result).astype(np.uint8)
    for i in range(len(contours)):
        a = cv2.contourArea(contours[i], True)
        # keep hole > 2000 pixels
        if abs(a) > 2000:
            cv2.drawContours(refine_hole_mask, contours, i, color=255, thickness=-1)
    return refine_hole_mask + arm_mask

def onnx_inference(session, lip_session, input_dir):
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229])
    ])
    dataset = SimpleFolderDataset(root=input_dir, input_size=[512, 512], transform=transform)
    dataloader = DataLoader(dataset)
    with torch.no_grad():
        for _, batch in enumerate(tqdm(dataloader)):
            image, meta = batch
            c = meta['center'].numpy()[0]
            s = meta['scale'].numpy()[0]
            w = meta['width'].numpy()[0]
            h = meta['height'].numpy()[0]
            output = session.run(None, {"input.1": image.numpy().astype(np.float32)})
            upsample = torch.nn.Upsample(size=[512, 512], mode='bilinear', align_corners=True)
            upsample_output = upsample(torch.from_numpy(output[1][0]).unsqueeze(0))
            upsample_output = upsample_output.squeeze()
            upsample_output = upsample_output.permute(1, 2, 0)  # CHW -> HWC
            logits_result = transform_logits(upsample_output.data.cpu().numpy(), c, s, w, h, input_size=[512, 512])
            parsing_result = np.argmax(logits_result, axis=2)
            parsing_result = np.pad(parsing_result, pad_width=1, mode='constant', constant_values=0)
            # try holefilling the clothes part
            arm_mask = (parsing_result == 14).astype(np.float32) \
                       + (parsing_result == 15).astype(np.float32)
            upper_cloth_mask = (parsing_result == 4).astype(np.float32) + arm_mask
            img = np.where(upper_cloth_mask, 255, 0)
            dst = hole_fill(img.astype(np.uint8))
            parsing_result_filled = dst / 255 * 4
            parsing_result_woarm = np.where(parsing_result_filled == 4, parsing_result_filled, parsing_result)
            # add back arm and refined hole between arm and cloth
            refine_hole_mask = refine_hole(parsing_result_filled.astype(np.uint8), parsing_result.astype(np.uint8),
                                           arm_mask.astype(np.uint8))
            parsing_result = np.where(refine_hole_mask, parsing_result, parsing_result_woarm)
            # remove padding
            parsing_result = parsing_result[1:-1, 1:-1]

        dataset_lip = SimpleFolderDataset(root=input_dir, input_size=[473, 473], transform=transform)
        dataloader_lip = DataLoader(dataset_lip)
        with torch.no_grad():
            for _, batch in enumerate(tqdm(dataloader_lip)):
                image, meta = batch
                c = meta['center'].numpy()[0]
                s = meta['scale'].numpy()[0]
                w = meta['width'].numpy()[0]
                h = meta['height'].numpy()[0]

                output_lip = lip_session.run(None, {"input.1": image.numpy().astype(np.float32)})
                upsample = torch.nn.Upsample(size=[473, 473], mode='bilinear', align_corners=True)
                upsample_output_lip = upsample(torch.from_numpy(output_lip[1][0]).unsqueeze(0))
                upsample_output_lip = upsample_output_lip.squeeze()
                upsample_output_lip = upsample_output_lip.permute(1, 2, 0)  # CHW -> HWC
                logits_result_lip = transform_logits(upsample_output_lip.data.cpu().numpy(), c, s, w, h,
                                                     input_size=[473, 473])
                parsing_result_lip = np.argmax(logits_result_lip, axis=2)
    # add neck parsing result
    neck_mask = np.logical_and(np.logical_not((parsing_result_lip == 13).astype(np.float32)),
                               (parsing_result == 11).astype(np.float32))
    parsing_result = np.where(neck_mask, 18, parsing_result)
    palette = get_palette(19)
    output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8))
    output_img.putpalette(palette)
    face_mask = torch.from_numpy((parsing_result == 11).astype(np.float32))

    # Colorize the parsing result
    colored_parsing = colorize_parsing(parsing_result)
    
    # Add numbers to the colorized image
    numbered_parsing = add_numbers_to_image(colored_parsing, parsing_result)
    
    # Save the numbered parsing result
    output_filename =  "colored_parsing.png"
    cv2.imwrite(output_filename, cv2.cvtColor(numbered_parsing, cv2.COLOR_RGB2BGR))

    return output_img, face_mask