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'
def initialize_model(self, path=None, train=True):
self.G = RFRNet()
self.optm_G = optim.Adam(self.G.parameters(), lr=2e-4)
self.print_networks(False)
if train:
self.lossNet = VGG16FeatureExtractor()
try:
start_iter = load_ckpt(path, [('generator', self.G)],
[('optimizer_G', self.optm_G)])
if train:
self.optm_G = optim.Adam(self.G.parameters(), lr=2e-4)
print('Model Initialized, iter: ', start_iter)
self.iter = start_iter
except:
print('No trained model, from start')
self.iter = 0
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')
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)],
[('optimizer', optimizer)])
for param_group in optimizer.param_groups:
param_group['lr'] = lr
print('Starting from iter ', start_iter)
for i in tqdm(range(start_iter, args.max_iter)):
model.train()
image, mask, gt = [x.to(device) for x in next(iterator_train)]
if args.mask_root is not None:
guide = image[:, 3:4, :, :]
image = image[:, 0:3, :, :]
output, _ = model(image, mask, guide)
else:
output, _ = model(image, mask)
loss_dict = criterion(image, mask, output, gt)
type=str,
default='../posewarp/data/hunter_256.jpeg')
parser.add_argument('--maskPath',
type=str,
default='../posewarp/data/hunter_256_IUV.npy')
#parser.add_argument('--maskPath', type=str, default='/home/cloudlet/vlr_824/project/backgroundInpainting/mask/003328.jpg')
parser.add_argument('--modelPath', type=str, default='./pretrainedModel.pth')
parser.add_argument('--image_size', type=int, default=256)
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 = customData(args.imPath, args.maskPath, img_transform,
mask_transform, 'val')
model = PConvUNet()
load_ckpt(args.modelPath, [('model', model)])
model.to(device)
model.eval()
evaluateCustom(model, dataset_val, device, 'result.jpg')
