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

OTVision/application/config.py

@@ -2,6 +2,12 @@ from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from pathlib import Path
 
+from OTVision.plugin.ffmpeg_video_writer import (
+    ConstantRateFactor,
+    EncodingSpeed,
+    VideoCodec,
+)
+
 AVAILABLE_WEIGHTS = "AVAILABLEWEIGHTS"
 CALIBRATIONS = "CALIBRATIONS"
 COL_WIDTH = "COLWIDTH"
@@ -58,6 +64,10 @@ LOG_DIR = "LOG_DIR"
 START_TIME = "START_TIME"
 DETECT_END = "DETECT_END"
 DETECT_START = "DETECT_START"
+WRITE_VIDEO = "WRITE_VIDEO"
+VIDEO_CODEC = "VIDEO_CODEC"
+ENCODING_SPEED = "ENCODING_SPEED"
+CRF = "CRF"
 DATETIME_FORMAT = "%Y-%m-%d_%H-%M-%S"
 DEFAULT_EXPECTED_DURATION: timedelta = timedelta(minutes=15)
 """Default length of a video is 15 minutes."""
@@ -289,6 +299,10 @@ class DetectConfig:
     start_time: datetime | None = None
     detect_start: int | None = None
     detect_end: int | None = None
+    write_video: bool = False
+    video_codec: VideoCodec = VideoCodec.H264_SOFTWARE
+    encoding_speed: EncodingSpeed = EncodingSpeed.FAST
+    crf: ConstantRateFactor = ConstantRateFactor.DEFAULT
 
     def to_dict(self) -> dict:
         expected_duration = (
@@ -306,6 +320,10 @@ class DetectConfig:
             START_TIME: self.start_time,
             DETECT_START: self.detect_start,
             DETECT_END: self.detect_end,
+            WRITE_VIDEO: self.write_video,
+            VIDEO_CODEC: self.video_codec.value,
+            ENCODING_SPEED: self.encoding_speed.value,
+            CRF: self.crf.name,
         }
 
 

OTVision/application/config_parser.py

@@ -5,12 +5,14 @@ from OTVision.application.config import (
     COL_WIDTH,
     CONF,
     CONVERT,
+    CRF,
     DATETIME_FORMAT,
     DEFAULT_FILETYPE,
     DELETE_INPUT,
     DETECT,
     DETECT_END,
     DETECT_START,
+    ENCODING_SPEED,
     EXPECTED_DURATION,
     FLUSH_BUFFER_SIZE,
     FONT,
@@ -51,8 +53,10 @@ from OTVision.application.config import (
     TRANSFORM,
     UNDISTORT,
     VID,
+    VIDEO_CODEC,
     WEIGHTS,
     WINDOW,
+    WRITE_VIDEO,
     YOLO,
     Config,
     ConvertConfig,
@@ -181,15 +185,19 @@ class ConfigParser:
 
         start_time = self._parse_start_time(data)
         return DetectConfig(
-            sources,
-            data.get(RUN_CHAINED, DetectConfig.run_chained),
-            yolo_config,
-            expected_duration,
-            data.get(OVERWRITE, DetectConfig.overwrite),
-            data.get(HALF_PRECISION, DetectConfig.half_precision),
-            start_time,
-            data.get(DETECT_START, DetectConfig.detect_start),
-            data.get(DETECT_END, DetectConfig.detect_end),
+            paths=sources,
+            run_chained=data.get(RUN_CHAINED, DetectConfig.run_chained),
+            yolo_config=yolo_config,
+            expected_duration=expected_duration,
+            overwrite=data.get(OVERWRITE, DetectConfig.overwrite),
+            half_precision=data.get(HALF_PRECISION, DetectConfig.half_precision),
+            start_time=start_time,
+            detect_start=data.get(DETECT_START, DetectConfig.detect_start),
+            detect_end=data.get(DETECT_END, DetectConfig.detect_end),
+            write_video=data.get(WRITE_VIDEO, DetectConfig.write_video),
+            video_codec=data.get(VIDEO_CODEC, DetectConfig.video_codec),
+            encoding_speed=data.get(ENCODING_SPEED, DetectConfig.encoding_speed),
+            crf=data.get(CRF, DetectConfig.crf),
         )
 
     def parse_yolo_config(self, data: dict) -> YoloConfig:

OTVision/application/detect/update_detect_config_with_cli_args.py

@@ -84,6 +84,22 @@ class UpdateDetectConfigWithCliArgs:
                 if cli_args.detect_end is not None
                 else detect_config.detect_end
             ),
+            write_video=(
+                cli_args.write_video
+                if cli_args.write_video is not None
+                else detect_config.write_video
+            ),
+            video_codec=(
+                cli_args.video_codec
+                if cli_args.video_codec is not None
+                else detect_config.video_codec
+            ),
+            encoding_speed=(
+                cli_args.encoding_speed
+                if cli_args.encoding_speed is not None
+                else detect_config.encoding_speed
+            ),
+            crf=cli_args.crf if cli_args.crf is not None else detect_config.crf,
         )
 
     def _update_log_config(self, config: Config, cli_args: DetectCliArgs) -> _LogConfig:

