title: Keras: Multi-label classification with U-net architcture for effective background/foreground image segmentation
body: <p>I have a series of cell images (+ masks) similar to 'D' in the following picture:</p>
  
  <p><a href="https://i.stack.imgur.com/43RM9.jpg" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/43RM9.jpg" alt="enter image description here"></a></p>
  
  <p>However, I'm trying to set up a two-label classification network using the Unet architecture where I explicitly label the following two components of an image:</p>
  
  <p>1) The Cells, and</p>
  
  <p>2) The Background (or the borders of cells)</p>
  
  <p>My reasoning for having a dedicated label for the background is as follows: my original images are 2000x3000 in size, and I want to chunk the images into the Unet as 512x512 blocks. However, when training the original Unet in this manner, I find that because a fair amount of the 512x512 blocks don't contain any cells (they're mostly Background, so the accompanying masks are empty), the loss for every epoch hovers at 0.4 and the network has a lot of trouble training. This is in spite of the presence of good 512x512 blocks that do contain good cell + mask layers.</p>
  
  <p>My thinking is that, to get the Unet system to work with 2 labels and segment the images properly, I need to provide two things:
  1) Another set of masks for the background itself so the network can positively identify the background.
  2) An output modifier for the Unet Keras code itself so it gives out 2 labels instead of one.</p>
  
  <p>However, please correct me if I am wrong on 1 assumption: the Unet architecture will work better to segment images if provided with a series of negative controls (background) as well as positive controls (the cells).</p>
  
  <p>Here's the Keras Model based off <a href="https://www.kaggle.com/keegil/keras-u-net-starter-lb-0-277" rel="nofollow noreferrer">the Kaggle starter tutorial model</a>:</p>
  
  <pre><code>inputs = Input((IMG_HEIGHT, IMG_WIDTH, IMG_CHANNELS))
      s = Lambda(lambda x: x / 255) (inputs)
  
      c1 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (s)
      c1 = Dropout(0.1) (c1)
      c1 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c1)
      p1 = MaxPooling2D((2, 2)) (c1)
  
      c2 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p1)
      c2 = Dropout(0.1) (c2)
      c2 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c2)
      p2 = MaxPooling2D((2, 2)) (c2)
  
      c3 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p2)
      c3 = Dropout(0.2) (c3)
      c3 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c3)
      p3 = MaxPooling2D((2, 2)) (c3)
  
      c4 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p3)
      c4 = Dropout(0.2) (c4)
      c4 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c4)
      p4 = MaxPooling2D(pool_size=(2, 2)) (c4)
  
      c5 = Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (p4)
      c5 = Dropout(0.3) (c5)
      c5 = Conv2D(256, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c5)
  
      u6 = Conv2DTranspose(128, (2, 2), strides=(2, 2), padding='same') (c5)
      u6 = concatenate([u6, c4])
      c6 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u6)
      c6 = Dropout(0.2) (c6)
      c6 = Conv2D(128, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c6)
  
      u7 = Conv2DTranspose(64, (2, 2), strides=(2, 2), padding='same') (c6)
      u7 = concatenate([u7, c3])
      c7 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u7)
      c7 = Dropout(0.2) (c7)
      c7 = Conv2D(64, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c7)
  
      u8 = Conv2DTranspose(32, (2, 2), strides=(2, 2), padding='same') (c7)
      u8 = concatenate([u8, c2])
      c8 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u8)
      c8 = Dropout(0.1) (c8)
      c8 = Conv2D(32, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c8)
  
      u9 = Conv2DTranspose(16, (2, 2), strides=(2, 2), padding='same') (c8)
      u9 = concatenate([u9, c1], axis=3)
      c9 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (u9)
      c9 = Dropout(0.1) (c9)
      c9 = Conv2D(16, (3, 3), activation='elu', kernel_initializer='he_normal', padding='same') (c9)
  
      outputs = Conv2D(1, (1, 1), activation='sigmoid') (c9)
  
      model = Model(inputs=[inputs], outputs=[outputs])
      model.compile(optimizer='adam', loss='binary_crossentropy', metrics=[mean_iou])
      model.summary()
  </code></pre>
  
  <p><strong>1) Getting the Unet to output 2 labels:</strong></p>
  
  <p>Should I change the output layer from </p>
  
  <pre><code>outputs = Conv2D(1, (1, 1), activation='sigmoid') (c9)
  </code></pre>
  
  <p>to:</p>
  
  <pre><code>outputs = Conv2D(2, (1, 1), activation='sigmoid') (c9)
  </code></pre>
  
  <p><strong>2) How do I structure the input array for the masks</strong></p>
  
  <p>For the mask labels, I have a 3D numpy array in the shape of: <code>(# images, height, width)</code> with <code>dtype=uint8</code>. Every mask is binarised to 0 or 255 for a single label (eg the cell). Do I create another mask array with the same shape and pass that into the network as another training layer, or do I label the background as a separate number in the original mask image? To simplify: how do I train 2 types of labels in the Unet architecture? </p>
  
  <p><strong>EDIT:</strong></p>
  
  <p>This is the code for my loss metric. It's a combination between dice coefficient and binary crossentropy:</p>
  
  <pre><code>def dice_coef(y_true, y_pred):
      smooth = 1.
      y_true_f = K.flatten(y_true)
      y_pred_f = K.flatten(y_pred)
      intersection = K.sum(y_true_f * y_pred_f)
      return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)
  
  def bce_dice_loss(y_true, y_pred):
      return 0.5 * binary_crossentropy(y_true, y_pred) - dice_coef(y_true, y_pred)
  </code></pre>