dnt 0.3.1.3__py3-none-any.whl → 0.3.1.7__py3-none-any.whl

dnt/filter/filter.py

@@ -8,7 +8,8 @@ class Filter:
         pass
     
     @staticmethod
-    def filter_iou(detections: pd.DataFrame, zones: geometry.multipolygon = None, class_list: list[int] = None, score_threshold: float = 0):
+    def filter_iou(detections: pd.DataFrame, zones: geometry.multipolygon = None, 
+                   class_list: list[int] = None, score_threshold: float = 0)->pd.DataFrame:
 
         detections = detections.loc[detections[6]>=score_threshold].copy()
 
@@ -232,7 +233,7 @@ class Filter:
     
     @staticmethod
     def filter_tracks_by_zones(tracks:pd.DataFrame, 
-                    zones: geometry.MultiPolygon = None,
+                    zones: list[Polygon] = None,
                     method: str = 'list', 
                     ref_point: str = 'bc',
                     offset: tuple = (0, 0),
@@ -244,10 +245,19 @@ class Filter:
         Inputs:
             tracks - tracks 
             zones - zones (polygon)
-            method - 'list' (default), 'filter', 'label'
+            method - 'list' (default) - List track ids within zones
+                     'filter' - filter tracks within zones
+                     'label' - label tracks with zone index
             ref_point - the reference point of a track bbox, 
-                        br - buttom_right, center_point, 
-                        left_up, right_up, left_buttom, right_buttom
+                        br - buttom_right, 
+                        bl - bottom_left
+                        bc - bottom_center
+                        cc - center_point, 
+                        cl - left_center,
+                        cr - right_center,
+                        tc - top_center,
+                        tl - top_left,
+                        tr - top_right,
             offset - the offset to ref_point, default is (0, 0)
             aggregate - combine outputs to one dataframe, add zone column
             zone_name - if aggregate, the field name of zone variable, default is 'zone' 
@@ -291,7 +301,7 @@ class Filter:
             pbar.set_description_str("Filtering zones {} of {}".format(video_index, video_tot))
         else:
             pbar.set_description_str("Filtering zones ")
-        
+
         for zone in zones:
             matched  = geo_tracks[geo_tracks.geometry.within(zone)]
             if len(matched)>0:
@@ -305,7 +315,7 @@ class Filter:
             for i in range(len(matched_ids)):
                 tracks.loc[tracks['track'].isin(matched_ids[i]), zone_name] = i
             if method == 'filter':
-                results = tracks[tracks['zone']!=-1].copy()
+                results = tracks[tracks[zone_name]!=-1].copy()
             else:
                 results = tracks
         else:
@@ -375,7 +385,8 @@ class Filter:
                     tolerance: int = 0,
                     bbox_size: int = 0,
                     force_line_indexes: list[int] = None,
-                    video_index:int = None, video_tot:int = None) -> pd.DataFrame:
+                    video_index:int = None,
+                    video_tot:int = None) -> pd.DataFrame:
         '''
         Filter tracks by lines
         Inputs:

dnt/label/labeler2.py

@@ -1,5 +1,6 @@
 import os, sys
 sys.path.append(os.path.dirname(__file__))
+sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 
 import cv2
 import numpy as np
@@ -7,37 +8,74 @@ import matplotlib.pyplot as plt
 from tqdm import tqdm
 import pandas as pd
 import itertools
-from ..shared.util import load_classes
+from shared.util import load_classes
 import random
+import subprocess
 
 class Labeler:
-    def __init__(self, method:str='opencv', compress_message:bool=False, nodraw_empty:bool=True):
-        
+    def __init__(self,
+                 method:str='opencv',
+                 encoder:str='libx264',
+                 preset:str='medium',
+                 crf:int=23,
+                 pix_fmt:str='bgr24',
+                 compress_message:bool=False,
+                 nodraw_empty:bool=True):
+        '''
+        Parameters:\n
+            - method: 'opencv' (default) - use opencv to draw labels),
+                      'ffmpeg' - use ffmpeg to draw labels
+            - encoder: 'libx264' (default) - use libx264 encoder for ffmpeg,
+                       'libx265' - use libx265 encoder for ffmpeg
+                       'h264_nvenc' - use h264_nvenc encoder for ffmpeg
+                       'hevc_nvenc' - use hevc_nvenc encoder for ffmpeg
+            - preset: 'medium' (default) - use medium preset for ffmpeg
+                        'slow' - use slow preset for ffmpeg
+                        'fast' - use fast preset for ffmpeg
+            - crf: 23 (default) - use 23 crf for ffmpeg, lower is better quality
+            - compress_message: False (default) - show compress message in progress bar
+            - nodraw_empty: True (default) - not draw empty frames
+        '''
         self.method = method
+        self.encoder = encoder
+        self.preset = preset
+        self.crf = crf
+        self.pix_fmt = pix_fmt
         self.compress_message=compress_message
         self.nodraw_empty = nodraw_empty
     
-    def draw(self, input_video:str, output_video:str, 
-            draws:pd.DataFrame = None, draw_file:str = None,
-            start_frame:int=None, end_frame:int=None, 
-            video_index:int=None, video_tot:int=None, verbose:bool=True):
+    def draw(self,
+             input_video:str,
+             output_video:str,
+             draws:pd.DataFrame=None,
+             draw_file:str=None,
+             start_frame:int=None,
+             end_frame:int=None,
+             video_index:int=None,
+             video_tot:int=None,
+             verbose:bool=True):
         '''
