kingkenche/dinov2_retrievalsv2

DINOv2 3D Object Retrieval Pipeline (colour-aware)

Extract DINOv2 embeddings from colourful multi-view renders of 3D objects, retrieve the best-matching ShapeNetSem objects for each voxel chunk, and visualise the results. Shape similarity is always the primary signal; colour only breaks ties within DINOv2's own shortlist.

---

File overview

dinov2_retrieval/
├── config.py            ← all configuration dataclasses
├── utils.py              ← mesh loading (.ply/.obj/.off), path scanning, embedding I/O
├── render_views.py       ← multi-view colour renderer (pyrender EGL, auto-fallback to Open3D)
├── embed_dinov2.py        ← DINOv2 embedding extractor
├── color_features.py      ← colour histogram descriptor + placeholder detection (NEW)
├── similarity.py          ← shape similarity, shape+colour reranked top-K retrieval
├── visualize.py           ← 3 visualisations (grid, heatmap, bar charts)
├── pipeline.py             ← end-to-end CLI pipeline
├── test_pipeline.py        ← smoke tests (no real data needed)
└── requirements.txt

Chunk meshes come from the companion Blender voxelisation script and may be .ply (preferred — carries per-vertex colour sampled from the original textured mesh) or .obj (textured or untextured), mixed freely in the same input directory. ShapeNetSem models are .obj with UV-textured or flat-material colour, as usual.

---

Setup

# 1. Create a virtual environment
python -m venv .venv && source .venv/bin/activate

# 2. Install dependencies
pip install -r requirements.txt

# 3. For GPU headless rendering (Linux + NVIDIA):
pip install pyopengl pyopengl-accelerate

# 4. Open3D is required as the automatic rendering fallback (used if
#    EGL/pyrender isn't available — see "Rendering backend" below):
pip install open3d

---

Input directory structure

Chunks (mesh files — .ply and/or .obj, mixed OK)

data/chunks/
├── chunk_0000.ply     ← per-vertex colour (from the Blender script)
├── chunk_0001.ply
├── chunk_0002.obj     ← also accepted
└── ...

ShapeNetSem (either layout works)

# Layout A — flat
data/ShapeNetSem/
└── <model_id>/
    ├── model.obj
    └── model.mtl

# Layout B — hierarchical (standard ShapeNetSem)
data/ShapeNetSem/
└── <synset_id>/
    └── <model_id>/
        ├── model_normalized.obj
        └── model_normalized.mtl

Pre-rendered images (optional — skip rendering step)

data/shapenet_renders/
└── <model_id>/
    ├── view_00.png
    ├── view_01.png
    └── ...   (any number of views)

---

Run

Basic (render from meshes, colour tiebreak ON by default)

python pipeline.py \
  --chunks_dir   data/chunks \
  --shapenet_dir data/ShapeNetSem \
  --output_dir   output \
  --top_k        5

Shape-only retrieval (disable the colour tiebreak)

python pipeline.py \
  --chunks_dir   data/chunks \
  --shapenet_dir data/ShapeNetSem \
  --no_colour_rerank

Tune how much colour is allowed to influence ranking

python pipeline.py \
  --chunks_dir      data/chunks \
  --shapenet_dir    data/ShapeNetSem \
  --shape_pool      30      # colour can only reorder within DINOv2's top 30
  --colour_weight   0.2     # 0 = shape only, 1 = colour only, within the pool

With pre-rendered ShapeNetSem images

python pipeline.py \
  --chunks_dir            data/chunks_tiger_textures_r20_v7 \
  --shapenet_dir          data/ShapeNetSem \
  --shapenet_renders_dir  dinov2_retrievalsv2/output_full/rendered/shapenet \
  --output_dir            output_texture \
  --top_k                 5

Best geometric accuracy (larger model + patch features)

python pipeline.py \
  --chunks_dir   data/chunks \
  --shapenet_dir data/ShapeNetSem \
  --model_name   facebook/dinov2-large \
  --feature_type cls_patch \
  --n_views      12 \
  --top_k        5

Explicit rendering backend (skip auto-detection)

python pipeline.py \
  --chunks_dir      data/chunks \
  --shapenet_dir    data/ShapeNetSem \
  --opengl_platform osmesa

---

Rendering backend

--opengl_platform defaults to auto:

1. Tries pyrender with EGL (GPU headless) once.
2. If that fails, falls back to the Open3D offscreen renderer.

We do not attempt a same-process pyrender EGL → OSMesa retry: PyOpenGL picks its backend from the PYOPENGL_PLATFORM environment variable at the time the OpenGL package is first imported, and that choice is cached for the rest of the process — retrying with a different env var value after a failed EGL attempt has no effect. Open3D is a separate library, so it's the only backend that reliably works as an in-process fallback. If you know your machine only has OSMesa (no GPU/EGL at all), set --opengl_platform osmesa explicitly instead of relying on auto.

