tool.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Sequence, Tuple

import numpy as np
import soundfile as sf
from PIL import Image, ImageOps
from scipy.signal import hilbert

VERSION = "1.0.0"
AUTHOR = "@uckix"

SYNC_HZ = 1200.0
VIS_START_HZ = 1900.0
VIS_BIT1_HZ = 1100.0
VIS_BIT0_HZ = 1300.0
BLACK_HZ = 1500.0
WHITE_HZ = 2300.0

IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".gif", ".webp", ".tiff"}
AUDIO_EXTS = {".wav", ".mp3", ".ogg", ".flac", ".m4a", ".aac", ".opus"}


@dataclass(frozen=True)
class Mode:
    name: str
    vis: int
    width: int
    height: int
    sync_ms: float
    scan_ms: float
    gap_ms: float
    order: Tuple[str, str, str]
    family: str

    @property
    def line_ms(self) -> float:
        return self.sync_ms + 3 * self.scan_ms + 3 * self.gap_ms


MODES: Dict[str, Mode] = {
    "MartinM1": Mode("MartinM1", 0x2C, 320, 256, 4.862, 146.432, 0.572, ("G", "B", "R"), "martin"),
    "MartinM2": Mode("MartinM2", 0x28, 160, 256, 4.862, 73.216, 0.572, ("G", "B", "R"), "martin"),
    "ScottieS1": Mode("ScottieS1", 0x3C, 320, 256, 9.0, 138.24, 1.5, ("R", "G", "B"), "scottie"),
    "ScottieS2": Mode("ScottieS2", 0x38, 160, 256, 9.0, 88.064, 1.5, ("R", "G", "B"), "scottie"),
    "ScottieDX": Mode("ScottieDX", 0x4C, 320, 256, 9.0, 345.6, 1.5, ("R", "G", "B"), "scottie"),
}
VIS_TO_MODE = {mode.vis: mode for mode in MODES.values()}
MODE_ALIASES = {name.lower(): name for name in MODES}

HELP_EPILOG = """
examples:
  decode with auto-detect:
    python tool.py -i input.wav -o out.png
    python tool.py -i input.mp3 -o out.png

  force decode mode if needed:
    python tool.py -i input.wav -o out.png -m ScottieS1

  encode with a chosen mode:
    python tool.py -i in.png -o out.wav -m ScottieS1
    python tool.py -i in.jpg -o out.ogg -m MartinM1
    python tool.py -i in.jpg -o out.mp3 -m ScottieDX

  list available modes:
    python tool.py --list-modes
""".strip()


def byte_to_freq(value: np.ndarray | float | int) -> np.ndarray | float:
    return BLACK_HZ + (np.asarray(value, dtype=np.float64) * (WHITE_HZ - BLACK_HZ) / 255.0)


def freq_to_byte(freq: np.ndarray | float | int) -> np.ndarray | float:
    arr = np.asarray(freq, dtype=np.float64)
    out = np.round((arr - BLACK_HZ) * 255.0 / (WHITE_HZ - BLACK_HZ))
    return np.clip(out, 0, 255)


def normalize_mode_name(name: str | None) -> str | None:
    if name is None:
        return None
    key = name.strip().lower()
    if key not in MODE_ALIASES:
        supported = ", ".join(MODES)
        raise RuntimeError(f"unsupported mode '{name}'. supported: {supported}")
    return MODE_ALIASES[key]


def check_ffmpeg() -> str:
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise RuntimeError("ffmpeg not found in PATH. install ffmpeg to handle mp3/ogg/audio conversion.")
    return ffmpeg


def run(cmd: Sequence[str]) -> None:
    proc = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    if proc.returncode != 0:
        raise RuntimeError(f"command failed: {' '.join(cmd)}\n{proc.stderr.strip()}")