-        General labeling function
-        Inputs:
+        General labeling function\n
+        Inputs:\n
                 draws: a DataFrame contains labeling information, if None, read label_file
                 label_file: a txt file with a header ['frame','type','coords','color','size','thick','desc']
                 input_video: raw video
                 output_video: labeled video
-                start_frame:
-                end_frame:
+                start_frame: starting frame
+                end_frame: ending frame
                 video_index: display video index in batch processing
                 video_tot: display total video number in batch processing
         '''
         if draws is not None:
             data = draws
         else:
-            data = pd.read_csv(draw_file, dtype={'frame':int, 'type':str, 'size':float, 'desc':str, 'thick':int}, 
-                    converters={'coords': lambda x:list(eval(x)), 'color': lambda x:eval(x)})
+            data = pd.read_csv(draw_file,
+                               dtype={'frame':int,
+                                      'type':str,
+                                      'size':float,
+                                      'desc':str,
+                                      'thick':int},
+                               converters={'coords': lambda x:list(eval(x)), 'color': lambda x:eval(x)})
 
         cmap = plt.get_cmap('tab20b')
         colors = [cmap(i)[:3] for i in np.linspace(0, 1, 20)]
@@ -55,8 +93,31 @@ class Labeler:
         fps = int(cap.get(cv2.CAP_PROP_FPS))
         width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
         height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-        writer = cv2.VideoWriter(output_video, fourcc, fps, (width, height))
+
+        if self.method == 'ffmpeg':
+            # FFmpeg command to write H.265 encoded video
+            ffmpeg_cmd = [
+                "ffmpeg",
+                "-y",  # Overwrite output file if it exists
+                "-f", "rawvideo",
+                "-vcodec", "rawvideo",
+                "-pix_fmt", self.pix_fmt,
+                "-s", f"{width}x{height}",
+                "-r", str(fps),
+                "-i", "-",  # Read input from stdin
+                "-c:v", self.encoder,  # H.265 codec
+                "-preset", self.preset,  # Adjust preset as needed (ultrafast, fast, medium, slow, etc.)
+                "-crf", str(self.crf),  # Constant Rate Factor (higher = more compression, lower = better quality)
+                output_video]
+
+            # Start FFmpeg process
+            process = subprocess.Popen(ffmpeg_cmd,
+                               stdin=subprocess.PIPE,
+                               stdout=subprocess.DEVNULL,
+                               stderr=subprocess.DEVNULL)
+        else:
+            fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+            writer = cv2.VideoWriter(output_video, fourcc, fps, (width, height))
 
         if verbose:
             pbar = tqdm(total=tot_frames, unit=" frames")
@@ -128,11 +189,11 @@ class Labeler:
                     color = element['color']
                     thick = element['thick']
                     cv2.polylines(frame, [np.array(coords)], isClosed=False, color=color, thickness=thick)
-
-            writer.write(frame)
-            #key = cv2.waitKey(1) & 0xFF
-            #if key == ord("q"):
-            #    break
+                
+            if self.method == 'ffmpeg':
+                process.stdin.write(frame.tobytes())
+            else:
+                writer.write(frame)
             
             if verbose:
                 pbar.update()
@@ -141,17 +202,31 @@ class Labeler:
             pbar.close()
         #cv2.destroyAllWindows()
         cap.release()
-        writer.release()      
-
-    def draw_track_clips(self, input_video:str, output_path:str,
-                    tracks:pd.DataFrame = None, track_file:str = None,
-                    method:str='all', random_number:int=10, track_ids:list=None,
-                    start_frame_offset:int=0, end_frame_offset:int=0,
-                    tail:int=0, prefix:bool=False,
-                    size:int=1, thick:int=1,
-                    verbose:bool=True):
+        if self.method == 'ffmpeg':
+            process.stdin.close()
+            process.wait()
+        else:
+            writer.release()
+
+    def draw_track_clips(self,
+                         input_video:str,
+                         output_path:str,
+                         tracks:pd.DataFrame=None,
+                         track_file:str = None,
+                         method:str='all',
+                         random_number:int=10,
+                         track_ids:list=None,
+                         start_frame_offset:int=0,
+                         end_frame_offset:int=0,
+                         tail:int=0,
+                         prefix:bool=False,
+                         size:int=1,
+                         thick:int=1,
+                         video_index:int=None,
+                         video_tot:int=None,
+                         verbose:bool=True):
         '''
