inference.py

import os
import sys
import io

import torch
import numpy as np
from omegaconf import OmegaConf
import PIL.Image
from PIL import Image
import rembg

from dva.ray_marcher import RayMarcher
from dva.io import load_from_config
from dva.utils import to_device
from dva.visualize import visualize_primvolume, visualize_video_primvolume
from models.diffusion import create_diffusion
import logging
from tqdm import tqdm

import mcubes
import xatlas
import nvdiffrast.torch as dr
import cv2
from scipy.ndimage import binary_dilation, binary_erosion
from sklearn.neighbors import NearestNeighbors
from utils.meshutils import clean_mesh, decimate_mesh
from utils.mesh import Mesh
logger = logging.getLogger("inference.py")


def remove_background(image: PIL.Image.Image,
    rembg_session = None,
    force: bool = False,
    **rembg_kwargs,
) -> PIL.Image.Image:
    do_remove = True
    if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
        do_remove = False
    do_remove = do_remove or force
    if do_remove:
        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
    return image

def resize_foreground(
    image: PIL.Image.Image,
    ratio: float,
) -> PIL.Image.Image:
    image = np.array(image)
    assert image.shape[-1] == 4
    alpha = np.where(image[..., 3] > 0)
    y1, y2, x1, x2 = (
        alpha[0].min(),
        alpha[0].max(),
        alpha[1].min(),
        alpha[1].max(),
    )
    # crop the foreground
    fg = image[y1:y2, x1:x2]
    # pad to square
    size = max(fg.shape[0], fg.shape[1])
    ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 2
    ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw0
    new_image = np.pad(
        fg,
        ((ph0, ph1), (pw0, pw1), (0, 0)),
        mode="constant",
        constant_values=((0, 0), (0, 0), (0, 0)),
    )

    # compute padding according to the ratio
    new_size = int(new_image.shape[0] / ratio)
    # pad to size, double side
    ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
    ph1, pw1 = new_size - size - ph0, new_size - size - pw0
    new_image = np.pad(
        new_image,
        ((ph0, ph1), (pw0, pw1), (0, 0)),
        mode="constant",
        constant_values=((0, 0), (0, 0), (0, 0)),
    )
    new_image = PIL.Image.fromarray(new_image)
    return new_image

