Fix sys.path and networkx dependency for Hugging Face

script.py

@@ -1,472 +1,130 @@
-### This is example of the script that will be run in the test environment.
-
-### You can change the rest of the code to define and test your solution.
-### However, you should not change the signature of the provided function.
-### The script saves "submission.json" file in the current directory.
-### You can use any additional files and subdirectories to organize your code.
+"""S23DR 2026 submission — sklearn edges + edge validation + improved heuristic."""
 
 from pathlib import Path
 from tqdm import tqdm
-import json
 import numpy as np
 from datasets import load_dataset
 from typing import Dict
-from joblib import Parallel, delayed
-from sklearn.cluster import DBSCAN
-from hoho2025 import example_solutions as hoho_example
-
-VERTEX_MERGE_EPS = 1.0
-EDGE_MIN_LENGTH = 0.5
-EDGE_MIN_SUPPORT_IMAGES = 1
-EDGE_MAX_ANGLE_DEG = 45.0
-VERTEX_MAX_COLMAP_DIST = 2.0
-VERTEX_MIN_EDGE_DEGREE = 1
-VERTEX_MIN_VIEW_COUNT = 1
-
-def empty_solution(sample):
-    '''Return a minimal valid solution, i.e. 2 vertices and 1 edge.'''
-    return np.zeros((2,3)), [(0, 1)]
-
-def _unit_vector(vector):
-    norm = np.linalg.norm(vector)
-    if norm == 0:
-        return None
-    return vector / norm
-
-def merge_vertices_dbscan(vertices_3d, edges, eps=VERTEX_MERGE_EPS):
-    """Cluster nearby 3D vertices, replace clusters with centroids, and remap edges."""
-    vertices_3d = np.asarray(vertices_3d, dtype=float)
-    if vertices_3d.ndim != 2 or vertices_3d.shape[1] != 3 or len(vertices_3d) == 0:
-        return vertices_3d.reshape((-1, 3)), []
-
-    clustering = DBSCAN(eps=eps, min_samples=1).fit(vertices_3d)
-    labels = clustering.labels_
-    unique_labels = np.unique(labels)
-    label_to_new_idx = {label: idx for idx, label in enumerate(unique_labels)}
-
-    merged_vertices = np.stack(
-        [vertices_3d[labels == label].mean(axis=0) for label in unique_labels],
-        axis=0,
-    )
-
-    remapped_edges = []
-    seen_edges = set()
-    for a, b in edges:
-        a = int(a)
-        b = int(b)
-        if a < 0 or b < 0 or a >= len(labels) or b >= len(labels):
-            continue
-
-        new_a = label_to_new_idx[labels[a]]
-        new_b = label_to_new_idx[labels[b]]
-        if new_a == new_b:
-            continue
-
-        edge = (min(new_a, new_b), max(new_a, new_b))
-        if edge in seen_edges:
-            continue
-
-        seen_edges.add(edge)
-        remapped_edges.append(edge)
-
-    return merged_vertices, remapped_edges
-
-def merge_vertices_dbscan_with_edge_stats(vert_edge_per_image, eps=VERTEX_MERGE_EPS):
-    """Merge multi-view vertices with DBSCAN and keep edge support/direction stats."""
-    all_vertices = []
-    old_edges = []
-    vertex_types = []
-    vertex_src_images = []
-    cur_start = 0
-
-    for img_idx, (vertices, connections, vertices_3d) in vert_edge_per_image.items():
-        vertices_3d = np.asarray(vertices_3d, dtype=float).reshape((-1, 3))
-        if len(vertices_3d) == 0:
-            continue
-
-        all_vertices.append(vertices_3d)
-        vertex_types.extend([int(v.get('type') == 'apex') for v in vertices])
-        vertex_src_images.extend([img_idx] * len(vertices_3d))
-
-        for a, b in connections:
-            a = int(a)
-            b = int(b)
-            if a < 0 or b < 0 or a >= len(vertices_3d) or b >= len(vertices_3d):
-                continue
-            old_edges.append((cur_start + a, cur_start + b, img_idx))
-
-        cur_start += len(vertices_3d)
-
-    if not all_vertices:
-        return np.empty((0, 3)), [], np.array([], dtype=int), {}, {}
-
-    all_vertices = np.concatenate(all_vertices, axis=0)
-    labels = DBSCAN(eps=eps, min_samples=1).fit(all_vertices).labels_
-
-    # Keep apex and non-apex vertices from collapsing into the same final corner.
-    vertex_types = np.asarray(vertex_types, dtype=int)
-    cluster_keys = [(int(label), int(vtype)) for label, vtype in zip(labels, vertex_types)]
-    unique_keys = sorted(set(cluster_keys))
-    key_to_new_idx = {key: idx for idx, key in enumerate(unique_keys)}
-    old_to_new_idx = np.array([key_to_new_idx[key] for key in cluster_keys], dtype=int)
-
-    merged_vertices = np.stack(
-        [
-            all_vertices[old_to_new_idx == idx].mean(axis=0)
-            for idx in range(len(unique_keys))
-        ],
-        axis=0,
-    )
-
-    vertex_src_images = np.asarray(vertex_src_images)
-    vertex_view_count = np.array(
-        [
-            len(set(vertex_src_images[old_to_new_idx == idx]))
-            for idx in range(len(unique_keys))
-        ],
-        dtype=int,
-    )
-
-    edge_image_sets = {}
-    edge_dirs = {}
-    for old_a, old_b, img_idx in old_edges:
-        new_a = int(old_to_new_idx[old_a])
-        new_b = int(old_to_new_idx[old_b])
-        if new_a == new_b:
-            continue
-
-        edge = (min(new_a, new_b), max(new_a, new_b))
-        direction = _unit_vector(all_vertices[old_b] - all_vertices[old_a])
-        if direction is None:
-            continue
-
-        final_direction = _unit_vector(merged_vertices[edge[1]] - merged_vertices[edge[0]])
-        if final_direction is None:
-            continue
-        if np.dot(direction, final_direction) < 0:
-            direction = -direction
-
-        edge_image_sets.setdefault(edge, set()).add(img_idx)
-        edge_dirs.setdefault(edge, []).append(direction)
-
-    edge_vote_count = {edge: len(imgs) for edge, imgs in edge_image_sets.items()}
-    edge_angle_ok = {
-        edge: edge_directions_are_consistent(dirs, EDGE_MAX_ANGLE_DEG)
-        for edge, dirs in edge_dirs.items()
-    }
-
-    return merged_vertices, list(edge_image_sets.keys()), vertex_view_count, edge_vote_count, edge_angle_ok
-
-def edge_directions_are_consistent(directions, max_angle_deg=EDGE_MAX_ANGLE_DEG):
-    """Check whether per-image 3D edge directions agree after sign alignment."""
-    if len(directions) <= 1:
-        return True
-
-    directions = np.asarray(directions, dtype=float)
-    mean_direction = _unit_vector(directions.mean(axis=0))
-    if mean_direction is None:
-        return False
-
-    min_cos = np.cos(np.deg2rad(max_angle_deg))
-    return bool(np.all(directions @ mean_direction >= min_cos))
-
-def filter_edges_by_geometry(vertices_3d, edges, edge_vote_count=None, edge_angle_ok=None,
-                             min_len=EDGE_MIN_LENGTH, min_support=EDGE_MIN_SUPPORT_IMAGES):
-    """Drop short, weakly supported, or direction-inconsistent edges."""
-    vertices_3d = np.asarray(vertices_3d, dtype=float)
-    if len(vertices_3d) == 0:
-        return []
-
-    filtered_edges = []
-    seen_edges = set()
-    for a, b in edges:
-        a = int(a)
-        b = int(b)
-        if a < 0 or b < 0 or a >= len(vertices_3d) or b >= len(vertices_3d) or a == b:
-            continue
-
-        edge = (min(a, b), max(a, b))
-        if edge in seen_edges:
-            continue
-
-        edge_length = np.linalg.norm(vertices_3d[edge[1]] - vertices_3d[edge[0]])
-        if edge_length < min_len:
-            continue
-
-        if edge_vote_count is not None and edge_vote_count.get(edge, 0) < min_support:
-            continue
-
-        if edge_angle_ok is not None and not edge_angle_ok.get(edge, True):
-            continue
-
-        seen_edges.add(edge)
-        filtered_edges.append(edge)
-
-    return filtered_edges
-
-def prune_bad_vertices(vertices_3d, edges, colmap_rec, vertex_view_count,
-                       max_colmap_dist=VERTEX_MAX_COLMAP_DIST,
-                       min_edge_degree=VERTEX_MIN_EDGE_DEGREE,
-                       min_view_count=VERTEX_MIN_VIEW_COUNT):
-    """Remove vertices with weak COLMAP, edge-degree, or multi-view support."""
-    vertices_3d = np.asarray(vertices_3d, dtype=float)
-    vertex_view_count = np.asarray(vertex_view_count, dtype=int)
-    if len(vertices_3d) == 0:
-        return vertices_3d.reshape((-1, 3)), []
-
-    valid_edges = []
-    for a, b in edges:
-        a = int(a)
-        b = int(b)
-        if 0 <= a < len(vertices_3d) and 0 <= b < len(vertices_3d) and a != b:
-            valid_edges.append((a, b))
-
-    colmap_mask = np.ones(len(vertices_3d), dtype=bool)
-    xyz_sfm = [point.xyz for point in colmap_rec.points3D.values()]
-    if xyz_sfm:
-        xyz_sfm = np.asarray(xyz_sfm, dtype=float)
-        diff = vertices_3d[:, None, :] - xyz_sfm[None, :, :]
-        min_dist = np.sqrt((diff ** 2).sum(axis=-1)).min(axis=1)
-        colmap_mask = min_dist <= max_colmap_dist
-
-    view_mask = vertex_view_count >= min_view_count
-    keep_mask = colmap_mask & view_mask
-
-    while True:
-        degree = np.zeros(len(vertices_3d), dtype=int)
-        for a, b in valid_edges:
-            if keep_mask[a] and keep_mask[b]:
-                degree[a] += 1
-                degree[b] += 1
-
-        next_keep_mask = keep_mask & (degree >= min_edge_degree)
-        if np.array_equal(next_keep_mask, keep_mask):
-            break
-        keep_mask = next_keep_mask
-
-    old_to_new = {}
-    new_vertices = []
-    for old_idx, keep in enumerate(keep_mask):
-        if keep:
-            old_to_new[old_idx] = len(new_vertices)
-            new_vertices.append(vertices_3d[old_idx])
-
-    new_edges = []
-    seen_edges = set()
-    for a, b in valid_edges:
-        if a not in old_to_new or b not in old_to_new:
-            continue
-        edge = tuple(sorted((old_to_new[a], old_to_new[b])))
-        if edge in seen_edges:
-            continue
-        seen_edges.add(edge)
-        new_edges.append(edge)
-
-    if not new_vertices:
-        return np.empty((0, 3)), []
-    return np.asarray(new_vertices), new_edges
-
-def predict_wireframe_filtered(sample, verbose=False):
-    """Baseline prediction with DBSCAN vertex merge plus stricter edge filtering."""
-    good_entry = hoho_example.convert_entry_to_human_readable(sample)
-    vert_edge_per_image = {}
-    colmap_rec = None
-
-    for i, (gest, depth, img_id, ade_seg) in enumerate(zip(
-        good_entry['gestalt'],
-        good_entry['depth'],
-        good_entry['image_ids'],
-        good_entry['ade'],
-    )):
-        if 'colmap' in good_entry:
-            colmap_rec = good_entry['colmap']
-        else:
-            colmap_rec = good_entry['colmap_binary']
-
-        depth_size = (np.array(depth).shape[1], np.array(depth).shape[0])
-        gest_seg_np = np.array(gest.resize(depth_size)).astype(np.uint8)
-        vertices, connections = hoho_example.get_vertices_and_edges_from_segmentation(
-            gest_seg_np, edge_th=10.
-        )
-
-        ade_seg_np = np.array(ade_seg.resize(depth_size)).astype(np.uint8)
-        vertices, connections = hoho_example.filter_vertices_by_background(
-            vertices, connections, ade_seg_np
-        )
-
-        if (len(vertices) < 2) or (len(connections) < 1):
-            if verbose:
-                print(f'Not enough vertices or connections found in image {i}, skipping.')
-            vert_edge_per_image[i] = [], [], np.empty((0, 3))
-            continue
-
-        vertices_3d = hoho_example.create_3d_wireframe_single_image(
-            vertices, connections, depth, colmap_rec, img_id, ade_seg, verbose=verbose
-        )
-        vert_edge_per_image[i] = vertices, connections, vertices_3d
-
-    all_3d_vertices, connections_3d, vertex_view_count, edge_vote_count, edge_angle_ok = (
-        merge_vertices_dbscan_with_edge_stats(vert_edge_per_image, VERTEX_MERGE_EPS)
-    )
-    connections_3d = filter_edges_by_geometry(
-        all_3d_vertices,
-        connections_3d,
-        edge_vote_count=edge_vote_count,
-        edge_angle_ok=edge_angle_ok,
-    )
-    if colmap_rec is None or len(all_3d_vertices) < 2 or len(connections_3d) < 1:
-        return empty_solution(sample)
-
-    all_3d_vertices_clean, connections_3d_clean = prune_bad_vertices(
-        all_3d_vertices, connections_3d, colmap_rec, vertex_view_count
-    )
-    all_3d_vertices_clean, connections_3d_clean = hoho_example.prune_not_connected(
-        all_3d_vertices_clean, connections_3d_clean, keep_largest=False
-    )
-
-    if (len(all_3d_vertices_clean) < 2) or len(connections_3d_clean) < 1:
-        if verbose:
-            print(f'Not enough vertices or connections in the 3D vertices')
-        return empty_solution(sample)
+import os
+import json
+import gc
+import subprocess
+import sys
 