-        Parameters:
+        Parameters:\n
             input_video: the raw video file
             outputh_path: the folder for outputing track clips
             tracks: the dataframe of tracks
@@ -184,7 +259,7 @@ class Labeler:
         else:
             track_ids = tracks['track'].unique().tolist()
 
-        pbar = tqdm(total=len(track_ids), desc='Labeling tracks ', unit='videos')
+        #pbar = tqdm(total=len(track_ids), desc='Labeling tracks ', unit='clips')
         for id in track_ids:
             selected_tracks = tracks[tracks['track']==id].copy()
             start_frame = max(selected_tracks['frame'].min() - start_frame_offset, 0)
@@ -195,20 +270,42 @@ class Labeler:
             else:
                 out_video = os.path.join(output_path, str(id)+'.mp4')
 
-            self.draw_tracks(input_video=input_video, output_video=out_video, tracks=selected_tracks, 
-                             start_frame=start_frame, end_frame=end_frame, verbose=False, tail=tail, thick=thick, size=size)
-            pbar.update()
-        pbar.close()
-
-    def draw_tracks(self, input_video:str, output_video:str,
-                    tracks:pd.DataFrame = None, track_file:str = None, label_file:str=None, 
-                    color = None, tail:int=0, thick:int=2, size:int=1,
-                    class_name = False, 
-                    start_frame:int=None, end_frame:int=None, 
-                    video_index:int=None, video_tot:int=None, verbose:bool=True):
+            self.draw_tracks(input_video=input_video,
+                             output_video=out_video,
+                             tracks=selected_tracks, 
+                             start_frame=start_frame,
+                             end_frame=end_frame, 
+                             verbose=verbose, 
+                             tail=tail, 
+                             thick=thick,
+                             size=size,
+                             video_index=video_index,
+                             video_tot=video_tot)
+            
+            #pbar.update()
+        #pbar.close()
+
+    def draw_tracks(self,
+                    input_video:str,
+                    output_video:str,
+                    tracks:pd.DataFrame=None,
+                    track_file:str=None,
+                    label_file:str=None,
+                    color=None,
+                    tail:int=0,
+                    thick:int=2,
+                    size:int=1,
+                    class_name=False,
+                    start_frame:int=None,
+                    end_frame:int=None,
+                    video_index:int=None,
+                    video_tot:int=None,
+                    verbose:bool=True):
 
         if tracks is None:
