luboslenco
1192 lines
38 KiB
C
1192 lines
38 KiB
C
/*
|
|
* iris_depth.c - Depth Anything 3 (DA3MONO-LARGE) monocular depth estimation
|
|
*
|
|
* Pure-C implementation of the Depth Anything 3 "mono" depth estimator
|
|
* (DA3MONO-LARGE). The network is a DINOv2 ViT-L/14 backbone followed by a
|
|
* DPT-style head, in the original (non-HuggingFace) checkpoint layout used by
|
|
* the Depth-Anything-3 release:
|
|
*
|
|
* backbone (DINOv2 ViT-L/14, "model.backbone.pretrained.*"):
|
|
* - patch embed: 14x14 stride-14 conv, 3 -> 1024, on a 518x518 input
|
|
* (37x37 = 1369 patches), prepend cls token, add position embeddings
|
|
* - 24 transformer blocks (pre-norm, fused-QKV MHSA + GELU MLP, LayerScale)
|
|
* - the outputs of blocks {4,11,17,23} are taken as multi-scale features,
|
|
* each passed through the final backbone LayerNorm
|
|
*
|
|
* head (DPT, "model.head.*"):
|
|
* - projects: drop cls, reshape to 37x37, 1x1 project to {256,512,1024,1024}
|
|
* channels, then resize_layers (transpose-conv x4 / x2, identity, conv
|
|
* stride2) to scales {148,74,37,19}
|
|
* - scratch.layer{1..4}_rn: 3x3 convs map each to 256 channels
|
|
* - scratch.refinenet{4..1}: feature-fusion (RefineNet) merges coarse->fine
|
|
* to 296x296
|
|
* - scratch.output_conv1: 3x3 conv 256->128, bilinear upsample to 518x518,
|
|
* output_conv2: 3x3 conv 128->32 + ReLU, 1x1 conv 32->1 + ReLU
|
|
*
|
|
* (The model also carries a parallel "sky_output_conv2" branch predicting a sky
|
|
* mask; it is not needed for the depth map and is ignored here.)
|
|
*
|
|
* The output is resized back to the input resolution and min/max normalized to
|
|
* a single-channel grayscale image (brighter = nearer).
|
|
*
|
|
* All weights are F32 in the memory-mapped safetensors file (zero-copy). The
|
|
* transformer linear/attention matrix multiplies are dispatched to the Metal
|
|
* or Vulkan GEMM backend when available; convolutions (small spatially) run on
|
|
* the CPU.
|
|
*
|
|
* Note: the official image processor resizes with bicubic resampling; here we
|
|
* use bilinear, which yields a visually equivalent relative-depth map.
|
|
*/
|
|
|
|
#include "iris_depth.h"
|
|
#include "iris_safetensors.h"
|
|
#include <math.h>
|
|
#include <stdarg.h>
|
|
#include <stdio.h>
|
|
#include <stdlib.h>
|
|
#include <string.h>
|
|
|
|
#if defined(USE_METAL)
|
|
#include "iris_metal.h"
|
|
#elif defined(USE_VULKAN)
|
|
#include "iris_vulkan.h"
|
|
#endif
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* Architecture constants (Depth Anything 3 mono - Large / ViT-L/14) */
|
|
/* ------------------------------------------------------------------ */
|
|
#define DA_PATCH 14
|
|
#define DA_IMG 518 /* model input is DA_IMG x DA_IMG */
|
|
#define DA_GRID (DA_IMG / DA_PATCH) /* 37 patches per side */
|
|
#define DA_NPATCH (DA_GRID * DA_GRID) /* 1369 */
|
|
#define DA_TOKENS (DA_NPATCH + 1) /* 1370 (+ cls) */
|
|
#define DA_HIDDEN 1024
|
|
#define DA_LAYERS 24
|
|
#define DA_HEADS 16
|
|
#define DA_HEAD_DIM (DA_HIDDEN / DA_HEADS) /* 64 */
|
|
#define DA_MLP 4096
|
|
#define DA_LN_EPS 1e-6f
|
|
#define DA_FUSION_CH 256
|
|
#define DA_HEAD_CH (DA_FUSION_CH / 2) /* head conv1 output (128) */
|
|
#define DA_HEAD_HIDDEN 32 /* head conv2 output */
|
|
|
|
/* Tensor-name roots for the DA3 checkpoint layout. */
|
|
#define BB "model.backbone.pretrained"
|
|
#define HEAD "model.head"
|
|
|
|
/* Backbone blocks whose outputs feed the head (0-indexed). */
|
|
static const int DA_OUT_LAYERS[4] = {4, 11, 17, 23};
|
|
/* Per-feature channel widths after the 1x1 projects. */
|
|
static const int DA_PROJ_CH[4] = {256, 512, 1024, 1024};
|
|
|
|
/* ImageNet normalization (matches the model's image processor). */
|
|
static const float DA_MEAN[3] = {0.485f, 0.456f, 0.406f};
|
|
static const float DA_STD[3] = {0.229f, 0.224f, 0.225f};
|
|
|
|
struct iris_depth {
|
|
safetensors_file_t *sf;
|
|
int tileable; /* remove the perspective tilt for seamless top-down maps */
|
|
};
|
|
|
|
/* ================================================================== */
|
|
/* Weight access (zero-copy; all tensors are F32 in the mmap'd file) */
|
|
/* ================================================================== */
|
|
|
|
static const float *da_tensor(iris_depth_t *m, const char *name) {
|
|
const safetensor_t *t = safetensors_find(m->sf, name);
|
|
if (!t) {
|
|
fprintf(stderr, "DepthAnything3: missing tensor '%s'\n", name);
|
|
return NULL;
|
|
}
|
|
if (t->dtype != DTYPE_F32) {
|
|
fprintf(stderr, "DepthAnything3: tensor '%s' is not F32\n", name);
|
|
return NULL;
|
|
}
|
|
return (const float *)safetensors_data(m->sf, t);
|
|
}
|
|
|
|
/* printf-style tensor lookup. */
|
|
static const float *da_tensorf(iris_depth_t *m, const char *fmt, ...) {
|
|
char name[256];
|
|
va_list ap;
|
|
va_start(ap, fmt);
|
|
vsnprintf(name, sizeof(name), fmt, ap);
|
|
va_end(ap);
|
|
return da_tensor(m, name);
|
|
}
|
|
|
|
/* ================================================================== */
|
|
/* GEMM helpers (Vulkan-accelerated when available) */
|
|
/* ================================================================== */
|
|
|
|
/* C[M,N] = scale * A[M,K] @ B[N,K]^T (row-wise dot products).
|
|
* cache_b: weight matrix is immutable -> cache it on the GPU by pointer. */
|
|
static void gemm_nt(int M, int N, int K, const float *A, int lda, const float *B, int ldb, float *C, int ldc, float scale, int cache_b) {
|
|
#if defined(USE_METAL)
|
|
if (iris_metal_available()) {
|
|
if (cache_b)
|
|
iris_metal_sgemm_cached(0, 1, M, N, K, scale, A, lda, B, ldb, 0.0f, C, ldc);
|
|
else
|
|
iris_metal_sgemm(0, 1, M, N, K, scale, A, lda, B, ldb, 0.0f, C, ldc);
|
|
return;
|
|
}
|
|
#elif defined(USE_VULKAN)
|
|
if (iris_vulkan_available()) {
|
|
if (cache_b)
|
|
iris_vulkan_sgemm_cached(0, 1, M, N, K, scale, A, lda, B, ldb, 0.0f, C, ldc);
|
|
else
|
|
iris_vulkan_sgemm(0, 1, M, N, K, scale, A, lda, B, ldb, 0.0f, C, ldc);
|
|
return;
|
|
}
|
|
#endif
|
|
(void)cache_b;
|
|
for (int i = 0; i < M; i++) {
|
|
const float *a = A + (size_t)i * lda;
|
|
float *c = C + (size_t)i * ldc;
|
|
for (int j = 0; j < N; j++) {
|
|
const float *b = B + (size_t)j * ldb;
|
|
float s = 0.0f;
|
|
for (int k = 0; k < K; k++)
|
|
s += a[k] * b[k];
|
|
c[j] = s * scale;
|
|
}
|
|
}
|
|
}
|
|
|
|
/* C[M,N] = A[M,K] @ B[K,N] (B not transposed). */
|
|
static void gemm_nn(int M, int N, int K, const float *A, int lda, const float *B, int ldb, float *C, int ldc) {
|
|
#if defined(USE_METAL)
|
|
if (iris_metal_available()) {
|
|
iris_metal_sgemm(0, 0, M, N, K, 1.0f, A, lda, B, ldb, 0.0f, C, ldc);
|
|
return;
|
|
}
|
|
#elif defined(USE_VULKAN)
|
|
if (iris_vulkan_available()) {
|
|
iris_vulkan_sgemm(0, 0, M, N, K, 1.0f, A, lda, B, ldb, 0.0f, C, ldc);
|
|
return;
|
|
}
|
|
#endif
|
|
for (int i = 0; i < M; i++) {
|
|
float *c = C + (size_t)i * ldc;
|
|
for (int j = 0; j < N; j++)
|
|
c[j] = 0.0f;
|
|
const float *a = A + (size_t)i * lda;
|
|
for (int k = 0; k < K; k++) {
|
|
float av = a[k];
|
|
const float *b = B + (size_t)k * ldb;
|
|
for (int j = 0; j < N; j++)
|
|
c[j] += av * b[j];
|
|
}
|
|
}
|
|
}
|
|
|
|
/* out[M,N] = x[M,K] @ W[N,K]^T + bias[N] (bias may be NULL). */
|
|
static void linear(const float *x, int M, int K, const float *W, const float *bias, int N, float *out) {
|
|
gemm_nt(M, N, K, x, K, W, K, out, N, 1.0f, 1);
|
|
if (bias) {
|
|
for (int i = 0; i < M; i++) {
|
|
float *o = out + (size_t)i * N;
|
|
for (int j = 0; j < N; j++)
|
|
o[j] += bias[j];
|
|
}
|
|
}
|
|
}
|
|
|
|
/* ================================================================== */
|
|
/* Elementwise / norm primitives */
|
|
/* ================================================================== */
|
|
|
|
/* Row-wise LayerNorm over `dim` (per token), affine weight+bias. */
|
|
static void layer_norm(const float *in, int rows, int dim, const float *gamma, const float *beta, float *out) {
|
|
for (int r = 0; r < rows; r++) {
|
|
const float *x = in + (size_t)r * dim;
|
|
float *o = out + (size_t)r * dim;
|
|
float mean = 0.0f;
|
|
for (int i = 0; i < dim; i++)
|
|
mean += x[i];
|
|
mean /= dim;
|
|
float var = 0.0f;
|
|
for (int i = 0; i < dim; i++) {
|
|
float d = x[i] - mean;
|
|
var += d * d;
|
|
}
|
|
var /= dim;
|
|
float inv = 1.0f / sqrtf(var + DA_LN_EPS);
|
|
for (int i = 0; i < dim; i++)
|
|
o[i] = (x[i] - mean) * inv * gamma[i] + beta[i];
|
|
}
|
|
}
|
|
|
|
/* Exact GELU (erf form), in place. */
|
|
static void gelu(float *buf, size_t n) {
|
|
const float inv_sqrt2 = 0.70710678118654752440f;
|
|
for (size_t i = 0; i < n; i++)
|
|
buf[i] = 0.5f * buf[i] * (1.0f + erff(buf[i] * inv_sqrt2));
|
|
}
|
|
|
|
static void relu_inplace(float *buf, size_t n) {
|
|
for (size_t i = 0; i < n; i++)
|
|
if (buf[i] < 0.0f)
|
|
buf[i] = 0.0f;
|
|
}
|
|
|
|
/* ================================================================== */
|
|
/* Convolutions (CPU) */
|
|
/* ================================================================== */
|
|
|
|
/* General 2D convolution, planar CHW.
|
|
* in: [Cin, H, W]
|
|
* weight: [Cout, Cin, kH, kW]
|
|
* bias: [Cout] or NULL
|
|
* out: [Cout, OH, OW] with OH=(H+2*pad-kH)/stride+1
|
|
*/
|
|
static void conv2d(const float *in, int Cin, int H, int W, const float *weight, const float *bias, int Cout, int kH, int kW, int stride, int pad, float *out) {
|
|
int OH = (H + 2 * pad - kH) / stride + 1;
|
|
int OW = (W + 2 * pad - kW) / stride + 1;
|
|
const size_t in_plane = (size_t)H * W;
|
|
const size_t out_plane = (size_t)OH * OW;
|
|
|
|
for (int oc = 0; oc < Cout; oc++) {
|
|
float *op = out + (size_t)oc * out_plane;
|
|
float b = bias ? bias[oc] : 0.0f;
|
|
for (size_t i = 0; i < out_plane; i++)
|
|
op[i] = b;
|
|
|
|
for (int ic = 0; ic < Cin; ic++) {
|
|
const float *ip = in + (size_t)ic * in_plane;
|
|
const float *wk = weight + (((size_t)oc * Cin + ic) * kH) * kW;
|
|
for (int ky = 0; ky < kH; ky++) {
|
|
for (int kx = 0; kx < kW; kx++) {
|
|
float wv = wk[ky * kW + kx];
|
|
if (wv == 0.0f)
|
|
continue;
|
|
for (int oy = 0; oy < OH; oy++) {
|
|
int sy = oy * stride + ky - pad;
|
|
if (sy < 0 || sy >= H)
|
|
continue;
|
|
float *orow = op + (size_t)oy * OW;
|
|
const float *irow = ip + (size_t)sy * W;
|
|
for (int ox = 0; ox < OW; ox++) {
|
|
int sx = ox * stride + kx - pad;
|
|
if (sx < 0 || sx >= W)
|
|
continue;
|
|
orow[ox] += wv * irow[sx];
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/* Transposed convolution with stride == kernel and no padding (the only
|
|
* configuration the DPT resize stage uses). Upsamples HxW -> (H*k)x(W*k).
|
|
* weight: [Cin, Cout, k, k] bias: [Cout] or NULL
|
|
* out[oc, oh, ow] = bias[oc] + sum_ic in[ic, oh/k, ow/k] * W[ic,oc, oh%k, ow%k]
|
|
*/
|
|
static void conv_transpose_block(const float *in, int Cin, int H, int W, const float *weight, const float *bias, int Cout, int k, float *out) {
|
|
int OH = H * k, OW = W * k;
|
|
const size_t in_plane = (size_t)H * W;
|
|
const size_t out_plane = (size_t)OH * OW;
|
|
|
|
for (int oc = 0; oc < Cout; oc++) {
|
|
float *op = out + (size_t)oc * out_plane;
|
|
float b = bias ? bias[oc] : 0.0f;
|
|
for (int oy = 0; oy < OH; oy++) {
|
|
int iy = oy / k, ky = oy % k;
|
|
for (int ox = 0; ox < OW; ox++) {
|
|
int ix = ox / k, kx = ox % k;
|
|
float s = b;
|
|
for (int ic = 0; ic < Cin; ic++) {
|
|
float iv = in[(size_t)ic * in_plane + (size_t)iy * W + ix];
|
|
float wv = weight[(((size_t)ic * Cout + oc) * k + ky) * k + kx];
|
|
s += iv * wv;
|
|
}
|
|
op[(size_t)oy * OW + ox] = s;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/* ================================================================== */
|
|
/* Bilinear resize (planar CHW) */
|
|
/* ================================================================== */
|
|
|
|
static void bilinear_chw(const float *in, int C, int H, int W, int OH, int OW, int align_corners, float *out) {
|
|
const size_t in_plane = (size_t)H * W;
|
|
const size_t out_plane = (size_t)OH * OW;
|
|
float sy = 0.0f, sx = 0.0f;
|
|
if (align_corners) {
|
|
sy = (OH > 1) ? (float)(H - 1) / (OH - 1) : 0.0f;
|
|
sx = (OW > 1) ? (float)(W - 1) / (OW - 1) : 0.0f;
|
|
}
|
|
else {
|
|
sy = (float)H / OH;
|
|
sx = (float)W / OW;
|
|
}
|
|
for (int oy = 0; oy < OH; oy++) {
|
|
float fy = align_corners ? oy * sy : (oy + 0.5f) * sy - 0.5f;
|
|
if (fy < 0.0f)
|
|
fy = 0.0f;
|
|
int y0 = (int)fy;
|
|
if (y0 > H - 1)
|
|
y0 = H - 1;
|
|
int y1 = (y0 + 1 < H) ? y0 + 1 : y0;
|
|
float wy = fy - y0;
|
|
for (int ox = 0; ox < OW; ox++) {
|
|
float fx = align_corners ? ox * sx : (ox + 0.5f) * sx - 0.5f;
|
|
if (fx < 0.0f)
|
|
fx = 0.0f;
|
|
int x0 = (int)fx;
|
|
if (x0 > W - 1)
|
|
x0 = W - 1;
|
|
int x1 = (x0 + 1 < W) ? x0 + 1 : x0;
|
|
float wx = fx - x0;
|
|
for (int c = 0; c < C; c++) {
|
|
const float *ip = in + (size_t)c * in_plane;
|
|
float v00 = ip[(size_t)y0 * W + x0];
|
|
float v01 = ip[(size_t)y0 * W + x1];
|
|
float v10 = ip[(size_t)y1 * W + x0];
|
|
float v11 = ip[(size_t)y1 * W + x1];
|
|
float top = v00 + (v01 - v00) * wx;
|
|
float bot = v10 + (v11 - v10) * wx;
|
|
out[(size_t)c * out_plane + (size_t)oy * OW + ox] = top + (bot - top) * wy;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/* ================================================================== */
|
|
/* Backbone (DINOv2 ViT-L/14) */
|
|
/* ================================================================== */
|
|
|
|
/* Multi-head self-attention for one transformer block, fused QKV.
|
|
* x is [DA_TOKENS, DA_HIDDEN]; result written to `out` (same shape). */
|
|
static int da_attention(iris_depth_t *m, int layer, const float *x, float *out, float *scratch) {
|
|
const int T = DA_TOKENS, Hd = DA_HIDDEN;
|
|
|
|
/* Fused QKV: weight [3*Hd, Hd], bias [3*Hd]; output projection. */
|
|
const float *wqkv = da_tensorf(m, BB ".blocks.%d.attn.qkv.weight", layer);
|
|
const float *bqkv = da_tensorf(m, BB ".blocks.%d.attn.qkv.bias", layer);
|
|
const float *wo = da_tensorf(m, BB ".blocks.%d.attn.proj.weight", layer);
|
|
const float *bo = da_tensorf(m, BB ".blocks.%d.attn.proj.bias", layer);
|
|
if (!wqkv || !bqkv || !wo || !bo)
|
|
return -1;
|
|
|
|
/* scratch layout: QKV (T*3*Hd), attn output A (T*Hd), scores S (T*T). */
|
|
float *QKV = scratch;
|
|
float *A = QKV + (size_t)T * 3 * Hd; /* attention output (concat heads) */
|
|
float *S = A + (size_t)T * Hd; /* per-head scores [T,T] */
|
|
|
|
/* QKV[T,3*Hd] = x @ Wqkv^T + bqkv. Within a row: [q(Hd) | k(Hd) | v(Hd)],
|
|
* each Hd block laid out as heads*head_dim. */
|
|
linear(x, T, Hd, wqkv, bqkv, 3 * Hd, QKV);
|
|
const float *Qbase = QKV;
|
|
const float *Kbase = QKV + Hd;
|
|
const float *Vbase = QKV + 2 * Hd;
|
|
|
|
const float scale = 1.0f / sqrtf((float)DA_HEAD_DIM);
|
|
for (int h = 0; h < DA_HEADS; h++) {
|
|
int off = h * DA_HEAD_DIM;
|
|
/* S[T,T] = scale * Qh @ Kh^T (strided per-head views, stride 3*Hd). */
|
|
gemm_nt(T, T, DA_HEAD_DIM, Qbase + off, 3 * Hd, Kbase + off, 3 * Hd, S, T, scale, 0);
|
|
/* row softmax */
|
|
for (int i = 0; i < T; i++) {
|
|
float *row = S + (size_t)i * T;
|
|
float mx = row[0];
|
|
for (int j = 1; j < T; j++)
|
|
if (row[j] > mx)
|
|
mx = row[j];
|
|
float sum = 0.0f;
|
|
for (int j = 0; j < T; j++) {
|
|
row[j] = expf(row[j] - mx);
|
|
sum += row[j];
|
|
}
|
|
float inv = 1.0f / sum;
|
|
for (int j = 0; j < T; j++)
|
|
row[j] *= inv;
|
|
}
|
|
/* Ah[T,head_dim] = S @ Vh -> write into A at head offset (stride Hd). */
|
|
gemm_nn(T, DA_HEAD_DIM, T, S, T, Vbase + off, 3 * Hd, A + off, Hd);
|
|
}
|
|
|
|
/* output projection: out = A @ Wo^T + bo */
|
|
linear(A, T, Hd, wo, bo, Hd, out);
|
|
return 0;
|
|
}
|
|
|
|
/* Run the backbone, capturing the four feature maps (after the final
|
|
* LayerNorm) into feats[k] = [DA_TOKENS, DA_HIDDEN]. Returns 0 on success. */
|
|
static int da_backbone(iris_depth_t *m, const float *pixels /* [3,518,518] */, float *feats[4]) {
|
|
const int T = DA_TOKENS, Hd = DA_HIDDEN;
|
|
|
|
/* ---- Patch embedding (14x14 stride-14 conv, 3 -> 1024) ---- */
|
|
const float *pw = da_tensor(m, BB ".patch_embed.proj.weight");
|
|
const float *pb = da_tensor(m, BB ".patch_embed.proj.bias");
|
|
const float *cls = da_tensor(m, BB ".cls_token");
|
|
const float *pos = da_tensor(m, BB ".pos_embed");
|
|
if (!pw || !pb || !cls || !pos)
|
|
return -1;
|
|
|
|
float *x = malloc((size_t)T * Hd * sizeof(float));
|
|
if (!x)
|
|
return -1;
|
|
|
|
/* token 0 = cls */
|
|
memcpy(x, cls, Hd * sizeof(float));
|
|
/* tokens 1.. = patches (row-major), conv with stride=kernel, no overlap */
|
|
for (int r = 0; r < DA_GRID; r++) {
|
|
for (int c = 0; c < DA_GRID; c++) {
|
|
float *tok = x + (size_t)(1 + r * DA_GRID + c) * Hd;
|
|
for (int oc = 0; oc < Hd; oc++) {
|
|
float s = pb[oc];
|
|
const float *wk = pw + (size_t)oc * 3 * DA_PATCH * DA_PATCH;
|
|
for (int ic = 0; ic < 3; ic++) {
|
|
const float *ip = pixels + (size_t)ic * DA_IMG * DA_IMG;
|
|
const float *wic = wk + (size_t)ic * DA_PATCH * DA_PATCH;
|
|
for (int ky = 0; ky < DA_PATCH; ky++) {
|
|
const float *irow = ip + (size_t)(r * DA_PATCH + ky) * DA_IMG + c * DA_PATCH;
|
|
const float *wrow = wic + ky * DA_PATCH;
|
|
for (int kx = 0; kx < DA_PATCH; kx++)
|
|
s += irow[kx] * wrow[kx];
|
|
}
|
|
}
|
|
tok[oc] = s;
|
|
}
|
|
}
|
|
}
|
|
/* add position embeddings (exact size match, no interpolation) */
|
|
for (size_t i = 0; i < (size_t)T * Hd; i++)
|
|
x[i] += pos[i];
|
|
|
|
/* ---- work buffers ---- */
|
|
float *norm = malloc((size_t)T * Hd * sizeof(float));
|
|
float *tmp = malloc((size_t)T * Hd * sizeof(float));
|
|
float *mlp = malloc((size_t)T * DA_MLP * sizeof(float));
|
|
/* attention scratch: QKV (3*T*Hd) + A (T*Hd) + scores (T*T) */
|
|
float *attn_scratch = malloc(((size_t)4 * T * Hd + (size_t)T * T) * sizeof(float));
|
|
if (!norm || !tmp || !mlp || !attn_scratch) {
|
|
free(x);
|
|
free(norm);
|
|
free(tmp);
|
|
free(mlp);
|
|
free(attn_scratch);
|
|
return -1;
|
|
}
|
|
|
|
int next_capture = 0;
|
|
int ok = 1;
|
|
|
|
for (int l = 0; l < DA_LAYERS && ok; l++) {
|
|
const float *n1w = da_tensorf(m, BB ".blocks.%d.norm1.weight", l);
|
|
const float *n1b = da_tensorf(m, BB ".blocks.%d.norm1.bias", l);
|
|
const float *n2w = da_tensorf(m, BB ".blocks.%d.norm2.weight", l);
|
|
const float *n2b = da_tensorf(m, BB ".blocks.%d.norm2.bias", l);
|
|
const float *ls1 = da_tensorf(m, BB ".blocks.%d.ls1.gamma", l);
|
|
const float *ls2 = da_tensorf(m, BB ".blocks.%d.ls2.gamma", l);
|
|
const float *f1w = da_tensorf(m, BB ".blocks.%d.mlp.fc1.weight", l);
|
|
const float *f1b = da_tensorf(m, BB ".blocks.%d.mlp.fc1.bias", l);
|
|
const float *f2w = da_tensorf(m, BB ".blocks.%d.mlp.fc2.weight", l);
|
|
const float *f2b = da_tensorf(m, BB ".blocks.%d.mlp.fc2.bias", l);
|
|
if (!n1w || !n1b || !n2w || !n2b || !ls1 || !ls2 || !f1w || !f1b || !f2w || !f2b) {
|
|
ok = 0;
|
|
break;
|
|
}
|
|
|
|
/* --- attention block --- */
|
|
layer_norm(x, T, Hd, n1w, n1b, norm);
|
|
if (da_attention(m, l, norm, tmp, attn_scratch) != 0) {
|
|
ok = 0;
|
|
break;
|
|
}
|
|
/* x = x + ls1 * attn */
|
|
for (int i = 0; i < T; i++) {
|
|
float *xr = x + (size_t)i * Hd;
|
|
const float *tr = tmp + (size_t)i * Hd;
|
|
for (int j = 0; j < Hd; j++)
|
|
xr[j] += ls1[j] * tr[j];
|
|
}
|
|
|
|
/* --- MLP block --- */
|
|
layer_norm(x, T, Hd, n2w, n2b, norm);
|
|
linear(norm, T, Hd, f1w, f1b, DA_MLP, mlp);
|
|
gelu(mlp, (size_t)T * DA_MLP);
|
|
linear(mlp, T, DA_MLP, f2w, f2b, Hd, tmp);
|
|
/* x = x + ls2 * mlp_out */
|
|
for (int i = 0; i < T; i++) {
|
|
float *xr = x + (size_t)i * Hd;
|
|
const float *tr = tmp + (size_t)i * Hd;
|
|
for (int j = 0; j < Hd; j++)
|
|
xr[j] += ls2[j] * tr[j];
|
|
}
|
|
|
|
/* capture selected block outputs (post final LayerNorm) */
|
|
if (next_capture < 4 && l == DA_OUT_LAYERS[next_capture]) {
|
|
const float *lnw = da_tensor(m, BB ".norm.weight");
|
|
const float *lnb = da_tensor(m, BB ".norm.bias");
|
|
if (!lnw || !lnb) {
|
|
ok = 0;
|
|
break;
|
|
}
|
|
layer_norm(x, T, Hd, lnw, lnb, feats[next_capture]);
|
|
next_capture++;
|
|
}
|
|
}
|
|
|
|
free(x);
|
|
free(norm);
|
|
free(tmp);
|
|
free(mlp);
|
|
free(attn_scratch);
|
|
return ok ? 0 : -1;
|
|
}
|
|
|
|
/* ================================================================== */
|
|
/* DPT head */
|
|
/* ================================================================== */
|
|
|
|
/* Reshape a captured feature [DA_TOKENS, DA_HIDDEN] (cls dropped) into a
|
|
* planar CHW map [DA_HIDDEN, DA_GRID, DA_GRID]. */
|
|
static void da_tokens_to_chw(const float *feat, float *out) {
|
|
for (int r = 0; r < DA_GRID; r++)
|
|
for (int c = 0; c < DA_GRID; c++) {
|
|
const float *tok = feat + (size_t)(1 + r * DA_GRID + c) * DA_HIDDEN;
|
|
for (int ch = 0; ch < DA_HIDDEN; ch++)
|
|
out[(size_t)ch * DA_GRID * DA_GRID + (size_t)r * DA_GRID + c] = tok[ch];
|
|
}
|
|
}
|
|
|
|
/* ResidualConvUnit (in place into `buf`):
|
|
* buf = buf + conv2(relu(conv1(relu(buf)))) , all 3x3 pad-1, C channels. */
|
|
static int da_res_conv_unit(iris_depth_t *m, const char *prefix, float *buf, int C, int H, int W, float *t1, float *t2) {
|
|
const size_t plane = (size_t)C * H * W;
|
|
const float *w1 = da_tensorf(m, "%s.conv1.weight", prefix);
|
|
const float *b1 = da_tensorf(m, "%s.conv1.bias", prefix);
|
|
const float *w2 = da_tensorf(m, "%s.conv2.weight", prefix);
|
|
const float *b2 = da_tensorf(m, "%s.conv2.bias", prefix);
|
|
if (!w1 || !b1 || !w2 || !b2)
|
|
return -1;
|
|
|
|
memcpy(t1, buf, plane * sizeof(float));
|
|
relu_inplace(t1, plane);
|
|
conv2d(t1, C, H, W, w1, b1, C, 3, 3, 1, 1, t2);
|
|
relu_inplace(t2, plane);
|
|
conv2d(t2, C, H, W, w2, b2, C, 3, 3, 1, 1, t1);
|
|
for (size_t i = 0; i < plane; i++)
|
|
buf[i] += t1[i];
|
|
return 0;
|
|
}
|
|
|
|
/* One RefineNet feature-fusion block (scratch.refinenet{N}).
|
|
* in: [C, H, W] current path (consumed)
|
|
* residual: [C, H, W] or NULL (skip connection at same resolution)
|
|
* out: [C, OH, OW] caller-allocated, OH/OW are the upsample target
|
|
* Returns 0 on success. */
|
|
static int da_refinenet(iris_depth_t *m, int n, float *in, const float *residual, int C, int H, int W, int OH, int OW, float *out) {
|
|
char prefix[64];
|
|
snprintf(prefix, sizeof(prefix), HEAD ".scratch.refinenet%d", n);
|
|
const size_t plane = (size_t)C * H * W;
|
|
|
|
float *t1 = malloc(plane * sizeof(float));
|
|
float *t2 = malloc(plane * sizeof(float));
|
|
float *res = malloc(plane * sizeof(float));
|
|
if (!t1 || !t2 || !res) {
|
|
free(t1);
|
|
free(t2);
|
|
free(res);
|
|
return -1;
|
|
}
|
|
|
|
int rc = 0;
|
|
if (residual) {
|
|
char rp[96];
|
|
memcpy(res, residual, plane * sizeof(float));
|
|
snprintf(rp, sizeof(rp), "%s.resConfUnit1", prefix);
|
|
if (da_res_conv_unit(m, rp, res, C, H, W, t1, t2) != 0) {
|
|
rc = -1;
|
|
goto done;
|
|
}
|
|
for (size_t i = 0; i < plane; i++)
|
|
in[i] += res[i];
|
|
}
|
|
|
|
{
|
|
char rp[96];
|
|
snprintf(rp, sizeof(rp), "%s.resConfUnit2", prefix);
|
|
if (da_res_conv_unit(m, rp, in, C, H, W, t1, t2) != 0) {
|
|
rc = -1;
|
|
goto done;
|
|
}
|
|
}
|
|
|
|
/* upsample to (OH,OW) (align_corners=True), then 1x1 out_conv (with bias). */
|
|
{
|
|
float *up = malloc((size_t)C * OH * OW * sizeof(float));
|
|
if (!up) {
|
|
rc = -1;
|
|
goto done;
|
|
}
|
|
bilinear_chw(in, C, H, W, OH, OW, 1, up);
|
|
const float *pw = da_tensorf(m, "%s.out_conv.weight", prefix);
|
|
const float *pb = da_tensorf(m, "%s.out_conv.bias", prefix);
|
|
if (!pw || !pb) {
|
|
free(up);
|
|
rc = -1;
|
|
goto done;
|
|
}
|
|
conv2d(up, C, OH, OW, pw, pb, C, 1, 1, 1, 0, out);
|
|
free(up);
|
|
}
|
|
|
|
done:
|
|
free(t1);
|
|
free(t2);
|
|
free(res);
|
|
return rc;
|
|
}
|
|
|
|
/* Run the DPT head. `feats` are the four captured backbone features.
|
|
* Writes a [DA_IMG, DA_IMG] depth map into `depth_out`. Returns 0 on success. */
|
|
static int da_head(iris_depth_t *m, float *feats[4], float *depth_out) {
|
|
int rc = -1;
|
|
/* Per-stage "_rn" features at 256 channels. */
|
|
float *rn[4] = {0};
|
|
int rdim[4][2]; /* [k] = {H, W} of the _rn output */
|
|
float *chw = malloc((size_t)DA_HIDDEN * DA_GRID * DA_GRID * sizeof(float));
|
|
if (!chw)
|
|
return -1;
|
|
|
|
for (int k = 0; k < 4; k++) {
|
|
da_tokens_to_chw(feats[k], chw);
|
|
|
|
/* projects.k : 1x1 projection 1024 -> DA_PROJ_CH[k] */
|
|
int pc = DA_PROJ_CH[k];
|
|
const float *prw = da_tensorf(m, HEAD ".projects.%d.weight", k);
|
|
const float *prb = da_tensorf(m, HEAD ".projects.%d.bias", k);
|
|
if (!prw || !prb)
|
|
goto done;
|
|
float *proj = malloc((size_t)pc * DA_GRID * DA_GRID * sizeof(float));
|
|
if (!proj)
|
|
goto done;
|
|
conv2d(chw, DA_HIDDEN, DA_GRID, DA_GRID, prw, prb, pc, 1, 1, 1, 0, proj);
|
|
|
|
/* resize_layers.k */
|
|
float *res = NULL;
|
|
int rh = 0, rw = 0;
|
|
if (k == 0 || k == 1) {
|
|
int s = (k == 0) ? 4 : 2;
|
|
const float *rwt = da_tensorf(m, HEAD ".resize_layers.%d.weight", k);
|
|
const float *rbs = da_tensorf(m, HEAD ".resize_layers.%d.bias", k);
|
|
if (!rwt || !rbs) {
|
|
free(proj);
|
|
goto done;
|
|
}
|
|
rh = DA_GRID * s;
|
|
rw = DA_GRID * s;
|
|
res = malloc((size_t)pc * rh * rw * sizeof(float));
|
|
if (!res) {
|
|
free(proj);
|
|
goto done;
|
|
}
|
|
conv_transpose_block(proj, pc, DA_GRID, DA_GRID, rwt, rbs, pc, s, res);
|
|
}
|
|
else if (k == 2) {
|
|
/* identity (no resize_layers.2 in the checkpoint) */
|
|
rh = DA_GRID;
|
|
rw = DA_GRID;
|
|
res = proj;
|
|
proj = NULL;
|
|
}
|
|
else { /* k == 3: 3x3 stride-2 pad-1 downsample */
|
|
const float *rwt = da_tensorf(m, HEAD ".resize_layers.%d.weight", k);
|
|
const float *rbs = da_tensorf(m, HEAD ".resize_layers.%d.bias", k);
|
|
if (!rwt || !rbs) {
|
|
free(proj);
|
|
goto done;
|
|
}
|
|
rh = (DA_GRID + 2 - 3) / 2 + 1;
|
|
rw = rh;
|
|
res = malloc((size_t)pc * rh * rw * sizeof(float));
|
|
if (!res) {
|
|
free(proj);
|
|
goto done;
|
|
}
|
|
conv2d(proj, pc, DA_GRID, DA_GRID, rwt, rbs, pc, 3, 3, 2, 1, res);
|
|
}
|
|
free(proj);
|
|
|
|
/* scratch.layer{k+1}_rn : 3x3 pad-1, no bias, pc -> 256 */
|
|
const float *cw = da_tensorf(m, HEAD ".scratch.layer%d_rn.weight", k + 1);
|
|
if (!cw) {
|
|
free(res);
|
|
goto done;
|
|
}
|
|
rn[k] = malloc((size_t)DA_FUSION_CH * rh * rw * sizeof(float));
|
|
if (!rn[k]) {
|
|
free(res);
|
|
goto done;
|
|
}
|
|
conv2d(res, pc, rh, rw, cw, NULL, DA_FUSION_CH, 3, 3, 1, 1, rn[k]);
|
|
free(res);
|
|
rdim[k][0] = rh;
|
|
rdim[k][1] = rw;
|
|
}
|
|
|
|
/* ---- feature fusion (coarse -> fine): refinenet 4,3,2,1 ---- */
|
|
{
|
|
float *path = NULL;
|
|
int ph = 0, pw = 0;
|
|
for (int step = 0; step < 4; step++) {
|
|
int k = 3 - step;
|
|
int n = 4 - step; /* refinenet index (4..1) */
|
|
int H = rdim[k][0], W = rdim[k][1];
|
|
/* target upsample size: next-finer feature's size, or x2 for the last. */
|
|
int OH, OW;
|
|
if (step < 3) {
|
|
OH = rdim[k - 1][0];
|
|
OW = rdim[k - 1][1];
|
|
}
|
|
else {
|
|
OH = H * 2;
|
|
OW = W * 2;
|
|
}
|
|
|
|
float *out = malloc((size_t)DA_FUSION_CH * OH * OW * sizeof(float));
|
|
if (!out) {
|
|
free(path);
|
|
goto done;
|
|
}
|
|
|
|
if (step == 0) {
|
|
/* deepest (refinenet4): no residual; rn[k] is the path input. */
|
|
if (da_refinenet(m, n, rn[k], NULL, DA_FUSION_CH, H, W, OH, OW, out) != 0) {
|
|
free(out);
|
|
free(path);
|
|
goto done;
|
|
}
|
|
}
|
|
else {
|
|
if (da_refinenet(m, n, path, rn[k], DA_FUSION_CH, ph, pw, OH, OW, out) != 0) {
|
|
free(out);
|
|
free(path);
|
|
goto done;
|
|
}
|
|
}
|
|
free(path);
|
|
path = out;
|
|
ph = OH;
|
|
pw = OW;
|
|
}
|
|
|
|
/* ---- output convs ---- */
|
|
const float *c1w = da_tensor(m, HEAD ".scratch.output_conv1.weight");
|
|
const float *c1b = da_tensor(m, HEAD ".scratch.output_conv1.bias");
|
|
const float *c2w = da_tensor(m, HEAD ".scratch.output_conv2.0.weight");
|
|
const float *c2b = da_tensor(m, HEAD ".scratch.output_conv2.0.bias");
|
|
const float *c3w = da_tensor(m, HEAD ".scratch.output_conv2.2.weight");
|
|
const float *c3b = da_tensor(m, HEAD ".scratch.output_conv2.2.bias");
|
|
if (!c1w || !c1b || !c2w || !c2b || !c3w || !c3b) {
|
|
free(path);
|
|
goto done;
|
|
}
|
|
|
|
/* output_conv1: DA_FUSION_CH -> DA_HEAD_CH, 3x3 pad-1 */
|
|
float *h1 = malloc((size_t)DA_HEAD_CH * ph * pw * sizeof(float));
|
|
if (!h1) {
|
|
free(path);
|
|
goto done;
|
|
}
|
|
conv2d(path, DA_FUSION_CH, ph, pw, c1w, c1b, DA_HEAD_CH, 3, 3, 1, 1, h1);
|
|
free(path);
|
|
|
|
/* upsample to (518,518), bilinear align_corners=True */
|
|
float *h1u = malloc((size_t)DA_HEAD_CH * DA_IMG * DA_IMG * sizeof(float));
|
|
if (!h1u) {
|
|
free(h1);
|
|
goto done;
|
|
}
|
|
bilinear_chw(h1, DA_HEAD_CH, ph, pw, DA_IMG, DA_IMG, 1, h1u);
|
|
free(h1);
|
|
|
|
/* output_conv2.0: DA_HEAD_CH -> DA_HEAD_HIDDEN, 3x3 pad-1, ReLU */
|
|
float *h2 = malloc((size_t)DA_HEAD_HIDDEN * DA_IMG * DA_IMG * sizeof(float));
|
|
if (!h2) {
|
|
free(h1u);
|
|
goto done;
|
|
}
|
|
conv2d(h1u, DA_HEAD_CH, DA_IMG, DA_IMG, c2w, c2b, DA_HEAD_HIDDEN, 3, 3, 1, 1, h2);
|
|
free(h1u);
|
|
relu_inplace(h2, (size_t)DA_HEAD_HIDDEN * DA_IMG * DA_IMG);
|
|
|
|
/* output_conv2.2: DA_HEAD_HIDDEN -> 1, 1x1 -> depth.
|
|
* (No terminal ReLU: the DA3 mono head emits a signed depth field.) */
|
|
conv2d(h2, DA_HEAD_HIDDEN, DA_IMG, DA_IMG, c3w, c3b, 1, 1, 1, 1, 0, depth_out);
|
|
free(h2);
|
|
rc = 0;
|
|
}
|
|
|
|
done:
|
|
for (int k = 0; k < 4; k++)
|
|
free(rn[k]);
|
|
free(chw);
|
|
return rc;
|
|
}
|
|
|
|
/* ================================================================== */
|
|
/* Tileable detrending */
|
|
/* ================================================================== */
|
|
|
|
/* Remove the dominant low-frequency tilt from a depth map so that the
|
|
* result tiles seamlessly. Monocular models assume a perspective camera and
|
|
* bake a smooth front-to-back ramp into the depth (darker far / brighter
|
|
* near). For a top-down, tileable texture that ramp is spurious and makes
|
|
* opposite edges mismatch. Under perspective the disparity of a flat ground
|
|
* plane is affine in pixel coordinates, so we least-squares fit a plane
|
|
* a*(x-cx) + b*(y-cy) + c and subtract its tilt, which equalizes opposite
|
|
* borders while preserving genuine surface relief. The coordinate grid is
|
|
* regular and centered, so the design matrix is orthogonal and the fit is a
|
|
* pair of independent 1D slopes. */
|
|
static void detrend_plane(float *img, int H, int W) {
|
|
if (H < 2 && W < 2)
|
|
return;
|
|
const double cx = (W - 1) * 0.5, cy = (H - 1) * 0.5;
|
|
double Sxx = 0.0, Syy = 0.0, Sxz = 0.0, Syz = 0.0;
|
|
for (int x = 0; x < W; x++) {
|
|
double dx = x - cx;
|
|
Sxx += dx * dx;
|
|
}
|
|
Sxx *= H;
|
|
for (int y = 0; y < H; y++) {
|
|
double dy = y - cy;
|
|
Syy += dy * dy;
|
|
}
|
|
Syy *= W;
|
|
for (int y = 0; y < H; y++) {
|
|
double dy = y - cy;
|
|
const float *row = img + (size_t)y * W;
|
|
for (int x = 0; x < W; x++) {
|
|
double v = row[x];
|
|
Sxz += v * (x - cx);
|
|
Syz += v * dy;
|
|
}
|
|
}
|
|
double a = (Sxx > 0.0) ? Sxz / Sxx : 0.0;
|
|
double b = (Syy > 0.0) ? Syz / Syy : 0.0;
|
|
for (int y = 0; y < H; y++) {
|
|
double dy = y - cy;
|
|
float *row = img + (size_t)y * W;
|
|
for (int x = 0; x < W; x++)
|
|
row[x] -= (float)(a * (x - cx) + b * dy);
|
|
}
|
|
}
|
|
|
|
/* ------------------------------------------------------------------ */
|
|
/* Seamless tiling via periodic+smooth decomposition (Moisan 2011) */
|
|
/* ------------------------------------------------------------------ */
|
|
/* Equalizing opposite edges with a separable per-row/per-column shift injects a
|
|
* visible streak on any row or column whose two edges differ sharply (e.g. black
|
|
* on one side, white on the other): that line gets shifted hard relative to its
|
|
* neighbours. The cure is to stop treating rows and columns independently and
|
|
* instead subtract the single smoothest 2D field that makes the image periodic.
|
|
* That field is Moisan's "smooth component": the solution of a periodic Poisson
|
|
* equation whose source is the jump between opposite borders. A localized edge
|
|
* mismatch is then diffused into a gentle 2D bump rather than smeared along a
|
|
* whole scanline, so no streak remains and the result tiles exactly. Solved in
|
|
* the Fourier domain, so it requires power-of-two dimensions (always the case
|
|
* for textures); other sizes are left unchanged. */
|
|
|
|
#ifndef IRIS_TWO_PI
|
|
#define IRIS_TWO_PI 6.28318530717958647692
|
|
#endif
|
|
|
|
static int is_pow2(int n) {
|
|
return n > 0 && (n & (n - 1)) == 0;
|
|
}
|
|
|
|
/* In-place iterative radix-2 Cooley-Tukey FFT. inv=0 forward (e^-i), inv=1
|
|
* inverse (e^+i, scaled by 1/n). n must be a power of two. */
|
|
static void fft1d(double *re, double *im, int n, int inv) {
|
|
for (int i = 1, j = 0; i < n; i++) {
|
|
int bit = n >> 1;
|
|
for (; j & bit; bit >>= 1)
|
|
j ^= bit;
|
|
j ^= bit;
|
|
if (i < j) {
|
|
double tr = re[i];
|
|
re[i] = re[j];
|
|
re[j] = tr;
|
|
double ti = im[i];
|
|
im[i] = im[j];
|
|
im[j] = ti;
|
|
}
|
|
}
|
|
for (int len = 2; len <= n; len <<= 1) {
|
|
double ang = IRIS_TWO_PI / len * (inv ? 1.0 : -1.0);
|
|
double wr = cos(ang), wi = sin(ang);
|
|
for (int i = 0; i < n; i += len) {
|
|
double cwr = 1.0, cwi = 0.0;
|
|
for (int k = 0; k < len / 2; k++) {
|
|
int a = i + k, b = i + k + len / 2;
|
|
double vr = re[b] * cwr - im[b] * cwi;
|
|
double vi = re[b] * cwi + im[b] * cwr;
|
|
re[b] = re[a] - vr;
|
|
im[b] = im[a] - vi;
|
|
re[a] += vr;
|
|
im[a] += vi;
|
|
double ncwr = cwr * wr - cwi * wi;
|
|
cwi = cwr * wi + cwi * wr;
|
|
cwr = ncwr;
|
|
}
|
|
}
|
|
}
|
|
if (inv) {
|
|
for (int i = 0; i < n; i++) {
|
|
re[i] /= n;
|
|
im[i] /= n;
|
|
}
|
|
}
|
|
}
|
|
|
|
/* 2D FFT: transform every row (length W), then every column (length H), using
|
|
* caller-provided column scratch of length >= H. */
|
|
static void fft2d(double *re, double *im, int H, int W, int inv, double *cr, double *ci) {
|
|
for (int i = 0; i < H; i++)
|
|
fft1d(re + (size_t)i * W, im + (size_t)i * W, W, inv);
|
|
for (int j = 0; j < W; j++) {
|
|
for (int i = 0; i < H; i++) {
|
|
cr[i] = re[(size_t)i * W + j];
|
|
ci[i] = im[(size_t)i * W + j];
|
|
}
|
|
fft1d(cr, ci, H, inv);
|
|
for (int i = 0; i < H; i++) {
|
|
re[(size_t)i * W + j] = cr[i];
|
|
im[(size_t)i * W + j] = ci[i];
|
|
}
|
|
}
|
|
}
|
|
|
|
/* Subtract the smooth component so img becomes seamlessly tileable. Only acts on
|
|
* power-of-two dimensions (textures always qualify); other sizes are untouched. */
|
|
static void make_seamless_poisson(float *img, int H, int W) {
|
|
if (!is_pow2(H) || !is_pow2(W) || H < 2 || W < 2)
|
|
return;
|
|
|
|
size_t n = (size_t)H * W;
|
|
double *re = calloc(n, sizeof(double));
|
|
double *im = calloc(n, sizeof(double));
|
|
double *cr = malloc((size_t)H * sizeof(double));
|
|
double *ci = malloc((size_t)H * sizeof(double));
|
|
if (!re || !im || !cr || !ci) {
|
|
free(re);
|
|
free(im);
|
|
free(cr);
|
|
free(ci);
|
|
return;
|
|
}
|
|
|
|
/* Boundary jump field: the wrap-around difference across each border,
|
|
* accumulated (corners receive both a row and a column contribution). */
|
|
for (int i = 0; i < H; i++) {
|
|
double *row = re + (size_t)i * W;
|
|
double l = img[(size_t)i * W + 0], r = img[(size_t)i * W + (W - 1)];
|
|
row[0] += r - l;
|
|
row[W - 1] += l - r;
|
|
}
|
|
for (int j = 0; j < W; j++) {
|
|
double t = img[(size_t)0 * W + j], b = img[(size_t)(H - 1) * W + j];
|
|
re[(size_t)0 * W + j] += b - t;
|
|
re[(size_t)(H - 1) * W + j] += t - b;
|
|
}
|
|
|
|
fft2d(re, im, H, W, 0, cr, ci);
|
|
|
|
/* Divide by the periodic-Laplacian eigenvalues to solve the Poisson eq;
|
|
* the DC term (constant offset) is undetermined, so pin it to zero. */
|
|
for (int q = 0; q < H; q++) {
|
|
double cq = 2.0 * cos(IRIS_TWO_PI * q / H);
|
|
for (int j = 0; j < W; j++) {
|
|
size_t idx = (size_t)q * W + j;
|
|
if (q == 0 && j == 0) {
|
|
re[idx] = 0.0;
|
|
im[idx] = 0.0;
|
|
continue;
|
|
}
|
|
double denom = cq + 2.0 * cos(IRIS_TWO_PI * j / W) - 4.0;
|
|
re[idx] /= denom;
|
|
im[idx] /= denom;
|
|
}
|
|
}
|
|
|
|
fft2d(re, im, H, W, 1, cr, ci);
|
|
|
|
for (size_t i = 0; i < n; i++)
|
|
img[i] -= (float)re[i];
|
|
|
|
free(re);
|
|
free(im);
|
|
free(cr);
|
|
free(ci);
|
|
}
|
|
|
|
/* ================================================================== */
|
|
/* Public API */
|
|
/* ================================================================== */
|
|
|
|
iris_depth_t *iris_depth_load(const char *path) {
|
|
safetensors_file_t *sf = safetensors_open(path);
|
|
if (!sf) {
|
|
fprintf(stderr, "DepthAnything3: failed to open '%s'\n", path);
|
|
return NULL;
|
|
}
|
|
if (!safetensors_find(sf, BB ".cls_token") || !safetensors_find(sf, HEAD ".scratch.output_conv2.2.weight")) {
|
|
fprintf(stderr, "DepthAnything3: '%s' does not look like a Depth Anything 3 mono model\n", path);
|
|
safetensors_close(sf);
|
|
return NULL;
|
|
}
|
|
iris_depth_t *m = calloc(1, sizeof(iris_depth_t));
|
|
if (!m) {
|
|
safetensors_close(sf);
|
|
return NULL;
|
|
}
|
|
m->sf = sf;
|
|
return m;
|
|
}
|
|
|
|
void iris_depth_set_tileable(iris_depth_t *m, int on) {
|
|
if (m)
|
|
m->tileable = on ? 1 : 0;
|
|
}
|
|
|
|
void iris_depth_free(iris_depth_t *m) {
|
|
if (!m)
|
|
return;
|
|
safetensors_close(m->sf);
|
|
free(m);
|
|
}
|
|
|
|
iris_image *iris_depth_estimate(iris_depth_t *m, const iris_image *input) {
|
|
if (!m || !input || !input->data)
|
|
return NULL;
|
|
|
|
int W = input->width, H = input->height, ic = input->channels;
|
|
|
|
/* ---- Preprocess: resize to 518x518 (bilinear) + ImageNet normalize ----
|
|
* Build a planar CHW float source at the original resolution, then resize. */
|
|
float *src = malloc((size_t)3 * H * W * sizeof(float));
|
|
if (!src)
|
|
return NULL;
|
|
const size_t src_plane = (size_t)H * W;
|
|
for (int y = 0; y < H; y++) {
|
|
for (int x = 0; x < W; x++) {
|
|
const uint8_t *px = input->data + ((size_t)y * W + x) * ic;
|
|
size_t pi = (size_t)y * W + x;
|
|
src[0 * src_plane + pi] = px[0] / 255.0f;
|
|
src[1 * src_plane + pi] = (ic > 1 ? px[1] : px[0]) / 255.0f;
|
|
src[2 * src_plane + pi] = (ic > 2 ? px[2] : px[0]) / 255.0f;
|
|
}
|
|
}
|
|
|
|
float *pixels = malloc((size_t)3 * DA_IMG * DA_IMG * sizeof(float));
|
|
if (!pixels) {
|
|
free(src);
|
|
return NULL;
|
|
}
|
|
bilinear_chw(src, 3, H, W, DA_IMG, DA_IMG, 0, pixels);
|
|
free(src);
|
|
|
|
/* normalize per channel */
|
|
for (int c = 0; c < 3; c++) {
|
|
float *p = pixels + (size_t)c * DA_IMG * DA_IMG;
|
|
float mean = DA_MEAN[c], istd = 1.0f / DA_STD[c];
|
|
for (size_t i = 0; i < (size_t)DA_IMG * DA_IMG; i++)
|
|
p[i] = (p[i] - mean) * istd;
|
|
}
|
|
|
|
/* ---- Backbone ---- */
|
|
float *feats[4] = {0};
|
|
for (int k = 0; k < 4; k++) {
|
|
feats[k] = malloc((size_t)DA_TOKENS * DA_HIDDEN * sizeof(float));
|
|
if (!feats[k]) {
|
|
for (int j = 0; j < k; j++)
|
|
free(feats[j]);
|
|
free(pixels);
|
|
return NULL;
|
|
}
|
|
}
|
|
|
|
#ifdef USE_VULKAN
|
|
/* Fold the whole forward into one GPU batch when resident ops exist. */
|
|
if (iris_vulkan_available())
|
|
iris_gpu_batch_begin();
|
|
#endif
|
|
|
|
int ok = (da_backbone(m, pixels, feats) == 0);
|
|
free(pixels);
|
|
|
|
float *depth = NULL;
|
|
if (ok) {
|
|
depth = malloc((size_t)DA_IMG * DA_IMG * sizeof(float));
|
|
if (!depth || da_head(m, feats, depth) != 0) {
|
|
free(depth);
|
|
depth = NULL;
|
|
ok = 0;
|
|
}
|
|
}
|
|
|
|
#ifdef USE_VULKAN
|
|
if (iris_vulkan_available())
|
|
iris_gpu_batch_end();
|
|
#endif
|
|
|
|
for (int k = 0; k < 4; k++)
|
|
free(feats[k]);
|
|
if (!ok) {
|
|
free(depth);
|
|
return NULL;
|
|
}
|
|
|
|
/* ---- Resize depth back to input resolution (bilinear) ---- */
|
|
float *depth_full = malloc((size_t)H * W * sizeof(float));
|
|
if (!depth_full) {
|
|
free(depth);
|
|
return NULL;
|
|
}
|
|
bilinear_chw(depth, 1, DA_IMG, DA_IMG, H, W, 0, depth_full);
|
|
free(depth);
|
|
|
|
/* ---- Optional: make the map tile seamlessly ----
|
|
* First remove the perspective tilt (keeps the border correction small),
|
|
* then feather opposite edges together so no seam remains. */
|
|
if (m->tileable) {
|
|
detrend_plane(depth_full, H, W);
|
|
make_seamless_poisson(depth_full, H, W);
|
|
}
|
|
|
|
/* ---- Min/max normalize to grayscale (inverted: brighter = farther) ---- */
|
|
float mn = depth_full[0], mx = depth_full[0];
|
|
for (size_t i = 1; i < (size_t)H * W; i++) {
|
|
if (depth_full[i] < mn)
|
|
mn = depth_full[i];
|
|
if (depth_full[i] > mx)
|
|
mx = depth_full[i];
|
|
}
|
|
float range = (mx - mn) > 1e-8f ? (mx - mn) : 1.0f;
|
|
|
|
iris_image *out = iris_image_create(W, H, 1);
|
|
if (!out) {
|
|
free(depth_full);
|
|
return NULL;
|
|
}
|
|
for (size_t i = 0; i < (size_t)H * W; i++) {
|
|
float v = 1.0f - (depth_full[i] - mn) / range;
|
|
int iv = (int)(v * 255.0f + 0.5f);
|
|
if (iv < 0)
|
|
iv = 0;
|
|
if (iv > 255)
|
|
iv = 255;
|
|
out->data[i] = (uint8_t)iv;
|
|
}
|
|
free(depth_full);
|
|
return out;
|
|
}
|