File size: 8,554 Bytes
ff07ed4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 |
import os
import io
import torch
import numpy as np
from termcolor import colored
try:
# os.environ["PYOPENGL_PLATFORM"] = "osmesa"
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):
# np.random.seed(131)
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)
# print(img.shape, vis_map.shape)
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)
# draw patches
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)
# visualize patches
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 = (139, 97, 233)
color = (173,178,241)
elif l == 1:
half_patch = patch_size
# color = (246, 222, 118)
color = (239,198,175)
elif l >= 2:
half_patch = patch_size*(2**(l-1))
# color = (0,0,0)
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:
# pass
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)
# img = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
return img
def render_mesh(height, width, meshes, face, cam_intrinsics, colors = None):
# renderer
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
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)
# mesh
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}')
# camera
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)
# render
rgb, depth = renderer.render(scene, flags=pyrender.RenderFlags.RGBA)
rgb = rgb[:,:,:3].astype(np.float32)
renderer.delete()
return rgb, depth
|