| import os, sys |
| import numpy as np |
| import cv2 |
| import pdb |
| import pickle |
| from mpl_toolkits import mplot3d |
| import matplotlib.pyplot as plt |
| import PIL.Image as pil |
| import matplotlib as mpl |
| import matplotlib.cm as cm |
|
|
|
|
| colorlib = [(0,0,255),(255,0,0),(0,255,0),(255,255,0),(0,255,255),(255,0,255),(0,0,0),(255,255,255)] |
|
|
| class Visualizer(object): |
| def __init__(self, loss_weights_dict, dump_dir=None): |
| self.loss_weights_dict = loss_weights_dict |
| |
| self.dump_dir = dump_dir |
|
|
| self.log_list = [] |
|
|
| def add_log_pack(self, log_pack): |
| self.log_list.append(log_pack) |
|
|
| def dump_log(self, fname=None): |
| if fname is None: |
| fname = self.dump_dir |
| with open(fname, 'wb') as f: |
| pickle.dump(self.log_list, f) |
|
|
| def print_loss(self, loss_pack, iter_=None): |
| loss_pixel = loss_pack['loss_pixel'].mean().detach().cpu().numpy() |
| loss_ssim = loss_pack['loss_ssim'].mean().detach().cpu().numpy() |
| loss_flow_smooth = loss_pack['loss_flow_smooth'].mean().detach().cpu().numpy() |
| loss_flow_consis = loss_pack['loss_flow_consis'].mean().detach().cpu().numpy() |
| if 'pt_depth_loss' in loss_pack.keys(): |
| loss_pt_depth = loss_pack['pt_depth_loss'].mean().detach().cpu().numpy() |
| loss_pj_depth = loss_pack['pj_depth_loss'].mean().detach().cpu().numpy() |
| loss_depth_smooth = loss_pack['depth_smooth_loss'].mean().detach().cpu().numpy() |
| str_= ('iter: {0}, loss_pixel: {1:.6f}, loss_ssim: {2:.6f}, loss_pt_depth: {3:.6f}, loss_pj_depth: {4:.6f}, loss_depth_smooth: {5:.6f}'.format(\ |
| iter_, loss_pixel, loss_ssim, loss_pt_depth, loss_pj_depth, loss_depth_smooth)) |
| |
| |
| |
| print(str_) |
| else: |
| print('iter: {4}, loss_pixel: {0:.6f}, loss_ssim: {1:.6f}, loss_flow_smooth: {2:.6f}, loss_flow_consis: {3:.6f}'.format(loss_pixel, loss_ssim, loss_flow_smooth, loss_flow_consis, iter_)) |
|
|
| class Visualizer_debug(): |
| def __init__(self, dump_dir=None, img1=None, img2=None): |
| self.dump_dir = dump_dir |
| self.img1 = img1 |
| self.img2 = img2 |
| |
| def draw_point_corres(self, batch_idx, match, name): |
| img1 = self.img1[batch_idx] |
| img2 = self.img2[batch_idx] |
| self.show_corres(img1, img2, match, name) |
| print("Correspondence Saved in " + self.dump_dir + '/' + name) |
|
|
| def draw_invalid_corres_ray(self, img1, img2, depth_match, point2d_1_coord, point2d_2_coord, point2d_1_depth, point2d_2_depth, P1, P2): |
| |
| idx = np.where(point2d_1_depth < 0)[0] |
| select_match = depth_match[:, idx] |
| self.show_corres(img1, img2, select_match) |
| pdb.set_trace() |
| |
| def draw_epipolar_line(self, batch_idx, match, F, name): |
| |
| img1 = self.img1[batch_idx] |
| img2 = self.img2[batch_idx] |
| self.show_epipolar_line(img1, img2, match, F, name) |
| print("Epipolar Lines Saved in " + self.dump_dir + '/' + name) |
|
|
| def show_corres(self, img1, img2, match, name): |
| |
| cv2.imwrite(os.path.join(self.dump_dir, name+'_img1_cor.png'), img1) |
| cv2.imwrite(os.path.join(self.dump_dir, name+'_img2_cor.png'), img2) |
| img1 = cv2.imread(os.path.join(self.dump_dir, name+'_img1_cor.png')) |
| img2 = cv2.imread(os.path.join(self.dump_dir, name+'_img2_cor.png')) |
| n = np.shape(match)[1] |
| for i in range(n): |
| x1,y1 = match[:2,i] |
| x2,y2 = match[2:,i] |
| |
| |
| cv2.circle(img1, (x1,y1), radius=1, color=colorlib[i%len(colorlib)], thickness=2) |
| cv2.circle(img2, (x2,y2), radius=1, color=colorlib[i%len(colorlib)], thickness=2) |
| cv2.imwrite(os.path.join(self.dump_dir, name+'_img1_cor.png'), img1) |
| cv2.imwrite(os.path.join(self.dump_dir, name+'_img2_cor.png'), img2) |
| |
| def show_mask(self, mask, name): |
| |
| mask = mask / np.max(mask) * 255.0 |
| cv2.imwrite(os.path.join(self.dump_dir, name+'.png'), mask) |
| |
| def save_img(self, img, name): |
| cv2.imwrite(os.path.join(self.dump_dir, name+'.png'), img) |
| |
| def save_depth_img(self, depth, name): |
| |
| minddepth = np.min(depth) |
| maxdepth = np.max(depth) |
| depth_nor = (depth-minddepth) / (maxdepth-minddepth) * 255.0 |
| depth_nor = depth_nor.astype(np.uint8) |
| cv2.imwrite(os.path.join(self.dump_dir, name+'_depth.png'), depth_nor) |
| |
| def save_disp_color_img(self, disp, name): |
| vmax = np.percentile(disp, 95) |
| normalizer = mpl.colors.Normalize(vmin=disp.min(), vmax=vmax) |
| mapper = cm.ScalarMappable(norm=normalizer, cmap='magma') |
| colormapped_im = (mapper.to_rgba(disp)[:,:,:3] * 255).astype(np.uint8) |
| im = pil.fromarray(colormapped_im) |
| |
| name_dest_im = os.path.join(self.dump_dir, name + '_depth.jpg') |
| im.save(name_dest_im) |
|
|
|
|
| def drawlines(self, img1, img2, lines, pts1, pts2): |
| ''' img1 - image on which we draw the epilines for the points in img2 |
| lines - corresponding epilines ''' |
| r,c, _ = img1.shape |
| for r,pt1,pt2 in zip(lines,pts1,pts2): |
| color = tuple(np.random.randint(0,255,3).tolist()) |
| x0,y0 = map(int, [0, -r[2]/r[1] ]) |
| x1,y1 = map(int, [c, -(r[2]+r[0]*c)/r[1] ]) |
| img1 = cv2.line(img1, (x0,y0), (x1,y1), color,1) |
| img1 = cv2.circle(img1,tuple(pt1),3,color,-1) |
| img2 = cv2.circle(img2,tuple(pt2),3,color,-1) |
| return img1,img2 |
| |
| def show_epipolar_line(self, img1, img2, match, F, name): |
| |
| pts1 = np.transpose(match[:2,:], [1,0]) |
| pts2 = np.transpose(match[2:,:], [1,0]) |
| lines1 = cv2.computeCorrespondEpilines(pts2.reshape(-1,1,2), 2,F) |
| lines1 = lines1.reshape(-1,3) |
| img5,img6 = self.drawlines(img1,img2,lines1,pts1,pts2) |
|
|
| |
| |
| lines2 = cv2.computeCorrespondEpilines(pts1.reshape(-1,1,2), 1,F) |
| lines2 = lines2.reshape(-1,3) |
| img3,img4 = self.drawlines(img2,img1,lines2,pts2,pts1) |
|
|
| cv2.imwrite(os.path.join(self.dump_dir, name+'_1eline.png'), img5) |
| cv2.imwrite(os.path.join(self.dump_dir, name+'_2eline.png'), img3) |
|
|
| return None |
|
|
| |
| def show_ray(self, ax, K, RT, point2d, cmap='Greens'): |
| K_inv = np.linalg.inv(K) |
| R, T = RT[:,:3], RT[:,3] |
| ray_direction = np.matmul(np.matmul(R.T, K_inv), np.array([point2d[0], point2d[1], 1])) |
| ray_direction = ray_direction / (np.linalg.norm(ray_direction, ord=2) + 1e-12) |
| ray_origin = (-1) * np.matmul(R.T, T) |
|
|
| scatters = [ray_origin + t * ray_direction for t in np.linspace(0.0, 100.0, 1000)] |
| scatters = np.stack(scatters, axis=0) |
| self.visualize_points(ax, scatters, cmap=cmap) |
| self.scatter_3d(ax, scatters[0], scatter_color='r') |
| return ray_direction |
|
|
| def visualize_points(self, ax, points, cmap=None): |
| |
| |
| ax.plot3D(points[:,0], points[:,1], points[:,2]) |
|
|
| def scatter_3d(self, ax, point, scatter_color='r'): |
| ax.scatter(point[0], point[1], point[2], c=scatter_color) |
|
|
| def visualize_two_rays(self, ax, match, P1, P2): |
| |
| K = P1[:,:3] |
| K_inv = np.linalg.inv(K) |
| RT1, RT2 = np.matmul(K_inv, P1), np.matmul(K_inv, P2) |
| x1, y1, x2, y2 = match |
| d1 = self.show_ray(ax, K, RT1, [x1, y1], cmap='Greens') |
| d2 = self.show_ray(ax, K, RT2, [x2, y2], cmap='Reds') |
| print(np.dot(d1.squeeze(), d2.squeeze())) |
| |
| if __name__ == '__main__': |
| img1 = cv2.imread('./vis/ga.png') |
| img2 = cv2.imread('./vis/gb.png') |
| match = np.load('./vis/gmatch.npy') |
| print(np.shape(img1)) |
| match = np.reshape(match, [4,-1]) |
| select_match = match[:,np.random.randint(200000, size=100)] |
| show_corres(img1, img2, select_match) |
|
|
