From eb385abd56adfebf5779479ec330254331aed426 Mon Sep 17 00:00:00 2001
From: Aaron Parecki <aaron@parecki.com>
Date: Fri, 21 Jun 2024 16:08:20 -0700
Subject: [PATCH] add support for RTSP feeds

adds `-rtsp` option that lets you count objects from an RTSP camera instead of the raspberry pi camera
---
 umt/umt_main.py  |  1 +
 umt/umt_utils.py | 18 ++++++++++++++++++
 2 files changed, 19 insertions(+)

diff --git a/umt/umt_main.py b/umt/umt_main.py
index 136c8d2..b1bd623 100644
--- a/umt/umt_main.py
+++ b/umt/umt_main.py
@@ -38,6 +38,7 @@ def main():
     parser.add_argument('-imageseq', dest='image_path', type=str, required=False, help='specify an image sequence')
     parser.add_argument('-video', dest='video_path', type=str, required=False, help='specify video file')
     parser.add_argument('-camera', dest='camera', default=False, action='store_true', help='specify this when using the rpi camera as the input')
+    parser.add_argument('-rtsp', dest='rtsp_url', type=str, required=False, help='specify RTSP camera URL')
     parser.add_argument('-threshold', dest='threshold', type=float, default=0.5, required=False, help='specify a custom inference threshold')
     parser.add_argument('-tpu', dest='tpu', required=False, default=False, action='store_true', help='add this when using a coral usb accelerator')
     parser.add_argument('-nframes', dest='nframes', type=int, required=False, default=10, help='specify nunber of frames to process')
diff --git a/umt/umt_utils.py b/umt/umt_utils.py
index 77a2cf0..4844cd2 100644
--- a/umt/umt_utils.py
+++ b/umt/umt_utils.py
@@ -74,6 +74,22 @@ def video_frame_gen(args):
         yield Image.fromarray(frame)
 
 
+def rtsp_frame_gen(args):
+
+    cap = cv2.VideoCapture(args.rtsp_url, cv2.CAP_FFMPEG)
+
+    while True:
+        if cv2.waitKey(1) & 0xFF == ord('q'): break
+
+        # pull frame from video stream
+        _, frame = cap.read()
+
+        # array to PIL image format
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+
+        yield Image.fromarray(frame)
+
+
 def initialize_img_source(args):
 
     # track objects from video file
@@ -85,6 +101,8 @@ def initialize_img_source(args):
     # track objects from camera source
     if args.camera: return camera_frame_gen
 
+    # track objects from RTSP camera feed
+    if args.rtsp_url: return rtsp_frame_gen
 
 def initialize_detector(args):