Adds inference script and utils

Adds a script for running inference on the dataset, including downloading the dataset from Hugging Face Hub if running locally.

It also adds utility functions for reading Colmap reconstructions and providing an empty solution in case of prediction errors.

Renames `test.py` to `train.py` and incorporates prediction and visualization logic.

Adds `.gitignore` file to exclude unnecessary files.

.gitignore

@@ -0,0 +1,4 @@
+
+.vscode/launch.json
+
+__pycache__/

predict.py

@@ -9,7 +9,7 @@ def convert_entry_to_human_readable(entry):
         if 'colmap' in k:
             out[k] = read_colmap_rec(v)
         elif k in ['wf_vertices', 'wf_edges', 'K', 'R', 't', 'depth']:
-            out[k] = np.array(v)
+            out[k] = v
         else:
             out[k]=v
     out['__key__'] = entry['order_id']
@@ -56,10 +56,9 @@ def predict_wireframe(entry) -> Tuple[np.ndarray, List[int]]:
     
     # Merge vertices from all images
     all_3d_vertices, connections_3d = merge_vertices_3d(vert_edge_per_image, 0.5)
-    
-    
     all_3d_vertices_clean, connections_3d_clean  = prune_not_connected(all_3d_vertices, connections_3d, keep_largest=False)
     all_3d_vertices_clean, connections_3d_clean  = prune_too_far(all_3d_vertices_clean, connections_3d_clean, colmap_rec, th = 4.0)
+
     
     if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1:
         print (f'Not enough vertices or connections in the 3D vertices')

script.py

@@ -0,0 +1,101 @@
+from pathlib import Path
+from tqdm import tqdm
+import pandas as pd
+import numpy as np
+from datasets import load_dataset
+from typing import Dict
+from joblib import Parallel, delayed
+import os
+import json
+import gc
+
+from utils import empty_solution
+from predict import predict_wireframe
+
+if __name__ == "__main__":
+    print ("------------ Loading dataset------------ ")
+    param_path = Path('params.json')
+    print(param_path)
+    with param_path.open() as f:
+        params = json.load(f)
+    print(params)
+    import os
+    
+    print('pwd:')
+    os.system('pwd')
+    print(os.system('ls -lahtr'))
+    print('/tmp/data/')
+    print(os.system('ls -lahtr /tmp/data/'))
+    print('/tmp/data/data')
+    print(os.system('ls -lahtrR /tmp/data/data'))
+    
+
+    data_path_test_server = Path('/tmp/data')
+    data_path_local = Path().home() / '.cache/huggingface/datasets/usm3d___hoho25k_test_x/'
+
+    if data_path_test_server.exists():
+        # data_path = data_path_test_server
+        TEST_ENV = True
+    else: 
+        # data_path = data_path_local
+        TEST_ENV = False
+        from huggingface_hub import snapshot_download 
+        _ = snapshot_download(
+            repo_id=params['dataset'],
+            local_dir="/tmp/data",
+            repo_type="dataset",
+        )
+    data_path = data_path_test_server
+    
+    
+    print(data_path)
+
+    # dataset = load_dataset(params['dataset'], trust_remote_code=True, use_auth_token=params['token'])
+    # data_files = {
+    #     "validation": [str(p) for p in [*data_path.rglob('*validation*.arrow')]+[*data_path.rglob('*public*/**/*.tar')]],
+    #     "test": [str(p) for p in [*data_path.rglob('*test*.arrow')]+[*data_path.rglob('*private*/**/*.tar')]],
+    # }
+    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 / 'hoho25k_test_x.py'),
+        data_files=data_files,
+        trust_remote_code=True,
+        writer_batch_size=100
+    )
+
+    print('load with webdataset')
+
+
+    print(dataset, flush=True)
+    
+    print('------------ Now you can do your solution ---------------')
+    solution = []
+
+    def process_sample(sample, i):
+        try:
+            pred_vertices, pred_edges = predict_wireframe(sample)
+        except:
+            pred_vertices, pred_edges = empty_solution()
+        if i %10 == 0:
+            gc.collect()
+        return {
+            'order_id': sample['order_id'], 
+            'wf_vertices': pred_vertices.tolist(),
+            'wf_edges': pred_edges
+        }
+    num_cores = 4
+    
+    for subset_name in dataset.keys():
+        print (f"Predicting {subset_name}")
+        for i, sample in enumerate(tqdm(dataset[subset_name])):
+            res = process_sample(sample, i)
+            solution.append(res)
+
+    print('------------ Saving results ---------------')
+    sub = pd.DataFrame(solution, columns=["order_id", "wf_vertices", "wf_edges"])
+    sub.to_parquet("submission.parquet")
+    print("------------ Done ------------ ")

test.py → train.py

@@ -7,7 +7,7 @@ import io
 import open3d as o3d
 
 from visu import plot_reconstruction_local, plot_wireframe_local, plot_bpo_cameras_from_entry_local, _plotly_rgb_to_normalized_o3d_color
-from utils import read_colmap_rec
+from utils import read_colmap_rec, empty_solution
 
 #from hoho2025.example_solutions import predict_wireframe
 from hoho2025.metric_helper import hss
@@ -19,18 +19,25 @@ scores_hss = []
 scores_f1 = []
 scores_iou = []
 
+show_visu = False
+
 idx = 0
 for a in ds['train']:
     colmap = read_colmap_rec(a['colmap_binary'])
-    pred_vertices, pred_edges = predict_wireframe(a)
 
-    pcd, geometries = plot_reconstruction_local(None, colmap, points=True, cameras=True, crop_outliers=True)
-    wireframe = plot_wireframe_local(None, a['wf_vertices'], a['wf_edges'], a['wf_classifications'])
-    wireframe2 = plot_wireframe_local(None, pred_vertices, pred_edges, None, color='rgb(255, 0, 0)')
-    bpo_cams = plot_bpo_cameras_from_entry_local(None, a)
+    try:
+        pred_vertices, pred_edges = predict_wireframe(a)
+    except:
+        pred_vertices, pred_edges = empty_solution()
+
+    if show_visu:
+        pcd, geometries = plot_reconstruction_local(None, colmap, points=True, cameras=True, crop_outliers=True)
+        wireframe = plot_wireframe_local(None, a['wf_vertices'], a['wf_edges'], a['wf_classifications'])
+        wireframe2 = plot_wireframe_local(None, pred_vertices, pred_edges, None, color='rgb(255, 0, 0)')
+        bpo_cams = plot_bpo_cameras_from_entry_local(None, a)
 
-    visu_all = [pcd] + geometries + wireframe + bpo_cams + wireframe2
-    o3d.visualization.draw_geometries(visu_all, window_name="3D Reconstruction")
+        visu_all = [pcd] + geometries + wireframe + bpo_cams + wireframe2
+        o3d.visualization.draw_geometries(visu_all, window_name="3D Reconstruction")
 
     score = hss(pred_vertices, pred_edges, a['wf_vertices'], a['wf_edges'], vert_thresh=0.5, edge_thresh=0.5)
     print(f"Score: {score}")
@@ -38,3 +45,9 @@ for a in ds['train']:
     scores_f1.append(score.f1)
     scores_iou.append(score.iou)
 
+for i in range(10):
+    print("END OF DATASET")
+print(f"Mean HSS: {np.mean(scores_hss):.4f}")
+print(f"Mean F1: {np.mean(scores_f1):.4f}")
+print(f"Mean IoU: {np.mean(scores_iou):.4f}")
+

utils.py

@@ -1,6 +1,8 @@
 import pycolmap
 import tempfile,zipfile
 import io
+import numpy as np
+from typing import Dict
 
 def read_colmap_rec(colmap_data):
     with tempfile.TemporaryDirectory() as tmpdir:
@@ -8,4 +10,22 @@ def read_colmap_rec(colmap_data):
             zf.extractall(tmpdir)  # unpacks cameras.txt, images.txt, etc. to tmpdir
         # Now parse with pycolmap
         rec = pycolmap.Reconstruction(tmpdir)
-        return rec
+        return rec
+    
+def empty_solution():
+    '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.'''
+    return np.zeros((2,3)), [(0, 1)]
+
+class Sample(Dict):
+    def pick_repr_data(self, x):
+        if hasattr(x, 'shape'):
+            return x.shape
+        if isinstance(x, (str, float, int)):
+            return x
+        if isinstance(x, list):
+            return [type(x[0])] if len(x) > 0 else []
+        return type(x)
+
+    def __repr__(self):
+        # return str({k: v.shape if hasattr(v, 'shape') else [type(v[0])] if isinstance(v, list) else type(v) for k,v in self.items()})
+        return str({k: self.pick_repr_data(v) for k,v in self.items()})