{"cells":[{"metadata":{},"cell_type":"markdown","source":"# Centernet \nIn this notebook, I am going to train a simple centernet based neural network to detect wheat head.\nThe network architecture and some parts are borrow from [this notebook](https://www.kaggle.com/kmat2019/centernet-keypoint-detector)\n\n**[Inference notebook here](https://www.kaggle.com/nvnnghia/keras-centernet-inference)**"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"\nimport cv2, os\nimport pandas as pd\nfrom PIL import Image, ImageDraw\nimport numpy as np \nfrom keras.utils import Sequence\nfrom imgaug import augmenters as iaa\nfrom sklearn.utils import shuffle\nfrom sklearn.model_selection import train_test_split\n\nfrom PIL import Image\nfrom keras.layers import Dense,Dropout, Conv2D,Conv2DTranspose, BatchNormalization, Activation,AveragePooling2D,GlobalAveragePooling2D, Input, Concatenate, MaxPool2D, Add, UpSampling2D, LeakyReLU,ZeroPadding2D\nfrom keras.models import Model\nfrom keras.objectives import mean_squared_error\nfrom keras import backend as K\nfrom keras.losses import binary_crossentropy\nfrom keras.callbacks import ModelCheckpoint, EarlyStopping, TensorBoard, ReduceLROnPlateau,LearningRateScheduler\nfrom keras.optimizers import Adam, RMSprop, SGD\nfrom matplotlib import pyplot as plt\nimport re","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Data preprocessing \nthis part is borrowed from this notebook [this notebook](https://www.kaggle.com/pestipeti/pytorch-starter-fasterrcnn-train)"},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"DIR_INPUT = '/kaggle/input/global-wheat-detection'\nDIR_TRAIN = f'{DIR_INPUT}/train'\nDIR_TEST = f'{DIR_INPUT}/test'","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df = pd.read_csv(f'{DIR_INPUT}/train.csv')\ntrain_df.shape","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train_df['x'] = -1\ntrain_df['y'] = -1\ntrain_df['w'] = -1\ntrain_df['h'] = -1\n\ndef expand_bbox(x):\n    r = np.array(re.findall(\"([0-9]+[.]?[0-9]*)\", x))\n    if len(r) == 0:\n        r = [-1, -1, -1, -1]\n    return r\n\ntrain_df[['x', 'y', 'w', 'h']] = np.stack(train_df['bbox'].apply(lambda x: expand_bbox(x)))\ntrain_df.drop(columns=['bbox'], inplace=True)\ntrain_df['x'] = train_df['x'].astype(np.float)\ntrain_df['y'] = train_df['y'].astype(np.float)\ntrain_df['w'] = train_df['w'].astype(np.float)\ntrain_df['h'] = train_df['h'].astype(np.float)\n\n\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"image_ids = train_df['image_id'].unique()\nvalid_ids = image_ids[-665:]\ntrain_ids = image_ids[:-665]\nvalid_df = train_df[train_df['image_id'].isin(valid_ids)]\ntrain_df = train_df[train_df['image_id'].isin(train_ids)]","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"###\ncategory_n=1\noutput_layer_n=category_n+4\n\n\n####FOR AUGMENTATION#####\n####FOR AUGMENTATION#####\nsometimes = lambda aug: iaa.Sometimes(0.5, aug)\nseq = iaa.Sequential([\n    sometimes(\n        iaa.OneOf([\n            iaa.Add((-10, 10), per_channel=0.5),\n            iaa.Multiply((0.9, 1.1), per_channel=0.5),\n            iaa.ContrastNormalization((0.9, 1.1), per_channel=0.5)\n        ])\n    ),\n    iaa.AdditiveGaussianNoise(scale=(0, 0.08*255)),\n    #sometimes(iaa.Rotate((-90, 90))),\n    sometimes(iaa.Fliplr(0.5)),\n    #sometimes(iaa.Crop(percent=(0, 0.2))),\n    sometimes(iaa.Flipud(0.5))\n    \n],random_order=False)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# data generator"},{"metadata":{"trusted":true},"cell_type":"code","source":"#####FOR DATASET######\nclass My_generator(Sequence):\n    def __init__(self, dataframe, image_dir, batch_size, input_size=512, is_train = False):\n        self.image_ids = dataframe['image_id'].unique()\n        self.df = dataframe\n        self.image_dir = image_dir\n        self.batch_size = batch_size\n        self.input_size = input_size\n        self.output_size = self.input_size//4\n        self.is_train = is_train\n        if self.is_train:\n            self.on_epoch_end()\n\n    def __len__(self):\n        return int(np.ceil(len(self.image_ids)/float(self.batch_size)))\n\n    def __getitem__(self, idx):\n        batch_x = self.image_ids[idx*self.batch_size:(idx+1)*self.batch_size]\n\n        if self.is_train:\n            return self.train_generator(batch_x)\n\n        return self.valid_generate(batch_x)\n\n    def on_epoch_end(self):\n        if(self.is_train):\n            self.image_ids= shuffle(self.image_ids)\n\n\n    def train_generator(self, batch_x):\n        batch_imgs = []\n        batch_segs = []\n        output_height,output_width=self.output_size,self.output_size\n        for image_id in batch_x:\n            records = self.df[self.df['image_id'] == image_id]\n\n            img = cv2.imread(f'{self.image_dir}/{image_id}.jpg')\n            im_h, im_w = img.shape[:2]\n            img = cv2.resize(img, (self.input_size, self.input_size))\n\n            boxes = records[['x', 'y', 'w', 'h']].values\n            #boxes[:, 2] = boxes[:, 0] + boxes[:, 2]\n            #boxes[:, 3] = boxes[:, 1] + boxes[:, 3]\n\n            # output_height,output_width=104,104\n\n            #PROCESS LABELS\n            output_layer=np.zeros((output_height,output_width,(output_layer_n+category_n)))\n\n            for box in boxes:\n                x, y, w, h = box\n                xc = x + 0.5*w\n                yc = y + 0.5*h\n                x_c, y_c, width, height = xc*output_height/im_w, yc*output_height/im_h, w*output_height/im_w, h*output_height/im_h\n                # print(x_c, y_c)\n\n                category=0 #not classify, just detect\n                heatmap=((np.exp(-(((np.arange(output_width)-x_c)/(width/10))**2)/2)).reshape(1,-1)\n                                    *(np.exp(-(((np.arange(output_height)-y_c)/(height/10))**2)/2)).reshape(-1,1))\n                output_layer[:,:,category]=np.maximum(output_layer[:,:,category],heatmap[:,:])\n\n                output_layer[int(y_c//1),int(x_c//1),category_n+category]=1\n                output_layer[int(y_c//1),int(x_c//1),2*category_n]=y_c%1#height offset\n                output_layer[int(y_c//1),int(x_c//1),2*category_n+1]=x_c%1\n                output_layer[int(y_c//1),int(x_c//1),2*category_n+2]=height/output_height\n                output_layer[int(y_c//1),int(x_c//1),2*category_n+3]=width/output_width\n\n\n            #images_aug, segmaps_aug = seq(images=[img], heatmaps=[output_layer.astype(np.float32)])\n\n            #batch_imgs.append(images_aug[0])\n            #batch_segs.append(segmaps_aug[0])\n            batch_imgs.append(img)\n            batch_segs.append(output_layer)\n\n\n        batch_imgs = np.array(batch_imgs, np.float32) /255\n        batch_segs = np.array(batch_segs, np.float32)\n\n        return batch_imgs, batch_segs\n\n\n    def valid_generate(self, batch_x, batch_y):\n        batch_imgs = []\n        batch_segs = []\n        output_height,output_width=self.output_size,self.output_size\n        for image_id in batch_x:\n            records = self.df[self.df['image_id'] == image_id]\n\n            img = cv2.imread(f'{self.image_dir}/{image_id}.jpg')\n            im_h, im_w = img.shape[:2]\n            img = cv2.resize(img, (self.input_size, self.input_size))\n\n            boxes = records[['x', 'y', 'w', 'h']].values\n\n            #PROCESS LABELS\n            output_layer=np.zeros((output_height,output_width,(output_layer_n+category_n)))\n\n            for box in boxes:\n                x, y, w, h = box\n                xc = x + 0.5*w\n                yc = y + 0.5*h\n                x_c, y_c, width, height = xc*output_height/im_w, yc*output_height/im_h, w*output_height/im_w, h*output_height/im_h\n                # print(x_c, y_c)\n\n                category=0 #not classify, just detect\n                heatmap=((np.exp(-(((np.arange(output_width)-x_c)/(width/10))**2)/2)).reshape(1,-1)\n                                    *(np.exp(-(((np.arange(output_height)-y_c)/(height/10))**2)/2)).reshape(-1,1))\n                output_layer[:,:,category]=np.maximum(output_layer[:,:,category],heatmap[:,:])\n\n                output_layer[int(y_c//1),int(x_c//1),category_n+category]=1\n                output_layer[int(y_c//1),int(x_c//1),2*category_n]=y_c%1#height offset\n                output_layer[int(y_c//1),int(x_c//1),2*category_n+1]=x_c%1\n                output_layer[int(y_c//1),int(x_c//1),2*category_n+2]=height/output_height\n                output_layer[int(y_c//1),int(x_c//1),2*category_n+3]=width/output_width\n\n            #images_aug, segmaps_aug = seq(images=[img], heatmaps=[output_layer.astype(np.float32)])\n\n\n            batch_imgs.append(img)\n            batch_segs.append(output_layer)\n\n\n        batch_imgs = np.array(batch_imgs, np.float32) /255\n        batch_segs = np.array(batch_segs, np.float32)\n\n        return batch_imgs, batch_segs","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Test generator"},{"metadata":{"trusted":true},"cell_type":"code","source":"def test_dataset():\n    \n    mygen = My_generator(train_df, DIR_TRAIN, batch_size=1, is_train=True, input_size = 416)\n\n    for count, (x,y) in enumerate(mygen):\n        # print(x.shape)\n        # print(y.shape)\n        x = x[0]\n        y= y[0]\n\n        points = np.argwhere(y[:,:,1] ==1)\n        for y1,x1 in points:\n            # print(x1,y1)\n            offsety = y[:,:,2][y1,x1]\n            offetx = y[:,:,3][y1,x1]\n            \n            h = y[:,:,4][y1,x1]*104\n            w = y[:,:,5][y1,x1]*104\n\n            x1, y1 = x1 + offetx, y1+offsety \n\n            xmin = int((x1-w/2)*4)\n            xmax = int((x1+w/2)*4)\n            ymin = int((y1-h/2)*4)\n            ymax = int((y1+h/2)*4)\n\n            cv2.rectangle(x, (xmin, ymin), (xmax, ymax), (0,255,255), 2)\n\n\n            cv2.circle(x, (int(x1*4),int(y1*4)), 5, (0,0,255), 2) \n\n        #cv2.imshow('djpg',y[:,:,1]*255)\n        #cv2.imshow('drawjpg',x)\n        fig, ax = plt.subplots(1, 1, figsize=(16, 8))\n\n        ax.set_axis_off()\n        ax.imshow(x)\n        break\n        \ntest_dataset()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# LOSSES FUNCTION#####"},{"metadata":{"trusted":true},"cell_type":"code","source":"\ndef all_loss(y_true, y_pred):\n    mask=K.sign(y_true[...,2*category_n+2])\n    N=K.sum(mask)\n    alpha=2.\n    beta=4.\n\n    heatmap_true_rate = K.flatten(y_true[...,:category_n])\n    heatmap_true = K.flatten(y_true[...,category_n:(2*category_n)])\n    heatmap_pred = K.flatten(y_pred[...,:category_n])\n    heatloss=-K.sum(heatmap_true*((1-heatmap_pred)**alpha)*K.log(heatmap_pred+1e-6)+(1-heatmap_true)*((1-heatmap_true_rate)**beta)*(heatmap_pred**alpha)*K.log(1-heatmap_pred+1e-6))\n    offsetloss=K.sum(K.abs(y_true[...,2*category_n]-y_pred[...,category_n]*mask)+K.abs(y_true[...,2*category_n+1]-y_pred[...,category_n+1]*mask))\n    sizeloss=K.sum(K.abs(y_true[...,2*category_n+2]-y_pred[...,category_n+2]*mask)+K.abs(y_true[...,2*category_n+3]-y_pred[...,category_n+3]*mask))\n    \n    all_loss=(heatloss+1.0*offsetloss+5.0*sizeloss)/N\n    return all_loss\n\ndef size_loss(y_true, y_pred):\n    mask=K.sign(y_true[...,2*category_n+2])\n    N=K.sum(mask)\n    sizeloss=K.sum(K.abs(y_true[...,2*category_n+2]-y_pred[...,category_n+2]*mask)+K.abs(y_true[...,2*category_n+3]-y_pred[...,category_n+3]*mask))\n    return (5*sizeloss)/N\n\ndef offset_loss(y_true, y_pred):\n    mask=K.sign(y_true[...,2*category_n+2])\n    N=K.sum(mask)\n    offsetloss=K.sum(K.abs(y_true[...,2*category_n]-y_pred[...,category_n]*mask)+K.abs(y_true[...,2*category_n+1]-y_pred[...,category_n+1]*mask))\n    return (offsetloss)/N\n  \ndef heatmap_loss(y_true, y_pred):\n    mask=K.sign(y_true[...,2*category_n+2])\n    N=K.sum(mask)\n    alpha=2.\n    beta=4.\n\n    heatmap_true_rate = K.flatten(y_true[...,:category_n])\n    heatmap_true = K.flatten(y_true[...,category_n:(2*category_n)])\n    heatmap_pred = K.flatten(y_pred[...,:category_n])\n    heatloss=-K.sum(heatmap_true*((1-heatmap_pred)**alpha)*K.log(heatmap_pred+1e-6)+(1-heatmap_true)*((1-heatmap_true_rate)**beta)*(heatmap_pred**alpha)*K.log(1-heatmap_pred+1e-6))\n    return heatloss/N","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Model definition"},{"metadata":{"trusted":true},"cell_type":"code","source":"##########MODEL#############\n\ndef aggregation_block(x_shallow, x_deep, deep_ch, out_ch):\n    x_deep= Conv2DTranspose(deep_ch, kernel_size=2, strides=2, padding='same', use_bias=False)(x_deep)\n    x_deep = BatchNormalization()(x_deep)   \n    x_deep = LeakyReLU(alpha=0.1)(x_deep)\n    x = Concatenate()([x_shallow, x_deep])\n    x=Conv2D(out_ch, kernel_size=1, strides=1, padding=\"same\")(x)\n    x = BatchNormalization()(x)   \n    x = LeakyReLU(alpha=0.1)(x)\n    return x\n  \n\n\ndef cbr(x, out_layer, kernel, stride):\n    x=Conv2D(out_layer, kernel_size=kernel, strides=stride, padding=\"same\")(x)\n    x = BatchNormalization()(x)\n    x = LeakyReLU(alpha=0.1)(x)\n    return x\n\ndef resblock(x_in,layer_n):\n    x=cbr(x_in,layer_n,3,1)\n    x=cbr(x,layer_n,3,1)\n    x=Add()([x,x_in])\n    return x  \n\n\n#I use the same network at CenterNet\ndef create_model(input_shape, aggregation=True):\n    input_layer = Input(input_shape)\n    \n    #resized input\n    input_layer_1=AveragePooling2D(2)(input_layer)\n    input_layer_2=AveragePooling2D(2)(input_layer_1)\n\n    #### ENCODER ####\n\n    x_0= cbr(input_layer, 16, 3, 2)#512->256\n    concat_1 = Concatenate()([x_0, input_layer_1])\n\n    x_1= cbr(concat_1, 32, 3, 2)#256->128\n    concat_2 = Concatenate()([x_1, input_layer_2])\n\n    x_2= cbr(concat_2, 64, 3, 2)#128->64\n    \n    x=cbr(x_2,64,3,1)\n    x=resblock(x,64)\n    x=resblock(x,64)\n    \n    x_3= cbr(x, 128, 3, 2)#64->32\n    x= cbr(x_3, 128, 3, 1)\n    x=resblock(x,128)\n    x=resblock(x,128)\n    x=resblock(x,128)\n    \n    x_4= cbr(x, 256, 3, 2)#32->16\n    x= cbr(x_4, 256, 3, 1)\n    x=resblock(x,256)\n    x=resblock(x,256)\n    x=resblock(x,256)\n    x=resblock(x,256)\n    x=resblock(x,256)\n \n    x_5= cbr(x, 512, 3, 2)#16->8\n    x= cbr(x_5, 512, 3, 1)\n    \n    x=resblock(x,512)\n    x=resblock(x,512)\n    x=resblock(x,512)\n    \n    #### DECODER ####\n    x_1= cbr(x_1, output_layer_n, 1, 1)\n    x_1 = aggregation_block(x_1, x_2, output_layer_n, output_layer_n)\n    x_2= cbr(x_2, output_layer_n, 1, 1)\n    x_2 = aggregation_block(x_2, x_3, output_layer_n, output_layer_n)\n    x_1 = aggregation_block(x_1, x_2, output_layer_n, output_layer_n)\n    x_3= cbr(x_3, output_layer_n, 1, 1)\n    x_3 = aggregation_block(x_3, x_4, output_layer_n, output_layer_n) \n    x_2 = aggregation_block(x_2, x_3, output_layer_n, output_layer_n)\n    x_1 = aggregation_block(x_1, x_2, output_layer_n, output_layer_n)\n\n    x_4= cbr(x_4, output_layer_n, 1, 1)\n\n    x=cbr(x, output_layer_n, 1, 1)\n    x= UpSampling2D(size=(2, 2))(x)#8->16 \n\n    x = Concatenate()([x, x_4])\n    x=cbr(x, output_layer_n, 3, 1)\n    x= UpSampling2D(size=(2, 2))(x)#16->32\n\n    x = Concatenate()([x, x_3])\n    x=cbr(x, output_layer_n, 3, 1)\n    x= UpSampling2D(size=(2, 2))(x)#32->64 \n\n    x = Concatenate()([x, x_2])\n    x=cbr(x, output_layer_n, 3, 1)\n    x= UpSampling2D(size=(2, 2))(x)#64->128 \n\n    x = Concatenate()([x, x_1])\n    x=Conv2D(output_layer_n, kernel_size=3, strides=1, padding=\"same\")(x)\n    out = Activation(\"sigmoid\")(x)\n    \n    model=Model(input_layer, out)\n    \n    return model","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# Training"},{"metadata":{"trusted":true},"cell_type":"code","source":"#####TRAIN##########\ndef lrs(epoch):\n    lr = 0.001\n    if epoch >= 20: lr = 0.0002\n    return lr\n    \ndef train(batch_size=32, input_size=512, n_epoch=30):\n    learning_rate=0.001\n\n    mygen = My_generator(train_df, DIR_TRAIN, batch_size=batch_size, is_train=True, input_size = input_size)\n    myval = My_generator(valid_df, DIR_TRAIN, batch_size=batch_size, is_train=True, input_size = input_size)\n\n\n    model=create_model(input_shape=(input_size,input_size,3))\n    \n    lr_schedule = LearningRateScheduler(lrs)\n\n\n    # EarlyStopping\n    early_stopping = EarlyStopping(monitor = 'val_loss', min_delta=0, patience = 60, verbose = 1)\n    # ModelCheckpoint\n    #weights_dir = './'\n\n    #if os.path.exists(weights_dir) == False:os.mkdir(weights_dir)\n    #model_checkpoint = ModelCheckpoint(weights_dir + \"{epoch:02d}-{val_loss:.3f}.hdf5\", monitor = 'val_loss', verbose = 1,\n    #                                      save_best_only = False, save_weights_only = True, period = 3)\n    # reduce learning rate\n    reduce_lr = ReduceLROnPlateau(monitor = 'val_loss', factor = 0.5, patience = 4, verbose = 1)\n    \n\n    print(model.summary())\n\n    model.compile(loss=all_loss, optimizer=Adam(lr=learning_rate), metrics=[heatmap_loss,size_loss,offset_loss])\n\n    hist = model.fit_generator(\n        mygen,\n        steps_per_epoch = len(train_df['image_id'].unique()) // batch_size,\n        epochs = n_epoch,\n        validation_data=myval,\n        validation_steps = len(valid_df['image_id'].unique()) // batch_size,\n        callbacks = [early_stopping, reduce_lr],\n        shuffle = True,\n        verbose = 1\n    )\n    model.save_weights('final_weights.h5')\n\n    #print(hist.history.keys())\n    #np.save('his.npy', hist.history)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"train(n_epoch=1)","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}