def load_audio_mono(path: Path, sample_rate: int = 48000) -> Tuple[np.ndarray, int]:
    if path.suffix.lower() == ".wav":
        data, sr = sf.read(str(path), always_2d=False)
        if data.ndim == 2:
            data = data.mean(axis=1)
        return data.astype(np.float64), sr

    ffmpeg = check_ffmpeg()
    with tempfile.TemporaryDirectory() as temp_dir:
        wav_path = Path(temp_dir) / "decoded.wav"
        run([ffmpeg, "-y", "-i", str(path), "-ac", "1", "-ar", str(sample_rate), "-vn", str(wav_path)])
        data, sr = sf.read(str(wav_path), always_2d=False)
        if data.ndim == 2:
            data = data.mean(axis=1)
        return data.astype(np.float64), sr


def save_audio(path: Path, audio: np.ndarray, sample_rate: int = 48000) -> None:
    audio = np.asarray(audio, dtype=np.float32)
    if path.suffix.lower() == ".wav":
        sf.write(str(path), audio, sample_rate, subtype="PCM_16")
        return

    ffmpeg = check_ffmpeg()
    with tempfile.TemporaryDirectory() as temp_dir:
        wav_path = Path(temp_dir) / "temp.wav"
        sf.write(str(wav_path), audio, sample_rate, subtype="PCM_16")
        cmd = [ffmpeg, "-y", "-i", str(wav_path)]
        if path.suffix.lower() == ".mp3":
            cmd += ["-b:a", "320k"]
        elif path.suffix.lower() == ".ogg":
            cmd += ["-c:a", "libvorbis", "-q:a", "8"]
        cmd += [str(path)]
        run(cmd)


def preprocess_audio(audio: np.ndarray) -> np.ndarray:
    audio = np.asarray(audio, dtype=np.float64)
    if audio.size == 0:
        raise ValueError("empty audio")
    audio = audio - np.mean(audio)
    peak = np.max(np.abs(audio))
    if peak > 0:
        audio = audio / peak
    return audio


def instantaneous_frequency(audio: np.ndarray, sample_rate: int) -> np.ndarray:
    analytic = hilbert(audio)
    phase = np.unwrap(np.angle(analytic))
    inst = np.diff(phase, prepend=phase[0]) * sample_rate / (2.0 * np.pi)
    inst = np.clip(inst, 200.0, 4000.0)
    window = max(3, int(round(sample_rate * 0.0008)))
    if window % 2 == 0:
        window += 1
    kernel = np.ones(window, dtype=np.float64) / window
    return np.convolve(inst, kernel, mode="same")


def build_prefix(arr: np.ndarray) -> np.ndarray:
    return np.concatenate([[0.0], np.cumsum(arr, dtype=np.float64)])


def mean_range(prefix: np.ndarray, start: int, end: int) -> float:
    start = max(0, min(start, len(prefix) - 1))
    end = max(start + 1, min(end, len(prefix) - 1))
    return (prefix[end] - prefix[start]) / (end - start)


def segment_mean(prefix: np.ndarray, sample_rate: int, t0: float, ms: float) -> float:
    start = int(round(t0 * sample_rate))
    end = int(round((t0 + ms / 1000.0) * sample_rate))
    return mean_range(prefix, start, end)


def find_header(inst: np.ndarray, sample_rate: int, search_seconds: float = 8.0) -> float:
    prefix = build_prefix(inst)
    limit = min(len(inst) / sample_rate, search_seconds)
    best_t = None
    best_score = float("inf")
    step = 0.002
    t = 0.0
    while t + 0.62 < limit:
        f1 = segment_mean(prefix, sample_rate, t + 0.12, 40.0)
        f2 = segment_mean(prefix, sample_rate, t + 0.300, 10.0)
        f3 = segment_mean(prefix, sample_rate, t + 0.43, 40.0)
        score = abs(f1 - VIS_START_HZ) + abs(f2 - SYNC_HZ) + abs(f3 - VIS_START_HZ)
        if score < best_score:
            best_score = score
            best_t = t
        t += step
    if best_t is None or best_score > 220.0:
        raise RuntimeError("could not find SSTV VIS header")
    return best_t


