def traditional_nms(self, boxes, masks, scores, iou_threshold=0.5, conf_thresh=0.05): import pyximport pyximport.install(setup_args={"include_dirs": np.get_include()}, reload_support=True) from utils.cython_nms import nms as cnms num_classes = scores.size(0) idx_lst = [] cls_lst = [] scr_lst = [] # Multiplying by max_size is necessary because of how cnms computes its area and intersections boxes = boxes * cfg.max_size for _cls in range(num_classes): cls_scores = scores[_cls, :] # threshold 넘는 confidence score 남김 conf_mask = cls_scores > conf_thresh # 점수 내림차순으로 정렬 idx = torch.arange(cls_scores.size(0), device=boxes.device) cls_scores = cls_scores[conf_mask] idx = idx[conf_mask] # 없으면 건너뛰기 if cls_scores.size(0) == 0: continue preds = torch.cat([boxes[conf_mask], cls_scores[:, None]], dim=1).cpu().numpy() # 예측에 대해 iou 계산 후 threshold 넘는 것 제거 keep = cnms(preds, iou_threshold) keep = torch.Tensor(keep, device=boxes.device).long() idx_lst.append(idx[keep]) cls_lst.append(keep * 0 + _cls) scr_lst.append(cls_scores[keep]) # 남은 것들 list에 추가 idx = torch.cat(idx_lst, dim=0) classes = torch.cat(cls_lst, dim=0) scores = torch.cat(scr_lst, dim=0) scores, idx2 = scores.sort(0, descending=True) idx2 = idx2[:cfg.max_num_detections] scores = scores[:cfg.max_num_detections] idx = idx[idx2] classes = classes[idx2] # Undo the multiplication above return boxes[idx] / cfg.max_size, masks[idx], classes, scores
def traditional_nms(self, boxes, masks, scores, iou_threshold=0.5, conf_thresh=0.05): import pyximport pyximport.install(setup_args={"include_dirs": np.get_include()}, reload_support=True) from utils.cython_nms import nms as cnms num_classes = scores.shape[0] idx_lst = [] cls_lst = [] scr_lst = [] # Multiplying by max_size is necessary because of how cnms computes its area and intersections boxes = boxes * cfg.max_size for _cls in range(num_classes): cls_scores = scores[_cls] conf_mask = cls_scores > conf_thresh idx = jt.arange(cls_scores.shape[0], device=boxes.device) cls_scores = cls_scores[conf_mask] idx = idx[conf_mask] if cls_scores.shape[0] == 0: continue preds = jt.contrib.concat([boxes[conf_mask], cls_scores[:, None]], dim=1).numpy() keep = cnms(preds, iou_threshold) keep = jt.array(keep).int64() idx_lst.append(idx[keep]) cls_lst.append(keep * 0 + _cls) scr_lst.append(cls_scores[keep]) idx = jt.contrib.concat(idx_lst, dim=0) classes = jt.contrib.concat(cls_lst, dim=0) scores = jt.contrib.concat(scr_lst, dim=0) scores, idx2 = scores.sort(0, descending=True) idx2 = idx2[:cfg.max_num_detections] scores = scores[:cfg.max_num_detections] idx = idx[idx2] classes = classes[idx2] # Undo the multiplication above return boxes[idx] / cfg.max_size, masks[idx], classes, scores
def traditional_nms_ours(self, boxes, masks, scores, iou_threshold=0.5, conf_thresh=0.05): import pyximport pyximport.install(setup_args={"include_dirs": np.get_include()}, reload_support=True) from utils.cython_nms import nms as cnms num_classes = scores.size(0) idx_lst = [] cls_lst = [] scr_lst = [] box_lst = [] mask_lst = [] # Multiplying by max_size is necessary because of how cnms computes its area and intersections boxes = boxes * cfg.max_size for _cls in range(num_classes): cls_scores = scores[_cls, :] _, id = cls_scores.sort(0, descending=True) id = id[:200].contiguous() cls_scores = cls_scores[id] idx = torch.arange(cls_scores.size(0), device=boxes.device) if cls_scores.size(0) == 0: continue preds = torch.cat([boxes[id], cls_scores[:, None]], dim=1).cpu().numpy() keep = cnms(preds, iou_threshold) keep = torch.Tensor(keep, device=boxes.device).long() m = (cls_scores[keep] > 0.01) idx_lst.append(idx[keep][m]) cls_lst.append(keep[m] * 0 + _cls) scr_lst.append(cls_scores[keep][m]) box_lst.append(boxes[id][keep][m]) mask_lst.append(masks[id][keep][m]) idx = torch.cat(idx_lst, dim=0) classes = torch.cat(cls_lst, dim=0) scores = torch.cat(scr_lst, dim=0) boxes = torch.cat(box_lst, dim=0) masks = torch.cat(mask_lst, dim=0) scores, idx2 = scores.sort(0, descending=True) idx2 = idx2[:cfg.max_num_detections] scores = scores[:cfg.max_num_detections] classes = classes[idx2] # Undo the multiplication above return boxes[idx2] / cfg.max_size, masks[idx2], classes, scores
def traditional_nms(self, boxes, masks, track, scores, iou_threshold=0.5, conf_thresh=0.05): num_classes = scores.size(0) idx_lst = [] cls_lst = [] scr_lst = [] # Multiplying by max_size is necessary because of how cnms computes its area and intersections boxes = boxes * cfg.max_size for _cls in range(num_classes): cls_scores = scores[_cls, :] conf_mask = cls_scores > conf_thresh idx = torch.arange(cls_scores.size(0), device=boxes.device) cls_scores = cls_scores[conf_mask] idx = idx[conf_mask] if cls_scores.size(0) == 0: continue preds = torch.cat([boxes[conf_mask], cls_scores[:, None]], dim=1).cpu().numpy() keep = cnms(preds, iou_threshold) keep = torch.Tensor(keep, device=boxes.device).long() idx_lst.append(idx[keep]) cls_lst.append(keep * 0 + _cls) scr_lst.append(cls_scores[keep]) idx = torch.cat(idx_lst, dim=0) classes = torch.cat(cls_lst, dim=0) scores = torch.cat(scr_lst, dim=0) scores, idx2 = scores.sort(0, descending=True) idx2 = idx2[:cfg.max_num_detections] scores = scores[:cfg.max_num_detections] idx = idx[idx2] classes = classes[idx2] # Undo the multiplication above return boxes[idx] / cfg.max_size, masks[idx], track[ idx], classes, scores