Ticket/425/motion/correction/tiffs

src/ophys_etl/modules/video/schemas.py

@@ -1,6 +1,7 @@
 import argschema
 import pathlib
 from marshmallow import post_load
+from marshmallow.validate import OneOf
 
 
 class VideoBaseSchema(argschema.ArgSchema):
@@ -16,8 +17,23 @@ class VideoBaseSchema(argschema.ArgSchema):
             required=False,
             allow_none=True,
             default=3,
-            description=("Radius of median filter kernel; "
-                         "if None, no median filter applied"))
+            description=("Size of spatial filter kernel; "
+                         "if None or zero, no spatial filter applied"))
+
+    kernel_type = argschema.fields.String(
+            required=False,
+            allow_none=False,
+            default='median',
+            validation=OneOf(('median', 'mean')),
+            description=("Type of spatial smoothing kernel to be "
+                         "applied to the video after downsampling"))
+
+    video_dtype = argschema.fields.String(
+            required=False,
+            allow_none=False,
+            default='uint8',
+            validation=OneOf(('uint8', 'uint16')),
+            description=("Type to which the output video is cast"))
 
     input_frame_rate_hz = argschema.fields.Float(
             required=True,
@@ -113,8 +129,10 @@ def check_quality(self, data, **kwargs):
     def check_output_path(self, data, **kwargs):
         output_path = pathlib.Path(data['output_path'])
         output_suffix = output_path.suffix
-        if output_suffix not in ('.mp4', '.avi'):
-            msg = "output_path must be an .mp4 or .avi file\n"
+        allowed = ('.mp4', '.avi', '.tiff', '.tif')
+        if output_suffix not in allowed:
+            msg = "output_path must have one of these extensions:\n"
+            msg += f"{allowed}\n"
             msg += f"you gave {data['output_path']}"
             raise ValueError(msg)
         return data

src/ophys_etl/modules/video/side_by_side_video.py

@@ -1,3 +1,5 @@
+from ophys_etl.modules.video.util_classes import VideoModuleMixin
+
 from ophys_etl.modules.video.utils import (
     create_side_by_side_video)
 
@@ -25,30 +27,29 @@ class SideBySideVideoSchema(VideoBaseSchema):
                          "in the right panel of the output video"))
 
 
-class SideBySideVideoGenerator(argschema.ArgSchemaParser):
+class SideBySideVideoGenerator(argschema.ArgSchemaParser,
+                               VideoModuleMixin):
 
     default_schema = SideBySideVideoSchema
 
     def run(self):
-        if self.args['upper_quantile'] is not None:
-            quantiles = (self.args['lower_quantile'],
-                         self.args['upper_quantile'])
-        else:
-            quantiles = None
+
+        supplemental_args = self._get_supplemental_args()
 
         create_side_by_side_video(
-            pathlib.Path(self.args['left_video_path']),
-            pathlib.Path(self.args['right_video_path']),
-            self.args['input_frame_rate_hz'],
-            pathlib.Path(self.args['output_path']),
-            self.args['output_frame_rate_hz'],
-            self.args['kernel_size'],
-            self.args['n_parallel_workers'],
+            left_video_path=pathlib.Path(self.args['left_video_path']),
+            right_video_path=pathlib.Path(self.args['right_video_path']),
+            input_hz=self.args['input_frame_rate_hz'],
+            output_path=pathlib.Path(self.args['output_path']),
+            output_hz=self.args['output_frame_rate_hz'],
+            spatial_filter=supplemental_args['spatial_filter'],
+            n_processors=self.args['n_parallel_workers'],
             quality=self.args['quality'],
-            quantiles=quantiles,
+            quantiles=supplemental_args['quantiles'],
             reticle=self.args['reticle'],
             speed_up_factor=self.args['speed_up_factor'],
-            tmp_dir=self.args['tmp_dir'])
+            tmp_dir=self.args['tmp_dir'],
+            video_dtype=supplemental_args['video_dtype'])
 
 
 if __name__ == "__main__":

src/ophys_etl/modules/video/single_video.py

@@ -1,3 +1,5 @@
+from ophys_etl.modules.video.util_classes import VideoModuleMixin
+
 from ophys_etl.modules.video.utils import (
     create_downsampled_video)
 
@@ -17,29 +19,27 @@ class VideoSchema(VideoBaseSchema):
            description="Path to the input video file")
 
 
-class VideoGenerator(argschema.ArgSchemaParser):
+class VideoGenerator(argschema.ArgSchemaParser,
+                     VideoModuleMixin):
 
     default_schema = VideoSchema
 
     def run(self):
-        if self.args['upper_quantile'] is not None:
-            quantiles = (self.args['lower_quantile'],
-                         self.args['upper_quantile'])
-        else:
-            quantiles = None
+        supplemental_args = self._get_supplemental_args()
 
         create_downsampled_video(
-            pathlib.Path(self.args['video_path']),
-            self.args['input_frame_rate_hz'],
-            pathlib.Path(self.args['output_path']),
-            self.args['output_frame_rate_hz'],
-            self.args['kernel_size'],
-            self.args['n_parallel_workers'],
+            input_path=pathlib.Path(self.args['video_path']),
+            input_hz=self.args['input_frame_rate_hz'],
+            output_path=pathlib.Path(self.args['output_path']),
+            output_hz=self.args['output_frame_rate_hz'],
+            spatial_filter=supplemental_args['spatial_filter'],
+            n_processors=self.args['n_parallel_workers'],
             quality=self.args['quality'],
-            quantiles=quantiles,
+            quantiles=supplemental_args['quantiles'],
             reticle=self.args['reticle'],
             speed_up_factor=self.args['speed_up_factor'],
-            tmp_dir=self.args['tmp_dir'])
+            tmp_dir=self.args['tmp_dir'],
+            video_dtype=supplemental_args['video_dtype'])
 
 
 if __name__ == "__main__":

src/ophys_etl/modules/video/util_classes.py

@@ -0,0 +1,44 @@
+from ophys_etl.modules.video.utils import (
+    apply_downsampled_mean_filter_to_video)
+from ophys_etl.modules.median_filtered_max_projection.utils import (
+    apply_median_filter_to_video)
+from functools import partial
+import numpy as np
+
+
+class VideoModuleMixin(object):
+
+    def _get_supplemental_args(self) -> dict:
+        """
+        Parse self.args and return a dict like
+        {'quantiles': optional tuple of quantiles to use in clipping video
+         'spatial_filter': optional callable spatial filter to apply to video
+         'video_dtype': dtype of video
+        """
+        if self.args['upper_quantile'] is not None:
+            quantiles = (self.args['lower_quantile'],
+                         self.args['upper_quantile'])
+        else:
+            quantiles = (0.0, 1.0)
+
+        use_kernel = False
+        if self.args['kernel_size'] is not None:
+            if self.args['kernel_size'] > 0:
+                use_kernel = True
+
+        if use_kernel:
+            if self.args['kernel_type'] == 'median':
+                spatial_filter = partial(apply_median_filter_to_video,
+                                         kernel_size=self.args['kernel_size'])
+            else:
+                spatial_filter = partial(
+                                     apply_downsampled_mean_filter_to_video,
+                                     kernel_size=self.args['kernel_size'])
+        else:
+            spatial_filter = None
+
+        video_dtype = np.dtype(self.args['video_dtype'])
+
+        return {'quantiles': quantiles,
+                'spatial_filter': spatial_filter,
+                'video_dtype': video_dtype}