File size: 34,589 Bytes
318d421
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
#!/usr/bin/env python

# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Unit tests for ``lerobot.datasets.video_utils`` encoding functions and ``lerobot.configs.video.VideoEncoderConfig`` config class."""

import json
from pathlib import Path

import numpy as np
import pytest

pytest.importorskip("av", reason="av is required (install lerobot[dataset])")

import av  # noqa: E402

from lerobot.configs import VALID_VIDEO_CODECS, DepthEncoderConfig, RGBEncoderConfig, VideoEncoderConfig
from lerobot.datasets.image_writer import write_image
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.pyav_utils import get_codec
from lerobot.datasets.utils import INFO_PATH
from lerobot.datasets.video_utils import (
    concatenate_video_files,
    encode_video_frames,
    get_video_info,
    reencode_video,
)
from tests.fixtures.constants import (
    DUMMY_DEPTH_FEATURES,
    DUMMY_DEPTH_KEY,
    DUMMY_DEPTH_VIDEO_INFO_FULL,
    DUMMY_VIDEO_FEATURES,
    DUMMY_VIDEO_INFO,
    DUMMY_VIDEO_KEY,
)
from tests.fixtures.dataset_factories import add_frames


# Per-codec skip markers β€” validation tests only fire when the codec is available
def _require_encoder(vcodec: str) -> pytest.MarkDecorator:
    """Skip the test if ``vcodec`` is not available in the local FFmpeg build."""
    return pytest.mark.skipif(get_codec(vcodec) is None, reason=f"{vcodec!r} not in local FFmpeg build")


require_libsvtav1 = _require_encoder("libsvtav1")
require_h264 = _require_encoder("h264")
require_hevc = _require_encoder("hevc")
require_videotoolbox = _require_encoder("h264_videotoolbox")
require_nvenc = _require_encoder("h264_nvenc")
require_vaapi = _require_encoder("h264_vaapi")
require_qsv = _require_encoder("h264_qsv")


TEST_ARTIFACTS_DIR = Path(__file__).parent.parent / "artifacts" / "encoded_videos"


def _write_color_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None:
    imgs_dir.mkdir(parents=True, exist_ok=True)
    for i in range(num_frames):
        arr = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
        write_image(arr, imgs_dir / f"frame-{i:06d}.png")


def _write_depth_frames(imgs_dir: Path, num_frames: int = 4, height: int = 64, width: int = 96) -> None:
    """Write synthetic uint16 depth TIFFs (millimetres) for depth encoder tests.

    Uses a smooth linear ramp + per-frame offset (not white noise) so HEVC Main 12
    on ``gray12le`` compresses well. Values span ~100 mm to 10 m, covering most
    of the default ``[DEPTH_MIN, DEPTH_MAX]`` metres range after
    ``quantize_depth(input_unit="auto"="mm")``.
    """
    imgs_dir.mkdir(parents=True, exist_ok=True)
    base = np.linspace(100.0, 10_000.0, height * width, dtype=np.float32).reshape(height, width)
    for i in range(num_frames):
        arr = (base + 50.0 * i).clip(0, 65535).astype(np.uint16)
        write_image(arr, imgs_dir / f"frame-{i:06d}.tiff")


def _encode_video(
    path: Path,
    num_frames: int = 4,
    fps: int = 30,
    cfg: VideoEncoderConfig | None = None,
    depth: bool = False,
) -> Path:
    """Write synthetic frames to a temp dir and encode them to ``path``.

    ``depth=False`` writes uint8 RGB PNG noise and encodes with ``cfg``
    (defaulting to the library default). ``depth=True`` writes synthetic uint16
    depth TIFFs and encodes with ``cfg`` or a default :class:`DepthEncoderConfig`
    (HEVC Main 12 / ``gray12le``).
    """
    imgs_dir = path.parent / f"imgs_{path.stem}"
    if depth:
        _write_depth_frames(imgs_dir, num_frames=num_frames)
        cfg = cfg or DepthEncoderConfig()
    else:
        _write_color_frames(imgs_dir, num_frames=num_frames)
    encode_video_frames(imgs_dir, path, fps=fps, video_encoder=cfg, overwrite=True)
    return path


def _read_feature_info(dataset: LeRobotDataset, key: str = DUMMY_VIDEO_KEY) -> dict:
    info = json.loads((dataset.root / INFO_PATH).read_text())
    return info["features"][key]["info"]


# ─── RGBEncoderConfig / codec options ──────────────────────────────


