bklronin
426 lines
15 KiB
Python
426 lines
15 KiB
Python
"""Mitsuba 3 photorealistic render backend.
|
|
|
|
Requires: ``pip install mitsuba``
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
import os
|
|
from typing import Callable, Optional
|
|
|
|
import numpy as np
|
|
|
|
from .render_backend import RenderBackend, RenderCamera, RenderMaterial, RenderSettings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class MitsubaBackend(RenderBackend):
|
|
"""Render backend using Mitsuba 3 path tracer."""
|
|
|
|
def name(self) -> str:
|
|
return "Mitsuba 3"
|
|
|
|
def is_available(self) -> bool:
|
|
import sys
|
|
import io
|
|
|
|
old_stderr = sys.stderr
|
|
sys.stderr = io.StringIO()
|
|
try:
|
|
import mitsuba # noqa: F401
|
|
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
finally:
|
|
sys.stderr = old_stderr
|
|
|
|
# ── Scene construction ──────────────────────────────────────────
|
|
|
|
def _base_scene_dict(
|
|
self,
|
|
camera: RenderCamera,
|
|
settings: RenderSettings,
|
|
first_mesh_path: Optional[str] = None,
|
|
) -> dict:
|
|
"""Return a scene dict with everything *except* the shape entries.
|
|
|
|
When *first_mesh_path* is given the ground plane / backdrop is sized
|
|
from its bounding box; otherwise a large default is used.
|
|
"""
|
|
import mitsuba as mi
|
|
|
|
lighting = settings.lighting
|
|
ground = settings.ground_plane
|
|
|
|
cam_to_world = mi.ScalarTransform4f.look_at(
|
|
origin=list(camera.origin),
|
|
target=list(camera.target),
|
|
up=list(camera.up),
|
|
)
|
|
|
|
scene: dict = {
|
|
"type": "scene",
|
|
"integrator": {"type": "path", "max_depth": settings.max_depth},
|
|
"sensor": {
|
|
"type": "perspective",
|
|
"fov": camera.fov,
|
|
"to_world": cam_to_world,
|
|
"film": {
|
|
"type": "hdrfilm",
|
|
"width": settings.width,
|
|
"height": settings.height,
|
|
"rfilter": {"type": "gaussian"},
|
|
},
|
|
"sampler": {
|
|
"type": "independent",
|
|
"sample_count": settings.spp,
|
|
},
|
|
},
|
|
"emitter": {
|
|
"type": "constant",
|
|
"radiance": {
|
|
"type": "rgb",
|
|
"value": [
|
|
lighting.ambient_intensity,
|
|
lighting.ambient_intensity * 0.97,
|
|
lighting.ambient_intensity * 0.94,
|
|
],
|
|
},
|
|
},
|
|
}
|
|
|
|
# ── 3-point lighting ──────────────────────────────────────────
|
|
key_rgb = [c * lighting.key_intensity for c in lighting.key_color]
|
|
scene["key_light"] = {
|
|
"type": "directional",
|
|
"to_world": mi.ScalarTransform4f.look_at(
|
|
origin=[1.0, -0.8, 1.2],
|
|
target=[0.0, 0.0, 0.0],
|
|
up=[0.0, 0.0, 1.0],
|
|
),
|
|
"irradiance": {"type": "rgb", "value": key_rgb},
|
|
}
|
|
|
|
fill_rgb = [c * lighting.fill_intensity for c in lighting.fill_color]
|
|
scene["fill_light"] = {
|
|
"type": "directional",
|
|
"to_world": mi.ScalarTransform4f.look_at(
|
|
origin=[-1.0, 0.6, 0.8],
|
|
target=[0.0, 0.0, 0.0],
|
|
up=[0.0, 0.0, 1.0],
|
|
),
|
|
"irradiance": {"type": "rgb", "value": fill_rgb},
|
|
}
|
|
|
|
rim_rgb = [c * lighting.rim_intensity for c in lighting.rim_color]
|
|
scene["rim_light"] = {
|
|
"type": "directional",
|
|
"to_world": mi.ScalarTransform4f.look_at(
|
|
origin=[-0.3, 1.2, -0.8],
|
|
target=[0.0, 0.0, 0.0],
|
|
up=[0.0, 0.0, 1.0],
|
|
),
|
|
"irradiance": {"type": "rgb", "value": rim_rgb},
|
|
}
|
|
|
|
# ── Ground plane / backdrop ───────────────────────────────────
|
|
if ground.enabled:
|
|
try:
|
|
ext = os.path.splitext(first_mesh_path)[1].lower() if first_mesh_path else ""
|
|
shape_type = "ply" if ext == ".ply" else "obj"
|
|
mesh_shape = mi.load_dict({"type": shape_type, "filename": first_mesh_path})
|
|
bbox = mesh_shape.bbox()
|
|
bbox_min, bbox_max = bbox[0], bbox[1]
|
|
model_height = bbox_max[2] - bbox_min[2]
|
|
ground_z = bbox_min[2] - 0.001 * model_height
|
|
dx = bbox_max[0] - bbox_min[0]
|
|
dy = bbox_max[1] - bbox_min[1]
|
|
dz = bbox_max[2] - bbox_min[2]
|
|
diag = float((dx * dx + dy * dy + dz * dz) ** 0.5)
|
|
except Exception:
|
|
ground_z = -ground.distance_below
|
|
diag = 1000.0
|
|
|
|
bsdf_ground = {
|
|
"type": "diffuse",
|
|
"reflectance": {"type": "rgb", "value": list(ground.color)},
|
|
}
|
|
|
|
if ground.curved_backdrop:
|
|
half_size = diag * 50.0
|
|
radius = diag * 3.0
|
|
cyl_height = diag * 20.0
|
|
scene["ground_floor"] = {
|
|
"type": "rectangle",
|
|
"to_world": mi.ScalarTransform4f.translate([0.0, 0.0, ground_z])
|
|
@ mi.ScalarTransform4f.scale([half_size, half_size, 1.0]),
|
|
"bsdf": bsdf_ground,
|
|
}
|
|
scene["ground_backdrop"] = {
|
|
"type": "cylinder",
|
|
"radius": radius,
|
|
"p0": [-cyl_height / 2, -radius, ground_z],
|
|
"p1": [cyl_height / 2, -radius, ground_z],
|
|
"to_world": mi.ScalarTransform4f.rotate([1, 0, 0], 90)
|
|
@ mi.ScalarTransform4f.translate([0.0, 0.0, -radius]),
|
|
"bsdf": bsdf_ground,
|
|
}
|
|
else:
|
|
half_size = diag * 50.0
|
|
scene["ground_plane"] = {
|
|
"type": "rectangle",
|
|
"to_world": mi.ScalarTransform4f.translate([0.0, 0.0, ground_z])
|
|
@ mi.ScalarTransform4f.scale([half_size, half_size, 1.0]),
|
|
"bsdf": bsdf_ground,
|
|
}
|
|
|
|
return scene
|
|
|
|
def _build_scene_dict(
|
|
self,
|
|
mesh_path: str,
|
|
material: RenderMaterial,
|
|
camera: RenderCamera,
|
|
settings: RenderSettings,
|
|
) -> dict:
|
|
"""Build a single-shape Mitsuba scene dictionary."""
|
|
ext = os.path.splitext(mesh_path)[1].lower()
|
|
shape_type = "ply" if ext == ".ply" else "obj"
|
|
|
|
scene = self._base_scene_dict(camera, settings, first_mesh_path=mesh_path)
|
|
scene["shape"] = {
|
|
"type": shape_type,
|
|
"filename": mesh_path,
|
|
"bsdf": self._make_bsdf(material),
|
|
}
|
|
return scene
|
|
|
|
def _build_assembly_scene_dict(
|
|
self,
|
|
parts: list,
|
|
camera: RenderCamera,
|
|
settings: RenderSettings,
|
|
) -> dict:
|
|
"""Build a multi-shape Mitsuba scene dictionary.
|
|
|
|
*parts* is a list of ``(mesh_path, RenderMaterial)`` tuples.
|
|
"""
|
|
first_path = parts[0][0] if parts else None
|
|
scene = self._base_scene_dict(camera, settings, first_mesh_path=first_path)
|
|
|
|
for i, (mesh_path, material) in enumerate(parts):
|
|
ext = os.path.splitext(mesh_path)[1].lower()
|
|
shape_type = "ply" if ext == ".ply" else "obj"
|
|
scene[f"shape_{i}"] = {
|
|
"type": shape_type,
|
|
"filename": mesh_path,
|
|
"bsdf": self._make_bsdf(material),
|
|
}
|
|
return scene
|
|
|
|
def _make_bsdf(self, material: RenderMaterial) -> dict:
|
|
"""Convert a RenderMaterial to a Mitsuba BSDF dict."""
|
|
mt = material.bsdf_type
|
|
|
|
if mt == "roughconductor":
|
|
# Use material="none" with specular_reflectance set to the
|
|
# metal color. The pip-installed Mitsuba doesn't include
|
|
# spectral metal data files (iron.spd, copper.spd, etc.).
|
|
return {
|
|
"type": "roughconductor",
|
|
"material": "none",
|
|
"alpha": max(material.roughness, 0.01),
|
|
"specular_reflectance": {
|
|
"type": "rgb",
|
|
"value": list(material.color),
|
|
},
|
|
}
|
|
|
|
if mt == "roughdielectric":
|
|
return {
|
|
"type": "roughdielectric",
|
|
"int_ior": material.int_ior,
|
|
"ext_ior": 1.0,
|
|
"alpha": max(material.roughness, 0.01),
|
|
}
|
|
|
|
if mt == "dielectric":
|
|
return {
|
|
"type": "dielectric",
|
|
"int_ior": material.int_ior,
|
|
"ext_ior": 1.0,
|
|
}
|
|
|
|
if mt == "plastic":
|
|
return {
|
|
"type": "plastic",
|
|
"diffuse_reflectance": {
|
|
"type": "rgb",
|
|
"value": list(material.color),
|
|
},
|
|
"int_ior": material.int_ior,
|
|
}
|
|
|
|
# Default: diffuse
|
|
return {
|
|
"type": "diffuse",
|
|
"reflectance": {
|
|
"type": "rgb",
|
|
"value": list(material.color),
|
|
},
|
|
}
|
|
|
|
# ── Rendering ───────────────────────────────────────────────────
|
|
|
|
def render(
|
|
self,
|
|
mesh_path: str,
|
|
material: RenderMaterial,
|
|
camera: RenderCamera,
|
|
settings: RenderSettings,
|
|
progress_callback: Optional[Callable[[float], None]] = None,
|
|
) -> np.ndarray:
|
|
"""Render a mesh file and return (H, W, 3) float32 RGB array."""
|
|
self._set_variant()
|
|
import mitsuba as mi
|
|
|
|
scene_dict = self._build_scene_dict(mesh_path, material, camera, settings)
|
|
scene = mi.load_dict(scene_dict)
|
|
|
|
logger.info(f"Rendering {settings.width}x{settings.height} @ {settings.spp} spp")
|
|
|
|
try:
|
|
image = mi.render(scene, spp=settings.spp, seed=int(settings.seed or 0))
|
|
except Exception as e:
|
|
logger.error(f"Mitsuba render failed: {e}")
|
|
raise
|
|
|
|
if progress_callback:
|
|
progress_callback(1.0)
|
|
|
|
arr = np.array(image, dtype=np.float32)
|
|
arr = np.clip(arr, 0.0, None)
|
|
arr = np.power(arr, 1.0 / 2.2)
|
|
arr = np.clip(arr, 0.0, 1.0)
|
|
return arr
|
|
|
|
def render_assembly(
|
|
self,
|
|
parts: list,
|
|
camera: RenderCamera,
|
|
settings: RenderSettings,
|
|
progress_callback: Optional[Callable[[float], None]] = None,
|
|
) -> np.ndarray:
|
|
"""Render multiple meshes with individual materials.
|
|
|
|
*parts* is a list of ``(mesh_path, RenderMaterial)`` tuples.
|
|
Returns (H, W, 3) float32 RGB array.
|
|
"""
|
|
self._set_variant()
|
|
import mitsuba as mi
|
|
|
|
scene_dict = self._build_assembly_scene_dict(parts, camera, settings)
|
|
scene = mi.load_dict(scene_dict)
|
|
|
|
logger.info(
|
|
f"Rendering assembly ({len(parts)} parts) "
|
|
f"{settings.width}x{settings.height} @ {settings.spp} spp"
|
|
)
|
|
|
|
try:
|
|
image = mi.render(scene, spp=settings.spp, seed=int(settings.seed or 0))
|
|
except Exception as e:
|
|
logger.error(f"Mitsuba assembly render failed: {e}")
|
|
raise
|
|
|
|
if progress_callback:
|
|
progress_callback(1.0)
|
|
|
|
arr = np.array(image, dtype=np.float32)
|
|
arr = np.clip(arr, 0.0, None)
|
|
arr = np.power(arr, 1.0 / 2.2)
|
|
arr = np.clip(arr, 0.0, 1.0)
|
|
return arr
|
|
|
|
def render_preview(
|
|
self,
|
|
mesh_path: str,
|
|
material: RenderMaterial,
|
|
camera: RenderCamera,
|
|
settings: RenderSettings,
|
|
) -> np.ndarray:
|
|
"""Quick low-quality preview (4x fewer spp)."""
|
|
preview_settings = RenderSettings(
|
|
width=settings.width // 2,
|
|
height=settings.height // 2,
|
|
spp=max(settings.spp // 4, 16),
|
|
max_depth=min(settings.max_depth, 4),
|
|
seed=settings.seed,
|
|
lighting=settings.lighting,
|
|
ground_plane=settings.ground_plane,
|
|
)
|
|
return self.render(mesh_path, material, camera, preview_settings)
|
|
|
|
# ── Export ──────────────────────────────────────────────────────
|
|
|
|
def export_image(self, image: np.ndarray, path: str) -> None:
|
|
"""Save a rendered image to PNG or EXR."""
|
|
from PIL import Image
|
|
|
|
ext = os.path.splitext(path)[1].lower()
|
|
|
|
if ext == ".exr":
|
|
# Save as EXR (HDR) — no tonemapping
|
|
try:
|
|
import OpenEXR # type: ignore[import-not-found]
|
|
import Imath # type: ignore[import-not-found]
|
|
|
|
h, w = image.shape[:2]
|
|
header = OpenEXR.Header(w, h)
|
|
header["channels"] = {
|
|
"R": Imath.PixelType(Imath.PixelType.FLOAT),
|
|
"G": Imath.PixelType(Imath.PixelType.FLOAT),
|
|
"B": Imath.PixelType(Imath.PixelType.FLOAT),
|
|
}
|
|
exr = OpenEXR.OutputFile(path, header)
|
|
exr.write(
|
|
{
|
|
"R": image[:, :, 0].tobytes(),
|
|
"G": image[:, :, 1].tobytes(),
|
|
"B": image[:, :, 2].tobytes(),
|
|
}
|
|
)
|
|
exr.close()
|
|
except ImportError:
|
|
# Fallback: save as 16-bit PNG
|
|
logger.warning("OpenEXR not available, saving as 16-bit PNG")
|
|
img = Image.fromarray((image * 65535).astype(np.uint16), "RGB")
|
|
img.save(path)
|
|
else:
|
|
# PNG / JPEG — already tonemapped
|
|
img = Image.fromarray((image * 255).astype(np.uint8), "RGB")
|
|
img.save(path)
|
|
|
|
logger.info(f"Exported render to {path}")
|
|
|
|
# ── Helpers ─────────────────────────────────────────────────────
|
|
|
|
def _set_variant(self) -> None:
|
|
"""Set the Mitsuba variant (called once)."""
|
|
import sys
|
|
import io
|
|
|
|
# Suppress the harmless "LLVM API initialization failed" warning
|
|
# that drjit emits on macOS ARM when scalar variant is used.
|
|
old_stderr = sys.stderr
|
|
sys.stderr = io.StringIO()
|
|
try:
|
|
import mitsuba as mi
|
|
|
|
mi.set_variant("scalar_rgb")
|
|
finally:
|
|
sys.stderr = old_stderr
|