| """Depth refinement module for S23DR 2026. | |
| Learn a per-scene scale + shift correction for MoGe v2 depth maps | |
| using COLMAP sparse depth as supervision. | |
| """ | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| from typing import Tuple, Optional | |
| def ransac_scale_fit(dense_depth, sparse_depth, house_mask=None, n_iterations=100, inlier_threshold=0.5, min_samples=3): | |
| mask = (sparse_depth > 0) & (dense_depth > 0) & (sparse_depth < 50) & (dense_depth < 50) | |
| if house_mask is not None: mask = mask & house_mask | |
| X, Y = dense_depth[mask], sparse_depth[mask] | |
| if len(X) < min_samples: | |
| return (np.median(Y / X), 0.0) if len(X) > 0 else (1.0, 0.0) | |
| best_scale, best_shift, best_inliers = 1.0, 0.0, 0 | |
| rng = np.random.RandomState(42) | |
| for _ in range(n_iterations): | |
| idx = rng.choice(len(X), min(min_samples, len(X)), replace=False) | |
| A = np.column_stack([X[idx], np.ones(len(idx))]) | |
| try: params, _, _, _ = np.linalg.lstsq(A, Y[idx], rcond=None); scale, shift = params[0], params[1] | |
| except: continue | |
| if scale <= 0: continue | |
| inliers = (np.abs(scale * X + shift - Y) < inlier_threshold).sum() | |
| if inliers > best_inliers: best_inliers = inliers; best_scale = scale; best_shift = shift | |
| inlier_mask = np.abs(best_scale * X + best_shift - Y) < inlier_threshold | |
| if inlier_mask.sum() >= min_samples: | |
| A = np.column_stack([X[inlier_mask], np.ones(inlier_mask.sum())]) | |
| try: params, _, _, _ = np.linalg.lstsq(A, Y[inlier_mask], rcond=None); best_scale, best_shift = params[0], params[1] | |
| except: pass | |
| return best_scale, best_shift | |
| def fit_depth_to_colmap(dense_depth, colmap_rec, img_id, ade_seg=None, method='ransac'): | |
| from hoho2025.example_solutions import get_sparse_depth, get_house_mask | |
| sparse_depth, found, col_img, proj_pts = get_sparse_depth(colmap_rec, img_id, dense_depth) | |
| if not found: return dense_depth, False | |
| house_mask = get_house_mask(ade_seg) if ade_seg is not None else None | |
| if method == 'ransac': | |
| scale, shift = ransac_scale_fit(dense_depth, sparse_depth, house_mask) | |
| fitted = scale * dense_depth + shift | |
| else: | |
| from hoho2025.example_solutions import fit_scale_robust_median | |
| _, fitted = fit_scale_robust_median(dense_depth, sparse_depth, house_mask) | |
| return np.clip(fitted, 0, None), True | |
| class DepthRefinerNet(nn.Module): | |
| def __init__(self, in_channels=5): | |
| super().__init__() | |
| self.net = nn.Sequential(nn.Conv2d(in_channels, 32, 3, padding=1), nn.ReLU(), nn.Conv2d(32, 32, 3, padding=1), nn.ReLU(), nn.Conv2d(32, 16, 3, padding=1), nn.ReLU(), nn.Conv2d(16, 2, 1)) | |
| def forward(self, x, dense_depth): | |
| correction = self.net(x) | |
| scale = torch.sigmoid(correction[:, 0:1]) * 4.0 | |
| shift = correction[:, 1:2] * 2.0 | |
| return torch.clamp(scale * dense_depth + shift, min=0) | |