CSD_MT_eval.py

import os
import torch
import cv2
import os.path as osp
import numpy as np
from PIL import Image
from CSD_MT.options import Options
from CSD_MT.model import CSD_MT
from faceutils.face_parsing.model import BiSeNet
import torchvision.transforms as transforms
import faceutils as futils

import warnings
warnings.filterwarnings("ignore", category=FutureWarning, module="torch")

# load face_parsing model
n_classes = 19
face_paseing_model = BiSeNet(n_classes=n_classes)
save_pth = osp.join('faceutils/face_parsing/res/cp', '79999_iter.pth')
face_paseing_model.load_state_dict(torch.load(save_pth,map_location='cpu'))
face_paseing_model.eval()

# load makeup transfer model
parser = Options()
opts = parser.parse()
makeup_model = CSD_MT(opts)
ep0, total_it = makeup_model.resume('CSD_MT/weights/CSD_MT.pth')
makeup_model.eval()

def crop_image(image):
    up_ratio = 0.2 / 0.85  # delta_size / face_size
    down_ratio = 0.15 / 0.85  # delta_size / face_size
    width_ratio = 0.2 / 0.85  # delta_size / face_size

    image = Image.fromarray(image)
    face = futils.dlib.detect(image)

    if not face:
        raise ValueError("No face !")

    face_on_image = face[0]

    image, face, crop_face = futils.dlib.crop(image, face_on_image, up_ratio, down_ratio, width_ratio)
    np_image = np.array(image)
    return np_image

def get_face_parsing(x):
    to_tensor = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
    ])
    with torch.no_grad():
        img = Image.fromarray(x)
        image = img.resize((512, 512), Image.BILINEAR)
        img = to_tensor(image)
        img = torch.unsqueeze(img, 0)
        out = face_paseing_model(img)[0]
        parsing = out.squeeze(0).cpu().numpy().argmax(0)
    return parsing


def split_parse(opts,parse):
    h, w = parse.shape
    result = np.zeros([h, w, opts.semantic_dim])
    result[:, :, 0][np.where(parse == 0)] = 1
    result[:, :, 0][np.where(parse == 16)] = 1
    result[:, :, 0][np.where(parse == 17)] = 1
    result[:, :, 0][np.where(parse == 18)] = 1
    result[:, :, 0][np.where(parse == 9)] = 1
    result[:, :, 1][np.where(parse == 1)] = 1
    result[:, :, 2][np.where(parse == 2)] = 1
    result[:, :, 2][np.where(parse == 3)] = 1
    result[:, :, 3][np.where(parse == 4)] = 1
    result[:, :, 3][np.where(parse == 5)] = 1
    result[:, :, 1][np.where(parse == 6)] = 1
    result[:, :, 4][np.where(parse == 7)] = 1
    result[:, :, 4][np.where(parse == 8)] = 1
    result[:, :, 5][np.where(parse == 10)] = 1
    result[:, :, 6][np.where(parse == 11)] = 1
    result[:, :, 7][np.where(parse == 12)] = 1
    result[:, :, 8][np.where(parse == 13)] = 1
    result[:, :, 9][np.where(parse == 14)] = 1
    result[:, :, 9][np.where(parse == 15)] = 1
    result = np.array(result)
    return result


def local_masks(opts,split_parse):
    h, w, c = split_parse.shape
    all_mask = np.zeros([h, w])
    all_mask[np.where(split_parse[:, :, 0] == 0)] = 1
    all_mask[np.where(split_parse[:, :, 3] == 1)] = 0
    all_mask[np.where(split_parse[:, :, 6] == 1)] = 0
    all_mask = np.expand_dims(all_mask, axis=2)  # Expansion of the dimension
    all_mask = np.concatenate((all_mask, all_mask, all_mask), axis=2)
    return all_mask



def load_data_from_image(non_makeup_img, makeup_img,opts):
    non_makeup_img=crop_image(non_makeup_img)
    makeup_img = crop_image(makeup_img)
    non_makeup_img=cv2.resize(non_makeup_img,(opts.resize_size,opts.resize_size))
    makeup_img = cv2.resize(makeup_img, (opts.resize_size, opts.resize_size))
    non_makeup_parse = get_face_parsing(non_makeup_img)
    non_makeup_parse = cv2.resize(non_makeup_parse, (opts.resize_size, opts.resize_size),interpolation=cv2.INTER_NEAREST)
    makeup_parse = get_face_parsing(makeup_img)
    makeup_parse = cv2.resize(makeup_parse, (opts.resize_size, opts.resize_size),interpolation=cv2.INTER_NEAREST)

    non_makeup_split_parse = split_parse(opts,non_makeup_parse)
    makeup_split_parse = split_parse(opts,makeup_parse)

    non_makeup_all_mask = local_masks(opts,
        non_makeup_split_parse)
    makeup_all_mask = local_masks(opts,
        makeup_split_parse)

    non_makeup_img = non_makeup_img / 127.5 - 1
    non_makeup_img = np.transpose(non_makeup_img, (2, 0, 1))
    non_makeup_split_parse = np.transpose(non_makeup_split_parse, (2, 0, 1))

    makeup_img = makeup_img / 127.5 - 1
    makeup_img = np.transpose(makeup_img, (2, 0, 1))
    makeup_split_parse = np.transpose(makeup_split_parse, (2, 0, 1))

    non_makeup_img=torch.from_numpy(non_makeup_img).type(torch.FloatTensor)
    non_makeup_img = torch.unsqueeze(non_makeup_img, 0)
    non_makeup_split_parse = torch.from_numpy(non_makeup_split_parse).type(torch.FloatTensor)
    non_makeup_split_parse = torch.unsqueeze(non_makeup_split_parse, 0)
    non_makeup_all_mask = np.transpose(non_makeup_all_mask, (2, 0, 1))

    makeup_img = torch.from_numpy(makeup_img).type(torch.FloatTensor)
    makeup_img = torch.unsqueeze(makeup_img, 0)
    makeup_split_parse = torch.from_numpy(makeup_split_parse).type(torch.FloatTensor)
    makeup_split_parse = torch.unsqueeze(makeup_split_parse, 0)
    makeup_all_mask = np.transpose(makeup_all_mask, (2, 0, 1))

    data = {'non_makeup_color_img': non_makeup_img,
            'non_makeup_split_parse':non_makeup_split_parse,
            'non_makeup_all_mask': torch.unsqueeze(torch.from_numpy(non_makeup_all_mask).type(torch.FloatTensor), 0),

            'makeup_color_img': makeup_img,
            'makeup_split_parse': makeup_split_parse,
            'makeup_all_mask': torch.unsqueeze(torch.from_numpy(makeup_all_mask).type(torch.FloatTensor), 0)
            }
    return data