OTVision/application/event/new_video_start.py

@@ -0,0 +1,9 @@
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class NewVideoStartEvent:
+    output: str
+    width: int
+    height: int
+    fps: float

OTVision/application/video/generate_video.py

@@ -0,0 +1,15 @@
+from OTVision.abstraction.pipes_and_filter import Filter
+from OTVision.domain.frame import Frame
+from OTVision.domain.input_source_detect import InputSourceDetect
+
+
+class GenerateVideo:
+    def __init__(
+        self, input_source: InputSourceDetect, video_writer: Filter[Frame, Frame]
+    ) -> None:
+        self._input_source = input_source
+        self._video_writer = video_writer
+
+    def generate(self) -> None:
+        for frame in self._video_writer.filter(self._input_source.produce()):
+            pass

OTVision/config.py

@@ -77,6 +77,7 @@ from OTVision.application.config import (
     VIDEOS,
     WEIGHTS,
     WINDOW,
+    WRITE_VIDEO,
     YOLO,
     Config,
 )
@@ -176,6 +177,7 @@ CONFIG[DETECT][OVERWRITE] = True
 CONFIG[DETECT][HALF_PRECISION] = False
 CONFIG[DETECT][DETECT_START] = None
 CONFIG[DETECT][DETECT_END] = None
+CONFIG[DETECT][WRITE_VIDEO] = False
 
 # TRACK
 CONFIG[TRACK] = {}

OTVision/detect/builder.py

@@ -4,7 +4,7 @@ from functools import cached_property
 
 from OTVision.abstraction.observer import Subject
 from OTVision.application.buffer import Buffer
-from OTVision.application.config import Config
+from OTVision.application.config import Config, DetectConfig
 from OTVision.application.config_parser import ConfigParser
 from OTVision.application.configure_logger import ConfigureLogger
 from OTVision.application.detect.current_object_detector import CurrentObjectDetector
@@ -12,9 +12,6 @@ from OTVision.application.detect.current_object_detector_metadata import (
     CurrentObjectDetectorMetadata,
 )
 from OTVision.application.detect.detected_frame_factory import DetectedFrameFactory
-from OTVision.application.detect.detected_frame_producer import (
-    SimpleDetectedFrameProducer,
-)
 from OTVision.application.detect.detection_file_save_path_provider import (
     DetectionFileSavePathProvider,
 )
@@ -34,6 +31,10 @@ from OTVision.detect.detected_frame_buffer import (
     DetectedFrameBufferEvent,
     FlushEvent,
 )
+from OTVision.detect.detected_frame_producer import (
+    DetectedFrameProducerFactory,
+    SimpleDetectedFrameProducer,
+)
 from OTVision.detect.otdet import OtdetBuilder
 from OTVision.detect.otdet_file_writer import OtdetFileWriter
 from OTVision.detect.plugin_av.rotate_frame import AvVideoFrameRotator
@@ -47,6 +48,7 @@ from OTVision.domain.frame import DetectedFrame
 from OTVision.domain.input_source_detect import InputSourceDetect
 from OTVision.domain.object_detection import ObjectDetectorFactory
 from OTVision.domain.serialization import Deserializer
