ClusterVQA/run_tests.py at master · pekkavaa/ClusterVQA · GitHub

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
import os.path
import sys
import time
import glob
import io
import vqa
from vqa import quantizer
from vqa.encoders import PickleEncoder, HiColorVQAEncoder, VizEncoder
from vqa.video_types import *
from vqa.audiotrack import load_raw_s16le_file
from pickle_decoder import decode_pickled_video

datadir = "testdata/logo"
tempdir = "testoutput"
frame_paths = list(sorted(glob.glob(os.path.join(datadir, "*.png"))))
test_fps = 15
cfg = Config()
cfg.num_blocks = 100
cfg.max_vectors_to_fit = 10000
cfg.min_keyframe_distance = 30
cfg.max_keyframe_distance = 30
cfg.scene_cut_error_limit = 1e9  # disable scene cut detection


def pickle_end_to_end_test():
    class MemoryPickler:
        def __init__(self):
            self.data = bytearray()

        def write(self, stuff):
            self.data += stuff

    memory = MemoryPickler()
    encoder = PickleEncoder(memory, None)

    start = time.time()
    quantizer.encode_video(frame_paths, test_fps, cfg, encoder)
    took = time.time() - start
    print(f"Got {len(memory.data)} bytes")
    assert len(memory.data) > 1000
    assert encoder.frame_idx == 35

    if not os.path.isdir(tempdir):
        os.mkdir(tempdir)

    inputstream = io.BytesIO(memory.data)
    decode_pickled_video(inputstream, tempdir)

    for name in ["frame0001.png", "frame0012.png", "frame0025.png", "frame0035.png"]:
        source = quantizer.load_float_rgb(os.path.join(datadir, name))
        result = quantizer.load_float_rgb(os.path.join(tempdir, name))
        error = ((source - result) ** 2).mean()
        print(f"{name} MSE: {error}")
        assert error < 0.001

    print(f"Encoding took {took:.3f} s, {took/len(frame_paths):.4f} s/frame")
    print("Pickle end to end test OK")


def visual_inspection_test():
    """Not an automated test. It just shows the test encoding results."""

    encoder = VizEncoder(None, framerange=range(len(frame_paths)))
    quantizer.encode_video(frame_paths, test_fps, cfg, encoder)


def vqa_smoke_test():
    if not os.path.isdir(tempdir):
        os.mkdir(tempdir)

    videofile = os.path.join(tempdir, "logo.vqa")
    track = load_raw_s16le_file(
        "testdata/bleeps_22050_pcm_s16le.raw", rate=22050, is_stereo=True
    )

    if os.path.exists(videofile):
        os.remove(videofile)

    with open(videofile, "w") as _:
        pass

    with open(videofile, "rb+") as f:
        start = time.time()
        quantizer.encode_video(frame_paths, test_fps, cfg, HiColorVQAEncoder(f, track))
        took = time.time() - start

    assert os.path.exists(videofile), "VQA file should exist"
    assert (
        os.path.getsize(videofile) > 1024 * 10
    ), "VQA file should be larger than 10 KiB"

    print("VQA Smoke Test OK")


suites = ["default", "all", "visual"]


def print_usage():
    print("run_tests.py [SUITE]")
    print("Where SUITE is one of ", suites)


if __name__ == "__main__":
    suite = "default"
    if len(sys.argv) > 1:
        suite = sys.argv[1]

    if suite not in suites:
        print_usage()
        sys.exit(1)

    if suite == "default" or suite == "all":
        pickle_end_to_end_test()
        vqa_smoke_test()

    if suite == "all":
        vqa.block_ops.test_block_ops()
        vqa.vptr_format.test_vptr_encoding()
        vqa.lcw.test_compression()
        vqa.lcw.test_decompression()
        vqa.audiotrack.test_pcm_loading()

    if suite == "visual":
        visual_inspection_test()

    print(f"OK")