FCN train

Author: YunYang1994

5-Image_Segmentation/FCN/README.md

@@ -12,7 +12,7 @@ $ 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.
 ```bashrc
-VOC           # path:  /home/yang/test/VOC/
+VOC           # path:  /home/yang/dataset/VOC
 ├── test
 |    └──VOCdevkit
 |       └──VOC2007 (from VOCtest_06-Nov-2007.tar)
@@ -21,9 +21,10 @@ VOC           # path:  /home/yang/test/VOC/
              └──VOC2007 (from VOCtrainval_06-Nov-2007.tar)
              └──VOC2012 (from VOCtrainval_11-May-2012.tar)
 ```
-Then you need to make some transformation.
+Finally you need to make some transformation and train it.
 ```bashrc
-$ python parser_voc.py --data_path /home/yang/test/VOC
+$ python parser_voc.py --data_path /home/yang/dataset/VOC
+$ python train.py
 ```
 
 |![image](https://user-images.githubusercontent.com/30433053/66732790-d4d56680-ee8f-11e9-9120-07b0e8aa53d4.jpg)|![image](https://user-images.githubusercontent.com/30433053/66732791-d69f2a00-ee8f-11e9-9c5d-16cc84bc7e9e.jpg)|![image](https://user-images.githubusercontent.com/30433053/66732795-da32b100-ee8f-11e9-9d85-f0ddba7a3ab1.jpg)|

5-Image_Segmentation/FCN/config.py

@@ -11,6 +11,8 @@
 #
 #================================================================
 
+import numpy as np
+
 classes = ['background','aeroplane','bicycle','bird','boat',
            'bottle','bus','car','cat','chair','cow','diningtable',
            'dog','horse','motorbike','person','potted plant',
@@ -22,4 +24,5 @@
             [64,128,128],[192,128,128],[0,64,0],[128,64,0],
             [0,192,0],[128,192,0],[0,64,128]]
 
-image_size = [224, 224]
+rgb_mean = np.array([0.485, 0.456, 0.406])
+rgb_std = np.array([0.229, 0.224, 0.225])

5-Image_Segmentation/FCN/fcn8s.py

@@ -84,7 +84,6 @@ def call(self, x, training=False):
         h = self.conv4_3(h)
         h = self.pool4(h)
         pool4 = h # 1/16
-        print(pool4.shape)
 
         h = self.conv5_1(h)
         h = self.conv5_2(h)
@@ -121,6 +120,5 @@ def call(self, x, training=False):
         h = self.upscore8(h)
         h = h[:, 31:31+x.shape[1], 31:31+x.shape[2], :] # channel last
 
-        return h
-
+        return tf.nn.softmax(h, axis=-1)
 

5-Image_Segmentation/FCN/parser_voc.py

@@ -17,14 +17,14 @@
 from scipy import misc
 
 VOC_path = "/home/yang/dataset/VOC"
-train_labels = "./data/train_labels"
 
 if not os.path.exists("./data"): os.mkdir("./data")
 if not os.path.exists("./data/train_labels"): os.mkdir("./data/train_labels")
+if not os.path.exists("./data/test_labels"): os.mkdir("./data/test_labels")
 
-train_image_write = open(os.path.join(os.getcwd(), "data/train_image.txt"), "w")
 
 for mode in ["train", "test"]:
+    image_write = open(os.path.join(os.getcwd(), "data/%s_image.txt" %mode), "w")
     for year in [2007, 2012]:
         if mode == "test" and year == 2012: continue
         train_label_folder = os.path.join(VOC_path, "%s/VOCdevkit/VOC%d/SegmentationClass" %(mode, year))
@@ -35,10 +35,10 @@
             label_name = train_label_image[:-4]
             image_path = os.path.join(train_image_folder, label_name + ".jpg")
             if not os.path.exists(image_path): continue
-            train_image_write.writelines(image_path+"\n")
+            image_write.writelines(image_path+"\n")
             label_path = os.path.join(train_label_folder, train_label_image)
             label_image = np.array(misc.imread(label_path))
-            write_label = open("./data/train_labels/"+label_name+".txt", 'w')
+            write_label = open(("./data/%s_labels/" % mode)+label_name+".txt", 'w')
             print("=> processing %s" %label_path)
             H, W, C = label_image.shape
             for i in range(H):

5-Image_Segmentation/FCN/test.py

@@ -11,4 +11,10 @@
 #
 #================================================================
 
+import numpy as np
+from fcn8s import FCN8s
 
+data = np.arange(224*224*3).reshape([1, 224,224,3]).astype(np.float)
+
+model = FCN8s()
+y = model(data)

5-Image_Segmentation/FCN/train.py

@@ -6,7 +6,7 @@
 #   Editor      : VIM
 #   File name   : train.py
 #   Author      : YunYang1994
-#   Created date: 2019-10-12 17:44:30
+#   Created date: 2019-10-14 19:12:36
 #   Description :
 #
 #================================================================
@@ -15,17 +15,15 @@
 import cv2
 import random
 import tensorflow as tf
-from config import colormap, classes
 import numpy as np
-from PIL import Image
+from fcn8s import FCN8s
 from scipy import misc
-
+from config import colormap, classes, rgb_mean, rgb_std
 
 
 def create_image_label_path_generator(images_filepath, labels_filepath):
     image_paths = open(images_filepath).readlines()
     all_label_txts = os.listdir(labels_filepath)
-    print(all_label_txts)
     image_label_paths = []
     for label_txt in all_label_txts:
         label_name = label_txt[:-4]
@@ -35,61 +33,88 @@ def create_image_label_path_generator(images_filepath, labels_filepath):
             image_name = image_path.split("/")[-1][:-4]
             if label_name == image_name:
                 image_label_paths.append((image_path, label_path))
-    print(image_label_paths)
     while True:
         random.shuffle(image_label_paths)
         for i in range(len(image_label_paths)):
             yield image_label_paths[i]
 
-image_label_path_generator = create_image_label_path_generator(
-    "./data/train_image.txt", "./data/train_labels"
-)
 
 def process_image_label(image_path, label_path):
     # image = misc.imread(image_path)
     image = cv2.imread(image_path)
     image = cv2.resize(image, (224, 224), interpolation=cv2.INTER_NEAREST)
     image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-    image = np.array(image)
+    # data augmentation here
+    # pass
+    # image transformation here
+    image = (image / 255. - rgb_mean) / rgb_std
 
     label = open(label_path).readlines()
     label = [np.array(line.rstrip().split(" ")) for line in label]
     label = np.array(label, dtype=np.int)
     label = cv2.resize(label, (224, 224), interpolation=cv2.INTER_NEAREST)
+    label = label.astype(np.int)
 
     return image, label
 
 
-
-def TrainGenerator(batch_size):
+def DataGenerator(train_image_txt, train_labels_dir, batch_size):
     """
     generate image and mask at the same time
     """
-    pass
-
-
-for epoch in range(4):
-    image_path, label_path = next(image_label_path_generator)
-    # print(image_path, label_path)
-    image, label = process_image_label(image_path, label_path)
-
+    image_label_path_generator = create_image_label_path_generator(
+        train_image_txt, train_labels_dir
+    )
+    while True:
+        images = np.zeros(shape=[batch_size, 224, 224, 3])
+        labels = np.zeros(shape=[batch_size, 224, 224], dtype=np.float)
+        for i in range(batch_size):
+            image_path, label_path = next(image_label_path_generator)
+            image, label = process_image_label(image_path, label_path)
+            images[i], labels[i] = image, label
+        yield images, labels
+
+
+def visual_result(image, label):
+    image = (image * rgb_std + rgb_mean) * 255
+    image, label = image.astype(np.int), label.astype(np.int)
     H, W, C = image.shape
     new_label = np.zeros(shape=[H, W, C])
     cls = []
     for i in range(H):
         for j in range(W):
-            new_label[i, j] = np.array(colormap[label[i,j]])
-            cls.append(label[i, j])
+            cls_idx = label[i, j]
+            new_label[i, j] = np.array(colormap[cls_idx])
+            cls.append(cls_idx)
 
-    show_image = 0.7*new_label + 0.3*image
-    write_image = np.zeros(shape=[224, 448, 3])
-    write_image[:, :224, :] = image
-    write_image[:, 224:, :] = show_image
+    # show_image = 0.7*new_label + 0.3*image
+    show_image = np.zeros(shape=[224, 448, 3])
     cls = set(cls)
-    # for x in cls:
-        # print(classes[x])
-    # misc.imshow(show_image)
-    misc.imshow(write_image)
-    misc.imsave("%d.jpg"%epoch, write_image)
+    for x in cls:
+        print(classes[x])
+    show_image[:, :224, :] = image
+    show_image[:, 224:, :] = new_label
+    misc.imshow(show_image)
+
+TrainSet = DataGenerator("./data/train_image.txt", "./data/train_labels", 2)
+TestSet  = DataGenerator("./data/test_image.txt", "./data/test_labels", 2)
+
+model = FCN8s()
+callback = tf.keras.callbacks.ModelCheckpoint("model.h5", verbose=1, save_weights_only=True)
+model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4),
+              callback=callback,
+              loss='sparse_categorical_crossentropy',
+              metrics=['accuracy'])
+model.fit_generator(TrainSet, steps_per_epoch=6000, epochs=30)
+model.save_weights("model.h5")
+
+# data = np.arange(224*224*3).reshape([1,224,224,3]).astype(np.float)
+# model(data)
+# model.load_weights("model.h5")
+
+for x, y in TrainSet:
+    result = model(x)
+    pred_label = tf.argmax(result, axis=-1)
+    visual_result(x[0], pred_label[0].numpy())