diff --git a/pyproject.toml b/pyproject.toml index a384a3d..1fe3f70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -dependencies = ["opencv-python>=4.5.0","scipy"] +dependencies = ["opencv-python>=4.5.0","opencv-contrib-python>=4.5.0","scipy"] [project.scripts] tsmark = "tsmark:main" diff --git a/setup.py b/setup.py index 7ffdb9d..a09201a 100644 --- a/setup.py +++ b/setup.py @@ -22,5 +22,5 @@ setup( "tsmark=tsmark:main", ], }, - install_requires=["opencv-python>=4.5.0", "scipy"], + install_requires=["opencv-python>=4.5.0","opencv-contrib-python>=4.5.0", "scipy"], ) diff --git a/tsmark/video_annotator.py b/tsmark/video_annotator.py index bf79e33..bd2baba 100755 --- a/tsmark/video_annotator.py +++ b/tsmark/video_annotator.py @@ -36,10 +36,15 @@ class Marker: self.crop = [(None, None), (None, None), None] self.crop_click = 0 self.point_click = 0 + self.point_tracking = 0 + self.point_tracking_length = 4 self.points = {} self.points_interpolated = {} self.point_index = None + self.message = None + self.message_timer = time.time() + self.forced_fps = opts.fps try: @@ -338,6 +343,7 @@ class Marker: pass def get_point(self, nr=None, index=None): + """[x,y,x2,y2, cx, cy, w, h]""" if nr is None: nr = self.nr if index is None: @@ -349,6 +355,8 @@ class Marker: *self.points[index][nr], int((self.points[index][nr][0] + self.points[index][nr][2]) / 2), int((self.points[index][nr][1] + self.points[index][nr][3]) / 2), + int(abs(self.points[index][nr][0] - self.points[index][nr][2])), + int(abs(self.points[index][nr][1] - self.points[index][nr][3])), ] return [None, None, None, None, None, None] @@ -446,6 +454,110 @@ class Marker: self.interpolate_points() + def modify_point_wh(self): + + if self.point_click == 0: + self.add_message("Not in point clicking mode") + return + if self.opts.output_points is None: + return + curr_point = self.get_point() + if curr_point[0] is None: + self.add_message("Not in point frame (green)") + return + + new_wh = abs(self.mouse_position[0] - curr_point[4]) + new_hh = abs(self.mouse_position[1] - curr_point[5]) + x1 = int(curr_point[4] - new_wh) + y1 = int(curr_point[5] - new_hh) + x2 = int(curr_point[4] + new_wh) + y2 = int(curr_point[5] + new_hh) + self.points[self.point_index][self.nr] = [x1, y1, x2, y2] + self.interpolate_points() + + def track_point(self): + + if self.point_click == 0: + self.add_message("Not in point clicking mode") + return + if self.opts.output_points is None: + return + + old_nr = self.nr + + curr_point = self.get_point() + if curr_point[0] is None: + self.add_message("Not in point frame (green)") + return + + max_frames = int(min(self.point_tracking_length * self.fps, self.frames - self.nr - 1)) + cv2.namedWindow("tsmark - tracker", flags=cv2.WINDOW_AUTOSIZE | cv2.WINDOW_KEEPRATIO | cv2.WINDOW_GUI_NORMAL) + tracker = cv2.TrackerKCF_create() + + self.video_reader.set(cv2.CAP_PROP_POS_FRAMES, self.nr) + ok, frame = self.video_reader.read() + frame = cv2.resize(frame.copy(), self.video_res) + bbox = tuple([*curr_point[0:2], *curr_point[6:8]]) + ok = tracker.init(frame, bbox) + tracked = {} + for i in range(max_frames): + # Read a new frame + ok, frame = self.video_reader.read() + frame = cv2.resize(frame.copy(), self.video_res) + if not ok: + break + + ok, bbox = tracker.update(frame) + # Draw bounding box + if ok: + # Tracking success + p1 = (int(bbox[0]), int(bbox[1])) + p2 = (int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])) + cv2.rectangle(frame, p1, p2, (255, 0, 0), 2, 1) + tracked[i] = bbox + cv2.putText(frame, "Tracking...", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2) + else: + # Tracking failure + cv2.putText( + frame, "Tracking failure detected", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2 + ) + + # Display result + cv2.imshow("tsmark - tracker", frame) + + # Exit if ESC pressed + if cv2.waitKey(1) & 0xFF == ord("q"): # if press SPACE bar + break + + done = False + while True: + if done: + break + self.video_reader.set(cv2.CAP_PROP_POS_FRAMES, self.nr + 1) + for i in range(max_frames): + show_time = time.time() + ok, frame = self.video_reader.read() + frame = cv2.resize(frame.copy(), self.video_res) + cv2.putText(frame, "Replay...", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2) + if i in tracked: + bbox = tracked[i] + p1 = (int(bbox[0]), int(bbox[1])) + p2 = (int(bbox[0] + bbox[2]), int(bbox[1] + bbox[3])) + cv2.rectangle(frame, p1, p2, (255, 0, 0), 2, 1) + cv2.imshow("tsmark - tracker", frame) + if cv2.waitKey(1) & 0xFF == ord("q"): # if press SPACE bar + done = True + break + + time_to_wait = self.viewer_spf - time.time() + show_time + if time_to_wait > 0: + time.sleep(time_to_wait) + + cv2.destroyWindow("tsmark - tracker") + + self.nr = old_nr - 1 + self.read_next = True + def interpolate_points(self): if not self.point_index in self.points_interpolated: @@ -526,6 +638,24 @@ class Marker: ) self.shadow_text(frame, formatted, (left, bottom), 1.1, 2, (255, 255, 255)) + def draw_message(self, frame): + + if self.message is None: + return + if time.time() - 5 > self.message_timer: + self.message = None + return + + left = 10 + bottom = 90 + + self.shadow_text(frame, self.message, (left, bottom), 0.9, 2, (255, 255, 255)) + + def add_message(self, new): + + self.message = new + self.message_timer = time.time() + def format_time(self, nframe): seconds = int(nframe / self.fps) @@ -797,6 +927,7 @@ class Marker: cv2.namedWindow("tsmark", flags=cv2.WINDOW_AUTOSIZE | cv2.WINDOW_KEEPRATIO | cv2.WINDOW_GUI_NORMAL) cv2.setMouseCallback("tsmark", self.mouse_click) digits_ords = [ord(str(x)) for x in range(10)] + FPS_modifier = 1 FPS_modifiers = [0.25, 1, 4] while self.video_reader.isOpened(): @@ -817,6 +948,8 @@ class Marker: self.draw_time(frame_visu) self.draw_bar(frame_visu) self.draw_label(frame_visu) + self.draw_message(frame_visu) + if self.show_help: self.draw_help(frame_visu) @@ -904,6 +1037,7 @@ class Marker: elif k & 0xFF == ord("f"): # modify FPS FPS_modifier = (FPS_modifier + 1) % len(FPS_modifiers) + self.add_message(f"Player speed {FPS_modifiers[FPS_modifier]}") elif k & 0xFF == ord("a"): # toggle crop offset self.crop_click = 0 if self.crop_click == 1 else 1 self.crop[2] = True @@ -931,6 +1065,11 @@ class Marker: else: self.point_index = chr(k2) + elif k & 0xFF == ord("t"): # tracking + self.track_point() + elif k & 0xFF == ord("e"): # point edit (width height) + self.modify_point_wh() + elif k & 0xFF == ord("x"): # toggle ts if self.point_click == 1: self.del_point(self.nr)