import math import os from dataclasses import dataclass, field from typing import List, Union import numpy as np import PIL.Image import torch import torch.nn.functional as F import trimesh from einops import rearrange from huggingface_hub import hf_hub_download from omegaconf import OmegaConf from PIL import Image from .models.isosurface import MarchingCubeHelper from .utils import ( BaseModule, ImagePreprocessor, find_class, get_spherical_cameras, scale_tensor, ) class TSR(BaseModule): @dataclass class Config(BaseModule.Config): cond_image_size: int image_tokenizer_cls: str image_tokenizer: dict tokenizer_cls: str tokenizer: dict backbone_cls: str backbone: dict post_processor_cls: str post_processor: dict decoder_cls: str decoder: dict renderer_cls: str renderer: dict cfg: Config @classmethod def from_pretrained( cls, pretrained_model_name_or_path: str, config_name: str, weight_name: str, token=None ): if os.path.isdir(pretrained_model_name_or_path): config_path = os.path.join(pretrained_model_name_or_path, config_name) weight_path = os.path.join(pretrained_model_name_or_path, weight_name) else: config_path = hf_hub_download( repo_id=pretrained_model_name_or_path, filename=config_name, token=token ) weight_path = hf_hub_download( repo_id=pretrained_model_name_or_path, filename=weight_name, token=token ) cfg = OmegaConf.load(config_path) OmegaConf.resolve(cfg) model = cls(cfg) ckpt = torch.load(weight_path, map_location="cpu") if "module" in list(ckpt["state_dict"].keys())[0]: state_dict = {key.replace('module.',''): item for key, item in ckpt["state_dict"].items()} else: state_dict = ckpt["state_dict"] missing_states = set(model.state_dict().keys()) - set(state_dict.keys()) if len(missing_states) > 0: warnings.warn("Missing keys ! : {}".format(missing_states)) model.load_state_dict(state_dict) return model def configure(self): self.image_tokenizer = find_class(self.cfg.image_tokenizer_cls)( self.cfg.image_tokenizer ) self.tokenizer = find_class(self.cfg.tokenizer_cls)(self.cfg.tokenizer) self.backbone = find_class(self.cfg.backbone_cls)(self.cfg.backbone) self.post_processor = find_class(self.cfg.post_processor_cls)( self.cfg.post_processor ) self.decoder = find_class(self.cfg.decoder_cls)(self.cfg.decoder) self.renderer = find_class(self.cfg.renderer_cls)(self.cfg.renderer) self.image_processor = ImagePreprocessor() self.isosurface_helper = None def forward( self, image: Union[ PIL.Image.Image, np.ndarray, torch.FloatTensor, List[PIL.Image.Image], List[np.ndarray], List[torch.FloatTensor], ], device: str, ) -> torch.FloatTensor: rgb_cond = self.image_processor(image, self.cfg.cond_image_size)[:, None].to( device ) batch_size = rgb_cond.shape[0] input_image_tokens: torch.Tensor = self.image_tokenizer( rearrange(rgb_cond, "B Nv H W C -> B Nv C H W", Nv=1), ) input_image_tokens = rearrange( input_image_tokens, "B Nv C Nt -> B (Nv Nt) C", Nv=1 ) tokens: torch.Tensor = self.tokenizer(batch_size) tokens = self.backbone( tokens, encoder_hidden_states=input_image_tokens, ) scene_codes = self.post_processor(self.tokenizer.detokenize(tokens)) return scene_codes def render( self, scene_codes, n_views: int, elevation_deg: float = 0.0, camera_distance: float = 1.9, fovy_deg: float = 40.0, height: int = 256, width: int = 256, return_type: str = "pil", ): rays_o, rays_d = get_spherical_cameras( n_views, elevation_deg, camera_distance, fovy_deg, height, width ) rays_o, rays_d = rays_o.to(scene_codes.device), rays_d.to(scene_codes.device) def process_output(image: torch.FloatTensor): if return_type == "pt": return image elif return_type == "np": return image.detach().cpu().numpy() elif return_type == "pil": return Image.fromarray( (image.detach().cpu().numpy() * 255.0).astype(np.uint8) ) else: raise NotImplementedError images = [] for scene_code in scene_codes: images_ = [] for i in range(n_views): with torch.no_grad(): image = self.renderer( self.decoder, scene_code, rays_o[i], rays_d[i] ) images_.append(process_output(image)) images.append(images_) return images def set_marching_cubes_resolution(self, resolution: int): if ( self.isosurface_helper is not None and self.isosurface_helper.resolution == resolution ): return self.isosurface_helper = MarchingCubeHelper(resolution) def extract_mesh(self, scene_codes, resolution: int = 256, threshold: float = 25.0): self.set_marching_cubes_resolution(resolution) meshes = [] for scene_code in scene_codes: with torch.no_grad(): density = self.renderer.query_triplane( self.decoder, scale_tensor( self.isosurface_helper.grid_vertices.to(scene_codes.device), self.isosurface_helper.points_range, (-self.renderer.cfg.radius, self.renderer.cfg.radius), ), scene_code, )["density_act"] v_pos, t_pos_idx = self.isosurface_helper(-(density - threshold)) v_pos = scale_tensor( v_pos, self.isosurface_helper.points_range, (-self.renderer.cfg.radius, self.renderer.cfg.radius), ) with torch.no_grad(): color = self.renderer.query_triplane( self.decoder, v_pos, scene_code, )["color"] mesh = trimesh.Trimesh( vertices=v_pos.cpu().numpy(), faces=t_pos_idx.cpu().numpy(), vertex_colors=color.cpu().numpy(), ) meshes.append(mesh) return meshes