[fbsync] Add video kernels to dispatcher tests (#6905)

Summary:
* add KernelInfo's for {five, ten}_crop_video

* add video kernels to dispatcher tests

Reviewed By: NicolasHug

Differential Revision: D41265199

fbshipit-source-id: 156dee1b71a2431f1a1f66a9c3779af9682b4be2

Author: Yosua Michael Maranatha

test/prototype_transforms_dispatcher_infos.py

@@ -140,6 +140,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.horizontal_flip,
         kernels={
             features.Image: F.horizontal_flip_image_tensor,
+            features.Video: F.horizontal_flip_video,
             features.BoundingBox: F.horizontal_flip_bounding_box,
             features.Mask: F.horizontal_flip_mask,
         },
@@ -149,6 +150,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.resize,
         kernels={
             features.Image: F.resize_image_tensor,
+            features.Video: F.resize_video,
             features.BoundingBox: F.resize_bounding_box,
             features.Mask: F.resize_mask,
         },
@@ -161,6 +163,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.affine,
         kernels={
             features.Image: F.affine_image_tensor,
+            features.Video: F.affine_video,
             features.BoundingBox: F.affine_bounding_box,
             features.Mask: F.affine_mask,
         },
@@ -177,6 +180,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.vertical_flip,
         kernels={
             features.Image: F.vertical_flip_image_tensor,
+            features.Video: F.vertical_flip_video,
             features.BoundingBox: F.vertical_flip_bounding_box,
             features.Mask: F.vertical_flip_mask,
         },
@@ -186,6 +190,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.rotate,
         kernels={
             features.Image: F.rotate_image_tensor,
+            features.Video: F.rotate_video,
             features.BoundingBox: F.rotate_bounding_box,
             features.Mask: F.rotate_mask,
         },
@@ -200,6 +205,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.crop,
         kernels={
             features.Image: F.crop_image_tensor,
+            features.Video: F.crop_video,
             features.BoundingBox: F.crop_bounding_box,
             features.Mask: F.crop_mask,
         },
@@ -209,6 +215,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.resized_crop,
         kernels={
             features.Image: F.resized_crop_image_tensor,
+            features.Video: F.resized_crop_video,
             features.BoundingBox: F.resized_crop_bounding_box,
             features.Mask: F.resized_crop_mask,
         },
@@ -218,6 +225,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.pad,
         kernels={
             features.Image: F.pad_image_tensor,
+            features.Video: F.pad_video,
             features.BoundingBox: F.pad_bounding_box,
             features.Mask: F.pad_mask,
         },
@@ -244,6 +252,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.perspective,
         kernels={
             features.Image: F.perspective_image_tensor,
+            features.Video: F.perspective_video,
             features.BoundingBox: F.perspective_bounding_box,
             features.Mask: F.perspective_mask,
         },
@@ -256,6 +265,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.elastic,
         kernels={
             features.Image: F.elastic_image_tensor,
+            features.Video: F.elastic_video,
             features.BoundingBox: F.elastic_bounding_box,
             features.Mask: F.elastic_mask,
         },
@@ -265,6 +275,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.center_crop,
         kernels={
             features.Image: F.center_crop_image_tensor,
+            features.Video: F.center_crop_video,
             features.BoundingBox: F.center_crop_bounding_box,
             features.Mask: F.center_crop_mask,
         },
@@ -277,6 +288,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.gaussian_blur,
         kernels={
             features.Image: F.gaussian_blur_image_tensor,
+            features.Video: F.gaussian_blur_video,
         },
         pil_kernel_info=PILKernelInfo(F.gaussian_blur_image_pil),
         test_marks=[
@@ -288,48 +300,55 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.equalize,
         kernels={
             features.Image: F.equalize_image_tensor,
+            features.Video: F.equalize_video,
         },
         pil_kernel_info=PILKernelInfo(F.equalize_image_pil, kernel_name="equalize_image_pil"),
     ),
     DispatcherInfo(
         F.invert,
         kernels={
             features.Image: F.invert_image_tensor,
+            features.Video: F.invert_video,
         },
         pil_kernel_info=PILKernelInfo(F.invert_image_pil, kernel_name="invert_image_pil"),
     ),
     DispatcherInfo(
         F.posterize,
         kernels={
             features.Image: F.posterize_image_tensor,
+            features.Video: F.posterize_video,
         },
         pil_kernel_info=PILKernelInfo(F.posterize_image_pil, kernel_name="posterize_image_pil"),
     ),
     DispatcherInfo(
         F.solarize,
         kernels={
             features.Image: F.solarize_image_tensor,
+            features.Video: F.solarize_video,
         },
         pil_kernel_info=PILKernelInfo(F.solarize_image_pil, kernel_name="solarize_image_pil"),
     ),
     DispatcherInfo(
         F.autocontrast,
         kernels={
             features.Image: F.autocontrast_image_tensor,
+            features.Video: F.autocontrast_video,
         },
         pil_kernel_info=PILKernelInfo(F.autocontrast_image_pil, kernel_name="autocontrast_image_pil"),
     ),
     DispatcherInfo(
         F.adjust_sharpness,
         kernels={
             features.Image: F.adjust_sharpness_image_tensor,
+            features.Video: F.adjust_sharpness_video,
         },
         pil_kernel_info=PILKernelInfo(F.adjust_sharpness_image_pil, kernel_name="adjust_sharpness_image_pil"),
     ),
     DispatcherInfo(
         F.erase,
         kernels={
             features.Image: F.erase_image_tensor,
+            features.Video: F.erase_video,
         },
         pil_kernel_info=PILKernelInfo(F.erase_image_pil),
         test_marks=[
@@ -340,41 +359,47 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.adjust_brightness,
         kernels={
             features.Image: F.adjust_brightness_image_tensor,
+            features.Video: F.adjust_brightness_video,
         },
         pil_kernel_info=PILKernelInfo(F.adjust_brightness_image_pil, kernel_name="adjust_brightness_image_pil"),
     ),
     DispatcherInfo(
         F.adjust_contrast,
         kernels={
             features.Image: F.adjust_contrast_image_tensor,
+            features.Video: F.adjust_contrast_video,
         },
         pil_kernel_info=PILKernelInfo(F.adjust_contrast_image_pil, kernel_name="adjust_contrast_image_pil"),
     ),
     DispatcherInfo(
         F.adjust_gamma,
         kernels={
             features.Image: F.adjust_gamma_image_tensor,
+            features.Video: F.adjust_gamma_video,
         },
         pil_kernel_info=PILKernelInfo(F.adjust_gamma_image_pil, kernel_name="adjust_gamma_image_pil"),
     ),
     DispatcherInfo(
         F.adjust_hue,
         kernels={
             features.Image: F.adjust_hue_image_tensor,
+            features.Video: F.adjust_hue_video,
         },
         pil_kernel_info=PILKernelInfo(F.adjust_hue_image_pil, kernel_name="adjust_hue_image_pil"),
     ),
     DispatcherInfo(
         F.adjust_saturation,
         kernels={
             features.Image: F.adjust_saturation_image_tensor,
+            features.Video: F.adjust_saturation_video,
         },
         pil_kernel_info=PILKernelInfo(F.adjust_saturation_image_pil, kernel_name="adjust_saturation_image_pil"),
     ),
     DispatcherInfo(
         F.five_crop,
         kernels={
             features.Image: F.five_crop_image_tensor,
+            features.Video: F.five_crop_video,
         },
         pil_kernel_info=PILKernelInfo(F.five_crop_image_pil),
         test_marks=[
@@ -386,6 +411,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.ten_crop,
         kernels={
             features.Image: F.ten_crop_image_tensor,
+            features.Video: F.ten_crop_video,
         },
         test_marks=[
             xfail_jit_python_scalar_arg("size"),
@@ -397,6 +423,7 @@ def fill_sequence_needs_broadcast(args_kwargs):
         F.normalize,
         kernels={
             features.Image: F.normalize_image_tensor,
+            features.Video: F.normalize_video,
         },
         test_marks=[
             skip_dispatch_feature,

test/prototype_transforms_kernel_infos.py

@@ -2020,6 +2020,12 @@ def reference_inputs_five_crop_image_tensor():
             yield ArgsKwargs(image_loader, size=size)
 
 
+def sample_inputs_five_crop_video():
+    size = _FIVE_TEN_CROP_SIZES[0]
+    for video_loader in make_video_loaders(sizes=[_get_five_ten_crop_spatial_size(size)]):
+        yield ArgsKwargs(video_loader, size=size)
+
+
 def sample_inputs_ten_crop_image_tensor():
     for size, vertical_flip in itertools.product(_FIVE_TEN_CROP_SIZES, [False, True]):
         for image_loader in make_image_loaders(
@@ -2036,30 +2042,45 @@ def reference_inputs_ten_crop_image_tensor():
             yield ArgsKwargs(image_loader, size=size, vertical_flip=vertical_flip)
 
 
+def sample_inputs_ten_crop_video():
+    size = _FIVE_TEN_CROP_SIZES[0]
+    for video_loader in make_video_loaders(sizes=[_get_five_ten_crop_spatial_size(size)]):
+        yield ArgsKwargs(video_loader, size=size)
+
+
+_common_five_ten_crop_marks = [
+    xfail_jit_python_scalar_arg("size"),
+    mark_framework_limitation(("TestKernels", "test_batched_vs_single"), "Custom batching needed."),
+]
+
 KERNEL_INFOS.extend(
     [
         KernelInfo(
             F.five_crop_image_tensor,
             sample_inputs_fn=sample_inputs_five_crop_image_tensor,
             reference_fn=pil_reference_wrapper(F.five_crop_image_pil),
             reference_inputs_fn=reference_inputs_five_crop_image_tensor,
-            test_marks=[
-                xfail_jit_python_scalar_arg("size"),
-                mark_framework_limitation(("TestKernels", "test_batched_vs_single"), "Custom batching needed."),
-            ],
+            test_marks=_common_five_ten_crop_marks,
             closeness_kwargs=DEFAULT_IMAGE_CLOSENESS_KWARGS,
         ),
+        KernelInfo(
+            F.five_crop_video,
+            sample_inputs_fn=sample_inputs_five_crop_video,
+            test_marks=_common_five_ten_crop_marks,
+        ),
         KernelInfo(
             F.ten_crop_image_tensor,
             sample_inputs_fn=sample_inputs_ten_crop_image_tensor,
             reference_fn=pil_reference_wrapper(F.ten_crop_image_pil),
             reference_inputs_fn=reference_inputs_ten_crop_image_tensor,
-            test_marks=[
-                xfail_jit_python_scalar_arg("size"),
-                mark_framework_limitation(("TestKernels", "test_batched_vs_single"), "Custom batching needed."),
-            ],
+            test_marks=_common_five_ten_crop_marks,
             closeness_kwargs=DEFAULT_IMAGE_CLOSENESS_KWARGS,
         ),
+        KernelInfo(
+            F.ten_crop_video,
+            sample_inputs_fn=sample_inputs_ten_crop_video,
+            test_marks=_common_five_ten_crop_marks,
+        ),
     ]
 )