from dataclasses import dataclass import cv2 import numpy as np import torch import torch.nn.functional as F from diffusers import DiffusionPipeline from diffusers.utils import BaseOutput def pad_camera_extrinsics_4x4(extrinsics): if extrinsics.shape[-2] == 4: return extrinsics padding = torch.tensor([[0, 0, 0, 1]]).to(extrinsics) if extrinsics.ndim == 3: padding = padding.unsqueeze(0).repeat(extrinsics.shape[0], 1, 1) extrinsics = torch.cat([extrinsics, padding], dim=-2) return extrinsics def center_looking_at_camera_pose( camera_position: torch.Tensor, look_at: torch.Tensor = None, up_world: torch.Tensor = None, ): if look_at is None: look_at = torch.tensor([0, 0, 0], dtype=torch.float32) if up_world is None: up_world = torch.tensor([0, 0, 1], dtype=torch.float32) if camera_position.ndim == 2: look_at = look_at.unsqueeze(0).repeat(camera_position.shape[0], 1) up_world = up_world.unsqueeze(0).repeat(camera_position.shape[0], 1) z_axis = camera_position - look_at z_axis = F.normalize(z_axis, dim=-1).float() x_axis = torch.linalg.cross(up_world, z_axis, dim=-1) x_axis = F.normalize(x_axis, dim=-1).float() y_axis = torch.linalg.cross(z_axis, x_axis, dim=-1) y_axis = F.normalize(y_axis, dim=-1).float() extrinsics = torch.stack([x_axis, y_axis, z_axis, camera_position], dim=-1) extrinsics = pad_camera_extrinsics_4x4(extrinsics) return extrinsics def spherical_camera_pose(azimuths: np.ndarray, elevations: np.ndarray, radius=2.5): azimuths = np.deg2rad(azimuths) elevations = np.deg2rad(elevations) xs = radius * np.cos(elevations) * np.cos(azimuths) ys = radius * np.cos(elevations) * np.sin(azimuths) zs = radius * np.sin(elevations) cam_locations = np.stack([xs, ys, zs], axis=-1) cam_locations = torch.from_numpy(cam_locations).float() c2ws = center_looking_at_camera_pose(cam_locations) return c2ws def FOV_to_intrinsics(fov, device="cpu"): focal_length = 0.5 / np.tan(np.deg2rad(fov) * 0.5) intrinsics = torch.tensor( [[focal_length, 0, 0.5], [0, focal_length, 0.5], [0, 0, 1]], device=device ) return intrinsics def get_zero123plus_input_cameras(batch_size=1, radius=4.0, fov=30.0): azimuths = np.array([30, 90, 150, 210, 270, 330]).astype(float) elevations = np.array([20, -10, 20, -10, 20, -10]).astype(float) c2ws = spherical_camera_pose(azimuths, elevations, radius) c2ws = c2ws.float().flatten(-2) Ks = FOV_to_intrinsics(fov).unsqueeze(0).repeat(6, 1, 1).float().flatten(-2) extrinsics = c2ws[:, :12] intrinsics = torch.stack([Ks[:, 0], Ks[:, 4], Ks[:, 2], Ks[:, 5]], dim=-1) cameras = torch.cat([extrinsics, intrinsics], dim=-1) return cameras.unsqueeze(0).repeat(batch_size, 1, 1) @dataclass class InstantMeshPipelineOutput(BaseOutput): vertices: np.ndarray faces: np.ndarray uvs: np.ndarray texture: np.ndarray class InstantMeshPipeline(DiffusionPipeline): def __init__(self, lrm): super().__init__() self.lrm = lrm self.register_modules(lrm=self.lrm) @torch.no_grad() def __call__(self, images: torch.Tensor): """if remove_bg: image = rembg.remove(image) image = np.array(image) alpha = np.where(image[..., 3] > 0) y1, y2, x1, x2 = ( alpha[0].min(), alpha[0].max(), alpha[1].min(), alpha[1].max(), ) fg = image[y1:y2, x1:x2] 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 image = np.pad( fg, ((ph0, ph1), (pw0, pw1), (0, 0)), mode="constant", constant_values=((0, 0), (0, 0), (0, 0)), ) new_size = int(image.shape[0] / 0.85) ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2 ph1, pw1 = new_size - size - ph0, new_size - size - pw0 image = np.pad( image, ((ph0, ph1), (pw0, pw1), (0, 0)), mode="constant", constant_values=((0, 0), (0, 0), (0, 0)), ) image = Image.fromarray(image) self.multi_view_diffusion = self.multi_view_diffusion.to(self._execution_device) images = self.multi_view_diffusion(image).images[0] images = np.asarray(images, dtype=np.float32) / 255.0 images = torch.from_numpy(images).permute(2, 0, 1).contiguous().float() n, m = 3, 2 c, h, w = images.shape images = ( images.view(c, n, h // n, m, w // m).permute(1, 3, 0, 2, 4).contiguous() ) images = images.view(n * m, c, h // n, w // m) images = images.unsqueeze(0) images = v2.functional.resize( images, 320, interpolation=3, antialias=True ).clamp(0, 1)""" self.lrm.init_flexicubes_geometry(self._execution_device, fovy=30.0) cameras = get_zero123plus_input_cameras().to(self._execution_device) planes = self.lrm.forward_planes(images, cameras) mesh_out = self.lrm.extract_mesh( planes, use_texture_map=True, texture_resolution=1024, ) vertices, vertex_indices, uvs, uv_indices, texture = mesh_out vertices = vertices.cpu().numpy() vertex_indices = vertex_indices.cpu().numpy() uvs = uvs.cpu().numpy() uv_indices = uv_indices.cpu().numpy() texture = texture.permute(1, 2, 0).cpu().numpy() vertex_indices_flat = vertex_indices.reshape(-1) uv_indices_flat = uv_indices.reshape(-1) vertex_uv_pairs = np.stack([vertex_indices_flat, uv_indices_flat], axis=1) unique_pairs, unique_indices = np.unique( vertex_uv_pairs, axis=0, return_inverse=True ) vertices = vertices[unique_pairs[:, 0]] uvs = uvs[unique_pairs[:, 1]] faces = unique_indices.reshape(-1, 3) lo, hi = 0, 1 img = np.asarray(texture, dtype=np.float32) img = (img - lo) * (255 / (hi - lo)) img = img.clip(0, 255) mask = np.sum(img.astype(np.float32), axis=-1, keepdims=True) mask = (mask <= 3.0).astype(np.float32) kernel = np.ones((3, 3), "uint8") dilate_img = cv2.dilate(img, kernel, iterations=1) img = img * (1 - mask) + dilate_img * mask img = img.clip(0, 255).astype(np.uint8) texture = np.ascontiguousarray(img[::-1, :, :]) return InstantMeshPipelineOutput( vertices=vertices, faces=faces, uvs=uvs, texture=texture, )