%matplotlib inline
#這是juoyter notebook的magic word˙

import matplotlib
import matplotlib.pyplot as plt
from IPython import display
"""
# Steel Defect 資料預處理

"""
import os
#判斷是否在jupyter notebook上
def is_in_ipython():
    "Is the code running in the ipython environment (jupyter including)"
    program_name = os.path.basename(os.getenv('_', ''))

    if ('jupyter-notebook' in program_name or # jupyter-notebook
        'ipython'          in program_name or # ipython
        'jupyter' in program_name or  # jupyter
        'JPY_PARENT_PID'   in os.environ):    # ipython-notebook
        return True
    else:
        return False


#判斷是否在colab上
def is_in_colab():
    if not is_in_ipython(): return False
    try:
        from google import colab
        return True
    except: return False

#判斷是否在kaggke_kernal上
def is_in_kaggle_kernal():
    if 'kaggle' in os.environ['PYTHONPATH']:
        return True
    else:
        return False

if is_in_colab():
    from google.colab import drive
    drive.mount('/content/gdrive')
os.environ['TRIDENT_BACKEND'] = 'pytorch'

if is_in_kaggle_kernal():
    os.environ['TRIDENT_HOME'] = './trident'
    
elif is_in_colab():
    os.environ['TRIDENT_HOME'] = '/content/gdrive/My Drive/trident'

#為確保安裝最新版 
!pip uninstall tridentx -y
!pip install ../input/trident/tridentx-0.7.4-py3-none-any.whl --upgrade
import json
import copy
import numpy as np
#調用trident api
import trident as T
from trident import *
from trident.models import resnet,efficientnet
import random
import glob
import pandas as pd
imgs=glob.glob('../input/severstal-steel-defect-detection/train_images/*jpg')
print(len(imgs))

df_train=pd.read_csv('../input/severstal-steel-defect-detection/train.csv')
print(df_train)

df_images_frequency=df_train['ImageId'].value_counts()
print(df_images_frequency)

palette = [(0, 0, 0),(256, 192, 0), (0, 192, 256), (128, 0, 256), (256,64,0)]

