def main():
args = tools.get_args(parser)
config = tools.get_config(args)
tools.init(config)
tb_logger, logger = tools.get_logger(config)
tools.check_dist_init(config, logger)
checkpoint = tools.get_checkpoint(config)
runner = tools.get_model(config, checkpoint)
loaders = tools.get_data_loader(config)
if dist.is_master():
logger.info(config)
if args.mode == 'train':
train(config, runner, loaders, checkpoint, tb_logger)
elif args.mode == 'evaluate':
evaluate(runner, loaders)
elif args.mode == 'calc_flops':
if dist.is_master():
flops = tools.get_model_flops(config, runner.get_model())
logger.info('flops: {}'.format(flops))
elif args.mode == 'calc_params':
if dist.is_master():
params = tools.get_model_parameters(runner.get_model())
logger.info('params: {}'.format(params))
else:
assert checkpoint is not None
from models.dmcp.utils import sample_model
sample_model(config, runner.get_model())
if dist.is_master():
logger.info('Done')
def main(argv):
(opt, args) = parser.parse_args(argv)
config = get_config(opt.config)
print(opt)
if opt.manualSeed is None:
opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if config['cuda']:
torch.cuda.manual_seed_all(opt.manualSeed)
torch.cuda.set_device(opt.gpu_ids)
cudnn.benchmark = True
transform = transforms.Compose([
transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
dataset = Aligned_Dataset(config['datapath'],
direction='AtoB',
transform=transform)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=int(2))
trainer = GAN_Trainer(config, dataloader)
trainer.train()
return
def main(argv):
(opt, args) = parser.parse_args(argv)
config = get_config(opt.config)
print(opt)
if opt.manualSeed is None:
opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if config['cuda']:
torch.cuda.manual_seed_all(opt.manualSeed)
torch.cuda.set_device(opt.gpu_ids)
cudnn.benchmark = True
transform = transforms.Compose([transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])
dataset = Aligned_Dataset(config['datapath'], subfolder='test', direction='AtoB', transform=transform)
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1,
shuffle=False, num_workers=int(2))
model_dir = '/media/scw4750/AIwalker/stackgan-like/checkpoints/generator_epoch_160.pkl'
trainer = GAN_Trainer(config, dataloader)
# load the model
trainer.G.load_state_dict(torch.load(model_dir))
trainer.test()
return
def main(argv):
(opt, args) = parser.parse_args(argv)
print(opt)
config = get_config(opt.config)
if opt.manualSeed is None:
opt.manualSeed = random.randint(1, 10000)
print('Random Seed: ', opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
torch.cuda.manual_seed_all(opt.manualSeed)
torch.cuda.set_device(opt.gpu_ids)
cudnn.benchmark = True
# loading data set
transform = transforms.Compose([transforms.Resize((config['fineSizeH'], config['fineSizeW'])),
transforms.ToTensor()])
dataset = Aligned_Dataset(config['dataPath'], direction='AtoB', transform=transform)
train_loader = torch.utils.data.DataLoader(dataset, batch_size=config['batchSize'],
shuffle=True, num_workers=int(4))
# setup model
trainer = trainer_gan(config, train_loader, resume_epoch=opt.resume_epoch)
if opt.cuda:
trainer.cuda()
if opt.resume_epoch:
trainer.resume()
# training
for epoch in range(opt.resume_epoch, config['nepoch']):
trainer.train(epoch)
trainer.update_learning_rate(epoch)
if epoch % 10 == 0:
trainer.save(epoch)
def main(argv):
(opt, args) = parser.parse_args(argv)
print(opt)
config = get_config(opt.config)
if opt.manualSeed is None:
opt.manualSeed = random.randint(1, 10000)
print('Random Seed: ', opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
torch.cuda.manual_seed_all(opt.manualSeed)
torch.cuda.set_device(opt.gpu_ids)
cudnn.benchmark = True
# loading data set
transform = transforms.Compose([transforms.Resize((config['fineSizeH'], config['fineSizeW'])),
transforms.ToTensor()])
dataset = Aligned_Dataset(config['dataPath'], subfolder='test', direction='AtoB', transform=transform)
test_loader = torch.utils.data.DataLoader(dataset, batch_size=1,
shuffle=False, num_workers=int(4))
# setup model
trainer = trainer_gan(config, test_loader, resume_epoch=opt.resume_epoch)
# load a model
trainer.netG.load_state_dict(torch.load(opt.modeldir))
if opt.cuda:
trainer.cuda()
# testing
trainer.test()
def main(argv):
(opt, args) = parser.parse_args(argv)
config = get_config(opt.config)
print(opt)
if opt.manualSeed is None:
opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if config['cuda']:
torch.cuda.manual_seed_all(opt.manualSeed)
torch.cuda.set_device(opt.gpu_ids)
cudnn.benchmark = True
transform = transforms.Compose([
transforms.Resize((512, 512)),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
dataset = Aligned_Dataset(config['datapath'],
subfolder='test',
direction='AtoB',
transform=transform)
dataloader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=False,
num_workers=int(2))
model_dir = '/media/scw4750/AIwalker/stackgan-like/checkpoints/generator_epoch_160.pkl'
trainer = GAN_Trainer(config, dataloader)
# load the model
trainer.G.load_state_dict(torch.load(model_dir))
trainer.test()
return
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
# datetime object containing current date and time
#
now = datetime.now()
dt_string = now.strftime("%d-%m-%Y_%H.%M.%S")
# Configure checkpoint path
checkpoint_path = os.path.join(
'ckpts',
config['dataset_name'] + '_' + dt_string,
config['mask_type'] + '_' + config['expname']
)
if not os.path.exists(checkpoint_path):
os.makedirs(checkpoint_path, exist_ok=True)
shutil.copy(args.config, os.path.join(checkpoint_path, os.path.basename(args.config)))
writer = SummaryWriter(logdir=checkpoint_path)
logger = get_logger(checkpoint_path) # get logger and configure it at the first call
logger.info("Arguments: {}".format(args))
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
logger.info("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
# Log the configuration
logger.info("Configuration: {}".format(config))
if args.distributed:
print("Distributed training...")
# train_distributed(config, logger, writer, checkpoint_path)
train_distributed_v2(config, logger, writer, checkpoint_path)
else:
train(config, logger, checkpoint_path)
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 main():
config = get_config(args.config)
if config['cuda']:
device = torch.device("cuda:{}".format(config['gpu_ids'][0]))
else:
device = torch.device("cpu")
trainer = Trainer(config)
trainer.load_state_dict(load_weights(args.model_path, device),
strict=False)
trainer.eval()
image = imageio.imread(args.image)
image = torch.FloatTensor(image).permute(2, 0, 1).unsqueeze(0).cuda()
mask = imageio.imread(args.mask)
mask = (torch.FloatTensor(mask[:, :, 0]) /
255).unsqueeze(0).unsqueeze(0).cuda()
x = (image / 127.5 - 1) * (1 - mask).cuda()
with torch.no_grad():
_, result, _ = trainer.netG(x, mask)
imageio.imwrite(args.output,
upcast(result[0].permute(1, 2, 0).detach().cpu().numpy()))
def main(argv):
(opt, args) = parser.parse_args(argv)
print(opt)
config = get_config(opt.config)
if opt.manualSeed is None:
opt.manualSeed = random.randint(1, 10000)
print('Random Seed: ', opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
torch.cuda.manual_seed_all(opt.manualSeed)
torch.cuda.set_device(opt.gpu_ids)
cudnn.benchmark = True
# loading data set
transform = transforms.Compose([
transforms.Resize((config['fineSize'], config['fineSize'])),
transforms.ToTensor(),
transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
dataset = Aligned_Dataset(config['dataPath'],
subfolder='test',
direction='AtoB',
transform=transform)
test_loader = torch.utils.data.DataLoader(dataset,
batch_size=1,
shuffle=True,
num_workers=int(4))
# setup model
trainer = trainer_gan(config, test_loader)
# load a model
trainer.netG.load_state_dict(torch.load(opt.modeldir))
if opt.cuda:
trainer.cuda()
# testing
trainer.test()
def generateInpaintedImage(args, netG, imagePath):
config = get_config(args.g_config)
occlusions = []
# 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)
try: # for unexpected error logging
with torch.no_grad(): # enter no grad context
if is_image_file(imagePath):
if args.mask and is_image_file(args.mask):
# Test a multiple masked image with a given mask
x = default_loader(imagePath)
x = transforms.Resize([512, 1024])(x)
mask = default_loader(args.mask)
mask = transforms.Resize(config['image_shape'][:-1])(mask)
mask = transforms.CenterCrop(
config['image_shape'][:-1])(mask)
mask = transforms.ToTensor()(mask)[0].unsqueeze(dim=0)
mask = mask.unsqueeze(dim=0)
w, h = x.size
first = x.crop((0, 0, w // 3, h))
second = x.crop((w // 3, 0, ((w // 3) * 2) + 2, h))
third = x.crop(((w // 3) * 2, 0, w, h))
for y in [first, second, third]:
y = transforms.CenterCrop(
config['image_shape'][:-1])(y)
y = transforms.ToTensor()(y)
y = normalize(y)
y = y * (1. - mask)
occlusions.append(y)
elif args.mask:
raise TypeError("{} is not an image file.".format(
args.mask))
default_image = default_loader(imagePath)
di_w, di_h = default_image.size
for idx, occlusion in enumerate(occlusions):
if cuda:
occlusion = occlusion.cuda()
mask = mask.cuda()
# Inference
x1, x2, offset_flow = netG(occlusion, mask)
inpainted_result = x2 * mask + occlusion * (1. - mask)
inp_hw = config['image_shape'][1]
if idx == 0:
offset = ((di_w // 3 - inp_hw) // 2,
(di_h - inp_hw) // 2)
elif idx == 1:
offset = ((di_w - inp_hw) // 2, (di_h - inp_hw) // 2)
elif idx == 2:
offset = ((((di_w - inp_hw) // 2) + (di_w // 3)),
(di_h - inp_hw) // 2)
grid = vutils.make_grid(inpainted_result, normalize=True)
# Add 0.5 after unnormalizing to [0, 255] to round to nearest integer
ndarr = grid.mul_(255).add_(0.5).clamp_(0, 255).permute(
1, 2, 0).to('cpu', torch.uint8).numpy()
im = Image.fromarray(ndarr)
im = transforms.CenterCrop(config['mask_shape'])(im)
im = transforms.Resize(config['image_shape'][:-1])(im)
default_image.paste(im, offset)
return default_image
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
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():
# 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
# ----- Directory to save checkpoint file
checkpoint_path = os.path.join('checkpoints_',
config['dataset_name'],
'parser' + '_' + config['expname'])
if not os.path.exists(checkpoint_path):
os.makedirs(checkpoint_path)
shutil.copy(args.config, os.path.join(checkpoint_path, os.path.basename(args.config)))
writer = SummaryWriter(log_dir=checkpoint_path)
logger = get_logger(checkpoint_path) # get logger and configure it at the first call
if not os.path.exists(config['output_test_dir']):
os.makedirs(config['output_test_dir'])
logger.info(f"Arguments: {args}")
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
logger.info(f"Random seed: {args.seed}")
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
logger.info(f"configuration : {config}")
# 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
logger.info(f"Training on dataset: {config['dataset_name']}")
train_dataset = Parse_Dataset(data_path=config['all_data_path'],
gt_path=config['gt_data_path'],
with_subfolder=config['data_with_subfolder'],
image_shape=config['image_shape'],
random_crop=config['random_crop'], return_name=True)
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
batch_size=config['batch_size'],
shuffle=True,
num_workers=config['num_workers'])
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]
trainer = Seg_Trainer(config)
logger.info(f"\n{trainer.netParser}")
# CUDA AVAILABLE
if cuda:
trainer = nn.parallel.DataParallel(trainer, device_ids=device_ids)
trainer_module = trainer.module
else:
trainer_module = trainer
# Get the resume iteration to restart training
#start_iteration = trainer_module.resume(config['resume']) if config['resume'] else 1
start_iteration = trainer_module.resume(config['resume'], config['resume_iter']) if config['resume'] else 1
# connect loaders to iter()
iterable_train_loader = iter(train_loader)
iterable_test_loader = iter(test_loader)
# learing rate
lr = config['lr']
print('Training Start.........')
for iteration in range(start_iteration, config['niter'] + 1):
#=============== TRAIN ===================
# ------ [ train batch loader ] ---------
try:
train_img_names, gt_images, gt_targets, orig_images = iterable_train_loader.next()
except StopIteration:
iterable_train_loader = iter(train_loader)
train_img_names, gt_images, gt_targets, orig_images = iterable_train_loader.next()
# ------ [ train batch ] ---------
if cuda:
orig_images = orig_images.cuda()
gt_images = gt_images.cuda()
# Forward
loss, predict = trainer(orig_images, gt_images)
if not loss.dim() == 0:
loss = torch.mean(loss)
# Backward (update optimizer)
trainer_module.optimizerSGD.zero_grad()
loss.backward()
trainer_module.optimizerSGD.step()
# [print loss] (in this, 1 print for 30 iteration)
if iteration % 30 == 0:
print("Epoch [%d/%d] Loss: %.10f lr:%.6f" % (iteration ,config['niter'], loss.data, lr))
#=============== TEST ===================
if iteration % 50 == 0:
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 = trainer.module.netParser(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()
cv2.imwrite(config['output_test_dir'] + test_img_names[test_idx] + '.png', test_sam)
# <learning rate up>
if iteration % 50000 == 0:
lr = lr * config['lr_decay']
for param_group in trainer_module.optimizerSGD.param_groups:
param_group['lr'] = lr
# save the model
if iteration % config['snapshot_save_iter'] == 0:
trainer_module.save_model(checkpoint_path, iteration)
except Exception as e: # for unexpected error logging
logger.error(f"{e}")
raise e
def main(argv):
(opt, args) = parser.parse_args(argv)
print(opt)
if opt.manualSeed is None:
opt.manualSeed = random.randint(1, 10000)
print("random seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
torch.cuda.manual_seed_all(opt.manualSeed)
torch.cuda.set_device(opt.gpu_ids)
cudnn.benchmark = True
config = get_config(opt.config)
# loading data sets
datasetA = DATASET(os.path.join(config['dataPath'], 'trainA'),
config['loadSize'], config['fineSize'], config['flip'])
datasetB = DATASET(os.path.join(config['dataPath'], 'trainB'),
config['loadSize'], config['fineSize'], config['flip'])
loader_A = torch.utils.data.DataLoader(dataset=datasetA,
batch_size=config['batchSize'],
shuffle=True,
num_workers=config['num_workers'])
loaderA = iter(loader_A)
loader_B = torch.utils.data.DataLoader(dataset=datasetB,
batch_size=config['batchSize'],
shuffle=True,
num_workers=config['num_workers'])
loaderB = iter(loader_B)
# define the trainer
trainer = ImGANTrainer(config)
print(trainer.netG_ab, trainer.netD_ab)
if opt.cuda:
trainer.cuda()
for iteration in range(1, config['niter'] + 1):
try:
imgA = loaderA.next()
imgB = loaderB.next()
except StopIteration:
loaderA, loaderB = iter(loader_A), iter(loader_B)
imgA = loaderA.next()
imgB = loaderB.next()
image_A = Variable(imgA.cuda())
image_B = Variable(imgB.cuda())
trainer.dis_upodate(image_A, image_B)
trainer.gen_update(image_A, image_B)
if iteration % 100 == 0:
trainer.get_current_losses()
losses = trainer.get_current_losses()
message = '([%d/%d][%d/%d]) ' % (iteration, config['niter'],
len(loader_A), len(loader_B))
for k, v in losses.items():
message += '%s: %.6f ' % (k, v)
print(message)
input_a, input_b, fake_ab, fake_ba, fake_bab, fake_aba = trainer.test(
image_A, image_B)
vutils.save_image(image_A.data,
'%s/realA_niter_%03d_1.png' %
(config['outf'], iteration),
normalize=True)
vutils.save_image(image_B.data,
'%s/realB_niter_%03d_1.png' %
(config['outf'], iteration),
normalize=True)
vutils.save_image(fake_aba.data,
'%s/recA_niter_%03d_1.png' %
(config['outf'], iteration),
normalize=True)
vutils.save_image(fake_bab.data,
'%s/recB_niter_%03d_1.png' %
(config['outf'], iteration),
normalize=True)
vutils.save_image(fake_ab.data,
'%s/AB_niter_%03d_1.png' %
(config['outf'], iteration),
normalize=True)
vutils.save_image(fake_ba.data,
'%s/BA_niter_%03d_1.png' %
(config['outf'], iteration),
normalize=True)
parser.add_argument('--seed', type=int, default=0, help='manual seed')
# parser.add_argument('--image', type=str)
# parser.add_argument('--mask', type=str, default='')
# parser.add_argument('--output', type=str, default='output.png')
# parser.add_argument('--flow', type=str, default='')
parser.add_argument('--checkpoint_path',
type=str,
default='checkpoints2/imagenet/hole_benchmark')
parser.add_argument('--iter', type=int,
default=0) # default means the latest iteration
parser.add_argument(
'--which_model',
type=str,
default='checkpoints2/imagenet/hole_benchmark/gen_00080000.pt')
args = parser.parse_args()
config = get_config(args.config)
# CUDA configuration
cuda = config['cuda']
if torch.cuda.device_count() > 0:
cuda = True # memory problem
else:
cuda = False
if not args.checkpoint_path:
args.checkpoint_path = os.path.join(
'checkpoints', config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
print("args.checkpoint_path = {}".format(args.checkpoint_path))
dataset_name = args.checkpoint_path.split("/")[1]
device_ids = config['gpu_ids']
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
# Configure checkpoint path
checkpoint_path = os.path.join(
'checkpoints', config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
if not os.path.exists(checkpoint_path):
os.makedirs(checkpoint_path)
shutil.copy(args.config,
os.path.join(checkpoint_path, os.path.basename(args.config)))
writer = SummaryWriter(log_dir=checkpoint_path)
logger = get_logger(
checkpoint_path) # get logger and configure it at the first call
logger.info(f"Arguments: {args}")
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
logger.info(f"Random seed: {args.seed}")
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
# Log the configuration
logger.info(f"Configuration: {config}")
try: # for unexpected error logging
# Load the dataset
logger.info(f"Training on dataset: {config['dataset_name']}")
train_dataset = Dataset(data_path=config['train_data_path'],
with_subfolder=config['data_with_subfolder'],
image_shape=config['image_shape'],
random_crop=config['random_crop'])
train_loader = torch.utils.data.DataLoader(
dataset=train_dataset,
batch_size=config['batch_size'],
shuffle=True,
num_workers=config['num_workers'])
# Define the trainer
trainer = Trainer(config)
logger.info(f"\n{trainer.netG}")
logger.info(f"\n{trainer.globalD}")
if cuda:
trainer = nn.parallel.DataParallel(trainer, device_ids=device_ids)
trainer_module = trainer.module
else:
trainer_module = trainer
# Get the resume iteration to restart training
start_iteration = trainer_module.resume(
config['checkpoint_dir'],
args.iteration) if config['resume'] else 1
iterable_train_loader = iter(train_loader)
time_count = time.time()
for iteration in range(start_iteration, config['niter'] + 1):
try:
ground_truth = iterable_train_loader.next()
except StopIteration:
iterable_train_loader = iter(train_loader)
ground_truth = iterable_train_loader.next()
# Prepare the inputs
x, mask = mask_image(ground_truth, config)
if cuda:
x = x.cuda()
mask = mask.cuda()
ground_truth = ground_truth.cuda()
###### Forward pass ######
losses, coarse_result, inpainted_result = trainer(
x, mask, ground_truth)
# Scalars from different devices are gathered into vectors
for k in losses.keys():
if not losses[k].dim() == 0:
losses[k] = torch.mean(losses[k])
###### Backward pass ######
# Update D
trainer_module.optimizer_d.zero_grad()
losses['d'] = losses['d_loss_rel'] + losses['d_loss_loren']
losses['d'].backward(retain_graph=True)
trainer_module.optimizer_d.step()
# Update G
trainer_module.optimizer_g.zero_grad()
losses['g'] = losses['l1'] * config['l1_loss_alpha'] + losses[
'g_loss_rel'] + losses['g_loss_loren']
losses['g'].backward()
trainer_module.optimizer_g.step()
# Log and visualization
log_losses = ['l1', 'g', 'd']
if iteration % config['print_iter'] == 0:
time_count = time.time() - time_count
speed = config['print_iter'] / time_count
speed_msg = 'speed: %.2f batches/s ' % speed
time_count = time.time()
message = 'Iter: [%d/%d] ' % (iteration, config['niter'])
for k in log_losses:
v = losses[k]
writer.add_scalar(k, v, iteration)
message += '%s: %.6f ' % (k, v)
message += speed_msg
logger.info(message)
def denorm(x):
out = (x + 1) / 2 # [-1,1] -> [0,1]
return out.clamp_(0, 1)
if iteration % (config['viz_iter']) == 0:
ims = torch.cat(
[x, coarse_result, inpainted_result, ground_truth], dim=3)
writer.add_image('raw_masked_coarse_refine', denorm(ims),
iteration)
# if iteration % (config['save_image']) ==0:
# viz_max_out = config['viz_max_out']
# if x.size(0) > viz_max_out:
# viz_images = torch.stack([x[:viz_max_out], coarse_result[:viz_max_out], inpainted_result[:viz_max_out],
# ], dim=1)
# else:
# viz_images = torch.stack([x, coarse_result, inpainted_result], dim=1)
# viz_images = viz_images.view(-1, *list(x.size())[1:])
# vutils.save_image(viz_images,
# '%s/niter_%03d.png' % (checkpoint_path, iteration),
# nrow=3 * 4,
# normalize=True)
# Save the model
if iteration % config['snapshot_save_iter'] == 0:
trainer_module.save_model(checkpoint_path)
except Exception as e: # for unexpected error logging
logger.error(f"{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
# Configure checkpoint path
checkpoint_path = os.path.join(
'checkpoints', config['dataset_name'],
config['mask_type'] + '_' + config['expname'])
if not os.path.exists(checkpoint_path):
os.makedirs(checkpoint_path)
shutil.copy(args.config,
os.path.join(checkpoint_path, os.path.basename(args.config)))
writer = SummaryWriter(logdir=checkpoint_path)
logger = get_logger(
checkpoint_path) # get logger and configure it at the first call
logger.info("Arguments: {}".format(args))
# Set random seed
if args.seed is None:
args.seed = random.randint(1, 10000)
logger.info("Random seed: {}".format(args.seed))
random.seed(args.seed)
torch.manual_seed(args.seed)
if cuda:
torch.cuda.manual_seed_all(args.seed)
# Log the configuration
logger.info("Configuration: {}".format(config))
try: # for unexpected error logging
# Load the dataset
logger.info("Training on dataset: {}".format(config['dataset_name']))
train_dataset = Dataset(data_path=config['train_data_path'],
with_subfolder=config['data_with_subfolder'],
image_shape=config['image_shape'],
random_crop=config['random_crop'])
# val_dataset = Dataset(data_path=config['val_data_path'],
# with_subfolder=config['data_with_subfolder'],
# image_size=config['image_size'],
# random_crop=config['random_crop'])
train_loader = torch.utils.data.DataLoader(
dataset=train_dataset,
batch_size=config['batch_size'],
shuffle=True,
num_workers=config['num_workers'])
# val_loader = torch.utils.data.DataLoader(dataset=val_dataset,
# batch_size=config['batch_size'],
# shuffle=False,
# num_workers=config['num_workers'])
# Define the trainer
trainer = Trainer(config)
logger.info("\n{}".format(trainer.netG))
logger.info("\n{}".format(trainer.localD))
logger.info("\n{}".format(trainer.globalD))
if cuda:
trainer = nn.parallel.DataParallel(trainer, device_ids=device_ids)
trainer_module = trainer.module
else:
trainer_module = trainer
# Get the resume iteration to restart training
start_iteration = trainer_module.resume(
"checkpoints/{}/hole_benchmark".format(config['dataset_name']),
iteration=config['resume']) if config['resume'] else 1
iterable_train_loader = iter(train_loader)
time_count = time.time()
for iteration in range(start_iteration, config['niter'] + 1):
try:
ground_truth = next(iterable_train_loader)
except StopIteration:
iterable_train_loader = iter(train_loader)
ground_truth = next(iterable_train_loader)
# Prepare the inputs
bboxes = random_bbox(config, batch_size=ground_truth.size(0))
x, mask = mask_image(ground_truth, bboxes, config)
if cuda:
x = x.cuda()
mask = mask.cuda()
ground_truth = ground_truth.cuda()
###### Forward pass ######
compute_g_loss = iteration % config['n_critic'] == 0
losses, inpainted_result, offset_flow = trainer(
x, bboxes, mask, ground_truth, compute_g_loss)
# Scalars from different devices are gathered into vectors
for k in losses.keys():
if not losses[k].dim() == 0:
losses[k] = torch.mean(losses[k])
###### Backward pass ######
# Update D
trainer_module.optimizer_d.zero_grad()
losses['d'] = losses[
'wgan_d'] + losses['wgan_gp'] * config['wgan_gp_lambda']
losses['d'].backward()
trainer_module.optimizer_d.step()
# Update G
if compute_g_loss:
trainer_module.optimizer_g.zero_grad()
losses['g'] = losses['l1'] * config['l1_loss_alpha'] \
+ losses['ae'] * config['ae_loss_alpha'] \
+ losses['wgan_g'] * config['gan_loss_alpha']
losses['g'].backward()
trainer_module.optimizer_g.step()
# Log and visualization
log_losses = ['l1', 'ae', 'wgan_g', 'wgan_d', 'wgan_gp', 'g', 'd']
if iteration % config['print_iter'] == 0:
time_count = time.time() - time_count
speed = config['print_iter'] / time_count
speed_msg = 'speed: %.2f batches/s ' % speed
time_count = time.time()
message = 'Iter: [%d/%d] ' % (iteration, config['niter'])
for k in log_losses:
v = losses.get(k, 0.)
writer.add_scalar(k, v, iteration)
message += '%s: %.6f ' % (k, v)
message += speed_msg
logger.info(message)
if iteration % (config['viz_iter']) == 0:
viz_max_out = config['viz_max_out']
if x.size(0) > viz_max_out:
viz_images = torch.stack([
x[:viz_max_out], inpainted_result[:viz_max_out],
offset_flow[:viz_max_out]
],
dim=1)
else:
viz_images = torch.stack(
[x, inpainted_result, offset_flow], dim=1)
viz_images = viz_images.view(-1, *list(x.size())[1:])
vutils.save_image(viz_images,
'%s/niter_%03d.png' %
(checkpoint_path, iteration),
nrow=3 * 4,
normalize=True)
# Save the model
if iteration % config['snapshot_save_iter'] == 0:
trainer_module.save_model(checkpoint_path, iteration)
except Exception as e: # for unexpected error logging
logger.error("{}".format(e))
raise e
path2best_model = os.path.join(weights_folder, f"{run_name}.pth")
log_metrics(metrics,
epoch,
"val",
logger,
config["log_mlflow"],
path2best_model=path2best_model)
if __name__ == '__main__':
from datetime import datetime
from utils.tools import get_config
run_name = datetime.now().strftime("%d_%m_%Y.%H_%M")
config = get_config()
log_mlflow = config.get("mlflow", None) is not None
config["log_mlflow"] = log_mlflow
config["run_name"] = run_name
setup_mlflow(config)
weights_folder = os.path.join(os.getcwd(), "weights",
config["description"]["model_name"],
config["description"]["task_type"])
log_folder = os.path.join(os.getcwd(), "logs",
config["description"]["model_name"],
config["description"]["task_type"], run_name)
os.makedirs(weights_folder, exist_ok=True)
os.makedirs(log_folder, exist_ok=True)
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
file = dataset_files(args.test_root, "*.jpg")
mask_file = dataset_files(args.mask_dir, "*.png")
for j in range(len(mask_file)):
for i in range(len(file)):
if is_image_file(file[i]):
if mask_file and is_image_file(mask_file[j]):
# Test a single masked image with a given mask
x = default_loader(file[i])
mask = default_loader(mask_file[j])
# x = cv2.cvtColor(cv2.imread(file[i]), cv2.COLOR_BGR2RGB)
# mask = cv2.cvtColor(cv2.imread(mask_file[j]), cv2.COLOR_BGR2RGB)
# x = cv2.resize(x, (config['image_shape'][0], config['image_shape'][1]))
# mask = cv2.resize(mask, (config['image_shape'][0], config['image_shape'][1]))
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)
# x_raw = x
mask = mask.unsqueeze(dim=0)
elif mask_file[j]:
raise TypeError("{} is not an image file.".format(
mask_file[j]))
else:
# Test a single ground-truth image with a random mask
ground_truth = default_loader(file[i])
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 = test_bbox(config,
batch_size=ground_truth.size(0),
t=50,
l=50)
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
g_checkpoint = torch.load(f'{checkpoint_path}/gen.pt')
netG.load_state_dict(g_checkpoint)
# model_iteration = int(last_model_name[-11:-3])
print("Model Resumed".format(checkpoint_path))
if cuda:
netG = nn.parallel.DataParallel(
netG, device_ids=device_ids)
x = x.cuda()
mask = mask.cuda()
# Inference
x1, x2 = netG(x, mask)
inpainted_result = x2 * mask + x * (1. - mask)
inpainted_result_cpu = torch.Tensor.cpu(
inpainted_result).detach().permute(0, 2, 3, 1)
inpainted_result_cpu = np.asarray(
inpainted_result_cpu[0])
inpainted_result_cpu = cv2.normalize(
inpainted_result_cpu, inpainted_result_cpu, 0, 255,
cv2.NORM_MINMAX)
# cat_result = torch.cat([x, inpainted_result, ground_truth], dim=3).cuda()
vutils.save_image(inpainted_result,
args.output_dir +
'output_{}/'.format(j + 1) +
'output_{}.png'.format(i),
padding=0,
normalize=True)
# cv2.imwrite(args.output_dir+ 'output_{}/'.format(j+1) + 'output_{}.png'.format(i), inpainted_result_cpu)
# cv2.cvtColor(inpainted_result_cpu, cv2.COLOR_BGR2RGB))
print("{}th image saved".format(i))
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
from utils.tools import get_config, random_bbox, mask_image, is_image_file, default_loader, normalize, get_model_list
import torch
import math
from multiprocessing.dummy import Pool
irange = range
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
config = get_config('configs/config.yaml')
# CUDA configuration
cuda = config['cuda']
print(' ------------- cuda -------------')
print(cuda)
device_ids = config['gpu_ids']
def make_grid(tensor, nrow=8, padding=2,
normalize=False, range=None, scale_each=False, pad_value=0):
"""Make a grid of images.
Args:
tensor (Tensor or list): 4D mini-batch Tensor of shape (B x C x H x W)
or a list of images all of the same size.
nrow (int, optional): Number of images displayed in each row of the grid.
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))
