def forward(self, im_data, batch=2, allInfer=False):
'''
@param im_data [B, 1, H, W]
@return im_topk [B, 1, H, W]
@return kpts [[N, 4] for B] (B, 0, H, W)
'''
if allInfer:
im_data = im_data
im_rawsc, _, _ = self.det(im_data)
im_score = self.det.process(im_rawsc)[0]
im_topk = topk_map(im_score, self.TOPK).permute(0, 3, 1, 2) # B, 1, H, W
kpts = im_topk.nonzero() # (B*topk, 4)
kpts = [kpts[kpts[:, 0] == idx, :] for idx in range(im_data.shape[0])] # [[N, 4] for B]
im_topk = im_topk.float()
else:
im_topK_ = []
kpts_ = []
for j in range(0, im_data.shape[0], batch):
im_data_clip = im_data[j:j+batch]
im_rawsc, _, _ = self.det(im_data_clip)
im_score = self.det.process(im_rawsc)[0]
im_topk = topk_map(im_score, self.TOPK).permute(0, 3, 1, 2) # B, 1, H, W
kpts = im_topk.nonzero() # (B*topk, 4)
kpts = [kpts[kpts[:, 0] == idx, :] for idx in range(im_data_clip.shape[0])] # [[N, 4] for B]
im_topk = im_topk.float()
im_topK_.append(im_topk)
kpts_ = kpts_ + kpts
kpts = kpts_
im_topk = torch.cat(im_topK_, 0)
return im_topk, kpts # B, 1, H, W; N, 4;
def inference(self, im_data, im_info, im_raw):
# import time
# start_time = time.time()
# for i in range(100):
# im_rawsc, im_scale, im_orint = self.det(im_data)
# end_time = time.time()
# print((end_time-start_time)/100)
im_rawsc, im_scale, im_orint = self.det(im_data)
im_score = self.det.process(im_rawsc)[0]
im_topk = topk_map(im_score, self.TOPK)
kpts = im_topk.nonzero() # (B*topk, 4)
cos, sim = im_orint.squeeze().chunk(chunks=2, dim=-1)
cos = cos.masked_select(im_topk.to(torch.bool)) # (B*topk)
sim = sim.masked_select(im_topk.to(torch.bool)) # (B*topk)
im_orint = torch.cat((cos.unsqueeze(-1), sim.unsqueeze(-1)), dim=-1)
im_patches = clip_patch(
kpts,
im_scale.masked_select(im_topk.to(torch.bool)),
im_orint,
im_info,
im_raw,
PSIZE=self.PSIZE,
) # (numkp, 1, 32, 32)
# start_time = time.time()
# for i in range(100):
# im_des = self.des(im_patches)
# end_time = time.time()
# print((end_time - start_time) / 100)
im_des = self.des(im_patches)
# import pdb
# pdb.set_trace()
# from matplotlib import pyplot as plt
# plt.imshow(scoremap_gaus[0, , :, :].cpu().detach().numpy())
return im_scale, kpts, im_des ,im_scale.masked_select(im_topk.to(torch.bool)), im_orint
def process(self, im1w_score):
"""
nms(n), topk(t), gaussian kernel(g) operation
:param im1w_score: warped score map
:return: processed score map, topk mask, topk value
"""
im1w_score = filter_border(im1w_score)
# apply nms to im1w_score
nms_mask = nms(im1w_score,
thresh=self.NMS_THRESH,
ksize=self.NMS_KSIZE)
im1w_score = im1w_score * nms_mask
topk_value = im1w_score
# apply topk to im1w_score
topk_mask = topk_map(im1w_score, self.TOPK)
im1w_score = topk_mask.to(torch.float) * im1w_score
# apply gaussian kernel to im1w_score
psf = get_gauss_filter_weight(
self.GAUSSIAN_KSIZE,
self.GAUSSIAN_SIGMA)[None, None, :, :].to(im1w_score.device)
# psf = im1w_score.new_tensor(
# get_gauss_filter_weight(self.GAUSSIAN_KSIZE, self.GAUSSIAN_SIGMA)[
# None, None, :, :
# ]
# )
im1w_score = F.conv2d(
input=im1w_score.permute(0, 3, 1, 2),
weight=psf,
stride=1,
padding=self.GAUSSIAN_KSIZE // 2,
).permute(0, 2, 3, 1) # (B, H, W, 1)
"""
apply tf.clamp to make sure all value in im1w_score isn't greater than 1
but this won't happend in correct way
"""
im1w_score = im1w_score.clamp(min=0.0, max=1.0)
return im1w_score, topk_mask, topk_value
def inference(self, im_data, im_info, im_raw):
im_rawsc, im_scale, im_orint = self.det(im_data)
im_score = self.det.process(im_rawsc)[0]
im_topk = topk_map(im_score, self.TOPK)
kpts = im_topk.nonzero() # (B*topk, 4)
cos, sim = im_orint.squeeze().chunk(chunks=2, dim=-1)
cos = cos.masked_select(im_topk) # (B*topk)
sim = sim.masked_select(im_topk) # (B*topk)
im_orint = torch.cat((cos.unsqueeze(-1), sim.unsqueeze(-1)), dim=-1)
im_patches = clip_patch(
kpts,
im_scale.masked_select(im_topk),
im_orint,
im_info,
im_raw,
PSIZE=self.PSIZE,
) # (numkp, 1, 32, 32)
im_des = self.des(im_patches)
return im_scale.masked_select(im_topk), kpts, im_des, im_orint
def process(self, im1w_score):
"""
nms(n), topk(t), gaussian kernel(g) operation
:param im1w_score: warped score map
:return: processed score map, topk mask, topk value
"""
im1w_score = filter_border(im1w_score)
# apply nms to im1w_score
#nms:non-maximum suppression非最大值抑制
nms_mask = nms(im1w_score,
thresh=self.NMS_THRESH,
ksize=self.NMS_KSIZE)
im1w_score = im1w_score * nms_mask
topk_value = im1w_score
# apply topk to im1w_score
topk_mask = topk_map(im1w_score, self.TOPK)
im1w_score = topk_mask.to(torch.float) * im1w_score
# apply gaussian kernel to im1w_score
psf = im1w_score.new_tensor(
get_gauss_filter_weight(self.GAUSSIAN_KSIZE,
self.GAUSSIAN_SIGMA)[None, None, :, :])
im1w_score = F.conv2d(
input=im1w_score.permute(0, 3, 1, 2),
weight=psf,
stride=1,
padding=self.GAUSSIAN_KSIZE // 2,
).permute(0, 2, 3, 1) # (B, H, W, 1)
#python pytorch.permute函数用于变换参数的维数，比如这里的F.conv2d的参数有4个，现在将里面参数的顺序调整为0 2 3 1
"""
apply tf.clamp to make sure all value in im1w_score isn't greater than 1
but this won't happend in correct way
"""
im1w_score = im1w_score.clamp(min=0.0, max=1.0)
#python torch clamp函数，用于将参数的取值范围限制在某个范围内，相当于把数值夹在某个区间上
#如果clamp函数里面参数小于等于最小值，那么这个参数设置为最小值；如果这个参数大于最大值，那么这个参数设置为最大值
#如果clamp函数里面参数介于最大值与最小值之间，那么就取这个参数 本身的值
return im1w_score, topk_mask, topk_value
def inference(im_data, im_info, im_raw):
score_maps, orint_maps = det(im_data)
im_rawsc, im_scale, im_orint = handle_det_out(score_maps, orint_maps,
det.scale_list,
det.score_com_strength,
det.scale_com_strength)
im_score = det.process(im_rawsc)[0]
im_topk = topk_map(im_score, TOPK)
kpts = im_topk.nonzero() # (B*topk, 4)
cos, sim = im_orint.squeeze().chunk(chunks=2, dim=-1)
cos = cos.masked_select(im_topk) # (B*topk)
sim = sim.masked_select(im_topk) # (B*topk)
im_orint = torch.cat((cos.unsqueeze(-1), sim.unsqueeze(-1)), dim=-1)
im_patches = clip_patch(
kpts,
im_scale.masked_select(im_topk),
im_orint,
im_info,
im_raw,
PSIZE=PSIZE,
) # (numkp, 1, 32, 32)
im_patches = des.input_norm(im_patches)
im_des = des(im_patches)
return im_scale, kpts, im_des
