Skip to content

Instantly share code, notes, and snippets.

@BIGBALLON
Last active November 4, 2022 10:48
Show Gist options
  • Save BIGBALLON/42de49b6b57a7d764b616292dab46f47 to your computer and use it in GitHub Desktop.
Save BIGBALLON/42de49b6b57a7d764b616292dab46f47 to your computer and use it in GitHub Desktop.
for EMAE figures
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment