Blame for vli/scene.py

"""Configuration side: once the serial ledger is resolved, any placement that
satisfies its relations is valid. The seed picks one configuration.
"""
import math

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt  # noqa: E402
import numpy as np  # noqa: E402
import trimesh  # noqa: E402
from mpl_toolkits.mplot3d.art3d import Poly3DCollection  # noqa: E402
from scipy.spatial import ConvexHull  # noqa: E402

from .simplify import base_color, decimate  # noqa: E402

GAP = 0.1
GLOW = [255, 214, 107, 255]
RENDER_FACE_BUDGET = 3000


def _half_width(mesh):
    return 0.5 * max(mesh.extents[0], mesh.extents[2])


def sun_direction(sun):
    el, az = math.radians(sun["elevation_deg"]), math.radians(sun["azimuth_deg"])
    return np.array([math.cos(el) * math.sin(az), math.sin(el), math.cos(el) * math.cos(az)])


def configure(ledger, meshes, seed=0):
    """Place every object so the ledger's relations hold. Returns a JSON-able dict."""
    rng = np.random.default_rng(seed)
    items = ledger["items"]
    relations = [i for i in items if i["kind"] == "relation"]
    lights = [i for i in items if i.get("emits_light")]
    objects = [i for i in items if i["kind"] == "object" and not i.get("emits_light")]

    place, x = {}, 0.0
    for it in objects:  # unrelated objects stand in a loose row
        hw = _half_width(meshes[it["id"]])
        place[it["id"]] = {"position": [x + hw, 0.0, float(rng.uniform(-0.5, 0.5) * hw)],
                           "yaw_deg": float(rng.uniform(-180, 180))}
        x += 2 * hw + 4 * GAP
    for p in place.values():
        p["position"][0] -= x / 2

    if lights:
        sun = {"id": lights[0]["id"], "elevation_deg": float(rng.uniform(25, 65)),
               "azimuth_deg": float(rng.uniform(-150, 150)),
               "shadows": any(r["relation"] == "lights" for r in relations)}
    else:  # nothing shines: soft daylight from no object (the ledger asks "what light?")
        sun = {"id": None, "elevation_deg": float(rng.uniform(35, 60)),
               "azimuth_deg": float(rng.uniform(-150, 150)), "shadows": True, "daylight": True}

    def height(oid):
        return float(meshes[oid].extents[1])

    for _ in range(2):  # second pass lets chains (a on b on c) settle
        for r in relations:
            a, b, rel = r["subject"], r["object"], r["relation"]
            if a not in place:
                continue
            pa = place[a]["position"]
            if b not in place:  # environment object: ground, sun, scene
                if rel == "on":
                    pa[1] = 0.0
                elif rel == "faces" and sun and b == sun["id"]:
                    place[a]["yaw_deg"] = sun["azimuth_deg"]
                continue
            pb = place[b]["position"]
            hwa, hwb = _half_width(meshes[a]), _half_width(meshes[b])
            if rel == "on":
                slack = max(0.0, hwb - hwa)
                ang = rng.uniform(0, 2 * math.pi)
                pa[0], pa[2] = pb[0] + slack * math.cos(ang) * 0.5, pb[2] + slack * math.sin(ang) * 0.5
                pa[1] = pb[1] + height(b)
            elif rel == "under":
                pa[0], pa[1], pa[2] = pb[0], 0.0, pb[2]
            elif rel in ("beside", "near"):
                side = rng.choice([-1, 1])
                gap = GAP if rel == "beside" else 3 * GAP + hwa
                pa[0], pa[1], pa[2] = pb[0] + side * (hwa + hwb + gap), pb[1], pb[2]
            elif rel in ("behind", "in_front_of"):
                sign = -1 if rel == "behind" else 1
                pa[0], pa[1], pa[2] = pb[0], pb[1], pb[2] + sign * (hwa + hwb + GAP)
            elif rel == "above":
                pa[0], pa[2] = pb[0], pb[2]
                pa[1] = pb[1] + height(b) + 0.3 * max(height(a), height(b))
            elif rel == "faces":
                place[a]["yaw_deg"] = math.degrees(math.atan2(pb[0] - pa[0], pb[2] - pa[2]))

    if lights:  # whatever shone sits in the sky along the light direction
        reach = max([abs(c) + _half_width(meshes[o]) for o, p in place.items()
                     for c in (p["position"][0], p["position"][2])] + [0.3])
        for it in lights:
            # shadows use the direction; distance only frames the view
            pos = sun_direction(sun) * (1.5 * reach + 2 * _half_width(meshes[it["id"]]))
            place[it["id"]] = {"position": [float(c) for c in pos], "yaw_deg": 0.0, "light": True}
    return {"seed": seed, "placements": place, "sun": sun}


def transform(p):
    t = trimesh.transformations.rotation_matrix(math.radians(p["yaw_deg"]), [0, 1, 0])
    t[:3, 3] = p["position"]
    return t


def _ground(ledger, placed):
    if not any(i["kind"] == "environment" and i["role"] == "plane" for i in ledger["items"]):
        return None
    lights = {i["id"] for i in ledger["items"] if i.get("emits_light")}
    reach = max([np.abs(m.bounds[:, [0, 2]]).max() for oid, m in placed.items() if oid not in lights] + [0.3])
    g = trimesh.creation.cylinder(radius=reach * 1.6, height=0.01, sections=48)
    g.apply_transform(trimesh.transformations.rotation_matrix(math.pi / 2, [1, 0, 0]))
    g.apply_translation([0, -0.005, 0])
    g.visual = trimesh.visual.ColorVisuals(g, face_colors=np.tile([77, 124, 58, 255], (len(g.faces), 1)))
    return g


