File size: 2,403 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 | from __future__ import annotations
import argparse
from pathlib import Path
from typing import Optional
import numpy as np
from infer.io import load_array, load_frames_from_dir, save_image, save_sequence_grid
from infer.predictor import Predictor
def _parse_args() -> argparse.Namespace:
ap = argparse.ArgumentParser(description="Black-box SavedModel inference: output predicted frames.")
ap.add_argument("--model_dir", type=str, default="savedmodel", help="Path to SavedModel directory")
ap.add_argument("--frames_dir", type=str, default=None, help="Directory containing input frames (images)")
ap.add_argument("--array", type=str, default=None, help="Path to .npy/.npz containing frames")
ap.add_argument(
"--pad_last_frame",
type=str,
default="none",
choices=["none", "zero", "one", "repeat"],
help="If model expects 4 frames but you provide 3, pad the last frame with: zero/one/repeat",
)
ap.add_argument("--out_dir", type=str, default="outputs", help="Output directory")
ap.add_argument("--save_sequence_grid", action="store_true", help="Save a grid of the predicted sequence")
ap.add_argument("--grid_cols", type=int, default=8, help="Columns for sequence grid")
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)
pred = Predictor(args.model_dir)
seq = pred.predict_sequence(frames, pad_last_frame=args.pad_last_frame) # [B,T,H,W,C]
if seq.ndim == 5:
seq0 = seq[0]
last = seq[0, -1]
elif seq.ndim == 4:
seq0 = seq
last = seq[-1]
else:
raise RuntimeError(f"Unexpected prediction shape: {seq.shape}")
out_dir = Path(args.out_dir)
out_dir.mkdir(parents=True, exist_ok=True)
save_image(out_dir / "pred_last.png", last)
if args.save_sequence_grid:
save_sequence_grid(out_dir / "pred_sequence_grid.png", seq0, cols=args.grid_cols)
print(f"Wrote: {out_dir / 'pred_last.png'}")
if args.save_sequence_grid:
print(f"Wrote: {out_dir / 'pred_sequence_grid.png'}")
if __name__ == "__main__":
main()
|