liblaf.melon.ext.wrap ¶

def annotate_landmarks(
    left: Any,
    right: Any,
    *,
    left_landmarks: Float[ArrayLike, "L 3"] | None = None,
    right_landmarks: Float[ArrayLike, "L 3"] | None = None,
) -> tuple[Float[np.ndarray, "L 3"], Float[np.ndarray, "L 3"]]:
    """Open a Wrap project for editing corresponding landmarks.

    The left mesh is pre-aligned to the right mesh with either landmark
    Procrustes fitting or Trimesh mesh registration. Returned left landmarks are
    transformed back into the original left-mesh coordinate system.

    Args:
        left: Left mesh.
        right: Right mesh.
        left_landmarks: Optional starting landmarks on `left`.
        right_landmarks: Optional starting landmarks on `right`.

    Returns:
        Updated `(left_landmarks, right_landmarks)`.

    Raises:
        subprocess.CalledProcessError: If `Wrap.sh` fails.
    """
    left: tm.Trimesh = io.as_trimesh(left, triangulate=True)
    right: tm.Trimesh = io.as_trimesh(right, triangulate=True)
    if left_landmarks is None:
        left_landmarks: Float[np.ndarray, "L 3"] = np.zeros((0, 3))
    if right_landmarks is None:
        right_landmarks: Float[np.ndarray, "L 3"] = np.zeros((0, 3))

    environment: j2.Environment = get_environment()
    template: j2.Template = environment.get_template("annotate-landmarks.wrap")

    if np.size(left_landmarks) > 0 and np.shape(left_landmarks) == np.shape(
        right_landmarks
    ):
        transform, transformed, cost = tm.registration.procrustes(
            left_landmarks, right_landmarks
        )
        logger.info("procrustes cost: %f", cost)
        left: tm.Trimesh = left.copy()
        left: tm.Trimesh = left.apply_transform(transform)
        left_landmarks: Float[np.ndarray, "L 3"] = transformed
    else:
        transform, cost = tm.registration.mesh_other(left, right, scale=True)
        logger.info("mesh_other cost: %f", cost)
        left: tm.Trimesh = left.copy()
        left: tm.Trimesh = left.apply_transform(transform)
        left_landmarks: Float[np.ndarray, "L 3"] = tm.transform_points(
            left_landmarks, transform
        )

    with tempfile.TemporaryDirectory() as tmpdir_:
        tmpdir: Path = Path(tmpdir_)
        project_file: Path = tmpdir / "annotate-landmarks.wrap"
        left_file: Path = tmpdir / "left.obj"
        right_file: Path = tmpdir / "right.obj"
        left_landmarks_file: Path = tmpdir / "left.landmarks.json"
        right_landmarks_file: Path = tmpdir / "right.landmarks.json"
        io.save(left, left_file)
        io.save(right, right_file)
        io.save_landmarks(left_landmarks, left_landmarks_file)
        io.save_landmarks(right_landmarks, right_landmarks_file)
        project: str = template.render(
            {
                "left": left_file,
                "right": right_file,
                "left_landmarks": left_landmarks_file,
                "right_landmarks": right_landmarks_file,
            }
        )
        project_file.write_text(project)
        subprocess.run(["Wrap.sh", project_file], check=True)
        left_landmarks: Float[np.ndarray, "L 3"] = io.load_landmarks(
            left_landmarks_file
        )
        right_landmarks: Float[np.ndarray, "L 3"] = io.load_landmarks(
            right_landmarks_file
        )

    left_landmarks: Float[np.ndarray, "L 3"] = tm.transform_points(
        left_landmarks, tm.transformations.inverse_matrix(transform)
    )
    return left_landmarks, right_landmarks

def delta_transfer(
    floating_neutral: pv.PolyData, ref_neutral: Any, ref_expression: Any
) -> pv.PolyData:
    environment: j2.Environment = get_environment()
    template: j2.Template = environment.get_template("delta-transfer.wrap")

    with tempfile.TemporaryDirectory() as tmpdir_:
        tmpdir: Path = Path(tmpdir_)
        project_path: Path = tmpdir / "delta-transfer.wrap"
        floating_neutral_path: Path = tmpdir / "floating-neutral.obj"
        ref_neutral_path: Path = tmpdir / "ref-neutral.obj"
        ref_expression_path: Path = tmpdir / "ref-expression.obj"
        output_path: Path = tmpdir / "output.obj"
        io.save(floating_neutral, floating_neutral_path)
        io.save(ref_neutral, ref_neutral_path)
        io.save(ref_expression, ref_expression_path)
        project: str = template.render(
            {
                "floating_neutral": floating_neutral_path,
                "ref_neutral": ref_neutral_path,
                "ref_expression": ref_expression_path,
                "output": output_path,
            }
        )
        project_path.write_text(project)
        subprocess.run(["WrapCmd.sh", "compute", project_path], check=True)
        output: pv.PolyData = io.load_polydata(output_path)
        output.copy_attributes(floating_neutral)
    return output

