tracker

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"

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"],
 )

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)