mri2avatar.py

import os
import pdb
import torch
import shutil
import logging
import argparse
import subprocess
import numpy as np
import nibabel as nib
from tqdm import tqdm
from pathlib import Path
from typing import Dict, List, Union, Optional
from skimage.filters import threshold_otsu
from scipy.ndimage import (
    generate_binary_structure, binary_fill_holes, 
    binary_dilation, binary_erosion, label
    )
from pytorch3d.io import (
    load_obj,
    load_objs_as_meshes,
    save_obj
)
from pytorch3d.renderer import (
    PerspectiveCameras,
    AmbientLights,
    RasterizationSettings,
    look_at_view_transform,
    MeshRendererWithFragments,
    MeshRasterizer,
    SoftPhongShader,
    HardFlatShader,
    TexturesVertex,
)
from pytorch3d.structures import Meshes
import imageio.v2 as imageio

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('utils')

if torch.cuda.is_available():
    DEVICE = torch.device("cuda:0")
    torch.cuda.set_device(DEVICE)
else:
    DEVICE = torch.device("cpu")

TEXT2TEX_DIR = "/home-local/software/Text2Tex"


def user_input():
    parser = argparse.ArgumentParser(description="Convert NIfTI to 3D mesh with texture.")
    parser.add_argument('--path_input', type=str, required=True, help="Path to input NIfTI file.")
    parser.add_argument('--outdir', type=str, required=True, help="Output directory.")
    parser.add_argument('--use_binary_mask', action='store_true', help=(
        "Generate binary mask for main object, using Otsu's thresholding and morphological operations, "
        "and then produce mesh from the mask, instead of the original image."
    ))
    parser.add_argument('--reduction_factor', type=float, default=0.1, help=(
        "Smaller values produce simpler mesh, resulting in smaller file size and faster rendering."
    ))
    parser.add_argument('--prompt', type=str, required=True, help="Text prompt for diffusion model to generate texture.")
    parser.add_argument("--device", type=str, choices=["a6000", "2080"], default="a6000", 
                        help="Type of GPU available. Will determine the quality of the texture.")
    parser.add_argument('--intermediate', action='store_true', help="Preserve intermediate files")
    parser.add_argument('--generate_gif', action='store_true', help="Generate rotating GIFs of the original (untextured) and final (textured) meshes.")
    args = parser.parse_args()
    return args


def standardize_nifti_for_meshing(
    path_input: str,
    path_output: str,
    dtype: Optional[np.dtype] = None,
    overwrite: bool = False,
):
    """ First step preprocessing to standardize the format and orientation of NIfTI file.
    """
    if not Path(path_input).exists():
        return f"Error: Input file not found: {path_input}"
    if Path(path_output).exists() and not overwrite:
        return f"Error: Output file already exists: {path_output}. Set overwrite=True to replace."
    
    output_dir = Path(path_output).parent
    output_dir.mkdir(parents=True, exist_ok=True)
    
    try:
        logger.info(f"Loading NIfTI file from {path_input}")
        img = nib.load(path_input)
        
        logger.info("Reorienting to closest canonical orientation")
        img = nib.as_closest_canonical(img)
        
        data = img.get_fdata().astype("float32")
        if dtype is not None:
            logger.info(f"Converting data type to {dtype}")
            data = data.astype(dtype)
        
        logger.info("Checking sform and qform codes")
        affine = img.affine
        img_fixed = nib.Nifti1Image(data, affine)
        
        if img.header['sform_code'] == 0:
            logger.info("Fixing sform code")
            img_fixed.header.set_sform(affine, code=1)
        else:
            img_fixed.header.set_sform(img.header.get_sform(), code=img.header['sform_code'].item())
                
        if img.header['qform_code'] == 0:
            logger.info("Fixing qform code")
            img_fixed.header.set_qform(affine, code=1)
        else:
            img_fixed.header.set_qform(img.header.get_qform(), code=img.header['qform_code'].item())
            
        nib.save(img_fixed, path_output)
        logger.info(f"Finished standardizing {path_input}")
        return f"Successfully saved standardized NIfTI to {path_output}"
    
    except Exception as e:
        logger.error(f"Error processing {path_input}: {str(e)}", exc_info=True)
        return f"An error occurred: {str(e)}"