def placed_meshes(ledger, meshes, config):
    out = {}
    for oid, p in config["placements"].items():
        m = meshes[oid].copy()
        m.apply_transform(transform(p))
        if p.get("light"):  # a light source glows: warm, flat colour
            m.visual = trimesh.visual.ColorVisuals(m, face_colors=np.tile(GLOW, (len(m.faces), 1)))
        out[oid] = m
    return out


def export_glb(ledger, meshes, config, path):
    placed = placed_meshes(ledger, meshes, config)
    scene = trimesh.Scene()
    for oid, m in placed.items():
        scene.add_geometry(m, node_name=oid)
    ground = _ground(ledger, placed)
    if ground is not None:
        scene.add_geometry(ground, node_name="ground")
    scene.export(path)
    return path


# ---------------- picture (headless, for people and for models) ----------------

def _sky(sun):
    if not sun:
        return "#2a3346"
    t = min(1.0, max(0.0, (sun["elevation_deg"] - 5) / 40))  # low sun -> warm sky
    low, high = np.array([0.95, 0.62, 0.38]), np.array([0.53, 0.71, 0.88])
    return tuple(low + (high - low) * t)


def _to_axes(v):
    """Y-up scene coords -> matplotlib Z-up axes (a proper rotation)."""
    return np.stack([v[..., 0], -v[..., 2], v[..., 1]], axis=-1)


def _shade(mesh, light, color, glow=False):
    if glow:
        k = np.ones(len(mesh.faces))
    else:  # two-sided: Objaverse meshes often have flipped faces
        k = 0.45 + 0.55 * np.abs(mesh.face_normals @ light)
    rgb = np.clip(color[:3] / 255.0 * k[:, None], 0, 1)
    return np.c_[rgb, np.ones(len(rgb))]


def render(meshes_by_id, path, sun=None, ground=None, title=None):
    light = sun_direction(sun) if sun else np.array([0.4, 0.8, 0.45]) / np.linalg.norm([0.4, 0.8, 0.45])
    fig = plt.figure(figsize=(6.4, 4.8), dpi=110)
    fig.patch.set_facecolor(_sky(sun))
    ax = fig.add_subplot(projection="3d", computed_zorder=False)
    ax.set_facecolor(_sky(sun))
    pts = []

    if ground is not None:
        v = ground.vertices
        top = v[(v[:, 1] > -0.005) & (np.hypot(v[:, 0], v[:, 2]) > 1e-6)]  # top cap rim
        ring = top[np.argsort(np.arctan2(top[:, 2], top[:, 0]))]
        ax.add_collection3d(Poly3DCollection([_to_axes(ring)], facecolors=[(0.30, 0.49, 0.23, 1)], zorder=0))
        pts.append(ring)

    if sun and sun.get("shadows") and light[1] > 0.05:
        for oid, m in meshes_by_id.items():
            if oid == sun["id"]:
                continue
            v = m.vertices if len(m.vertices) < 4000 else m.vertices[:: len(m.vertices) // 4000 + 1]
            flat = v - np.outer(v[:, 1] / light[1], light)
            try:
                hull = ConvexHull(flat[:, [0, 2]])
            except Exception:
                continue
            poly = flat[hull.vertices]
            poly[:, 1] = 0.002
            ax.add_collection3d(Poly3DCollection([_to_axes(poly)], facecolors=[(0, 0, 0, 0.35)], zorder=1))

    tris, colors = [], []
    for oid, m in meshes_by_id.items():
        color = base_color(m)
        r = decimate(m, RENDER_FACE_BUDGET) if len(m.faces) > RENDER_FACE_BUDGET else m
        tris.append(r.triangles)
        colors.append(_shade(r, light, color, glow=bool(sun) and oid == sun["id"]))
        pts.append(r.vertices)
    if tris:
        ax.add_collection3d(Poly3DCollection(_to_axes(np.concatenate(tris)), facecolors=np.concatenate(colors),
                                             linewidths=0, zorder=2))

    allp = _to_axes(np.concatenate(pts)) if pts else np.zeros((1, 3))
    lo, hi = allp.min(axis=0), allp.max(axis=0)
    span = max((hi - lo).max(), 1e-3)
    mid = (lo + hi) / 2
    ax.set_xlim(mid[0] - span / 2, mid[0] + span / 2)
    ax.set_ylim(mid[1] - span / 2, mid[1] + span / 2)
    ax.set_zlim(max(lo[2], mid[2] - span / 2), max(lo[2], mid[2] - span / 2) + span)
    ax.set_box_aspect((1, 1, 1))
    ax.view_init(elev=22, azim=-60)
    ax.set_axis_off()
    if title:
        ax.set_title(title, color="white", fontsize=10)
    fig.savefig(path, facecolor=fig.get_facecolor(), bbox_inches="tight")
    plt.close(fig)
    return path


def render_scene(ledger, meshes, config, path):
    placed = placed_meshes(ledger, meshes, config)
    return render(placed, path, sun=config["sun"], ground=_ground(ledger, placed))