def decode_vis(inst: np.ndarray, sample_rate: int, header_t: float) -> Tuple[int, int, int]:
    prefix = build_prefix(inst)
    vis_start = header_t + 0.610
    start_bit = segment_mean(prefix, sample_rate, vis_start, 30.0)
    bits: List[int] = []
    for i in range(7):
        freq = segment_mean(prefix, sample_rate, vis_start + (i + 1) * 0.030, 30.0)
        bit = 1 if abs(freq - VIS_BIT1_HZ) < abs(freq - VIS_BIT0_HZ) else 0
        bits.append(bit)
    parity_freq = segment_mean(prefix, sample_rate, vis_start + 8 * 0.030, 30.0)
    parity_bit = 1 if abs(parity_freq - VIS_BIT1_HZ) < abs(parity_freq - VIS_BIT0_HZ) else 0
    stop_bit = segment_mean(prefix, sample_rate, vis_start + 9 * 0.030, 30.0)

    value = 0
    for i, bit in enumerate(bits):
        value |= (bit & 1) << i
    return value, int(round(start_bit)), int(round(stop_bit))


def find_sync_near(
    err1200_prefix: np.ndarray,
    sample_rate: int,
    expected_t: float,
    sync_ms: float,
    window_ms: float = 18.0,
) -> float:
    best_t = expected_t
    best_score = float("inf")
    offset = -window_ms
    while offset <= window_ms:
        t = expected_t + offset / 1000.0
        start = int(round(t * sample_rate))
        end = int(round((t + sync_ms / 1000.0) * sample_rate))
        if start < 0 or end >= len(err1200_prefix):
            offset += 0.5
            continue
        score = mean_range(err1200_prefix, start, end) + abs(offset) * 0.2
        if score < best_score:
            best_score = score
            best_t = t
        offset += 0.5
    return best_t


def sample_channel(inst_prefix: np.ndarray, sample_rate: int, start_t: float, scan_ms: float, width: int) -> np.ndarray:
    seg_edges = np.linspace(start_t, start_t + scan_ms / 1000.0, width + 1)
    starts = np.round(seg_edges[:-1] * sample_rate).astype(int)
    ends = np.round(seg_edges[1:] * sample_rate).astype(int)
    starts = np.clip(starts, 0, len(inst_prefix) - 2)
    ends = np.clip(ends, starts + 1, len(inst_prefix) - 1)
    means = (inst_prefix[ends] - inst_prefix[starts]) / (ends - starts)
    return freq_to_byte(means).astype(np.uint8)


