wip tracking uniformization with main repo

tracking/postprocess_and_count_tracks.py

@@ -3,6 +3,7 @@ import argparse
 from scipy.signal import convolve
 from tracking.utils import write_tracking_results_to_file, read_tracking_results
 from collections import defaultdict
+from statistics import mode
 
 import json
 
@@ -21,14 +22,20 @@ def filter_tracks(tracklets, kappa, tau):
         for det in dets:
             results.append((det[0], tracker_nb, det[1], det[2], det[3], det[4]))
 
-    # for tracker_nb, associated_detections in enumerate(tracks):
-    #     for det 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 process_class_and_confidences(class_confs):
+    ''' Finds the majority and most confident class from list [(classid, conf), ...]
+    '''
+    d = defaultdict(lambda: (0, 0.0))
+    for (cls, conf) in class_confs:
+        d[cls] = (d[cls][0] + 1, d[cls][1] + conf)
+    best_class = sorted(d.items(), key=lambda v: v[1][0]+v[1][1])[-1]
+    return best_class[0], round(best_class[1][1]/best_class[1][0],2)
+
+
 def postprocess_for_api(results, class_dict=defaultdict(lambda: "fragment")):
     """ Converts tracking results into json object for API
     """
@@ -47,11 +54,18 @@ def postprocess_for_api(results, class_dict=defaultdict(lambda: "fragment")):
             result_list.append({"label":classname,
                                 "id": id,
                                 "frame_to_box": {str(frame_number): box},
-                                "frame_to_conf": {str(frame_number): conf}})
+                                "frame_to_class_conf": {str(frame_number): (res[5], conf)}})
         # otherwise, retrieve the jsonline and append the box
         else:
             result_list[id_list[id]]["frame_to_box"][str(frame_number)] = box
-            result_list[id_list[id]]["frame_to_conf"][str(frame_number)] = conf
+            result_list[id_list[id]]["frame_to_class_conf"][str(frame_number)] = (res[5], conf)
+
+    # Finally, collapse the confidence and class
+    for res in result_list:
+        classid, avg_conf = process_class_and_confidences(res.pop("frame_to_class_conf").values())
+        res["avg_conf"] = avg_conf
+        # update the label
+        res["label"] = class_dict[classid]
 
     return {"detected_trash": result_list}
 
@@ -63,10 +77,11 @@ def count_objects(input_json, class_dict):
         results[trash["label"]] += 1
         total += 1
 
-    if total==0: 
+    if total==0:
         total = 1
     return {k+f": {str(v)}":v/total for k,v in results.items()}
 
+
 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, ...)

tracking/track_video.py

@@ -2,56 +2,22 @@ import cv2
 import numpy as np
 import os
 from detection.detect import detect
-from tracking.utils import in_frame, init_trackers
+from tracking.utils import in_frame
 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 init_trackers(engine, detections, confs, labels, frame_nb, state_variance, observation_variance, delta):
+    """ Initializes the trackers based on detections
+    """
+    trackers = []
+    for detection, conf, label in zip(detections, confs, labels):
+        tracker_for_detection = engine(frame_nb, detection, conf, label, state_variance, observation_variance, delta)
+        trackers.append(tracker_for_detection)
+    return trackers
 
 
 def build_confidence_function_for_trackers(trackers, flow01):
@@ -63,19 +29,23 @@ def build_confidence_function_for_trackers(trackers, flow01):
             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):
+
+def associate_detections_to_trackers(detections_for_frame, confs, labels, 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 detection_nb, (detection, conf, label) in enumerate(zip(detections_for_frame, confs, labels)):
             for tracker_id, confidence_function in enumerate(confidence_functions):
                 score = confidence_function(detection)
+                cls_score = trackers[tracker_id].cls_score_function(conf, label)
+                if cls_score < 0.5:
+                    score = score * 0.1 # if wrong class, reduce the score, to tweak
                 if score > confidence_threshold:
-                    cost_matrix[detection_nb,tracker_id] = score
+                    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)
+                    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]
 
