Merge pull request #38 from groundlight/tyler/allow-specifying-stream…

…-type

Sometimes stream incorrectly determines the correct way to grab frames from the stream. For example, some internet .mjpg streams should be processed via FileStreamFrameGrabber but instead it is routed to InternetFrameGrabber due. This allows the user to manually specify how their stream should be processed.

README.md

@@ -5,7 +5,7 @@ process frames from a video file, device, or stream.
 
 ## Download
 
-This application is easy to use on any system with Docker installed. 
+This application is easy to use on any system with Docker installed.
 
 ```shell
 $ docker pull groundlight/stream
@@ -29,6 +29,7 @@ options:
   -f, --fps=FPS          number of frames to capture per second. 0 to use maximum rate possible. [default: 5]
   -h, --help             show this message.
   -s, --stream=STREAM    id, filename or URL of a video stream (e.g. rtsp://host:port/script?params OR video.mp4 OR *.jpg) [default: 0]
+  -x, --streamtype=TYPE  (optional) type of stream. One of [device, directory, rtsp, youtube, file, image_url] [default: auto-infer]
   -t, --token=TOKEN      API token to authenticate with the Groundlight API
   -v, --verbose          enable debug logs
   -w, --width=WIDTH      resize images to w pixels wide (and scale height proportionately if not set explicitly)

src/grabber.py

@@ -2,7 +2,6 @@
 import logging
 import os
 import random
-import re
 import time
 import urllib
 from abc import ABCMeta, abstractmethod
@@ -18,28 +17,30 @@
 
 class FrameGrabber(metaclass=ABCMeta):
     @staticmethod
-    def create_grabber(stream=None, **kwargs):
+    def create_grabber(stream=None, stream_type=None, **kwargs):
         logger.debug(f"Input {stream=} (type {type(stream)}")
-        if type(stream) == int:
+        if (type(stream) == int and not streamtype) or stream_type == "device":
             logger.debug("Looking for camera {stream=}")
             return DeviceFrameGrabber(stream=stream)
-        elif (type(stream) == str) and (stream.find("*") != -1):
+        elif ((type(stream) == str) and (stream.find("*") != -1) and not stream_type) or stream_type == "directory":
             logger.debug(f"Found wildcard file {stream=}")
             return DirectoryFrameGrabber(stream=stream)
-        elif (type(stream) == str) and (stream[:4] == "rtsp"):
+        elif ((type(stream) == str) and (stream[:4] == "rtsp") and not stream_type) or stream_type == "rtsp":
             logger.debug(f"found rtsp stream {stream=}")
             return RTSPFrameGrabber(stream=stream)
-        elif (type(stream) == str) and (stream.find("youtube.com") > 0):
+        elif (
+            (type(stream) == str) and (stream.find("youtube.com") > 0) and not stream_type
+        ) or stream_type == "youtube":
             logger.debug(f"found youtube stream {stream=}")
             return YouTubeFrameGrabber(stream=stream)
-        elif (type(stream) == str) and Path(stream).is_file():
+        elif ((type(stream) == str) and Path(stream).is_file() and not stream_type) or stream_type == "file":
             logger.debug(f"found filename stream {stream=}")
             return FileStreamFrameGrabber(stream=stream, **kwargs)
-        elif (type(stream) == str) and (stream[:4] == "http"):
+        elif ((type(stream) == str) and (stream[:4] == "http") and not stream_type) or stream_type == "image_url":
             logger.debug(f"found image url {stream=}")
             return ImageURLFrameGrabber(url=stream, **kwargs)
         else:
-            raise ValueError(f"cannot create a frame grabber from {stream=}")
+            raise ValueError(f"cannot create a frame grabber from {stream=} {stream_type=}")
 
     @abstractmethod
     def grab():

src/stream.py

@@ -9,6 +9,7 @@
   -f, --fps=FPS          number of frames to capture per second. 0 to use maximum rate possible. [default: 5]
   -h, --help             show this message.
   -s, --stream=STREAM    id, filename or URL of a video stream (e.g. rtsp://host:port/script?params OR movie.mp4 OR *.jpg) [default: 0]
+  -x, --streamtype=TYPE  (optional) type of stream. One of [device, directory, rtsp, youtube, file, image_url] [default: auto-infer]
   -t, --token=TOKEN      api token to authenticate with the groundlight api
   -v, --verbose          enable debug logs
   -w, --width=WIDTH      resize images to w pixels wide (and scale height proportionately if not set explicitly)
@@ -169,10 +170,17 @@ def main():
     DETECTOR = args["--detector"]
 
     STREAM = args["--stream"]
-    try:
-        STREAM = int(STREAM)
-    except ValueError as e:
-        logger.debug(f"{STREAM=} is not an int.  Treating as a filename or url.")
+    STREAM_TYPE = args.get("--streamtype")
+    if STREAM_TYPE is None:
+        try:
+            STREAM = int(STREAM)
+        except ValueError as e:
+            logger.debug(f"{STREAM=} is not an int.  Treating as a filename or url.")
+    else:
+        STREAM_TYPE = STREAM_TYPE.lower()
+        if STREAM_TYPE not in ["device", "directory", "rtsp", "youtube", "file", "image_url"]:
+            raise ValueError(f"Invalid stream type {STREAM_TYPE=}")
+        logger.debug(f"{STREAM_TYPE=}")
 
     FPS = args["--fps"]
     try:
@@ -215,7 +223,7 @@ def main():
 
     logger.debug(f"creating groundlight client with {ENDPOINT=} and {TOKEN=}")
     gl = Groundlight(endpoint=ENDPOINT, api_token=TOKEN)
-    grabber = FrameGrabber.create_grabber(stream=STREAM, fps_target=FPS)
+    grabber = FrameGrabber.create_grabber(stream=STREAM, stream_type=STREAM_TYPE, fps_target=FPS)
     q = Queue()
     tc = ThreadControl()
     if motion_detect: