예제 #1
0
    def update(self, image, tlwhs, pids):
        self.frame_id += 1

        # set features
        tlbrs = [tlwh2tlbr(tlwh) for tlwh in tlwhs]
        to_remove = []

        for i in range(len(tlbrs)):
            tlbr = tlbrs[i]
            if (tlbr[1] < 0 and tlbr[3] < 0) or (tlbr[0] < 0 and tlbr[2] < 0):
                to_remove.append(i)
        pids = [pids[i] for i in range(len(tlbrs)) if i not in to_remove]
        tlbrs = [tlbrs[i] for i in range(len(tlbrs)) if i not in to_remove]

        features = extract_reid_features(self.reid_model, image, tlbrs)
        features = features.cpu().numpy()
        pids = np.array(pids, ndmin=2).transpose()
        if not features.any():
            return
        line = np.concatenate([
            np.ones_like(pids) * self.cam, pids,
            np.ones_like(pids) * self.frame_id, features
        ],
                              axis=1)
        self.lines = np.concatenate([self.lines, line], axis=0)
예제 #2
0
    def update(self, image, tlwhs, det_scores=None):
        self.frame_id += 1

        activated_starcks = []
        refind_stracks = []
        lost_stracks = []
        removed_stracks = []

        """step 1: prediction"""
        for strack in itertools.chain(self.tracked_stracks, self.lost_stracks):
            strack.predict()

        """step 2: scoring and selection"""
        if det_scores is None:
            det_scores = np.ones(len(tlwhs), dtype=float)
        detections = [STrack(tlwh, score, from_det=True) for tlwh, score in zip(tlwhs, det_scores)]

        if self.classifier is None:
            pred_dets = []
        else:
            self.classifier.update(image)

            n_dets = len(tlwhs)
            if self.use_tracking:
                tracks = [STrack(t.self_tracking(image), t.tracklet_score(), from_det=False)
                          for t in itertools.chain(self.tracked_stracks, self.lost_stracks) if t.is_activated]
                detections.extend(tracks)
            rois = np.asarray([d.tlbr for d in detections], dtype=np.float32)

            cls_scores = self.classifier.predict(rois)
            scores = np.asarray([d.score for d in detections], dtype=np.float)
            scores[0:n_dets] = 1.
            scores = scores * cls_scores
            # nms
            if len(detections) > 0:
                keep = nms_detections(rois, scores.reshape(-1), nms_thresh=0.3)
                mask = np.zeros(len(rois), dtype=np.bool)
                mask[keep] = True
                keep = np.where(mask & (scores >= self.min_cls_score))[0]
                detections = [detections[i] for i in keep]
                scores = scores[keep]
                for d, score in zip(detections, scores):
                    d.score = score
            pred_dets = [d for d in detections if not d.from_det]
            detections = [d for d in detections if d.from_det]

        # set features
        tlbrs = [det.tlbr for det in detections]
        features = extract_reid_features(self.reid_model, image, tlbrs)
        features = features.cpu().numpy()
        for i, det in enumerate(detections):
            det.set_feature(features[i])

        """step 3: association for tracked"""
        # matching for tracked targets
        unconfirmed = []
        tracked_stracks = []  # type: list[STrack]
        for track in self.tracked_stracks:
            if not track.is_activated:
                unconfirmed.append(track)
            else:
                tracked_stracks.append(track)

        dists = matching.nearest_reid_distance(tracked_stracks, detections, metric='euclidean')
        dists = matching.gate_cost_matrix(self.kalman_filter, dists, tracked_stracks, detections)
        matches, u_track, u_detection = matching.linear_assignment(dists, thresh=self.min_ap_dist)
        for itracked, idet in matches:
            tracked_stracks[itracked].update(detections[idet], self.frame_id, image)

        # matching for missing targets
        detections = [detections[i] for i in u_detection]
        dists = matching.nearest_reid_distance(self.lost_stracks, detections, metric='euclidean')
        dists = matching.gate_cost_matrix(self.kalman_filter, dists, self.lost_stracks, detections)
        matches, u_lost, u_detection = matching.linear_assignment(dists, thresh=self.min_ap_dist)
        for ilost, idet in matches:
            track = self.lost_stracks[ilost]  # type: STrack
            det = detections[idet]
            track.re_activate(det, self.frame_id, image, new_id=not self.use_refind)
            refind_stracks.append(track)

        # remaining tracked
        # tracked
        len_det = len(u_detection)
        detections = [detections[i] for i in u_detection] + pred_dets
        r_tracked_stracks = [tracked_stracks[i] for i in u_track]
        dists = matching.iou_distance(r_tracked_stracks, detections)
        matches, u_track, u_detection = matching.linear_assignment(dists, thresh=0.7)
        for itracked, idet in matches:
            r_tracked_stracks[itracked].update(detections[idet], self.frame_id, image, update_feature=True)
        for it in u_track:
            track = r_tracked_stracks[it]
            track.mark_lost()
            lost_stracks.append(track)

        # unconfirmed
        detections = [detections[i] for i in u_detection if i < len_det]
        dists = matching.iou_distance(unconfirmed, detections)
        matches, u_unconfirmed, u_detection = matching.linear_assignment(dists, thresh=0.7)
        for itracked, idet in matches:
            unconfirmed[itracked].update(detections[idet], self.frame_id, image, update_feature=True)
        for it in u_unconfirmed:
            track = unconfirmed[it]
            track.mark_removed()
            removed_stracks.append(track)

        """step 4: init new stracks"""
        for inew in u_detection:
            track = detections[inew]
            if not track.from_det or track.score < 0.6:
                continue
            track.activate(self.kalman_filter, self.frame_id, image)
            activated_starcks.append(track)

        """step 6: update state"""
        for track in self.lost_stracks:
            if self.frame_id - track.end_frame > self.max_time_lost:
                track.mark_removed()
                removed_stracks.append(track)

        self.tracked_stracks = [t for t in self.tracked_stracks if t.state == TrackState.Tracked]
        self.lost_stracks = [t for t in self.lost_stracks if t.state == TrackState.Lost]  # type: list[STrack]
        self.tracked_stracks.extend(activated_starcks)
        self.tracked_stracks.extend(refind_stracks)
        self.lost_stracks.extend(lost_stracks)
        self.removed_stracks.extend(removed_stracks)

        # output_stracks = self.tracked_stracks + self.lost_stracks

        # get scores of lost tracks
        rois = np.asarray([t.tlbr for t in self.lost_stracks], dtype=np.float32)
        lost_cls_scores = self.classifier.predict(rois)
        out_lost_stracks = [t for i, t in enumerate(self.lost_stracks)
                            if lost_cls_scores[i] > 0.3 and self.frame_id - t.end_frame <= 4]
        output_tracked_stracks = [track for track in self.tracked_stracks if track.is_activated]

        output_stracks = output_tracked_stracks + out_lost_stracks

        logger.debug('===========Frame {}=========='.format(self.frame_id))
        logger.debug('Activated: {}'.format([track.track_id for track in activated_starcks]))
        logger.debug('Refind: {}'.format([track.track_id for track in refind_stracks]))
        logger.debug('Lost: {}'.format([track.track_id for track in lost_stracks]))
        logger.debug('Removed: {}'.format([track.track_id for track in removed_stracks]))

        return output_stracks
