noaa-goes-visualization/defect_audit.py

#!/usr/bin/env python
"""Build a hand-confirmed set of REAL archive defects, the ground truth detection lacks.

Every detector score before this came from damage the bench injected itself, and
the bench's notion of a clean frame came from the same detectors being scored.  The
real 2024 render showed what that hid: saturated bars, bright blocks, ghost and
displaced discs, and short-exposure static, all judged good.  This builds the
missing reference set in three steps:

    defect_audit.py mine    --start 2023-01-01 --end 2025-04-01 --out AUDIT/candidates.json
    defect_audit.py sheets  --candidates AUDIT/candidates.json [--corrections C.txt] --out-dir AUDIT/sheets
    defect_audit.py apply   --candidates AUDIT/candidates.json --corrections C.txt --out labels/real_defects.json

``mine`` gathers candidates from several independent sources, so that no single
detector's blind spot decides what gets looked at:

* one-frame IMG_SDEV spikes against both neighbours (how the bars were found);
* frames whose header says no long exposure went into them (N_LONG=0);
* frames the legacy pixel filter rejected;
* frames around the era's largest flares, the hardest *clean* frames there are;
* random frames, stratified by satellite, band and year.

Each candidate is measured -- far-field brightness and grain, the current verdict --
and given a *proposed* category.  ``sheets`` draws numbered contact sheets grouped
by category for a person to confirm; corrections are plain ``ID CATEGORY`` lines.
``apply`` refuses to write labels while anything is still only proposed as a
generic far-field defect, so every defect in the set was classified by eye.

Read-only against the archive and the index.
"""

import argparse
import collections
import datetime
import json
import os
import sys
from multiprocessing import Pool

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from bench import parse_date
from suvi import db, detectors, fitsio, paths

#: Categories a confirmed label may carry.
#: ``gain`` is a whole-frame brightness jump against both neighbours with no flare
#: to explain it -- invisible in any one frame, obvious in the sequence.
DEFECTS = ("bar", "block", "static", "ghost", "displaced", "gain", "dark", "other_defect")
CLEAN = ("normal", "flare_normal")
#: Neither clean nor visibly defective: the header is wrong (e.g. WAVELNTH names
#: another band) while the pixels look right.  Reported apart, scored in neither.
NEUTRAL = ("header_only",)
CATEGORIES = DEFECTS + CLEAN + NEUTRAL
#: Proposed for far-field defects until someone says which kind; never a label.
PENDING = "far_field"

#: A frame's IMG_SDEV must exceed both neighbours' by this factor to be a spike.
SPIKE_RATIO = 1.6
#: Neighbours further apart than this (4 slots) say nothing about a one-frame spike.
SPIKE_MAX_GAP = 4 * paths.CADENCE
#: Far-field brightness, as a multiple of the band's typical level, from which a
#: frame is proposed as a far-field defect.  Normal and X-class flare frames reached
#: at most 3.2x in the 2024 calibration sample; defects started at 4x.
FAR_PROPOSE = 4.0
#: Far-field grain, as a multiple of the band's typical level, proposed as static.
NOISE_PROPOSE = 10.0
#: Minutes either side of a flare peak whose frames are flare controls.
FLARE_WINDOW = 40 * 60


# ---------------------------------------------------------------- candidate sources


def header_spikes(series, ratio=SPIKE_RATIO, max_gap=SPIKE_MAX_GAP):
    """Times whose IMG_SDEV jumps above both neighbours by `ratio` and comes back.

    `series` is time-sorted (t_start, img_sdev, n_long) for one satellite and band.
    N_LONG=0 frames are left to their own source; missing or non-positive spreads
    and neighbours more than `max_gap` away cannot establish a spike.
    """
    hits = []
    for i in range(1, len(series) - 1):
        (t0, s0, _), (t, s, n_long), (t1, s1, _) = series[i - 1], series[i], series[i + 1]
        if n_long == 0 or t - t0 > max_gap or t1 - t > max_gap:
            continue
        if not all(v is not None and v > 0 for v in (s0, s, s1)):
            continue
        if min(s / s0, s / s1) > ratio:
            hits.append(t)
    return hits