class TestCodecOptions:
    @require_libsvtav1
    def test_libsvtav1_defaults(self):
        cfg = RGBEncoderConfig()
        opts = cfg.get_codec_options()
        assert opts["g"] == 2
        assert opts["crf"] == 30
        assert opts["preset"] == 12

    @require_libsvtav1
    def test_libsvtav1_custom_preset(self):
        cfg = RGBEncoderConfig(preset=8)
        assert cfg.get_codec_options()["preset"] == 8

    @require_h264
    def test_h264_options(self):
        cfg = RGBEncoderConfig(vcodec="h264", g=10, crf=23, preset=None)
        opts = cfg.get_codec_options()
        assert opts["g"] == 10
        assert opts["crf"] == 23
        assert "preset" not in opts

    @require_videotoolbox
    def test_videotoolbox_options(self):
        cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", g=2, crf=30, preset=None)
        opts = cfg.get_codec_options()
        assert opts["g"] == 2
        assert opts["q:v"] == 40
        assert "crf" not in opts

    @require_nvenc
    def test_nvenc_options(self):
        cfg = RGBEncoderConfig(vcodec="h264_nvenc", g=2, crf=25, preset=None)
        opts = cfg.get_codec_options()
        assert opts["rc"] == 0
        assert opts["qp"] == 25
        assert "crf" not in opts
        assert opts["g"] == 2

    @require_vaapi
    def test_vaapi_options(self):
        cfg = RGBEncoderConfig(vcodec="h264_vaapi", crf=28, preset=None)
        assert cfg.get_codec_options()["qp"] == 28

    @require_qsv
    def test_qsv_options(self):
        cfg = RGBEncoderConfig(vcodec="h264_qsv", crf=25, preset=None)
        assert cfg.get_codec_options()["global_quality"] == 25

    @require_h264
    def test_no_g_no_crf(self):
        cfg = RGBEncoderConfig(vcodec="h264", g=None, crf=None, preset=None)
        opts = cfg.get_codec_options()
        assert "g" not in opts
        assert "crf" not in opts

    @require_libsvtav1
    def test_encoder_threads_libsvtav1(self):
        cfg = RGBEncoderConfig(fast_decode=0)
        opts = cfg.get_codec_options(encoder_threads=4)
        assert "lp=4" in opts.get("svtav1-params", "")

    @require_h264
    def test_encoder_threads_h264(self):
        cfg = RGBEncoderConfig(vcodec="h264", preset=None)
        assert cfg.get_codec_options(encoder_threads=2)["threads"] == 2

    @require_libsvtav1
    def test_fast_decode_libsvtav1(self):
        cfg = RGBEncoderConfig(fast_decode=1)
        opts = cfg.get_codec_options()
        assert "fast-decode=1" in opts.get("svtav1-params", "")

    @require_libsvtav1
    def test_libsvtav1_fast_decode_clamped_to_svt_range(self):
        """Out-of-range fast_decode is clamped to [0, 2] in svtav1-params (SVT-AV1 FastDecode)."""
        cfg = RGBEncoderConfig(fast_decode=100)
        assert "fast-decode=2" in cfg.get_codec_options().get("svtav1-params", "")
        cfg_neg = RGBEncoderConfig(fast_decode=-5)
        assert "fast-decode=0" in cfg_neg.get_codec_options().get("svtav1-params", "")

    @require_h264
    def test_fast_decode_h264(self):
        cfg = RGBEncoderConfig(vcodec="h264", fast_decode=1, preset=None)
        assert cfg.get_codec_options()["tune"] == "fastdecode"

    @require_libsvtav1
    def test_pix_fmt_unsupported_raises(self):
        """Passing an unsupported pix_fmt is a hard error."""
        with pytest.raises(ValueError, match="pix_fmt"):
            RGBEncoderConfig(pix_fmt="yuv444p")  # libsvtav1 only supports yuv420p variants

    @require_libsvtav1
    @require_h264
    def test_preset_default_behaviour(self):
        """Empty constructor picks preset=12 (libsvtav1 path); other codecs stay None."""
        assert RGBEncoderConfig().preset == 12
        assert RGBEncoderConfig(vcodec="libsvtav1").preset == 12
        assert RGBEncoderConfig(vcodec="h264").preset is None
        assert RGBEncoderConfig(vcodec="h264", preset=None).preset is None

    @require_h264
    def test_preset_string_on_h264(self):
        """h264 accepts string presets and forwards them to FFmpeg."""
        cfg = RGBEncoderConfig(vcodec="h264", preset="slow")
        assert cfg.get_codec_options()["preset"] == "slow"

    @require_videotoolbox
    def test_preset_on_videotoolbox_not_set(self):
        """videotoolbox has no preset option at all."""
        cfg = RGBEncoderConfig(vcodec="h264_videotoolbox", preset="slow")
        assert "preset" not in cfg.get_codec_options()

    @require_libsvtav1
    def test_libsvtav1_preset_out_of_range_raises(self):
        """libsvtav1 preset must sit in [-2, 13] as exposed by PyAV."""
        with pytest.raises(ValueError, match="out of range"):
            RGBEncoderConfig(vcodec="libsvtav1", preset=100)
        with pytest.raises(ValueError, match="out of range"):
            RGBEncoderConfig(vcodec="libsvtav1", preset=-3)

    @require_libsvtav1
    def test_libsvtav1_crf_out_of_range_raises(self):
        """libsvtav1 crf must sit in [0, 63]."""
        with pytest.raises(ValueError, match="crf.*out of range"):
            RGBEncoderConfig(vcodec="libsvtav1", crf=64)

    @require_libsvtav1
    def test_libsvtav1_crf_rejects_python_float(self):
        """libsvtav1 exposes ``crf`` as an INT AVOption; Python float must not pass validation."""
        with pytest.raises(ValueError, match="float values are not allowed"):
            RGBEncoderConfig(vcodec="libsvtav1", crf=2.5)

    @require_libsvtav1
    def test_libsvtav1_extra_crf_rejects_fractional_string(self):
        """INT options reject fractional values even when supplied only via ``extra_options``."""
        with pytest.raises(ValueError, match="float values are not allowed"):
            RGBEncoderConfig(
                vcodec="libsvtav1",
                crf=None,
                extra_options={"crf": "2.5"},
            )

    @require_libsvtav1
    def test_libsvtav1_extra_crf_rejects_float(self):
        with pytest.raises(ValueError, match="float values are not allowed"):
            RGBEncoderConfig(
                vcodec="libsvtav1",
                crf=None,
                extra_options={"crf": 2.5},
            )

    @require_h264
    def test_h264_crf_accepts_float_and_int(self):
        """x264 exposes crf as a FLOAT option, so both int and float are accepted."""
        assert RGBEncoderConfig(vcodec="h264", crf=23).get_codec_options()["crf"] == 23
        assert RGBEncoderConfig(vcodec="h264", crf=23.5).get_codec_options()["crf"] == 23.5

    @require_libsvtav1
    def test_validate_is_rerunnable(self):
        """After mutating a field, validate() re-checks and surfaces new issues."""
        cfg = RGBEncoderConfig(vcodec="libsvtav1")
        cfg.preset = 100  # now out of range
        with pytest.raises(ValueError, match="out of range"):
            cfg.validate()


