-
Notifications
You must be signed in to change notification settings - Fork 0
/
make_raw_images_from_directory.py
61 lines (47 loc) · 1.96 KB
/
make_raw_images_from_directory.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import argparse
import glob
import numpy
import os
from PIL import Image
import sys
from paint_transfer.config import Config
from paint_transfer.data_process import BaseDataProcess
from paint_transfer import dataset
from paint_transfer.utility.image import array_to_image, save_images, save_tiled_image
parser = argparse.ArgumentParser()
parser.add_argument('--config_json_path')
parser.add_argument('--image_directory', required=True)
parser.add_argument('--path_save_directory')
args = parser.parse_args()
config_json_path = args.config_json_path
path_save = args.path_save_directory
os.makedirs(path_save, exist_ok=True)
config = Config(config_json_path)
class CropTopImageSquareProcess(BaseDataProcess):
def __call__(self, image: Image.Image, test):
width = image.size[0]
image = image.crop((0, 0, width, width))
return image
dataset.data_process._process.insert(3, CropTopImageSquareProcess())
paths = glob.glob(os.path.join(args.image_directory, '*'))
num_image = len(paths)
dataset_test = dataset.DataProcessDataset(paths, data_process=dataset.data_process, test=True)
# save images
images = []
raw_lines = []
for i, data in enumerate(dataset_test[:num_image]):
image = data['target'][numpy.newaxis]
image = array_to_image(image)
images += image
raw_line = data['raw_line'][numpy.newaxis]
raw_line = 1 - raw_line
raw_line = numpy.repeat(raw_line, 3, axis=1)
raw_line = array_to_image(raw_line, minmax=(0, 1))
raw_lines += raw_line
save_images(images, path_save, '')
save_images(raw_lines, path_save, 'raw_line_')
# save tiled images
paths_input = [os.path.join(path_save, '{}.png'.format(i)) for i in range(num_image)]
save_tiled_image(paths_input, os.path.join(path_save, 'tile{}.png'.format(num_image)), col=num_image)
paths_input = [os.path.join(path_save, 'raw_line_{}.png'.format(i)) for i in range(num_image)]
save_tiled_image(paths_input, os.path.join(path_save, 'tile_raw_line_{}.png'.format(num_image)), col=num_image)