def extract_texmesh(args, model, output_path, device):
    # Prepare directory
    ins_dir = output_path

    # Get SDFs
    with torch.no_grad():
        xx = torch.linspace(-1, 1, args.mc_resolution, device=device)
        pts = torch.stack(torch.meshgrid(xx, xx, xx, indexing='ij'), dim=-1).reshape(-1,3)
        chunks = torch.split(pts, args.batch_size)
        dists = []
        for chunk_pts in tqdm(chunks):
            preds = model(chunk_pts)
            dists.append(preds['sdf'].detach())
    dists = torch.cat(dists, dim=0)
    grid = dists.reshape(args.mc_resolution, args.mc_resolution, args.mc_resolution)

    # Meshify
    vertices, triangles = mcubes.marching_cubes(grid.cpu().numpy(), 0.0)
    
    # Resize + recenter
    b_min_np = np.array([-1., -1., -1.])
    b_max_np = np.array([ 1.,  1.,  1.])
    vertices = vertices / (args.mc_resolution - 1.0) * (b_max_np - b_min_np) + b_min_np

    vertices, triangles = clean_mesh(vertices, triangles, min_f=8, min_d=5, repair=True, remesh=False)

    if args.decimate > 0 and triangles.shape[0] > args.decimate:
        vertices, triangles = decimate_mesh(vertices, triangles, args.decimate, remesh=args.remesh)

    h0 = 1024
    w0 = 1024
    ssaa = 1
    fp16 = True
    glctx = dr.RasterizeGLContext(output_db=False)
    v_np = vertices.astype(np.float32)
    f_np = triangles.astype(np.int64)
    v = torch.from_numpy(vertices).float().contiguous().to(device)
    f = torch.from_numpy(triangles.astype(np.int64)).int().contiguous().to(device)
    print(f'[INFO] running xatlas to unwrap UVs for mesh: v={v_np.shape} f={f_np.shape}')
    # unwrap uv in contracted space
    atlas = xatlas.Atlas()
    atlas.add_mesh(v_np, f_np)
    chart_options = xatlas.ChartOptions()
    chart_options.max_iterations = 0 # disable merge_chart for faster unwrap...
    pack_options = xatlas.PackOptions()
    # pack_options.blockAlign = True
    # pack_options.bruteForce = False
    atlas.generate(chart_options=chart_options, pack_options=pack_options)
    vmapping, ft_np, vt_np = atlas[0] # [N], [M, 3], [N, 2]

    vt = torch.from_numpy(vt_np.astype(np.float32)).float().contiguous().to(device)
    ft = torch.from_numpy(ft_np.astype(np.int64)).int().contiguous().to(device)
    uv = vt * 2.0 - 1.0 # uvs to range [-1, 1]
    uv = torch.cat((uv, torch.zeros_like(uv[..., :1]), torch.ones_like(uv[..., :1])), dim=-1) # [N, 4]

    if ssaa > 1:
        h = int(h0 * ssaa)
        w = int(w0 * ssaa)
    else:
        h, w = h0, w0

    rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), ft, (h, w)) # [1, h, w, 4]
    xyzs, _ = dr.interpolate(v.unsqueeze(0), rast, f) # [1, h, w, 3]
    mask, _ = dr.interpolate(torch.ones_like(v[:, :1]).unsqueeze(0), rast, f) # [1, h, w, 1]
    # masked query 
    xyzs = xyzs.view(-1, 3)
    mask = (mask > 0).view(-1)
    feats = torch.zeros(h * w, 6, device=device, dtype=torch.float32)

    if mask.any():
        xyzs = xyzs[mask] # [M, 3]
        # batched inference to avoid OOM
        all_feats = []
        head = 0
        chunk_size = args.batch_size
        while head < xyzs.shape[0]:
            tail = min(head + chunk_size, xyzs.shape[0])
            with torch.cuda.amp.autocast(enabled=fp16):
                preds = model(xyzs[head:tail])
                # [R, G, B, NA, roughness, metallic]
                all_feats.append(torch.concat([preds['tex'].float(), torch.zeros_like(preds['tex'])[..., 0:1].float(), preds['mat'].float()], dim=-1))
            head += chunk_size
        feats[mask] = torch.cat(all_feats, dim=0)
    feats = feats.view(h, w, -1) # 6 channels
    mask = mask.view(h, w)
    # quantize [0.0, 1.0] to [0, 255]
    feats = feats.cpu().numpy()
    feats = (feats * 255)

    ### NN search as a queer antialiasing ...
    mask = mask.cpu().numpy()
    inpaint_region = binary_dilation(mask, iterations=32) # pad width
    inpaint_region[mask] = 0
    search_region = mask.copy()
    not_search_region = binary_erosion(search_region, iterations=3)
    search_region[not_search_region] = 0
    search_coords = np.stack(np.nonzero(search_region), axis=-1)
    inpaint_coords = np.stack(np.nonzero(inpaint_region), axis=-1)
    knn = NearestNeighbors(n_neighbors=1, algorithm='kd_tree').fit(search_coords)
    _, indices = knn.kneighbors(inpaint_coords)
    feats[tuple(inpaint_coords.T)] = feats[tuple(search_coords[indices[:, 0]].T)]
    # do ssaa after the NN search, in numpy
    feats0 = cv2.cvtColor(feats[..., :3].astype(np.uint8), cv2.COLOR_RGB2BGR) # albedo
    feats1 = cv2.cvtColor(feats[..., 3:].astype(np.uint8), cv2.COLOR_RGB2BGR) # visibility features
    if ssaa > 1:
        feats0 = cv2.resize(feats0, (w0, h0), interpolation=cv2.INTER_LINEAR)
        feats1 = cv2.resize(feats1, (w0, h0), interpolation=cv2.INTER_LINEAR)

    cv2.imwrite(os.path.join(ins_dir, f'texture.jpg'), feats0)
    cv2.imwrite(os.path.join(ins_dir, f'roughness_metallic.jpg'), feats1)

    target_mesh = Mesh(v=torch.from_numpy(v_np).contiguous(), f=torch.from_numpy(f_np).contiguous(), ft=ft.contiguous(), vt=torch.from_numpy(vt_np).contiguous(), albedo=torch.from_numpy(feats[..., :3]) / 255, metallicRoughness=torch.from_numpy(feats[..., 3:]) / 255)
    target_mesh.write(os.path.join(ins_dir, f'pbr_mesh.glb'))

