Fix memory leak

fastmot/feature_extractor.py

@@ -21,6 +21,9 @@ def __init__(self, config):
         self.embeddings = []
         self.num_features = 0
 
+    def __del__(self):
+        self.pool.close()
+
     def __call__(self, frame, detections):
         self.extract_async(frame, detections)
         return self.postprocess()

fastmot/mot.py

@@ -1,5 +1,6 @@
 from enum import Enum
 import logging
+import pycuda.driver as cuda
 import cv2
 
 from .detector import SSDDetector, YOLODetector, PublicDetector
@@ -44,6 +45,9 @@ def __init__(self, size, config, draw=False, verbose=False):
         self.detector_type = DetectorType[config['detector_type']]
         self.detector_frame_skip = config['detector_frame_skip']
 
+        # create CUDA context for tensorRT
+        self.cfx = cuda.Device(0).make_context()
+
         LOGGER.info('Loading detector model...')
         if self.detector_type == DetectorType.SSD:
             self.detector = SSDDetector(self.size, config['ssd_detector'])
@@ -58,6 +62,9 @@ def __init__(self, size, config, draw=False, verbose=False):
         self.tracker = MultiTracker(self.size, self.extractor.metric, config['multi_tracker'])
         self.frame_count = 0
 
+    def __del__(self):
+        self.cfx.pop()
+
     @property
     def visible_tracks(self):
         # retrieve confirmed and active tracks from the tracker
@@ -83,6 +90,7 @@ def step(self, frame):
         frame : ndarray
             The next frame.
         """
+        self.cfx.push()
         detections = []
         if self.frame_count == 0:
             detections = self.detector(frame)
@@ -112,6 +120,7 @@ def step(self, frame):
         if self.draw:
             self._draw(frame, detections)
         self.frame_count += 1
+        self.cfx.pop()
 
     @staticmethod
     def print_timing_info():