Plugin support

tsmark/__init__.py

@@ -2,7 +2,7 @@ import argparse
 
 from tsmark.video_annotator import Marker
 
-VERSION = "0.7.0"
+VERSION = "0.7.1"
 
 
 def get_options():
@@ -105,7 +105,14 @@ def get_options():
         required=False,
         help="Run FFMPEG commands instead of just printing them",
     )
-
+    parser.add_argument(
+        "--plugin",
+        action="store",
+        dest="plugin",
+        default=None,
+        type=str,
+        help="Name of plugin class to use, ex. myplugin:ClassName. Plugins loaded from ~/.config/tsmark/plugins/ (try hello:World)",
+    )
     parser.add_argument("--version", action="version", version=VERSION)
     parser.add_argument(action="store", dest="video")
     return parser.parse_args()

tsmark/video_annotator.py

@@ -1,3 +1,4 @@
+import importlib.util
 import json
 import os
 import shlex
@@ -16,6 +17,7 @@ AUDIO_DISCARD = "-an".split(" ")
 VIDEO_COMPRESS = "-c:v mpeg2video -q:v 3 -g 1".split(" ")
 VIDEO_COPY = "-c:v copy".split(" ")
 EXT_COMPRESS = ".mpeg"
+PLUGIN_FOLDER = os.path.expanduser("~/.config/tsmark/plugins")
 
 
 class Marker:
@@ -52,6 +54,7 @@ class Marker:
         self.forced_fps = opts.fps
 
         try:
+            self.load_plugin()
             self.open()
             self.calculate_res()
             self.parse_timestamps()
@@ -166,6 +169,17 @@ class Marker:
                 (84, 255, 63),
                 1,
             )
+
+        if self.point_click == 1 and self.point_index in self.points:
+            bar_middle = int((self.bar_top + self.bar_bottom) / 2)
+            for ts in self.points[self.point_index]:
+                p_pos = int(self.bar_start + ts / self.frames * (self.bar_end - self.bar_start))
+                cv2.circle(frame, (p_pos, bar_middle), 3, (32, 32, 32), -1)
+            for ts in self.points[self.point_index]:
+                p_pos = int(self.bar_start + ts / self.frames * (self.bar_end - self.bar_start))
+                color = (60, 205, 60) if self.points[self.point_index][ts]["visible"] == 1 else (96, 96, 96)
+                cv2.circle(frame, (p_pos, bar_middle), 1, color, -1)
+
         cv2.line(
             frame,
             (bar_position, self.bar_top),
@@ -558,6 +572,48 @@ class Marker:
             self.nr = max(tracker_gui.points) - 1
             self.read_next = True
 
+    def load_plugin(self):
+
+        self.plugin = None
+        if self.opts.plugin:
+            if not os.path.exists(os.path.join(PLUGIN_FOLDER, "hello.py")):
+                os.makedirs(PLUGIN_FOLDER, exist_ok=True)
+                with open(os.path.join(PLUGIN_FOLDER, "hello.py"), "wt") as fp:
+                    fp.write(
+                        """import cv2
+import numpy as np
+class World:
+    def __init__(self, tsmark):
+        self.tsmark = tsmark
+        self.window_name = "tsmark - plugin"
+        self.tsmark.paused = True
+        cv2.namedWindow(self.window_name, flags=cv2.WINDOW_AUTOSIZE | cv2.WINDOW_KEEPRATIO | cv2.WINDOW_GUI_NORMAL)
+        frame = cv2.resize(np.zeros((16, 16, 3), dtype=np.uint8), self.tsmark.video_res)
+        self.tsmark.shadow_text(frame, "Hello World! press q to exit.", (100, 80), 0.75, 2, (255, 255, 255))
+        cv2.imshow(self.window_name, frame)
+        while True:
+            k = cv2.waitKey(10)
+            # break 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
+        cv2.destroyWindow(self.window_name)
+        return
+"""
+                    )
+
+            plugin_file, plugin_class = self.opts.plugin.split(":", 1)
+            plugin_path = os.path.join(PLUGIN_FOLDER, plugin_file + ".py")
+            module_spec = importlib.util.spec_from_file_location("Plugin", plugin_path)
+            loaded_plugin = importlib.util.module_from_spec(module_spec)
+            module_spec.loader.exec_module(loaded_plugin)
+            plugin_class = getattr(loaded_plugin, plugin_class)
+            self.plugin = plugin_class
+
+    def launch_plugin(self):
+
+        if self.plugin:
+            self.plugin(self)
+
     def interpolate_points(self):
         """types:
         key:  user clicked / accepted frame
@@ -1163,6 +1219,8 @@ class Marker:
                         else:
                             pass
 
+                elif k & 0xFF == ord("m"):  # launch plugin module
+                    self.launch_plugin()
                 elif k & 0xFF == ord("t"):  # tracking
                     self.track_point()
                 elif k & 0xFF == ord("e"):  # point edit (width height)