Skip to content

Instantly share code, notes, and snippets.

@harry7557558

harry7557558/colmap_export.py Secret

Last active July 9, 2026 00:55
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save harry7557558/1e31891d1fc8050bfdf771ef78f48c1c to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save harry7557558/1e31891d1fc8050bfdf771ef78f48c1c to your computer and use it in GitHub Desktop.
Download ZIP
Raw
colmap_export.py
#!/usr/bin/env python3
"""Shared COLMAP export helpers for the render_*.py fractal renderers.
The renderers use an OpenGL-style camera: +X right, +Y up, camera looks
along -Z. COLMAP stores world->camera poses in an OpenCV-style frame:
+X right, +Y down, +Z forward. ``c2w_gl_to_colmap_w2c`` converts between
the two; getting this flip wrong is what makes cameras appear inverted in
the COLMAP GUI.
Because the renderers ray-trace an exact implicit surface, the initial
point cloud is sampled directly from primary-ray hits (ground truth)
instead of being triangulated by COLMAP from image features.
"""
import numpy as np
from pathlib import Path
def rotmat_to_quat(R: np.ndarray) -> np.ndarray:
"""Convert 3x3 rotation matrix to quaternion in COLMAP order [w, x, y, z]."""
m = R
tr = float(m[0, 0] + m[1, 1] + m[2, 2])
if tr > 0:
S = np.sqrt(tr + 1.0) * 2.0
qw = 0.25 * S
qx = (m[2, 1] - m[1, 2]) / S
qy = (m[0, 2] - m[2, 0]) / S
qz = (m[1, 0] - m[0, 1]) / S
elif (m[0, 0] > m[1, 1]) and (m[0, 0] > m[2, 2]):
S = np.sqrt(1.0 + float(m[0, 0] - m[1, 1] - m[2, 2])) * 2.0
qw = (m[2, 1] - m[1, 2]) / S
qx = 0.25 * S
qy = (m[0, 1] + m[1, 0]) / S
qz = (m[0, 2] + m[2, 0]) / S
elif m[1, 1] > m[2, 2]:
S = np.sqrt(1.0 + float(m[1, 1] - m[0, 0] - m[2, 2])) * 2.0
qw = (m[0, 2] - m[2, 0]) / S
qx = (m[0, 1] + m[1, 0]) / S
qy = 0.25 * S
qz = (m[1, 2] + m[2, 1]) / S
else:
S = np.sqrt(1.0 + float(m[2, 2] - m[0, 0] - m[1, 1])) * 2.0
qw = (m[1, 0] - m[0, 1]) / S
qx = (m[0, 2] + m[2, 0]) / S
qy = (m[1, 2] + m[2, 1]) / S
qz = 0.25 * S
q = np.array([qw, qx, qy, qz], dtype=np.float64)
return q / np.linalg.norm(q)
# Flips the camera-space +Y (up -> down) and +Z (backward -> forward) axes.
GL_TO_CV = np.diag([1.0, -1.0, -1.0])
def c2w_gl_to_colmap_w2c(c2w: np.ndarray):
"""OpenGL camera-to-world matrix -> COLMAP world-to-camera (R, t)."""
R_c2w_cv = c2w[:3, :3] @ GL_TO_CV
R = R_c2w_cv.T
t = -R @ c2w[:3, 3]
return R, t
def sample_surface_points(hits: np.ndarray, img8: np.ndarray, n_points: int,
rng: np.random.Generator):
"""Pick up to n_points surface points from a primary-ray hit map.
hits: (th, tw, 4) array of [hit_flag, x, y, z] per traced pixel.
img8: (H, W, 3) uint8 rendered image used to color the points.
Returns (points_xyz, points_rgb, pixels_xy) where pixels_xy are the
pixel-center coordinates in full-resolution image space (for
reprojection checks).
"""
th, tw = hits.shape[:2]
ys, xs = np.nonzero(hits[:, :, 0] > 0.5)
if ys.size == 0:
return (np.zeros((0, 3)), np.zeros((0, 3), np.uint8), np.zeros((0, 2)))
if ys.size > n_points:
sel = rng.choice(ys.size, size=n_points, replace=False)
ys, xs = ys[sel], xs[sel]
pts = hits[ys, xs, 1:4].astype(np.float64)
H, W = img8.shape[:2]
pix = np.stack([(xs + 0.5) * (W / tw), (ys + 0.5) * (H / th)], axis=1)
ix = np.clip(pix[:, 0].astype(np.int64), 0, W - 1)
iy = np.clip(pix[:, 1].astype(np.int64), 0, H - 1)
rgb = img8[iy, ix]
return pts, rgb, pix
def reprojection_errors(points_xyz: np.ndarray, pixels_xy: np.ndarray,
R: np.ndarray, t: np.ndarray,
fx: float, fy: float, cx: float, cy: float) -> np.ndarray:
"""Project world points with a COLMAP pose; return per-point pixel error.
Points behind the camera get infinite error, so any convention mistake
(axis flip, wrong inverse) shows up as a huge value instead of passing
silently.
"""
cam = points_xyz @ R.T + t
z = cam[:, 2]
valid = z > 1e-9
z_safe = np.where(valid, z, 1.0)
u = fx * cam[:, 0] / z_safe + cx
v = fy * cam[:, 1] / z_safe + cy
err = np.hypot(u - pixels_xy[:, 0], v - pixels_xy[:, 1])
err[~valid] = np.inf
return err
def write_colmap_text(dataset_dir, cameras, images_meta,
points_xyz=None, points_rgb=None) -> None:
"""Write sparse/0/{cameras,images,points3D}.txt under dataset_dir.
cameras: list of (camera_id, model, width, height, params)
images_meta: list of (image_id, q_wxyz, t, camera_id, name); name must be
relative to the images/ directory (no "images/" prefix).
points_xyz/points_rgb: (N, 3) float / (N, 3) uint8 initial point cloud.
"""
sparse0 = Path(dataset_dir) / "sparse" / "0"
sparse0.mkdir(parents=True, exist_ok=True)
stale = sorted(p.name for p in sparse0.glob("*.bin"))
if stale:
print(f"WARNING: {sparse0} already contains a binary model "
f"({', '.join(stale)}). Most loaders prefer *.bin over *.txt, "
f"so delete or move those files or they will shadow this export.")
if points_xyz is None:
points_xyz = np.zeros((0, 3))
if points_rgb is None:
points_rgb = np.full((len(points_xyz), 3), 128, np.uint8)
with (sparse0 / "cameras.txt").open("w", encoding="utf-8") as f:
f.write("# Camera list with one line of data per camera:\n")
f.write("# CAMERA_ID, MODEL, WIDTH, HEIGHT, PARAMS[]\n")
f.write("# Number of cameras: {}\n".format(len(cameras)))
for cam_id, model, width, height, params in cameras:
pstr = " ".join("{:.10f}".format(float(p)) for p in params)
f.write(f"{cam_id} {model} {width} {height} {pstr}\n")
with (sparse0 / "images.txt").open("w", encoding="utf-8") as f:
f.write("# Image list with two lines of data per image:\n")
f.write("# IMAGE_ID, QW, QX, QY, QZ, TX, TY, TZ, CAMERA_ID, NAME\n")
f.write("# POINTS2D[] as (X, Y, POINT3D_ID)\n")
f.write("# Number of images: {}\n".format(len(images_meta)))
for image_id, q, t, cam_id, name in images_meta:
qstr = " ".join("{:.10f}".format(float(v)) for v in q)
tstr = " ".join("{:.10f}".format(float(v)) for v in t)
f.write(f"{image_id} {qstr} {tstr} {cam_id} {name}\n")
f.write("\n")
with (sparse0 / "points3D.txt").open("w", encoding="utf-8") as f:
f.write("# 3D point list with one line of data per point:\n")
f.write("# POINT3D_ID, X, Y, Z, R, G, B, ERROR, TRACK[] as (IMAGE_ID, POINT2D_IDX)\n")
f.write("# Number of points: {}\n".format(len(points_xyz)))
for j, (p, c) in enumerate(zip(points_xyz, points_rgb)):
f.write(f"{j + 1} {p[0]:.8f} {p[1]:.8f} {p[2]:.8f} "
f"{int(c[0])} {int(c[1])} {int(c[2])} 0\n")
Raw
render_mandelbulb.py
#!/usr/bin/env python3
import argparse
import math
import os
from pathlib import Path
from typing import Tuple
import numpy as np
from PIL import Image
from colmap_export import (rotmat_to_quat, c2w_gl_to_colmap_w2c,
sample_surface_points, reprojection_errors,
write_colmap_text)
# -----------------------------
# Utilities
# -----------------------------
def clamp(x, lo, hi):
return max(lo, min(hi, x))
def normalize_np(v: np.ndarray, eps: float = 1e-12) -> np.ndarray:
n = np.linalg.norm(v)
if n < eps:
return v
return v / n
def look_at_c2w(eye: np.ndarray, target: np.ndarray, up: np.ndarray) -> np.ndarray:
"""Return camera-to-world 4x4 matrix, right-handed, camera looks along -Z."""
forward = normalize_np(target - eye)
right = normalize_np(np.cross(forward, up))
true_up = np.cross(right, forward)
c2w = np.eye(4, dtype=np.float64)
# Camera axes in world space:
# +X = right, +Y = true_up, +Z = -forward (because camera looks down -Z)
c2w[:3, 0] = right
c2w[:3, 1] = true_up
c2w[:3, 2] = -forward
c2w[:3, 3] = eye
return c2w
def c2w_to_w2c(c2w: np.ndarray) -> np.ndarray:
return np.linalg.inv(c2w)
def srgb_encode(x: np.ndarray) -> np.ndarray:
x = np.clip(x, 0.0, 1.0)
a = 0.055
out = np.where(x <= 0.0031308, 12.92 * x, (1.0 + a) * np.power(x, 1.0 / 2.4) - a)
return np.clip(out, 0.0, 1.0)
# -----------------------------
# Taichi setup
# -----------------------------
import taichi as ti
try:
ti.init(arch=ti.cuda, default_fp=ti.f32)
except Exception:
ti.init(arch=ti.gpu, default_fp=ti.f32)
except Exception:
ti.init(arch=ti.cpu, default_fp=ti.f32)
vec3 = ti.types.vector(3, ti.f32)
# Global config fields (shape=())
width_f = ti.field(dtype=ti.i32, shape=())
height_f = ti.field(dtype=ti.i32, shape=())
fov_y_deg_f = ti.field(dtype=ti.f32, shape=())
camera_radius_f = ti.field(dtype=ti.f32, shape=())
camera_elev_deg_f = ti.field(dtype=ti.f32, shape=())
camera_target = ti.Vector.field(3, dtype=ti.f32, shape=())
camera_pos = ti.Vector.field(3, dtype=ti.f32, shape=())
camera_right = ti.Vector.field(3, dtype=ti.f32, shape=())
camera_up = ti.Vector.field(3, dtype=ti.f32, shape=())
camera_forward = ti.Vector.field(3, dtype=ti.f32, shape=())
img = ti.Vector.field(3, dtype=ti.f32, shape=())
# Per-pixel primary-ray hits: [hit_flag, x, y, z]. Traced at a reduced
# resolution to sample the initial point cloud; reallocated in main().
hit_map = ti.Vector.field(4, dtype=ti.f32, shape=())
# Small constants
PI = 3.1415926535897932384626433832795
@ti.func
def colormap_turbo(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(-266.03287948, -63.82163439, -97.24397473) # k=7
c = c * t + vec3(874.85901369, 202.98596195, 447.99287080) # k=6
c = c * t + vec3(-1061.31232675, -245.26823636, -766.82678386) # k=5
c = c * t + vec3(550.44382083, 149.40813531, 604.07329733) # k=4
c = c * t + vec3(-90.46877071, -55.09180383, -203.76964931) # k=3
c = c * t + vec3(-10.15061040, 9.64581379, 9.26380342) # k=2
c = c * t + vec3(2.94711014, 2.06520532, 6.33792818) # k=1
c = c * t + vec3(0.14637796, 0.08594198, 0.23431523) # k=0
return c
@ti.func
def colormap_viridis(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(2.89151588, -0.17213704, 0.17750210) # k=3
c = c * t + vec3(-2.14339482, -0.16754949, -1.71784815) # k=2
c = c * t + vec3(0.04824199, 1.23905227, 1.23133794) # k=1
c = c * t + vec3(0.29387218, 0.01752027, 0.34360395) # k=0
return c
@ti.func
def colormap_magma(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(1.69082060, -2.36381252, 0.75796302) # k=4
c = c * t + vec3(-5.49016037, 5.48987916, 3.89534203) # k=3
c = c * t + vec3(4.30604404, -2.93736985, -7.52578597) # k=2
c = c * t + vec3(0.45897387, 0.82685306, 3.73227094) # k=1
c = c * t + vec3(-0.00157294, -0.01042786, -0.05490989) # k=0
return c
@ti.func
def colormap_inferno(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(1.66959821, -2.07978236, 7.14515755) # k=4
c = c * t + vec3(-5.11462265, 4.46548484, -6.74156149) # k=3
c = c * t + vec3(3.64171238, -1.86975118, -2.30104808) # k=2
c = c * t + vec3(0.74909880, 0.51010061, 2.59796038) # k=1
c = c * t + vec3(-0.01087222, -0.00208239, 0.00011545) # k=0
return c
@ti.func
def colormap_cividis(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(0.59905315, 0.14530728, -0.89703048) # k=3
c = c * t + vec3(-0.95976947, -0.07321042, 0.66155846) # k=2
c = c * t + vec3(1.43486697, 0.69925032, 0.05514949) # k=1
c = c * t + vec3(-0.05784352, 0.13417334, 0.38674306) # k=0
return c
@ti.func
def saturate(x):
return ti.min(ti.max(x, 0.0), 1.0)
@ti.func
def hash_u32(x):
x ^= x >> 16
x *= ti.u32(0x7feb352d)
x ^= x >> 15
x *= ti.u32(0x846ca68b)
x ^= x >> 16
return x
@ti.func
def rand01(seed):
seed = hash_u32(seed)
return ti.cast(seed & 0x00FFFFFF, ti.f32) / ti.cast(0x01000000, ti.f32)
@ti.func
def rand2(seed):
return ti.Vector([rand01(seed), rand01(seed ^ ti.u32(0x9e3779b9))])
@ti.func
def mix(a, b, t):
return a * (1.0 - t) + b * t
SURFACE_EPS = 1e-5
BOUNDING_RADIUS = 2.0
@ti.func
def mandelbulb_implicit(p: vec3) -> ti.types.vector(2, ti.f32):
z = p
dr = 1.0
r = 0.0
power = 2
min_radius = 1e9
for i in range(64):
r = z.norm()
min_radius = ti.min(min_radius, r)
if r > 64.0:
break
r = ti.max(r, 1e-8)
theta = ti.acos(ti.max(-1.0, ti.min(1.0, z.z / r)))
phi = ti.atan2(z.y, z.x)
zr = ti.pow(r, power)
dr = ti.pow(r, power - 1.0) * power * dr + 1.0
theta *= power
phi *= power
sin_t = ti.sin(theta)
z1 = zr * vec3([sin_t * ti.cos(phi), sin_t * ti.sin(phi), ti.cos(theta)]) + p
r1 = z1.norm()
z = z1
r = z.norm()
r = ti.max(r, 1e-8)
de = 0.5 * ti.log(r) * r / dr
#de = ti.abs(de)
return ti.Vector([de, min_radius])
@ti.func
def sphere_clip(ro: vec3, rd: vec3, radius: ti.f32) -> ti.types.vector(3, ti.f32):
b = ro.dot(rd)
c = ro.dot(ro) - radius * radius
disc = b * b - c
result = ti.Vector([0.0, 0.0, 0.0])
if disc >= 0.0:
sqrt_disc = ti.sqrt(disc)
result = ti.Vector([1.0, -b - sqrt_disc, -b + sqrt_disc])
return result
@ti.func
def estimate_normal(p: vec3) -> vec3:
e = 5e-5
ex = vec3([e, 0.0, 0.0])
ey = vec3([0.0, e, 0.0])
ez = vec3([0.0, 0.0, e])
dx1 = mandelbulb_implicit(p + ex)[0]
dx0 = mandelbulb_implicit(p - ex)[0]
dy1 = mandelbulb_implicit(p + ey)[0]
dy0 = mandelbulb_implicit(p - ey)[0]
dz1 = mandelbulb_implicit(p + ez)[0]
dz0 = mandelbulb_implicit(p - ez)[0]
n = vec3([dx1 - dx0, dy1 - dy0, dz1 - dz0])
return n / e
@ti.func
def random_hemisphere_cosine(n: vec3, seed: ti.u32) -> vec3:
r = rand2(seed)
phi = 2.0 * PI * r.x
cos_theta = ti.sqrt(1.0 - r.y)
sin_theta = ti.sqrt(r.y)
local = vec3([ti.cos(phi) * sin_theta, ti.sin(phi) * sin_theta, cos_theta])
# Build orthonormal basis
tangent = ti.Vector([0.0, 1.0, 0.0])
if ti.abs(n.y) > 0.999:
tangent = ti.Vector([1.0, 0.0, 0.0])
tangent = (tangent - n * tangent.dot(n)).normalized()
bitangent = n.cross(tangent)
return (tangent * local.x + bitangent * local.y + n * local.z).normalized()
@ti.func
def sky_radiance(rd: vec3) -> vec3:
# t = 0.5 * (rd.y + 1.0)
t = max(rd.normalized().y, 0.0)
sky_bottom = ti.Vector([0.7, 0.75, 0.8]) ** 2.2
sky_top = ti.Vector([1.2, 1.4, 1.8]) ** 2.2
col = mix(sky_bottom, sky_top, t)
return col
@ti.func
def trace_scene(ro: vec3, rd: vec3, seed: ti.u32) -> ti.types.vector(5, ti.f32):
# based on https://raw.githubusercontent.com/harry7557558/spirulae/refs/heads/master/implicit3-rt/frag-render.glsl
rd_n = rd.normalized()
result = ti.Vector([0.0, 0.0, ro.x, ro.y, ro.z])
clip = sphere_clip(ro, rd_n, BOUNDING_RADIUS)
t0c = ti.max(clip[1], 0.0)
t1c = clip[2]
if clip[0] >= 0.5 and t1c >= t0c:
step_sz = 500.0*SURFACE_EPS
t = t0c
dt = step_sz
v = 0.0
v0 = 0.0
v00 = 0.0
v1 = 0.0
dt0 = 0.0
dt00 = 0.0
is_bisecting = 0
prev_t = t0c
t_hi = 0.0
tfar = t1c
found = 0
result_t = 0.0
for iteration in range(16384):
if found == 1:
break
if t > tfar:
break
val_mr = mandelbulb_implicit(ro + rd_n * t) - 20.0*SURFACE_EPS
v = val_mr[0]
if is_bisecting == 1:
if t_hi - prev_t <= 1e-5:
dv = v1 - v0
if ti.abs(dv) > 1e-12:
result_t = prev_t - v0 / dv * (t_hi - prev_t)
else:
result_t = 0.5 * (prev_t + t_hi)
found = 1
else:
if v * v0 < 0.0:
t_hi = t
v1 = v
else:
prev_t = t
v0 = v
t = 0.5 * (prev_t + t_hi)
elif v * v0 < 0.0:
is_bisecting = 1
t_hi = t
v1 = v
t = 0.5 * (prev_t + t_hi)
else:
if ti.math.isnan(dt0) or dt0 <= 0.0:
v00 = v
v0 = v
dt0 = 0.0
dt00 = 0.0
g = 0.0
if dt0 > 0.0:
if dt00 > 0.0:
g = (v00 * dt0 / (dt00 * (dt0 + dt00))
- v0 * (dt0 + dt00) / (dt0 * dt00)
+ v * (2.0 * dt0 + dt00) / (dt0 * (dt0 + dt00)))
else:
g = (v - v0) / dt0
dt00 = dt0
dt0 = dt
prev_t = t
v00 = v0
v0 = v
ddt = ti.abs(v / g)
if ddt > step_sz or ti.math.isnan(ddt) or ti.math.isinf(ddt):
dt = step_sz
else:
dt_cand = ti.min(ddt - step_sz, tfar - prev_t - 0.01 * step_sz)
dt = ti.max(0.05 * step_sz, ti.min(dt_cand, step_sz))
if iteration < 2:
dt *= rand01(seed + iteration)
t += dt
if found == 1:
p = ro + rd_n * result_t
result = ti.Vector([1.0, result_t, p.x, p.y, p.z])
return result
@ti.func
def trace_scene_slow(ro: vec3, rd: vec3, seed: ti.u32) -> ti.types.vector(5, ti.f32):
"""
Sphere tracing.
Returns [hit_flag, t, p.x, p.y, p.z] with hit_flag in {0,1}.
"""
t = 0.0
hit = 0.0
p = ro
for i in range(16384):
p = ro + rd * t
sdf = 0.25 * mandelbulb_implicit(p)[0]
if sdf < SURFACE_EPS:
hit = ti.f32(i)
break
t += sdf
if t > 100.0:
break
return ti.Vector([hit, t, p.x, p.y, p.z])
@ti.func
def direct_light(n: vec3, p: vec3, view_dir: vec3, seed: ti.u32) -> vec3:
# One sun light with a hard shadow ray (cheap, but useful for contrast).
ldir = -ti.Vector([-0.4, 0.8, -0.2]).normalized()
ndotl = saturate(n.dot(ldir))
return_val = ti.Vector([0.0, 0.0, 0.0])
if ndotl > 0.0:
# Shadow ray
shadow_ro = p + n * (SURFACE_EPS * 4.0)
shadow_hit = trace_scene(shadow_ro, ldir, seed)
visible = 1.0 if shadow_hit[0] == 0.0 else 0.0
# Simple specular-like glint to help structure read well.
h = (ldir + view_dir).normalized()
spec = ti.pow(saturate(n.dot(h)), 96.0) * 0.15
sun_color = 1.0 * ti.Vector([1.0, 0.95, 0.9])
return_val = visible * (ndotl + spec) * sun_color
return return_val
@ti.func
def shade_background(rd: vec3, throughput: vec3, seed: ti.u32) -> vec3:
# Add a subtle fog / aerial perspective based on ray direction.
c = sky_radiance(rd)
return throughput * c
@ti.func
def render_path(ro: vec3, rd: vec3, px_seed: ti.u32) -> vec3:
throughput = ti.Vector([1.0, 1.0, 1.0])
radiance = ti.Vector([0.0, 0.0, 0.0])
cur_ro = ro
cur_rd = rd
for bounce in range(4):
hit = trace_scene(cur_ro, cur_rd, px_seed + bounce * 131 + 17)
if hit[0] == 0.0:
if bounce > 0:
radiance += shade_background(cur_rd, throughput, px_seed + bounce * 131)
break
t = hit[1]
p = vec3([hit[2], hit[3], hit[4]])
n = estimate_normal(p)
n_norm = n.norm()
n = n.normalized()
# Material: mostly diffuse with a tiny amount of view-dependent sheen.
# albedo = (0.5 + 0.35 * ti.sin(1.0 * PI * n) + 0.1 * ti.sin(4.0 * PI * n))
albedo = colormap_magma(0.5 + 0.5 * ti.cos(ti.log(0.25*n_norm)))
# albedo = ti.max(1.0 - 1.5 * (1.0 - albedo), 0.0)
# albedo = albedo ** 2.0
# albedo = albedo + (ti.Vector([0.1, 0.0, 0.2]) - albedo) * (1.0 - ti.exp(-0.01 * hit[0]))
# albedo = albedo + (ti.Vector([0.0, 0.0, 0.0]) - albedo) * (1.0 - ti.exp(-0.1 / n_norm))
albedo = 0.9 * albedo ** 2.2
radiance += throughput * albedo * direct_light(n, p, -cur_rd, px_seed + bounce * 977 + 123)
new_dir = random_hemisphere_cosine(n, px_seed + bounce * 977 + 17)
cur_ro = p + n * (SURFACE_EPS * 8.0)
cur_rd = new_dir
throughput *= albedo
# Russian roulette
if bounce >= 2:
p_rr = ti.min(0.95, ti.max(throughput.x, ti.max(throughput.y, throughput.z)))
if rand01(px_seed + bounce * 31337 + 123) > p_rr:
break
throughput /= p_rr
# A tiny ambient term to avoid dead-black crevices.
# radiance += throughput * ti.Vector([0.01, 0.01, 0.015])
return radiance
@ti.func
def is_finite(x):
return (not ti.math.isinf(x)) and (not ti.math.isnan(x))
@ti.kernel
def render_kernel(view_idx: ti.i32, frame_seed: ti.i32, spp: ti.i32):
w = width_f[None]
h = height_f[None]
fov = fov_y_deg_f[None] * PI / 180.0
aspect = ti.cast(w, ti.f32) / ti.cast(h, ti.f32)
tan_half = ti.tan(0.5 * fov)
cam_pos = camera_pos[None]
cam_right = camera_right[None]
cam_up = camera_up[None]
cam_forward = camera_forward[None]
for y, x in img:
color = ti.Vector([0.0, 0.0, 0.0])
total = 0.0
for s in range(spp):
seed = ti.u32(frame_seed * 1315423911 + view_idx * ti.u32(2654435761) + y * 92821 + x * 68917 + s * 97531)
# Jittered pixel center.
j = rand2(seed)
px = (ti.cast(x, ti.f32) + j.x) / ti.cast(w, ti.f32)
py = (ti.cast(y, ti.f32) + j.y) / ti.cast(h, ti.f32)
ndc_x = (2.0 * px - 1.0) * aspect * tan_half
ndc_y = (1.0 - 2.0 * py) * tan_half
rd = (cam_forward + cam_right * ndc_x + cam_up * ndc_y).normalized()
ro = cam_pos
color_d = render_path(ro, rd, seed)
if is_finite(color_d.norm()):
color += color_d
total += 1.0
color /= ti.max(total, 1.0)
img[y, x] = color
@ti.kernel
def primary_hit_kernel(w: ti.i32, h: ti.i32):
fov = fov_y_deg_f[None] * PI / 180.0
aspect = ti.cast(w, ti.f32) / ti.cast(h, ti.f32)
tan_half = ti.tan(0.5 * fov)
cam_pos = camera_pos[None]
cam_right = camera_right[None]
cam_up = camera_up[None]
cam_forward = camera_forward[None]
for y, x in hit_map:
px = (ti.cast(x, ti.f32) + 0.5) / ti.cast(w, ti.f32)
py = (ti.cast(y, ti.f32) + 0.5) / ti.cast(h, ti.f32)
ndc_x = (2.0 * px - 1.0) * aspect * tan_half
ndc_y = (1.0 - 2.0 * py) * tan_half
rd = (cam_forward + cam_right * ndc_x + cam_up * ndc_y).normalized()
hit = trace_scene(cam_pos, rd, ti.u32(y * 92821 + x * 68917 + 1))
hit_map[y, x] = ti.Vector([hit[0], hit[2], hit[3], hit[4]])
def configure_scene(args) -> None:
width_f[None] = args.width
height_f[None] = args.height
fov_y_deg_f[None] = args.fov_y_deg
camera_radius_f[None] = args.camera_radius
camera_elev_deg_f[None] = args.camera_elevation_deg
camera_target[None] = ti.Vector([args.target_x, args.target_y, args.target_z])
def set_camera_from_pose(c2w: np.ndarray) -> None:
camera_pos[None] = ti.Vector(c2w[:3, 3].astype(np.float32))
camera_right[None] = ti.Vector(c2w[:3, 0].astype(np.float32))
camera_up[None] = ti.Vector(c2w[:3, 1].astype(np.float32))
# Camera forward in world space is -Z axis of c2w
camera_forward[None] = ti.Vector((-c2w[:3, 2]).astype(np.float32))
def generate_orbit_pose(view_idx: int, num_views: int, args) -> Tuple[np.ndarray, np.ndarray]:
# --- Fibonacci sphere formula ---
i = view_idx
N = num_views
# y coordinate goes from +1 to -1
y = 1.0 - 2.0 * (i + 0.5) / N
radius_xy = math.sqrt(max(0.0, 1.0 - y * y))
golden_angle = math.pi * (3.0 - math.sqrt(5.0))
theta = golden_angle * i
x = radius_xy * math.cos(theta)
z = radius_xy * math.sin(theta)
# --- Camera radius ---
r = args.camera_radius
# Direction = normalized Fibonacci point
direction = np.array([x, y, z], dtype=np.float64)
# --- Camera target ---
target = np.array([args.target_x, args.target_y, args.target_z], dtype=np.float64)
# --- Compute camera position ---
eye = target + r * direction
# --- Up vector (can adjust if needed based on hemisphere) ---
up = np.array([0.0, 1.0, 0.0], dtype=np.float64)
# --- Build transforms ---
c2w = look_at_c2w(eye, target, up)
w2c = c2w_to_w2c(c2w)
return c2w, w2c
def save_image(path: Path, linear_rgb: np.ndarray) -> np.ndarray:
srgb = srgb_encode(linear_rgb)
img8 = (np.clip(srgb, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
Image.fromarray(img8).save(path)
return img8
def main():
parser = argparse.ArgumentParser(description="Render a Mandelbulb dataset in COLMAP format.")
parser.add_argument("--out_dir", type=str, default="./mandelbulb_colmap")
parser.add_argument("--width", type=int, default=3840)
parser.add_argument("--height", type=int, default=3840)
parser.add_argument("--views", type=int, default=200)
parser.add_argument("--fov_y_deg", type=float, default=50.0)
parser.add_argument("--spp", type=int, default=256)
parser.add_argument("--camera_radius", type=float, default=5.0)
parser.add_argument("--camera_elevation_deg", type=float, default=20.0)
parser.add_argument("--camera_azimuth_offset_deg", type=float, default=0.0)
parser.add_argument("--target_x", type=float, default=0.0)
parser.add_argument("--target_y", type=float, default=0.0)
parser.add_argument("--target_z", type=float, default=0.0)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--num_points", type=int, default=100000,
help="target size of the initial point cloud")
parser.add_argument("--point_trace_res", type=int, default=800,
help="max side of the ray grid used to sample surface points")
parser.add_argument("--skip_render", action="store_true",
help="reuse existing images/ (pass the same camera args as the "
"original render) and only rebuild sparse/0")
args = parser.parse_args()
np.random.seed(args.seed)
rng = np.random.default_rng(args.seed)
configure_scene(args)
global img, hit_map
img = ti.Vector.field(3, dtype=ti.f32, shape=(args.height, args.width))
scale = min(1.0, args.point_trace_res / max(args.width, args.height))
trace_w = max(1, round(args.width * scale))
trace_h = max(1, round(args.height * scale))
hit_map = ti.Vector.field(4, dtype=ti.f32, shape=(trace_h, trace_w))
out_dir = Path(args.out_dir)
images_dir = out_dir / "images"
images_dir.mkdir(parents=True, exist_ok=True)
# COLMAP camera model: SIMPLE_PINHOLE with fx = fy = f, cx = w/2, cy = h/2.
# fov_y_deg is the vertical FOV, so f is derived from the image height.
f = 0.5 * args.height / math.tan(0.5 * math.radians(args.fov_y_deg))
camera_id = 1
cameras = [(camera_id, "SIMPLE_PINHOLE", args.width, args.height, [f, args.width * 0.5, args.height * 0.5])]
images_meta = []
pts_per_view = max(1, -(-args.num_points // args.views))
all_pts, all_rgb = [], []
max_reproj = 0.0
print(f"Rendering {args.views} views to {out_dir} ...")
for i in range(args.views):
c2w, _ = generate_orbit_pose(i, args.views, args)
set_camera_from_pose(c2w)
img_name = f"{i:06d}.jpg"
img_path = images_dir / img_name
if args.skip_render:
if not img_path.is_file():
raise SystemExit(f"--skip_render: missing {img_path}")
with Image.open(img_path) as im:
img8 = np.asarray(im.convert("RGB"))
else:
render_kernel(i, args.seed, args.spp)
img8 = save_image(img_path, img.to_numpy())
# Sample ground-truth surface points for the initial point cloud.
primary_hit_kernel(trace_w, trace_h)
pts, rgb, pix = sample_surface_points(hit_map.to_numpy(), img8,
pts_per_view, rng)
all_pts.append(pts)
all_rgb.append(rgb)
R, t = c2w_gl_to_colmap_w2c(c2w)
q = rotmat_to_quat(R)
images_meta.append((i + 1, q, t, camera_id, img_name))
if len(pts):
err = reprojection_errors(pts, pix, R, t, f, f,
args.width * 0.5, args.height * 0.5)
max_reproj = max(max_reproj, float(err.max()))
print(f" wrote {img_name} ({len(pts)} surface points)")
points_xyz = np.concatenate(all_pts) if all_pts else np.zeros((0, 3))
points_rgb = np.concatenate(all_rgb) if all_rgb else np.zeros((0, 3), np.uint8)
write_colmap_text(out_dir, cameras, images_meta, points_xyz, points_rgb)
print(f"Wrote sparse/0: {len(images_meta)} images, {len(points_xyz)} points")
print(f"Max reprojection error of sampled surface points: {max_reproj:.3f} px")
if max_reproj > 2.0:
print("WARNING: large reprojection error; exported poses may be "
"inconsistent with the rendered images.")
print("Done.")
if __name__ == "__main__":
main()
Raw
render_mandeltorus.py
#!/usr/bin/env python3
import argparse
import math
import os
from pathlib import Path
from typing import Tuple
import numpy as np
from PIL import Image
from colmap_export import (rotmat_to_quat, c2w_gl_to_colmap_w2c,
sample_surface_points, reprojection_errors,
write_colmap_text)
# -----------------------------
# Utilities
# -----------------------------
def clamp(x, lo, hi):
return max(lo, min(hi, x))
def normalize_np(v: np.ndarray, eps: float = 1e-12) -> np.ndarray:
n = np.linalg.norm(v)
if n < eps:
return v
return v / n
def look_at_c2w(eye: np.ndarray, target: np.ndarray, up: np.ndarray) -> np.ndarray:
"""Return camera-to-world 4x4 matrix, right-handed, camera looks along -Z."""
forward = normalize_np(target - eye)
right = normalize_np(np.cross(forward, up))
true_up = np.cross(right, forward)
c2w = np.eye(4, dtype=np.float64)
# Camera axes in world space:
# +X = right, +Y = true_up, +Z = -forward (because camera looks down -Z)
c2w[:3, 0] = right
c2w[:3, 1] = true_up
c2w[:3, 2] = -forward
c2w[:3, 3] = eye
return c2w
def c2w_to_w2c(c2w: np.ndarray) -> np.ndarray:
return np.linalg.inv(c2w)
def srgb_encode(x: np.ndarray) -> np.ndarray:
x = np.clip(x, 0.0, 1.0)
a = 0.055
out = np.where(x <= 0.0031308, 12.92 * x, (1.0 + a) * np.power(x, 1.0 / 2.4) - a)
return np.clip(out, 0.0, 1.0)
# -----------------------------
# Taichi setup
# -----------------------------
import taichi as ti
try:
ti.init(arch=ti.cuda, default_fp=ti.f32)
except Exception:
ti.init(arch=ti.gpu, default_fp=ti.f32)
except Exception:
ti.init(arch=ti.cpu, default_fp=ti.f32)
vec3 = ti.types.vector(3, ti.f32)
# Global config fields (shape=())
width_f = ti.field(dtype=ti.i32, shape=())
height_f = ti.field(dtype=ti.i32, shape=())
fov_y_deg_f = ti.field(dtype=ti.f32, shape=())
camera_radius_f = ti.field(dtype=ti.f32, shape=())
camera_elev_deg_f = ti.field(dtype=ti.f32, shape=())
camera_target = ti.Vector.field(3, dtype=ti.f32, shape=())
camera_pos = ti.Vector.field(3, dtype=ti.f32, shape=())
camera_right = ti.Vector.field(3, dtype=ti.f32, shape=())
camera_up = ti.Vector.field(3, dtype=ti.f32, shape=())
camera_forward = ti.Vector.field(3, dtype=ti.f32, shape=())
img = ti.Vector.field(3, dtype=ti.f32, shape=())
# Per-pixel primary-ray hits: [hit_flag, x, y, z]. Traced at a reduced
# resolution to sample the initial point cloud; reallocated in main().
hit_map = ti.Vector.field(4, dtype=ti.f32, shape=())
# Small constants
PI = 3.1415926535897932384626433832795
@ti.func
def colormap_turbo(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(-266.03287948, -63.82163439, -97.24397473) # k=7
c = c * t + vec3(874.85901369, 202.98596195, 447.99287080) # k=6
c = c * t + vec3(-1061.31232675, -245.26823636, -766.82678386) # k=5
c = c * t + vec3(550.44382083, 149.40813531, 604.07329733) # k=4
c = c * t + vec3(-90.46877071, -55.09180383, -203.76964931) # k=3
c = c * t + vec3(-10.15061040, 9.64581379, 9.26380342) # k=2
c = c * t + vec3(2.94711014, 2.06520532, 6.33792818) # k=1
c = c * t + vec3(0.14637796, 0.08594198, 0.23431523) # k=0
return c
@ti.func
def colormap_viridis(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(2.89151588, -0.17213704, 0.17750210) # k=3
c = c * t + vec3(-2.14339482, -0.16754949, -1.71784815) # k=2
c = c * t + vec3(0.04824199, 1.23905227, 1.23133794) # k=1
c = c * t + vec3(0.29387218, 0.01752027, 0.34360395) # k=0
return c
@ti.func
def colormap_magma(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(1.69082060, -2.36381252, 0.75796302) # k=4
c = c * t + vec3(-5.49016037, 5.48987916, 3.89534203) # k=3
c = c * t + vec3(4.30604404, -2.93736985, -7.52578597) # k=2
c = c * t + vec3(0.45897387, 0.82685306, 3.73227094) # k=1
c = c * t + vec3(-0.00157294, -0.01042786, -0.05490989) # k=0
return c
@ti.func
def colormap_inferno(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(1.66959821, -2.07978236, 7.14515755) # k=4
c = c * t + vec3(-5.11462265, 4.46548484, -6.74156149) # k=3
c = c * t + vec3(3.64171238, -1.86975118, -2.30104808) # k=2
c = c * t + vec3(0.74909880, 0.51010061, 2.59796038) # k=1
c = c * t + vec3(-0.01087222, -0.00208239, 0.00011545) # k=0
return c
@ti.func
def colormap_cividis(t: ti.f32) -> vec3:
c = vec3(0.0, 0.0, 0.0)
c = c * t + vec3(0.59905315, 0.14530728, -0.89703048) # k=3
c = c * t + vec3(-0.95976947, -0.07321042, 0.66155846) # k=2
c = c * t + vec3(1.43486697, 0.69925032, 0.05514949) # k=1
c = c * t + vec3(-0.05784352, 0.13417334, 0.38674306) # k=0
return c
@ti.func
def saturate(x):
return ti.min(ti.max(x, 0.0), 1.0)
@ti.func
def hypot(x, y):
return ti.sqrt(x * x + y * y)
@ti.func
def hash_u32(x):
x ^= x >> 16
x *= ti.u32(0x7feb352d)
x ^= x >> 15
x *= ti.u32(0x846ca68b)
x ^= x >> 16
return x
@ti.func
def rand01(seed):
seed = hash_u32(seed)
return ti.cast(seed & 0x00FFFFFF, ti.f32) / ti.cast(0x01000000, ti.f32)
@ti.func
def rand2(seed):
return ti.Vector([rand01(seed), rand01(seed ^ ti.u32(0x9e3779b9))])
@ti.func
def mix(a, b, t):
return a * (1.0 - t) + b * t
SURFACE_EPS = 1e-4
BOUNDING_RADIUS = 2.0
@ti.func
def mandeltorus_implicit(p: vec3) -> ti.types.vector(2, ti.f32):
p *= 2.0
z = p
dr = 1.0
r = 0.0
min_radius = 1e9
k = 6.0
m = 6.0
n = 6.0
R = 1.5
# k = 4.0
# m = 12.0
# n = 8.0
# R = 1.5
niter = 24
for i in range(niter):
r = z.norm()
min_radius = ti.min(min_radius, r)
if r > 64.0:
niter = i
break
r = ti.max(r, 1e-8)
r0 = hypot(hypot(z.x, z.y) - R, z.z)
rk = r0 ** k
ma = m * ti.atan2(z.y, z.x)
nb = n * ti.atan2(hypot(z.x, z.y) - R, z.z)
z1 = vec3([
ti.cos(ma) * (R + rk * ti.sin(nb)),
ti.sin(ma) * (R + rk * ti.sin(nb)),
rk * ti.cos(nb)
]) + p
rk = z1.norm()
z = z1
dr = ti.pow(r0, k - 1.0) * k * dr + 1.0 # approximation
r = z.norm()
r = ti.max(r, 1e-8)
de = 0.5 * ti.log(r) * r / dr
# de = ti.log(r + 1.0) / (0.375 * niter) - R * (ti.exp(-k) + k - 1.0) / (1.0 + ti.exp(-k))
#de = ti.abs(de)
return ti.Vector([de, min_radius])
@ti.func
def sphere_clip(ro: vec3, rd: vec3, radius: ti.f32) -> ti.types.vector(3, ti.f32):
b = ro.dot(rd)
c = ro.dot(ro) - radius * radius
disc = b * b - c
result = ti.Vector([0.0, 0.0, 0.0])
if disc >= 0.0:
sqrt_disc = ti.sqrt(disc)
result = ti.Vector([1.0, -b - sqrt_disc, -b + sqrt_disc])
return result
@ti.func
def estimate_normal(p: vec3) -> vec3:
e = 5e-5
ex = vec3([e, 0.0, 0.0])
ey = vec3([0.0, e, 0.0])
ez = vec3([0.0, 0.0, e])
dx1 = mandeltorus_implicit(p + ex)[0]
dx0 = mandeltorus_implicit(p - ex)[0]
dy1 = mandeltorus_implicit(p + ey)[0]
dy0 = mandeltorus_implicit(p - ey)[0]
dz1 = mandeltorus_implicit(p + ez)[0]
dz0 = mandeltorus_implicit(p - ez)[0]
n = vec3([dx1 - dx0, dy1 - dy0, dz1 - dz0])
return n / e
@ti.func
def random_hemisphere_cosine(n: vec3, seed: ti.u32) -> vec3:
r = rand2(seed)
phi = 2.0 * PI * r.x
cos_theta = ti.sqrt(1.0 - r.y)
sin_theta = ti.sqrt(r.y)
local = vec3([ti.cos(phi) * sin_theta, ti.sin(phi) * sin_theta, cos_theta])
# Build orthonormal basis
tangent = ti.Vector([0.0, 1.0, 0.0])
if ti.abs(n.y) > 0.999:
tangent = ti.Vector([1.0, 0.0, 0.0])
tangent = (tangent - n * tangent.dot(n)).normalized()
bitangent = n.cross(tangent)
return (tangent * local.x + bitangent * local.y + n * local.z).normalized()
@ti.func
def sky_radiance(rd: vec3) -> vec3:
# t = 0.5 * (rd.y + 1.0)
t = max(rd.normalized().y, 0.0)
sky_bottom = ti.Vector([0.7, 0.75, 0.8]) ** 2.2
sky_top = ti.Vector([1.2, 1.4, 1.8]) ** 2.2
col = mix(sky_bottom, sky_top, t)
return col
@ti.func
def trace_scene(ro: vec3, rd: vec3, seed: ti.u32) -> ti.types.vector(5, ti.f32):
# based on https://raw.githubusercontent.com/harry7557558/spirulae/refs/heads/master/implicit3-rt/frag-render.glsl
rd_n = rd.normalized()
result = ti.Vector([0.0, 0.0, ro.x, ro.y, ro.z])
clip = sphere_clip(ro, rd_n, BOUNDING_RADIUS)
t0c = ti.max(clip[1], 0.0)
t1c = clip[2]
if clip[0] >= 0.5 and t1c >= t0c:
step_sz = 500.0*SURFACE_EPS
t = t0c
dt = step_sz
v = 0.0
v0 = 0.0
v00 = 0.0
v1 = 0.0
dt0 = 0.0
dt00 = 0.0
is_bisecting = 0
prev_t = t0c
t_hi = 0.0
tfar = t1c
found = 0
result_t = 0.0
for iteration in range(16384):
if found == 1:
break
if t > tfar:
break
val_mr = mandeltorus_implicit(ro + rd_n * t) - 20.0*SURFACE_EPS
v = val_mr[0]
if is_bisecting == 1:
if t_hi - prev_t <= 1e-5:
dv = v1 - v0
if ti.abs(dv) > 1e-12:
result_t = prev_t - v0 / dv * (t_hi - prev_t)
else:
result_t = 0.5 * (prev_t + t_hi)
found = 1
else:
if v * v0 < 0.0:
t_hi = t
v1 = v
else:
prev_t = t
v0 = v
t = 0.5 * (prev_t + t_hi)
elif v * v0 < 0.0:
is_bisecting = 1
t_hi = t
v1 = v
t = 0.5 * (prev_t + t_hi)
else:
if ti.math.isnan(dt0) or dt0 <= 0.0:
v00 = v
v0 = v
dt0 = 0.0
dt00 = 0.0
g = 0.0
if dt0 > 0.0:
if dt00 > 0.0:
g = (v00 * dt0 / (dt00 * (dt0 + dt00))
- v0 * (dt0 + dt00) / (dt0 * dt00)
+ v * (2.0 * dt0 + dt00) / (dt0 * (dt0 + dt00)))
else:
g = (v - v0) / dt0
dt00 = dt0
dt0 = dt
prev_t = t
v00 = v0
v0 = v
ddt = ti.abs(v / g)
if ddt > step_sz or ti.math.isnan(ddt) or ti.math.isinf(ddt):
dt = step_sz
else:
dt_cand = ti.min(ddt - step_sz, tfar - prev_t - 0.01 * step_sz)
dt = ti.max(0.05 * step_sz, ti.min(dt_cand, step_sz))
if iteration < 2:
dt *= rand01(seed + iteration)
t += dt
if found == 1:
p = ro + rd_n * result_t
result = ti.Vector([1.0, result_t, p.x, p.y, p.z])
return result
@ti.func
def direct_light(n: vec3, p: vec3, view_dir: vec3, seed: ti.u32) -> vec3:
# One sun light with a hard shadow ray (cheap, but useful for contrast).
ldir = -ti.Vector([-0.4, 0.8, -0.2]).normalized()
ndotl = saturate(n.dot(ldir))
return_val = ti.Vector([0.0, 0.0, 0.0])
if ndotl > 0.0:
# Shadow ray
shadow_ro = p + n * (SURFACE_EPS * 4.0)
shadow_hit = trace_scene(shadow_ro, ldir, seed)
visible = 1.0 if shadow_hit[0] == 0.0 else 0.0
# Simple specular-like glint to help structure read well.
h = (ldir + view_dir).normalized()
spec = ti.pow(saturate(n.dot(h)), 96.0) * 0.15
sun_color = 1.0 * ti.Vector([1.0, 0.95, 0.9])
return_val = visible * (ndotl + spec) * sun_color
return return_val
@ti.func
def shade_background(rd: vec3, throughput: vec3, seed: ti.u32) -> vec3:
# Add a subtle fog / aerial perspective based on ray direction.
c = sky_radiance(rd)
return throughput * c
@ti.func
def render_path(ro: vec3, rd: vec3, px_seed: ti.u32) -> vec3:
throughput = ti.Vector([1.0, 1.0, 1.0])
radiance = ti.Vector([0.0, 0.0, 0.0])
cur_ro = ro
cur_rd = rd
for bounce in range(4):
hit = trace_scene(cur_ro, cur_rd, px_seed + bounce * 131 + 17)
if hit[0] == 0.0:
if bounce > 0:
radiance += shade_background(cur_rd, throughput, px_seed + bounce * 131)
break
t = hit[1]
p = vec3([hit[2], hit[3], hit[4]])
n = estimate_normal(p)
n_norm = n.norm()
n = n.normalized()
# Material: mostly diffuse with a tiny amount of view-dependent sheen.
# albedo = (0.5 + 0.35 * ti.sin(1.0 * PI * n) + 0.1 * ti.sin(4.0 * PI * n))
albedo = colormap_turbo(0.5 + 0.5 * ti.cos(ti.log(1.0*n_norm)))
# albedo = ti.max(1.0 - 1.5 * (1.0 - albedo), 0.0)
# albedo = albedo ** 2.0
# albedo = albedo + (ti.Vector([0.1, 0.0, 0.2]) - albedo) * (1.0 - ti.exp(-0.01 * hit[0]))
# albedo = albedo + (ti.Vector([0.0, 0.0, 0.0]) - albedo) * (1.0 - ti.exp(-0.1 / n_norm))
albedo = 0.9 * albedo ** 2.2
radiance += throughput * albedo * direct_light(n, p, -cur_rd, px_seed + bounce * 977 + 123)
new_dir = random_hemisphere_cosine(n, px_seed + bounce * 977 + 17)
cur_ro = p + n * (SURFACE_EPS * 8.0)
cur_rd = new_dir
throughput *= albedo
# Russian roulette
if bounce >= 2:
p_rr = ti.min(0.95, ti.max(throughput.x, ti.max(throughput.y, throughput.z)))
if rand01(px_seed + bounce * 31337 + 123) > p_rr:
break
throughput /= p_rr
# A tiny ambient term to avoid dead-black crevices.
# radiance += throughput * ti.Vector([0.01, 0.01, 0.015])
return radiance
@ti.func
def is_finite(x):
return (not ti.math.isinf(x)) and (not ti.math.isnan(x))
@ti.kernel
def render_kernel(view_idx: ti.i32, frame_seed: ti.i32, spp: ti.i32):
w = width_f[None]
h = height_f[None]
fov = fov_y_deg_f[None] * PI / 180.0
aspect = ti.cast(w, ti.f32) / ti.cast(h, ti.f32)
tan_half = ti.tan(0.5 * fov)
cam_pos = camera_pos[None]
cam_right = camera_right[None]
cam_up = camera_up[None]
cam_forward = camera_forward[None]
for y, x in img:
color = ti.Vector([0.0, 0.0, 0.0])
total = 0.0
for s in range(spp):
seed = ti.u32(frame_seed * 1315423911 + view_idx * ti.u32(2654435761) + y * 92821 + x * 68917 + s * 97531)
# Jittered pixel center.
j = rand2(seed)
px = (ti.cast(x, ti.f32) + j.x) / ti.cast(w, ti.f32)
py = (ti.cast(y, ti.f32) + j.y) / ti.cast(h, ti.f32)
ndc_x = (2.0 * px - 1.0) * aspect * tan_half
ndc_y = (1.0 - 2.0 * py) * tan_half
rd = (cam_forward + cam_right * ndc_x + cam_up * ndc_y).normalized()
ro = cam_pos
color_d = render_path(ro, rd, seed)
if is_finite(color_d.norm()):
color += color_d
total += 1.0
color /= ti.max(total, 1.0)
img[y, x] = color
@ti.kernel
def primary_hit_kernel(w: ti.i32, h: ti.i32):
fov = fov_y_deg_f[None] * PI / 180.0
aspect = ti.cast(w, ti.f32) / ti.cast(h, ti.f32)
tan_half = ti.tan(0.5 * fov)
cam_pos = camera_pos[None]
cam_right = camera_right[None]
cam_up = camera_up[None]
cam_forward = camera_forward[None]
for y, x in hit_map:
px = (ti.cast(x, ti.f32) + 0.5) / ti.cast(w, ti.f32)
py = (ti.cast(y, ti.f32) + 0.5) / ti.cast(h, ti.f32)
ndc_x = (2.0 * px - 1.0) * aspect * tan_half
ndc_y = (1.0 - 2.0 * py) * tan_half
rd = (cam_forward + cam_right * ndc_x + cam_up * ndc_y).normalized()
hit = trace_scene(cam_pos, rd, ti.u32(y * 92821 + x * 68917 + 1))
hit_map[y, x] = ti.Vector([hit[0], hit[2], hit[3], hit[4]])
def configure_scene(args) -> None:
width_f[None] = args.width
height_f[None] = args.height
fov_y_deg_f[None] = args.fov_y_deg
camera_radius_f[None] = args.camera_radius
camera_elev_deg_f[None] = args.camera_elevation_deg
camera_target[None] = ti.Vector([args.target_x, args.target_y, args.target_z])
def set_camera_from_pose(c2w: np.ndarray) -> None:
camera_pos[None] = ti.Vector(c2w[:3, 3].astype(np.float32))
camera_right[None] = ti.Vector(c2w[:3, 0].astype(np.float32))
camera_up[None] = ti.Vector(c2w[:3, 1].astype(np.float32))
# Camera forward in world space is -Z axis of c2w
camera_forward[None] = ti.Vector((-c2w[:3, 2]).astype(np.float32))
def generate_orbit_pose(view_idx: int, num_views: int, args) -> Tuple[np.ndarray, np.ndarray]:
# --- Fibonacci sphere formula ---
i = view_idx
N = num_views
# y coordinate goes from +1 to -1
y = 1.0 - 2.0 * (i + 0.5) / N
radius_xy = math.sqrt(max(0.0, 1.0 - y * y))
golden_angle = math.pi * (3.0 - math.sqrt(5.0))
theta = golden_angle * i
x = radius_xy * math.cos(theta)
z = radius_xy * math.sin(theta)
# --- Camera radius ---
r = args.camera_radius
# Direction = normalized Fibonacci point
direction = np.array([x, y, z], dtype=np.float64)
# --- Camera target ---
target = np.array([args.target_x, args.target_y, args.target_z], dtype=np.float64)
# --- Compute camera position ---
eye = target + r * direction
# --- Up vector (can adjust if needed based on hemisphere) ---
up = np.array([0.0, 1.0, 0.0], dtype=np.float64)
# --- Build transforms ---
c2w = look_at_c2w(eye, target, up)
w2c = c2w_to_w2c(c2w)
return c2w, w2c
def save_image(path: Path, linear_rgb: np.ndarray) -> np.ndarray:
srgb = srgb_encode(linear_rgb)
img8 = (np.clip(srgb, 0.0, 1.0) * 255.0 + 0.5).astype(np.uint8)
Image.fromarray(img8).save(path)
return img8
def main():
parser = argparse.ArgumentParser(description="Render a 3D fractal.")
parser.add_argument("--out_dir", type=str, default="./mandeltorus_colmap")
parser.add_argument("--width", type=int, default=3840)
parser.add_argument("--height", type=int, default=3840)
parser.add_argument("--views", type=int, default=200)
parser.add_argument("--fov_y_deg", type=float, default=50.0)
parser.add_argument("--spp", type=int, default=256)
parser.add_argument("--camera_radius", type=float, default=5.0)
parser.add_argument("--camera_elevation_deg", type=float, default=20.0)
parser.add_argument("--camera_azimuth_offset_deg", type=float, default=0.0)
parser.add_argument("--target_x", type=float, default=0.0)
parser.add_argument("--target_y", type=float, default=0.0)
parser.add_argument("--target_z", type=float, default=0.0)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--num_points", type=int, default=100000,
help="target size of the initial point cloud")
parser.add_argument("--point_trace_res", type=int, default=800,
help="max side of the ray grid used to sample surface points")
parser.add_argument("--skip_render", action="store_true",
help="reuse existing images/ (pass the same camera args as the "
"original render) and only rebuild sparse/0")
args = parser.parse_args()
np.random.seed(args.seed)
rng = np.random.default_rng(args.seed)
configure_scene(args)
global img, hit_map
img = ti.Vector.field(3, dtype=ti.f32, shape=(args.height, args.width))
scale = min(1.0, args.point_trace_res / max(args.width, args.height))
trace_w = max(1, round(args.width * scale))
trace_h = max(1, round(args.height * scale))
hit_map = ti.Vector.field(4, dtype=ti.f32, shape=(trace_h, trace_w))
out_dir = Path(args.out_dir)
images_dir = out_dir / "images"
images_dir.mkdir(parents=True, exist_ok=True)
# COLMAP camera model: SIMPLE_PINHOLE with fx = fy = f, cx = w/2, cy = h/2.
# fov_y_deg is the vertical FOV, so f is derived from the image height.
f = 0.5 * args.height / math.tan(0.5 * math.radians(args.fov_y_deg))
camera_id = 1
cameras = [(camera_id, "SIMPLE_PINHOLE", args.width, args.height, [f, args.width * 0.5, args.height * 0.5])]
images_meta = []
pts_per_view = max(1, -(-args.num_points // args.views))
all_pts, all_rgb = [], []
max_reproj = 0.0
print(f"Rendering {args.views} views to {out_dir} ...")
for i in range(args.views):
c2w, _ = generate_orbit_pose(i, args.views, args)
set_camera_from_pose(c2w)
img_name = f"{i:06d}.jpg"
img_path = images_dir / img_name
if args.skip_render:
if not img_path.is_file():
raise SystemExit(f"--skip_render: missing {img_path}")
with Image.open(img_path) as im:
img8 = np.asarray(im.convert("RGB"))
else:
render_kernel(i, args.seed, args.spp)
img8 = save_image(img_path, img.to_numpy())
# Sample ground-truth surface points for the initial point cloud.
primary_hit_kernel(trace_w, trace_h)
pts, rgb, pix = sample_surface_points(hit_map.to_numpy(), img8,
pts_per_view, rng)
all_pts.append(pts)
all_rgb.append(rgb)
R, t = c2w_gl_to_colmap_w2c(c2w)
q = rotmat_to_quat(R)
images_meta.append((i + 1, q, t, camera_id, img_name))
if len(pts):
err = reprojection_errors(pts, pix, R, t, f, f,
args.width * 0.5, args.height * 0.5)
max_reproj = max(max_reproj, float(err.max()))
print(f" wrote {img_name} ({len(pts)} surface points)")
points_xyz = np.concatenate(all_pts) if all_pts else np.zeros((0, 3))
points_rgb = np.concatenate(all_rgb) if all_rgb else np.zeros((0, 3), np.uint8)
write_colmap_text(out_dir, cameras, images_meta, points_xyz, points_rgb)
print(f"Wrote sparse/0: {len(images_meta)} images, {len(points_xyz)} points")
print(f"Max reprojection error of sampled surface points: {max_reproj:.3f} px")
if max_reproj > 2.0:
print("WARNING: large reprojection error; exported poses may be "
"inconsistent with the rendered images.")
print("Done.")
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment