| |
| |
|
|
| import sys |
| import os |
| import subprocess |
| sys.path.insert(0, os.path.dirname(__file__)) |
|
|
| |
| def _pip(*pkgs): |
| subprocess.check_call([sys.executable, "-m", "pip", "install", "-q"] + list(pkgs)) |
|
|
| for pkg, mod in [ |
| ("torch==2.4.0", "torch"), |
| ("numpy<2", "numpy"), |
| ("scipy", "scipy"), |
| ("plyfile", "plyfile"), |
| ("pycolmap", "pycolmap"), |
| ("tqdm", "tqdm"), |
| ("joblib", "joblib"), |
| ("huggingface-hub", "huggingface_hub"), |
| ("datasets==3.6.0", "datasets"), |
| ]: |
| try: |
| __import__(mod.split("=")[0]) |
| except ImportError: |
| _pip(pkg) |
|
|
| import json |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from datasets import load_dataset |
| from joblib import Parallel, delayed |
| from tqdm import tqdm |
|
|
| from s23dr_2026.model import get_model, load_checkpoint_compat |
| from s23dr_2026.inference import predict_wireframe_v2 |
| from s23dr_2026.scene import Scene |
| from s23dr_2026.utils import set_random_seed |
|
|
| CHECKPOINT = Path(__file__).parent / "wireframe_detr_cdn_multiscale_384d_128q.pth" |
| TTA_ROTATIONS = 1 |
| NUM_POINTS = 7168 |
| SCORE_THRESH = 0.9 |
| MERGE_DIST = 0.5 |
|
|
|
|
| def empty_solution(sample): |
| return np.zeros((2, 3)), [(0, 1)], sample["order_id"] |
|
|
|
|
| def predict_wireframe_safely(sample, model): |
| try: |
| scene = Scene(sample) |
| verts, edges = predict_wireframe_v2( |
| scene, model, "cuda", |
| pt_type="colmap", |
| num_points=NUM_POINTS, |
| score_threshold=SCORE_THRESH, |
| merge_distance_threshold=MERGE_DIST, |
| ) |
| if len(edges) == 0: |
| verts, edges, _ = empty_solution(sample) |
| except Exception as e: |
| print(f"Failed ({sample.get('order_id', '?')}): {e} — empty solution") |
| verts, edges, _ = empty_solution(sample) |
| edges = [(int(a), int(b)) for a, b in edges] |
| return verts, edges, sample["order_id"] |
|
|
|
|
| if __name__ == "__main__": |
| print("------------ Loading dataset ------------") |
| with open("params.json") as f: |
| params = json.load(f) |
| print(params) |
|
|
| data_path = Path("/tmp/data") |
| if not data_path.exists(): |
| from huggingface_hub import snapshot_download |
| snapshot_download(repo_id=params["dataset"], local_dir=str(data_path), 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")], |
| } |
| print(data_files) |
| 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("------------ Loading model ------------") |
| set_random_seed(0) |
| checkpoint = torch.load(CHECKPOINT, map_location="cuda") |
| model = get_model(checkpoint, num_classes=1) |
| load_checkpoint_compat(model, checkpoint) |
| model.to("cuda").eval() |
| print(f"Model loaded from {CHECKPOINT}") |
| print(f"Inference: TTA={TTA_ROTATIONS}, {NUM_POINTS} pts, threshold={SCORE_THRESH}") |
|
|
| print("------------ Predicting ------------") |
| solution = [] |
| for subset_name in dataset: |
| print(f"Predicting on {subset_name}") |
| preds = Parallel(n_jobs=1, prefer="processes")( |
| delayed(predict_wireframe_safely)(sample, model) |
| for sample in tqdm(dataset[subset_name]) |
| ) |
| for verts, edges, order_id in preds: |
| print(f"{order_id}: {len(verts)} verts, {len(edges)} edges") |
| solution.append({ |
| "order_id": order_id, |
| "wf_vertices": verts.tolist(), |
| "wf_edges": edges, |
| }) |
|
|
| print("------------ Saving ------------") |
| output_path = Path(params.get("output_path", ".")) |
| output_path.mkdir(parents=True, exist_ok=True) |
| for save_path in [Path("submission.json"), output_path / "submission.json"]: |
| save_path.parent.mkdir(parents=True, exist_ok=True) |
| with open(save_path, "w") as f: |
| json.dump(solution, f) |
| print(f"Saved → {save_path}") |
|
|
| print("------------ Done ------------") |
|
|
