visualizer.py

import numpy as np
import torch

import os
import io
import base64

import numpy as np
import pandas as pd
import cv2
import matplotlib.pyplot as plt
import matplotlib.colors as colors

import time
from IPython.display import clear_output
from IPython.display import HTML


def show_data_augmentations(img_tensor: torch.Tensor,
                            mask_tensor: torch.Tensor,
                            mean: tuple = (0.485, 0.456, 0.406),
                            std: tuple = (0.229, 0.224, 0.225),
                            labels: list=["image", "lung", "heart", "trachea"]):
    
    img = img_tensor.numpy().transpose(1, 2, 0)
    img = (img * std + mean).astype("float32")
    img = np.clip(img, 0, 1)
    mask = mask_tensor.numpy().transpose(1, 2, 0)
    data_to_plot = [img, *[mask[:,:, i] for i in range(3)]]

    fig, axes = plt.subplots(nrows=1, ncols=4, figsize=(15, 8))
    for i, ax in enumerate(axes):
        ax.imshow(data_to_plot[i])
        ax.set_title(labels[i])

    plt.show()
    
    
def show_video(video_path: str):
    """
    show video in jupyter notebook, agent interaction in environment.
    Takes - path to video file.
    Returns - html video player in jupyter notebook.
    """  
    video = io.open(video_path, 'r+b').read()
    encoded = base64.b64encode(video)

    return HTML(data='''<video alt="test" controls>
    <source src="data:video/mp4;base64,{0}" type="video/mp4" /> </video>'''
    .format(encoded.decode('ascii')))
  
  
def get_color_info(classes: list = ['lung', 'heart', 'trachea']):

    def get_color(cmap):
        new_cmap = colors.LinearSegmentedColormap.from_list(
            (cmap.name, 0.0, 0.0,),
            cmap(np.linspace(0.0, 0.0, 100))
            )
        return new_cmap

    colormaps = [plt.get_cmap(cmap_name) for cmap_name in 
                ['cool', 'autumn', 'autumn_r']
                ]
    arr = np.linspace(0, 50, 100).reshape((10, 10))
    fig, ax = plt.subplots(1, 3, figsize=(15, 5))
    for i in range(3):
        cmap = get_color(colormaps[i])
        ax[i].axis('off')
        ax[i].imshow(arr, cmap=cmap)
        ax[i].set_title(classes[i], fontsize=15)

    fig.suptitle("Color definition", fontsize=20, y=0.99)
    fig.savefig(f"color_definition.png", bbox_inches='tight', 
                pad_inches=0.2, dpi=100, format="png")
    
    plt.savefig(f"color_definition.svg", bbox_inches='tight',
                pad_inches=0.2, dpi=100, format="svg")
    plt.show()