split/test_picky.py at main · Lumenman/split · 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
"""Self-check for picky.py: sine-trace image -> WAVs with sane sizes."""
import os
import wave
from pathlib import Path

import numpy as np
from PIL import Image

import picky


def make_image(path, w=200, h=64):
    img = np.zeros((h, w), np.uint8)
    r = (h / 2 + (h / 3) * np.sin(np.linspace(0, 8 * np.pi, w))).astype(int)
    img[r, np.arange(w)] = 255
    Image.fromarray(img).save(path)


def main():
    tmp = Path(__file__).parent / "selftest_tmp"
    tmp.mkdir(exist_ok=True)
    os.chdir(tmp)

    make_image("trace.tif")
    r = picky.process_single("trace.tif", "w", 0, 1.0, "d", 1, 1.0, 0.95,
                             "all", 44100, 24)
    assert len(r["W"]) == 200 and len(r["T"]) == 199
    assert r["W"].max() > r["W"].min(), "flat output"
    pdp = sorted(Path(".").glob("trace_*_pdp.wav"))
    assert pdp, "no pdp wav written"
    with wave.open(str(pdp[-1])) as f:
        assert f.getnframes() == 200
        assert f.getsampwidth() == 3 and f.getframerate() == 44100

    make_image("trace2.tif")
    picky.process_batch(["trace.tif", "trace2.tif"], "x", 50, 1.0, "d", 1, 1.0,
                        0.95, "all", 44100, 24)
    clk = sorted(Path(".").glob("Batch_*_clk.wav"))
    bpdp = sorted(Path(".").glob("Batch_*_pdp.wav"))
    assert clk and bpdp
    with wave.open(str(clk[-1])) as fa, wave.open(str(bpdp[-1])) as fb:
        assert fa.getnframes() == fb.getnframes() == 3000 + 200 + 200 + 3000

    # alpha is not a channel (Octave's imread returns it separately)
    g = np.asarray(Image.open("trace.tif"))
    rgb = np.repeat(g[:, :, None], 3, axis=2)
    Image.fromarray(rgb, "RGB").save("rgb.tif")
    Image.fromarray(np.dstack([rgb, np.full(g.shape, 128, np.uint8)]), "RGBA").save("rgba.tif")
    a = picky.process_single("rgb.tif", "-", 0, 1.0, "d", 1, 1.0, 0.95, "all", 44100, 24)
    b = picky.process_single("rgba.tif", "-", 0, 1.0, "d", 1, 1.0, 0.95, "all", 44100, 24)
    assert np.array_equal(a["W"], b["W"]), "alpha channel leaked into the sum"

    # outputs land next to the source image even when the cwd is elsewhere
    os.chdir(tmp.parent)
    picky.process_single(str(tmp / "trace.tif"), "w", 0, 1.0, "d", 1, 1.0, 0.95,
                         "pdp", 44100, 24)
    assert not list(Path(".").glob("trace_*.wav")), "output leaked into the cwd"
    assert list(tmp.glob("trace_*_pdp.wav"))

    print("OK")


if __name__ == "__main__":
    main()