visualization.py

import torch
import numpy as np
import matplotlib.pyplot as plt
from se3_helpers import get_T_CO_init_and_gt
from renderer import render_scene
from models import fetch_network
from rotation_representation import calculate_T_CO_pred
import os
import torch
from data_loaders import *
from parser_config import get_dict_from_cli
from loss import compute_ADD_L1_loss
from image_dataloaders import get_dataloader
from PIL import Image
from scipy.signal import convolve2d
import cv2


def create_silhouette(img, silhouette_only=False):
    gray = np.mean(img, axis=2)
    sil = np.where(gray>0, 1.0, 0.0)
    filt = np.ones((5,5))
    res = convolve2d(sil, filt, mode='same')
    res = np.clip(res, 0.0, 1.0)
    sil = res-sil
    if silhouette_only:
        return sil
    else:
        img[sil>0] = (1.0,0.0, 0.0)
        return img


def blend_imgs(im1, im2, alpha):
    im1 = Image.fromarray(np.uint8(im1*255.0))
    im2 = Image.fromarray(np.uint8(im2*255.0))
    return np.asarray(Image.blend(im1,im2,alpha))/255.0

def save_img(img, resize_to, save_dir, filename):
    save_path = os.path.join(save_dir, filename)
    img_cv2 = cv2.resize(cv2.cvtColor(np.uint8(img*255), cv2.COLOR_RGB2BGR), resize_to)
    cv2.imwrite(save_path, img_cv2)



def create_interpolated_image_fig(init_imgs, gt_imgs, pred_imgs_sequence, gt_imgs_raster=None,  save_dir=None, show_fig=False, train_examples="", train_val_or_test =None):
    print("init_imgs", type(init_imgs), init_imgs.shape)
    print("gt imgs", type(gt_imgs), gt_imgs.shape)
    print("pred_imgs_seq", type(pred_imgs_sequence), np.array(pred_imgs_sequence).shape)
    print("gt_imgs_raster", type(gt_imgs_raster), gt_imgs_raster.shape)
    print("save_dir:", save_dir)
    BLEND_ALPHA = 0.4
    IMG_SIZE = (200,200)

    save_figure = (save_dir is not None)

    batch_size = len(init_imgs)
    iter_num = len(pred_imgs_sequence)
    fig, ax = plt.subplots(batch_size, iter_num+2, squeeze=False)
    fig.set_size_inches(18.5, 10.5, forward=True)
    if(save_dir is not None):
        assert(train_val_or_test is not None)
        save_subdir = os.path.join(save_dir, train_val_or_test)
        print("save subdir", save_subdir)
        os.makedirs(save_subdir, exist_ok=True)
        save_path = os.path.join(save_subdir, "viz_at_train_ex"+str(train_examples)+".png")
        print(save_path)
        save_full_img_subdir = os.path.join(save_subdir, "train_ex"+str(train_examples))
        os.makedirs(save_full_img_subdir, exist_ok=True)
        
    #ax = fig.add_subplot(batch_size, iter_num+1)
    for i in range(batch_size):
        init_img = init_imgs[i]
        gt_img = gt_imgs[i]
        ax[i,0].imshow(gt_img)
        ax[i,0].axis('off')
        ax[i,0].set_title("Real img")
        if(gt_imgs_raster is not None):
            gt_raster = gt_imgs_raster[i]
            sil = create_silhouette(gt_raster, silhouette_only=True)
            gt_img[sil>0] = (0.0,1.0, 0.0)
        initial = blend_imgs(create_silhouette(init_img), gt_img, BLEND_ALPHA)
        
        if save_figure:
            save_full_img_subdir_bsz = os.path.join(save_full_img_subdir, "example"+str(i))
            os.makedirs(save_full_img_subdir_bsz, exist_ok=True)
            save_img(gt_img, IMG_SIZE, save_full_img_subdir_bsz, "0-gt.png")
            save_img(initial, IMG_SIZE, save_full_img_subdir_bsz, "1-init.png")
        ax[i,1].imshow(initial)
        ax[i,1].axis('off')
        ax[i,1].set_title("Green: GT. Red: init")
        for j in range(iter_num):
            gt_img = gt_imgs[i]
            pred_img = pred_imgs_sequence[j][i]
            blend_pred = blend_imgs(create_silhouette(pred_img), gt_img, BLEND_ALPHA)
            if save_figure:
                save_img(blend_pred, IMG_SIZE, save_full_img_subdir_bsz, f'{j+2}-pred-{j}.png')
            ax[i,j+2].imshow(blend_pred)
            ax[i,j+2].axis('off')
            ax[i,j+2].set_title(f'Green: GT. Red: pred {j}')

    if show_fig:
        plt.show()
    if(save_figure):
        print("Saving visualization:", save_path)
        plt.savefig(save_path,dpi=150, bbox_inches='tight')
    plt.close()


def combine_imgs(img1, img2):
    gs1 = np.mean(img1, axis=2)
    gs2 = np.mean(img2, axis=2)
    img = np.zeros((gs1.shape[0], gs1.shape[1], 3))
    img[:,:,0] = gs1
    img[:,:,1] = gs2
    return img