def extract_eye_mask(parsing, expansion=20, upward_bias=20, side_bias=20):
    # 눈 영역 마스크 생성
    
    eye_mask = np.zeros_like(parsing, dtype=np.uint8)
    eye_mask[np.where(parsing == 4)] = 1  # 왼쪽 눈
    eye_mask[np.where(parsing == 5)] = 1  # 오른쪽 눈

    # 기본 확장
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (expansion, expansion))
    expanded_mask = cv2.dilate(eye_mask, kernel, iterations=1)

    upward_mask = np.zeros_like(expanded_mask)
    upward_mask[:-upward_bias, :] = expanded_mask[upward_bias:, :]

    left_mask = np.zeros_like(expanded_mask)
    left_mask[:, :-side_bias] = expanded_mask[:, side_bias:]

    right_mask = np.zeros_like(expanded_mask)
    right_mask[:, side_bias:] = expanded_mask[:, :-side_bias]

    final_mask = np.clip(expanded_mask + upward_mask + left_mask + right_mask, 0, 1)

    return final_mask

def extract_eyebrow_mask(parsing):
    # 눈썹 마스크 생성

    eyebrow_mask = np.zeros_like(parsing, dtype=np.uint8)
    eyebrow_mask[np.where(parsing == 2)] = 1  # 왼쪽 눈썹
    eyebrow_mask[np.where(parsing == 3)] = 1  # 오른쪽 눈썹
    return eyebrow_mask

def extract_lips_mask(parsing):
    # 입술 마스크 생성

    lips_mask = np.zeros_like(parsing, dtype=np.uint8)
    lips_mask[np.where(parsing == 12)] = 1  # 윗입술
    lips_mask[np.where(parsing == 13)] = 1  # 아랫입술
    return lips_mask


def makeup_transfer256(non_makeup_image, makeup_image, alpha_values, regions):
    """
    메이크업 전이 함수: 영역별로 다른 alpha 값을 사용하여 특정 영역에 필터 적용.
    """
    # 메이크업 전이 수행
    data = load_data_from_image(non_makeup_image, makeup_image, opts=opts)
    with torch.no_grad():
        transfer_tensor = makeup_model.test_pair(data)
        transfer_img = transfer_tensor[0].cpu().float().numpy()
        transfer_img = np.transpose((transfer_img / 2 + 0.5) * 255., (1, 2, 0))
        transfer_img = np.clip(transfer_img, 0, 255).astype(np.uint8)

    # 원본 이미지 크기에 맞게 리사이즈
    target_size = (non_makeup_image.shape[1], non_makeup_image.shape[0])
    transfer_img = cv2.resize(transfer_img, target_size, interpolation=cv2.INTER_LINEAR)

    # 얼굴 파싱 및 영역별 마스크 생성
    non_makeup_parse = get_face_parsing(non_makeup_image)
    masks = {
        "eye": extract_eye_mask(non_makeup_parse),
        "eyebrow": extract_eyebrow_mask(non_makeup_parse),
        "lip": extract_lips_mask(non_makeup_parse),
    }

    # 결과 이미지 생성
    result_image = non_makeup_image.astype(np.float32)
    transfer_img = transfer_img.astype(np.float32)

    # 선택된 영역에만 메이크업 적용
    for region in regions:
        mask = masks.get(region, None)
        if mask is not None:
            mask = cv2.resize(mask, target_size, interpolation=cv2.INTER_NEAREST)
            mask = cv2.GaussianBlur(mask.astype(np.float32), (9, 9), 0)
            mask = mask / mask.max()

            alpha = alpha_values.get(region, 1)  # 해당 영역의 alpha 값 가져오기
            for c in range(3):  # RGB 채널별 적용
                result_image[:, :, c] = result_image[:, :, c] * (1 - alpha * mask) + transfer_img[:, :, c] * (
                    alpha * mask
                )

    # 전체 영역에 대한 처리 (regions="all")
    if "all" in regions:
        alpha = alpha_values.get("all", 1)
        for c in range(3):
            result_image[:, :, c] = result_image[:, :, c] * (1 - alpha) + transfer_img[:, :, c] * alpha

    return result_image.astype(np.uint8)