-    connections_3d_clean = [(int(a), int(b)) for a, b in connections_3d_clean]
-    return all_3d_vertices_clean, connections_3d_clean
+# Ensure local imports work regardless of how script.py is invoked
+current_dir = str(Path(__file__).parent.absolute())
+if current_dir not in sys.path:
+    sys.path.insert(0, current_dir)
 
-def predict_wireframe_safely(sample):
+def install_if_missing(package):
     try:
-        pred_vertices, pred_edges = predict_wireframe_filtered(sample)
-    except Exception as e:
-        print (f"Failed due to {e},  returning empty solution")
-        pred_vertices, pred_edges = empty_solution(sample)
-    pred_edges  = [(int(a), int(b)) for a, b in pred_edges] # to remove possible np.int64
-    return pred_vertices, pred_edges, sample['order_id']
+        __import__(package)
+    except ImportError:
+        subprocess.check_call([sys.executable, "-m", "pip", "install", package])
 
-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)
 
+class Sample(Dict):
     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()})
-
-def load_competition_dataset(params):
-    """
-    Loads dataset both:
-    1. Locally from public parquet files.
-    2. In official competition environment from /tmp/data.
-    """
-    import os
-
-    data_path = Path("/tmp/data")
-
-    print("------------ Dataset path check ------------")
-    print("pwd:")
-    os.system("pwd")
-
-    print("/tmp/data:")
-    os.system("ls -lahtr /tmp/data || true")
+        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()})
 
