🚧 update ffmeg and pyrender

This commit is contained in:
Qing Shuai 2022-08-21 16:02:48 +08:00
parent 050cb209d1
commit d534ba41fc
3 changed files with 315 additions and 0 deletions

View File

@ -0,0 +1,79 @@
'''
@ Date: 2021-11-27 16:50:33
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2022-04-13 18:19:03
@ FilePath: /EasyMocapPublic/easymocap/visualize/ffmpeg_wrapper.py
'''
import shutil
import os
from os.path import join
from glob import glob
from tqdm import tqdm
class VideoMaker:
def __init__(self, restart=True, fps_in=50, fps_out=50, remove_images=False,
reorder=False,
ext='.jpg', debug=False) -> None:
self.restart = ' -y' if restart else ''
self.fps_in = ' -r {}'.format(fps_in)
self.remove_images = remove_images
cmd = ' -pix_fmt yuv420p -vcodec libx264'
cmd += ' -r {}'.format(fps_out)
if ext == '.png':
cmd += ' -profile:v main'
self.cmd = cmd
self.ext = ext
self.shell = 'ffmpeg{restart}{fps_in} -i {path}/%06d{ext} -vf scale="2*ceil(iw/2):2*ceil(ih/2)"{cmd} {path}.mp4'
if not debug:
self.shell += ' -loglevel quiet'
self.reorder = reorder
def make_video(self, path):
imgnames = sorted(glob(join(path, '*'+self.ext)))
if len(imgnames) == 0:
print('[ffmpeg] No images in folder {}'.format(path))
return 0
firstname = imgnames[0]
index = os.path.basename(firstname).replace(self.ext, '')
if index.isdigit():
index = int(index)
if index != 0:
self.reorder = True
if self.reorder:
tmpdir = '/tmp/ffmpeg-tmp'
shutil.rmtree(tmpdir)
os.makedirs(tmpdir, exist_ok=True)
for nf, imgname in tqdm(enumerate(imgnames), desc='copy to /tmp'):
tmpname = join(tmpdir, '{:06d}{}'.format(nf, self.ext))
shutil.copyfile(imgname, tmpname)
path_ori = path
path = tmpdir
cmd = self.shell.format(
restart=self.restart,
fps_in=self.fps_in,
cmd=self.cmd,
path=path,
ext=self.ext
)
print(cmd)
os.system(cmd)
if self.reorder:
shutil.copy(path+'.mp4', path_ori+'.mp4')
if self.remove_images:
shutil.rmtree(path)
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('path', type=str)
parser.add_argument('--ext', type=str, default='.jpg')
parser.add_argument('--fps', type=int, default=50)
parser.add_argument('--remove', action='store_true')
parser.add_argument('--debug', action='store_true')
parser.add_argument('--reorder', action='store_true')
args = parser.parse_args()
video_maker = VideoMaker(
restart=True, fps_in=args.fps, fps_out=args.fps, remove_images=args.remove, ext=args.ext,
reorder=args.reorder,
debug=args.debug)
video_maker.make_video(args.path)

View File

@ -0,0 +1,43 @@
'''
@ Date: 2021-05-13 14:34:27
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2021-05-13 14:37:24
@ FilePath: /EasyMocap/easymocap/visualize/pyrender_flags.py
'''
from pyrender import RenderFlags
render_flags_default = {
'flip_wireframe': False,
'all_wireframe': False,
'all_solid': True,
'shadows': False, # TODO:bug exists in shadow mode
'vertex_normals': False,
'face_normals': False,
'cull_faces': True, # set to False
'point_size': 1.0,
'rgba':True
}
def get_flags(flags):
render_flags = render_flags_default.copy()
render_flags.update(flags)
flags = RenderFlags.NONE
if render_flags['flip_wireframe']:
flags |= RenderFlags.FLIP_WIREFRAME
elif render_flags['all_wireframe']:
flags |= RenderFlags.ALL_WIREFRAME
elif render_flags['all_solid']:
flags |= RenderFlags.ALL_SOLID
if render_flags['shadows']:
flags |= RenderFlags.SHADOWS_DIRECTIONAL | RenderFlags.SHADOWS_SPOT
if render_flags['vertex_normals']:
flags |= RenderFlags.VERTEX_NORMALS
if render_flags['face_normals']:
flags |= RenderFlags.FACE_NORMALS
if not render_flags['cull_faces']:
flags |= RenderFlags.SKIP_CULL_FACES
if render_flags['rgba']:
flags |= RenderFlags.RGBA
return flags