def save_main_object_binary_mask(
    path_input: str,
    path_output: str,
    iterations: int = 10,
    dilate: int = 0,
    pad: bool = True,
    pad_value: Union[int, float, str] = 0,
):
    """ Given a NIfTI file of an image, save the binary mask of the main object.
    Otsu's thresholding and morphological operations are used to generate the mask.
    
    Args:
        path_input (str): path to the input NIfTI file.
        path_output (str): path to save the output binary mask NIfTI file.
        iterations (int): number of iterations for dilation and erosion. Large values
            are good for filling holes and gaps, after which the largest connected component
            will cover the entire object, but also smooth out the fine structures.
        dilate (int): number of additional dilation iterations to apply on the final mask.
        pad (bool): whether to pad the image before morphological operations to avoid border effects.
        pad_value (int, float, str): value to use for padding.
    """
    img = nib.load(path_input)
    data = img.get_fdata()

    # Otsu's thresholding
    threshold = threshold_otsu(data)
    mask = data > threshold
    
    # Dilation + Fill holes + Erosion
    struct_element = generate_binary_structure(rank=data.ndim, connectivity=1)
    if pad:
        pad_width = [(iterations, iterations) for _ in range(data.ndim)]
        if isinstance(pad_value, str):
            if pad_value.lower() == 'min':
                pad_value = data.min()
            elif pad_value.lower() == 'max':
                pad_value = data.max()
            else:
                raise ValueError("pad_value supports 'min', 'max', or numeric values only.")
        else:
            pad_value = float(pad_value)
        mask = np.pad(mask, pad_width, mode='constant', constant_values=pad_value)

    mask = binary_dilation(mask, structure=struct_element, iterations=iterations)
    mask = binary_fill_holes(mask)
    mask = binary_erosion(mask, structure=struct_element, iterations=iterations)

    # Keep the largest connected component
    labels, _ = label(mask)
    sizes = np.bincount(labels.ravel())
    sizes[0] = 0
    idx_largest = sizes.argmax()
    mask = labels == idx_largest
    
    # Dilate the mask if specified
    if dilate > 0:
        mask = binary_dilation(mask, structure=struct_element, iterations=dilate)

    # Undo padding
    if pad:
        slices = tuple(slice(iterations, -iterations) for _ in range(data.ndim))
        mask = mask[slices]
        
    mask_img = nib.Nifti1Image(mask.astype(np.uint8), img.affine, img.header)
    nib.save(mask_img, path_output)
    logger.info(f"Saved binary mask to {path_output}")


def init_mesh(path_input: str, device=DEVICE):
    # https://github.com/daveredrum/Text2Tex/blob/main/scripts/normalize_mesh.py
    logger.info(f"Loading mesh: {path_input}")
    verts, faces, aux = load_obj(path_input, device=device)
    mesh = load_objs_as_meshes([path_input], device=device)
    return mesh, verts, faces, aux


def normalize_mesh(mesh):
    # https://github.com/daveredrum/Text2Tex/blob/main/scripts/normalize_mesh.py
    bbox = mesh.get_bounding_boxes()
    num_verts = mesh.verts_packed().shape[0]

    # move mesh to origin
    mesh_center = bbox.mean(dim=2).repeat(num_verts, 1)
    mesh = mesh.offset_verts(-mesh_center)

    # scale
    lens = bbox[0, :, 1] - bbox[0, :, 0]
    max_len = lens.max()
    scale = 1 / max_len
    scale = scale.unsqueeze(0).repeat(num_verts)
    mesh.scale_verts_(scale)

    return mesh.verts_packed()


def apply_mesh_normalization(
    path_input: str,
    path_output: str,
):
    mesh, verts, faces, aux = init_mesh(path_input)
    verts = normalize_mesh(mesh)
    save_obj(
        path_output,
        verts=verts,
        faces=faces.verts_idx,
        decimal_places=5,
    )
    logger.info(f"Saved normalized mesh to {path_output}")


