Python nms_wrapper.nms方法代码示例

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def _keep_top_k(self, boxes, box_inds, end_at, top_k, thresh):
        '''Set dynamic class threshold based on average detections per class per frame
        (before nms) and keep boxes above this threshold.
        '''
        box_list = boxes[:end_at].tolist()
        box_list = [box for box in box_list if box is not None]
        if len(box_list)==0:
            return
        X = torch.cat(box_list, dim=0)
        if X.size(0) == 0:
            return

        scores,_ = torch.sort(X[:,4], descending=True)
        # set threshold for this class
        thresh = scores[min(scores.numel(), top_k)]
        for image_index in range(end_at):
            if boxes[image_index] is not None and boxes[image_index].size(0)>0:
                bbox = boxes[image_index]
                keep = torch.nonzero(bbox[:,4]>=thresh).view(-1)
                if keep.numel()==0:
                    continue
                boxes[image_index] = bbox[keep]
                box_inds[image_index] = box_inds[image_index][keep]
        return boxes, box_inds, thresh 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def apply_nms(all_boxes, thresh):
    """Apply non-maximum suppression to all predicted boxes output by the
    test_net method.
    """
    num_classes = len(all_boxes)
    num_images = len(all_boxes[0])
    nms_boxes = [[[] for _ in range(num_images)]
                 for _ in range(num_classes)]
    for cls_ind in range(num_classes):
        for im_ind in range(num_images):
            dets = all_boxes[cls_ind][im_ind]
            if dets == []:
                continue
            keep = nms(torch.FloatTensor(dets).cuda(), thresh)
            if len(keep) == 0:
                continue
            nms_boxes[cls_ind][im_ind] = dets[keep, :].copy()
    return nms_boxes 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def __init__(self, pred_boxes, scores, pred_trk_boxes, classes):
        print("Starting post-processing.")
        self.pred_boxes = pred_boxes
        self.scores = scores
        self.pred_trk_boxes = pred_trk_boxes
        self.num_frames = pred_boxes.size(0)+1
        self.num_classes = len(classes)
        self.classes = classes
        self.top_scores = torch.cuda.FloatTensor(self.num_classes,1).zero_() 
        #self._aboxes = [[[] for _ in xrange(self.num_frames)] for _ in xrange(self.num_classes)]
        self._aboxes = np.ndarray(shape=(self.num_classes, self.num_frames), dtype=np.object)
        #self._box_inds = [[[] for _ in xrange(self.num_frames)] for _ in xrange(self.num_classes)]
        self._box_inds = np.ndarray(shape=(self.num_classes, self.num_frames), dtype=np.object)
        #self._track_boxes = [[[] for _ in xrange(self.num_frames)] for _ in xrange(self.num_classes)]
        self._track_boxes = np.ndarray((self.num_classes,self.num_frames), dtype=np.object)
        #self._paths = [[] for _ in xrange(self.num_classes)]
        self._paths = np.ndarray((self.num_classes,), dtype=np.object)
        #self._ascores_track =  [[] for _ in xrange(self.num_frames)]  
        self._ascores_track = np.ndarray(shape=(self.num_frames,), dtype=np.object)
        #self._aboxes_track = [[[] for _ in xrange(2)] for _ in xrange(self.num_frames)]
        self._aboxes_track = np.ndarray(shape=(self.num_frames, 2), dtype=np.object)
        self.CONF_THRESH = torch.cuda.FloatTensor(self.num_classes)
        self.CONF_THRESH.fill_(-1*float('Inf'))
        self.max_per_image = 400 # keep at most max_per_image dets per class b/4 nms
        self.max_per_set = 160*self.num_frames
        self.class_agnostic = True

        self._process_frame_pairs()
        
        for j in range(1, self.num_classes):
            self._aboxes[j], self._box_inds[j], self.CONF_THRESH[j] = \
                    self._keep_top_k(self._aboxes[j], self._box_inds[j], 
                            self.num_frames, self.max_per_set, self.CONF_THRESH[j]) 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def _make_tubes(self, frameBoxes, max_per_image, box_voting, tracks_cell):
        '''Build tubes for cth class.
        '''

        tracks = {'boxes': None, 'scores': None, 'c': None}
        tracks['boxes'] = tracks_cell[0]
        tracks['scores'] = tracks_cell[1]
        tracks['c'] = tracks_cell[2]

        nms_thresh = 0.3

        object_frames_boxes = np.ndarray((len(frameBoxes),), dtype=np.object)
        object_frames_scores =  np.ndarray((len(frameBoxes),), dtype=np.object)
        object_frames_boxes_idx = np.ndarray((len(frameBoxes),), dtype=np.object)
        object_frames_trackedboxes = np.ndarray((len(frameBoxes),), dtype=np.object)
        
        # Iterate over the non-empty frames
        for f in range(len(frameBoxes)-1):
            # boxes in frame f
            boxes = frameBoxes[f]
            if box_voting: # TODO
                raise NotImplementedError
            else:
                nms_idx = nms(boxes[:, :5].clone(), nms_thresh).long().view(-1)
                if nms_idx.numel() > max_per_image:
                    nms_idx = nms_idx[:max_per_image]
                boxes = boxes[nms_idx]
                object_frames_boxes[f] = boxes[:, :4]
                object_frames_scores[f] = boxes[:, 4]
                object_frames_boxes_idx[f] = torch.arange(boxes.size(0)).cuda()
                if tracks['boxes'] is not None and tracks['boxes'][f, 0] is not None:
                    object_frames_trackedboxes[f] = tracks['boxes'][f, :]

        paths = self._zero_jump_link(object_frames_boxes, object_frames_scores,
                                     object_frames_boxes_idx, object_frames_trackedboxes)

        return paths 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def _extract_frames(self, v_path):
        '''Extract all fromes from @v_path

        :param v_path: full path to video
        :return: list of full paths to extracted video frames
        '''
        print("Extracting frames from {}".format(v_path))
        # Store frames in tmp dir by default
        tmp_dir = os.path.join('/tmp', os.path.basename(v_path))
        # If path exists, delete it
        if os.path.isdir(tmp_dir):
            shutil.rmtree(tmp_dir)
        os.mkdir(tmp_dir)
        # create directory to dump output to
        save_dir = tmp_dir.replace('.mp4', '') + "_processed"
        # clear contents of output directory before saving
        if os.path.isdir(save_dir):
            shutil.rmtree(save_dir)
        print("Saving to {}".format(save_dir))
        
        self._input_dir = tmp_dir
        self._output_dir = save_dir
        # TODO Make fps configurable at command line
        cmd = "ffmpeg -i %s -vf fps=10 %s" % (v_path,os.path.join(tmp_dir,'%09d.png'))
        # execute ffmpeg cmd
        subprocess.call(cmd,shell=True)
        # set frame paths of the current video
        self._frame_paths = sorted(glob.glob("%s/*.png" % tmp_dir))
        self._n_frames = len(self._frame_paths)
        self._max_per_set = 160*self._n_frames # average 160 dets per class per frame before nms
        print("Found {} frames".format(self._n_frames))
        return 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Train a Fast R-CNN network')
    parser.add_argument('exp_name', type=str, default=None, help='experiment name')
    parser.add_argument('--dataset', dest='dataset',
                        help='training dataset',
                        default='pascal_voc', type=str)
    parser.add_argument('--cfg', dest='cfg_file',
                        help='optional config file',
                        default='cfgs/vgg16.yml', type=str)
    parser.add_argument('--net', dest='net',
                        help='detnet59, etc',
                        default='detnet59', type=str)
    parser.add_argument('--set', dest='set_cfgs',
                        help='set config keys', default=None,
                        nargs=argparse.REMAINDER)
    parser.add_argument('--load_dir', dest='load_dir',
                        help='directory to load models', default="weights")
    parser.add_argument('--image_dir', dest='image_dir',
                        help='directory to load images', default="demo_images/",
                        type=str)
    parser.add_argument('--result_dir', dest='result_dir', help='directory to save visual result', default="vis_results/",
                        type=str)
    parser.add_argument('--cuda', dest='cuda',
                        help='whether use CUDA',
                        action='store_true')
    parser.add_argument('--checksession', dest='checksession',
                        help='checksession to load model',
                        default=4, type=int)
    parser.add_argument('--checkepoch', dest='checkepoch',
                        help='checkepoch to load network',
                        default=6, type=int)
    parser.add_argument('--checkpoint', dest='checkpoint',
                        help='checkpoint to load network',
                        default=10000, type=int)
    parser.add_argument('--soft_nms', help='whether use soft nms', action='store_true')
    parser.add_argument('--cascade', help='whether use cascade', action='store_true')
    parser.add_argument('--cag', dest='class_agnostic',
                        help='whether perform class_agnostic bbox regression',
                        action='store_true')

    args = parser.parse_args()
    return args 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Train a Fast R-CNN network')
    parser.add_argument('exp_name', type=str, default=None, help='experiment name')
    parser.add_argument('--dataset', dest='dataset',
                        help='training dataset',
                        default='pascal_voc', type=str)
    parser.add_argument('--cfg', dest='cfg_file',
                        help='optional config file',
                        default='cfgs/vgg16.yml', type=str)
    parser.add_argument('--net', dest='net',
                        help='detnet59, etc',
                        default='detnet59', type=str)
    parser.add_argument('--set', dest='set_cfgs',
                        help='set config keys', default=None,
                        nargs=argparse.REMAINDER)
    parser.add_argument('--load_dir', dest='load_dir',
                        help='directory to load models', default="/srv/share/jyang375/models",
                        type=str)
    parser.add_argument('--cuda', dest='cuda',
                        help='whether use CUDA',
                        action='store_true')
    parser.add_argument('--ls', dest='large_scale',
                        help='whether use large imag scale',
                        action='store_true')
    parser.add_argument('--mGPUs', dest='mGPUs',
                        help='whether use multiple GPUs',
                        action='store_true')
    parser.add_argument('--cag', dest='class_agnostic',
                        help='whether perform class_agnostic bbox regression',
                        action='store_true')
    parser.add_argument('--parallel_type', dest='parallel_type',
                        help='which part of model to parallel, 0: all, 1: model before roi pooling',
                        default=0, type=int)
    parser.add_argument('--checksession', dest='checksession',
                        help='checksession to load model',
                        default=1, type=int)
    parser.add_argument('--checkepoch', dest='checkepoch',
                        help='checkepoch to load network',
                        default=1, type=int)
    parser.add_argument('--checkpoint', dest='checkpoint',
                        help='checkpoint to load network',
                        default=10021, type=int)
    parser.add_argument('--bs', dest='batch_size',
                        help='batch_size',
                        default=1, type=int)
    parser.add_argument('--vis', dest='vis',
                        help='visualization mode',
                        action='store_true')
    parser.add_argument('--soft_nms', help='whether use soft nms', action='store_true')
    parser.add_argument('--cascade', help='whether use soft nms', action='store_true')
    args = parser.parse_args()
    return args 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def visualize_without_paths(video_dataset, pred_boxes, scores, pred_trk_boxes, det_classes):
    print("Visualizing...")
    list_im = video_dataset._frame_paths 

    num_classes = len(det_classes)
    num_frames = len(list_im)
    try:
        font = ImageFont.truetype('arial.ttf', 24)
    except IOError:
        font = ImageFont.load_default()

    for i_frame in range(num_frames-1):
        print('frame: {}/{}'.format(i_frame, num_frames))
        fig, ax = plt.subplots(figsize=(12, 12))
        img_path = list_im[i_frame]
        img = cv2.imread(img_path)
        img = img[:,:,(2,1,0)]
        disp_image = Image.fromarray(np.uint8(img))
        for cls_ind in range(1, num_classes):
            ax.imshow(disp_image, aspect='equal')
            class_name = det_classes[cls_ind]
            keep = torch.nonzero(scores[i_frame][0][:, cls_ind]>CONF_THRESH).view(-1)
            if keep.numel()==0:
                # no detections above threshold for this class
                continue
            cls_scores = scores[i_frame][0][keep][:, cls_ind]
            _, order = torch.sort(cls_scores, 0, True)
            cls_boxes = pred_boxes[i_frame][0][keep, :]
            cls_dets = torch.cat([cls_boxes, cls_scores.contiguous().view(-1,1)], dim=1)
            cls_dets = cls_dets[order]
            keep = nms(cls_dets, 0.3)
            cls_dets = cls_dets[keep.view(-1).long()]
            for ibox in range(cls_dets.size(0)):
                bbox = cls_dets[ibox, :4].cpu().numpy().flatten()
                score = cls_dets[ibox, 4]
                ax.add_patch(
                        plt.Rectangle((bbox[0], bbox[1]),
                            bbox[2] - bbox[0],
                            bbox[3] - bbox[1], fill=False,
                            edgecolor=COLOR_WHEEL[cls_ind], linewidth=3.5)
                        )
                ax.text(bbox[0], bbox[1] - 2,
                        '{:s} {:.3f}'.format(class_name, score),
                        bbox=dict(facecolor=COLOR_WHEEL[cls_ind], alpha=0.5),
                        fontsize=14, color='white')

        # Save image with bboxes overlaid
        plt.axis('off')
        plt.tight_layout()
        #plt.show()
        if not os.path.exists(video_dataset._output_dir):
            os.makedirs(video_dataset._output_dir)
        plt.savefig(os.path.join(video_dataset._output_dir, os.path.basename(img_path)))
        plt.clf()
        plt.close('all') 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def detect(net, im_path,device, thresh=0.5, visualize=False, timers=None, pyramid=False, visualization_folder=None):
    """
    Main module to detect faces
    :param net: The trained network
    :param im_path: The path to the image
    :param device: GPU or CPU device to be used
    :param thresh: Detection with a less score than thresh are ignored
    :param visualize: Whether to visualize the detections
    :param timers: Timers for calculating detect time (if None new timers would be created)
    :param pyramid: Whether to use pyramid during inference
    :param visualization_folder: If set the visualizations would be saved in this folder (if visualize=True)
    :return: cls_dets (bounding boxes concatenated with scores) and the timers
    """

    if not timers:
        timers = {'detect': Timer(),
                  'misc': Timer()}

    im = cv2.imread(im_path)
    imfname = os.path.basename(im_path)
    sys.stdout.flush()
    timers['detect'].tic()

    if not pyramid:
        im_scale = _compute_scaling_factor(im.shape, cfg.TEST.SCALES[0], cfg.TEST.MAX_SIZE)
        im_blob = _get_image_blob(im, [im_scale])[0]
        ssh_rois = forward(net,im_blob,im_scale,device,thresh)

    else :
        assert False, 'not implement'

    timers['detect'].toc()
    timers['misc'].tic()

    nms_keep = nms(ssh_rois, cfg.TEST.RPN_NMS_THRESH)
    cls_dets = ssh_rois[nms_keep, :]

    if visualize:
        plt_name = os.path.splitext(imfname)[0] + '_detections_{}'.format("SSH pytorch")
        visusalize_detections(im, cls_dets, plt_name=plt_name, visualization_folder=visualization_folder)
    timers['misc'].toc()
    return cls_dets, timers 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Train a Fast R-CNN network')
    parser.add_argument('exp_name', type=str, default=None, help='experiment name')
    parser.add_argument('--dataset', dest='dataset',
                        help='training dataset',
                        default='pascal_voc', type=str)
    parser.add_argument('--cfg', dest='cfg_file',
                        help='optional config file',
                        default='cfgs/vgg16.yml', type=str)
    parser.add_argument('--net', dest='net',
                        help='detnet59, etc',
                        default='detnet59', type=str)
    parser.add_argument('--set', dest='set_cfgs',
                        help='set config keys', default=None,
                        nargs=argparse.REMAINDER)
    parser.add_argument('--load_dir', dest='load_dir',
                        help='directory to load models', default="weights")
    parser.add_argument('--image_dir', dest='image_dir',
                        help='directory to load images', default="demo_images/",
                        type=str)
    parser.add_argument('--result_dir', dest='result_dir', help='directory to save visual result',
                        default="vis_results/",
                        type=str)
    parser.add_argument('--cuda', dest='cuda',
                        help='whether use CUDA',
                        action='store_true')
    parser.add_argument('--checksession', dest='checksession',
                        help='checksession to load model',
                        default=4, type=int)
    parser.add_argument('--checkepoch', dest='checkepoch',
                        help='checkepoch to load network',
                        default=6, type=int)
    parser.add_argument('--checkpoint', dest='checkpoint',
                        help='checkpoint to load network',
                        default=10000, type=int)
    parser.add_argument('--soft_nms', help='whether use soft nms', action='store_true')
    parser.add_argument('--cag', dest='class_agnostic',
                        help='whether perform class_agnostic bbox regression',
                        action='store_true')

    args = parser.parse_args()
    return args 