+from OTVision.domain.video_writer import VideoWriter
 from OTVision.plugin.yaml_serialization import YamlDeserializer
 
 
@@ -152,9 +154,17 @@ class DetectBuilder(ABC):
     @cached_property
     def detected_frame_producer(self) -> DetectedFrameProducer:
         return SimpleDetectedFrameProducer(
+            producer_factory=self.detected_frame_producer_factory,
+        )
+
+    @cached_property
+    def detected_frame_producer_factory(self) -> DetectedFrameProducerFactory:
+        return DetectedFrameProducerFactory(
             input_source=self.input_source,
+            video_writer_filter=self.video_file_writer,
             detection_filter=self.current_object_detector,
             detected_frame_buffer=self.detected_frame_buffer,
+            get_current_config=self.get_current_config,
         )
 
     @cached_property
@@ -165,6 +175,10 @@ class DetectBuilder(ABC):
     def yaml_deserializer(self) -> Deserializer:
         return YamlDeserializer()
 
+    @property
+    def detect_config(self) -> DetectConfig:
+        return self.current_config.get().detect
+
     def __init__(self, argv: list[str] | None = None) -> None:
         self.argv = argv
 
@@ -173,6 +187,15 @@ class DetectBuilder(ABC):
     def input_source(self) -> InputSourceDetect:
         raise NotImplementedError
 
+    @property
+    @abstractmethod
+    def video_file_writer(self) -> VideoWriter:
+        raise NotImplementedError
+
+    @abstractmethod
+    def register_observers(self) -> None:
+        raise NotImplementedError
+
     def build(self) -> OTVisionVideoDetect:
         self.register_observers()
         self._preload_object_detection_model()
@@ -181,7 +204,3 @@ class DetectBuilder(ABC):
     def _preload_object_detection_model(self) -> None:
         model = self.current_object_detector.get()
         model.preload()
-
-    def register_observers(self) -> None:
-        self.input_source.register(self.detected_frame_buffer.on_flush)
-        self.detected_frame_buffer.register(self.otdet_file_writer.write)

OTVision/detect/cli.py

@@ -6,6 +6,11 @@ from OTVision.application.config import DATETIME_FORMAT
 from OTVision.domain.cli import CliParseError, DetectCliArgs, DetectCliParser
 from OTVision.helpers.files import check_if_all_paths_exist
 from OTVision.helpers.log import DEFAULT_LOG_FILE, VALID_LOG_LEVELS
+from OTVision.plugin.ffmpeg_video_writer import (
+    ConstantRateFactor,
+    EncodingSpeed,
+    VideoCodec,
+)
 
 
 class ArgparseDetectCliParser(DetectCliParser):
@@ -128,6 +133,34 @@ class ArgparseDetectCliParser(DetectCliParser):
             help="Specify end of detection in seconds.",
             required=False,
         )
+        self._parser.add_argument(
+            "--write-video",
+            default=None,
+            action="store_true",
+            help="Write video to output folder. Not supported on Windows.",
+            required=False,
+        )
+        self._parser.add_argument(
+            "--video-codec",
+            default=None,
+            choices=VideoCodec.as_list(),
+            help="Video codec for video writer. Default is 'libx264'",
+            required=False,
+        )
+        self._parser.add_argument(
+            "--encoding-speed",
+            default=None,
+            choices=EncodingSpeed.as_list(),
+            help="Encoding speed for video writer. Default is 'fast'",
+            required=False,
+        )
+        self._parser.add_argument(
+            "--crf",
+            default=None,
+            choices=ConstantRateFactor.as_list(),
+            help="Constant rate factor for video writer. Default is 'DEFAULT'",
+            required=False,
+        )
 
     def parse(self) -> DetectCliArgs:
         args = self._parser.parse_args(self._argv)
@@ -151,11 +184,21 @@ class ArgparseDetectCliParser(DetectCliParser):
             detect_start=(
                 int(args.detect_start) if args.detect_start is not None else None
             ),
-            detect_end=int(args.detect_end) if args.detect_end is not None else None,
+            detect_end=(int(args.detect_end) if args.detect_end is not None else None),
             logfile=Path(args.logfile),
             log_level_console=args.log_level_console,
             log_level_file=args.log_level_file,
             logfile_overwrite=args.logfile_overwrite,