예제 #3
0
    def update(self, image, tlwhs=None, det_scores=None):
        self.frame_id += 1

        activated_starcks = []
        refind_stracks = []
        lost_stracks = []
        removed_stracks = []

        # <custom> Detect
        detects_candidate = self.models["personDetect"].Detect(image)

        if len(detects_candidate) != 0:
            detects = detects_candidate[
                detects_candidate[:, 4] >= self.configs["detect_threshold"]]

            # Re Detect
            if self.configs.get(
                    "redetect") is not None and self.configs["redetect"]:
                detects_refind = detects_candidate[
                    (detects_candidate[:,
                                       4] < self.configs["detect_threshold"])
                    & (detects_candidate[:, 4] >
                       self.configs["redetect_target_min"]) &
                    ((detects_candidate[:, 2] - detects_candidate[:, 0]) *
                     (detects_candidate[:, 3] - detects_candidate[:, 1]) <
                     480000)]
                for detect in detects_refind:
                    # Super Resolution & Deblurring
                    frame_redetect = image[int(detect[1]):int(detect[3]),
                                           int(detect[0]):int(detect[2])]
                    frame_redetect = self.models[
                        "superResolution"].SuperResolution(frame_redetect)

                    # Re Detect
                    redetects_candidate = self.models["personDetect"].Detect(
                        frame_redetect)
                    if len(redetects_candidate) == 0:
                        continue
                    redetects = redetects_candidate[
                        redetects_candidate[:, 4] >
                        self.configs["redetect_threshold"]]
                    if len(redetects) == 0:
                        continue
                    redetects[:, 0:4] = redetects[:, 0:4] / 4
                    redetects[:, 0:4] = redetects[:, 0:4] + [
                        detect[0], detect[1], detect[0], detect[1]
                    ]
                    detects = np.vstack((detects, redetects))
        else:
            detects = []

        if len(detects) != 0:
            if self.configs.get("recovery_body") is not None and self.configs[
                    "recovery_body"]:
                for detect in detects:
                    crop_x1, crop_y1, crop_x2, crop_y2 = [
                        int(detect[0]),
                        int(detect[1]),
                        int(detect[2]),
                        int(detect[3])
                    ]
                    print((crop_x2 - crop_x1), (crop_y2 - crop_y1))
                    if (crop_x2 - crop_x1) * (crop_y2 - crop_y1) > 480000:
                        continue
                    frame_crop = image[crop_y1:crop_y2, crop_x1:crop_x2]
                    if cv2.Laplacian(frame_crop, cv2.CV_64F).var(
                    ) < self.configs["recovery_body_threshold"]:
                        if self.configs["recovery_body_method"] == "sr":
                            frame_crop = self.models[
                                "superResolution"].SuperResolution(frame_crop)
                            image[crop_y1:crop_y2,
                                  crop_x1:crop_x2] = cv2.resize(
                                      frame_crop,
                                      (crop_x2 - crop_x1, crop_y2 - crop_y1),
                                      interpolation=cv2.INTER_AREA)
                            print("SR Recovery")
                        elif self.configs["recovery_body_method"] == "deblur":
                            frame_crop = self.models["deblur"].Deblur(
                                frame_crop)
                            image[
                                crop_y1:crop_y2, crop_x1:
                                crop_x2] = frame_crop  #cv2.resize(frame_crop, (crop_x2-crop_x1, crop_y2-crop_y1), interpolation=cv2.INTER_AREA)
                            print("Deblur Recovery")
                        else:
                            print("Recovery Parameter Error!")

        if len(detects) != 0:
            if self.configs.get("recovery_face") is not None and self.configs[
                    "recovery_face"]:
                ""

        if len(detects) != 0:
            # xyxy -> xywh
            tlwhs = detects.copy()[:, :4]
            tlwhs[:, 2:4] = tlwhs[:, 2:4] - tlwhs[:, 0:2]
            #tlwhs = tlwhs.tolist()
        else:
            tlwhs = []
        """step 1: prediction"""
        for strack in itertools.chain(self.tracked_stracks, self.lost_stracks):
            strack.predict()
        """step 2: scoring and selection"""
        if det_scores is None:
            det_scores = np.ones(len(tlwhs), dtype=float)
        detections = [
            STrack(tlwh, score, from_det=True)
            for tlwh, score in zip(tlwhs, det_scores)
        ]

        if self.classifier is None:
            pred_dets = []
        else:
            self.classifier.update(image)

            n_dets = len(tlwhs)
            if self.use_tracking:
                tracks = [
                    STrack(t.self_tracking(image),
                           t.tracklet_score(),
                           from_det=False) for t in itertools.chain(
                               self.tracked_stracks, self.lost_stracks)
                    if t.is_activated
                ]
                detections.extend(tracks)
            rois = np.asarray([d.tlbr for d in detections], dtype=np.float32)

            cls_scores = self.classifier.predict(rois)
            scores = np.asarray([d.score for d in detections], dtype=np.float)
            scores[0:n_dets] = 1.
            scores = scores * cls_scores
            # nms
            if len(detections) > 0:
                keep = nms_detections(rois, scores.reshape(-1), nms_thresh=0.3)
                mask = np.zeros(len(rois), dtype=np.bool)
                mask[keep] = True
                keep = np.where(mask & (scores >= self.min_cls_score))[0]
                detections = [detections[i] for i in keep]
                scores = scores[keep]
                for d, score in zip(detections, scores):
                    d.score = score
            pred_dets = [d for d in detections if not d.from_det]
            detections = [d for d in detections if d.from_det]

        # set features
        tlbrs = [det.tlbr for det in detections]
        features = extract_reid_features(self.reid_model, image, tlbrs)
        features = features.cpu().numpy()
        for i, det in enumerate(detections):
            det.set_feature(features[i])

        # <custom>set face feature
        detects_face = []
        if (self.configs.get("faceRecognition_1") is not None
                and self.configs["faceRecognition_1"]) or (
                    self.configs.get("faceRecognition_2") is not None
                    and self.configs["faceRecognition_2"]):
            detects_face = self.models["faceDetect"].Detect(image)
            for det in detections:
                person_x1, person_y1, person_x2, person_y2 = det.tlbr[:4]
                faces = []
                for detect_face in detects_face:
                    face_x1, face_y1, face_x2, face_y2 = detect_face[:4]
                    if person_x1 < face_x1 and person_y1 < face_y1 and person_x2 > face_x2 and person_y2 > face_x2:
                        faces.append(detect_face)
                if len(faces) == 1:
                    #cv2.imshow('face',image[int(faces[0][1]):int(faces[0][3]), int(faces[0][0]):int(faces[0][2])])
                    #if cv2.waitKey(1) & 0xFF == ord('q'):
                    #return

                    faceFeature = self.models["faceEmbed"].faceEmbedding(
                        image[int(faces[0][1]):int(faces[0][3]),
                              int(faces[0][0]):int(faces[0][2])])
                    det.set_faceFeature(faceFeature)
        """step 3: association for tracked"""
        # matching for tracked targets
        unconfirmed = []
        tracked_stracks = []  # type: list[STrack]
        for track in self.tracked_stracks:
            if not track.is_activated:
                unconfirmed.append(track)
            else:
                tracked_stracks.append(track)

        dists = matching.nearest_reid_distance(tracked_stracks,
                                               detections,
                                               metric='euclidean')
        dists = matching.gate_cost_matrix(self.kalman_filter, dists,
                                          tracked_stracks, detections)
        matches, u_track, u_detection = matching.linear_assignment(
            dists, thresh=self.min_ap_dist)
        for itracked, idet in matches:
            tracked_stracks[itracked].update(detections[idet], self.frame_id,
                                             image)

        # <custom> face matching for tracked targets
        if self.configs.get("faceRecognition_1") is not None and self.configs[
                "faceRecognition_1"]:
            '''
            detections : 전체 탐지 후보(수정되어도됨)
            tracked_stracks : 전체 트랙 후보(수정되면안됨)
            '''
            cost_matrix = np.zeros((len(u_track), len(u_detection)),
                                   dtype=np.float)
            for idx_idx_track, idx_track in enumerate(u_track):
                if len(tracked_stracks[idx_track].faceFeatures) == 0:
                    cost_matrix[idx_idx_track, :] = 3
                    continue
                for idx_idx_detect, idx_detect in enumerate(u_detection):
                    if detections[idx_detect].curr_faceFeature is None:
                        cost_matrix[idx_idx_track, idx_idx_detect] = 3
                        continue

                    if self.configs["faceRecognition_1_method"] == "last":
                        cost_matrix[
                            idx_idx_track, idx_idx_detect] = np.linalg.norm(
                                tracked_stracks[idx_track].faceFeatures[-1] -
                                detections[idx_detect].curr_faceFeature)
                    if self.configs["faceRecognition_1_method"] == "min":
                        min = self.configs["faceRecognition_1_threshold"] + 1e-5
                        for faceFeature in tracked_stracks[
                                idx_track].faceFeatures:
                            distance = np.linalg.norm(
                                faceFeature -
                                detections[idx_detect].curr_faceFeature)
                            if min > distance:
                                min = distance
                        cost_matrix[idx_idx_track, idx_idx_detect] = min
                    if self.configs["faceRecognition_1_method"] == "mean":
                        mean = 0
                        for faceFeature in tracked_stracks[
                                idx_track].faceFeatures:
                            mean += np.linalg.norm(
                                faceFeature -
                                detections[idx_detect].curr_faceFeature)
                        mean = mean / len(
                            tracked_stracks[idx_track].faceFeatures)
                        cost_matrix[idx_idx_track, idx_idx_detect] = mean

            cost_matrix[
                cost_matrix >
                self.configs["faceRecognition_1_threshold"]] = self.configs[
                    "faceRecognition_1_threshold"] + 1e-5  # 너무 outsider는 threshold 만큼으로 고정
            row_indices, col_indices = linear_sum_assignment(cost_matrix)

            unmatched_tracks_remove = []
            unmatched_detections_remove = []
            for row, col in zip(
                    row_indices,
                    col_indices):  # for row, col in indices: 에서 교체됨
                if cost_matrix[
                        row,
                        col] < self.configs["faceRecognition_1_threshold"]:
                    tracked_stracks[u_track[row]].update(
                        detections[u_detection[col]], self.frame_id, image)
                    unmatched_tracks_remove.append(u_track[row])
                    unmatched_detections_remove.append(u_detection[col])

            u_track = [
                track for track in u_track
                if track not in unmatched_tracks_remove
            ]
            u_detection = [
                detect for detect in u_detection
                if detect not in unmatched_detections_remove
            ]

        # matching for missing targets
        detections = [detections[i] for i in u_detection]
        dists = matching.nearest_reid_distance(self.lost_stracks,
                                               detections,
                                               metric='euclidean')
        dists = matching.gate_cost_matrix(self.kalman_filter, dists,
                                          self.lost_stracks, detections)
        matches, u_lost, u_detection = matching.linear_assignment(
            dists, thresh=self.min_ap_dist)
        for ilost, idet in matches:
            track = self.lost_stracks[ilost]  # type: STrack
            det = detections[idet]
            track.re_activate(det,
                              self.frame_id,
                              image,
                              new_id=not self.use_refind)
            refind_stracks.append(track)

        # <custom> face matching for missing targets
        if self.configs.get("faceRecognition_2") is not None and self.configs[
                "faceRecognition_2"]:
            cost_matrix = np.zeros((len(u_lost), len(u_detection)),
                                   dtype=np.float)

            for idx_idx_track, idx_track in enumerate(u_lost):
                if len(self.lost_stracks[idx_track].faceFeatures) == 0:
                    cost_matrix[idx_idx_track, :] = 3
                    continue
                for idx_idx_detect, idx_detect in enumerate(u_detection):
                    if detections[idx_detect].curr_faceFeature is None:
                        cost_matrix[idx_idx_track, idx_idx_detect] = 3
                        continue

                    if self.configs["faceRecognition_2_method"] == "last":
                        cost_matrix[
                            idx_idx_track, idx_idx_detect] = np.linalg.norm(
                                self.lost_stracks[idx_track].faceFeatures[-1] -
                                detections[idx_detect].curr_faceFeature)
                    if self.configs["faceRecognition_2_method"] == "min":
                        min = self.configs["faceRecognition_2_threshold"] + 1e-5
                        for faceFeature in self.lost_stracks[
                                idx_track].faceFeatures:
                            distance = np.linalg.norm(
                                faceFeature -
                                detections[idx_detect].curr_faceFeature)
                            if min > distance:
                                min = distance
                        cost_matrix[idx_idx_track, idx_idx_detect] = min
                    if self.configs["faceRecognition_2_method"] == "mean":
                        mean = 0
                        for faceFeature in self.lost_stracks[
                                idx_track].faceFeatures:
                            mean += np.linalg.norm(
                                faceFeature -
                                detections[idx_detect].curr_faceFeature)
                        mean = mean / len(
                            self.lost_stracks[idx_track].faceFeatures)
                        cost_matrix[idx_idx_track, idx_idx_detect] = mean

            cost_matrix[
                cost_matrix >
                self.configs["faceRecognition_2_threshold"]] = self.configs[
                    "faceRecognition_2_threshold"] + 1e-5  # 너무 outsider는 threshold 만큼으로 고정
            row_indices, col_indices = linear_sum_assignment(cost_matrix)

            unmatched_tracks_remove = []
            unmatched_detections_remove = []
            for row, col in zip(
                    row_indices,
                    col_indices):  # for row, col in indices: 에서 교체됨
                if cost_matrix[
                        row,
                        col] < self.configs["faceRecognition_2_threshold"]:
                    #self.lost_stracks[u_lost[row]].update(detections[u_detection[col]], self.frame_id, image)
                    unmatched_tracks_remove.append(u_lost[row])
                    unmatched_detections_remove.append(u_detection[col])

                    track = self.lost_stracks[u_lost[row]]  # type: STrack
                    det = detections[u_detection[col]]
                    track.re_activate(det,
                                      self.frame_id,
                                      image,
                                      new_id=not self.use_refind)
                    refind_stracks.append(track)

            u_lost = [
                track for track in u_lost
                if track not in unmatched_tracks_remove
            ]
            u_detection = [
                detect for detect in u_detection
                if detect not in unmatched_detections_remove
            ]

        # remaining tracked
        # tracked
        len_det = len(u_detection)
        detections = [detections[i] for i in u_detection] + pred_dets
        r_tracked_stracks = [tracked_stracks[i] for i in u_track]
        dists = matching.iou_distance(r_tracked_stracks, detections)
        matches, u_track, u_detection = matching.linear_assignment(dists,
                                                                   thresh=0.7)
        for itracked, idet in matches:
            r_tracked_stracks[itracked].update(detections[idet],
                                               self.frame_id,
                                               image,
                                               update_feature=True)
        for it in u_track:
            track = r_tracked_stracks[it]
            track.mark_lost()
            lost_stracks.append(track)

        # unconfirmed
        detections = [detections[i] for i in u_detection if i < len_det]
        dists = matching.iou_distance(unconfirmed, detections)
        matches, u_unconfirmed, u_detection = matching.linear_assignment(
            dists, thresh=0.7)
        for itracked, idet in matches:
            unconfirmed[itracked].update(detections[idet],
                                         self.frame_id,
                                         image,
                                         update_feature=True)
        for it in u_unconfirmed:
            track = unconfirmed[it]
            track.mark_removed()
            removed_stracks.append(track)
        """step 4: init new stracks"""
        for inew in u_detection:
            track = detections[inew]
            if not track.from_det or track.score < 0.6:
                continue
            track.activate(self.kalman_filter, self.frame_id, image)
            activated_starcks.append(track)
        """step 6: update state"""
        for track in self.lost_stracks:
            if self.frame_id - track.end_frame > self.max_time_lost:
                track.mark_removed()
                removed_stracks.append(track)

        self.tracked_stracks = [
            t for t in self.tracked_stracks if t.state == TrackState.Tracked
        ]
        self.lost_stracks = [
            t for t in self.lost_stracks if t.state == TrackState.Lost
        ]  # type: list[STrack]
        self.tracked_stracks.extend(activated_starcks)
        self.tracked_stracks.extend(refind_stracks)
        self.lost_stracks.extend(lost_stracks)
        self.removed_stracks.extend(removed_stracks)

        # output_stracks = self.tracked_stracks + self.lost_stracks

        # get scores of lost tracks
        rois = np.asarray([t.tlbr for t in self.lost_stracks],
                          dtype=np.float32)
        lost_cls_scores = self.classifier.predict(rois)
        out_lost_stracks = [
            t for i, t in enumerate(self.lost_stracks)
            if lost_cls_scores[i] > 0.3 and self.frame_id - t.end_frame <= 4
        ]
        output_tracked_stracks = [
            track for track in self.tracked_stracks if track.is_activated
        ]

        output_stracks = output_tracked_stracks + out_lost_stracks

        logger.debug('===========Frame {}=========='.format(self.frame_id))
        logger.debug('Activated: {}'.format(
            [track.track_id for track in activated_starcks]))
        logger.debug('Refind: {}'.format(
            [track.track_id for track in refind_stracks]))
        logger.debug('Lost: {}'.format(
            [track.track_id for track in lost_stracks]))
        logger.debug('Removed: {}'.format(
            [track.track_id for track in removed_stracks]))

        return output_stracks