File size: 5,450 Bytes
637209e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 | import torch
import torchvision.transforms.v2.functional as tvF
from torchcodec.decoders import VideoDecoder
from transformers.image_processing_utils import BatchFeature
from transformers.image_utils import PILImageResampling
from transformers.processing_utils import Unpack, VideosKwargs
from transformers.video_processing_utils import BaseVideoProcessor, VideoMetadata
from transformers.video_utils import VideoInput
class MSPVisualVideoProcessor(BaseVideoProcessor):
resample = PILImageResampling.BILINEAR
def __init__(
self,
do_convert_rgb_to_grayscale: bool = True,
do_rescale: bool = True,
rescale_factor: float = 1 / 255.0,
image_mean=0.421,
image_std=0.165,
do_normalize: bool = True,
do_resize: bool = True,
size: dict[str, int] = {"height": 96, "width": 96},
do_center_crop: bool = True,
crop_size: dict[str, int] = {"height": 88, "width": 88},
**kwargs: Unpack[VideosKwargs],
):
super().__init__(
do_rescale=do_rescale,
rescale_factor=rescale_factor,
image_mean=image_mean,
image_std=image_std,
do_normalize=do_normalize,
do_resize=do_resize,
size=size,
do_center_crop=do_center_crop,
crop_size=crop_size,
**kwargs,
)
self.do_convert_rgb_to_grayscale = do_convert_rgb_to_grayscale
def sample_frames(
self,
metadata: VideoMetadata,
num_frames: int | None = None,
fps: int | float | None = None,
**kwargs,
):
if num_frames:
total_frames = metadata.total_num_frames
num_frames = num_frames if num_frames is not None else self.num_frames
assert num_frames is not None, (
"`num_frames` must be specified if `fixed_len_video == True`"
)
frame_idxs = [
int(i * (total_frames - 1) / (num_frames - 1))
for i in range(num_frames)
]
return torch.tensor(frame_idxs)
else:
return super().sample_frames(metadata, num_frames, fps, **kwargs)
def _load_video(self, src: str | bytes) -> torch.Tensor:
"""
Load video from a file path or bytes and return as a 4D torch.Tensor.
Args:
src (str | bytes): Path to the video file or bytes of the video file.
Returns:
torch.Tensor: Loaded video as a 4D tensor (num_frames, height, width, num_channels).
"""
vd = VideoDecoder(src)
video = vd.get_frames_in_range(0, vd.metadata.num_frames).data
return video
def __call__(
self, videos: VideoInput | str | list[str] | bytes | list[bytes], **kwargs
):
"""Overrides the __call__ method to handle video input as file paths or bytes."""
if isinstance(videos, (str, bytes)):
videos = self._load_video(videos)
elif isinstance(videos, list) and isinstance(videos[0], (str, bytes)):
videos = [self._load_video(v) for v in videos]
# remove kwargs not in VideosKwargs
# for key in list(kwargs.keys()):
# if key not in VideosKwargs.__optional_keys__:
# kwargs.pop(key, None)
return super().__call__(videos, **kwargs)
def convert_rgb_to_grayscale(self, video: torch.Tensor) -> torch.Tensor:
"""
Convert a video to grayscale.
"""
video = tvF.rgb_to_grayscale(video)
return video
def _preprocess(
self,
videos: VideoInput,
**kwargs: Unpack[VideosKwargs],
) -> BatchFeature:
"""
Preprocesses a video or a batch of videos.
Args:
videos (VideoInput): Video to preprocess.
See `VideoInput` for details.
**kwargs: Additional keyword arguments.
Returns:
BatchFeature: A BatchFeature with the following fields:
- pixel_values_videos: Pixel values to be fed to a model, of shape (batch_size,num_channels, num_frames, height, width).
- padding_mask_videos (optional): Mask to be used for padding, of shape (batch_size, num_frames).
"""
# Always set `return_tensors` to `None` since it won't pad variable length videos
# We'll handle this after we call the parent' method
return_tensors = kwargs.pop("return_tensors", None)
result = super()._preprocess(videos, **kwargs)
pixels = result.pixel_values_videos
if self.do_convert_rgb_to_grayscale:
pixels = [self.convert_rgb_to_grayscale(video) for video in pixels]
data = {"pixel_values_videos": pixels}
if return_tensors:
lengths = torch.tensor([video.size(0) for video in pixels])
pixels = torch.nn.utils.rnn.pad_sequence(
pixels, batch_first=True, padding_value=0.0
)
data["pixel_values_videos"] = pixels
if lengths.unique().size(0) > 1:
mask = torch.arange(lengths.max())[None] < lengths[:, None]
data["padding_mask_videos"] = mask
# pixel_values_videos shape [batch_size, num_channels, num_frames, height, width]
data["pixel_values_videos"] = data["pixel_values_videos"].permute(0, 2, 1, 3, 4)
return BatchFeature(data=data, tensor_type=return_tensors)
|