class TestExtraOptions:
    @require_libsvtav1
    def test_default_is_empty_dict(self):
        cfg = RGBEncoderConfig()
        assert cfg.extra_options == {}

    @require_libsvtav1
    def test_unknown_key_passes_through(self):
        """Keys not published as AVOptions are forwarded to FFmpeg."""
        cfg = RGBEncoderConfig(extra_options={"totally_made_up_option": "value"})
        assert cfg.extra_options == {"totally_made_up_option": "value"}

    @require_libsvtav1
    def test_numeric_value_in_range_ok(self):
        """libsvtav1 exposes ``qp`` as INT in [0, 63]."""
        cfg = RGBEncoderConfig(extra_options={"qp": 30})
        assert cfg.extra_options == {"qp": 30}

    @require_libsvtav1
    def test_numeric_out_of_range_raises(self):
        with pytest.raises(ValueError, match=r"qp=.*out of range"):
            RGBEncoderConfig(extra_options={"qp": 999})

    @require_libsvtav1
    def test_numeric_string_accepted_in_range(self):
        """Numeric strings are accepted for numeric options (mirrors FFmpeg)."""
        cfg = RGBEncoderConfig(extra_options={"qp": "18"})
        assert cfg.extra_options == {"qp": "18"}

    @require_libsvtav1
    def test_numeric_string_out_of_range_raises(self):
        with pytest.raises(ValueError, match=r"qp=.*out of range"):
            RGBEncoderConfig(extra_options={"qp": "999"})

    @require_libsvtav1
    def test_non_numeric_string_on_numeric_option_raises(self):
        with pytest.raises(ValueError, match=r"qp=.*not numeric"):
            RGBEncoderConfig(extra_options={"qp": "medium"})

    @require_libsvtav1
    def test_bool_on_numeric_option_raises(self):
        """``bool`` is explicitly rejected for numeric options."""
        with pytest.raises(ValueError, match=r"qp=.*not numeric"):
            RGBEncoderConfig(extra_options={"qp": True})

    @require_h264
    def test_string_option_passes_through_unchecked(self):
        """String-typed AVOptions are NOT enum-checked (too many accept freeform)."""
        cfg = RGBEncoderConfig(vcodec="h264", preset=None, extra_options={"tune": "some-future-tune"})
        assert cfg.extra_options == {"tune": "some-future-tune"}

    @require_libsvtav1
    def test_merged_into_codec_options_and_stringified(self):
        """Typed merge by default; ``as_strings=True`` matches FFmpeg option dict."""
        cfg = RGBEncoderConfig(extra_options={"qp": 20})
        opts = cfg.get_codec_options()
        assert opts["qp"] == 20
        assert isinstance(opts["qp"], int)
        str_opts = cfg.get_codec_options(as_strings=True)
        assert str_opts["qp"] == "20"
        assert all(isinstance(v, str) for v in str_opts.values())

    @require_libsvtav1
    def test_structured_fields_win_on_collision(self):
        """A colliding extra_options key is discarded; the structured field wins."""
        cfg = RGBEncoderConfig(crf=30, extra_options={"crf": 18})
        assert cfg.get_codec_options()["crf"] == 30


