customize_mot_model.md

Customize MOT Models

We basically categorize model components into 5 types.

• tracker: the component that associate the objects across the video with the cues extracted by the components below.
• detector: usually a detector to detect objects from the input image, e.g., Faster R-CNN.
• motion: the component to compute motion information between consecutive frames, e.g., KalmanFilter.
• reid: usually an independent ReID model to extract the feature embeddings from the cropped image, e.g., BaseReID.
• track_head: the component to extract tracking cues but share the same backbone with the detector, e.g., a embedding head or a regression head.

Add a new tracker

1. Define a tracker (e.g. MyTracker)

Create a new file mmtrack/models/mot/trackers/my_tracker.py.

We implement a BaseTracker that provide basic APIs to maintain the tracks across the video. We recommend to inherit the new tracker from it. The users may refer to the documentations of BaseTracker for the details.

from mmtrack.models import TRACKERS
from .base_tracker import BaseTracker

@TRACKERS.register_module()
class MyTracker(BaseTracker):

    def __init__(self,
                 arg1,
                 arg2,
                 *args,
                 **kwargs):
        super().__init__(*args, **kwargs)
        pass

    def track(self, inputs):
        # implementation is ignored
        pass

2. Import the module

You can either add the following line to mmtrack/models/mot/trackers/__init__.py,

from .my_tracker import MyTracker

or alternatively add

custom_imports = dict(
    imports=['mmtrack.models.mot.trackers.my_tracker.py'],
    allow_failed_imports=False)

to the config file and avoid modifying the original code.

3. Modify the config file

tracker=dict(
    type='MyTracker',
    arg1=xxx,
    arg2=xxx)

Add a new detector

Please refer to tutorial in mmdetection for developping a new detector.

Add a new motion model

1. Define a motion model (e.g. MyFlowNet)

Create a new file mmtrack/models/motion/my_flownet.py.

You can inherit the motion model from nn.Module if it is a deep learning module and from object if not.

from ..builder import MOTION

@MOTION.register_module()
class MyFlowNet(nn.Module):

    def __init__(self,
                arg1,
                arg2):
        pass

    def forward(self, inputs):
        # implementation is ignored
        pass

2. Import the module

You can either add the following line to mmtrack/models/motion/__init__.py,

from .my_flownet import MyFlowNet

or alternatively add

custom_imports = dict(
    imports=['mmtrack.models.motion.my_flownet.py'],
    allow_failed_imports=False)

to the config file and avoid modifying the original code.

3. Modify the config file

motion=dict(
    type='MyFlowNet',
    arg1=xxx,
    arg2=xxx)

Add a new reid model

1. Define a reid model (e.g. MyReID)

Create a new file mmtrack/models/reid/my_reid.py.

from ..builder import REID

@REID.register_module()
class MyReID(nn.Module):

    def __init__(self,
                arg1,
                arg2):
        pass

    def forward(self, inputs):
        # implementation is ignored
        pass

2. Import the module

You can either add the following line to mmtrack/models/reid/__init__.py,

from .my_reid import MyReID

or alternatively add

custom_imports = dict(
    imports=['mmtrack.models.reid.my_reid.py'],
    allow_failed_imports=False)

to the config file and avoid modifying the original code.

3. Modify the config file

reid=dict(
    type='MyReID',
    arg1=xxx,
    arg2=xxx)

Add a new track head

1. Define a head (e.g. MyHead)

Create a new file mmtrack/models/track_heads/my_head.py.

from mmdet.models import HEADS

@HEADS.register_module()
class MyHead(nn.Module):

    def __init__(self,
                arg1,
                arg2):
        pass

    def forward(self, inputs):
        # implementation is ignored
        pass

2. Import the module

You can either add the following line to mmtrack/models/track_heads/__init__.py,

from .my_head import MyHead

or alternatively add

custom_imports = dict(
    imports=['mmtrack.models.track_heads.my_head.py'],
    allow_failed_imports=False)

to the config file and avoid modifying the original code.

3. Modify the config file

track_head=dict(
    type='MyHead',
    arg1=xxx,
    arg2=xxx)

Add a new loss

Assume you want to add a new loss as MyLoss, for bounding box regression. To add a new loss function, the users need implement it in mmtrack/models/losses/my_loss.py. The decorator weighted_loss enable the loss to be weighted for each element.

import torch
import torch.nn as nn

from ..builder import LOSSES
from .utils import weighted_loss

@weighted_loss
def my_loss(pred, target):
    assert pred.size() == target.size() and target.numel() > 0
    loss = torch.abs(pred - target)
    return loss

@LOSSES.register_module()
class MyLoss(nn.Module):

    def __init__(self, reduction='mean', loss_weight=1.0):
        super(MyLoss, self).__init__()
        self.reduction = reduction
        self.loss_weight = loss_weight

    def forward(self,
                pred,
                target,
                weight=None,
                avg_factor=None,
                reduction_override=None):
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (
            reduction_override if reduction_override else self.reduction)
        loss_bbox = self.loss_weight * my_loss(
            pred, target, weight, reduction=reduction, avg_factor=avg_factor)
        return loss_bbox

Then the users need to add it in the mmtrack/models/losses/__init__.py.

from .my_loss import MyLoss, my_loss

Alternatively, you can add

custom_imports=dict(
    imports=['mmtrack.models.losses.my_loss'])

to the config file and achieve the same goal.

To use it, modify the loss_xxx field. Since MyLoss is for regression, you need to modify the loss_bbox field in the head.

loss_bbox=dict(type='MyLoss', loss_weight=1.0))