code/hamer/rendering.py · changh95/hamer-p150 at main

Add files using upload-large-folder tool

2ed2aa1 verified about 2 months ago

5.85 kB

# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
# SPDX-License-Identifier: Apache-2.0
"""Render the predicted MANO mesh on the *original* image.
Projection convention derived from ``hamer/utils/renderer.py`` upstream.
The upstream applies a 180° X-rotation to vertices (trimesh) *and* then
renders with pyrender/OpenGL which does a y-flip (NDC y-up → image y-down).
Those two y-negations cancel.  Our equivalent — without pyrender — is to
project the MANO vertices **directly** without any rotation:
    proj_x = focal * (vx + tx) / (vz + tz) + W/2
    proj_y = focal * (vy + ty) / (vz + tz) + H/2
where (tx, ty, tz) = cam_crop_to_full(pred_cam, ...).  Applying only the
rotation (without the matching y-flip) inverts every vertex's y relative
to the hand center, making a downward-facing palm appear upward.
Triangles are rasterized as alpha-blended polygons with simple diffuse
shading from face-normal · view-vector — no ``pyrender`` / ``trimesh``
dependency.  The shading is enough to make finger curl readable instead of
the flat silhouette we used to ship.
"""
from __future__ import annotations

from pathlib import Path
from typing import Optional, Tuple

import numpy as np
import torch

# matplotlib is imported lazily inside ``overlay_mesh`` so the projection
# helpers (``cam_crop_to_full`` / ``project_mesh``) stay importable in a
# runtime that ships no matplotlib (the serving image).

FOCAL_LENGTH = 5000.0
MODEL_IMAGE_SIZE = 256


def cam_crop_to_full(
    cam_bbox: np.ndarray,
    box_center: np.ndarray,
    box_size: float,
    img_w: int,
    img_h: int,
    focal: float,
) -> np.ndarray:
    """Replicates ``hamer.utils.renderer.cam_crop_to_full`` on numpy.
    Returns the 3-D camera translation ``(tx, ty, tz)`` in original-image
    pixel-scaled world coordinates.
    """
    s, tx_b, ty_b = float(cam_bbox[0]), float(cam_bbox[1]), float(cam_bbox[2])
    cx, cy = float(box_center[0]), float(box_center[1])
    bs = box_size * s + 1e-9
    tz = 2.0 * focal / bs
    tx = 2.0 * (cx - img_w / 2.0) / bs + tx_b
    ty = 2.0 * (cy - img_h / 2.0) / bs + ty_b
    return np.array([tx, ty, tz], dtype=np.float64)


def project_mesh(
    verts: torch.Tensor,
    cam_bbox: torch.Tensor,
    box_center: Tuple[float, float],
    box_size: float,
    img_w: int,
    img_h: int,
    is_right: bool = True,
) -> Tuple[np.ndarray, np.ndarray]:
    """Project MANO vertices to ORIGINAL-image pixel coords.
    For a LEFT hand the upstream pipeline runs the (right-hand) model on a
    horizontally-mirrored input crop, then flips ``cam.x`` and ``verts.x``
    back before rendering — see ``demo.py`` lines 136-138 / 178-179.
    Reproduce that here.
    Returns ``(verts_2d, verts_camera_z)`` — the second array is per-vertex
    camera-space z (positive = depth in front of the camera) used for
    back-to-front face sorting in rendering.
    """
    focal = FOCAL_LENGTH / MODEL_IMAGE_SIZE * max(img_w, img_h)
    cam_bbox_np = cam_bbox.detach().cpu().numpy().copy()
    if not is_right:
        cam_bbox_np[1] *= -1.0
    cam_t = cam_crop_to_full(cam_bbox_np, np.array(box_center), box_size, img_w, img_h, focal)

    v = verts.detach().cpu().numpy().astype(np.float64).copy()
    if not is_right:
        v[:, 0] *= -1.0
    # Direct projection — no rotation applied (upstream's 180° X rotation and
    # pyrender's NDC y-flip cancel; applying only the rotation inverts y).
    v_cam = v + cam_t
    z = v_cam[:, 2]
    proj = focal * v_cam[:, :2] / z[:, None]
    proj[:, 0] += img_w / 2.0
    proj[:, 1] += img_h / 2.0
    return proj, z


