Refactoring and improvements

.gitignore

@@ -1,2 +1,3 @@
 env
 *.swp
+*.pyc

audiobuffer.py

@@ -3,6 +3,7 @@
 import numpy as np
 from pyloudnorm import Meter
 from av import AudioFrame
+from logger import log
 
 BUFFER_LENGTH_SECONDS = 1
 SAMPLE_RATE = 48000
@@ -57,11 +58,11 @@ def lufs(self):
         """ Returns the LUFS value for the given channel. """
 
         if (len(self.left_buffer) < BUFFER_SIZE or len(self.right_buffer) < BUFFER_SIZE):
-            print(
+            log.info(
                 f"Buffer {len(self.left_buffer)} too small to calculate LUFS")
             return 0
 
-        print("Calculating LUFS")
+        log.debug("Calculating LUFS")
         left_array = np.array(self.left_buffer)
         right_array = np.array(self.right_buffer)
         return self.meter.integrated_loudness(np.array([left_array, right_array]).T)

logger.py

@@ -0,0 +1,54 @@
+import structlog
+import logging
+import os
+import sys
+
+if sys.stdout.isatty():
+    structlog.configure(
+        processors=[
+            structlog.stdlib.add_logger_name,
+            structlog.stdlib.add_log_level,
+            structlog.stdlib.PositionalArgumentsFormatter(),
+            structlog.processors.TimeStamper(fmt="%Y-%m-%d %H:%M.%S"),
+            structlog.processors.StackInfoRenderer(),
+            structlog.dev.ConsoleRenderer()
+        ],
+        context_class=dict,
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        wrapper_class=structlog.stdlib.BoundLogger,
+        cache_logger_on_first_use=True,
+    )
+else:
+    structlog.configure(
+        processors=[
+            structlog.stdlib.filter_by_level,
+            structlog.processors.StackInfoRenderer(),
+            structlog.processors.format_exc_info,
+            structlog.processors.JSONRenderer()
+        ],
+        context_class=dict,
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        wrapper_class=structlog.stdlib.BoundLogger,
+        cache_logger_on_first_use=True,
+    )
+
+log = structlog.get_logger('ts_probe')
+
+# Get log level from environment variable
+
+
+def get_log_level():
+    VALID_LOG_LEVELS = ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']
+    log_level = os.environ.get('LOG_LEVEL', None)
+
+    if log_level is None:
+        return logging.INFO
+
+    if log_level.upper() not in VALID_LOG_LEVELS:
+        log.error("Valid levels are: %s" % VALID_LOG_LEVELS)
+        raise ValueError("Invalid log level: %s" % log_level)
+
+    return logging.getLevelName(log_level.upper())
+
+
+log.setLevel(get_log_level())

requirements.txt

@@ -3,3 +3,5 @@ av
 numpy
 prometheus_client
 pyloudnorm
+structlog
+rich

ts_probe.py

@@ -7,34 +7,34 @@
 from audiobuffer import AudioBuffer
 from videobuffer import VideoBuffer
 from prometheus import Prometheus
+from logger import log
 
 prom = Prometheus()
+DEFAULT_VIDEO_URL = 'http://simula.frikanalen.no:9094/frikanalen.ts'
+PROMETHEUS_PORT = 8000
+BUFFER_SIZE = 100
 
 # MPEG-2 transport stream URL
-url = os.environ.get(
-    'VIDEO_URL', 'http://simula.frikanalen.no:9094/frikanalen.ts')
-
+url = os.environ.get('VIDEO_URL', None)
 
-# Parameters for circular buffer
-BUFFER_SIZE = 100
-video_brightness_buffer = deque(maxlen=BUFFER_SIZE)
-audio_amplitude_buffer = deque(maxlen=BUFFER_SIZE)
-motion_buffer = deque(maxlen=BUFFER_SIZE)
+if url is None:
+    url = DEFAULT_VIDEO_URL
+    log.warning("No video URL specified, using default: %s" % url)
 
+audio_buffer = AudioBuffer()
+video_buffer = VideoBuffer()
+padding_packet_count: int = 0
+total_bytes_received: int = 0
+START_TIME = time.time()
 
 # Start Prometheus HTTP server
-prom.listen(8000)
+prom.listen(PROMETHEUS_PORT)
 
+log.info("Opening stream: %s" % url)
 stream = av.open(url)
 
-padding_packet_count: int = 0
-START_TIME = time.time()
-total_bytes_received: int = 0
-
-audio_buffer = AudioBuffer()
-video_buffer = VideoBuffer()
-
 while True:
+    log.info("Stream is open")
     try:
         for frame in stream.decode():
             if isinstance(frame, VideoFrame):
@@ -49,6 +49,14 @@
                 prom.audio_amplitude_dbfs_gauge.set(audio_buffer.dbfs(0))
 
     except av.AVError as e:
-        print(e)
+        log.error(e)
         prom.decode_error_count.inc()
         stream = av.open(url)
+
+    except KeyboardInterrupt:
+        log.info("Keyboard interrupt")
+        break
+
+    except Exception as e:
+        log.error(e)
+        break

videobuffer.py

@@ -1,46 +1,38 @@
 """ Video analysis """
-import logging
+from logger import log
 from collections import deque
 import numpy as np
+import numpy.typing as npt
 from av import VideoFrame
+from itertools import islice
 
-log = logging.getLogger(__name__)
+BUFFER_SIZE = 10
+FRAME_RATE = 50
 
 
-def _calculate_motion(buffer: deque[VideoFrame]):
+def _motion(buffer: deque[npt.NDArray[np.uint8]]):
     """ Returns the average motion between the given frames. """
     if len(buffer) < 2:
         log.warning("Buffer too small to calculate motion")
         return 0
 
-    prev_frame = buffer[0]
-    abs_diff = 0
-
-    for frame in buffer:
-        abs_diff += np.abs(prev_frame.to_ndarray().astype(
-            int) - frame.to_ndarray().astype(int))
-
-        prev_frame = frame
-
-    return np.mean(abs_diff/len(buffer)) / 255
+    frame_diffs = [
+        np.mean(abs(buffer[i].astype(np.int16) - frame).astype(np.uint8))
+        for i, frame in enumerate(islice(buffer, 1, None))
+    ]
 
+    # Return the mean of the absolute differences
+    return np.mean(frame_diffs) / 255
 
-# Function to calculate average video brightness
 
-
-def calculate_average_brightness(buffer: deque[VideoFrame]):
+def _avg_brightness(buffer: deque[npt.NDArray[np.uint8]]):
     """ Returns the average brightness of the given frames. """
 
-    # Return 0 and log a warning if the deque is empty
     if len(buffer) == 0:
         log.warning("Buffer empty, returning 0")
         return 0
 
-    return np.mean([frame.planes[0] for frame in buffer])
-
-
-BUFFER_SIZE = 10
-FRAME_RATE = 50
+    return np.mean(buffer) / 255
 
 
 class VideoBuffer():
@@ -50,14 +42,15 @@ class VideoBuffer():
 
     def append(self, frame: VideoFrame):
         """ Appends a frame to the buffer. """
-        self.video_buffer.append(frame.reformat(format='gray8'))
+        self.video_buffer.append(frame.reformat(
+            format='gray8').to_ndarray().astype(np.uint8))
 
     @property
     def avg_brightness(self):
         """ Returns the average brightness of the frames in the buffer. """
-        return calculate_average_brightness(self.video_buffer)
+        return _avg_brightness(self.video_buffer)
 
     @property
     def motion(self):
         """ Returns the average motion between the frames in the buffer. """
-        return _calculate_motion(self.video_buffer)
+        return _motion(self.video_buffer)