Formatting

trap/tools.py

@@ -1,4 +1,5 @@
 from argparse import Namespace
+import math
 from pathlib import Path
 import pickle
 from tempfile import mktemp
@@ -37,16 +38,18 @@ class FrameGenerator():
         self.video_nr = None
         self.frame_count = None
         self.frame_idx = None
+        self.n = 0
 
     def __iter__(self):
-        n = 0
         for video_nr, video_path in enumerate(self.video_srcs):
             self.video_path = video_path
-            self.video_nr = video_path
+            self.video_nr = video_nr
             logger.info(f"Play from '{str(video_path)}'")
             video = cv2.VideoCapture(str(video_path))
             fps = video.get(cv2.CAP_PROP_FPS)
             self.frame_count = video.get(cv2.CAP_PROP_FRAME_COUNT)
+            if self.frame_count < 0:
+                self.frame_count = math.inf
             self.frame_idx = 0
             if self.config.video_offset:
                 logger.info(f"Start at frame {self.config.video_offset}")
@@ -56,16 +59,29 @@ class FrameGenerator():
             while True:
                 ret, img = video.read()
                 self.frame_idx+=1
-                n+=1
+                self.n+=1
 
                 # seek to 0 if video has finished. Infinite loop
                 if not ret:
                     # now loading multiple files        
                     break
 
-                frame = Frame(index=n, img=img, H=self.config.H, camera=self.config.camera)
+                frame = Frame(index=self.n, img=img, H=self.config.H, camera=self.config.camera)
                 yield frame
 
+def marquee_string(string: str, window: int, i: int):
+    if window > len(string):
+        return string
+    
+    # too_much = len(string) - window
+    # offset = i % too_much
+    # return string[offset:offset+window]
+    
+    too_much = len(string) - window
+    offset = i % (too_much*2)
+    if offset > too_much:
+        offset = too_much - (offset-too_much)
+    return string[offset:offset+window]
 
 def tracker_preprocess():
 
@@ -89,7 +105,7 @@ def tracker_preprocess():
                 total += len(detections)
                 # detections = _yolov8_track(frame, model, imgsz=1440, classes=[0])
 
-                bar.set_description(f"[{frames.video_nr}/{len(frames.video_srcs)}] [{frames.frame_idx}/{frames.frame_count}] {str(frames.video_path)} -- Detections {len(detections)}: {[d.track_id for d in detections]} (so far {total})")
+                bar.set_description(f"{frames.video_nr}/{len(frames.video_srcs)} [{frames.frame_idx}/{frames.frame_count}] {marquee_string(str(frames.video_path), 10, frames.n//2)} | dets {len(detections)}: {[d.track_id for d in detections]} (∑{total})")
             
                 for detection in detections:
                     track = tracks[detection.track_id]