def _assign_priors(self, gt_boxes, gt_labels, corner_form_priors,
iou_threshold):
"""Assign ground truth boxes and targets to priors.
Args:
gt_boxes (num_targets, 4): ground truth boxes.
gt_labels (num_targets): labels of targets.
priors (num_priors, 4): corner form priors
Returns:
boxes (num_priors, 4): real values for priors.
labels (num_priros): labels for priors.
"""
# size: num_priors x num_targets
ious = iou_xywh(corner_form_priors, gt_boxes, type='N2N')
# size: num_priors
best_target_per_prior, best_target_per_prior_index = ious.max(1)
# size: num_targets
best_prior_per_target, best_prior_per_target_index = ious.max(0)
for target_index, prior_index in enumerate(best_prior_per_target_index):
best_target_per_prior_index[prior_index] = target_index
# 2.0 is used to make sure every target has a prior assigned
best_target_per_prior.index_fill_(0, best_prior_per_target_index, 2)
# size: num_priors
labels = gt_labels[best_target_per_prior_index]
labels[best_target_per_prior < iou_threshold] = 0 # the backgournd id
boxes = gt_boxes[best_target_per_prior_index]
return boxes, labels
def _ignore(self, pre_loc, lab_boxes):
# pylint: disable=no-self-use
"""
Count the ignore mask.
:param pre_loc:
:param lab_boxes:
:return:
"""
batch = pre_loc.shape[0]
_ignore_mask = []
for i in range(batch):
b_box = lab_boxes[i]
b_box = torch.stack(b_box, 0)
iou = iou_xywh(pre_loc[i], b_box, type='N2N_yolo')
iou_max = torch.max(iou, 0)
# print(iou_max[0][15, 31:34])
ignore_iou = torch.lt(iou_max[0], 0.6)
# print(ignore_iou[15, 31:34])
_ignore_mask.append(ignore_iou)
ignore_mask = torch.stack(_ignore_mask, 0)
# print(ignore_mask[0, 15, 31:34])
return ignore_mask
def NMS_Greedy(self, pre_score, pre_class, pre_loc):
"""
Nms.
:param pre_score: in the shape of [box_number, class_number]
:param pre_loc: int the shape of [box_number, 4] 4 means [x1, y1, x2, y2]
:param score_thresh:score_thresh
:param iou_thresh:iou_thresh
:return: labels_out
"""
# print('using Greedy NMS')
score_sort = pre_score.sort(descending=True) # sort the scores.
score_idx = score_sort[1][
score_sort[0] > self.score_thresh] # find the scores>0.7(thresh)
keep = []
for i in self.class_range: # with different classess.
a = pre_class[score_idx] == i # each class for NMS.
order_index = score_idx[a] # get the index of orders
while order_index.shape[0] > 0: # deal with all the boxes.
max_one = order_index[0].item(
) # get index of the max score box.
box_head = pre_loc[max_one] # get the score of it
box_others = pre_loc[order_index[1:]] # the rest boxes.
ious = iou_xywh(
box_head, box_others, type='N21'
) # count the ious between the max one and the others
rest = torch.lt(ious, self.iou_thresh).squeeze(
) # find the boxes of iou<0.5(thresh), discard the iou>0.5.
order_index = order_index[1:][
rest] # get the new index of the rest boxes, except the max one.
keep.append(max_one)
return keep
def NMS_old_edition(self, pre_score_raw, pre_loc, score_thresh,
iou_thresh):
"""
Nms.
:param pre_score: in the shape of [box_number, class_number]
:param pre_loc: int the shape of [box_number, 4] 4 means [x1, y1, x2, y2]
:param score_thresh:score_thresh
:param iou_thresh:iou_thresh
:return: labels_out
"""
pre_score_raw = pre_score_raw.max(-1)
pre_score = pre_score_raw[0]
pre_class = pre_score_raw[1]
idx = pre_score > score_thresh # Return the index of the max pre_score
print('max score:', torch.max(pre_score).item())
pre_score_max = pre_score[idx]
# print(pre_score[torch.gt(pre_score, pre_score_thre)])
print('Num of box:', len(pre_score_max))
if len(pre_score_max) > 1:
sor = pre_score_max.sort(descending=True)[1]
_idx = idx.nonzero()[sor]
_leng = len(_idx)
for i in range(_leng):
for j in range(i + 1, _leng):
if pre_class[_idx[i]] != pre_class[
_idx[j]]: # diffent classes
continue
if pre_score[_idx[i]] < score_thresh or \
pre_score[_idx[j]] < score_thresh:
continue # get out of the likely anchors which have been counted
box1 = pre_loc[_idx[i]].unsqueeze(0)
box2 = pre_loc[_idx[j]].unsqueeze(0)
iou_ = iou_xywh(box1, box2)
if iou_ > iou_thresh:
pre_score[_idx[j]] = 0.0
_idx = (pre_score > score_thresh
).nonzero() # Return the index of the max pre_score
labels_out = []
for keep_idx in _idx:
box = pre_loc[keep_idx]
box = box.squeeze()
boxx1 = box[0] - box[2] / 2
boxy1 = box[1] - box[3] / 2
boxx2 = box[0] + box[2] / 2
boxy2 = box[1] + box[3] / 2
box_out = [boxx1, boxy1, boxx2, boxy2]
if min(box_out) <= 0:
print('error!')
continue
pre_score_out = pre_score[keep_idx].item()
class_out = pre_class[keep_idx].item()
labels_out.append([pre_score_out, class_out, box_out])
return labels_out
def _reshape_labels(self, labels, grid_xy, shape, f_id):
"""
Reshape the labels.
:param labels: labels from training data
:param grid_xy: the matrix of the grid numbers
:return: labels_obj, labels_cls, lab_loc_xy, lab_loc_wh, labels_boxes, area_scal
"""
mask = np.arange(self.anc_num) + self.anc_num * f_id
anchors = self.anchors[mask]
labels_obj = torch.zeros([self.batch_size, shape[0], shape[1],
self.anc_num, 1]).to(self.cfg.TRAIN.DEVICE)
labels_loc = torch.zeros([self.batch_size, shape[0], shape[1],
self.anc_num, 4]).to(self.cfg.TRAIN.DEVICE)
labels_cls = torch.zeros([self.batch_size, shape[0], shape[1],
self.anc_num, self.cls_num]).to(self.cfg.TRAIN.DEVICE)
labels_boxes = []
for batch_idx, labs in enumerate(labels):
lab_boxes = []
for lab in labs:
lab = torch.Tensor(lab).to(self.cfg.TRAIN.DEVICE)
box_xy = (lab[1:3] + lab[3:5]) / 2
box_wh = (lab[3:5] - lab[1:3])
boxes = torch.cat([box_xy, box_wh])
box_center = (box_xy * torch.Tensor([shape[1], shape[0]]).to(self.cfg.TRAIN.DEVICE)).long()
anc = torch.cat([torch.zeros_like(anchors), anchors], 1)
box_iou = torch.cat([torch.Tensor([0, 0]).to(self.cfg.TRAIN.DEVICE), box_wh])
iou = iou_xywh(anc.to(self.cfg.TRAIN.DEVICE), box_iou.to(self.cfg.TRAIN.DEVICE), type='N21')
iou_max = torch.max(iou, 0)
anc_idx = iou_max[1].item()
lab_boxes.append(boxes)
# print(box_center, boxes)
labels_obj[batch_idx, box_center[1], box_center[0], anc_idx, 0] = 1
labels_loc[batch_idx, box_center[1], box_center[0], anc_idx] = boxes
labels_cls[batch_idx, box_center[1], box_center[0], anc_idx, lab[0].long()] = 1
labels_boxes.append(lab_boxes)
lab_loc_xy = labels_loc[..., 0:2]
lab_loc_wh = labels_loc[..., 2:4]
# count the scale of w*h, in order to count area_scal*wh
if self.multiply_area_scale:
area_scale = (1 - torch.sqrt(lab_loc_wh[..., 0] * lab_loc_wh[..., 1])).unsqueeze(-1).expand_as(lab_loc_wh)
else:
area_scale = 1.0
# print(area_scal.shape, area_scal[0, 14, 39, 12:])
lab_loc_xy = lab_loc_xy * torch.Tensor([shape[1], shape[0]]).to(self.cfg.TRAIN.DEVICE) - grid_xy
anchor_ch = anchors.view(1, 1, 1, self.anc_num, 2).expand(1, shape[0], shape[1], self.anc_num, 2).to(self.cfg.TRAIN.DEVICE)
lab_loc_wh = lab_loc_wh / anchor_ch
lab_loc_wh = torch.log(torch.clamp(lab_loc_wh, 1e-9, 1e9))
return labels_obj, labels_cls, lab_loc_xy, lab_loc_wh, labels_boxes, area_scale
def Soft_NMS(pre_score_raw, pre_loc, score_thresh, theta=0.5):
"""
Nms.
:param pre_score: in the shape of [box_number, class_number]
:param pre_loc: int the shape of [box_number, 4] 4 means [x1, y1, x2, y2]
:param score_thresh:score_thresh
:param iou_thresh:iou_thresh
:return: labels_out
"""
print('using Soft NMS')
class_num = pre_score_raw.shape[1]
pre_score_raw = pre_score_raw.max(-1)
pre_score = pre_score_raw[0]
pre_class = pre_score_raw[1]
score_sort = pre_score.sort(descending=True)
score_idx = score_sort[1][score_sort[0] > score_thresh]
keep = []
for i in range(class_num): # with different classess.
a = pre_class[score_idx] == i
order_index = score_idx[a]
while order_index.shape[0] > 0:
max_one = order_index[0].item()
keep.append(max_one)
box_head = pre_loc[max_one]
box_others = pre_loc[order_index[1:]]
score_others = pre_score[order_index[1:]]
# print(score_others)
ious = iou_xywh(box_head, box_others).reshape(1, -1)
# print('iou', ious)
soft_score = score_others * torch.exp(
-pow(ious, 2) / theta).squeeze() # s = s*e^(-iou^2 / theta)
# print('s',soft_score)
rest = torch.gt(soft_score, score_thresh)
order_index = order_index[1:][rest]
new_index = pre_score[order_index].sort(descending=True)
order_index = order_index[new_index[1]]
# ########### NMS UP ####################
labels_out = nms2labels(keep, pre_score, pre_loc, pre_class)
return labels_out
def NMS(pre_score_raw, pre_loc, score_thresh, iou_thresh=0.5):
"""
Nms.
:param pre_score: in the shape of [box_number, class_number]
:param pre_loc: int the shape of [box_number, 4] 4 means [x1, y1, x2, y2]
:param score_thresh:score_thresh
:param iou_thresh:iou_thresh
:return: labels_out
"""
print('using Greedy NMS')
class_num = pre_score_raw.shape[1] # get the numbers of classes.
pre_score_raw = pre_score_raw.max(
-1) # get the max score of the scores of classes.
pre_score = pre_score_raw[0] # score out
pre_class = pre_score_raw[1] # idx of score: is class
score_sort = pre_score.sort(descending=True) # sort the scores.
score_idx = score_sort[1][score_sort[0] >
score_thresh] # find the scores>0.7(thresh)
keep = []
for i in range(class_num): # with different classess.
a = pre_class[score_idx] == i # each class for NMS.
order_index = score_idx[a] # get the index of orders
while order_index.shape[0] > 0: # deal with all the boxes.
max_one = order_index[0].item() # get index of the max score box.
box_head = pre_loc[max_one] # get the score of it
box_others = pre_loc[order_index[1:]] # the rest boxes.
ious = iou_xywh(
box_head, box_others
) # count the ious between the max one and the others
rest = torch.lt(ious, iou_thresh).squeeze(
) # find the boxes of iou<0.5(thresh), discard the iou>0.5.
order_index = order_index[1:][
rest] # get the new index of the rest boxes, except the max one.
keep.append(max_one)
# ########### NMS UP ####################
labels_out = nms2labels(keep, pre_score, pre_loc, pre_class)
return labels_out
def NMS_Soft(self, pre_score, pre_class, pre_loc):
"""
Nms.
:param pre_score: in the shape of [box_number, class_number]
:param pre_loc: int the shape of [box_number, 4] 4 means [x1, y1, x2, y2]
:param score_thresh:score_thresh
:param iou_thresh:iou_thresh
:return: labels_out
"""
# print('using Soft NMS')
score_sort = pre_score.sort(descending=True)
score_idx = score_sort[1][score_sort[0] > self.score_thresh]
keep = []
for i in self.class_range: # with different classess.
a = pre_class[score_idx] == i
order_index = score_idx[a]
while order_index.shape[0] > 0:
max_one = order_index[0].item()
keep.append(max_one)
box_head = pre_loc[max_one]
box_others = pre_loc[order_index[1:]]
score_others = pre_score[order_index[1:]]
# print(score_others)
ious = iou_xywh(box_head, box_others,
type='N21').reshape(1, -1)
# print('iou', ious)
soft_score = score_others * torch.exp(
-pow(ious, 2) /
self.theta).squeeze() # s = s*e^(-iou^2 / theta)
# print('s',soft_score)
rest = torch.gt(soft_score, self.score_thresh)
order_index = order_index[1:][rest]
new_index = pre_score[order_index].sort(descending=True)
order_index = order_index[new_index[1]]
# ########### NMS UP ####################
return keep