-    print("/tmp/data/data:")
-    os.system("ls -lahtr /tmp/data/data || true")
 
-    # Case 1: local debugging with public parquet dataset
-    parquet_dir = data_path / "data"
-    train_parquet = list(parquet_dir.glob("train-*.parquet"))
-    val_parquet = list(parquet_dir.glob("validation-*.parquet"))
+def empty_solution():
+    return np.zeros((2, 3)), [(0, 1)]
 
-    if len(train_parquet) > 0 or len(val_parquet) > 0:
-        print("Loading local/public parquet dataset")
 
-        data_files = {}
-
-        if len(train_parquet) > 0:
-            data_files["train"] = str(parquet_dir / "train-*.parquet")
-
-        if len(val_parquet) > 0:
-            data_files["validation"] = str(parquet_dir / "validation-*.parquet")
-
-        dataset = load_dataset("parquet", data_files=data_files)
-        return dataset
-
-    # Case 2: official test environment with custom dataset script
-    dataset_script_candidates = list(data_path.glob("*.py"))
-
-    if len(dataset_script_candidates) > 0:
-        dataset_script = dataset_script_candidates[0]
-        print(f"Loading official dataset script: {dataset_script}")
-
-        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:", data_files)
-
-        dataset = load_dataset(
-            str(dataset_script),
-            data_files=data_files,
-            trust_remote_code=True,
-            writer_batch_size=100,
-        )
-
-        return dataset
-
-    # Case 3: fallback download for local run
-    print("No local /tmp/data files found. Trying Hugging Face download.")
-
-    from huggingface_hub import snapshot_download
-
-    snapshot_download(
-        repo_id=params["dataset"],
-        local_dir="/tmp/data",
-        repo_type="dataset",
-        token=params.get("token", None),
-    )
+def process_sample(sample, i, sklearn_model=None):
+    try:
+        from sklearn_submission import predict_wireframe_sklearn
+        pred_vertices, pred_edges = predict_wireframe_sklearn(sample, sklearn_model)
+    except Exception as e:
+        if i < 5:
+            print(f"  Sample {i} sklearn failed: {e}", flush=True)
+        try:
+            from hoho2025.example_solutions import predict_wireframe
+            pred_vertices, pred_edges = predict_wireframe(sample)
+        except Exception:
+            pred_vertices, pred_edges = empty_solution()
+    if i % 10 == 0:
+        gc.collect()
+    return {
+        'order_id': sample['order_id'],
+        'wf_vertices': np.array(pred_vertices).tolist(),
+        'wf_edges': [list(e) for e in pred_edges],
+    }
 
