major commit with first working version

app.py

@@ -1,7 +1,157 @@
 import gradio as gr
+import os
+import os.path as op
+from urllib.request import urlretrieve
 
-def greet(name):
-    return "Surfrider " + name + "!!"
+import json
+import multiprocessing
+from typing import Dict, List, Tuple
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+import datetime
+import logging
+
+# imports for tracking
+import torch
+import cv2
+import numpy as np
+from tools.video_readers import IterableFrameReader, SimpleVideoReader
+from detection.mobilenet import get_mobilenet_v3_small
+
+from detection.detect import detect
+from tracking.postprocess_and_count_tracks import filter_tracks, postprocess_for_api
+from tracking.utils import get_detections_for_video, write_tracking_results_to_file, read_tracking_results, gather_tracklets, generate_video_with_annotations
+from tracking.track_video import track_video
+from tracking.trackers import get_tracker
+
+logger = logging.getLogger()
+logger.setLevel(logging.DEBUG)
+ch = logging.StreamHandler()
+ch.setLevel(logging.DEBUG)
+formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+ch.setFormatter(formatter)
+logger.addHandler(ch)
+
+
+# download model
+FILE_MODEL = "mobilenet_v3_pretrained.pth"
+URL_MODEL = "https://partage.imt.fr/index.php/s/sJi22N6gedN6T4q/download"
+model_filename = op.realpath('./models/' + FILE_MODEL)
+if not op.exists(model_filename):
+    logger.info('---Downloading model...')
+    urlretrieve(URL_MODEL, model_filename)
+else:
+    logger.info('---Model already downloaded.')
+
+# load model
+print("loading model")
+logger.info('---Loading model...')
+
+model = get_mobilenet_v3_small(num_layers=0, heads={'hm': 1}, head_conv=256)
+checkpoint = torch.load(model_filename, map_location="cpu")
+model.load_state_dict(checkpoint['model'], strict=True)
+
+class DotDict(dict):
+    """dot.notation access to dictionary attributes"""
+    __getattr__ = dict.get
+    __setattr__ = dict.__setitem__
+    __delattr__ = dict.__delitem__
+
+config_track = DotDict({
+    "confidence_threshold": 0.5,
+    "detection_threshold": 0.3,
+    "downsampling_factor": 4,
+    "noise_covariances_path": "data/tracking_parameters",
+    "output_shape": (960,544),
+    "skip_frames": 3, #3
+    "arch": "mobilenet_v3_small",
+    "device": "cpu",
+    "detection_batch_size": 1,
+    "display": 0,
+    "kappa": 4, #4
+    "tau": 3, #4
+    "max_length": 240
+})
+
+def track(args):
+    device = torch.device("cpu")
+
+    engine = get_tracker('EKF')
+
+    detector = lambda frame: detect(frame, threshold=args.detection_threshold, model=model)
+
+    transition_variance = np.load(os.path.join(args.noise_covariances_path, 'transition_variance.npy'))
+    observation_variance = np.load(os.path.join(args.noise_covariances_path, 'observation_variance.npy'))
+
+    logger.info(f'---Processing {args.video_path}')
+    reader = IterableFrameReader(video_filename=args.video_path,
+                                 skip_frames=args.skip_frames,
+                                 output_shape=args.output_shape,
+                                 progress_bar=True,
+                                 preload=False,
+                                 max_frame=args.max_length)
+
+
+    input_shape = reader.input_shape
+    output_shape = reader.output_shape
+    ratio_y = input_shape[0] / (output_shape[0] // args.downsampling_factor)
+    ratio_x = input_shape[1] / (output_shape[1] // args.downsampling_factor)
+
+    logger.info('---Detecting...')
+    detections = get_detections_for_video(reader, detector, batch_size=args.detection_batch_size, device=device)
+
+    logger.info('---Tracking...')
+    display = None
+    results = track_video(reader, iter(detections), args, engine, transition_variance, observation_variance, display)
+
+    # store unfiltered results
+    datestr = datetime.datetime.now().strftime('%Y%m%d%H%M%S%f')
+    output_filename = os.path.splitext(args.video_path)[0] + "_" + datestr + '_unfiltered.txt'
+    write_tracking_results_to_file(results, ratio_x=ratio_x, ratio_y=ratio_y, output_filename=output_filename)
+    logger.info('---Filtering...')
+
+    # read from the file
+    results = read_tracking_results(output_filename)
+    filtered_results = filter_tracks(results, config_track.kappa, config_track.tau)
+    # store filtered results
+    output_filename = os.path.splitext(args.video_path)[0] + "_" + datestr + '_filtered.txt'
+    write_tracking_results_to_file(filtered_results, ratio_x=ratio_x, ratio_y=ratio_y, output_filename=output_filename)
+
+    return filtered_results
+
+
+def run_model(video_path, skip, tau, kappa, maxframes):
+    logger.info('---video filename: '+ video_path)
+
+    # launch the tracking
+    config_track.video_path = video_path
+    config_track.skip_frames = int(skip)
+    config_track.tau = int(tau)
+    config_track.kappa = int(kappa)
+    config_track.max_length = int(maxframes)
+
+    output_path = "/tmp/video_out.mp4"
+    filtered_results = track(config_track)
+
+    # postprocess
+    logger.info('---Postprocessing...')
+    output_json = postprocess_for_api(filtered_results)
+    #output_json={"detected_trash":[]}
+    # build video output
+    logger.info('---Generating new video...')
+    video = SimpleVideoReader(video_path, skip_frames=0)
+    generate_video_with_annotations(video, output_json, output_path, config_track.skip_frames, logger)
+
+    return output_path #, output_json
+
+video_in = gr.inputs.Video(type="mp4", source="upload", label="Video Upload", optional=False)
+skip_slider = gr.inputs.Slider(minimum=0, maximum=15, step=1, default=3, label="skip frames")
+tau_slider = gr.inputs.Slider(minimum=1, maximum=7, step=1, default=3, label="tau")
+kappa_slider = gr.inputs.Slider(minimum=1, maximum=7, step=1, default=4, label="kappa")
+maxframes_num = gr.inputs.Number(default=240, label="max frames at 24fps")
+
+
+gr.Interface(fn=run_model, inputs=[video_in, skip_slider, tau_slider, kappa_slider, maxframes_num],
+             outputs="playable_video",
+             title="Surfnet demo",
+             description="Upload a video, only the first seconds will be processed (use maxframe to change that)",
+             allow_screenshot=False).launch(debug=True)

detection/detect.py

@@ -0,0 +1,18 @@
+import torch
+from time import time
+
+
+def nms(heat, kernel=3):
+    pad = (kernel - 1) // 2
+
+    hmax = torch.nn.functional.max_pool2d(
+        heat, (kernel, kernel), stride=1, padding=pad)
+    keep = (hmax == heat).float()
+    return heat * keep
+
+def detect(preprocessed_frames, threshold, model):
+
+    batch_result = torch.sigmoid(model(preprocessed_frames)[-1]['hm'])
+    batch_peaks = nms(batch_result).gt(threshold).squeeze(dim=1)
+    detections = [torch.nonzero(peaks).cpu().numpy()[:,::-1] for peaks in batch_peaks]
+    return detections

detection/mobilenet.py

@@ -0,0 +1,141 @@
+# ------------------------------------------------------------------------------
+# Copyright (c) Microsoft
+# Licensed under the MIT License.
+# Written by Bin Xiao (Bin.Xiao@microsoft.com)
+# Modified by Xingyi Zhou
+# ------------------------------------------------------------------------------
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+import torch
+import torch.nn as nn
+import torchvision.models as models
+
+BN_MOMENTUM = 0.1
+
+class MobiletNetHM(nn.Module):
+
+    def __init__(self, heads, head_conv, **kwargs):
+        self.inplanes = 576
+        self.deconv_with_bias = False
+        self.heads = heads
+
+        super(MobiletNetHM, self).__init__()
+        self.features = models.mobilenet_v3_small(pretrained=True).features
+
+        # used for deconv layers
+        self.deconv_layers = self._make_deconv_layer(
+            3,
+            [256, 256, 256],
+            [4, 4, 4],
+        )
+        # self.final_layer = []
+
+        for head in sorted(self.heads):
+          num_output = self.heads[head]
+          if head_conv > 0:
+            fc = nn.Sequential(
+                nn.Conv2d(256, head_conv,
+                  kernel_size=3, padding=1, bias=True),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(head_conv, num_output, 
+                  kernel_size=1, stride=1, padding=0))
+          else:
+            fc = nn.Conv2d(
+              in_channels=256,
+              out_channels=num_output,
+              kernel_size=1,
+              stride=1,
+              padding=0
+          )
+          self.__setattr__(head, fc)
+
+        # self.final_layer = nn.ModuleList(self.final_layer)
+
+    def _get_deconv_cfg(self, deconv_kernel, index):
+        if deconv_kernel == 4:
+            padding = 1
+            output_padding = 0
+        elif deconv_kernel == 3:
+            padding = 1
+            output_padding = 1
+        elif deconv_kernel == 2:
+            padding = 0
+            output_padding = 0
+
+        return deconv_kernel, padding, output_padding
+
+    def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
+        assert num_layers == len(num_filters), \
+            'ERROR: num_deconv_layers is different len(num_deconv_filters)'
+        assert num_layers == len(num_kernels), \
+            'ERROR: num_deconv_layers is different len(num_deconv_filters)'
+
+        layers = []
+        for i in range(num_layers):
+            kernel, padding, output_padding = \
+                self._get_deconv_cfg(num_kernels[i], i)
+
+            planes = num_filters[i]
+            layers.append(
+                nn.ConvTranspose2d(
+                    in_channels=self.inplanes,
+                    out_channels=planes,
+                    kernel_size=kernel,
+                    stride=2,
+                    padding=padding,
+                    output_padding=output_padding,
+                    bias=self.deconv_with_bias))
+            layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
+            layers.append(nn.ReLU(inplace=True))
+            self.inplanes = planes
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.features(x)
+        x = self.deconv_layers(x)
+        ret = {}
+        for head in self.heads:
+            ret[head] = self.__getattr__(head)(x)
+        return [ret]
+
+    def init_weights(self):
+        # print('=> init resnet deconv weights from normal distribution')
+        for _, m in self.deconv_layers.named_modules():
+            if isinstance(m, nn.ConvTranspose2d):
+                # print('=> init {}.weight as normal(0, 0.001)'.format(name))
+                # print('=> init {}.bias as 0'.format(name))
+                nn.init.normal_(m.weight, std=0.001)
+                if self.deconv_with_bias:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                # print('=> init {}.weight as 1'.format(name))
+                # print('=> init {}.bias as 0'.format(name))
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+        # print('=> init final conv weights from normal distribution')
+        for head in self.heads:
+            final_layer = self.__getattr__(head)
+            for i, m in enumerate(final_layer.modules()):
+                if isinstance(m, nn.Conv2d):
+                    # nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+                    # print('=> init {}.weight as normal(0, 0.001)'.format(name))
+                    # print('=> init {}.bias as 0'.format(name))
+                    if m.weight.shape[0] == self.heads[head]:
+                        if 'hm' in head:
+                            nn.init.constant_(m.bias, -2.19)
+                        else:
+                            nn.init.normal_(m.weight, std=0.001)
+                            nn.init.constant_(m.bias, 0)
+
+
+def get_mobilenet_v3_small(num_layers, heads, head_conv):
+
+  model = MobiletNetHM(heads, head_conv=head_conv)
+  model.init_weights()
+  return model

detection/transforms.py

@@ -0,0 +1,135 @@
+import imgaug as ia
+from imgaug import augmenters as iaa 
+from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage
+import numpy as np 
+from tools.misc import blob_for_bbox
+import torch 
+ia.seed(1)
+from torchvision.transforms import functional as F
+import torchvision.transforms as T
+import cv2
+
+class Compose(object):
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, target):
+        for t in self.transforms:
+            image, target = t(image, target)
+        return image, target
+
+class ToTensorBboxes(object):
+
+    def __init__(self, num_classes, downsampling_factor):
+        self.num_classes = num_classes
+        self.downsampling_factor = downsampling_factor
+
+    def __call__(self, image, bboxes):
+        h,w = image.shape[:-1]
+        image = F.to_tensor(image)
+        if self.downsampling_factor is not None: 
+            blobs = np.zeros(shape=(self.num_classes + 2, h // self.downsampling_factor, w // self.downsampling_factor))
+        else: 
+            blobs = np.zeros(shape=(self.num_classes + 2, h, w))
+
+        for bbox_imgaug in bboxes:
+            cat = bbox_imgaug.label-1
+            bbox = [bbox_imgaug.x1, bbox_imgaug.y1, bbox_imgaug.width, bbox_imgaug.height]
+
+            new_blobs, ct_int = blob_for_bbox(bbox,  blobs[cat], self.downsampling_factor)
+            blobs[cat] = new_blobs
+            if ct_int is not None: 
+                ct_x, ct_y = ct_int
+                if ct_x < blobs.shape[2] and ct_y < blobs.shape[1]:
+                    blobs[-2, ct_y, ct_x] = bbox[3]
+                    blobs[-1, ct_y, ct_x] = bbox[2]
+                # else:
+                #     import matplotlib.pyplot as plt 
+                #     print(ct_x, ct_y)
+                #     fig, ax = plt.subplots(1,1,figsize=(20,20))
+                #     ax.imshow(blobs[cat])
+                #     import pickle
+                #     with open('verbose.pickle','wb') as f: 
+                #         pickle.dump((fig,ax),f)
+                #     plt.close()
+                
+
+        target = torch.from_numpy(blobs)
+        return image, target
+
+class Normalize(object):
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, image, target):
+        image = F.normalize(image, mean=self.mean, std=self.std)
+        return image, target
+
+class TrainTransforms:
+    def __init__(self, base_size, crop_size, num_classes, downsampling_factor, hflip_prob=0.5, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
+        self.num_classes = num_classes
+        self.downsampling_factor = downsampling_factor
+        self.base_size = base_size
+        self.crop_height, self.crop_width = crop_size
+        self.hflip_prob = hflip_prob
+        self.random_size_range = (int(self.base_size),int(2.0*self.base_size))
+        self.seq = iaa.Sequential([
+            iaa.Resize({"height": self.random_size_range, "width": "keep-aspect-ratio"}),
+            iaa.Fliplr(p=self.hflip_prob),
+            iaa.PadToFixedSize(width=self.crop_width, height=self.crop_height),
+            iaa.CropToFixedSize(width=self.crop_width, height=self.crop_height)
+        ])
+        self.last_transforms = Compose([ToTensorBboxes(num_classes, downsampling_factor),
+                                        Normalize(mean=mean,std=std)])
+
+
+        
+    def __call__(self, img, target):
+        
+        bboxes_imgaug = [BoundingBox(x1=bbox[0], y1=bbox[1], x2=bbox[0]+bbox[2], y2=bbox[1]+bbox[3], label=cat) \
+            for bbox, cat in zip(target['bboxes'],target['cats'])]
+        bboxes = BoundingBoxesOnImage(bboxes_imgaug, shape=img.shape)
+
+        img, bboxes_imgaug = self.seq(image=img, bounding_boxes=bboxes)
+        return self.last_transforms(img, bboxes_imgaug)
+
+class ValTransforms:
+    def __init__(self, base_size, crop_size, num_classes, downsampling_factor, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
+        self.num_classes = num_classes
+        self.downsampling_factor = downsampling_factor
+        self.base_size = base_size
+        self.crop_height, self.crop_width = crop_size
+        self.seq = iaa.Sequential([
+            iaa.Resize({"height": int(self.base_size), "width": "keep-aspect-ratio"}),
+            # iaa.Rotate((-45,45)),
+            iaa.CenterPadToFixedSize(width=self.crop_width, height=self.crop_height),
+            iaa.CenterCropToFixedSize(width=self.crop_width, height=self.crop_height)
+        ])
+        self.last_transforms = Compose([ToTensorBboxes(num_classes, downsampling_factor),
+                                        Normalize(mean=mean,std=std)])
+
+
+        
+    def __call__(self, img, target):
+        
+        bboxes_imgaug = [BoundingBox(x1=bbox[0], y1=bbox[1], x2=bbox[0]+bbox[2], y2=bbox[1]+bbox[3], label=cat) \
+            for bbox, cat in zip(target['bboxes'],target['cats'])]
+        bboxes = BoundingBoxesOnImage(bboxes_imgaug, shape=img.shape)
+
+        img, bboxes_imgaug = self.seq(image=img, bounding_boxes=bboxes)
+        return self.last_transforms(img, bboxes_imgaug)
+
+class TransformFrames:
+    def __init__(self):
+        transforms = []
+
+        transforms.append(T.Lambda(lambda img: cv2.cvtColor(img, cv2.COLOR_BGR2RGB)))
+        transforms.append(T.ToTensor())
+        transforms.append(T.Normalize(mean=[0.485, 0.456, 0.406],
+                                    std=[0.229, 0.224, 0.225]))
+
+        self.transforms = T.Compose(transforms)
+    
+    def __call__(self, img):
+        return self.transforms(img)

models/.gitignore

@@ -0,0 +1,5 @@
+# Ignore everything in this directory
+*
+# Except this file
+!.gitignore
+!download_pretrained_base.sh

packages.txt

@@ -0,0 +1,3 @@
+ffmpeg
+x264
+libx264-dev

requirements.txt

@@ -0,0 +1,13 @@
+scikit-image
+opencv-python
+pytorch
+torchvision
+pycocotools
+debugpy
+scipy
+tqdm
+tensorboard
+imgaug
+psycopg2-binary
+moviepy
+git+https://github.com/pykalman/pykalman.git

tools/misc.py

@@ -0,0 +1,201 @@
+from matplotlib.pyplot import grid
+import numpy as np
+import math
+import cv2
+import torch
+import torchvision.transforms.functional as F
+
+class ResizeForCenterNet(object):
+    def __init__(self, fix_res=False):
+        self.fix_res = fix_res
+
+    def __call__(self, image):
+        if self.fix_res:
+            new_h = 512
+            new_w = 512
+        else:
+            w, h = image.size
+            new_h = (h | 31) + 1
+            new_w = (w | 31) + 1
+        image = F.resize(image, (new_h, new_w))
+        return image
+
+def gaussian_radius(det_size, min_overlap=0.7):
+    height, width = det_size
+
+    a1  = 1
+    b1  = (height + width)
+    c1  = width * height * (1 - min_overlap) / (1 + min_overlap)
+    sq1 = np.sqrt(b1 ** 2 - 4 * a1 * c1)
+    r1  = (b1 + sq1) / 2
+
+    a2  = 4
+    b2  = 2 * (height + width)
+    c2  = (1 - min_overlap) * width * height
+    sq2 = np.sqrt(b2 ** 2 - 4 * a2 * c2)
+    r2  = (b2 + sq2) / 2
+
+    a3  = 4 * min_overlap
+    b3  = -2 * min_overlap * (height + width)
+    c3  = (min_overlap - 1) * width * height
+    sq3 = np.sqrt(b3 ** 2 - 4 * a3 * c3)
+    r3  = (b3 + sq3) / 2
+    return min(r1, r2, r3)
+
+def gaussian2D(shape, sigma=1):
+    m, n = [(ss - 1.) / 2. for ss in shape]
+    y, x = np.ogrid[-m:m+1,-n:n+1]
+
+    h = np.exp(-(x * x + y * y) / (2 * sigma * sigma))
+    h[h < np.finfo(h.dtype).eps * h.max()] = 0
+    return h
+
+def draw_umich_gaussian(heatmap, center, radius, k=1):
+    diameter = 2 * radius + 1
+    gaussian = gaussian2D((diameter, diameter), sigma=diameter / 6)
+
+    x, y = int(center[0]), int(center[1])
+
+    height, width = heatmap.shape[0:2]
+
+    left, right = min(x, radius), min(width - x, radius + 1)
+    top, bottom = min(y, radius), min(height - y, radius + 1)
+
+    masked_heatmap  = heatmap[y - top:y + bottom, x - left:x + right]
+    masked_gaussian = gaussian[radius - top:radius + bottom, radius - left:radius + right]
+    if min(masked_gaussian.shape) > 0 and min(masked_heatmap.shape) > 0: # TODO debug
+        np.maximum(masked_heatmap, masked_gaussian * k, out=masked_heatmap)
+
+    return heatmap
+
+def blob_for_bbox(bbox, heatmap, downsampling_factor=None):
+    if downsampling_factor is not None:
+        left, top, w, h = [bbox_coord // downsampling_factor for bbox_coord in bbox]
+    else:
+        left, top, w, h = [bbox_coord for bbox_coord in bbox]
+
+    right, bottom = left+w, top+h
+    ct_int = None
+    if h > 0 and w > 0:
+        radius = gaussian_radius((math.ceil(h), math.ceil(w)))
+        radius = max(0, int(radius))
+        ct = np.array([(left + right) / 2, (top + bottom) / 2], dtype=np.float32)
+        ct_int = ct.astype(np.int32)
+        heatmap = draw_umich_gaussian(heatmap, ct_int, radius)
+    return heatmap, ct_int
+
+def pre_process_centernet(image, meta=None, fix_res=True):
+    scale = 1.0
+    mean = [0.408, 0.447, 0.47]
+    std = [0.289, 0.274, 0.278]
+    height, width = image.shape[0:2]
+    new_height = int(height * scale)
+    new_width  = int(width * scale)
+    if fix_res:
+        inp_height, inp_width = 512, 512
+        c = np.array([new_width / 2., new_height / 2.], dtype=np.float32)
+        s = max(height, width) * 1.0
+    else:
+        inp_height = (new_height | 31) + 1
+        inp_width = (new_width | 31) + 1
+        c = np.array([new_width // 2, new_height // 2], dtype=np.float32)
+        s = np.array([inp_width, inp_height], dtype=np.float32)
+
+    trans_input = get_affine_transform(c, s, 0, [inp_width, inp_height])
+    resized_image = cv2.resize(image, (new_width, new_height))
+    inp_image = cv2.warpAffine(
+        resized_image, trans_input, (inp_width, inp_height),
+        flags=cv2.INTER_LINEAR)
+    inp_image = ((inp_image / 255. - mean) / std).astype(np.float32)
+
+    images = inp_image.transpose(2, 0, 1).reshape(1, 3, inp_height, inp_width)
+    # if self.opt.flip_test:
+    #     images = np.concatenate((images, images[:, :, :, ::-1]), axis=0)
+    images = torch.from_numpy(images)
+    # meta = {'c': c, 's': s,
+    #         'out_height': inp_height // self.opt.down_ratio,
+    #         'out_width': inp_width // self.opt.down_ratio}
+    return images.squeeze() #, meta
+
+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 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
+
+def get_3rd_point(a, b):
+    direct = a - b
+    return b + np.array([-direct[1], direct[0]], dtype=np.float32)
+
+def load_checkpoint(model, trained_model_weights_filename):
+    checkpoint = torch.load(trained_model_weights_filename, map_location='cpu')
+    model.load_state_dict(checkpoint['model'])
+    return model
+
+def load_model(arch, model_weights, device):
+
+    if model_weights is None:
+        if arch == 'mobilenet_v3_small':
+            model_weights = 'models/mobilenet_v3_pretrained.pth'
+            arch = 'mobilenetv3small'
+        elif arch == 'res_18':
+            model_weights = 'models/res18_pretrained.pth'
+        elif arch == 'dla_34':
+            model_weights = 'models/dla_34_pretrained.pth'
+
+    heads = {'hm':1} if arch != 'dla_34' else {'hm':1, 'wh':2}
+
+    model = create_base(arch, heads=heads, head_conv=256).to(device)
+    model = load_checkpoint(model, model_weights)
+    for param in model.parameters():
+        param.requires_grad = False
+    model.eval()
+
+    return model
+
+def _calculate_euclidean_similarity(distances, zero_distance):
+    """ Calculates the euclidean distance between two sets of detections, and then converts this into a similarity
+    measure with values between 0 and 1 using the following formula: sim = max(0, 1 - dist/zero_distance).
+    The default zero_distance of 2.0, corresponds to the default used in MOT15_3D, such that a 0.5 similarity
+    threshold corresponds to a 1m distance threshold for TPs.
+    """
+    sim = np.maximum(0, 1 - distances/zero_distance)
+    return sim

tools/optical_flow.py

@@ -0,0 +1,68 @@
+import cv2
+import numpy as np
+
+# params for ShiTomasi corner detection
+feature_params = dict( maxCorners = 100,
+                       qualityLevel = 0.3,
+                       minDistance = 2,
+                       blockSize = 7 )
+# Parameters for lucas kanade optical flow
+lk_params = dict( winSize  = (15,15),
+                  maxLevel = 20,
+                  criteria = (cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03))
+
+def viz_dense_flow(frame_reader, nb_frames_to_process):
+
+    ret, frame1 = frame_reader.read_frame()
+    prvs = cv2.cvtColor(frame1,cv2.COLOR_BGR2GRAY)
+    hsv = np.zeros_like(frame1)
+    hsv[...,1] = 0
+
+    while(frame_reader.nb_frames_read < nb_frames_to_process):  
+        ret, frame2 = frame_reader.read_frame()
+        next = cv2.cvtColor(frame2,cv2.COLOR_BGR2GRAY)
+        flow = cv2.calcOpticalFlowFarneback(prvs,next, None, 0.5, 3, 15, 3, 5, 1.2, 0)
+        mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1])
+        hsv[...,0] = ang*180/np.pi/2
+        hsv[...,2] = cv2.normalize(mag,None,0,255,cv2.NORM_MINMAX)
+        bgr = cv2.cvtColor(hsv,cv2.COLOR_HSV2BGR)
+        cv2.imshow('frame',np.concatenate([frame2, bgr], axis=0))
+        k = cv2.waitKey(0) & 0xff
+        prvs = next
+    cv2.destroyAllWindows()  
+
+def flow_opencv_dense(img, img2):
+
+    prvs = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    next = cv2.cvtColor(img2,cv2.COLOR_BGR2GRAY)
+
+    flow = cv2.calcOpticalFlowFarneback(prvs,next, None, 0.5, 3, 15, 3, 5, 1.2, 0)
+
+    return flow
+
+def flow_opencv_sparse(img, img2, p0):
+    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    img2_gray = cv2.cvtColor(img2, cv2.COLOR_BGR2GRAY)
+
+    p1, st, err = cv2.calcOpticalFlowPyrLK(img_gray, img2_gray, p0, None, **lk_params)
+
+    # Select good points
+    good_new = p1[st==1]
+
+    return good_new
+
+def dense_flow_norm(dense_flow):
+    v, u = dense_flow[...,1], dense_flow[...,0]
+    return np.sqrt(u ** 2 + v ** 2)
+
+def compute_flow(frame0, frame1, downsampling_factor):
+
+    if downsampling_factor > 1:
+        h, w = frame0.shape[:-1]
+        new_h = h // downsampling_factor
+        new_w = w // downsampling_factor
+
+        frame0 = cv2.resize(frame0, (new_w, new_h))
+        frame1 = cv2.resize(frame1, (new_w, new_h))
+
+    return flow_opencv_dense(frame0, frame1)

tools/video_readers.py

@@ -0,0 +1,193 @@
+import cv2
+import torch
+from tqdm import tqdm
+from itertools import cycle
+class AdvancedFrameReader:
+    def __init__(self, video_name, read_every, rescale_factor, init_time_min, init_time_s):
+
+        self.cap = cv2.VideoCapture(video_name)
+
+        self.original_width = self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)
+        self.original_height = self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
+
+        if rescale_factor != 1:
+            self.set_rescale_factor(rescale_factor)
+        else:
+            self.original_shape_mode = True
+
+        self.init_rescale_factor = rescale_factor
+
+
+        self.frame_skip = read_every -  1
+        self.fps = self.cap.get(cv2.CAP_PROP_FPS)/read_every
+        print(f'Reading at {self.fps:.2f} fps')
+
+        self.set_time_position(init_time_min, init_time_s)
+        self.init_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
+        self.total_num_frames = self.cap.get(cv2.CAP_PROP_FRAME_COUNT)
+
+
+    def post_process(self, ret, frame):
+        if ret:
+            if self.original_shape_mode:
+                return ret, frame
+            else:
+                return ret, cv2.resize(frame, self.new_shape)
+        else:
+            return ret, []
+
+    def skip(self):
+        if not self.nb_frames_read:
+            return
+        else:
+            for _ in range(self.frame_skip):
+                self.cap.read()
+
+    def read_frame(self):
+
+        self.skip()
+        ret, frame = self.cap.read()
+        self.nb_frames_read += 1
+        return self.post_process(ret, frame)
+
+    def set_time_position(self, time_min, time_s):
+        time = 60 * time_min  + time_s
+        self.cap.set(cv2.CAP_PROP_POS_MSEC, 1000 * time)
+        self.init_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
+        print('Reading from {}min{}sec'.format(time_min,time_s))
+        self.nb_frames_read = 0
+
+    def reset_init_frame(self):
+        self.cap.set(cv2.CAP_PROP_POS_FRAMES,int(self.init_frame))
+        self.nb_frames_read = 0
+
+    def set_init_frame(self, init_frame):
+        self.cap.set(cv2.CAP_PROP_POS_FRAMES, int(init_frame))
+        self.init_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
+        self.nb_frames_read = 0
+
+    def set_rescale_factor(self, rescale_factor):
+        width = int(self.original_width/rescale_factor)
+        height = int(self.original_height/rescale_factor)
+        self.new_shape  = (width, height)
+        self.original_shape_mode = False
+        print('Reading in {}x{}'.format(width, height))
+
+    def set_original_shape_mode(self, mode):
+        self.original_shape_mode = mode
+
+    def reset_init_rescale_factor(self):
+        self.set_rescale_factor(self.init_rescale_factor)
+
+class IterableFrameReader:
+
+    def __init__(self, video_filename, skip_frames=0, output_shape=None, progress_bar=False, preload=False, max_frame=0):
+
+        self.video = cv2.VideoCapture(video_filename)
+        self.input_shape = (self.video.get(cv2.CAP_PROP_FRAME_WIDTH), self.video.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        self.skip_frames = skip_frames
+        self.preload = preload
+        self.total_num_frames = self.video.get(cv2.CAP_PROP_FRAME_COUNT)
+        self.max_num_frames = min(max_frame, self.total_num_frames) if max_frame!=0 else self.total_num_frames
+        self.counter = 0
+
+        if output_shape is None:
+            w, h = self.input_shape
+            new_h = (h | 31) + 1
+            new_w = (w | 31) + 1
+            self.output_shape = (new_w, new_h)
+        else:
+            self.output_shape = output_shape
+
+        self.fps = self.video.get(cv2.CAP_PROP_FPS) / (self.skip_frames+1)
+
+        print(f'Reading video at {self.fps}fps.')
+        if progress_bar:
+            self.progress_bar = tqdm(total=int(self.max_num_frames/(self.skip_frames+1)), leave=True)
+            self.progress_bar_update = self.progress_bar.update
+        else:
+            self.progress_bar_update = lambda: None
+
+        if self.preload:
+            print('Preloading frames in RAM...')
+            self.frames = self._load_all_frames()
+
+
+    def _load_all_frames(self):
+        frames = []
+        while True:
+            ret, frame = self._read_frame()
+            if ret:
+                frames.append(frame)
+            else: break
+        if self.progress_bar: self.progress_bar.reset()
+        return frames
+
+    def __next__(self):
+        self.counter+=1
+        if self.preload:
+            if self.counter < len(self.frames):
+                frame = self.frames[self.counter]
+                self.progress_bar_update()
+                return frame
+        else:
+            if self.counter < self.max_num_frames:
+                ret, frame = self._read_frame()
+                if ret:
+                    return frame
+        print("here")
+
+        self.counter=0
+        self.video.set(cv2.CAP_PROP_POS_FRAMES, 0)
+        if self.progress_bar: self.progress_bar.reset()
+        raise StopIteration
+
+    def _read_frame(self):
+        ret, frame = self.video.read()
+        self._skip_frames()
+        if ret:
+            self.progress_bar_update()
+            frame =  cv2.resize(frame, self.output_shape)
+        return ret, frame
+
+    def __iter__(self):
+        return self
+
+    def _skip_frames(self):
+        for _ in range(self.skip_frames):
+            self.counter+=1
+            self.video.read()
+
+
+class SimpleVideoReader:
+    def __init__(self, video_filename, skip_frames=0):
+        self.skip_frames = skip_frames
+        self.video = cv2.VideoCapture(video_filename)
+        self.shape = (int(self.video.get(cv2.CAP_PROP_FRAME_WIDTH)), int(self.video.get(cv2.CAP_PROP_FRAME_HEIGHT)))
+        self.fps = self.video.get(cv2.CAP_PROP_FPS) / (skip_frames+1)
+        self.frame_nb = 0
+        self.num_frames = self.video.get(cv2.CAP_PROP_FRAME_COUNT)
+
+    def read(self):
+        ret, frame = self.video.read()
+        self.frame_nb+=1
+        self._skip_frames()
+        return ret, frame, self.frame_nb-1
+
+    def set_frame(self, frame_nb_to_set):
+        self.video.set(cv2.CAP_PROP_POS_FRAMES, frame_nb_to_set)
+        self.frame_nb = frame_nb_to_set
+
+    def _skip_frames(self):
+        for _ in range(self.skip_frames):
+            self.video.read()
+
+class TorchIterableFromReader(torch.utils.data.IterableDataset):
+
+    def __init__(self, reader, transforms):
+        self.transforms = transforms
+        self.reader = reader
+
+    def __iter__(self):
+        for frame in self.reader:
+            yield self.transforms(frame)

tracking/postprocess_and_count_tracks.py

@@ -0,0 +1,119 @@
+import numpy as np
+import argparse
+from scipy.signal import convolve
+from tracking.utils import write_tracking_results_to_file, read_tracking_results
+
+import json
+
+
+def filter_tracks(tracklets, kappa, tau):
+    """ filters the tracks depending on params
+    kappa: size of the moving average window
+    tau: minimum length of tracklet
+    returns raw filtered tracks
+    """
+    if not kappa == 1:
+        tracks = filter_by_nb_consecutive_obs(tracklets, kappa, tau)
+    else: tracks = tracklets
+    results = []
+    for tracker_nb, associated_detections in enumerate(tracks):
+        for associated_detection in associated_detections:
+            results.append((associated_detection[0], tracker_nb, associated_detection[1], associated_detection[2]))
+
+    results = sorted(results, key=lambda x: x[0])
+    return results
+
+
+def postprocess_for_api(results):
+    """ Converts tracking results into json object for API
+    """
+    result_list = []
+    id_list = {}
+
+    for res in results:
+        frame_number = res[0]
+        box = [round(res[2],1), round(res[3],1), round(res[2],1), round(res[3],1)]
+        id = res[1]
+        # if the id is not already is the results, add a new jsonline
+        if id not in id_list:
+            id_list[id] = len(result_list)
+            result_list.append({"label":"fragments",
+                                "id": id,
+                                "frame_to_box": {str(frame_number): box}})
+        # otherwise, retrieve the jsonline and append the box
+        else:
+            result_list[id_list[id]]["frame_to_box"][str(frame_number)] = box
+    return {"detected_trash": result_list}
+
+
+def write(results, output_name):
+    """ Writes the results in two files:
+    - tracking in a Mathis format xxx_track.txt (frame, id, box_x, box_y, ...)
+    - the number of detected objects in a separate file xxx_count.txt
+    """
+    output_tracks_filename = output_name.split('.')[0]+'_tracks.txt'
+    write_tracking_results_to_file(results, ratio_x=1, ratio_y=1,
+                                   output_filename=output_tracks_filename)
+
+    with open(output_name.split('.')[0]+'_count.txt','w') as out_file:
+        if len(results):
+            out_file.write(f'{max(result[1]+1 for result in results)}')
+        else:
+            out_file.write('0')
+
+
+def threshold(tracklets, tau):
+    return [tracklet for tracklet in tracklets if len(tracklet) > tau]
+
+
+def compute_moving_average(tracklet, kappa):
+    if len(tracklet)==0 or len(tracklet[0])==0:
+        return tracklet
+    pad = (kappa-1)//2
+    observation_points = np.zeros(tracklet[-1][0] - tracklet[0][0] + 1)
+    first_frame_id = tracklet[0][0] - 1
+    for observation in tracklet:
+        frame_id = observation[0] - 1
+        observation_points[frame_id - first_frame_id] = 1
+    density_fill = convolve(observation_points, np.ones(kappa)/kappa, mode='same')
+    if pad>0 and len(observation_points) >= kappa:
+        density_fill[:pad] = density_fill[pad:2*pad]
+        density_fill[-pad:] = density_fill[-2*pad:-pad]
+    density_fill = observation_points * density_fill
+
+    return  density_fill[density_fill > 0]
+
+
+def filter_by_nb_consecutive_obs(tracklets, kappa, tau):
+
+    new_tracklets = []
+
+    for tracklet in tracklets:
+        new_tracklet = []
+        density_fill = compute_moving_average(tracklet, kappa=kappa)
+        for (observation, density_fill_value) in zip(tracklet, density_fill):
+            if density_fill_value > 0.6:
+                new_tracklet.append(observation)
+        new_tracklets.append(new_tracklet)
+
+    return threshold(new_tracklets, tau)
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--input_file',type=str)
+    parser.add_argument('--output_name',type=str)
+    parser.add_argument('--kappa',type=int)
+    parser.add_argument('--tau',type=int)
+    parser.add_argument('--output_type',type=str,default="api")
+    args = parser.parse_args()
+
+    tracklets = read_tracking_results(args.input_file)
+    filtered_results = filter_tracks(tracklets, args.kappa, args.tau)
+    if args.output_type == "api":
+        output = postprocess_for_api(filtered_results)
+        with open(args.output_name, 'w') as f:
+            json.dump(output, f)
+    else:
+        write(filtered_results, args.output_name)

tracking/track_video.py

@@ -0,0 +1,152 @@
+import cv2
+import numpy as np
+import os
+from detection.detect import detect
+from tracking.utils import in_frame, init_trackers
+from tools.optical_flow import compute_flow
+from tracking.trackers import get_tracker
+import matplotlib.pyplot as plt
+from scipy.spatial.distance import euclidean
+from scipy.optimize import linear_sum_assignment
+import torch
+
+
+class Display:
+    def __init__(self, on, interactive=True):
+        self.on = on
+        self.fig, self.ax = plt.subplots()
+        self.interactive = interactive
+        if interactive:
+            plt.ion()
+        self.colors =  plt.rcParams['axes.prop_cycle'].by_key()['color']
+        self.legends = []
+        self.plot_count = 0
+
+    def display(self, trackers):
+
+        something_to_show = False
+        for tracker_nb, tracker in enumerate(trackers):
+            if tracker.enabled:
+                tracker.fill_display(self, tracker_nb)
+                something_to_show = True
+
+        self.ax.imshow(self.latest_frame_to_show)
+
+        if len(self.latest_detections):
+            self.ax.scatter(self.latest_detections[:, 0], self.latest_detections[:, 1], c='r', s=40)
+
+        if something_to_show:
+            self.ax.xaxis.tick_top()
+            plt.legend(handles=self.legends)
+            self.fig.canvas.draw()
+            if self.interactive:
+                plt.show()
+                while not plt.waitforbuttonpress():
+                    continue
+            else:
+                plt.savefig(os.path.join('plots',str(self.plot_count)))
+            self.ax.cla()
+            self.legends = []
+            self.plot_count+=1
+
+    def update_detections_and_frame(self, latest_detections, frame):
+        self.latest_detections = latest_detections
+        self.latest_frame_to_show = cv2.cvtColor(cv2.resize(frame, self.display_shape), cv2.COLOR_BGR2RGB)
+
+
+def build_confidence_function_for_trackers(trackers, flow01):
+    tracker_nbs = []
+    confidence_functions = []
+    for tracker_nb, tracker in enumerate(trackers):
+        if tracker.enabled:
+            tracker_nbs.append(tracker_nb)
+            confidence_functions.append(tracker.build_confidence_function(flow01))
+    return tracker_nbs, confidence_functions
+
+def associate_detections_to_trackers(detections_for_frame, trackers, flow01, confidence_threshold):
+    tracker_nbs, confidence_functions = build_confidence_function_for_trackers(trackers, flow01)
+    assigned_trackers = [None]*len(detections_for_frame)
+    if len(tracker_nbs):
+        cost_matrix = np.zeros(shape=(len(detections_for_frame),len(tracker_nbs)))
+        for detection_nb, detection in enumerate(detections_for_frame):
+            for tracker_id, confidence_function in enumerate(confidence_functions):
+                score = confidence_function(detection)
+                if score > confidence_threshold:
+                    cost_matrix[detection_nb,tracker_id] = score
+                else:
+                    cost_matrix[detection_nb,tracker_id] = 0
+        row_inds, col_inds = linear_sum_assignment(cost_matrix,maximize=True)
+        for row_ind, col_ind in zip(row_inds, col_inds):
+            if cost_matrix[row_ind,col_ind] > confidence_threshold: assigned_trackers[row_ind] = tracker_nbs[col_ind]
+
+    return assigned_trackers
+
+def track_video(reader, detections, args, engine, transition_variance, observation_variance, display):
+    init = False
+    trackers = dict()
+    frame_nb = 0
+    frame0 = next(reader)
+    detections_for_frame = next(detections)
+
+    max_distance = euclidean(reader.output_shape, np.array([0,0]))
+    delta = 0.05*max_distance
+
+    if display is not None and display.on:
+
+        display.display_shape = (reader.output_shape[0] // args.downsampling_factor, reader.output_shape[1] // args.downsampling_factor)
+        display.update_detections_and_frame(detections_for_frame, frame0)
+
+    if len(detections_for_frame):
+        trackers = init_trackers(engine, detections_for_frame, frame_nb, transition_variance, observation_variance, delta)
+        init = True
+
+    if display is not None and display.on: display.display(trackers)
+
+    for frame_nb, (frame1, detections_for_frame) in enumerate(zip(reader, detections), start=1):
+
+        if display is not None and display.on:
+            display.update_detections_and_frame(detections_for_frame, frame1)
+
+        if not init:
+            if len(detections_for_frame):
+                trackers = init_trackers(engine, detections_for_frame, frame_nb, transition_variance, observation_variance, delta)
+                init = True
+
+        else:
+
+            new_trackers = []
+            flow01 = compute_flow(frame0, frame1, args.downsampling_factor)
+
+            if len(detections_for_frame):
+
+                assigned_trackers = associate_detections_to_trackers(detections_for_frame, trackers,
+                                                                     flow01, args.confidence_threshold)
+
+                for detection, assigned_tracker in zip(detections_for_frame, assigned_trackers):
+                    if in_frame(detection, flow01.shape[:-1]):
+                        if assigned_tracker is None :
+                            new_trackers.append(engine(frame_nb, detection, transition_variance, observation_variance, delta))
+                        else:
+                            trackers[assigned_tracker].update(detection, flow01, frame_nb)
+
+            for tracker in trackers:
+                tracker.update_status(flow01)
+
+            if len(new_trackers):
+                trackers.extend(new_trackers)
+
+        if display is not None and display.on:
+            display.display(trackers)
+        frame0 = frame1.copy()
+
+
+    results = []
+    tracklets = [tracker.tracklet for tracker in trackers]
+
+    for tracker_nb, associated_detections in enumerate(tracklets):
+        for associated_detection in associated_detections:
+            results.append((associated_detection[0], tracker_nb, associated_detection[1][0], associated_detection[1][1]))
+
+    results = sorted(results, key=lambda x: x[0])
+
+    return results

tracking/trackers.py

@@ -0,0 +1,269 @@
+import numpy as np
+from scipy.stats import multivariate_normal
+from tracking.utils import in_frame, exp_and_normalise, GaussianMixture
+from pykalman import KalmanFilter, AdditiveUnscentedKalmanFilter
+import matplotlib.patches as mpatches
+
+class Tracker:
+
+    def __init__(self, frame_nb, X0, transition_variance, observation_variance, delta):
+
+        self.transition_covariance = np.diag(transition_variance)
+        self.observation_covariance = np.diag(observation_variance)
+        self.updated = False
+        self.steps_since_last_observation = 0
+        self.enabled = True
+        self.tracklet = [(frame_nb, X0)]
+        self.delta = delta
+
+    def store_observation(self, observation, frame_nb):
+        self.tracklet.append((frame_nb, observation))
+        self.updated = True
+
+    def update_status(self, flow):
+        if self.enabled and not self.updated:
+            self.steps_since_last_observation += 1
+            self.enabled = self.update(None, flow)
+        else:
+            self.steps_since_last_observation = 0
+        self.updated = False
+
+    def build_confidence_function(self, flow):
+
+        def confidence_from_multivariate_distribution(coord, distribution):
+            delta = self.delta
+            x = coord[0]
+            y = coord[1]
+            right_top = np.array([x+delta, y+delta])
+            left_low = np.array([x-delta, y-delta])
+            right_low = np.array([x+delta, y-delta])
+            left_top = np.array([x-delta, y+delta])
+
+            return distribution.cdf(right_top) \
+                - distribution.cdf(right_low) \
+                - distribution.cdf(left_top) \
+                + distribution.cdf(left_low)
+
+        distribution = self.predictive_distribution(flow)
+        
+        return lambda coord: confidence_from_multivariate_distribution(coord, distribution)
+    
+    def get_display_colors(self, display, tracker_nb):
+        colors = display.colors
+        color = colors[tracker_nb % len(colors)]
+        display.legends.append(mpatches.Patch(color=color, label=len(self.tracklet)))
+        return colors[tracker_nb % len(colors)]
+
+class SMC(Tracker):     
+    def set_param(param):
+        SMC.n_particles = int(param)
+
+    def __init__(self, frame_nb, X0, transition_variance, observation_variance, delta):
+        super().__init__(frame_nb, X0, transition_variance, observation_variance, delta)
+
+        self.particles = multivariate_normal(
+            X0, cov=self.observation_covariance).rvs(SMC.n_particles)
+        self.normalized_weights = np.ones(SMC.n_particles)/SMC.n_particles
+
+    def update(self, observation, flow, frame_nb=None):
+        if observation is not None: self.store_observation(observation, frame_nb)
+        self.resample()
+        enabled = self.move_particles(flow)
+        if observation is not None: 
+            self.importance_reweighting(observation)
+        else: 
+            self.normalized_weights = np.ones(
+                len(self.particles))/len(self.particles)
+
+        return enabled 
+
+    def state_transition(self, state, flow):
+
+        mean = state + \
+            flow[max(0, int(state[1])),
+                 max(0, int(state[0])), :]
+        cov = np.diag(self.transition_covariance)
+        return multivariate_normal(mean, cov)
+    
+    def observation(self, state):
+
+        return multivariate_normal(state, self.observation_covariance)
+
+    def move_particles(self, flow):
+        new_particles = []
+        for particle in self.particles:
+            new_particle = self.state_transition(particle, flow).rvs(1)
+            if in_frame(new_particle, flow.shape[:-1]):
+                new_particles.append(new_particle)
+        if len(new_particles):
+            self.particles = np.array(new_particles)
+            enabled = True
+        else:
+            enabled = False
+        
+        return enabled
+
+    def importance_reweighting(self, observation):
+        log_weights_unnormalized = np.zeros(len(self.particles))
+        for particle_nb, particle in enumerate(self.particles):
+            log_weights_unnormalized[particle_nb] = self.observation(
+                particle).logpdf(observation)
+        self.normalized_weights = exp_and_normalise(log_weights_unnormalized)
+
+    def resample(self):
+        resampling_indices = np.random.choice(
+            a=len(self.particles), p=self.normalized_weights, size=len(self.particles))
+        self.particles = self.particles[resampling_indices]
+
+    def predictive_distribution(self, flow, nb_new_particles=5):
+        new_particles = []
+        new_weights = []
+
+        for particle, normalized_weight in zip(self.particles, self.normalized_weights):
+            new_particles_for_particle = self.state_transition(
+                particle, flow).rvs(nb_new_particles)
+
+            new_particles_for_particle = [
+                particle for particle in new_particles_for_particle if in_frame(particle, flow.shape[:-1])]
+
+            if len(new_particles_for_particle):
+                new_particles.extend(new_particles_for_particle)
+                new_weights.extend([normalized_weight/len(new_particles_for_particle)] *
+                                len(new_particles_for_particle))
+
+        new_particles = np.array(new_particles)
+
+
+
+        return GaussianMixture(new_particles, self.observation_covariance, new_weights)
+
+    def fill_display(self, display, tracker_nb):
+        color = self.get_display_colors(display, tracker_nb)
+        display.ax.scatter(self.particles[:,0], self.particles[:,1], s=5, c=color)
+
+class EKF(Tracker): 
+
+    def __init__(self, frame_nb, X0, transition_variance, observation_variance, delta):
+            super().__init__(frame_nb, X0, transition_variance, observation_variance, delta)
+            self.filter = KalmanFilter(initial_state_mean=X0, 
+                                       initial_state_covariance=self.observation_covariance, 
+                                       transition_covariance=self.transition_covariance,
+                                       observation_matrices=np.eye(2),
+                                       observation_covariance=self.observation_covariance)
+
+            self.filtered_state_mean = X0
+            self.filtered_state_covariance = self.observation_covariance
+
+    def get_update_parameters(self, flow):
+
+        flow_value = flow[int(self.filtered_state_mean[1]),int(self.filtered_state_mean[0]), :]
+
+
+        grad_flow_value = np.array([np.gradient(flow[:,:,0]),np.gradient(flow[:,:,1])])[:,:,int(self.filtered_state_mean[1]),int(self.filtered_state_mean[0])]
+        return np.eye(2) + grad_flow_value, flow_value - grad_flow_value.dot(self.filtered_state_mean) 
+
+        
+    def EKF_step(self, observation, flow):
+        transition_matrix, transition_offset = self.get_update_parameters(flow)
+
+        return self.filter.filter_update(self.filtered_state_mean, 
+                                        self.filtered_state_covariance, 
+                                        transition_matrix=transition_matrix,
+                                        transition_offset=transition_offset,
+                                        observation=observation)
+                                                                                             
+    def update(self, observation, flow, frame_nb=None):
+        if observation is not None: self.store_observation(observation, frame_nb)
+
+        self.filtered_state_mean, self.filtered_state_covariance = self.EKF_step(observation, flow)
+
+        enabled=False if not in_frame(self.filtered_state_mean,flow.shape[:-1]) else True
+
+        return enabled
+
+    def predictive_distribution(self, flow):
+
+        filtered_state_mean, filtered_state_covariance = self.EKF_step(None, flow)
+
+        distribution = multivariate_normal(filtered_state_mean, filtered_state_covariance + self.observation_covariance)
+
+        return distribution
+    
+    def fill_display(self, display, tracker_nb):
+        yy, xx = np.mgrid[0:display.display_shape[1]:1, 0:display.display_shape[0]:1]
+        pos = np.dstack((xx, yy))    
+        distribution = multivariate_normal(self.filtered_state_mean, self.filtered_state_covariance)
+
+        color = self.get_display_colors(display, tracker_nb)
+        cs = display.ax.contour(distribution.pdf(pos), colors=color)
+        display.ax.clabel(cs, inline=True, fontsize='large')
+        display.ax.scatter(self.filtered_state_mean[0], self.filtered_state_mean[1], color=color, marker="x", s=100)
+
+class UKF(Tracker):
+
+    def __init__(self, frame_nb, X0, transition_variance, observation_variance, delta):
+            super().__init__(frame_nb, X0, transition_variance, observation_variance, delta)
+            self.filter = AdditiveUnscentedKalmanFilter(initial_state_mean=X0, 
+                                        initial_state_covariance=self.observation_covariance, 
+                                        observation_functions = lambda z: np.eye(2).dot(z),
+                                        transition_covariance=self.transition_covariance,
+                                        observation_covariance=self.observation_covariance)
+
+            self.filtered_state_mean = X0
+            self.filtered_state_covariance = self.observation_covariance
+
+    def UKF_step(self, observation, flow):
+        return self.filter.filter_update(self.filtered_state_mean, 
+                                        self.filtered_state_covariance, 
+                                        transition_function=lambda x: x + flow[int(x[1]),int(x[0]),:],
+                                        observation=observation)
+
+                                                                                             
+    def update(self, observation, flow, frame_nb=None):
+        if observation is not None: self.store_observation(observation, frame_nb)
+
+        self.filtered_state_mean, self.filtered_state_covariance = self.UKF_step(observation, flow)
+
+        enabled=False if not in_frame(self.filtered_state_mean,flow.shape[:-1]) else True
+
+        return enabled
+
+    def predictive_distribution(self, flow):
+
+        filtered_state_mean, filtered_state_covariance = self.UKF_step(None, flow)
+
+        distribution = multivariate_normal(filtered_state_mean, filtered_state_covariance + self.observation_covariance)
+
+        return distribution
+    
+    def fill_display(self, display, tracker_nb):
+        yy, xx = np.mgrid[0:display.display_shape[1]:1, 0:display.display_shape[0]:1]
+        pos = np.dstack((xx, yy))    
+        distribution = multivariate_normal(self.filtered_state_mean, self.filtered_state_covariance)
+
+        color = self.get_display_colors(display, tracker_nb)
+        cs = display.ax.contour(distribution.pdf(pos), colors=color)
+        display.ax.clabel(cs, inline=True, fontsize='large')
+        display.ax.scatter(self.filtered_state_mean[0], self.filtered_state_mean[1], color=color, marker="x", s=100)
+
+trackers = {'EKF': EKF,
+           'SMC': SMC,
+           'UKF': UKF} 
+
+def get_tracker(algorithm_and_params):
+    print(f'{algorithm_and_params} will be used for tracking.')
+
+    splitted_name = algorithm_and_params.split('_')
+    if len(splitted_name) > 1: 
+        algorithm_name, param = splitted_name
+        tracker = trackers[algorithm_name]
+        tracker.set_param(param)
+
+    else: 
+        algorithm_name = splitted_name[0]
+        tracker = trackers[algorithm_name]
+    
+    return tracker
+
+
+

tracking/utils.py

@@ -0,0 +1,221 @@
+from scipy.stats import multivariate_normal
+import numpy as np
+import os
+import cv2
+from torch.utils.data import DataLoader
+import torch
+from tools.video_readers import TorchIterableFromReader
+from time import time
+from detection.transforms import TransformFrames
+from collections import defaultdict
+from moviepy.editor import ImageSequenceClip
+from skimage.transform import downscale_local_mean
+
+class GaussianMixture(object):
+    def __init__(self, means, covariance, weights):
+        self.components = [multivariate_normal(
+            mean=mean, cov=covariance) for mean in means]
+        self.weights = weights
+
+    def pdf(self, x):
+        result = 0
+        for weight, component in zip(self.weights, self.components):
+            result += weight*component.pdf(x)
+        return result
+
+    def logpdf(self, x):
+        return np.log(self.pdf(x))
+
+    def cdf(self, x):
+        result = 0
+        for weight, component in zip(self.weights, self.components):
+            result += weight*component.cdf(x)
+        return result
+
+def init_trackers(engine, detections, frame_nb, state_variance, observation_variance, delta):
+    trackers = []
+
+    for detection in detections:
+        tracker_for_detection = engine(frame_nb, detection, state_variance, observation_variance, delta)
+        trackers.append(tracker_for_detection)
+
+    return trackers
+
+def exp_and_normalise(lw):
+    w = np.exp(lw - lw.max())
+    return w / w.sum()
+
+def in_frame(position, shape, border=0.02):
+
+
+    shape_x = shape[1]
+    shape_y = shape[0]
+    x = position[0]
+    y = position[1]
+
+    return x > border*shape_x and x < (1-border)*shape_x and y > border*shape_y and y < (1-border)*shape_y
+
+def gather_filenames_for_video_in_annotations(video, images, data_dir):
+    images_for_video = [image for image in images
+                        if image['video_id'] == video['id']]
+    images_for_video = sorted(
+        images_for_video, key=lambda image: image['frame_id'])
+
+    return [os.path.join(data_dir, image['file_name'])
+                 for image in images_for_video]
+
+def get_detections_for_video(reader, detector, batch_size=16, device=None):
+
+    detections = []
+    dataset = TorchIterableFromReader(reader, TransformFrames())
+    loader = DataLoader(dataset, batch_size=batch_size)
+    average_times = []
+    with torch.no_grad():
+        for preprocessed_frames in loader:
+            time0 = time()
+            detections_for_frames = detector(preprocessed_frames.to(device))
+            average_times.append(time() - time0)
+            for detections_for_frame in detections_for_frames:
+                if len(detections_for_frame): detections.append(detections_for_frame)
+                else: detections.append(np.array([]))
+    print(f'Frame-wise inference time: {batch_size/np.mean(average_times)} fps')
+    return detections
+
+def generate_video_with_annotations(video, output_detected, output_filename, skip_frames, logger):
+    fps = 24
+    logger.info("---intepreting json")
+    results = defaultdict(list)
+    for trash in output_detected["detected_trash"]:
+        for k, v in trash["frame_to_box"].items():
+            frame_nb = int(k) - 1
+            object_nb = trash["id"] + 1
+            center_x = v[0]
+            center_y = v[1]
+            results[frame_nb * (skip_frames+1)].append((object_nb, center_x, center_y))
+            # append next skip_frames
+            if str(frame_nb + 2) in trash["frame_to_box"]:
+                next_trash = trash["frame_to_box"][str(frame_nb + 2)]
+                next_x = next_trash[0]
+                next_y = next_trash[1]
+                for i in range(1, skip_frames+1):
+                    new_x = center_x + (next_x - center_x) * i/(skip_frames+1)
+                    new_y = center_y + (next_y - center_y) * i/(skip_frames+1)
+                    results[frame_nb * (skip_frames+1) + i].append((object_nb, new_x, new_y))
+    logger.info("---writing video")
+
+    #fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+    # writer = cv2.VideoWriter(filename=output_filename,
+                                    #apiPreference=cv2.CAP_FFMPEG,
+    #                                fourcc=fourcc,
+    #                                fps=fps,
+    #                                frameSize=video.shape)
+
+    font = cv2.FONT_HERSHEY_SIMPLEX
+    ret, frame, frame_nb = video.read()
+    frames = []
+    while ret:
+        detections_for_frame = results[frame_nb]
+        for detection in detections_for_frame:
+            cv2.putText(frame, '{}'.format(detection[0]), (int(detection[1]), int(detection[2])+5), font, 2, (0, 0, 255), 3, cv2.LINE_AA)
+
+        #writer.write(frame)
+        frame = downscale_local_mean(frame, (4,4,1))
+        frames.append(frame[:,:,::-1])
+
+        ret, frame, frame_nb = video.read()
+        if frame_nb > 24 * 10:
+            break
+    #writer.release()
+
+    # Test with another library
+
+    clip = ImageSequenceClip(sequence=frames, fps=fps)
+    clip.write_videofile(output_filename, fps=fps)
+    del frames
+
+
+    logger.info("---finished writing video")
+
+
+def resize_external_detections(detections, ratio):
+
+    for detection_nb in range(len(detections)):
+        detection = detections[detection_nb]
+        if len(detection):
+            detection = np.array(detection)[:,:-1]
+            detection[:,0] = (detection[:,0] + detection[:,2])/2
+            detection[:,1] = (detection[:,1] + detection[:,3])/2
+            detections[detection_nb] = detection[:,:2]/ratio
+    return detections
+
+
+def write_tracking_results_to_file(results, ratio_x, ratio_y, output_filename):
+    """ writes the output result of a tracking the following format:
+    - frame
+    - id
+    - x_tl, y_tl, w=0, h=0
+    - 4x unused=-1
+    """
+    with open(output_filename, 'w') as output_file:
+        for result in results:
+            output_file.write('{},{},{},{},{},{},{},{},{},{}\n'.format(result[0]+1,
+                                                                result[1]+1,
+                                                                ratio_x * result[2],
+                                                                ratio_y * result[3],
+                                                                0,
+                                                                0,
+                                                                -1,-1,-1,-1))
+
+
+def read_tracking_results(input_file):
+    """ read the input filename and interpret it as tracklets
+    i.e. lists of lists
+    """
+    raw_results = np.loadtxt(input_file, delimiter=',')
+    if raw_results.ndim == 1: raw_results = np.expand_dims(raw_results,axis=0)
+    tracklets = defaultdict(list)
+    for result in raw_results:
+        frame_id = int(result[0])
+        track_id = int(result[1])
+        left, top, width, height = result[2:6]
+        center_x = left + width/2
+        center_y = top + height/2
+        tracklets[track_id].append((frame_id, center_x, center_y))
+
+    tracklets = list(tracklets.values())
+    return tracklets
+
+def gather_tracklets(tracklist):
+    """ Converts a list of flat tracklets into a list of lists
+    """
+    tracklets = defaultdict(list)
+    for track in tracklist:
+        frame_id = track[0]
+        track_id = track[1]
+        center_x = track[2]
+        center_y = track[3]
+        tracklets[track_id].append((frame_id, center_x, center_y))
+
+    tracklets = list(tracklets.values())
+    return tracklets
+
+class FramesWithInfo:
+    def __init__(self, frames, output_shape=None):
+        self.frames = frames
+        if output_shape is None:
+            self.output_shape = frames[0].shape[:-1][::-1]
+        else: self.output_shape = output_shape
+        self.end = len(frames)
+        self.read_head = 0
+
+    def __next__(self):
+        if self.read_head < self.end:
+            frame = self.frames[self.read_head]
+            self.read_head+=1
+            return frame
+
+        else:
+            raise StopIteration
+
+    def __iter__(self):
+        return self