# 需要导入模块: from model.nms import nms_wrapper [as 别名]
# 或者: from model.nms.nms_wrapper import nms [as 别名]
def parse_args():
    """
    Parse input arguments
    """
    parser = argparse.ArgumentParser(description='Train a Fast R-CNN network')
    parser.add_argument('exp_name', type=str, default=None, help='experiment name')
    parser.add_argument('--dataset', dest='dataset',
                        help='training dataset',
                        default='pascal_voc', type=str)
    parser.add_argument('--cfg', dest='cfg_file',
                        help='optional config file',
                        default='cfgs/vgg16.yml', type=str)
    parser.add_argument('--net', dest='net',
                        help='detnet59, etc',
                        default='detnet59', type=str)
    parser.add_argument('--set', dest='set_cfgs',
                        help='set config keys', default=None,
                        nargs=argparse.REMAINDER)
    parser.add_argument('--load_dir', dest='load_dir',
                        help='directory to load models', default="/srv/share/jyang375/models",
                        type=str)
    parser.add_argument('--cuda', dest='cuda',
                        help='whether use CUDA',
                        action='store_true')
    parser.add_argument('--ls', dest='large_scale',
                        help='whether use large imag scale',
                        action='store_true')
    parser.add_argument('--mGPUs', dest='mGPUs',
                        help='whether use multiple GPUs',
                        action='store_true')
    parser.add_argument('--cag', dest='class_agnostic',
                        help='whether perform class_agnostic bbox regression',
                        action='store_true')
    parser.add_argument('--parallel_type', dest='parallel_type',
                        help='which part of model to parallel, 0: all, 1: model before roi pooling',
                        default=0, type=int)
    parser.add_argument('--checksession', dest='checksession',
                        help='checksession to load model',
                        default=1, type=int)
    parser.add_argument('--checkepoch', dest='checkepoch',
                        help='checkepoch to load network',
                        default=1, type=int)
    parser.add_argument('--checkpoint', dest='checkpoint',
                        help='checkpoint to load network',
                        default=10021, type=int)
    parser.add_argument('--bs', dest='batch_size',
                        help='batch_size',
                        default=1, type=int)
    parser.add_argument('--vis', dest='vis',
                        help='visualization mode',
                        action='store_true')
    parser.add_argument('--soft_nms', help='whether use soft nms', action='store_true')
    args = parser.parse_args()
    return args