UnOpticalFlow/core/visualize/visualizer.py

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.use_flow_error = (self.loss_weights_dict['flow_error'] > 0)
        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))
            #if self.use_flow_error:
            #    loss_flow_error = loss_pack['flow_error'].mean().detach().cpu().numpy()
            #    str_ = str_ + ', loss_flow_error: {0:.6f}'.format(loss_flow_error)
            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):
        # img: [H, W, 3] match: [4, n] point2d_coord: [n, 2] P: [3, 4]
        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):
        # img: [H, W, 3] match: [4,n] F: [3,3]
        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):
        # img: [H, W, 3] match: [4, n]
        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]
            #print((x1, y1))
            #print((x2, y2))
            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: [H, W, 1]
        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):
        # depth: [h,w,1]
        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):
        # img: [H,W,3] match: [4,n] F: [3,3]
        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)

        # Find epilines corresponding to points in left image (first image) and
        # drawing its lines on right image
        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], c=points[:,2], cmap=cmap)
        # ax.plot3D(points[:,0], points[:,1], points[:,2], c=points[:,2])
        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):
        # match: [4] P: [3,4]
        K = P1[:,:3] # the first P1 has identity rotation matrix and zero translation.
        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)