diff --git a/tsmark/__init__.py b/tsmark/__init__.py index 685e01a..87b4c5e 100644 --- a/tsmark/__init__.py +++ b/tsmark/__init__.py @@ -40,6 +40,15 @@ def get_options(): required=False, help="Load points from a JSON file", ) + parser.add_argument( + "--max-track", + action="store", + dest="max_track", + type=float, + default=4, + required=False, + help="Length of tracking segment in seconds: %(default)s", + ) parser.add_argument( "--fps", action="store", diff --git a/tsmark/video_annotator.py b/tsmark/video_annotator.py index 8db451b..6fadeba 100755 --- a/tsmark/video_annotator.py +++ b/tsmark/video_annotator.py @@ -3,6 +3,7 @@ import os import shlex import subprocess import sys +import threading import time import cv2 @@ -37,7 +38,7 @@ class Marker: self.crop_click = 0 self.point_click = 0 self.point_tracking = 0 - self.point_tracking_length = 4 + self.point_tracking_length = float(self.opts.max_track) self.points = {} self.points_interpolated = {} self.point_index = None @@ -45,6 +46,9 @@ class Marker: self.message = None self.message_timer = time.time() + self.autosave_interval = 60 + self.autosave_timer = time.time() + self.forced_fps = opts.fps try: @@ -234,19 +238,22 @@ class Marker: continue current = self.get_interpolated_point(index=index) - if current[5] == 0: + if current["type"] in ("pre", "post"): continue color = (0, 192, 192) - if current[5] == 2: + if current["type"] == "key": color = (60, 205, 60) - if current[5] == 1: + if current["type"] == "interp": color = (192, 0, 192) - cv2.circle(frame, (current[6], current[7]), 10, (0, 0, 0), 2) - cv2.circle(frame, (current[6], current[7]), 10, color, 1) + if current["visible"] == 0: + color = (96, 96, 96) + continue + cv2.circle(frame, (current["cx"], current["cy"]), 10, (0, 0, 0), 2) + cv2.circle(frame, (current["cx"], current["cy"]), 10, color, 1) self.shadow_text( frame, index, - (current[6], current[7]), + (current["cx"], current["cy"]), 0.5, 1, color, @@ -261,89 +268,109 @@ class Marker: frame, (0, self.mouse_position[1]), (self.video_res[0], self.mouse_position[1]), (128, 128, 128), 1 ) # Show current track - x, y = [self.video_res[0] - 120, 50] + x, y = [20, 70] self.shadow_text( frame, - "Points: " + str(self.point_index), + "P:" + str(self.point_index), (x, y), 0.5, 1, - (0, 192, 192), + (255, 255, 255), ) try: - current = self.get_interpolated_point() # self.points_interpolated[self.point_index][self.nr] - color = (0, 192, 192) - if current[5] == 2: - color = (60, 205, 60) - if current[5] == 1: - color = (192, 0, 192) - cv2.rectangle( - frame, - (current[1], current[2]), - (current[3], current[4]), - color, - 2, - ) - cv2.circle(frame, (current[6], current[7]), 10, color, 1) + current = self.get_interpolated_point() + if current["type"] is not None: + color = (0, 192, 192) + if current["type"] == "key": + color = (60, 205, 60) + if current["type"] == "interp": + color = (192, 0, 192) + if current["visible"] == 0: + color = (96, 96, 96) - history = list(range(max(1, int(self.nr - self.viewer_fps)), self.nr + 1)) - for p in history: - current = self.get_interpolated_point(p) - past = self.get_interpolated_point(p - 1) - cv2.line( + cv2.rectangle( frame, - (past[6], past[7]), - (current[6], current[7]), - (192, 0, 192), - 1, + (current["x0"], current["y0"]), + (current["x1"], current["y1"]), + color, + 2, ) + cv2.circle(frame, (current["cx"], current["cy"]), 10, color, 1) + + history = [] + for p in range(max(1, int(self.nr - self.viewer_fps)), self.nr + 1): + po = self.get_interpolated_point(p) + history.append([po["cx"], po["cy"]]) + history = np.array(history, np.int32).reshape((-1, 1, 2)) + cv2.polylines(frame, [history], False, (192, 0, 192), 1) + except KeyError: + print(current, self.nr) pass except IndexError: print(current, self.nr) pass try: - # ~ point_keys = list(sorted(self.points[self.point_index].keys())) - # current = self.points[self.point_index][self.nr] - current = self.get_point() # self.points_interpolated[self.point_index][self.nr] - color = (60, 205, 60) - cv2.circle(frame, (current[4], current[5]), 13, color, 2) + current = self.get_point() + if current["x0"] is not None: + cv2.circle(frame, (current["cx"], current["cy"]), 13, (60, 205, 60), 2) except KeyError: pass except IndexError: - # ~ print(self.points[self.point_index]) - # ~ print(self.nr) + print(self.points[self.point_index]) + print(self.nr) pass def scan_point(self, direction): + def set_nr(ts): + self.nr = ts - 1 + self.read_next = True + try: + if direction == "first": + return set_nr(min(list(self.points[self.point_index].keys()))) + + if direction == "last": + return set_nr(max(list(self.points[self.point_index].keys()))) if direction == "next": for ts in sorted(list(self.points[self.point_index].keys())): if ts > self.nr: - self.nr = ts - 1 - self.read_next = True - return + return set_nr(ts) if direction == "previous": for ts in reversed(sorted(list(self.points[self.point_index].keys()))): if ts < self.nr - 1: - self.nr = ts - 1 - self.read_next = True - return + return set_nr(ts) + except Exception: pass - def del_point(self, ts): + def toggle_point(self, ts): try: - del self.points[self.point_index][ts] + if ts in self.points[self.point_index]: + # Remove point + del self.points[self.point_index][ts] + else: + # Introduce point from interpolated + ip = self.get_interpolated_point() + if ip["type"] is None: + return + {"x0": None, "y0": None, "x1": None, "y1": None, "cx": None, "cy": None, "visible": 0, "type": None} + self.points[self.point_index][self.nr] = { + "x0": ip["x0"], + "y0": ip["y0"], + "x1": ip["x1"], + "y1": ip["y1"], + "visible": 1, + } self.interpolate_points() except Exception: pass def get_point(self, nr=None, index=None): - """[x,y,x2,y2, cx, cy, w, h]""" + """{x0,y0,x1,y1, cx, cy, w, h, visible}""" if nr is None: nr = self.nr if index is None: @@ -351,105 +378,130 @@ class Marker: if index in self.points: if nr in self.points[index]: - return [ - *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])), - ] + value = self.points[index][nr].copy() + value.update( + { + "cx": int((value["x0"] + value["x1"]) / 2), + "cy": int((value["y0"] + value["y1"]) / 2), + "w": int(abs(value["x0"] - value["x1"])), + "h": int(abs(value["y0"] - value["y1"])), + } + ) + return value - return [None, None, None, None, None, None] + return { + "x0": None, + "y0": None, + "x1": None, + "y1": None, + "cx": None, + "cy": None, + "w": None, + "h": None, + "visible": 0, + } def get_interpolated_point(self, nr=None, index=None): + """{x0,y0,x1,y1, cx, cy, visible,type}""" if nr is None: nr = self.nr if index is None: index = self.point_index - if index in self.points: + if index in self.points_interpolated: if nr in self.points_interpolated[index]: - return [ - *self.points_interpolated[index][nr], - int((self.points_interpolated[index][nr][1] + self.points_interpolated[index][nr][3]) / 2), - int((self.points_interpolated[index][nr][2] + self.points_interpolated[index][nr][4]) / 2), - ] + value = self.points_interpolated[index][nr].copy() + value.update( + { + "cx": int((value["x0"] + value["x1"]) / 2), + "cy": int((value["y0"] + value["y1"]) / 2), + } + ) + return value - return [None, None, None, None, None, None, None, None] + return {"x0": None, "y0": None, "x1": None, "y1": None, "cx": None, "cy": None, "visible": 0, "type": None} def modify_point(self, position, x, y): """position: tl topleft, br bottomright, c center""" - if position == "tl": - ix = 0 - iy = 1 - if position == "br": - ix = 2 - iy = 3 - if not self.point_index in self.points: self.points[self.point_index] = {} + if not self.nr in self.points[self.point_index]: if len(self.points[self.point_index]) > 0: keys = sorted(list(self.points[self.point_index].keys())) - if self.nr > keys[-1]: + if self.nr > keys[-1]: # last point if at end of track last_p = self.points[self.point_index][keys[-1]] - elif self.nr < keys[0]: + elif self.nr < keys[0]: # first point if before track last_p = self.points[self.point_index][keys[0]] - else: + else: # previous point if in the middle of track prev_key = keys[0] for key in keys: if key > self.nr: last_p = self.points[self.point_index][prev_key] break prev_key = key - w = abs(last_p[2] - last_p[0]) - h = abs(last_p[3] - last_p[1]) + w = abs(last_p["x1"] - last_p["x0"]) + h = abs(last_p["y1"] - last_p["y0"]) else: w = 50 h = 50 if position == "tl": - self.points[self.point_index][self.nr] = [ - x, - y, - min(self.video_res[0] - 1, x + w), - min(self.video_res[1] - 1, y + h), - ] + self.points[self.point_index][self.nr] = { + "x0": x, + "y0": y, + "x1": min(self.video_res[0] - 1, x + w), + "y1": min(self.video_res[1] - 1, y + h), + "visible": 1, + } if position == "br": - self.points[self.point_index][self.nr] = [max(0, x - w), max(0, y - h), x, y] + self.points[self.point_index][self.nr] = { + "x0": max(0, x - w), + "y0": max(0, y - h), + "x1": x, + "y1": y, + "visible": 1, + } if position == "c": - self.points[self.point_index][self.nr] = [ - max(0, int(x - w / 2)), - max(0, int(y - h / 2)), - min(self.video_res[0] - 1, int(x + w / 2)), - min(self.video_res[1] - 1, int(y + h / 2)), - ] + self.points[self.point_index][self.nr] = { + "x0": max(0, int(x - w / 2)), + "y0": max(0, int(y - h / 2)), + "x1": min(self.video_res[0] - 1, int(x + w / 2)), + "y1": min(self.video_res[1] - 1, int(y + h / 2)), + "visible": 1, + } else: # not a new point + self.points[self.point_index][self.nr]["visible"] = 1 if position == "c": current = self.points[self.point_index][self.nr] - w = abs(current[2] - current[0]) - h = abs(current[3] - current[1]) - self.points[self.point_index][self.nr] = [ - max(0, int(x - w / 2)), - max(0, int(y - h / 2)), - min(self.video_res[0] - 1, int(x + w / 2)), - min(self.video_res[1] - 1, int(y + h / 2)), - ] - else: - self.points[self.point_index][self.nr][ix] = x - self.points[self.point_index][self.nr][iy] = y + w = abs(current["x1"] - current["x0"]) + h = abs(current["y1"] - current["y0"]) + self.points[self.point_index][self.nr] = { + "x0": max(0, int(x - w / 2)), + "y0": max(0, int(y - h / 2)), + "x1": min(self.video_res[0] - 1, int(x + w / 2)), + "y1": min(self.video_res[1] - 1, int(y + h / 2)), + "visible": 1, + } + elif position == "tl": + self.points[self.point_index][self.nr]["x0"] = x + self.points[self.point_index][self.nr]["y0"] = y - if self.points[self.point_index][self.nr][0] > self.points[self.point_index][self.nr][2]: - self.points[self.point_index][self.nr][2], self.points[self.point_index][self.nr][0] = ( - self.points[self.point_index][self.nr][0], - self.points[self.point_index][self.nr][2], + elif position == "br": + self.points[self.point_index][self.nr]["x1"] = x + self.points[self.point_index][self.nr]["y1"] = y + + if self.points[self.point_index][self.nr]["x0"] > self.points[self.point_index][self.nr]["x1"]: + self.points[self.point_index][self.nr]["x1"], self.points[self.point_index][self.nr]["x0"] = ( + self.points[self.point_index][self.nr]["x0"], + self.points[self.point_index][self.nr]["x1"], ) - if self.points[self.point_index][self.nr][1] > self.points[self.point_index][self.nr][3]: - self.points[self.point_index][self.nr][3], self.points[self.point_index][self.nr][1] = ( - self.points[self.point_index][self.nr][1], - self.points[self.point_index][self.nr][3], + if self.points[self.point_index][self.nr]["y0"] > self.points[self.point_index][self.nr]["y1"]: + self.points[self.point_index][self.nr]["y1"], self.points[self.point_index][self.nr]["y0"] = ( + self.points[self.point_index][self.nr]["y0"], + self.points[self.point_index][self.nr]["y1"], ) self.interpolate_points() @@ -462,17 +514,32 @@ class Marker: if self.opts.output_points is None: return curr_point = self.get_point() - if curr_point[0] is None: + if curr_point["x0"] 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] + new_wh = abs(self.mouse_position[0] - curr_point["cx"]) + new_hh = abs(self.mouse_position[1] - curr_point["cy"]) + self.points[self.point_index][self.nr]["x0"] = int(curr_point["cx"] - new_wh) + self.points[self.point_index][self.nr]["y0"] = int(curr_point["cy"] - new_hh) + self.points[self.point_index][self.nr]["x1"] = int(curr_point["cx"] + new_wh) + self.points[self.point_index][self.nr]["y1"] = int(curr_point["cy"] + new_hh) + self.points[self.point_index][self.nr]["visible"] = 1 + + self.interpolate_points() + + def toggle_point_visibility(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["x0"] is None: + self.add_message("Not in point frame (green)") + return + self.points[self.point_index][self.nr]["visible"] = 1 - self.points[self.point_index][self.nr]["visible"] self.interpolate_points() def track_point(self): @@ -488,60 +555,77 @@ class Marker: for nr in tracker_gui.points: self.points[self.point_index][nr] = tracker_gui.points[nr] self.interpolate_points() - self.nr = max(tracker_gui.points)-1 + self.nr = max(tracker_gui.points) - 1 self.read_next = True def interpolate_points(self): + """types: + key: user clicked / accepted frame + interp: interpolated frame + pre: before any keyframes + post: after any keyframes + """ + + def i_point(x0=None, y0=None, x1=None, y1=None, t=None, visible=None): + return {"x0": x0, "y0": y0, "x1": x1, "y1": y1, "type": t, "visible": visible} + + def point2array(p): + + return [p["x0"], p["y0"], p["x1"], p["y1"]] if not self.point_index in self.points_interpolated: - self.points_interpolated[self.point_index] = {key: [] for key in range(self.frames)} + self.points_interpolated[self.point_index] = {key: {} for key in range(self.frames)} - if len(self.points[self.point_index]) == 1: + if len(self.points[self.point_index]) == 1: # only one point added key = list(self.points[self.point_index].keys())[0] - x, y, x2, y2 = self.points[self.point_index][key] + vals = self.points[self.point_index][key] + # x, y, x2, y2 = self.points[self.point_index][key] for key in range(self.frames): - self.points_interpolated[self.point_index][key] = [False, int(x), int(y), x2, y2, 0] - self.points_interpolated[self.point_index][self.nr][5] = 2 + self.points_interpolated[self.point_index][key] = i_point() + self.points_interpolated[self.point_index][key].update(vals) + self.points_interpolated[self.point_index][key]["type"] = "pre" + + self.points_interpolated[self.point_index][self.nr]["type"] = "key" else: # more points point_keys = list(sorted(list(self.points[self.point_index].keys()))) - point_values = [self.points[self.point_index][k] for k in point_keys] + point_values = [point2array(self.points[self.point_index][k]) for k in point_keys] xyxy = np.array(point_values).T spline = PchipInterpolator(point_keys, xyxy, axis=1) start_key = min(point_keys) end_key = max(point_keys) + 1 t2 = np.arange(start_key, end_key) + # Pre points for key in range(0, start_key): - self.points_interpolated[self.point_index][key] = [ - False, - self.points[self.point_index][start_key][0], - self.points[self.point_index][start_key][1], - self.points[self.point_index][start_key][2], - self.points[self.point_index][start_key][3], - 0, - ] + self.points_interpolated[self.point_index][key]["type"] = "pre" + self.points_interpolated[self.point_index][key].update(self.points[self.point_index][start_key]) # interpolated points + visible = self.points[self.point_index][start_key]["visible"] for row in np.vstack((t2, spline(t2))).T: - self.points_interpolated[self.point_index][row[0]] = [ - True, - int(row[1]), - int(row[2]), - int(row[3]), - int(row[4]), - 1, - ] + if row[0] in point_keys: + visible = self.points[self.point_index][row[0]]["visible"] + self.points_interpolated[self.point_index][row[0]] = { + "type": "interp", + "x0": int(row[1]), + "y0": int(row[2]), + "x1": int(row[3]), + "y1": int(row[4]), + "visible": visible, + } + + # post points for key in range(end_key, self.frames + 1): - self.points_interpolated[self.point_index][key] = [ - False, - int(row[1]), - int(row[2]), - int(row[3]), - int(row[4]), - 3, - ] + self.points_interpolated[self.point_index][key] = { + "type": "post", + "x0": int(row[1]), + "y0": int(row[2]), + "x1": int(row[3]), + "y1": int(row[4]), + "visible": visible, + } # clicked points (not necessary, could determine at draw time!) for key in point_keys: - self.points_interpolated[self.point_index][key][5] = 2 + self.points_interpolated[self.point_index][key]["type"] = "key" def draw_help(self, frame): @@ -598,7 +682,7 @@ class Marker: def get_help(self): return """Keyboard help: - (Note: after mouse click, arrows stop working due to unknown bug: use j,l,i,k) + Arrows, PgUp, PgDn, Home, End or click mouse in position bar j l i k [ ] jump in video position @@ -610,11 +694,27 @@ class Marker: space or click video pause a and s modify crop offset or size - p toggle bounding box drawing. follow with any key as index. left/middle/right mouse sets box position + f toggle 0.25x 1x or 4x FPS v toggle HUD h toggle help q or esc quit + + Bounding box editor: + p toggle bounding box drawing. follow with any key as index. + o toggle object is occluded + x toggle (delete) key frame + mouse left: set top-left corner of box + mouse middle: set center of box + mouse right: set lower right corner of box + e set width/height of box symmetric around center + z c Home End move between key-frames + t start optical flow tracker + Color codes: + green keyframe + purple interpolated frame + gray object is occluded + """ def mouse_click(self, event, x, y, flags, param): @@ -713,10 +813,13 @@ class Marker: self.point_index = index self.points[index] = {int(k): v for k, v in self.points[index].items()} for key in self.points[index]: - self.points[index][key] = [ - *self.original_to_visual((self.points[index][key][0], self.points[index][key][1])), - *self.original_to_visual((self.points[index][key][2], self.points[index][key][3])), - ] + self.points[index][key]["x0"], self.points[index][key]["y0"] = self.original_to_visual( + (self.points[index][key]["x0"], self.points[index][key]["y0"]) + ) + self.points[index][key]["x1"], self.points[index][key]["y1"] = self.original_to_visual( + (self.points[index][key]["x1"], self.points[index][key]["y1"]) + ) + self.interpolate_points() print(f"Loaded points with index: {index}") self.point_index = None @@ -795,10 +898,13 @@ class Marker: for index in self.points.keys(): points[index] = {} for key in sorted(self.points[index].keys()): - points[index][key] = [ - *self.visual_to_original((self.points[index][key][0], self.points[index][key][1])), - *self.visual_to_original((self.points[index][key][2], self.points[index][key][3])), - ] + points[index][key] = self.points[index][key].copy() + points[index][key]["x0"], points[index][key]["y0"] = self.visual_to_original( + (self.points[index][key]["x0"], self.points[index][key]["y0"]) + ) + points[index][key]["x1"], points[index][key]["y1"] = self.visual_to_original( + (self.points[index][key]["x1"], self.points[index][key]["y1"]) + ) with open(self.opts.output_points, "wt") as fp: json.dump(points, fp, indent=2) @@ -863,13 +969,14 @@ class Marker: FPS_modifier = 1 FPS_modifiers = [0.25, 1, 4] + read_fails = 0 while self.video_reader.isOpened(): show_time = time.time() if (not self.paused) or self.read_next: ret, frame = self.video_reader.read() if ret == True: - - draw_wait = 200 if self.paused or ( self.paused and self.point_click == 0 ) else 1 + read_fails = 0 + draw_wait = 200 if self.paused or (self.paused and self.point_click == 0) else 1 if (not self.paused) or self.read_next: self.read_next = False @@ -890,6 +997,7 @@ class Marker: break cv2.imshow("tsmark", frame_visu) k = cv2.waitKey(draw_wait) + if k & 0xFF == ord("q") or k & 0xFF == 27: break elif k & 0xFF == 32: # space @@ -897,13 +1005,19 @@ class Marker: # Movement ================= elif k & 0xFF == 80: # home key - self.nr = -1 - self.read_next = True + if self.point_click == 1: + self.scan_point("first") + else: + self.nr = -1 + self.read_next = True elif k & 0xFF == 87: # end key - self.nr = self.frames - 1 - self.paused = True - self.read_next = True + if self.point_click == 1: + self.scan_point("last") + else: + self.nr = self.frames - 1 + self.paused = True + self.read_next = True elif k & 0xFF == 85 or k & 0xFF == ord("]"): # pg up self.nr = int((nr_time + self.hugestep) * self.fps) - 1 @@ -977,6 +1091,10 @@ class Marker: elif k & 0xFF == ord("s"): # toggle crop size self.crop_click = 0 if self.crop_click == 2 else 2 self.crop[2] = True + elif k & 0xFF == ord("o"): # toggle point visibility (occlusion) + if self.opts.output_points is None: + continue + self.toggle_point_visibility() elif k & 0xFF == ord("p"): # toggle points if self.opts.output_points is None: continue @@ -985,10 +1103,10 @@ class Marker: self.shadow_text( frame_visu, "Enter point index", - (self.video_res[0] - 200, 50), - 0.5, - 1, - (0, 192, 192), + (20, 70), + 0.9, + 2, + (255, 255, 255), ) cv2.imshow("tsmark", frame_visu) @@ -1004,10 +1122,10 @@ class Marker: (20, 70), 0.9, 2, - (255,255,255), + (255, 255, 255), ) cv2.imshow("tsmark", frame_visu) - entered_chars="" + entered_chars = "" while True: frame_query = frame_visu.copy() self.shadow_text( @@ -1016,7 +1134,7 @@ class Marker: (20, 100), 0.9, 2, - (255,255,255), + (255, 255, 255), ) cv2.imshow("tsmark", frame_query) del frame_query @@ -1025,7 +1143,7 @@ class Marker: break elif k2 & 0xFF == ord("g") or k2 & 0xFF == 13: try: - self.nr = int(entered_chars) -1 + self.nr = int(entered_chars) - 1 except ValueError: try: self.nr = self.parse_time(entered_chars) @@ -1034,7 +1152,7 @@ class Marker: break self.read_next = True break - elif k2 & 0xFF == 8: # backspace + elif k2 & 0xFF == 8: # backspace entered_chars = entered_chars[0:-1] elif k2 & 0xFF in digits_ords: entered_chars += str(digits_ords.index(k2 & 0xFF)) @@ -1045,7 +1163,6 @@ class Marker: else: pass - elif k & 0xFF == ord("t"): # tracking self.track_point() elif k & 0xFF == ord("e"): # point edit (width height) @@ -1053,7 +1170,7 @@ class Marker: elif k & 0xFF == ord("x"): # toggle ts if self.point_click == 1: - self.del_point(self.nr) + self.toggle_point(self.nr) else: self.toggle_stamp() @@ -1078,11 +1195,23 @@ class Marker: time.sleep(time_to_wait) else: - self.nr = self.frames - 2 + self.nr = self.frames - 2 - read_fails self.video_reader.set(cv2.CAP_PROP_POS_FRAMES, self.nr) + read_fails += 1 + if read_fails > self.frames: + self.nr = 0 + self.open() self.paused = True self.read_next = True + if time.time() > self.autosave_timer + self.autosave_interval: + self.autosave_timer = time.time() + try: + print("Autosave timestamps / points") + self.save_timestamps() + except Exception as e: + print(e) + self.video_reader.release() cv2.destroyAllWindows() self.print_timestamps() @@ -1093,28 +1222,41 @@ class TrackerGUI: def __init__(self, marker): self.marker = marker + self.points = {} + try: + cv2.TrackerKCF_create() + except AttributeError: + marker.add_message("Tracking failed: missing opencv contrib") + return + self.start() def start(self): - old_nr = self.marker.nr curr_point = self.marker.get_point() - if curr_point[0] is None: + if curr_point["x0"] is None: self.marker.add_message("Not in point frame (green)") return - max_frames = int(min(self.marker.point_tracking_length * self.marker.fps, self.marker.frames - self.marker.nr - 1)) + max_frames = int( + min(self.marker.point_tracking_length * self.marker.fps, self.marker.frames - self.marker.nr - 1) + ) cv2.namedWindow("tsmark - tracker", flags=cv2.WINDOW_AUTOSIZE | cv2.WINDOW_KEEPRATIO | cv2.WINDOW_GUI_NORMAL) tracker = cv2.TrackerKCF_create() self.marker.video_reader.set(cv2.CAP_PROP_POS_FRAMES, self.marker.nr) + # TODO: track using original video resolution! ok, frame = self.marker.video_reader.read() frame = cv2.resize(frame.copy(), self.marker.video_res) - bbox = tuple([*curr_point[0:2], *curr_point[6:8]]) + bbox = tuple([curr_point["x0"], curr_point["y0"], curr_point["w"], curr_point["h"]]) ok = tracker.init(frame, bbox) + visu_interval = 0.2 + show_time = 0 + show_message = "" tracked = {} + tracked[0] = [*bbox, 1] for i in range(max_frames): # Read a new frame ok, frame = self.marker.video_reader.read() @@ -1123,26 +1265,29 @@ class TrackerGUI: break ok, bbox = tracker.update(frame) - # Draw bounding box + # ~ print(f"Tracking... ({i}/{max_frames})") 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,1] - cv2.putText(frame, "Tracking...", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2) + tracked[i + 1] = [*bbox, 1] + show_message = f"Tracking... ({i}/{max_frames})" else: # Tracking failure - cv2.putText( - frame, "Tracking failure detected", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2 - ) + show_message = f"Tracking failure detected ({i}/{max_frames})" + bbox = None - # Display result - cv2.imshow("tsmark - tracker", frame) - - # Exit if ESC pressed - if cv2.waitKey(1) & 0xFF == ord("q"): # if press SPACE bar - break + if time.time() > show_time + visu_interval: + # Display result + if bbox is not None: + 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) + self.marker.shadow_text(frame, show_message, (100, 80), 0.75, 2, (255, 255, 255)) + cv2.imshow("tsmark - tracker", frame) + show_time = time.time() + k = cv2.waitKey(1) + # break tracking if ESC pressed, q, space or enter + if k & 0xFF == ord("q") or k & 0xFF == 32 or k & 0xFF == 27 or k & 0xFF == 13: + break done = False paused = False @@ -1151,8 +1296,8 @@ class TrackerGUI: while True: if done: break - self.marker.video_reader.set(cv2.CAP_PROP_POS_FRAMES, self.marker.nr + 1) - i = 0 + self.marker.video_reader.set(cv2.CAP_PROP_POS_FRAMES, self.marker.nr) + i = -1 while True: show_time = time.time() if done: @@ -1165,20 +1310,25 @@ class TrackerGUI: frame_copy = frame.copy() i += 1 seek = False - - cv2.putText(frame, f"Replay... {i}", (100, 80), cv2.FONT_HERSHEY_SIMPLEX, 0.75, (0, 0, 255), 2) + self.marker.shadow_text(frame, f"Accept? ({i+1}/{max_frames})", (100, 80), 0.75, 2, (255, 255, 255)) 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])) - color = (0,255,0) if cut_after > i else (0, 192, 192) + color = (0, 255, 0) if cut_after > i else (0, 192, 192) thicc = 2 if cut_after > i else 1 cv2.rectangle(frame, p1, p2, color, thicc, 1) cv2.imshow("tsmark - tracker", frame) - k = cv2.waitKey(1) - if k & 0xFF == ord("q"): # if press SPACE bar + # speed up fps by 2 + time_to_wait = self.marker.viewer_spf / 2 - time.time() + show_time + k = cv2.waitKey(max(1, int(time_to_wait * 1000))) + if k & 0xFF == ord("q") or k & 0xFF == 13: # accept with q or enter done = True break + if k & 0xFF == 27: # decline with escape + done = True + cut_after = 0 + break elif k & 0xFF == 32: # space paused = not paused # Movement ================= @@ -1189,48 +1339,40 @@ class TrackerGUI: i -= int(self.marker.fps) + 1 seek = True # Move by frame - elif k & 0xFF == ord("."): + elif k & 0xFF == ord(".") or k & 0xFF == ord("c"): paused = True seek = True - elif k & 0xFF == ord(","): + elif k & 0xFF == ord(",") or k & 0xFF == ord("z"): paused = True i -= 2 seek = True elif k & 0xFF == ord("x"): cut_after = i - #if i in tracked: - # tracked[i][4] = 1 - tracked[i][4] + # TODO: ord("h") for help! - time_to_wait = self.marker.viewer_spf - time.time() + show_time - if time_to_wait > 0: - time.sleep(time_to_wait) - - if i >= max_frames: - i = max_frames-1 + if i >= max_frames - 1: + i = max_frames - 2 paused = True seek = True - if i<0: - i=0 + if i < 0: + i = -1 paused = True seek = True if seek: - self.marker.video_reader.set(cv2.CAP_PROP_POS_FRAMES, self.marker.nr + 1 + i) - + self.marker.video_reader.set(cv2.CAP_PROP_POS_FRAMES, self.marker.nr + i + 1) cv2.destroyWindow("tsmark - tracker") - self.marker.nr = old_nr - 1 self.marker.read_next = True self.points = {} for i in sorted(list(tracked.keys())): if i >= cut_after: continue - self.points[self.marker.nr+1+i] = [ - tracked[i][0], - tracked[i][1], - tracked[i][0]+tracked[i][2], - tracked[i][1]+tracked[i][3], - ] - - + self.points[self.marker.nr + i + 1] = { + "x0": tracked[i][0], + "y0": tracked[i][1], + "x1": tracked[i][0] + tracked[i][2], + "y1": tracked[i][1] + tracked[i][3], + "visible": 1, + }