-    return load_competition_dataset(params)
 
-import json
 if __name__ == "__main__":
-    print ("------------ Loading dataset------------ ")
+    print("------------ Loading dataset ------------", flush=True)
     param_path = Path('params.json')
-    print(param_path)
-    with param_path.open() as f:
-        params = json.load(f)
-    safe_params = dict(params)
-    if "token" in safe_params:
-        safe_params["token"] = "hf_******"
-
-    print(safe_params)
+    if param_path.exists():
+        with param_path.open() as f:
+            params = json.load(f)
+    else:
+        params = {"dataset": "usm3d/hoho22k_2026_trainval", "output_path": "."}
+
+    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")
+
+    os.system(f'ls -lahtrR {data_path}')
+
+    data_files = {}
+    public_tars = sorted([str(p) for p in data_path.rglob('*public*/**/*.tar')])
+    private_tars = sorted([str(p) for p in data_path.rglob('*private*/**/*.tar')])
+    if public_tars:
+        data_files["validation"] = public_tars
+    if private_tars:
+        data_files["test"] = private_tars
+
+    loading_scripts = sorted(data_path.rglob('*.py'))
+    loading_script = str(loading_scripts[0]) if loading_scripts else str(data_path)
+
+    dataset = load_dataset(
+        loading_script, data_files=data_files,
+        trust_remote_code=True, writer_batch_size=100,
+    )
+    print(f"Dataset: {dataset}", flush=True)
 
