LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 verages
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

evaluate.py

@@ -0,0 +1,72 @@
+# -*- coding: utf-8 -*-
+# @Brief: 测试miou脚本
+import numpy as np
+import os
+import core.config as cfg
+from core.VOCdataset import VOCDataset
+from core.metrics import get_confusion_matrix_and_miou
+from nets.UNet import *
+from tqdm import tqdm
+import tensorflow as tf
+
+
+def evaluate(model, val_file_path, num_classes):
+    """
+    评价SegNet网络指标，主要是测试miou
+    :param model: 模型对象
+    :param val_file_path: 验证集文件路径
+    :param num_classes: 分类数量
+    :return: None
+    """
+    val_dataset = VOCDataset(val_file_path, batch_size=1)
+    val_dataset = val_dataset.tf_dataset()
+    val_dataset = iter(val_dataset)
+
+    f = open(val_file_path, mode='r')
+    images = f.readlines()
+    num_sample = len(images)
+
+    sum_confusion_matrix = np.zeros((num_classes, num_classes), dtype=np.int32)
+    process_bar = tqdm(range(num_sample), ncols=100, unit="step")
+
+    for i in process_bar:
+        image, y_true = next(val_dataset)
+
+        y_pred = model.predict(image)
+        y_pred = tf.nn.softmax(y_pred)
+
+        y_pred = np.argmax(y_pred, axis=-1).astype(np.uint8)
+        y_pred = np.squeeze(y_pred, axis=0)
+        y_true = np.squeeze(y_true, axis=0).astype(np.uint8)
+
+        confusion_matrix, miou = get_confusion_matrix_and_miou(y_true, y_pred, num_classes=21)
+        sum_confusion_matrix += confusion_matrix
+
+        process_bar.set_postfix(image_id=images[i].strip(), miou="{:.4f}".format(miou))
+
+    intersection = np.diag(sum_confusion_matrix)
+    union = np.sum(sum_confusion_matrix, axis=0) + np.sum(sum_confusion_matrix, axis=1) - intersection
+
+    iou = intersection / union
+    iou = np.nan_to_num(iou)    # 避免计算iou时出现nan
+
+    meanIOU = np.mean(iou)
+    object_meanIOU = np.mean(iou[1:])
+    print("-"*80)
+    print("Total MIOU: {:.4f}".format(meanIOU))
+    print("Object MIOU: {:.4f}".format(object_meanIOU))
+    print('pixel acc: {:.4f}'.format(np.sum(intersection)/np.sum(sum_confusion_matrix)))
+    print("IOU: ", iou)
+
+
+if __name__ == '__main__':
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+    gpus = tf.config.experimental.list_physical_devices("GPU")
+    if gpus:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+
+    model = UNet(cfg.input_shape, cfg.num_classes)
+    model.load_weights("./weights/unet_weights.h5")
+
+    evaluate(model, cfg.val_txt_path, cfg.num_classes)

predict.py

@@ -0,0 +1,54 @@
+# -*- coding: utf-8 -*-
+# @Brief: 预测脚本
+import core.config as cfg
+from core.dataset import resize_image_with_pad
+from core.metrics import *
+from nets.UNet import *
+import tensorflow as tf
+import os
+from PIL import Image
+import cv2 as cv
+import numpy as np
+
+
+def inference(model, image):
+    """
+    前向推理
+    :param model: 模型对象
+    :param image:  输入图像
+    :return:
+    """
+    image = cv.cvtColor(image, cv.COLOR_BGR2RGB)
+    image = resize_image_with_pad(image, target_size=cfg.input_shape[:2])
+    image = np.expand_dims(image, axis=0)
+    image /= 255.
+
+    pred_mask = model.predict(image)
+    pred_mask = tf.nn.softmax(pred_mask)
+
+    pred_mask = np.squeeze(pred_mask)
+    pred_mask = np.argmax(pred_mask, axis=-1).astype(np.uint8)
+
+    return pred_mask
+
+
+if __name__ == '__main__':
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+    gpus = tf.config.experimental.list_physical_devices("GPU")
+    if gpus:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+
+    model = UNet(cfg.input_shape, cfg.num_classes)
+    model.load_weights("./weights/unet_weights.h5")
+
+    image = cv.imread("D:/Code/Data/VOC2012/JPEGImages/2007_000033.jpg")
+    mask = Image.open("D:/Code/Data/VOC2012/SegmentationClass/2007_000033.png")
+    palette = mask.palette
+
+    result = inference(model, image)
+
+    result = Image.fromarray(result, mode='P')
+    result.palette = palette
+    result.show()

requirements.txt

@@ -0,0 +1,5 @@
+tensorflow-gpu==2.3.0
+opencv-python
+numpy
+tqdm
+pillow

train.py

@@ -0,0 +1,73 @@
+# -*- coding: utf-8 -*-
+# @Brief: 训练脚本
+from tensorflow.keras import optimizers, callbacks, utils, applications
+from core.VOCdataset import VOCDataset
+from nets.UNet import *
+from core.losses import *
+from core.metrics import *
+from core.callback import *
+import core.config as cfg
+from evaluate import evaluate
+import tensorflow as tf
+import os
+import cv2 as cv
+
+
+def train_by_fit(model, epochs, train_gen, test_gen, train_steps, test_steps):
+    """
+    fit方式训练
+    :param model: 训练模型
+    :param epochs: 训练轮数
+    :param train_gen: 训练集生成器
+    :param test_gen: 测试集生成器
+    :param train_steps: 训练次数
+    :param test_steps: 测试次数
+    :return: None
+    """
+
+    cbk = [
+        callbacks.ModelCheckpoint(
+            './weights/epoch={epoch:02d}_val_loss={val_loss:.04f}_miou={val_object_miou:.04f}.h5',
+            save_weights_only=True),
+    ]
+
+    learning_rate = CosineAnnealingLRScheduler(epochs, train_steps, 1e-4, 1e-6, warmth_rate=0.1)
+    optimizer = optimizers.Adam(learning_rate)
+    lr_info = print_lr(optimizer)
+
+    model.compile(optimizer=optimizer,
+                  loss=crossentropy_with_logits,
+                  metrics=[object_accuracy, object_miou, lr_info])
+
+    model.fit(train_gen,
+              steps_per_epoch=train_steps,
+              validation_data=test_gen,
+              validation_steps=test_steps,
+              epochs=epochs,
+              callbacks=cbk)
+
+
+if __name__ == '__main__':
+    os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+
+    gpus = tf.config.experimental.list_physical_devices("GPU")
+    if gpus:
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+
+    if not os.path.exists("weights"):
+        os.mkdir("weights")
+
+    model = UNet(cfg.input_shape, cfg.num_classes)
+    model.summary()
+
+    train_dataset = VOCDataset(cfg.train_txt_path, batch_size=cfg.batch_size, aug=True)
+    test_dataset = VOCDataset(cfg.val_txt_path, batch_size=cfg.batch_size)
+
+    train_steps = len(train_dataset) // cfg.batch_size
+    test_steps = len(test_dataset) // cfg.batch_size
+
+    train_gen = train_dataset.tf_dataset()
+    test_gen = test_dataset.tf_dataset()
+
+    train_by_fit(model, cfg.epochs, train_gen, test_gen, train_steps, test_steps)