+            write_video=args.write_video,
+            video_codec=(
+                VideoCodec(args.video_codec) if args.video_codec is not None else None
+            ),
+            encoding_speed=(
+                EncodingSpeed(args.encoding_speed)
+                if args.encoding_speed is not None
+                else None
+            ),
+            crf=ConstantRateFactor[args.crf] if args.crf is not None else None,
         )
 
     def _parse_start_time(self, start_time: str | None) -> datetime | None:

OTVision/detect/detected_frame_producer.py

@@ -0,0 +1,14 @@
+from typing import Generator
+
+from OTVision.detect.detected_frame_producer_factory import DetectedFrameProducerFactory
+from OTVision.domain.detect_producer_consumer import DetectedFrameProducer
+from OTVision.domain.frame import DetectedFrame
+
+
+class SimpleDetectedFrameProducer(DetectedFrameProducer):
+
+    def __init__(self, producer_factory: DetectedFrameProducerFactory) -> None:
+        self._producer_factory = producer_factory
+
+    def produce(self) -> Generator[DetectedFrame, None, None]:
+        return self._producer_factory.create()

OTVision/detect/detected_frame_producer_factory.py

@@ -0,0 +1,43 @@
+from typing import Generator
+
+from OTVision.abstraction.pipes_and_filter import Filter
+from OTVision.application.get_current_config import GetCurrentConfig
+from OTVision.domain.frame import DetectedFrame, Frame
+from OTVision.domain.input_source_detect import InputSourceDetect
+
+
+class DetectedFrameProducerFactory:
+    def __init__(
+        self,
+        input_source: InputSourceDetect,
+        video_writer_filter: Filter[Frame, Frame],
+        detection_filter: Filter[Frame, DetectedFrame],
+        detected_frame_buffer: Filter[DetectedFrame, DetectedFrame],
+        get_current_config: GetCurrentConfig,
+    ) -> None:
+        self._input_source = input_source
+        self._video_writer_filter = video_writer_filter
+        self._detection_filter = detection_filter
+        self._detected_frame_buffer = detected_frame_buffer
+        self._get_current_config = get_current_config
+
+    def create(self) -> Generator[DetectedFrame, None, None]:
+        if self._get_current_config.get().detect.write_video:
+            return self.__create_with_video_writer()
+        return self.__create_without_video_writer()
+
+    def __create_without_video_writer(
+        self,
+    ) -> Generator[DetectedFrame, None, None]:
+        return self._detected_frame_buffer.filter(
+            self._detection_filter.filter(self._input_source.produce())
+        )
+
+    def __create_with_video_writer(
+        self,
+    ) -> Generator[DetectedFrame, None, None]:
+        return self._detected_frame_buffer.filter(
+            self._detection_filter.filter(
+                self._video_writer_filter.filter(self._input_source.produce())
+            )
+        )

OTVision/detect/file_based_detect_builder.py

@@ -1,19 +1,55 @@
 from functools import cached_property
 
 from OTVision.abstraction.observer import Subject
+from OTVision.application.event.new_video_start import NewVideoStartEvent
 from OTVision.detect.builder import DetectBuilder
 from OTVision.detect.detected_frame_buffer import FlushEvent
 from OTVision.detect.video_input_source import VideoSource
-from OTVision.domain.input_source_detect import InputSourceDetect
+from OTVision.domain.video_writer import VideoWriter
+from OTVision.plugin.ffmpeg_video_writer import (
+    FfmpegVideoWriter,
+    PixelFormat,
+    VideoFormat,
+    append_save_suffix_to_save_location,
+)
 
 
 class FileBasedDetectBuilder(DetectBuilder):
+
     @cached_property
-    def input_source(self) -> InputSourceDetect:
+    def input_source(self) -> VideoSource:
         return VideoSource(
-            subject=Subject[FlushEvent](),
+            subject_flush=Subject[FlushEvent](),
+            subject_new_video_start=Subject[NewVideoStartEvent](),
             get_current_config=self.get_current_config,
             frame_rotator=self.frame_rotator,
             timestamper_factory=self.timestamper_factory,
             save_path_provider=self.detection_file_save_path_provider,
         )
