{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":52279,"databundleVersionId":5822112,"sourceType":"competition"}],"dockerImageVersionId":30673,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Buckit Net: How far down can you sample? \nBiModal UNet Consistent with Kernel Integrated Targets\n\n\n(please note, this is a first draft and the whole thing is extremely disorganized. I'm just trying to get this thing running, atm)\nThis notebook is not sufficient to offer a submission for this challenge. ","metadata":{}},{"cell_type":"markdown","source":"# Introduction\n\nThis notebook illustrates Buckit Net (Bimodal Unet consistent with Kernel Integrated Targets). This model is designed for semantic segmentation tasks, and is unique because it allows a \"full resolution\" UNet to be trained with significantly down-sampled target data. In this notebook we'll train a UNet on (512, 512) sized input images, generate (512, 512, 1) segmentation mask predictions, but train using (32, 32, 2) sized target data.\n\n\nIn many semantic segmentation tasks, it is common to down sample both input images and target mask data, and train a UNet at this lower resolution. Afterwards, the models' predictions are typically up sampled to full resolution. Buckit net takes a different approach. We keep the UNet at full resolution, but apply a statistical head to both the predictions of the UNet, as well as our target data. In some sense, this is like evaluating the error between the model's prediction blurred, and the target data blurred.\n\n\nHere, our statistical head is the concatenation of a Max Pooling operation, and a Min Pooling operation. In this notebook each pooling operation has a pooling size of 32, and a stride of 16. Applied to the original (512, 512, 1) sized labeled mask data, this yields a (32, 32, 2) down sampled target data. We also use a custom loss function, based on Binary Crossentropy, to handle the 2-channel output.","metadata":{}},{"cell_type":"markdown","source":"# Motivation\n\nIn computer vision, high resolution data can be burdensome. There are some hard practical issues, such as VRAM necessary to process high resolution images. There are also subtler issues, such as concerns of overfitting, as well as how consistent human labelers can draw particles in labeled data. Its common in many semantic segmentation tasks to down sample both input images as well as labeled data, so that a UNet can be applied simplistically. In my estimation, it doesn't seem reasonable to me to throw away hard-earned high resolution data. As such, Buckit Net is designed to try and capture the benefit of down sampling labeled data, while not throwing away information from the input images.\n\n\nBuckit Net does not up sample prediction masks, and does not use full resolution labeled data. Instead, Buckit Net applies a Statistical Head to the full-resolution target data, and applies that same Statistical Head to the full-resolution UNet output during training. In some sense, instead of setting up our loss as BCE(y_true, y_pred), we instead have BCE(S(y_true), S(y_pred)). This is somewhat like taking a loss based on your predicted mask blurred, compared to your target data blurred. In this way, Buckit Net is trained to be consistent with target data, not exactly match it. I think one way to motivate the need for Consistency with Targets, rather than exact matching, is to consider a dataset which has two copies of an image, but with two slightly different masks. it seems to me, this will lead to a minimum error floor, which could lead to instability especially with adaptive and momentum based optimizers.\n","metadata":{}},{"cell_type":"markdown","source":"# Data Preprocessing\n\nTo illustrate how this model works, we'll be using the HuBMAP - Hacking the Human Vascular System challenge. We won't be using all of this data, only the vascular mask data, as well as using the best annotated dataset, which is marked as dataset one. Once we've loaded the dataset into a numpy array format, we'll go through other basic steps quickly, including train-validation split, writing a data generator, writing some callbacks to visualize predictions, and then training the model.\n\nAfter the model has been trained we'll briefly look at some sample predictions, and talk about the details behind some design decisions. \n","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport json\nimport os\nimport tifffile as tiff\nfrom PIL import Image\nimport cv2\n\n\n#This is data loading is terribly written, sorry. This was translated from script, and compute hours are fininte. \n\nsource_dir = \"/kaggle/input/hubmap-hacking-the-human-vasculature/\"\ndata_dir = \"/kaggle/working/\"\ntarget_dir = data_dir + \"target/\"\ntrain_dir = data_dir + \"train/\"\n\nif not os.path.exists(target_dir):\n    os.makedirs(target_dir)\n    \nif not os.path.exists(train_dir):\n    os.makedirs(train_dir)\n\n    #This is a super hacky way to avoid re-doing all the \n    images = []\n    targets = []\n    filenames = []\n\n    #first, get target data\n    with open(source_dir+\"polygons.jsonl\",'r') as target_file:\n\n        json_labels = [json.loads(line) for line in target_file]\n\n\n\n    sample_names = []\n\n    for sample in json_labels:\n        label_stack = []\n        filename = sample['id']\n\n        for anno in sample['annotations']:\n\n            if anno['type'] == \"blood_vessel\":\n                lines = np.asarray(anno['coordinates'])\n                label_slice = np.zeros((512,512,1),dtype=np.uint8)\n                cv2.fillPoly(label_slice, [lines], 1)\n\n                label_stack.append(label_slice)\n        sample_names.append(filename)\n        if label_stack==[]:\n            label_stack = np.zeros((1,512,512,1),dtype=bool)\n\n        np.save(data_dir+\"target/\"+filename+\".npy\", np.asarray(label_stack,dtype=bool))\n\n\n    np.save(data_dir+\"filenames.npy\", np.asarray(sample_names))\n    for name in sample_names:\n        with tiff.TiffFile(source_dir+\"/train/\"+name+\".tif\") as tif:\n                    img = tif.asarray()\n                    np.save(train_dir+name+'.npy',img)","metadata":{"_kg_hide-input":true,"scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Sorry for how badly that's written. This was not originally trained in a notebook env. Will fix maybe. ","metadata":{}},{"cell_type":"markdown","source":"Here we just open up the right filenames in the right dataset, and get a simple array of images, and their corresponding masks. ","metadata":{}},{"cell_type":"code","source":"x_meta = pd.read_csv(source_dir+\"tile_meta.csv\")#ONLY DATASET 1 HAS EXPERT ANNOTATIONS. DS2 has sparse annotations, others do not have any!\n\nx_filenames = x_meta['id'][x_meta['dataset']==1]\nx_filenames = [str(name) for name in x_filenames]\n\nx_images = np.asarray([np.load(data_dir+\"train/\"+name+\".npy\") for name in x_filenames], dtype=np.uint8)\n\ny_masks = np.asarray([np.max(np.load(data_dir+\"target/\"+name+'.npy'),axis=0) for name in x_filenames], dtype=bool)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Sample Data - Full Resolution ","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport matplotlib\n\nplt.imshow(x_images[0])\nplt.show()\n\nplt.imshow(y_masks[0])\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train-Validation Split","metadata":{}},{"cell_type":"code","source":"#note we're just using this tool to do an easy train-val split, we're not actually doing a full k-fold validation. \nfrom sklearn.model_selection import StratifiedKFold as KFold\nkfold_dummy = np.ones((len(x_filenames),))\nkfold = KFold(10, shuffle=True) #90-10 validation split\nfor train_ix, val_ix in kfold.split(kfold_dummy,kfold_dummy):\n    continue\n    \n\n\nfilenames_train = np.array(x_filenames)[train_ix.astype(int)]\nfilenames_val = np.array(x_filenames)[val_ix.astype(int)]\n\nx_train = x_images[train_ix]\nx_val = x_images[val_ix]\n\n\ny_train_mask = y_masks[train_ix]\ny_val_mask = y_masks[val_ix]\n\n#note we are calling full resolution labels \"masks\", while the down-sampled versions are called \"target\"","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Define the Statistical Head\n\nHere we are going to define the Statistical Head. This will be applied to the labeled mask data directly to get the down-sampled targets. The layers are basically copy pasted into the main model definition, as well. ","metadata":{}},{"cell_type":"code","source":"import tensorflow as tf\nfrom tensorflow.keras.layers import Input\nfrom tensorflow.keras.models import Model\nfrom tensorflow.keras.layers import Activation, Lambda, Concatenate\nfrom tensorflow.keras.layers import Conv2D, Conv2DTranspose, GlobalMaxPooling2D, MaxPooling2D\nfrom tensorflow.keras.layers import BatchNormalization as BN\nfrom tensorflow.keras.losses import BinaryCrossentropy\nfrom tensorflow.keras.callbacks import Callback\n\nfrom tensorflow.keras.optimizers import SGD, RMSprop, Adam, Nadam\nimport tensorflow.keras.backend as K\n\ndef build_attention_head_model(ws_size):\n    window_size, stride = ws_size\n    inputs = Input(shape=(512, 512, 1))\n    x = Lambda(lambda x: K.cast_to_floatx(x))(inputs)\n    \n    y_max = MaxPooling2D(window_size, stride, padding=\"same\")(x)\n    \n    #to do min pooling, we just do min = -max(-x)\n    y_neg = Lambda(lambda x: -x)(x)\n    y_neg_max = MaxPooling2D(window_size,stride, padding=\"same\")(y_neg)\n    y_min_proper = Lambda(lambda x: -x)(y_neg_max)\n\n    y = Concatenate()([y_max, y_min_proper])  \n     \n    model = Model(inputs=inputs, outputs=y)\n    \n    opt = SGD(learning_rate=1e-3)#this doesn't matter\n    model.compile(loss='binary_crossentropy',optimizer=opt, metrics=['Accuracy'])#this doesn't matter\n    return model\n\n\n\ndownsample_params = (32, 16)#width, stride\nattention_head_model = build_attention_head_model(downsample_params)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Labeled Data Down Sampled\n\nWith the attention head model initialized, we can just pass the mask data to this model, and it'll give us the down sampled target data. We'll split this into training and validation sets, and print out an example of the target data","metadata":{}},{"cell_type":"code","source":"\ny_targets = attention_head_model.predict(y_masks)\n\ny_train_target = y_targets[train_ix]\ny_val_target = y_targets[val_ix]\n\n#lets fine one that actually has 2nd channel values...\nviz_ix = np.where(y_targets[:,:,:,1] == 1)[0][0]\n\nplt.imshow(y_masks[viz_ix])\nplt.show()\nplt.imshow(y_targets[viz_ix,:,:,0])\nplt.show()\nplt.imshow(y_targets[viz_ix,:,:,1])\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Defining Model Functions\n\nNow we've largely loaded data. We're going to build out the model functions, then split the data into a train and validation set, and then build some data generators before we train the model. ","metadata":{}},{"cell_type":"markdown","source":"## Custom Loss and Metrics\n","metadata":{}},{"cell_type":"markdown","source":"Before we define the model, we'll need to define that custom loss function I mentioned before. This is not a complicated adjustment, but why we do it is somewhat subtle. Our output is of shape (32, 32, 2), and has two channels. We apply a non-reducing Binary Cross Entropy function, giving us a (32, 32, 2) sized error array. We do not average the two channels, instead we apply a \"reduce_max\" function along the last axis. This means that for each Max-Min pair of losses, we only take the highest value loss. This reduces our error array to size (32, 32, 1). This remaining array is averaged into a scalar batch error. \n\n\nWhy we do this is a little complicated, but it is primarily to avoid excessive reinforcement epoch after epoch. In a region where the label is (1,1), both the maximum and minimum pixel value will be selected. Both the pixels will be trained to increase their score. When the minimum pixel value is increased, it is somewhat likely to not remain the minimum value after some time. However, the maximum pixel value is very likely to be chosen, reinforced, and chosen again; forever. There is a similar issue with regions where the label is (0,0), and the minimum valued pixel in that region. \n\nBy selecting the Larger Error Value in each region- we can avoid this issue. In a (1,1) labeled region, the error associated with the Minimum Pixel Value will always be larger, because both targets are 1, and the minimum value is further from the target than the maximum value is. Put more simply, in a region labeled (1, 1), the larger error is always the minimum activation. In a region labeled (0, 0), the Maximum pixel value is always the highest error. This effectively stops both of our worst-case reinforcement concerns. There is nothing really special that happens with the (1, 0) labeled regions, and there are no (0, 1) labeled regions. \n\nplease note, by convention the labels are listed as (Maximum, Minimum)","metadata":{}},{"cell_type":"code","source":"def binary_maxmin_crossentropy(y_true, y_pred):\n    y_true = tf.keras.backend.cast_to_floatx(y_true)\n    \n    #NOTE: K.binary_crossentropy (K = tf.keras.backend) is important to use, its shape preserving. \n    #using tf.keras.losses.BinaryCrossentropy is wrong, it will automatically average the last dim. \n    \n    maxmin_error = K.binary_crossentropy(y_true, y_pred, from_logits=True)#yield(batch, height, width, 2)\n    \n    pixel_error = tf.reduce_max(maxmin_error, axis=3, keepdims=False)#returns (batch, 32, 32)\n    batch_error = tf.reduce_mean(pixel_error, axis=(0,1,2),keepdims=False)#returns (batch,)\n    return batch_error\n\n\n#these are two Accuracy metrics which are used. \ndef bucket_logit_accuracy(y_true, y_pred):#each should be shape (batch, 32, 32, 2)\n    y_rund = tf.math.greater_equal(y_pred,tf.constant([0.0]))#y_pred is in logits\n    y_rund = tf.keras.backend.cast_to_floatx(y_rund)\n    y_true = tf.keras.backend.cast_to_floatx(y_true)\n\n    ch_accuracy = tf.math.equal(y_true,y_rund) #accuracy of Max and Min individually. Both must be true to be accurate. \n    pixel_accuracy = tf.keras.backend.cast_to_floatx(tf.reduce_all(ch_accuracy,axis=3)) #k.all applies to both min and max buckets. Per pixel, accuracy is either 100% or 0%\n    sample_accuracy = tf.math.reduce_mean(pixel_accuracy, axis=(0,1,2))\n    \n    #batch_accuracy = tf.reduce_mean(image_accuracy, axis=0)\n    \n    return sample_accuracy\n\n\ndef bucket_accuracy(y_true, y_pred):#each should be shape (batch, 32, 32, 2)\n    y_rund = tf.math.greater_equal(y_pred,tf.constant([0.5]))#y_pred is in logits\n    y_rund = tf.keras.backend.cast_to_floatx(y_rund)\n    y_true = tf.keras.backend.cast_to_floatx(y_true)\n\n    \n    pixel_accuracy = tf.keras.backend.cast_to_floatx(tf.reduce_all(tf.math.equal(y_true,y_rund),axis=3)) #k.all applies to both min and max buckets. There is no 50%; it is 0 or 100 per region. \n    sample_accuracy = tf.math.reduce_mean(pixel_accuracy, axis=(1,2))\n    \n    #batch_accuracy = tf.reduce_mean(image_accuracy, axis=0)\n    \n    return sample_accuracy","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model Definition \n\nNext, we'll define the largely-standard UNet. We have a slightly odd way of normalizing and activating the UNet blocks, but nothing too odd about the UNet part of the model, itself. Please note the way we organize Batch Norm, Pooling Operations, and Relu Activation. Also note, the same Statistical Head is used here, but we just re-write the pooling operations to keep things straight forward. Tensorflow is weird sometimes. \n","metadata":{}},{"cell_type":"code","source":"def build_model(ws_size, mode):\n    ws_size = (32, 16)\n    #policy = tf.keras.mixed_precision.Policy('mixed_float16')#fp 16 instead of fp32 for some savings idk\n    policy = \"mixed_bfloat16\"\n    inputs = Input(shape=(512, 512, 3))\n    \n    act1 = \"relu\"\n    ini1 = \"he_normal\"\n    \n    act2 = \"relu\"\n    ini2 = \"he_normal\"\n    \n    #layer sizes\n    l = [32, 32, 64, 64, 128]\n    \n\n\n\n    \n    x = Lambda(lambda x: x/255.0,dtype=policy)(inputs)\n    x = Lambda(lambda x: 2.0*(x-0.5),dtype=policy)(x)#i very much prefer this static solution instead of BN'ing the images directly\n    \n    #note, we do Conv > {branch}> Max Pooling > Relu ->\n    #                                > Relu ->\n    \n    #we do not want to apply BN to the output of relu. \n    #we would also like the Encoder and Decoder to share the same activations\n    #BN after Max Pooling seems OK, but idk its unstable sometimes. \n    #therefore, this\n    \n    #Encode\n    x_512 = Conv2D(l[0], 3, padding='same',kernel_initializer=ini1,activation=act1,dtype=policy)(x)\n    x_512 = Conv2D(l[0], 3, padding='same',kernel_initializer=ini1,activation='linear',dtype=policy)(x_512)\n\n    x_256 = MaxPooling2D(2, padding='same',dtype=policy)(x_512)\n    x_256 = Activation(act1, dtype=policy)(x_256)\n    x_256 = Conv2D(l[1], 3, padding='same',kernel_initializer=ini1,activation=act1,dtype=policy)(x_256)\n    x_256 = Conv2D(l[1], 3, padding='same',kernel_initializer=ini1,activation='linear',dtype=policy)(x_256)\n\n    x_128 = MaxPooling2D(2, padding='same',dtype=policy)(x_256)\n    x_128 = Activation(act1, dtype=policy)(x_128)\n    x_128 = Conv2D(l[2], 3, padding='same',kernel_initializer=ini1,activation=act1,dtype=policy)(x_128)\n    x_128 = Conv2D(l[2], 3, padding='same',kernel_initializer=ini1,activation='linear',dtype=policy)(x_128)\n\n    x_64 = MaxPooling2D(2, padding='same',dtype=policy)(x_128)\n    x_64 = Activation(act1, dtype=policy)(x_64)\n    x_64 = Conv2D(l[3], 3, padding='same',kernel_initializer=ini1,activation=act1,dtype=policy)(x_64)\n    x_64 = Conv2D(l[3], 3, padding='same',kernel_initializer=ini1,activation='linear',dtype=policy)(x_64)\n\n    x_32 = MaxPooling2D(2, padding='same',dtype=policy)(x_64)\n    x_32 = Activation(act1, dtype=policy)(x_32)\n    x_32 = Conv2D(l[4], 3, padding='same',kernel_initializer=ini1,activation=act1,dtype=policy)(x_32)\n    x_32 = Conv2D(l[4], 3, padding='same',kernel_initializer=ini1,activation=act1,dtype=policy)(x_32)\n    #Bottleneck\n\n\n    #Decode\n    ux_64 = Conv2DTranspose(l[3], 2, 2, padding=\"same\",kernel_initializer=ini2,activation='linear',dtype=policy)(x_32)\n    ux_64 = Concatenate()([ux_64,x_64])\n    ux_64 = Activation(act2, dtype=policy)(ux_64)\n    ux_64 = Conv2D(l[3], 3, padding='same',kernel_initializer=ini2,activation=act2,dtype=policy)(ux_64)\n    ux_64 = Conv2D(l[3], 3, padding='same',kernel_initializer=ini2,activation=act2,dtype=policy)(ux_64)\n\n\n    ux_128 = Conv2DTranspose(l[2], 2, 2, padding=\"same\",kernel_initializer=ini2,activation='linear',dtype=policy)(ux_64)\n    ux_128 = Concatenate()([ux_128,x_128])\n    ux_128 = Activation(act2, dtype=policy)(ux_128)\n    ux_128 = Conv2D(l[2], 3, padding='same',kernel_initializer=ini2,activation=act2,dtype=policy)(ux_128)\n    ux_128 = Conv2D(l[2], 3, padding='same',kernel_initializer=ini2,activation=act2,dtype=policy)(ux_128)\n\n\n    ux_256 = Conv2DTranspose(l[1], 2, 2, padding=\"same\",kernel_initializer=ini2,activation='linear',dtype=policy)(ux_128)\n    ux_256 = Concatenate()([ux_256,x_256])\n    ux_256 = Activation(act2, dtype=policy)(ux_256)\n    ux_256 = Conv2D(l[1], 3, padding='same',kernel_initializer=ini2,activation=act2,dtype=policy)(ux_256)\n    ux_256 = Conv2D(l[1], 3, padding='same',kernel_initializer=ini2,activation=act2,dtype=policy)(ux_256)\n\n    \n    ux_512 = Conv2DTranspose(l[0], 2, 2, padding=\"same\",kernel_initializer=ini2,activation='linear',dtype=policy)(ux_256)\n    ux_512 = Concatenate()([ux_512,x_512])\n    ux_512 = Activation(act2, dtype=policy)(ux_512)\n    ux_512 = Conv2D(l[0], 3, padding='same',kernel_initializer=ini2,activation=act2,dtype=policy)(ux_512)\n    ux_512 = Conv2D(l[0], 3, padding='same',kernel_initializer=ini2,activation=act2,dtype=policy)(ux_512)\n    bias_constant = tf.keras.initializers.constant(-1.0)\n\n    pixel_logits = Conv2D(1, 1, activation=\"linear\", bias_initializer=bias_constant)(ux_512)#This bias is initialized to -1. Most pixels are negative, and initialization assumes a balanced target. This should help reduce the \"shock\" of large early gradients. \n\n    pixel_map = Activation(\"sigmoid\", name='map_output')(pixel_logits)\n    #sigmoid output for Map Mode\n    \n    \n    #this is where the standard UNet typically ends. Below is the Stastical Head. This is used during training, so it take the Logit Outputs instead of the sigmoid activated ones\n    #i probably could have just used the earlier defined attention head but whatever \n    y_region_max = MaxPooling2D(ws_size[0], ws_size[1], padding=\"same\")(pixel_logits)\n    \n    y_region_min = Lambda(lambda x: -x)(pixel_logits)\n    y_region_min_a = MaxPooling2D(downsample_params[0],downsample_params[1],padding=\"same\")(y_region_min)\n    y_region_min_proper = Lambda(lambda x: -x)(y_region_min_a)\n\n    \n    downsample_map = Concatenate()([y_region_max, y_region_min_proper])\n    \n    \n    \n    #To \"remove\" the statisical head and view the full-resolution predictions, we really just build the model twice in \"train\" mode and \"map\" mode, and pass the trained model's weights to the map model. \n\n    if mode==\"train\":\n        model = Model(inputs=inputs, outputs=downsample_map)\n    if mode==\"map\":\n        model = Model(inputs=inputs, outputs=pixel_map)\n        \n        \n\n    opt = SGD(1e-2)#this is replaced later\n    if mode==\"train\":\n        model.compile(loss=binary_maxmin_crossentropy,optimizer=opt, metrics=[bucket_logit_accuracy])\n        \n\n    if mode==\"map\":\n        model.compile(loss=binary_maxmin_crossentropy,optimizer=opt, metrics=[bucket_accuracy])\n        \n        \n    return model","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"That's the model. Its just a UNet with an attention head added, sometimes. ","metadata":{}},{"cell_type":"code","source":"model = build_model(downsample_params, \"train\")\nmap_model = build_model(downsample_params, \"map\")\nprint(model.summary())\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data Generator\n\nNext we need to define the Data Generator to work on all of the files we loaded before. This should be a relatively slim data generator which just pulls items from active RAM. It includes shuffle and some very simple augmentations. ","metadata":{}},{"cell_type":"code","source":"#from tensorflow.keras.utils import Sequence\nfrom tensorflow.keras.utils import PyDataset #i guess this was a tf update at some point\n\nimport random\n\ndef augment_img(img, target):\n    #assumes img is in form [width,height,3]\n    #assumes target is in form [width,height,2]\n    \n    if random.getrandbits(1):\n        img = np.flip(img,axis=-2)\n        target = np.flip(target,axis=-2)\n        \n    if random.getrandbits(1):\n        img = np.flip(img,axis=-3)\n        target = np.flip(target,axis=-3)\n        \n    r = np.random.randint(0,3)\n    img = np.rot90(img, r, axes=(0,1))\n    target = np.rot90(target, r, axes=(0,1))\n    return img, target\n\n\ndef no_augment_img(img, target):\n    \n    return img, target\n\nclass Augment_Generator(PyDataset):\n    def __init__(self, x_filenames, x_samples, y_samples, batch_size, shuffle, augment, mode, **kwargs):\n        super().__init__(**kwargs)\n\n        self.x_samples = x_samples\n        self.filenames = x_filenames\n        self.y_samples = y_samples\n        \n        self.shuffle = shuffle\n        self.augment = augment\n        self.mode = mode\n    \n        self.batch_size = batch_size\n        self.n_samples=len(self.x_samples)\n        self.len = int(np.ceil(self.n_samples/self.batch_size))\n\n        self.shuffle_ix = list(np.arange(self.n_samples))\n\n        \n        if self.shuffle:\n            np.random.shuffle(self.shuffle_ix)\n            \n        if self.augment:\n            self.img_fn = augment_img\n        else:\n            self.img_fn = no_augment_img\n            \n        \n    def __len__(self):\n        return self.len\n    \n    def __getitem__(self, idx):\n        batch_ix = self.shuffle_ix[idx*self.batch_size:(idx+1)*self.batch_size]\n        \n        #x_data = [np.load(self.data_dir+\"train/\"+str(self.filenames[i])+\".npy\") for i in ix]#with out a very fast drive, this isnt' advisible. \n        #y_data = [np.max(np.load(self.data_dir+\"target/\"+str(self.filenames[i])+\".npy\"),axis=(0,1,2),keepdims=False) for i in ix]\n        x_batch = [self.x_samples[ix] for ix in batch_ix]\n        y_batch = [self.y_samples[ix] for ix in batch_ix]\n\n        \n        #x_data = [self.img_fn(img) for img in x_data]#augment needs to do both x and y\n        x_batch = np.asarray(x_batch)\n        y_batch = np.asarray(y_batch)\n        \n        #i know I'm bad at this\n        augmented = [self.img_fn(x_batch[i], y_batch[i]) for i in range(len(batch_ix))]\n        x_batch = np.asarray([augmented[i][0] for i in range(len(augmented))])\n        y_batch = np.asarray([augmented[i][1] for i in range(len(augmented))])\n        \n        assert x_batch.shape[1:] == (512, 512, 3)#don't ask\n        assert y_batch.shape[1:] == (32, 32, 2)\n        \n        return x_batch, y_batch\n        \n          \n    def on_epoch_end(self):\n        if self.shuffle:\n\n            np.random.shuffle(self.shuffle_ix)\n            ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"With the generator defined we'll initialize one for both the Training and Testing dataset. ","metadata":{}},{"cell_type":"code","source":"batch_size = 8\ntrain_generator = Augment_Generator(filenames_train, x_train, y_train_target, batch_size, True, True, 'Train')\nval_generator = Augment_Generator(filenames_val, x_val, y_val_target, batch_size, False, False, 'Validate')\n\ntrain_steps = train_generator.len\nval_steps = val_generator.len\n","metadata":{"scrolled":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Testing Generator Output","metadata":{}},{"cell_type":"code","source":"#Illustrating the generrator's output\nplt.imshow(val_generator.__getitem__(val_steps-1)[0][0])\nplt.show()\nplt.imshow(val_generator.__getitem__(val_steps-1)[1][0][:,:,0])\nplt.show()\nplt.imshow(val_generator.__getitem__(val_steps-1)[1][0][:,:,1])\nplt.show()","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"These are actually kind of an incomprehensible mess. They work, don't get me wrong. But I really should have some graph titles and whatnot. Oh well.\n\n\nThere are two main functions in here, and some callback wrappers. Basically, we have a function to plot our Loss Graphs over time. Also, we have a function which uses the Map Model to show us the full resolution predictions every few epochs, too. This kind of monitoring is something I think you should have for ANY kind of CV task.\n\n\nHere, we're going to grab 3 samples from the Training Set and 3 samples from the Validation Set. During training, we'll have the model predict on these 6 samples and display the results so we can monitor how the model is learning.","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nimport matplotlib\nimport numpy as np\nimport pylab as pl\n\n\n####================Callbacks and stuff for monitoring progress\nclass PlotLoss(Callback):\n    def __init__(self, freq):\n        self.freq = freq      \n        self.logs = []\n        \n        self.loss = []\n        self.accuracy = []\n        \n        self.bias = []\n        self.prelu = []\n\n        self.val_loss = []\n        self.val_accuracy = []\n        \n        \n    def on_epoch_end(self, epoch, logs={}):\n            #self.logs.append(logs)\n            \n            self.loss.append(logs.get('loss'))\n            self.val_loss.append(logs.get('val_loss'))\n            \n            self.accuracy.append(logs.get('bucket_logit_accuracy'))\n            self.val_accuracy.append(logs.get('val_bucket_logit_accuracy'))\n            \n            bias = self.model.get_weights()[-1][0]\n\n            if np.shape(bias) == np.shape(1):\n                self.bias.append(bias)\n\n            nbias = [n - self.bias[0] for n in self.bias]\n            if (epoch+1)%self.freq==0:\n\n                fig, ax = plt.subplots(1,3, figsize=(9,5))\n\n                ax[0].plot(self.loss)\n                ax[0].plot(self.val_loss)\n\n                ax[1].plot(self.accuracy)\n                ax[1].plot(self.val_accuracy)            \n\n                ax[2].plot(nbias)\n\n\n                plt.show()            \n            return\n   \n\n#takes in 6 images and activation maps, reutrns a 3x2 grid of images with box highlighting   \ndef plot_boxes(images, targets, maps, k, dims):\n    map_size=dims[0]\n    r_size = dims[1]\n    r_stride = dims[2]\n    cms = pl.cm.YlOrBr\n\n    fig, ax = plt.subplots(2, 3, figsize=(9,6))\n    for j , (img, target, mask) in enumerate(zip(images, targets, maps)):\n        mask = mask.flatten()\n        j_a = (j-j%3)//3\n        j_b = j%3\n        \n        img = np.asarray(img,dtype=np.float32)/255\n        img = img[:,:,0]#to grey, just grab R channel now 512 x 512\n        img = np.stack([img, img, img], axis=-1) #now 512x512x3\n        \n\n        \n        target = np.max(target,axis=0)[:,:,0]\n        target = np.asarray(target, dtype=np.float32)\n        \n        \n        \n        \n        img[:,:,0] = img[:,:,0]-target\n        img[:,:,1] = img[:,:,1]-target\n        img[:,:,2] = img[:,:,2]+target#we want pure blue outlines\n        \n        img = np.clip(img,0,1)\n        \n        \n        ax[j_a][j_b].imshow(img,cmap=plt.get_cmap('gray'))\n        normal = pl.Normalize(0., 1.)\n\n        for i in range(k):\n            a_max = np.argmax(mask)\n            x_region = a_max%map_size#note, numpy addresses work on arr[y, x, z], compared to image coordinates\n            y_region = (a_max-(a_max%map_size))//map_size\n            prob = mask[a_max]\n            if prob<0.5:break\n            x_pixel = r_stride*(x_region)\n            y_pixel = r_stride*(y_region)\n   \n            color = cms(normal(prob**2))\n            \n            #rectangle expects bottom left coordinates. We've generated Top Right\n            if i==0: \n                line_width = 3.0\n            else: \n                line_width = 1.5\n                \n            rectangle = matplotlib.patches.Rectangle((x_pixel, y_pixel), r_size, r_size, edgecolor=color,facecolor=\"none\",lw=line_width)\n            ax[j_a][j_b].add_patch(rectangle)\n            font = {'color':color,'size':'large'}\n            ax[j_a][j_b].text(x_pixel-3,y_pixel-6,s=\"{0:.3f}\".format(prob), fontdict=font)\n\n            \n            mask[a_max]= 0#to help get the next most maximum\n    plt.subplots_adjust(wspace=0, hspace=0)\n    plt.show()\n            \n            \n        \n#call back to display box-highlighted images using earlier Fn\nclass PlotBoxes(Callback):\n    def __init__(self, map_model, test_images,test_targets, dims,freq):\n        self.map_model = map_model\n        self.test_images = test_images\n        self.test_targets = test_targets\n        self.freq = freq\n        self.dims=dims\n    def on_epoch_end(self, epoch, logs={}):\n        if (epoch+1)%self.freq==0:    \n            self.map_model.set_weights(self.model.get_weights())\n            test_maps = self.map_model.predict(self.test_images)\n            plot_boxes(self.test_images, self.test_targets, test_maps, 5, self.dims)\n        return\n            \ndef PlotImages(img_stack):\n    for l in range(np.shape(img_stack)[1]):\n        fig, ax = plt.subplots(1,4, figsize=(14,4))\n        for ch in range(4):\n            ax[ch].imshow(img_stack[0,l,:,:,ch], cmap='gray')\n        plt.show()\n\n\n\n\n\n\nclass PlotMaps(Callback):\n    def __init__(self, map_model, test_images, test_targets, freq):\n        self.map_model = map_model\n        self.test_images = test_images\n        self.test_targets = test_targets\n        self.freq = freq\n    def on_epoch_end(self, epoch, logs={}):\n        if (epoch+1)%self.freq==0:    \n            self.map_model.set_weights(self.model.get_weights())\n            pred_maps = self.map_model.predict(self.test_images)\n            plot_maps(self.test_images, self.test_targets, pred_maps)\n\n\n#takes in 6 images and activation maps, reutrns a 3x2 grid of images with box highlighting   \ndef plot_maps(images, targets, predictions):\n    mapA = predictions[0]\n    mapB = predictions[1]\n    \n    fig, ax = plt.subplots(3, 3, figsize=(18,24))\n    for j in range(np.shape(images)[0]):\n        img = images[j]\n        target = targets[j]\n        pred = predictions[j][:,:,0]\n\n        \n        img = np.asarray(img,dtype=np.float32)/255\n        img[:,:,1] = img[:,:,0]#to grey, just grab R channel now 512 x 512\n        img[:,:,2] = img[:,:,0]#to grey, just grab R channel now 512 x 512\n\n        \n        #First image has target data\n        img1 = np.zeros_like(img) #now 512x512x3\n        #target = np.max(target,axis=0)\n        target = np.asarray(target[:,:,0], dtype=np.float32)\n        \n        img1[:,:,0] = img[:,:,0]-target\n        img1[:,:,1] = img[:,:,1]-target\n        img1[:,:,2] = img[:,:,2]+target#we want pure blue objects\n        img1 = np.clip(img1,0,1)\n\n\n\n        #second image has pred data\n        img2 = np.zeros_like(img) #now 512x512x3\n        \n        pred_mask = np.asarray(pred>0.5,dtype=np.float32)\n        \n        img2[:,:,0] = img[:,:,0]+pred_mask\n        img2[:,:,1] = img[:,:,1]-pred_mask\n        img2[:,:,2] = img[:,:,2]-pred_mask\n        img2 = np.clip(img2,0,1)\n\n\n        \n        #Third Image only shows activations and agreement\n        img3 = np.ones_like(img)\n        \n        \n        img3[:,:,0] = (1-target)*pred_mask#Solid Red on FP\n        img3[:,:,1] = target*(pred_mask)#solid Green on TP\n        img3[:,:,2] = target*(1-pred_mask)#solid Blue on targets\n        \n        img3[:,:,0] += (1-target)*pred*(1-pred_mask)#Variable Red in Negative Space\n        img3[:,:,1] += target*pred*(1-pred_mask)#variable green on FN\n\n        img3 = np.clip(img3, 0, 1)\n        \n        ax[j][0].imshow(img1)\n        ax[j][1].imshow(img2)\n        ax[j][2].imshow(img3)\n\n        \n\n    plt.subplots_adjust(wspace=0, hspace=0)\n    plt.show()\n            ","metadata":{"_kg_hide-input":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Here, we're going to grab 3 samples from the Training Set and 3 samples from the Validaiton Set. During trainig, we'll have the model predict on these 6 samples and display the results so we can monitor how the model is learning. ","metadata":{}},{"cell_type":"code","source":"#we really don't need to do this from the filenames but whatever who cares\n\nviz_images_train = np.asarray([np.load(data_dir+\"train/\"+name+\".npy\") for name in filenames_train[:3]])\nviz_images_val = np.asarray([np.load(data_dir+\"train/\"+name+\".npy\") for name in filenames_val[:3]])\n\nviz_masks_train = np.asarray([np.max(np.load(data_dir+\"target/\"+name+\".npy\"),axis=0) for name in filenames_train[:3]])\nviz_masks_val = np.asarray([np.max(np.load(data_dir+\"target/\"+name+\".npy\"),axis=0) for name in filenames_val[:3]])\n\n\nplot_samples_train = PlotMaps(map_model, viz_images_train, viz_masks_train, freq=10)\nplot_samples_val = PlotMaps(map_model, viz_images_val, viz_masks_val, freq=10)\n\nplot_loss = PlotLoss(freq=10)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model Training\nWe are finally ready to train the model. During training, every 10 epochs, we'll display the UNet's output with out the attention head attached. We train for a while time, so the output is scrollable. If you're wondering, the 3rd graph after the loss and accuracy is the last layer's Bias unit's change (its initialized at 1, but this is subtracted out of each entry)","metadata":{}},{"cell_type":"code","source":"model.optimizer = RMSprop(1e-4)\n\nmodel.fit(train_generator, batch_size=batch_size,\n          #steps_per_epoch=train_steps,\n          validation_data=val_generator, #validation_steps=val_steps,\n          epochs=500, callbacks=[plot_loss, plot_samples_train, plot_samples_val])\n","metadata":{"scrolled":true,"_kg_hide-output":true,"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# An extremely brief review of predictions\nHere is the manual process for accessing the full-resolution predictions","metadata":{}},{"cell_type":"code","source":"map_model = build_model(downsample_params, \"map\")\nmap_model.set_weights(model.get_weights())\n\nviz_pred_train = map_model.predict(viz_images_train)\nviz_pred_val = map_model.predict(viz_images_val)\n\nprint(\"train set predictions\")\nplot_maps(viz_images_train, viz_masks_train, viz_pred_train)\nprint(\"validation set predictions\")\nplot_maps(viz_images_val, viz_masks_val, viz_pred_val)\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}