def rotate_verts(verts, plane, deg):
    # https://github.com/daveredrum/Text2Tex/blob/main/scripts/rotate_mesh.py
    theta = deg * np.pi / 180
    A = torch.tensor([
        [np.cos(theta), -np.sin(theta)],
        [np.sin(theta), np.cos(theta)]
    ]).float()

    if plane == "xy":
        xy = torch.stack([verts[:, 0], verts[:, 1]], dim=1).float()
        xy = torch.matmul(A, xy.T).T
        verts = torch.stack([xy[:, 0], xy[:, 1], verts[:, 2]], dim=1)
    elif plane == "xz":
        xz = torch.stack([verts[:, 0], verts[:, 2]], dim=1).float()
        xz = torch.matmul(A, xz.T).T
        verts = torch.stack([xz[:, 0], verts[:, 1], xz[:, 1]], dim=1)
    else:
        yz = torch.stack([verts[:, 1], verts[:, 2]], dim=1).float()
        yz = torch.matmul(A, yz.T).T
        verts = torch.stack([verts[:, 0], yz[:, 0], yz[:, 1]], dim=1)
    
    logger.info(f"Rotated mesh {deg} degrees around {plane} plane")
    return verts


def apply_mesh_rotation(
    path_input: str,
    path_output: str,
    plane: List[str] = ["yz", "xz"],
    deg: List[float] = [270, 180],
):
    assert len(plane) == len(deg), "plane and deg must have the same length"
    verts, faces, aux = load_obj(path_input)
    for p, d in zip(plane, deg):
        verts = rotate_verts(verts, p, d)
        
    save_obj(
        path_output,
        verts=verts,
        faces=faces.verts_idx,
        decimal_places=5,
    )
    logger.info(f"Saved rotated mesh to {path_output}")


def run_text2tex(
    path_input: str,
    path_output_dir: str,
    prompt: str,
    device: str,
    text2tex_dir=TEXT2TEX_DIR
):
    path_input = Path(path_input).absolute()
    path_output_dir = Path(path_output_dir).absolute()
    path_output_dir.mkdir(parents=True, exist_ok=True)

    command = (
        f"python {text2tex_dir}/scripts/generate_texture.py "
        f"--input_dir {path_input.parent} "
        f"--output_dir {path_output_dir} "
        f"--obj_name {path_input.stem} "
        f"--obj_file {path_input.name} "
        f'--prompt "{prompt}" '
        f"--add_view_to_prompt --ddim_steps 50 --new_strength 1 --update_strength 0.3 "
        f"--view_threshold 0.1 --blend 0 --dist 1 --num_viewpoints 36 --viewpoint_mode predefined "
        f"--use_principle --update_steps 20 --update_mode heuristic --seed 42 --post_process "
        f'--device "{device}" '
        f"--use_objaverse"
    )
    logger.info(f"Running Text2Tex with command: {command}")
    subprocess.run(command, shell=True, check=True, cwd=text2tex_dir)


def generate_gif(
    path_obj: str,
    path_gif: str,
    num_views: int = 32,
    dist: float = 1.5,
    elev: float = 15.0,
    azim_start: float = 180.0,
    image_size: int = 1024,
    center_crop: int = 512,
    duration: float = 1.5,
    shader_type: str = "textured",
    device=DEVICE,
):
    """Renders a mesh from multiple viewpoints and saves it as a GIF."""
    Path(path_gif).parent.mkdir(parents=True, exist_ok=True)
    logger.info(f"Generating GIF: {path_gif}")

    # Load the mesh
    verts, faces, aux = load_obj(path_obj, device=device)
    if shader_type == "textured":
        mesh = load_objs_as_meshes([path_obj], device=device)
    else:
        face_verts = verts[faces.verts_idx]
        v0, v1, v2 = face_verts[:, 0], face_verts[:, 1], face_verts[:, 2]
        face_normals = torch.cross(v1 - v0, v2 - v0, dim=1)
        face_normals = face_normals / face_normals.norm(dim=1, keepdim=True)

        vertex_normals = torch.zeros_like(verts)
        vertex_counts = torch.zeros((verts.shape[0], 1), device=device)
        
        for i in range(3):
            vertex_indices = faces.verts_idx[:, i]
            vertex_normals.index_add_(0, vertex_indices, face_normals)
            vertex_counts.index_add_(0, vertex_indices, torch.ones((len(vertex_indices), 1), device=device))
            
        vertex_normals = vertex_normals / torch.clamp(vertex_counts, min=1)
        vertex_normals = vertex_normals / torch.clamp(vertex_normals.norm(dim=1, keepdim=True), min=1e-6)
        vertex_colors = (vertex_normals + 1) / 2.0
        textures = TexturesVertex(verts_features=vertex_colors.unsqueeze(0))
        mesh = Meshes(verts=[verts], faces=[faces.verts_idx], textures=textures)

    # center the mesh
    verts_packed = mesh.verts_packed()
    center = verts_packed.mean(0)
    mesh.offset_verts_(-center)

    images = []
    for i in tqdm(range(num_views), desc="Rendering frames"):
        azim = azim_start + i * (360.0 / num_views)
        R, T = look_at_view_transform(dist=dist, elev=elev, azim=azim)
        cameras = PerspectiveCameras(device=device, R=R, T=T)
        lights = AmbientLights(device=device)

        if shader_type == "textured":
            shader = SoftPhongShader(device=device, cameras=cameras, lights=lights)
        else:
            shader = HardFlatShader(device=device, cameras=cameras, lights=lights)

        renderer = MeshRendererWithFragments(
            rasterizer=MeshRasterizer(
                cameras=cameras,
                raster_settings=RasterizationSettings(image_size=image_size, faces_per_pixel=1)
            ),
            shader=shader
        )
        
        rendered_images, _ = renderer(mesh)
        
        # Crop the image
        assert center_crop <= image_size, "center_crop should not be greater than image_size"
        start_idx = (image_size - center_crop) // 2
        img_np = rendered_images[0, start_idx:start_idx+center_crop, start_idx:start_idx+center_crop, :3].cpu().numpy()
        
        img_uint8 = (img_np * 255).astype(np.uint8)
        images.append(img_uint8)
    imageio.mimsave(path_gif, images, duration=duration, loop=0)