def create_rgb_overlapped_image_fig(init_imgs, gt_imgs, pred_imgs_sequence, save_dir, train_val_or_test, show_fig, save_fig):
    print("save_dir:", save_dir)
    batch_size = len(init_imgs)
    iter_num = len(pred_imgs_sequence)
    fig, ax = plt.subplots(batch_size, iter_num+1)
    fig.set_size_inches(18.5, 10.5, forward=True)
    #ax = fig.add_subplot(batch_size, iter_num+1)
    for i in range(batch_size):
        init_img = init_imgs[i]
        gt_img = gt_imgs[i]
        ax[i,0].imshow(combine_imgs(init_img, gt_img))
        ax[i,0].axis('off')
        ax[i,0].set_title("Green: GT. Red: init")
        for j in range(iter_num):
            gt_img = gt_imgs[i]
            init_img = init_imgs[i]
            pred_img = pred_imgs_sequence[j][i]
            ax[i,j+1].imshow(combine_imgs(pred_img, gt_img))
            ax[i,j+1].axis('off')
            ax[i,j+1].set_title(f'Green: GT. Red: pred {j}')
    if show_fig:
        plt.show()
    if save_fig:
        plt.savefig(save_path,dpi=300, bbox_inches='tight')
    plt.close()

def visualize_examples(model, config, train_val_or_test, show_fig=False, save_dir=None, n_train_examples=""):

    assert (train_val_or_test=="test" or train_val_or_test=="train" or train_val_or_test=='val')

    batch_size = 5
    device = config["train_params"]["device"]
    iter_num = 5

    scene_config = config["scene_config"]
    ds_name = config["dataset_config"]["model3d_dataset"]
    img_size= config["camera_intrinsics"]["image_resolution"]
    print("img size", img_size)
    rot_repr = config["network"]["rotation_representation"]
    use_norm_depth = config["advanced"]["use_normalized_depth"]
    use_par_render = config["scene_config"]["use_parallel_rendering"]
    if train_val_or_test == 'train':
        ds_conf = config["dataset_config"]
    elif train_val_or_test=='val':
        ds_conf = config["val_dataset_config"]
    elif(train_val_or_test=='test'):
        ds_conf = config["test_dataset_config"]

    #model = fetch_network(model_name, rot_repr, use_norm_depth, use_pretrained=True, pretrained_path=model_load_path)
    model.eval()
    #model = model.to(device)


    data_loader = get_dataloader(ds_conf, batch_size, train_val_or_test)
    init_imgs, gt_imgs, T_CO_init, T_CO_gt, mesh_verts, mesh_paths, depths, cam_mats = next(iter(data_loader))
    T_CO_init = T_CO_init.numpy()
    T_CO_gt = T_CO_gt.numpy()
    gt_imgs = gt_imgs.numpy()
    init_imgs = init_imgs.numpy()
    print(cam_mats.numpy())
    depths = depths.numpy()
    gt_imgs_raster, _ = render_batch(T_CO_gt, mesh_paths, cam_mats.numpy(), img_size, use_par_render)

    T_CO_pred = T_CO_init
    pred_imgs = init_imgs
    pred_imgs_sequence = []
    T_CO_gt = torch.tensor(T_CO_gt).to(device)
    with torch.no_grad():
        for i in range(iter_num):
            model_input = prepare_model_input(pred_imgs, gt_imgs, depths, use_norm_depth).to(device)
            #cam_mats = get_camera_mat_tensor(cam_intrinsics, batch_size).to(device)
            #mesh_verts = sample_verts_to_batch(mesh_paths, num_sample_verts).to(device)

            T_CO_pred = torch.tensor(T_CO_pred).to(device)
            model_output = model(model_input)
            T_CO_pred = calculate_T_CO_pred(model_output, T_CO_pred, rot_repr, cam_mats)
            T_CO_pred = T_CO_pred.detach().cpu().numpy()
            pred_imgs, depths = render_batch(T_CO_pred, mesh_paths, cam_mats.numpy(), img_size, use_par_render)
            if not use_norm_depth: norm_depth=None
            
            pred_imgs_sequence.append(pred_imgs)

    #create_rgb_overlapped_image_fig(init_imgs, gt_imgs, pred_imgs_sequence, save_path, show_fig, save_fig)
    print("Creating images")
    create_interpolated_image_fig(init_imgs, gt_imgs, pred_imgs_sequence, gt_imgs_raster, save_dir, show_fig, n_train_examples, train_val_or_test)



if __name__ == '__main__':
    config = get_dict_from_cli()
    model_name = config["network"]["backend_network"]
    rotation_repr = config["network"]["rotation_representation"]
    use_pretrained = config["model_io"]["use_pretrained_model"]
    model_save_dir = config["model_io"]["model_save_dir"]
    model_save_name = config["model_io"]["model_save_name"]
    train_ds = config["dataset_config"]["img_dataset"]
    pretrained_path = os.path.join(model_save_dir, model_save_name)
    print("Loading pretrained from", pretrained_path)
    use_norm_depth = config["advanced"]["use_normalized_depth"]
    model = fetch_network(model_name, rotation_repr, use_norm_depth, True, pretrained_path).to('cuda')
    visualize_examples(model, config, "test", show_fig=True)