Coverage for orchestrant/streaming/generator.py: 29%

1"""Helpers for streaming encoded frames.""" 

2 

3from __future__ import annotations 

4 

5import threading 

6import time 

7from typing import TYPE_CHECKING 

8 

9import cv2 

10from loguru import logger 

11 

12 

13if TYPE_CHECKING: 

14 from collections.abc import Iterator 

15 

16 from orchestrant.streaming.capture import FrameCapture 

17 

18 

19# Re-armed, never rebound, so no function needs a `global` statement. 

20_shutdown_event = threading.Event() 

21 

22 

23def init_shutdown_event() -> None: 

24 """Re-arm the shutdown event for a fresh streaming session.""" 

25 _shutdown_event.clear() 

26 

27 

28def request_shutdown() -> None: 

29 """Request shutdown of all active frame generators.""" 

30 _shutdown_event.set() 

31 

32 

33def is_shutdown_requested() -> bool: 

34 """Check if shutdown has been requested.""" 

35 return _shutdown_event.is_set() 

36 

37 

38def gen_frames( 

39 frame_capture: FrameCapture, 

40 jpeg_quality: int = 30, 

41 wait_for_frame: float = 0.1, 

42 wait_on_empty: float = 0.5, 

43) -> Iterator[bytes]: 

44 """Yield MJPEG multipart chunks until shutdown is requested or capture stops.""" 

45 logger.info("Starting video stream...") 

46 

47 init_shutdown_event() 

48 

49 while frame_capture.frame_queue.empty(): 

50 if is_shutdown_requested(): 

51 logger.info("Shutdown requested during initial frame wait") 

52 return 

53 logger.debug("Waiting for the first frame...") 

54 time.sleep(wait_for_frame) 

55 

56 while not is_shutdown_requested(): 

57 if is_shutdown_requested(): 

58 logger.info("Shutdown requested, stopping stream") 

59 break 

60 

61 frame = frame_capture.get_frame() 

62 if frame is None: 

63 logger.warning("No frame available; skipping frame.") 

64 time.sleep(wait_on_empty) 

65 continue 

66 

67 ret, buffer = cv2.imencode( 

68 ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, jpeg_quality] 

69 ) 

70 if not ret: 

71 logger.warning("Frame encoding failed; skipping frame...") 

72 continue 

73 frame_bytes = buffer.tobytes() 

74 

75 yield ( 

76 b"--frame\r\nContent-Type: image/jpeg\r\n\r\n" + frame_bytes + b"\r\n\r\n" 

77 ) 

78 

79 logger.info("Video stream stopped")