+
+    @cached_property
+    def video_file_writer(self) -> VideoWriter:
+        # Using save_location_strategy=keep_original_save_location is not supported for
+        # file-based detection. Otherwise, we would be overwriting the input source that
+        # we are reading from.
+        return FfmpegVideoWriter(
+            save_location_strategy=append_save_suffix_to_save_location,
+            encoding_speed=self.detect_config.encoding_speed,
+            input_format=VideoFormat.RAW,
+            output_format=VideoFormat.MP4,
+            input_pixel_format=PixelFormat.RGB24,
+            output_pixel_format=PixelFormat.YUV420P,
+            output_video_codec=self.detect_config.video_codec,
+            constant_rate_factor=self.detect_config.crf,
+        )
+
+    def register_observers(self) -> None:
+        if self.detect_config.write_video:
+            self.input_source.subject_new_video_start.register(
+                self.video_file_writer.notify_on_new_video_start
+            )
+            self.input_source.subject_flush.register(
+                self.video_file_writer.notify_on_flush_event
+            )
+        self.input_source.subject_flush.register(self.detected_frame_buffer.on_flush)
+        self.detected_frame_buffer.register(self.otdet_file_writer.write)

OTVision/detect/rtsp_based_detect_builder.py

@@ -2,11 +2,18 @@ from functools import cached_property
 
 from OTVision.abstraction.observer import Subject
 from OTVision.application.config import StreamConfig
+from OTVision.application.event.new_video_start import NewVideoStartEvent
 from OTVision.detect.builder import DetectBuilder
 from OTVision.detect.detected_frame_buffer import FlushEvent
 from OTVision.detect.rtsp_input_source import Counter, RtspInputSource
-from OTVision.domain.input_source_detect import InputSourceDetect
 from OTVision.domain.time import CurrentDatetimeProvider, DatetimeProvider
+from OTVision.domain.video_writer import VideoWriter
+from OTVision.plugin.ffmpeg_video_writer import (
+    FfmpegVideoWriter,
+    PixelFormat,
+    VideoFormat,
+    keep_original_save_location,
+)
 
 FLUSH_BUFFER_SIZE = 18000
 FLUSH_BUFFER_SIZE = 1200
@@ -24,9 +31,10 @@ class RtspBasedDetectBuilder(DetectBuilder):
         return config.stream
 
     @cached_property
-    def input_source(self) -> InputSourceDetect:
+    def input_source(self) -> RtspInputSource:
         return RtspInputSource(
-            subject=Subject[FlushEvent](),
+            subject_flush=Subject[FlushEvent](),
+            subject_new_video_start=Subject[NewVideoStartEvent](),
             datetime_provider=self.datetime_provider,
             frame_counter=Counter(),
             get_current_config=self.get_current_config,
@@ -35,3 +43,27 @@ class RtspBasedDetectBuilder(DetectBuilder):
     @cached_property
     def datetime_provider(self) -> DatetimeProvider:
         return CurrentDatetimeProvider()
+
+    @cached_property
+    def video_file_writer(self) -> VideoWriter:
+        return FfmpegVideoWriter(
+            save_location_strategy=keep_original_save_location,
+            encoding_speed=self.detect_config.encoding_speed,
+            input_format=VideoFormat.RAW,
+            output_format=VideoFormat.MP4,
+            input_pixel_format=PixelFormat.RGB24,
+            output_pixel_format=PixelFormat.YUV420P,
+            output_video_codec=self.detect_config.video_codec,
+            constant_rate_factor=self.detect_config.crf,
+        )
+
+    def register_observers(self) -> None:
+        if self.detect_config.write_video:
+            self.input_source.subject_new_video_start.register(
+                self.video_file_writer.notify_on_new_video_start
+            )
+            self.input_source.subject_flush.register(
+                self.video_file_writer.notify_on_flush_event
+            )
+        self.input_source.subject_flush.register(self.detected_frame_buffer.on_flush)
+        self.detected_frame_buffer.register(self.otdet_file_writer.write)