def flare_peaks(times, means, count, separation=2 * 86400, smooth=5, min_rise=2.0):
    """Up to `count` of the largest sustained brightenings, at least `separation` apart.

    A running median over `smooth` frames keeps a one-frame defect -- which lifts
    IMG_MEAN as a flare does -- from being mistaken for a flare peak.  A peak must
    reach `min_rise` times the series median, so a quiet stretch yields fewer
    peaks rather than background frames posing as flares.
    """
    if len(times) < smooth:
        return []
    values = np.asarray(means, dtype=np.float64)
    half = smooth // 2
    padded = np.pad(values, half, mode="edge")
    windows = np.lib.stride_tricks.sliding_window_view(padded, smooth)
    sustained = np.median(windows, axis=1)
    floor = min_rise * float(np.median(sustained))
    peaks = []
    for index in np.argsort(-sustained):
        if sustained[index] < floor:
            break
        when = times[index]
        if all(abs(when - other) >= separation for other in peaks):
            peaks.append(int(when))
            if len(peaks) == count:
                break
    return sorted(peaks)


def stratified_sample(rows, key, per_stratum, rng):
    """At most `per_stratum` rows from each stratum, chosen without replacement."""
    groups = collections.defaultdict(list)
    for row in rows:
        groups[key(row)].append(row)
    chosen = []
    for stratum in sorted(groups):
        members = groups[stratum]
        take = min(per_stratum, len(members))
        chosen.extend(members[i] for i in sorted(rng.choice(len(members), take, replace=False)))
    return chosen


# ------------------------------------------------------------------- measurement


def measure(job):
    """Far-field brightness and grain plus the current verdict for one frame."""
    satellite, wavelength, when, path = job
    fitsio.quiet_astropy()
    values, _ = fitsio.scan_header(path)
    image, error = fitsio.read_image(path)
    verdict = detectors.judge_file((satellite, wavelength, when), path)
    radius = (values or {}).get("diam_sun")
    radius = radius / 2.0 if radius else None
    return {
        "far_field": detectors.far_field(image, radius) if image is not None else None,
        "noise": detectors.far_field_noise(image, radius) if image is not None else None,
        "n_long": (values or {}).get("n_long"),
        "verdict": verdict.verdict,
        "reason": verdict.reason if image is not None else (verdict.reason or error),
    }


def band_baselines(candidates):
    """Typical far-field brightness and grain per band, from the random controls."""
    out = {}
    for wavelength in paths.WAVELENGTHS:
        pool = [c for c in candidates if "random" in c["sources"]
                and c["wavelength"] == wavelength and c["far_field"] is not None
                and c.get("n_long") != 0]
        if pool:
            out[wavelength] = {
                "far_field": float(np.median([c["far_field"] for c in pool])),
                "noise": float(np.median([c["noise"] for c in pool])),
            }
    return out


def propose(candidate, baselines):
    """A first-guess category for a person to confirm or correct."""
    base = baselines.get(candidate["wavelength"])
    far, noise = candidate.get("far_field"), candidate.get("noise")
    if candidate.get("n_long") == 0 or (
            base and noise is not None and base["noise"] > 0
            and noise > NOISE_PROPOSE * base["noise"]):
        return "static"
    if base and far is not None and base["far_field"] > 0 and far > FAR_PROPOSE * base["far_field"]:
        return PENDING
    if candidate.get("verdict") == detectors.BAD:
        return "dark" if "ECLIPSE" in (candidate.get("reason") or "") else "other_defect"
    return "flare_normal" if "flare" in candidate["sources"] else "normal"


# --------------------------------------------------------------------------- mine


def _rows(conn, sql, params):
    return [tuple(row) for row in conn.execute(sql, params)]


