video.py 10.4 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import base64
4
import math
5
from abc import abstractmethod
6
from functools import partial
7
8
from io import BytesIO
from pathlib import Path
9
from typing import Any
10
11

import numpy as np
12
import numpy.typing as npt
13
from PIL import Image
14

15
from vllm import envs
16
from vllm.logger import init_logger
17

18
19
from .base import MediaIO
from .image import ImageMediaIO
20

21
22
logger = init_logger(__name__)

23
24
25
26

def resize_video(frames: npt.NDArray, size: tuple[int, int]) -> npt.NDArray:
    num_frames, _, _, channels = frames.shape
    new_height, new_width = size
27
28
29
    resized_frames = np.empty(
        (num_frames, new_height, new_width, channels), dtype=frames.dtype
    )
30
31
    # lazy import cv2 to avoid bothering users who only use text models
    import cv2
32

33
34
35
36
37
38
39
40
41
42
43
44
45
46
    for i, frame in enumerate(frames):
        resized_frame = cv2.resize(frame, (new_width, new_height))
        resized_frames[i] = resized_frame
    return resized_frames


def rescale_video_size(frames: npt.NDArray, size_factor: float) -> npt.NDArray:
    _, height, width, _ = frames.shape
    new_height = int(height * size_factor)
    new_width = int(width * size_factor)

    return resize_video(frames, (new_height, new_width))


47
def sample_frames_from_video(frames: npt.NDArray, num_frames: int) -> npt.NDArray:
48
49
50
51
52
53
54
    total_frames = frames.shape[0]
    if num_frames == -1:
        return frames

    frame_indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
    sampled_frames = frames[frame_indices, ...]
    return sampled_frames
55
56


57
58
class VideoLoader:
    @classmethod
59
    @abstractmethod
60
61
62
    def load_bytes(
        cls, data: bytes, num_frames: int = -1, **kwargs
    ) -> tuple[npt.NDArray, dict[str, Any]]:
63
64
65
        raise NotImplementedError


66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
class VideoLoaderRegistry:
    def __init__(self) -> None:
        self.name2class: dict[str, type] = {}

    def register(self, name: str):
        def wrap(cls_to_register):
            self.name2class[name] = cls_to_register
            return cls_to_register

        return wrap

    @staticmethod
    def load(cls_name: str) -> VideoLoader:
        cls = VIDEO_LOADER_REGISTRY.name2class.get(cls_name)
        assert cls is not None, f"VideoLoader class {cls_name} not found"
        return cls()


VIDEO_LOADER_REGISTRY = VideoLoaderRegistry()


@VIDEO_LOADER_REGISTRY.register("opencv")
88
89
90
91
92
93
94
95
96
97
class OpenCVVideoBackend(VideoLoader):
    def get_cv2_video_api(self):
        import cv2.videoio_registry as vr

        api_pref = None
        for backend in vr.getStreamBufferedBackends():
            if not vr.hasBackend(backend):
                continue
            if not vr.isBackendBuiltIn(backend):
                _, abi, api = vr.getStreamBufferedBackendPluginVersion(backend)
98
                if abi < 1 or (abi == 1 and api < 2):
99
100
101
102
103
104
                    continue
            api_pref = backend
            break
        return api_pref

    @classmethod
105
106
107
108
    def load_bytes(
        cls,
        data: bytes,
        num_frames: int = -1,
109
        fps: int = -1,
110
111
        **kwargs,
    ) -> tuple[npt.NDArray, dict[str, Any]]:
112
113
114
115
116
117
118
119
        import cv2

        backend = cls().get_cv2_video_api()
        cap = cv2.VideoCapture(BytesIO(data), backend, [])
        if not cap.isOpened():
            raise ValueError("Could not open video stream")

        total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
120
121
122
        original_fps = cap.get(cv2.CAP_PROP_FPS)
        duration = total_frames_num / original_fps if original_fps > 0 else 0

123
124
125
126
127
128
129
130
131
132
133
        # resample video to target num_frames and fps
        # - the minimum of the two will be used
        num_frames_to_sample = total_frames_num
        if num_frames > 0:
            num_frames_to_sample = min(num_frames, total_frames_num)
        if fps > 0:
            num_frames_to_sample = min(num_frames_to_sample, math.floor(duration * fps))
        num_frames_to_sample = max(1, num_frames_to_sample)  # at least one sample

        if num_frames_to_sample == total_frames_num:
            frame_idx = list(range(0, num_frames_to_sample))
134
        else:
135
            uniform_sampled_frames = np.linspace(
136
                0, total_frames_num - 1, num_frames_to_sample, dtype=int
137
            )
138
139
140
141
142
143
144
            frame_idx = uniform_sampled_frames.tolist()

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
        frames = np.empty((len(frame_idx), height, width, 3), dtype=np.uint8)

        i = 0
145
        for idx in range(max(frame_idx) + 1):
146
            ok = cap.grab()
147
148
            if not ok:
                break
149
            if idx in frame_idx:
150
151
152
153
                ret, frame = cap.retrieve()
                if ret:
                    frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    i += 1
154

155
156
        assert i == num_frames_to_sample, (
            f"Expected reading {num_frames_to_sample} frames, "
157
158
            f"but only loaded {i} frames from video."
        )
159