class TestEncoderDetection:
    @require_h264
    def test_explicit_codec_kept_when_available(self):
        cfg = RGBEncoderConfig(vcodec="h264")
        assert cfg.vcodec == "h264"

    @require_videotoolbox
    def test_auto_picks_videotoolbox_when_available(self):
        """``h264_videotoolbox`` sits at the top of ``HW_VIDEO_CODECS`` so it wins when present."""
        cfg = RGBEncoderConfig(vcodec="auto")
        assert cfg.vcodec == "h264_videotoolbox"

    def test_invalid_codec_raises(self):
        with pytest.raises(ValueError, match="Invalid vcodec"):
            RGBEncoderConfig(vcodec="not_a_real_codec")

    def test_hw_encoder_names_listed_as_valid(self):
        assert "auto" in VALID_VIDEO_CODECS
        assert "h264_videotoolbox" in VALID_VIDEO_CODECS
        assert "h264_nvenc" in VALID_VIDEO_CODECS


class TestGetVideoInfo:
    def test_returns_all_stream_fields(self):
        info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4")

        assert info["video.height"] == 64
        assert info["video.width"] == 96
        assert info["video.pix_fmt"] == "yuv420p"
        assert info["video.fps"] == 30
        assert info["video.channels"] == 3
        assert info["is_depth_map"] is False
        assert info["has_audio"] is False
        assert "video.g" not in info
        assert "video.crf" not in info
        assert "video.preset" not in info

    @require_libsvtav1
    def test_merges_encoder_config_as_video_prefixed_entries(self):
        cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)

        info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg)

        assert info["video.g"] == 2
        assert info["video.crf"] == 30
        assert info["video.preset"] == 12
        assert info["video.fast_decode"] == 0
        assert info["video.video_backend"] == "pyav"
        assert info["video.extra_options"] == {}

    @require_libsvtav1
    def test_stream_derived_keys_take_precedence_over_config(self):
        cfg = RGBEncoderConfig(vcodec="libsvtav1", pix_fmt="yuv420p")

        info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=cfg)

        assert info["video.codec"]  # populated from stream, not from config's vcodec
        assert info["video.pix_fmt"] == "yuv420p"

    def test_depth_encoder_config_sets_is_depth_map_true(self):
        """A ``DepthEncoderConfig`` causes ``get_video_info`` to mark the stream as depth."""
        info = get_video_info(TEST_ARTIFACTS_DIR / "clip_4frames.mp4", video_encoder=DepthEncoderConfig())
        assert info["is_depth_map"] is True


