def predict(image, mask, root_path, AI_directory_path, model_type="life"): device = torch.device('cuda') size = (256, 256) img_transform = transforms.Compose([ transforms.Resize(size=size), transforms.ToTensor(), transforms.Normalize(mean=opt.MEAN, std=opt.STD) ]) mask_transform = transforms.Compose( [transforms.Resize(size=size), transforms.ToTensor()]) dataset_val = Places2(root_path, image, mask, img_transform, mask_transform) model = PConvUNet().to(device) load_ckpt(AI_directory_path, [('model', model)]) model.eval() evaluate(model, dataset_val, device, image.split('.')[0] + 'result.jpg') return image.split('.')[0] + 'result.jpg'
parser = argparse.ArgumentParser() # training options parser.add_argument('--root', type=str, default='./data') parser.add_argument('--snapshot', type=str, default='') parser.add_argument('--image_size', type=int, default=256) parser.add_argument('--mask_root', type=str, default='./mask') args = parser.parse_args() device = torch.device('cuda') size = (args.image_size, args.image_size) img_transform = transforms.Compose([ transforms.Resize(size=size), transforms.ToTensor(), transforms.Normalize(mean=opt.MEAN, std=opt.STD) ]) mask_transform = transforms.Compose( [transforms.Resize(size=size), transforms.ToTensor()]) dataset_val = Places2(args.root, args.mask_root, img_transform, mask_transform, 'val') model = PConvUNet().to(device) load_ckpt(args.snapshot, [('model', model)]) model.eval() evaluate(model, dataset_val, device, 'result.jpg')
auto_resize=not random_masks) for train_flag in [True, False] ] print('train size:', len(dataset_train)) print('val size:', len(dataset_val)) print('dataset shapes:') for tensor in dataset_train[0]: print(tensor.shape) iterator_train = iter( data.DataLoader(dataset_train, batch_size=args.batch_size, sampler=InfiniteSampler(len(dataset_train)), num_workers=args.n_threads)) print(len(dataset_train)) model = PConvUNet(input_guides=1 if use_depth else 0).to(device) if args.finetune: lr = args.lr_finetune model.freeze_enc_bn = True else: lr = args.lr start_iter = 0 optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr) criterion = InpaintingLoss(VGG16FeatureExtractor()).to(device) if args.resume: start_iter = load_ckpt(args.resume, [('model', model)],
#total_white = total_white /(len(dataset_train) * 64 *64) iterator_train = iter( data.DataLoader(dataset_train, batch_size=args.batch_size, sampler=InfiniteSampler(len(dataset_train)), num_workers=args.n_threads)) iterator_val = iter( data.DataLoader(dataset_val, batch_size=args.batch_size, sampler=InfiniteSampler(len(dataset_val)), num_workers=args.n_threads)) print(len(dataset_train)) model = PConvUNet(layer_size=3).to(device) if args.finetune: lr = args.lr_finetune model.freeze_enc_bn = True else: lr = args.lr start_iter = 0 optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr) criterion = InpaintingLoss(VGG16FeatureExtractor()).to(device) if args.resume: start_iter = load_ckpt(args.resume, [('model', model)],
[transforms.Resize(size=size), transforms.ToTensor()]) dataset_train = Places2(args.root, args.mask_root, img_tf, mask_tf, 'train') dataset_val = Places2(args.root, args.mask_root, img_tf, mask_tf, 'val') # dataset_train = DIV2K(args.root, args.mask_root, img_tf, mask_tf, 'train') # dataset_val = DIV2K(args.root, args.mask_root, img_tf, mask_tf, 'val') iterator_train = iter( data.DataLoader(dataset_train, batch_size=args.batch_size, sampler=InfiniteSampler(len(dataset_train)), num_workers=args.n_threads)) print(len(dataset_train)) model = PConvUNet() if torch.cuda.device_count() > 1: print("Use", torch.cuda.device_count(), "GPUs!") model = nn.DataParallel(model) model = model.to(device) if args.finetune: lr = args.lr_finetune model.freeze_enc_bn = True else: lr = args.lr start_iter = 0 optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()),
print("set dataset") dataset_train = Places2(args.root, args.mask_root, img_tf, mask_tf, 'train') dataset_val = Places2(args.root, args.mask_root, img_tf, mask_tf, 'val') print("prepare data iterator") iterator_train = iter( data.DataLoader(dataset_train, batch_size=args.batch_size, sampler=InfiniteSampler(len(dataset_train)), num_workers=args.n_threads)) print(len(dataset_train)) print("set model") #model = PConvUNet().to(device) model = PConvUNet().to("cuda") model = torch.nn.DataParallel( model, device_ids=None) #device_ids=None means using all devices if args.finetune: lr = args.lr_finetune model.freeze_enc_bn = True else: lr = args.lr start_iter = 0 optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr) criterion = InpaintingLoss(VGG16FeatureExtractor()).to(device)