def resume(self, checkpoint_dir, iteration=0, test=False):
# Load generators
last_model_name = get_model_list(checkpoint_dir,
"gen",
iteration=iteration)
self.netG.load_state_dict(torch.load(last_model_name))
iteration = int(last_model_name[-11:-3])
if not test:
# Load discriminators
last_model_name = get_model_list(checkpoint_dir,
"dis",
iteration=iteration)
state_dict = torch.load(last_model_name)
self.localD.load_state_dict(state_dict['localD'])
self.globalD.load_state_dict(state_dict['globalD'])
# Load optimizers
state_dict = torch.load(
os.path.join(checkpoint_dir, 'optimizer.pt'))
self.optimizer_d.load_state_dict(state_dict['dis'])
self.optimizer_g.load_state_dict(state_dict['gen'])
print("Resume from {} at iteration {}".format(checkpoint_dir,
iteration))
logger.info("Resume from {} at iteration {}".format(
checkpoint_dir, iteration))
return iteration
def resume(self, checkpoint_dir, iteration=0, test=False):
# Load generators
last_model_name = get_model_list(checkpoint_dir,
"gen",
iteration=iteration)
mp_loc = 'cpu' if not self.config['cuda'] else None
self.netG.load_state_dict(
torch.load(last_model_name, map_location=mp_loc))
iteration = int(last_model_name[-11:-3])
print("last model name is {}".format(last_model_name))
if not test:
# Load discriminators
last_model_name = get_model_list(checkpoint_dir,
"dis",
iteration=iteration)
# this function returns the path including filename of the latest checkpoint
state_dict = torch.load(last_model_name, map_location=mp_loc)
self.localD.load_state_dict(state_dict['localD'])
self.globalD.load_state_dict(state_dict['globalD'])
# Load optimizers
state_dict = torch.load(os.path.join(checkpoint_dir,
'optimizer.pt'),
map_location=mp_loc)
self.optimizer_d.load_state_dict(state_dict['dis'])
self.optimizer_g.load_state_dict(state_dict['gen'])
print("Resume from {} at iteration {}".format(checkpoint_dir,
iteration))
logger.info("Resume from {} at iteration {}".format(
checkpoint_dir, iteration))
return iteration
def loadGenerator(args):
config = get_config(args.g_config)
# CUDA configuration
cuda = config['cuda']
device_ids = config['gpu_ids']
if cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
str(i) for i in device_ids)
device_ids = list(range(len(device_ids)))
config['gpu_ids'] = device_ids
cudnn.benchmark = True
# Set checkpoint path
if not args.checkpoint_path:
checkpoint_path = os.path.join(
'checkpoints', config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
else:
checkpoint_path = args.checkpoint_path
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids).cuda()
# Resume weight
last_model_name = get_model_list(checkpoint_path,
"gen",
iteration=args.iter)
model_iteration = int(last_model_name[-11:-3])
netG.load_state_dict(torch.load(last_model_name))
print("Configuration: {}".format(config))
print("Resume from {} at iteration {}".format(checkpoint_path,
model_iteration))
if cuda:
netG = nn.parallel.DataParallel(netG, device_ids=device_ids)
return netG
def resume(self, checkpoint_dir, iteration=0, test=False):
# Load generators
last_model_name = get_model_list(checkpoint_dir,
"parser",
iteration=iteration)
self.netParser.load_state_dict(torch.load(last_model_name))
iteration = int(last_model_name[-11:-3])
'''
if not test:
# Load optimizers
#last_opt_name = torch.load(os.path.join(checkpoint_dir, 'parser_optimizer.pt'))
opt_state_dict = torch.load(os.path.join(checkpoint_dir, 'parser_optimizer.pt'))
self.optimizerSGD.load_state_dict(opt_state_dict)
'''
print("Resume from {} at iteration {}".format(checkpoint_dir,
iteration))
logger.info("Resume from {} at iteration {}".format(
checkpoint_dir, iteration))
return iteration
def generate(img, img_mask_path, model_path):
with torch.no_grad(): # enter no grad context
if img_mask_path and is_image_file(img_mask_path):
# Test a single masked image with a given mask
x = Image.fromarray(img)
mask = default_loader(img_mask_path)
x = transforms.Resize(config['image_shape'][:-1])(x)
x = transforms.CenterCrop(config['image_shape'][:-1])(x)
mask = transforms.Resize(config['image_shape'][:-1])(mask)
mask = transforms.CenterCrop(config['image_shape'][:-1])(mask)
x = transforms.ToTensor()(x)
mask = transforms.ToTensor()(mask)[0].unsqueeze(dim=0)
x = normalize(x)
x = x * (1. - mask)
x = x.unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
elif img_mask_path:
raise TypeError("{} is not an image file.".format(img_mask_path))
else:
# Test a single ground-truth image with a random mask
#ground_truth = default_loader(img_path)
ground_truth = img
ground_truth = transforms.Resize(config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.CenterCrop(config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.ToTensor()(ground_truth)
ground_truth = normalize(ground_truth)
ground_truth = ground_truth.unsqueeze(dim=0)
bboxes = random_bbox(config, batch_size=ground_truth.size(0))
x, mask = mask_image(ground_truth, bboxes, config)
# Set checkpoint path
if not model_path:
checkpoint_path = os.path.join('checkpoints',
config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
else:
checkpoint_path = model_path
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids)
# Resume weight
last_model_name = get_model_list(checkpoint_path, "gen", iteration=0)
if cuda:
netG.load_state_dict(torch.load(last_model_name))
else:
netG.load_state_dict(torch.load(last_model_name, map_location='cpu'))
model_iteration = int(last_model_name[-11:-3])
print("Resume from {} at iteration {}".format(checkpoint_path, model_iteration))
if cuda:
netG = nn.parallel.DataParallel(netG, device_ids=device_ids)
x = x.cuda()
mask = mask.cuda()
# Inference
x1, x2, offset_flow = netG(x, mask)
inpainted_result = x2 * mask + x * (1. - mask)
inpainted_result = from_torch_img_to_numpy(inpainted_result, 'output.png', padding=0, normalize=True)
return inpainted_result
def main():
args = parser.parse_args()
config = get_config(args.config)
# CUDA configuration
cuda = config['cuda']
device_ids = config['gpu_ids']
if cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
str(i) for i in device_ids)
device_ids = list(range(len(device_ids)))
config['gpu_ids'] = device_ids
cudnn.benchmark = True
print("Arguments: {}".format(args))
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
print("Configuration: {}".format(config))
try: # for unexpected error logging
with torch.no_grad(): # enter no grad context
if is_image_file(args.image):
if args.mask and is_image_file(args.mask):
# Test a single masked image with a given mask
x = default_loader(args.image)
mask = default_loader(args.mask)
x = transforms.Resize(config['image_shape'][:-1])(x)
x = transforms.CenterCrop(config['image_shape'][:-1])(x)
mask = transforms.Resize(config['image_shape'][:-1])(mask)
mask = transforms.CenterCrop(
config['image_shape'][:-1])(mask)
x = transforms.ToTensor()(x)
mask = transforms.ToTensor()(mask)[0].unsqueeze(dim=0)
x = normalize(x)
x = x * (1. - mask)
x = x.unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
elif args.mask:
raise TypeError(
"{} is not an image file.".format(args.mask))
else:
# Test a single ground-truth image with a random mask
ground_truth = default_loader(args.image)
ground_truth = transforms.Resize(
config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.CenterCrop(
config['image_shape'][:-1])(ground_truth)
ground_truth = transforms.ToTensor()(ground_truth)
ground_truth = normalize(ground_truth)
ground_truth = ground_truth.unsqueeze(dim=0)
bboxes = random_bbox(
config, batch_size=ground_truth.size(0))
x, mask = mask_image(ground_truth, bboxes, config)
# Set checkpoint path
if not args.checkpoint_path:
checkpoint_path = os.path.join('checkpoints',
config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
else:
checkpoint_path = args.checkpoint_path
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids)
# Resume weight
last_model_name = get_model_list(
checkpoint_path, "gen", iteration=args.iter)
netG.load_state_dict(torch.load(last_model_name))
model_iteration = int(last_model_name[-11:-3])
print("Resume from {} at iteration {}".format(
checkpoint_path, model_iteration))
if cuda:
netG = nn.parallel.DataParallel(
netG, device_ids=device_ids)
x = x.cuda()
mask = mask.cuda()
# Inference
x1, x2, offset_flow = netG(x, mask)
inpainted_result = x2 * mask + x * (1. - mask)
vutils.save_image(inpainted_result, args.output,
padding=0, normalize=True)
print("Saved the inpainted result to {}".format(args.output))
if args.flow:
vutils.save_image(offset_flow, args.flow,
padding=0, normalize=True)
print("Saved offset flow to {}".format(args.flow))
else:
raise TypeError("{} is not an image file.".format)
# exit no grad context
except Exception as e: # for unexpected error logging
print("Error: {}".format(e))
raise e
def main():
args = parser.parse_args()
config = get_config(args.config)
# CUDA configuration
cuda = config['cuda']
device_ids = config['gpu_ids']
if cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
str(i) for i in device_ids)
device_ids = list(range(len(device_ids)))
config['gpu_ids'] = device_ids
cudnn.benchmark = True
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
chunker = ImageChunker(config['image_shape'][0], config['image_shape'][1],
args.overlap)
try: # for unexpected error logging
with torch.no_grad(): # enter no grad context
if is_image_file(args.image):
print("Loading image...")
imgs, masks = [], []
img_ori = default_loader(args.image)
img_w, img_h = img_ori.size
# Load mask txt file
fname = args.image.replace('.jpg', '.txt')
bboxes, _ = load_bbox_txt(fname, img_w, img_h)
mask_ori = create_mask(bboxes, img_w, img_h)
chunked_images = chunker.dimension_preprocess(
np.array(deepcopy(img_ori)))
chunked_masks = chunker.dimension_preprocess(
np.array(deepcopy(mask_ori)))
for (x, msk) in zip(chunked_images, chunked_masks):
x = transforms.ToTensor()(x)
mask = transforms.ToTensor()(msk)[0].unsqueeze(dim=0)
# x = normalize(x)
x = x * (1. - mask)
x = x.unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
imgs.append(x)
masks.append(mask)
# Set checkpoint path
if not args.checkpoint_path:
checkpoint_path = os.path.join(
'checkpoints', config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
else:
checkpoint_path = args.checkpoint_path
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids)
# Resume weight
last_model_name = get_model_list(checkpoint_path,
"gen",
iteration=args.iter)
netG.load_state_dict(torch.load(last_model_name))
model_iteration = int(last_model_name[-11:-3])
print("Resume from {} at iteration {}".format(
checkpoint_path, model_iteration))
pred_imgs = []
for (x, mask) in zip(imgs, masks):
if torch.max(mask) == 1:
if cuda:
netG = nn.parallel.DataParallel(
netG, device_ids=device_ids)
x = x.cuda()
mask = mask.cuda()
# Inference
x1, x2, offset_flow = netG(x, mask)
inpainted_result = x2 * mask + x * (1. - mask)
inpainted_result = inpainted_result.squeeze(
dim=0).permute(1, 2, 0).cpu()
pred_imgs.append(inpainted_result.numpy())
else:
pred_imgs.append(
x.squeeze(dim=0).permute(1, 2, 0).numpy())
pred_imgs = np.asarray(pred_imgs, dtype=np.float32)
reconstructed_image = chunker.dimension_postprocess(
pred_imgs, np.array(img_ori))
# plt.imshow(reconstructed_image); plt.show()
reconstructed_image = torch.tensor(
reconstructed_image).permute(2, 0, 1).unsqueeze(dim=0)
vutils.save_image(reconstructed_image,
args.output,
padding=0,
normalize=True)
print("Saved the inpainted result to {}".format(args.output))
if args.flow:
vutils.save_image(offset_flow,
args.flow,
padding=0,
normalize=True)
print("Saved offset flow to {}".format(args.flow))
else:
raise TypeError("{} is not an image file.".format)
# exit no grad context
except Exception as e: # for unexpected error logging
print("Error: {}".format(e))
raise e
def main():
# if you want to see output images with colored segmentation?
VIEW_COLORED_SEGMENTATION = False
# Config file reading
args = parser.parse_args()
config = get_config(args.config)
# ------ CUDA configuration
cuda = config['cuda']
device_ids = config['gpu_ids']
if cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(
str(i) for i in device_ids)
device_ids = list(range(len(device_ids)))
config['gpu_ids'] = device_ids
cudnn.benchmark = True
if not os.path.exists(config['output_test_dir']):
os.makedirs(config['output_test_dir'])
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
# first, read images and pick labels with same name
# we will train all images from HQ dataset
# ---------- train and test dataset&loader
try: # for unexpected error logging
# Load the dataset
print("Inference on dataset: {" + config['dataset_name'] + "}")
test_dataset = Test_Dataset(
data_path=config['test_data_path'],
with_subfolder=config['data_with_subfolder'],
image_shape=config['image_shape'],
random_crop=config['random_crop'],
return_name=True)
test_loader = torch.utils.data.DataLoader(
dataset=test_dataset,
batch_size=config['batch_size'],
shuffle=False,
num_workers=config['num_workers'])
# [Trainer] (in test, not use trainer class directly)
netG = Parser(config, cuda, device_ids)
# Get the resume iteration to restart training
#================== <LOAD CHECKPOINT FILE starting with parser*.pt> ============================
last_checkpoint_file = get_model_list(config['resume'], "parser",
config['resume_iter'])
netG.load_state_dict(torch.load(last_checkpoint_file))
print("Resume from {}".format(config['resume']))
# CUDA AVAILABLE
if cuda:
netG = nn.parallel.DataParallel(netG, device_ids=device_ids)
# connect loaders to iter()
iterable_test_loader = iter(test_loader)
# learing rate
#lr = config['lr']
print('Inference Start.........')
start_iteration = 0
# =============== TEST ===================
for iteration in range(start_iteration, config['niter'] + 1):
print('ITERATION {}..... [{}/{}]'.format(
iteration, iteration * config['batch_size'],
int(len(test_dataset.samples) / config['batch_size'])))
try:
test_img_names, test_orig_images = iterable_test_loader.next()
except StopIteration:
iterable_test_loader = iter(test_loader)
test_img_names, test_orig_images = iterable_test_loader.next()
if cuda:
test_orig_images = test_orig_images.cuda()
# <predict test set>
test_predict = netG(test_orig_images)
for test_idx in range(test_orig_images.shape[0]):
pred_out = torch.argmax(test_predict[test_idx], dim=0)
test_sam = pred_out.cpu().numpy()
if VIEW_COLORED_SEGMENTATION:
decoded = decode_segmap(test_sam)
misc.imsave(
config['output_test_dir'] + test_img_names[test_idx] +
'.png', decoded)
else:
cv2.imwrite(
config['output_test_dir'] + test_img_names[test_idx] +
'.png', test_sam)
except Exception as e: # for unexpected error logging (set
print("{e}")
raise e
if args.seed is None:
args.seed = random.randint(1, 10000)
print("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
# print("Configuration: {}".format(config))
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids)
# Resume weight
# if cuda:
# netG.cuda()
last_model_name = get_model_list(args.checkpoint_path,
"gen",
iteration=args.iter)
last_model_name = args.which_model
# last_model_name = args.which_model
# print("loading model from here --------------> {}".format(last_model_name))
# if not cuda:
# netG.load_state_dict(torch.load(last_model_name, map_location='cpu'))
# else:
# netG.load_state_dict(torch.load(last_model_name))
def init_weights(m):
if type(m) == nn.Linear:
torch.nn.init.xavier_uniform(m.weight)
m.bias.data.fill_(0.01)
def main():
args = parser.parse_args()
config = get_config(args.config)
# CUDA configuration
cuda = config['cuda']
device_ids = config['gpu_ids']
if cuda:
os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(i) for i in device_ids)
device_ids = list(range(len(device_ids)))
config['gpu_ids'] = device_ids
cudnn.benchmark = True
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
# print("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
t0 = time.time()
dataset = datasets.LoadImages(args.image)
chunker = ImageChunker(config['image_shape'][0],
config['image_shape'][1],
args.overlap)
try: # for unexpected error logging
with torch.no_grad(): # enter no grad context
# Set checkpoint path
if not args.checkpoint_path:
checkpoint_path = os.path.join('checkpoints', config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
else:
checkpoint_path = args.checkpoint_path
last_model_name = get_model_list(checkpoint_path, "gen", iteration=args.iter)
prev_fname = ''
vid_writer = None
for fpath, img_ori, vid_cap in dataset :
imgs, masks = [], []
if prev_fname == fpath :
frame += 1 # increase frame number if still on the same file
else :
frame = 0 # start frame number
_, img_h, img_w = img_ori.shape
txtfile = pathlib.Path(fpath).with_suffix('.txt') # Load mask txt file
txtfile = os.path.join(args.output, str(txtfile).split('/')[-1])
if os.path.exists(txtfile) :
bboxes, bframes = load_bbox_txt(txtfile, img_w, img_h)
assert len(bboxes) == len(bframes)
idx = [ii for ii, val in enumerate(bframes) if val==frame]
bndbxs = [bboxes[ii] for ii in idx]
img_ori = np.moveaxis(img_ori, 0, -1)
if len(bndbxs) > 0 : # if any logo detected
mask_ori = create_mask(bndbxs, img_w, img_h)
# fig, axes = plt.subplots(1,2); axes[0].imshow(img_ori[0]); axes[1].imshow(mask_ori); plt.show()
chunked_images = chunker.dimension_preprocess(np.array(deepcopy(img_ori)))
chunked_masks = chunker.dimension_preprocess(np.array(deepcopy(mask_ori)))
for (x, msk) in zip(chunked_images, chunked_masks) :
x = transforms.ToTensor()(x)
mask = transforms.ToTensor()(msk)[0].unsqueeze(dim=0)
# x = normalize(x)
x = x * (1. - mask)
x = x.unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
imgs.append(x)
masks.append(mask)
# Define the trainer
netG = Generator(config['netG'], cuda, device_ids)
netG.load_state_dict(torch.load(last_model_name))
model_iteration = int(last_model_name[-11:-3])
# print("Resume from {} at iteration {}".format(checkpoint_path, model_iteration))
pred_imgs = []
for (x, mask) in zip(imgs, masks) :
if torch.max(mask) == 1 :
if cuda:
netG = nn.parallel.DataParallel(netG, device_ids=device_ids)
x = x.cuda()
mask = mask.cuda()
# Inference
x1, x2, offset_flow = netG(x, mask)
inpainted_result = x2 * mask + x * (1. - mask)
inpainted_result = inpainted_result.squeeze(dim=0).permute(1,2,0).cpu()
pred_imgs.append(inpainted_result.numpy())
else :
pred_imgs.append(x.squeeze(dim=0).permute(1,2,0).numpy())
pred_imgs = np.asarray(pred_imgs, dtype=np.float32)
reconstructed_image = chunker.dimension_postprocess(pred_imgs, np.array(img_ori))
reconstructed_image = np.uint8(reconstructed_image[:, :, ::-1]*255) # BGR to RGB, and rescaling
else : # no logo detected
reconstructed_image = img_ori[:, :, ::-1]
# Save results (image with detections)
outname = fpath.split('/')[-1]
outname = outname.split('.')[0] + '-inp.' + outname.split('.')[-1]
outpath = os.path.join(args.output, outname)
if dataset.mode == 'images':
cv2.imwrite(outpath, reconstructed_image)
print("Saved the inpainted image to {}".format(outpath))
else :
if fpath != prev_fname: # new video
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
print("Saved the inpainted video to {}".format(outpath))
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
vid_writer = cv2.VideoWriter(outpath, cv2.VideoWriter_fourcc(*args.fourcc), fps, (w, h))
vid_writer.write(reconstructed_image)
prev_fname = fpath
# exit no grad context
except Exception as err: # for unexpected error logging
print("Error: {}".format(err))
pass
print('Inpainting: (%.3fs)' % (time.time() - t0))