-            tracks = pd.read_csv(track_file, header=None, dtype={0:int, 1:int, 2:int, 3:int, 4:int, 5:int, 6:float, 7:int, 8:int, 9:int})
+            tracks = pd.read_csv(track_file,
+                                 header=None,
+                                 dtype={0:int, 1:int, 2:int, 3:int, 4:int, 5:int, 6:float, 7:int, 8:int, 9:int})
         tracks.columns = ['frame', 'track', 'x', 'y', 'w', 'h', 'score', 'cls', 'r3', 'r4']
 
         cmap = plt.get_cmap('tab20b')
@@ -223,22 +320,20 @@ class Labeler:
         if end_frame is None:
             end_frame = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))-1
 
-        tot_frames = end_frame - start_frame + 1       
+        selected_tracks = tracks.loc[(tracks['frame']>=start_frame) & (tracks['frame']<=end_frame)].copy()
 
-        if verbose:
-            pbar = tqdm(total=tot_frames, unit=" frames")
-            if self.compress_message:
-                    pbar.set_description_str("Generating labels")
+        pbar_desc = ""   
+        if self.compress_message:
+            pbar_desc = "Generating labels"
+        else:
+            if video_index and video_tot:
+                pbar_desc = "Generating labels {} of {}".format(video_index, video_tot)
             else:
-                if video_index and video_tot:
-                    pbar.set_description_str("Generating labels {} of {}".format(video_index, video_tot))
-                else:
-                    pbar.set_description_str("Generating labels {} ".format(input_video))
-
-        selected_tracks = tracks.loc[(tracks['frame']>=start_frame) & (tracks['frame']<=end_frame)].copy()
+                pbar_desc = "Generating labels {} ".format(input_video)
 
+        pbar = tqdm(total=len(selected_tracks), unit=" frames", desc=pbar_desc)
         results = []
-        for index, track in selected_tracks.iterrows():
+        for _, track in selected_tracks.iterrows():
 
             if color is None:
                 final_color = colors[int(track['track']) % len(colors)]
@@ -265,7 +360,7 @@ class Labeler:
 
             if verbose:
                 pbar.update()
-        
+
         if verbose:
             pbar.close()
 
@@ -308,8 +403,9 @@ class Labeler:
             end_frame = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))-1
 
         tot_frames = end_frame - start_frame + 1       
+        selected_dets = dets.loc[(dets[0]>=start_frame) & (dets[0]<=end_frame)].copy()
 
-        pbar = tqdm(total=len(dets), unit=" dets")
+        pbar = tqdm(total=len(selected_dets), unit=" dets")
         if self.compress_message:
                 pbar.set_description_str("Generating labels")
         else:
@@ -318,8 +414,6 @@ class Labeler:
             else:
                 pbar.set_description_str("Generating labels {} ".format(input_video))
 
-        selected_dets = dets.loc[(dets[0]>=start_frame) & (dets[0]<=end_frame)].copy()
-
         results = []
         for index, det in selected_dets.iterrows():
 
@@ -492,11 +586,20 @@ class LabelGenerator():
 
     
 if __name__=='__main__':
