forked from DmitryUlyanov/deep-image-prior
-
Notifications
You must be signed in to change notification settings - Fork 0
/
generate_inpainting_dataset.py
28 lines (22 loc) · 1010 Bytes
/
generate_inpainting_dataset.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
import shutil
from PIL import Image
from utils.inpainting_utils import get_random_img_mask, gen_text_mask
import glob
import tqdm
import os
import numpy as np
SOURCE_DATASET = './data/sr_datasets/Set5/images/' #'./data/denoising_dataset/' #
TARGET_DATASET = ['./data/inpainting_scribble_dataset/', './data/inpainting_text_dataset',
'./data/inpainitng_scribble_set5/'][2]
os.makedirs(TARGET_DATASET, exist_ok=True)
source_imgs = sorted(glob.glob(SOURCE_DATASET + '/*.png'))
for src_img_path in tqdm.tqdm(source_imgs[1:2]):
src_img_name = os.path.basename(src_img_path)
dst_img_path = os.path.join(TARGET_DATASET, src_img_name)
shutil.copy(src_img_path, dst_img_path)
img_pil = Image.open(src_img_path)
img_np = np.array(img_pil)
curr_mask = get_random_img_mask(img_np)
# curr_mask = gen_text_mask(img_np)
curr_mask_pil = Image.fromarray((curr_mask).astype(np.uint8))
# curr_mask_pil.save(os.path.join(TARGET_DATASET, 'mask_{}'.format(src_img_name)))