def main(config):
    logging.basicConfig(level=logging.INFO)
    ddim_steps = config.inference.ddim
    if ddim_steps > 0:
        use_ddim = True
    else:
        use_ddim = False
    cfg_scale = config.inference.get("cfg", 0.0)

    inference_dir = f"{config.output_dir}/inference_folder"
    os.makedirs(inference_dir, exist_ok=True)

    amp = False
    precision = config.inference.get("precision", 'fp16')
    if precision == 'tf32':
        precision_dtype = torch.float32
    elif precision == 'fp16':
        amp = True
        precision_dtype = torch.float16
    else:
       raise NotImplementedError("{} precision is not supported".format(precision))

    device = torch.device(f"cuda:{0}")
    seed = config.inference.seed
    torch.manual_seed(seed)
    torch.cuda.set_device(device)
    
    model = load_from_config(config.model.generator)
    vae = load_from_config(config.model.vae)
    conditioner = load_from_config(config.model.conditioner)
    vae_state_dict = torch.load(config.model.vae_checkpoint_path, map_location='cpu')
    vae.load_state_dict(vae_state_dict['model_state_dict'])
    
    if config.checkpoint_path:
        state_dict = torch.load(config.checkpoint_path, map_location='cpu')
        model.load_state_dict(state_dict['ema'])
    vae = vae.to(device)
    conditioner = conditioner.to(device)
    model = model.to(device)
    config.diffusion.pop("timestep_respacing")
    if use_ddim:
        respacing = "ddim{}".format(ddim_steps)
    else:
        respacing = ""
    diffusion = create_diffusion(timestep_respacing=respacing, **config.diffusion)  # default: 1000 steps, linear noise schedule
    if use_ddim:
        sample_fn = diffusion.ddim_sample_loop_progressive
    else:
        sample_fn = diffusion.p_sample_loop_progressive
    
    if cfg_scale > 0:
        fwd_fn = model.forward_with_cfg
    else:
        fwd_fn = model.forward

    rm = RayMarcher(
        config.image_height,
        config.image_width,
        **config.rm,
    ).to(device)

    perchannel_norm = False
    if "latent_mean" in config.model:
        latent_mean = torch.Tensor(config.model.latent_mean)[None, None, :].to(device)
        latent_std = torch.Tensor(config.model.latent_std)[None, None, :].to(device)
        assert latent_mean.shape[-1] == config.model.generator.in_channels
        perchannel_norm = True

    model.eval()
    examples_dir = config.inference.input_dir
    img_list = os.listdir(examples_dir)
    rembg_session = rembg.new_session()
    logger.info(f"Starting Inference...")
    for img_path in img_list:
        full_img_path = os.path.join(examples_dir, img_path)
        img_name = img_path[:-4]
        current_output_dir = os.path.join(inference_dir, img_name)
        os.makedirs(current_output_dir, exist_ok=True)
        input_image = Image.open(full_img_path)
        input_image = remove_background(input_image, rembg_session)
        input_image = resize_foreground(input_image, 0.85)
        raw_image = np.array(input_image)
        mask = (raw_image[..., -1][..., None] > 0) * 1
        raw_image = raw_image[..., :3] * mask
        input_cond = torch.from_numpy(np.array(raw_image)[None, ...]).to(device)
        with torch.no_grad():
            latent = torch.randn(1, config.model.num_prims, 1, 4, 4, 4)
            batch = {}
            inf_bs = 1
            inf_x = torch.randn(inf_bs, config.model.num_prims, 68).to(device)
            y = conditioner.encoder(input_cond)
            model_kwargs = dict(y=y[:inf_bs, ...], precision_dtype=precision_dtype, enable_amp=amp)
            if cfg_scale > 0:
                model_kwargs['cfg_scale'] = cfg_scale
            sampled_count = -1
            for samples in sample_fn(fwd_fn, inf_x.shape, inf_x, clip_denoised=False, model_kwargs=model_kwargs, progress=True, device=device
            ):
                sampled_count += 1
                if not (sampled_count % 10 == 0 or sampled_count == diffusion.num_timesteps - 1):
                    continue
                else:
                    recon_param = samples["sample"].reshape(inf_bs, config.model.num_prims, -1)
                    if perchannel_norm:
                        recon_param = recon_param / config.model.latent_nf * latent_std + latent_mean
                    recon_srt_param = recon_param[:, :, 0:4]
                    recon_feat_param = recon_param[:, :, 4:] # [8, 2048, 64]
                    recon_feat_param_list = []
                    # one-by-one to avoid oom
                    for inf_bidx in range(inf_bs):
                        if not perchannel_norm:
                            decoded = vae.decode(recon_feat_param[inf_bidx, ...].reshape(1*config.model.num_prims, *latent.shape[-4:]) / config.model.latent_nf)
                        else:
                            decoded = vae.decode(recon_feat_param[inf_bidx, ...].reshape(1*config.model.num_prims, *latent.shape[-4:]))
                        recon_feat_param_list.append(decoded.detach())
                    recon_feat_param = torch.concat(recon_feat_param_list, dim=0)
                    # invert normalization
                    if not perchannel_norm:
                        recon_srt_param[:, :, 0:1] = (recon_srt_param[:, :, 0:1] / 10) + 0.05
                    recon_feat_param[:, 0:1, ...] /= 5.
                    recon_feat_param[:, 1:, ...] = (recon_feat_param[:, 1:, ...] + 1) / 2.
                    recon_feat_param = recon_feat_param.reshape(inf_bs, config.model.num_prims, -1)
                    recon_param = torch.concat([recon_srt_param, recon_feat_param], dim=-1)
                    visualize_primvolume("{}/dstep{:04d}_recon.jpg".format(current_output_dir, sampled_count), batch, recon_param, rm, device)
            visualize_video_primvolume(current_output_dir, batch, recon_param, 60, rm, device)
            prim_params = {'srt_param': recon_srt_param[0].detach().cpu(), 'feat_param': recon_feat_param[0].detach().cpu()}
            torch.save({'model_state_dict': prim_params}, "{}/denoised.pt".format(current_output_dir))

    if config.inference.export_glb:
        logger.info(f"Starting GLB Mesh Extraction...")
        config.model.pop("vae")
        config.model.pop("vae_checkpoint_path")
        config.model.pop("conditioner")
        config.model.pop("generator")
        config.model.pop("latent_nf")
        config.model.pop("latent_mean")
        config.model.pop("latent_std")
        model_primx = load_from_config(config.model)
        for img_path in img_list:
            img_name = img_path[:-4]
            output_path = os.path.join(inference_dir, img_name)
            denoise_param_path = os.path.join(inference_dir, img_name, 'denoised.pt')
            ckpt_weight = torch.load(denoise_param_path, map_location='cpu')['model_state_dict']
            model_primx.load_state_dict(ckpt_weight)
            model_primx.to(device)
            model_primx.eval()
            with torch.no_grad():
                model_primx.srt_param[:, 1:4] *= 0.85
                extract_texmesh(config.inference, model_primx, output_path, device)

if __name__ == "__main__":
    torch.backends.cudnn.benchmark = True
    # manually enable tf32 to get speedup on A100 GPUs
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True
    os.environ["CC"] = "/mnt/lustre/share/gcc/gcc-8.5.0/bin/gcc"
    os.environ["CPP"] = "/mnt/lustre/share/gcc/gcc-8.5.0/bin/g++"
    os.environ["CXX"] = "/mnt/lustre/share/gcc/gcc-8.5.0/bin/g++"
    # set config
    config = OmegaConf.load(str(sys.argv[1]))
    config_cli = OmegaConf.from_cli(args_list=sys.argv[2:])
    if config_cli:
        logger.info("overriding with following values from args:")
        logger.info(OmegaConf.to_yaml(config_cli))
        config = OmegaConf.merge(config, config_cli)

    main(config)