test_preproc.py

"""
test_preproc.py — Test image_preproc modes on PYNQ (no camera needed)

Generates a synthetic YUYV test frame, encodes it in each mode,
saves 4 JPEGs so you can visually compare the effects.

Run on PYNQ:
    python3 test_preproc.py
"""

import numpy as np
from mjpeg_driver import MjpegEncoder

BITSTREAM = "mjpeg_enc.bit"
WIDTH, HEIGHT = 320, 240


def make_test_yuyv(width, height):
    """Generate a YUYV test pattern with soft edges, noise, and color regions.

    Contains low-frequency content (sine blobs, smooth gradients) so that
    sharpening mode produces a visible difference.
    """
    n_pairs = width // 2
    frame = np.zeros((height, n_pairs, 4), dtype=np.uint8)

    # Precompute a soft luma field: sine blobs + gradient + mild noise
    rng = np.random.RandomState(42)
    noise = rng.randint(-8, 9, size=(height, width), dtype=np.int16)

    for row in range(height):
        for col in range(n_pairs):
            x0 = col * 2
            x1 = x0 + 1

            # Smooth luma: overlapping sine blobs (low freq content)
            def luma(x, y):
                # vertical gradient
                base = 40 + int(180 * y / height)
                # horizontal sine (period ~160px)
                blob1 = int(30 * np.sin(2 * np.pi * x / 160)
                            * np.sin(2 * np.pi * y / 120))
                # diagonal sine (period ~200px)
                blob2 = int(20 * np.sin(2 * np.pi * (x + y) / 200))
                val = base + blob1 + blob2 + int(noise[y, x])
                return max(16, min(235, val))

            y0 = luma(x0, row)
            y1 = luma(x1, row)

            # Cb/Cr: 4 color zones with smooth transitions at boundaries
            zone_f = x0 * 4.0 / width  # 0.0 .. 4.0
            zone = int(zone_f)
            colors = [
                (128, 128),  # grey
                (90, 240),   # red-ish
                (240, 90),   # blue-ish
                (60, 170),   # green-ish
            ]
            u, v = colors[min(zone, 3)]

            frame[row, col] = [y0, u, y1, v]
    return frame.tobytes()


def main():
    print("Generating test YUYV frame...")
    yuyv = make_test_yuyv(WIDTH, HEIGHT)

    print(f"Loading overlay: {BITSTREAM}")
    enc = MjpegEncoder(BITSTREAM, WIDTH, HEIGHT, quality=75)

    modes = ["normal", "sharpen", "blue", "binarize"]

    for i, name in enumerate(modes):
        enc.set_mode(i)
        jpeg = enc.encode_yuyv(yuyv)
        fname = f"test_{name}.jpg"
        with open(fname, "wb") as f:
            f.write(jpeg)
        print(f"  -> {fname}: {len(jpeg)} bytes")

    # Test button GPIO read
    try:
        btns = enc.ol.btns_gpio
        val = btns.read(0x0)
        print(f"\nButton GPIO read: 0x{val:02x} (BTN3..BTN0 = {val:04b})")
    except Exception as e:
        print(f"\nButton GPIO read failed: {e}")

    enc.set_mode(0)  # reset to normal
    enc.close()
    print("\nDone! Check test_normal.jpg, test_sharpen.jpg, test_blue.jpg, test_red.jpg")


if __name__ == "__main__":
    main()