File size: 3,270 Bytes
8b78124 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | from __future__ import annotations
import argparse
import csv
import json
from pathlib import Path
from typing import Any, Dict, List, Optional
import numpy as np
from infer.io import load_array, load_frames_from_dir
from infer.predictor import Predictor
from infer.anomaly import compute_anomaly_from_4frames
def _parse_args() -> argparse.Namespace:
ap = argparse.ArgumentParser(description="Compute per-window anomaly scores (ahat_error and 1-SSIM vectors).")
ap.add_argument("--model_dir", type=str, default="savedmodel")
ap.add_argument("--frames_dir", type=str, default=None)
ap.add_argument("--array", type=str, default=None)
ap.add_argument("--out_json", type=str, default="anomaly.json")
ap.add_argument("--out_csv", type=str, default="anomaly.csv")
ap.add_argument("--start", type=int, default=0, help="Start frame index (inclusive) on the full sequence")
ap.add_argument("--end", type=int, default=-1, help="End frame index (exclusive); -1 means to the end")
return ap.parse_args()
def main() -> None:
args = _parse_args()
if (args.frames_dir is None) == (args.array is None):
raise SystemExit("Provide exactly one of --frames_dir or --array")
if args.frames_dir is not None:
frames = load_frames_from_dir(args.frames_dir) # [T,H,W,C] 0..255
else:
frames = load_array(args.array)
if frames.ndim == 5:
if frames.shape[0] != 1:
raise SystemExit("Batch arrays are not supported in CLI; please provide a single sequence.")
frames = frames[0]
if frames.ndim != 4 or frames.shape[0] < 4:
raise SystemExit(f"Expected frames with shape [T,H,W,C] and T>=4, got {frames.shape}")
T = frames.shape[0]
start = max(0, int(args.start))
end = T if int(args.end) == -1 else min(T, int(args.end))
if end - start < 4:
raise SystemExit("Range too small; need at least 4 frames.")
pred = Predictor(args.model_dir)
if pred.sequence_len is not None and pred.sequence_len != 4:
raise SystemExit(f"This tool assumes model sequence_len=4, got {pred.sequence_len}")
rows: List[Dict[str, Any]] = []
for t in range(start + 3, end):
win = frames[t - 3 : t + 1] # [4,H,W,C] where win[3] is GT frame4
res = compute_anomaly_from_4frames(pred, win)
rows.append(
{
"t": int(t),
"ahat_error": res.ahat_error.tolist(),
"dissim_1mssim": res.dissim_1mssim.tolist(),
}
)
out_json = Path(args.out_json)
out_json.parent.mkdir(parents=True, exist_ok=True)
out_json.write_text(json.dumps({"items": rows}, ensure_ascii=False, indent=2), encoding="utf-8")
out_csv = Path(args.out_csv)
out_csv.parent.mkdir(parents=True, exist_ok=True)
with out_csv.open("w", newline="", encoding="utf-8") as f:
w = csv.writer(f)
w.writerow(["t", "ahat_error_1", "ahat_error_2", "ahat_error_3", "ahat_error_4", "dissim_1", "dissim_2", "dissim_3", "dissim_4"])
for r in rows:
a = r["ahat_error"]
d = r["dissim_1mssim"]
w.writerow([r["t"], *a, *d])
print(f"Wrote: {out_json}")
print(f"Wrote: {out_csv}")
if __name__ == "__main__":
main()
|