def train_joint(config, output_dir, args):
assert 'train_iter' in config
# config
# from utils.utils import pltImshow
# from utils.utils import saveImg
print("Start")
torch.set_default_tensor_type(torch.FloatTensor)
task = config['data']['dataset']
print("Start 1")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info('train on device: %s', device)
print("Start 2:", config['front_end_model'])
with open(os.path.join(output_dir, 'config.yml'), 'w') as f:
yaml.dump(config, f, default_flow_style=False)
# writer = SummaryWriter(getWriterPath(task=args.command, date=True))
writer = SummaryWriter(
getWriterPath(task=args.command, exper_name=args.exper_name,
date=True))
print("Start 3")
## save data
save_path = get_save_path(output_dir)
print("Save path: ", save_path)
# data loading
# data = dataLoader(config, dataset='syn', warp_input=True)
print("Start 4: ", task)
data = dataLoader(config, dataset=task, warp_input=True)
print("Start 5")
train_loader, val_loader = data['train_loader'], data['val_loader']
print("Mid ")
datasize(train_loader, config, tag='train')
datasize(val_loader, config, tag='val')
# init the training agent using config file
# from train_model_frontend import Train_model_frontend
from utils.loader import get_module
train_model_frontend = get_module('', config['front_end_model'])
print("Front end is: ", train_model_frontend)
train_agent = train_model_frontend(config,
save_path=save_path,
device=device)
# writer from tensorboard
train_agent.writer = writer
# feed the data into the agent
train_agent.train_loader = train_loader
train_agent.val_loader = val_loader
# load model initiates the model and load the pretrained model (if any)
train_agent.loadModel()
train_agent.dataParallel()
try:
# train function takes care of training and evaluation
train_agent.train()
except KeyboardInterrupt:
print("press ctrl + c, save model!")
train_agent.saveModel()
pass
def simple_render(config, output_dir, args):
from utils.tools import squeezeToNumpy
import matplotlib.pyplot as plt
from utils.draw import plot_imgs
from utils.loader import get_save_path
savepath = get_save_path(output_dir)
savepath = savepath / "../rendered"
os.makedirs(savepath, exist_ok=True)
# Prepare for data loader
renderloader = renderLoader(config,
dataset=config['data']['dataset'],
warp_input=True)
steps = config['rendering']['steps']
# Denormalization & Save
whole_step = 0
for batch_idx, (img0, img1, mat) in tqdm(enumerate(renderloader),
desc='step',
total=steps):
whole_step += 1
img_0 = 255 * squeezeToNumpy(img0)[[2, 1, 0], :, :].transpose(
(1, 2, 0))
img_1 = 255 * squeezeToNumpy(img1)[[2, 1, 0], :, :].transpose(
(1, 2, 0))
plot_imgs([img_0.astype(np.uint8),
img_1.astype(np.uint8)],
titles=['img1', 'img2'],
dpi=200)
plt.title(str(batch_idx))
plt.tight_layout()
plt.savefig(savepath / (str(batch_idx) + '.png'),
dpi=300,
bbox_inches='tight')
if whole_step >= steps:
break
def export_descriptor(config, output_dir, args):
"""
# input 2 images, output keypoints and correspondence
save prediction:
pred:
'image': np(320,240)
'prob' (keypoints): np (N1, 2)
'desc': np (N2, 256)
'warped_image': np(320,240)
'warped_prob' (keypoints): np (N2, 2)
'warped_desc': np (N2, 256)
'homography': np (3,3)
'matches': np [N3, 4]
"""
from utils.loader import get_save_path
from utils.var_dim import squeezeToNumpy
# basic settings
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
logging.info("train on device: %s", device)
with open(os.path.join(output_dir, "config.yml"), "w") as f:
yaml.dump(config, f, default_flow_style=False)
writer = SummaryWriter(getWriterPath(task=args.command, date=True))
save_path = get_save_path(output_dir)
save_output = save_path / "../predictions"
os.makedirs(save_output, exist_ok=True)
## parameters
outputMatches = True
subpixel = config["model"]["subpixel"]["enable"]
patch_size = config["model"]["subpixel"]["patch_size"]
# data loading
from utils.loader import dataLoader_test as dataLoader
task = config["data"]["dataset"]
data = dataLoader(config, dataset=task)
test_set, test_loader = data["test_set"], data["test_loader"]
from utils.print_tool import datasize
datasize(test_loader, config, tag="test")
# model loading
from utils.loader import get_module
Val_model_heatmap = get_module("", config["front_end_model"])
## load pretrained
val_agent = Val_model_heatmap(config["model"], device=device)
val_agent.loadModel()
## tracker
tracker = PointTracker(max_length=2, nn_thresh=val_agent.nn_thresh)
###### check!!!
count = 0
for i, sample in tqdm(enumerate(test_loader)):
img_0, img_1 = sample["image"], sample["warped_image"]
# first image, no matches
# img = img_0
def get_pts_desc_from_agent(val_agent, img, device="cpu"):
"""
pts: list [numpy (3, N)]
desc: list [numpy (256, N)]
"""
heatmap_batch = val_agent.run(
img.to(device)
) # heatmap: numpy [batch, 1, H, W]
# heatmap to pts
pts = val_agent.heatmap_to_pts()
# print("pts: ", pts)
if subpixel:
pts = val_agent.soft_argmax_points(pts, patch_size=patch_size)
# heatmap, pts to desc
desc_sparse = val_agent.desc_to_sparseDesc()
# print("pts[0]: ", pts[0].shape, ", desc_sparse[0]: ", desc_sparse[0].shape)
# print("pts[0]: ", pts[0].shape)
outs = {"pts": pts[0], "desc": desc_sparse[0]}
return outs
def transpose_np_dict(outs):
for entry in list(outs):
outs[entry] = outs[entry].transpose()
outs = get_pts_desc_from_agent(val_agent, img_0, device=device)
pts, desc = outs["pts"], outs["desc"] # pts: np [3, N]
if outputMatches == True:
tracker.update(pts, desc)
# save keypoints
pred = {"image": squeezeToNumpy(img_0)}
pred.update({"prob": pts.transpose(), "desc": desc.transpose()})
# second image, output matches
outs = get_pts_desc_from_agent(val_agent, img_1, device=device)
pts, desc = outs["pts"], outs["desc"]
if outputMatches == True:
tracker.update(pts, desc)
pred.update({"warped_image": squeezeToNumpy(img_1)})
# print("total points: ", pts.shape)
pred.update(
{
"warped_prob": pts.transpose(),
"warped_desc": desc.transpose(),
"homography": squeezeToNumpy(sample["homography"]),
}
)
if outputMatches == True:
matches = tracker.get_matches()
print("matches: ", matches.transpose().shape)
pred.update({"matches": matches.transpose()})
print("pts: ", pts.shape, ", desc: ", desc.shape)
# clean last descriptor
tracker.clear_desc()
filename = str(count)
path = Path(save_output, "{}.npz".format(filename))
np.savez_compressed(path, **pred)
# print("save: ", path)
count += 1
print("output pairs: ", count)
def train_joint(config, output_dir, args):
assert 'train_iter' in config
# config
# from utils.utils import pltImshow
# from utils.utils import saveImg
torch.set_default_tensor_type(torch.FloatTensor)
task = config['data']['dataset']
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
logging.info('train on device: %s', device)
with open(os.path.join(output_dir, 'config.yml'), 'w') as f:
yaml.dump(config, f, default_flow_style=False)
# writer = SummaryWriter(getWriterPath(task=args.command, date=True))
writer = SummaryWriter(getWriterPath(task=args.command,
exper_name=args.exper_name, date=True))
## save data
save_path = get_save_path(output_dir)
# let us create Dataloaders
def _create_loader(dataset, bs=8, shuffle=False, n_workers=8):
return torch.utils.data.DataLoader(dataset,
batch_size=bs,
shuffle=shuffle,
pin_memory=True,
num_workers=n_workers,
)
splits = ["train", "val"]
data_loaders = {}
# create the dataset and dataloader classes
for split in splits:
dataset = PhototourismSuperpoint(split=split, **config["data"])
data_loaders[split] = _create_loader(dataset,
config["model"]["batch_size"],
split == "train",
config["training"]["workers_train"])
# data loading
train_loader, val_loader = data_loaders["train"], data_loaders["val"]
datasize(train_loader, config, tag='train')
datasize(val_loader, config, tag='val')
# init the training agent using config file
# from train_model_frontend import Train_model_frontend
from utils.loader import get_module
train_agent = TrainModelHeatmapMy(config,
save_path=save_path,
device=device)
# writer from tensorboard
train_agent.writer = writer
# feed the data into the agent
train_agent.train_loader = train_loader
train_agent.val_loader = val_loader
# load model initiates the model and load the pretrained model (if any)
train_agent.loadModel()
train_agent.dataParallel()
try:
# train function takes care of training and evaluation
train_agent.train_my()
except KeyboardInterrupt:
print ("press ctrl + c, save model!")
train_agent.saveModel()
pass
def val_feature(config, output_dir, args):
"""
1) input 2 images, output keypoints and correspondence
:param config:
:param output_dir:
:param args:
:return:
"""
# config
# device = torch.device("cpu")
from superpoint.utils.var_dim import squeezeToNumpy
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
logging.info("train on device: %s", device)
with open(os.path.join(output_dir, "config_sp.yml"), "w") as f:
yaml.dump(config, f, default_flow_style=False)
writer = SummaryWriter(getWriterPath(task=args.command, date=True))
if config["training"]["reproduce"]:
logging.info("reproduce = True")
torch.manual_seed(0)
np.random.seed(0)
print(
f"test random # : np({np.random.rand(1)}), torch({torch.rand(1)})")
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
## save data
from pathlib import Path
# save_path = save_path_formatter(config, output_dir)
from utils.loader import get_save_path
save_path = get_save_path(output_dir)
save_output = save_path / "../predictions"
os.makedirs(save_output, exist_ok=True)
# data loading
from utils.loader import dataLoader as dataLoader
# task = config['data']['dataset']
val = 'test' if args.test else 'val'
val_shuffle = False if args.test else True
train = False if args.test else True
data = dataLoader(config,
dataset=config["data"]["dataset"],
train=train,
warp_input=True,
val=val,
val_shuffle=val_shuffle)
train_loader, test_loader = data["train_loader"], data["val_loader"]
# data = dataLoader(config, dataset=config["data"]["dataset"], warp_input=True)
# train_loader, test_loader = data["train_loader"], data["val_loader"]
datasize(test_loader, config, tag="test")
# model loading
from utils.loader import get_module
Val_model_heatmap = get_module("", config["front_end_model"])
## load pretrained
val_agent = Val_model_heatmap(config["model"], device=device)
val_agent.loadModel()
## tracker
tracker = PointTracker(max_length=2, nn_thresh=val_agent.nn_thresh)
###### check!!!
outputMatches = True
count = 0
subpixel = config["model"]["subpixel"]["enable"]
patch_size = config["model"]["subpixel"]["patch_size"]
# print("Do subpixel!!!")
rand_noise = config["model"]["rand_noise"]["enable"]
from utils.eval_tools import Result_processor
result_dict_entry = [
"epi_dist_mean_gt",
"num_prob",
"num_warped_prob",
"num_matches",
"mscores",
]
result_processor = Result_processor(result_dict_entry)
for i, sample in tqdm(enumerate(test_loader)):
if config['training']['val_interval'] == -1:
pass
elif i > config['training']['val_interval']:
break
# imgs_grey_float = [img_grey.float().cuda() / 255.0 for img_grey in imgs_grey]
# img_0, img_1 = sample['imgs_grey'][0].unsqueeze(1), sample['imgs_grey'][1].unsqueeze(1)
img_0, img_1 = (
process_grey2tensor(sample["imgs_grey"][0], device=device),
process_grey2tensor(sample["imgs_grey"][1], device=device),
)
# first image, no matches
# img = img_0
outs = get_pts_desc_from_agent(val_agent,
img_0,
subpixel,
patch_size,
device=device)
pts, desc = outs["pts"], outs["desc"] # pts: np [3, N]
logging.info(f"pts: {pts.shape}, desc: {desc.shape}")
if rand_noise:
sigma = config["model"]["rand_noise"]["sigma"]
noise = np.random.normal(loc=0.0, scale=sigma, size=pts.shape)
pts[:2, :] += noise[:2, :]
print("pts: ", pts[:, :3])
if outputMatches == True:
tracker.update(pts, desc)
# save keypoints
pred = {"image": squeezeToNumpy(img_0)}
pred.update({"prob": pts.transpose(), "desc": desc.transpose()})
pred.update({"num_prob": pts.shape[1]})
logging.debug(f"1 -- pts: {pts.shape}, desc: {desc.shape}")
# second image, output matches
outs = get_pts_desc_from_agent(val_agent,
img_1,
subpixel,
patch_size,
device=device)
pts, desc = outs["pts"], outs["desc"]
# print(f"1 -- pts: {pts[:,:5]}, desc: {desc[:10, :5]}")
# print(f"2 -- pts: {pts[:,:5]}, desc: {desc[:10, :5]}")
if rand_noise:
sigma = config["model"]["rand_noise"]["sigma"]
noise = np.random.normal(loc=0.0, scale=sigma, size=pts.shape)
pts[:2, :] += noise[:2, :]
print("pts: ", pts[:, :3])
if outputMatches == True:
tracker.update(pts, desc)
pred.update({"warped_image": squeezeToNumpy(img_1)})
# print("total points: ", pts.shape)
pred.update({
"warped_prob": pts.transpose(),
"warped_desc": desc.transpose(),
# "homography": squeezeToNumpy(sample["homography"]),
})
logging.debug(f"2 -- pts: {pts.shape}")
pred.update({"num_warped_prob": pts.shape[1]})
# if subpixel:
# pts = subpixel_fix(pts)
if outputMatches == True:
matches = tracker.get_matches()
result = epi_dist_from_matches(matches.transpose()[np.newaxis,
...],
sample,
device,
five_point=False)
epi_dist_mean_gt = result["epi_dist_mean_gt"]
logging.debug(f"epi_dist_mean_gt: {epi_dist_mean_gt.shape}")
logging.info(
f"matches: {matches.transpose().shape}, num_prob: {pred['num_prob']}, num_warped_prob: {pred['num_warped_prob']}"
)
pred.update({"matches": matches.transpose()})
pred.update({"num_matches": matches.shape[1]})
pred.update({"epi_dist_mean_gt": epi_dist_mean_gt})
mscores = tracker.get_mscores()
# logging.info(f"tracker.get_mscores(): {mscores.shape}")
pred.update({"mscores": mscores})
"""
pred:
'image': np(320,240)
'prob' (keypoints): np (N1, 2)
'desc': np (N2, 256)
'warped_image': np(320,240)
'warped_prob' (keypoints): np (N2, 2)
'warped_desc': np (N2, 256)
'homography': np (3,3)
"""
# clean last descriptor
tracker.clear_desc()
# save data
from pathlib import Path
filename = str(count)
path = Path(save_output, "{}.npz".format(filename))
np.savez_compressed(path, **pred)
# print("save: ", path)
count += 1
# process results
result_processor.load_result(pred)
params = {"inlier_ratio": config['evaluation']['inlier_thd']}
# result_processor.output_result(['inlier_ratio'], **params)
result_processor.inlier_ratio("epi_dist_mean_gt",
params["inlier_ratio"],
if_print=True)
result_processor.save_result(Path(save_output, "result_dict_all.npz"),
"result_dict_all")
print(f"exp path: {save_output}, output pairs: {count}")
def simple_export(config, output_dir, args):
"""
# input 2 images, output keypoints and correspondence
save prediction:
pred:
'image': np(320,240)
'prob' (keypoints): np (N1, 2)
'desc': np (N1, 256)
'warped_image': np(320,240)
'warped_prob' (keypoints): np (N2, 2)
'warped_desc': np (N2, 256)
'homography': np (3,3)
'matches': np [N3, 4]
"""
from utils.loader import get_save_path
from utils.tools import squeezeToNumpy
# basic settings
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
logging.info("train on device: %s", device)
with open(os.path.join(output_dir, "config.yml"), "w") as f:
yaml.dump(config, f, default_flow_style=False)
save_path = get_save_path(output_dir)
save_output = save_path / "../predictions"
os.makedirs(save_output, exist_ok=True)
model_path = os.path.join(output_dir, 'checkpoints', args.model_name)
## parameters
outputMatches = True
# data loading
from utils.loader import testLoader as dataLoader
task = config["data"]["dataset"]
data = dataLoader(config, dataset=task)
test_set, test_loader = data["test_set"], data["test_loader"]
# model loading
# from utils.loader import get_module
# Val_model = get_module("", config["front_end_model"])
## load pretrained
# val_agent = Val_model()
val_agent = UnSuperPoint()
val_agent.load_state_dict(torch.load(model_path))
val_agent.to(val_agent.dev)
val_agent.train(False)
val_agent.task = 'test'
## tracker
tracker = PointTracker(max_length=2, nn_thresh=val_agent.nn_thresh)
###### check!!!
count = 0
total = len(test_loader)
for i, sample in tqdm(enumerate(test_loader), desc='scene', total=total):
img_0, img_1 = sample["image"], sample["warped_image"]
# first image, no matches
# img = img_0
@torch.no_grad()
def get_pts_desc_from_agent(val_agent, img, device="cpu"):
"""
pts: list [numpy (3, N)]
desc: list [numpy (256, N)]
"""
score, point, desc = val_agent.forward(
img.to(device)) # heatmap: numpy [batch, 1, H, W]
point = val_agent.get_position(point)
# heatmap, pts to pts, desc
pts, desc = val_agent.getPtsDescFromHeatmap(
point[0].cpu().numpy(), score[0].cpu().numpy(),
desc[0].cpu().numpy())
outs = {"pts": pts, "desc": desc}
return outs
def transpose_np_dict(outs):
for entry in list(outs):
outs[entry] = outs[entry].transpose()
outs = get_pts_desc_from_agent(val_agent, img_0, device=device)
pts, desc = outs["pts"], outs["desc"] # pts: np [3, N]
tqdm.write("pts A: {}, desc A: {}".format(pts.shape, desc.shape))
if outputMatches == True:
tracker.update(pts, desc)
# save keypoints
img_0 = squeezeToNumpy(img_0).transpose((1, 2, 0))
pred = {"image": ((img_0 / 0.225 + 0.5) * 255).astype(np.uint8)}
pred.update({"prob": pts.transpose(), "desc": desc.transpose()})
# second image, output matches
outs = get_pts_desc_from_agent(val_agent, img_1, device=device)
pts, desc = outs["pts"], outs["desc"]
tqdm.write("pts B: {}, desc B: {}".format(pts.shape, desc.shape))
if outputMatches == True:
tracker.update(pts, desc)
img_1 = squeezeToNumpy(img_1).transpose((1, 2, 0))
pred.update(
{"warped_image": ((img_1 / 0.225 + 0.5) * 255).astype(np.uint8)})
# print("total points: ", pts.shape)
pred.update({
"warped_prob": pts.transpose(),
"warped_desc": desc.transpose(),
"homography": squeezeToNumpy(sample["homography"]),
})
if outputMatches == True:
matches = tracker.get_matches()
tqdm.write("matches: {}".format(matches.transpose().shape))
pred.update({"matches": matches.transpose()})
# clean last descriptor
tracker.clear_desc()
filename = str(count)
path = Path(save_output, "{}.npz".format(filename))
np.savez_compressed(path, **pred)
# print("save: ", path)
count += 1
print("output pairs: ", count)