"""
## Importing Library
"""
from ipywidgets import interact, interactive, IntSlider, ToggleButtons
import numpy as np
import pandas as pd
import random, re, math
import matplotlib.pyplot as plt
import seaborn as sns
"""
## Function and pathes for loading data
"""
DATA_DIR = "/kaggle/input/covid-segmentation/"
rad=DATA_DIR + "images_radiopedia.npy", DATA_DIR + "masks_radiopedia.npy"
med=DATA_DIR + "images_medseg.npy", DATA_DIR + "masks_medseg.npy"
test=DATA_DIR + "test_images_medseg.npy"

def load_case(image_file, mask_file=None):
    image =np.load(image_file).astype(np.float32)
    if not mask_file is None:
        mask =np.load(mask_file).astype(np.int8)
        return image, mask
    return image
"""
## Visualization functions
"""
"""
### Let's look at some images and their masks using widget interact 
"""
def explore_image(slices):
    plt.figure(figsize=(10, 10))
    plt.imshow(load_case(med[0],med[1])[0][slices, :, :, 0], cmap='gray');
    plt.imshow(load_case(med[0],med[1])[1][slices, :, :, 0].squeeze(), alpha=0.3, cmap='Reds')
    plt.imshow(load_case(med[0],med[1])[1][slices, :, :, 1].squeeze(), alpha=0.2, cmap='Blues')
    plt.title('Explore slices of Images_medseg \n Blue color is Consolidations \n Red color is Ground Glass  ', fontsize=20)
    plt.axis('off')
    return slices


interact(explore_image, slices=(0, load_case(med[0],med[1])[0].shape[0] - 1));
def visualize(batch_1,batch_2,pred=False):
   
    classes_dict = {'Ground_Glass': 0.,
                    'Consolidations': 1.,
                    'Lungs_Other': 2.,
                    'Background': 3. }
                                   
    
    
    fix, ax = plt.subplots(2, 5, figsize=(60, 40))
    
    for i in range(2):
        
        
        ax[i,0].set_xticks([]) 
        ax[i,0].set_yticks([])
        if i==0:
            ax[i,0].imshow(batch_1[0][:,:,0], cmap='Greys')
            ax[i,0].set_ylabel("Ground Truth ", fontsize=55) if pred==True else ax[i,0].set_ylabel("From batch_1 ", fontsize=55)
            for j in range(4):
                mask_f = batch_1[1][:,:,j]
                img_label = list(classes_dict.keys())[j]

                ax[i, j + 1].imshow(mask_f, vmin=0, vmax=1,cmap='cividis')
                ax[i, j + 1].set_xticks([]) 
                ax[i, j + 1].set_yticks([]) 
                ax[i][j + 1].set_title(f"{img_label} ", fontsize=55)
 
        else:
            ax[i,0].imshow(batch_2[0][:,:,0], cmap='Greys')
            ax[i,0].set_ylabel("Prediction", fontsize=55) if pred==True else ax[i,0].set_ylabel("From batch_2 ", fontsize=55)
            for j in range(4):
                mask_f = batch_2[1][:,:,j]
                img_label = list(classes_dict.keys())[j]

                ax[i, j + 1].imshow(mask_f, vmin=0, vmax=1,cmap='cividis')
                ax[i, j + 1].set_xticks([]) 
                ax[i, j + 1].set_yticks([]) 
 
                    

    plt.tight_layout()
    plt.show()
    
maxval = 99
i = np.random.randint(0, maxval)
visualize(batch_1=(load_case(med[0])[i], load_case(med[1])[i]),batch_2=(load_case(rad[0])[i], load_case(rad[1])[i]),pred=False)
def count_classes(image,mask):
    ground_glass=[]
    consolidations=[]
    other=[]
    background=[]
    
    for i in range(image.shape[0]):
        if mask[i,:,:,0].sum()!=0:
            ground_glass.append(image[i])
        if mask[i,:,:,1].sum()!=0:
            consolidations.append(image[i])
        if mask[i,:,:,2].sum()!=0:
            other.append(image[i])
        if mask[i,:,:,3].sum()!=0:
            background.append(image[i])
            
    x=['ground_glass','consolidations','other','background']
    y=[len(ground_glass),len(consolidations),len(other),len(background)]
    plt.figure(figsize=(12,5))
    plt.title('The number of each category segmentaion in images', fontsize=20)
    sns.set_context('paper',font_scale=2)
    sns.barplot(x=x,y=y,palette='coolwarm')
    
    print(f'The number of images that has: \n Ground_Glass Segmentation:{y[0]} \
    \n Consolidations Segmentation:{y[1]} \n Lung_Others Segmentation:{y[2]} \n Background Segmentation:{y[3]}')
        
                
    return y
y=count_classes(load_case(med[0],med[1])[0],load_case(med[0],med[1])[1])
v=count_classes(load_case(rad[0],rad[1])[0],load_case(rad[0],rad[1])[1])