-    print("------------ Loading dataset ------------")
-    dataset = load_competition_dataset(params)
+    # Try to load sklearn model
+    sklearn_model = None
+    try:
+        install_if_missing('scikit-learn')
+        install_if_missing('networkx')
+        import pickle
+        model_path = Path(__file__).parent / 'sklearn_edge.pkl'
+        print(f"Looking for sklearn model at: {model_path} (exists={model_path.exists()})", flush=True)
+        if model_path.exists():
+            with open(model_path, 'rb') as f:
+                sklearn_model = pickle.load(f)
+            print("Loaded sklearn edge model OK", flush=True)
+        else:
+            print("sklearn model not found — using heuristic + edge validation only", flush=True)
+    except Exception as e:
+        print(f"sklearn failed: {e} — using heuristic + edge validation only", flush=True)
 
-    print(dataset, flush=True)
-    
-    print('------------ Now you can do your solution ---------------')
+    print("------------ Running predictions ---------------", flush=True)
     solution = []
-    for subset_name in dataset:
-        print (f"Predicitng on {subset_name}")
-        preds = Parallel(n_jobs=-1, prefer="processes")(
-            delayed(predict_wireframe_safely)(a) for a in tqdm(dataset[subset_name])
-        )
-        print ("Converting")
-        for p in preds:
-            pred_vertices, pred_edges, order_id = p
-            print (f'{order_id}: {len(pred_vertices)} verts, {len(pred_edges)} edges')
-            solution.append({
-                            'order_id': order_id, 
-                            'wf_vertices': pred_vertices.tolist(),
-                            'wf_edges': pred_edges
-                        })
-    print('------------ Saving results ---------------')
-    with open("submission.json", "w") as f:
+    for subset_name in dataset.keys():
+        print(f"Predicting {subset_name}", flush=True)
+        for i, sample in enumerate(tqdm(dataset[subset_name])):
+            res = process_sample(sample, i, sklearn_model)
+            solution.append(res)
+            if i % 50 == 0:
+                print(f"  Processed {i} samples", flush=True)
+
+    print("------------ Saving results ---------------", flush=True)
+    output_path = Path(params.get('output_path', '.'))
+
+    with open(output_path / "submission.json", 'w') as f:
         json.dump(solution, f)
+    print(f"Saved {len(solution)} predictions to submission.json", flush=True)
+
+    try:
+        import pandas as pd
+        sub = pd.DataFrame(solution, columns=["order_id", "wf_vertices", "wf_edges"])
+        sub.to_parquet(output_path / "submission.parquet")
+    except Exception:
+        pass
 
-    print("------------ Done ------------ ")
+    print("------------ Done ------------", flush=True)

sklearn_submission.py

@@ -3,6 +3,12 @@
 import numpy as np
 import cv2
 from typing import Tuple, List
+import sys
+from pathlib import Path
+
+_cur_dir = str(Path(__file__).parent.absolute())
+if _cur_dir not in sys.path:
+    sys.path.insert(0, _cur_dir)
 
 from hoho2025.example_solutions import (
     convert_entry_to_human_readable, empty_solution,