class TestEncodeVideoFrames:
    @require_libsvtav1
    def test_produces_readable_mp4(self, tmp_path):
        video_path = _encode_video(tmp_path / "out.mp4")

        assert video_path.exists()
        info = get_video_info(video_path)
        assert info["video.height"] == 64
        assert info["video.width"] == 96

    @require_libsvtav1
    def test_frame_count_and_duration_match_input(self, tmp_path):
        num_frames = 10
        fps = 30
        video_path = _encode_video(tmp_path / "out.mp4", num_frames=num_frames, fps=fps)

        with av.open(str(video_path)) as container:
            stream = container.streams.video[0]
            actual_frames = sum(1 for _ in container.decode(stream))
            duration = (
                float(stream.duration * stream.time_base)
                if stream.duration is not None
                else float(container.duration / av.time_base)
            )

        assert actual_frames == num_frames
        assert abs(duration - num_frames / fps) < 0.1

    def test_overwrite_false_skips_existing_file(self, tmp_path):
        imgs_dir = tmp_path / "imgs"
        _write_color_frames(imgs_dir)
        video_path = tmp_path / "out.mp4"
        sentinel = b"pre-existing content"
        video_path.write_bytes(sentinel)

        encode_video_frames(imgs_dir, video_path, fps=30, overwrite=False)

        assert video_path.read_bytes() == sentinel

    @require_libsvtav1
    def test_overwrite_true_replaces_existing_file(self, tmp_path):
        imgs_dir = tmp_path / "imgs"
        _write_color_frames(imgs_dir)
        video_path = tmp_path / "out.mp4"
        video_path.write_bytes(b"stale content")

        encode_video_frames(imgs_dir, video_path, fps=30, overwrite=True)

        info = get_video_info(video_path)
        assert info["video.height"] == 64

    @require_libsvtav1
    def test_custom_encoder_config_fields_stored_in_info(self, tmp_path):
        """All stream-derived and encoder config fields are present after encoding."""
        cfg = RGBEncoderConfig(vcodec="libsvtav1", g=4, crf=25, preset=10)
        video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg)

        info = get_video_info(video_path, video_encoder=cfg)

        # Stream-derived
        assert info["video.height"] == 64
        assert info["video.width"] == 96
        assert info["video.channels"] == 3
        assert info["video.codec"] == "av1"
        assert info["video.pix_fmt"] == "yuv420p"
        assert info["video.fps"] == 30
        assert info["is_depth_map"] is False
        assert info["has_audio"] is False
        # Encoder config
        assert info["video.g"] == 4
        assert info["video.crf"] == 25
        assert info["video.preset"] == 10
        assert info["video.fast_decode"] == 0
        assert info["video.video_backend"] == "pyav"
        assert info["video.extra_options"] == {}


class TestReencodeVideo:
    @require_libsvtav1
    @require_h264
    def test_reencode_video(self, tmp_path):
        src = TEST_ARTIFACTS_DIR / "clip_4frames.mp4"
        out = tmp_path / "reencoded.mp4"
        cfg = RGBEncoderConfig(vcodec="h264", g=6, crf=23, pix_fmt="yuv444p")
        reencode_video(src, out, video_encoder=cfg, overwrite=True)

        assert out.exists()
        with av.open(str(out)) as container:
            n_frames = sum(1 for _ in container.decode(video=0))
        assert n_frames == 4

        info = get_video_info(out, video_encoder=cfg)
        assert info["video.codec"] == "h264"
        assert info["video.pix_fmt"] == "yuv444p"
        assert info["video.height"] == 64
        assert info["video.width"] == 96
        assert info["video.fps"] == 30
        assert info["video.g"] == 6
        assert info["video.crf"] == 23

    @require_h264
    def test_reencode_video_trim_window(self, tmp_path):
        src = TEST_ARTIFACTS_DIR / "clip_6frames.mp4"
        out = tmp_path / "trim_window.mp4"
        cfg = RGBEncoderConfig(vcodec="h264")
        reencode_video(src, out, video_encoder=cfg, start_time_s=0.05, end_time_s=0.12, overwrite=True)

        with av.open(str(out)) as container:
            frames = list(container.decode(video=0))
        # Only the frames at 0.067 and 0.1 s fall inside [0.05, 0.12).
        assert len(frames) == 2
        assert frames[0].time == pytest.approx(0.0, abs=1e-3)


