1#!/usr/bin/env python3
2# SPDX-License-Identifier: Apache-2.0
3# -----------------------------------------------------------------------------
4# Copyright 2019-2023 Arm Limited
5#
6# Licensed under the Apache License, Version 2.0 (the "License"); you may not
7# use this file except in compliance with the License. You may obtain a copy
8# of the License at:
9#
10#     http://www.apache.org/licenses/LICENSE-2.0
11#
12# Unless required by applicable law or agreed to in writing, software
13# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
14# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
15# License for the specific language governing permissions and limitations
16# under the License.
17# -----------------------------------------------------------------------------
18"""
19The image test runner is used for image quality and performance testing.
20
21It is designed to process directories of arbitrary test images, using the
22directory structure and path naming conventions to self-describe how each image
23is to be compressed. Some built-in test sets are provided in the ./Test/Images
24directory, and others can be downloaded by running the astc_test_image_dl
25script.
26
27Attributes:
28    RESULT_THRESHOLD_WARN: The result threshold (dB) for getting a WARN.
29    RESULT_THRESHOLD_FAIL: The result threshold (dB) for getting a FAIL.
30    TEST_BLOCK_SIZES: The block sizes we can test. This is a subset of the
31        block sizes supported by ASTC, simply to keep test run times
32        manageable.
33"""
34
35import argparse
36import os
37import platform
38import sys
39
40import testlib.encoder as te
41import testlib.testset as tts
42import testlib.resultset as trs
43
44# Require bit exact with reference scores
45RESULT_THRESHOLD_WARN = -0.00
46RESULT_THRESHOLD_FAIL = -0.00
47RESULT_THRESHOLD_3D_FAIL = -0.00
48
49
50TEST_BLOCK_SIZES = ["4x4", "5x5", "6x6", "8x8", "12x12", "3x3x3", "6x6x6"]
51
52TEST_QUALITIES = ["fastest", "fast", "medium", "thorough", "verythorough", "exhaustive"]
53
54
55def is_3d(blockSize):
56    """
57    Is the given block size a 3D block type?
58
59    Args:
60        blockSize (str): The block size.
61
62    Returns:
63        bool: ``True`` if the block string is a 3D block size, ``False`` if 2D.
64    """
65    return blockSize.count("x") == 2
66
67
68def count_test_set(testSet, blockSizes):
69    """
70    Count the number of test executions needed for a test set.
71
72    Args:
73        testSet (TestSet): The test set to run.
74        blockSizes (list(str)): The block sizes to run.
75
76    Returns:
77        int: The number of test executions needed.
78    """
79    count = 0
80    for blkSz in blockSizes:
81        for image in testSet.tests:
82            # 3D block sizes require 3D images
83            if is_3d(blkSz) != image.is3D:
84                continue
85
86            count += 1
87
88    return count
89
90
91def determine_result(image, reference, result):
92    """
93    Determine a test result against a reference and thresholds.
94
95    Args:
96        image (TestImage): The image being compressed.
97        reference (Record): The reference result to compare against.
98        result (Record): The test result.
99
100    Returns:
101        Result: The result code.
102    """
103    dPSNR = result.psnr - reference.psnr
104
105    if (dPSNR < RESULT_THRESHOLD_FAIL) and (not image.is3D):
106        return trs.Result.FAIL
107
108    if (dPSNR < RESULT_THRESHOLD_3D_FAIL) and image.is3D:
109        return trs.Result.FAIL
110
111    if dPSNR < RESULT_THRESHOLD_WARN:
112        return trs.Result.WARN
113
114    return trs.Result.PASS
115
116
117def format_solo_result(image, result):
118    """
119    Format a metrics string for a single (no compare) result.
120
121    Args:
122        image (TestImage): The image being tested.
123        result (Record): The test result.
124
125    Returns:
126        str: The metrics string.
127    """
128    name = "%5s %s" % (result.blkSz, result.name)
129    tPSNR = "%2.3f dB" % result.psnr
130    tTTime = "%.3f s" % result.tTime
131    tCTime = "%.3f s" % result.cTime
132    tCMTS = "%.3f MT/s" % result.cRate
133
134    return "%-32s | %8s | %9s | %9s | %11s" % \
135        (name, tPSNR, tTTime, tCTime, tCMTS)
136
137
138def format_result(image, reference, result):
139    """
140    Format a metrics string for a comparison result.
141
142    Args:
143        image (TestImage): The image being tested.
144        reference (Record): The reference result to compare against.
145        result (Record): The test result.
146
147    Returns:
148        str: The metrics string.
149    """
150    dPSNR = result.psnr - reference.psnr
151    sTTime = reference.tTime / result.tTime
152    sCTime = reference.cTime / result.cTime
153
154    name  = "%5s %s" % (result.blkSz, result.name)
155    tPSNR = "%2.3f dB (% 1.3f dB)" % (result.psnr, dPSNR)
156    tTTime = "%.3f s (%1.2fx)" % (result.tTime, sTTime)
157    tCTime = "%.3f s (%1.2fx)" % (result.cTime, sCTime)
158    tCMTS = "%.3f MT/s" % (result.cRate)
159    result = determine_result(image, reference, result)
160
161    return "%-32s | %22s | %15s | %15s | %11s | %s" % \
162           (name, tPSNR, tTTime, tCTime, tCMTS, result.name)
163
164
165def run_test_set(encoder, testRef, testSet, quality, blockSizes, testRuns,
166                 keepOutput, threads):
167    """
168    Execute all tests in the test set.
169
170    Args:
171        encoder (EncoderBase): The encoder to use.
172        testRef (ResultSet): The test reference results.
173        testSet (TestSet): The test set.
174        quality (str): The quality level to execute the test against.
175        blockSizes (list(str)): The block sizes to execute each test against.
176        testRuns (int): The number of test repeats to run for each image test.
177        keepOutput (bool): Should the test preserve output images? This is
178            only a hint and discarding output may be ignored if the encoder
179            version used can't do it natively.
180        threads (int or None): The thread count to use.
181
182    Returns:
183        ResultSet: The test results.
184    """
185    resultSet = trs.ResultSet(testSet.name)
186
187    curCount = 0
188    maxCount = count_test_set(testSet, blockSizes)
189
190    dat = (testSet.name, encoder.name, quality)
191    title = "Test Set: %s / Encoder: %s -%s" % dat
192    print(title)
193    print("=" * len(title))
194
195    for blkSz in blockSizes:
196        for image in testSet.tests:
197            # 3D block sizes require 3D images
198            if is_3d(blkSz) != image.is3D:
199                continue
200
201            curCount += 1
202
203            dat = (curCount, maxCount, blkSz, image.testFile)
204            print("Running %u/%u %s %s ... " % dat, end='', flush=True)
205            res = encoder.run_test(image, blkSz, "-%s" % quality, testRuns,
206                                   keepOutput, threads)
207            res = trs.Record(blkSz, image.testFile, res[0], res[1], res[2], res[3])
208            resultSet.add_record(res)
209
210            if testRef:
211                refResult = testRef.get_matching_record(res)
212                res.set_status(determine_result(image, refResult, res))
213
214                res.tTimeRel = refResult.tTime / res.tTime
215                res.cTimeRel = refResult.cTime / res.cTime
216                res.psnrRel = res.psnr - refResult.psnr
217
218                res = format_result(image, refResult, res)
219            else:
220                res = format_solo_result(image, res)
221
222            print("\r[%3u] %s" % (curCount, res))
223
224    return resultSet
225
226
227def get_encoder_params(encoderName, referenceName, imageSet):
228    """
229    The the encoder and image set parameters for a test run.
230
231    Args:
232        encoderName (str): The encoder name.
233        referenceName (str): The reference encoder name.
234        imageSet (str): The test image set.
235
236    Returns:
237        tuple(EncoderBase, str, str, str): The test parameters for the
238        requested encoder and test set. An instance of the encoder wrapper
239        class, the output data name, the output result directory, and the
240        reference to use.
241    """
242    # 1.7 variants
243    if encoderName == "ref-1.7":
244        encoder = te.Encoder1_7()
245        name = "reference-1.7"
246        outDir = "Test/Images/%s" % imageSet
247        refName = None
248        return (encoder, name, outDir, refName)
249
250    if encoderName.startswith("ref"):
251        _, version, simd = encoderName.split("-")
252
253        # 2.x, 3.x, and 4.x variants
254        compatible2xPrefixes = ["2.", "3.", "4."]
255        if any(True for x in compatible2xPrefixes if version.startswith(x)):
256            encoder = te.Encoder2xRel(version, simd)
257            name = f"reference-{version}-{simd}"
258            outDir = "Test/Images/%s" % imageSet
259            refName = None
260            return (encoder, name, outDir, refName)
261
262        # Latest main
263        if version == "main":
264            encoder = te.Encoder2x(simd)
265            name = f"reference-{version}-{simd}"
266            outDir = "Test/Images/%s" % imageSet
267            refName = None
268            return (encoder, name, outDir, refName)
269
270        assert False, f"Encoder {encoderName} not recognized"
271
272    encoder = te.Encoder2x(encoderName)
273    name = "develop-%s" % encoderName
274    outDir = "TestOutput/%s" % imageSet
275    refName = referenceName.replace("ref", "reference")
276    return (encoder, name, outDir, refName)
277
278
279def parse_command_line():
280    """
281    Parse the command line.
282
283    Returns:
284        Namespace: The parsed command line container.
285    """
286    parser = argparse.ArgumentParser()
287
288    # All reference encoders
289    refcoders = ["ref-1.7",
290                 "ref-2.5-neon", "ref-2.5-sse2", "ref-2.5-sse4.1", "ref-2.5-avx2",
291                 "ref-3.7-neon", "ref-3.7-sse2", "ref-3.7-sse4.1", "ref-3.7-avx2",
292                 "ref-4.4-neon", "ref-4.4-sse2", "ref-4.4-sse4.1", "ref-4.4-avx2",
293                 "ref-4.5-neon", "ref-4.5-sse2", "ref-4.5-sse4.1", "ref-4.5-avx2",
294                 "ref-main-neon", "ref-main-sse2", "ref-main-sse4.1", "ref-main-avx2"]
295
296    # All test encoders
297    testcoders = ["none", "neon", "sse2", "sse4.1", "avx2", "native", "universal"]
298    testcodersAArch64 = ["neon"]
299    testcodersX86 = ["sse2", "sse4.1", "avx2"]
300
301    coders = refcoders + testcoders + ["all-aarch64", "all-x86"]
302
303    parser.add_argument("--encoder", dest="encoders", default="avx2",
304                        choices=coders, help="test encoder variant")
305
306    parser.add_argument("--reference", dest="reference", default="ref-main-avx2",
307                        choices=refcoders, help="reference encoder variant")
308
309    astcProfile = ["ldr", "ldrs", "hdr", "all"]
310    parser.add_argument("--color-profile", dest="profiles", default="all",
311                        choices=astcProfile, help="test color profile")
312
313    imgFormat = ["l", "xy", "rgb", "rgba", "all"]
314    parser.add_argument("--color-format", dest="formats", default="all",
315                        choices=imgFormat, help="test color format")
316
317    choices = list(TEST_BLOCK_SIZES) + ["all"]
318    parser.add_argument("--block-size", dest="blockSizes",
319                        action="append", choices=choices,
320                        help="test block size")
321
322    testDir = os.path.dirname(__file__)
323    testDir = os.path.join(testDir, "Images")
324    testSets = []
325    for path in os.listdir(testDir):
326        fqPath = os.path.join(testDir, path)
327        if os.path.isdir(fqPath):
328            testSets.append(path)
329    testSets.append("all")
330
331    parser.add_argument("--test-set", dest="testSets", default="Small",
332                        choices=testSets, help="test image test set")
333
334    parser.add_argument("--test-image", dest="testImage", default=None,
335                        help="select a specific test image from the test set")
336
337    choices = list(TEST_QUALITIES) + ["all", "all+"]
338    parser.add_argument("--test-quality", dest="testQual", default="thorough",
339                        choices=choices, help="select a specific test quality")
340
341    parser.add_argument("--repeats", dest="testRepeats", default=1,
342                        type=int, help="test iteration count")
343
344    parser.add_argument("--keep-output", dest="keepOutput", default=False,
345                        action="store_true", help="keep image output")
346
347    parser.add_argument("-j", dest="threads", default=None,
348                        type=int, help="thread count")
349
350
351    args = parser.parse_args()
352
353    # Turn things into canonical format lists
354    if args.encoders == "all-aarch64":
355        args.encoders = testcodersAArch64
356    elif args.encoders == "all-x86":
357        args.encoders = testcodersX86
358    else:
359        args.encoders = [args.encoders]
360
361    if args.testQual == "all+":
362        args.testQual = TEST_QUALITIES
363    elif args.testQual == "all":
364        args.testQual = TEST_QUALITIES
365        args.testQual.remove("verythorough")
366        args.testQual.remove("exhaustive")
367    else:
368        args.testQual = [args.testQual]
369
370    if not args.blockSizes or ("all" in args.blockSizes):
371        args.blockSizes = TEST_BLOCK_SIZES
372
373    args.testSets = testSets[:-1] if args.testSets == "all" \
374        else [args.testSets]
375
376    args.profiles = astcProfile[:-1] if args.profiles == "all" \
377        else [args.profiles]
378
379    args.formats = imgFormat[:-1] if args.formats == "all" \
380        else [args.formats]
381
382    return args
383
384
385def main():
386    """
387    The main function.
388
389    Returns:
390        int: The process return code.
391    """
392    # Parse command lines
393    args = parse_command_line()
394
395    testSetCount = 0
396    worstResult = trs.Result.NOTRUN
397
398    for quality in args.testQual:
399        for imageSet in args.testSets:
400            for encoderName in args.encoders:
401                (encoder, name, outDir, refName) = \
402                    get_encoder_params(encoderName, args.reference, imageSet)
403
404                testDir = "Test/Images/%s" % imageSet
405                testRes = "%s/astc_%s_%s_results.csv" % (outDir, name, quality)
406
407                testRef = None
408                if refName:
409                    dat = (testDir, refName, quality)
410                    testRefPath = "%s/astc_%s_%s_results.csv" % dat
411                    testRef = trs.ResultSet(imageSet)
412                    testRef.load_from_file(testRefPath)
413
414                testSetCount += 1
415                testSet = tts.TestSet(imageSet, testDir,
416                                      args.profiles, args.formats, args.testImage)
417
418                resultSet = run_test_set(encoder, testRef, testSet, quality,
419                                         args.blockSizes, args.testRepeats,
420                                         args.keepOutput, args.threads)
421
422                resultSet.save_to_file(testRes)
423
424                if refName:
425                    summary = resultSet.get_results_summary()
426                    worstResult = max(summary.get_worst_result(), worstResult)
427                    print(summary)
428
429        if (testSetCount > 1) and (worstResult != trs.Result.NOTRUN):
430            print("OVERALL STATUS: %s" % worstResult.name)
431
432    if worstResult == trs.Result.FAIL:
433        return 1
434
435    return 0
436
437
438if __name__ == "__main__":
439    sys.exit(main())
440
441
442if __name__ == "__main__":
443    sys.exit(main())
444