- added renderer

- Added undo
This commit is contained in:
bklronin
2026-07-12 22:21:43 +02:00
parent 210e3cfb5d
commit 9f1387fe68
8 changed files with 3419 additions and 0 deletions
+345
View File
@@ -0,0 +1,345 @@
"""Mitsuba 3 photorealistic render backend.
Requires: ``pip install mitsuba``
"""
from __future__ import annotations
import logging
import os
import tempfile
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 _build_scene_dict(
self,
mesh_path: str,
material: RenderMaterial,
camera: RenderCamera,
settings: RenderSettings,
) -> dict:
"""Build a Mitsuba scene dictionary from our data classes.
Uses a 3-point lighting setup (key + fill + rim) plus an ambient
environment emitter for soft fill, giving well-balanced shading on
all faces of the model. Lighting intensities and colors come from
``settings.lighting``; ground plane comes from ``settings.ground_plane``.
"""
import mitsuba as mi
lighting = settings.lighting
ground = settings.ground_plane
# Map our BSDF types to Mitsuba BSDF dicts
bsdf = self._make_bsdf(material)
# Determine mesh file type from extension
ext = os.path.splitext(mesh_path)[1].lower()
shape_type = "ply" if ext == ".ply" else "obj"
# Build camera-to-world transform using the Python API
cam_to_world = mi.ScalarTransform4f.look_at(
origin=list(camera.origin),
target=list(camera.target),
up=list(camera.up),
)
scene = {
"type": "scene",
# Integrator
"integrator": {
"type": "path",
"max_depth": settings.max_depth,
},
# Camera
"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,
},
},
# Ambient environment fill
"emitter": {
"type": "constant",
"radiance": {
"type": "rgb",
"value": [
lighting.ambient_intensity,
lighting.ambient_intensity * 0.97,
lighting.ambient_intensity * 0.94,
],
},
},
# Shape
"shape": {
"type": shape_type,
"filename": mesh_path,
"bsdf": bsdf,
},
}
# ── 3-point lighting (colors and intensities from config) ───
key_rgb = [
c * lighting.key_intensity for c in lighting.key_color
]
key_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],
)
scene["key_light"] = {
"type": "directional",
"to_world": key_to_world,
"irradiance": {"type": "rgb", "value": key_rgb},
}
fill_rgb = [
c * lighting.fill_intensity for c in lighting.fill_color
]
fill_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],
)
scene["fill_light"] = {
"type": "directional",
"to_world": fill_to_world,
"irradiance": {"type": "rgb", "value": fill_rgb},
}
rim_rgb = [
c * lighting.rim_intensity for c in lighting.rim_color
]
rim_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],
)
scene["rim_light"] = {
"type": "directional",
"to_world": rim_to_world,
"irradiance": {"type": "rgb", "value": rim_rgb},
}
# ── Ground plane (optional) ─────────────────────────────────
if ground.enabled:
scene["ground_plane"] = {
"type": "rectangle",
"size": [500.0, 500.0],
"to_world": mi.ScalarTransform4f.translate(
[0.0, 0.0, -ground.distance_below]
)
@ mi.ScalarTransform4f.rotate_about_z(90),
"bsdf": {
"type": "diffuse",
"reflectance": {
"type": "rgb",
"value": list(ground.color),
},
},
}
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"
)
# Render
image = mi.render(scene, spp=settings.spp, seed=settings.seed or 0)
if progress_callback:
progress_callback(1.0)
# Convert to numpy (H, W, 3)
arr = np.array(image, dtype=np.float32)
# Apply approximate sRGB tonemapping
arr = np.clip(arr, 0.0, None)
arr = np.power(arr, 1.0 / 2.2) # gamma
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,
)
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
import Imath
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