class TestConcatenateVideoFiles:
    def test_two_clips_frame_count(self, tmp_path):
        """Output frame count equals the sum of the two input frame counts."""
        out = tmp_path / "out.mp4"
        concatenate_video_files(
            [TEST_ARTIFACTS_DIR / "clip_6frames.mp4", TEST_ARTIFACTS_DIR / "clip_4frames.mp4"], out
        )

        with av.open(str(out)) as container:
            total = sum(1 for _ in container.decode(video=0))
        assert total == 10

    def test_three_clips_frame_count(self, tmp_path):
        out = tmp_path / "out.mp4"
        clip = TEST_ARTIFACTS_DIR / "clip_5frames.mp4"
        concatenate_video_files([clip, clip, clip], out)

        with av.open(str(out)) as container:
            total = sum(1 for _ in container.decode(video=0))
        assert total == 15

    @require_libsvtav1
    def test_geometry_preserved(self, tmp_path):
        """Output resolution, fps, codec and pixel format must match the inputs."""
        out = tmp_path / "out.mp4"
        concatenate_video_files(
            [TEST_ARTIFACTS_DIR / "clip_4frames.mp4", TEST_ARTIFACTS_DIR / "clip_4frames.mp4"], out
        )

        info = get_video_info(out)
        assert info["video.height"] == 64
        assert info["video.width"] == 96
        assert info["video.fps"] == 30
        assert info["video.codec"] == "av1"
        assert info["video.pix_fmt"] == "yuv420p"

    def test_compatibility_check_raises_on_different_codec(self, tmp_path):
        with pytest.raises(ValueError):
            concatenate_video_files(
                [TEST_ARTIFACTS_DIR / "clip_4frames.mp4", TEST_ARTIFACTS_DIR / "clip_h264.mp4"],
                tmp_path / "out.mp4",
                compatibility_check=True,
            )

    def test_compatibility_check_raises_on_different_resolution(self, tmp_path):
        with pytest.raises(ValueError):
            concatenate_video_files(
                [TEST_ARTIFACTS_DIR / "clip_4frames.mp4", TEST_ARTIFACTS_DIR / "clip_32x48.mp4"],
                tmp_path / "out.mp4",
                compatibility_check=True,
            )


class TestEncoderConfigPersistence:
    """Encoder config must be stored as ``video.<field>`` entries in
    ``info["features"][key]["info"]`` when the first episode is saved.
    """

    @require_libsvtav1
    def test_first_episode_save_persists_encoder_config(self, tmp_path, empty_lerobot_dataset_factory):
        cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
        dataset = empty_lerobot_dataset_factory(
            root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg
        )

        add_frames(dataset, num_frames=4)
        dataset.save_episode()
        dataset.finalize()

        info = _read_feature_info(dataset)

        assert info["video.height"] == 64
        assert info["video.width"] == 96
        assert info["video.fps"] == 30
        assert info["video.g"] == 2
        assert info["video.crf"] == 30
        assert info["video.preset"] == 12
        assert info["video.fast_decode"] == 0
        assert info["video.video_backend"] == "pyav"
        assert info["video.extra_options"] == {}

    @require_libsvtav1
    def test_second_episode_does_not_overwrite_encoder_fields(self, tmp_path, empty_lerobot_dataset_factory):
        cfg = RGBEncoderConfig(vcodec="libsvtav1", g=2, crf=30, preset=12)
        dataset = empty_lerobot_dataset_factory(
            root=tmp_path / "ds", features=DUMMY_VIDEO_FEATURES, use_videos=True, rgb_encoder=cfg
        )

        add_frames(dataset, num_frames=4)
        dataset.save_episode()
        first_info = dict(_read_feature_info(dataset))

        add_frames(dataset, num_frames=4)
        dataset.save_episode()
        dataset.finalize()

        assert _read_feature_info(dataset) == first_info


class TestFromVideoInfo:
    """``RGBEncoderConfig.from_video_info`` reconstructs an encoder config
    from the ``video.*`` keys persisted in a dataset's ``info.json``.
    """

    @require_libsvtav1
    def test_reconstructs_from_dummy_video_info(self):
        cfg = RGBEncoderConfig.from_video_info(DUMMY_VIDEO_INFO)

        # Canonical stream codec ``"av1"`` is aliased to the encoder name.
        assert cfg.vcodec == "libsvtav1"
        assert cfg.pix_fmt == DUMMY_VIDEO_INFO["video.pix_fmt"]
        assert cfg.g == DUMMY_VIDEO_INFO["video.g"]
        assert cfg.crf == DUMMY_VIDEO_INFO["video.crf"]
        assert cfg.preset == DUMMY_VIDEO_INFO["video.preset"]
        assert cfg.fast_decode == DUMMY_VIDEO_INFO["video.fast_decode"]
        assert cfg.video_backend == DUMMY_VIDEO_INFO["video.video_backend"]
        # ``{}`` placeholder (typical after a merge with disagreeing sources)
        # must not leak into the reconstructed config.
        assert cfg.extra_options == RGBEncoderConfig().extra_options


