Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- (base) mona@mona:~/research/3danimals/SMALViewer$ cat smal/smal3d_renderer.py
- import sys, os
- sys.path.append(os.path.dirname(sys.path[0]))
- import torch
- import torch.nn as nn
- import neural_renderer as nr
- import torch.nn.functional as F
- from SMPL.smal_torch_batch import SMALModel
- from smal.joint_catalog import SMALJointInfo
- import pickle as pkl
- import numpy as np
- import cv2
- import matplotlib.pyplot as plt
- class SMAL3DRenderer(nn.Module):
- def __init__(self, image_size, z_distance = 2.5, elevation = 89.9, azimuth = 0.0):
- super(SMAL3DRenderer, self).__init__()
- self.smal_model = SMALModel()
- self.image_size = image_size
- self.smal_info = SMALJointInfo()
- self.renderer = nr.Renderer(camera_mode='look_at')
- self.renderer.eye = nr.get_points_from_angles(z_distance, elevation, azimuth)
- self.renderer.image_size = image_size
- self.renderer.light_intensity_ambient = 1.0
- with open("smal/dog_texture.pkl", 'rb') as f:
- self.textures = pkl.load(f).cuda()
- def forward(self, batch_params):
- batch_size = batch_params['betas'].shape[0]
- verts, joints_3d = self.smal_model(
- batch_params['betas'],
- torch.cat((batch_params['global_rotation'], batch_params['joint_rotations']), dim = 1),
- batch_params['trans'])
- faces = self.smal_model.faces.unsqueeze(0).expand(batch_size, -1, -1)
- textures = self.textures.unsqueeze(0).expand(batch_size, -1, -1, -1, -1, -1)
- rendered_joints = self.renderer.render_points(joints_3d[:, self.smal_info.include_classes])
- rendered_silhouettes = self.renderer.render_silhouettes(verts, faces)
- rendered_silhouettes = rendered_silhouettes.unsqueeze(1)
- rendered_images = self.renderer.render(verts, faces, textures)
- rendered_images = torch.clamp(rendered_images[0], 0.0, 1.0)
- return rendered_images, rendered_silhouettes, rendered_joints, verts, joints_3d(base)
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement