Merge pull request #137 from AllenNeuralDynamics/refactor-update-ffmp…

…eg-default

Update ffmpeg string and improve documentation

docs/api.base.rst

@@ -1,14 +1,6 @@
 base
 -------------
 
-.. toctree::
-   :maxdepth: 2
-
-   api.base/session
-   api.base/rig
-   api.base/task_logic
-
-
 .. automodule:: aind_behavior_services.base
    :members:
    :undoc-members:

docs/api.rig.rst

@@ -0,0 +1,7 @@
+rig
+-------------------------
+
+.. automodule:: aind_behavior_services.rig
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api.rst

@@ -5,6 +5,9 @@ API
    :maxdepth: 2
 
    api.base
-   api.calibration
+   api.rig
+   api.task_logic
+   api.session
    api.data_types
+   api.calibration
    api.utils

docs/api.session.rst

@@ -0,0 +1,10 @@
+session
+--------------------------
+
+.. image:: _static/AindBehaviorSessionModel.svg
+   :target: _static/AindBehaviorSessionModel.svg
+
+.. automodule:: aind_behavior_services.session
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/api.task_logic.rst

@@ -0,0 +1,7 @@
+task_logic
+--------------------------------
+
+.. automodule:: aind_behavior_services.task_logic
+   :members:
+   :undoc-members:
+   :show-inheritance:

docs/articles/data_formats/video.rst

@@ -62,24 +62,18 @@ or
 
 - Use mp4 container
 - Acquire without any gamma correction
-- Use ``ffmpeg`` with the following encoding codec string for online encoding (optimized for compression quality and speed):
+- Use ``ffmpeg`` with the encoding codec string for online encoding (optimized for compression quality and speed) set to the default values defined in :py:attr:`~aind_behavior_services.rig.FFMPEG_OUTPUT_8BIT` and :py:attr:`~aind_behavior_services.rig.FFMPEG_INPUT`:
 
-  .. note::
-     This pipeline has been designed and tested with monochrome videos with the raw pixel format ``gray``. For color videos, the arguments might need to be altered to match the color space of the input.
-
-  - output arguments: ``-vf "scale=out_color_matrix=bt709:out_range=full,format=bgr24,scale=out_range=full" -c:v h264_nvenc -pix_fmt yuv420p -color_range full -colorspace bt709 -color_trc linear -tune hq -preset p4 -rc vbr -cq 12 -b:v 0M -metadata author="Allen Institute for Neural Dynamics" -maxrate 700M -bufsize 350M``
-  - input_arguments: ``-colorspace bt709 -color_primaries bt709 -color_range full -color_trc linear``
-
-and the following encoding codec string for offline re-encoding (optimized for quality and size):
+The following encoding codec string can be used for offline re-encoding (optimized for quality and size):
 
 - output arguments: ``-vf "scale=out_color_matrix=bt709:out_range=full:sws_dither=none,format=yuv420p10le,colorspace=ispace=bt709:all=bt709:dither=none,scale=out_range=tv:sws_dither=none,format=yuv420p" -c:v libx264 -preset veryslow -crf 18 -pix_fmt yuv420p -metadata author="Allen Institute for Neural Dynamics" -movflags +faststart+write_colr``
 
 .. warning::
 
-  For higher bit depth (more than 8 bit) recordings, change the output arguments of the online, first stage, encoding to be as follows:
-    - output arguments: ``-vf "scale=out_color_matrix=bt709:out_range=full,format=rgb48le,scale=out_range=full" -c:v h264_nvenc -pix_fmt yuv420p -color_range full -colorspace bt709 -color_trc linear -tune hq -preset p4 -rc vbr -cq 12 -b:v 0M -metadata author="Allen Institute for Neural Dynamics" -maxrate 700M -bufsize 350M``
+  This pipeline has been designed and tested with monochrome videos with the raw pixel format ``gray``. For color videos, the arguments might need to be altered to match the color space of the input.
 
-This is almost the same, except the intermediate color representation is 48 bits per pixel instead of 24.
+  For higher bit depth (more than 8 bit) recordings, change the output arguments of the online, first stage, encoding to :py:attr:`~aind_behavior_services.rig.FFMPEG_OUTPUT_16BIT`
+  This is almost the same, except the intermediate color representation is 48 bits per pixel instead of 24.
 
 Application notes
 #####################

src/aind_behavior_services/rig/__init__.py

@@ -2,7 +2,7 @@
 
 import os
 from enum import Enum, IntEnum, auto
-from typing import Annotated, Dict, Generic, List, Literal, Optional, TypeVar, Union
+from typing import Annotated, Any, Dict, Generic, List, Literal, Optional, TypeVar, Union
 
 from pydantic import BaseModel, Field, field_validator
 from typing_extensions import TypeAliasType
@@ -16,16 +16,56 @@ class Device(BaseModel):
     calibration: Optional[BaseModel] = Field(default=None, description="Calibration")
 
 
