membranorama-like with rastering

triangle_raster.py

#!/usr/bin/env python3

import numpy as np
import pandas as pd
import napari
import mrcfile
from morphosamplers.sampler import (
    generate_1d_grid,
    place_sampling_grids,
    sample_volume_at_coordinates,
)
from scipy.spatial.transform import Rotation
from rich.progress import track
from rich import print


# vert = np.load("simple_vert.npy")
# faces = np.load("simple_faces.npy")
# vol = np.load("simple_label.npy").astype(np.float32)

# obj_path = 'TS_004_dose-filt_lp50_bin8_membrain_model_28885faces.obj'
# mesh_px_size = 1
# tomo_path = 'TS_004_dose-filt_lp50_bin8.rec'

obj_path = "T17S1C3M4.obj"
mesh_px_size = 14.08
tomo_path = "tomo17_load1G5L3_bin4_denoised_ctfcorr_scaled3.rec"

# obj_path = 'tomo_17_M10_grow1_1_mesh_data.obj'
# mesh_px_size = 1
# tomo_path = 'tomo_17_M10_grow1_1_mesh_data.mrc'

df = pd.read_csv(obj_path, sep=" ", names=["type", "a", "b", "c"])
vert = df[df["type"] == "v"][["a", "b", "c"]].to_numpy() / mesh_px_size
faces = df[df["type"] == "f"][["a", "b", "c"]].to_numpy().astype(int) - 1
with mrcfile.open(tomo_path, permissive=True) as mrc:
    vol = mrc.data.T.astype(np.float32)

v = napari.Viewer(ndisplay=3)
v.add_image(vol, blending="translucent", name="volume")
# v.add_points(
#     vert,
#     face_color="red",
#     size=0.2,
#     features={"idx": np.arange(len(vert))},
#     text={"string": "idx", "color": "white", "size": 20},
#     name='vertices'
# )
# v.add_surface((vert, faces), name='surface')

edges = vert[faces] - np.roll(vert[faces], 1, axis=1)
edge_length = np.linalg.norm(edges, axis=2).mean().round().astype(int)
edge_length = max(edge_length, 2)  # fails with 1x1 texture
thickness = 5

tex_shape = (edge_length, edge_length)
base_grid = generate_1d_grid(thickness, grid_spacing=1)

raster = np.empty((len(faces), *tex_shape, 3))
raster_flat = raster.reshape(-1, 3)

tex_full_width = np.ceil(np.sqrt(len(faces))).astype(int)
tex_full_shape = np.array(tex_shape) * tex_full_width
texture = np.full(tex_full_shape, np.nan, dtype=np.float32)
texcoords = []
orientations = []


for f_idx, face in track(list(enumerate(faces)), description="Processing face coords"):
    pts = vert[face]
    istep = (pts[1] - pts[0]) / (edge_length - 1)
    jstep = (pts[2] - pts[0]) / (edge_length - 1)

    raster[f_idx, :] = pts[0]

    for i in range(edge_length):
        for j in range(edge_length):
            raster[f_idx, i, j] += i * istep + j * jstep

    # v.add_points(raster_flat, size=0.2, face_color='yellow')

    normal = np.cross(istep, jstep)
    normal /= np.linalg.norm(normal)
    z = [0, 0, 1]
    ori = Rotation.align_vectors(normal, z)[0]
    # norm_computed = ori.apply([0, 0, 1])
    # norm_vecs = np.array([np.mean(pts, axis=0), norm_computed])
    # v.add_vectors(norm_vecs, edge_width=0.1)

    orientations.append(ori)

    # sampled = sample_volume_at_coordinates(vol, grids).astype(np.float32)
    # texel = sampled.mean(axis=1).reshape(tex_shape)

    I = f_idx // tex_full_width
    J = f_idx % tex_full_width
    i, j = np.array(tex_shape) * (I, J)
    # texture[i:i+edge_length, j:j+edge_length] = texel

    texcoords_face = np.array(
        [
            [i, j],
            [i + edge_length - 1, j],
            [i, j + edge_length - 1],
        ]
    )
    # +0.5 cause uv coords are "corner based" and not centered on voxel
    texcoords_face = (texcoords_face + 0.5) / tex_shape / tex_full_width
    texcoords.append(texcoords_face)

    # print(f"======== {f_idx=} ========")
    # print(f"{I=}, {J=}, {i=}, {j=}")
    # print(f"{texcoords_face=}")


# currently place_sampling_grids is not handling stacked dims
print("Generating full sample grid")
grids = place_sampling_grids(
    base_grid, raster_flat, np.repeat(orientations, np.prod(tex_shape))
)
print("Sampling")
sampled = sample_volume_at_coordinates(vol, grids, interpolation_order=1)
print("Sampled")
# v3 = napari.Viewer()
# v3.add_image(sampled)

# v.add_points(
#     grids.reshape(-1, 3),
#     size=0.05,
#     face_color="blue",
#     antialiasing=0,
#     features={"idx": np.arange(len(grids.reshape(-1, 3)))},
#     text={"string": "idx", "color": "white", "size": 20},
# )

sampled = sampled.mean(axis=-1)


for f_idx, face_sample in track(
    list(enumerate(np.split(sampled, len(faces)))), description="Generating texture"
):
    texel = face_sample.reshape(tex_shape)

    I = f_idx // tex_full_width
    J = f_idx % tex_full_width
    i, j = np.array(tex_shape) * (I, J)
    texture[i : i + edge_length, j : j + edge_length] = texel

texture -= np.nanmin(texture)
texture /= np.nanmax(texture)

texcoords = np.concatenate(texcoords)

v2 = napari.Viewer()
v2.add_image(texture, name="texture atlas").bounding_box.visible = True
# -0.5 cause in napari (unlike opengl uv coords) we're pixel centered
# v2.add_points(
#     (texcoords * tex_full_shape) - 0.5,
#     size=1,
#     face_color="red",
#     features={"idx": faces.reshape(-1)},
#     text={"string": "idx", "color": "white", "size": 20},
#     name='texcoords',
# )

vert_unwrapped = vert[faces].reshape(-1, 3)
faces_unwrapped = np.arange(len(vert_unwrapped)).reshape(-1, 3)
# TODO: why the .T[::-1] ???
v.add_surface(
    (vert_unwrapped, faces_unwrapped),
    texture=texture.T[::-1],
    texcoords=texcoords,
    name="textured surface",
)