示例#1
0
文件: utils.py 项目: wjfwzzc/Models
    def predict(self, **inputs):
        """
        Args:
            val_func(callable): model inference function

        Returns:
            results boxes: detection model output
        """
        box_cls, box_delta = self.pred_func(**inputs)
        # box_cls, box_delta = self.model(**inputs)
        box_cls, box_delta = box_cls.numpy(), box_delta.numpy()
        dtboxes_all = list()
        all_inds = np.where(box_cls > self.model.cfg.test_cls_threshold)

        for c in range(self.model.cfg.num_classes):
            inds = np.where(all_inds[1] == c)[0]
            inds = all_inds[0][inds]
            scores = box_cls[inds, c]
            if self.model.cfg.class_aware_box:
                bboxes = box_delta[inds, c, :]
            else:
                bboxes = box_delta[inds, :]

            dtboxes = np.hstack(
                (bboxes, scores[:, np.newaxis])).astype(np.float32)

            if dtboxes.size > 0:
                if self.model.cfg.test_nms == -1:
                    keep = dtboxes[:, 4].argsort()[::-1]
                else:
                    assert 0 < self.model.cfg.test_nms <= 1.0
                    keep = py_cpu_nms(dtboxes, self.model.cfg.test_nms)
                dtboxes = np.hstack(
                    (dtboxes[keep], np.full(
                        (len(keep), 1), c,
                        dtype=np.float32))).astype(np.float32)
                dtboxes_all.extend(dtboxes)

        if len(dtboxes_all) > self.model.cfg.test_max_boxes_per_image:
            dtboxes_all = sorted(
                dtboxes_all, reverse=True,
                key=lambda i: i[4])[:self.model.cfg.test_max_boxes_per_image]

        dtboxes_all = np.array(dtboxes_all, dtype=np.float)
        return dtboxes_all
示例#2
0
    def predict(self, val_func):
        """
        Args:
            val_func(callable): model inference function

        Returns:
            results boxes: detection model output
        """
        model = self.model

        box_cls, box_delta = val_func()
        box_cls, box_delta = box_cls.numpy(), box_delta.numpy()
        dtboxes_all = list()
        all_inds = np.where(box_cls > model.cfg.test_cls_threshold)

        for c in range(0, model.cfg.num_classes):
            inds = np.where(all_inds[1] == c)[0]
            inds = all_inds[0][inds]
            scores = box_cls[inds, c]
            if model.cfg.class_aware_box:
                bboxes = box_delta[inds, c, :]
            else:
                bboxes = box_delta[inds, :]

            dtboxes = np.hstack(
                (bboxes, scores[:, np.newaxis])).astype(np.float32)

            if dtboxes.size > 0:
                keep = py_cpu_nms(dtboxes, model.cfg.test_nms)
                dtboxes = np.hstack(
                    (dtboxes[keep], np.ones(
                        (len(keep), 1), np.float32) * c)).astype(np.float32)
                dtboxes_all.extend(dtboxes)

        if len(dtboxes_all) > model.cfg.test_max_boxes_per_image:
            dtboxes_all = sorted(
                dtboxes_all, reverse=True,
                key=lambda i: i[4])[:model.cfg.test_max_boxes_per_image]

        dtboxes_all = np.array(dtboxes_all, dtype=np.float)
        return dtboxes_all
示例#3
0
    def detect_persons(self, image):

        data, im_info = DetEvaluator.process_inputs(
            image.copy(),
            self.detector.cfg.test_image_short_size,
            self.detector.cfg.test_image_max_size,
        )

        evaluator = DetEvaluator(self.detector)

        det_res = evaluator.predict(image=mge.tensor(data),
                                    im_info=mge.tensor(im_info))

        persons = []
        for d in det_res:
            cls_id = int(d[5] + 1)
            if cls_id == 1:
                bbox = d[:5]
                persons.append(bbox)
        persons = np.array(persons).reshape(-1, 5)
        keep = py_cpu_nms(persons, cfg.nms_thr)
        return persons[keep]