Spaces:
Runtime error
Runtime error
File size: 2,618 Bytes
da48dbe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import torch
import numpy as np
import neural_renderer as nr
from core import path_config
from models import SMPL
class PartRenderer():
"""Renderer used to render segmentation masks and part segmentations.
Internally it uses the Neural 3D Mesh Renderer
"""
def __init__(self, focal_length=5000., render_res=224):
# Parameters for rendering
self.focal_length = focal_length
self.render_res = render_res
# We use Neural 3D mesh renderer for rendering masks and part segmentations
self.neural_renderer = nr.Renderer(dist_coeffs=None, orig_size=self.render_res,
image_size=render_res,
light_intensity_ambient=1,
light_intensity_directional=0,
anti_aliasing=False)
self.faces = torch.from_numpy(SMPL(path_config.SMPL_MODEL_DIR).faces.astype(np.int32)).cuda()
textures = np.load(path_config.VERTEX_TEXTURE_FILE)
self.textures = torch.from_numpy(textures).cuda().float()
self.cube_parts = torch.cuda.FloatTensor(np.load(path_config.CUBE_PARTS_FILE))
def get_parts(self, parts, mask):
"""Process renderer part image to get body part indices."""
bn,c,h,w = parts.shape
mask = mask.view(-1,1)
parts_index = torch.floor(100*parts.permute(0,2,3,1).contiguous().view(-1,3)).long()
parts = self.cube_parts[parts_index[:,0], parts_index[:,1], parts_index[:,2], None]
parts *= mask
parts = parts.view(bn,h,w).long()
return parts
def __call__(self, vertices, camera):
"""Wrapper function for rendering process."""
# Estimate camera parameters given a fixed focal length
cam_t = torch.stack([camera[:,1], camera[:,2], 2*self.focal_length/(self.render_res * camera[:,0] +1e-9)],dim=-1)
batch_size = vertices.shape[0]
K = torch.eye(3, device=vertices.device)
K[0,0] = self.focal_length
K[1,1] = self.focal_length
K[2,2] = 1
K[0,2] = self.render_res / 2.
K[1,2] = self.render_res / 2.
K = K[None, :, :].expand(batch_size, -1, -1)
R = torch.eye(3, device=vertices.device)[None, :, :].expand(batch_size, -1, -1)
faces = self.faces[None, :, :].expand(batch_size, -1, -1)
parts, _, mask = self.neural_renderer(vertices, faces, textures=self.textures.expand(batch_size, -1, -1, -1, -1, -1), K=K, R=R, t=cam_t.unsqueeze(1))
parts = self.get_parts(parts, mask)
return mask, parts |