def get_config(): parser = argparse.ArgumentParser() parser.add_argument('--model', required=True) parser.add_argument('--mode',type=str,choices=['train','predict'],required=True) parser.add_argument('--config-path',type=str,default='config/cfg.json') parser.add_argument('--gpu-ids',type=int,nargs='+',default=0,dest='gpu_ids') parser.add_argument('--state',type=int,default=1,dest='state') parser.add_argument('-e', '--epochs', type=int, default=5, help='Number of epochs', dest='epochs') parser.add_argument('-b', '--batch-size', type=int, default=2, help='Batch size', dest='batch_size') parser.add_argument('-l', '--learning-rate', type=float, default=0.1, help='Learning rate', dest='lr') parser.add_argument('-p', '--port', type=int,default=10001, help='Visualization port', dest='port') parser.add_argument('-w', '--worker-num', type=int, default=8, help='Dataloader worker number', dest='num_workers') parser.add_argument('-c', '--class-num', type=int, default=2, help='class number', dest='class_num') parser.add_argument('-v', '--valid-percent', type=get_range_limited_float_type(0,100), default=10.0, help='Percent of the data that is used as validation (0-100)', dest='valid_percent') args = parser.parse_args() assert os.path.exists(args.config_path),'config json not exists' with open(args.config_path,'r') as f: config = json.load(f) for arg in vars(args): config[arg]=getattr(args,arg) if isinstance(config['gpu_ids'],int): config['gpu_ids'] = [config['gpu_ids']] config['gpu_ids'] = list(set(config['gpu_ids'])) config['device'] = 'cuda' if torch.cuda.is_available() else 'cpu' if config['device'] == 'cuda': gpu_num = torch.cuda.device_count() assert len(config['gpu_ids'])!=0,'unexpected gpu number' for gpu_id in config['gpu_ids']: assert gpu_id>=0 and gpu_id<gpu_num,'invalid gpu id input' config['input_transform'] = Compose([ # CenterCrop(256), ToTensor(), # The mean and std result from data statistics of ImageNet dataset, you should fill corresponding mean and std here. # Normalize([.485, .456, .406], [.229, .224, .225]), ]) config['target_transform'] = Compose([ # CenterCrop(256), ToLabel(), # Relabel(255, 21), ]) return config
tt = time.time() all_E = Infer_MO(all_F, all_M, num_frames, num_objects, scales=[0.5, 0.75, 1.0]) print('{} | num_objects: {}, FPS: {}'.format( seq_name, num_objects, num_frames / (time.time() - tt))) # Save results for quantitative eval ###################### if MO: folder = 'results/MO' else: folder = 'results/SO' test_path = os.path.join(folder, seq_name) if not os.path.exists(test_path): os.makedirs(test_path) for f in range(num_frames): E = all_E[0, :, f].numpy() # make hard label E = ToLabel(E) (lh, uh), (lw, uw) = info['pad'] E = E[lh[0]:-uh[0], lw[0]:-uw[0]] img_E = Image.fromarray(E) img_E.putpalette(palette) img_E.save(os.path.join(test_path, '{:05d}.png'.format(f)))
os.makedirs(os.path.join(args.name + '_results', 'Colorization')) if not os.path.exists(os.path.join(args.name + '_results', 'log')): os.makedirs(os.path.join(args.name + '_results', 'log')) # Loss plot logger = Logger(os.path.join(args.name + '_results', 'log')) image_transform = transforms.Compose([ transforms.CenterCrop(args.input_size), transforms.ToTensor(), transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)) ]) target_transform = transforms.Compose([ transforms.CenterCrop(args.input_size), ToLabel(), Relabel(255, 21), ]) train_loader_src = DataLoader(VOC12(args.src_data, image_transform, target_transform), num_workers=4, batch_size=args.batch_size, shuffle=True) train_loader_tgt = DataLoader(VOC12(args.tgt_data, image_transform, target_transform), num_workers=4, batch_size=args.batch_size, shuffle=True) test_loader_src = data_load('./data/test_data/', 'test',