import os
import cv2


original_image_path = '/home/vybt/Downloads/olk/rmbg_tiennv3/tiennv3_rmbg'
resize_padding_output = '/home/vybt/Downloads/olk/rmbg_tiennv3/deep_lab_tiennv3_resize_lookbook'
where_to_resize_from = '/home/vybt/Downloads/olk/rmbg_tiennv3/deep_lab_tiennv3_lookbook'

def resize_and_padding(im, desired_size):
    """
    :param im: image need to be resize
    :param desired_size: desired size
    :return: image has desired size with black border
    """
    old_size = im.shape[:2]  # old_size is in (height, width) format
    ratio = float(desired_size) / max(old_size)
    new_size = tuple([int(x * ratio) for x in old_size])

    im = cv2.resize(im, (new_size[1], new_size[0]))

    delta_w = desired_size - new_size[1]
    delta_h = desired_size - new_size[0]
    top, bottom = delta_h // 2, delta_h - (delta_h // 2)
    left, right = delta_w // 2, delta_w - (delta_w // 2)

    color = [0, 0, 0]
    new_im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT,
                                value=color)
    old_size_keeping = (old_size[0], old_size[1], top, left, bottom, right)
    return new_im, old_size_keeping


def recovery_size(image, old_size):
    """
    :param image: padding resized image
    :param old_size: size of original image before padding resize
    :return: image without black border
    """
    h, w, top, left, bottom, right = old_size

    cropped = image[top: image.shape[0] - bottom, left:image.shape[1] - right]
    cropped_resize = cv2.resize(cropped, (w, h))

    return cropped_resize


list_file = os.listdir(where_to_resize_from)
for file_name in list_file:

    print(os.path.join(original_image_path, file_name))
    original_image = cv2.imread(os.path.join(original_image_path, file_name.replace('.png', '.jpg')))
    print(original_image.shape)
    output, old_size = resize_and_padding(original_image.copy(), 513)
    deep_lab_image_path = os.path.join(where_to_resize_from, file_name)
    deep_lab_image = cv2.imread(deep_lab_image_path)
    deep_lab_image = recovery_size(deep_lab_image, old_size)
    cv2.imwrite(os.path.join(resize_padding_output, file_name.replace('.png', '.jpg')), deep_lab_image)