View File

@ -0,0 +1,193 @@
'''
@ Date: 2021-05-13 14:20:13
@ Author: Qing Shuai
@ LastEditors: Qing Shuai
@ LastEditTime: 2022-06-10 22:47:42
@ FilePath: /EasyMocapPublic/easymocap/visualize/pyrender_wrapper.py
'''
import pyrender
import numpy as np
import trimesh
import cv2
from .pyrender_flags import get_flags
from ..mytools.vis_base import get_rgb
def offscree_render(renderer, scene, img, flags):
rend_rgba, rend_depth = renderer.render(scene, flags=flags)
assert rend_depth.max() < 65, 'depth should less than 65.536: {}'.format(rend_depth.max())
rend_depth = (rend_depth * 1000).astype(np.uint16)
if rend_rgba.shape[2] == 3: # fail to generate transparent channel
valid_mask = (rend_depth > 0)[:, :, None]
rend_rgba = np.dstack((rend_rgba, (valid_mask*255).astype(np.uint8)))
rend_rgba = rend_rgba[..., [2, 1, 0, 3]]
if False:
rend_cat = cv2.addWeighted(
cv2.bitwise_and(img, 255 - rend_rgba[:, :, 3:4].repeat(3, 2)), 1,
cv2.bitwise_and(rend_rgba[:, :, :3], rend_rgba[:, :, 3:4].repeat(3, 2)), 1, 0)
else:
rend_cat = img.copy()
rend_cat[rend_rgba[:,:,-1]==255] = rend_rgba[:,:,:3][rend_rgba[:,:,-1]==255]
return rend_rgba, rend_depth, rend_cat
class Renderer:
def __init__(self, bg_color=[1.0, 1.0, 1.0, 0.0], ambient_light=[0.5, 0.5, 0.5], flags={}) -> None:
self.bg_color = bg_color
self.ambient_light = ambient_light
self.renderer = pyrender.OffscreenRenderer(1024, 1024)
self.flags = get_flags(flags)
@staticmethod
def add_light(scene, camera=None):
# Use 3 directional lights
# Create light source
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=3)
light_forward = np.eye(4)
# here the location of the light is set to be the origin
# and this location doesn't affect the render results
scene.add(light, pose=light_forward)
light_z = np.eye(4)
light_z[:3, :3] = cv2.Rodrigues(np.array([np.pi, 0, 0]))[0]
# if camera is not None:
# light_z[:3, :3] = camera['R'] @ light_z[:3, :3]
scene.add(light, pose=light_z)
def __call__(self, render_data, images, cameras, extra_mesh=[],
ret_image=False, ret_depth=False, ret_color=False, ret_mask=False, ret_all=True):
if isinstance(images, np.ndarray) and isinstance(cameras, dict):
images, cameras = [images], [cameras]
assert isinstance(cameras, list)
rot = trimesh.transformations.rotation_matrix(
np.radians(180), [1, 0, 0])
output_images, output_colors, output_depths = [], [], []
for nv, img in enumerate(images):
cam = cameras[nv]
K, R, T = cam['K'], cam['R'], cam['T']
self.renderer.viewport_height = img.shape[0]
self.renderer.viewport_width = img.shape[1]
scene = pyrender.Scene(bg_color=self.bg_color,
ambient_light=self.ambient_light)
for iextra, _mesh in enumerate(extra_mesh):
mesh = _mesh.copy()
trans_cam = np.eye(4)
trans_cam[:3, :3] = R
trans_cam[:3, 3:] = T
mesh.apply_transform(trans_cam)
mesh.apply_transform(rot)
# mesh.vertices = np.asarray(mesh.vertices) @ R.T + T.T
mesh_ = pyrender.Mesh.from_trimesh(mesh)
scene.add(mesh_, 'extra{}'.format(iextra))
for trackId, data in render_data.items():
vert = data['vertices'].copy()
faces = data['faces']
vert = vert @ R.T + T.T
if 'colors' not in data.keys():
# 如果使用了vid这个键那么可视化的颜色使用vid的颜色
if False:
col = get_rgb(data.get('vid', trackId))
else:
col = get_colors(data.get('vid', trackId))
mesh = trimesh.Trimesh(vert, faces, process=False)
mesh.apply_transform(rot)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
roughnessFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=col)
# material = pyrender.material.SpecularGlossinessMaterial(
# diffuseFactor=1.0, glossinessFactor=0.0
# )
mesh = pyrender.Mesh.from_trimesh(mesh, material=material, smooth=True)
else:
mesh = trimesh.Trimesh(vert, faces, vertex_colors=data['colors'], process=False)
mesh.apply_transform(rot)
mesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
scene.add(mesh, data['name'])
camera_pose = np.eye(4)
camera = pyrender.camera.IntrinsicsCamera(fx=K[0, 0], fy=K[1, 1], cx=K[0, 2], cy=K[1, 2])
scene.add(camera, pose=camera_pose)
self.add_light(scene, camera=cam)
# pyrender.Viewer(scene, use_raymond_lighting=True)
rend_rgba, rend_depth, rend_cat = offscree_render(self.renderer, scene, img, self.flags)
output_colors.append(rend_rgba)
output_depths.append(rend_depth)
output_images.append(rend_cat)
res = None
if ret_depth:
res = output_depths
elif ret_color:
res = output_colors
elif ret_mask:
res = [val[:, :, 3] for val in output_colors]
elif ret_image:
res = output_images
else:
res = output_colors, output_depths, output_images
return res
def render_image(self, render_data, images, cameras, extra_mesh,
**kwargs):
return self.__call__(render_data, images, cameras, extra_mesh,
ret_all=True, **kwargs)
def plot_meshes(img, meshes, K, R, T, mode='image'):
renderer = Renderer()
out = renderer.render_image(meshes, img, {'K': K, 'R': R, 'T': T}, [])
if mode == 'image':
return out[2][0]
elif mode == 'mask':
return out[0][0][..., -1]
elif mode == 'hstack':
return np.hstack([img, out[0][0][:, :, :3]])
elif mode == 'left':
out = out[0][0]
rend_rgba = np.roll(out, out.shape[1]//10, axis=1)
rend_cat = img.copy()
rend_cat[rend_rgba[:,:,-1]==255] = rend_rgba[:,:,:3][rend_rgba[:,:,-1]==255]
return rend_cat
# 这个顺序是BGR的。虽然render的使用的是RGB的但是由于和图像拼接了所以又变成BGR的了
colors = [
# (0.5, 0.2, 0.2, 1.), # Defalut BGR
(.5, .5, .7, 1.), # Pink BGR
(.44, .50, .98, 1.), # Red
(.7, .7, .6, 1.), # Neutral
(.5, .5, .7, 1.), # Blue
(.5, .55, .3, 1.), # capsule
(.3, .5, .55, 1.), # Yellow
# (.6, .6, .6, 1.), # gray
(.9, 1., 1., 1.),
(0.95, 0.74, 0.65, 1.),
(.9, .7, .7, 1.)
]
colors_table = {
# colorblind/print/copy safe:
'_blue': [0.65098039, 0.74117647, 0.85882353],
'_pink': [.9, .7, .7],
'_mint': [ 166/255., 229/255., 204/255.],
'_mint2': [ 202/255., 229/255., 223/255.],
'_green': [ 153/255., 216/255., 201/255.],
'_green2': [ 171/255., 221/255., 164/255.],
'_red': [ 251/255., 128/255., 114/255.],
'_orange': [ 253/255., 174/255., 97/255.],
'_yellow': [ 250/255., 230/255., 154/255.],
'r':[255/255,0,0],
'g':[0,255/255,0],
'b':[0,0,255/255],
'k':[0,0,0],
'y':[255/255,255/255,0],
'purple':[128/255,0,128/255]
}
def get_colors(pid):
if isinstance(pid, int):
return colors[pid % len(colors)]
elif isinstance(pid, str):
return colors_table[pid]
elif isinstance(pid, list) or isinstance(pid, tuple):
if len(pid) == 3:
pid = (pid[0], pid[1], pid[2], 1.)
assert len(pid) == 4
return pid