Refactor Rtsp client

RtspClient.py

@@ -1,85 +1,47 @@
-import socket
-
+import os
 import cv2
-import numpy
-import socks
 
 
 class RtspClient:
 
-    def __init__(self, ip, username, password, port=554, profile="main", **kwargs):
+    def __init__(self, ip, username, password, port=554, profile="main", use_udp=True, **kwargs):
         """
 
-        :param ip:
-        :param username:
-        :param password:
-        :param port: rtsp port
+        :param ip: Camera IP
+        :param username: Camera Username
+        :param password: Camera User Password
+        :param port: RTSP port
         :param profile: "main" or "sub"
+        :param use_upd: True to use UDP, False to use TCP
         :param proxies: {"host": "localhost", "port": 8000}
         """
+        capture_options = 'rtsp_transport;'
         self.ip = ip
         self.username = username
         self.password = password
         self.port = port
-        self.sockt = None
-        self.url = "rtsp://" + self.username + ":" + self.password + "@" + self.ip + ":" + str(
-            self.port) + "//h264Preview_01_" + profile
         self.proxy = kwargs.get("proxies")
+        self.url = "rtsp://" + self.username + ":" + self.password + "@" + \
+            self.ip + ":" + str(self.port) + "//h264Preview_01_" + profile
+        if use_udp:
+            capture_options = capture_options + 'udp'
+        else:
+            capture_options = capture_options + 'tcp'
 
-    def __enter__(self):
-        self.sockt = self.connect()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self.sockt.close()
-
-    def connect(self) -> socket:
-        try:
-            sockt = socks.socksocket(socket.AF_INET, socket.SOCK_STREAM)
-            sockt.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-            if self.proxy is not None:
-                sockt.set_proxy(socks.SOCKS5, self.proxy["host"], self.proxy["port"])
-            sockt.connect((self.ip, self.port))
-            return sockt
-        except Exception as e:
-            print(e)
-
-    def get_frame(self) -> bytearray:
-        try:
-            self.sockt.send(str.encode(self.url))
-            data = b''
-            while True:
-                try:
-                    r = self.sockt.recv(90456)
-                    if len(r) == 0:
-                        break
-                    a = r.find(b'END!')
-                    if a != -1:
-                        data += r[:a]
-                        break
-                    data += r
-                except Exception as e:
-                    print(e)
-                    continue
-            nparr = numpy.fromstring(data, numpy.uint8)
-            frame = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
-            return frame
-        except Exception as e:
-            print(e)
+        os.environ["OPENCV_FFMPEG_CAPTURE_OPTIONS"] = capture_options
 
     def preview(self):
         """ Blocking function. Opens OpenCV window to display stream. """
-        self.connect()
-        win_name = 'RTSP'
-        cv2.namedWindow(win_name, cv2.WINDOW_AUTOSIZE)
-        cv2.moveWindow(win_name, 20, 20)
+        win_name = self.ip
+        cap = cv2.VideoCapture(self.url, cv2.CAP_FFMPEG)
+        ret, frame = cap.read()
 
-        while True:
-            cv2.imshow(win_name, self.get_frame())
-            # if self._latest is not None:
-            #    cv2.imshow(win_name,self._latest)
-            if cv2.waitKey(25) & 0xFF == ord('q'):
+        while ret:
+            cv2.imshow(win_name, frame)
+
+            ret, frame = cap.read()
+            if (cv2.waitKey(1) & 0xFF == ord('q')):
                 break
-        cv2.waitKey()
+
+        cap.release()
         cv2.destroyAllWindows()
-        cv2.waitKey()

api/recording.py

@@ -30,16 +30,15 @@ class RecordingAPIMixin:
     ###########
     # RTSP Stream
     ###########
-    def open_video_stream(self, profile: str = "main") -> Image:
+    def open_video_stream(self, profile: str = "main", proxies=None) -> None:
         """
-        profile is "main" or "sub"
-        https://support.reolink.com/hc/en-us/articles/360007010473-How-to-Live-View-Reolink-Cameras-via-VLC-Media-Player
-        :param profile:
-        :return:
+        'https://support.reolink.com/hc/en-us/articles/360007010473-How-to-Live-View-Reolink-Cameras-via-VLC-Media-Player'
+        :param profile: profile is "main" or "sub"
+        :param proxies: Default is none, example: {"host": "localhost", "port": 8000}
         """
-        with RtspClient(ip=self.ip, username=self.username, password=self.password,
-                        proxies={"host": "127.0.0.1", "port": 8000}) as rtsp_client:
-            rtsp_client.preview()
+        rtsp_client = RtspClient(
+            ip=self.ip, username=self.username, password=self.password, proxies=proxies)
+        rtsp_client.preview()
 
     def get_snap(self, timeout: int = 3, proxies=None) -> Image or None:
         """