@@ -83,8 +53,7 @@ def associate_detections_to_trackers(detections_for_frame, trackers, flow01, con
 
 
 def interpret_detection(detections_for_frame, downsampling_factor, is_yolo=False):
-    """
-    normalizes the detections depending whether they come from centernet or yolo
+    """ normalizes the detections depending whether they come from centernet or yolo
     """
     if not is_yolo:
         confs = [0.0]*len(detections_for_frame)
@@ -112,7 +81,6 @@ def track_video(reader, detections, args, engine, transition_variance, observati
     delta = 0.005*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)
 
@@ -132,15 +100,12 @@ def track_video(reader, detections, args, engine, transition_variance, observati
             if len(detections_for_frame):
                 trackers = init_trackers(engine, detections_for_frame, confs, labels, 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,
+                assigned_trackers = associate_detections_to_trackers(detections_for_frame, confs, labels, trackers,
                                                                      flow01, args.confidence_threshold)
 
                 for detection, conf, label, assigned_tracker in zip(detections_for_frame, confs, labels, assigned_trackers):

tracking/trackers.py

@@ -8,7 +8,6 @@ import matplotlib.patches as mpatches
 class Tracker:
 
     def __init__(self, frame_nb, X0, confidence, class_id, transition_variance, observation_variance, delta):
-
         self.transition_covariance = np.diag(transition_variance)
         self.observation_covariance = np.diag(observation_variance)
         self.updated = False
@@ -49,12 +48,20 @@ class Tracker:
 
         return lambda coord: confidence_from_multivariate_distribution(coord, distribution)
 
+    def cls_score_function(self, conf, label):
+        """ generates a score based on classes associated with observation in this tracker
+        """
+        class_conf = sum([tr[2] for tr in self.tracklet if tr[3]==label])
+        other_conf = sum([tr[2] for tr in self.tracklet])
+        return (class_conf+conf) / (other_conf+conf)
+
     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)

tracking/utils.py

@@ -10,8 +10,10 @@ from detection.transforms import TransformFrames
 from collections import defaultdict
 from skimage.transform import downscale_local_mean
 from skvideo.io import FFmpegWriter
+import matplotlib.pyplot as plt
 from tracking.gps import get_media_duration
 
+
 class GaussianMixture(object):
     def __init__(self, means, covariance, weights):
         self.components = [multivariate_normal(
@@ -33,22 +35,13 @@ class GaussianMixture(object):
             result += weight*component.cdf(x)
         return result
 
-def init_trackers(engine, detections, confs, labels, frame_nb, state_variance, observation_variance, delta):
-    trackers = []
-
-    for detection, conf, label in zip(detections, confs, labels):
-        tracker_for_detection = engine(frame_nb, detection, conf, label, 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):
-
 
+def in_frame(position, shape, border=0.02):
     shape_x = shape[1]
     shape_y = shape[0]
     x = position[0]
@@ -56,6 +49,7 @@ def in_frame(position, shape, border=0.02):
 
     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']]
@@ -65,8 +59,8 @@ def gather_filenames_for_video_in_annotations(video, images, data_dir):
     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):
 
+def get_detections_for_video(reader, detector, batch_size=16, device=None):
     detections = []
     dataset = TorchIterableFromReader(reader, TransformFrames())
     loader = DataLoader(dataset, batch_size=batch_size)
@@ -244,6 +238,7 @@ def gather_tracklets(tracklist):
     tracklets = list(tracklets.values())
     return tracklets
 
+
 class FramesWithInfo:
     def __init__(self, frames, output_shape=None):
         self.frames = frames
@@ -264,3 +259,47 @@ class FramesWithInfo:
 
     def __iter__(self):
         return self
+
+class Display:
+    """ Display tracking
+    """
+    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)