{"cells":[{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","collapsed":true,"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"!mkdir /kaggle/working/train_images\n!unzip /kaggle/input/data-science-bowl-2018/stage1_train.zip -d /kaggle/working/train_images","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import os\nfrom tqdm import tqdm\nimport numpy as np\nfrom skimage.io import imread\nfrom skimage.transform import resize\nfrom skimage.segmentation import find_boundaries\nfrom joblib import Parallel, delayed\nimport h5py\nimport matplotlib.pyplot as plt\n%matplotlib inline","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"op = os.path\nTRAIN_ROOT = \"/kaggle/working/train_images/\"\nOUTSIZE = 512, 512\nw0 = 10\nsigma = 5","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def iter_train_imageids():\n    \"\"\"Iterate over image IDs in the training set.\"\"\"\n    for imageid in tqdm(os.listdir(TRAIN_ROOT)):\n        if op.isdir(op.join(TRAIN_ROOT, imageid)):\n            yield imageid\n\n\ndef load_train_image(imageid, as_gray=False):\n    \"\"\"Load a given image ID as a NumPy array, optionally as grayscale.\"\"\"\n    impath = op.join(TRAIN_ROOT, imageid, 'images', imageid + '.png')\n    return imread(impath, as_gray=as_gray)\n\n\ndef load_masks(imageid):\n    \"\"\"Given an image ID, load its masks as a 3-D NumPy array of shape H x W x N;\n    Where:\n        H: height\n        W: width\n        N: number of masks\n    Each H x W slice of the resulting array is one binary mask.\"\"\"\n    maskdir = op.join(TRAIN_ROOT, imageid, 'masks')\n    masks = []\n    for maskpath in os.listdir(maskdir):\n        masks.append(imread(op.join(maskdir, maskpath), as_gray=True))\n    return np.array(masks)\n\n\ndef resize_image(imageid, outsize=(512, 512)):\n    \"\"\"Load an image and resize it.\"\"\"\n    image = load_train_image(imageid)[..., :3]\n    # image = rgba2rgb(image)\n    image = resize(image, OUTSIZE, preserve_range=True) / 255\n    masks = load_masks(imageid)\n    masks = (masks.sum(0) > 0).astype(int)\n    masks = resize(masks, OUTSIZE, preserve_range=True)\n    masks = (masks > 0).astype(int)\n    return image, masks, imageid\n\n\ndef resize_parallel(imageids, outsize=(512, 512), x_fname=\"XTrain.h5\", y_fname=\"ymasks.h5\"):\n    \"\"\"Resize images in parallel, and write them and their masks to .h5 files.\"\"\"\n    res = Parallel(n_jobs=-1)(delayed(resize_image)(imageid) for imageid in imageids)\n    images, masks, imageid = zip(*res)\n    fx = h5py.File(x_fname, 'w')\n    fy = h5py.File(y_fname, 'w')\n    for x, y, imageid in zip(images, masks, imageid):\n        fx.create_dataset(imageid, data=x)\n        fy.create_dataset(imageid, data=y)\n    fx.close()\n    fy.close()\n\n\ndef weight_map(imageid, outsize=(512, 512), show=False):\n    \"\"\"Compute the weight map for a given Image ID, as described in Ronneberger et al.\n    (https://arxiv.org/pdf/1505.04597.pdf)\n    \"\"\"\n    masks = load_masks(imageid)\n    nrows, ncols = masks.shape[1:]\n    masks = (masks > 0).astype(int)\n    distMap = np.zeros((nrows * ncols, masks.shape[0]))\n    X1, Y1 = np.meshgrid(np.arange(nrows), np.arange(ncols))\n    X1, Y1 = np.c_[X1.ravel(), Y1.ravel()].T\n    for i, mask in enumerate(masks):\n        bounds = find_boundaries(mask, mode='inner')\n        X2, Y2 = np.nonzero(bounds)\n        xSum = (X2.reshape(-1, 1) - X1.reshape(1, -1)) ** 2\n        ySum = (Y2.reshape(-1, 1) - Y1.reshape(1, -1)) ** 2\n        distMap[:, i] = np.sqrt(xSum + ySum).min(axis=0)\n    ix = np.arange(distMap.shape[0])\n    if distMap.shape[1] == 1:\n        d1 = distMap.ravel()\n        border_loss_map = w0 * np.exp((-1 * (d1) ** 2) / (2 * (sigma ** 2)))\n    else:\n        if distMap.shape[1] == 2:\n            d1_ix, d2_ix = np.argpartition(distMap, 1, axis=1)[:, :2].T\n        else:\n            d1_ix, d2_ix = np.argpartition(distMap, 2, axis=1)[:, :2].T\n        d1 = distMap[ix, d1_ix]\n        d2 = distMap[ix, d2_ix]\n        border_loss_map = w0 * np.exp((-1 * (d1 + d2) ** 2) / (2 * (sigma ** 2)))\n    xBLoss = np.zeros((nrows, ncols))\n    xBLoss[X1, Y1] = border_loss_map\n    # class weight map\n    loss = np.zeros((nrows, ncols))\n    w_1 = 1 - masks.sum() / loss.size\n    w_0 = 1 - w_1\n    loss[masks.sum(0) == 1] = w_1\n    loss[masks.sum(0) == 0] = w_0\n    ZZ = xBLoss + loss\n    ZZ = resize(ZZ, outsize, preserve_range=True)\n    if show:\n        plt.imshow(ZZ)\n        plt.colorbar()\n        plt.axis('off')\n    return ZZ, imageid\n\n\ndef weight_map_parallel(imageids, masks_fout=\"wmasks.h5\"):\n    res = Parallel(n_jobs=-1, verbose=1)(delayed(weight_map)(imageid) for imageid in imageids)\n    with h5py.File(masks_fout, 'w') as fout:\n        for mask, imageid in res:\n            fout.create_dataset(imageid, data=mask)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"collapsed":true},"cell_type":"code","source":"import random\nids = list(iter_train_imageids())\n\ntrain_ids = []\ntest_ids = []\n\n\ndef sample(acc):\n    choice = random.choice(ids)\n    acc.append(choice)\n    ids.remove(choice)\n\n# Pick 80 images for training\n[sample(train_ids) for i in range(80)]\n\n# And 20 for validation\n[sample(test_ids) for i in range(20)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Write the .h5 files for train & validation datasets\n# (The h5 files are just for my personal convenience, we don't have to create them.)\nresize_parallel(train_ids, x_fname=\"xTrain.h5\", y_fname=\"yTrain.h5\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"resize_parallel(test_ids, x_fname=\"xTest.h5\", y_fname=\"yTest.h5\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"weight_map_parallel(train_ids, masks_fout=\"weightTrain.h5\")","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"xtrain = []  # Collect the training image array\nwtrain = []  # Collect the weight mask arrays\nytrain = []  # Collect the masks\n\nxtest = []  # And same for the testing set\nytest = []\n\nwith h5py.File('/kaggle/working/xTrain.h5', 'r') as fin:\n    for imid in train_ids:\n        xtrain.append(fin[imid][:])\n\nwith h5py.File('/kaggle/working/xTest.h5', 'r') as fin:\n    for imid in test_ids:\n        xtest.append(fin[imid][:])\n\nwith h5py.File('/kaggle/working/yTrain.h5', 'r') as fin:\n    for imid in train_ids:\n        ytrain.append(fin[imid][:])\n\nwith h5py.File('/kaggle/working/yTest.h5', 'r') as fin:\n    for imid in test_ids:\n        ytest.append(fin[imid][:])\n\nwith h5py.File('/kaggle/working/weightTrain.h5', 'r') as fin:\n    for imid in train_ids:\n        wtrain.append(fin[imid][:])\n\n\nxtrain = np.array(xtrain)\nytrain = np.array(ytrain)\nwtrain = np.array(wtrain)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import tensorflow as tf\nL = tf.keras.layers\nK = tf.keras.backend\nModel = tf.keras.models.Model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# Tensors for the model to work with\nycat = tf.keras.utils.to_categorical(ytrain)\nwmap = np.zeros((80, 512, 512, 2), dtype=np.float32)\nwmap[..., 0] = wtrain\nwmap[..., 1] = wtrain","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"_epsilon = tf.convert_to_tensor(K.epsilon(), np.float32)\n\n\ndef my_loss(target, output):\n    \"\"\"\n    A custom function defined to simply sum the pixelwise loss.\n    This function doesn't compute the crossentropy loss, since that is made a\n    part of the model's computational graph itself.\n    Parameters\n    ----------\n    target : tf.tensor\n        A tensor corresponding to the true labels of an image.\n    output : tf.tensor\n        Model output\n    Returns\n    -------\n    tf.tensor\n        A tensor holding the aggregated loss.\n    \"\"\"\n    return - tf.reduce_sum(target * output,\n                           len(output.get_shape()) - 1)\n\n\ndef make_weighted_loss_unet(input_shape, n_classes):\n    # two inputs, one for the image and one for the weight maps\n    ip = tf.keras.Input(shape=input_shape, name=\"image_input\")\n    # the shape of the weight maps has to be such that it can be element-wise\n    # multiplied to the softmax output.\n    weight_ip = tf.keras.Input(shape=input_shape[:2] + (n_classes,))\n\n    # adding the layers\n    conv1 = L.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(ip)\n    conv1 = L.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1)\n    conv1 = L.Dropout(0.1)(conv1)\n    mpool1 = L.MaxPool2D()(conv1)\n\n    conv2 = L.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(mpool1)\n    conv2 = L.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)\n    conv2 = L.Dropout(0.2)(conv2)\n    mpool2 = L.MaxPool2D()(conv2)\n\n    conv3 = L.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(mpool2)\n    conv3 = L.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)\n    conv3 = L.Dropout(0.3)(conv3)\n    mpool3 = L.MaxPool2D()(conv3)\n\n    conv4 = L.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(mpool3)\n    conv4 = L.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)\n    conv4 = L.Dropout(0.4)(conv4)\n    mpool4 = L.MaxPool2D()(conv4)\n\n    conv5 = L.Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(mpool4)\n    conv5 = L.Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5)\n    conv5 = L.Dropout(0.5)(conv5)\n\n    up6 = L.Conv2DTranspose(512, 2, strides=2, kernel_initializer='he_normal', padding='same')(conv5)\n    conv6 = L.Concatenate()([up6, conv4])\n    conv6 = L.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)\n    conv6 = L.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)\n    conv6 = L.Dropout(0.4)(conv6)\n\n    up7 = L.Conv2DTranspose(256, 2, strides=2, kernel_initializer='he_normal', padding='same')(conv6)\n    conv7 = L.Concatenate()([up7, conv3])\n    conv7 = L.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv7)\n    conv7 = L.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv7)\n    conv7 = L.Dropout(0.3)(conv7)\n\n    up8 = L.Conv2DTranspose(128, 2, strides=2, kernel_initializer='he_normal', padding='same')(conv7)\n    conv8 = L.Concatenate()([up8, conv2])\n    conv8 = L.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv8)\n    conv8 = L.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv8)\n    conv8 = L.Dropout(0.2)(conv8)\n\n    up9 = L.Conv2DTranspose(64, 2, strides=2, kernel_initializer='he_normal', padding='same')(conv8)\n    conv9 = L.Concatenate()([up9, conv1])\n    conv9 = L.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)\n    conv9 = L.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)\n    conv9 = L.Dropout(0.1)(conv9)\n\n    c10 = L.Conv2D(n_classes, 1, activation='softmax', kernel_initializer='he_normal', name=\"unet-activation\")(conv9)\n\n    # Add a few non trainable layers to mimic the computation of the crossentropy\n    # loss, so that the actual loss function just has to peform the\n    # aggregation.\n    c11 = L.Lambda(lambda x: x / tf.reduce_sum(x, len(x.get_shape()) - 1, True))(c10)\n    c11 = L.Lambda(lambda x: tf.clip_by_value(x, _epsilon, 1. - _epsilon))(c11)\n    c11 = L.Lambda(lambda x: K.log(x))(c11)\n    weighted_sm = L.multiply([c11, weight_ip])\n\n    model = Model(inputs=[ip, weight_ip], outputs=[weighted_sm])\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"model = make_weighted_loss_unet((512, 512, 3), 2)\nadam = tf.keras.optimizers.Adam()\nmodel.compile(adam, loss=my_loss)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true,"collapsed":true},"cell_type":"code","source":"history = model.fit([xtrain, wmap], ycat, batch_size=1, epochs=100)\nplt.plot(history.history['loss'])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_input = model.get_layer('image_input').input\nsoftmax_output = model.get_layer('unet-activation').output\npredictor = K.function([image_input], [softmax_output])","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"testImage = xtest[0]\ntestLabel = ytest[0]\nfig, (ax1, ax2) = plt.subplots(nrows=1, ncols=2)\nax1.imshow(testImage, cmap=plt.cm.gray)\nax2.imshow(testLabel, cmap=plt.cm.gray)\nax1.axis('off')\n_ = ax2.axis('off')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"yhat = predictor([testImage.reshape((1, 512, 512, 3))])[0]\nyhat = np.argmax(yhat[0], axis=-1)\nplt.imshow(yhat)","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat":4,"nbformat_minor":4}