| import os |
| import io |
| import torch |
| import numpy as np |
| from termcolor import colored |
| try: |
| |
| os.environ["PYOPENGL_PLATFORM"] = "egl" |
| import pyrender |
| except: |
| print(colored('pyrender is not correctly imported.', 'red')) |
| import matplotlib |
| from matplotlib import colormaps |
| from matplotlib.colors import LightSource |
| import matplotlib.pyplot as plt |
| import math |
| import cv2 |
| import trimesh |
| from sklearn.decomposition import PCA |
| from scipy.spatial.transform import Rotation as R |
| import torchvision |
| from .transforms import adjust_colors |
|
|
| BASE_COLORS = np.loadtxt(os.path.abspath(os.path.join(__file__, "../colors.txt")), skiprows=0)/255. |
| BASE_COLORS = adjust_colors(BASE_COLORS, |
| saturation_threshold = 0.3, |
| brightness_threshold = 0.8) |
|
|
|
|
| def get_colors_rgb(size): |
| |
| return BASE_COLORS[np.random.choice(BASE_COLORS.shape[0], size=size, replace=False)] |
|
|
| def tensor_to_BGR(img_tensor): |
| img = img_tensor.numpy()*255 |
| img = img.astype(np.uint8).transpose((1,2,0))[:,:,::-1].copy() |
| return img |
|
|
| def pad_img(img, pad_size = None, pad_color_offset = 127): |
| if not isinstance(img, np.ndarray): |
| img = tensor_to_BGR(img.detach().cpu()) |
| if pad_size is None: |
| pad_size = max(img.shape[0],img.shape[1]) |
|
|
| pad = np.zeros((pad_size,pad_size,img.shape[-1]), dtype=img.dtype) + pad_color_offset |
| pad[:img.shape[0], :img.shape[1]] = img.copy() |
| return pad |
|
|
|
|
| def vis_scale_img(img, scale_map, conf_thresh = 0.3, patch_size=14): |
| cmap = plt.get_cmap('coolwarm') |
|
|
| vis_map = np.zeros((scale_map.shape[0]*patch_size, scale_map.shape[1]*patch_size, 3), dtype=np.uint8) |
| loc_i, loc_j = torch.where(scale_map[:,:,0] > conf_thresh) |
| for (i, j) in zip(loc_i, loc_j): |
| scale = round(math.sqrt(scale_map[i,j,1].item()),2) |
| vis_map[i*patch_size: (i+1)*patch_size, j*patch_size: (j+1)*patch_size] = (np.array(cmap(1-scale)[:3][::-1])*255).astype(np.uint8) |
| vis_map = pad_img(vis_map, pad_color_offset=0) |
| img = pad_img(img) |
| |
| assert img.shape == vis_map.shape |
|
|
| white_img = 0.6*img + 0.4*np.array((255,255,255)) |
|
|
| valid_mask = (vis_map > 0) |
| visible_weight = 0.8 |
| img = vis_map * valid_mask * visible_weight +\ |
| img * valid_mask * (1-visible_weight)+\ |
| white_img * (1-valid_mask) |
|
|
| |
| loc_i, loc_j = torch.where(scale_map[:,:,0]+1) |
| for (i, j) in zip(loc_i.tolist(), loc_j.tolist()): |
| cv2.rectangle(img, (j*patch_size, i*patch_size), ((j+1)*patch_size, (i+1)*patch_size), |
| color=(255,255,255), thickness = 2 ) |
|
|
| return img |
|
|
| def vis_meshes_img(img, verts, smpl_faces, cam_intrinsics, colors = None, padding = True): |
| if not isinstance(img, np.ndarray): |
| img = tensor_to_BGR(img.detach().cpu()) |
|
|
| if padding: |
| pad_size = max(img.shape[0],img.shape[1]) |
| img = pad_img(img, pad_size) |
|
|
| if colors is not None: |
| assert len(colors) == len(verts) |
|
|
| if len(cam_intrinsics.flatten()) == 9: |
| cam_intrinsics = cam_intrinsics.reshape(3,3) |
| rgb, depth = render_mesh(img.shape[0],img.shape[1],verts,smpl_faces,cam_intrinsics,colors) |
| valid_mask = (depth > 0)[:,:,None] |
| visible_weight = 1. |
| rendered_img = rgb[:,:,::-1] * valid_mask * visible_weight +\ |
| img * valid_mask * (1-visible_weight)+\ |
| img * (1-valid_mask) |
| else: |
| rendered_img = img |
| for i, cam_int in enumerate(cam_intrinsics): |
| rgb, depth = render_mesh(img.shape[0],img.shape[1],[verts[i]],smpl_faces,cam_int,colors) |
| valid_mask = (depth > 0)[:,:,None] |
| visible_weight = 0.8 |
| rendered_img = rgb[:,:,::-1] * valid_mask * visible_weight +\ |
| rendered_img * valid_mask * (1-visible_weight)+\ |
| rendered_img * (1-valid_mask) |
| rendered_img = rendered_img.astype(np.uint8) |
|
|
| return rendered_img |
|
|
|
|
| def vis_joints_img(img, j2ds): |
| pass |
|
|
| def vis_sat(img, input_size, patch_size, sat_dict, bid, padding=True): |
| if not isinstance(img, np.ndarray): |
| img = tensor_to_BGR(img.detach().cpu()) |
|
|
| assert max(img.shape[0], img.shape[1]) == input_size |
| if padding: |
| img = pad_img(img, input_size) |
|
|
| |
| pos_y, pos_x = sat_dict['pos_y'][bid], sat_dict['pos_x'][bid] |
| pos_y = (pos_y * input_size).detach().int().cpu().numpy() |
| pos_x = (pos_x * input_size).detach().int().cpu().numpy() |
|
|
| lvls = sat_dict['lvl'] |
| if lvls is None: |
| lvl = np.zeros(len(pos_x),dtype=int) |
| else: |
| lvl = lvls[bid].detach().int().cpu().numpy() |
|
|
| for (cx, cy, l) in zip(pos_x, pos_y, lvl): |
| if l == 0: |
| half_patch = patch_size//2 |
| |
| color = (173,178,241) |
| elif l == 1: |
| half_patch = patch_size |
| |
| color = (239,198,175) |
| elif l >= 2: |
| half_patch = patch_size*(2**(l-1)) |
| |
| color = (255, 255, 255) |
| else: |
| raise NotImplementedError |
|
|
| x1, x2 = cx - half_patch, cx + half_patch |
| y1, y2 = cy - half_patch, cy + half_patch |
|
|
| if l>0: |
| k = 7*l |
| img[y1:y2,x1:x2] = 0.5*cv2.blur(img[y1:y2,x1:x2].copy(),(k,k)) + 0.5*np.array(color) |
| else: |
| |
| img[y1:y2,x1:x2] = 0.5*img[y1:y2,x1:x2].copy() + 0.5*np.array(color) |
|
|
|
|
| for (cx, cy, l) in zip(pos_x, pos_y, lvl): |
| if l == 0: |
| half_patch = patch_size//2 |
| color = (139, 97, 233) |
| elif l == 1: |
| half_patch = patch_size |
| color = (246, 222, 118) |
| elif l >= 2: |
| half_patch = patch_size*(2**(l-1)) |
| color = (255, 255, 255) |
| else: |
| raise NotImplementedError |
| |
| x1, x2 = cx - half_patch, cx + half_patch |
| y1, y2 = cy - half_patch, cy + half_patch |
|
|
| cv2.rectangle(img, (x1, y1), (x2, y2), |
| color=(255,255,255), thickness = 2 ) |
| |
|
|
| return img |
|
|
| def vis_boxes(img, boxes, padding=True, color = (0,0,255)): |
| if not isinstance(img, np.ndarray): |
| img = tensor_to_BGR(img.detach().cpu()) |
| if padding: |
| pad_size = max(img.shape[0],img.shape[1]) |
| img = pad_img(img, pad_size) |
| |
| for bbox in boxes: |
| bbox = bbox.int().tolist() |
| cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), |
| color=color, thickness = 2 ) |
| |
| return img |
|
|
|
|
| def get_img_from_fig(fig, dpi=120): |
| buf = io.BytesIO() |
| fig.savefig(buf, format="png", dpi=dpi, transparent=False, bbox_inches="tight", pad_inches=0) |
| buf.seek(0) |
| img_arr = np.frombuffer(buf.getvalue(), dtype=np.uint8) |
| buf.close() |
| img = cv2.imdecode(img_arr, 1) |
| |
| return img |
|
|
|
|
| def render_mesh(height, width, meshes, face, cam_intrinsics, colors = None): |
| |
| |
| scene = pyrender.Scene(ambient_light=(0.3, 0.3, 0.3)) |
| renderer = pyrender.OffscreenRenderer(viewport_width=width, viewport_height=height, point_size=1.0) |
|
|
| |
| light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=0.8) |
| light_pose = np.eye(4) |
| light_pose[:3, 3] = np.array([0, -1, 1]) |
| scene.add(light, pose=light_pose) |
| light_pose[:3, 3] = np.array([0, 1, 1]) |
| scene.add(light, pose=light_pose) |
| light_pose[:3, 3] = np.array([1, 1, 2]) |
| scene.add(light, pose=light_pose) |
|
|
| |
| if colors is None: |
| colors = get_colors_rgb(len(meshes)) |
|
|
| for i, mesh in enumerate(meshes): |
| mesh = trimesh.Trimesh(mesh, face) |
| rot = trimesh.transformations.rotation_matrix(np.radians(180), [1, 0, 0]) |
| mesh.apply_transform(rot) |
| material = pyrender.MetallicRoughnessMaterial(metallicFactor=0.0, alphaMode='OPAQUE', baseColorFactor=(*colors[i], 1.0)) |
| mesh = pyrender.Mesh.from_trimesh(mesh, material=material, smooth=True) |
|
|
| scene.add(mesh, f'mesh_{i}') |
|
|
|
|
| |
| f=np.array([cam_intrinsics[0,0],cam_intrinsics[1,1]]) |
| c=cam_intrinsics[0:2,2] |
| camera = pyrender.camera.IntrinsicsCamera(fx=f[0], fy=f[1], cx=c[0], cy=c[1]) |
| scene.add(camera) |
|
|
| |
| rgb, depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA) |
| rgb = rgb[:,:,:3].astype(np.float32) |
| renderer.delete() |
| return rgb, depth |
|
|
