bad render result

[changfali]

hello， I use nvdiffrast to render a mesh, but got a bad result: could you help me with this?

This is my render code: `class MeshRenderer(nn.Module): def init(self, rasterize_fov, znear=0.1, zfar=10, rasterize_size=224): super(MeshRenderer, self).init()

    x = np.tan(np.deg2rad(rasterize_fov * 0.5)) * znear
    self.ndc_proj = torch.tensor(ndc_projection(x=x, n=znear, f=zfar)).matmul(
        torch.diag(torch.tensor([1.0, -1, -1, 1]))
    )
    self.rasterize_size = rasterize_size
    self.glctx = None

def forward(self, vertex, tri, feat=None):
    """
    Return:
        mask               -- torch.tensor, size (B, 1, H, W)
        depth              -- torch.tensor, size (B, 1, H, W)
        features(optional) -- torch.tensor, size (B, C, H, W) if feat is not None

    Parameters:
        vertex          -- torch.tensor, size (B, N, 3)
        tri             -- torch.tensor, size (B, M, 3) or (M, 3), triangles
        feat(optional)  -- torch.tensor, size (B, C), features
    """
    device = vertex.device
    rsize = int(self.rasterize_size)
    ndc_proj = self.ndc_proj.to(device)
    verts_proj = to_image(vertex)
    # trans to homogeneous coordinates of 3d vertices, the direction of y is the same as v
    if vertex.shape[-1] == 3:
        vertex = torch.cat([vertex, torch.ones([*vertex.shape[:2], 1]).to(device)], dim=-1)
        vertex[..., 1] = -vertex[..., 1]

    vertex_ndc = vertex @ ndc_proj.t()
    if self.glctx is None:
        self.glctx = dr.RasterizeCudaContext(device=device)

    ranges = None
    if isinstance(tri, List) or len(tri.shape) == 3:
        vum = vertex_ndc.shape[1]
        fnum = torch.tensor([f.shape[0] for f in tri]).unsqueeze(1).to(device)

        print("fnum shape:{}".format(fnum.shape))

        fstartidx = torch.cumsum(fnum, dim=0) - fnum
        ranges = torch.cat([fstartidx, fnum], axis=1).type(torch.int32).cpu()
        for i in range(tri.shape[0]):
            tri[i] = tri[i] + i * vum
        vertex_ndc = torch.cat(vertex_ndc, dim=0)
        tri = torch.cat(tri, dim=0)

    # for range_mode vetex: [B*N, 4], tri: [B*M, 3], for instance_mode vetex: [B, N, 4], tri: [M, 3]
    tri = tri.type(torch.int32).contiguous()
    rast_out, _ = dr.rasterize(self.glctx, vertex_ndc.contiguous(), tri, resolution=[rsize, rsize], ranges=ranges)

    depth, _ = dr.interpolate(vertex.reshape([-1, 4])[..., 2].unsqueeze(1).contiguous(), rast_out, tri)
    depth = depth.permute(0, 3, 1, 2)
    mask = (rast_out[..., 3] > 0).float().unsqueeze(1)
    depth = mask * depth

    image = None

    verts_x = verts_proj[0, :, 0]
    verts_y = 224 - verts_proj[0, :, 1]
    verts_int = torch.ceil(verts_proj[0]).long()  # (n, 2)
    verts_xr_int = verts_int[:, 0].clamp(1, 224 - 1)
    verts_yt_int = 224 - verts_int[:, 1].clamp(2, 224)
    verts_right_float = verts_xr_int - verts_x
    verts_left_float = 1 - verts_right_float
    verts_top_float = verts_y - verts_yt_int
    verts_bottom_float = 1 - verts_top_float

    rast_lt = rast_out[0, verts_yt_int, verts_xr_int - 1, 3]
    rast_lb = rast_out[0, verts_yt_int + 1, verts_xr_int - 1, 3]
    rast_rt = rast_out[0, verts_yt_int, verts_xr_int, 3]
    rast_rb = rast_out[0, verts_yt_int + 1, verts_xr_int, 3]

    occ_feat = (
        (rast_lt > 0) * 1.0 * (verts_left_float + verts_top_float)
        + (rast_lb > 0) * 1.0 * (verts_left_float + verts_bottom_float)
        + (rast_rt > 0) * 1.0 * (verts_right_float + verts_top_float)
        + (rast_rb > 0) * 1.0 * (verts_right_float + verts_bottom_float)
    )
    occ_feat = occ_feat[None, :, None] / 4.0

    # occ_feat = torch.ones([1, vertex.shape[1], 1], dtype=torch.float32).to(vertex.device)
    occ, _ = dr.interpolate(occ_feat, rast_out, tri)
    occ = occ.permute(0, 3, 1, 2)
    # occ = mask * occ

    if feat is not None:
        image, _ = dr.interpolate(feat, rast_out, tri)
        image = image.permute(0, 3, 1, 2)
        image = mask * image

    return mask, depth, image, occ

def render_uv_texture(self, vertex, tri, uv, uv_texture):
    """
    Return:
        mask               -- torch.tensor, size (B, 1, H, W)
        depth              -- torch.tensor, size (B, 1, H, W)
        features(optional) -- torch.tensor, size (B, C, H, W) if feat is not None

    Parameters:
        vertex          -- torch.tensor, size (B, N, 3)
        tri             -- torch.tensor, size (M, 3), triangles
        uv                -- torch.tensor, size (B,N, 2),  uv mapping
        uv_texture   -- torch.tensor, size (B,C,H,W,C) texture map
    """
    device = vertex.device
    rsize = int(self.rasterize_size)
    ndc_proj = self.ndc_proj.to(device)
    # trans to homogeneous coordinates of 3d vertices, the direction of y is the same as v
    if vertex.shape[-1] == 3:
        vertex = torch.cat([vertex, torch.ones([*vertex.shape[:2], 1]).to(device)], dim=-1)
        vertex[..., 1] = -vertex[..., 1]

    vertex_ndc = vertex @ ndc_proj.t()
    if self.glctx is None:
        self.glctx = dr.RasterizeCudaContext(device=device)
        # print("create glctx on device cuda:%d" % device.index)

    # print('vertex_ndc shape:{}'.format(vertex_ndc.shape))  # Size([1, 35709, 4])
    # print('tri shape:{}'.format(tri.shape))  # Size([70789, 3])

    ranges = None
    if isinstance(tri, List) or len(tri.shape) == 3:
        vum = vertex_ndc.shape[1]
        fnum = torch.tensor([f.shape[0] for f in tri]).unsqueeze(1).to(device)

        print("fnum shape:{}".format(fnum.shape))

        fstartidx = torch.cumsum(fnum, dim=0) - fnum
        ranges = torch.cat([fstartidx, fnum], axis=1).type(torch.int32).cpu()
        for i in range(tri.shape[0]):
            tri[i] = tri[i] + i * vum
        vertex_ndc = torch.cat(vertex_ndc, dim=0)
        tri = torch.cat(tri, dim=0)

    # for range_mode vetex: [B*N, 4], tri: [B*M, 3], for instance_mode vetex: [B, N, 4], tri: [M, 3]
    tri = tri.type(torch.int32).contiguous()
    rast_out, _ = dr.rasterize(self.glctx, vertex_ndc.contiguous(), tri, resolution=[rsize, rsize], ranges=ranges)

    depth, _ = dr.interpolate(vertex.reshape([-1, 4])[..., 2].unsqueeze(1).contiguous(), rast_out, tri)
    depth = depth.permute(0, 3, 1, 2)
    mask = (rast_out[..., 3] > 0).float().unsqueeze(1)
    depth = mask * depth
    uv_ = uv.clone()
    uv_[..., -1] = 1.0 - uv_[..., -1]

    rast_out, rast_db = dr.rasterize(
        self.glctx, vertex_ndc.contiguous(), tri, resolution=[rsize, rsize], ranges=ranges
    )

    interp_out, uv_da = dr.interpolate(uv_, rast_out, tri, rast_db, diff_attrs="all")

    # tex = torch.zeros((1, 128*5//4, 128, 3), dtype=torch.float32)

    uv_texture = uv_texture.permute(0, 2, 3, 1).contiguous()
    img = dr.texture(uv_texture, interp_out, filter_mode="linear")  # , uv_da)
    img = img * torch.clamp(rast_out[..., -1:], 0, 1)  # Mask out background.

    image = img.permute(0, 3, 1, 2)

    return mask, depth, image

def pred_texture(self, vertex, tri, uv, target_img, tex_size=1024):
    """
    Return:
        mask               -- torch.tensor, size (B, 1, H, W)
        depth              -- torch.tensor, size (B, 1, H, W)
        features(optional) -- torch.tensor, size (B, C, H, W) if feat is not None

    Parameters:
        vertex          -- torch.tensor, size (B, N, 3)
        tri             -- torch.tensor, size (B, M, 3) or (M, 3), triangles
        uv                -- torch.tensor, size (B,N, 2),  uv mapping
        base_tex   -- torch.tensor, size (B,H,W,C)
    """

    device = vertex.device
    rsize = int(self.rasterize_size)
    ndc_proj = self.ndc_proj.to(device)
    # trans to homogeneous coordinates of 3d vertices, the direction of y is the same as v
    if vertex.shape[-1] == 3:
        vertex = torch.cat([vertex, torch.ones([*vertex.shape[:2], 1]).to(device)], dim=-1)
        vertex[..., 1] = -vertex[..., 1]

    vertex_ndc = vertex @ ndc_proj.t()
    if self.glctx is None:
        self.glctx = dr.RasterizeCudaContext(device=device)
        # print("create glctx on device cuda:%d" % device.index)

    # print('vertex_ndc shape:{}'.format(vertex_ndc.shape))  # Size([1, 35709, 4])
    # print('tri shape:{}'.format(tri.shape))  # Size([70789, 3])

    ranges = None
    if isinstance(tri, List) or len(tri.shape) == 3:
        vum = vertex_ndc.shape[1]
        fnum = torch.tensor([f.shape[0] for f in tri]).unsqueeze(1).to(device)

        print("fnum shape:{}".format(fnum.shape))

        fstartidx = torch.cumsum(fnum, dim=0) - fnum
        ranges = torch.cat([fstartidx, fnum], axis=1).type(torch.int32).cpu()
        for i in range(tri.shape[0]):
            tri[i] = tri[i] + i * vum
        vertex_ndc = torch.cat(vertex_ndc, dim=0)
        tri = torch.cat(tri, dim=0)

    # for range_mode vetex: [B*N, 4], tri: [B*M, 3], for instance_mode vetex: [B, N, 4], tri: [M, 3]
    tri = tri.type(torch.int32).contiguous()
    rast_out, _ = dr.rasterize(self.glctx, vertex_ndc.contiguous(), tri, resolution=[rsize, rsize], ranges=ranges)

    depth, _ = dr.interpolate(vertex.reshape([-1, 4])[..., 2].unsqueeze(1).contiguous(), rast_out, tri)
    depth = depth.permute(0, 3, 1, 2)
    mask = (rast_out[..., 3] > 0).float().unsqueeze(1)
    depth = mask * depth
    uv[..., -1] = 1.0 - uv[..., -1]

    rast_out, rast_db = dr.rasterize(
        self.glctx, vertex_ndc.contiguous(), tri, resolution=[rsize, rsize], ranges=ranges
    )

    interp_out, uv_da = dr.interpolate(uv, rast_out, tri, rast_db, diff_attrs="all")

    mask_3c = mask.permute(0, 2, 3, 1)
    mask_3c = torch.cat((mask_3c, mask_3c, mask_3c), dim=-1)
    maskout_img = mask_3c * target_img
    mean_color = torch.sum(maskout_img, dim=(1, 2))
    valid_pixel_count = torch.sum(mask)
    mean_color = mean_color / valid_pixel_count

    tex = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
    tex[:, :, :, 0] = mean_color[0, 0]
    tex[:, :, :, 1] = mean_color[0, 1]
    tex[:, :, :, 2] = mean_color[0, 2]
    tex = tex.cuda()

    tex_resolution_list = []
    cur_tex_size = 64
    while cur_tex_size <= tex_size:
        tex_resolution_list.append(cur_tex_size)
        if cur_tex_size == tex_size:
            break
        else:
            cur_tex_size = min(cur_tex_size * 2, tex_size)

    tex_mask = torch.zeros((1, tex_size, tex_size, 3), dtype=torch.float32)
    # tex_mask = torch.zeros((1, 2048, 2048, 3), dtype=torch.float32)
    tex_mask[:, :, :, 1] = 1.0
    tex_mask = tex_mask.cuda()
    tex_mask.requires_grad = True
    tex_mask = tex_mask.contiguous()

    criterionTV = TVLoss()

    for tex_resolution in tex_resolution_list:
        tex = tex.detach()
        tex = tex.permute(0, 3, 1, 2)
        tex = F.interpolate(tex, (tex_resolution, tex_resolution))
        # tex = F.interpolate(tex, (tex_resolution, tex_resolution))
        tex = tex.permute(0, 2, 3, 1).contiguous()

        tex.requires_grad = True

        optim = torch.optim.Adam([tex], lr=1e-2)

        texture_opt_iters = 100

        if tex_resolution == tex_size:
            optim_mask = torch.optim.Adam([tex_mask], lr=1e-2)

        for i in range(int(texture_opt_iters)):
            if tex_resolution == tex_size:
                optim_mask.zero_grad()
                rendered = dr.texture(tex_mask, interp_out, filter_mode="linear")  # , uv_da)
                rendered = rendered * torch.clamp(rast_out[..., -1:], 0, 1)  # Mask out background.
                tex_loss = torch.mean((target_img - rendered) ** 2)

                tex_loss.backward()
                optim_mask.step()

            optim.zero_grad()

            img = dr.texture(tex, interp_out, filter_mode="linear")  # , uv_da)
            img = img * torch.clamp(rast_out[..., -1:], 0, 1)  # Mask out background.
            recon_loss = torch.mean((target_img - img) ** 2)

            if tex_resolution < tex_size:
                tv_loss = criterionTV(tex.permute(0, 3, 1, 2))

                total_loss = recon_loss + tv_loss * 0.01
            else:
                total_loss = recon_loss

            total_loss.backward()
            optim.step()

    # tex_map = tex[0].detach().cpu().numpy()[...,::-1] * 255.0

    image = img.permute(0, 3, 1, 2)

    tex_mask = tex_mask[0].detach().cpu().numpy() * 255.0
    tex_mask = (
        np.where(tex_mask[..., 1] > 250, 1.0, 0.0)
        * np.where(tex_mask[..., 0] < 10, 1.0, 0)
        * np.where(tex_mask[..., 2] < 10, 1.0, 0)
    )
    tex_mask = 1.0 - tex_mask

    return mask, depth, image, tex.detach(), tex_mask

`

