def get_matches(deses_SP):
from models.model_wrap import PointTracker
tracker = PointTracker(max_length=2, nn_thresh=1.2)
f = lambda x: x.cpu().detach().numpy()
# tracker = PointTracker(max_length=2, nn_thresh=1.2)
# print("deses_SP[1]: ", deses_SP[1].shape)
matching_mask = tracker.nn_match_two_way(f(deses_SP[0]).T, f(deses_SP[1]).T, nn_thresh=1.2)
return matching_mask
def getMatches(data):
from models.model_wrap import PointTracker
desc = data['desc']
warped_desc = data['warped_desc']
nn_thresh = 1.2
print("nn threshold: ", nn_thresh)
tracker = PointTracker(max_length=2, nn_thresh=nn_thresh)
# matches = tracker.nn_match_two_way(desc, warped_desc, nn_)
tracker.update(keypoints.T, desc.T)
tracker.update(warped_keypoints.T, warped_desc.T)
matches = tracker.get_matches().T
mscores = tracker.get_mscores().T
# mAP
# matches = data['matches']
print("matches: ", matches.shape)
print("mscores: ", mscores.shape)
try:
print("mscore max: ", mscores.max(axis=0))
print("mscore min: ", mscores.min(axis=0))
except ValueError:
pass
return matches, mscores
def dump_SP_match_idx(delta_ij, N_frames, dump_dir, save_npy, nn_threshes):
for nn_thresh, name in zip(nn_threshes, ['good', 'all']):
SP_matcher = PointTracker(max_length=2, nn_thresh=nn_thresh)
for ii in tqdm(range(N_frames-delta_ij)):
jj = ii + delta_ij
SP_kps_ii, SP_des_ii = load_SP(dump_dir, '%06d'%ii, ext='.npy' if save_npy else '.h5')
SP_kps_jj, SP_des_jj = load_SP(dump_dir, '%06d'%jj, ext='.npy' if save_npy else '.h5')
matches, scores = get_SP_match_idx_pair(SP_matcher, SP_kps_ii, SP_kps_jj, SP_des_ii, SP_des_jj)
dump_ij_match_quality_file = dump_dir/'SP_ij_match_quality_{}-{}'.format(ii, jj)
if save_npy:
# print(matches.shape, scores.shape)
match_quality = np.hstack((matches, scores)) # [[x1, y1, x2, y2, dist_good, ratio_good]]
np.save(dump_ij_match_quality_file+'_%s.npy'%name, match_quality)
else:
pass
def load_net_SP(self, name="net_SP"):
config = self.config
device = self.device
SP_params = {
"out_num_points": 2000,
"patch_size": 5,
"device": device,
"nms_dist": 4,
"conf_thresh": 0.015,
}
from models.model_utils import SuperPointNet_process
from models.model_wrap import PointTracker
from models.SuperPointNet_gauss2 import SuperPointNet_gauss2
from train_good_corr_4_vals_goodF_baseline import prepare_model
SP_processer = SuperPointNet_process(**SP_params)
SP_tracker = PointTracker(max_length=2, nn_thresh=1.2)
net_SP = SuperPointNet_gauss2()
net_SP, optimizer_SP, n_iter_SP, n_iter_val_SP = prepare_model(
config,
net_SP,
device,
n_iter=0,
n_iter_val=0,
net_postfix="_SP",
train=False,
)
logging.info("+++[Train]+++ training superpoint")
## put to class
self.net_SP_helper = {
"SP_processer": SP_processer,
"SP_tracker": SP_tracker
}
self.net_dict[name] = net_SP
pass
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