from pathlib import Path

import cv2
import numpy as np
import pandas as pd


def maskblur(mask, kernel_size, sigma=1):
    mask_blur = cv2.GaussianBlur(mask, (kernel_size, kernel_size), sigma)
    return mask_blur


def erosion(mask, kernel_size):
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    erosion_image = cv2.erode(mask, kernel, iterations=1)  # // make erosion image
    return erosion_image


def dilate(mask, kernel_size):
    kernel = np.ones((kernel_size, kernel_size), np.uint8)
    erosion_image = cv2.dilate(mask, kernel, iterations=1)  # // make erosion image
    return erosion_image


def resize_adapt(model_out, crop_region):
    def inter_alg(target_size, img):
        if isinstance(target_size, tuple):
            w, h = target_size
        else:
            w, h = target_size, target_size
        return inter_alg_(w, h, img)

    def inter_alg_(w, h, img):
        if w * h < img.shape[0] * img.shape[1]:
            return cv2.INTER_AREA
        else:
            return cv2.INTER_CUBIC

    x1, y1, x2, y2 = crop_region
    h, w = y2 - y1 + 1, x2 - x1 + 1
    sz = model_out.shape[0]  # h,w 동일하다.
    if h == sz and w == sz:
        return model_out

    r = max(h, w) / sz
    max_hw = max(h, w)
    temp_ = cv2.resize(model_out, (max_hw, max_hw), inter_alg(max_hw, model_out))
    temp_ = temp_[
        (max_hw - h) // 2 : (max_hw - h) // 2 + h,
        (max_hw - w) // 2 : (max_hw - w) // 2 + w,
    ]
    return temp_


