diff --git a/docker-compose.yml b/docker-compose.yml index 126c0d717bd3952d20e029fa00303d0dd407a6f4..a30c938e55e25a07c0a74ed7de1b8383da8e86ff 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -128,7 +128,7 @@ services: --username "${DEVBEH_KAFKA_USERNAME:-}" --password "${DEVBEH_KAFKA_PASSWORD:-}" --batch_size ${DEVBEH_TRACKING_BATCH_SIZE:-1} - --frames_commit_latency ${DEVBEH_TRACKING_FRAMES_COMMIT_LATENCY:-100} + --max_frames_lost ${DEVBEH_MAX_FRAMES_LOST:-50} --name ${DEVBEH_TRACKING_INSTANCE_NAME:-tracking}' networks: kafka_net: diff --git a/test/tracking_test/docker-compose.yml b/test/tracking_test/docker-compose.yml index e191d271a7a74e28ec3aeb05842235fbcfa8a507..97d0424b34b4987669b4a1cf6e9fdfdc8f15c315 100644 --- a/test/tracking_test/docker-compose.yml +++ b/test/tracking_test/docker-compose.yml @@ -57,7 +57,6 @@ services: command: 'python run.py tracking --log_dir /opt/tracking/logs --fairmot_weights_url ${DEVBEH_FAIRMOT_WEIGHTS_URL:-https://drive.google.com/u/0/uc?export=download&confirm=VHcw&id=1iqRQjsG9BawIl8SlFomMg5iwkb6nqSpi} --use_fairmot - --frames_commit_latency 100 --bootstrap_servers kafka:29097' networks: tracking_net: diff --git a/tracking/docker-build-context/fairmot/src/lib/tracker/multitracker.py b/tracking/docker-build-context/fairmot/src/lib/tracker/multitracker.py index 26ee99ca35251ac116467e36740de6e56091b47c..c1e78dec45ae3424f438c402aa5a4cc3a952bb45 100644 --- a/tracking/docker-build-context/fairmot/src/lib/tracker/multitracker.py +++ b/tracking/docker-build-context/fairmot/src/lib/tracker/multitracker.py @@ -1,3 +1,4 @@ +import logging import numpy as np from collections import deque import itertools @@ -22,7 +23,7 @@ from models.utils import _tranpose_and_gather_feat class STrack(BaseTrack): shared_kalman = KalmanFilter() - def __init__(self, tlwh, score, temp_feat, buffer_size=30): + def __init__(self, tlwh, score, temp_feat, next_id_supplier, buffer_size=30): # wait activate self._tlwh = np.asarray(tlwh, dtype=np.float) @@ -38,6 +39,7 @@ class STrack(BaseTrack): self.update_features(temp_feat) self.features = deque([], maxlen=buffer_size) self.alpha = 0.9 + self.next_id_supplier = next_id_supplier def update_features(self, feat): feat /= np.linalg.norm(feat) @@ -71,7 +73,7 @@ class STrack(BaseTrack): def activate(self, kalman_filter, frame_id): """Start a new tracklet""" self.kalman_filter = kalman_filter - self.track_id = self.next_id() + self.track_id = self.next_id_supplier() self.mean, self.covariance = self.kalman_filter.initiate(self.tlwh_to_xyah(self._tlwh)) self.tracklet_len = 0 @@ -94,7 +96,7 @@ class STrack(BaseTrack): self.is_activated = True self.frame_id = frame_id if new_id: - self.track_id = self.next_id() + self.track_id = self.next_id_supplier() self.score = new_track.score self.score_list.append(self.score) @@ -192,9 +194,15 @@ class JDETracker(object): self.max_per_image = opt.K self.mean = np.array(opt.mean, dtype=np.float32).reshape(1, 1, 3) self.std = np.array(opt.std, dtype=np.float32).reshape(1, 1, 3) + self.id_counter = 1 self.kalman_filter = KalmanFilter() + def next_id(self): + t_id = self.id_counter + self.id_counter += 1 + return t_id + def post_process(self, dets, meta): dets = dets.detach().cpu().numpy() dets = dets.reshape(1, -1, dets.shape[2]) @@ -221,7 +229,7 @@ class JDETracker(object): results[j] = results[j][keep_inds] return results - def update(self, im_blob, img0): + def update(self, im_blob, img0, fairmot_detection_info, use_detections): self.frame_id += 1 activated_starcks = [] refind_stracks = [] @@ -244,11 +252,33 @@ class JDETracker(object): hm = output['hm'].sigmoid_() wh = output['wh'] id_feature = output['id'] + # logging.info("id_feature " + str(id_feature.shape) + " " + str(id_feature)) id_feature = F.normalize(id_feature, dim=1) reg = output['reg'] if self.opt.reg_offset else None - dets, inds = mot_decode(hm, wh, reg=reg, ltrb=self.opt.ltrb, K=self.opt.K) + if not use_detections: + dets, inds = mot_decode(hm, wh, reg=reg, ltrb=self.opt.ltrb, K=self.opt.K) + else: + dets = np.zeros((fairmot_detection_info.shape[0], 6), dtype=np.float32) + dets[:, :4] = fairmot_detection_info[:, :4] / 4 + dets[:, 4] = fairmot_detection_info[:, 4] + # 0 - класс человека, но, кажется, он не используется + dets[:, 5] = np.zeros((fairmot_detection_info.shape[0]), dtype=np.float32) + dets = np.expand_dims(dets, axis=0) + dets = torch.from_numpy(dets).cuda() + + inds = np.zeros((1, dets.shape[1]), dtype=int) + for i in range(dets.shape[1]): + c_x = int((dets[0, i, 2] + dets[0, i, 0]) / 2) + c_y = int((dets[0, i, 3] + dets[0, i, 1]) / 2) + # Индексы преобразованы в одно число, width - 1088/4=272 + inds[0, i] = c_y * 272 + c_x + inds = torch.from_numpy(inds).cuda() + + # logging.info("inds " + str(inds.shape) + " " + str(inds)) + # logging.info("dets " + str(dets.shape) + " "+ str(dets)) id_feature = _tranpose_and_gather_feat(id_feature, inds) + # logging.info("id_feature2 " + str(id_feature.shape) + " " + str(id_feature)) id_feature = id_feature.squeeze(0) id_feature = id_feature.cpu().numpy() @@ -265,6 +295,7 @@ class JDETracker(object): dets = dets[remain_inds] id_feature = id_feature[remain_inds] + # logging.info("remain ids " + str(remain_inds)) # vis ''' for i in range(0, dets.shape[0]): @@ -279,7 +310,7 @@ class JDETracker(object): if len(dets) > 0: '''Detections''' - detections = [STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f, 30) for + detections = [STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f, self.next_id, 30) for (tlbrs, f) in zip(dets[:, :5], id_feature)] else: detections = [] @@ -332,7 +363,7 @@ class JDETracker(object): # association the untrack to the low score detections if len(dets_second) > 0: '''Detections''' - detections_second = [STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f, 30) for + detections_second = [STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f, self.next_id, 30) for (tlbrs, f) in zip(dets_second[:, :5], id_feature_second)] else: detections_second = [] @@ -401,6 +432,8 @@ class JDETracker(object): logger.debug('Lost: {}'.format([track.track_id for track in lost_stracks])) logger.debug('Removed: {}'.format([track.track_id for track in removed_stracks])) + # logging.info("Out count: " + str(len(output_stracks))) + return output_stracks diff --git a/tracking/run.py b/tracking/run.py index e4ab30b68930de1ec70b26e95d28022954397c85..2ecd4bbb9d454924dce7dae22489958f0219477f 100644 --- a/tracking/run.py +++ b/tracking/run.py @@ -22,22 +22,23 @@ PERSON_CLASS = 1 # noinspection DuplicatedCode def process_tracking_requests(requests: Iterable[PerceptionRequest], results: Iterable[PerceptionResponse], - tracker_factory, fairmot_tracker_factory, use_fairmot, trackers): + tracker_factory, fairmot_tracker_factory, use_fairmot, use_detections, trackers): for req, res in zip(requests, results): if req.finished: - track(req.manager_id, req.video_id, tracker_factory, fairmot_tracker_factory, use_fairmot, trackers, None, - None) + track(req.manager_id, req.video_id, tracker_factory, fairmot_tracker_factory, use_fairmot, use_detections, + trackers, None, None) else: image = np.fromstring(req.tracking.image, np.uint8) image = cv2.imdecode(image, cv2.IMREAD_ANYCOLOR) logging.info('Processing image of shape %s', image.shape) tracking_entries = track(req.manager_id, req.video_id, tracker_factory, fairmot_tracker_factory, - use_fairmot, trackers, image, req.tracking.entries) + use_fairmot, use_detections, trackers, image, req.tracking.entries) for tracking_entry in tracking_entries: res.tracking.entries.append(tracking_entry) -def track(manager_id, video_id, tracker_factory, fairmot_tracker_factory, use_fairmot, trackers, image, detections): +def track(manager_id, video_id, tracker_factory, fairmot_tracker_factory, use_fairmot, use_detections, trackers, image, + detections): """ Обрабатывает очередной кадр видеозаписи, сопоставляя объекты с прошлым кадром. :param manager_id: идентификатор менеджера, передающего видепоток @@ -45,6 +46,8 @@ def track(manager_id, video_id, tracker_factory, fairmot_tracker_factory, use_fa :param tracker_factory: фабрика для создания трекеров :param fairmot_tracker_factory: фабрика для создания трекеров fairmot или None :param use_fairmot: следует ли использовать fairmot для отслеживания людей + :param use_detections: следует ли использовать рамки, переданные в качестве входных данных, даже + при использовании fairmot :param trackers: хранит трекеры для видеозаписей {(manager_id, video_id): tracker} :param image: изображение или None, если видеопоток завершился :param detections: набор DetectionEntry @@ -100,17 +103,26 @@ def track(manager_id, video_id, tracker_factory, fairmot_tracker_factory, use_fa if use_fairmot: img0 = image - img, _, _, _ = letterbox(img0) + img, ratio, dw, dh = letterbox(img0) + border_top = round(dh - 0.1) + border_left = round(dw - 0.1) + + # Correcting detected boxes + for i in range(fairmot_detection_info.shape[0]): + fairmot_detection_info[i, 0] = round(fairmot_detection_info[i, 0] * ratio) + border_left + fairmot_detection_info[i, 1] = round(fairmot_detection_info[i, 1] * ratio) + border_top + fairmot_detection_info[i, 2] = round(fairmot_detection_info[i, 2] * ratio) + border_left + fairmot_detection_info[i, 3] = round(fairmot_detection_info[i, 3] * ratio) + border_top # Normalize RGB img = img[:, :, ::-1].transpose(2, 0, 1) img = np.ascontiguousarray(img, dtype=np.float32) img /= 255.0 - blob = torch.from_numpy(img).cuda().unsqueeze(0) fairmot_tracker: JDETracker = trackers[tracker_key][1] - fairmot_tracks: Iterable[FairmotSTrack] = fairmot_tracker.update(blob, img0) + fairmot_tracks: Iterable[FairmotSTrack] = fairmot_tracker.update(blob, img0, fairmot_detection_info, + use_detections) for t in fairmot_tracks: tlbr = t.tlbr @@ -140,21 +152,22 @@ def correct_track_id(id, use_fairmot): return id * 2 + 1 if use_fairmot else id -def get_tracker_factory(fps): - return lambda: create_tracker(fps) +def get_tracker_factory(max_frames_lost): + return lambda: create_tracker(max_frames_lost) -def create_tracker(fps): +def create_tracker(max_frames_lost): args = mock.Mock() # todo: пробовать разные параметры args.track_thresh = 0.7 - args.track_buffer = 150 + fps = 30 # чтобы компенсировать деление на 30 в трэкере + args.track_buffer = max_frames_lost args.mot20 = False # todo: попробовать True тоже args.match_thresh = 0.9 return BYTETracker(args, fps) -def get_fairmot_tracker_factory(fairmot_weights_path, fps): +def get_fairmot_tracker_factory(fairmot_weights_path, max_frames_lost): device = torch.device('cuda') ltrb = True reid_dim = 128 @@ -171,10 +184,10 @@ def get_fairmot_tracker_factory(fairmot_weights_path, fps): model = model.to(device) model.eval() - return lambda: create_fairmot_tracker(model, fps) + return lambda: create_fairmot_tracker(model, max_frames_lost) -def create_fairmot_tracker(model, fps): +def create_fairmot_tracker(model, max_frames_lost): opt = mock.Mock() # mean и std не используются opt.mean = [0.408, 0.447, 0.470] @@ -184,7 +197,7 @@ def create_fairmot_tracker(model, fps): # 0.4 - дефолт. Попробовать opt.conf_thres = 0.6 # todo: попробовать другие параметры (30 - default) - opt.track_buffer = 150 + opt.track_buffer = max_frames_lost opt.reid_dim = 128 opt.num_classes = 1 opt.nID = 14455 @@ -202,6 +215,8 @@ def create_fairmot_tracker(model, fps): opt.output_res = max(opt.output_h, opt.output_w) opt.match_thres = 0.4 + # чтобы компенсировать деление на 30 в трэкере + fps=30 return JDETracker(opt, model, frame_rate=fps) @@ -211,13 +226,16 @@ def run(parser: argparse.ArgumentParser): parser.add_argument('--log_dir', help='logs dir path', required=True) parser.add_argument('--logging_level', dest='logging_level', choices=logging_levels.keys(), default='INFO', help='logging level. One of ' + str(logging_levels.keys())) - parser.add_argument('--fairmot_weights_path', default='',# todo + parser.add_argument('--fairmot_weights_path', default='', help='path to the model weights file. Default: ' 'PROJECT_DIR_PATH/tracking/models/fairmot_dla34.pth') parser.add_argument('--fairmot_weights_url', default='', help='url to the model weights file that ' 'should be downloaded and placed to the ' 'fairmot_weights_path if it does not already exist') parser.add_argument('--use_fairmot', action='store_true', help='enables application of fairmot for people tracking') + parser.add_argument('--use_detections', action='store_true', help='the option is used together with the use_fairmot' + ' option. It enables the use of provided ' + 'detections instead of results of the fairmot') parser.add_argument('--bootstrap_servers', type=str, required=True, help='comma separated kafka bootstrap servers. Example: kafka1.local:9092,kafka2.local:9092') parser.add_argument('--username', type=str, default='', help='kafka SASL username or empty str') @@ -226,8 +244,8 @@ def run(parser: argparse.ArgumentParser): parser.add_argument('--group_id', default='tracking', help='kafka clients group id') parser.add_argument('--result_topic_prefix', default='tracking-results-', help='prefix of the result topic names') parser.add_argument('--batch_size', type=int, default=1, help='batch size for frames processing') - parser.add_argument('--frames_commit_latency', type=int, default=10, help='learning interval in frames') - parser.add_argument('--fps', type=int, default=1, help='fps') + parser.add_argument('--max_frames_lost', type=int, default=50, help='max number of frames a person can be missed ' + 'and reidentified later') parser.add_argument('--commit_period_ms', type=int, default=1000, help='period for committing processed and delivered request offsets') args = parser.parse_args() @@ -248,8 +266,8 @@ def run(parser: argparse.ArgumentParser): logging.info('Fairmot model weights download finished') logging.info("Initializing tracking models") - tracker_factory = get_tracker_factory(args.fps) - fairmot_tracker_factory = get_fairmot_tracker_factory(fairmot_weights_path, args.fps) if args.use_fairmot else None + tracker_factory = get_tracker_factory(args.max_frames_lost) + fairmot_tracker_factory = get_fairmot_tracker_factory(fairmot_weights_path, args.max_frames_lost) if args.use_fairmot else None logging.info("Model was initialized") trackers = {} @@ -260,11 +278,12 @@ def run(parser: argparse.ArgumentParser): username=args.username, password=args.password, topic_name=args.topic, group_id=args.group_id, result_topic_prefix=args.result_topic_prefix, req_batch_size=args.batch_size, commit_period_ms=args.commit_period_ms, - stream_max_idle_time_ms=600_000, frames_commit_latency=args.frames_commit_latency, + stream_max_idle_time_ms=600_000, frames_commit_latency=args.max_frames_lost, request_handler=lambda req, res: process_tracking_requests(req, res, tracker_factory, fairmot_tracker_factory, args.use_fairmot, + args.use_detections, trackers)) logging.info("Starting tracking service '%s'", args.name)