Coverage for orchestr_ant_ion / streaming / generator.py: 0%

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 orchestr_ant_ion.streaming.capture import FrameCapture 

17 

18 

19_shutdown_event: threading.Event | None = None 

20 

21 

22def init_shutdown_event() -> None: 

23 """Initialize the global shutdown event for graceful termination.""" 

24 global _shutdown_event 

25 _shutdown_event = threading.Event() 

26 

27 

28def request_shutdown() -> None: 

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

30 global _shutdown_event 

31 if _shutdown_event is not None: 

32 _shutdown_event.set() 

33 

34 

35def is_shutdown_requested() -> bool: 

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

37 global _shutdown_event 

38 return _shutdown_event is not None and _shutdown_event.is_set() 

39 

40 

41def gen_frames( 

42 frame_capture: FrameCapture, 

43 jpeg_quality: int = 30, 

44 wait_for_frame: float = 0.1, 

45 wait_on_empty: float = 0.5, 

46) -> Iterator[bytes]: 

47 """Yield MJPEG frame chunks suitable for multipart responses. 

48 

49 This generator will continue yielding frames until shutdown is requested 

50 via request_shutdown() or the frame capture stops. 

51 

52 Args: 

53 frame_capture: The frame capture instance to read from. 

54 jpeg_quality: JPEG encoding quality (1-100). 

55 wait_for_frame: Time to wait for initial frame in seconds. 

56 wait_on_empty: Time to wait when no frame is available in seconds. 

57 

58 Yields: 

59 MJPEG frame chunks ready for multipart responses. 

60 """ 

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

62 

63 init_shutdown_event() 

64 

65 while frame_capture.frame_queue.empty(): 

66 if is_shutdown_requested(): 

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

68 return 

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

70 time.sleep(wait_for_frame) 

71 

72 while not is_shutdown_requested(): 

73 if is_shutdown_requested(): 

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

75 break 

76 

77 frame = frame_capture.get_frame() 

78 if frame is None: 

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

80 time.sleep(wait_on_empty) 

81 continue 

82 

83 ret, buffer = cv2.imencode( 

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

85 ) 

86 if not ret: 

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

88 continue 

89 frame_bytes = buffer.tobytes() 

90 

91 yield ( 

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

93 ) 

94 

95 logger.info("Video stream stopped")