def gather(conn, start, end, rng, satellites=(16, 18), random_per_stratum=45,
           nlong_per_stratum=15, legacy_per_band=50, flare_count=15, flare_frames=600):
    """Candidate frames by source: {(satellite, wavelength, t_start): (path, sources)}."""
    found = {}

    def add(rows, source):
        for satellite, wavelength, when, path in rows:
            key = (satellite, wavelength, when)
            found.setdefault(key, [path, []])[1].append(source)

    frame = ("SELECT f.satellite, f.wavelength, f.t_start, f.path FROM frame f "
             "JOIN header h ON h.frame_id = f.id WHERE f.t_start >= ? AND f.t_start < ? "
             "AND f.size_bytes >= 100000 AND f.satellite IN (%s)" % ",".join("?" * len(satellites)))
    params = (start, end, *satellites)

    for satellite in satellites:
        for wavelength in paths.WAVELENGTHS:
            rows = _rows(conn, "SELECT f.t_start, h.img_sdev, h.n_long, f.path FROM frame f "
                         "JOIN header h ON h.frame_id = f.id WHERE f.satellite = ? AND "
                         "f.wavelength = ? AND f.t_start >= ? AND f.t_start < ? "
                         "AND f.size_bytes >= 100000 AND h.read_ok = 1 ORDER BY f.t_start",
                         (satellite, wavelength, start, end))
            spikes = set(header_spikes([r[:3] for r in rows]))
            add([(satellite, wavelength, r[0], r[3]) for r in rows if r[0] in spikes], "spike")

    nlong = _rows(conn, frame + " AND h.n_long = 0", params)
    add(stratified_sample(nlong, lambda r: (r[0], r[1]), nlong_per_stratum, rng), "n_long0")

    legacy = _rows(conn, frame.replace("FROM frame f ", "FROM frame f JOIN detection d ON "
                   "d.frame_id = f.id JOIN detector_run r ON r.id = d.run_id ")
                   + " AND r.name = 'legacy_filter' AND d.verdict = 'bad' AND h.n_long >= 1",
                   params)
    add(stratified_sample(legacy, lambda r: r[1], legacy_per_band, rng), "legacy")

    year = lambda r: datetime.datetime.fromtimestamp(r[2], datetime.timezone.utc).year
    everything = conn.execute(
        "SELECT f.satellite, f.wavelength, f.t_start, f.path FROM frame f JOIN header h "
        "ON h.frame_id = f.id WHERE f.t_start >= ? AND f.t_start < ? AND f.size_bytes >= 100000 "
        "AND h.n_long >= 1 AND f.satellite IN (%s) AND (f.id %% 97) = ?" % ",".join("?" * len(satellites)),
        (start, end, *satellites, int(rng.integers(97))))
    add(stratified_sample([tuple(r) for r in everything], lambda r: (r[0], r[1], year(r)),
                          random_per_stratum, rng), "random")

    series = _rows(conn, "SELECT f.t_start, h.img_mean FROM frame f JOIN header h ON "
                   "h.frame_id = f.id WHERE f.satellite = ? AND f.wavelength = 131 AND "
                   "f.t_start >= ? AND f.t_start < ? AND h.img_mean IS NOT NULL "
                   "ORDER BY f.t_start", (satellites[0], start, end))
    peaks = flare_peaks([r[0] for r in series], [r[1] for r in series], flare_count)
    around = []
    for peak in peaks:
        around += _rows(conn, frame + " AND f.t_start >= ? AND f.t_start < ?",
                        params + (peak - FLARE_WINDOW, peak + FLARE_WINDOW))
    if around:
        picks = rng.choice(len(around), min(flare_frames, len(around)), replace=False)
        add([around[i] for i in sorted(picks)], "flare")
    return found, peaks


def cmd_mine(args):
    rng = np.random.default_rng(args.seed)
    conn = db.connect(args.db, readonly=True)
    root = paths.data_root()
    found, peaks = gather(conn, args.start, args.end, rng)
    counts = collections.Counter(s for _, sources in found.values() for s in sources)
    print(f"{len(found)} candidate frames: {dict(counts)}; flare peaks "
          + ", ".join(datetime.datetime.fromtimestamp(p, datetime.timezone.utc)
                      .strftime("%Y-%m-%d %H:%M") for p in peaks), flush=True)

    keys = sorted(found)
    jobs = [(s, w, t, paths.abspath(found[(s, w, t)][0], root)) for s, w, t in keys]
    candidates = []
    with Pool(args.workers, initializer=fitsio.quiet_astropy) as pool:
        for (s, w, t), measured in zip(keys, pool.imap(measure, jobs, 8)):
            candidates.append({"satellite": s, "wavelength": w, "t_start": t,
                               "path": found[(s, w, t)][0], "sources": found[(s, w, t)][1],
                               **measured})
    baselines = band_baselines(candidates)
    for candidate in candidates:
        candidate["proposed"] = propose(candidate, baselines)

    # Most spikes are real solar activity.  Keep all the ones that look defective,
    # and enough of the rest to measure false positives on hard clean frames.
    rng_thin = np.random.default_rng(args.seed + 1)
    spike_clean = [i for i, c in enumerate(candidates)
                   if c["sources"] == ["spike"] and c["proposed"] in CLEAN]
    if len(spike_clean) > args.spike_clean:
        drop = set(rng_thin.choice(spike_clean, len(spike_clean) - args.spike_clean,
                                   replace=False).tolist())
        candidates = [c for i, c in enumerate(candidates) if i not in drop]

    candidates.sort(key=lambda c: (CATEGORIES.index(c["proposed"])
                                   if c["proposed"] in CATEGORIES else -1, c["t_start"]))
    for index, candidate in enumerate(candidates, 1):
        candidate["id"] = index
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    with open(args.out, "w") as handle:
        json.dump({"start": args.start, "end": args.end, "baselines": baselines,
                   "flare_peaks": peaks, "candidates": candidates}, handle, indent=1)
    print(f"wrote {len(candidates)} candidates: "
          f"{dict(collections.Counter(c['proposed'] for c in candidates))}")
    return 0


