def createNirDataloader(nir_path,
rgb_path,
label_path,
batch_size=16,
use_cuda=True):
'''
创建dataloader
:param image_path:
:param label_path:
:param batch_size:
:param use_cuda:
:return:
'''
dataset = NirRgbData(nir_path,
rgb_path,
label_path,
transform=NormalizeImageDict(
["nir_image", "rgb_image"]))
dataloader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=4,
pin_memory=True)
pair_generator = NirRgbTnsPair(use_cuda=use_cuda)
return dataloader, pair_generator
def createDataloader(image_path, label_path, batch_size=16, use_cuda=True):
'''
创建dataloader
:param image_path:
:param label_path:
:param batch_size:
:param use_cuda:
:return:
'''
#dataset = SinglechannelData(image_path,label_path,transform=NormalizeImage(normalize_range=True, normalize_img=False))
dataset = TestDataset(image_path,
label_path,
transform=NormalizeImageDict(["image"]))
dataloader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=4,
pin_memory=True)
#pair_generator = SingleChannelPairTnf(use_cuda=use_cuda)
pair_generator = RandomTnsPair(use_cuda=use_cuda)
return dataloader, pair_generator
def main(args):
# checkpoint_path = "/home/zale/project/registration_cnn_ntg/trained_weight/voc2011_multi_gpu/checkpoint_voc2011_multi_gpu_paper_NTG_resnet101.pth.tar"
# checkpoint_path = "/home/zale/project/registration_cnn_ntg/trained_weight/coco2017_multi_gpu/checkpoint_coco2017_multi_gpu_paper30_NTG_resnet101.pth.tar"
#args.training_image_path = '/home/zale/datasets/vocdata/VOC_train_2011/VOCdevkit/VOC2011/JPEGImages'
# args.training_image_path = '/media/disk2/zale/datasets/coco2017/train2017'
checkpoint_path = "/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011_multi_gpu/checkpoint_voc2011_multi_gpu_three_channel_paper_origin_NTG_resnet101.pth.tar"
args.training_image_path = '/home/zlk/datasets/vocdata/VOC_train_2011/VOCdevkit/VOC2011/JPEGImages'
random_seed = 10021
init_seeds(random_seed + random.randint(0, 10000))
mixed_precision = True
utils.init_distributed_mode(args)
print(args)
#device,local_rank = torch_util.select_device(multi_process =True,apex=mixed_precision)
device = torch.device(args.device)
use_cuda = True
# Data loading code
print("Loading data")
RandomTnsDataset = RandomTnsData(args.training_image_path,
cache_images=False,
paper_affine_generator=True,
transform=NormalizeImageDict(["image"]))
# train_dataloader = DataLoader(RandomTnsDataset, batch_size=args.batch_size, shuffle=True, num_workers=4,
# pin_memory=True)
print("Creating data loaders")
if args.distributed:
train_sampler = torch.utils.data.distributed.DistributedSampler(
RandomTnsDataset)
# test_sampler = torch.utils.data.distributed.DistributedSampler(dataset_test)
else:
train_sampler = torch.utils.data.RandomSampler(RandomTnsDataset)
# test_sampler = torch.utils.data.SequentialSampler(dataset_test)
# train_batch_sampler = torch.utils.data.BatchSampler(
# train_sampler, args.batch_size, drop_last=True)
data_loader = DataLoader(RandomTnsDataset,
sampler=train_sampler,
num_workers=4,
shuffle=(train_sampler is None),
pin_memory=False,
batch_size=args.batch_size)
# data_loader_test = torch.utils.data.DataLoader(
# dataset_test, batch_size=1,
# sampler=test_sampler, num_workers=args.workers,
# collate_fn=utils.collate_fn)
print("Creating model")
model = CNNRegistration(use_cuda=use_cuda)
model.to(device)
# 优化器 和scheduler
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.Adam(params, lr=args.lr)
# 学习率小于1e-6 ntg损失下降很慢，所以最小设置为1e-6
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer, T_max=args.lr_max_iter, eta_min=1e-6)
# if mixed_precision:
# model,optimizer = amp.initialize(model,optimizer,opt_level='O1',verbosity=0)
model_without_ddp = model
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu])
model_without_ddp = model.module
minium_loss, saved_epoch = load_checkpoint(model_without_ddp, optimizer,
lr_scheduler, checkpoint_path,
args.rank)
vis_env = "multi_gpu_rgb_train_paper_30"
loss = NTGLoss()
pair_generator = RandomTnsPair(use_cuda=use_cuda)
gridGen = AffineGridGen()
vis = VisdomHelper(env_name=vis_env)
print('Starting training...')
start_time = time.time()
draw_test_loss = False
log_interval = 20
for epoch in range(saved_epoch, args.num_epochs):
start_time = time.time()
if args.distributed:
train_sampler.set_epoch(epoch)
train_loss = train(epoch,
model,
loss,
optimizer,
data_loader,
pair_generator,
gridGen,
vis,
use_cuda=use_cuda,
log_interval=log_interval,
lr_scheduler=lr_scheduler,
rank=args.rank)
if draw_test_loss:
#test_loss = test(model,loss,test_dataloader,pair_generator,gridGen,use_cuda=use_cuda)
#vis.drawBothLoss(epoch,train_loss,test_loss,'loss_table')
pass
else:
vis.drawLoss(epoch, train_loss)
end_time = time.time()
print("epoch:", str(end_time - start_time), '秒')
is_best = train_loss < minium_loss
minium_loss = min(train_loss, minium_loss)
state_dict = model_without_ddp.state_dict()
if is_main_process():
save_checkpoint(
{
'epoch': epoch + 1,
'args': args,
# 'state_dict': model.state_dict(),
'state_dict': state_dict,
'minium_loss': minium_loss,
'model_loss': train_loss,
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
},
is_best,
checkpoint_path)
def start_train(training_path,test_image_path,load_from,out_path,vis_env,paper_affine_generator = False,
random_seed=666,log_interval=100,multi_gpu=True,use_cuda=True):
init_seeds(random_seed+random.randint(0,10000))
device,local_rank = torch_util.select_device(multi_process =multi_gpu,apex=mixed_precision)
# args.batch_size = args.batch_size * torch.cuda.device_count()
args.batch_size = 16
args.lr_scheduler = True
draw_test_loss = False
print(args.batch_size)
print("创建模型中")
model = CNNRegistration(use_cuda=use_cuda)
model = model.to(device)
# 优化器 和scheduler
optimizer = optim.Adam(model.FeatureRegression.parameters(), lr=args.lr)
if args.lr_scheduler:
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer,
T_max=args.lr_max_iter,
eta_min=1e-7)
else:
scheduler = False
print("加载权重")
minium_loss,saved_epoch= load_checkpoint(model,optimizer,load_from,0)
# Mixed precision training https://github.com/NVIDIA/apex
if mixed_precision:
model,optimizer = amp.initialize(model,optimizer,opt_level='01',verbosity=0)
if multi_gpu:
model = nn.DataParallel(model)
loss = NTGLoss()
pair_generator = RandomTnsPair(use_cuda=use_cuda)
gridGen = AffineGridGen()
vis = VisdomHelper(env_name=vis_env)
print("创建dataloader")
RandomTnsDataset = RandomTnsData(training_path, cache_images=False,paper_affine_generator = paper_affine_generator,
transform=NormalizeImageDict(["image"]))
train_dataloader = DataLoader(RandomTnsDataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
if draw_test_loss:
testDataset = RandomTnsData(test_image_path, cache_images=False, paper_affine_generator=paper_affine_generator,
transform=NormalizeImageDict(["image"]))
test_dataloader = DataLoader(testDataset, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=False)
print('Starting training...')
for epoch in range(saved_epoch, args.num_epochs):
start_time = time.time()
train_loss = train(epoch, model, loss, optimizer, train_dataloader, pair_generator, gridGen, vis,
use_cuda=use_cuda, log_interval=log_interval,scheduler = scheduler)
if draw_test_loss:
test_loss = test(model,loss,test_dataloader,pair_generator,gridGen,use_cuda=use_cuda)
vis.drawBothLoss(epoch,train_loss,test_loss,'loss_table')
else:
vis.drawLoss(epoch,train_loss)
end_time = time.time()
print("epoch:", str(end_time - start_time),'秒')
is_best = train_loss < minium_loss
minium_loss = min(train_loss, minium_loss)
state_dict = model.module.state_dict() if multi_gpu else model.state_dict()
save_checkpoint({
'epoch': epoch + 1,
'args': args,
#'state_dict': model.state_dict(),
'state_dict': state_dict,
'minium_loss': minium_loss,
'model_loss':train_loss,
'optimizer': optimizer.state_dict(),
}, is_best, out_path)
def main():
print("eval pf dataset")
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
# ntg_checkpoint_path = "/home/zlk/project/registration_cnn_ntg/trained_weight/output/voc2012_coco2014_NTG_resnet101.pth.tar"
# ntg_checkpoint_path = "/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011/checkpoint_voc2011_NTG_resnet101.pth.tar"
# ntg_checkpoint_path = "/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011/checkpoint_voc2011_20r_NTG_resnet101.pth.tar"
# ntg_checkpoint_path = '/home/zlk/project/registration_cnn_ntg/trained_weight/three_channel/checkpoint_NTG_resnet101.pth.tar'
small_aff_ntg_checkpoint_path = '/home/zlk/project/registration_cnn_ntg/trained_weight/three_channel/coco2014_small_aff_checkpoint_NTG_resnet101.pth.tar'
ntg_checkpoint_path = '/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011/best_checkpoint_voc2011_three_channel_paper_NTG_resnet101.pth.tar'
# ntg_checkpoint_path = '/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011_paper_affine/best_checkpoint_voc2011_NTG_resnet101.pth.tar'
#ntg_checkpoint_path = "/home/zlk/project/registration_cnn_ntg/trained_weight/voc2011/checkpoint_voc2011_30r_NTG_resnet101.pth.tar"
# image_path = '../datasets/row_data/VOC/3
# label_path = '../datasets/row_data/label_file/aff_param2.csv'
#image_path = '../datasets/row_data/COCO/'
#label_path = '../datasets/row_data/label_file/aff_param_coco.csv'
pf_data_path = 'datasets/row_data/pf_data'
batch_size = 128
# 加载模型
use_cuda = torch.cuda.is_available()
ntg_model = createModel(ntg_checkpoint_path, use_cuda=use_cuda)
small_aff_ntg_model = createModel(small_aff_ntg_checkpoint_path,
use_cuda=use_cuda)
dataset = PFDataset(
csv_file=os.path.join(pf_data_path, 'test_pairs_pf.csv'),
training_image_path=pf_data_path,
transform=NormalizeImageDict(['source_image', 'target_image']))
dataloader = DataLoader(dataset,
batch_size=batch_size,
shuffle=False,
num_workers=4)
batchTensorToVars = BatchTensorToVars(use_cuda=use_cuda)
pt = PointTnf(use_cuda=use_cuda)
print('Computing PCK...')
total_correct_points_aff = 0
ntg_total_correct_points_aff = 0
cnn_ntg_total_correct_points_aff = 0
total_correct_points_tps = 0
total_correct_points_aff_tps = 0
total_points = 0
ntg_total_points = 0
cnn_ntg_total_points = 0
for i, batch in enumerate(dataloader):
batch = batchTensorToVars(batch)
source_im_size = batch['source_im_size']
target_im_size = batch['target_im_size']
source_points = batch['source_points']
target_points = batch['target_points']
source_image_batch = batch['source_image']
target_image_batch = batch['target_image']
# warp points with estimated transformations
target_points_norm = PointsToUnitCoords(target_points, target_im_size)
theta_estimate_batch = ntg_model(batch)
#warped_image_batch = affine_transform_pytorch(source_image_batch, theta_estimate_batch)
#batch['source_image'] = warped_image_batch
#theta_estimate_batch = small_aff_ntg_model(batch)
# 将pytorch的变换参数转为opencv的变换参数
#theta_opencv = theta2param(theta_estimate_batch.view(-1, 2, 3), 240, 240, use_cuda=use_cuda)
# P5使用传统NTG方法进行优化cnn的结果
#cnn_ntg_param_batch = estimate_param_batch(source_image_batch, target_image_batch, theta_opencv,itermax = 600)
#theta_pytorch = param2theta(cnn_ntg_param_batch.view(-1, 2, 3),240,240,use_cuda=use_cuda)
# theta_opencv = theta2param(theta_estimate_batch.view(-1, 2, 3), 240, 240, use_cuda=use_cuda)
# with torch.no_grad():
# ntg_param_batch = estimate_aff_param_iterator(source_image_batch[:, 0, :, :].unsqueeze(1),
# target_image_batch[:, 0, :, :].unsqueeze(1),
# None, use_cuda=use_cuda, itermax=600)
#
# cnn_ntg_param_batch = estimate_aff_param_iterator(source_image_batch[:, 0, :, :].unsqueeze(1),
# target_image_batch[:, 0, :, :].unsqueeze(1),
# theta_opencv, use_cuda=use_cuda, itermax=600)
#
# ntg_param_pytorch_batch = param2theta(ntg_param_batch,240, 240, use_cuda=use_cuda)
# cnn_ntg_param_pytorch_batch = param2theta(cnn_ntg_param_batch,240, 240, use_cuda=use_cuda)
warped_points_aff_norm = pt.affPointTnf(theta_estimate_batch,
target_points_norm)
warped_points_aff = PointsToPixelCoords(warped_points_aff_norm,
source_im_size)
# ntg_warped_points_aff_norm = pt.affPointTnf(ntg_param_pytorch_batch, target_points_norm)
# ntg_warped_points_aff = PointsToPixelCoords(ntg_warped_points_aff_norm, source_im_size)
#
# cnn_ntg_warped_points_aff_norm = pt.affPointTnf(cnn_ntg_param_pytorch_batch, target_points_norm)
# cnn_ntg_warped_points_aff = PointsToPixelCoords(cnn_ntg_warped_points_aff_norm, source_im_size)
L_pck = batch['L_pck'].data
correct_points_aff, num_points = correct_keypoints(
source_points.data, warped_points_aff.data, L_pck)
# ntg_correct_points_aff, ntg_num_points = correct_keypoints(source_points.data,
# ntg_warped_points_aff.data, L_pck)
# cnn_ntg_correct_points_aff, cnn_ntg_num_points = correct_keypoints(source_points.data,
# cnn_ntg_warped_points_aff.data, L_pck)
total_correct_points_aff += correct_points_aff
total_points += num_points
# ntg_total_correct_points_aff += ntg_correct_points_aff
# ntg_total_points += ntg_num_points
#
# cnn_ntg_total_correct_points_aff += cnn_ntg_correct_points_aff
# cnn_ntg_total_points += cnn_ntg_num_points
print('Batch: [{}/{} ({:.0f}%)]'.format(i, len(dataloader),
100. * i / len(dataloader)))
total_correct_points_aff = total_correct_points_aff.__float__()
# ntg_total_correct_points_aff = ntg_total_correct_points_aff.__float__()
# cnn_ntg_total_correct_points_aff = cnn_ntg_total_correct_points_aff.__float__()
PCK_aff = total_correct_points_aff / total_points
# ntg_PCK_aff=ntg_total_correct_points_aff/ntg_total_points
# cnn_ntg_PCK_aff=cnn_ntg_total_correct_points_aff/cnn_ntg_total_points
print('PCK affine:', PCK_aff)
# print('ntg_PCK affine:',ntg_PCK_aff)
# print('cnn_ntg_PCK affine:',cnn_ntg_PCK_aff)
print('Done!')