def get_face_mask(
    img_size, df_fan_row, blur_ratio=0.3, dilate_ratio=0.2, erosion_ratio=0
):
    assert blur_ratio >= 0 and blur_ratio <= 1
    assert erosion_ratio >= 0 and erosion_ratio <= 1
    assert dilate_ratio >= 0 and dilate_ratio <= 1

    def _masking(img, pts, value):
        img = cv2.fillPoly(img, [pts], value)
        return img

    def _get_face_pts_n_box(img_size, df_fan_row):
        box = df_fan_row["cropped_box"]
        pts2d = df_fan_row["pts2d"] - np.array([box[0], box[1]])

        if isinstance(df_fan_row["cropped_size"], float):
            cropped_size = df_fan_row["cropped_size"]
        else:
            cropped_size = df_fan_row["cropped_size"][0]
        ratio = img_size[0] / cropped_size
        pts2d = pts2d * ratio
        xs, ys = pts2d[:, 0], pts2d[:, 1]
        l, t, r, b = min(xs), min(ys), max(xs), max(ys)
        return np.concatenate([pts2d[0:17, :], pts2d[17:27, :][::-1]]).astype(
            np.int32
        ), (l, t, r, b)

    if df_fan_row["pts2d"] is None:
        mask = np.zeros((img_size[1], img_size[0]), dtype=np.uint8)
        if len(mask.shape) == 2:
            mask = np.expand_dims(mask, axis=2)
        return {"crop": mask, "origin": 1 - mask}

    pts, box = _get_face_pts_n_box(img_size, df_fan_row)
    h = max(box[2] - box[0], box[3] - box[1])
    mask = np.zeros((img_size[1], img_size[0]), dtype=np.uint8)
    mask = _masking(mask, pts, (255))
    if dilate_ratio != 0:
        mask = dilate(mask, int(h * dilate_ratio) // 2 * 2 + 1)
    if erosion_ratio != 0:
        mask = erosion(mask, int(h * erosion_ratio) // 2 * 2 + 1)
    if blur_ratio != 0:
        blur_kernel_size = int(h * blur_ratio) // 2 * 2 + 1
        mask = maskblur(mask, blur_kernel_size, 0)
    mask = mask / 255
    if len(mask.shape) == 2:
        mask = np.expand_dims(mask, axis=2)
    return {"crop": mask, "origin": 1 - mask}


def cromakey_green(img):
    r = img[:, :, 0]
    g = img[:, :, 1]
    b = img[:, :, 2]
    g_alpha = g > 50
    r_alpha = (g * 1.0) > r
    b_alpha = (g * 0.7) > b
    alpha = g_alpha & (r_alpha & b_alpha)
    alpha = (1 - alpha) * 255
    alpha = alpha.astype(np.uint8)

    alpha = maskblur(alpha, kernel_size=13)
    alpha[np.where(alpha > 100)] = 255
    alpha = erosion(alpha, kernel_size=5)

    if len(alpha.shape) == 2:
        alpha2 = np.expand_dims(alpha, axis=2)
    else:
        alpha2 = alpha
    new = np.concatenate((img, alpha2), axis=2)
    return new


def cromakey_green_binary(img):
    img = cromakey_green(img)
    alpha = img[:, :, 3]
    alpha[np.where(alpha <= 128)] = 0
    alpha[np.where(alpha > 128)] = 1


def cromakey_green_hunet_lmy(img):
    r = img[:, :, 0]
    g = img[:, :, 1]
    b = img[:, :, 2]
    g_alpha = g > 70
    r_alpha = g > r
    b_alpha = (g * 0.8) > b
    alpha = g_alpha & (r_alpha & b_alpha)
    alpha = (1 - alpha) * 255
    alpha = alpha.astype(np.uint8)

    alpha = maskblur(alpha, kernel_size=11)
    alpha[np.where(alpha > 100)] = 255
    alpha = maskblur(alpha, kernel_size=3)
    alpha = erosion(alpha, kernel_size=3)

    if len(alpha.shape) == 2:
        alpha2 = np.expand_dims(alpha, axis=2)
    else:
        alpha2 = alpha
    new = np.concatenate((img, alpha2), axis=2)
    return new


# ybm 영상용 크로마키 함수
def cromakey_green_ybm_front(img):
    r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]

    g_alpha = g > 70
    # r_alpha = (g * 0.7) > r
    # b_alpha = (g * 0.7) > b
    r_alpha = g > r
    b_alpha = (g * 0.9) > b

    alpha = g_alpha & (r_alpha & b_alpha)
    alpha = (1 - alpha) * 255
    alpha = alpha.astype(np.uint8)

    alpha = maskblur(alpha, kernel_size=11)
    alpha[np.where(alpha > 100)] = 255
    alpha = maskblur(alpha, kernel_size=3)
    alpha = maskblur(alpha, kernel_size=3)

    grey_alpha = alpha < 255
    g[grey_alpha] = r[grey_alpha] * 0.8

    if len(alpha.shape) == 2:
        alpha2 = np.expand_dims(alpha, axis=2)
    else:
        alpha2 = alpha
    new = np.concatenate((img, alpha2), axis=2)
    return new


# ybm 영상용 크로마키 함수
def cromakey_green_ybm_side(img):
    img = img.copy()
    r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]
    g_alpha = g > 50
    r_alpha = g > r
    b_alpha = (g * 0.9) > b

    alpha = g_alpha & (r_alpha & b_alpha)
    alpha = (1 - alpha) * 255
    alpha = alpha.astype(np.uint8)

    alpha = maskblur(alpha, kernel_size=11)
    alpha[np.where(alpha > 100)] = 255
    alpha = maskblur(alpha, kernel_size=3)
    alpha = maskblur(alpha, kernel_size=3)

    grey_alpha = alpha < 255
    g[grey_alpha] = r[grey_alpha] * 0.8

    if len(alpha.shape) == 2:
        alpha2 = np.expand_dims(alpha, axis=2)
    else:
        alpha2 = alpha
    new = np.concatenate((img, alpha2), axis=2)
    return new


# devin 영상용 크로마키 함수
def cromakey_green_devin_side(img):
    img = img.copy()
    r, g, b = img[:, :, 0], img[:, :, 1], img[:, :, 2]
    g_alpha = g > 70
    r_alpha = (g * 0.8) > r
    # r_alpha = g > r
    b_alpha = (g * 0.9) > b

    alpha = g_alpha & (r_alpha & b_alpha)
    alpha = (1 - alpha) * 255
    alpha = alpha.astype(np.uint8)

    alpha = maskblur(alpha, kernel_size=7, sigma=3)
    alpha[np.where(alpha < 150)] = 0
    alpha = maskblur(alpha, kernel_size=5, sigma=2)

    if len(alpha.shape) == 2:
        alpha = np.expand_dims(alpha, axis=2)

    new = np.concatenate((img, alpha), axis=2)
    return new


def get_cromakey_func(args):
    if "cromakey" not in args.keys():
        return cromakey_green_hunet_lmy
    if "cromakey_green_ybm_front" == args.cromakey:
        return cromakey_green_ybm_front
    if "cromakey_green_ybm_side" == args.cromakey:
        return cromakey_green_ybm_side
    if "cromakey_green_devin_side" == args.cromakey:
        return cromakey_green_devin_side

    raise "cromakey not found"