# ----------------------------------------------------------------- labels & sheets


def parse_corrections(text, known_ids):
    """``ID CATEGORY`` lines -> {id: category}.  '#' starts a comment.

    Unknown ids and categories are errors rather than skipped: a typo must not
    silently leave a frame at its proposed category.
    """
    out = {}
    for number, line in enumerate(text.splitlines(), 1):
        line = line.split("#", 1)[0].strip()
        if not line:
            continue
        parts = line.split()
        if len(parts) != 2 or not parts[0].isdigit():
            raise ValueError(f"line {number}: expected 'ID CATEGORY', got {line!r}")
        ident, category = int(parts[0]), parts[1]
        if ident not in known_ids:
            raise ValueError(f"line {number}: no candidate {ident}")
        if category not in CATEGORIES:
            raise ValueError(f"line {number}: unknown category {category!r}; "
                             f"expected one of {', '.join(CATEGORIES)}")
        out[ident] = category
    return out


def current_labels(candidates, corrections):
    """Each candidate's category after corrections: {id: category}."""
    return {c["id"]: corrections.get(c["id"], c["proposed"]) for c in candidates}


def _load(args):
    with open(args.candidates) as handle:
        audit = json.load(handle)
    corrections = {}
    if args.corrections:
        with open(args.corrections) as handle:
            corrections = parse_corrections(handle.read(),
                                             {c["id"] for c in audit["candidates"]})
    return audit, corrections


def thumbnail(candidate, root, baselines, size=256):
    """The band's own asinh quick-look, captioned with id, frame and ratios."""
    import cv2 as cv
    from PIL import Image, ImageDraw

    image, _ = fitsio.read_image(paths.abspath(candidate["path"], root))
    base = baselines.get(str(candidate["wavelength"])) or baselines.get(candidate["wavelength"])
    if image is None:
        tile = Image.new("RGB", (size, size), (60, 0, 0))
    else:
        scale = 2.0 * (base["far_field"] if base else 0.05)
        x = np.arcsinh(np.clip(np.nan_to_num(image), 0, None) / scale)
        x = np.clip(x / max(np.percentile(x, 99.5), 1e-9), 0, 1)
        small = cv.resize(x.astype(np.float32), (size, size), interpolation=cv.INTER_AREA)
        tile = Image.fromarray((small * 255).astype(np.uint8)).convert("RGB")
    when = datetime.datetime.fromtimestamp(candidate["t_start"], datetime.timezone.utc)
    draw = ImageDraw.Draw(tile)
    draw.text((3, 2), f"#{candidate['id']} g{candidate['satellite']} {candidate['wavelength']}",
              fill=(255, 255, 0))
    draw.text((3, 14), when.strftime("%Y-%m-%d %H:%M"), fill=(255, 255, 0))
    if base and candidate.get("far_field") is not None:
        ratio = candidate["far_field"] / base["far_field"]
        noise = (candidate["noise"] or 0) / base["noise"] if base["noise"] else 0
        draw.text((3, size - 13), f"far x{ratio:.1f} grain x{noise:.0f} {candidate['verdict']}",
                  fill=(0, 255, 255))
    return tile