Every mesh is rendered with a single lighting pass (moderate ambient + restrained directional light) that serves both DINOv2 embedding extraction and the visualisation grids — tuned to avoid over/under-saturating true surface colour without needing a separate unlit render.

---

How colour affects retrieval

DINOv2 is trained with strong colour-jitter augmentation, which makes it fairly colour-invariant by design — it mostly attends to shape and structure. Two shape-similar but differently-coloured ShapeNetSem objects can therefore score nearly identically under DINOv2 alone.

This pipeline keeps DINOv2 shape similarity as the primary ranking signal, and adds an explicit RGB-histogram colour similarity (color_features.py) that can only re-order candidates DINOv2 already shortlisted — it never pulls in a shape-dissimilar object just because its colour matches:

for each chunk:
    1. pool = top `shape_pool` objects by DINOv2 cosine similarity alone
    2. for each object in pool:
           final_score = (1 - colour_weight) * shape_sim + colour_weight * colour_sim
           (colour_weight ignored, i.e. shape-only, for objects whose
            render looks like a colourless placeholder — missing texture,
            or an unresolved-appearance chunk)
    3. return the top-K of `pool` by final_score

Defaults: shape_pool=20, colour_weight=0.15 — shape decides who's in the running, colour only nudges the final order. Set --no_colour_rerank to fall back to pure DINOv2 ranking.

---

Output structure

output/
├── rendered/
│   ├── chunks/
│   │   └── <chunk_name>/
│   │       ├── view_00.png  ...  view_11.png   ← colourful (vertex colour / texture)
│   └── shapenet/
│       └── <model_id>/
│           └── view_00.png  ...  view_11.png   ← colourful
│
├── embeddings/
│   ├── chunk_embeddings.npy    ← (N_chunks, D) float32
│   ├── chunk_names.json
│   ├── shapenet_embeddings.npy ← (N_objects, D) float32
│   └── shapenet_names.json
│
├── results/
│   ├── similarity_matrix.npy   ← (N_chunks, N_objects) float32, DINOv2 shape similarity
│   ├── chunk_names.json
│   ├── object_names.json
│   └── top_k_results.json      ← per-chunk top-K matches, final (shape+colour) scores,
│                                  and whether colour re-ranking was applied
│
└── visualizations/
    ├── topk_grids/
    │   └── topk_grid_<chunk_name>.png   ← query + K matches side by side
    ├── similarity_heatmap.png            ← DINOv2 shape similarity, colour-coded
    └── bar_charts/
        ├── barchart_<chunk_name>.png
        └── barcharts_all_chunks.png      ← combined multi-panel

---

Run the smoke tests (no real data needed)

python test_pipeline.py

Tests:

• test_utils → mesh loading (.obj + .ply), vertex-colour preservation, normalisation, viewpoints
• test_similarity → cosine sim matrix + shape-only top-K
• test_colour_features_and_reranking → colour histograms, placeholder detection, shape+colour reranked retrieval (checks colour never pulls a match outside the shape-similarity pool)
• test_visualize_synthetic → all 3 visualisations with fake colourful images
• test_full_pipeline_dry_run → end-to-end with synthetic renders, including colour re-ranking, save/reload
• test_embed_dinov2_synthetic → DINOv2 forward pass (downloads model)

---

DINOv2 model options

Model	Embedding dim	Speed	Geometry quality
facebook/dinov2-small	384	Fast	Good
facebook/dinov2-base	768	Medium	Better ✓ default
facebook/dinov2-large	1024	Slow	Best
facebook/dinov2-giant	1536	Slowest	Best

Feature type:

• cls → CLS token only (fast, recommended)
• patch_mean → mean of all patch tokens (richer local geometry)
• cls_patch → concat both (best accuracy, 2× dimension)

---

Key flags

Flag	Default	Notes
--model_name	facebook/dinov2-base	Larger = better geometry
--feature_type	cls	cls / patch_mean / cls_patch
--n_views	12	More views = more stable embedding
--top_k	5	Matches returned per chunk
--opengl_platform	auto	auto (EGL→Open3D) / egl / osmesa
--no_colour_rerank	off (colour ON by default)	Disable colour tiebreak, pure DINOv2 ranking
--shape_pool	20	Candidate pool colour is allowed to re-order
--colour_weight	0.15	Colour blend weight within the pool
--colour_hist_bins	8	Histogram bins per RGB channel
--placeholder_std_threshold	6.0	Below this render std-dev = colourless, shape-only
--force_recompute	False	Rerun even if .npy cache exists
--max_shapenet_objects	None	Cap objects for quick testing
--no_vis	False	Skip visualisation step