Coverage for tests/test_video_utils.py: 99%
366 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
3from __future__ import annotations
5import os
6import logging
7import pytest
8import aiofiles
9import aiofiles.os
11import numpy as np
13from typing import List
14from typing import Optional
16from PIL import Image
18from media_utils import get_frame_with_text
19from media_utils import video_frames_to_base64
20from media_utils import base64_to_video_frames
21from media_utils import chunk_video_binary
22from media_utils import save_video_frames
23from media_utils import save_video_audio
24from media_utils import add_text_to_frame
25from media_utils import concatenate_videos
26from media_utils import change_video_fps
27from media_utils import get_video_file_info
28from media_utils import get_video_duration
29from media_utils import get_video_frames
30from media_utils import get_video_frame
31from media_utils import get_video_fps
32from media_utils import get_video_num_frames
33from media_utils import get_video_with_text
34from media_utils import get_video_frames_at_fps
35from media_utils import get_video_size
36from media_utils import get_font_size
37from media_utils import get_audio_duration
38from media_utils import get_ffmpeg_version
40from file_utils import read_file_base64
43@pytest.mark.asyncio
44async def test_video() -> None:
45 video_frames: List[np.ndarray] = [
46 get_frame_with_text(100, 60, f"frame{frame_id:02d}", output_type="numpy") # type: ignore[misc]
47 for frame_id in range(24)
48 ]
49 video_path = await save_video_frames(video_frames, fps=24)
50 video_file_info = get_video_file_info(video_path)
52 video_overall_info = video_file_info.get("overall")
53 assert video_overall_info is not None
54 assert video_overall_info["num_bytes"] > 1000
55 assert video_overall_info["duration_seconds"] == 1.0
57 video_info = video_file_info.get("video")
58 assert video_info is not None
59 assert video_info.get("duration_seconds") == 1.0
60 assert video_info.get("width") == 100
61 assert video_info.get("height") == 60
62 assert video_info.get("num_frames") == 24
63 assert video_info.get("fps") == 24
64 assert video_info.get("codec") == "h264"
66 pix_fmt = video_info.get("pix_fmt")
67 logging.info(f"Pixel format: {pix_fmt}")
69 del video_frames
70 del video_path
71 del video_info
74@pytest.mark.asyncio
75async def test_base64() -> None:
76 """Test video frames to/from base64 conversion."""
77 video_frames: List[np.ndarray] = [
78 get_frame_with_text(80, 64, f"frame{frame_id:02d}", output_type="pil") # type: ignore[misc]
79 for frame_id in range(12)
80 ]
81 video_base64 = video_frames_to_base64(video_frames) # type: ignore[arg-type]
82 video_frames_2 = base64_to_video_frames(video_base64)
83 assert len(video_frames) == len(video_frames_2)
84 for frame1, frame2 in zip(video_frames, video_frames_2):
85 assert frame1.size == frame2.size
86 assert get_video_size(video_frames_2) == 184320
88 with pytest.raises(TypeError):
89 video_frames_to_base64(12345) # type: ignore[arg-type]
90 with pytest.raises(ValueError):
91 video_frames_to_base64([])
92 with pytest.raises(TypeError):
93 base64_to_video_frames(12345) # type: ignore[arg-type]
95 assert get_video_size(None) == 0 # type: ignore[arg-type]
96 assert get_video_size([]) == 0
98 del video_frames
99 del video_frames_2
102def test_get_video_frames_at_fps() -> None:
103 """Test getting a video frames at a specific FPS."""
104 video_frames: list[Image.Image] = [
105 get_frame_with_text(128, 64, f"frame{frame_id:02d}", output_type="pil") # type: ignore[misc]
106 for frame_id in range(10)
107 ]
108 new_video_frames = get_video_frames_at_fps(video_frames, src_fps=30, dst_fps=30)
109 assert len(new_video_frames) == 10
110 new_video_frames = get_video_frames_at_fps(video_frames, src_fps=30, dst_fps=24)
111 assert len(new_video_frames) == 8 # 10 * 24 / 30 = 8
112 assert get_video_frames_at_fps([], src_fps=24) == []
113 assert get_video_frames_at_fps(None, dst_fps=16) == [] # type: ignore[arg-type]
115 with pytest.raises(ValueError):
116 get_video_frames_at_fps(video_frames, src_fps=0, dst_fps=0)
117 with pytest.raises(ValueError):
118 get_video_frames_at_fps(video_frames, src_fps=30, dst_fps=-16)
120 del video_frames
121 del new_video_frames
124@pytest.mark.asyncio
125async def test_chunk() -> None:
126 """Test chunking a video binary."""
127 video_binary = await get_video_with_text(50, 30, "Test, video", duration_seconds=3.0, fps=30)
128 video_binary_chunked = chunk_video_binary(video_binary, 0.5, 1.5)
129 duration = await get_video_duration(video_binary_chunked)
130 assert duration == 1.0
132 video_binary_chunked = chunk_video_binary(video_binary, 0.0, 1.5)
133 duration = await get_video_duration(video_binary_chunked)
134 assert duration == 1.5
136 video_binary_chunked = chunk_video_binary(video_binary_chunked, 0.5, 2.5)
137 duration = await get_video_duration(video_binary_chunked)
138 assert duration == 1.0
140 with pytest.raises(ValueError, match="end_seconds must be greater than start_seconds"):
141 chunk_video_binary(video_binary, 0.5, 0.4)
142 with pytest.raises(ValueError, match="must be non-negative"):
143 chunk_video_binary(video_binary, -0.5, 0.4)
144 with pytest.raises(TypeError):
145 chunk_video_binary(12345, 0.5, 1.5) # type: ignore[arg-type]
146 with pytest.raises(ValueError, match="Video input is empty"):
147 chunk_video_binary(None, 0.5, 1.5) # type: ignore[arg-type]
149 del video_binary
152@pytest.mark.asyncio
153async def test_get_frame() -> None:
154 """Test getting a specific frame from a video binary."""
155 video_frames = [
156 get_frame_with_text(160, 90, f"frame{frame_id:02d}", output_type="pil")
157 for frame_id in range(12)
158 ]
159 video_path = await save_video_frames(video_frames, fps=24)
160 async with aiofiles.open(video_path, "rb") as f:
161 video_binary = await f.read()
162 frame = await get_video_frame(video_binary, 5)
163 assert frame.size == (160, 90)
164 duration_seconds = await get_video_duration(video_binary)
165 assert duration_seconds == 0.5
166 duration_seconds = await get_video_duration(video_path)
167 assert duration_seconds == 0.5
169 with pytest.raises(TypeError):
170 await get_video_frame(video_path, 0) # type: ignore[arg-type]
171 with pytest.raises(ValueError):
172 await get_video_frame(video_binary, 512)
173 with pytest.raises(ValueError):
174 await get_video_frame(video_binary, -1)
175 with pytest.raises(ValueError):
176 await get_video_frame(b"", 5)
178 await aiofiles.os.remove(video_path)
179 del video_frames
180 del video_binary
183def test_add_text() -> None:
184 """Test adding text to video frames."""
185 NUM_FRAMES = 12
186 video_frames = [
187 get_frame_with_text(64, 64, "", output_type="pil")
188 for frame_id in range(NUM_FRAMES)
189 ]
190 assert len(video_frames) == NUM_FRAMES
191 for frame in video_frames:
192 assert frame.size == (64, 64)
193 assert np.abs(np.array(frame)).sum() == 0 # All black
195 video_frames = [
196 add_text_to_frame(frame, text="Sample Text", position="top-left")
197 for frame in video_frames
198 ]
199 assert len(video_frames) == NUM_FRAMES
200 for frame in video_frames:
201 assert frame.size == (64, 64)
202 assert np.abs(np.array(frame)).sum() > 0 # Not all black
204 # Centered text
205 video_frames = [
206 add_text_to_frame(frame, text="Sample Text 2", font_color="blue", position="center")
207 for frame in video_frames
208 ]
209 assert len(video_frames) == NUM_FRAMES
210 for frame in video_frames:
211 assert frame.size == (64, 64)
212 assert np.abs(np.array(frame)).sum() > 0 # Not all black
214 # Specific text position
215 font_size = get_font_size(64, 64)
216 video_frames = [
217 add_text_to_frame(frame, text="Sample Text 3", font_size=font_size, position=(0, 0))
218 for frame in video_frames
219 ]
220 assert len(video_frames) == NUM_FRAMES
221 for frame in video_frames:
222 assert frame.size == (64, 64)
223 assert np.abs(np.array(frame)).sum() > 0 # Not all black
225 del video_frames
228@pytest.mark.asyncio
229async def test_concatenate() -> None:
230 """Test concatenating multiple video binaries."""
231 video_frames1 = [
232 get_frame_with_text(120, 60, f"frame1_{frame_id:02d}", output_type="numpy")
233 for frame_id in range(5)
234 ]
235 video_frames2 = [
236 get_frame_with_text(120, 60, f"frame2_{frame_id:02d}", output_type="numpy")
237 for frame_id in range(10)
238 ]
239 video_path = await save_video_frames(video_frames1, fps=30)
240 async with aiofiles.open(video_path, 'rb') as file:
241 video_binary1 = await file.read()
242 video_path = await save_video_frames(video_frames2, fps=30)
243 async with aiofiles.open(video_path, 'rb') as file:
244 video_binary2 = await file.read()
245 video_binary = await concatenate_videos([
246 video_binary1,
247 video_binary2
248 ])
249 num_video_frames = await get_video_num_frames(video_binary)
250 assert num_video_frames == 15 # 5 + 10
252 video_frames = await get_video_frames(video_binary)
253 assert len(video_frames) == 15
254 video_frame = await get_video_frame(video_binary, 5)
255 assert video_frame.size == (120, 60)
257 video_frames = await get_video_frames(None)
258 assert video_frames == []
259 with pytest.raises(ValueError):
260 await concatenate_videos([])
261 with pytest.raises(TypeError, match="Video input 0 is not bytes or str"):
262 await concatenate_videos([video_frames1, video_frames2]) # type: ignore[arg-type]
263 with pytest.raises(FileNotFoundError):
264 await concatenate_videos(["abc", "def"])
265 with pytest.raises(ValueError, match="One of the inputs is corrupted"):
266 await concatenate_videos([b"abc", b"def"])
268 with pytest.raises(FileNotFoundError, match="Video file does not exist"):
269 await get_video_num_frames("abc")
270 with pytest.raises(ValueError, match="The video binary is corrupted or has an unsupported format"):
271 await get_video_num_frames(b"abc")
273 await aiofiles.os.remove(video_path)
275 del video_frames1
276 del video_frames2
277 del video_binary1
278 del video_binary2
279 del video_binary
282@pytest.mark.asyncio
283async def test_get_video_frames_empty() -> None:
284 video_frames = await get_video_frames(None)
285 assert video_frames == []
287 video_frames = await get_video_frames(b"")
288 assert video_frames == []
290 video_frames = await get_video_frames([]) # type: ignore[arg-type]
291 assert video_frames == []
294@pytest.mark.asyncio
295@pytest.mark.parametrize(
296 "num_frames,fps",
297 [
298 (5, 30.0), # 5 frames at 30 fps
299 (11, 23.0), # 11 frames at 23 fps
300 (17, 23.0), # 17 frames at 23 fps
301 (23, 23.0), # 23 frames at 23 fps
302 (32, 23.0), # 32 frames at 23 fps
303 (33, 23.0), # 33 frames at 23 fps
304 (81, 23.0), # 81 frames at 23 fps
305 (5, 16.0), # 5 frames at 16 fps
306 ],
307)
308async def test_get_video_frames(num_frames: int, fps: float) -> None:
309 width, height = 120, 60
310 video_frames = [
311 get_frame_with_text(width, height, f"frame{frame_id:02d}", output_type="numpy")
312 for frame_id in range(num_frames)
313 ]
314 video_path = await save_video_frames(
315 video_frames,
316 fps=fps)
317 async with aiofiles.open(video_path, "rb") as file:
318 video_binary = await file.read()
319 video_frames_out = await get_video_frames(video_binary)
320 assert len(video_frames_out) == num_frames
322 video_file_info = get_video_file_info(video_binary)
323 assert "video" in video_file_info
324 video_info = video_file_info["video"]
325 assert video_info is not None
326 assert video_info.get("width") == width
327 assert video_info.get("height") == height
328 assert video_info.get("num_frames") == num_frames
329 assert video_info.get("num_frames") == len(video_frames_out)
330 assert video_info.get("fps") == fps
331 duration_seconds = video_info.get("duration_seconds")
332 assert duration_seconds is not None
333 assert abs(duration_seconds - (num_frames / fps)) < 0.01
334 video_overall_info = video_file_info.get("overall")
335 assert video_overall_info is not None
336 num_bytes = video_overall_info.get("num_bytes")
337 assert num_bytes is not None
338 assert num_bytes > 1000
340 os.remove(video_path)
341 del video_frames
344@pytest.mark.asyncio
345async def test_save_video_frames() -> None:
346 NUM_FRAMES = 81
347 WIDTH, HEIGHT = 160, 100
348 FPS = 23.0
350 video_frames = [
351 get_frame_with_text(WIDTH, HEIGHT, f"frame{frame_id:02d}", output_type="numpy")
352 for frame_id in range(NUM_FRAMES)
353 ]
355 video_path = await save_video_frames(video_frames, fps=FPS)
356 video_file_info = get_video_file_info(video_path)
357 video_info = video_file_info.get("video")
358 assert video_info is not None
359 assert video_info.get("fps") == FPS
360 assert video_info.get("num_frames") == 81
361 assert video_info.get("duration_seconds") == 3.521739 # 81 / 23
362 os.remove(video_path)
364 """
365 # Test with trimming
366 video_path = await save_video_frames(video_frames, fps=FPS, time_in_seconds=2.0)
367 video_info = get_video_file_info(video_path)
368 assert video_info["video"]["fps"] == FPS
369 assert video_info["video"]["num_frames"] == 46 # 2.0 * 23
370 assert video_info["video"]["duration_seconds"] == 2.0
371 os.remove(video_path)
373 # Test with expansion adds empty frames
374 video_path = await save_video_frames(video_frames, fps=FPS, time_in_seconds=5.0)
375 video_info = get_video_file_info(video_path)
376 assert video_info["video"]["fps"] == FPS
377 assert video_info["video"]["num_frames"] == 115
378 assert video_info["video"]["duration_seconds"] == 5.0
379 os.remove(video_path)
380 """
382 with pytest.raises(ValueError):
383 await save_video_frames(None, fps=FPS) # type: ignore[arg-type]
384 with pytest.raises(ValueError):
385 await save_video_frames([], fps=FPS)
386 with pytest.raises(TypeError):
387 await save_video_frames(b"BLAH", fps=FPS) # type: ignore[arg-type]
390def assert_approx(
391 a: float,
392 b: float,
393 tol: float = 1e-3,
394 msg: Optional[str] = None
395) -> None:
396 if msg:
397 assert abs(a - b) < tol, f"{a} !~= {b}: {msg}"
398 else:
399 assert abs(a - b) < tol, f"{a} !~= {b}"
402@pytest.mark.asyncio
403async def test_save_video_frames_audio() -> None:
404 # Get the audio
405 audio_path = "tests/data/audio_4675.wav"
406 audio_base64 = await read_file_base64(audio_path)
407 audio_duration_secs = get_audio_duration(audio_base64)
408 assert audio_duration_secs == 4.675
410 # Get video only with the video frames
411 NUM_FRAMES = 113
412 WIDTH, HEIGHT = 180, 100
413 FPS = 23.0
414 video_frames: list[Image.Image] = [
415 get_frame_with_text( # type: ignore[misc]
416 WIDTH, HEIGHT,
417 text=f"frame{frame_id:02d}",
418 font_size=24,
419 output_type="pil")
420 for frame_id in range(NUM_FRAMES)
421 ]
422 assert len(video_frames) == NUM_FRAMES
424 # Video without audio
425 video_path = await save_video_frames(
426 video_frames=video_frames,
427 fps=FPS)
428 video_file_info = get_video_file_info(video_path)
429 video_overall_info = video_file_info.get("overall")
430 assert video_overall_info is not None
431 num_bytes = video_overall_info.get("num_bytes")
432 assert num_bytes is not None
433 assert num_bytes > 1000
434 duration_seconds = video_overall_info.get("duration_seconds")
435 assert duration_seconds is not None
436 assert_approx(duration_seconds, 4.913)
438 assert "video" in video_file_info
439 video_info = video_file_info["video"]
440 assert video_info is not None
441 assert video_info.get("fps") == FPS, f"Video info: {video_file_info}"
442 assert video_info.get("num_frames") == 113, f"Video info: {video_file_info}"
443 assert video_info.get("width") == 180, f"Video info: {video_file_info}"
444 assert video_info.get("height") == 100, f"Video info: {video_file_info}"
445 video_duration_seconds = video_info.get("duration_seconds")
446 assert video_duration_seconds is not None
447 assert_approx(video_duration_seconds, 4.913) # 113 / 23
449 assert "audio" not in video_file_info
451 os.remove(video_path)
453 # Video + audio (video longer than audio)
454 video_path = await save_video_audio(
455 video_content=video_frames,
456 audio_path=audio_path,
457 fps=FPS)
458 video_file_info = get_video_file_info(video_path)
459 assert video_file_info is not None
460 video_overall_info = video_file_info.get("overall")
461 assert video_overall_info is not None
462 num_bytes = video_overall_info.get("num_bytes")
463 assert num_bytes is not None
464 assert num_bytes > 1000
465 duration_seconds = video_overall_info.get("duration_seconds")
466 assert duration_seconds is not None
467 assert_approx(duration_seconds, 4.913)
469 assert "video" in video_file_info
470 video_info = video_file_info["video"]
471 assert video_info is not None
472 assert video_info.get("fps") == FPS, f"Video info: {video_file_info}"
473 # TODO the following works with ffmpeg 4.2 but not with 6.0
474 """
475 assert video_info.get("num_frames") == 113, f"Video info: {video_file_info}"
476 assert_approx(
477 video_info.get("duration_seconds"), 4.913,
478 msg=f"Video info: {video_file_info}")
479 """
480 logging.warning(f"Video info: {video_file_info}")
481 num_frames = video_info.get("num_frames")
482 assert num_frames is not None
483 assert num_frames > 100
484 video_duration_seconds = video_info.get("duration_seconds")
485 assert video_duration_seconds is not None
486 assert video_duration_seconds > 4.0
488 assert "audio" in video_file_info
489 audio_info = video_file_info["audio"]
490 assert audio_info is not None
491 audio_duration_seconds = audio_info.get("duration_seconds")
492 assert audio_duration_seconds is not None
493 assert audio_duration_seconds > 0
494 """
495 # assert_approx(audio_info["duration_seconds"], 4.913)
496 assert_approx(
497 audio_duration_seconds, 4.864,
498 msg=f"Video info: {video_file_info}") # TODO this is not good
499 """
500 logging.info(f"Audio info: {audio_info}")
501 assert audio_duration_seconds > 4.0
503 os.remove(video_path)
506@pytest.mark.asyncio
507async def test_ffmpeg_version() -> None:
508 """Test getting the FFmpeg version."""
509 ffmpeg_version = get_ffmpeg_version()
510 logging.info(f"FFmpeg version: {get_ffmpeg_version()}")
511 assert ffmpeg_version is not None
514@pytest.mark.asyncio
515async def test_change_fps() -> None:
516 """Test changing the FPS of a video binary."""
517 NUM_FRAMES = 30
518 FPS = 30
519 video_frames = [
520 get_frame_with_text(320, 240, f"frame{frame_id:02d}", output_type="numpy")
521 for frame_id in range(NUM_FRAMES)
522 ]
523 video_path = await save_video_frames(video_frames, fps=FPS)
524 async with aiofiles.open(video_path, 'rb') as f:
525 video_binary = await f.read()
526 video_file_info = get_video_file_info(video_path)
527 assert "video" in video_file_info
528 video_info = video_file_info["video"]
529 assert video_info.get("fps") == FPS
530 assert video_info.get("num_frames") == 30
531 assert video_info.get("duration_seconds") == 1.0
532 video_num_frames = await get_video_num_frames(video_binary)
533 assert video_num_frames == 30
535 # Change FPS to 15
536 new_video_binary = change_video_fps(video_binary, 15)
537 new_video_size = len(new_video_binary)
538 assert new_video_size > 100
539 video_file_info = get_video_file_info(new_video_binary)
540 assert "video" in video_file_info
541 video_info = video_file_info["video"]
542 assert video_info.get("fps") == 15
543 assert video_info.get("num_frames") == 15
544 assert video_info.get("duration_seconds") == 1.0
545 assert get_video_fps(new_video_binary) == 15
547 with pytest.raises(TypeError):
548 change_video_fps(video_path, 15) # type: ignore[arg-type]
549 with pytest.raises(ValueError):
550 change_video_fps(video_binary, 0)
552 with pytest.raises(ValueError, match="No video frames provided."):
553 await save_video_frames([], fps=30)
555 with pytest.raises(ValueError, match="No video frames provided."):
556 await save_video_frames(None, fps=30) # type: ignore[arg-type]
558 with pytest.raises(TypeError):
559 get_video_fps(12345) # type: ignore[arg-type]
560 with pytest.raises(ValueError):
561 get_video_fps(b"")
563 os.remove(video_path)
564 del video_frames
565 del video_binary