# ─── Depth-specific encoding tests ────────────────────────────────────


class TestEncodeDepthVideoFrames:
    """Depth mirror of :class:`TestEncodeVideoFrames`.

    Exercises ``encode_video_frames`` end-to-end through
    :class:`DepthEncoderConfig` (HEVC Main 12 / ``gray12le``) on synthetic
    uint16 depth TIFFs.
    """

    @require_hevc
    def test_produces_readable_file(self, tmp_path):
        video_path = _encode_video(tmp_path / "out.mp4", depth=True)

        assert video_path.exists()
        info = get_video_info(video_path, video_encoder=DepthEncoderConfig())
        assert info["video.height"] == 64
        assert info["video.width"] == 96
        assert info["video.codec"] == "hevc"
        assert info["video.pix_fmt"] == "gray12le"
        assert info["video.channels"] == 1
        assert info["is_depth_map"] is True

    @require_hevc
    def test_frame_count_and_duration_match_input(self, tmp_path):
        num_frames = 10
        fps = 30
        video_path = _encode_video(tmp_path / "out.mp4", num_frames=num_frames, fps=fps, depth=True)

        with av.open(str(video_path)) as container:
            stream = container.streams.video[0]
            actual_frames = sum(1 for _ in container.decode(stream))
            duration = (
                float(stream.duration * stream.time_base)
                if stream.duration is not None
                else float(container.duration / av.time_base)
            )

        assert actual_frames == num_frames
        assert abs(duration - num_frames / fps) < 0.1

    def test_overwrite_false_skips_existing_file(self, tmp_path):
        """Codec-agnostic: file-system semantics must hold even without an HEVC encoder."""
        imgs_dir = tmp_path / "imgs"
        _write_depth_frames(imgs_dir)
        video_path = tmp_path / "out.mp4"
        sentinel = b"pre-existing depth content"
        video_path.write_bytes(sentinel)

        encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=False)

        assert video_path.read_bytes() == sentinel

    @require_hevc
    def test_overwrite_true_replaces_existing_file(self, tmp_path):
        imgs_dir = tmp_path / "imgs"
        _write_depth_frames(imgs_dir)
        video_path = tmp_path / "out.mp4"
        video_path.write_bytes(b"stale content")

        encode_video_frames(imgs_dir, video_path, fps=30, video_encoder=DepthEncoderConfig(), overwrite=True)

        info = get_video_info(video_path, video_encoder=DepthEncoderConfig())
        assert info["video.height"] == 64
        assert info["video.pix_fmt"] == "gray12le"
        assert info["is_depth_map"] is True

    @require_hevc
    def test_custom_encoder_config_fields_stored_in_info(self, tmp_path):
        """All stream-derived and depth-encoder config fields are present after encoding."""
        cfg = DepthEncoderConfig(
            vcodec="hevc",
            pix_fmt="gray12le",
            g=4,
            crf=25,
            extra_options={},
            depth_min=0.05,
            depth_max=8.0,
            shift=2.5,
            use_log=False,
        )
        video_path = _encode_video(tmp_path / "out.mp4", num_frames=4, fps=30, cfg=cfg, depth=True)

        info = get_video_info(video_path, video_encoder=cfg)

        # Stream-derived
        assert info["video.height"] == 64
        assert info["video.width"] == 96
        assert info["video.channels"] == 1
        assert info["video.codec"] == "hevc"
        assert info["video.pix_fmt"] == "gray12le"
        assert info["video.fps"] == 30
        assert info["is_depth_map"] is True
        assert info["has_audio"] is False
        # Base encoder config
        assert info["video.g"] == 4
        assert info["video.crf"] == 25
        assert info["video.fast_decode"] == 0
        assert info["video.video_backend"] == "pyav"
        assert info["video.extra_options"] == {}
        # Depth-specific tuning
        assert info["video.depth_min"] == 0.05
        assert info["video.depth_max"] == 8.0
        assert info["video.shift"] == 2.5
        assert info["video.use_log"] is False


