323 lines
14 KiB
Python
323 lines
14 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
"""Render a filler over real archive data: no synthetic damage, no truth.
|
||
|
|
|
||
|
|
SUVI_MODEL=.../v8_best.pt render_real.py --start 2024-01-01 --end 2025-01-01 \\
|
||
|
|
--out /media/shared/NOAA_GOES/renders/real2024_v8
|
||
|
|
|
||
|
|
The bench measures the model against damage it injected and truth it kept. This
|
||
|
|
is the other thing: the archive as it really is, where the only repairs are for
|
||
|
|
frames the archive itself lost or spoiled, and nothing exists to score them
|
||
|
|
against. The output is the product -- one composite per 4-minute slot, with bad
|
||
|
|
bands reconstructed by the model -- not a number.
|
||
|
|
|
||
|
|
What counts as bad is the **stored production judgment** (``filter_FITS.py
|
||
|
|
judge``: header, lunar, disc, field and clean-neighbour checks), the one run every
|
||
|
|
consumer reads, scored on hand-confirmed real defects. An earlier version judged
|
||
|
|
frames here with header_v1 + disc_v1 alone; that pair passed 41% of the real
|
||
|
|
defects in the audit -- bars, ghost discs, static -- and the renderer composited
|
||
|
|
them, which is what the first 2024 video showed. A day whose frames the run has
|
||
|
|
not judged is refused rather than rendered on a weaker test.
|
||
|
|
|
||
|
|
``--filler optical_flow`` repairs each bad band from its nearest clean frames
|
||
|
|
either side (the classical repair, the same as ``bench.py render --variant new``)
|
||
|
|
instead of the model: the baseline a learned filler is compared against on real
|
||
|
|
data.
|
||
|
|
|
||
|
|
Work proceeds a day at a time. Each day is seen through a window reaching two
|
||
|
|
days back and two forward, which is enough for the model's +/-128-slot stack and
|
||
|
|
guarantees that any gap overlapping the day is seen at its true length up to
|
||
|
|
the 24-hour fill policy, so a long outage is left black rather than invented.
|
||
|
|
Gaps are judged band by band: a slot is only out of policy when some band has
|
||
|
|
gone unobserved by *both* satellites for longer than that. GOES-16 lacking
|
||
|
|
94/131 while GOES-18 lacks 304 is not an outage -- every band is observed, and
|
||
|
|
the model takes each band from whichever frames hold it. Every slot on the
|
||
|
|
cadence timeline is composited, not just slots the archive holds: a short outage
|
||
|
|
on both satellites is still in policy and still gets filled.
|
||
|
|
|
||
|
|
Resumable: a day is marked done only once all its composites are written, and
|
||
|
|
composites are never rewritten.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import datetime
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
|
|
|
||
|
|
import bench
|
||
|
|
from suvi import availability, cases, dataset, db, detectors, fillers, fitsio, paths
|
||
|
|
|
||
|
|
DAY = 86400
|
||
|
|
#: Days of context either side of the day being rendered. Must cover the
|
||
|
|
#: model's stack (+/-128 slots is ~8.5 h) and the dual-gap policy (24 h).
|
||
|
|
BEFORE_DAYS = 2
|
||
|
|
AFTER_DAYS = 2
|
||
|
|
|
||
|
|
|
||
|
|
def day_starts(start, end):
|
||
|
|
"""UTC midnights in [start, end)."""
|
||
|
|
return list(range(start - start % DAY, end, DAY))
|
||
|
|
|
||
|
|
|
||
|
|
def window_for(day):
|
||
|
|
"""The context window a day is rendered through, as [t0, t1)."""
|
||
|
|
return day - BEFORE_DAYS * DAY, day + (1 + AFTER_DAYS) * DAY
|
||
|
|
|
||
|
|
|
||
|
|
def day_name(day):
|
||
|
|
return datetime.datetime.fromtimestamp(day, datetime.timezone.utc).strftime("%Y-%m-%d")
|
||
|
|
|
||
|
|
|
||
|
|
#: Fillers a render can use; each writes ``<out>/<filler>/goes<satellite>``.
|
||
|
|
FILLERS = ("learned", "optical_flow")
|
||
|
|
#: Anchor frames the optical-flow filler keeps between fills. Consecutive bad
|
||
|
|
#: slots share anchors; the bound keeps a bar-storm day from holding hundreds.
|
||
|
|
FLOW_CACHE = 64
|
||
|
|
|
||
|
|
|
||
|
|
class State:
|
||
|
|
"""Resume state under ``<out>/_state``: which (day, satellite) are finished
|
||
|
|
for each filler, so two fillers rendered into one root never share markers."""
|
||
|
|
|
||
|
|
def __init__(self, out, filler="learned"):
|
||
|
|
self.done = os.path.join(out, "_state", "done", filler)
|
||
|
|
os.makedirs(self.done, exist_ok=True)
|
||
|
|
|
||
|
|
def is_done(self, day, satellite):
|
||
|
|
return os.path.exists(os.path.join(self.done, f"{day_name(day)}_g{satellite}"))
|
||
|
|
|
||
|
|
def mark_done(self, day, satellite, summary):
|
||
|
|
with open(os.path.join(self.done, f"{day_name(day)}_g{satellite}"), "w") as handle:
|
||
|
|
json.dump(summary, handle)
|
||
|
|
|
||
|
|
|
||
|
|
def frame_rows(conn, t0, t1, satellites, run_id):
|
||
|
|
"""Index rows for [t0, t1) with the run's verdict:
|
||
|
|
(satellite, wavelength, t_start, path, size_bytes, verdict or None)."""
|
||
|
|
marks = ",".join("?" * len(satellites))
|
||
|
|
return [
|
||
|
|
(row["satellite"], row["wavelength"], row["t_start"], row["path"],
|
||
|
|
row["size_bytes"], row["verdict"])
|
||
|
|
for row in conn.execute(
|
||
|
|
f"SELECT f.satellite, f.wavelength, f.t_start, f.path, f.size_bytes, d.verdict "
|
||
|
|
f"FROM frame f LEFT JOIN detection d ON d.frame_id = f.id AND d.run_id = ? "
|
||
|
|
f"WHERE f.t_start >= ? AND f.t_start < ? AND f.satellite IN ({marks})",
|
||
|
|
(run_id, t0, t1, *satellites),
|
||
|
|
)
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def archive_map(rows, root, wavelengths):
|
||
|
|
"""slot -> absolute path, leaving out stub downloads (they hold no image)."""
|
||
|
|
out = {}
|
||
|
|
for satellite, wavelength, t_start, relpath, size, *_ in rows:
|
||
|
|
if wavelength not in wavelengths:
|
||
|
|
continue
|
||
|
|
if size is not None and size < dataset.MIN_FRAME_BYTES:
|
||
|
|
continue
|
||
|
|
out[(satellite, wavelength, t_start)] = paths.abspath(relpath, root)
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def judged_bad(rows, archive):
|
||
|
|
"""(bad slots, unjudged slots) among the frames in `archive`."""
|
||
|
|
bad, unjudged = set(), set()
|
||
|
|
for satellite, wavelength, t_start, _, _, verdict in rows:
|
||
|
|
slot = (satellite, wavelength, t_start)
|
||
|
|
if slot not in archive:
|
||
|
|
continue
|
||
|
|
if verdict is None:
|
||
|
|
unjudged.add(slot)
|
||
|
|
elif verdict == detectors.BAD:
|
||
|
|
bad.add(slot)
|
||
|
|
return bad, unjudged
|
||
|
|
|
||
|
|
|
||
|
|
class FlowFiller:
|
||
|
|
"""Per-band optical flow between the band's nearest clean frames either side.
|
||
|
|
|
||
|
|
The classical repair the learned model is compared against, using the bench's
|
||
|
|
own `new` variant so the two cannot differ in how they choose the anchors.
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(self, overlay, wavelengths, bad):
|
||
|
|
self.overlay = overlay
|
||
|
|
self.wavelengths = tuple(wavelengths)
|
||
|
|
self.bad = bad
|
||
|
|
self.cache = {}
|
||
|
|
|
||
|
|
def fill(self, satellite, when, bands):
|
||
|
|
"""{band index: array} for `bands`, or None if any of them has no anchor."""
|
||
|
|
if len(self.cache) > FLOW_CACHE:
|
||
|
|
self.cache.clear()
|
||
|
|
out = {}
|
||
|
|
for index in bands:
|
||
|
|
slot = (satellite, self.wavelengths[index], when)
|
||
|
|
image, ok = bench._replacement_band("new", self.overlay, slot, self.bad,
|
||
|
|
self.cache, self.wavelengths)
|
||
|
|
if not ok:
|
||
|
|
return None
|
||
|
|
out[index] = image
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def make_repair(filler, overlay, satellites, wavelengths, bad):
|
||
|
|
"""``repair(satellite, when, bands)`` -> indexable by band, or None."""
|
||
|
|
if filler == "optical_flow":
|
||
|
|
return FlowFiller(overlay, wavelengths, bad).fill
|
||
|
|
# The model reconstructs all six bands jointly, whichever were asked for; a band
|
||
|
|
# no frame in its stack holds comes back NaN, and the slot is unfillable.
|
||
|
|
joint = bench._LearnedFiller(overlay, satellites, wavelengths, bad)
|
||
|
|
|
||
|
|
def repair(satellite, when, bands):
|
||
|
|
filled = joint.fill(satellite, when)
|
||
|
|
if filled is None or any(np.isnan(filled[band]).all() for band in bands):
|
||
|
|
return None
|
||
|
|
return filled
|
||
|
|
|
||
|
|
return repair
|
||
|
|
|
||
|
|
|
||
|
|
def bands_to_fill(overlay, bad, satellite, when, wavelengths):
|
||
|
|
"""Indices of the bands at (satellite, when) that the model must supply."""
|
||
|
|
return [index for index, wavelength in enumerate(wavelengths)
|
||
|
|
if (satellite, wavelength, when) in bad
|
||
|
|
or overlay.path((satellite, wavelength, when)) is None]
|
||
|
|
|
||
|
|
|
||
|
|
def render_day(day, satellite, satellites, wavelengths, rows, root, state, out,
|
||
|
|
workers, batch, max_dual_gap, filler="learned"):
|
||
|
|
"""Composite every slot of one day for one satellite. Returns a summary."""
|
||
|
|
t0, t1 = window_for(day)
|
||
|
|
archive = archive_map(rows, root, wavelengths)
|
||
|
|
bad, unjudged = judged_bad(rows, archive)
|
||
|
|
if unjudged:
|
||
|
|
raise SystemExit(f"{day_name(day)}: {len(unjudged)} frames in the context "
|
||
|
|
f"window have no stored verdict; run `filter_FITS.py "
|
||
|
|
f"judge` over {day_name(t0)}..{day_name(t1)} first")
|
||
|
|
|
||
|
|
overlay = cases.Overlay(archive=archive)
|
||
|
|
unfillable = bench.unfillable_by_band(overlay, satellites, wavelengths, bad,
|
||
|
|
cases.timeline(t0, t1), max_dual_gap)
|
||
|
|
repair = make_repair(filler, overlay, satellites, wavelengths, bad)
|
||
|
|
|
||
|
|
out_dir = os.path.join(out, filler, f"goes{satellite}")
|
||
|
|
os.makedirs(out_dir, exist_ok=True)
|
||
|
|
summary = {"written": 0, "filled": 0, "policy_gap": 0, "unfillable": 0,
|
||
|
|
"existing": 0, "bad_bands": 0}
|
||
|
|
jobs = []
|
||
|
|
for when in cases.timeline(day, day + DAY):
|
||
|
|
target = os.path.join(out_dir, f"Composite-{int(when)}.jpg")
|
||
|
|
if os.path.exists(target):
|
||
|
|
summary["existing"] += 1 # composites are immutable
|
||
|
|
continue
|
||
|
|
if when in unfillable:
|
||
|
|
summary["policy_gap"] += 1 # the gap stays a gap: black frame
|
||
|
|
continue
|
||
|
|
missing = bands_to_fill(overlay, bad, satellite, when, wavelengths)
|
||
|
|
filled = repair(satellite, when, missing) if missing else {}
|
||
|
|
if filled is None:
|
||
|
|
summary["unfillable"] += 1 # no usable source for a missing band
|
||
|
|
continue
|
||
|
|
arrays = []
|
||
|
|
for index, wavelength in enumerate(wavelengths):
|
||
|
|
if index in missing:
|
||
|
|
arrays.append(filled[index])
|
||
|
|
continue
|
||
|
|
image, _ = fitsio.read_image(overlay.path((satellite, wavelength, when)))
|
||
|
|
if image is None:
|
||
|
|
# Judged good but unreadable now: the filler supplies it instead.
|
||
|
|
extra = repair(satellite, when, [index])
|
||
|
|
if extra is None:
|
||
|
|
arrays = None
|
||
|
|
break
|
||
|
|
image = extra[index]
|
||
|
|
arrays.append(image)
|
||
|
|
if arrays is None:
|
||
|
|
summary["unfillable"] += 1
|
||
|
|
continue
|
||
|
|
if missing:
|
||
|
|
summary["filled"] += 1
|
||
|
|
summary["bad_bands"] += len(missing)
|
||
|
|
jobs.append((arrays, when, target))
|
||
|
|
if len(jobs) >= batch:
|
||
|
|
bench._flush_renders(jobs, workers)
|
||
|
|
summary["written"] += len(jobs)
|
||
|
|
jobs.clear()
|
||
|
|
if jobs:
|
||
|
|
bench._flush_renders(jobs, workers)
|
||
|
|
summary["written"] += len(jobs)
|
||
|
|
return summary
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv=None):
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__,
|
||
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
|
|
parser.add_argument("--start", type=bench.parse_date, required=True)
|
||
|
|
parser.add_argument("--end", type=bench.parse_date, required=True,
|
||
|
|
help="exclusive")
|
||
|
|
parser.add_argument("--satellites", default="16,18",
|
||
|
|
help="the pair the model fuses; both are used as context")
|
||
|
|
parser.add_argument("--render", default=None,
|
||
|
|
help="satellites to write composites for, in order "
|
||
|
|
"(default: all of --satellites)")
|
||
|
|
parser.add_argument("--out", required=True)
|
||
|
|
parser.add_argument("--workers", type=int, default=4)
|
||
|
|
parser.add_argument("--batch", type=int, default=12,
|
||
|
|
help="composites per worker dispatch; bounds peak memory")
|
||
|
|
parser.add_argument("--max-dual-gap", type=int, default=availability.MAX_DUAL_GAP)
|
||
|
|
parser.add_argument("--filler", choices=FILLERS, default="learned",
|
||
|
|
help="what repairs bad bands: the model (SUVI_MODEL) or "
|
||
|
|
"per-band optical flow")
|
||
|
|
parser.add_argument("--run", type=int, default=None,
|
||
|
|
help="judgment run id (default: the run made under the current rules)")
|
||
|
|
parser.add_argument("--db", default=None)
|
||
|
|
args = parser.parse_args(argv)
|
||
|
|
|
||
|
|
if args.filler == "learned" and fillers.load_learned() is None:
|
||
|
|
raise SystemExit(f"needs a checkpoint: set {fillers.LEARNED_CHECKPOINT_ENV}")
|
||
|
|
fitsio.quiet_astropy()
|
||
|
|
satellites = tuple(bench.parse_int_list(args.satellites))
|
||
|
|
render = tuple(bench.parse_int_list(args.render)) if args.render else satellites
|
||
|
|
wavelengths = tuple(paths.WAVELENGTHS)
|
||
|
|
conn = db.connect(args.db, readonly=True)
|
||
|
|
run_id = args.run or detectors.judgment_run(conn)
|
||
|
|
if run_id is None:
|
||
|
|
raise SystemExit("No production judgment under the current rules; run "
|
||
|
|
"`filter_FITS.py judge` over the dates (and context) first.")
|
||
|
|
root = paths.data_root()
|
||
|
|
state = State(args.out, args.filler)
|
||
|
|
days = day_starts(args.start, args.end)
|
||
|
|
print(f"real run: {day_name(days[0])} .. {day_name(days[-1])}, "
|
||
|
|
f"{len(days)} days, rendering {render} with {args.filler}, "
|
||
|
|
f"context {satellites}, judgment run {run_id}", flush=True)
|
||
|
|
|
||
|
|
started = time.time()
|
||
|
|
for satellite in render:
|
||
|
|
for index, day in enumerate(days, 1):
|
||
|
|
if state.is_done(day, satellite):
|
||
|
|
continue
|
||
|
|
t0, t1 = window_for(day)
|
||
|
|
rows = frame_rows(conn, t0, t1, satellites, run_id)
|
||
|
|
began = time.time()
|
||
|
|
summary = render_day(day, satellite, satellites, wavelengths, rows, root,
|
||
|
|
state, args.out, args.workers, args.batch,
|
||
|
|
args.max_dual_gap, args.filler)
|
||
|
|
state.mark_done(day, satellite, summary)
|
||
|
|
print(f" g{satellite} [{index:>3}/{len(days)}] {day_name(day)}: "
|
||
|
|
f"{summary['written']} written ({summary['filled']} filled, "
|
||
|
|
f"{summary['bad_bands']} bands), {summary['policy_gap']} left "
|
||
|
|
f"black by policy, {summary['unfillable']} unfillable, "
|
||
|
|
f"{time.time() - began:.0f}s", flush=True)
|
||
|
|
print(f"done in {(time.time() - started) / 3600:.1f} h")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|