diff --git a/trap/tools.py b/trap/tools.py index 0aa5614..f34555b 100644 --- a/trap/tools.py +++ b/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]