Spaces:
Sleeping
Sleeping
| 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.RasterizeCudaContext() | |
| 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) | |