[Fix] Fix endless poll from ffmpeg subprocess (#2)

1. Updated the read_img_from_rtsp_ffmpeg function in the rtsp_reader.py file to handle exceptions when the RTSP stream shape is obtained but no data is available. This prevents the code from getting stuck in an endless loop.
Added a timeout of 10 seconds for the communicate() method in the ffmpeg subprocess to prevent it from hanging indefinitely.
2. Moved the main program logic to the main.py file for better organization and readability.
3. Made various changes in the stream_helper.py file, including:
Initialized the state attributes.
Implemented a missing count mechanism to track the number of times the cat is missing from the scene.
Switched to the default scene if the cat is missing for a specified tolerance limit.
Improved error logging and scene switching logic.

cat_stream/rtsp_reader.py

@@ -85,7 +85,13 @@ def read_img_from_rtsp_ffmpeg(
     process = subprocess.Popen(
         ffmpeg_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
     # read str from stderr
-    out, err = process.communicate()
+    try:
+        out, err = process.communicate(timeout=10)
+    except subprocess.TimeoutExpired:
+        process.kill()
+        msg = 'Timeout when reading the frame from the RTSP stream.'
+        logger.error(msg)
+        return None
     err_str = err.decode('utf-8')
     resolution_match = re.search(r', (\d+)x(\d+),', err_str)
     if resolution_match:
@@ -100,7 +106,12 @@ def read_img_from_rtsp_ffmpeg(
     raw_frame = out
     # Convert the raw frame data to a numpy array
     frame_array = np.frombuffer(raw_frame, dtype=np.uint8)
-    frame = frame_array.reshape((height, width, 3))
+    try:
+        frame = frame_array.reshape((height, width, 3))
+    except ValueError:
+        msg = 'Failed to reshape the frame data to an image.'
+        logger.error(msg)
+        return None
     # Check if the FFmpeg process has finished
     if process.poll() is not None:
         pass

cat_stream/stream_helper.py

@@ -82,10 +82,11 @@ def __init__(self,
             self.default_scene_name = first_scene_name
         # init state attr
         self.current_scene_name = self.obs_client.get_current_scene_name()
-        self.missing_count = 0
+        self.missing_count = 2
 
     def start(self) -> None:
         """Start the main loop of the stream helper."""
+        self.logger.info('Start the main loop of the stream helper.')
         while True:
             lasted_frames = dict()
             for scene_name, scene_dict in self.obs_scenes.items():
@@ -108,12 +109,17 @@ def start(self) -> None:
             if len(cat_seen_scenes) <= 0:
                 self.missing_count += 1
                 if self.missing_count >= self.missing_tolerance:
-                    msg = f'Cat missing for {self.missing_count} times, ' +\
-                        f'switch to default scene {self.default_scene_name}.'
-                    self.logger.info(msg)
-                    self.obs_client.set_current_scene(self.default_scene_name)
-                    self.current_scene_name = self.default_scene_name
+                    tgt_scene_name = self.default_scene_name
+                    missing_count_tmp = self.missing_count
                     self.missing_count = 0
+                    if self.current_scene_name != tgt_scene_name:
+                        msg = f'Cat missing for {missing_count_tmp} times, ' +\
+                            'switch to default scene ' +\
+                            f'{self.default_scene_name}.'
+                        self.logger.info(msg)
+                        self.obs_client.set_current_scene(
+                            self.default_scene_name)
+                        self.current_scene_name = self.default_scene_name
             else:
                 self.missing_count = 0
                 tgt_scene_name = cat_seen_scenes[0]

cat_stream/version.py

@@ -1,4 +1,4 @@
-__version__ = '0.1.0'
+__version__ = '0.1.1'
 
 
 def parse_version_info(version_str):

configs/default_config.py

@@ -0,0 +1,13 @@
+type = 'StreamHelper'
+obs_ws_url = 'ws://192.168.111.111:4455'
+obs_ws_pwd = 'qqdIciF8FPNfzrqX'
+cat_det_cfg = dict(type='YOLOv5CatDetection', device='cpu')
+rtsp_reader_backend = 'ffmpeg'
+obs_scenes = dict(
+    balcony=dict(default=True, media_source='MeizuX'),
+    kitchen_lo=dict(media_source='Mi8SE'),
+    kitchen_hi=dict(media_source='Redmi4'),
+    entry_hi=dict(media_source='MiMax'),
+    entry_lo=dict(media_source='Nokia6'))
+detect_interval = 10
+missing_tolerance = 6

configs/yolov5_ffmpeg_5scenes.py

@@ -0,0 +1,13 @@
+type = 'StreamHelper'
+obs_ws_url = 'ws://192.168.111.111:4455'
+obs_ws_pwd = 'qqdIciF8FPNfzrqX'
+cat_det_cfg = dict(type='YOLOv5CatDetection', device='cpu')
+rtsp_reader_backend = 'ffmpeg'
+obs_scenes = dict(
+    balcony=dict(default=True, media_source='MeizuX'),
+    kitchen_lo=dict(media_source='Mi8SE'),
+    kitchen_hi=dict(media_source='Redmi4'),
+    entry_hi=dict(media_source='MiMax'),
+    entry_lo=dict(media_source='Nokia6'))
+detect_interval = 10
+missing_tolerance = 6

tools/main.py → main.py

@@ -30,7 +30,7 @@ def setup_parser():
         '--config_path',
         help='Path to the configuration file.',
         type=str,
-        default='configs/yolov5_ffmpeg_3scenes.py')
+        default='configs/default_config.py')
     parser.add_argument(
         '--obs_ws_pwd', help='Password of the OBS websocket.', required=False)
     return parser.parse_args()