def test_dataset_desc_file(): ddf = 'data/fake_datasets.yml' datasets = load_datasets(ddf) assert len(datasets) is 9 assert len(datasets.NONE.train.hr.compile()) is 0 assert len(datasets.NORMAL.train.hr.compile()) is 7 assert len(datasets.NORMAL.val.hr.compile()) is 5 assert len(datasets.NORMAL.test.hr.compile()) is 1 assert len(datasets.PAIR.train.hr.compile()) is 2 assert len(datasets.PAIR.train.lr.compile()) is 2 assert len(datasets.VIDEOPAIR.train.hr.compile()) is 1 assert len(datasets.VIDEOPAIR.train.lr.compile()) is 1 assert len(datasets.VIDEOPAIR.val.hr.compile()) is 1 assert len(datasets.VIDEOPAIR.val.lr.compile()) is 1 assert len(datasets.VIDEOPAIR.test.hr.compile()) is 1 assert len(datasets.VIDEOPAIR.test.lr.compile()) is 1 assert len(datasets.FOO.test.hr.compile()) is 2 assert len(datasets.BAR.test.hr.compile()) is 5 assert datasets.VIDEOPAIR.train.hr.as_video assert datasets.XIUXIAN.test.hr.as_video raw = load_datasets(ddf, 'RAW') assert len(raw.train.hr.compile()) is 1 assert len(raw.val.hr.compile()) is 1 assert len(raw.test.hr.compile()) is 1 assert raw.train.hr.as_video
def test_dataset_desc_file(self): ddf = 'data/fake_datasets.yml' datasets = load_datasets(ddf) self.assertEqual(len(datasets), 9) self.assertEqual(len(datasets.NONE.train.hr.compile()), 0) self.assertEqual(len(datasets.NORMAL.train.hr.compile()), 7) self.assertEqual(len(datasets.NORMAL.val.hr.compile()), 5) self.assertEqual(len(datasets.NORMAL.test.hr.compile()), 1) self.assertEqual(len(datasets.PAIR.train.hr.compile()), 2) self.assertEqual(len(datasets.PAIR.train.lr.compile()), 2) self.assertEqual(len(datasets.VIDEOPAIR.train.hr.compile()), 1) self.assertEqual(len(datasets.VIDEOPAIR.train.lr.compile()), 1) self.assertEqual(len(datasets.VIDEOPAIR.val.hr.compile()), 1) self.assertEqual(len(datasets.VIDEOPAIR.val.lr.compile()), 1) self.assertEqual(len(datasets.VIDEOPAIR.test.hr.compile()), 1) self.assertEqual(len(datasets.VIDEOPAIR.test.lr.compile()), 1) self.assertEqual(len(datasets.FOO.test.hr.compile()), 2) self.assertEqual(len(datasets.BAR.test.hr.compile()), 5) self.assertTrue(datasets.VIDEOPAIR.train.hr.as_video) self.assertTrue(datasets.XIUXIAN.test.hr.as_video) raw = load_datasets(ddf, 'RAW') self.assertEqual(len(raw.train.hr.compile()), 1) self.assertEqual(len(raw.val.hr.compile()), 1) self.assertEqual(len(raw.test.hr.compile()), 1) self.assertTrue(raw.train.hr.as_video)
def fetch_datasets(data_config_file, opt): all_datasets = load_datasets(data_config_file) dataset = all_datasets[opt.dataset.upper()] if opt.test: test_data = all_datasets[opt.test.upper()] else: test_data = dataset if opt.infer: infer_dir = Path(opt.infer) if infer_dir.exists(): # infer files in this directory if infer_dir.is_file(): images = [str(infer_dir)] else: images = list(infer_dir.glob('*')) if not images: images = infer_dir.iterdir() infer_data = Dataset(infer=images, mode='pil-image1', modcrop=False) else: infer_data = all_datasets[opt.infer.upper()] else: infer_data = test_data # TODO temp use, delete if not need if opt.cifar: cifar_data, cifar_test = tf.keras.datasets.cifar10.load_data() dataset = Dataset(**dataset) dataset.mode = 'numpy' dataset.train = [cifar_data[0]] dataset.val = [cifar_test[0]] return dataset, test_data, infer_data
def fetch_datasets(data_config_file, opt): all_datasets = load_datasets(data_config_file) dataset = all_datasets[opt.dataset.upper()] if opt.test: test_data = all_datasets[opt.test.upper()] else: test_data = dataset if opt.infer: infer_dir = Path(opt.infer) if infer_dir.exists(): # infer files in this directory if infer_dir.is_file(): images = [str(infer_dir)] else: images = list(infer_dir.glob('*')) if not images: images = infer_dir.iterdir() infer_data = Dataset(infer=images, mode='pil-image1', modcrop=False) else: infer_data = all_datasets[opt.infer.upper()] else: infer_data = test_data if opt.mode: dataset.mode = opt.mode test_data.mode = opt.mode infer_data.mode = opt.mode return dataset, test_data, infer_data
def main(): model = InformationDistillationNetwork(3, rgb_input=False).compile() dataset = load_datasets('../Data/datasets.json')['91-IMAGE'] dataset.setattr(patch_size=48, depth=1, random=True, max_patches=64 * 300) env = Environment(model, f'{model.name}/save', f'{model.name}/log') env.fit(64, 100, dataset, restart=False, learning_rate=1e-5) env.output_callbacks = [save_image(f'{model.name}/test')] env.test(dataset)
def main(): model = ResidualDenseNetwork(3, rdb_blocks=10, rdb_conv=6, rgb_input=False).compile() dataset = load_datasets('../Data/datasets.json')['BSD'] dataset.setattr(patch_size=96, depth=1, random=True, max_patches=64 * 1) env = Environment(model, f'{model.name}/save', f'{model.name}/log') env.fit(64, 1, dataset, restart=True) env.output_callbacks = [save_image(f'{model.name}/test')] env.test(dataset)
def main(): flags, args = parser.parse_known_args() opt = Config() for pair in flags._get_kwargs(): opt.setdefault(*pair) data_config_file = Path(flags.data_config) if not data_config_file.exists(): raise RuntimeError("dataset config file doesn't exist!") for _ext in ('json', 'yaml', 'yml'): # for compat # apply a 2-stage (or master-slave) configuration, master can be # override by slave model_config_root = Path('Parameters/root.{}'.format(_ext)) if opt.p: model_config_file = Path(opt.p) else: model_config_file = Path('Parameters/{}.{}'.format( opt.model, _ext)) if model_config_root.exists(): opt.update(Config(str(model_config_root))) if model_config_file.exists(): opt.update(Config(str(model_config_file))) model_params = opt.get(opt.model, {}) opt.update(model_params) suppress_opt_by_args(model_params, *args) model = get_model(flags.model)(**model_params) if flags.cuda: model.cuda() root = f'{flags.save_dir}/{flags.model}' if flags.comment: root += '_' + flags.comment verbosity = logging.DEBUG if flags.verbose else logging.INFO trainer = model.trainer datasets = load_datasets(data_config_file) dataset = datasets[flags.dataset.upper()] train_config = Config(crop=opt.train_data_crop, feature_callbacks=[], label_callbacks=[], convert_to='rgb', **opt) if opt.channel == 1: train_config.convert_to = 'gray' if opt.lr_decay: train_config.lr_schedule = lr_decay(lr=opt.lr, **opt.lr_decay) train_config.random_val = not opt.traced_val train_config.cuda = flags.cuda if opt.verbose: dump(opt) with trainer(model, root, verbosity, opt.pth) as t: if opt.seed is not None: t.set_seed(opt.seed) tloader = QuickLoader(dataset, 'train', train_config, True, flags.thread) vloader = QuickLoader(dataset, 'val', train_config, False, flags.thread, batch=1, crop=opt.val_data_crop, steps_per_epoch=opt.val_num) t.fit([tloader, vloader], train_config) if opt.export: t.export(opt.export)
def main(*args, **kwargs): args = argparse.ArgumentParser() args.add_argument('name', type=str, choices=list_supported_models(), help='the model name can be found in model_alias.py') args.add_argument('--scale', type=int, default=3, help='scale factor, default 3') args.add_argument('--channel', type=int, default=1, help='image channels, default 1') args.add_argument('--dataconfig', type=str, default='../Data/datasets.json', help='the path to dataset config json file') args.add_argument( '--dataset', type=str, default='91-IMAGE', help= 'specified dataset name(as described in config file, default 91-image') args.add_argument('--batch', type=int, default=64, help='training batch size, default 64') args.add_argument('--epochs', type=int, default=200, help='training epochs, default 200') args.add_argument( '--patch_size', type=int, default=48, help= 'patch size of cropped training and validating sub-images, default 48') args.add_argument( '--strides', type=int, default=48, help='crop stride if random_patches is set 0, default 48') args.add_argument('--depth', type=int, default=1, help='image1 depth used for video sources, default 1') args.add_argument( '--random_patches', type=int, default=0, help= 'if set more than 0, use random crop to generate `random_patches` sub-image1 batches' ) args.add_argument('--retrain', type=int, default=0, help='retrain the model from scratch, default 0') args.add_argument('--lr', type=float, default=1e-4, help='initial learning rate, default 1e-4') args.add_argument('--lr_decay', type=float, default=1, help='learning rate decay, default 1') args.add_argument('--lr_decay_step', type=int, default=1000, help='learning rate decay step') args.add_argument( '--add_noise', type=float, default=None, help='if not None, add noise with given stddev to input features') args.add_argument( '--add_random_noise', type=list, default=None, help= 'if not None, add random noise with given stddev bound [low, high, step=1]' ) args.add_argument( '--test', type=str, default=None, help='specify a dataset used to test, or use --dataset values if None') args.add_argument('--predict', type=str, default=None, help='evaluate model on given files') args.add_argument( '--savedir', type=str, default='../Results', help='directory to save model checkpoints, default ../Results') args.add_argument('--output_color', type=str, default='RGB', choices=('RGB', 'L', 'GRAY', 'Y'), help='output color mode, default RGB') args.add_argument('--output_index', type=int, default=-1, help='access index of model outputs to save, default -1') args.add_argument( '--export_pb', type=str, default=None, help= 'if not None, specify the path that export trained model into pb format' ) args.add_argument('--comment', type=str, default=None) args.add_argument('--custom_feature_cb', type=str, default=None) args = args.parse_args() model_args = json.load(open(f'parameters/{args.name}.json', mode='r')) model = get_model(args.name)(scale=args.scale, channel=args.channel, **model_args) dataset = load_datasets(args.dataconfig)[args.dataset.upper()] dataset.setattr(patch_size=args.patch_size, strides=args.strides, depth=args.depth) if args.random_patches: dataset.setattr(random=True, max_patches=args.batch * args.random_patches) save_root = f'{args.savedir}/{model.name}_sc{args.scale}_c{args.channel}' if args.comment: save_root += '_' + args.comment with Environment(model, f'{save_root}/save', f'{save_root}/log', feature_index=model.feature_index, label_index=model.label_index) as env: if args.add_noise: env.feature_callbacks = [add_noise(args.add_noise)] if args.add_random_noise: env.feature_callbacks = [add_random_noise(*args.add_random_noise)] if args.custom_feature_cb: func = args.custom_feature_cb.split(' ') for f_name in func: env.feature_callbacks += [globals()[f_name]] fit_fn = partial(env.fit, args.batch, args.epochs, dataset, restart=args.retrain, learning_rate=args.lr, learning_rate_schedule=lr_decay( 'stair', args.lr, decay_step=args.lr_decay_step, decay_rate=args.lr_decay)) if model.channel > 1: fit_fn(convert_to='RGB') test_format = 'RGB' else: fit_fn(convery_to='GRAY') # use callback to generate colored images from grayscale ones # all models inputs is gray image1 however test_format = 'YUV' env.feature_callbacks += [to_gray()] env.label_callbacks = [to_gray()] if args.output_color == 'RGB': env.output_callbacks += [to_rgb()] if args.test: test_set = load_datasets(args.dataconfig)[args.test.upper()] test_set.setattr(patch_size=args.patch_size, strides=args.strides, depth=args.depth) else: test_set = dataset env.output_callbacks += [ save_image(f'{save_root}/test', args.output_index) ] env.test(test_set, convert_to=test_format) # load image1 with 3 channels if args.predict: pth = Path(args.predict) if not pth.exists(): raise ValueError('[Error] File path does not exist') images = pth.rglob('*') env.fi, fi_old = 0, env.fi # upscale directly env.output_callbacks[-1] = save_image(f'{save_root}/output', args.output_index) env.predict(images, convert_to=test_format) env.fi = fi_old if args.export_pb: model = get_model(args.name)(scale=args.scale, rgb_input=True) with Environment(model, f'{save_root}/save', f'{save_root}/log', feature_index=model.feature_index, label_index=model.label_index) as env: env.export(args.export_pb)
from VSR.Framework.Callbacks import * from VSR.Models.DnCnn import DnCNN from pathlib import Path def add_noise(low=0, high=55): import numpy as np def feature_callback(feature): sigma = np.random.randint(low, high) # sigma = 25 noise = np.random.normal(0, sigma, size=feature.shape) return feature + noise return feature_callback if __name__ == '__main__': model = DnCNN(layers=20, rgb_input=False).compile() dataset = load_datasets('../Data/datasets.json')['BSD'] dataset.setattr(patch_size=40, strides=40, depth=1) env = Environment(model, f'{model.name}/save', f'{model.name}/log', feature_index=0, label_index=0) env.feature_callbacks = [add_noise(0, 55)] env.fit(128, 100, dataset, restart=False) env.feature_callbacks = [] env.output_callbacks = [lambda output, **kwargs: output[0], reduce_residual()] env.output_callbacks += [save_image(model.name)] env.predict(list(Path('srcnn').glob('*.png'))) # env.test(dataset) env.export(f'{model.name}')
pre-process images in dataset - cal mean and covariance - ~random crop~ """ import numpy as np import tensorflow as tf from VSR.DataLoader.Dataset import load_datasets from VSR.DataLoader.Loader import QuickLoader, BasicLoader from VSR.Util import FID from VSR.Util.ImageProcess import imresize, array_to_img try: DATASETS = load_datasets('./Data/datasets.json') except FileNotFoundError: DATASETS = load_datasets('../Data/datasets.json') FLAGS = tf.flags.FLAGS SAVE = {} def main(*args): for name in FLAGS.dataset: name = name.upper() d = DATASETS.get(name) if not d: tf.logging.error('Could not find ' + name) return
from VSR.Framework.Envrionment import Environment from VSR.DataLoader.Dataset import load_datasets from VSR.Models.Vespcn import VESPCN if __name__ == '__main__': model = VESPCN(scale=3, depth=3) model.compile() datasets = load_datasets('../Data/datasets.json') env = Environment(model, './vespcn/save', './vespcn/log') env.fit(64, 100, datasets['MCL-V'], patch_size=48, strides=48, depth=3) env.test(datasets['MCL-V'], depth=3)
from VSR.Models.SrGan import SRGAN from VSR.DataLoader.Dataset import load_datasets from VSR.Framework.Envrionment import Environment from VSR.Framework.Callbacks import * if __name__ == '__main__': dataset = load_datasets('../Data/datasets.json')['DIV2K'] dataset.setattr(patch_size=96, strides=96, random=True, max_patches=16 * 100) model = SRGAN(scale=4, glayers=16, dlayers=8, vgg_layer=[2, 2], init=True) with Environment(model, f'../Results/{model.name}/save', f'../Results/{model.name}/log') as env: env.fit(16, 100, dataset, learning_rate=1e-4) env.output_callbacks = [ save_image(f'../Results/{model.name}/test_init') ] env.test(dataset) model = SRGAN(scale=4, glayers=16, dlayers=8, vgg_layer=[2, 2], init=False) with Environment(model, f'../Results/{model.name}/save', f'../Results/{model.name}/log') as env: env.fit(32, 200, dataset, learning_rate_schedule=lr_decay('stair', 0.001, decay_step=1000, decay_rate=0.96)) env.feature_callbacks = [to_gray()]
def main(*args, **kwargs): args = argparse.ArgumentParser() args.add_argument('name', type=str, choices=list_supported_models(), help='the model name can be found in model_alias.py') # basic options args.add_argument('--scale', type=int, default=4, help='scale factor, default 4') args.add_argument('--channel', type=int, default=3, help='image channels, default 3') args.add_argument('--batch', type=int, default=16, help='training batch size, default 16') args.add_argument('--epochs', type=int, default=200, help='training epochs, default 200') args.add_argument('--steps_per_epoch', type=int, default=200, help='training steps each epoch, default 200') args.add_argument('--retrain', type=bool, default=False, help='retrain the model from scratch') # dataset options args.add_argument('--dataconfig', type=str, default='../Data/datasets.json', help='the path to dataset config json file') args.add_argument( '--dataset', type=str, default='91-IMAGE', help= 'specified dataset name(as described in config file, default 91-image') args.add_argument( '--patch_size', type=int, default=48, help= 'patch size of cropped training and validating sub-images, default 48') args.add_argument('--depth', type=int, default=1, help='image depth used for video sources, default 1') args.add_argument( '--parallel', type=int, default=1, help='number of cores used to load training sets in parallel') args.add_argument('--memory', type=str, default=None, help='limit the memory usage. i.e. 4GB, 100MB') # learning options args.add_argument('--lr', type=float, default=1e-4, help='initial learning rate, default 1e-4') args.add_argument('--lr_decay', type=float, default=1, help='learning rate decay, default 1') args.add_argument('--lr_decay_step', type=int, default=1000, help='learning rate decay step') # output options args.add_argument( '--test', type=str, default=None, help='specify a dataset used to test, or use --dataset values if None') args.add_argument('--predict', type=str, default=None, help='evaluate model on given files') args.add_argument( '--savedir', type=str, default='../Results', help='directory to save model checkpoints, default ../Results') args.add_argument('--output_color', type=str, default='RGB', choices=('RGB', 'L', 'GRAY', 'Y'), help='output color mode, default RGB') args.add_argument('--output_index', type=int, default=-1, help='access index of model outputs to save, default -1') args.add_argument( '--export_pb', type=str, default=None, help= 'if not None, specify the path that export trained model into pb format' ) args.add_argument( '--comment', type=str, default=None, help='add a suffix to output dir to distinguish each experiments') # callbacks args.add_argument( '--add_noise', type=float, default=None, help='if not None, add noise with given stddev to input features') args.add_argument( '--add_random_noise', type=list, default=None, help= 'if not None, add random noise with given stddev bound [low, high, step=1]' ) args.add_argument('--custom_feature_cb', type=str, default=None, help='customized callbacks, defined in `custom_api.py`') args = args.parse_args() if Path(f'parameters/{args.name}.json').exists(): model_args = json.load(open(f'parameters/{args.name}.json', mode='r')) else: print( f'[warning] model parameter file not found, use default parameters' ) model_args = dict() model = get_model(args.name)(scale=args.scale, channel=args.channel, **model_args) dataset = load_datasets(args.dataconfig)[args.dataset.upper()] dataset.setattr(patch_size=args.patch_size, depth=args.depth) save_root = f'{args.savedir}/{model.name}_sc{args.scale}_c{args.channel}' if args.comment: save_root += '_' + args.comment with Environment(model, f'{save_root}/save', f'{save_root}/log', feature_index=model.feature_index, label_index=model.label_index) as env: if args.add_noise: env.feature_callbacks = [add_noise(args.add_noise)] if args.add_random_noise: env.feature_callbacks = [add_random_noise(*args.add_random_noise)] if args.custom_feature_cb: func = args.custom_feature_cb.split(' ') for f_name in func: env.feature_callbacks += [globals()[f_name]] fit_fn = partial(env.fit, args.batch, args.epochs, args.steps_per_epoch, dataset, augmentation=True, restart=args.retrain, learning_rate=args.lr, learning_rate_schedule=lr_decay( 'stair', args.lr, decay_step=args.lr_decay_step, decay_rate=args.lr_decay), parallel=args.parallel, memory_usage=args.memory) if model.channel > 1: fit_fn(convert_to='RGB') test_format = 'RGB' else: fit_fn(convery_to='GRAY') # use callback to generate colored images from grayscale ones # all models inputs is gray image however test_format = 'YUV' env.feature_callbacks += [to_gray()] env.label_callbacks = [to_gray()] if args.output_color == 'RGB': env.output_callbacks += [to_rgb()] if args.test: test_set = load_datasets(args.dataconfig)[args.test.upper()] test_set.setattr(patch_size=args.patch_size, depth=args.depth) else: test_set = dataset env.output_callbacks += [ save_image(f'{save_root}/test', args.output_index) ] env.test(test_set, convert_to=test_format) # load image with 3 channels if args.predict: pth = Path(args.predict) if not pth.exists(): raise ValueError('[Error] File path does not exist') if pth.is_dir(): images = list(pth.glob('*')) if not images: images = pth.iterdir() elif pth.is_file(): images = pth env.fi, fi_old = 0, env.fi # upscale directly env.output_callbacks[-1] = save_image(f'{save_root}/output', args.output_index) env.predict(images, convert_to=test_format, depth=args.depth) env.fi = fi_old elif test_set.pred: env.fi, fi_old = 0, env.fi # upscale directly env.output_callbacks[-1] = save_image(f'{save_root}/output', args.output_index) env.predict(test_set.pred, convert_to=test_format, depth=args.depth) env.fi = fi_old if args.export_pb: env.export(args.export_pb)
Calculate metrics for outputs and labels: - PSNR - SSIM """ from pathlib import Path import numpy as np import tensorflow as tf from VSR.DataLoader.Dataset import load_datasets, Dataset from VSR.DataLoader.Loader import BasicLoader from VSR.Util.ImageProcess import rgb_to_yuv from VSR.Util.Config import Config try: DATASETS = load_datasets('./Data/datasets.yaml') except FileNotFoundError: DATASETS = load_datasets('../Data/datasets.yaml') tf.flags.DEFINE_string("input_dir", None, "images to test") tf.flags.DEFINE_string("dataset", None, "dataset to compare") tf.flags.DEFINE_bool("no_ssim", False, "disable ssim metric") tf.flags.DEFINE_bool("no_psnr", False, "disable psnr metric") tf.flags.DEFINE_bool("l_only", False, "compute luminance only") tf.flags.DEFINE_string( "l_standard", "matlab", "yuv convertion standard, either 'bt601', 'bt709' or 'matlab'") tf.flags.DEFINE_integer("shave", 0, "shave border pixels") tf.flags.DEFINE_integer("clip", -1, "depth of a clip, -1 includes all frames") tf.flags.DEFINE_integer("offset", 0, "skip n files in the dataset")
def main(): flags, args = parser.parse_known_args() opt = Config() for pair in flags._get_kwargs(): opt.setdefault(*pair) data_config_file = Path(flags.data_config) if not data_config_file.exists(): raise RuntimeError("dataset config file doesn't exist!") for _ext in ('json', 'yaml', 'yml'): # for compat # apply a 2-stage (or master-slave) configuration, master can be # override by slave model_config_root = Path('Parameters/root.{}'.format(_ext)) if opt.p: model_config_file = Path(opt.p) else: model_config_file = Path('Parameters/{}.{}'.format(opt.model, _ext)) if model_config_root.exists(): opt.update(Config(str(model_config_root))) if model_config_file.exists(): opt.update(Config(str(model_config_file))) model_params = opt.get(opt.model, {}) suppress_opt_by_args(model_params, *args) opt.update(model_params) model = get_model(flags.model)(**model_params) if flags.cuda: model.cuda() root = f'{flags.save_dir}/{flags.model}' if flags.comment: root += '_' + flags.comment verbosity = logging.DEBUG if flags.verbose else logging.INFO trainer = model.trainer datasets = load_datasets(data_config_file) try: test_datas = [datasets[t.upper()] for t in flags.test] run_benchmark = True except KeyError: test_datas = [] for pattern in flags.test: test_data = Dataset(test=_glob_absolute_pattern(pattern), mode='pil-image1', modcrop=False) father = Path(flags.test) while not father.is_dir(): if father.parent == father: break father = father.parent test_data.name = father.stem test_datas.append(test_data) run_benchmark = False if opt.verbose: dump(opt) for test_data in test_datas: loader_config = Config(convert_to='rgb', feature_callbacks=[], label_callbacks=[], output_callbacks=[], **opt) loader_config.batch = 1 loader_config.subdir = test_data.name loader_config.output_callbacks += [ save_image(root, flags.output_index, flags.auto_rename)] if opt.channel == 1: loader_config.convert_to = 'gray' with trainer(model, root, verbosity, flags.pth) as t: if flags.seed is not None: t.set_seed(flags.seed) loader = QuickLoader(test_data, 'test', loader_config, n_threads=flags.thread) loader_config.epoch = flags.epoch if run_benchmark: t.benchmark(loader, loader_config) else: t.infer(loader, loader_config)
from VSR.Util.Utility import Vgg from VSR.DataLoader.Dataset import load_datasets from VSR.DataLoader.Loader import BatchLoader import tensorflow as tf import numpy as np data = load_datasets('../Data/datasets.json')['91-IMAGE'] loader = BatchLoader(1, data, 'test', convert_to_gray=True, crop=False, scale=1) m = Vgg(input_shape=[None, None, 3], type='vgg19') with tf.Session() as sess: # m = Vgg(input_shape=[None, None, 3], type='vgg19') for hr, lr in loader: y = m(hr, [2, 3], [2, 3]) break tf.reset_default_graph() m = Vgg(input_shape=[None, None, 3], type='vgg19') with tf.Session() as sess: # m = Vgg(input_shape=[None, None, 3], type='vgg19') for hr, lr in loader: y = m(hr, [2, 3], [2, 3]) break # Error for not reusing VGG variables
import os if not os.getcwd().endswith('UTest'): os.chdir('UTest') from VSR.DataLoader.Dataset import _glob_absolute_pattern, load_datasets DATASETS = load_datasets('./data/fake_datasets.yml') def test_glob_absolute_pattern(): URL = './data/set5_x2' node = _glob_absolute_pattern(URL) assert len(node) == 5 assert node[0].match('img_001_SRF_2_LR.png') assert node[1].match('img_002_SRF_2_LR.png') assert node[2].match('img_003_SRF_2_LR.png') assert node[3].match('img_004_SRF_2_LR.png') assert node[4].match('img_005_SRF_2_LR.png') URL = './data/flying_chair/**/*.flo' node = _glob_absolute_pattern(URL) assert len(node) == 1 assert node[0].match('0000.flo') URL = './data/**/*.png' node = _glob_absolute_pattern(URL) assert len(node) == 10 def test_existence(): _K = DATASETS.keys()
def main(*args, **kwargs): args = argparse.ArgumentParser() args.add_argument('name', type=str, choices=list_supported_models(), help='the model name can be found in model_alias.py') args.add_argument('--scale', type=int, default=3, help='scale factor') args.add_argument('--dataconfig', type=str, default='../Data/datasets.json', help='the path to dataset config json file') args.add_argument('--dataset', type=str, default='91-IMAGE', help='specified dataset name(as described in config file') args.add_argument('--batch', type=int, default=64, help='training batch size') args.add_argument('--epochs', type=int, default=200, help='training epochs') args.add_argument('--patch_size', type=int, default=48, help='patch size of cropped training and validating sub-images') args.add_argument('--strides', type=int, default=48, help='crop stride if random_patches is set 0') args.add_argument('--depth', type=int, default=1, help='image1 depth used for video sources') args.add_argument('--random_patches', type=int, default=0, help='if set more than 0, use random crop to generate `random_patches` sub-image1 batches') args.add_argument('--retrain', type=int, default=0, help='retrain the model from scratch') args.add_argument('--lr', type=float, default=1e-4, help='initial learning rate') args.add_argument('--lr_decay', type=float, default=1, help='learning rate decay') args.add_argument('--lr_decay_step', type=int, default=1000, help='learning rate decay step') args.add_argument('--add_noise', type=float, default=None, help='if not None, add noise with given stddev to input features') args.add_argument('--add_random_noise', type=list, default=None, help='if not None, add random noise with given stddev bound [low, high, step=1]') args.add_argument('--test', type=bool, default=True, help='test model and save tested images') args.add_argument('--savedir', type=str, default='../Results', help='directory to save model checkpoints') args.add_argument('--output_color', type=str, default='RGB', choices=('RGB', 'L', 'GRAY', 'Y'), help='output color mode') args.add_argument('--export_pb', type=str, default=None, help='if not None, specify the path that export trained model into pb format') args = args.parse_args() model_args = json.load(open(f'parameters/{args.name}.json', mode='r')) model = get_model(args.name)(scale=args.scale, **model_args) dataset = load_datasets(args.dataconfig)[args.dataset.upper()] dataset.setattr(patch_size=args.patch_size, strides=args.strides, depth=args.depth) if args.random_patches: dataset.setattr(random=True, max_patches=args.batch * args.random_patches) with Environment(model, f'{args.savedir}/{model.name}/save', f'{args.savedir}/{model.name}/log', feature_index=model.feature_index, label_index=model.label_index) as env: if args.add_noise: env.feature_callbacks = [add_noise(args.add_noise)] if args.add_random_noise: env.feature_callbacks = [add_random_noise(*args.add_random_noise)] env.fit(args.batch, args.epochs, dataset, restart=args.retrain, learning_rate=args.lr, learning_rate_schedule=lr_decay('stair', args.lr, decay_step=args.lr_decay_step, decay_rate=args.lr_decay)) if args.test: # use callback to generate colored images from grayscale ones # all models inputs is gray image1 however env.feature_callbacks += [to_gray()] env.label_callbacks = [to_gray()] if args.output_color == 'RGB': env.output_callbacks += [to_rgb()] env.output_callbacks += [save_image(f'../Results/{model.name}/test')] env.test(dataset, convert_to_gray=False) # load image1 with 3 channels if args.export_pb: model = get_model(args.name)(scale=args.scale, rgb_input=True) with Environment(model, f'../Results/{model.name}/save', f'../Results/{model.name}/log') as env: env.export(args.export_pb)