Video: add screenshot function

labgrid/driver/httpvideodriver.py

@@ -1,3 +1,4 @@
+import pathlib
 import subprocess
 import sys
 from urllib.parse import urlsplit
@@ -20,15 +21,21 @@ class HTTPVideoDriver(Driver, VideoProtocol):
     def get_qualities(self):
         return ("high", [("high", None)])
 
-    @Driver.check_active
-    def stream(self, quality_hint=None):
+    def _get_default_port(self):
+        default_port = None
         s = urlsplit(self.video.url)
         if s.scheme == "http":
             default_port = 80
         elif s.scheme == "https":
             default_port = 443
         else:
             print(f"Unknown scheme: {s.scheme}", file=sys.stderr)
+            return default_port
+
+    @Driver.check_active
+    def stream(self, quality_hint=None):
+        default_port = self._get_default_port()
+        if default_port is None:
             return
 
         url = proxymanager.get_url(self.video.url, default_port=default_port)
@@ -47,3 +54,28 @@ def stream(self, quality_hint=None):
 
         sub = subprocess.run(pipeline)
         return sub.returncode
+
+    @Driver.check_active
+    def screenshot(self, filename):
+        default_port = self._get_default_port()
+        if default_port is None:
+            return
+
+        filepath = pathlib.Path(filename)
+        url = proxymanager.get_url(self.video.url, default_port=default_port)
+        pipeline = [
+            "gst-launch-1.0",
+            "souphttpsrc",
+            "num-buffers=75",
+            f"location={url}",
+            "!",
+            "decodebin",
+            "!",
+            "jpegenc",
+            "!",
+            "filesink",
+            f"location={filepath.absolute()}",
+        ]
+
+        sub = subprocess.run(pipeline)
+        return sub.returncode

labgrid/driver/usbvideodriver.py

@@ -1,3 +1,4 @@
+import pathlib
 import subprocess
 
 import attr
@@ -170,3 +171,45 @@ def stream(self, caps_hint=None, controls=None):
 
         rx.communicate()
         tx.communicate()
+
+    @Driver.check_active
+    def screenshot(self, filename, caps_hint=None, controls=None):
+        assert isinstance(filename, str)
+
+        filepath = pathlib.Path(filename)
+        caps = self.select_caps(caps_hint)
+        pipeline = self.get_pipeline(self.video.path, caps, controls)
+
+        tx_cmd = self.video.command_prefix + ["gst-launch-1.0", "-q"]
+        tx_cmd += pipeline.split()
+        rx_cmd = ["gst-launch-1.0", "fdsrc", " num-buffers=75", "!", "jpegdec", "!", "jpegenc", "!", "filesink", f"location={filepath.absolute()}"]
+
+        tx = subprocess.Popen(
+            tx_cmd,
+            stdin=subprocess.DEVNULL,
+            stdout=subprocess.PIPE,
+        )
+        rx = subprocess.Popen(
+            rx_cmd,
+            stdin=tx.stdout,
+            stdout=subprocess.DEVNULL,
+        )
+
+        # wait until one subprocess has terminated
+        while True:
+            try:
+                tx.wait(timeout=0.1)
+                break
+            except subprocess.TimeoutExpired:
+                pass
+            try:
+                rx.wait(timeout=0.1)
+                break
+            except subprocess.TimeoutExpired:
+                pass
+
+        rx.terminate()
+        tx.terminate()
+
+        rx.communicate()
+        tx.communicate()

labgrid/protocol/videoprotocol.py

@@ -9,3 +9,7 @@ def get_qualities(self):
     @abc.abstractmethod
     def stream(self, quality_hint=None):
         raise NotImplementedError
+
+    @abc.abstractmethod
+    def screenshot(self, filename):
+        raise NotImplementedError