File size: 5,493 Bytes
d1999d9 | 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 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 | """S23DR 2026 submission: rayv9_learnt_baseline_snap
Pipeline per scene:
1. v1d heatmap -> spatial NMS -> rays
2. rays + SfM -> voxel volume -> RayVoxelTransformer -> hybrid NMS -> vertices
3. fuse_and_sample -> learned baseline -> wireframe
4. midpoint-snap baseline verts to v9, append unmatched v9 verts
Val HSS (100 scenes): baseline=0.352 snap=0.411 (+0.059, wins=85/100)
Vertex F1@0.5=0.494 F1@1.0=0.685
Usage:
# local smoke-test on training split (n scenes)
python script.py --mode local --n_scenes 10
# competition submission (reads params.json, writes submission.json)
python script.py
"""
import argparse
import json
import time
from pathlib import Path
import numpy as np
import torch
import v9_inference as v9
import baseline_inference as bl
from snap import snap_midpoint_plus_unmatched
from s23dr_2026_example.point_fusion import FuserConfig
SCRIPT_DIR = Path(__file__).resolve().parent
V1D_CKPT = SCRIPT_DIR / "v1d_checkpoint.pt"
V9_CKPT = SCRIPT_DIR / "v9_checkpoint.pt"
BASELINE_CKPT = SCRIPT_DIR / "baseline_checkpoint.pt"
def empty_solution():
return np.zeros((2, 3)), [(0, 1)]
def run(dataset, v1d_model, v9_model, v9_img_size, bl_model, device, n_scenes=None):
cfg = FuserConfig()
rng = np.random.RandomState(2718)
solution = []
processed = 0
t0 = time.time()
for subset_name in dataset:
print(f"\nProcessing {subset_name}...", flush=True)
for sample in dataset[subset_name]:
if n_scenes is not None and processed >= n_scenes:
break
order_id = sample["order_id"]
v9_verts = np.zeros((0, 3))
try:
v9_verts = v9.predict_vertices(
sample, v1d_model, v9_model, v9_img_size, device)
except Exception as e:
print(f" v9 failed {order_id}: {e}", flush=True)
bl_result = None
try:
fused = bl.fuse_and_sample(sample, cfg, rng)
if fused is not None:
bl_result = bl.predict(fused, bl_model, device)
except Exception as e:
print(f" baseline failed {order_id}: {e}", flush=True)
if bl_result is None:
pred_v, pred_e = empty_solution()
else:
pred_v, pred_e = snap_midpoint_plus_unmatched(bl_result[0], bl_result[1], v9_verts)
solution.append({
"order_id": order_id,
"wf_vertices": pred_v.tolist(),
"wf_edges": [(int(a), int(b)) for a, b in pred_e],
})
processed += 1
elapsed = time.time() - t0
print(f" [{processed}] {order_id} "
f"v9={len(v9_verts)} bl={'ok' if bl_result else 'fail'} "
f"{elapsed:.0f}s elapsed", flush=True)
return solution
def load_models(device):
print("Loading v1d...", flush=True)
v1d_model = v9.load_v1d(V1D_CKPT, device)
print("Loading v9...", flush=True)
v9_model, v9_img_size = v9.load_v9(V9_CKPT, device)
print("Loading baseline...", flush=True)
bl_model = bl.load_model(BASELINE_CKPT, device)
return v1d_model, v9_model, v9_img_size, bl_model
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--mode", choices=["local", "submit"], default="submit",
help="local: stream from hoho22k_2026_trainval training split; "
"submit: read params.json and use test data")
parser.add_argument("--n_scenes", type=int, default=None,
help="cap number of scenes (local mode)")
args = parser.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}", flush=True)
v1d_model, v9_model, v9_img_size, bl_model = load_models(device)
from datasets import load_dataset
if args.mode == "local":
print("Mode: local (hoho22k_2026_trainval / train split)", flush=True)
dataset = load_dataset("usm3d/hoho22k_2026_trainval", split="train",
streaming=True, trust_remote_code=True)
# Wrap in dict to match the multi-subset loop
dataset = {"train": dataset}
else:
print("Mode: submit", flush=True)
with open("params.json") as f:
params = json.load(f)
print(f"Competition: {params.get('competition_id', '?')}", flush=True)
data_path = Path("/tmp/data")
if not data_path.exists():
from huggingface_hub import snapshot_download
snapshot_download(repo_id=params["dataset"], local_dir="/tmp/data",
repo_type="dataset")
data_files = {
"validation": [str(p) for p in data_path.rglob("*public*/**/*.tar")],
"test": [str(p) for p in data_path.rglob("*private*/**/*.tar")],
}
dataset = load_dataset(
str(data_path / "hoho22k_2026_test_x_anon.py"),
data_files=data_files,
trust_remote_code=True,
writer_batch_size=100,
)
print(f"Loaded: {dataset}", flush=True)
solution = run(dataset, v1d_model, v9_model, v9_img_size, bl_model,
device, n_scenes=args.n_scenes)
with open("submission.json", "w") as f:
json.dump(solution, f)
print(f"\nSaved submission.json ({len(solution)} entries)", flush=True)
|