k20hcmus/FishEye8K

1import numpy as np2import torch3 4from .deep.feature_extractor import Extractor5from .sort.nn_matching import NearestNeighborDistanceMetric6from .sort.detection import Detection7from .sort.tracker import Tracker8 9 10__all__ = ['DeepSort']11 12 13class DeepSort(object):14    def __init__(self, model_path, max_dist=0.2, min_confidence=0.3, nms_max_overlap=1.0, max_iou_distance=0.7, max_age=70, n_init=3, nn_budget=100, use_cuda=True):15        self.min_confidence = min_confidence16        self.nms_max_overlap = nms_max_overlap17 18        self.extractor = Extractor(model_path, use_cuda=use_cuda)19 20        max_cosine_distance = max_dist21        metric = NearestNeighborDistanceMetric(22            "cosine", max_cosine_distance, nn_budget)23        self.tracker = Tracker(24            metric, max_iou_distance=max_iou_distance, max_age=max_age, n_init=n_init)25 26    def update(self, bbox_xywh, confidences, oids, ori_img):27        self.height, self.width = ori_img.shape[:2]28        # generate detections29        features = self._get_features(bbox_xywh, ori_img)30        bbox_tlwh = self._xywh_to_tlwh(bbox_xywh)31        detections = [Detection(bbox_tlwh[i], conf, features[i],oid) for i, (conf,oid) in enumerate(zip(confidences,oids)) if conf > self.min_confidence]32 33        # run on non-maximum supression34        boxes = np.array([d.tlwh for d in detections])35        scores = np.array([d.confidence for d in detections])36 37        # update tracker38        self.tracker.predict()39        self.tracker.update(detections)40 41        # output bbox identities42        outputs = []43        for track in self.tracker.tracks:44            if not track.is_confirmed() or track.time_since_update > 1:45                continue46            box = track.to_tlwh()47            x1, y1, x2, y2 = self._tlwh_to_xyxy(box)48            track_id = track.track_id49            track_oid = track.oid50            outputs.append(np.array([x1, y1, x2, y2, track_id, track_oid], dtype=np.int64))51        if len(outputs) > 0:52            outputs = np.stack(outputs, axis=0)53        return outputs54 55    """56    TODO:57        Convert bbox from xc_yc_w_h to xtl_ytl_w_h58    Thanks JieChen91@github.com for reporting this bug!59    """60    @staticmethod61    def _xywh_to_tlwh(bbox_xywh):62        if isinstance(bbox_xywh, np.ndarray):63            bbox_tlwh = bbox_xywh.copy()64        elif isinstance(bbox_xywh, torch.Tensor):65            bbox_tlwh = bbox_xywh.clone()66        bbox_tlwh[:, 0] = bbox_xywh[:, 0] - bbox_xywh[:, 2] / 2.67        bbox_tlwh[:, 1] = bbox_xywh[:, 1] - bbox_xywh[:, 3] / 2.68        return bbox_tlwh69 70    def _xywh_to_xyxy(self, bbox_xywh):71        x, y, w, h = bbox_xywh72        x1 = max(int(x - w / 2), 0)73        x2 = min(int(x + w / 2), self.width - 1)74        y1 = max(int(y - h / 2), 0)75        y2 = min(int(y + h / 2), self.height - 1)76        return x1, y1, x2, y277 78    def _tlwh_to_xyxy(self, bbox_tlwh):79        """80        TODO:81            Convert bbox from xtl_ytl_w_h to xc_yc_w_h82        Thanks JieChen91@github.com for reporting this bug!83        """84        x, y, w, h = bbox_tlwh85        x1 = max(int(x), 0)86        x2 = min(int(x+w), self.width - 1)87        y1 = max(int(y), 0)88        y2 = min(int(y+h), self.height - 1)89        return x1, y1, x2, y290 91    def increment_ages(self):92        self.tracker.increment_ages()93 94    def _xyxy_to_tlwh(self, bbox_xyxy):95        x1, y1, x2, y2 = bbox_xyxy96 97        t = x198        l = y199        w = int(x2 - x1)100        h = int(y2 - y1)101        return t, l, w, h102 103    def _get_features(self, bbox_xywh, ori_img):104        im_crops = []105        for box in bbox_xywh:106            x1, y1, x2, y2 = self._xywh_to_xyxy(box)107            im = ori_img[y1:y2, x1:x2]108            im_crops.append(im)109        if im_crops:110            features = self.extractor(im_crops)111        else:112            features = np.array([])113        return features114