jnhg_door/stream.py

import os
import cv2
import time
import requests
from logger import logger


class StreamCapture:
    def __init__(self, ip, channel, reconnect_interval=5):
        self.channel            = channel
        self.ip                 = ip
        self.reconnect_interval = reconnect_interval
        self.cap                = None
        self.stream_url         = None

    def get_stream_url(self):
        data = {
            "channel": str(self.channel),
            "ip":self.ip
        }
        url = requests.post(url='http://172.19.152.231/open/api/operate/previewURLs',data=data).json()['msg']
        self.stream_url = url
        return url
    
    def get_frame(self):
        if self.cap is None or not self.cap.isOpened():
            return None, False
        ret, frame = self.cap.read()

        if not ret:
            return None, False
        return frame, ret

    def reconnect(self):
        stream_url = self.get_stream_url()
        logger.info(f"Reconnect {stream_url} ...")
        self.cap.release()
        self.cap = cv2.VideoCapture(stream_url)

        if not self.cap.isOpened():
            logger.info("Reconnect failed")
            return False
        logger.info("Reconnect success")
        return True

    def close(self):
        self.cap.release()

    def __call__(self):
        stream_url = self.get_stream_url()
        logger.info(f"Connect {stream_url} ...")
        self.cap = cv2.VideoCapture(stream_url)
        if not self.cap.isOpened():
            logger.info("Connect failed")
            yield None, False
        logger.info("Connect sucess")
        
        while True:
            frame, sucess = self.get_frame()
            if not sucess:
                logger.info("Disconnect ...")
                while not self.reconnect():
                    time.sleep(self.reconnect_interval)
            yield frame, sucess
        self.cap.release()

if __name__ == "__main__":
    stream = StreamCapture()
    for frame, ret in stream():
        if not ret:
            print("Cannot Get Image")
            continue