-    video_file = "/mnt/d/videos/hfst/Standard_SCU7WH_2022-09-16_0630.02.001.mp4"    
-    iou_file = "/mnt/d/videos/hfst/Standard_SCU7WH_2022-09-16_0630.02.001_iou.txt"
-    track_file = "/mnt/d/videos/hfst/tracks/Standard_SCU7WH_2022-09-16_0630.02.001_track.txt"
-    label_video = "/mnt/d/videos/hfst/labels/Standard_SCU7WH_2022-09-16_0630.02.001_track.mp4"
-    label_file = "/mnt/d/videos/hfst/tracks/Standard_SCU7WH_2022-09-16_0630.02.001_label.txt"
-
-    labeler = Labeler(video_file, zoom_factor=0.5, nodraw_empty=True, label_fields=[6])
-    labeler.draw(label_file, video_file, label_video)
+    
+    video_file = "/mnt/d/videos/sample/traffic.mp4"    
+    iou_file = "/mnt/d/videos/sample/dets/traffic_iou.txt"
+    track_file = "/mnt/d/videos/sample/tracks/traffic_track.txt"
+    label_video = "/mnt/d/videos/sample/labels/traffic_track-ffmpeg.mp4"
+    label_file = "/mnt/d/videos/sample/labels/traffic_track.txt"
+
+    labeler = Labeler(method='ffmpeg',
+                      encoder='hevc_nvenc',
+                      preset='medium',
+                      crf=23,
+                      compress_message=True)
+    labeler.draw_tracks(input_video=video_file,
+                        output_video=label_video,
+                        track_file=track_file)
+    
+    

dnt/shared/synhcro.py

@@ -66,7 +66,7 @@ class Synchronizer():
         milliseconds_per_frame = 1/fps * 1000 
 
         dif_frame = frame - ref_frame
