iqa 1.2.1

A single, ergonomic API over the patchwork of visual quality assessment metrics in the Rust ecosystem.
Documentation 
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

#!/usr/bin/env python3
#
# gen-msssim-goldens.py — a no-dependency cross-check of the MS-SSIM goldens.
#
# The AUTHORITATIVE goldens come from `scripts/gen-msssim-goldens.sh`, which runs
# the original reference `msssim.m` (Wang, Simoncelli & Bovik, "Multiscale
# Structural Similarity for Image Quality Assessment", 2003) under Octave. This
# script is an INDEPENDENT NumPy reimplementation of that same five-scale
# algorithm, written from the paper rather than from the Rust: a different
# language, a separate convolution, and the textbook E[x^2]-E[x]^2 variance form
# (the crate uses the deviation form). It exists so the goldens can be re-derived
# without Octave and so a third implementation has to agree — it reproduces the
# reference values exactly.
#
# It reads the committed grayscale PGM fixtures and prints the Rust `GOLDENS`
# table to paste into tests/ms_ssim_reference.rs. The fixtures are the single
# source of truth: this script and the Rust test read the identical .pgm files.
#
#   python3 scripts/gen-msssim-goldens.py
#
# Regenerate the fixtures first if they changed:
#   cargo test --features ms-ssim --test ms_ssim_reference \
#       write_ms_ssim_fixtures -- --ignored

import os
import sys

import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

WINDOW = 11
SIGMA = 1.5
K1, K2 = 0.01, 0.03
WEIGHTS = (0.0448, 0.2856, 0.3001, 0.2363, 0.1333)

FIXTURE_DIR = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "tests",
    "fixtures",
    "ms_ssim",
)

# (case label, reference fixture, distorted fixture) — keep in lockstep with the
# CASES table in tests/ms_ssim_reference.rs.
CASES = [
    ("gradient_lo", "gradient_ref.pgm", "gradient_lo.pgm"),
    ("gradient_hi", "gradient_ref.pgm", "gradient_hi.pgm"),
    ("solid_dist", "solid_ref.pgm", "solid_dist.pgm"),
    ("texture_dist", "texture_ref.pgm", "texture_dist.pgm"),
]


def read_pgm(path):
    """Reads a binary (P5) PGM into an (h, w) float64 array."""
    with open(path, "rb") as f:
        data = f.read()
    if data[:2] != b"P5":
        raise ValueError(f"{path}: not a binary PGM")
    idx = 2
    fields = []
    while len(fields) < 3:
        while idx < len(data) and data[idx] in b" \t\n\r":
            idx += 1
        if data[idx : idx + 1] == b"#":
            while data[idx] not in b"\n":
                idx += 1
            continue
        start = idx
        while data[idx] not in b" \t\n\r":
            idx += 1
        fields.append(int(data[start:idx]))
    w, h, _maxv = fields
    idx += 1  # exactly one whitespace byte follows maxval
    buf = np.frombuffer(data[idx : idx + w * h], dtype=np.uint8).astype(np.float64)
    return buf.reshape(h, w)


def gaussian_window():
    c = (WINDOW - 1) / 2.0
    ax = np.arange(WINDOW) - c
    xx, yy = np.meshgrid(ax, ax)
    w = np.exp(-(xx**2 + yy**2) / (2 * SIGMA**2))
    return w / w.sum()


def scale_means(a, b, win, c1, c2):
    """Mean full-SSIM and mean contrast-structure over all valid windows."""
    wa = sliding_window_view(a, (WINDOW, WINDOW))
    wb = sliding_window_view(b, (WINDOW, WINDOW))
    mu_a = np.einsum("ijkl,kl->ij", wa, win)
    mu_b = np.einsum("ijkl,kl->ij", wb, win)
    aa = np.einsum("ijkl,kl->ij", wa * wa, win)
    bb = np.einsum("ijkl,kl->ij", wb * wb, win)
    ab = np.einsum("ijkl,kl->ij", wa * wb, win)
    var_a = aa - mu_a * mu_a
    var_b = bb - mu_b * mu_b
    cov = ab - mu_a * mu_b
    luminance = (2 * mu_a * mu_b + c1) / (mu_a * mu_a + mu_b * mu_b + c1)
    cs = (2 * cov + c2) / (var_a + var_b + c2)
    return float((luminance * cs).mean()), float(cs.mean())


def downsample(img):
    """Average each aligned 2x2 block, dropping any odd trailing row/column."""
    h, w = img.shape
    h2, w2 = (h // 2) * 2, (w // 2) * 2
    img = img[:h2, :w2]
    return (
        img[0::2, 0::2] + img[1::2, 0::2] + img[0::2, 1::2] + img[1::2, 1::2]
    ) / 4.0


def msssim(a, b):
    c1 = (K1 * 255.0) ** 2
    c2 = (K2 * 255.0) ** 2
    win = gaussian_window()
    n = len(WEIGHTS)
    product = 1.0
    for i in range(n):
        full, cs = scale_means(a, b, win, c1, c2)
        term = full if i == n - 1 else cs
        product *= max(term, 1e-12) ** WEIGHTS[i]
        if i < n - 1:
            a, b = downsample(a), downsample(b)
    return product


def main():
    print('// Independent NumPy MS-SSIM reference, scripts/gen-msssim-goldens.py.')
    print("const GOLDENS: &[(&str, f64)] = &[")
    for label, ref, dist in CASES:
        a = read_pgm(os.path.join(FIXTURE_DIR, ref))
        b = read_pgm(os.path.join(FIXTURE_DIR, dist))
        print(f'    ("{label}", {msssim(a, b):.6f}),')
    print("];")


if __name__ == "__main__":
    sys.exit(main())