unet

5-Image_Segmentation/Unet/README.md

@@ -1,28 +1,16 @@
-![image](https://user-images.githubusercontent.com/30433053/65319751-ee59fa00-dbd2-11e9-9065-10629dfd6d2f.png)
 
-# Train VOC dataset
-Download VOC PASCAL trainval and test data
-```bashrc
-$ wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar
-$ wget http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
-$ wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar
-```
-Extract all of these tars into one directory and rename them, which should have the following basic structure.
+# U-Net: Convolutional Networks for Biomedical Image Segmentation.
 
-```bashrc
-VOC           # path:  /home/yang/dataset/VOC
-├── test
-|    └──VOCdevkit
-|       └──VOC2007 (from VOCtest_06-Nov-2007.tar)
-└── train
-     └──VOCdevkit
-             └──VOC2007 (from VOCtrainval_06-Nov-2007.tar)
-                     └──VOC2012 (from VOCtrainval_11-May-2012.tar)
-```
+![image](https://lmb.informatik.uni-freiburg.de/Publications/2015/RFB15a/u-net-architecture.png)
 
-Then you need to parser VOC dataset and train it
 
 ```bashrc
-$ python parser_voc.py --voc_path /home/yang/dataset/VOC
+$ git clone https://github.com/YunYang1994/membrane.git
 $ python train.py
-```
+$ python test.py
+```
+
+| input | result |
+|---|:---:|
+|![image](./results/origin_0.png)|![image](./results/0.png)|
+|![image](./results/origin_1.png)|![image](./results/1.png)|

5-Image_Segmentation/Unet/Unet.py

@@ -12,57 +12,13 @@
 #================================================================
 
 from tensorflow.keras.models import Model
+from tensorflow.keras.optimizers import Adam
 from tensorflow.keras.layers import Conv2D, Input, MaxPooling2D, Dropout, concatenate, UpSampling2D
 
 
-# def Unet(num_class, image_size):
-
-    # inputs = Input(shape=[image_size, image_size, 3])
-    # conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(inputs)
-    # conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv1)
-    # pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
-    # conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool1)
-    # conv2 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv2)
-    # pool2 = MaxPooling2D(pool_size=(2, 2))(conv2)
-    # conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool2)
-    # conv3 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv3)
-    # pool3 = MaxPooling2D(pool_size=(2, 2))(conv3)
-    # conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool3)
-    # conv4 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv4)
-    # drop4 = Dropout(0.5)(conv4)
-    # pool4 = MaxPooling2D(pool_size=(2, 2))(drop4)
-
-    # conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(pool4)
-    # conv5 = Conv2D(1024, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv5)
-    # drop5 = Dropout(0.5)(conv5)
-
-    # up6 = Conv2D(512, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(drop5))
-    # merge6 = concatenate([drop4,up6], axis = 3)
-    # conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge6)
-    # conv6 = Conv2D(512, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv6)
-
-    # up7 = Conv2D(256, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv6))
-    # merge7 = concatenate([conv3,up7], axis = 3)
-    # conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge7)
-    # conv7 = Conv2D(256, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv7)
-
-    # up8 = Conv2D(128, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv7))
-    # merge8 = concatenate([conv2,up8], axis = 3)
-    # conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge8)
-    # conv8 = Conv2D(128, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv8)
-
-    # up9 = Conv2D(64, 2, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(UpSampling2D(size = (2,2))(conv8))
-    # merge9 = concatenate([conv1,up9], axis = 3)
-    # conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(merge9)
-    # conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
-    # conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same', kernel_initializer = 'he_normal')(conv9)
-    # conv10 = Conv2D(num_class, 1)(conv9)
-
-    # return Model(inputs = inputs, outputs = conv10)
-
 def Unet(num_class, image_size):
 
-    inputs = Input(shape=[image_size, image_size, 3])
+    inputs = Input(shape=[image_size, image_size, 1])
     conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same')(inputs)
     conv1 = Conv2D(64, 3, activation = 'relu', padding = 'same')(conv1)
     pool1 = MaxPooling2D(pool_size=(2, 2))(conv1)
@@ -101,6 +57,8 @@ def Unet(num_class, image_size):
     conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same')(merge9)
     conv9 = Conv2D(64, 3, activation = 'relu', padding = 'same')(conv9)
     conv9 = Conv2D(2, 3, activation = 'relu', padding = 'same')(conv9)
-    conv10 = Conv2D(num_class, 1)(conv9)
+    conv10 = Conv2D(num_class, 1, activation = 'sigmoid')(conv9)
+    model = Model(inputs = inputs, outputs = conv10)
+    model.compile(optimizer = Adam(lr = 1e-4), loss = 'binary_crossentropy', metrics = ['accuracy'])
 
-    return Model(inputs = inputs, outputs = conv10)
+    return model

5-Image_Segmentation/Unet/test.py

@@ -6,50 +6,73 @@
 #   Editor      : VIM
 #   File name   : test.py
 #   Author      : YunYang1994
-#   Created date: 2019-09-23 14:57:12
+#   Created date: 2019-10-05 12:55:29
 #   Description :
 #
 #================================================================
 
-import os
 import cv2
-import json
 import numpy as np
-import tensorflow as tf
 from Unet import Unet
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
 
+def testGenerator(batch_size):
+    """
+    generate image and mask at the same time
+    use the same seed for image_datagen and mask_datagen
+    to ensure the transformation for image and mask is the same
+    """
+    aug_dict = dict(horizontal_flip=False,fill_mode='nearest')
 
-batch_size = 2
-image_size = 512
-alpha = 0.3
-model = Unet(21, image_size)
-model.load_weights("Unet.h5")
-
-image_paths = open("./data/test_image.txt").readlines()
-label_paths = open("./data/test_label.txt").readlines()
-voc_colormap = json.load(open("./data/voc_colormap.json"))
-if not os.path.exists("./images"): os.mkdir("./images")
-
-class_name = sorted(voc_colormap.keys())
-colormap = [voc_colormap[cls] for cls in class_name]
-
-while len(image_paths):
-    output_mask = np.zeros([image_size, image_size, 3], np.uint8)
-    image_path = image_paths.pop().rstrip()
-    image_name = image_path.split("/")[-1]
-    image = cv2.imread(image_path)
-    image = cv2.resize(image, dsize=(image_size, image_size), interpolation=cv2.INTER_NEAREST)
-    batch_image = np.expand_dims(image / 255., 0)
-
-    pred_mask = model.predict(batch_image)[0]
-    pred_mask = tf.nn.softmax(pred_mask, axis=-1)
-    pred_mask = tf.argmax(pred_mask, axis=-1).numpy()
-    for i in range(image_size):
-        for j in range(image_size):
-            cls_idx = pred_mask[i][j]
-            color = colormap[cls_idx]
-            output_mask[i][j] = color
-    image = (1-alpha) * output_mask + alpha * image
-    cv2.imwrite("./images/"+image_name, image)
+    image_datagen = ImageDataGenerator(**aug_dict)
+    mask_datagen = ImageDataGenerator(**aug_dict)
+    image_generator = image_datagen.flow_from_directory(
+        "membrane/test",
+        classes=["images"],
+        color_mode = "grayscale",
+        target_size = (256, 256),
+        class_mode = None,
+        batch_size = batch_size, seed=1)
+
+    mask_generator = mask_datagen.flow_from_directory(
+        "membrane/test",
+        classes=["labels"],
+        color_mode = "grayscale",
+        target_size = (256, 256),
+        class_mode = None,
+        batch_size = batch_size, seed=1)
+
+    train_generator = zip(image_generator, mask_generator)
+    for (img,mask) in train_generator:
+        img = img / 255.
+        mask = mask / 255.
+        mask[mask > 0.5] = 1
+        mask[mask <= 0.5] = 0
+        yield (img, mask[0])
+
+testSet = testGenerator(batch_size=1)
+alpha   = 0.3
+model   = Unet(1, image_size=256)
+model.load_weights("model.hf5")
+
+for idx, (img, mask) in enumerate(testSet):
+    oring_img = img[0]
+    pred_mask = model.predict(img)[0]
+    pred_mask[pred_mask > 0.5] = 1
+    pred_mask[pred_mask <= 0.5] = 0
+    img = cv2.cvtColor(img[0], cv2.COLOR_GRAY2RGB)
+    H, W, C = img.shape
+    for i in range(H):
+        for j in range(W):
+            if pred_mask[i][j][0] <= 0.5:
+                img[i][j] = (1-alpha)*img[i][j]*255 + alpha*np.array([0, 0, 255])
+            else:
+                img[i][j] = img[i][j]*255
+    image_accuracy = np.mean(mask == pred_mask)
+    image_path = "./results/"+str(idx)+".png"
+    print("=> accuracy: %.4f, saving %s" %(image_accuracy, image_path))
+    cv2.imwrite(image_path, img)
+    cv2.imwrite("./results/origin_%d.png" %idx, oring_img*255)
+    if idx == 1: break
 
 

5-Image_Segmentation/Unet/train.py

@@ -11,66 +11,57 @@
 #
 #================================================================
 
-import os
-import cv2
-import json
-import shutil
-import numpy as np
-import tensorflow as tf
 from Unet import Unet
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
 
-image_size = 512
-epochs = 50
-lr = 0.0005
-batch_size = 2
-model = Unet(21, image_size)
-logdir = "./log"
-global_steps = 0
-optimizer = tf.keras.optimizers.Adam(lr)
+def adjustData(img,mask):
+    pass
 
-voc_colormap = json.load(open("./data/voc_colormap.json"))
+def trainGenerator(batch_size):
+    """
+    generate image and mask at the same time
+    use the same seed for image_datagen and mask_datagen
+    to ensure the transformation for image and mask is the same
+    """
+    aug_dict = dict(rotation_range=0.2,
+                        width_shift_range=0.05,
+                        height_shift_range=0.05,
+                        shear_range=0.05,
+                        zoom_range=0.05,
+                        horizontal_flip=True,
+                        fill_mode='nearest')
+    aug_dict = dict(horizontal_flip=True,
+                        fill_mode='nearest')
 
-if os.path.exists(logdir): shutil.rmtree(logdir)
-writer = tf.summary.create_file_writer(logdir)
+    image_datagen = ImageDataGenerator(**aug_dict)
+    mask_datagen = ImageDataGenerator(**aug_dict)
+    image_generator = image_datagen.flow_from_directory(
+        "membrane/train",
+        classes=["images"],
+        color_mode = "grayscale",
+        target_size = (256, 256),
+        class_mode = None,
+        batch_size = batch_size, seed=1)
 
-class_name = sorted(voc_colormap.keys())
-colormap = [voc_colormap[cls] for cls in class_name]
+    mask_generator = mask_datagen.flow_from_directory(
+        "membrane/train",
+        classes=["labels"],
+        color_mode = "grayscale",
+        target_size = (256, 256),
+        class_mode = None,
+        batch_size = batch_size, seed=1)
 
-for epoch in range(epochs):
-    image_paths = open("./data/train_image.txt").readlines()
-    label_paths = open("./data/train_label.txt").readlines()
-    batch_image = np.zeros(shape=[batch_size, image_size, image_size, 3], dtype=np.float32)
-    batch_label = np.zeros(shape=[batch_size, image_size, image_size, 21], dtype=np.float32)
-    while len(image_paths):
-        global_steps += 1
-        for i in range(batch_size):
-            image = cv2.imread(image_paths.pop().rstrip())
-            image = cv2.resize(image, dsize=(image_size, image_size), interpolation=cv2.INTER_NEAREST)
-            batch_image[i] = image / 255.
-
-            label_image = cv2.imread(label_paths.pop().rstrip())
-            label_image = cv2.resize(label_image, dsize=(image_size, image_size), interpolation=cv2.INTER_NEAREST)
-
-            H,W,C = label_image.shape
-            for x in range(H):
-                for y in range(W):
-                    pixel_color = image[x][y].tolist()
-                    if pixel_color in colormap:
-                        cls_idx = colormap.index(pixel_color)
-                        batch_label[i][x][y][cls_idx] = 1.
-
-        with tf.GradientTape() as tape:
-            pred_result = model(batch_image, training=True)
-            loss = tf.nn.softmax_cross_entropy_with_logits(logits=pred_result, labels=batch_label)
-            loss = tf.reduce_mean(tf.reduce_sum(loss, axis=[1,2]))
-            gradients = tape.gradient(loss, model.trainable_variables)
-            optimizer.apply_gradients(zip(gradients, model.trainable_variables))
-            # writing summary data
-        with writer.as_default():
-            tf.summary.scalar("loss", loss, step=global_steps)
-            print("=> Epoch: %2d, global_steps: %5d loss: %.6f" %(epoch+1, global_steps, loss.numpy()))
-        writer.flush()
-    model.save_weights("Unet.h5")
+    train_generator = zip(image_generator, mask_generator)
+    for (img,mask) in train_generator:
+        img = img / 255.
+        mask = mask / 255.
+        mask[mask > 0.5] = 1
+        mask[mask <= 0.5] = 0
+        yield (img,mask)
 
+trainset = trainGenerator(batch_size=2)
+model = Unet(1, image_size=256)
+model.fit_generator(trainset,steps_per_epoch=2000,epochs=5)
+model.save_weights("model.h5")