gluing point tracking tool

tsmark/video_annotator.py

@@ -5,6 +5,8 @@ import sys
 import time
 
 import cv2
+import numpy as np
+from scipy.interpolate import PchipInterpolator
 
 AUDIO_COMPRESS = "-c:a libmp3lame -ar 44100 -b:a 192k".split(" ")
 AUDIO_COPY = "-c:a copy".split(" ")
@@ -31,6 +33,11 @@ class Marker:
         self.min_res = (512, None)
         self.crop = [(None, None), (None, None), None]
         self.crop_click = 0
+        self.point_click = 0
+        self.points = {}
+        self.points_interpolated = {}
+        self.point_index = None
+
         self.forced_fps = opts.fps
 
         try:
@@ -219,6 +226,118 @@ class Marker:
                 (0, 192, 192),
             )
 
+    def draw_points(self, frame):
+
+        if self.point_click == 0:
+            return
+
+        x, y = [self.video_res[0] - 120, 50]
+        self.shadow_text(
+            frame,
+            "Points: " + str(self.point_index),
+            (x, y),
+            0.5,
+            1,
+            (0, 192, 192),
+        )
+        try:
+            current = self.points_interpolated[self.point_index][self.nr]
+            color = (0, 192, 192)
+            if current[3] == 2:
+                color = (60, 205, 60)
+            if current[3] == 1:
+                color = (192, 0, 192)
+            cv2.circle(frame, (current[1], current[2]), 10, color, 1)
+            history = list(range(max(1, int(self.nr - self.viewer_fps)), self.nr + 1))
+            for p in history:
+                self.points_interpolated[self.point_index][p - 1][1:2]
+                cv2.line(
+                    frame,
+                    tuple(self.points_interpolated[self.point_index][p - 1][1:3]),
+                    tuple(self.points_interpolated[self.point_index][p][1:3]),
+                    (192, 0, 192),
+                    1,
+                )
+
+        except KeyError:
+            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]
+            color = (60, 205, 60)
+            cv2.circle(frame, (current[0], current[1]), 13, color, 2)
+        except KeyError:
+            pass
+        except IndexError:
+            print(self.points[self.point_index])
+            print(self.nr)
+            pass
+
+    def scan_point(self, direction):
+        try:
+
+            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
+
+            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
+        except Exception:
+            pass
+
+    def del_point(self, ts):
+        try:
+            del self.points[self.point_index][ts]
+            self.interpolate_points()
+        except Exception:
+            pass
+
+    def interpolate_points(self):
+
+        if not self.point_index in self.points_interpolated:
+            self.points_interpolated[self.point_index] = {key: [] for key in range(self.frames)}
+
+        if len(self.points[self.point_index]) == 1:
+            key = list(self.points[self.point_index].keys())[0]
+            x, y = self.points[self.point_index][key]
+            for key in range(self.frames):
+                self.points_interpolated[self.point_index][key] = [False, int(x), int(y), 0]
+            self.points_interpolated[self.point_index][self.nr][3] = 2
+
+        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]
+            xy = np.array(point_values).T
+            spline = PchipInterpolator(point_keys, xy, axis=1)
+            start_key = min(point_keys)
+            end_key = max(point_keys) + 1
+            t2 = np.arange(start_key, end_key)
+            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],
+                    0,
+                ]
+            # interpolated points
+            for row in np.vstack((t2, spline(t2))).T:
+                self.points_interpolated[self.point_index][row[0]] = [True, int(row[1]), int(row[2]), 1]
+            for key in range(end_key, self.frames + 1):
+                self.points_interpolated[self.point_index][key] = [False, int(row[1]), int(row[2]), 0]
+            # clicked points (not necessary, could determine at draw time!)
+            for key in point_keys:
+                self.points_interpolated[self.point_index][key][3] = 2
+
     def draw_help(self, frame):
 
         bottom = 80
@@ -296,6 +415,14 @@ class Marker:
             return
 
         if event == cv2.EVENT_LBUTTONDOWN:
+            if self.point_click == 1:
+                if not self.point_index in self.points:
+                    self.points[self.point_index] = {}
+                self.points[self.point_index][self.nr] = (int(x), int(y))
+                self.interpolate_points()
+
+                return
+
             if in_bar:
                 click_relative = (x - self.bar_start) / (self.bar_end - self.bar_start)
                 self.nr = int(click_relative * self.frames)
@@ -304,6 +431,13 @@ class Marker:
                 self.paused = not self.paused
 
         if event == cv2.EVENT_LBUTTONDBLCLK:
+            if self.point_click == 1:
+                if not self.point_index in self.points:
+                    return
+                if self.nr in self.points[self.point_index]:
+                    del self.points[self.point_index][self.nr]
+                return
+
             if not in_bar:
                 self.toggle_stamp()
             # doubleclick (toggle?)
@@ -485,6 +619,7 @@ class Marker:
                     self.read_next = False
                 frame_visu = cv2.resize(frame.copy(), self.video_res)
                 self.draw_crop(frame_visu)
+                self.draw_points(frame_visu)
                 nr_time = self.nr / self.fps
                 if self.show_info:
                     self.draw_time(frame_visu)
@@ -550,17 +685,23 @@ class Marker:
                     self.read_next = True
 
                 elif k & 0xFF == ord("z"):  # move to previous ts
-                    for ts in reversed(sorted(self.stamps)):
-                        if ts < self.nr - 1:
-                            self.nr = ts - 1
-                            self.read_next = True
-                            break
-                elif k & 0xFF == ord("c"):  # move to previous ts
-                    for ts in sorted(self.stamps):
-                        if ts > self.nr:
-                            self.nr = ts - 1
-                            self.read_next = True
-                            break
+                    if self.point_click == 1:
+                        self.scan_point("previous")
+                    else:
+                        for ts in reversed(sorted(self.stamps)):
+                            if ts < self.nr - 1:
+                                self.nr = ts - 1
+                                self.read_next = True
+                                break
+                elif k & 0xFF == ord("c"):  # move to next ts
+                    if self.point_click == 1:
+                        self.scan_point("next")
+                    else:
+                        for ts in sorted(self.stamps):
+                            if ts > self.nr:
+                                self.nr = ts - 1
+                                self.read_next = True
+                                break
 
                 # Move by number
                 elif k & 0xFF in digits_ords:
@@ -577,8 +718,30 @@ 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("p"):  # toggle points
+                    self.point_click = 0 if self.point_click == 1 else 1
+                    if self.point_click == 1:
+                        self.shadow_text(
+                            frame_visu,
+                            "Enter point index",
+                            (self.video_res[0] - 200, 50),
+                            0.5,
+                            1,
+                            (0, 192, 192),
+                        )
+                        cv2.imshow("tsmark", frame_visu)
+
+                        k2 = cv2.waitKey(0)
+                        if k2 & 0xFF == ord("q") or k2 & 0xFF == 27:
+                            self.point_click = 0
+                        else:
+                            self.point_index = chr(k2)
+
                 elif k & 0xFF == ord("x"):  # toggle ts
-                    self.toggle_stamp()
+                    if self.point_click == 1:
+                        self.del_point(self.nr)
+                    else:
+                        self.toggle_stamp()
 
                 elif k & 0xFF == ord("v"):
                     self.show_info = not self.show_info