-        return round(ref_time + dif_frame * milliseconds_per_frame)
+        return round(ref_time + float(dif_frame) * milliseconds_per_frame)
 
     @staticmethod
     def add_unix_time_to_frames(frames:pd.DataFrame, ref_frame:int, ref_time:int, ref_timezone:str='US/Eastern', 

dnt/track/dsort/deep_sort/deep_sort.py

@@ -34,13 +34,13 @@ class DeepSort(object):
         metric = NearestNeighborDistanceMetric("cosine", self.max_cosine_distance, self.nn_budget)
         self.tracker = Tracker(metric, max_iou_distance=self.max_iou_distance, max_age=self.max_age, n_init=self.n_init)
 
-    def update(self, bbox_xywh, confidences, ori_img):
+    def update(self, bbox_xywh, confidences, classes, ori_img):
         self.height, self.width = ori_img.shape[:2]
 
         # generate detections
         features = self._get_features(bbox_xywh, ori_img)   # extract features for bboxes
         bbox_tlwh = self._xywh_to_tlwh(bbox_xywh)           # convert bbox from xc_yc_w_h to left top width height/width
-        detections = [Detection(bbox_tlwh[i], conf, features[i]) 
+        detections = [Detection(bbox_tlwh[i], conf, classes[i], features[i]) 
                         for i, conf in enumerate(confidences) if conf>self.min_confidence] # ignore low confidence bboxes
 
         # run on non-maximum supression
@@ -61,7 +61,8 @@ class DeepSort(object):
             box = track.to_tlwh()
             x1,y1,x2,y2 = self._tlwh_to_xyxy(box)
             track_id = track.track_id
-            outputs.append(np.array([x1,y1,x2,y2,track_id], dtype=int))
+            track_cls = track.cls
+            outputs.append(np.array([x1,y1,x2,y2,track_id,track_cls], dtype=int))
         if len(outputs) > 0:
             outputs = np.stack(outputs,axis=0)
         return outputs

dnt/track/dsort/deep_sort/sort/detection.py

@@ -26,10 +26,11 @@ class Detection(object):
 
     """
 
-    def __init__(self, tlwh, confidence, feature):
+    def __init__(self, tlwh, confidence, cls, feature):
         self.tlwh = np.asarray(tlwh, dtype=float)
         self.confidence = float(confidence)
         self.feature = np.asarray(feature, dtype=float)
+        self.cls = cls
 
     def to_tlbr(self):
         """Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,

dnt/track/dsort/deep_sort/sort/track.py

@@ -64,7 +64,7 @@ class Track:
     """
 
     def __init__(self, mean, covariance, track_id, n_init, max_age,
-                 feature=None):
+                 feature=None, cls=None):
         self.mean = mean
         self.covariance = covariance
         self.track_id = track_id
@@ -73,6 +73,7 @@ class Track:
         self.time_since_update = 0
 
         self.state = TrackState.Tentative
+        self.cls = cls
         self.features = []
         if feature is not None:
             self.features.append(feature)

dnt/track/dsort/deep_sort/sort/tracker.py

@@ -134,5 +134,5 @@ class Tracker:
         mean, covariance = self.kf.initiate(detection.to_xyah())
         self.tracks.append(Track(
             mean, covariance, self._next_id, self.n_init, self.max_age,
-            detection.feature))
+            detection.feature, detection.cls))
         self._next_id += 1

dnt/track/dsort/dsort.py

@@ -4,15 +4,31 @@ import torch
 import numpy as np
 import pandas as pd
 from tqdm import tqdm
-
-sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
-
 from deep_sort import DeepSort
 from config import Config
 
-def track(video_file:str, det_file:str, out_file:str = None, device:str = 'auto', half:bool = False,
-            cfg:dict = Config.get_cfg_dsort('default'), video_index:int = None, total_videos:int = None):
-  
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+
+def track(video_file:str,
+          det_file:str,
+          out_file:str=None,
+          device:str='auto',
+          half:bool=False,
+          cfg:dict=Config.get_cfg_dsort('default'),
+          video_index:int=None,
+          total_videos:int=None):
+    '''
+    Track objects in a video using Deep SORT.
+    Args:
+        video_file (str): Path to the input video file.
+        det_file (str): Path to the detection results file.
+        out_file (str): Path to save the tracking results.
+        device (str): Device to run the model on ('cpu' or 'cuda').
+        half (bool): Whether to use half precision.
+        cfg (dict): Configuration dictionary for Deep SORT.
+        video_index (int): Index of the current video in a batch.
+        total_videos (int): Total number of videos in a batch.
+    '''
     #device = torch.device('cuda') if (torch.cuda.is_available() and gpu) else torch.device('cpu')
     cap = cv2.VideoCapture(video_file)
     if not cap.isOpened():
@@ -50,11 +66,18 @@ def track(video_file:str, det_file:str, out_file:str = None, device:str = 'auto'
             bbox_xywh = np.array(bbox_xywh)
             conf_score = np.array(frame_dets[:,6])
             classes = np.array(frame_dets[:,7])
-            outputs = deepsort.update(bbox_xywh, conf_score, im)
+            outputs = deepsort.update(bbox_xywh, conf_score, classes, im)
 
             if len(outputs) > 0:
                 for output in outputs:
-                    results.append([pos_frame, output[4], output[0], output[1], output[2]-output[0], output[3]-output[1], -1, -1, -1, -1])
+                    results.append([pos_frame,
+                                    output[4],
+                                    output[0],
+                                    output[1],
+                                    output[2]-output[0],
+                                    output[3]-output[1],
+                                    output[5],
+                                    -1, -1, -1])
 
         pbar.update()
 
@@ -65,15 +88,9 @@ def track(video_file:str, det_file:str, out_file:str = None, device:str = 'auto'
         df.to_csv(out_file, index=False, header=None)
 
 if __name__ == "__main__":
-
-    '''
-    video_file = "/mnt/d/videos/ped2stage/videos/gh1293.mp4"
-    iou_file = "/mnt/d/videos/ped2stage/dets/gh1293_iou_ped.txt"
-    out_file = "/mnt/d/videos/ped2stage/tracks/gh1293_ped_track_2.txt"
-    '''
-    video_file = "/mnt/d/videos/samples/ped_veh.mp4"
-    iou_file = "/mnt/d/videos/samples/dets/ped_veh_iou.txt"
-    out_file = "/mnt/d/videos/samples/tracks/ped_veh_track.txt"
+    video_file = "/mnt/d/videos/sample/traffic.mp4"
+    iou_file = "/mnt/d/videos/sample/dets/traffic_det.txt"
+    out_file = "/mnt/d/videos/sample/tracks/traffic_track.txt"
 
     track(video_file, iou_file, out_file)
     

dnt/track/re_class.py

@@ -5,8 +5,25 @@ from ..detect import Detector
 from ..engine.iob import iobs
 
 class ReClass:
-    def __init__(self, num_frames:int=25, threshold:float=0.75, model:str='rtdetr', device:str='auto',
-                 default_class:int=0, match_class:list=[1, 36]) -> None:
+    def __init__(self,
+                 num_frames:int=25,
+                 threshold:float=0.75,
+                 model:str='rtdetr',
+                 weights:str='x',
+                 device:str='auto',
+                 default_class:int=0,
+                 match_class:list=[1, 36]) -> None:
+        """
+        Re-classify tracks based on detection results
+        Parameters:
+            num_frames: Number of frames to consider for re-classification, default 25
+            threshold: Threshold for matching, default 0.75
+            model: Detection model to use, default 'rtdetr'
+            weights: Weights for the detection model, default 'x'
+            device: Device to use for detection, default 'auto'
+            default_class: Default class to assign if no match found, default 0 (pedestrian)
+            match_class: List of classes to match, default [1, 36] (bicycle, skateboard/scooter)
+        """
         self.detector = Detector(model=model, device=device)
         self.num_frames = num_frames
         self.threshold = threshold
@@ -36,8 +53,12 @@ class ReClass:
 
         return hit, avg_score
 
-    def re_classify(self, tracks:pd.DataFrame, input_video:str, track_ids:list=None, 
-                 out_file:str=None, verbose:bool=True)->pd.DataFrame:
+    def re_classify(self,
+                    tracks:pd.DataFrame,
+                    input_video:str,
+                    track_ids:list=None,
+                    out_file:str=None,
+                    verbose:bool=True)->pd.DataFrame:
         """
         Re-classify tracks
         Parameters:
@@ -55,7 +76,8 @@ class ReClass:
             track_ids = tracks['track'].unique().tolist()
 
         results = []
-        pbar = tqdm(total=len(track_ids), unit='track', desc='Re-classifying tracks')
+        if verbose:
+            pbar = tqdm(total=len(track_ids), unit='track', desc='Re-classifying tracks')
         for track_id in track_ids:
             
             target_track = tracks[tracks['track'] == track_id].copy()
@@ -85,7 +107,8 @@ class ReClass:
             results.append([track_id, cls, round(avg_score, 2)])
             if verbose:
                 pbar.update()
-        pbar.close()
+        if verbose:
+            pbar.close()
 
         df = pd.DataFrame(results, columns=['track', 'cls', 'avg_score'])
         if out_file:

dnt/track/sort/sort.py

@@ -202,7 +202,7 @@ class Sort(object):
   def update(self, dets=np.empty((0, 5))):
     """
     Params:
-      dets - a numpy array of detections in the format [[x1,y1,x2,y2,score],[x1,y1,x2,y2,score],...]
+      dets - a numpy array of detections in the format [[x1,y1,x2,y2,score]
     Requires: this method must be called once for each frame even with empty detections (use np.empty((0, 5)) for frames without detections).
     Returns the a similar array, where the last column is the object ID.
 
@@ -246,9 +246,9 @@ class Sort(object):
 
 def track(det_file, out_file, max_age=1, min_inits=3, iou_threshold=0.3, video_index = None, total_videos = None):
   
-  tracker = Sort(max_age=max_age, 
-                       min_hits=min_inits,
-                       iou_threshold=iou_threshold)
+  tracker = Sort(max_age=max_age,
+                 min_hits=min_inits,
+                 iou_threshold=iou_threshold)
   dets = np.loadtxt(det_file, delimiter=',')
   start_frame = int(dets[:,0].min())
   end_frame = int(dets[:,0].max())
@@ -265,7 +265,6 @@ def track(det_file, out_file, max_age=1, min_inits=3, iou_threshold=0.3, video_i
     if len(dets_frame)>0:
       dets_revised[:,[0,1]] = dets_frame[:,[2,3]]
       dets_revised[:,2] = dets_frame[:,2] + dets_frame[:,4]
-
       dets_revised[:,3] = dets_frame[:,3] + dets_frame[:,5]
       dets_revised[:,4] = dets_frame[:,6] / 100