Update docs and docstrings

Author: prouast

README.md

@@ -22,6 +22,7 @@ The library provides:
 - **High-Fidelity Accuracy:** A simple interface to the VitalLens API for state-of-the-art estimation (heart rate, respiratory rate, HRV).
 - **Local Fallbacks:** Implementations of classic rPPG algorithms (`pos`, `chrom`, `g`) for local, API-free processing.
 - **Flexible Input:** Support for video files and in-memory `np.ndarray`.
+- **Real-time Streaming:** `stream()` context manager for low-latency live inference.
 - **Face Detection:** Integrated fast face detection and ROI management.
 
 Using a different language? Check out our [JavaScript client](https://github.com/Rouast-Labs/vitallens.js) and [iOS SDK](https://github.com/Rouast-Labs/vitallens-ios).
@@ -53,7 +54,7 @@ print("Heart Rate:", results[0]['vitals']['heart_rate']['value'])
 To get improved accuracy and advanced metrics like **Respiratory Rate** and **HRV**, use the `vitallens` method. You can get a free key from the [API Dashboard](https://www.rouast.com/api).
 
 ```python
-from vitallens import VitalLens, Method
+from vitallens import VitalLens
 
 # Automatically selects the best model for your plan
 vl = VitalLens(method="vitallens", api_key="YOUR_API_KEY")
@@ -68,6 +69,23 @@ print(f"Respiratory Rate: {vitals['respiratory_rate']['value']:.1f} rpm")
 if 'hrv_sdnn' in vitals:
     print(f"HRV (SDNN):       {vitals['hrv_sdnn']['value']:.1f} ms")
 ```
+
+### Real-time Streaming
+
+Process live frames from a webcam or stream.
+
+```python
+import time
+from vitallens import VitalLens
+
+# Process live frames
+vl = VitalLens(method="vitallens", api_key="YOUR_API_KEY")
+
+with vl.stream() as session:
+    # In your capture loop (e.g., OpenCV)
+    session.push(frame, timestamp=time.time())
+    results = session.get_result(block=False)
+```
 <!-- mkdocs-end -->
 
 ## Documentation

examples/README.md

@@ -97,6 +97,51 @@ video_arr = np.array(frames)
 results = vl(video_arr, fps=fps)
 ```
 
+### Real-time Streaming
+
+For live feeds or webcams. There are two ways to handle results:
+
+#### Polling (Non-blocking)
+
+Best for applications with their own main loop (e.g., OpenCV display).
+
+```python
+import time
+from vitallens import VitalLens
+
+vl = VitalLens(method="vitallens", api_key="YOUR_API_KEY")
+
+with vl.stream() as session:
+    while True:
+        frame, ts = get_frame() # Your capture logic
+        session.push(frame, timestamp=ts)
+        
+        # Check for results whenever you want
+        results = session.get_result(block=False)
+        if results:
+            print(results[0]['vitals']['heart_rate']['value'])
+```
+
+#### Callback
+
+Best for event-driven applications. The callback is triggered automatically as soon as inference finishes.
+
+```python
+import time
+from vitallens import VitalLens
+
+def my_callback(results):
+    print(f"Callback received HR: {results[0]['vitals']['heart_rate']['value']}")
+
+vl = VitalLens(method="vitallens", api_key="YOUR_API_KEY")
+
+with vl.stream(on_result=my_callback) as session:
+    while True:
+        frame, ts = get_frame()
+        session.push(frame, timestamp=ts)
+        # No need to call get_result()
+```
+
 ## Running with Docker
 
 If you encounter dependency issues (e.g., with `onnxruntime` or `ffmpeg`), you can run the example scripts inside our Docker container.

vitallens/client.py

@@ -250,8 +250,36 @@ def __call__(
     return results
 
   def stream(self, on_result: Callable = None) -> StreamSession:
-    """
-    Returns a StreamSession context manager for real-time video processing.
+    """Returns a context manager for real-time vital sign estimation.
+
+    This method creates a `StreamSession` that manages background inference threads,
+    sliding window buffers, and signal state via `vitallens-core`. It is designed 
+    for low-latency applications like webcam feeds.
+
+    Usage:
+      ```python
+      with vl.stream() as session:
+          session.push(frame, timestamp)
+          results = session.get_result(block=False)
+      ```
+
+    Args:
+      on_result: An optional callback function triggered automatically whenever 
+        new inference results are available. The function should accept one 
+        argument (the results list).
+
+    Returns:
+      session: A `StreamSession` context manager. The session object provides:
+        * `push(frame, timestamp, face=None)`: Ingests a new RGB frame. 
+          `timestamp` should be in seconds (e.g., `time.time()`).
+        * `get_result(block=False, timeout=None)`: Pulls the latest analysis 
+          results from the queue.
+        * `current_face`: The most recently detected face coordinates.
+
+    Note:
+      The results returned by `get_result()` or the callback follow the same 
+      format as `__call__`, but represent the physiological state of the 
+      current sliding window rather than a global file average.
     """
     return StreamSession(
       rppg_method=self.rppg,

vitallens/methods/simple_rppg_method.py

@@ -123,7 +123,19 @@ def infer_stream(
       fps: float,
       state
     ):
-    """TODO"""
+    """Estimate pulse signal from a sequence of frames in a streaming context.
+
+    Args:
+      frames: The input video frames of shape (n_frames, h, w, 3).
+      fps: The sampling frequency of the input frames.
+      state: The internal state of the rPPG method (unused for simple methods).
+    Returns:
+      Tuple of
+        - sig_dict: A dictionary of the estimated signals.
+        - conf_dict: A dictionary of the estimated confidences.
+        - live_out: Dummy live confidence estimation (set to always 1). Shape (n_frames,)
+        - state: The updated internal state of the rPPG method (None).
+    """
     sig = self.algorithm(frames, fps)
     sig_dict = {'ppg_waveform': sig}
     conf_dict = {'ppg_waveform': np.ones_like(sig)}

vitallens/methods/vitallens.py

@@ -223,7 +223,19 @@ def stack_dicts(dict_list):
     return sig_dict, conf_dict, live
 
   def infer_stream(self, frames: np.ndarray, fps: float, state=None):
-    """TODO"""
+    """Estimate vitals from a sequence of frames using the VitalLens streaming API.
+
+    Args:
+      frames: The input video frames of shape (n_frames, h, w, 3).
+      fps: The sampling frequency of the input frames.
+      state: The internal state of the rPPG method used to maintain temporal continuity.
+    Returns:
+      Tuple of
+        - sig_dict: A dictionary of the estimated signals.
+        - conf_dict: A dictionary of the estimated confidences.
+        - live: The face live confidence. Shape (n_frames,)
+        - new_state: The updated internal state of the rPPG method.
+    """
     headers = {
       "Content-Type": "application/octet-stream",
       "X-Encoding": "gzip"

vitallens/stream.py

@@ -39,6 +39,14 @@ class InferenceContext:
 
 class FrameBuffer:
   def __init__(self, buffer_id: str, roi: vc.Rect, max_capacity: int, timestamp: float):
+    """Initialize a thread-safe frame buffer for a specific region of interest.
+
+    Args:
+      buffer_id: Unique identifier for the buffer.
+      roi: The region of interest (ROI) associated with this buffer.
+      max_capacity: The maximum number of frames the buffer can hold.
+      timestamp: The creation timestamp of the buffer.
+    """
     self.id = buffer_id
     self.roi = roi
     self.created_at = timestamp
@@ -51,13 +59,27 @@ def count(self) -> int:
     return len(self.buffer)
 
   def append(self, frame: np.ndarray, context: InferenceContext):
+    """Add a frame and its metadata to the buffer, maintaining max capacity.
+
+    Args:
+      frame: The parsed RGB image data.
+      context: Metadata including timestamp and ROI.
+    """
     self.buffer.append((frame, context))
     self.last_seen = context.timestamp
     if len(self.buffer) > self.max_capacity:
       overflow = len(self.buffer) - self.max_capacity
       self.buffer = self.buffer[overflow:]
 
   def execute(self, take_count: int, keep_count: int) -> list:
+    """Extract frames for inference and manage buffer overlap.
+
+    Args:
+      take_count: Number of frames to extract for the current batch.
+      keep_count: Number of frames to retain for the next sliding window.
+    Returns:
+      payload: A list of (frame, context) tuples, or None if insufficient frames.
+    """
     if take_count <= 0 or len(self.buffer) < take_count:
       return None
     payload = self.buffer[:take_count]
@@ -66,8 +88,13 @@ def execute(self, take_count: int, keep_count: int) -> list:
     return payload\
 
 class BufferManager:
-  """Thread-safe buffer manager."""
+  """Thread-safe manager for multiple streaming frame buffers."""
   def __init__(self, buffer_config: vc.BufferConfig):
+    """Initialize the manager with specific buffering and overlap constraints.
+
+    Args:
+      buffer_config: Configuration for stream window sizes and overlaps.
+    """
     self.buffer_planner = vc.BufferPlanner(buffer_config)
     self.buffer_config = buffer_config
     self.buffers = {}
@@ -84,6 +111,14 @@ def _get_active_metadata(self):
     ]
 
   def register_target(self, target_roi: vc.Rect, timestamp: float):
+    """Register or update a target ROI in the buffer planner.
+
+    Args:
+      target_roi: The latest detected face ROI to track.
+      timestamp: The current frame timestamp.
+    Returns:
+      buffer_id: The ID of the buffer assigned to this target, or None.
+    """
     with self.lock:
       self.current_timestamp = max(self.current_timestamp, timestamp)
       action = self.buffer_planner.evaluate_target(target_roi, timestamp, self._get_active_metadata())
@@ -100,12 +135,26 @@ def register_target(self, target_roi: vc.Rect, timestamp: float):
       return None
 
   def append(self, buffer_id: str, frame: np.ndarray, context: InferenceContext):
+    """Route a frame to a specific active buffer.
+
+    Args:
+      buffer_id: The identifier for the target buffer.
+      frame: The parsed RGB image data.
+      context: Metadata including timestamp and ROI.
+    """
     with self.lock:
       self.current_timestamp = max(self.current_timestamp, context.timestamp)
       if buffer_id in self.buffers:
         self.buffers[buffer_id].append(frame, context)
 
   def poll(self, flush: bool = False):
+    """Poll the planner for the next required inference command.
+
+    Args:
+      flush: If True, forces the planner to return commands for remaining data.
+    Returns:
+      command: A vc.InferenceCommand containing the buffer ID and frame counts.
+    """
     with self.lock:
       has_state = self.state is not None
       plan = self.buffer_planner.poll(
@@ -121,30 +170,62 @@ def poll(self, flush: bool = False):
       return plan.command
 
   def execute(self, command: vc.InferenceCommand) -> list:
+    """Execute a specific inference command by extracting data from the relevant buffer.
+
+    Args:
+      command: The command containing buffer ID and frame counts.
+    Returns:
+      window: A list of frames and contexts ready for inference.
+    """
     with self.lock:
       if command.buffer_id in self.buffers:
         return self.buffers[command.buffer_id].execute(command.take_count, command.keep_count)
       return None
 
   def get_state(self):
+    """Retrieve the current physiological state of the rPPG model.
+
+    Returns:
+      state: The model's internal state vector or None.
+    """
     with self.lock:
       return self.state
 
   def update_state(self, new_state):
+    """Update the internal state of the rPPG model after an inference step.
+
+    Args:
+      new_state: The updated state vector returned by the API.
+    """
     with self.lock:
       self.state = new_state
 
   def get_all_buffers(self):
+    """Get a list of all currently active buffers.
+
+    Returns:
+      buffers: A list of (buffer_id, roi) tuples.
+    """
     with self.lock:
       return [(buf.id, buf.roi) for buf in self.buffers.values()]
 
   def reset(self):
+    """Clear all active buffers and reset the model state."""
     with self.lock:
       self.buffers.clear()
       self.state = None
 
 class StreamSession:
+  """Context manager handling background inference and buffer synchronization for live streams."""
   def __init__(self, rppg_method, face_detector=None, fdet_fs=1.0, on_result: Optional[Callable] = None):
+    """Initialize the streaming session and start the background inference thread.
+
+    Args:
+      rppg_method: The method instance (e.g., VitalLensRPPGMethod) to use.
+      face_detector: Optional instance of FaceDetector for automatic tracking.
+      fdet_fs: Frequency [Hz] at which to run automatic face detection.
+      on_result: Optional callback for asynchronous result handling.
+    """
     self.rppg = rppg_method
     self.face_detector = face_detector
     self.fdet_fs = fdet_fs
@@ -167,12 +248,19 @@ def __exit__(self, exc_type, exc_val, exc_tb):
     self.close()
 
   def close(self):
+    """Stop the background thread and shut down the streaming session."""
     self.running = False
     if self.thread.is_alive():
       self.thread.join(timeout=2.0)
 
   def push(self, frame: np.ndarray, timestamp: float, face: np.ndarray = None):
-    """Pushes a single frame to the ingestion pipeline on the main thread."""
+    """Push a single frame and its timestamp into the streaming pipeline.
+
+    Args:
+      frame: The input video frame of shape (h, w, 3) in RGB format.
+      timestamp: The absolute timestamp of the frame in seconds.
+      face: Optional pre-detected face box [x0, y0, x1, y1].
+    """
     if not self.running: return
     # Throttled face detection
     min_interval = 1.0 / self.rppg.fps_target
@@ -235,14 +323,21 @@ def push(self, frame: np.ndarray, timestamp: float, face: np.ndarray = None):
       self.buffer_manager.append(buf_id, payload, ctx)
 
   def get_result(self, block=False, timeout=None):
-    """Pulls the latest result from the queue."""
+    """Pull the latest inference results from the session queue.
+
+    Args:
+      block: Whether to block until a result is available.
+      timeout: Maximum time to block if block is True.
+    Returns:
+      results: A list of analysis results for detected faces.
+    """
     try:
       return self.result_queue.get(block=block, timeout=timeout)
     except queue.Empty:
       return None
 
   def _inference_loop(self):
-    """Background thread executing inference without blocking the webcam."""
+    """Background loop that polls for ready buffers and executes model inference."""
     while self.running:
       command = self.buffer_manager.poll()
       if not command:
@@ -279,7 +374,13 @@ def _inference_loop(self):
         self.vc_session.reset()
 
   def _format_result(self, session_result):
-    # TODO: Is there a smarter way to do this?
+    """Format the raw session result into the standard vitallens-python result dictionary.
+
+    Args:
+      session_result: The vc.SessionResult object from vitallens-core.
+    Returns:
+      res_dict: A list containing a formatted results dictionary for the face.
+    """
     face_dict = {
       'coordinates': [], 
       'confidence': [],