bandsaunter/bandsaunter/images.py

"""Pictures off the air, and the file they get written to.

Three of the things a receiver can hear are images rather than sound or
bits: weather satellites in the 137 MHz band sending APT, amateurs sending
SSTV on 14.230 MHz and 144.5 MHz, and the marine weather fax stations that
have been sending charts on shortwave since before any of this was digital.
All three are analogue -- brightness is a frequency, and a picture is a very
long single tone that wobbles -- so all three come out of the same place:
the audio the scanner already records.

This module is what they have in common.  The decoders themselves are in
:mod:`sstv`, :mod:`apt` and :mod:`fax`.

PNG is written here rather than pulled in from Pillow.  A scanner that
cannot start because an imaging library is missing is worse than one that
cannot draw, the format's own container is a dozen lines of zlib and
struct, and this way the pictures work on a machine with nothing installed
but numpy.
"""

from __future__ import annotations

import struct
import zlib
from dataclasses import dataclass, field
from pathlib import Path

import numpy as np
from scipy import signal as sps

__all__ = ["ImageDecode", "write_png", "instantaneous_frequency",
           "tone_amplitude", "resample_to", "PNG_SIGNATURE"]

PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"


# ---------------------------------------------------------------------------
# What a decoder gives back
# ---------------------------------------------------------------------------

@dataclass
class ImageDecode:
    """One picture recovered from a capture."""

    ok: bool = False
    kind: str = ""                # APT / SSTV / HF fax
    mode: str = ""                # NOAA APT, Martin M1, 120 lpm IOC 576, ...
    width: int = 0
    height: int = 0
    pixels: np.ndarray | None = field(default=None, repr=False)
    path: str = ""                # where it was written, once it has been
    complete: bool = True         # False when the capture ran out partway
    seconds: float = 0.0          # of signal that went into it
    confidence: float = 0.0
    channels: dict[str, str] = field(default_factory=dict)   # extra images
    note: str = ""

    @property
    def name(self) -> str:
        return f"{self.kind} ({self.mode})" if self.mode else self.kind

    def summary(self) -> str:
        if not self.ok:
            return self.note or "no image decoded"
        size = f"{self.width}x{self.height}"
        bits = [self.name, size]
        if not self.complete:
            bits.append("partial")
        if self.path:
            bits.append(Path(self.path).name)
        return "  ".join(bits)

    def save(self, path) -> str:
        """Write the picture, and remember where it went."""
        if self.pixels is None:
            return ""
        written = write_png(path, self.pixels)
        self.path = str(written)
        return self.path


# ---------------------------------------------------------------------------
# PNG
# ---------------------------------------------------------------------------

def _chunk(tag: bytes, data: bytes) -> bytes:
    return (struct.pack(">I", len(data)) + tag + data
            + struct.pack(">I", zlib.crc32(tag + data) & 0xFFFFFFFF))


def write_png(path, pixels: np.ndarray) -> Path:
    """Write a greyscale or RGB array as a PNG.

    ``pixels`` is ``(height, width)`` for grey or ``(height, width, 3)`` for
    colour, in any numeric type; it is clipped into a byte rather than scaled,
    because a decoder that has already decided what is black and what is white
    should not have that decision quietly changed here.
    """
    data = np.asarray(pixels)
    if data.ndim == 2:
        colour, depth = 0, 1
    elif data.ndim == 3 and data.shape[2] == 3:
        colour, depth = 2, 3
    else:
        raise ValueError(f"cannot write an array of shape {data.shape} as PNG")
    if data.dtype != np.uint8:
        data = np.clip(data, 0, 255).astype(np.uint8)
    height, width = data.shape[0], data.shape[1]

    raw = bytearray()
    flat = data.reshape(height, width * depth)
    for row in flat:
        raw.append(0)                      # filter: none
        raw.extend(row.tobytes())

    out = bytearray(PNG_SIGNATURE)
    out += _chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, colour,
                                       0, 0, 0))
    out += _chunk(b"IDAT", zlib.compress(bytes(raw), 6))
    out += _chunk(b"IEND", b"")

    path = Path(path)
    tmp = path.with_suffix(path.suffix + ".tmp")
    tmp.write_bytes(bytes(out))
    tmp.replace(path)                      # never a half-written picture
    return path


