PSKernel.py

import gc
import math
import os
import pdb
from glob import glob

import numpy as np
import pandas as pd
import pydicom
import pysnooper
import tensorflow
from matplotlib import pyplot as plt
from sklearn.model_selection import KFold
from tensorflow.keras.layers import (
    Conv2D,
    Conv2DTranspose,
    Dropout,
    Input,
    MaxPooling2D,
    concatenate,
)
from tensorflow.keras.models import Model
from tensorflow.python.ops import math_ops
from tqdm import tqdm

import utils
from kernel import KaggleKernel


class PS(KaggleKernel):
    def __init__(self, *args, **kargs):
        super(PS, self).__init__(*args, **kargs)
        self._PS_init_()

    def _PS_init_(self):
        self.use_tf_data = False
        self.developing = True
        # self.train_dev_mask = None  # for train dev split
        self.DATA_PATH_BASE = "../input/siim-train-test"
        self._im_chan = 1

        self.BS = 16
        try:
            self.ds  # if has this attribute, no need to overwrite
        except AttributeError:
            self.ds = None  # for train
            self.ds_len = 0
            self.dev_ds = None
            self.dev_ds_len = 0
            self.tf_data_handler = None

    def analyze_data(self):
        pass

    @staticmethod
    def _PS_data_preprocess_np(fns, df, TARGET_COLUMN, im_height, im_width, im_chan):
        X_train = np.zeros(
            (len(fns), im_height, im_width, im_chan), dtype=np.uint8)
        Y_train = np.zeros((len(fns), im_height, im_width, 1), dtype=np.uint8)
        print("Getting train images and masks ... ")
        # sys.stdout.flush()
        for n, _id in tqdm(enumerate(fns), total=len(fns)):
            dataset = pydicom.read_file(_id)
            _id_keystr = _id.split("/")[-1][:-4]
            X_train[n] = np.expand_dims(dataset.pixel_array, axis=2)
            try:
                mask_data = df.loc[_id_keystr, TARGET_COLUMN]

                if "-1" in mask_data:
                    Y_train[n] = np.zeros((1024, 1024, 1))
                else:
                    if type(mask_data) == str:
                        Y_train[n] = np.expand_dims(
                            utils.rle2mask(
                                df.loc[_id_keystr, TARGET_COLUMN], 1024, 1024
                            ).T,
                            axis=2,
                        )
                    else:
                        Y_train[n] = np.zeros((1024, 1024, 1))
                        for x in mask_data:
                            Y_train[n] = Y_train[n] + np.expand_dims(
                                utils.rle2mask(x, 1024, 1024).T, axis=2
                            )
            except KeyError:
                print(
                    f"Key {_id.split('/')[-1][:-4]} without mask, assuming healthy patient."
                )
                # Assume missing masks are empty masks.
                Y_train[n] = np.zeros((1024, 1024, 1))

        print("Done data preprocessing as numpy array!")

        return X_train, Y_train

    @staticmethod
    def _PS_data_preprocess(fns, df, tf=False):
        """
        for tf=True, need to use TF2.0

        :param fns:
        :param df:
        :param tf:
        :return:
        """
        TARGET_COLUMN = " EncodedPixels"
        im_height = 1024
        im_width = 1024
        im_chan = 1
        # Get train images and masks
        if not tf:
            return PS._PS_data_preprocess_np(
                fns, df, TARGET_COLUMN, im_height, im_width, im_chan
            )
        else:
            return PS._PS_data_preprocess_tf(
                fns, df, TARGET_COLUMN, im_height, im_width, im_chan
            )

    @staticmethod
    def _PS_data_preprocess_tf(fns, df, TARGET_COLUMN, im_height, im_width, im_chan):
        tf_data_handler = utils.PS_TF_DataHandler()
        return tf_data_handler.to_tf_from_disk(
            fns, df, TARGET_COLUMN, im_height, im_width, im_chan
        )

    @staticmethod
    def img_build_patches(imgs):
        im_height = 128
        im_width = 128
        imgs = imgs.reshape((-1, im_height, im_width, 1))
        return imgs

    def prepare_train_dev_data(self):
        if self.use_tf_data:
            self._prepare_train_data_as_tf()
        else:
            self._prepare_train_data_as_np()

    def _get_train_data(self):
        train_data_wildcard = self.DATA_PATH_BASE + "/dicom-images-train/*/*/*.dcm"
        if self.developing:
            train_fns = sorted(glob(train_data_wildcard))[:100]
        else:
            train_fns = sorted(glob(train_data_wildcard))
        return train_fns

    # just prepare train/dev here together
    def _prepare_train_data_as_tf(self):
        train_fns = self._get_train_data()

        self.logger.debug(f"train & dev counts: {len(train_fns)}")
        df_full = pd.read_csv(
            os.path.join(self.DATA_PATH_BASE, "train-rle.csv"), index_col="ImageId"
        )
        self.ds = PS._PS_data_preprocess(train_fns, df_full, tf=True)

    # just prepare train/dev here together
    def _prepare_train_data_as_np(self):
        # getting path of all the train and test images

        train_fns = self._get_train_data()

        self.logger.debug(f"train & dev counts: {len(train_fns)}")
        df_full = pd.read_csv(
            os.path.join(self.DATA_PATH_BASE, "train-rle.csv"), index_col="ImageId"
        )

        train_fns_splits = []
        idx = 0
        LEN = 1024
        for idx in range(len(train_fns) // LEN + 1):
            train_fns_splits.append(train_fns[idx * LEN: idx * LEN + LEN])

        for i, fns in enumerate(train_fns_splits):
            images, mask_e = PS._PS_data_preprocess(fns, df_full)

            # just use sklearn split to get id and it is fine. For text thing,
            splits = list(
                KFold(n_splits=5, random_state=2019,
                      shuffle=True).split(images)
            )
            # just do 1 fold, later we can add them all back
            tr_ind, val_ind = splits[0]

            self.train_X = images[tr_ind]
            self.train_Y = mask_e[tr_ind]

            self.dev_X = images[val_ind]
            self.dev_Y = mask_e[val_ind]

            self.logger.debug(
                f"size: train_X {self.train_X.shape}, train_Y {self.train_Y.shape}"
            )

            self.logger.debug(sorted(glob("*pkl")))  # list file now
            # self.save_data_tf(file_name=f"train_dev.{i}.tfrec")
            # del self.train_X # not del if not save to tfrec
            # del self.train_Y
            # del self.dev_X
            # del self.dev_Y
            # del images
            # del mask_e
            # gc.collect()
            # self.logger.debug(sorted(glob("*")))

    def _clone_data(self, src):
        self.model = src.model
        self.model_metrics = src.model_metrics
        self.model_loss = src.model_loss

        self.train_X = src.train_X
        self.train_Y = src.train_Y

        self.dev_X = src.dev_X
        self.dev_Y = src.dev_Y

        self.test_X = src.test_X

        self.result_analyzer = src.result_analyzer  # for analyze the result

        self._stage = src._stage

    def save_data_tf(self, file_name="train_dev.tfrec"):
        assert self.train_X is not None

        all_X = np.concatenate((self.train_X, self.dev_X), axis=0)
        all_Y = np.concatenate((self.train_Y, self.dev_Y), axis=0)

        ds = utils.PS_TF_DataHandler.get_train_dataset(all_X, all_Y)

        # def train_input_fn_bt(features, labels, batch_size, cv, split_id=None, n_splits=None, ds=None):
        assert len(all_X) == len(all_Y)
        try:
            self.BS
        except AttributeError:
            self._PS_init_()
        split = 5
        # self.ds = utils.PS_TF_DataHandler.train_input_fn_bt(None, None, BS, cv=True, split_id=0, n_splits=split, ds=ds, ds_len=math.floor(len(all_X)*(1-1/split)))
        self.ds = ds
        self.ds_len = math.floor(len(all_X) * (1 - 1 / split))
        # self.dev_ds = utils.PS_TF_DataHandler.train_input_fn_bt(None, None, BS, cv=True, cv_train=False, split_id=0, n_splits=split, ds=ds, ds_len=math.floor(len(all_X) * (1 / split)))
        self.dev_ds_len = math.floor(len(all_X) * (1 / split))

        self.logger.debug(f"dev set counts: {self.dev_ds_len} / {len(all_X)}")

        utils.PS_TF_DataHandler.to_tfrecord(self.ds, file_name=file_name)

    def after_prepare_data_hook(self):
        self.analyze_data()

    def train_model(self):
        # self.model.fit(self.train_X, self.train_Y, validation_data=(self.dev_X, self.dev_Y), batch_size=256, epochs=5)  # can run now...
        v_s = math.floor(self.dev_ds_len / self.BS)

        def split_train_dev(ds, n_splits, split_id):
            ds_shards = [ds.shard(n_splits, i) for i in range(n_splits)]

            shards_cross = [
                ds_shards[val_id] for val_id in range(n_splits) if val_id != split_id
            ]

            ds_train = shards_cross[0]
            for t in shards_cross[1:]:
                ds_train = ds.concatenate(t)
            return ds, ds.shard(n_splits, split_id)

        def mask_to_binary(a, b, threshold=0.5):
            # b = tf.transpose(b, [1,0,2])
            threshold = math_ops.cast(threshold, b.dtype)
            b = math_ops.cast(b > threshold, b.dtype)
            return a, b

        def ds_prepare():
            ds = self.ds
            ds = ds.map(mask_to_binary)
            ds, ds_dev = split_train_dev(ds, 5, 0)
            ds = ds.shuffle(buffer_size=1024).repeat().batch(
                self.BS).prefetch(10)
            ds_dev = ds_dev.batch(self.BS).prefetch(10)
            self.ds = ds
            self.dev_ds = ds_dev

        ds_prepare()
        self.model.fit(
            self.ds,
            steps_per_epoch=1024,
            epochs=1,
            validation_data=self.dev_ds,
            validation_steps=v_s if v_s > 0 else 1,
            verbose=1,
        )

    def prepare_test_data(self):
        test_data_wildcard = self.DATA_PATH_BASE + "/dicom-images-test/*/*/*.dcm"
        if self.developing:
            test_fns = sorted(glob(test_data_wildcard))[:100]
        else:
            test_fns = sorted(glob(test_data_wildcard))
        self.logger.debug(f"test counts: {len(test_fns)}")

    def _build_model(self, input_layer, start_neurons):
        # ref: https://www.kaggle.com/phoenigs/u-net-dropout-augmentation-stratification
        # 128 -> 64
        conv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(
            input_layer
        )
        conv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(
            conv1
        )
        pool1 = MaxPooling2D((2, 2))(conv1)
        pool1 = Dropout(0.25)(pool1)

        # 64 -> 32
        conv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(
            pool1
        )
        conv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(
            conv2
        )
        pool2 = MaxPooling2D((2, 2))(conv2)
        pool2 = Dropout(0.5)(pool2)

        # 32 -> 16
        conv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(
            pool2
        )
        conv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(
            conv3
        )
        pool3 = MaxPooling2D((2, 2))(conv3)
        pool3 = Dropout(0.5)(pool3)

        # 16 -> 8
        conv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(
            pool3
        )
        conv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(
            conv4
        )
        pool4 = MaxPooling2D((2, 2))(conv4)
        pool4 = Dropout(0.5)(pool4)

        # Middle
        convm = Conv2D(start_neurons * 16, (3, 3), activation="relu", padding="same")(
            pool4
        )
        convm = Conv2D(start_neurons * 16, (3, 3), activation="relu", padding="same")(
            convm
        )

        # 8 -> 16
        deconv4 = Conv2DTranspose(
            start_neurons * 8, (3, 3), strides=(2, 2), padding="same"
        )(convm)
        uconv4 = concatenate([deconv4, conv4])
        uconv4 = Dropout(0.5)(uconv4)
        uconv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(
            uconv4
        )
        uconv4 = Conv2D(start_neurons * 8, (3, 3), activation="relu", padding="same")(
            uconv4
        )

        # 16 -> 32
        deconv3 = Conv2DTranspose(
            start_neurons * 4, (3, 3), strides=(2, 2), padding="same"
        )(uconv4)
        uconv3 = concatenate([deconv3, conv3])
        uconv3 = Dropout(0.5)(uconv3)
        uconv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(
            uconv3
        )
        uconv3 = Conv2D(start_neurons * 4, (3, 3), activation="relu", padding="same")(
            uconv3
        )

        # 32 -> 64
        deconv2 = Conv2DTranspose(
            start_neurons * 2, (3, 3), strides=(2, 2), padding="same"
        )(uconv3)
        uconv2 = concatenate([deconv2, conv2])
        uconv2 = Dropout(0.5)(uconv2)
        uconv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(
            uconv2
        )
        uconv2 = Conv2D(start_neurons * 2, (3, 3), activation="relu", padding="same")(
            uconv2
        )

        # 64 -> 128
        deconv1 = Conv2DTranspose(
            start_neurons * 1, (3, 3), strides=(2, 2), padding="same"
        )(uconv2)
        uconv1 = concatenate([deconv1, conv1])
        uconv1 = Dropout(0.5)(uconv1)
        uconv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(
            uconv1
        )
        uconv1 = Conv2D(start_neurons * 1, (3, 3), activation="relu", padding="same")(
            uconv1
        )

        # uconv1 = Dropout(0.5)(uconv1)
        output_layer = Conv2D(1, (1, 1), padding="same",
                              activation="sigmoid")(uconv1)

        return output_layer

    def build_and_set_model(self):

        img_size = 256
        input_layer = Input((img_size, img_size, 1))
        output_layer = self._build_model(input_layer, 16)
        model = Model(input_layer, output_layer)
        model.compile(
            loss="binary_crossentropy",
            optimizer="adam",
            metrics=["accuracy", utils.dice_coef],
        )

        self.model = model

    def set_result_analyzer(self):
        self.result_analyzer = PS_result_analyzer()

    @staticmethod
    def _check_image_data(ds):
        cnt = 0
        for image, mask in ds:
            print(cnt)
            cnt += 1
            m = mask.numpy()
            if m.max() <= 0:
                continue
            m = np.reshape(m, (1024, 1024))
            img = image.numpy()
            img = np.reshape(img, (1024, 1024))
            # plt.imshow(img)
            plt.imshow(m)
            break

    def _recover_from_tf(self):
        self.ds = utils.PS_TF_DataHandler.from_tfrecord()
        PS._check_image_data(self.ds)


class PS_result_analyzer:  # todo maybe add other good analyze functions
    def dev_set_performance(self, y_true, y_pred):
        return utils.dice_coef(y_true, y_pred)