From b2a78b5ad08b7de0d4e7dd20ec600348c50abd65 Mon Sep 17 00:00:00 2001 From: "lanjinpeng.ljp" Date: Mon, 9 Jan 2023 07:11:15 +0800 Subject: [PATCH] =?UTF-8?q?=E6=94=AF=E6=8C=81=E8=A7=86=E9=A2=91=E5=A4=9A?= =?UTF-8?q?=E7=9B=AE=E6=A0=87=E8=B7=9F=E8=B8=AA=E6=A8=A1=E5=9E=8B=20=20=20?= =?UTF-8?q?=20=20=20=20=20=20Link:=20https://code.alibaba-inc.com/Ali-MaaS?= =?UTF-8?q?/MaaS-lib/codereview/11249098?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- data/test/videos/MOT17-03-partial.mp4 | 3 + modelscope/metainfo.py | 1 + .../video_multi_object_tracking/__init__.py | 22 + .../models/__init__.py | 0 .../models/common.py | 104 +++++ .../models/decode.py | 73 +++ .../models/model.py | 52 +++ .../models/yolo.py | 149 +++++++ .../tracker/__init__.py | 0 .../tracker/basetrack.py | 55 +++ .../tracker/matching.py | 95 ++++ .../tracker/multitracker.py | 418 ++++++++++++++++++ .../utils/__init__.py | 0 .../utils/image.py | 73 +++ .../utils/kalman_filter.py | 264 +++++++++++ .../utils/utils.py | 208 +++++++++ .../utils/visualization.py | 85 ++++ modelscope/outputs/outputs.py | 13 + modelscope/pipeline_inputs.py | 2 + modelscope/pipelines/builder.py | 3 + .../video_multi_object_tracking_pipeline.py | 81 ++++ modelscope/utils/constant.py | 1 + requirements/cv.txt | 1 + .../test_video_multi_object_tracking.py | 39 ++ 24 files changed, 1742 insertions(+) create mode 100644 data/test/videos/MOT17-03-partial.mp4 create mode 100644 modelscope/models/cv/video_multi_object_tracking/__init__.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/models/__init__.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/models/common.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/models/decode.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/models/model.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/models/yolo.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/tracker/__init__.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/tracker/basetrack.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/tracker/matching.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/tracker/multitracker.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/utils/__init__.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/utils/image.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/utils/kalman_filter.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/utils/utils.py create mode 100644 modelscope/models/cv/video_multi_object_tracking/utils/visualization.py create mode 100644 modelscope/pipelines/cv/video_multi_object_tracking_pipeline.py create mode 100644 tests/pipelines/test_video_multi_object_tracking.py diff --git a/data/test/videos/MOT17-03-partial.mp4 b/data/test/videos/MOT17-03-partial.mp4 new file mode 100644 index 00000000..ed1ee7ba --- /dev/null +++ b/data/test/videos/MOT17-03-partial.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:ca26615762e3f4ccca53a020efe73c3cf3598edc68bb68b5555c24e815718336 +size 3151767 diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index b59ec351..c940b2ba 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -244,6 +244,7 @@ class Pipelines(object): crowd_counting = 'hrnet-crowd-counting' action_detection = 'ResNetC3D-action-detection' video_single_object_tracking = 'ostrack-vitb-video-single-object-tracking' + video_multi_object_tracking = 'video-multi-object-tracking' image_panoptic_segmentation = 'image-panoptic-segmentation' image_panoptic_segmentation_easycv = 'image-panoptic-segmentation-easycv' video_summarization = 'googlenet_pgl_video_summarization' diff --git a/modelscope/models/cv/video_multi_object_tracking/__init__.py b/modelscope/models/cv/video_multi_object_tracking/__init__.py new file mode 100644 index 00000000..985a8da5 --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/__init__.py @@ -0,0 +1,22 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +from typing import TYPE_CHECKING + +from modelscope.utils.import_utils import LazyImportModule + +if TYPE_CHECKING: + from .video_multi_object_tracking import VideoMultiObjectTracking + +else: + _import_structure = { + 'video_multi_object_tracking': ['VideoMultiObjectTracking'], + } + + import sys + + sys.modules[__name__] = LazyImportModule( + __name__, + globals()['__file__'], + _import_structure, + module_spec=__spec__, + extra_objects={}, + ) diff --git a/modelscope/models/cv/video_multi_object_tracking/models/__init__.py b/modelscope/models/cv/video_multi_object_tracking/models/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/video_multi_object_tracking/models/common.py b/modelscope/models/cv/video_multi_object_tracking/models/common.py new file mode 100644 index 00000000..e5eaa686 --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/models/common.py @@ -0,0 +1,104 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT + +import torch +import torch.nn as nn + + +def autopad(k, p=None): + if p is None: + p = k // 2 if isinstance(k, int) else [x // 2 for x in k] + return p + + +class Conv(nn.Module): + + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): + super(Conv, self).__init__() + self.conv = nn.Conv2d( + c1, c2, k, s, autopad(k, p), groups=g, bias=False) + self.bn = nn.BatchNorm2d(c2) + self.act = nn.SiLU() if act is True else ( + act if isinstance(act, nn.Module) else nn.Identity()) + + def forward(self, x): + return self.act(self.bn(self.conv(x))) + + def fuseforward(self, x): + return self.act(self.conv(x)) + + +class Bottleneck(nn.Module): + + def __init__(self, c1, c2, shortcut=True, g=1, e=0.5): + super(Bottleneck, self).__init__() + c_ = int(c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = Conv(c_, c2, 3, 1, g=g) + self.add = shortcut and c1 == c2 + + def forward(self, x): + return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x)) + + +class C3(nn.Module): + """ + CSP Bottleneck with 3 convolutions + """ + + def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5): + super(C3, self).__init__() + c_ = int(c2 * e) # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = Conv(c1, c_, 1, 1) + self.cv3 = Conv(2 * c_, c2, 1) + self.m = nn.Sequential( + *[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)]) + + def forward(self, x): + return self.cv3(torch.cat((self.m(self.cv1(x)), self.cv2(x)), dim=1)) + + +class SPP(nn.Module): + """ + Spatial pyramid pooling layer used in YOLOv3-SPP + """ + + def __init__(self, c1, c2, k=(5, 9, 13)): + super(SPP, self).__init__() + c_ = c1 // 2 # hidden channels + self.cv1 = Conv(c1, c_, 1, 1) + self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1) + self.m = nn.ModuleList( + [nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k]) + + def forward(self, x): + x = self.cv1(x) + return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) + + +class Focus(nn.Module): + """ + Focus wh information into c-space + """ + + def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): + super(Focus, self).__init__() + self.conv = Conv(c1 * 4, c2, k, s, p, g, act) + + def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) + return self.conv( + torch.cat([ + x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], + x[..., 1::2, 1::2] + ], 1)) + + +class Concat(nn.Module): + + def __init__(self, dimension=1): + super(Concat, self).__init__() + self.d = dimension + + def forward(self, x): + return torch.cat(x, self.d) diff --git a/modelscope/models/cv/video_multi_object_tracking/models/decode.py b/modelscope/models/cv/video_multi_object_tracking/models/decode.py new file mode 100644 index 00000000..4c2ad91e --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/models/decode.py @@ -0,0 +1,73 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +import torch +import torch.nn as nn + +from modelscope.models.cv.video_multi_object_tracking.utils.utils import ( + _gather_feat, _tranpose_and_gather_feat) + + +def _nms(heat, kernel=3): + pad = (kernel - 1) // 2 + + hmax = nn.functional.max_pool2d( + heat, (kernel, kernel), stride=1, padding=pad) + keep = (hmax == heat).float() + return heat * keep + + +def _topk(scores, K=40): + batch, cat, height, width = scores.size() + + topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K) + + topk_inds = topk_inds % (height * width) + topk_ys = torch.true_divide(topk_inds, width).int().float() + topk_xs = (topk_inds % width).int().float() + + topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K) + topk_clses = torch.true_divide(topk_ind, K).int() + topk_inds = _gather_feat(topk_inds.view(batch, -1, 1), + topk_ind).view(batch, K) + topk_ys = _gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K) + topk_xs = _gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K) + + return topk_score, topk_inds, topk_clses, topk_ys, topk_xs + + +def mot_decode(heat, wh, reg=None, ltrb=False, K=100): + batch, cat, height, width = heat.size() + + heat = _nms(heat) + + scores, inds, clses, ys, xs = _topk(heat, K=K) + if reg is not None: + reg = _tranpose_and_gather_feat(reg, inds) + reg = reg.view(batch, K, 2) + xs = xs.view(batch, K, 1) + reg[:, :, 0:1] + ys = ys.view(batch, K, 1) + reg[:, :, 1:2] + else: + xs = xs.view(batch, K, 1) + 0.5 + ys = ys.view(batch, K, 1) + 0.5 + wh = _tranpose_and_gather_feat(wh, inds) + if ltrb: + wh = wh.view(batch, K, 4) + else: + wh = wh.view(batch, K, 2) + clses = clses.view(batch, K, 1).float() + scores = scores.view(batch, K, 1) + if ltrb: + a = xs - wh[..., 0:1] + b = ys - wh[..., 1:2] + c = xs + wh[..., 2:3] + d = ys + wh[..., 3:4] + bboxes = torch.cat([a, b, c, d], dim=2) + else: + a = xs - wh[..., 0:1] / 2 + b = ys - wh[..., 1:2] / 2 + c = xs + wh[..., 0:1] / 2 + d = ys + wh[..., 1:2] / 2 + bboxes = torch.cat([a, b, c, d], dim=2) + detections = torch.cat([bboxes, scores, clses], dim=2) + + return detections, inds diff --git a/modelscope/models/cv/video_multi_object_tracking/models/model.py b/modelscope/models/cv/video_multi_object_tracking/models/model.py new file mode 100644 index 00000000..255e92de --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/models/model.py @@ -0,0 +1,52 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +import torch + +from modelscope.utils.logger import get_logger +from .yolo import get_pose_net as get_pose_net_yolo + +logger = get_logger() + +_model_factory = {'yolo': get_pose_net_yolo} + + +def create_model(arch, heads, head_conv): + num_layers = int(arch[arch.find('_') + 1:]) if '_' in arch else 0 + arch = arch[:arch.find('_')] if '_' in arch else arch + get_model = _model_factory[arch] + model = get_model(num_layers=num_layers, heads=heads, head_conv=head_conv) + return model + + +def load_model(model, model_path): + checkpoint = torch.load( + model_path, map_location=lambda storage, loc: storage) + state_dict_ = checkpoint['state_dict'] + state_dict = {} + + # convert data_parallal to model + for k in state_dict_: + if k.startswith('module') and not k.startswith('module_list'): + state_dict[k[7:]] = state_dict_[k] + else: + state_dict[k] = state_dict_[k] + model_state_dict = model.state_dict() + + # check loaded parameters and created model parameters + msg = 'If you see this, your model does not fully load the ' + \ + 'pre-trained weight. Please make sure ' + \ + 'you have correctly specified --arch xxx ' + \ + 'or set the correct --num_classes for your own dataset.' + for k in state_dict: + if k in model_state_dict: + if state_dict[k].shape != model_state_dict[k].shape: + state_dict[k] = model_state_dict[k] + else: + logger.info('Drop parameter {}.'.format(k) + msg) + for k in model_state_dict: + if not (k in state_dict): + logger.info('No param {}.'.format(k) + msg) + state_dict[k] = model_state_dict[k] + model.load_state_dict(state_dict, strict=False) + + return model diff --git a/modelscope/models/cv/video_multi_object_tracking/models/yolo.py b/modelscope/models/cv/video_multi_object_tracking/models/yolo.py new file mode 100644 index 00000000..df66060b --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/models/yolo.py @@ -0,0 +1,149 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +import math +from copy import deepcopy + +import torch.nn as nn + +from modelscope.models.base import TorchModel +from modelscope.utils.logger import get_logger +from .common import C3, SPP, Concat, Conv, Focus + +logger = get_logger() + +backbone_param = { + 'nc': + 80, + 'depth_multiple': + 0.33, + 'width_multiple': + 0.5, + 'backbone': [[-1, 1, 'Focus', [64, 3]], [-1, 1, 'Conv', [128, 3, 2]], + [-1, 3, 'C3', [128]], [-1, 1, 'Conv', [256, 3, 2]], + [-1, 9, 'C3', [256]], [-1, 1, 'Conv', [512, 3, 2]], + [-1, 9, 'C3', [512]], [-1, 1, 'Conv', [1024, 3, 2]], + [-1, 1, 'SPP', [1024, [5, 9, 13]]], + [-1, 3, 'C3', [1024, False]], [-1, 1, 'Conv', [512, 1, 1]], + [-1, 1, 'nn.Upsample', ['None', 2, 'nearest']], + [[-1, 6], 1, 'Concat', [1]], [-1, 3, 'C3', [512, False]], + [-1, 1, 'Conv', [256, 1, 1]], + [-1, 1, 'nn.Upsample', ['None', 2, 'nearest']], + [[-1, 4], 1, 'Concat', [1]], [-1, 3, 'C3', [256, False]], + [-1, 1, 'Conv', [128, 1, 1]], + [-1, 1, 'nn.Upsample', ['None', 2, 'nearest']], + [[-1, 2], 1, 'Concat', [1]], [-1, 3, 'C3', [128, False]]] +} + + +def fill_fc_weights(layers): + for m in layers.modules(): + if isinstance(m, nn.Conv2d): + if m.bias is not None: + nn.init.constant_(m.bias, 0) + + +class Model(nn.Module): + + def __init__(self, config=backbone_param, ch=3, nc=None, anchors=None): + super(Model, self).__init__() + self.yaml = config # model dict + + # Define model + ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels + if nc and nc != self.yaml['nc']: + self.yaml['nc'] = nc # override yaml value + self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) + self.names = [str(i) for i in range(self.yaml['nc'])] + self.inplace = self.yaml.get('inplace', True) + + def forward(self, x, augment=False, profile=False): + return self.forward_once(x, profile) + + def forward_once(self, x, profile=False): + y = [] + for m in self.model: + if m.f != -1: # if not from previous layer + x = y[m.f] if isinstance( + m.f, int) else [x if j == -1 else y[j] for j in m.f] + + x = m(x) # run + y.append(x if m.i in self.save else None) + + return x + + +def parse_model(d, ch): + gd, gw = d['depth_multiple'], d['width_multiple'] + + layers, save, c2 = [], [], ch[-1] + for i, (f, n, m, args) in enumerate(d['backbone']): + m = eval(m) if isinstance(m, str) else m + for j, a in enumerate(args): + try: + args[j] = eval(a) if isinstance(a, str) else a + except Exception: + pass + + n = max(round(n * gd), 1) if n > 1 else n + if m in [Conv, SPP, Focus, C3]: + c1, c2 = ch[f], args[0] + c2 = make_divisible(c2 * gw, 8) + + args = [c1, c2, *args[1:]] + if m in [C3]: + args.insert(2, n) + n = 1 + elif m is nn.BatchNorm2d: + args = [ch[f]] + elif m is Concat: + c2 = sum([ch[x] for x in f]) + else: + c2 = ch[f] + + m_ = nn.Sequential(*[m(*args) + for _ in range(n)]) if n > 1 else m(*args) + t = str(m)[8:-2].replace('__main__.', '') + np = sum([x.numel() for x in m_.parameters()]) + m_.i, m_.f, m_.type, m_.np = i, f, t, np + save.extend(x % i for x in ([f] if isinstance(f, int) else f) + if x != -1) + layers.append(m_) + if i == 0: + ch = [] + ch.append(c2) + return nn.Sequential(*layers), sorted(save) + + +class PoseYOLO(TorchModel): + + def __init__(self, heads): + self.heads = heads + super(PoseYOLO, self).__init__() + self.backbone = Model() + for head in sorted(self.heads): + num_output = self.heads[head] + fc = nn.Sequential( + nn.Conv2d(64, 64, kernel_size=3, padding=1, bias=True), + nn.SiLU(), + nn.Conv2d(64, num_output, kernel_size=1, stride=1, padding=0)) + self.__setattr__(head, fc) + if 'hm' in head: + fc[-1].bias.data.fill_(-2.19) + else: + fill_fc_weights(fc) + + def forward(self, x): + x = self.backbone(x) + ret = {} + for head in self.heads: + ret[head] = self.__getattr__(head)(x) + return [ret] + + +def get_pose_net(num_layers, heads, head_conv): + model = PoseYOLO(heads) + return model + + +def make_divisible(x, divisor): + return math.ceil(x / divisor) * divisor diff --git a/modelscope/models/cv/video_multi_object_tracking/tracker/__init__.py b/modelscope/models/cv/video_multi_object_tracking/tracker/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/video_multi_object_tracking/tracker/basetrack.py b/modelscope/models/cv/video_multi_object_tracking/tracker/basetrack.py new file mode 100644 index 00000000..f8763308 --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/tracker/basetrack.py @@ -0,0 +1,55 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +from collections import OrderedDict + +import numpy as np + + +class TrackState(object): + New = 0 + Tracked = 1 + Lost = 2 + Removed = 3 + + +class BaseTrack(object): + _count = 0 + + track_id = 0 + is_activated = False + state = TrackState.New + + history = OrderedDict() + features = [] + curr_feature = None + score = 0 + start_frame = 0 + frame_id = 0 + time_since_update = 0 + + # multi-camera + location = (np.inf, np.inf) + + @property + def end_frame(self): + return self.frame_id + + @staticmethod + def next_id(): + BaseTrack._count += 1 + return BaseTrack._count + + def activate(self, *args): + raise NotImplementedError + + def predict(self): + raise NotImplementedError + + def update(self, *args, **kwargs): + raise NotImplementedError + + def mark_lost(self): + self.state = TrackState.Lost + + def mark_removed(self): + self.state = TrackState.Removed diff --git a/modelscope/models/cv/video_multi_object_tracking/tracker/matching.py b/modelscope/models/cv/video_multi_object_tracking/tracker/matching.py new file mode 100644 index 00000000..45d2f5c0 --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/tracker/matching.py @@ -0,0 +1,95 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +import lap +import numpy as np +from scipy.spatial.distance import cdist + +from modelscope.models.cv.video_multi_object_tracking.utils import \ + kalman_filter +from modelscope.models.cv.video_multi_object_tracking.utils.utils import \ + bbox_iou + + +def linear_assignment(cost_matrix, thresh): + if cost_matrix.size == 0: + return np.empty((0, 2), + dtype=int), tuple(range(cost_matrix.shape[0])), tuple( + range(cost_matrix.shape[1])) + matches, unmatched_a, unmatched_b = [], [], [] + cost, x, y = lap.lapjv(cost_matrix, extend_cost=True, cost_limit=thresh) + for ix, mx in enumerate(x): + if mx >= 0: + matches.append([ix, mx]) + unmatched_a = np.where(x < 0)[0] + unmatched_b = np.where(y < 0)[0] + matches = np.asarray(matches) + return matches, unmatched_a, unmatched_b + + +def ious(atlbrs, btlbrs): + ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float) + if ious.size == 0: + return ious + + ious = bbox_iou(atlbrs, btlbrs, True).numpy() + + return ious + + +def iou_distance(atracks, btracks): + if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or ( + len(btracks) > 0 and isinstance(btracks[0], np.ndarray)): + atlbrs = atracks + btlbrs = btracks + else: + atlbrs = [track.tlbr for track in atracks] + btlbrs = [track.tlbr for track in btracks] + _ious = ious(atlbrs, btlbrs) + cost_matrix = 1 - _ious + + return cost_matrix + + +def embedding_distance(tracks, detections, metric='cosine'): + """ + Args: + tracks: list[STrack] + detections: list[BaseTrack] + metric: str + Returns: + cost_matrix: np.ndarray + """ + + cost_matrix = np.zeros((len(tracks), len(detections)), dtype=np.float) + if cost_matrix.size == 0: + return cost_matrix + det_features = np.asarray([track.curr_feat for track in detections], + dtype=np.float) + track_features = np.asarray([track.smooth_feat for track in tracks], + dtype=np.float) + cost_matrix = np.maximum(0.0, cdist(track_features, det_features, metric)) + return cost_matrix + + +def fuse_motion(kf, + cost_matrix, + tracks, + detections, + only_position=False, + lambda_=0.98): + if cost_matrix.size == 0: + return cost_matrix + gating_dim = 2 if only_position else 4 + gating_threshold = kalman_filter.chi2inv95[gating_dim] + measurements = np.asarray([det.to_xyah() for det in detections]) + for row, track in enumerate(tracks): + gating_distance = kf.gating_distance( + track.mean, + track.covariance, + measurements, + only_position, + metric='maha') + cost_matrix[row, gating_distance > gating_threshold] = np.inf + cost_matrix[row] = lambda_ * cost_matrix[row] + ( + 1 - lambda_) * gating_distance + return cost_matrix diff --git a/modelscope/models/cv/video_multi_object_tracking/tracker/multitracker.py b/modelscope/models/cv/video_multi_object_tracking/tracker/multitracker.py new file mode 100644 index 00000000..1dc3297f --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/tracker/multitracker.py @@ -0,0 +1,418 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +from collections import deque + +import numpy as np +import torch +import torch.nn.functional as F + +from modelscope.models.cv.video_multi_object_tracking.models.decode import \ + mot_decode +from modelscope.models.cv.video_multi_object_tracking.models.model import ( + create_model, load_model) +from modelscope.models.cv.video_multi_object_tracking.tracker import matching +from modelscope.models.cv.video_multi_object_tracking.tracker.basetrack import ( + BaseTrack, TrackState) +from modelscope.models.cv.video_multi_object_tracking.utils.kalman_filter import \ + KalmanFilter +from modelscope.models.cv.video_multi_object_tracking.utils.utils import ( + _tranpose_and_gather_feat, ctdet_post_process) +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +class STrack(BaseTrack): + shared_kalman = KalmanFilter() + + def __init__(self, tlwh, score, temp_feat, buffer_size=30): + + # wait activate + self._tlwh = np.asarray(tlwh, dtype=np.float) + self.kalman_filter = None + self.mean, self.covariance = None, None + self.is_activated = False + + self.score = score + self.tracklet_len = 0 + + self.smooth_feat = None + self.update_features(temp_feat) + self.features = deque([], maxlen=buffer_size) + self.alpha = 0.9 + + def update_features(self, feat): + feat /= np.linalg.norm(feat) + self.curr_feat = feat + if self.smooth_feat is None: + self.smooth_feat = feat + else: + self.smooth_feat = self.alpha * self.smooth_feat + ( + 1 - self.alpha) * feat + self.features.append(feat) + self.smooth_feat /= np.linalg.norm(self.smooth_feat) + + def predict(self): + mean_state = self.mean.copy() + if self.state != TrackState.Tracked: + mean_state[7] = 0 + self.mean, self.covariance = self.kalman_filter.predict( + mean_state, self.covariance) + + @staticmethod + def multi_predict(stracks): + if len(stracks) > 0: + multi_mean = np.asarray([st.mean.copy() for st in stracks]) + multi_covariance = np.asarray([st.covariance for st in stracks]) + for i, st in enumerate(stracks): + if st.state != TrackState.Tracked: + multi_mean[i][7] = 0 + multi_mean, multi_covariance = STrack.shared_kalman.multi_predict( + multi_mean, multi_covariance) + for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)): + stracks[i].mean = mean + stracks[i].covariance = cov + + def activate(self, kalman_filter, frame_id): + self.kalman_filter = kalman_filter + self.track_id = self.next_id() + self.mean, self.covariance = self.kalman_filter.initiate( + self.tlwh_to_xyah(self._tlwh)) + + self.tracklet_len = 0 + self.state = TrackState.Tracked + if frame_id == 1: + self.is_activated = True + self.frame_id = frame_id + self.start_frame = frame_id + + def re_activate(self, new_track, frame_id, new_id=False): + self.mean, self.covariance = self.kalman_filter.update( + self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh)) + + self.update_features(new_track.curr_feat) + self.tracklet_len = 0 + self.state = TrackState.Tracked + self.is_activated = True + self.frame_id = frame_id + if new_id: + self.track_id = self.next_id() + + def update(self, new_track, frame_id, update_feature=True): + """ + Update a matched track + Args: + new_track: STrack + frame_id: int + update_feature: bool + """ + self.frame_id = frame_id + self.tracklet_len += 1 + + new_tlwh = new_track.tlwh + self.mean, self.covariance = self.kalman_filter.update( + self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh)) + self.state = TrackState.Tracked + self.is_activated = True + + self.score = new_track.score + if update_feature: + self.update_features(new_track.curr_feat) + + @property + def tlwh(self): + """Get current position in bounding box format `(top left x, top left y, + width, height)`. + """ + if self.mean is None: + return self._tlwh.copy() + ret = self.mean[:4].copy() + ret[2] *= ret[3] + ret[:2] -= ret[2:] / 2 + return ret + + @property + def tlbr(self): + """Convert bounding box to format `(min x, min y, max x, max y)`, i.e., + `(top left, bottom right)`. + """ + ret = self.tlwh.copy() + ret[2:] += ret[:2] + return ret + + @staticmethod + def tlwh_to_xyah(tlwh): + """Convert bounding box to format `(center x, center y, aspect ratio, + height)`, where the aspect ratio is `width / height`. + """ + ret = np.asarray(tlwh).copy() + ret[:2] += ret[2:] / 2 + ret[2] /= ret[3] + return ret + + def to_xyah(self): + return self.tlwh_to_xyah(self.tlwh) + + @staticmethod + def tlbr_to_tlwh(tlbr): + ret = np.asarray(tlbr).copy() + ret[2:] -= ret[:2] + return ret + + @staticmethod + def tlwh_to_tlbr(tlwh): + ret = np.asarray(tlwh).copy() + ret[2:] += ret[:2] + return ret + + def __repr__(self): + return 'OT_{}_({}-{})'.format(self.track_id, self.start_frame, + self.end_frame) + + +class JDETracker(object): + + def __init__(self, opt, model_path, device): + self.opt = opt + self.device = device + self.model = create_model(opt.arch, opt.heads, opt.head_conv) + self.model = load_model(self.model, model_path) + self.model = self.model.to(device) + self.model.eval() + + self.tracked_stracks = [] # type: list[STrack] + self.lost_stracks = [] # type: list[STrack] + self.removed_stracks = [] # type: list[STrack] + + self.frame_id = 0 + self.det_thresh = opt.conf_thres + 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.kalman_filter = KalmanFilter() + + def set_buffer_len(self, frame_rate): + self.buffer_size = int(frame_rate / 30.0 * self.opt.track_buffer) + self.max_time_lost = self.buffer_size + + def post_process(self, dets, meta): + dets = dets.detach().cpu().numpy() + dets = dets.reshape(1, -1, dets.shape[2]) + dets = ctdet_post_process(dets.copy(), [meta['c']], [meta['s']], + meta['out_height'], meta['out_width'], + self.opt.num_classes) + for j in range(1, self.opt.num_classes + 1): + dets[0][j] = np.array(dets[0][j], dtype=np.float32).reshape(-1, 5) + return dets[0] + + def merge_outputs(self, detections): + results = {} + for j in range(1, self.opt.num_classes + 1): + results[j] = np.concatenate( + [detection[j] for detection in detections], + axis=0).astype(np.float32) + + scores = np.hstack( + [results[j][:, 4] for j in range(1, self.opt.num_classes + 1)]) + if len(scores) > self.max_per_image: + kth = len(scores) - self.max_per_image + thresh = np.partition(scores, kth)[kth] + for j in range(1, self.opt.num_classes + 1): + keep_inds = (results[j][:, 4] >= thresh) + results[j] = results[j][keep_inds] + return results + + def update(self, im_blob, img0): + self.frame_id += 1 + activated_starcks = [] + refind_stracks = [] + lost_stracks = [] + removed_stracks = [] + + width = img0.shape[1] + height = img0.shape[0] + inp_height = im_blob.shape[2] + inp_width = im_blob.shape[3] + if self.device.type == 'cuda': + im_blob = im_blob.cuda() + c = np.array([width / 2., height / 2.], dtype=np.float32) + s = max(float(inp_width) / float(inp_height) * height, width) * 1.0 + meta = { + 'c': c, + 's': s, + 'out_height': inp_height // self.opt.down_ratio, + 'out_width': inp_width // self.opt.down_ratio + } + + # Step 1: Network forward, get detections & embeddings + with torch.no_grad(): + output = self.model(im_blob)[-1] + hm = output['hm'].sigmoid_() + wh = output['wh'] + id_feature = output['id'] + 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) + id_feature = _tranpose_and_gather_feat(id_feature, inds) + id_feature = id_feature.squeeze(0) + id_feature = id_feature.cpu().numpy() + + dets = self.post_process(dets, meta) + dets = self.merge_outputs([dets])[1] + + remain_inds = dets[:, 4] > self.opt.conf_thres + dets = dets[remain_inds] + id_feature = id_feature[remain_inds] + + if len(dets) > 0: + '''Detections''' + detections = [ + STrack(STrack.tlbr_to_tlwh(tlbrs[:4]), tlbrs[4], f, 30) + for (tlbrs, f) in zip(dets[:, :5], id_feature) + ] + else: + detections = [] + + # Add newly detected tracklets to tracked_stracks + 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) + + # Step 2: First association, with embedding + strack_pool = joint_stracks(tracked_stracks, self.lost_stracks) + STrack.multi_predict(strack_pool) + dists = matching.embedding_distance(strack_pool, detections) + dists = matching.fuse_motion(self.kalman_filter, dists, strack_pool, + detections) + matches, u_track, u_detection = matching.linear_assignment( + dists, thresh=0.4) + + for itracked, idet in matches: + track = strack_pool[itracked] + det = detections[idet] + if track.state == TrackState.Tracked: + track.update(detections[idet], self.frame_id) + activated_starcks.append(track) + else: + track.re_activate(det, self.frame_id, new_id=False) + refind_stracks.append(track) + + # Step 3: Second association, with IOU + detections = [detections[i] for i in u_detection] + r_tracked_stracks = [ + strack_pool[i] for i in u_track + if strack_pool[i].state == TrackState.Tracked + ] + dists = matching.iou_distance(r_tracked_stracks, detections) + matches, u_track, u_detection = matching.linear_assignment( + dists, thresh=0.5) + + for itracked, idet in matches: + track = r_tracked_stracks[itracked] + det = detections[idet] + if track.state == TrackState.Tracked: + track.update(det, self.frame_id) + activated_starcks.append(track) + else: + track.re_activate(det, self.frame_id, new_id=False) + refind_stracks.append(track) + + for it in u_track: + track = r_tracked_stracks[it] + if not track.state == TrackState.Lost: + track.mark_lost() + lost_stracks.append(track) + + detections = [detections[i] for i in u_detection] + 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) + activated_starcks.append(unconfirmed[itracked]) + 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 track.score < self.det_thresh: + continue + track.activate(self.kalman_filter, self.frame_id) + activated_starcks.append(track) + # Step 5: 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.tracked_stracks = joint_stracks(self.tracked_stracks, + activated_starcks) + self.tracked_stracks = joint_stracks(self.tracked_stracks, + refind_stracks) + self.lost_stracks = sub_stracks(self.lost_stracks, + self.tracked_stracks) + self.lost_stracks.extend(lost_stracks) + self.lost_stracks = sub_stracks(self.lost_stracks, + self.removed_stracks) + self.removed_stracks.extend(removed_stracks) + self.tracked_stracks, self.lost_stracks = remove_duplicate_stracks( + self.tracked_stracks, self.lost_stracks) + output_stracks = [ + track for track in self.tracked_stracks if track.is_activated + ] + + return output_stracks + + +def joint_stracks(tlista, tlistb): + exists = {} + res = [] + for t in tlista: + exists[t.track_id] = 1 + res.append(t) + for t in tlistb: + tid = t.track_id + if not exists.get(tid, 0): + exists[tid] = 1 + res.append(t) + return res + + +def sub_stracks(tlista, tlistb): + stracks = {} + for t in tlista: + stracks[t.track_id] = t + for t in tlistb: + tid = t.track_id + if stracks.get(tid, 0): + del stracks[tid] + return list(stracks.values()) + + +def remove_duplicate_stracks(stracksa, stracksb): + pdist = matching.iou_distance(stracksa, stracksb) + pairs = np.where(pdist < 0.15) + dupa, dupb = list(), list() + for p, q in zip(*pairs): + timep = stracksa[p].frame_id - stracksa[p].start_frame + timeq = stracksb[q].frame_id - stracksb[q].start_frame + if timep > timeq: + dupb.append(q) + else: + dupa.append(p) + resa = [t for i, t in enumerate(stracksa) if i not in dupa] + resb = [t for i, t in enumerate(stracksb) if i not in dupb] + return resa, resb diff --git a/modelscope/models/cv/video_multi_object_tracking/utils/__init__.py b/modelscope/models/cv/video_multi_object_tracking/utils/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/modelscope/models/cv/video_multi_object_tracking/utils/image.py b/modelscope/models/cv/video_multi_object_tracking/utils/image.py new file mode 100644 index 00000000..1e8cf76b --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/utils/image.py @@ -0,0 +1,73 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +import cv2 +import numpy as np + + +def flip(img): + return img[:, :, ::-1].copy() + + +def transform_preds(coords, center, scale, output_size): + target_coords = np.zeros(coords.shape) + trans = get_affine_transform(center, scale, 0, output_size, inv=1) + for p in range(coords.shape[0]): + target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans) + return target_coords + + +def get_affine_transform(center, + scale, + rot, + output_size, + shift=np.array([0, 0], dtype=np.float32), + inv=0): + if not isinstance(scale, np.ndarray) and not isinstance(scale, list): + scale = np.array([scale, scale], dtype=np.float32) + + scale_tmp = scale + src_w = scale_tmp[0] + dst_w = output_size[0] + dst_h = output_size[1] + + rot_rad = np.pi * rot / 180 + src_dir = get_dir([0, src_w * -0.5], rot_rad) + dst_dir = np.array([0, dst_w * -0.5], np.float32) + + src = np.zeros((3, 2), dtype=np.float32) + dst = np.zeros((3, 2), dtype=np.float32) + src[0, :] = center + scale_tmp * shift + src[1, :] = center + src_dir + scale_tmp * shift + dst[0, :] = [dst_w * 0.5, dst_h * 0.5] + dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5], np.float32) + dst_dir + + src[2:, :] = get_3rd_point(src[0, :], src[1, :]) + dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :]) + + if inv: + trans = cv2.getAffineTransform(np.float32(dst), np.float32(src)) + else: + trans = cv2.getAffineTransform(np.float32(src), np.float32(dst)) + + return trans + + +def affine_transform(pt, t): + new_pt = np.array([pt[0], pt[1], 1.], dtype=np.float32).T + new_pt = np.dot(t, new_pt) + return new_pt[:2] + + +def get_3rd_point(a, b): + direct = a - b + return b + np.array([-direct[1], direct[0]], dtype=np.float32) + + +def get_dir(src_point, rot_rad): + sn, cs = np.sin(rot_rad), np.cos(rot_rad) + + src_result = [0, 0] + src_result[0] = src_point[0] * cs - src_point[1] * sn + src_result[1] = src_point[0] * sn + src_point[1] * cs + + return src_result diff --git a/modelscope/models/cv/video_multi_object_tracking/utils/kalman_filter.py b/modelscope/models/cv/video_multi_object_tracking/utils/kalman_filter.py new file mode 100644 index 00000000..c9cb0b41 --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/utils/kalman_filter.py @@ -0,0 +1,264 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +import numpy as np +import scipy.linalg + +chi2inv95 = { + 1: 3.8415, + 2: 5.9915, + 3: 7.8147, + 4: 9.4877, + 5: 11.070, + 6: 12.592, + 7: 14.067, + 8: 15.507, + 9: 16.919 +} + + +class KalmanFilter(object): + """ + A simple Kalman filter for tracking bounding boxes in image space. + + The 8-dimensional state space + + x, y, a, h, vx, vy, va, vh + + contains the bounding box center position (x, y), aspect ratio a, height h, + and their respective velocities. + + Object motion follows a constant velocity model. The bounding box location + (x, y, a, h) is taken as direct observation of the state space (linear + observation model). + + """ + + def __init__(self): + ndim, dt = 4, 1. + + # Create Kalman filter model matrices. + self._motion_mat = np.eye(2 * ndim, 2 * ndim) + for i in range(ndim): + self._motion_mat[i, ndim + i] = dt + self._update_mat = np.eye(ndim, 2 * ndim) + + # Motion and observation uncertainty are chosen relative to the current + # state estimate. These weights control the amount of uncertainty in + # the model. This is a bit hacky. + self._std_weight_position = 1. / 20 + self._std_weight_velocity = 1. / 160 + + def initiate(self, measurement): + """Create track from unassociated measurement. + + Args: + measurement : ndarray + Bounding box coordinates (x, y, a, h) with center position (x, y), + aspect ratio a, and height h. + + Returns: + (ndarray, ndarray): Returns the mean vector (8 dimensional) and covariance matrix + (8x8 dimensional) of the new track. Unobserved velocities are initialized + to 0 mean. + + """ + mean_pos = measurement + mean_vel = np.zeros_like(mean_pos) + mean = np.r_[mean_pos, mean_vel] + + std = [ + 2 * self._std_weight_position * measurement[3], + 2 * self._std_weight_position * measurement[3], 1e-2, + 2 * self._std_weight_position * measurement[3], + 10 * self._std_weight_velocity * measurement[3], + 10 * self._std_weight_velocity * measurement[3], 1e-5, + 10 * self._std_weight_velocity * measurement[3] + ] + covariance = np.diag(np.square(std)) + return mean, covariance + + def predict(self, mean, covariance): + """Run Kalman filter prediction step. + + Args: + mean : ndarray + The 8 dimensional mean vector of the object state at the previous + time step. + covariance : ndarray + The 8x8 dimensional covariance matrix of the object state at the + previous time step. + + Returns: + (ndarray, ndarray) + Returns the mean vector and covariance matrix of the predicted + state. Unobserved velocities are initialized to 0 mean. + + """ + std_pos = [ + self._std_weight_position * mean[3], + self._std_weight_position * mean[3], 1e-2, + self._std_weight_position * mean[3] + ] + std_vel = [ + self._std_weight_velocity * mean[3], + self._std_weight_velocity * mean[3], 1e-5, + self._std_weight_velocity * mean[3] + ] + motion_cov = np.diag(np.square(np.r_[std_pos, std_vel])) + + mean = np.dot(mean, self._motion_mat.T) + covariance = np.linalg.multi_dot( + (self._motion_mat, covariance, self._motion_mat.T)) + motion_cov + + return mean, covariance + + def project(self, mean, covariance): + """Project state distribution to measurement space. + + Args: + mean : ndarray + The state's mean vector (8 dimensional array). + covariance : ndarray + The state's covariance matrix (8x8 dimensional). + + Returns: + (ndarray, ndarray) + Returns the projected mean and covariance matrix of the given state + estimate. + + """ + std = [ + self._std_weight_position * mean[3], + self._std_weight_position * mean[3], 1e-1, + self._std_weight_position * mean[3] + ] + innovation_cov = np.diag(np.square(std)) + + mean = np.dot(self._update_mat, mean) + covariance = np.linalg.multi_dot( + (self._update_mat, covariance, self._update_mat.T)) + return mean, covariance + innovation_cov + + def multi_predict(self, mean, covariance): + """Run Kalman filter prediction step (Vectorized version). + + Args: + mean : ndarray + The Nx8 dimensional mean matrix of the object states at the previous + time step. + covariance : ndarray + The Nx8x8 dimensional covariance matrics of the object states at the + previous time step. + + Returns: + (ndarray, ndarray) + Returns the mean vector and covariance matrix of the predicted + state. Unobserved velocities are initialized to 0 mean. + """ + std_pos = [ + self._std_weight_position * mean[:, 3], + self._std_weight_position * mean[:, 3], + 1e-2 * np.ones_like(mean[:, 3]), + self._std_weight_position * mean[:, 3] + ] + std_vel = [ + self._std_weight_velocity * mean[:, 3], + self._std_weight_velocity * mean[:, 3], + 1e-5 * np.ones_like(mean[:, 3]), + self._std_weight_velocity * mean[:, 3] + ] + sqr = np.square(np.r_[std_pos, std_vel]).T + + motion_cov = [] + for i in range(len(mean)): + motion_cov.append(np.diag(sqr[i])) + motion_cov = np.asarray(motion_cov) + + mean = np.dot(mean, self._motion_mat.T) + left = np.dot(self._motion_mat, covariance).transpose((1, 0, 2)) + covariance = np.dot(left, self._motion_mat.T) + motion_cov + + return mean, covariance + + def update(self, mean, covariance, measurement): + """Run Kalman filter correction step. + + Args: + mean : ndarray + The predicted state's mean vector (8 dimensional). + covariance : ndarray + The state's covariance matrix (8x8 dimensional). + measurement : ndarray + The 4 dimensional measurement vector (x, y, a, h), where (x, y) + is the center position, a the aspect ratio, and h the height of the + bounding box. + + Returns: + (ndarray, ndarray) + Returns the measurement-corrected state distribution. + + """ + projected_mean, projected_cov = self.project(mean, covariance) + + chol_factor, lower = scipy.linalg.cho_factor( + projected_cov, lower=True, check_finite=False) + kalman_gain = scipy.linalg.cho_solve((chol_factor, lower), + np.dot(covariance, + self._update_mat.T).T, + check_finite=False).T + innovation = measurement - projected_mean + + new_mean = mean + np.dot(innovation, kalman_gain.T) + new_covariance = covariance - np.linalg.multi_dot( + (kalman_gain, projected_cov, kalman_gain.T)) + return new_mean, new_covariance + + def gating_distance(self, + mean, + covariance, + measurements, + only_position=False, + metric='maha'): + """Compute gating distance between state distribution and measurements. + A suitable distance threshold can be obtained from `chi2inv95`. If + `only_position` is False, the chi-square distribution has 4 degrees of + freedom, otherwise 2. + + Args: + mean : ndarray + Mean vector over the state distribution (8 dimensional). + covariance : ndarray + Covariance of the state distribution (8x8 dimensional). + measurements : ndarray + An Nx4 dimensional matrix of N measurements, each in + format (x, y, a, h) where (x, y) is the bounding box center + position, a the aspect ratio, and h the height. + only_position : Optional[bool] + If True, distance computation is done with respect to the bounding + box center position only. + + Returns: + an array of length N, where the i-th element contains the + squared Mahalanobis distance between (mean, covariance) and + `measurements[i]`. + """ + mean, covariance = self.project(mean, covariance) + if only_position: + mean, covariance = mean[:2], covariance[:2, :2] + measurements = measurements[:, :2] + + d = measurements - mean + if metric == 'gaussian': + return np.sum(d * d, axis=1) + elif metric == 'maha': + cholesky_factor = np.linalg.cholesky(covariance) + z = scipy.linalg.solve_triangular( + cholesky_factor, + d.T, + lower=True, + check_finite=False, + overwrite_b=True) + squared_maha = np.sum(z * z, axis=0) + return squared_maha + else: + raise ValueError('invalid distance metric') diff --git a/modelscope/models/cv/video_multi_object_tracking/utils/utils.py b/modelscope/models/cv/video_multi_object_tracking/utils/utils.py new file mode 100644 index 00000000..e961b720 --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/utils/utils.py @@ -0,0 +1,208 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +import cv2 +import numpy as np +import torch + +from .image import transform_preds + + +def xyxy2xywh(x): + # Convert bounding box format from [x1, y1, x2, y2] to [x, y, w, h] + y = torch.zeros(x.shape) if x.dtype is torch.float32 else np.zeros(x.shape) + y[:, 0] = (x[:, 0] + x[:, 2]) / 2 + y[:, 1] = (x[:, 1] + x[:, 3]) / 2 + y[:, 2] = x[:, 2] - x[:, 0] + y[:, 3] = x[:, 3] - x[:, 1] + return y + + +def xywh2xyxy(x): + # Convert bounding box format from [x, y, w, h] to [x1, y1, x2, y2] + y = torch.zeros(x.shape) if x.dtype is torch.float32 else np.zeros(x.shape) + y[:, 0] = (x[:, 0] - x[:, 2] / 2) + y[:, 1] = (x[:, 1] - x[:, 3] / 2) + y[:, 2] = (x[:, 0] + x[:, 2] / 2) + y[:, 3] = (x[:, 1] + x[:, 3] / 2) + return y + + +def bbox_iou(box1, box2, x1y1x2y2=False): + """ + Returns the IoU of two bounding boxes + """ + N, M = len(box1), len(box2) + box1 = torch.from_numpy(np.stack(box1)) + box2 = torch.from_numpy(np.stack(box2)) + if x1y1x2y2: + # Get the coordinates of bounding boxes + b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, + 2], box1[:, + 3] + b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, + 2], box2[:, + 3] + else: + # Transform from center and width to exact coordinates + b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2 + b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2 + b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2 + b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2 + + # get the coordinates of the intersection rectangle + inter_rect_x1 = torch.max(b1_x1.unsqueeze(1), b2_x1) + inter_rect_y1 = torch.max(b1_y1.unsqueeze(1), b2_y1) + inter_rect_x2 = torch.min(b1_x2.unsqueeze(1), b2_x2) + inter_rect_y2 = torch.min(b1_y2.unsqueeze(1), b2_y2) + # Intersection area + inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, 0) * torch.clamp( + inter_rect_y2 - inter_rect_y1, 0) + # Union Area + b1_area = ((b1_x2 - b1_x1) * (b1_y2 - b1_y1)) + b1_area = ((b1_x2 - b1_x1) * (b1_y2 - b1_y1)).view(-1, 1).expand(N, M) + b2_area = ((b2_x2 - b2_x1) * (b2_y2 - b2_y1)).view(1, -1).expand(N, M) + + return inter_area / (b1_area + b2_area - inter_area + 1e-16) + + +class LoadVideo: # for inference + + def __init__(self, path, img_size=(1088, 608)): + self.cap = cv2.VideoCapture(path) + self.frame_rate = int(round(self.cap.get(cv2.CAP_PROP_FPS))) + self.vw = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + self.vh = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + self.vn = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + self.width = img_size[0] + self.height = img_size[1] + self.count = 0 + + self.w, self.h = 1920, 1080 + print('Length of the video: {:d} frames'.format(self.vn)) + + def get_size(self, vw, vh, dw, dh): + wa, ha = float(dw) / vw, float(dh) / vh + a = min(wa, ha) + return int(vw * a), int(vh * a) + + def __iter__(self): + self.count = -1 + return self + + def __next__(self): + self.count += 1 + if self.count == len(self): + raise StopIteration + # Read image + res, img0 = self.cap.read() # BGR + assert img0 is not None, 'Failed to load frame {:d}'.format(self.count) + img0 = cv2.resize(img0, (self.w, self.h)) + + # Padded resize + shape = [self.height, self.width] + img, ratio, pad = letterbox(img0, shape) + + # Normalize RGB + img = img[:, :, ::-1].transpose(2, 0, 1) + img = np.ascontiguousarray(img, dtype=np.float32) + img /= 255.0 + + return self.count, img, img0 + + def __len__(self): + return self.vn # number of files + + +def letterbox(img, + new_shape=(640, 640), + color=(114, 114, 114), + auto=True, + scaleFill=False, + scaleup=True): + # Resize image to a 32-pixel-multiple rectangle https://github.com/ultralytics/yolov3/issues/232 + shape = img.shape[:2] # current shape [height, width] + if isinstance(new_shape, int): + new_shape = (new_shape, new_shape) + + # Scale ratio (new / old) + r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) + if not scaleup: # only scale down, do not scale up (for better test mAP) + r = min(r, 1.0) + + # Compute padding + ratio = r, r # width, height ratios + new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) + dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[ + 1] # wh padding + if auto: # minimum rectangle + dw, dh = np.mod(dw, 64), np.mod(dh, 64) # wh padding + elif scaleFill: # stretch + dw, dh = 0.0, 0.0 + new_unpad = (new_shape[1], new_shape[0]) + ratio = new_shape[1] / shape[1], new_shape[0] / shape[ + 0] # width, height ratios + + dw /= 2 # divide padding into 2 sides + dh /= 2 + + if shape[::-1] != new_unpad: # resize + img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) + top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) + left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) + img = cv2.copyMakeBorder( + img, top, bottom, left, right, cv2.BORDER_CONSTANT, + value=color) # add border + return img, ratio, (dw, dh) + + +def _gather_feat(feat, ind, mask=None): + dim = feat.size(2) + ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim) + feat = feat.gather(1, ind) + if mask is not None: + mask = mask.unsqueeze(2).expand_as(feat) + feat = feat[mask] + feat = feat.view(-1, dim) + return feat + + +def _tranpose_and_gather_feat(feat, ind): + feat = feat.permute(0, 2, 3, 1).contiguous() + feat = feat.view(feat.size(0), -1, feat.size(3)) + feat = _gather_feat(feat, ind) + return feat + + +class cfg_opt: + K = 500 + arch = 'yolo' + conf_thres = 0.4 + down_ratio = 4 + head_conv = 256 + heads = {'hm': 1, 'wh': 4, 'id': 64, 'reg': 2} + img_size = (1088, 608) + ltrb = True + mean = [0.408, 0.447, 0.47] + min_box_area = 100 + num_classes = 1 + reg_offset = True + reid_dim = 64 + std = [0.289, 0.274, 0.278] + track_buffer = 30 + + +def ctdet_post_process(dets, c, s, h, w, num_classes): + ret = [] + for i in range(dets.shape[0]): + top_preds = {} + dets[i, :, :2] = transform_preds(dets[i, :, 0:2], c[i], s[i], (w, h)) + dets[i, :, 2:4] = transform_preds(dets[i, :, 2:4], c[i], s[i], (w, h)) + classes = dets[i, :, -1] + for j in range(num_classes): + inds = (classes == j) + det4 = dets[i, inds, :4].astype(np.float32) + det5 = dets[i, inds, 4:5].astype(np.float32) + top_preds[j + 1] = np.concatenate([det4, det5], axis=1).tolist() + ret.append(top_preds) + return ret diff --git a/modelscope/models/cv/video_multi_object_tracking/utils/visualization.py b/modelscope/models/cv/video_multi_object_tracking/utils/visualization.py new file mode 100644 index 00000000..8ed7b601 --- /dev/null +++ b/modelscope/models/cv/video_multi_object_tracking/utils/visualization.py @@ -0,0 +1,85 @@ +# The implementation is adopted from FairMOT, +# made publicly available under the MIT License at https://github.com/ifzhang/FairMOT +import cv2 +import numpy as np + + +def get_color(idx): + idx = idx * 3 + color = ((37 * idx) % 255, (17 * idx) % 255, (29 * idx) % 255) + + return color + + +def plot_tracking(image, + tlwhs, + obj_ids, + scores=None, + frame_id=0, + fps=0., + ids2=None): + im = np.ascontiguousarray(np.copy(image)) + text_scale = max(1, image.shape[1] / 1600.) + text_thickness = 2 + line_thickness = max(1, int(image.shape[1] / 500.)) + + cv2.putText( + im, + 'frame: %d fps: %.2f num: %d' % (frame_id, fps, len(tlwhs)), + (0, int(15 * text_scale)), + cv2.FONT_HERSHEY_PLAIN, + text_scale, (0, 0, 255), + thickness=2) + + for i, tlwh in enumerate(tlwhs): + x1, y1, w, h = tlwh + intbox = tuple(map(int, (x1, y1, x1 + w, y1 + h))) + obj_id = int(obj_ids[i]) + id_text = '{}'.format(int(obj_id)) + if ids2 is not None: + id_text = id_text + ', {}'.format(int(ids2[i])) + color = get_color(abs(obj_id)) + cv2.rectangle( + im, + intbox[0:2], + intbox[2:4], + color=color, + thickness=line_thickness) + cv2.putText( + im, + id_text, (intbox[0], intbox[1] + 30), + cv2.FONT_HERSHEY_PLAIN, + text_scale, (0, 0, 255), + thickness=text_thickness) + return im + + +def show_multi_object_tracking_result(video_in_path, bboxes, video_save_path): + cap = cv2.VideoCapture(video_in_path) + frame_idx = 0 + while (cap.isOpened()): + frame_idx += 1 + success, frame = cap.read() + if not success: + if frame_idx == 1: + raise Exception(video_in_path, + ' can not be correctly decoded by OpenCV.') + else: + break + cur_frame_boxes = [] + cur_obj_ids = [] + for box in bboxes: + if box[0] == frame_idx: + cur_frame_boxes.append( + [box[2], box[3], box[4] - box[2], box[5] - box[3]]) + cur_obj_ids.append(box[1]) + if frame_idx == 1: + size = (frame.shape[1], frame.shape[0]) + fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G') + video_writer = cv2.VideoWriter(video_save_path, fourcc, + cap.get(cv2.CAP_PROP_FPS), size, + True) + frame = plot_tracking(frame, cur_frame_boxes, cur_obj_ids, frame_idx) + video_writer.write(frame) + video_writer.release + cap.release() diff --git a/modelscope/outputs/outputs.py b/modelscope/outputs/outputs.py index 3fbaba68..dc384d91 100644 --- a/modelscope/outputs/outputs.py +++ b/modelscope/outputs/outputs.py @@ -386,6 +386,19 @@ TASK_OUTPUTS = { OutputKeys.BOXES, OutputKeys.TIMESTAMPS ], + # video multi object tracking result for single video + # { + # "boxes": [ + # [frame_num, obj_id, x1, y1, x2, y2], + # [frame_num, obj_id, x1, y1, x2, y2], + # [frame_num, obj_id, x1, y1, x2, y2], + # ], + # "timestamps": ["hh:mm:ss", "hh:mm:ss", "hh:mm:ss"] + # } + Tasks.video_multi_object_tracking: [ + OutputKeys.BOXES, OutputKeys.TIMESTAMPS + ], + # live category recognition result for single video # { # "scores": [0.885272, 0.014790631, 0.014558001], diff --git a/modelscope/pipeline_inputs.py b/modelscope/pipeline_inputs.py index a064ffd7..e83c2b52 100644 --- a/modelscope/pipeline_inputs.py +++ b/modelscope/pipeline_inputs.py @@ -131,6 +131,8 @@ TASK_INPUTS = { Tasks.hand_2d_keypoints: InputType.IMAGE, Tasks.video_single_object_tracking: (InputType.VIDEO, InputType.BOX), + Tasks.video_multi_object_tracking: + InputType.VIDEO, Tasks.video_category: InputType.VIDEO, Tasks.product_retrieval_embedding: diff --git a/modelscope/pipelines/builder.py b/modelscope/pipelines/builder.py index ae5b6603..48ba54ed 100644 --- a/modelscope/pipelines/builder.py +++ b/modelscope/pipelines/builder.py @@ -248,6 +248,9 @@ DEFAULT_MODEL_FOR_PIPELINE = { Tasks.video_object_segmentation: (Pipelines.video_object_segmentation, 'damo/cv_rdevos_video-object-segmentation'), + Tasks.video_multi_object_tracking: ( + Pipelines.video_multi_object_tracking, + 'damo/cv_yolov5_video-multi-object-tracking_fairmot'), Tasks.image_multi_view_depth_estimation: ( Pipelines.image_multi_view_depth_estimation, 'damo/cv_casmvs_multi-view-depth-estimation_general'), diff --git a/modelscope/pipelines/cv/video_multi_object_tracking_pipeline.py b/modelscope/pipelines/cv/video_multi_object_tracking_pipeline.py new file mode 100644 index 00000000..0f02413c --- /dev/null +++ b/modelscope/pipelines/cv/video_multi_object_tracking_pipeline.py @@ -0,0 +1,81 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os.path as osp +from typing import Any, Dict + +import torch + +from modelscope.metainfo import Pipelines +from modelscope.models.cv.video_multi_object_tracking.tracker.multitracker import \ + JDETracker +from modelscope.models.cv.video_multi_object_tracking.utils.utils import ( + LoadVideo, cfg_opt) +from modelscope.models.cv.video_single_object_tracking.utils.utils import \ + timestamp_format +from modelscope.outputs import OutputKeys +from modelscope.pipelines.base import Input, Pipeline +from modelscope.pipelines.builder import PIPELINES +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +@PIPELINES.register_module( + Tasks.video_multi_object_tracking, + module_name=Pipelines.video_multi_object_tracking) +class VideoMultiObjectTrackingPipeline(Pipeline): + + def __init__(self, model: str, **kwargs): + """ + use `model` to create a multi object tracking pipeline + Args: + model: model id on modelscope hub. + """ + super().__init__(model=model, **kwargs) + ckpt_path = osp.join(model, ModelFile.TORCH_MODEL_BIN_FILE) + logger.info(f'loading model from {ckpt_path}') + opt = cfg_opt() + self.opt = opt + self.tracker = JDETracker(opt, ckpt_path, self.device) + logger.info('init tracker done') + + def preprocess(self, input) -> Input: + self.video_path = input[0] + return input + + def forward(self, input: Input) -> Dict[str, Any]: + dataloader = LoadVideo(input, self.opt.img_size) + self.tracker.set_buffer_len(dataloader.frame_rate) + + results = [] + output_timestamps = [] + frame_id = 0 + for i, (path, img, img0) in enumerate(dataloader): + output_timestamps.append( + timestamp_format(seconds=frame_id / dataloader.frame_rate)) + blob = torch.from_numpy(img).unsqueeze(0) + online_targets = self.tracker.update(blob, img0) + online_tlwhs = [] + online_ids = [] + for t in online_targets: + tlwh = t.tlwh + tid = t.track_id + vertical = tlwh[2] / tlwh[3] > 1.6 + if tlwh[2] * tlwh[3] > self.opt.min_box_area and not vertical: + online_tlwhs.append([ + tlwh[0], tlwh[1], tlwh[0] + tlwh[2], tlwh[1] + tlwh[3] + ]) + online_ids.append(tid) + results.append([ + frame_id + 1, tid, tlwh[0], tlwh[1], tlwh[0] + tlwh[2], + tlwh[1] + tlwh[3] + ]) + frame_id += 1 + + return { + OutputKeys.BOXES: results, + OutputKeys.TIMESTAMPS: output_timestamps + } + + def postprocess(self, inputs: Dict[str, Any]) -> Dict[str, Any]: + return inputs diff --git a/modelscope/utils/constant.py b/modelscope/utils/constant.py index d50e28f7..37737aab 100644 --- a/modelscope/utils/constant.py +++ b/modelscope/utils/constant.py @@ -102,6 +102,7 @@ class CVTasks(object): # reid and tracking video_single_object_tracking = 'video-single-object-tracking' + video_multi_object_tracking = 'video-multi-object-tracking' video_summarization = 'video-summarization' image_reid_person = 'image-reid-person' diff --git a/requirements/cv.txt b/requirements/cv.txt index 08256e53..7ea00b1e 100644 --- a/requirements/cv.txt +++ b/requirements/cv.txt @@ -12,6 +12,7 @@ imageio>=2.9.0 imageio-ffmpeg>=0.4.2 imgaug>=0.4.0 kornia>=0.5.0 +lap lmdb lpips ml_collections diff --git a/tests/pipelines/test_video_multi_object_tracking.py b/tests/pipelines/test_video_multi_object_tracking.py new file mode 100644 index 00000000..eb37ffd0 --- /dev/null +++ b/tests/pipelines/test_video_multi_object_tracking.py @@ -0,0 +1,39 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import unittest + +from modelscope.outputs import OutputKeys +from modelscope.pipelines import pipeline +from modelscope.utils.constant import Tasks +from modelscope.utils.demo_utils import DemoCompatibilityCheck +from modelscope.utils.test_utils import test_level + + +class MultiObjectTracking(unittest.TestCase, DemoCompatibilityCheck): + + def setUp(self) -> None: + self.task = Tasks.video_multi_object_tracking + self.model_id = 'damo/cv_yolov5_video-multi-object-tracking_fairmot' + + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_end2end(self): + video_multi_object_tracking = pipeline( + Tasks.video_multi_object_tracking, model=self.model_id) + video_path = 'data/test/videos/MOT17-03-partial.mp4' + result = video_multi_object_tracking(video_path) + print('result is : ', result[OutputKeys.BOXES]) + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_run_modelhub_default_model(self): + video_multi_object_tracking = pipeline( + Tasks.video_multi_object_tracking) + video_path = 'data/test/videos/MOT17-03-partial.mp4' + result = video_multi_object_tracking(video_path) + print('result is : ', result[OutputKeys.BOXES]) + + @unittest.skip('demo compatibility test is only enabled on a needed-basis') + def test_demo_compatibility(self): + self.compatibility_check() + + +if __name__ == '__main__': + unittest.main()