# ---------------------------------------------------------------------------
# The measurements every one of these decoders needs
# ---------------------------------------------------------------------------

def instantaneous_frequency(audio: np.ndarray, rate: float,
                            low: float = 800.0,
                            high: float = 2800.0) -> np.ndarray:
    """The frequency of an audio tone, sample by sample, in hertz.

    All three picture modes encode brightness as a frequency, so this is the
    whole of the demodulation: band-pass to the range the tone lives in, take
    the analytic signal, and differentiate its phase.
    """
    audio = np.asarray(audio, dtype=np.float64).ravel()
    if audio.size < 64:
        return np.zeros(0)
    nyquist = rate / 2.0
    lo = max(1e-4, low / nyquist)
    hi = min(0.999, high / nyquist)
    if hi <= lo:
        return np.zeros(0)
    taps = sps.firwin(129, [lo, hi], pass_zero=False)
    filtered = sps.lfilter(taps, [1.0], audio)
    # Undo the filter's delay.  A linear-phase FIR of 129 taps holds its
    # output back by 64 samples, which is four milliseconds at 16 kHz -- and
    # four milliseconds is most of an SSTV sync pulse and a fifteen-pixel
    # shift across a Martin M2 line.  It cancels only if every time in the
    # decode is measured from this same trace, and the line timings are not:
    # they come from the mode specification, in real seconds.
    delay = 64
    filtered = np.concatenate((filtered[delay:],
                               np.zeros(delay, dtype=filtered.dtype)))
    analytic = sps.hilbert(filtered)
    phase = np.unwrap(np.angle(analytic))
    freq = np.diff(phase) * rate / (2.0 * np.pi)
    # The first samples are the filter filling up and mean nothing.
    freq[:65] = freq[65] if freq.size > 65 else 0.0
    return freq


def tone_amplitude(audio: np.ndarray, rate: float, centre: float,
                   width: float = 1200.0) -> np.ndarray:
    """The envelope of an amplitude-modulated subcarrier.

    APT is a 2400 Hz tone whose loudness is the brightness of the picture, so
    for that one the frequency is fixed and it is the envelope that carries
    everything.
    """
    audio = np.asarray(audio, dtype=np.float64).ravel()
    if audio.size < 64:
        return np.zeros(0)
    nyquist = rate / 2.0
    lo = max(1e-4, (centre - width) / nyquist)
    hi = min(0.999, (centre + width) / nyquist)
    if hi <= lo:
        return np.abs(sps.hilbert(audio))
    taps = sps.firwin(129, [lo, hi], pass_zero=False)
    return np.abs(sps.hilbert(sps.lfilter(taps, [1.0], audio)))


def resample_to(values: np.ndarray, count: int) -> np.ndarray:
    """Stretch or squash a run of samples to exactly ``count`` of them.

    Linear rather than a proper resampler: a picture line is a few hundred
    samples wide and the difference is invisible, while a polyphase filter
    per line over a fifteen-minute pass is not.
    """
    values = np.asarray(values, dtype=np.float64)
    if values.size == 0 or count <= 0:
        return np.zeros(max(0, count))
    if values.size == count:
        return values
    source = np.linspace(0.0, 1.0, values.size)
    target = np.linspace(0.0, 1.0, count)
    return np.interp(target, source, values)


def stretch(values: np.ndarray, low: float, high: float) -> np.ndarray:
    """Map a frequency range onto 0-255, clipped at both ends."""
    if high <= low:
        return np.zeros_like(values, dtype=np.uint8)
    scaled = (np.asarray(values, dtype=np.float64) - low) / (high - low)
    return np.clip(scaled * 255.0, 0, 255).astype(np.uint8)