def decode_sstv(input_path: Path, output_path: Path, forced_mode: str | None = None) -> Mode:
    audio, sample_rate = load_audio_mono(input_path)
    audio = preprocess_audio(audio)
    inst = instantaneous_frequency(audio, sample_rate)

    if forced_mode is None:
        header_t = find_header(inst, sample_rate)
        vis, _, _ = decode_vis(inst, sample_rate, header_t)
        if vis not in VIS_TO_MODE:
            supported = ", ".join(MODES)
            raise RuntimeError(f"unsupported or unknown VIS code: {vis}. try --mode with one of: {supported}")
        mode = VIS_TO_MODE[vis]
        image_start = header_t + 0.910
    else:
        mode = MODES[normalize_mode_name(forced_mode)]
        try:
            header_t = find_header(inst, sample_rate)
            image_start = header_t + 0.910
        except Exception:
            image_start = 0.0

    inst_prefix = build_prefix(inst)
    err1200_prefix = build_prefix(np.abs(inst - SYNC_HZ))
    channels: Dict[str, np.ndarray] = {
        "R": np.zeros((mode.height, mode.width), dtype=np.uint8),
        "G": np.zeros((mode.height, mode.width), dtype=np.uint8),
        "B": np.zeros((mode.height, mode.width), dtype=np.uint8),
    }

    def build_sync_sequence(first_expected: float, first_window_ms: float) -> Tuple[List[float], float]:
        sync_times: List[float] = []
        current_sync = find_sync_near(err1200_prefix, sample_rate, first_expected, mode.sync_ms, first_window_ms)
        sync_times.append(current_sync)
        quick_score = mean_range(
            err1200_prefix,
            int(round(current_sync * sample_rate)),
            int(round((current_sync + mode.sync_ms / 1000.0) * sample_rate)),
        )
        for _ in range(1, min(mode.height, 8)):
            expected_next = current_sync + mode.line_ms / 1000.0
            next_sync = find_sync_near(err1200_prefix, sample_rate, expected_next, mode.sync_ms, 20.0)
            sync_times.append(next_sync)
            quick_score += abs((next_sync - current_sync) - mode.line_ms / 1000.0) * 1000.0
            current_sync = next_sync
        while len(sync_times) < mode.height:
            current_sync = find_sync_near(
                err1200_prefix,
                sample_rate,
                current_sync + mode.line_ms / 1000.0,
                mode.sync_ms,
                20.0,
            )
            sync_times.append(current_sync)
        return sync_times, quick_score

    if mode.family == "scottie":
        candidate_specs = [
            (image_start, 80.0),
            (image_start + 2.0 * (mode.scan_ms + mode.gap_ms) / 1000.0, 80.0),
        ]
    else:
        candidate_specs = [(image_start, 40.0)]

    best_syncs: List[float] | None = None
    best_score = float("inf")
    for first_expected, first_window_ms in candidate_specs:
        sync_times, score = build_sync_sequence(first_expected, first_window_ms)
        if score < best_score:
            best_score = score
            best_syncs = sync_times

    for line, sync_t in enumerate(best_syncs or []):
        t = sync_t + mode.sync_ms / 1000.0 + mode.gap_ms / 1000.0
        for idx, channel_name in enumerate(mode.order):
            channels[channel_name][line] = sample_channel(inst_prefix, sample_rate, t, mode.scan_ms, mode.width)
            t += mode.scan_ms / 1000.0
            if idx < len(mode.order) - 1:
                t += mode.gap_ms / 1000.0

    rgb = np.dstack([channels["R"], channels["G"], channels["B"]])
    image = Image.fromarray(rgb, mode="RGB")
    output_path.parent.mkdir(parents=True, exist_ok=True)
    if output_path.suffix.lower() in {".jpg", ".jpeg"}:
        image.save(output_path, quality=95, subsampling=0)
    else:
        image.save(output_path)
    return mode


def resize_image_for_mode(image: Image.Image, mode: Mode) -> Image.Image:
    return ImageOps.fit(image.convert("RGB"), (mode.width, mode.height), method=Image.Resampling.LANCZOS)


def vis_segments(vis: int) -> List[Tuple[float, float]]:
    segments: List[Tuple[float, float]] = [
        (VIS_START_HZ, 300.0),
        (SYNC_HZ, 10.0),
        (VIS_START_HZ, 300.0),
        (SYNC_HZ, 30.0),
    ]
    ones = 0
    value = vis
    for _ in range(7):
        bit = value & 1
        value >>= 1
        ones += bit
        segments.append((VIS_BIT1_HZ if bit else VIS_BIT0_HZ, 30.0))
    parity_bit = 1 if (ones % 2 == 1) else 0
    segments.append((VIS_BIT1_HZ if parity_bit else VIS_BIT0_HZ, 30.0))
    segments.append((SYNC_HZ, 30.0))
    return segments


def image_line_segments(line_rgb: np.ndarray, mode: Mode) -> List[Tuple[np.ndarray | float, float]]:
    segments: List[Tuple[np.ndarray | float, float]] = [(SYNC_HZ, mode.sync_ms), (BLACK_HZ, mode.gap_ms)]
    color_map = {"R": line_rgb[:, 0], "G": line_rgb[:, 1], "B": line_rgb[:, 2]}
    pixel_ms = mode.scan_ms / mode.width
    for idx, channel_name in enumerate(mode.order):
        segments.append((byte_to_freq(color_map[channel_name]), pixel_ms))
        if idx < len(mode.order) - 1:
            segments.append((BLACK_HZ, mode.gap_ms))
    return segments


