237 lines
10 KiB
Python
237 lines
10 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
"""Gate 0: measure the ceiling of the aligned-candidate design before any training.
|
||
|
|
|
||
|
|
ceiling.py --window wide --per-stratum 12
|
||
|
|
|
||
|
|
The new learned filler can only ever output a per-pixel convex combination of its
|
||
|
|
*aligned* candidates, times a gain bounded to [0.74, 1.35]. That is a deliberate
|
||
|
|
structural cage -- it makes the collapse mode of the previous architecture
|
||
|
|
unreachable -- but it also means the achievable quality is capped by what the
|
||
|
|
candidates contain. This script measures that cap on the real holdout window, on
|
||
|
|
the CPU, in under an hour, so a training run that cannot possibly win is never
|
||
|
|
started.
|
||
|
|
|
||
|
|
For each simulated gap length and counterpart stratum it scores blends of the same
|
||
|
|
aligned candidates the model would see:
|
||
|
|
|
||
|
|
prior softmax of the fixed blend prior == an untrained model, exactly
|
||
|
|
uniform equal weight over valid candidates == the floor selection must beat
|
||
|
|
best1 the single best candidate per band == what perfect *frame* selection earns
|
||
|
|
|
||
|
|
If these sit far below the hand-written baselines (optical_flow 44.79 dB at gap 1 /
|
||
|
|
34.47 at gap 300; fixed crosssat ~39.9 flat), the candidate set or the alignment is
|
||
|
|
at fault and no training run can fix it -- which is exactly how the rotation-sign
|
||
|
|
defect was caught: plain solar_rotation outscored every blend of what should have
|
||
|
|
been the same warped frames.
|
||
|
|
|
||
|
|
Gaps are simulated by hiding the target satellite's frames inside a centred window
|
||
|
|
of the given length, exactly as a real outage hides them; `absent` strata hide the
|
||
|
|
counterpart satellite entirely. Targets never leak: the hidden frames resolve to
|
||
|
|
`missing`, not to suspiciously pristine `suspect` pixels.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import sys
|
||
|
|
import time
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
import bench
|
||
|
|
from suvi import cases, db, fillers, metrics, paths, samples, vfs
|
||
|
|
|
||
|
|
#: Simulated gap lengths, matching the published baseline table's columns.
|
||
|
|
GAPS = (1, 3, 30, 100, 300)
|
||
|
|
#: Margin (slots) a target keeps from the window edges, so anchors always exist.
|
||
|
|
MARGIN = 310
|
||
|
|
|
||
|
|
|
||
|
|
def band_psnr(shown, shown_truth):
|
||
|
|
mse = float(np.mean((shown - shown_truth) ** 2))
|
||
|
|
return float("inf") if mse == 0 else float(10.0 * np.log10(1.0 / mse))
|
||
|
|
|
||
|
|
|
||
|
|
def blend_psnr(candidates_coded, weights, truth_display):
|
||
|
|
"""Display PSNR of a weighted blend, composed in coded space like the model."""
|
||
|
|
composite = np.tensordot(weights, candidates_coded, axes=(0, 0))
|
||
|
|
radiance = samples.decode_from_model(composite)
|
||
|
|
scores = [band_psnr(metrics.to_display(radiance[i], wavelength), truth_display[i])
|
||
|
|
for i, wavelength in enumerate(paths.WAVELENGTHS)]
|
||
|
|
return float(np.mean(scores))
|
||
|
|
|
||
|
|
|
||
|
|
def prior_weights(condition, valid, decay, suspect_penalty=4.0):
|
||
|
|
distance = condition[:, 5] * 6.0
|
||
|
|
logits = -decay * distance - suspect_penalty * condition[:, 2]
|
||
|
|
logits = np.where(valid > 0.5, logits, -np.inf)
|
||
|
|
weights = np.exp(logits - logits.max())
|
||
|
|
return weights / weights.sum()
|
||
|
|
|
||
|
|
|
||
|
|
def score_target(aligned, condition, truth, decays, net=None):
|
||
|
|
"""All the blends for one target. Returns {name: mean-over-band PSNR}."""
|
||
|
|
candidates = aligned[0].numpy() # (S, 6, H, W) coded
|
||
|
|
cond = condition[0].numpy()
|
||
|
|
valid = np.clip(cond[:, 0] + cond[:, 2], 0, 1)
|
||
|
|
usable = valid > 0.5
|
||
|
|
|
||
|
|
truth_display = [metrics.to_display(truth[i], wavelength)
|
||
|
|
for i, wavelength in enumerate(paths.WAVELENGTHS)]
|
||
|
|
|
||
|
|
out = {}
|
||
|
|
if net is not None:
|
||
|
|
import torch
|
||
|
|
with torch.no_grad():
|
||
|
|
predicted = net(aligned, condition)[0].numpy()
|
||
|
|
radiance = samples.decode_from_model(predicted)
|
||
|
|
out["model"] = float(np.mean([
|
||
|
|
band_psnr(metrics.to_display(radiance[i], wavelength), truth_display[i])
|
||
|
|
for i, wavelength in enumerate(paths.WAVELENGTHS)
|
||
|
|
]))
|
||
|
|
uniform = usable / usable.sum()
|
||
|
|
out["uniform"] = blend_psnr(candidates, uniform, truth_display)
|
||
|
|
for decay in decays:
|
||
|
|
out[f"prior[{decay:g}]"] = blend_psnr(
|
||
|
|
candidates, prior_weights(cond, valid, decay), truth_display
|
||
|
|
)
|
||
|
|
|
||
|
|
best_single = []
|
||
|
|
for band, wavelength in enumerate(paths.WAVELENGTHS):
|
||
|
|
shown = np.stack([
|
||
|
|
metrics.to_display(samples.decode_from_model(candidates[s, band]), wavelength)
|
||
|
|
for s in range(candidates.shape[0])
|
||
|
|
])
|
||
|
|
best_single.append(max(
|
||
|
|
band_psnr(shown[s], truth_display[band])
|
||
|
|
for s in range(candidates.shape[0]) if usable[s]
|
||
|
|
))
|
||
|
|
out["best1"] = float(np.mean(best_single))
|
||
|
|
return out
|
||
|
|
|
||
|
|
|
||
|
|
def pick_targets(eligible, count):
|
||
|
|
"""`count` times spread evenly over the eligible list."""
|
||
|
|
if len(eligible) <= count:
|
||
|
|
return list(eligible)
|
||
|
|
stride = len(eligible) / count
|
||
|
|
return [eligible[int(i * stride)] for i in range(count)]
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv=None):
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__,
|
||
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||
|
|
parser.add_argument("--window", default="wide")
|
||
|
|
parser.add_argument("--per-stratum", type=int, default=12)
|
||
|
|
parser.add_argument("--gaps", default=",".join(str(g) for g in GAPS))
|
||
|
|
parser.add_argument("--decays", default="1.5",
|
||
|
|
help="comma-separated prior temporal decays to evaluate")
|
||
|
|
parser.add_argument("--model", default=None,
|
||
|
|
help="also score this trained checkpoint, as a 'model' column")
|
||
|
|
args = parser.parse_args(argv)
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
gaps = tuple(int(g) for g in args.gaps.split(","))
|
||
|
|
decays = tuple(float(d) for d in args.decays.split(","))
|
||
|
|
|
||
|
|
net = None
|
||
|
|
if args.model:
|
||
|
|
loaded = fillers.load_learned(args.model, device="cpu")
|
||
|
|
if loaded is None:
|
||
|
|
raise SystemExit(f"could not load checkpoint {args.model}")
|
||
|
|
net = loaded[0]
|
||
|
|
|
||
|
|
conn = db.connect(readonly=True)
|
||
|
|
window = bench._load_window(conn, args.window)
|
||
|
|
satellites = tuple(json.loads(window["satellites"]))
|
||
|
|
wavelengths = tuple(json.loads(window["wavelengths"]))
|
||
|
|
print(f"window '{args.window}': {window['t_start']} .. {window['t_end']}, "
|
||
|
|
f"satellites {satellites}")
|
||
|
|
|
||
|
|
archive = cases.scan_window(paths.data_root(), satellites, wavelengths,
|
||
|
|
window["t_start"], window["t_end"])
|
||
|
|
overlay = cases.Overlay(archive=archive)
|
||
|
|
reliever = vfs.Reliever(label="ceiling")
|
||
|
|
reader = bench._band_reader(overlay, wavelengths, set(), limit=64,
|
||
|
|
reliever=reliever)
|
||
|
|
clean_times = bench._clean_times(overlay, satellites, wavelengths, set())
|
||
|
|
paired = sorted(set(clean_times[satellites[0]]) & set(clean_times[satellites[1]]))
|
||
|
|
print("clean six-band times: "
|
||
|
|
+ ", ".join(f"G{s}: {len(clean_times[s])}" for s in satellites)
|
||
|
|
+ f", paired: {len(paired)}")
|
||
|
|
|
||
|
|
results = {}
|
||
|
|
started = time.time()
|
||
|
|
for gap in gaps:
|
||
|
|
blocked_offsets = range(-(gap // 2), gap - gap // 2)
|
||
|
|
for counterpart in (True, False):
|
||
|
|
label = (gap, "present" if counterpart else "absent")
|
||
|
|
stratum = []
|
||
|
|
for satellite in satellites:
|
||
|
|
other = next(s for s in satellites if s != satellite)
|
||
|
|
eligible = [
|
||
|
|
t for t in clean_times[satellite][MARGIN:-MARGIN or None]
|
||
|
|
if not counterpart or t in clean_times[other]
|
||
|
|
]
|
||
|
|
for target in pick_targets(eligible, args.per_stratum // 2 or 1):
|
||
|
|
blocked = {target + k * paths.CADENCE for k in blocked_offsets}
|
||
|
|
|
||
|
|
def read(source, when):
|
||
|
|
if source == satellite and when in blocked:
|
||
|
|
return None # inside the simulated gap
|
||
|
|
if not counterpart and source == other:
|
||
|
|
return None # counterpart outage
|
||
|
|
return reader(source, when)
|
||
|
|
|
||
|
|
bad = {(satellite, wl, t) for t in blocked for wl in wavelengths}
|
||
|
|
filtered = {
|
||
|
|
satellite: [t for t in clean_times[satellite]
|
||
|
|
if t not in blocked],
|
||
|
|
other: clean_times[other] if counterpart else [],
|
||
|
|
}
|
||
|
|
stack = bench._build_stack(overlay, satellite, target, satellites,
|
||
|
|
wavelengths, bad, read, filtered)
|
||
|
|
calibration = None
|
||
|
|
if counterpart:
|
||
|
|
usable_pairs = [t for t in paired if t not in blocked]
|
||
|
|
if usable_pairs:
|
||
|
|
pair_time = bench._closest(usable_pairs, target)
|
||
|
|
pair = (reader(other, pair_time),
|
||
|
|
reader(satellite, pair_time))
|
||
|
|
if pair[0] is not None and pair[1] is not None:
|
||
|
|
calibration = pair
|
||
|
|
|
||
|
|
context = fillers.FillContext(stack=stack, calibration=calibration)
|
||
|
|
assembled = fillers.assemble_stack(context, torch,
|
||
|
|
torch.device("cpu"))
|
||
|
|
truth = reader(satellite, target)
|
||
|
|
if assembled is None or truth is None:
|
||
|
|
continue
|
||
|
|
stratum.append(score_target(*assembled, truth, decays, net=net))
|
||
|
|
results[label] = stratum
|
||
|
|
mean = {k: float(np.mean([s[k] for s in stratum])) for k in stratum[0]} \
|
||
|
|
if stratum else {}
|
||
|
|
print(f" gap {gap:>3} {label[1]:>8} n={len(stratum):>2} "
|
||
|
|
+ " ".join(f"{k} {v:6.2f}" for k, v in mean.items()),
|
||
|
|
flush=True)
|
||
|
|
|
||
|
|
print(f"\n{time.time() - started:.0f} s. Reference: optical_flow "
|
||
|
|
"44.79/44.36/40.72/36.85/34.47 dB at gaps 1/3/30/100/300; "
|
||
|
|
"fixed crosssat ~39.5-40.9 flat.")
|
||
|
|
columns = sorted({k for stratum in results.values() for s in stratum for k in s})
|
||
|
|
print(f"\n{'gap':>4} {'counterpart':>11} {'n':>3} "
|
||
|
|
+ " ".join(f"{c:>12}" for c in columns))
|
||
|
|
for (gap, counterpart), stratum in results.items():
|
||
|
|
if not stratum:
|
||
|
|
print(f"{gap:>4} {counterpart:>11} 0")
|
||
|
|
continue
|
||
|
|
row = " ".join(
|
||
|
|
f"{float(np.mean([s[c] for s in stratum])):>12.2f}" for c in columns
|
||
|
|
)
|
||
|
|
print(f"{gap:>4} {counterpart:>11} {len(stratum):>3} {row}")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(main())
|