支持视频多目标跟踪模型

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11249098

data/test/videos/MOT17-03-partial.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ca26615762e3f4ccca53a020efe73c3cf3598edc68bb68b5555c24e815718336
+size 3151767

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'

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={},
+    )

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)

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

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

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

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

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

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

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

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')

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

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()

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],

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:

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'),

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

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'
 

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

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()