`import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import cv2 from tqdm import tqdm from torch_scatter import scatter_add from pytorch3d.io import load_obj from pytorch3d.renderer import look_at_view_transform from pytorch3d.structures import Meshes from OFScript.utils import read_poly_from_obj, read_from_obj from HRN.renderer import MeshRenderer, set_rasterizer, SRenderY

class SH: def init(self): self.a = [np.pi, 2 np.pi / np.sqrt(3.0), 2 np.pi / np.sqrt(8.0)] self.c = [1 / np.sqrt(4 np.pi), np.sqrt(3.0) / np.sqrt(4 np.pi), 3 np.sqrt(5.0) / np.sqrt(12 np.pi)]

class Metahuman_render(nn.Module): def init(self, focal=2000, center=256): super(Metahuman_render, self).init() self.device = "cuda" self.focal = focal self.center = center

    # aligned to bfm camera
    fov = 2 * np.arctan(self.center / self.focal) * 180 / np.pi  # 112, 1015
    self.renderer = MeshRenderer(
        rasterize_fov=fov,  # 12.59363743796881
        znear=5.0,  # 5.0
        zfar=15.0,  # 15.0
        rasterize_size=self.center * 2,  # 224
    )

    set_rasterizer("pytorch3d")
    verts, faces, aux = load_obj("data/OFdata/test_render/test8.obj")
    # mesh_template = {"vertices": verts.numpy(), "faces": faces.verts_idx.numpy() + 1}
    # meta_params = {"template_bfm_mesh": mesh_template, "bfm_uvs2": aux.verts_uvs.numpy()}
    # self.render_p3d = SRenderY(
    #     self.center * 2, uv_size=self.center * 2, rasterizer_type="pytorch3d", meta_params=meta_params
    # ).to(self.device)

    self.faces = faces.verts_idx.to(self.device)
    self.uvs = aux.verts_uvs.to(self.device)
    print(self.uvs.shape, verts.shape)

    # trans form meta_norm to bfm
    self.R_m2b = (
        torch.tensor(
            [
                [9.9997e-01, 7.6289e-04, -7.1648e-03],
                [-1.3363e-03, 9.9676e-01, -8.0375e-02],
                [7.0803e-03, 8.0383e-02, 9.9674e-01],
            ]
        )
        .view(1, 3, 3)
        .to(self.device)
    )
    self.T_m2b = torch.tensor([[0.0070, -0.3729, -0.5469]]).view(1, 1, 3).to(self.device)
    self.s_m2b = 2.7609

    self.persc_proj = torch.tensor(self.perspective_projection(self.focal, self.center))[None, :, :].to(self.device)

    # self.gamma = torch.tensor(
    #     [
    #         [
    #             0.0643,
    #             -0.0028,
    #             0.0560,
    #             0.0012,
    #             0.0231,
    #             -0.0286,
    #             -0.0125,
    #             -0.0125,
    #             0.0753,
    #             0.0869,
    #             0.0175,
    #             0.0763,
    #             0.0044,
    #             0.0171,
    #             -0.0307,
    #             -0.0170,
    #             -0.0167,
    #             0.1061,
    #             0.0498,
    #             0.0154,
    #             0.0349,
    #             0.0009,
    #             0.0117,
    #             -0.0326,
    #             -0.0086,
    #             -0.0189,
    #             0.0968,
    #         ]
    #     ]
    # ).to(self.device)

    self.gamma = torch.randn(27) * 0.1
    print(self.gamma)
    self.SH = SH()
    self.init_lit = torch.tensor([0.8, 0, 0, 0, 0, 0, 0, 0, 0]).view(1, 1, -1).to(self.device)

def perspective_projection(self, focal, center):
    # return p.T (N, 3) @ (3, 3)
    return np.array([focal, 0, center, 0, focal, center, 0, 0, 1]).reshape([3, 3]).astype(np.float32).transpose()

def transform(self, X, R=None, T=None, s=None):
    """
    X: [bs, N, 3]
    R: [bs, 3, 3] | None
    T: [bs, 1, 3] | None
    s: [bs, 1, 1] | None
    """
    bs = X.shape[0]
    if s is None:
        s = 1.0
    if R is None:
        if isinstance(X, np.ndarray):
            R = np.eye(3).reshape(bs, 3, 3)
        else:
            R = torch.eye(3).view(bs, 3, 3).to(X.device)
    if T is None:
        if isinstance(X, np.ndarray):
            T = np.zeros((bs, 1, 3))
        else:
            T = torch.zeros((bs, 1, 3)).to(X.device)
    return s * (X @ R) + T

def to_image(self, face_shape):
    """
    Return:
        face_proj        -- torch.tensor, size (B, N, 2), y direction is opposite to v direction

    Parameters:
        face_shape       -- torch.tensor, size (B, N, 3)
    """
    # to image_plane
    face_proj = face_shape @ self.persc_proj
    face_proj = face_proj[..., :2] / face_proj[..., 2:]

    return face_proj

def getVertexNormals(self, vertex, faces):
    f = faces.permute(1, 0)
    v = vertex

    vec1 = v[:, f[1]] - v[:, f[0]]
    vec2 = v[:, f[2]] - v[:, f[0]]
    face_norm = F.normalize(vec1.cross(vec2), p=2, dim=-1)  # [B, F, 3]
    idx = torch.cat([f[0], f[1], f[2]], dim=0)
    face_norm = face_norm.repeat(1, 3, 1)
    norm = scatter_add(face_norm, idx, dim=1, dim_size=v.size(1))

    return F.normalize(norm, p=2, dim=-1)

def compute_color_map(self, face_texture_uv, face_norm_uv, gamma):
    face_texture_uv = face_texture_uv.permute(0, 2, 3, 1).contiguous()  # (B, 256, 256, 3)
    face_norm_uv = face_norm_uv.permute(0, 2, 3, 1).contiguous()  # (B, 256, 256, 3)
    size = face_texture_uv.shape[1]

    batch_size = gamma.shape[0]
    a, c = self.SH.a, self.SH.c
    gamma = gamma.reshape([batch_size, 3, 9])
    gamma = gamma + self.init_lit
    gamma = gamma.permute(0, 2, 1)  # (B, 9, 3)
    Y = torch.cat(
        [
            a[0] * c[0] * torch.ones_like(face_norm_uv[..., :1]).to(self.device),
            -a[1] * c[1] * face_norm_uv[..., 1:2],
            a[1] * c[1] * face_norm_uv[..., 2:],
            -a[1] * c[1] * face_norm_uv[..., :1],
            a[2] * c[2] * face_norm_uv[..., :1] * face_norm_uv[..., 1:2],
            -a[2] * c[2] * face_norm_uv[..., 1:2] * face_norm_uv[..., 2:],
            0.5 * a[2] * c[2] / np.sqrt(3.0) * (3 * face_norm_uv[..., 2:] ** 2 - 1),
            -a[2] * c[2] * face_norm_uv[..., :1] * face_norm_uv[..., 2:],
            0.5 * a[2] * c[2] * (face_norm_uv[..., :1] ** 2 - face_norm_uv[..., 1:2] ** 2),
        ],
        dim=-1,
    )  # (B, 256, 256, 9)
    Y = Y.reshape(batch_size, -1, 9)
    r = Y @ gamma[..., :1]
    g = Y @ gamma[..., 1:2]
    b = Y @ gamma[..., 2:]
    face_color = (
        torch.cat([r, g, b], dim=-1).reshape(batch_size, size, size, 3) * face_texture_uv
    )  # (B, 256, 256, 3)
    face_color = face_color.permute(0, 3, 1, 2).contiguous()  # (B, 3, 256, 256)
    return face_color.clip(0, 1)

def forward(self, vertex, tex_img):
    vertex = self.transform(vertex, self.R_m2b, self.T_m2b, self.s_m2b)
    for i in tqdm(range(10)):
        R, T = look_at_view_transform(10, torch.rand(1) * 40 - 20, torch.rand(1) * 60 - 30, device=self.device)
        vertex_ = self.transform(vertex, R, T)

        # vn = self.getVertexNormals(vertex, self.faces)

        # face_albedo_map = torch.nn.functional.interpolate(
        #     tex_img, [self.render_p3d.uv_size, self.render_p3d.uv_size], mode="bilinear"
        # )
        # vn_uv = self.render_p3d.world2uv(vn)

        # gamma = torch.randn((vertex.shape[0], 1, 9)).to(self.device) * 0.1
        # gamma = gamma.repeat(1, 3, 1).view(-1, 27)

        # face_color_map = self.compute_color_map(face_albedo_map, vn_uv, gamma)  # 添加S

        pred_mask, _, pred_face = self.renderer.render_uv_texture(
            vertex_, self.faces, self.uvs, tex_img
        )  # uvs会被改变
        # pred_mask, _, pred_face = self.renderer.render_uv_texture(vertex_, faces, uvs, tex_img)  # uvs会被改变

        # cv2.imwrite(
        #     "test/pred_face{}.png".format(i),
        #     (face_color_map.cpu().numpy()[0] * 255).astype(np.uint8).transpose(1, 2, 0)[:, :, ::-1],
        # )

        cv2.imwrite(
            "test/pred_face{}_1.png".format(i),
            (pred_face.cpu().numpy()[0] * 255).astype(np.uint8).transpose(1, 2, 0)[:, :, ::-1],
        )
    # np.save("bfm.npy", vertex_.cpu().detach().numpy()[0])

verts, faces, aux = load_obj("data/OFdata/test_render/test1.obj", device="cuda")

mesh = read_poly_from_obj("data/OFdata/test_render/test8.obj") mesh.normalization() verts = torch.tensor(mesh.vertex)[None, :, :].cuda()

img = cv2.imread("data/OFdata/test_render/chenruifan_00.png") img = cv2.resize(img, (512,512))

print(img.shape, img.max())

tex_img = ( torch.tensor(img / 255.0)[None, :, :, [2, 1, 0]] .permute(0, 3, 1, 2) .cuda() .float() ) renderer = Metahuman_render() renderer(verts, tex_img) `

[yltqy]

你好，请问你解决这个问题了没，求教