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
if __name__ == '__main__':
# filename = 'configs/magicpoint_shapes_subpix.yaml'
filename = 'configs/magicpoint_repeatability.yaml'
import yaml
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_default_tensor_type(torch.FloatTensor)
with open(filename, 'r') as f:
config = yaml.load(f)
task = config['data']['dataset']
# data loading
from utils.loader import dataLoader_test as dataLoader
data = dataLoader(config, dataset='hpatches')
test_set, test_loader = data['test_set'], data['test_loader']
# take one sample
for i, sample in tqdm(enumerate(test_loader)):
if i > 1: break
val_agent = Val_model_subpixel(config['subpixel'], device=device)
val_agent.loadModel()
# points from heatmap
img = sample['image']
print("image: ", img.shape)
points = torch.tensor([[1, 2], [3, 4]])
def points_to_4d(points):
num_of_points = points.shape[0]
# load config
filename = "configs/superpoint_coco_train_heatmap.yaml"
import yaml
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_default_tensor_type(torch.FloatTensor)
with open(filename, "r") as f:
config = yaml.load(f)
from utils.loader import dataLoader as dataLoader
# data = dataLoader(config, dataset='hpatches')
task = config["data"]["dataset"]
data = dataLoader(config, dataset=task, warp_input=True)
# test_set, test_loader = data['test_set'], data['test_loader']
train_loader, val_loader = data["train_loader"], data["val_loader"]
# model_fe = Train_model_frontend(config)
# print('==> Successfully loaded pre-trained network.')
train_agent = Train_model_heatmap(config, device=device)
train_agent.train_loader = train_loader
# train_agent.val_loader = val_loader
train_agent.loadModel()
train_agent.dataParallel()
train_agent.train()
if __name__ == '__main__':
# filename = 'configs/magicpoint_shapes_subpix.yaml'
filename = 'configs/magicpoint_shapes_pair.yaml'
import yaml
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
torch.set_default_tensor_type(torch.FloatTensor)
with open(filename, 'r') as f:
config = yaml.load(f)
task = config['data']['dataset']
# data loading
from utils.loader import dataLoader_test as dataLoader
data = dataLoader(config, dataset='SyntheticDataset_gaussian')
test_set, test_loader = data['test_set'], data['test_loader']
# load frontend
val_agent = Val_model_heatmap(config['model'], device=device)
# take one sample
for i, sample in tqdm(enumerate(test_loader)):
if i > 10: break
val_agent.loadModel()
# points from heatmap
img = sample['image']
img = np.transpose(img, (0, 3, 1, 2))
heatmap_batch = val_agent.run(
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 export_detector_homoAdapt_gpu(config, output_dir, args):
"""
input 1 images, output pseudo ground truth by homography adaptation.
Save labels:
pred:
'prob' (keypoints): np (N1, 3)
"""
from utils.utils import pltImshow
from utils.utils import saveImg
from utils.draw import draw_keypoints
# basic setting
task = config["data"]["dataset"]
export_task = config["data"]["export_folder"]
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, exper_name=args.exper_name, date=True)
)
## parameters
nms_dist = config["model"]["nms"] # 4
top_k = config["model"]["top_k"]
homoAdapt_iter = config["data"]["homography_adaptation"]["num"]
conf_thresh = config["model"]["detection_threshold"]
nn_thresh = 0.7
outputMatches = True
count = 0
max_length = 5
output_images = args.outputImg
check_exist = True
## save data
save_path = Path(output_dir)
save_output = save_path
save_output = save_output / "predictions" / export_task
save_path = save_path / "checkpoints"
logging.info("=> will save everything to {}".format(save_path))
os.makedirs(save_path, exist_ok=True)
os.makedirs(save_output, exist_ok=True)
# data loading
from utils.loader import dataLoader_test as dataLoader
data = dataLoader(config, dataset=task, export_task=export_task)
print("Data is: ",data)
test_set, test_loader = data["test_set"], data["test_loader"]
print("Size test: ",len(test_set))
print("Size loader: ",len(test_loader))
# model loading
## load pretrained
try:
path = config["pretrained"]
print("==> Loading pre-trained network.")
print("path: ", path)
# This class runs the SuperPoint network and processes its outputs.
fe = SuperPointFrontend_torch(
config=config,
weights_path=path,
nms_dist=nms_dist,
conf_thresh=conf_thresh,
nn_thresh=nn_thresh,
cuda=False,
device=device,
)
print("==> Successfully loaded pre-trained network.")
fe.net_parallel()
print(path)
# save to files
save_file = save_output / "export.txt"
with open(save_file, "a") as myfile:
myfile.write("load model: " + path + "\n")
except Exception:
print(f"load model: {path} failed! ")
raise
def load_as_float(path):
return imread(path).astype(np.float32) / 255
print("Tracker")
tracker = PointTracker(max_length, nn_thresh=fe.nn_thresh)
with open(save_file, "a") as myfile:
myfile.write("homography adaptation: " + str(homoAdapt_iter) + "\n")
print("Load save file")
'''
print(len(test_loader))
for i,sample in enumerate(test_loader):
print("Hello world")
print("Img: ",sample["image"].size())
print("Name: ",test_set[i]["name"])
print("valid mask: ",test_set[i]["valid_mask"].size())
print("valid img_2D: ",test_set[i]["image_2D"].size())
print("valid mask: ",test_set[i]["valid_mask"].size())
print("homograpgy: ",test_set[i]["homographies"].size())
print("inverse: ",test_set[i]["inv_homographies"].size())
print("scene name: ",test_set[i]["scene_name"])
print()
'''
## loop through all images
for i, sample in tqdm(enumerate(test_loader)):
img, mask_2D = sample["image"], sample["valid_mask"]
img = img.transpose(0, 1)
img_2D = sample["image_2D"].numpy().squeeze()
mask_2D = mask_2D.transpose(0, 1)
inv_homographies, homographies = (
sample["homographies"],
sample["inv_homographies"],
)
img, mask_2D, homographies, inv_homographies = (
img.to(device),
mask_2D.to(device),
homographies.to(device),
inv_homographies.to(device),
)
# sample = test_set[i]
name = sample["name"][0]
logging.info(f"name: {name}")
if check_exist:
p = Path(save_output, "{}.npz".format(name))
if p.exists():
logging.info("file %s exists. skip the sample.", name)
continue
print("Pass img to network")
# pass through network
heatmap = fe.run(img, onlyHeatmap=True, train=False)
outputs = combine_heatmap(heatmap, inv_homographies, mask_2D, device=device)
pts = fe.getPtsFromHeatmap(outputs.detach().cpu().squeeze()) # (x,y, prob)
# subpixel prediction
if config["model"]["subpixel"]["enable"]:
fe.heatmap = outputs # tensor [batch, 1, H, W]
print("outputs: ", outputs.shape)
print("pts: ", pts.shape)
pts = fe.soft_argmax_points([pts])
pts = pts[0]
## top K points
pts = pts.transpose()
print("total points: ", pts.shape)
print("pts: ", pts[:5])
if top_k:
if pts.shape[0] > top_k:
pts = pts[:top_k, :]
print("topK filter: ", pts.shape)
## save keypoints
pred = {}
pred.update({"pts": pts})
## - make directories
filename = str(name)
if task == "Kitti" or "Kitti_inh":
scene_name = sample["scene_name"][0]
os.makedirs(Path(save_output, scene_name), exist_ok=True)
path = Path(save_output, "{}.npz".format(filename))
np.savez_compressed(path, **pred)
## output images for visualization labels
if output_images:
img_pts = draw_keypoints(img_2D * 255, pts.transpose())
f = save_output / (str(count) + ".png")
if task == "Coco" or "Kitti":
f = save_output / (name + ".png")
saveImg(img_pts, str(f))
count += 1
print("output pseudo ground truth: ", count)
save_file = save_output / "export.txt"
with open(save_file, "a") as myfile:
myfile.write("Homography adaptation: " + str(homoAdapt_iter) + "\n")
myfile.write("output pairs: " + str(count) + "\n")
pass
def export_descriptor(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")
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 data
from pathlib import Path
# save_path = save_path_formatter(config, output_dir)
save_path = Path(output_dir)
save_output = save_path
save_output = save_output / 'predictions'
save_path = save_path / 'checkpoints'
logging.info('=> will save everything to {}'.format(save_path))
os.makedirs(save_path, exist_ok=True)
os.makedirs(save_output, exist_ok=True)
# data loading
from utils.loader import dataLoader_test as dataLoader
data = dataLoader(config, dataset='hpatches')
test_set, test_loader = data['test_set'], data['test_loader']
from utils.print_tool import datasize
datasize(test_loader, config, tag='test')
from imageio import imread
def load_as_float(path):
return imread(path).astype(np.float32) / 255
def squeezeToNumpy(tensor_arr):
return tensor_arr.detach().cpu().numpy().squeeze()
outputMatches = True
count = 0
max_length = 5
method = config['model']['method']
# tracker = PointTracker(max_length, nn_thresh=fe.nn_thresh)
# for sample in tqdm(enumerate(test_loader)):
for i, sample in tqdm(enumerate(test_loader)):
img_0, img_1 = sample['image'], sample['warped_image']
imgs_np, imgs_fil = [], []
# first image, no matches
imgs_np.append(img_0.numpy().squeeze())
imgs_np.append(img_1.numpy().squeeze())
# H, W = img.shape[1], img.shape[2]
# img = img.view(1,1,H,W)
##### add opencv functions here #####
def classicalDetectors(image, method='sift'):
"""
# sift keyframe detectors and descriptors
"""
image = image*255
round_method = False
if round_method == True:
from models.classical_detectors_descriptors import classical_detector_descriptor # with quantization
points, desc = classical_detector_descriptor(image, **{'method': method})
y, x = np.where(points)
# pnts = np.stack((y, x), axis=1)
pnts = np.stack((x, y), axis=1) # should be (x, y)
## collect descriptros
desc = desc[y, x, :]
else:
# sift with subpixel accuracy
from models.classical_detectors_descriptors import SIFT_det as classical_detector_descriptor
pnts, desc = classical_detector_descriptor(image, image)
print("desc shape: ", desc.shape)
return pnts, desc
pts_list = []
pts, desc_1 = classicalDetectors(imgs_np[0], method=method)
pts_list.append(pts)
print("total points: ", pts.shape)
'''
pts: list [numpy (N, 2)]
desc: list [numpy (N, 128)]
'''
# save keypoints
pred = {}
pred.update({
# 'image_fil': imgs_fil[0],
'image': imgs_np[0],
})
pred.update({'prob': pts,
'desc': desc_1})
# second image, output matches
pred.update({
'warped_image': imgs_np[1],
# 'warped_image_fil': imgs_fil[1],
})
pts, desc_2 = classicalDetectors(imgs_np[1], method=method)
pts_list.append(pts)
# if outputMatches == True:
# tracker.update(pts, desc)
# pred.update({'matches': matches.transpose()})
print("total points: ", pts.shape)
pred.update({'warped_prob': pts,
'warped_desc': desc_2,
'homography': squeezeToNumpy(sample['homography'])
})
## get matches
data = get_sift_match(sift_kps_ii=pts_list[0], sift_des_ii=desc_1, sift_kps_jj=pts_list[1], sift_des_jj=desc_2, if_BF_matcher=True)
matches = data['match_quality_good']
print(f"matches: {matches.shape}")
matches_all = data['match_quality_all']
pred.update({
'matches': matches,
'matches_all': matches_all
})
# clean last descriptor
'''
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)
'''
# save data
from pathlib import Path
filename = str(count)
path = Path(save_output, '{}.npz'.format(filename))
np.savez_compressed(path, **pred)
count += 1
print("output pairs: ", count)
save_file = save_output / "export.txt"
with open(save_file, "a") as myfile:
myfile.write("output pairs: " + str(count) + '\n')
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 train_good(config, output_dir, args):
"""
# training script, controlled by config file and args
# work with Train_model_pipeline.py
params:
config: config file path, contain the settings
output_dir: the path for results
args: some setting
"""
# config
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, exper_name=args.exper_name,
date=True))
logging.info(f"config: {config}")
## reproducibility
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 paths
save_path = Path(output_dir)
save_path = save_path / "checkpoints"
logging.info("+++[Train]+++ will save everything to {}".format(save_path))
os.makedirs(save_path, exist_ok=True)
# data loading
assert (config["data"]["sequence_length"] == 2
), "Sorry dude, we are only supporting two-frame setting for now."
# assert (config['data']['read_what']['with_X'] and config['model']['batch_size']==1) or (not config['data']['read_what']['with_X']), 'We are not suppoting batching lidar Xs with batch_size>1 yet!'
config["data"]["read_what"]["with_quality"] = config["model"]["if_quality"]
val = "test" if args.test else "val" # use test or val data
val_shuffle = False if args.test else True # not sorting when doing testing
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, val_loader = data["train_loader"], data["val_loader"]
logging.info(
"+++[Dataset]+++ train split size %d in %d batches, val split size %d in %d batches"
% (
len(train_loader) * config["data"]["batch_size"],
len(train_loader),
len(val_loader) * config["data"]["batch_size"],
len(val_loader),
))
# model loading
# model_params = {'depth': config['model']['depth'], 'clamp_at':config['model']['clamp_at'], 'in_channels': 4, 'num_seg_classes': 1, 'with_transform': False, 'with_instance_norm': True}
if config["model"]["if_SP"]:
config["model"]["quality_size"] = 1
img_zoom_xy = (
config["data"]["preprocessing"]["resize"][1] /
config["data"]["image"]["size"][1],
config["data"]["preprocessing"]["resize"][0] /
config["data"]["image"]["size"][0],
)
model_params = {
"depth": config["model"]["depth"],
"img_zoom_xy": img_zoom_xy,
"image_size": config["data"]["image"]["size"],
"quality_size": config["model"]["quality_size"],
"if_quality": config["model"]["if_quality"],
"if_img_des_to_pointnet": config["model"]["if_img_des_to_pointnet"],
"if_goodCorresArch": config["model"]["if_goodCorresArch"],
"if_img_feat": config["model"]["if_img_feat"],
"if_cpu_svd": config["model"]["if_cpu_svd"],
"if_learn_offsets": config["model"]["if_learn_offsets"],
"if_tri_depth": config["model"]["if_tri_depth"],
"if_sample_loss": config["model"]["if_sample_loss"],
}
## load model and weights - deep fundametal network (deepF)
net = modelLoader(config["model"]["name"], **model_params)
print(f"deepF net: {net}")
n_iter = 0
n_iter_val = 0 + n_iter
## load model and weights - superpoint (sp)
if config["model"]["if_SP"]:
SP_params = {
"out_num_points": 2000, ### no use
"patch_size": 5,
"device": device,
"nms_dist": 4,
"conf_thresh": 0.015,
"nn_thresh": 0.7,
}
params = config["training"].get("SP_params", None)
if params is not None:
SP_params.update(params)
else:
logging.warning(f"use default Superpoint Parameters")
# for e in list(params):
# if e != 'device':
# params[e] = float(params[e])
logging.info(f"SP_params: {SP_params}")
# checkpoint_path_SP = "logs/superpoint_coco_heat2_0/checkpoints/superPointNet_170000_checkpoint.pth.tar"
checkpoint_path_SP = config["training"]["pretrained_SP"]
checkpoint_mode_SP = "" if checkpoint_path_SP[-3:] == "pth" else "full"
# helping modules
SP_processer = SuperPointNet_process(**SP_params)
SP_tracker = PointTracker(max_length=2, nn_thresh=params["nn_thresh"])
# load the network
net_SP = SuperPointNet_gauss2()
n_iter_SP = 0
n_iter_val_SP = 0 + n_iter_SP
## load pretrained and create optimizer
net_SP, optimizer_SP, n_iter_SP, n_iter_val_SP = prepare_model(
config,
net_SP,
device,
n_iter_SP,
n_iter_val_SP,
net_postfix="_SP")
# net_SP = nn.DataParallel(net_SP) # AttributeError: 'DataParallel' object has no attribute 'process_output'
if config["training"].get("train_SP", True):
logging.info("+++[Train]+++ training superpoint")
else:
logging.info("+++[Train]+++ superpoint is used but not trained")
## load pretrained and create optimizer
net, optimizer, n_iter, n_iter_val = prepare_model(config,
net,
device,
n_iter,
n_iter_val,
net_postfix="")
if config["training"].get("train", True):
logging.info("+++[Train]+++ training deepF model")
else:
logging.info("+++[Train]+++ deepF model is used but not trained")
epoch = n_iter // len(train_loader)
# set up train_agent
train_agent = Train_model_pipeline(config,
save_path=save_path,
args=args,
device=device)
train_agent.writer = writer
# feed the data into the agent
train_agent.train_loader = train_loader
train_agent.val_loader = val_loader
train_agent.set_params(n_iter=n_iter, epoch=epoch, n_iter_val=n_iter_val)
if not config["model"]["if_SP"]:
net_SP = None
optimizer_SP = None
SP_processer = None
SP_tracker = None
train_agent.set_nets(net, net_SP=net_SP)
train_agent.set_optimizers(optimizer, optimizer_SP)
train_agent.set_SP_helpers(SP_processer, SP_tracker)
while True:
# Train for one epoch; val occasionally
epoch_loss, _, n_iter, n_iter_val = train_agent.train_epoch(train=True)
save_file = save_path / "training.txt"
# saveLoss(save_file, n_iter, epoch_loss)
if n_iter > config["training"]["train_iter"]:
break
print("Finished Training")
def simple_train(config, output_dir, args):
batch_size = config['training']['batch_size_train']
epochs = config['training']['epoch_train']
learning_rate = config['training']['learning_rate']
savepath = os.path.join(output_dir, 'checkpoints')
os.makedirs(savepath, exist_ok=True)
with open(os.path.join(output_dir, 'config.yml'), 'w') as f:
yaml.dump(config, f, default_flow_style=False)
# Prepare for data loader
data = dataLoader(config,
dataset=config['data']['dataset'],
warp_input=True)
trainloader, valloader = data['train_loader'], data['val_loader']
datasize(trainloader, config, tag='train')
datasize(valloader, config, tag='val')
# Prepare for model
model = UnSuperPoint(config)
model.train()
dev = torch.device("cuda") if torch.cuda.is_available() else torch.device(
'cpu')
model.to(dev)
# Prepare for tensorboard writer
model.writer = SummaryWriter(
getWriterPath(task=args.command,
exper_name=args.export_name,
date=True))
# Prepare for optimizer
# model.optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9)
model.optimizer = optim.Adam(model.parameters(),
lr=0.001,
betas=(0.9, 0.999))
whole_step = 0
total = len(trainloader)
try:
if config['data']['curriculum']:
maxX = trainloader.dataset.config['homographies'][
'perspective_amplitude_x']
maxY = trainloader.dataset.config['homographies'][
'perspective_amplitude_y']
for epoch in tqdm(range(1, epochs + 1), desc='epoch'):
if config['data']['curriculum']:
trainloader.dataset.config['homographies'][
'perspective_amplitude_x'] = epoch * maxX / epochs
trainloader.dataset.config['homographies'][
'perspective_amplitude_y'] = epoch * maxY / epochs
tqdm.write("Max x Perspective: {:.6f}".format(
trainloader.dataset.config['homographies']
['perspective_amplitude_x']))
tqdm.write("Max y Perspective: {:.6f}".format(
trainloader.dataset.config['homographies']
['perspective_amplitude_y']))
for batch_idx, (img0, img1, mat) in tqdm(enumerate(trainloader),
desc='step',
total=total):
whole_step += 1
model.step = whole_step
loss = model.train_val_step(img0, img1, mat, 'train')
tqdm.write('Loss: {:.6f}'.format(loss))
if whole_step % config['save_interval'] == 0:
torch.save(
model.state_dict(),
os.path.join(
savepath, config['model']['name'] +
'_{}.pkl'.format(whole_step)))
if args.eval and whole_step % config[
'validation_interval'] == 0:
for j, (img0, img1, mat) in enumerate(valloader):
model.train_val_step(img0, img1, mat, 'valid')
if j > config['training'].get("step_val", 4):
break
torch.save(
model.state_dict(),
os.path.join(
savepath,
config['model']['name'] + '_{}.pkl'.format(whole_step)))
except KeyboardInterrupt:
print("press ctrl + c, save model!")
torch.save(
model.state_dict(),
os.path.join(
savepath,
config['model']['name'] + '_{}.pkl'.format(whole_step)))
pass
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)
