- added renderer
- Added undo
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user