def enc2mask(enc, shape=(1600,256),fill_value=0):
    img = np.zeros(shape[0]*shape[1], dtype=np.uint8)
    s = enc.split()
    for i in range(len(s)//2):
        start = int(s[2*i]) - 1
        length = int(s[2*i+1])
        img[start:start+length] = int(fill_value)
    return img.reshape(shape).T


def mask2enc(mask, n=4):
    pixels = mask.T.flatten()
    encs = []
    for i in range(1,n+1):
        p = (pixels == i).astype(np.int8)
        if p.sum() == 0: encs.append('')
        else:
            p = np.concatenate([[0], p, [0]])
            runs = np.where(p[1:] != p[:-1])[0] + 1
            runs[1::2] -= runs[::2]
            encs.append(' '.join(str(x) for x in runs))
    return encs


def label2color(label_mask,palette):
    num_classes = len(palette)

    color_label= np.zeros((*label_mask.shape,3)).astype(np.int64)
    for i in range(num_classes):
        color_label[label_mask==i]=palette[i]
    return color_label
df_train['mask'] = df_train.apply(lambda row: enc2mask(enc=row.EncodedPixels,fill_value=row.ClassId),axis=1)
print(df_train)
for i in range(10):
    exsample_mask=df_train['mask'].iloc[i]
    print(exsample_mask.shape)
    print(exsample_mask.max())

exsample_mask=df_train['mask'].iloc[3]
exsample_image=image2array('../input/severstal-steel-defect-detection/train_images/'+df_train['ImageId'].iloc[3])
print(exsample_mask.max())
display.display(array2image(exsample_mask))

is_mask=np.expand_dims(np.greater(exsample_mask,0).astype(np.float32),-1)
color_mask=label2color(exsample_mask,palette)


display.display(array2image(label2color(exsample_mask,palette)))
display.display(array2image(exsample_image))

display.display(array2image(0.5*exsample_image+0.5*(1-is_mask)*exsample_image+0.5*is_mask*color_mask))

masked_dict=OrderedDict()

for index, row in df_train.iterrows():
    img_key='../input/severstal-steel-defect-detection/train_images/'+row['ImageId']
    if img_key not in masked_dict:
        masked_dict[img_key]=row['mask']
    else:
        masked_dict[img_key]=masked_dict[img_key]+row['mask']

print(len(masked_dict))
"""
你如果試圖要把所有mask生成出來，很快你會超過kaggle記憶體上限造成notebook重啟，而若是一張一張存檔來調用看起來可行，可是會佔用不少
"""
class MyMaskDataset(MaskDataset):
    def __init__(self, masks, class_names=None, symbol="mask", **kwargs):
        super().__init__(masks,class_names=class_names, symbol=symbol, object_type=ObjectType.label_mask, **kwargs)
     
    def __getitem__(self, index: int):
        img_id = self.items[index]  # self.pop(index)
        if img_id in masked_dict:
            return masked_dict[img_id].astype(np.int64)
        else:
            return np.zeros((256,1600,3),dtype=np.int64)
        
#如果是mask非零則ok，否則只有20%機率取用
sample_filter=lambda x:x[-1].max()>0 or random.random()>0.9

ds1=ImageDataset(list(masked_dict.keys()),symbol='image')
ds2=MyMaskDataset(list(masked_dict.keys()),symbol='mask')

#設定調色盤
for i in range(5):
    ds2.palette[i] =palette[i]

data_provider=DataProvider(traindata=Iterator(data=ds1,label=ds2,sample_filter=sample_filter))
data_provider.paired_transform_funcs=[
    RandomTransformAffine(rotation_range=5, zoom_range=0.00, shift_range=0.00, shear_range=0.1, random_flip=0.15 ,border_mode='zero'),
    RandomRescaleCrop((224,224),scale=(0.8,1.2))]

data_provider.image_transform_funcs=[
                     AddNoise(0.01),
                     RandomAdjustGamma(gamma_range=(0.6,1.5)),
                     RandomAdjustContrast(value_range=(0.6, 1.5)),
                     RandomAdjustHue(value_range=(-0.5, 0.5)),
                     Normalize(127.5,127.5)]
        

img_data,mask_data=data_provider.next()
print(mask_data.shape)
print(mask_data.max())


%%time
data_provider.preview_images()
from trident.models import efficientnet,deeplab
backbond_net=efficientnet.EfficientNetB0(pretrained=True,input_shape=(3,224,224))
backbond=backbond_net.model
backbond.trainable=False
deeplabv3=deeplab.DeeplabV3_plus(backbond,atrous_rates=(6,12,18,24),num_filters=256,classes=5)
deeplabv3.load_model('../input/steeldefect/Models/deeplabv3.pth')
deeplabv3.summary()

from trident.models import densenet
tiramisu=densenet.DenseNetFcn(blocks=(4, 5, 6, 7, 8),input_shape=(3,224,224),growth_rate=16, initial_filters=32,num_classes=5)
tiramisu.load_model('../input/steeldefect/Models/tiramisu_1.pth')
tiramisu.summary()
def draw_seg_image(training_context):
    data_feed = training_context['data_feed']
    data = training_context['train_data']
    model = training_context['current_model']
    output_data=data[data_feed['output']]
    target_data=to_numpy(data['mask'])
    input_data=to_numpy(data['image'])
    output_data=np.argmax(to_numpy(output_data),1)
    tile_images_list=[]
    input_arr = []
    target_arr=[]
    output_arr=[]
    for i in range(len(output_data)):
        input_arr.append(image_backend_adaption(data_provider.reverse_image_transform(input_data[i])))
        target_arr.append(label2color(target_data[i],palette))
        output_arr.append(label2color(output_data[i],palette))
    tile_images_list.append(input_arr)
    tile_images_list.append(target_arr)
    tile_images_list.append(output_arr)
    fig = tile_rgb_images(*tile_images_list, save_path='Results/segtile_image_{0}.png', imshow=True)
    plt.close()
        
        
    
deeplabv3.with_optimizer(optimizer=AdaBelief,lr=1e-3,betas=(0.9, 0.999),gradient_centralization='all')\
    .with_loss(DiceLoss(ignore_index=0),loss_weight=2)\
    .with_loss(CrossEntropyLoss(auto_balance=True)) \
    .with_loss(FocalLoss())\
    .with_metric(pixel_accuracy,name='pixel_accuracy',print_only=True)\
    .with_metric(iou,name='iou')\
    .with_regularizer('l2',reg_weight=1e-5)\
    .with_model_save_path('Models/deeplabv3.pth') \
    .trigger_when(when='on_batch_end',frequency=50,action=draw_seg_image)\
    .with_automatic_mixed_precision_training()

tiramisu.with_optimizer(optimizer=AdaBelief,lr=1e-3,betas=(0.9, 0.999),gradient_centralization='all')\
    .with_loss(DiceLoss(ignore_index=0),loss_weight=2)\
    .with_loss(CrossEntropyLoss(auto_balance=True)) \
    .with_loss(FocalLoss()) \
    .with_metric(pixel_accuracy,name='pixel_accuracy',print_only=True)\
    .with_metric(iou,name='iou')\
    .with_regularizer('l2',reg_weight=1e-5)\
    .with_model_save_path('Models/tiramisu_1.pth') \
    .trigger_when(when='on_batch_end',frequency=50,action=draw_seg_image)\
    .with_automatic_mixed_precision_training()

plan=TrainingPlan()\
    .add_training_item(deeplabv3,name='deeplabv3')\
    .add_training_item(tiramisu,name='tiramisu')\
    .with_data_loader(data_provider)\
    .repeat_epochs(30)\
    .with_batch_size(32)\
    .print_progress_scheduling(10,unit='batch')\
    .display_loss_metric_curve_scheduling(frequency=100,unit='batch',imshow=True)\
    .save_model_scheduling(20,unit='batch')\


plan.start_now()