def synthesize_segments(segments: Sequence[Tuple[np.ndarray | float, float]], sample_rate: int) -> np.ndarray:
    parts: List[np.ndarray] = []
    phase = 0.0
    for freq, ms in segments:
        if isinstance(freq, np.ndarray):
            for value in freq:
                count = max(1, int(round(sample_rate * (ms / 1000.0))))
                t = np.arange(count, dtype=np.float64)
                part = np.sin(phase + 2.0 * np.pi * float(value) * t / sample_rate)
                phase = (phase + 2.0 * np.pi * float(value) * count / sample_rate) % (2.0 * np.pi)
                parts.append(part.astype(np.float32))
            continue

        count = max(1, int(round(sample_rate * (ms / 1000.0))))
        t = np.arange(count, dtype=np.float64)
        part = np.sin(phase + 2.0 * np.pi * float(freq) * t / sample_rate)
        phase = (phase + 2.0 * np.pi * float(freq) * count / sample_rate) % (2.0 * np.pi)
        parts.append(part.astype(np.float32))

    audio = np.concatenate(parts) if parts else np.zeros(0, dtype=np.float32)
    if audio.size:
        audio *= 0.8 / max(1e-9, np.max(np.abs(audio)))
    return audio


def encode_sstv(input_path: Path, output_path: Path, mode_name: str, sample_rate: int = 48000) -> Mode:
    normalized_mode = normalize_mode_name(mode_name)
    mode = MODES[normalized_mode]
    image = resize_image_for_mode(Image.open(input_path), mode)
    rgb = np.asarray(image, dtype=np.uint8)

    segments: List[Tuple[np.ndarray | float, float]] = []
    segments.extend(vis_segments(mode.vis))
    for line in rgb:
        segments.extend(image_line_segments(line, mode))

    save_audio(output_path, synthesize_segments(segments, sample_rate), sample_rate)
    return mode


def infer_direction(input_path: Path, output_path: Path, explicit_encode: bool, explicit_decode: bool) -> str:
    if explicit_encode and explicit_decode:
        raise RuntimeError("use either --encode or --decode, not both")
    if explicit_encode:
        return "encode"
    if explicit_decode:
        return "decode"

    in_ext = input_path.suffix.lower()
    out_ext = output_path.suffix.lower()
    if in_ext in AUDIO_EXTS and out_ext in IMAGE_EXTS:
        return "decode"
    if in_ext in IMAGE_EXTS and out_ext in AUDIO_EXTS:
        return "encode"
    raise RuntimeError(
        "could not infer direction. use audio->image for decode, image->audio for encode, or pass --encode/--decode explicitly"
    )


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="single-file SSTV encoder/decoder with VIS auto-detect",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=HELP_EPILOG,
    )
    parser.add_argument("-i", "--input", help="input file path")
    parser.add_argument("-o", "--output", help="output file path")
    parser.add_argument("-m", "--mode", help="mode name for encode, or forced mode for decode")
    parser.add_argument("--encode", action="store_true", help="force encode mode")
    parser.add_argument("--decode", action="store_true", help="force decode mode")
    parser.add_argument("--rate", type=int, default=48000, help="audio sample rate for encoding or conversion")
    parser.add_argument("--list-modes", action="store_true", help="print supported modes and exit")
    parser.add_argument("--version", action="store_true", help="print version and exit")
    return parser


def print_modes() -> None:
    for mode in MODES.values():
        print(f"{mode.name:10} vis={mode.vis:<3} size={mode.width}x{mode.height} family={mode.family}")


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()

    if args.version:
        print(f"sstv-tool {VERSION} by {AUTHOR}")
        return 0

    if args.list_modes:
        print_modes()
        return 0

    if not args.input or not args.output:
        parser.error("-i/--input and -o/--output are required unless you use --list-modes or --version")

    input_path = Path(args.input)
    output_path = Path(args.output)

    try:
        direction = infer_direction(input_path, output_path, args.encode, args.decode)
        if direction == "decode":
            mode = decode_sstv(input_path, output_path, forced_mode=args.mode)
            print(f"decoded: mode={mode.name} vis={mode.vis} -> {output_path}")
        else:
            if not args.mode:
                raise RuntimeError("encoding requires --mode. use --list-modes to see supported modes")
            mode = encode_sstv(input_path, output_path, args.mode, sample_rate=args.rate)
            print(f"encoded: mode={mode.name} vis={mode.vis} -> {output_path}")
        return 0
    except Exception as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())