poisson_blend.py

""" 
    Example usage:
        python poisson_blend.py -x 1500 -y 2130 -n  3 --cw 168 --ch 168
"""
import numpy as np
import argparse
import random
import cv2

SRC_PATH = '/Users/jackyyeh/kohler/kavii/Fall24/Vitreous_Defects_Wisconsin/glaze_off/IMG_5493.JPG'
TGT_PATH = '/Users/jackyyeh/kohler/kavii/Fall24/Vitreous_Defects_Wisconsin/glaze/IMG_5491.JPG'


def imshow(window_name, img, binary=False):
    if binary:
        _, img = cv2.threshold(img, 0, 255, cv2.THRESH_BINARY)
    cv2.imshow(window_name, img)
    cv2.waitKey(0)


def random_select_coord(mask):
    """
    Randomly selects a coordinate from the non-zero (foreground) region of a binary mask.

    Args:
        mask (numpy.ndarray): The binary mask where the foreground pixels are non-zero.

    Returns:
        tuple: A randomly selected coordinate (x, y).
    """
    # Find all non-zero coordinates (foreground pixels)
    non_zero_coords = np.column_stack(np.where(mask > 0))

    # If there are no non-zero coordinates, return None
    assert len(non_zero_coords) != 0

    # Randomly select one of the coordinates
    selected_coord = tuple(
        non_zero_coords[random.randint(0, len(non_zero_coords) - 1)])
    selected_coord = (selected_coord[1], selected_coord[0])
    return selected_coord


def rotate(angle, mask):
    # Rotate the mask
    # Rotate around the center of the scaled mask
    center = (mask.shape[1] // 2, mask.shape[0] // 2)
    rotation_matrix = cv2.getRotationMatrix2D(center, angle, 1.0)

    # Perform the rotation
    rotated_mask = cv2.warpAffine(
        mask, rotation_matrix, (mask.shape[1], mask.shape[0]))
    return rotated_mask


def scale(img):
    scale_factor = random.uniform(0.5, 2.0)
    scale_crop_src = cv2.resize(
        img, None, fx=scale_factor, fy=scale_factor, interpolation=cv2.INTER_LINEAR)
    return scale_crop_src


def jitter(img):
    # Apply a simple jitter effect by shifting the pixels randomly
    jitter_amount = 5  # Max number of pixels to shift in any direction
    jittered_region = np.zeros_like(img)
    for i in range(img.shape[0]):
        for j in range(img.shape[1]):
            offset_i = min(max(i + np.random.randint(-jitter_amount,
                           jitter_amount + 1), 0), img.shape[0] - 1)
            offset_j = min(max(j + np.random.randint(-jitter_amount,
                           jitter_amount + 1), 0), img.shape[1] - 1)
            jittered_region[i, j] = img[offset_i, offset_j]
    return img


def synthesize(args):
    tl = [args.tl_x, args.tl_y]
    br = (tl[0]+args.crop_width, tl[1]+args.crop_height)

    # Load the source and target images
    src = cv2.imread(SRC_PATH)
    target = cv2.imread(TGT_PATH)

    # Apply seamless cloning
    valid_mask = np.zeros(src.shape[:2])
    valid_mask[args.crop_height:-args.crop_height,
               args.crop_width:-args.crop_width] = 1
    output = target.copy()
    cmp_output = target.copy()
    for i in range(args.num_synthesize):

        # 1) crop path from src image
        crop_src = src[tl[1]:br[1], tl[0]:br[0]]

        # 2) do some augment for cropped patch (scale, rotate, jitter)
        # jitter_src = jitter(crop_src)
        scale_crop_src = scale(crop_src)
        mask = np.zeros(scale_crop_src.shape[:2], dtype=np.uint8) + 255
        angle = random.randint(0, 180)
        rt_src = rotate(angle, scale_crop_src)
        rt_mask = rotate(angle, mask)

        # 3) select center on target image to paste cropped patch
        center = random_select_coord(valid_mask)
        print(center)

        # 4) seamlessClone
        output = cv2.seamlessClone(
            rt_src, output, rt_mask, center, cv2.NORMAL_CLONE)

        # 5) [optional] synthesize image using naitve copy and paste (only for comparison use)
        scale_height, scale_width = scale_crop_src.shape[:2]
        tl_clone = (center[0] - scale_width // 2,
                    (center[1] - scale_height // 2))
        br_clone = (center[0] + scale_width // 2,
                    (center[1] + scale_height // 2))
        cmp_output[tl_clone[1]:tl_clone[1]+rt_src.shape[0],
                   tl_clone[0]:tl_clone[0]+rt_src.shape[1]] = rt_src

        # cv2.rectangle(output, tl_clone, br_clone, (0, 255, 0), thickness=2)

    cv2.rectangle(src, tl, br, (255, 0, 0), thickness=2)

    # Display the result
    tmp_mask = cmp_output > 0
    cmp_output = cmp_output * tmp_mask + target * (1-tmp_mask)
    cmp_output = cmp_output.astype(np.uint8)

    vis = cv2.hconcat([src, output, cmp_output])
    imshow('Seamless Clone Result', vis)


def main():
    seed = 42
    random.seed(seed)
    np.random.seed(seed)

    parser = argparse.ArgumentParser()
    parser.add_argument('--num_synthesize', '-n', type=int, default=10)
    parser.add_argument('--tl_x', '-x', type=int, required=True)
    parser.add_argument('--tl_y', '-y', type=int, required=True)
    parser.add_argument('--crop_width', '--cw', type=int, required=True)
    parser.add_argument('--crop_height', '--ch', type=int, required=True)

    args = parser.parse_args()

    print('======= Initial arguments =======')
    for key, val in vars(args).items():
        print(f"name: {key} => {val}")
    print()

    synthesize(args)


if __name__ == "__main__":
    main()