OTVision 0.6.7__py3-none-any.whl → 0.6.8__py3-none-any.whl

OTVision/detect/rtsp_input_source.py

@@ -1,6 +1,8 @@
+import socket
 from datetime import datetime, timedelta
 from time import sleep
 from typing import Generator
+from urllib.parse import urlparse
 
 from cv2 import (
     CAP_PROP_FRAME_HEIGHT,
@@ -19,6 +21,7 @@ from OTVision.application.config import (
     StreamConfig,
 )
 from OTVision.application.configure_logger import logger
+from OTVision.application.event.new_video_start import NewVideoStartEvent
 from OTVision.application.get_current_config import GetCurrentConfig
 from OTVision.detect.detected_frame_buffer import FlushEvent
 from OTVision.domain.frame import Frame
@@ -26,7 +29,7 @@ from OTVision.domain.input_source_detect import InputSourceDetect
 from OTVision.domain.time import DatetimeProvider
 
 RTSP_URL = "rtsp://127.0.0.1:8554/test"
-RETRY_SECONDS = 1
+RETRY_SECONDS = 5
 DEFAULT_READ_FAIL_THRESHOLD = 5
 
 
@@ -79,13 +82,16 @@ class RtspInputSource(InputSourceDetect):
 
     def __init__(
         self,
-        subject: Subject[FlushEvent],
+        subject_flush: Subject[FlushEvent],
+        subject_new_video_start: Subject[NewVideoStartEvent],
         datetime_provider: DatetimeProvider,
         frame_counter: Counter,
         get_current_config: GetCurrentConfig,
         read_fail_threshold: int = DEFAULT_READ_FAIL_THRESHOLD,
     ) -> None:
-        super().__init__(subject)
+
+        self.subject_flush = subject_flush
+        self.subject_new_video_start = subject_new_video_start
         self._datetime_provider = datetime_provider
         self._stop_capture = False
         self._frame_counter = frame_counter
@@ -123,41 +129,50 @@ class RtspInputSource(InputSourceDetect):
     def produce(self) -> Generator[Frame, None, None]:
         self._stream_start_time = self._datetime_provider.provide()
         self._current_video_start_time = self._stream_start_time
-        while not self.should_stop():
-            if (frame := self._read_next_frame()) is not None:
-                self._frame_counter.increment()
-                occurrence = self._datetime_provider.provide()
-
-                if self._outdated:
-                    self._current_video_start_time = occurrence
-                    self._outdated = False
-
-                yield Frame(
-                    data=convert_frame_to_rgb(frame),  # YOLO expects RGB
-                    frame=self.current_frame_number,
-                    source=self.rtsp_url,
-                    output=self.create_output(),
-                    occurrence=occurrence,
-                )
-                if self.flush_condition_met():
-                    self._notify()
-                    self._outdated = True
-                    self._frame_counter.reset()
-
-        self._notify()
+        try:
+            while not self.should_stop():
+                if (frame := self._read_next_frame()) is not None:
+                    self._frame_counter.increment()
+                    occurrence = self._datetime_provider.provide()
+
+                    if self._outdated:
+                        self._current_video_start_time = occurrence
+                        self._outdated = False
+                        self._notify_new_video_start_observers()
+
+                    yield Frame(
+                        data=convert_frame_to_rgb(frame),  # YOLO expects RGB
+                        frame=self.current_frame_number,
+                        source=self.rtsp_url,
+                        output=self.create_output(),
+                        occurrence=occurrence,
+                    )
+                    if self.flush_condition_met():
+                        self._notify_flush_observers()
+                        self._outdated = True
+                        self._frame_counter.reset()
+            self._notify_flush_observers()
+        except InvalidRtspUrlError as cause:
+            logger().error(cause)
 
     def _init_video_capture(self, source: str) -> VideoCapture:
+        self._wait_for_connection(source)
+
         cap = VideoCapture(source)
         while not self.should_stop() and not cap.isOpened():
-            logger().warning(
-                f"Couldn't open the RTSP stream: {source}. "
-                f"Trying again in {RETRY_SECONDS}s..."
-            )
-            sleep(RETRY_SECONDS)
             cap.release()
+            self._wait_for_connection(source)
             cap = VideoCapture(source)
         return cap
 
+    def _wait_for_connection(self, connection: str) -> None:
+        while not self.should_stop() and not is_connection_available(connection):
+            logger().debug(
+                f"Couldn't open the RTSP stream: {connection}. "
+                f"Trying again in {RETRY_SECONDS}s..."
+            )
+            sleep(RETRY_SECONDS)
+
     def _read_next_frame(self) -> ndarray | None:
         successful, frame = self._video_capture.read()
         if successful:
@@ -189,9 +204,9 @@ class RtspInputSource(InputSourceDetect):
     def flush_condition_met(self) -> bool:
         return self.current_frame_number % self.flush_buffer_size == 0
 
-    def _notify(self) -> None:
-        frame_width = int(self._video_capture.get(CAP_PROP_FRAME_WIDTH))
-        frame_height = int(self._video_capture.get(CAP_PROP_FRAME_HEIGHT))
+    def _notify_flush_observers(self) -> None:
+        frame_width = self._get_width()
+        frame_height = self._get_height()
         frames = (
             self.flush_buffer_size
             if self.current_frame_number % self.flush_buffer_size == 0
@@ -199,7 +214,7 @@ class RtspInputSource(InputSourceDetect):
         )
         duration = timedelta(seconds=round(frames / self.fps))
         output = self.create_output()
-        self._subject.notify(
+        self.subject_flush.notify(
             FlushEvent.create(
                 source=self.rtsp_url,
                 output=output,
@@ -211,6 +226,21 @@ class RtspInputSource(InputSourceDetect):
             )
         )
 
+    def _get_width(self) -> int:
+        return int(self._video_capture.get(CAP_PROP_FRAME_WIDTH))
+
+    def _get_height(self) -> int:
+        return int(self._video_capture.get(CAP_PROP_FRAME_HEIGHT))
+
+    def _notify_new_video_start_observers(self) -> None:
+        event = NewVideoStartEvent(
+            output=self.create_output(),
+            width=self._get_width(),
+            height=self._get_height(),
+            fps=self.fps,
+        )
+        self.subject_new_video_start.notify(event)
+
     def create_output(self) -> str:
         output_filename = (
             f"{self.stream_config.name}_FR{round(self.fps)}"
@@ -221,3 +251,58 @@ class RtspInputSource(InputSourceDetect):
 
 def convert_frame_to_rgb(frame: ndarray) -> ndarray:
     return cvtColor(frame, COLOR_BGR2RGB)
+
+
+class InvalidRtspUrlError(Exception):
+    """Raised when the RTSP URL is invalid."""
+
+
+def is_connection_available(rtsp_url: str) -> bool:
+    """
+    Check if RTSP connection is available by sending a DESCRIBE request.
+
+    Args:
+        rtsp_url: The RTSP URL to check
+
+    Returns:
+        bool: True if stream is available, False otherwise
+    """
+    try:
+        parsed = urlparse(rtsp_url)
+        if parsed.hostname is None and parsed.port is None:
+            raise InvalidRtspUrlError(
+                f"Invalid RTSP URL: {rtsp_url}. Missing hostname or port."
+            )
+
+        host = parsed.hostname
+        port = parsed.port
+
+        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        sock.settimeout(5)
+
+        if sock.connect_ex((host, port)) != 0:
+            sock.close()
+            return False
+
+        # Send RTSP DESCRIBE request to get stream info
+        rtsp_request = (
+            f"DESCRIBE {rtsp_url} RTSP/1.0\r\n"
+            f"CSeq: 1\r\n"
+            f"Accept: application/sdp\r\n\r\n"
+        )
+        sock.send(rtsp_request.encode())
+
+        # Read response
+        response = sock.recv(4096).decode()
+        sock.close()
+
+        # Check if we got a valid RTSP response with SDP content
+        return (
+            response.startswith("RTSP/1.0 200 OK")
+            and "application/sdp" in response
+            and "m=video" in response
+        )
+    except InvalidRtspUrlError:
+        raise
+    except Exception:
+        return False

OTVision/detect/video_input_source.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from typing import Generator
 
 import av
+from av.container.input import InputContainer
 from tqdm import tqdm
 
 from OTVision.abstraction.observer import Subject
@@ -12,6 +13,7 @@ from OTVision.application.detect.detection_file_save_path_provider import (
     DetectionFileSavePathProvider,
 )
 from OTVision.application.detect.timestamper import Timestamper
+from OTVision.application.event.new_video_start import NewVideoStartEvent
 from OTVision.application.get_current_config import GetCurrentConfig
 from OTVision.detect.detected_frame_buffer import FlushEvent
 from OTVision.detect.plugin_av.rotate_frame import AvVideoFrameRotator
@@ -38,6 +40,9 @@ class VideoSource(InputSourceDetect):
     and selective frame processing based on configuration parameters.
 
     Args:
+        subject_flush: (Subject[FlushEvent]): Subject for notifying about flush events.
+        subject_new_video_start (Subject[NewVideoStartEvent): Subject for notifying
+            about new video start events.
         get_current_config (GetCurrentConfig): Use case to retrieve current
             configuration.
         frame_rotator (AvVideoFrameRotator): Use to rotate video frames.
@@ -56,13 +61,15 @@ class VideoSource(InputSourceDetect):
 
     def __init__(
         self,
-        subject: Subject[FlushEvent],
+        subject_flush: Subject[FlushEvent],
+        subject_new_video_start: Subject[NewVideoStartEvent],
         get_current_config: GetCurrentConfig,
         frame_rotator: AvVideoFrameRotator,
         timestamper_factory: TimestamperFactory,
         save_path_provider: DetectionFileSavePathProvider,
     ) -> None:
-        super().__init__(subject)
+        self.subject_flush = subject_flush
+        self.subject_new_video_start = subject_new_video_start
         self._frame_rotator = frame_rotator
         self._get_current_config = get_current_config
         self._timestamper_factory = timestamper_factory
@@ -97,13 +104,14 @@ class VideoSource(InputSourceDetect):
                 expected_duration=self._current_config.detect.expected_duration,
             )
             video_fps = get_fps(video_file)
+            self.notify_new_video_start_observers(video_file, video_fps)
             detect_start = self.__get_detect_start_in_frames(video_fps)
             detect_end = self.__get_detect_end_in_frames(video_fps)
             counter = 0
             try:
                 with av.open(str(video_file.absolute())) as container:
                     container.streams.video[0].thread_type = "AUTO"
-                    side_data = container.streams.video[0].side_data
+                    side_data = self._extract_side_data(container)
                     for frame_number, frame in enumerate(
                         container.decode(video=0), start=1
                     ):
@@ -129,10 +137,20 @@ class VideoSource(InputSourceDetect):
                                 }
                             )
                         counter += 1
-                self.notify_observers(video_file, video_fps)
+                self.notify_flush_event_observers(video_file, video_fps)
             except Exception as e:
                 log.error(f"Error processing {video_file}", exc_info=e)
 
+    def _extract_side_data(self, container: InputContainer) -> dict:
+        try:
+            return container.streams.video[0].side_data
+        except AttributeError:
+            log.warning(
+                "No side_data found in video stream. "
+                "Existing rotation will not be applied."
+            )
+            return {}
+
     def __collect_files_to_detect(self) -> list[Path]:
         filetypes = self._current_config.filetypes.video_filetypes.to_list()
         video_files = get_files(
@@ -169,7 +187,9 @@ class VideoSource(InputSourceDetect):
             return False
         return True
 
-    def notify_observers(self, current_video_file: Path, video_fps: float) -> None:
+    def notify_flush_event_observers(
+        self, current_video_file: Path, video_fps: float
+    ) -> None:
         if expected_duration := self._current_config.detect.expected_duration:
             duration = expected_duration
         else:
@@ -180,7 +200,7 @@ class VideoSource(InputSourceDetect):
             current_video_file, start_time=self._start_time
         )
 
-        self._subject.notify(
+        self.subject_flush.notify(
             FlushEvent.create(
                 source=str(current_video_file),
                 output=str(current_video_file),
@@ -192,6 +212,18 @@ class VideoSource(InputSourceDetect):
             )
         )
 
+    def notify_new_video_start_observers(
+        self, current_video_file: Path, video_fps: float
+    ) -> None:
+        width, height = get_video_dimensions(current_video_file)
+        event = NewVideoStartEvent(
+            output=str(current_video_file),
+            width=width,
+            height=height,
+            fps=video_fps,
+        )
+        self.subject_new_video_start.notify(event)
+
     def __get_detect_start_in_frames(self, video_fps: float) -> int:
         detect_start = convert_seconds_to_frames(
             self._current_config.detect.detect_start, video_fps

OTVision/detect/yolo.py

@@ -142,9 +142,17 @@ class YoloDetector(ObjectDetector, Filter[Frame, DetectedFrame]):
     def detect(
         self, frames: Generator[Frame, None, None]
     ) -> Generator[DetectedFrame, None, None]:
-        for frame in tqdm(frames, desc="Detected frames", unit=" frames"):
+        for frame in tqdm(
+            frames,
+            desc="Detected frames",
+            unit=" frames",
+            disable=self.disable_tqdm_logging(),
+        ):
             yield self._predict(frame)
 
+    def disable_tqdm_logging(self) -> bool:
+        return log.level > logging.INFO
+
     def _predict(self, frame: Frame) -> DetectedFrame:
         if frame[FrameKeys.data] is None:
             return self._create_empty_detection(frame)

OTVision/domain/cli.py

@@ -3,6 +3,12 @@ from dataclasses import dataclass
 from datetime import datetime, timedelta
 from pathlib import Path
 
+from OTVision.plugin.ffmpeg_video_writer import (
+    ConstantRateFactor,
+    EncodingSpeed,
+    VideoCodec,
+)
+
 
 class CliArgs(ABC):
     @abstractmethod
@@ -32,6 +38,10 @@ class DetectCliArgs(CliArgs):
     start_time: datetime | None = None
     detect_start: int | None = None
     detect_end: int | None = None
+    write_video: bool | None = None
+    video_codec: VideoCodec | None = None
+    encoding_speed: EncodingSpeed | None = None
+    crf: ConstantRateFactor | None = None
 
     def get_config_file(self) -> Path | None:
         return self.config_file

OTVision/domain/input_source_detect.py

@@ -1,12 +1,10 @@
 from abc import ABC, abstractmethod
 from typing import Generator
 
-from OTVision.abstraction.observer import Observable
-from OTVision.detect.detected_frame_buffer import FlushEvent
 from OTVision.domain.frame import Frame
 
 
-class InputSourceDetect(Observable[FlushEvent], ABC):
+class InputSourceDetect(ABC):
     """Interface for input sources that generate frames and notify about flush events.
 
     This class combines the Observable pattern for flush events with frame generation

OTVision/domain/time.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from datetime import datetime, timezone
+from datetime import datetime
 
 
 class DatetimeProvider(ABC):
@@ -10,4 +10,4 @@ class DatetimeProvider(ABC):
 
 class CurrentDatetimeProvider(DatetimeProvider):
     def provide(self) -> datetime:
-        return datetime.now(tz=timezone.utc)
+        return datetime.now()

OTVision/domain/video_writer.py

@@ -0,0 +1,30 @@
+from abc import ABC, abstractmethod
+
+from numpy import ndarray
+
+from OTVision.abstraction.pipes_and_filter import Filter
+from OTVision.application.event.new_video_start import NewVideoStartEvent
+from OTVision.detect.detected_frame_buffer import FlushEvent
+from OTVision.domain.frame import Frame
+
+
+class VideoWriter(Filter[Frame, Frame], ABC):
+    @abstractmethod
+    def write(self, image: ndarray) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def open(self, output: str, width: int, height: int, fps: float) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def close(self) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def notify_on_flush_event(self, event: FlushEvent) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def notify_on_new_video_start(self, event: NewVideoStartEvent) -> None:
+        raise NotImplementedError