def fast_wrapping(
    floating: Any,
    fixed: Any,
    *,
    floating_landmarks: Float[ArrayLike, "L 3"] | None = None,
    fixed_landmarks: Float[ArrayLike, "L 3"] | None = None,
    free_polygons_floating: Integer[ArrayLike, " F"] | None = None,
) -> pv.PolyData:
    """Run Faceform Wrap fast wrapping between floating and fixed surfaces.

    Args:
        floating: Mesh to deform.
        fixed: Target mesh.
        floating_landmarks: Optional landmarks on the floating mesh.
        fixed_landmarks: Optional corresponding landmarks on the fixed mesh.
        free_polygons_floating: Optional floating-mesh polygons left free during
            wrapping.

    Returns:
        Wrapped floating mesh with attributes copied from the aligned input.

    Raises:
        subprocess.CalledProcessError: If `WrapCmd.sh compute` fails.
    """
    if floating_landmarks is None:
        floating_landmarks: Float[np.ndarray, "L 3"] = np.empty((0, 3))
    if fixed_landmarks is None:
        fixed_landmarks: Float[np.ndarray, "L 3"] = np.empty((0, 3))
    if free_polygons_floating is None:
        free_polygons_floating: Integer[np.ndarray, " F"] = np.empty((0,), dtype=int)

    environment: j2.Environment = get_environment()
    template: j2.Template = environment.get_template("fast-wrapping.wrap")

    floating: pv.PolyData = io.as_polydata(floating)
    if np.size(floating_landmarks):
        matrix, transformed, cost = tm.registration.procrustes(
            floating_landmarks, fixed_landmarks
        )
        logger.info("procrustes cost: %f", cost)
        floating: pv.PolyData = floating.transform(matrix, inplace=False)
        floating_landmarks: Float[np.ndarray, "L 3"] = transformed

    with tempfile.TemporaryDirectory() as tmpdir_:
        tmpdir: Path = Path(tmpdir_)
        project_file: Path = tmpdir / "fast-wrapping.wrap"
        floating_path: Path = tmpdir / "floating.obj"
        fixed_path: Path = tmpdir / "fixed.obj"
        floating_landmarks_path: Path = tmpdir / "floating-landmarks.json"
        fixed_landmarks_path: Path = tmpdir / "fixed-landmarks.json"
        free_polygons_floating_path: Path = tmpdir / "free-polygons-floating.json"
        output_path: Path = tmpdir / "output.obj"
        io.save(floating, floating_path)
        io.save(fixed, fixed_path)
        io.save_landmarks(floating_landmarks, floating_landmarks_path)
        io.save_landmarks(fixed_landmarks, fixed_landmarks_path)
        io.save_polygons(free_polygons_floating, free_polygons_floating_path)
        project: str = template.render(
            {
                "floating": floating_path,
                "fixed": fixed_path,
                "floating_landmarks": floating_landmarks_path,
                "fixed_landmarks": fixed_landmarks_path,
                "free_polygons_floating": free_polygons_floating_path,
                "output": output_path,
            }
        )
        project_file.write_text(project)
        subprocess.run(["WrapCmd.sh", "compute", project_file], check=True)
        result: pv.PolyData = io.load_polydata(output_path)

    result.copy_attributes(floating)
    return result

@functools.cache
def get_environment() -> j2.Environment:
    """Return the cached Jinja environment for Wrap project templates."""
    return j2.Environment(
        undefined=j2.StrictUndefined,
        autoescape=j2.select_autoescape(),
        loader=j2.PackageLoader("liblaf.melon.ext.wrap"),
    )

def select_polygons(
    mesh: Any, polygons: Integer[ArrayLike, " N"] | None = None
) -> Integer[np.ndarray, " N"]:
    """Open a Wrap project that lets the user edit selected polygons.

    Args:
        mesh: Mesh to show in Wrap.
        polygons: Initial selected polygon indices.

    Returns:
        Polygon indices saved by Wrap.

    Raises:
        subprocess.CalledProcessError: If `Wrap.sh` fails.
    """
    if polygons is None:
        polygons: Integer[np.ndarray, " N"] = np.empty((0,), np.int32)

    environment: j2.Environment = get_environment()
    template: j2.Template = environment.get_template("select-polygons.wrap")

    with tempfile.TemporaryDirectory() as tmpdir_:
        tmpdir: Path = Path(tmpdir_)
        project_path: Path = tmpdir / "select-polygons.wrap"
        mesh_path: Path = tmpdir / "mesh.obj"
        polygons_path: Path = tmpdir / "polygons.json"
        io.save(mesh, mesh_path)
        io.save_polygons(polygons, polygons_path)
        project: str = template.render(
            {
                "mesh": mesh_path,
                "polygons": polygons_path,
            }
        )
        project_path.write_text(project)
        subprocess.run(["Wrap.sh", project_path], check=True)
        polygons: Integer[np.ndarray, " N"] = io.load_polygons(polygons_path)
    return polygons