def mri2avatar(args):
    path_input = Path(args.path_input)
    assert path_input.exists(), f"Input file does not exist: {path_input}"
    outdir = Path(args.outdir)
    tmpdir = outdir / "tmp"
    tmpdir.mkdir(parents=True, exist_ok=True)
    
    logger.info("====> Step 1: Standardizing NIfTI file")
    path_standardized = tmpdir / "standardized.nii.gz"
    standardize_nifti_for_meshing(
        path_input=str(path_input),
        path_output=str(path_standardized),
        overwrite=True,
    )
    
    if args.use_binary_mask:
        logger.info("====> Optional Step: Generating binary mask for the main object")
        path_mask = tmpdir / "standardized_mask.nii.gz"
        save_main_object_binary_mask(
            path_input=str(path_standardized),
            path_output=str(path_mask),
            pad=True,
            pad_value='min',
        )
        path_nifti = path_mask
        logger.info(f"====> Use binary mask for meshing: {path_mask}")
    else:
        path_nifti = path_standardized
        logger.info(f"====> Use standardized image for meshing: {path_standardized}")

    logger.info("====> Step 2: Running nii2mesh to generate initial mesh")
    path_init_mesh = tmpdir / "mesh.obj"
    command = f"nii2mesh {path_nifti} -p 1 -l 1 -r {args.reduction_factor} {path_init_mesh}"
    subprocess.run(command, shell=True, check=True)

    logger.info("====> Step 3: Preprocessing mesh (normalization and rotation)")
    path_normalized = tmpdir / "mesh_normalized.obj"
    apply_mesh_normalization(path_init_mesh, path_normalized)
    path_rotated = tmpdir / "mesh_normalized_rotated.obj"
    apply_mesh_rotation(path_normalized, path_rotated)

    logger.info("====> Step 4: Generating texture using Text2Tex")
    path_final_mesh = outdir / "42-p36-h20-1.0-0.3-0.1" / "update" / "mesh" / "19_post.obj"
    if not path_final_mesh.exists():
        run_text2tex(path_input=path_rotated, path_output_dir=outdir,
                     prompt=args.prompt, device=args.device)
    
    logger.info("====> Step 5: Generating GIF")
    generate_gif(
        path_obj=path_rotated,
        path_gif=str(outdir / "gif" / "untextured.gif"),
        shader_type="untextured",
    )
    generate_gif(
        path_obj=path_final_mesh,
        path_gif=str(outdir / "gif" / "textured.gif"),
        shader_type="textured",
    )
    
    logger.info(f"====> Finished! Check output directory: {outdir}")
    if not args.intermediate:
        logger.info("Cleaning up intermediate files...")
        shutil.rmtree(tmpdir)


if __name__ == "__main__":
    args = user_input()
    mri2avatar(args)