Previewer of 3D Scroll ink predictions

prediction_uv.py

import numpy as np
import zarr
from matplotlib import pyplot as plt
from PIL import Image
from tqdm import tqdm
import open3d as o3d
Image.MAX_IMAGE_PIXELS = None

# Taken from ThaumatoAnakalyptor by Julian Schiliger
def orient_uvs(vertices):
        # Rotate vertices and calculate the needed area
        vertices[:, 0] = 1.0 - vertices[:, 0]
        u_range = np.max(vertices[:, 0]) - np.min(vertices[:, 0])
        v_range = np.max(vertices[:, 1]) - np.min(vertices[:, 1])
        u_longer_v = u_range > v_range
        u_return = vertices[:, 0]
        v_return = vertices[:, 1]
        area_return = u_range * v_range
        for angle in range(-70, 70, 5):
            u_prime = vertices[:, 0] * np.cos(np.deg2rad(angle)) - vertices[:, 1] * np.sin(np.deg2rad(angle))
            v_prime = vertices[:, 0] * np.sin(np.deg2rad(angle)) + vertices[:, 1] * np.cos(np.deg2rad(angle))
            u_prime_range = np.max(u_prime) - np.min(u_prime)
            v_prime_range = np.max(v_prime) - np.min(v_prime)
            if u_prime_range < v_prime_range and u_longer_v:
                continue
            elif u_prime_range > v_prime_range and not u_longer_v:
                continue
            area = u_prime_range * v_prime_range
            if area < area_return:
                u_return = u_prime
                v_return = v_prime
                area_return = area
        return np.stack((u_return, v_return), axis=-1) 


def scatter_render(segment_id, mesh_path, image_path, predictions_path, output_image, marker_size=2, invert_axis=0):
    mesh = o3d.io.read_triangle_mesh(mesh_path)
    predictions = zarr.open(predictions_path, mode='r')
    V = np.asarray(mesh.vertices)
    N = np.asarray(mesh.vertex_normals)
    N /= np.linalg.norm(N, axis=1).reshape(-1,1)
    UV = np.asarray(mesh.triangle_uvs)  # Ensure your mesh has UV coordinates
    F = np.asarray(mesh.triangles)
    with Image.open(image_path) as img:
            # Get dimensions
            y_size, x_size = img.size
    UV[:,0] *= y_size
    UV[:,1] *= x_size
    uv = np.zeros((V.shape[0], 2), dtype=np.float64)
    uvs = UV.reshape((F.shape[0], F.shape[1], 2))

    for t in range(F.shape[0]):
        for v in range(F.shape[1]):
            uv[F[t,v]] = uvs[t,v]

    uv = orient_uvs(uv)
    newV = np.copy(V)
    for layers in tqdm(range(0,2)):
        Vlayer = V + ((-1)**layers) * N
        newV = np.vstack((newV, Vlayer))

    colors = np.zeros(newV.shape[0])

    roundV = np.round(newV).astype(np.int_)

    needed_cubes = set()
    for i in tqdm(range(roundV.shape[0])):
        x, y, z = roundV[i]
        x_cell = x//256 + 1
        y_cell = y//256 + 1
        z_cell = z//256 + 1
        needed_cubes.add((y_cell,x_cell,z_cell))

    cube_dict = {item: [] for item in needed_cubes}
    for i in tqdm(range(roundV.shape[0])):
        x, y, z = roundV[i]
        x_cell = x//256 + 1
        y_cell = y//256 + 1
        z_cell = z//256 + 1
        cube_dict[(y_cell,x_cell,z_cell)].append([i,x,y,z])


    for _, values in tqdm(cube_dict.items()):
        indices = []
        coordinates = []
        for element in values:
            indices.append(element[0])
            coordinates.append([element[1],element[2],element[3]])
        indices = np.array(indices)
        coordinates = np.array(coordinates).reshape(-1,3)
        colors[indices] = predictions[coordinates[:,1],coordinates[:,0],coordinates[:,2]]   #yx swap! ok

    #postprocess
    size = 3
    newV = newV.reshape(size,V.shape[0],3)
    colors = np.array(colors).reshape(size,V.shape[0])
    colors_avg = np.prod(colors, axis=0)
    colors_avg /= np.max(colors_avg)
    uv -= np.min(uv, axis=0)

    
    plt.figure(figsize=(np.max(uv,axis=0)[0]/500,np.max(uv,axis=0)[1]/500))
    plt.title(f"{segment_id}")
    if invert_axis == 1:
        plt.gca().invert_xaxis()
    plt.scatter(uv[:,0], uv[:,1], c=colors_avg[:], marker=',', s=marker_size, cmap='gray') # change s size if too sparse
    plt.savefig(output_image)


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Visualize 3D scroll ink predictions")

    # Add positional arguments
    parser.add_argument(
        "segment_id",
        type=int,
        help="ID of the mesh"
    )

    parser.add_argument(
        "mesh_path",
        type=str,
        help="Path to mesh .obj"
    )

    parser.add_argument(
        "image_path",
        type=str,
        help="Path to mesh mask .png"
    )

    parser.add_argument(
        "predictions_path",
        type=str,
        help="Path to the .zarr containing the 3D predictions",
    )

    parser.add_argument(
        "--output",
        type=str,
        help="Path to the .png rendered predictions",
        default="./output.png"
    )

    parser.add_argument(
        "--s",
        type=float,
        help="Size of the marker",
        default=2
    )

    parser.add_argument(
        "--i",
        type=int,
        help="Invert axis of the plot? 0 or 1",
        default=0
    )

    # Parse the arguments
    args = parser.parse_args()


    scatter_render(args.segment_id, args.mesh_path, args.image_path, args.predictions_path, args.output, args.s, args.i)

Usage:

python prediction_uv.py segment_id path_to_obj path_to_mask path_to_predictions_zarr --o output.png

Example usage:

python prediction_uv.py 753 "D:/vesuvius/Scroll1.volpkg/working/GP/753/753.obj" "D:/vesuvius/Scroll1.volpkg/working/GP/753/753.png" "D:/vesuvius/3d_predictions_scroll1.zarr"