infer.py

"""CLI for argus-3d.

Usage:
  python infer.py detect <image> --K K.json [--depth depth.png]
  python infer.py perceive <image> --K K.json
"""
import argparse
import json
import sys

import numpy as np
from PIL import Image

from argus_3d import Argus3D


def _load_K(path: str) -> np.ndarray:
    with open(path, "r") as f:
        return np.array(json.load(f), dtype=np.float64)


def _load_depth(path: str) -> np.ndarray:
    arr = np.array(Image.open(path))
    if arr.dtype == np.uint16:
        return arr.astype(np.float32) / 1000.0
    return arr.astype(np.float32)


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    p_det = sub.add_parser("detect")
    p_det.add_argument("image")
    p_det.add_argument("--K", required=True)
    p_det.add_argument("--depth", default=None)
    p_det.add_argument("--repo", default="phanerozoic/argus-3d")
    p_det.add_argument("--device", default="cuda")
    p_per = sub.add_parser("perceive")
    p_per.add_argument("image")
    p_per.add_argument("--K", required=True)
    p_per.add_argument("--depth", default=None)
    p_per.add_argument("--repo", default="phanerozoic/argus-3d")
    p_per.add_argument("--device", default="cuda")
    args = ap.parse_args()

    model = Argus3D.from_pretrained(args.repo, device=args.device)
    K = _load_K(args.K)
    depth = _load_depth(args.depth) if args.depth else None

    if args.cmd == "detect":
        boxes = model.detect(args.image, K, depth=depth)
        out = [
            {
                "box": [b.cx, b.cy, b.cz, b.w, b.h, b.d, b.theta],
                "score": b.score,
                "n_inliers": b.n_inliers,
            }
            for b in boxes
        ]
        json.dump(out, sys.stdout, indent=2)
    else:
        out = model.perceive(args.image, K, depth=depth)
        out["fg_score_map"] = out["fg_score_map"].tolist()
        out["depth_map"] = out["depth_map"].tolist()
        out["boxes"] = [
            {
                "box": [b.cx, b.cy, b.cz, b.w, b.h, b.d, b.theta],
                "score": b.score,
                "n_inliers": b.n_inliers,
            }
            for b in out["boxes"]
        ]
        json.dump(out, sys.stdout, indent=2)


if __name__ == "__main__":
    main()