160
161
162
163
        # Use transformers transformers.video_utils.VideoMetadata format
        # NOTE(Isotr0py): For models like Qwen3-VL/GLM4.5V, this metadata
        # can cause incorrect timestamp calculation without num_frames=-1.
        metadata = {
164
165
            "total_num_frames": total_frames_num,
            "fps": original_fps,
166
167
            "duration": duration,
            "video_backend": "opencv",
168
            "frames_indices": list(frame_idx),
169
170
            # extra field used to control hf processor's video
            # sampling behavior
171
            "do_sample_frames": num_frames_to_sample == total_frames_num,
172
173
        }

174
175
176
177
178
179
180
181
182
183
        return frames, metadata


@VIDEO_LOADER_REGISTRY.register("opencv_dynamic")
class OpenCVDynamicVideoBackend(OpenCVVideoBackend):
    @classmethod
    def load_bytes(
        cls,
        data: bytes,
        num_frames: int = -1,
184
        fps: int = 2,
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
        max_duration: int = 300,
        **kwargs,
    ) -> tuple[npt.NDArray, dict[str, Any]]:
        import cv2

        backend = cls().get_cv2_video_api()
        cap = cv2.VideoCapture(BytesIO(data), backend, [])
        if not cap.isOpened():
            raise ValueError("Could not open video stream")

        total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        original_fps = cap.get(cv2.CAP_PROP_FPS)
        duration = total_frames_num / original_fps if original_fps > 0 else 0

        # resample video to target num_frames
        max_frame_idx = total_frames_num - 1
        duration = duration or round(max_frame_idx / original_fps) + 1

        # Refer to:
        # https://github.com/huggingface/transformers/blob/v4.55.4/src/transformers/models/glm4v/video_processing_glm4v.py#L103-L140
205
        frame_indices: range | list[int]
206
        if duration <= max_duration:
207
            n = int(math.floor(duration * fps))
208
209
210
211
212
213
            frame_indices = sorted(
                {
                    min(max_frame_idx, int(math.ceil(i * original_fps / fps)))
                    for i in range(n)
                }
            )
214
        else:
215
            num_samples = int(max_duration * fps)
216
217
218
            if num_samples >= total_frames_num:
                frame_indices = range(total_frames_num)
            else:
219
220
221
222
223
224
225
                target_seconds = np.linspace(0, duration, num_samples, endpoint=True)
                frame_indices = sorted(
                    {
                        min(max_frame_idx, int(math.ceil(t * original_fps)))
                        for t in target_seconds
                    }
                )
226
227
228

        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
229
        frames = np.empty((len(frame_indices), height, width, 3), dtype=np.uint8)
230
231
232
233
234
235
236
237
238
239
240
241
242
243

        i = 0
        for idx in range(total_frames_num):
            ok = cap.grab()
            if not ok:
                break
            if idx in frame_indices:
                ret, frame = cap.retrieve()
                if ret:
                    frames[i] = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    i += 1

        assert i == len(frame_indices), (
            f"Expected reading {len(frame_indices)} frames, "
244
245
            f"but only loaded {i} frames from video."
        )
246

247
248
249
250
251
252
253
254
255
256
        # Use transformers transformers.video_utils.VideoMetadata format
        metadata = {
            "total_num_frames": total_frames_num,
            "fps": original_fps,
            "duration": duration,
            "video_backend": "opencv_dynamic",
            "frames_indices": list(frame_indices),
            "do_sample_frames": False,
        }

257
        return frames, metadata
258
259


260
261
262
263
264
class VideoMediaIO(MediaIO[npt.NDArray]):
    def __init__(
        self,
        image_io: ImageMediaIO,
        num_frames: int = 32,
265
        **kwargs,
266
267
268
269
270
    ) -> None:
        super().__init__()

        self.image_io = image_io
        self.num_frames = num_frames
271
272
273
274
275
276
        # `kwargs` contains custom arguments from
        # --media-io-kwargs for this modality.
        # They can be passed to the underlying
        # media loaders (e.g. custom implementations)
        # for flexible control.
        self.kwargs = kwargs
277
278
        video_loader_backend = envs.VLLM_VIDEO_LOADER_BACKEND
        self.video_loader = VIDEO_LOADER_REGISTRY.load(video_loader_backend)
279

280
    def load_bytes(self, data: bytes) -> tuple[npt.NDArray, dict[str, Any]]:
281
282
283
        return self.video_loader.load_bytes(
            data, num_frames=self.num_frames, **self.kwargs
        )
284

285
286
287
    def load_base64(
        self, media_type: str, data: str
    ) -> tuple[npt.NDArray, dict[str, Any]]:
288
289
290
291
292
293
        if media_type.lower() == "video/jpeg":
            load_frame = partial(
                self.image_io.load_base64,
                "image/jpeg",
            )

294
295
296
            return np.stack(
                [np.asarray(load_frame(frame_data)) for frame_data in data.split(",")]
            ), {}
297
298
299

        return self.load_bytes(base64.b64decode(data))

300
    def load_file(self, filepath: Path) -> tuple[npt.NDArray, dict[str, Any]]:
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        with filepath.open("rb") as f:
            data = f.read()

        return self.load_bytes(data)

    def encode_base64(
        self,
        media: npt.NDArray,
        *,
        video_format: str = "JPEG",
    ) -> str:
        video = media

        if video_format == "JPEG":
            encode_frame = partial(
                self.image_io.encode_base64,
                image_format=video_format,
            )

320
            return ",".join(encode_frame(Image.fromarray(frame)) for frame in video)
321
322
323

        msg = "Only JPEG format is supported for now."
        raise NotImplementedError(msg)