File size: 4,176 Bytes
f258387 | 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 | """Video Depth Anything preprocessing: the per-frame transform of the official ``infer_video_depth`` (Resize lower_bound
to 518 keeping the aspect ratio, sides multiples of 14, cv2 bicubic; ImageNet normalisation; vendored in ``vda_dpt.py``),
including its rule for wide videos (aspect ratio > 1.78: smaller input size), behind the transformers image-processor
interface. Frames are RGB uint8 arrays (H, W, 3), as the official ``read_video_frames`` returns them.
``post_process_depth_estimation`` resizes to the given sizes as ``infer_video_depth`` does (float32, bilinear,
align_corners=True)."""
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torchvision.transforms import Compose
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_processing_utils import BaseImageProcessor
from .vda_dpt import NormalizeImage, PrepareForNet, Resize
class VideoDepthAnythingImageProcessor(BaseImageProcessor):
model_input_names = ["pixel_values"]
def __init__(self, input_size: int = 518, **kwargs):
super().__init__(**kwargs)
self.input_size = input_size
def preprocess(self, videos, input_size: int | None = None, return_tensors=None, **kwargs) -> BatchFeature:
"""videos: one video or a list of videos of the same size; a video is an array (T, H, W, 3) or a list of
frames (RGB arrays or PIL images). pixel_values: (B, T, 3, H', W')."""
is_frame = lambda x: isinstance(x, Image.Image) or isinstance(x, np.ndarray) and x.ndim == 3 # noqa: E731
if isinstance(videos, np.ndarray) and videos.ndim == 4 or is_frame(videos[0]):
videos = [videos]
size = input_size or self.input_size
batch = []
for video in videos:
frames = [np.asarray(f.convert("RGB")) if isinstance(f, Image.Image) else f for f in video]
frame_height, frame_width = frames[0].shape[:2]
# infer_video_depth: wide videos get a smaller input size
input_size_video = size
ratio = max(frame_height, frame_width) / min(frame_height, frame_width)
if ratio > 1.78:
input_size_video = int(input_size_video * 1.777 / ratio)
input_size_video = round(input_size_video / 14) * 14
transform = Compose([
Resize(width=input_size_video, height=input_size_video, resize_target=False, keep_aspect_ratio=True,
ensure_multiple_of=14, resize_method="lower_bound", image_interpolation_method=cv2.INTER_CUBIC),
NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
PrepareForNet(),
])
batch.append(torch.stack([torch.from_numpy(transform({"image": f.astype(np.float32) / 255.0})["image"]) for f in frames]))
if len({t.shape for t in batch}) > 1:
raise ValueError("videos resize to different shapes; process them one at a time")
return BatchFeature({"pixel_values": torch.stack(batch)}, tensor_type=return_tensors)
def post_process_depth_estimation(self, outputs, target_sizes=None) -> list:
"""``DepthEstimatorOutput`` (predicted_depth (B, T, H, W)) -> ``[{"predicted_depth": (T, h, w)}]`` per video;
target_sizes: (height, width) per video (e.g. the frame size), None keeps the model resolution."""
predicted_depth = outputs.predicted_depth
if target_sizes is not None and len(predicted_depth) != len(target_sizes):
raise ValueError("Make sure that you pass in as many target sizes as the batch dimension of the predicted depth")
results = []
for depth, size in zip(predicted_depth, target_sizes if target_sizes is not None else [None] * len(predicted_depth)):
if size is not None: # infer_video_depth: depth.to(float32 input dtype), then upsample
depth = F.interpolate(depth.float().unsqueeze(1), size=(int(size[0]), int(size[1])), mode="bilinear", align_corners=True)[:, 0]
results.append({"predicted_depth": depth})
return results
__all__ = ["VideoDepthAnythingImageProcessor"]
|