def overlay_mesh(
    ax,
    image_uint8: np.ndarray,
    verts_2d: np.ndarray,
    verts_z: np.ndarray,
    faces: np.ndarray,
    color=(0.20, 0.85, 1.00),
    alpha: float = 0.65,
    is_right: bool = True,
) -> None:
    """Rasterize the projected mesh on ``image_uint8`` with face-normal
    diffuse shading + back-to-front depth ordering."""
    from matplotlib.collections import PolyCollection  # noqa: WPS433 (visualisation only)

    h, w = image_uint8.shape[:2]
    ax.imshow(image_uint8)
    ax.set_xlim(0, w)
    ax.set_ylim(h, 0)
    ax.set_aspect("equal")
    ax.axis("off")

    tri_z = verts_z[faces].mean(axis=1)               # (F,) camera z per face
    tris2d = verts_2d[faces]                          # (F, 3, 2)

    # Cull triangles whose mean z is behind the camera or way off-screen.
    inside = (
        (tris2d[..., 0] > -50).all(-1)
        & (tris2d[..., 0] < w + 50).all(-1)
        & (tris2d[..., 1] > -50).all(-1)
        & (tris2d[..., 1] < h + 50).all(-1)
        & (tri_z > 0.05)
    )
    tris2d, tri_z = tris2d[inside], tri_z[inside]

    # Winding: without the 180° X rotation, projecting v directly means
    # vy > 0 maps to larger proj_y (lower in image).  A right-hand
    # front-facing triangle that is CCW in 3D world projects to CW in
    # y-down image (cross_z > 0).  The x-flip for left hands reverses
    # winding, so front-facing left-hand triangles have cross_z < 0.
    e1 = tris2d[:, 1] - tris2d[:, 0]
    e2 = tris2d[:, 2] - tris2d[:, 0]
    cross_z = e1[:, 0] * e2[:, 1] - e1[:, 1] * e2[:, 0]
    winding_sign = 1.0 if is_right else -1.0
    facing = np.clip(winding_sign * cross_z / (np.abs(cross_z).max() + 1e-9), -1, 1)
    diffuse = 0.55 + 0.45 * facing

    # Painter's order: render farthest faces first.
    order = np.argsort(-tri_z)
    tris2d = tris2d[order]
    diffuse = diffuse[order]

    facecolors = np.empty((len(tris2d), 4), dtype=np.float32)
    facecolors[:, 0] = color[0] * diffuse
    facecolors[:, 1] = color[1] * diffuse
    facecolors[:, 2] = color[2] * diffuse
    facecolors[:, 3] = alpha

    coll = PolyCollection(
        tris2d,
        facecolors=facecolors,
        edgecolors=(0.05, 0.10, 0.15, 0.06),
        linewidths=0.25,
    )
    ax.add_collection(coll)

	# SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
	# SPDX-License-Identifier: Apache-2.0
	"""Render the predicted MANO mesh on the original image.

	Projection convention derived from ``hamer/utils/renderer.py`` upstream.
	The upstream applies a 180° X-rotation to vertices (trimesh) and then
	renders with pyrender/OpenGL which does a y-flip (NDC y-up → image y-down).
	Those two y-negations cancel. Our equivalent — without pyrender — is to
	project the MANO vertices directly without any rotation:

	proj_x = focal * (vx + tx) / (vz + tz) + W/2
	proj_y = focal * (vy + ty) / (vz + tz) + H/2

	where (tx, ty, tz) = cam_crop_to_full(pred_cam, ...). Applying only the
	rotation (without the matching y-flip) inverts every vertex's y relative
	to the hand center, making a downward-facing palm appear upward.

	Triangles are rasterized as alpha-blended polygons with simple diffuse
	shading from face-normal · view-vector — no ``pyrender`` / ``trimesh``
	dependency. The shading is enough to make finger curl readable instead of
	the flat silhouette we used to ship.
	"""
	from __future__ import annotations

	from pathlib import Path
	from typing import Optional, Tuple

	import numpy as np
	import torch

	# matplotlib is imported lazily inside ``overlay_mesh`` so the projection
	# helpers (``cam_crop_to_full`` / ``project_mesh``) stay importable in a
	# runtime that ships no matplotlib (the serving image).

	FOCAL_LENGTH = 5000.0
	MODEL_IMAGE_SIZE = 256


	def cam_crop_to_full(
	cam_bbox: np.ndarray,
	box_center: np.ndarray,
	box_size: float,
	img_w: int,
	img_h: int,
	focal: float,
	) -> np.ndarray:
	"""Replicates ``hamer.utils.renderer.cam_crop_to_full`` on numpy.

	Returns the 3-D camera translation ``(tx, ty, tz)`` in original-image
	pixel-scaled world coordinates.
	"""
	s, tx_b, ty_b = float(cam_bbox[0]), float(cam_bbox[1]), float(cam_bbox[2])
	cx, cy = float(box_center[0]), float(box_center[1])
	bs = box_size * s + 1e-9
	tz = 2.0 * focal / bs
	tx = 2.0 * (cx - img_w / 2.0) / bs + tx_b
	ty = 2.0 * (cy - img_h / 2.0) / bs + ty_b
	return np.array([tx, ty, tz], dtype=np.float64)


	def project_mesh(
	verts: torch.Tensor,
	cam_bbox: torch.Tensor,
	box_center: Tuple[float, float],
	box_size: float,
	img_w: int,
	img_h: int,
	is_right: bool = True,
	) -> Tuple[np.ndarray, np.ndarray]:
	"""Project MANO vertices to ORIGINAL-image pixel coords.

	For a LEFT hand the upstream pipeline runs the (right-hand) model on a
	horizontally-mirrored input crop, then flips ``cam.x`` and ``verts.x``
	back before rendering — see ``demo.py`` lines 136-138 / 178-179.
	Reproduce that here.

	Returns ``(verts_2d, verts_camera_z)`` — the second array is per-vertex
	camera-space z (positive = depth in front of the camera) used for
	back-to-front face sorting in rendering.
	"""
	focal = FOCAL_LENGTH / MODEL_IMAGE_SIZE * max(img_w, img_h)
	cam_bbox_np = cam_bbox.detach().cpu().numpy().copy()
	if not is_right:
	cam_bbox_np[1] *= -1.0
	cam_t = cam_crop_to_full(cam_bbox_np, np.array(box_center), box_size, img_w, img_h, focal)

	v = verts.detach().cpu().numpy().astype(np.float64).copy()
	if not is_right:
	v[:, 0] *= -1.0
	# Direct projection — no rotation applied (upstream's 180° X rotation and
	# pyrender's NDC y-flip cancel; applying only the rotation inverts y).
	v_cam = v + cam_t
	z = v_cam[:, 2]
	proj = focal * v_cam[:, :2] / z[:, None]
	proj[:, 0] += img_w / 2.0
	proj[:, 1] += img_h / 2.0
	return proj, z


	def overlay_mesh(
	ax,
	image_uint8: np.ndarray,
	verts_2d: np.ndarray,
	verts_z: np.ndarray,
	faces: np.ndarray,
	color=(0.20, 0.85, 1.00),
	alpha: float = 0.65,
	is_right: bool = True,
	) -> None:
	"""Rasterize the projected mesh on ``image_uint8`` with face-normal
	diffuse shading + back-to-front depth ordering."""
	from matplotlib.collections import PolyCollection # noqa: WPS433 (visualisation only)

	h, w = image_uint8.shape[:2]
	ax.imshow(image_uint8)
	ax.set_xlim(0, w)
	ax.set_ylim(h, 0)
	ax.set_aspect("equal")
	ax.axis("off")

	tri_z = verts_z[faces].mean(axis=1) # (F,) camera z per face
	tris2d = verts_2d[faces] # (F, 3, 2)

	# Cull triangles whose mean z is behind the camera or way off-screen.
	inside = (
	(tris2d[..., 0] > -50).all(-1)
	& (tris2d[..., 0] < w + 50).all(-1)
	& (tris2d[..., 1] > -50).all(-1)
	& (tris2d[..., 1] < h + 50).all(-1)
	& (tri_z > 0.05)
	)
	tris2d, tri_z = tris2d[inside], tri_z[inside]

	# Winding: without the 180° X rotation, projecting v directly means
	# vy > 0 maps to larger proj_y (lower in image). A right-hand
	# front-facing triangle that is CCW in 3D world projects to CW in
	# y-down image (cross_z > 0). The x-flip for left hands reverses
	# winding, so front-facing left-hand triangles have cross_z < 0.
	e1 = tris2d[:, 1] - tris2d[:, 0]
	e2 = tris2d[:, 2] - tris2d[:, 0]
	cross_z = e1[:, 0] * e2[:, 1] - e1[:, 1] * e2[:, 0]
	winding_sign = 1.0 if is_right else -1.0
	facing = np.clip(winding_sign * cross_z / (np.abs(cross_z).max() + 1e-9), -1, 1)
	diffuse = 0.55 + 0.45 * facing

	# Painter's order: render farthest faces first.
	order = np.argsort(-tri_z)
	tris2d = tris2d[order]
	diffuse = diffuse[order]

	facecolors = np.empty((len(tris2d), 4), dtype=np.float32)
	facecolors[:, 0] = color[0] * diffuse
	facecolors[:, 1] = color[1] * diffuse
	facecolors[:, 2] = color[2] * diffuse
	facecolors[:, 3] = alpha

	coll = PolyCollection(
	tris2d,
	facecolors=facecolors,
	edgecolors=(0.05, 0.10, 0.15, 0.06),
	linewidths=0.25,
	)
	ax.add_collection(coll)