Update for communication disconnection

people-counter/camera/detect.py

@@ -131,15 +131,25 @@ def run(
 
     # Dataloader
     bs = 1  # batch_size
-    if webcam:
-        view_img = check_imshow(warn=True)
-        dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
-        bs = len(dataset)
-    elif screenshot:
-        dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
-    else:
-        dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
-    vid_path, vid_writer = [None] * bs, [None] * bs
+    try:
+        if webcam:
+            view_img = check_imshow(warn=True)
+            dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
+            bs = len(dataset)
+        elif screenshot:
+            dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)
+        else:
+            dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)
+        vid_path, vid_writer = [None] * bs, [None] * bs
+
+        if dataset is None:
+            raise ValueError("Failed to load dataset")
+    except Exception as e:
+        LOGGER.error(f"Error loading stream: {e}")
+        if mqtt_server and mqtt_topic:
+            mqtt_client = MqttClient(mqtt_server, mqtt_port, mqtt_topic)
+            mqtt_client.publish("-1")  # Send -1 when an error occurs
+        sys.exit(1)
 
     # Initialize MQTT client
     if mqtt_server != '' and mqtt_topic != '':