for EMAE figures

EMAE_crop.py

import math
import os
import random
from itertools import product

import numpy as np
from imgaug import augmenters as iaa
from PIL import Image

IMG_SIZE = 600
PATCH_SIZE = 100
PATCH_NUM = 36  # 6*6 = 36 patchs, 100*100 each patch
PATCH_LINE = int(math.sqrt(PATCH_NUM))
PADDING = 10
FILENAME = "n04612504_5785.JPEG"
aug_seq = iaa.Sequential(
    [
        iaa.GaussianBlur(sigma=(0, 2.0)),
        iaa.Superpixels(p_replace=(0.1, 0.5), n_segments=(8, 32)),
        iaa.Emboss(alpha=(0.2, 0.5), strength=(0.5, 1.5)),
    ]
)


def tile(filename, resize_hw, d, outname=None, aug=False):
    img_list = []
    img = Image.open(filename)
    img = img.resize(resize_hw)
    w, h = img.size
    img.save(f"{w}x{h}-{filename}")
    grid = product(range(0, h - h % d, d), range(0, w - w % d, d))
    for i, j in grid:
        box = (j, i, j + d, i + d)
        img_crop = img.crop(box)
        img_list.append(img_crop)
        if outname:
            if aug:
                img_tile = np.array(img_crop)
                img_tile = aug_seq(image=img_tile)
                img_tile = Image.fromarray(img_tile)
            else:
                img_tile = img_crop
            img_tile.save(f"{outname}_{i//d}_{j//d}.jpg")
    return img_list


def concat_images_group(images):
    width = max(image.width for image in images)
    height = sum(image.height + PADDING for image in images) - PADDING
    composite = Image.new("RGB", (width, height), color=(255, 255, 255))
    y = 0
    for image in images:
        composite.paste(image, (0, y))
        y += image.height + PADDING
    return composite


def concat_images_all(images):

    padding_num = PATCH_LINE - 1
    width = max(image.width for image in images) * PATCH_LINE + PADDING * padding_num
    height = max(image.height for image in images) * PATCH_LINE + PADDING * padding_num
    composite_all = Image.new("RGB", (width, height), color=(255, 255, 255))
    y = 0
    for iw in range(PATCH_LINE):
        x = 0
        for iy in range(PATCH_LINE):
            patch_idx = iw * PATCH_LINE + iy
            img = images[patch_idx]
            composite_all.paste(img, (x, y))
            x += img.width + PADDING
        y += img.height + PADDING
    return composite_all


def concat_images_mask(images, patch_index, aug=False):
    img_each_line = int(math.sqrt(PATCH_NUM))
    padding_num = img_each_line - 1
    width = max(image.width for image in images) * img_each_line + PADDING * padding_num
    height = (
        max(image.height for image in images) * img_each_line + PADDING * padding_num
    )
    composite_input = Image.new("RGB", (width, height), color=(255, 255, 255))
    composite_target = Image.new("RGB", (width, height), color=(255, 255, 255))
    y = 0
    for iw in range(img_each_line):
        x = 0
        for iy in range(img_each_line):
            patch_idx = iw * img_each_line + iy
            img = images[patch_idx]
            mask_patch = Image.new("RGB", (img.width, img.height), color=(60, 60, 60))
            if patch_idx in patch_index:
                composite_target.paste(mask_patch, (x, y))
                composite_input.paste(img, (x, y))
            else:
                if aug:
                    img = np.array(img)
                    img = aug_seq(image=img)
                    img = Image.fromarray(img)
                composite_target.paste(img, (x, y))
                composite_input.paste(mask_patch, (x, y))
            x += img.width + PADDING
        y += img.height + PADDING
    return composite_input, composite_target


def gen_fig1():
    img_list = tile(FILENAME, (IMG_SIZE, IMG_SIZE), PATCH_SIZE)

    # x = [i for i in range(PATCH_NUM)]
    # random.shuffle(x)
    # path_num_each_group = PATCH_NUM // 4
    # patch_index = [
    #     x[i : i + path_num_each_group] for i in range(0, len(x), path_num_each_group)
    # ]
    patch_index = [
        [6, 10, 16, 25, 4, 9, 30, 8, 5],
        [2, 14, 32, 18, 11, 24, 22, 23, 27],
        [21, 15, 31, 1, 26, 29, 20, 13, 19],
        [7, 12, 35, 3, 33, 28, 0, 17, 34],
    ]
    cnt = 0
    for pi in patch_index:
        print(pi)
        composite_group = [img_list[idx] for idx in pi]
        composite = concat_images_group(composite_group)
        composite.save(f"out_{cnt}_vis_patch.jpg")
        composite_input, composite_target = concat_images_mask(img_list, pi, aug=True)
        composite_input.save(f"out_{cnt}_patch_input.jpg")
        composite_target.save(f"out_{cnt}_patch_target.jpg")
        cnt += 1

    composite = concat_images_all(img_list)
    composite.save(f"out_all_patch_no_mask.jpg")


def gen_fig2():
    tile(FILENAME, (IMG_SIZE, IMG_SIZE), PATCH_SIZE, "part1", aug=True)
    tile(FILENAME, (IMG_SIZE, IMG_SIZE), PATCH_SIZE, "part2", aug=False)


if __name__ == "__main__":
    gen_fig1()
    gen_fig2()