Delete freesplatter/utils/mesh_renderer.py

freesplatter/utils/mesh_renderer.py

@@ -1,608 +0,0 @@
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import nvdiffrast.torch as dr
-
-
-def get_ray_directions(h, w, intrinsics, norm=False, device=None):
-    """
-    Args:
-        h (int)
-        w (int)
-        intrinsics: (*, 4), in [fx, fy, cx, cy]
-
-    Returns:
-        directions: (*, h, w, 3), the direction of the rays in camera coordinate
-    """
-    batch_size = intrinsics.shape[:-1]
-    x = torch.linspace(0.5, w - 0.5, w, device=device)
-    y = torch.linspace(0.5, h - 0.5, h, device=device)
-    # (*, h, w, 2)
-    directions_xy = torch.stack(
-        [((x - intrinsics[..., 2:3]) / intrinsics[..., 0:1])[..., None, :].expand(*batch_size, h, w),
-         ((y - intrinsics[..., 3:4]) / intrinsics[..., 1:2])[..., :, None].expand(*batch_size, h, w)], dim=-1)
-    # (*, h, w, 3)
-    directions = F.pad(directions_xy, [0, 1], mode='constant', value=1.0)
-    if norm:
-        directions = F.normalize(directions, dim=-1)
-    return directions
-
-
-def edge_dilation(img, mask, radius=3, iter=7):
-    """
-    Args:
-        img (torch.Tensor): (n, c, h, w)
-        mask (torch.Tensor): (n, 1, h, w)
-        radius (float): Radius of dilation.
-
-    Returns:
-        torch.Tensor: Dilated image.
-    """
-    n, c, h, w = img.size()
-    int_radius = round(radius)
-    kernel_size = int(int_radius * 2 + 1)
-    distance1d_sq = torch.linspace(-int_radius, int_radius, kernel_size, dtype=img.dtype, device=img.device).square()
-    kernel_distance = (distance1d_sq.reshape(1, -1) + distance1d_sq.reshape(-1, 1)).sqrt()
-    kernel_neg_distance = kernel_distance.max() - kernel_distance + 1
-
-    for _ in range(iter):
-
-        mask_out = F.max_pool2d(mask, kernel_size, stride=1, padding=int_radius)
-        do_fill_mask = ((mask_out - mask) > 0.5).squeeze(1)
-        # (num_fill, 3) in [ind_n, ind_h, ind_w]
-        do_fill = do_fill_mask.nonzero()
-
-        # unfold the image and mask
-        mask_unfold = F.unfold(mask, kernel_size, padding=int_radius).reshape(
-            n, kernel_size * kernel_size, h, w).permute(0, 2, 3, 1)
-
-        fill_ind = (mask_unfold[do_fill_mask] * kernel_neg_distance.flatten()).argmax(dim=-1)
-        do_fill_h = do_fill[:, 1] + fill_ind // kernel_size - int_radius
-        do_fill_w = do_fill[:, 2] + fill_ind % kernel_size - int_radius
-
-        img_out = img.clone()
-        img_out[do_fill[:, 0], :, do_fill[:, 1], do_fill[:, 2]] = img[
-            do_fill[:, 0], :, do_fill_h, do_fill_w]
-
-        img = img_out
-        mask = mask_out
-
-    return img
-
-
-def depth_to_normal(depth, directions, format='opengl'):
-    """
-    Args:
-        depth: shape (*, h, w), inverse depth defined as 1 / z
-        directions: shape (*, h, w, 3), unnormalized ray directions, under OpenCV coordinate system
-
-    Returns:
-        out_normal: shape (*, h, w, 3), in range [0, 1]
-    """
-    out_xyz = directions / depth.unsqueeze(-1).clamp(min=1e-6)
-    dx = out_xyz[..., :, 1:, :] - out_xyz[..., :, :-1, :]
-    dy = out_xyz[..., 1:, :, :] - out_xyz[..., :-1, :, :]
-    right = F.pad(dx, (0, 0, 0, 1, 0, 0), mode='replicate')
-    up = F.pad(-dy, (0, 0, 0, 0, 1, 0), mode='replicate')
-    left = F.pad(-dx, (0, 0, 1, 0, 0, 0), mode='replicate')
-    down = F.pad(dy, (0, 0, 0, 0, 0, 1), mode='replicate')
-    out_normal = F.normalize(
-        F.normalize(torch.cross(right, up, dim=-1), dim=-1)
-        + F.normalize(torch.cross(up, left, dim=-1), dim=-1)
-        + F.normalize(torch.cross(left, down, dim=-1), dim=-1)
-        + F.normalize(torch.cross(down, right, dim=-1), dim=-1),
-        dim=-1)
-    if format == 'opengl':
-        out_normal[..., 1:3] = -out_normal[..., 1:3]  # to opengl coord
-    elif format == 'opencv':
-        out_normal = out_normal
-    else:
-        raise ValueError('format should be opengl or opencv')
-    out_normal = out_normal / 2 + 0.5
-    return out_normal
-
-
-def make_divisible(x, m=8):
-    return int(math.ceil(x / m) * m)
-
-
-def interpolate_hwc(x, scale_factor, mode='area'):
-    batch_dim = x.shape[:-3]
-    y = x.reshape(batch_dim.numel(), *x.shape[-3:]).permute(0, 3, 1, 2)
-    y = F.interpolate(y, scale_factor=scale_factor, mode=mode).permute(0, 2, 3, 1)
-    return y.reshape(*batch_dim, *y.shape[1:])
-
-
-def compute_edge_to_face_mapping(attr_idx):
-    with torch.no_grad():
-        # Get unique edges
-        # Create all edges, packed by triangle
-        all_edges = torch.cat((
-            torch.stack((attr_idx[:, 0], attr_idx[:, 1]), dim=-1),
-            torch.stack((attr_idx[:, 1], attr_idx[:, 2]), dim=-1),
-            torch.stack((attr_idx[:, 2], attr_idx[:, 0]), dim=-1),
-        ), dim=-1).view(-1, 2)
-
-        # Swap edge order so min index is always first
-        order = (all_edges[:, 0] > all_edges[:, 1]).long().unsqueeze(dim=1)
-        sorted_edges = torch.cat((
-            torch.gather(all_edges, 1, order),
-            torch.gather(all_edges, 1, 1 - order)
-        ), dim=-1)
-
-        # Elliminate duplicates and return inverse mapping
-        unique_edges, idx_map = torch.unique(sorted_edges, dim=0, return_inverse=True)
-
-        tris = torch.arange(attr_idx.shape[0]).repeat_interleave(3).cuda()
-
-        tris_per_edge = torch.zeros((unique_edges.shape[0], 2), dtype=torch.int64).cuda()
-
-        # Compute edge to face table
-        mask0 = order[:,0] == 0
-        mask1 = order[:,0] == 1
-        tris_per_edge[idx_map[mask0], 0] = tris[mask0]
-        tris_per_edge[idx_map[mask1], 1] = tris[mask1]
-
-        return tris_per_edge
-
-
-@torch.cuda.amp.autocast(enabled=False)
-def normal_consistency(face_normals, t_pos_idx):
-
-    tris_per_edge = compute_edge_to_face_mapping(t_pos_idx)
-
-    # Fetch normals for both faces sharind an edge
-    n0 = face_normals[tris_per_edge[:, 0], :]
-    n1 = face_normals[tris_per_edge[:, 1], :]
-
-    # Compute error metric based on normal difference
-    term = torch.clamp(torch.sum(n0 * n1, -1, keepdim=True), min=-1.0, max=1.0)
-    term = (1.0 - term)
-
-    return torch.mean(torch.abs(term))
-
-
-def laplacian_uniform(verts, faces):
-
-    V = verts.shape[0]
-    F = faces.shape[0]
-
-    # Neighbor indices
-    ii = faces[:, [1, 2, 0]].flatten()
-    jj = faces[:, [2, 0, 1]].flatten()
-    adj = torch.stack([torch.cat([ii, jj]), torch.cat([jj, ii])], dim=0).unique(dim=1)
-    adj_values = torch.ones(adj.shape[1], device=verts.device, dtype=torch.float)
-
-    # Diagonal indices
-    diag_idx = adj[0]
-
-    # Build the sparse matrix
-    idx = torch.cat((adj, torch.stack((diag_idx, diag_idx), dim=0)), dim=1)
-    values = torch.cat((-adj_values, adj_values))
-
-    # The coalesce operation sums the duplicate indices, resulting in the
-    # correct diagonal
-    return torch.sparse_coo_tensor(idx, values, (V,V)).coalesce()
-
-
-@torch.cuda.amp.autocast(enabled=False)
-def laplacian_smooth_loss(verts, faces):
-    with torch.no_grad():
-        L = laplacian_uniform(verts, faces.long())
-    loss = L.mm(verts)
-    loss = loss.norm(dim=1)
-    loss = loss.mean()
-    return loss
-
-
-class DMTet:
-
-    def __init__(self, device):
-        self.device = device
-        self.triangle_table = torch.tensor([
-            [-1, -1, -1, -1, -1, -1],
-            [1, 0, 2, -1, -1, -1],
-            [4, 0, 3, -1, -1, -1],
-            [1, 4, 2, 1, 3, 4],
-            [3, 1, 5, -1, -1, -1],
-            [2, 3, 0, 2, 5, 3],
-            [1, 4, 0, 1, 5, 4],
-            [4, 2, 5, -1, -1, -1],
-            [4, 5, 2, -1, -1, -1],
-            [4, 1, 0, 4, 5, 1],
-            [3, 2, 0, 3, 5, 2],
-            [1, 3, 5, -1, -1, -1],
-            [4, 1, 2, 4, 3, 1],
-            [3, 0, 4, -1, -1, -1],
-            [2, 0, 1, -1, -1, -1],
-            [-1, -1, -1, -1, -1, -1]
-        ], dtype=torch.long, device=device)
-        self.num_triangles_table = torch.tensor([0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long,
-                                                device=device)
-        self.base_tet_edges = torch.tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long, device=device)
-
-    def sort_edges(self, edges_ex2):
-        with torch.no_grad():
-            order = (edges_ex2[:, 0] > edges_ex2[:, 1]).long()
-            order = order.unsqueeze(dim=1)
-
-            a = torch.gather(input=edges_ex2, index=order, dim=1)
-            b = torch.gather(input=edges_ex2, index=1 - order, dim=1)
-
-        return torch.stack([a, b], -1)
-
-    def __call__(self, pos_nx3, sdf_n, tet_fx4):
-        # pos_nx3: [N, 3]
-        # sdf_n:   [N]
-        # tet_fx4: [F, 4]
-
-        with torch.no_grad():
-            occ_n = sdf_n > 0
-            occ_fx4 = occ_n[tet_fx4.reshape(-1)].reshape(-1, 4)
-            occ_sum = torch.sum(occ_fx4, -1)  # [F,]
-            valid_tets = (occ_sum > 0) & (occ_sum < 4)
-            # occ_sum = occ_sum[valid_tets]
-
-            # find all vertices
-            all_edges = tet_fx4[valid_tets][:, self.base_tet_edges].reshape(-1, 2)
-            all_edges = self.sort_edges(all_edges)
-            unique_edges, idx_map = torch.unique(all_edges, dim=0, return_inverse=True)
-
-            unique_edges = unique_edges.long()
-            mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1
-            mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device=self.device) * -1
-            mapping[mask_edges] = torch.arange(mask_edges.sum(), dtype=torch.long, device=self.device)
-            idx_map = mapping[idx_map]  # map edges to verts
-
-            interp_v = unique_edges[mask_edges]
-
-        edges_to_interp = pos_nx3[interp_v.reshape(-1)].reshape(-1, 2, 3)
-        edges_to_interp_sdf = sdf_n[interp_v.reshape(-1)].reshape(-1, 2, 1)
-        edges_to_interp_sdf[:, -1] *= -1
-
-        denominator = edges_to_interp_sdf.sum(1, keepdim=True)
-
-        edges_to_interp_sdf = torch.flip(edges_to_interp_sdf, [1]) / denominator
-        verts = (edges_to_interp * edges_to_interp_sdf).sum(1)
-
-        idx_map = idx_map.reshape(-1, 6)
-
-        v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=self.device))
-        tetindex = (occ_fx4[valid_tets] * v_id.unsqueeze(0)).sum(-1)
-        num_triangles = self.num_triangles_table[tetindex]
-
-        # Generate triangle indices
-        faces = torch.cat((
-            torch.gather(input=idx_map[num_triangles == 1], dim=1,
-                         index=self.triangle_table[tetindex[num_triangles == 1]][:, :3]).reshape(-1, 3),
-            torch.gather(input=idx_map[num_triangles == 2], dim=1,
-                         index=self.triangle_table[tetindex[num_triangles == 2]][:, :6]).reshape(-1, 3),
-        ), dim=0)
-
-        return verts, faces
-
-
-class MeshRenderer(nn.Module):
-    def __init__(self,
-                 near=0.1,
-                 far=10,
-                 ssaa=1,
-                 texture_filter='linear-mipmap-linear',
-                 opengl=False,
-                 device='cuda'):
-        super().__init__()
-        self.near = near
-        self.far = far
-        assert isinstance(ssaa, int) and ssaa >= 1
-        self.ssaa = ssaa
-        self.texture_filter = texture_filter
-        self.glctx = dr.RasterizeGLContext(output_db=False)
-
-    def forward(self, meshes, poses, intrinsics, h, w, shading_fun=None,
-                dilate_edges=0, normal_bg=[0.5, 0.5, 1.0], aa=True, render_vc=False):
-        """
-        Args:
-            meshes (list[Mesh]): list of Mesh objects
-            poses: Shape (num_scenes, num_images, 3, 4)
-            intrinsics: Shape (num_scenes, num_images, 4) in [fx, fy, cx, cy]
-        """
-        num_scenes, num_images, _, _ = poses.size()
-
-        if self.ssaa > 1:
-            h = h * self.ssaa
-            w = w * self.ssaa
-            intrinsics = intrinsics * self.ssaa
-
-        r_mat_c2w = torch.cat(
-            [poses[..., :3, :1], -poses[..., :3, 1:3]], dim=-1)  # opencv to opengl conversion
-
-        proj = poses.new_zeros([num_scenes, num_images, 4, 4])
-        proj[..., 0, 0] = 2 * intrinsics[..., 0] / w
-        proj[..., 0, 2] = -2 * intrinsics[..., 2] / w + 1
-        proj[..., 1, 1] = -2 * intrinsics[..., 1] / h
-        proj[..., 1, 2] = -2 * intrinsics[..., 3] / h + 1
-        proj[..., 2, 2] = -(self.far + self.near) / (self.far - self.near)
-        proj[..., 2, 3] = -(2 * self.far * self.near) / (self.far - self.near)
-        proj[..., 3, 2] = -1
-
-        # (num_scenes, (num_images, num_vertices, 3))
-        v_cam = [(mesh.v - poses[i, :, :3, 3].unsqueeze(-2)) @ r_mat_c2w[i] for i, mesh in enumerate(meshes)]
-        # (num_scenes, (num_images, num_vertices, 4))
-        v_clip = [F.pad(v, pad=(0, 1), mode='constant', value=1.0) @ proj[i].transpose(-1, -2) for i, v in enumerate(v_cam)]
-
-        if num_scenes == 1:
-            # (num_images, h, w, 4) in [u, v, z/w, triangle_id] & (num_images, h, w, 4 or 0)
-            rast, rast_db = dr.rasterize(
-                self.glctx, v_clip[0], meshes[0].f, (h, w), grad_db=torch.is_grad_enabled())
-
-            fg = (rast[..., 3] > 0).unsqueeze(0)  # (num_scenes, num_images, h, w)
-            alpha = fg.float().unsqueeze(-1)
-
-            depth = 1 / dr.interpolate(
-                -v_cam[0][..., 2:3].contiguous(), rast, meshes[0].f)[0].reshape(num_scenes, num_images, h, w)
-            depth.masked_fill_(~fg, 0)
-
-            normal = dr.interpolate(
-                meshes[0].vn.unsqueeze(0).contiguous(), rast, meshes[0].fn)[0].reshape(num_scenes, num_images, h, w, 3)
-            normal = F.normalize(normal, dim=-1)
-            # (num_scenes, num_images, h, w, 3) = (num_scenes, num_images, h, w, 3) @ (num_scenes, num_images, 1, 3, 3)
-            rot_normal = (normal @ r_mat_c2w.unsqueeze(2)) / 2 + 0.5
-            rot_normal[~fg] = rot_normal.new_tensor(normal_bg)
-
-            if meshes[0].vt is not None and meshes[0].albedo is not None:
-                # (num_images, h, w, 2) & (num_images, h, w, 4)
-                texc, texc_db = dr.interpolate(
-                    meshes[0].vt.unsqueeze(0).contiguous(), rast, meshes[0].ft, rast_db=rast_db, diff_attrs='all')
-                # (num_scenes, num_images, h, w, 3)
-                albedo = dr.texture(
-                    meshes[0].albedo.unsqueeze(0)[..., :3].contiguous(), texc, uv_da=texc_db, filter_mode=self.texture_filter).unsqueeze(0)
-                albedo[~fg] = 0
-            elif meshes[0].vc is not None:
-                rgba = dr.interpolate(
-                    meshes[0].vc.contiguous(), rast, meshes[0].f)[0].reshape(num_scenes, num_images, h, w, 4)
-                alpha = alpha * rgba[..., 3:4]
-                albedo = rgba[..., :3] * alpha
-            else:
-                albedo = torch.zeros_like(rot_normal)
-
-            prev_grad_enabled = torch.is_grad_enabled()
-            torch.set_grad_enabled(True)
-            if shading_fun is not None:
-                xyz = dr.interpolate(
-                    meshes[0].v.unsqueeze(0).contiguous(), rast, meshes[0].f)[0].reshape(num_scenes, num_images, h, w, 3)
-                rgb_reshade = shading_fun(
-                    world_pos=xyz[fg],
-                    albedo=albedo[fg],
-                    world_normal=normal[fg],
-                    fg_mask=fg)
-                albedo = torch.zeros_like(albedo)
-                albedo[fg] = rgb_reshade
-
-            # (num_scenes, num_images, h, w, 4)
-            rgba = torch.cat([albedo, alpha], dim=-1)
-
-            if dilate_edges > 0:
-                rgba = rgba.reshape(num_scenes * num_images, h, w, 4).permute(0, 3, 1, 2)
-                rgba = edge_dilation(rgba, rgba[:, 3:], dilate_edges)
-                rgba = rgba.permute(0, 2, 3, 1).reshape(num_scenes, num_images, h, w, 4)
-
-            if aa:
-                rgba, depth, rot_normal = dr.antialias(
-                    torch.cat([rgba, depth.unsqueeze(-1), rot_normal], dim=-1).squeeze(0),
-                    rast, v_clip[0], meshes[0].f).unsqueeze(0).split([4, 1, 3], dim=-1)
-                depth = depth.squeeze(-1)
-
-        else:  # concat and range mode
-            # v_cat = []
-            v_clip_cat = []
-            v_cam_cat = []
-            vn_cat = []
-            vt_cat = []
-            f_cat = []
-            fn_cat = []
-            ft_cat = []
-            v_count = 0
-            vn_count = 0
-            vt_count = 0
-            f_count = 0
-            f_ranges = []
-            for i, mesh in enumerate(meshes):
-                num_v = v_clip[i].size(1)
-                num_vn = mesh.vn.size(0)
-                num_vt = mesh.vt.size(0)
-                # v_cat.append(mesh.v.unsqueeze(0).expand(num_images, -1, -1).reshape(num_images * num_v, 3))
-                v_clip_cat.append(v_clip[i].reshape(num_images * num_v, 4))
-                v_cam_cat.append(v_cam[i].reshape(num_images * num_v, 3))
-                vn_cat.append(mesh.vn.unsqueeze(0).expand(num_images, -1, -1).reshape(num_images * num_vn, 3))
-                vt_cat.append(mesh.vt.unsqueeze(0).expand(num_images, -1, -1).reshape(num_images * num_vt, 2))
-                for _ in range(num_images):
-                    f_cat.append(mesh.f + v_count)
-                    fn_cat.append(mesh.fn + vn_count)
-                    ft_cat.append(mesh.ft + vt_count)
-                    v_count += num_v
-                    vn_count += num_vn
-                    vt_count += num_vt
-                    f_ranges.append([f_count, mesh.f.size(0)])
-                    f_count += mesh.f.size(0)
-            # v_cat = torch.cat(v_cat, dim=0)
-            v_clip_cat = torch.cat(v_clip_cat, dim=0)
-            v_cam_cat = torch.cat(v_cam_cat, dim=0)
-            vn_cat = torch.cat(vn_cat, dim=0)
-            f_cat = torch.cat(f_cat, dim=0)
-            f_ranges = torch.tensor(f_ranges, device=poses.device, dtype=torch.int32)
-            # (num_scenes * num_images, h, w, 4) in [u, v, z/w, triangle_id] & (num_scenes * num_images, h, w, 4 or 0)
-            rast, rast_db = dr.rasterize(
-                self.glctx, v_clip_cat, f_cat, (h, w), ranges=f_ranges, grad_db=torch.is_grad_enabled())
-
-            fg = (rast[..., 3] > 0).reshape(num_scenes, num_images, h, w)
-
-            depth = 1 / dr.interpolate(
-                -v_cam_cat[..., 2:3].contiguous(), rast, f_cat)[0].reshape(num_scenes, num_images, h, w)
-            depth.masked_fill_(~fg, 0)
-
-            normal = dr.interpolate(
-                vn_cat, rast, fn_cat)[0].reshape(num_scenes, num_images, h, w, 3)
-            normal = F.normalize(normal, dim=-1)
-            # (num_scenes, num_images, h, w, 3) = (num_scenes, num_images, h, w, 3) @ (num_scenes, num_images, 1, 3, 3)
-            rot_normal = (normal @ r_mat_c2w.unsqueeze(2)) / 2 + 0.5
-            rot_normal[~fg] = rot_normal.new_tensor(normal_bg)
-
-            # (num_scenes * num_images, h, w, 2) & (num_scenes * num_images, h, w, 4)
-            texc, texc_db = dr.interpolate(
-                vt_cat, rast, ft_cat, rast_db=rast_db, diff_attrs='all')
-            albedo = dr.texture(
-                torch.cat([mesh.albedo.unsqueeze(0)[..., :3].expand(num_images, -1, -1, -1) for mesh in meshes], dim=0),
-                texc, uv_da=texc_db, filter_mode=self.texture_filter
-            ).reshape(num_scenes, num_images, h, w, 3)
-
-            prev_grad_enabled = torch.is_grad_enabled()
-            torch.set_grad_enabled(True)
-            if shading_fun is not None:
-                raise NotImplementedError
-
-            # (num_scenes, num_images, h, w, 4)
-            rgba = torch.cat([albedo, fg.float().unsqueeze(-1)], dim=-1)
-
-            if dilate_edges > 0:
-                rgba = rgba.reshape(num_scenes * num_images, h, w, 4).permute(0, 3, 1, 2)
-                rgba = edge_dilation(rgba, rgba[:, 3:], dilate_edges)
-                rgba = rgba.permute(0, 2, 3, 1).reshape(num_scenes, num_images, h, w, 4)
-
-            if aa:
-                # Todo: depth/normal antialiasing
-                rgba = dr.antialias(
-                    rgba.reshape(num_scenes * num_images, h, w, 4), rast, v_clip_cat, f_cat
-                ).reshape(num_scenes, num_images, h, w, 4)
-
-        if self.ssaa > 1:
-            rgba = interpolate_hwc(rgba, 1 / self.ssaa)
-            depth = interpolate_hwc(depth.unsqueeze(-1), 1 / self.ssaa).squeeze(-1)
-            rot_normal = interpolate_hwc(rot_normal, 1 / self.ssaa)
-
-        results = dict(
-            rgba=rgba,
-            depth=depth,
-            normal=rot_normal)
-
-        torch.set_grad_enabled(prev_grad_enabled)
-
-        return results
-
-    def bake_xyz_shading_fun(self, meshes, shading_fun, map_size=1024, force_auto_uv=False):
-        assert len(meshes) == 1, 'only support one mesh'
-        mesh = meshes[0]
-
-        if mesh.vt is None or force_auto_uv:
-            mesh.auto_uv()
-        assert len(mesh.ft) == len(mesh.f)
-
-        vt_clip = torch.cat([mesh.vt * 2 - 1, mesh.vt.new_tensor([[0., 1.]]).expand(mesh.vt.size(0), -1)], dim=-1)
-
-        rast = dr.rasterize(self.glctx, vt_clip[None], mesh.ft, (map_size, map_size), grad_db=False)[0]
-        valid = (rast[..., 3] > 0).reshape(map_size, map_size)
-
-        xyz = dr.interpolate(mesh.v[None], rast, mesh.f)[0].reshape(map_size, map_size, 3)
-        rgb_reshade = shading_fun(world_pos=xyz[valid])
-        new_albedo_map = xyz.new_zeros((map_size, map_size, 3))
-        new_albedo_map[valid] = rgb_reshade
-        torch.cuda.empty_cache()
-        new_albedo_map = edge_dilation(
-            new_albedo_map.permute(2, 0, 1)[None], valid[None, None].float(),
-        ).squeeze(0).permute(1, 2, 0)
-        mesh.albedo = torch.cat(
-            [new_albedo_map.clamp(min=0, max=1),
-             torch.ones_like(new_albedo_map[..., :1])], dim=-1)
-
-        mesh.textureless = False
-        return [mesh]
-
-    def bake_multiview(self, meshes, images, alphas, poses, intrinsics, map_size=1024, cos_weight_pow=4.0):
-        assert len(meshes) == 1, 'only support one mesh'
-        mesh = meshes[0]
-        images = images[0]  # (n, h, w, 3)
-        alphas = alphas[0]  # (n, h, w, 1)
-        n, h, w, _ = images.size()
-
-        r_mat_c2w = torch.cat(
-            [poses[..., :3, :1], -poses[..., :3, 1:3]], dim=-1)[0]  # opencv to opengl conversion
-
-        proj = poses.new_zeros([n, 4, 4])
-        proj[..., 0, 0] = 2 * intrinsics[..., 0] / w
-        proj[..., 0, 2] = -2 * intrinsics[..., 2] / w + 1
-        proj[..., 1, 1] = -2 * intrinsics[..., 1] / h
-        proj[..., 1, 2] = -2 * intrinsics[..., 3] / h + 1
-        proj[..., 2, 2] = -(self.far + self.near) / (self.far - self.near)
-        proj[..., 2, 3] = -(2 * self.far * self.near) / (self.far - self.near)
-        proj[..., 3, 2] = -1
-
-        # (num_images, num_vertices, 3)
-        v_cam = (mesh.v.detach() - poses[0, :, :3, 3].unsqueeze(-2)) @ r_mat_c2w
-        # (num_images, num_vertices, 4)
-        v_clip = F.pad(v_cam, pad=(0, 1), mode='constant', value=1.0) @ proj.transpose(-1, -2)
-
-        rast, rast_db = dr.rasterize(self.glctx, v_clip, mesh.f, (h, w), grad_db=False)
-        texc, texc_db = dr.interpolate(
-            mesh.vt.unsqueeze(0).contiguous(), rast, mesh.ft, rast_db=rast_db, diff_attrs='all')
-
-        with torch.enable_grad():
-            dummy_maps = torch.ones((n, map_size, map_size, 1), device=images.device, dtype=images.dtype).requires_grad_(True)
-            # (num_images, h, w, 1)
-            albedo = dr.texture(
-                dummy_maps, texc, uv_da=texc_db, filter_mode=self.texture_filter)
-            visibility_grad = torch.autograd.grad(albedo.sum(), dummy_maps, create_graph=False)[0]
-
-        fg = rast[..., 3] > 0  # (num_images, h, w)
-        depth = 1 / dr.interpolate(
-            -v_cam[..., 2:3].contiguous(), rast, mesh.f)[0].reshape(n, h, w)
-        depth.masked_fill_(~fg, 0)
-
-        # # save all the depth maps for visualization debug
-        # import matplotlib.pyplot as plt
-        # for i in range(n):
-        #     plt.imshow(depth[i].cpu().numpy())
-        #     plt.savefig(f'depth_{i}.png')
-        # # also save the alphas
-        # for i in range(n):
-        #     plt.imshow(alphas[i].cpu().numpy())
-        #     plt.savefig(f'alpha_{i}.png')
-
-        directions = get_ray_directions(
-            h, w, intrinsics.squeeze(0), norm=True, device=intrinsics.device)
-
-        normals_opencv = depth_to_normal(
-            depth, directions, format='opencv') * 2 - 1
-        normals_cos_weight = (normals_opencv[..., None, :] @ directions[..., :, None]).squeeze(-1).neg().clamp(min=0)
-
-        img_space_weight = (normals_cos_weight ** cos_weight_pow) * alphas
-        img_space_weight = -F.max_pool2d(  # alleviate edge effect
-            -img_space_weight.permute(0, 3, 1, 2), 5, stride=1, padding=2).permute(0, 2, 3, 1)
-
-        # bake texture
-        vt_clip = torch.cat([mesh.vt * 2 - 1, mesh.vt.new_tensor([[0., 1.]]).expand(mesh.vt.size(0), -1)], dim=-1)
-
-        rast, rast_db = dr.rasterize(self.glctx, vt_clip[None], mesh.ft, (map_size, map_size), grad_db=False)
-        valid = (rast[..., 3] > 0).reshape(map_size, map_size)
-        rast = rast.expand(n, -1, -1, -1)
-        rast_db = rast_db.expand(n, -1, -1, -1)
-        v_img = v_clip[..., :2] / v_clip[..., 3:] * 0.5 + 0.5
-        # print(v_img.min(), v_img.max())
-        texc, texc_db = dr.interpolate(
-            v_img.contiguous(), rast.contiguous(), mesh.f, rast_db=rast_db.contiguous(), diff_attrs='all')
-        # (n, map_size, map_size, 4)
-        tex = dr.texture(
-            torch.cat([images, img_space_weight], dim=-1), texc, uv_da=texc_db, filter_mode=self.texture_filter)
-
-        weight = tex[..., 3:4] * visibility_grad
-
-        new_albedo_map = (tex[..., :3] * weight).sum(dim=0) / weight.sum(dim=0).clamp(min=1e-6)
-
-        new_albedo_map = edge_dilation(
-            new_albedo_map.permute(2, 0, 1)[None], valid[None, None].float(),
-        ).squeeze(0).permute(1, 2, 0)
-        mesh.albedo = torch.cat(
-            [new_albedo_map.clamp(min=0, max=1),
-             torch.ones_like(new_albedo_map[..., :1])], dim=-1)
-
-        mesh.textureless = False
-        return [mesh]