Added evaluation to tracking 01 notebook

Author: karineip

scenarios/tracking/01_training_introduction.ipynb

@@ -1,11 +1,14 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # Licensed under the MIT License.
 
+from collections import OrderedDict
+import numpy as np
 import os
 import os.path as osp
-from typing import Dict
+from typing import Dict, List
 from torch.utils.data import DataLoader
 from torchvision.transforms import transforms as T
+from .bbox import TrackingBbox
 from .references.fairmot.datasets.dataset.jde import JointDataset
 from .opts import opts
 from ..common.gpu import db_num_workers
@@ -54,3 +57,31 @@ def _init_dataloaders(self) -> None:
             pin_memory=True,
             drop_last=True,
         )
+        
+def boxes_to_mot(results: Dict[int, List[TrackingBbox]]) -> None:
+    """
+    Save the predicted tracks to csv file in MOT challenge format ["frame", "id", "left", "top", "width", "height",]
+    
+    Args:
+        results: dictionary mapping frame id to a list of predicted TrackingBboxes
+        txt_path: path to which results are saved in csv file
+    
+    """   
+    # convert results to dataframe in MOT challenge format
+    preds = OrderedDict(sorted(results.items()))
+    bboxes = [
+        [
+            bb.frame_id,
+            bb.track_id,
+            bb.top,
+            bb.left,
+            bb.bottom - bb.top,
+            bb.right - bb.left,
+            1, -1, -1, -1,
+        ]
+        for _, v in preds.items()
+        for bb in v
+    ]
+    bboxes_formatted = np.array(bboxes)
+    
+    return bboxes_formatted

utils_cv/tracking/dataset.py

@@ -2,22 +2,26 @@
 # Licensed under the MIT License.
 
 import argparse
-from collections import OrderedDict
+from collections import OrderedDict, defaultdict
 from copy import deepcopy
 import glob
 import requests
 import os
 import os.path as osp
+import tempfile #KIP
 from typing import Dict, List, Tuple
 
+
 import torch
 import torch.cuda as cuda
 import torch.nn as nn
 from torch.utils.data import DataLoader
 
 import cv2
+import numpy as np
 import pandas as pd
 import matplotlib.pyplot as plt
+import motmetrics as mm
 
 from .references.fairmot.datasets.dataset.jde import LoadImages, LoadVideo
 from .references.fairmot.models.model import (
@@ -26,10 +30,11 @@
     save_model,
 )
 from .references.fairmot.tracker.multitracker import JDETracker
+from .references.fairmot.tracking_utils.evaluation import Evaluator
 from .references.fairmot.trains.train_factory import train_factory
 
 from .bbox import TrackingBbox
-from .dataset import TrackingDataset
+from .dataset import TrackingDataset, boxes_to_mot
 from .opts import opts
 from .plot import draw_boxes, assign_colors
 from ..common.gpu import torch_device
