diff --git a/tsmark/video_annotator.py b/tsmark/video_annotator.py index 2e14a39..0b0bd73 100755 --- a/tsmark/video_annotator.py +++ b/tsmark/video_annotator.py @@ -305,9 +305,7 @@ class Marker: ) # Show current track intrp_str = ( - "" - if not self.points_interpolation_enabled - else " i*" if True in self.points_interpolation_required.values() else " i" + "" if not self.points_interpolation_enabled else " i*" if self.get_interpolation_required() else " i" ) self.shadow_text( frame, @@ -342,17 +340,13 @@ class Marker: thick_y1 = current[nearest_wall] if "y" in nearest_wall else current["y1"] cv2.line(frame, (thick_x0, thick_y0), (thick_x1, thick_y1), color, 5) - if self.mouse_flags["shift"]: - key_combos = (("x0", "y0"), ("x0", "y1"), ("x1", "y1"), ("x1", "y0")) - dists = [ - ( - (px, py), - abs(self.mouse_position[0] - current[px]) + abs(self.mouse_position[1] - current[py]), - ) - for i, (px, py) in enumerate(key_combos) - ] - dists.sort(key=lambda d: d[1]) - cv2.circle(frame, (current[dists[0][0][0]], current[dists[0][0][1]]), 5, color, -1) + else: + if self.mouse_flags["shift"]: + corners = (("x0", "y1"), ("x1", "y0")) + else: + corners = (("x0", "y0"), ("x1", "y1")) + for corner in corners: + cv2.circle(frame, (current[corner[0]], current[corner[1]]), 3, color, -1) except (KeyError, IndexError, TypeError): # print(self.get_interpolated_point(), self.nr) @@ -410,6 +404,7 @@ class Marker: if ts in self.points[self.point_index]: # Remove point del self.points[self.point_index][ts] + self.interpolate_set(self.point_index, ts) else: # Introduce point from interpolated ip = self.get_interpolated_point() @@ -422,7 +417,7 @@ class Marker: "y1": ip["y1"], "visible": ip["visible"], } - self.interpolate_set(self.point_index) + self.interpolate_set(self.point_index, self.nr) except Exception: pass @@ -540,7 +535,7 @@ class Marker: } def modify_point(self, position, x, y): - """position: tl topleft, br bottomright, c center, n nearest, snap to nearest wall""" + """position: tl topleft, br bottomright, tr,bl, c center, n nearest, snap to nearest wall""" def get_points_by_nearest(last_p, x, y, w=None, h=None): # Modifies in place! @@ -615,6 +610,22 @@ class Marker: "y1": y, "visible": visibility, } + if position == "tr": + self.points[self.point_index][self.nr] = { + "x0": max(0, x - w), + "y0": y, + "x1": x, + "y1": min(self.video_res[1] - 1, y + h), + "visible": visibility, + } + if position == "bl": + self.points[self.point_index][self.nr] = { + "x0": x, + "y0": max(0, y - h), + "x1": min(self.video_res[0] - 1, x + w), + "y1": y, + "visible": visibility, + } if position == "c": self.points[self.point_index][self.nr] = { "x0": max(0, int(x - w / 2)), @@ -643,10 +654,15 @@ class Marker: elif position == "tl": self.points[self.point_index][self.nr]["x0"] = x self.points[self.point_index][self.nr]["y0"] = y - elif position == "br": self.points[self.point_index][self.nr]["x1"] = x self.points[self.point_index][self.nr]["y1"] = y + elif position == "tr": + self.points[self.point_index][self.nr]["x1"] = x + self.points[self.point_index][self.nr]["y0"] = y + elif position == "bl": + self.points[self.point_index][self.nr]["x0"] = x + self.points[self.point_index][self.nr]["y1"] = y elif position == "n": self.points[self.point_index][self.nr] = get_points_by_nearest( @@ -674,7 +690,7 @@ class Marker: self.points[self.point_index][self.nr]["y0"], self.points[self.point_index][self.nr]["y1"], ) - self.interpolate_set(self.point_index) + self.interpolate_set(self.point_index, self.nr) def modify_point_wh(self): @@ -695,7 +711,7 @@ class Marker: 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"] = POINT_VISIBILITY[0] - self.interpolate_set(self.point_index) + self.interpolate_set(self.point_index, self.nr) def toggle_point_visibility(self): @@ -722,7 +738,7 @@ class Marker: except Exception as e: print(e) pass - self.interpolate_set(self.point_index) + self.interpolate_set(self.point_index, self.nr) def track_point(self): @@ -738,7 +754,7 @@ class Marker: if len(tracker_gui.points) > 0: for nr in tracker_gui.points: self.points[self.point_index][nr] = tracker_gui.points[nr] - self.interpolate_set(self.point_index) + self.interpolate_set(self.point_index, nr) self.nr = max(tracker_gui.points) - 1 self.read_next = True @@ -787,11 +803,26 @@ class World: if self.plugin: self.plugin() - def interpolate_set(self, point_index=None): + def get_interpolation_required(self, point_index=None): + + if point_index is None: + for key in self.points_interpolation_required: + if len(self.points_interpolation_required[key]) > 0: + return True + else: + if len(self.points_interpolation_required[point_index]) > 0: + return True + return False + + def interpolate_set(self, point_index=None, nr=None): if point_index is None: point_index = self.point_index - self.points_interpolation_required[point_index] = True + if nr is None: + nr = self.nr + if not point_index in self.points_interpolation_required: + self.points_interpolation_required[point_index] = [] + self.points_interpolation_required[point_index].append(nr) if not self.interpolate_thread_alive(): print("ERROR: Interpolator thread is not running!") self.interpolate_thread_start() @@ -803,14 +834,11 @@ class World: pre: before any keyframes post: after any keyframes """ + if point_index is None: + point_index = self.point_index + if not point_index in self.points: + return try: - if not point_index in self.points: - return - - if point_index is None: - point_index = self.point_index - - self.points_interpolation_required[point_index] = False def i_point(x0=None, y0=None, x1=None, y1=None, t=None, visible=None, age=None): return {"x0": x0, "y0": y0, "x1": x1, "y1": y1, "type": t, "visible": visible, "age": age} @@ -818,8 +846,23 @@ class World: def point2array(p): return [p["x0"], p["y0"], p["x1"], p["y1"]] + def tryadd(s, key, l): + try: + s.add(l[key]) + except IndexError: + pass + + if point_index in self.points_interpolation_required: + seed_keys = ( + [x for x in self.points_interpolation_required[point_index]] + if len(self.points_interpolation_required[point_index]) > 0 + else [] + ) + else: + seed_keys = [] + if not point_index in self.points_interpolated: - self.points_interpolated[point_index] = {key: {} for key in range(self.frames)} + self.points_interpolated[point_index] = {key: {} for key in range(self.frames + 1)} new_points = {k: v for k, v in self.points_interpolated[point_index].items()} @@ -835,19 +878,55 @@ class World: self.points_interpolated[point_index] = new_points else: # more points - point_keys = list(sorted(list(self.points[point_index].keys()))) - point_values = [point2array(self.points[point_index][k]) for k in point_keys] - xyxy = np.array(point_values).T + global_point_keys = list(sorted(list(self.points[point_index].keys()))) + + if len(seed_keys) > 0: + point_keys = set() + min_seed = min(seed_keys) + max_seed = max(seed_keys) + min_found = False + max_found = False + for pi, nr in enumerate(global_point_keys): + if not min_found: + min_index = pi + if not max_found: + max_index = pi + if nr >= min_seed: + min_found = True + if nr <= max_seed: + tryadd(point_keys, pi, global_point_keys) + + if nr > max_seed: + max_found = True + break + if min_found: + for shift in range(-3, 0): + tryadd(point_keys, min_index + shift, global_point_keys) + if max_found: + for shift in range(2): + tryadd(point_keys, max_index + shift, global_point_keys) + else: + for shift in range(-3, -1): + tryadd(point_keys, len(global_point_keys) + shift, global_point_keys) + + point_keys = list(sorted(list(point_keys))) + else: + point_keys = global_point_keys + + print(point_index, seed_keys, point_keys) + + xyxy = np.array([point2array(self.points[point_index][k]) for k in point_keys]).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) + start_key = min(global_point_keys) + end_key = max(global_point_keys) + 1 + # Pre points for key in range(0, start_key): new_points[key]["type"] = "pre" new_points[key].update(self.points[point_index][start_key]) # interpolated points visible = self.points[point_index][start_key]["visible"] + t2 = np.arange(min(point_keys), max(point_keys) + 1) for row in np.vstack((t2, spline(t2))).T: if row[0] in point_keys: visible = self.points[point_index][row[0]]["visible"] @@ -861,15 +940,10 @@ class World: } # post points + global_last = self.points[point_index][max(global_point_keys)] for key in range(end_key, self.frames + 1): - new_points[key] = { - "type": "post", - "x0": int(row[1]), - "y0": int(row[2]), - "x1": int(row[3]), - "y1": int(row[4]), - "visible": visible, - } + new_points[key]["type"] = "post" + new_points[key].update(global_last) # clicked points (not necessary, could determine at draw time!) for key in point_keys: new_points[key]["type"] = "key" @@ -882,8 +956,15 @@ class World: new_points[key]["age"] = age self.points_interpolated[point_index] = new_points + if point_index in self.points_interpolation_required: + self.points_interpolation_required[point_index] = [ + x for x in self.points_interpolation_required[point_index] if x not in seed_keys + ] except Exception as e: print(f"Interpolation error: {e}") + print(new_points) + print(point_keys) + raise (e) def interpolate_thread_start(self): @@ -901,7 +982,7 @@ class World: self.points_interpolation_thread_exit = True for point_index in self.points_interpolation_required: - if self.points_interpolation_required[point_index]: + if self.get_interpolation_required(point_index): self.interpolate_points(point_index) def interpolate_thread_alive(self): @@ -922,7 +1003,7 @@ class World: continue for point_index in self.points_interpolation_required: - if self.points_interpolation_required[point_index]: + if self.get_interpolation_required(point_index): self.interpolate_points(point_index) def toggle_interpolation(self, value=None): @@ -1035,9 +1116,9 @@ class World: x toggle (delete) key frame r convert interpolated points to points (no undo!) u toggle automatic interpolation - mouse left: set top-left corner of box. shift: modify nearest corner, ctrl: side to image edge + mouse left: set top-left corner of box. shift: bottom-left, ctrl: side to image edge mouse middle: set center of box - mouse right: set lower right corner of box + mouse right: set lower right corner of box. shift: top-right e set width/height of box symmetric around center z c Home End move between key-frames t start optical flow tracker @@ -1082,11 +1163,14 @@ class World: if self.mouse_flags["ctrl"]: self.modify_point("snap", int(x), int(y)) elif self.mouse_flags["shift"]: - self.modify_point("n", int(x), int(y)) + self.modify_point("bl", int(x), int(y)) else: self.modify_point("tl", int(x), int(y)) elif event == cv2.EVENT_RBUTTONDOWN: - self.modify_point("br", int(x), int(y)) + if self.mouse_flags["shift"]: + self.modify_point("tr", int(x), int(y)) + else: + self.modify_point("br", int(x), int(y)) elif event == cv2.EVENT_MBUTTONDOWN: self.modify_point("c", int(x), int(y)) @@ -1170,7 +1254,7 @@ class World: ) if not self.points[index][key].get("visible", "NA") in POINT_VISIBILITY: self.points[index][key]["visible"] = POINT_VISIBILITY[0] - self.interpolate_set(index) + self.interpolate_set(index, key) print(f"Loaded points with index: {index}") self.point_index = None