def _sheet(job):
    from PIL import Image

    path, members, root, baselines, columns = job
    fitsio.quiet_astropy()
    rows = (len(members) + columns - 1) // columns
    sheet = Image.new("RGB", (256 * columns, 256 * rows))
    for index, candidate in enumerate(members):
        sheet.paste(thumbnail(candidate, root, baselines), (256 * (index % columns),
                                                             256 * (index // columns)))
    sheet.save(path, quality=85)
    return path


def cmd_sheets(args):
    audit, corrections = _load(args)
    labels = current_labels(audit["candidates"], corrections)
    root = paths.data_root()
    os.makedirs(args.out_dir, exist_ok=True)
    groups = collections.defaultdict(list)
    for candidate in audit["candidates"]:
        groups[labels[candidate["id"]]].append(candidate)
    per_sheet = args.columns * args.rows
    jobs, index_lines = [], []
    for category in sorted(groups, key=lambda c: (c not in DEFECTS + (PENDING,), c)):
        members = groups[category]
        for n in range(0, len(members), per_sheet):
            chunk = members[n:n + per_sheet]
            name = f"{category}_{n // per_sheet + 1:02d}.jpg"
            jobs.append((os.path.join(args.out_dir, name), chunk, root,
                         audit["baselines"], args.columns))
            index_lines.append(f"{name}: ids {chunk[0]['id']}..{chunk[-1]['id']} "
                               f"({len(chunk)} frames)")
    with Pool(args.workers) as pool:
        for path in pool.imap_unordered(_sheet, jobs):
            print("  wrote", os.path.basename(path), flush=True)
    with open(os.path.join(args.out_dir, "index.txt"), "w") as handle:
        handle.write("\n".join(index_lines) + "\n")
    print(f"{len(jobs)} sheets: " + ", ".join(f"{k} {len(v)}" for k, v in sorted(groups.items())))
    return 0


def cmd_apply(args):
    audit, corrections = _load(args)
    labels = current_labels(audit["candidates"], corrections)
    pending = sorted(i for i, category in labels.items() if category not in CATEGORIES)
    if pending:
        raise SystemExit(f"{len(pending)} candidates still only proposed as {PENDING!r} "
                         f"(e.g. {pending[:10]}); classify them first")
    records = [{"satellite": c["satellite"], "wavelength": c["wavelength"],
                "t_start": c["t_start"], "category": labels[c["id"]],
                "sources": c["sources"]} for c in audit["candidates"]]
    records.sort(key=lambda r: (r["satellite"], r["wavelength"], r["t_start"]))
    os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
    with open(args.out, "w") as handle:
        json.dump({"source": os.path.abspath(args.candidates), "labels": records},
                  handle, indent=1)
    print(f"wrote {len(records)} labels: "
          f"{dict(collections.Counter(r['category'] for r in records))}")
    return 0


def main(argv=None):
    parser = argparse.ArgumentParser(description=__doc__,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = parser.add_subparsers(dest="command", required=True)

    mine = sub.add_parser("mine", help="gather and measure candidate frames")
    mine.add_argument("--start", type=parse_date, required=True)
    mine.add_argument("--end", type=parse_date, required=True)
    mine.add_argument("--out", required=True)
    mine.add_argument("--workers", type=int, default=4)
    mine.add_argument("--seed", type=int, default=1)
    mine.add_argument("--spike-clean", type=int, default=300,
                      help="clean-looking spike frames to keep as hard controls")
    mine.add_argument("--db", default=None)
    mine.set_defaults(func=cmd_mine)

    for name, func, helptext in (("sheets", cmd_sheets, "draw contact sheets by category"),
                                 ("apply", cmd_apply, "write confirmed labels")):
        command = sub.add_parser(name, help=helptext)
        command.add_argument("--candidates", required=True)
        command.add_argument("--corrections", default=None)
        command.set_defaults(func=func)
    sub.choices["sheets"].add_argument("--out-dir", required=True)
    sub.choices["sheets"].add_argument("--columns", type=int, default=8)
    sub.choices["sheets"].add_argument("--rows", type=int, default=6)
    sub.choices["sheets"].add_argument("--workers", type=int, default=4)
    sub.choices["apply"].add_argument("--out", required=True)

    args = parser.parse_args(argv)
    return args.func(args)


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