Rotating rendering utility for M3D-VTON

render_trunable.py

import io
import os
import torch
from skimage.io import imread
import numpy as np
import cv2
from tqdm import tqdm

# Util function for loading meshes
from pytorch3d.io import load_objs_as_meshes

# Data structures and functions for rendering
from pytorch3d.structures import Meshes
from pytorch3d.renderer import (
    look_at_view_transform,
    OpenGLOrthographicCameras, 
    FoVPerspectiveCameras,
    PointLights, 
    DirectionalLights, 
    Materials, 
    RasterizationSettings, 
    MeshRenderer, 
    MeshRasterizer,  
    HardPhongShader,
    SoftPhongShader,
    TexturesVertex
)

def set_renderer():
    # Setup
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)

    # Initialize an OpenGL perspective camera.
    # R, T = look_at_view_transform(2.0, 0, 180)
    # R, T = look_at_view_transform(1.0, 150, 150)
    R, T = look_at_view_transform(dist=2.7, elev=-150, azim=10) # for m3dvton
    
    # cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
    cameras = OpenGLOrthographicCameras(device=device, R=R, T=T)

    raster_settings = RasterizationSettings(
        image_size=800, 
        blur_radius=0.0, 
        faces_per_pixel=1, 
        bin_size = None, 
        max_faces_per_bin = None
    )

    

    lights = PointLights(device=device, location=((2.0, 2.0, 2.0),))

    renderer = MeshRenderer(
        rasterizer=MeshRasterizer(
            cameras=cameras, 
            raster_settings=raster_settings
        ),
        shader=HardPhongShader(
            device=device, 
            cameras=cameras,
            lights=lights
        )
    )
    return renderer

def get_verts_rgb_colors(obj_path):
  rgb_colors = []

  f = open(obj_path)
  lines = f.readlines()
  for line in lines:
    ls = line.split(' ')
    if len(ls) == 7:
      rgb_colors.append(ls[-3:])

  return np.array(rgb_colors, dtype='float32')[None, :, :]

def generate_video_from_obj(obj_path, video_path, renderer, render_texture=True):
    # input_image = cv2.imread(image_path)
    # input_image = input_image[:,:input_image.shape[1]//3]
    # input_image = cv2.resize(input_image, (512,512))

    # Setup
    device = torch.device("cuda:0")
    torch.cuda.set_device(device)

    # Load obj file
    verts_rgb_colors = get_verts_rgb_colors(obj_path)
    verts_rgb_colors = torch.from_numpy(verts_rgb_colors).to(device)
    if render_texture:
        textures = TexturesVertex(verts_features=verts_rgb_colors)
    else:
        wo_textures = TexturesVertex(verts_features=torch.ones_like(verts_rgb_colors)*0.75)

    # Load obj
    mesh = load_objs_as_meshes([obj_path], device=device)

    # Change specular color to green and change material shininess 
    materials = Materials(
        device=device,
        ambient_color = [[1.5, 1.5, 1.5]], 
        diffuse_color = [[1.0, 1.0, 1.0]], 
        specular_color=[[0.0, 0.0, 0.0]],
        shininess=64 
    )

    # Set mesh
    vers = mesh._verts_list
    faces = mesh._faces_list
    if render_texture:
        mesh_w_tex = Meshes(vers, faces, textures)
    else:
        mesh_wo_tex = Meshes(vers, faces, wo_textures)

    # create VideoWriter
    # fourcc = cv2. VideoWriter_fourcc(*'MP4V')
    # fourcc = cv2. VideoWriter_fourcc('M', 'J', 'P', 'G')
    # out = cv2.VideoWriter(video_path, fourcc, 20.0, (512,512))

    for i in tqdm(range(90)):
        R, T = look_at_view_transform(1.8, 0, i*4, device=device)
        if render_texture:
            images_w_tex = renderer(mesh_w_tex, R=R, T=T, materials=materials)
            images_w_tex = np.clip(images_w_tex[0, ..., :3].cpu().numpy(), 0.0, 1.0)[:, :, ::-1] * 255
            # image = np.concatenate([input_image, images_w_tex], axis=1)
            cv2.imwrite(video_path[:-9]+str(i)+'.png', images_w_tex)
        else:
            images_wo_tex = renderer(mesh_wo_tex, R=R, T=T)
            images_wo_tex = np.clip(images_wo_tex[0, ..., :3].cpu().numpy(), 0.0, 1.0)[:, :, ::-1] * 255
            cv2.imwrite(os.path.join('/public/zhaofuwei/Projects/M3D-VTON/video', str(i)+'.png'), images_wo_tex)
        
        # out.write(images_w_tex.astype('uint8'))
    # out.release()

if __name__ == '__main__':
    os.environ["CUDA_VISIBLE_DEVICES"]="0"
    obj_root = '/public/zhaofuwei/Projects/M3D-VTON/data/mesh'
    video_root = '/public/zhaofuwei/Projects/M3D-VTON/render'
    for objname in sorted(os.listdir(obj_root)):
        obj_path = os.path.join(obj_root, objname)
        video_path = os.path.join(video_root, objname[:-4]) + '/m3dvton/video.mp4'
        os.makedirs(os.path.join(video_root, objname[:-4]) + '/m3dvton', exist_ok=True)
        renderer = set_renderer()
        generate_video_from_obj(obj_path, video_path, renderer, render_texture=True)