+FFMPEG_OUTPUT_8BIT = '-vf "scale=out_color_matrix=bt709:out_range=full,format=bgr24,scale=out_range=full" -c:v h264_nvenc -pix_fmt yuv420p -color_range full -colorspace bt709 -color_trc linear -tune hq -preset p4 -rc vbr -cq 12 -b:v 0M -metadata author="Allen Institute for Neural Dynamics" -maxrate 700M -bufsize 350M'
+""" Default output arguments for 8-bit video encoding """
+
+FFMPEG_OUTPUT_16BIT = '-vf "scale=out_color_matrix=bt709:out_range=full,format=rgb48le,scale=out_range=full" -c:v hevc_nvenc -pix_fmt p010le -color_range full -colorspace bt709 -color_trc linear -tune hq -preset p4 -rc vbr -cq 12 -b:v 0M -metadata author="Allen Institute for Neural Dynamics" -maxrate 700M -bufsize 350M'
+""" Default output arguments for 16-bit video encoding """
+
+FFMPEG_INPUT = "-colorspace bt709 -color_primaries bt709 -color_range full -color_trc linear"
+""" Default input arguments """
+
+
+class VideoWriterFfmpegFactory:
+    def __init__(self, bit_depth: Literal[8, 16] = 8, video_writer_ffmpeg_kwargs: Dict[str, Any] = None):
+        self._bit_depth = bit_depth
+        self.video_writer_ffmpeg_kwargs = video_writer_ffmpeg_kwargs or {}
+        self._output_arguments: str
+        self._input_arguments: str
+        self._solve_strings()
+
+    def _solve_strings(self):
+        if self._bit_depth == 8:
+            self._output_arguments = FFMPEG_OUTPUT_8BIT
+        elif self._bit_depth == 16:
+            self._output_arguments = FFMPEG_OUTPUT_16BIT
+        else:
+            raise ValueError(f"Bit depth {self._bit_depth} not supported")
+        self._input_arguments = FFMPEG_INPUT
+
+    def construct_video_writer_ffmpeg(self) -> VideoWriterFfmpeg:
+        return VideoWriterFfmpeg(
+            output_arguments=self._output_arguments,
+            input_arguments=self._input_arguments,
+            **self.video_writer_ffmpeg_kwargs,
+        )
+
+    def update_video_writer_ffmpeg_kwargs(self, video_writer: VideoWriterFfmpeg):
+        return video_writer.model_copy(
+            update={"output_arguments": self._output_arguments, "input_arguments": self._input_arguments}
+        )
+
+
 class VideoWriterFfmpeg(BaseModel):
     video_writer_type: Literal["FFMPEG"] = Field(default="FFMPEG")
     frame_rate: int = Field(default=30, ge=0, description="Encoding frame rate")
     container_extension: str = Field(default="mp4", description="Container extension")
     output_arguments: str = Field(
-        default='-vf "scale=out_color_matrix=bt709:out_range=full,format=bgr24,scale=out_range=full" -c:v h264_nvenc -pix_fmt yuv420p -color_range full -colorspace bt709 -color_trc linear -tune hq -preset p4 -rc vbr -cq 12 -b:v 0M -metadata author="Allen Institute for Neural Dynamics" -maxrate 700M -bufsize 350M',  # E501
+        default=FFMPEG_OUTPUT_8BIT,
         description="Output arguments",
     )
     input_arguments: str = Field(
-        default="-colorspace bt709 -color_primaries bt709 -color_range full -color_trc linear",
+        default=FFMPEG_INPUT,
         description="Input arguments",
     )
 

tests/test_rig.py

@@ -0,0 +1,54 @@
+import unittest
+
+from aind_behavior_services.rig import (
+    FFMPEG_INPUT,
+    FFMPEG_OUTPUT_8BIT,
+    FFMPEG_OUTPUT_16BIT,
+    VideoWriterFfmpeg,
+    VideoWriterFfmpegFactory,
+)
+
+
+class TestVideoWriterFfmpegFactory(unittest.TestCase):
+    def test_initialization(self):
+        factory = VideoWriterFfmpegFactory(bit_depth=8)
+        self.assertEqual(factory._bit_depth, 8)
+        self.assertEqual(factory.video_writer_ffmpeg_kwargs, {})
+
+        factory = VideoWriterFfmpegFactory(bit_depth=16, video_writer_ffmpeg_kwargs={"frame_rate": 60})
+        self.assertEqual(factory._bit_depth, 16)
+        self.assertEqual(factory.video_writer_ffmpeg_kwargs, {"frame_rate": 60})
+
+    def test_solve_strings_8bit(self):
+        factory = VideoWriterFfmpegFactory(bit_depth=8)
+        self.assertEqual(factory._output_arguments, FFMPEG_OUTPUT_8BIT)
+        self.assertEqual(factory._input_arguments, FFMPEG_INPUT)
+
+    def test_solve_strings_16bit(self):
+        factory = VideoWriterFfmpegFactory(bit_depth=16)
+        self.assertEqual(factory._output_arguments, FFMPEG_OUTPUT_16BIT)
+        self.assertEqual(factory._input_arguments, FFMPEG_INPUT)
+
+    def test_construct_video_writer_ffmpeg(self):
+        factory = VideoWriterFfmpegFactory(bit_depth=8)
+        video_writer = factory.construct_video_writer_ffmpeg()
+        self.assertIsInstance(video_writer, VideoWriterFfmpeg)
+        self.assertEqual(video_writer.output_arguments, factory._output_arguments)
+        self.assertEqual(video_writer.input_arguments, factory._input_arguments)
+
+    def test_update_video_writer_ffmpeg_kwargs(self):
+        factory = VideoWriterFfmpegFactory(bit_depth=8)
+        video_writer = factory.construct_video_writer_ffmpeg()
+        updated_video_writer = factory.update_video_writer_ffmpeg_kwargs(video_writer)
+        self.assertEqual(updated_video_writer.output_arguments, factory._output_arguments)
+        self.assertEqual(updated_video_writer.input_arguments, factory._input_arguments)
+
+    def test_video_writer_ffmpeg_obj_equality(self):
+        factory = VideoWriterFfmpegFactory(bit_depth=8)
+        video_writer = VideoWriterFfmpeg(output_arguments=FFMPEG_OUTPUT_8BIT, input_arguments=FFMPEG_INPUT)
+        video_writer_from_factory = factory.construct_video_writer_ffmpeg()
+        self.assertEqual(video_writer, video_writer_from_factory)
+
+
+if __name__ == "__main__":
+    unittest.main()