class TestDepthEncoderConfigPersistence:
    """Depth mirror of :class:`TestEncoderConfigPersistence`.

    ``DepthEncoderConfig`` must be stored as ``video.<field>`` entries
    (including the depth-specific ``depth_min`` / ``depth_max`` / ``shift`` /
    ``use_log``) under ``info["features"][<depth_key>]["info"]`` when the
    first episode is saved.
    """

    @require_hevc
    def test_first_episode_save_persists_depth_encoder_config(self, tmp_path, empty_lerobot_dataset_factory):
        cfg = DepthEncoderConfig(
            vcodec="hevc",
            pix_fmt="gray12le",
            g=2,
            crf=30,
            extra_options={},
            depth_min=0.05,
            depth_max=8.0,
            shift=2.5,
            use_log=False,
        )
        dataset = empty_lerobot_dataset_factory(
            root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg
        )

        add_frames(dataset, num_frames=4)
        dataset.save_episode()
        dataset.finalize()

        info = _read_feature_info(dataset, key=DUMMY_DEPTH_KEY)

        # Stream-derived
        assert info["video.height"] == 64
        assert info["video.width"] == 96
        assert info["video.fps"] == 30
        assert info["video.codec"] == "hevc"
        assert info["video.pix_fmt"] == "gray12le"
        assert info["is_depth_map"] is True
        # Base encoder config
        assert info["video.g"] == 2
        assert info["video.crf"] == 30
        assert info["video.fast_decode"] == 0
        assert info["video.video_backend"] == "pyav"
        assert info["video.extra_options"] == {}
        # Depth-specific tuning
        assert info["video.depth_min"] == 0.05
        assert info["video.depth_max"] == 8.0
        assert info["video.shift"] == 2.5
        assert info["video.use_log"] is False

    @require_hevc
    def test_second_episode_does_not_overwrite_depth_encoder_fields(
        self, tmp_path, empty_lerobot_dataset_factory
    ):
        cfg = DepthEncoderConfig(
            vcodec="hevc",
            pix_fmt="gray12le",
            g=2,
            crf=30,
            depth_min=0.05,
            depth_max=8.0,
            shift=2.5,
            use_log=False,
        )
        dataset = empty_lerobot_dataset_factory(
            root=tmp_path / "ds", features=DUMMY_DEPTH_FEATURES, use_videos=True, depth_encoder=cfg
        )

        add_frames(dataset, num_frames=4)
        dataset.save_episode()
        first_info = dict(_read_feature_info(dataset, key=DUMMY_DEPTH_KEY))

        add_frames(dataset, num_frames=4)
        dataset.save_episode()
        dataset.finalize()

        assert _read_feature_info(dataset, key=DUMMY_DEPTH_KEY) == first_info


class TestDepthFromVideoInfo:
    """``DepthEncoderConfig.from_video_info`` reconstructs a depth encoder
    config from the ``video.*`` keys persisted in a dataset's ``info.json``.

    Depth mirror of :class:`TestFromVideoInfo`.
    """

    @require_hevc
    def test_reconstructs_from_dummy_depth_video_info(self):
        cfg = DepthEncoderConfig.from_video_info(DUMMY_DEPTH_VIDEO_INFO_FULL)

        # No alias for ``"hevc"``; the canonical stream codec is reused as-is.
        assert cfg.vcodec == "hevc"
        assert cfg.pix_fmt == DUMMY_DEPTH_VIDEO_INFO_FULL["video.pix_fmt"]
        assert cfg.g == DUMMY_DEPTH_VIDEO_INFO_FULL["video.g"]
        assert cfg.crf == DUMMY_DEPTH_VIDEO_INFO_FULL["video.crf"]
        assert cfg.fast_decode == DUMMY_DEPTH_VIDEO_INFO_FULL["video.fast_decode"]
        assert cfg.video_backend == DUMMY_DEPTH_VIDEO_INFO_FULL["video.video_backend"]
        # ``{}`` placeholder (typical after a merge with disagreeing sources)
        # must not leak into the reconstructed config.
        assert cfg.extra_options == DepthEncoderConfig().extra_options
        # Depth-specific tuning round-trips through ``info.json``.
        assert cfg.depth_min == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_min"]
        assert cfg.depth_max == DUMMY_DEPTH_VIDEO_INFO_FULL["video.depth_max"]
        assert cfg.shift == DUMMY_DEPTH_VIDEO_INFO_FULL["video.shift"]
        assert cfg.use_log == DUMMY_DEPTH_VIDEO_INFO_FULL["video.use_log"]