def __init__(self, do_cuda=True):
self.lock = RLock()
self.opts = rfnetOptions(do_cuda)
print(self.opts)
print('rfnet init')
#print(f"{gct()} model init")
print("model init")
det = RFDetSO(
cfg.TRAIN.score_com_strength,
cfg.TRAIN.scale_com_strength,
cfg.TRAIN.NMS_THRESH,
cfg.TRAIN.NMS_KSIZE,
cfg.TRAIN.TOPK,
cfg.MODEL.GAUSSIAN_KSIZE,
cfg.MODEL.GAUSSIAN_SIGMA,
cfg.MODEL.KSIZE,
cfg.MODEL.padding,
cfg.MODEL.dilation,
cfg.MODEL.scale_list,
)
des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH)
model = RFNetSO(
det, des, cfg.LOSS.SCORE, cfg.LOSS.PAIR, cfg.PATCH.SIZE, cfg.TRAIN.TOPK
)
device = torch.device("cuda")
model = model.to(device)
#resume = args.resume
#resume = "/home/cviss3/PycharmProjects/gensynth_dev_env/pyslam/thirdparty/rfnet/runs/10_24_09_25/model/e121_NN_0.480_NNT_0.655_NNDR_0.813_MeanMS_0.649.pth.tar"
resume = "/content/RFnetpyslam/thirdparty/rfnet/runs/10_24_09_25/model/e121_NN_0.480_NNT_0.655_NNDR_0.813_MeanMS_0.649.pth.tar"
print('==> Loading pre-trained network.')
checkpoint = torch.load(resume)
model.load_state_dict(checkpoint["state_dict"])
self.fe = model
print('==> Successfully loaded pre-trained network.')
self.device = device
self.pts = []
self.kps = []
self.des = []
self.img = []
self.heatmap = []
self.frame = None
self.frameFloat = None
self.keypoint_size = 20 # just a representative size for visualization and in order to convert extracted points to cv2.KeyPoint
###############################################################################
print(f"{gct()} : Build the model")
det = RFDetSO(
cfg.TRAIN.score_com_strength,
cfg.TRAIN.scale_com_strength,
cfg.TRAIN.NMS_THRESH,
cfg.TRAIN.NMS_KSIZE,
cfg.TRAIN.TOPK,
cfg.MODEL.GAUSSIAN_KSIZE,
cfg.MODEL.GAUSSIAN_SIGMA,
cfg.MODEL.KSIZE,
cfg.MODEL.padding,
cfg.MODEL.dilation,
cfg.MODEL.scale_list,
)
des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH)
model = RFNetSO(det, des, cfg.LOSS.SCORE, cfg.LOSS.PAIR, cfg.PATCH.SIZE,
cfg.TRAIN.TOPK)
if mgpu:
model = torch.nn.DataParallel(model)
model = model.to(device=device)
###############################################################################
# Load train data
###############################################################################
PPT = [cfg.PROJ.TRAIN_PPT, (cfg.PROJ.TRAIN_PPT + cfg.PROJ.EVAL_PPT)]
print(f"{gct()} : Loading traning data")
train_data = DataLoader(
HpatchDataset(
data_type="train",
# cfg.MODEL.scale_list,
# cfg.LOSS.SCORE,
# cfg.LOSS.PAIR,
# )
#
# device = torch.device("cuda:0")
# model = det.to(device)
# checkpoint = torch.load(resume)
# det.load_state_dict(checkpoint["state_dict"])
#
# img = torch.from_numpy(img.transpose((2, 0, 1)))[None, :].to(
# device, dtype=torch.float
# )
# score_maps, orint_maps = det(img)
#
# traced_script_module = torch.jit.trace(det, img)
# traced_script_module.save("det.pt")
resume = "/home/wang/workspace/RFNET/runs/07_25_14_05/model/e001_NN_0.283_NNT_0.337_NNDR_0.692_MeanMS_0.437_des.pth.tar"
des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH)
device = torch.device("cuda")
des = des.to(device).eval()
checkpoint = torch.load(resume)
des.load_state_dict(checkpoint["state_dict"])
device = torch.device("cuda:0")
im_patches = torch.randn((512, 1, 32, 32)).to(device)
im_des = des(im_patches)
traced_script_module2 = torch.jit.trace(des, im_patches)
traced_script_module2.save("des.pt")
det = RFDet(
cfg.TRAIN.score_com_strength,
cfg.TRAIN.scale_com_strength,
cfg.TRAIN.NMS_THRESH,
cfg.TRAIN.NMS_KSIZE,
cfg.TRAIN.TOPK,
cfg.MODEL.GAUSSIAN_KSIZE,
cfg.MODEL.GAUSSIAN_SIGMA,
cfg.MODEL.KSIZE,
cfg.MODEL.padding,
cfg.MODEL.dilation,
cfg.MODEL.scale_list,
cfg.LOSS.SCORE,
cfg.LOSS.PAIR,
)
des = HardNetNeiMask(cfg.HARDNET.MARGIN, cfg.MODEL.COO_THRSH)
if mgpu:
det = torch.nn.DataParallel(det)
des = torch.nn.DataParallel(des)
det = det.to(device=device)
des = des.to(device=device)
###############################################################################
# Load train data
###############################################################################
PPT = [cfg.PROJ.TRAIN_PPT, (cfg.PROJ.TRAIN_PPT + cfg.PROJ.EVAL_PPT)]
print(f"{gct()} : Loading traning data")
train_data = DataLoader(
HpatchDataset(