def compose_default_(model_out, org_image_with_alpha, mask, **kwargs):
    # 1. 마스크 섞기 : 원래 비디오의 투명값과 계산한 마스크를 섞는다.
    mask = mask[:, :, 0]
    mask[np.where(mask > 0)] = 1  # 마스크 영역을 128 -> 1 로 만든다.
    model_out[:, :, 3] = (
        org_image_with_alpha[:, :, 3] * (1 - mask) + model_out[:, :, 3] * mask
    )

    # 2. 섞인 마스크가 좀 자연스럽게 섞이도록 함.
    model_out[:, :, 3] = maskblur(model_out[:, :, 3], kernel_size=3, sigma=1)
    return model_out


def compose_devin_(model_out, org_image_with_alpha, mask, debug=False, **kwargs):
    mask = mask[:, :, 0]
    mask[np.where(mask > 0)] = 1  # 마스크 영역을 128 -> 1 로 만든다.
    mask = mask.astype(np.float32)

    # 1. 기존마스크와 경계가 잘 안보이도록 마스크를 부드럽게 만든다.
    kernel_size = int(mask.shape[0] * 0.03) // 2 * 2 + 1  # 이미지 크기의 3% 정도 마스크를 확장한다.
    if debug:
        print(
            f"## compose_devin_: kernel_size:{kernel_size}, mask_height:{mask.shape[0]}"
        )
    if kernel_size >= 3:
        mask = dilate(mask, kernel_size=kernel_size)
        mask = maskblur(mask, kernel_size=kernel_size, sigma=kernel_size // 2)
        mask = maskblur(mask, kernel_size=kernel_size, sigma=kernel_size // 2)
        mask = erosion(mask, kernel_size=3)  # 1pixel 만 줄임

    # 2. 마스크 섞기 : 원래 비디오의 투명값과 계산한 마스크를 섞는다.
    model_out[:, :, 3] = (
        org_image_with_alpha[:, :, 3] * (1 - mask) + model_out[:, :, 3] * mask
    )

    # 3. 섞인 마스크가 부드럽게 한번더 블러를 한다.
    model_out[:, :, 3] = maskblur(model_out[:, :, 3], kernel_size=3, sigma=1)

    return model_out


def get_compose_mask_func(args):
    if "cromakey" in args.keys():
        if "cromakey_green_devin_side" == args.cromakey:
            return compose_devin_
    if "compose" in args.keys():
        if "compose_smooth" == args.compose:
            return compose_devin_
    return compose_default_


def get_keying_func(template):
    cromakey_func = get_cromakey_func(template.model.args)
    compose_func = get_compose_mask_func(template.model.args)

    def keying_(pred, idx, box=None):
        model_out, mask, alpha = pred["pred"], pred["mask"], pred["img_gt_with_alpha"]

        if pred["filename"].endswith("_no.jpg") or pred["filename"].endswith("_no.png"):
            return alpha[:, :, [2, 1, 0, 3]]

        if (
            alpha.shape[0] != mask.shape[0]
            or alpha.shape[1] != mask.shape[1]
            or alpha.shape[0] != model_out.shape[0]
            or alpha.shape[1] != model_out.shape[1]
        ):
            raise Exception(
                f"not matched keying shape. "
                f"alpha: {alpha.shape[0]}, {alpha.shape[1]}, {alpha.shape[2]}, "
                f"mask: {mask.shape[0]}, {mask.shape[1]}, "
                f"model_out: {model_out.shape[0]}, {model_out.shape[1]}"
            )

        if box is not None:
            model_h = model_out.shape[0]
            box_h = box[3] - box[1]
            if box_h > model_h:
                model_out = resize_adapt(model_out, box)
                mask = resize_adapt(mask, box)
                alpha = resize_adapt(alpha, box)
        model_out = cromakey_func(model_out)
        model_out = compose_func(
            model_out=model_out, org_image_with_alpha=alpha, mask=mask
        )

        return model_out

    return keying_


def get_box_mask(width, height, config, verbose=False):
    def get_mask_(
        width, height, gradation_width, gradation_bottom=None, box_mask_erosion=None
    ):
        mask = np.ones((height, width, 1))
        r = list(range(0, gradation_width, 1))
        for s, e in zip(r, r[1:]):
            g = s / gradation_width
            # print(f'---- s:{s}, e:{e}, g:{g}')
            mask[s:e, s : width - s, :] = g
            mask[height - e : height - s, s : width - s, :] = g
            mask[s : height - s, s:e, :] = g
            mask[s : height - s, width - e : width - s, :] = g
        if gradation_bottom is not None:
            r = list(range(0, gradation_bottom, 1))
            for s, e in zip(r, r[1:]):
                g = s / gradation_bottom
                mask[height - e : height - s, s : width - s, :] = g
        if box_mask_erosion is not None:
            mask = erosion(mask, box_mask_erosion * 2 + 1)
            if len(mask.shape) == 2:
                mask = np.expand_dims(mask, 2)  # mask shape ex: (352,352,1)

        return mask

    gradation_width = int(height * 0.1)
    gradation_bottom = (
        int(height * config["gradation_bottom"])
        if "gradation_bottom" in config.keys()
        else None
    )
    box_mask_erosion = (
        int(height * config["box_mask_erosion"])
        if "box_mask_erosion" in config.keys()
        else None
    )
    # if verbose:
    #    print('gradation_width : ', gradation_width)
    #    print('gradation_bottom : ', gradation_bottom)
    #    print('box_mask_erosion : ', box_mask_erosion)
    mask = get_mask_(width, height, gradation_width, gradation_bottom, box_mask_erosion)
    mask_crop = mask
    mask_origin = 1 - mask
    return {"crop": mask_crop, "origin": mask_origin}


def get_compose_func_without_keying_move(template, ratio, verbose=False):
    args = template.model.args
    df = pd.read_pickle(
        f"{template.crop_mp4_dir}/{Path(template.template_video_path).stem}_000/df_fan.pickle"
    )
    df = df.set_index("frame_idx")
    move_head_box_size = (
        (df.loc[0]["cropped_box"][2] - df.loc[0]["cropped_box"][0] - 20) // 10 * 10
    )

    def resize_and_scale(model_out, head_box_idx):
        # ratio 1.0 에 맞는 크기로 resize 하고,
        # 원래 영상에서 10의 배수에 해당하는 위치로 (head_box, model_out) 모두 잘라낸다.
        head_box = df["cropped_box"][head_box_idx]
        if ratio == 1.0:
            return model_out, head_box

        # 일단 원래 크기로 만든다.
        model_out = resize_adapt(model_out, head_box)

        # 원래 크기에서의 박스에서 10의 배수에 해당하는 좌표를 찾는다.
        l, t = (np.array(head_box[:2]) + 9) // 10 * 10
        new_head_box = np.array(
            [l, t, l + move_head_box_size - 1, t + move_head_box_size - 1]
        )  # 양쪽포함이라서 1을 빼준다.

        # 10의 배수에 맞춰서 이미지를 잘라낸다.
        diff_box = new_head_box - head_box
        new_model_out = model_out[diff_box[1] : diff_box[3], diff_box[0] : diff_box[2]]
        # if verbose and head_box_idx == 0:
        #     print('org head_box:', head_box, ', new_head_box:', new_head_box)
        #     print('alpah2.shape:', model_out.shape, ', new_model_out:', new_model_out.shape)
        if (
            new_model_out.shape[0] % 10 != 0 or new_model_out.shape[1] % 10 != 0
        ):  # 크기는 10의 배수여야 한다.
            raise Exception(f"new_model_out.shape % 10 != 0, {new_model_out.shape}")

        # ratio에 맞는 크기로 변경한다.
        x1, y1, _, _ = np.round(new_head_box * ratio).astype(np.uint8)
        # 양쪽포함이라서 -1을 해준다.
        new_head_box = (
            x1,
            y1,
            x1 + int(move_head_box_size * ratio) - 1,
            y1 + int(move_head_box_size * ratio) - 1,
        )
        new_model_out = resize_adapt(new_model_out, new_head_box)

        # if verbose and head_box_idx == 0:
        #     print('org head_box:', head_box, ', new_head_box:', new_head_box)
        #     print('alpah2.shape:', model_out.shape, ', new_model_out:', new_model_out.shape)

        return new_model_out, new_head_box

    def compose_one(model_out, full_img, head_box_idx):
        model_out, box = resize_and_scale(model_out, head_box_idx)
        x1, y1, x2, y2 = box
        img = resize_adapt(model_out, (x1, y1, x2, y2))
        if (
            "compose" in template.config.keys()
            and template.config.compose == "face_only"
        ):
            row = df.loc[head_box_idx]
            mask_box = get_face_mask(
                (img.shape[1], img.shape[0]), row, **get_compose_option(template.config)
            )
        else:
            mask_box = get_box_mask(
                x2 - x1 + 1, y2 - y1 + 1, config=args, verbose=verbose
            )

        if y2 - y1 + 1 != img.shape[0] or x2 - x1 + 1 != img.shape[1]:
            raise Exception(
                f"not matched compose shape. x2-x1+1: {x2 - x1 + 1}, y2-y1+1:{y2 - y1 + 1}, img: {img.shape[1]}, {img.shape[0]}"
            )

        # Compose the image
        if full_img.shape[2] == 3:
            alpha = np.zeros_like(full_img[:, :, :1])
            alpha.fill(255)
            full_img = np.concatenate([full_img, alpha], axis=2)

        out_memory = full_img.copy()

        alpha = img[:, :, 3]
        alpha = cv2.merge([alpha, alpha, alpha])

        back = out_memory[y1 : y2 + 1, x1 : x2 + 1].copy()
        front = img[:, :, 0:3]

        img = np.concatenate(
            [np.where(alpha < (255, 255, 255), back[:, :, :3], front), back[:, :, 3:]],
            axis=2,
        )
        out_memory[y1 : y2 + 1, x1 : x2 + 1] = (
            full_img[y1 : y2 + 1, x1 : x2 + 1] * mask_box["origin"]
            + img * mask_box["crop"]
        )
        return out_memory

    return compose_one


def get_compose_func_without_keying_default(template, ratio, verbose=False):
    args = template.model.args
    df = pd.read_pickle(
        f"{template.crop_mp4_dir}/{Path(template.template_video_path).stem}_000/df_fan.pickle"
    )
    # sz = df['cropped_size'].values[0]
    # 원래 4k 템플릿에서 축소된 비율만큼 cropped_box 크기를 줄여준다.
    x1, y1, x2, y2 = np.round(np.array(df["cropped_box"].values[0]) * ratio).astype(
        np.uint8
    )
    del df
    mask_box = get_box_mask(x2 - x1 + 1, y2 - y1 + 1, config=args, verbose=verbose)
    img_size = args.img_size
    if verbose:
        print("croped size: ", x2 - x1 + 1, y2 - y1 + 1)
        print("croped region(x1,y1,x2,y2): ", x1, y1, x2, y2)

    def compose_one(model_out, full_img, _):
        img = resize_adapt(model_out, (x1, y1, x2, y2))
        if y2 - y1 + 1 != img.shape[0] or x2 - x1 + 1 != img.shape[1]:
            raise Exception(
                f"not matched compose shape. x2-x1+1: {x2 - x1 + 1}, y2-y1+1:{y2 - y1 + 1}, img: {img.shape[1]}, {img.shape[0]}"
            )

        # 붙여넣기
        if full_img.shape[2] == 3:
            alpha = np.zeros_like(full_img[:, :, :1])
            alpha.fill(255)
            full_img = np.concatenate([full_img, alpha], axis=2)

        out_memory = full_img.copy()

        alpha = img[:, :, 3]
        alpha = cv2.merge([alpha, alpha, alpha])

        back = out_memory[y1 : y2 + 1, x1 : x2 + 1].copy()
        front = img[:, :, 0:3]

        img = np.concatenate(
            [np.where(alpha < (255, 255, 255), back[:, :, :3], front), back[:, :, 3:]],
            axis=2,
        )

        out_memory[y1 : y2 + 1, x1 : x2 + 1] = (
            full_img[y1 : y2 + 1, x1 : x2 + 1] * mask_box["origin"]
            + img * mask_box["crop"]
        )
        return out_memory

    return compose_one


def get_compose_option(config):
    blur_ratio = 0.3
    dilate_ratio = 0.2
    erosion_ratio = 0.0
    if "compose_args" in config.keys():
        if "blur_ratio" in config.compose_args.keys():
            blur_ratio = config.compose_args.blur_ratio
        if "dilate_ratio" in config.compose_args.keys():
            dilate_ratio = config.compose_args.dilate_ratio
        if "erosion_ratio" in config.compose_args.keys():
            erosion_ratio = config.compose_args.erosion_ratio
    return {
        "blur_ratio": blur_ratio,
        "dilate_ratio": dilate_ratio,
        "erosion_ratio": erosion_ratio,
    }


def get_compose_func_without_keying_face_only(template, ratio, verbose=False):
    df = pd.read_pickle(
        f"{template.crop_mp4_dir}/{Path(template.template_video_path).stem}_000/df_fan.pickle"
    )
    x1, y1, x2, y2 = np.round(np.array(df["cropped_box"].values[0]) * ratio).astype(
        np.uint8
    )

    df = df.set_index("frame_idx")
    if verbose:
        print("get_compose_option")
        print(get_compose_option(template.config))

    def compose_one(model_out, full_img, head_box_idx):
        try:
            row = df.loc[head_box_idx]
        except Exception as e:
            print("exception get_compose_func_without_keying_face_only", e)
            raise Exception("exception get_compose_func_without_keying_face_only", e)
        img = resize_adapt(model_out, (x1, y1, x2, y2))
        if y2 - y1 + 1 != img.shape[0] or x2 - x1 + 1 != img.shape[1]:
            raise Exception(
                f"not matched compose shape. x2-x1+1: {x2 - x1 + 1}, y2-y1+1:{y2 - y1 + 1}, img: {img.shape[1]}, {img.shape[0]}"
            )

        mask_box = get_face_mask(
            (img.shape[1], img.shape[0]), row, **get_compose_option(template.config)
        )
        # 붙여넣기
        out_memory = full_img.copy()
        out_memory[y1 : y2 + 1, x1 : x2 + 1] = (
            full_img[y1 : y2 + 1, x1 : x2 + 1] * mask_box["origin"]
            + img * mask_box["crop"]
        )
        return out_memory

    return compose_one


# template video 의 frame 과 model inference 결과를 합성하는 함수를 리턴한다.
# params
#     ratio : 템플릿 scale 비율.
#             1.0: 템플릿 크기 그대로
#             0.5: width, height 를 절반으로 줄인 크기
def get_compose_func_without_keying(template, ratio, verbose=False):
    if "move" in template.config.keys() and template.config.move:
        return get_compose_func_without_keying_move(
            template=template, ratio=ratio, verbose=verbose
        )
    if "compose" in template.config.keys() and template.config.compose == "face_only":
        return get_compose_func_without_keying_face_only(
            template=template, ratio=ratio, verbose=verbose
        )

    return get_compose_func_without_keying_default(
        template=template, ratio=ratio, verbose=verbose
    )


def compose_direct(box, model_args, ratio, model_out, full_img):
    x1, y1, x2, y2 = box
    mask_box = get_box_mask(x2 - x1 + 1, y2 - y1 + 1, config=model_args)
    img_size = model_args.img_size

    img = resize_adapt(model_out, (x1, y1, x2, y2))
    if y2 - y1 + 1 != img.shape[0] or x2 - x1 + 1 != img.shape[1]:
        raise Exception(
            f"not matched compose shape. x2-x1+1: {x2 - x1 + 1}, y2-y1+1:{y2 - y1 + 1}, img: {img.shape[1]}, {img.shape[0]}"
        )

    # 붙여넣기
    out_memory = full_img.copy()
    out_memory[y1 : y2 + 1, x1 : x2 + 1] = (
        full_img[y1 : y2 + 1, x1 : x2 + 1] * mask_box["origin"] + img * mask_box["crop"]
    )
    return out_memory


def keying_direct(model_args, pred, box=None):
    cromakey_func = get_cromakey_func(model_args)
    compose_func = get_compose_mask_func(model_args)

    model_out, mask, alpha = pred["pred"], pred["mask"], pred["img_gt_with_alpha"]

    if pred["filename"].endswith("_no.jpg") or pred["filename"].endswith("_no.png"):
        return alpha[:, :, [2, 1, 0, 3]]

    if (
        alpha.shape[0] != mask.shape[0]
        or alpha.shape[1] != mask.shape[1]
        or alpha.shape[0] != model_out.shape[0]
        or alpha.shape[1] != model_out.shape[1]
        or alpha.shape[2] != 4
    ):
        raise Exception(
            f"not matched keying shape. "
            f"alpha: {alpha.shape[0]}, {alpha.shape[1]}, {alpha.shape[2]}, "
            f"mask: {mask.shape[0]}, {mask.shape[1]}, "
            f"model_out: {model_out.shape[0]}, {model_out.shape[1]}"
        )

    if box is not None:
        model_h = model_out.shape[0]
        box_h = box[3] - box[1]
        if box_h > model_h:
            model_out = resize_adapt(model_out, box)
            mask = resize_adapt(mask, box)
            alpha = resize_adapt(alpha, box)
    model_out = cromakey_func(model_out)
    model_out = compose_func(model_out=model_out, org_image_with_alpha=alpha, mask=mask)

    return model_out