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
torch.backends.cudnn.deterministic = True
random.seed(seed)
torch.manual_seed(seed)
np.random.seed(seed)
###############################################################################
# Build the model
###############################################################################
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