diff options
| -rw-r--r-- | augment.py | 139 | ||||
| -rwxr-xr-x | build_dataset.pl | 13 | ||||
| -rw-r--r-- | data/recursive_dataset.py | 2 | ||||
| -rw-r--r-- | data/sequence_dataset.py | 12 | ||||
| -rw-r--r-- | live.py | 7 | ||||
| -rw-r--r-- | options/dataset_options.py | 16 |
6 files changed, 176 insertions, 13 deletions
diff --git a/augment.py b/augment.py new file mode 100644 index 0000000..5edbc78 --- /dev/null +++ b/augment.py @@ -0,0 +1,139 @@ +### Copyright (C) 2017 NVIDIA Corporation. All rights reserved. +### Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode). +import os +import sys +sys.path.append(os.path.join(os.path.dirname(os.path.realpath(__file__)), '../live-cortex/rpc/')) +from collections import OrderedDict +from options.test_options import TestOptions +from options.dataset_options import DatasetOptions +from data.data_loader import CreateDataLoader +from models.models import create_model +import util.util as util +from util.visualizer import Visualizer +from util import html +import torch +from run_engine import run_trt_engine, run_onnx +from datetime import datetime +from PIL import Image, ImageOps +from shutil import copyfile, rmtree +from random import randint + +from img_ops import read_sequence + +import torch.utils.data as data +from PIL import Image +import torchvision.transforms as transforms + +def get_transform(opt, method=Image.BICUBIC, normalize=True): + transform_list = [] + + base = float(2 ** opt.n_downsample_global) + if opt.netG == 'local': + base *= (2 ** opt.n_local_enhancers) + transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method))) + + transform_list += [transforms.ToTensor()] + + if normalize: + transform_list += [transforms.Normalize((0.5, 0.5, 0.5), + (0.5, 0.5, 0.5))] + return transforms.Compose(transform_list) + +def normalize(): + return transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) + +def __make_power_2(img, base, method=Image.BICUBIC): + ow, oh = img.size + h = int(round(oh / base) * base) + w = int(round(ow / base) * base) + if (h == oh) and (w == ow): + return img + return img.resize((w, h), method) + + +opt = TestOptions().parse(save=False) +data_opt = DatasetOptions().parse(opt.unknown) +opt.nThreads = 1 # test code only supports nThreads = 1 +opt.batchSize = 1 # test code only supports batchSize = 1 +opt.serial_batches = True # no shuffle +opt.no_flip = True # no flip +if data_opt.tag == '': + d = datetime.now() + tag = data_opt.tag = "{}_{}".format( + opt.name, + # opt.experiment, + d.strftime('%Y%m%d%H%M') + ) +else: + tag = data_opt.tag + +opt.render_dir = render_dir = opt.results_dir + opt.name + "/" + tag + "/" + +print('tag:', tag) +print('render_dir:', render_dir) +util.mkdir(render_dir) + +data_loader = CreateDataLoader(opt) +dataset = data_loader.load_data() + +if not opt.engine and not opt.onnx: + model = create_model(opt) + if opt.data_type == 16: + model.half() + elif opt.data_type == 8: + model.type(torch.uint8) + if opt.verbose: + print(model) + +sequence = read_sequence(data_opt.sequence_name, '') +print("Got sequence {}, {} images".format(data_opt.sequence, len(sequence))) +_len = len(sequence) - data_opt.augment_take + +if _len <= 0: + print("Got empty sequence...") + data_opt.processing = False + rpc_client.send_status('processing', False) + sys.exit(1) + +transform = get_transform(opt) + +for m in range(data_opt.augment_take): + i = randint(0, _len) + index = i + + for n in range(data_opt.augment_make): + index = i + n + if n == 0: + A_path = sequence[i] + A = Image.open(A_path) + A_tensor = transform(A.convert('RGB')) + else: + A_path = os.path.join(self.opt.render_dir, "recur_{:05d}_{:05d}.png".format(m, index)) + A = Image.open(A_path) + A_tensor = transform(A.convert('RGB')) + B_path = sequence[index+1] + inst_tensor = 0 + + input_dict = {'label': A_tensor, 'inst': inst_tensor} + + if opt.data_type == 16: + data['label'] = data['label'].half() + data['inst'] = data['inst'].half() + elif opt.data_type == 8: + data['label'] = data['label'].uint8() + data['inst'] = data['inst'].uint8() + minibatch = 1 + generated = model.inference(data['label'], data['inst']) + + tmp_path = os.path.join(opt.render_dir, "recur_{:05d}_{:05d}_tmp.png".format(m, index+1)) + next_path = os.path.join(opt.render_dir, "recur_{:05d}_{:05d}.png".format(m, index+1)) + print('process image... %i' % index) + + im = util.tensor2im(generated.data[0]) + image_pil = Image.fromarray(im, mode='RGB') + image_pil.save(tmp_path) + os.rename(tmp_path, next_path) + + os.symlink(next_path, os.path.join("./datasets/", data_opt.sequence, "train_A", "recur_{:05d}_{:05d}.png".format(m, index+1))) + os.symlink(sequence[i+1], os.path.join("./datasets/", data_opt.sequence, "train_B", "recur_{:05d}_{:05d}.png".format(m, index+1))) + diff --git a/build_dataset.pl b/build_dataset.pl index f007038..f88b818 100755 --- a/build_dataset.pl +++ b/build_dataset.pl @@ -39,12 +39,19 @@ my $dir; my $src; my $dst; for ($i = 0; $i < $count; $i++) { - if (($i % $test_split) != ($test_split-1)) { - $dir = $thumbs_dir . "train_"; - } else { + if (($i % $test_split) = ($test_split-1)) { $dir = $thumbs_dir . "test_"; + + $src = $images_dir . $images[$i]; + $dst = $dir . sprintf("A/frame_%05d.png", $i); + system("ln", "-s", $src, $dst); + + $src = $images_dir . $images[$i+1]; + $dst = $dir . sprintf("B/frame_%05d.png", $i); + system("ln", "-s", $src, $dst); } + $dir = $thumbs_dir . "train_"; $src = $images_dir . $images[$i]; $dst = $dir . sprintf("A/frame_%05d.png", $i); system("ln", "-s", $src, $dst); diff --git a/data/recursive_dataset.py b/data/recursive_dataset.py index 40b7ebf..dda9c95 100644 --- a/data/recursive_dataset.py +++ b/data/recursive_dataset.py @@ -13,7 +13,7 @@ class RecursiveDataset(BaseDataset): ### input A (label maps) self.dataset_size = 1000000 - + def __getitem__(self, index): ### input A (label maps) A_path = os.path.join(self.opt.render_dir, "frame_{:05d}.png".format(index)) diff --git a/data/sequence_dataset.py b/data/sequence_dataset.py index 3eaa12b..c3c7d44 100644 --- a/data/sequence_dataset.py +++ b/data/sequence_dataset.py @@ -8,18 +8,18 @@ from PIL import Image class SequenceDataset(BaseDataset): def initialize(self, opt): self.opt = opt - self.root = opt.dataroot + self.root = opt.dataroot ### input A (label maps) self.dir_A = opt.dataroot self.A_paths = sorted(make_dataset(self.dir_A)) - self.dataset_size = len(self.A_paths) - - def __getitem__(self, index): + self.dataset_size = len(self.A_paths) + + def __getitem__(self, index): ### input A (label maps) - A_path = self.A_paths[index] - A = Image.open(A_path) + A_path = self.A_paths[index] + A = Image.open(A_path) params = get_params(self.opt, A.size) if self.opt.label_nc == 0: transform_A = get_transform(self.opt, params) @@ -24,8 +24,8 @@ import gevent from time import sleep from shutil import copyfile, rmtree -from img_ops import process_image, mix_next_image -from listener import Listener, read_sequence +from img_ops import read_sequence, process_image, mix_next_image +from listener import Listener module_name = 'pix2pixhd' @@ -134,7 +134,8 @@ def process_live_input(opt, data_opt, rpc_client): if data_opt.pause: data_opt.pause = False break - gevent.sleep(data_opt.frame_delay) + if data_opt.frame_delay > 0: + gevent.sleep(data_opt.frame_delay) data_opt.processing = False rpc_client.send_status('processing', False) diff --git a/options/dataset_options.py b/options/dataset_options.py index d9c3bf7..29e07c6 100644 --- a/options/dataset_options.py +++ b/options/dataset_options.py @@ -293,6 +293,22 @@ class DatasetOptions(BaseOptions): help='amount of processed image (clahe, poster, etc) to use in feeder step' ) + ### DATA AUGMENTATION + + self.parser.add_argument( + '--augment-take', + default=100, + type=int, + help='number of random images to take' + ) + + self.parser.add_argument( + '--augment-make', + default=15, + type=int, + help='number of recursive images to generate' + ) + ### GRAYSCALE self.parser.add_argument( |