@@ -211,7 +216,10 @@ def fit(
         Trainer = train_factory[opt_fit.task]
         trainer = Trainer(opt_fit.opt, self.model, self.optimizer)
         trainer.set_device(opt_fit.gpus, opt_fit.chunk_sizes, opt_fit.device)
-
+        
+        # initialize loss vars
+        self.losses_dict = defaultdict(list)
+        
         # training loop
         for epoch in range(
             start_epoch + 1, start_epoch + opt_fit.num_epochs + 1
@@ -229,10 +237,41 @@ def fit(
                 lr = opt_fit.lr * (0.1 ** (opt_fit.lr_step.index(epoch) + 1))
                 for param_group in optimizer.param_groups:
                     param_group["lr"] = lr
+                    
+            # store losses in each epoch                   
+            for k, v in log_dict_train.items():
+                if k in ['loss', 'hm_loss', 'wh_loss', 'off_loss', 'id_loss']:
+                    self.losses_dict[k].append(v)
 
         # save after training because at inference-time FairMOT src reads model weights from disk
         self.save(self.model_path)
 
+    def plot_training_losses(self, figsize: Tuple[int, int] = (10, 5))->None: 
+        '''
+        Plots training loss from calling `fit`  
+        
+        Args:
+            figsize (optional): width and height wanted for figure of training-loss plot
+        
+        '''
+        fig = plt.figure(figsize=figsize)
+        ax1 = fig.add_subplot(1, 1, 1)
+        
+        ax1.set_xlim([0, len(self.losses_dict['loss']) - 1])
+        ax1.set_xticks(range(0, len(self.losses_dict['loss'])))        
+        ax1.set_xlabel("epochs")
+        ax1.set_ylabel("losses")
+        
+        ax1.plot(self.losses_dict['loss'], c="r", label='loss')
+        ax1.plot(self.losses_dict['hm_loss'], c="y", label='hm_loss')
+        ax1.plot(self.losses_dict['wh_loss'], c="g", label='wh_loss')
+        ax1.plot(self.losses_dict['off_loss'], c="b", label='off_loss')
+        ax1.plot(self.losses_dict['id_loss'], c="m", label='id_loss')
+
+        plt.legend(loc='upper right')
+        fig.suptitle("Training losses over epochs")
+        
+    
     def save(self, path) -> None:
         """
         Save the model to a specified path.
@@ -243,9 +282,46 @@ def save(self, path) -> None:
         save_model(path, self.epoch, self.model, self.optimizer)
         print(f"Model saved to {path}")
 
-    def evaluate(self, results, gt) -> pd.DataFrame:
-        pass
-
+    def evaluate(self,
+                 results: Dict[int, List[TrackingBbox]],
+                 gt_root_path: str) -> str:
+        
+        """ eval code that calls on 'motmetrics' package in referenced FairMOT script, to produce MOT metrics on inference, given ground-truth.
+        Args:
+            results: prediction results from predict() function, i.e. Dict[int, List[TrackingBbox]] 
+            gt_root_path: path of dataset containing GT annotations in MOTchallenge format (xywh)
+        Returns:
+            strsummary: str output by method in 'motmetrics' package, containing metrics scores        
+        """
+       
+        #Implementation inspired from code found here: https://github.com/ifzhang/FairMOT/blob/master/src/track.py
+        evaluator = Evaluator(gt_root_path, "single_vid", "mot")
+        
+        with tempfile.TemporaryDirectory() as tmpdir1:
+            os.makedirs(osp.join(tmpdir1,'results'))
+            result_filename = osp.join(tmpdir1,'results', 'results.txt')
+          
+            # Save results im MOT format for evaluation            
+            bboxes_mot = boxes_to_mot(results)            
+            np.savetxt(result_filename, bboxes_mot, delimiter=",", fmt="%s")
+
+            # Run evaluation using pymotmetrics package
+            accs=[evaluator.eval_file(result_filename)]
+                           
+        # get summary
+        metrics = mm.metrics.motchallenge_metrics
+        mh = mm.metrics.create()
+       
+        summary = Evaluator.get_summary(accs, ("single_vid",), metrics)
+        strsummary = mm.io.render_summary(
+            summary,
+            formatters=mh.formatters,
+            namemap=mm.io.motchallenge_metric_names
+        )
+        print(strsummary)        
+        
+        return strsummary
+    
     def predict(
         self,
         im_or_video_path: str,

utils_cv/tracking/model.py

@@ -4,7 +4,7 @@
 import motmetrics as mm
 mm.lap.default_solver = 'lap'
 
-from tracking_utils.io import read_results, unzip_objs
+from .io import read_results, unzip_objs #EDITED
 
 
 class Evaluator(object):
@@ -76,10 +76,10 @@ def eval_frame(self, frame_id, trk_tlwhs, trk_ids, rtn_events=False):
         return events
 
     def eval_file(self, filename):
-        self.reset_accumulator()
-
+        self.reset_accumulator()        
         result_frame_dict = read_results(filename, self.data_type, is_gt=False)
         frames = sorted(list(set(self.gt_frame_dict.keys()) | set(result_frame_dict.keys())))
+        
         for frame_id in frames:
             trk_objs = result_frame_dict.get(frame_id, [])
             trk_tlwhs, trk_ids = unzip_objs(trk_objs)[:2]

utils_cv/tracking/plot.py

@@ -2,7 +2,7 @@
 from typing import Dict
 import numpy as np
 
-from tracking_utils.log import logger
+from .log import logger #EDITED
 
 
 def write_results(filename, results_dict: Dict, data_type: str):
@@ -61,16 +61,18 @@ def read_results(filename, data_type: str, is_gt=False, is_ignore=False):
 
 
 def read_mot_results(filename, is_gt, is_ignore):
+   
     valid_labels = {1}
     ignore_labels = {2, 7, 8, 12}
     results_dict = dict()
-    if os.path.isfile(filename):
-        with open(filename, 'r') as f:
-            for line in f.readlines():
-                linelist = line.split(',')
-                if len(linelist) < 7:
+    if os.path.isfile(filename):        
+        with open(filename, 'r') as f:            
+            for line in f.readlines():                
+                linelist = line.split(',')                
+                if len(linelist) < 7:                   
                     continue
                 fid = int(linelist[0])
+                
                 if fid < 1:
                     continue
                 results_dict.setdefault(fid, list())
@@ -82,6 +84,7 @@ def read_mot_results(filename, is_gt, is_ignore):
                         if mark == 0 or label not in valid_labels:
                             continue
                     score = 1
+                    
                 elif is_ignore:
                     if 'MOT16-' in filename or 'MOT17-' in filename:
                         label = int(float(linelist[7]))
@@ -91,14 +94,15 @@ def read_mot_results(filename, is_gt, is_ignore):
                     else:
                         continue
                     score = 1
+                    
                 else:
                     score = float(linelist[6])
-
+                    
+                
                 tlwh = tuple(map(float, linelist[2:6]))
-                target_id = int(linelist[1])
-
+                target_id = int(linelist[1])                
                 results_dict[fid].append((tlwh, target_id, score))
-
+                
     return results_dict