diff --git a/mogen/smpl/render_mesh.py b/mogen/smpl/render_mesh.py index fc52ff4..e6855e7 100644 --- a/mogen/smpl/render_mesh.py +++ b/mogen/smpl/render_mesh.py @@ -11,12 +11,15 @@ if os.name == 'posix' and "DISPLAY" not in os.environ: import torch from mogen.smpl.simplify_loc2rot import joints2smpl import pyrender +from pyrender.shader_program import ShaderProgramCache from shapely import geometry import trimesh from pyrender.constants import RenderFlags from comfy.model_management import get_torch_device from tqdm import tqdm +shader_dir = os.path.join(os.path.dirname(__file__), 'shaders') + class WeakPerspectiveCamera(pyrender.Camera): def __init__(self, scale, @@ -156,7 +159,7 @@ def render(motions): out = np.stack(vid, axis=0) return out -def render_from_smpl(thetas, yfov, move_x, move_y, move_z, draw_platform=True, depth_only=False, smpl_model_path=None): +def render_from_smpl(thetas, yfov, move_x, move_y, move_z, draw_platform=True, depth_only=False, normals=False, smpl_model_path=None): rot2xyz = Rotation2xyz(device=get_torch_device(), smpl_model_path=smpl_model_path) faces = rot2xyz.smpl_model.faces @@ -262,12 +265,15 @@ def render_from_smpl(thetas, yfov, move_x, move_y, move_z, draw_platform=True, d depth = r.render(scene, flags=RenderFlags.DEPTH_ONLY) color = np.zeros([960, 960, 3]) else: - color, depth = r.render(scene, flags=RenderFlags.RGBA) + if normals: + r._renderer._program_cache = ShaderProgramCache(shader_dir=shader_dir) + color, depth = r.render(scene, flags=RenderFlags.RGBA) + # Image.fromarray(color).save(outdir+name+'_'+str(i)+'.png') vid.append(color) vid_depth.append(depth) - r.delete() + r = None return np.stack(vid, axis=0), np.stack(vid_depth, axis=0) \ No newline at end of file diff --git a/mogen/smpl/shaders/mesh.frag b/mogen/smpl/shaders/mesh.frag new file mode 100644 index 0000000..2bb50a4 --- /dev/null +++ b/mogen/smpl/shaders/mesh.frag @@ -0,0 +1,13 @@ +#version 330 core + +in vec3 frag_position; +in vec3 frag_normal; + +out vec4 frag_color; + +void main() +{ + vec3 normal = normalize(frag_normal); + + frag_color = vec4(normal * 0.5 + 0.5, 1.0); +} \ No newline at end of file diff --git a/mogen/smpl/shaders/mesh.vert b/mogen/smpl/shaders/mesh.vert new file mode 100644 index 0000000..cacafc8 --- /dev/null +++ b/mogen/smpl/shaders/mesh.vert @@ -0,0 +1,24 @@ +#version 330 core + +// Vertex Attributes +layout(location = 0) in vec3 position; +layout(location = NORMAL_LOC) in vec3 normal; +layout(location = INST_M_LOC) in mat4 inst_m; + +// Uniforms +uniform mat4 M; +uniform mat4 V; +uniform mat4 P; + +// Outputs +out vec3 frag_position; +out vec3 frag_normal; + +void main() +{ + gl_Position = P * V * M * inst_m * vec4(position, 1); + frag_position = vec3(M * inst_m * vec4(position, 1.0)); + + mat4 N = transpose(inverse(M * inst_m)); + frag_normal = normalize(vec3(N * vec4(normal, 0.0))); +} \ No newline at end of file diff --git a/smpl_nodes.py b/smpl_nodes.py index 10af14e..0355aad 100644 --- a/smpl_nodes.py +++ b/smpl_nodes.py @@ -90,18 +90,21 @@ class RenderSMPLMesh: "move_y": ("FLOAT", {"default": 0,"min": -500, "max": 500, "step": 0.1}), "move_z": ("FLOAT", {"default": 0,"min": -500, "max": 500, "step": 0.1}), "background_hex_color": ("STRING", {"default": "#FFFFFF", "mutiline": False}) + }, + "optional": { + "normals": ("BOOLEAN", {"default": False}), } } - RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_TYPES = ("IMAGE", "IMAGE", "MASK") RETURN_NAMES = ("IMAGE", "DEPTH_MAP") CATEGORY = "MotionDiff/smpl" FUNCTION = "render" - def render(self, smpl, yfov, move_x, move_y, move_z, draw_platform, depth_only, background_hex_color): + def render(self, smpl, yfov, move_x, move_y, move_z, draw_platform, depth_only, background_hex_color, normals=False): smpl_model_path, thetas, _ = smpl color_frames, depth_frames = render_from_smpl( thetas.to(get_torch_device()), - yfov, move_x, move_y, move_z, draw_platform,depth_only, + yfov, move_x, move_y, move_z, draw_platform,depth_only, normals, smpl_model_path=smpl_model_path ) bg_color = ImageColor.getcolor(background_hex_color, "RGB") @@ -112,7 +115,10 @@ class RenderSMPLMesh: (color_frames[..., 2] == 1.) ] color_frames[..., :3][white_mask] = torch.Tensor(bg_color) - + white_mask_tensor = torch.stack(white_mask, dim=0) + white_mask_tensor = white_mask_tensor.float() / white_mask_tensor.max() + white_mask_tensor = 1.0 - white_mask_tensor.permute(1, 2, 3, 0).squeeze(dim=-1) + print(white_mask_tensor.shape) #Normalize to [0, 1] normalized_depth = (depth_frames - depth_frames.min()) / (depth_frames.max() - depth_frames.min()) #Pyrender's depths are the distance in meters to the camera, which is the inverse of depths in normal context @@ -121,7 +127,7 @@ class RenderSMPLMesh: #https://github.com/Fannovel16/comfyui_controlnet_aux/blob/main/src/controlnet_aux/util.py#L24 depth_frames = [torch.from_numpy(np.concatenate([x, x, x], axis=2)) for x in normalized_depth[..., None]] depth_frames = torch.stack(depth_frames, dim=0) - return (color_frames, depth_frames,) + return (color_frames, depth_frames, white_mask_tensor,) class SMPLLoader: @classmethod