├── SegNet ├── models │ └── README.md ├── pretrained │ └── README.md ├── networks │ ├── README.md │ ├── vggnet.py │ └── segnet.py ├── dataset │ ├── README.md │ └── make_txt.py ├── README.md ├── predict.py └── train.py ├── SegNet-MobileNet-Tensorflow2.0 ├── img_out │ └── README ├── models │ └── README ├── pretrained │ └── README ├── img_test │ └── README.md ├── networks │ ├── README.md │ ├── segnet.py │ └── mobilenet.py ├── datasets │ ├── README │ └── make_txt.py ├── README.md ├── write_vedio.py ├── predict.py ├── video_predict.py └── train.py └── README.md /SegNet/models/README.md: -------------------------------------------------------------------------------- 1 | ## 存放生成的模型权重文件 2 | -------------------------------------------------------------------------------- /SegNet/pretrained/README.md: -------------------------------------------------------------------------------- 1 | ## 存放预训练模型 2 | -------------------------------------------------------------------------------- /SegNet/networks/README.md: -------------------------------------------------------------------------------- 1 | ## 用于存放搭建的网络结构文件 2 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/img_out/README: -------------------------------------------------------------------------------- 1 | ## 存放测试的结果 2 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/models/README: -------------------------------------------------------------------------------- 1 | 存放训练生成的模型文件 2 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/pretrained/README: -------------------------------------------------------------------------------- 1 | 存放预训练模型 2 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/img_test/README.md: -------------------------------------------------------------------------------- 1 | ## 本文件夹用来存放待测试的图片 2 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/networks/README.md: -------------------------------------------------------------------------------- 1 | 语义分割网络SegNet, 骨干网络：MobileNet 2 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/datasets/README: -------------------------------------------------------------------------------- 1 | ## 存放数据集 2 | --jpg 原图 3 | --png 标签 4 | 5 | --make_txt 使原图与标签对应 6 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/README.md: -------------------------------------------------------------------------------- 1 | ## 语义分割网络SegNet-MobileNet 2 | - TensorFlow2.0 3 | - Python3.7 4 | - Ubantu16.4 & Win10 5 | -------------------------------------------------------------------------------- /SegNet/dataset/README.md: -------------------------------------------------------------------------------- 1 | ## 数据集 2 | 本项目采用的是自制数据集，具体的制作步骤可参见：[labelme安装以及使用教程——自制语义分割数据集（保姆级示范）](https://blog.csdn.net/wjinjie/article/details/106735141) 3 | -------------------------------------------------------------------------------- /SegNet/dataset/make_txt.py: -------------------------------------------------------------------------------- 1 | # coding:utf-8 2 | import os 3 | 4 | imgs_path = '/home/fmc/WX/Segmentation/SegNet-tf2/dataset/jpg' # 图片文件存放地址 5 | for files in os.listdir(imgs_path): 6 | print(files) 7 | image_name = files + ';' + files[:-4] + '.png' 8 | 9 | with open("train.txt", "a") as f: 10 | f.write(str(image_name) + '\n') 11 | f.close() -------------------------------------------------------------------------------- /SegNet/README.md: -------------------------------------------------------------------------------- 1 | ## 中文博客 2 | [深度学习语义分割实战（一）：SegNet](https://blog.csdn.net/wjinjie/article/details/107338470) 3 | 4 | ## 项目环境 5 | Anaconda3 6 | python：3.6&3.7 7 | TensorFlow：2.0.0 8 | Pycharm：2018&2019&2020 9 | Cuda: 10.0 10 | CuDNN：7.6.0&7.6.4 11 | 12 | ## 数据集 13 | 自制语义分割数据集：[labelme安装以及使用教程——自制语义分割数据集（保姆级示范）](https://blog.csdn.net/wjinjie/article/details/106735141) 14 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/datasets/make_txt.py: -------------------------------------------------------------------------------- 1 | # coding:utf-8 2 | import os 3 | 4 | imgs_path = '/home/fmc/WX/Segmentation/SegNet-Mobile-tf2/dataset/jpg' # 图片文件存放地址 5 | for files in os.listdir(imgs_path): 6 | print(files) 7 | image_name = files + ';' + files[:-4] + '.png' 8 | 9 | with open("train.txt", "a") as f: 10 | f.write(str(image_name) + '\n') 11 | f.close() -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/write_vedio.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | from networks.segnet import SegNet_Mobile 4 | from PIL import Image 5 | import numpy as np 6 | import random 7 | import copy 8 | import os 9 | import time 10 | 11 | cap = cv2.VideoCapture('/home/fmc/WX/Segmentation/SegNet-Mobile-tf2/1.mp4') 12 | 13 | fourcc = cv2.VideoWriter_fourcc(*'XVID') 14 | out = cv2.VideoWriter('output.mp4', fourcc, 20.0, (960, 544)) 15 | 16 | while(cap.isOpened()): 17 | ret, img = cap.read() 18 | out.write(img) 19 | cv2.imshow('image', img) 20 | k = cv2.waitKey(30) 21 | # q键退出 22 | if (k & 0xff == ord('q')): 23 | break 24 | 25 | cap.release() 26 | out.relaase() 27 | cv2.destroyAllWindows() -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | ## TensorFLow2-Segmentation 2 | 本专栏我将搭建一些经典的语义分割网路，用于对场景中的目标进行分割，具体算法原理和实现过程可参见我的中文博客： 3 | - [TensorFlow2深度学习实战（十三）：语义分割网络 SegNet 源码解析](https://ai-wx.blog.csdn.net/article/details/107338470) 4 | - [TensorFlow2深度学习实战（十二）：语义分割网络 SegNet 论文解读](https://ai-wx.blog.csdn.net/article/details/106732783) 5 | 6 | ## Tasks 7 | - SegNet(已完成) 8 | - SegNet-MobileNet(已完成) 9 | 10 | ## Environment 11 | - Anaconda3 12 | - python：3.6/3.7 13 | - TensorFlow：2.0.0 14 | - Pycharm：2018/2019/2020 15 | - Cuda: 10.0 16 | - CuDNN：7.6.0/7.6.4 17 | 18 | ## Dataset Annotation 19 | - 自制语义分割数据集：[labelme安装以及使用教程——自制语义分割数据集（保姆级示范）](https://blog.csdn.net/wjinjie/article/details/106735141) 20 | 21 | ## Demo Effect 22 | ![在这里插入图片描述](https://img-blog.csdnimg.cn/20200714190057663.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3dqaW5qaWU=,size_16,color_FFFFFF,t_70#pic_center) 23 | -------------------------------------------------------------------------------- /SegNet/predict.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from networks.segnet import SegNet 3 | from PIL import Image 4 | import numpy as np 5 | import random 6 | import copy 7 | import os 8 | 9 | class_colors = [[0,0,0],[0,255,0]] 10 | NCLASSES = 2 11 | HEIGHT = 416 12 | WIDTH = 416 13 | 14 | model = SegNet(n_classes=NCLASSES, input_height=HEIGHT, input_width=WIDTH) 15 | model.load_weights("models/last1.h5") 16 | 17 | imgs = os.listdir("./img_test/") 18 | 19 | for jpg in imgs: 20 | 21 | img = Image.open("./img_test/"+jpg) 22 | old_img = copy.deepcopy(img) 23 | orininal_h = np.array(img).shape[0] 24 | orininal_w = np.array(img).shape[1] 25 | 26 | img = img.resize((WIDTH,HEIGHT)) 27 | img = np.array(img) 28 | img = img/255 29 | img = img.reshape(-1,HEIGHT,WIDTH,3) 30 | pr = model.predict(img)[0] 31 | 32 | pr = pr.reshape((int(HEIGHT/2), int(WIDTH/2), NCLASSES)).argmax(axis=-1) 33 | 34 | seg_img = np.zeros((int(HEIGHT/2), int(WIDTH/2), 3)) 35 | colors = class_colors 36 | 37 | for c in range(NCLASSES): 38 | # seg_img[:,:,0] += ((pr[:,: ] == c )*( colors[c][0] )).astype('uint8') 39 | seg_img[:,:,1] += ((pr[:,: ] == c )*( colors[c][1] )).astype('uint8') 40 | # seg_img[:,:,2] += ((pr[:,: ] == c )*( colors[c][2] )).astype('uint8') 41 | 42 | # Image.fromarray将数组转换成image格式 43 | seg_img = Image.fromarray(np.uint8(seg_img)).resize((orininal_w, orininal_h)) 44 | # 将两张图片合成一张图片 45 | image = Image.blend(old_img, seg_img, 0.3) 46 | image.save("./img_out/"+jpg) 47 | 48 | 49 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/predict.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from networks.segnet import SegNet_Mobile 3 | from PIL import Image 4 | import numpy as np 5 | import random 6 | import copy 7 | import os 8 | 9 | class_colors = [[0,0,0], [0,255,0]] 10 | NCLASSES = 2 11 | HEIGHT = 416 12 | WIDTH = 416 13 | 14 | model = SegNet_Mobile(n_classes=NCLASSES, input_height=HEIGHT, input_width=WIDTH) 15 | model.load_weights("models/last1.h5") 16 | 17 | imgs = os.listdir("./img_test/") 18 | 19 | for jpg in imgs: 20 | 21 | img = Image.open("./img_test/"+jpg) 22 | old_img = copy.deepcopy(img) 23 | orininal_h = np.array(img).shape[0] 24 | orininal_w = np.array(img).shape[1] 25 | 26 | img = img.resize((WIDTH,HEIGHT)) 27 | img = np.array(img) 28 | img = img/255 29 | img = img.reshape(-1,HEIGHT,WIDTH,3) 30 | pr = model.predict(img)[0] 31 | 32 | pr = pr.reshape((int(HEIGHT/2), int(WIDTH/2), NCLASSES)).argmax(axis=-1) 33 | 34 | seg_img = np.zeros((int(HEIGHT/2), int(WIDTH/2), 3)) 35 | colors = class_colors 36 | 37 | for c in range(NCLASSES): 38 | # seg_img[:,:,0] += ((pr[:,: ] == c )*( colors[c][0] )).astype('uint8') 39 | seg_img[:,:,1] += ((pr[:,: ] == c )*( colors[c][1] )).astype('uint8') 40 | # seg_img[:,:,2] += ((pr[:,: ] == c )*( colors[c][2] )).astype('uint8') 41 | 42 | # Image.fromarray将数组转换成image格式 43 | seg_img = Image.fromarray(np.uint8(seg_img)).resize((orininal_w, orininal_h)) 44 | # 将两张图片合成一张图片 45 | image = Image.blend(old_img, seg_img, 0.3) 46 | image.save("./img_out/"+jpg) 47 | 48 | 49 | -------------------------------------------------------------------------------- /SegNet/networks/vggnet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.keras import models, layers, activations 3 | 4 | def vggnet_encoder(input_height=416, input_width=416, pretrained='imagenet'): 5 | 6 | img_input = tf.keras.Input(shape=(input_height, input_width, 3)) 7 | 8 | # 416,416,3 -> 208,208,64 9 | x = layers.Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv1')(img_input) 10 | x = layers.Conv2D(64, (3, 3), activation='relu', padding='same', name='block1_conv2')(x) 11 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block1_pool')(x) 12 | f1 = x 13 | 14 | # 208,208,64 -> 128,128,128 15 | x = layers.Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv1')(x) 16 | x = layers.Conv2D(128, (3, 3), activation='relu', padding='same', name='block2_conv2')(x) 17 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block2_pool')(x) 18 | f2 = x 19 | 20 | # 104,104,128 -> 52,52,256 21 | x = layers.Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv1')(x) 22 | x = layers.Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv2')(x) 23 | x = layers.Conv2D(256, (3, 3), activation='relu', padding='same', name='block3_conv3')(x) 24 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block3_pool')(x) 25 | f3 = x 26 | 27 | # 52,52,256 -> 26,26,512 28 | x = layers.Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv1')(x) 29 | x = layers.Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv2')(x) 30 | x = layers.Conv2D(512, (3, 3), activation='relu', padding='same', name='block4_conv3')(x) 31 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block4_pool')(x) 32 | f4 = x 33 | 34 | # 26,26,512 -> 13,13,512 35 | x = layers.Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv1')(x) 36 | x = layers.Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv2')(x) 37 | x = layers.Conv2D(512, (3, 3), activation='relu', padding='same', name='block5_conv3')(x) 38 | x = layers.MaxPooling2D((2, 2), strides=(2, 2), name='block5_pool')(x) 39 | f5 = x 40 | 41 | return img_input, [f1, f2, f3, f4, f5] 42 | -------------------------------------------------------------------------------- /SegNet/networks/segnet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.keras import layers, models 3 | from networks.vggnet import vggnet_encoder 4 | 5 | # 对应的输入输出形状：batch_size, height, width, channels 6 | IMAGE_ORDERING = 'channels_last' 7 | 8 | # 解码器 9 | def decoder(feature_input, n_classes, n_upSample): 10 | # feature_input是vggnet第四个卷积块的输出特征矩阵 11 | # 26,26,512 12 | output = (layers.ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(feature_input) 13 | output = (layers.Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING))(output) 14 | output = (layers.BatchNormalization())(output) 15 | 16 | # 进行一次UpSampling2D，此时hw变为原来的1/8 17 | # 52,52,256 18 | output = (layers.UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(output) 19 | output = (layers.ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(output) 20 | output = (layers.Conv2D(256, (3, 3), padding='valid', data_format=IMAGE_ORDERING))(output) 21 | output = (layers.BatchNormalization())(output) 22 | 23 | # 进行一次UpSampling2D，此时hw变为原来的1/4 24 | # 104,104,128 25 | for _ in range(n_upSample - 2): 26 | output = (layers.UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(output) 27 | output = (layers.ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(output) 28 | output = (layers.Conv2D(128, (3, 3), padding='valid', data_format=IMAGE_ORDERING))(output) 29 | output = (layers.BatchNormalization())(output) 30 | 31 | # 进行一次UpSampling2D，此时hw变为原来的1/2 32 | # 208,208,64 33 | output = (layers.UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(output) 34 | output = (layers.ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(output) 35 | output = (layers.Conv2D(64, (3, 3), padding='valid', data_format=IMAGE_ORDERING))(output) 36 | output = (layers.BatchNormalization())(output) 37 | 38 | # 此时输出为h_input/2,w_input/2,nclasses 39 | # 208,208,2 40 | output = layers.Conv2D(n_classes, (3, 3), padding='same', data_format=IMAGE_ORDERING)(output) 41 | 42 | return output 43 | 44 | 45 | # 语义分割网络SegNet 46 | def SegNet(input_height=416, input_width=416, n_classes=2, n_upSample=3, encoder_level=3): 47 | 48 | img_input, features = vggnet_encoder(input_height=input_height, input_width=input_width) 49 | feature = features[encoder_level] # (26,26,512) 50 | output = decoder(feature, n_classes, n_upSample) 51 | 52 | # 将结果进行reshape 53 | output = tf.reshape(output, (-1, int(input_height / 2) * int(input_width / 2), 2)) 54 | output = layers.Softmax()(output) 55 | 56 | model = tf.keras.Model(img_input, output) 57 | 58 | return model 59 | 60 | 61 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/networks/segnet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.keras import layers, models 3 | from networks.mobilenet import mobilenet_encoder 4 | 5 | # 对应的输入输出形状：batch_size, height, width, channels 6 | IMAGE_ORDERING = 'channels_last' 7 | 8 | # 解码器 9 | def decoder(feature_input, n_classes, n_upSample): 10 | # feature_input是vggnet第四个卷积块的输出特征矩阵 11 | # 26,26,512 12 | output = (layers.ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(feature_input) 13 | output = (layers.Conv2D(512, (3, 3), padding='valid', data_format=IMAGE_ORDERING))(output) 14 | output = (layers.BatchNormalization())(output) 15 | 16 | # 进行一次UpSampling2D，此时hw变为原来的1/8 17 | # 52,52,256 18 | output = (layers.UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(output) 19 | output = (layers.ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(output) 20 | output = (layers.Conv2D(256, (3, 3), padding='valid', data_format=IMAGE_ORDERING))(output) 21 | output = (layers.BatchNormalization())(output) 22 | 23 | # 进行一次UpSampling2D，此时hw变为原来的1/4 24 | # 104,104,128 25 | for _ in range(n_upSample - 2): 26 | output = (layers.UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(output) 27 | output = (layers.ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(output) 28 | output = (layers.Conv2D(128, (3, 3), padding='valid', data_format=IMAGE_ORDERING))(output) 29 | output = (layers.BatchNormalization())(output) 30 | 31 | # 进行一次UpSampling2D，此时hw变为原来的1/2 32 | # 208,208,64 33 | output = (layers.UpSampling2D((2, 2), data_format=IMAGE_ORDERING))(output) 34 | output = (layers.ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING))(output) 35 | output = (layers.Conv2D(64, (3, 3), padding='valid', data_format=IMAGE_ORDERING))(output) 36 | output = (layers.BatchNormalization())(output) 37 | 38 | # 此时输出为h_input/2,w_input/2,nclasses 39 | # 208,208,2 40 | output = layers.Conv2D(n_classes, (3, 3), padding='same', data_format=IMAGE_ORDERING)(output) 41 | 42 | return output 43 | 44 | 45 | # 语义分割网络SegNet 46 | def SegNet_Mobile(input_height=416, input_width=416, n_classes=2, n_upSample=3, encoder_level=3): 47 | 48 | img_input, features = mobilenet_encoder(input_height=input_height, input_width=input_width) 49 | feature = features[encoder_level] # (26,26,512) 50 | output = decoder(feature, n_classes, n_upSample) 51 | 52 | # 将结果进行reshape 53 | output = tf.reshape(output, (-1, int(input_height / 2) * int(input_width / 2), 2)) 54 | output = layers.Softmax()(output) 55 | 56 | model = tf.keras.Model(img_input, output) 57 | 58 | return model 59 | 60 | 61 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/video_predict.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | from networks.segnet import SegNet_Mobile 4 | from PIL import Image 5 | import numpy as np 6 | import random 7 | import copy 8 | import os 9 | import time 10 | 11 | class_colors = [[0,0,0], [0,255,0]] 12 | NCLASSES = 2 13 | HEIGHT = 416 14 | WIDTH = 416 15 | 16 | model = SegNet_Mobile(n_classes=NCLASSES, input_height=HEIGHT, input_width=WIDTH) 17 | model.load_weights("models/last1.h5") 18 | 19 | cap = cv2.VideoCapture('/home/fmc/WX/Segmentation/SegNet-Mobile-tf2/1.mp4') 20 | fourcc = cv2.VideoWriter_fourcc(*'XVID') 21 | out = cv2.VideoWriter('output.mp4', fourcc, 30.0, (960, 544)) 22 | fps = 0.0 23 | 24 | while(cap.isOpened()): 25 | 26 | t1 = time.time() 27 | ret, img = cap.read() 28 | 29 | ############################# 30 | # 格式转变，BGRtoRGB 31 | img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) 32 | # 转变成Image 33 | img = Image.fromarray(np.uint8(img)) 34 | 35 | old_img = copy.deepcopy(img) 36 | orininal_h = np.array(img).shape[0] 37 | orininal_w = np.array(img).shape[1] 38 | 39 | img = img.resize((WIDTH, HEIGHT)) 40 | img = np.array(img) 41 | 42 | img = img / 255 43 | img = img.reshape(-1, HEIGHT, WIDTH, 3) 44 | pr = model.predict(img)[0] 45 | 46 | pr = pr.reshape((int(HEIGHT / 2), int(WIDTH / 2), NCLASSES)).argmax(axis=-1) 47 | 48 | seg_img = np.zeros((int(HEIGHT / 2), int(WIDTH / 2), 3)) 49 | colors = class_colors 50 | 51 | for c in range(NCLASSES): 52 | seg_img[:, :, 0] += ((pr[:, :] == c) * (colors[c][0])).astype('uint8') 53 | seg_img[:, :, 1] += ((pr[:, :] == c) * (colors[c][1])).astype('uint8') 54 | seg_img[:, :, 2] += ((pr[:, :] == c) * (colors[c][2])).astype('uint8') 55 | 56 | seg_img = Image.fromarray(np.uint8(seg_img)).resize((orininal_w, orininal_h)) 57 | 58 | img = Image.blend(old_img, seg_img, 0.3) 59 | 60 | img = np.array(img) 61 | fps = (fps + (1. / (time.time() - t1))) / 2 62 | print("fps= %.2f" % (fps)) 63 | img = cv2.putText(img, "fps= %.2f" % (fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2) 64 | ############################# 65 | out.write(img) 66 | cv2.imshow('image', img) 67 | k = cv2.waitKey(10) 68 | #q键退出 69 | if (k & 0xff == ord('q')): 70 | break 71 | 72 | cap.release() 73 | out.release() 74 | cv2.destroyAllWindows() 75 | 76 | """ 77 | old_img = copy.deepcopy(img) 78 | orininal_h = np.array(img).shape[0] 79 | orininal_w = np.array(img).shape[1] 80 | 81 | img = img.resize((WIDTH, HEIGHT)) 82 | img = np.array(img) 83 | print(img) 84 | img = img / 255 85 | img = img.reshape(-1, HEIGHT, WIDTH, 3) 86 | pr = model.predict(img)[0] 87 | 88 | pr = pr.reshape((int(HEIGHT / 2), int(WIDTH / 2), NCLASSES)).argmax(axis=-1) 89 | 90 | seg_img = np.zeros((int(HEIGHT / 2), int(WIDTH / 2), 3)) 91 | colors = class_colors 92 | 93 | for c in range(NCLASSES): 94 | seg_img[:, :, 0] += ((pr[:, :] == c) * (colors[c][0])).astype('uint8') 95 | seg_img[:, :, 1] += ((pr[:, :] == c) * (colors[c][1])).astype('uint8') 96 | seg_img[:, :, 2] += ((pr[:, :] == c) * (colors[c][2])).astype('uint8') 97 | 98 | seg_img = Image.fromarray(np.uint8(seg_img)).resize((orininal_w, orininal_h)) 99 | 100 | image = Image.blend(old_img, seg_img, 0.3) 101 | 102 | """ -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/networks/mobilenet.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.keras import models, layers, activations, Input 3 | 4 | # 对应的输入输出形状：batch_size, height, width, channels 5 | IMAGE_ORDERING = 'channels_last' 6 | 7 | def relu6(x): 8 | return activations.relu(x, max_value=6) 9 | 10 | def conv_block(inputs, filters, alpha, kernel=(3,3), strides=(1,1)): 11 | 12 | channel_axis = 1 if IMAGE_ORDERING == 'channels_first' else -1 13 | filters = int(filters * alpha) # juanjihe 14 | 15 | x = layers.ZeroPadding2D(padding=(1, 1), data_format=IMAGE_ORDERING, name='conv1_pad')(inputs) 16 | x = layers.Conv2D(filters, kernel, data_format=IMAGE_ORDERING, padding='valid', use_bias=False, strides=strides, name='conv1')(x) 17 | x = layers.BatchNormalization(axis=channel_axis, name='conv1_bn')(x) 18 | y = layers.Activation(relu6, name='conv1_relu')(x) 19 | 20 | return y 21 | 22 | def depthwise_conv_block(inputs, pointwise_conv_filters, alpha, depth_multiplier=1, strides=(1,1), block_id=1): 23 | 24 | channel_axis = 1 if IMAGE_ORDERING == 'channels_first' else -1 25 | pointwise_conv_filters = int(pointwise_conv_filters * alpha) 26 | 27 | x = layers.ZeroPadding2D((1, 1), data_format=IMAGE_ORDERING, name='conv_pad_%d' % block_id)(inputs) 28 | x = layers.DepthwiseConv2D((3, 3), data_format=IMAGE_ORDERING, padding='valid', depth_multiplier=depth_multiplier, strides=strides, use_bias=False, name='conv_dw_%d' % block_id)(x) 29 | x = layers.BatchNormalization(axis=channel_axis, name='conv_dw_%d_bn' % block_id)(x) 30 | x = layers.Activation(relu6, name='conv_dw_%d_relu' % block_id)(x) 31 | 32 | x = layers.Conv2D(pointwise_conv_filters, (1,1), data_format=IMAGE_ORDERING, padding='same', use_bias=False, strides=(1, 1), name='conv_pw_%d' % block_id)(x) 33 | x = layers.BatchNormalization(axis=channel_axis, name='conv_pw_%d_bn' % block_id)(x) 34 | y = layers.Activation(relu6, name='conv_pw_%d_relu' % block_id)(x) 35 | 36 | return y 37 | 38 | def mobilenet_encoder(input_height=224, input_width=224, pretrained='imagenet'): 39 | 40 | alpha = 1.0 41 | depth_multiplier = 1 42 | dropout = 1e-3 43 | 44 | # (416, 416, 3) 45 | img_input = Input(shape=(input_height, input_width, 3)) 46 | 47 | # (208, 208, 3) 48 | x = conv_block(img_input, 32, alpha, strides=(2, 2)) 49 | x = depthwise_conv_block(x, 64, alpha, depth_multiplier, block_id=1) 50 | f1 = x 51 | 52 | # (104, 104, 3) 53 | x = depthwise_conv_block(x, 128, alpha, depth_multiplier, strides=(2, 2), block_id=2) 54 | x = depthwise_conv_block(x, 128, alpha, depth_multiplier, block_id=3) 55 | f2 = x 56 | 57 | # (52, 52, 3) 58 | x = depthwise_conv_block(x, 256, alpha, depth_multiplier, strides=(2, 2), block_id=4) 59 | x = depthwise_conv_block(x, 256, alpha, depth_multiplier, block_id=5) 60 | f3 = x 61 | 62 | # (26, 26, 3) 63 | x = depthwise_conv_block(x, 512, alpha, depth_multiplier, strides=(2, 2), block_id=6) 64 | x = depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=7) 65 | x = depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=8) 66 | x = depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=9) 67 | x = depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=10) 68 | x = depthwise_conv_block(x, 512, alpha, depth_multiplier, block_id=11) 69 | f4 = x 70 | 71 | x = depthwise_conv_block(x, 1024, alpha, depth_multiplier, strides=(2, 2), block_id=12) 72 | x = depthwise_conv_block(x, 1024, alpha, depth_multiplier, block_id=13) 73 | f5 = x 74 | 75 | return img_input, [f1, f2, f3, f4, f5] 76 | -------------------------------------------------------------------------------- /SegNet/train.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.keras import optimizers, losses 3 | from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping 4 | from PIL import Image 5 | import numpy as np 6 | from networks.segnet import SegNet 7 | NCLASSES = 2 8 | HEIGHT = 416 9 | WIDTH = 416 10 | 11 | def generate_arrays_from_file(lines, batch_size=4): 12 | # 获取总长度 13 | n = len(lines) 14 | i = 0 15 | while 1: 16 | X_train = [] 17 | Y_train = [] 18 | # 获取一个batch_size大小的数据 19 | for _ in range(batch_size): 20 | if i == 0: 21 | np.random.shuffle(lines) # 打乱300行 22 | name = lines[i].split(';')[0] # 训练图片名称 23 | # 从文件中读取图像 24 | 25 | img = Image.open("/home/fmc/WX/Segmentation/SegNet-tf2/dataset/jpg" + '/' + name) 26 | img = img.resize((WIDTH, HEIGHT)) # 调整为416*416 27 | img = np.array(img) 28 | img = img/255 29 | X_train.append(img) # 把训练集图片存放在X_train中 30 | 31 | name = (lines[i].split(';')[1]).replace("\n", "") # 图片标签的名称 32 | # 从文件中读取标签图像 33 | img = Image.open("/home/fmc/WX/Segmentation/SegNet-tf2/dataset/png" + '/' + name) 34 | img = img.resize((int(WIDTH/2), int(HEIGHT/2))) # 调整大小为208*208 35 | img = np.array(img) 36 | seg_labels = np.zeros((int(HEIGHT/2), int(WIDTH/2), NCLASSES)) # 初始化一个全零的208*208*2的张量 37 | for c in range(NCLASSES): 38 | # seg_labels[: , : , c ] = (img[:,:,0] == c ).astype(int) 39 | seg_labels[:, :, c] = (img[:, :] == c).astype(int) # 热独码：208*208的img标签->208*208*2的seg_labels标签 40 | # (208,208,2)->(208*208,2) 41 | seg_labels = np.reshape(seg_labels, (-1, NCLASSES)) 42 | Y_train.append(seg_labels) 43 | # 读完一个周期后重新开始 44 | i = (i+1) % n 45 | yield (np.array(X_train), np.array(Y_train)) 46 | 47 | def loss_function(y_true, y_pred): 48 | loss = losses.categorical_crossentropy(y_true, y_pred) 49 | return loss 50 | 51 | if __name__ == "__main__": 52 | 53 | model = SegNet(n_classes=NCLASSES, input_height=HEIGHT, input_width=WIDTH) 54 | weights_path = "/home/fmc/WX/Segmentation/SegNet-tf2/pretrained/vgg16_weights_tf_dim_ordering_tf_kernels_notop.h5" 55 | model.load_weights(weights_path, by_name=True) 56 | 57 | with open("/home/fmc/WX/Segmentation/SegNet-tf2/dataset/train.txt", "r") as f: 58 | lines = f.readlines() 59 | 60 | # 打乱的数据更有利于训练 61 | np.random.seed(10101) 62 | np.random.shuffle(lines) 63 | np.random.seed(None) 64 | # 90%用于训练，10%用于估计。 65 | num_val = int(len(lines) * 0.1) 66 | num_train = len(lines) - num_val 67 | 68 | # 保存的方式，3代保存一次 69 | log_dir = "models/" 70 | checkpoint_period = ModelCheckpoint( 71 | log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5', 72 | monitor='val_loss', 73 | save_weights_only=True, 74 | save_best_only=True, 75 | period=6 76 | # save_freq=3 77 | ) 78 | # 学习率下降的方式，val_loss3次不下降就下降学习率继续训练 79 | reduce_lr = ReduceLROnPlateau( 80 | monitor='val_loss', 81 | factor=0.5, 82 | patience=3, 83 | verbose=2 84 | ) 85 | # 是否需要早停，当val_loss连续10个epochs不下降的时候意味着模型基本训练完毕，可以停止 86 | early_stopping = EarlyStopping( 87 | monitor='val_loss', 88 | min_delta=0, 89 | patience=10, 90 | verbose=2 91 | ) 92 | 93 | model.compile(loss=loss_function, # 交叉熵损失函数 94 | optimizer=optimizers.Adam(lr=1e-3), # 优化器 95 | metrics=['accuracy']) # 评价标准 96 | batch_size = 4 97 | print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size)) 98 | 99 | # 开始训练 100 | model.fit_generator(generate_arrays_from_file(lines[:num_train], batch_size), # 训练集 101 | steps_per_epoch=max(1, num_train // batch_size), # 每一个epos的steps数 102 | validation_data=generate_arrays_from_file(lines[num_train:], batch_size), # 验证集 103 | validation_steps=max(1, num_val // batch_size), 104 | epochs=50, 105 | initial_epoch=0, 106 | callbacks=[checkpoint_period, reduce_lr, early_stopping]) # 回调 107 | 108 | model.save_weights(log_dir + 'last1.h5') # 保存最后一次模型权重 109 | -------------------------------------------------------------------------------- /SegNet-MobileNet-Tensorflow2.0/train.py: -------------------------------------------------------------------------------- 1 | import tensorflow as tf 2 | from tensorflow.keras import optimizers, losses 3 | from tensorflow.keras.callbacks import TensorBoard, ModelCheckpoint, ReduceLROnPlateau, EarlyStopping 4 | from PIL import Image 5 | import numpy as np 6 | from networks.segnet import SegNet_Mobile 7 | NCLASSES = 2 8 | HEIGHT = 416 9 | WIDTH = 416 10 | 11 | def generate_arrays_from_file(lines, batch_size=4): 12 | # 获取总长度 13 | n = len(lines) 14 | i = 0 15 | while 1: 16 | X_train = [] 17 | Y_train = [] 18 | # 获取一个batch_size大小的数据 19 | for _ in range(batch_size): 20 | if i == 0: 21 | np.random.shuffle(lines) # 打乱300行 22 | name = lines[i].split(';')[0] # 训练图片名称 23 | # 从文件中读取图像 24 | 25 | img = Image.open("/home/fmc/WX/Segmentation/SegNet-Mobile-tf2/dataset/jpg" + '/' + name) 26 | img = img.resize((WIDTH, HEIGHT)) # 调整为416*416 27 | img = np.array(img) 28 | img = img/255 29 | X_train.append(img) # 把训练集图片存放在X_train中 30 | 31 | name = (lines[i].split(';')[1]).replace("\n", "") # 图片标签的名称 32 | # 从文件中读取标签图像 33 | img = Image.open("/home/fmc/WX/Segmentation/SegNet-Mobile-tf2/dataset/png" + '/' + name) 34 | img = img.resize((int(WIDTH/2), int(HEIGHT/2))) # 调整大小为208*208 35 | img = np.array(img) 36 | seg_labels = np.zeros((int(HEIGHT/2), int(WIDTH/2), NCLASSES)) # 初始化一个全零的208*208*2的张量 37 | for c in range(NCLASSES): 38 | # seg_labels[: , : , c ] = (img[:,:,0] == c ).astype(int) 39 | seg_labels[:, :, c] = (img[:, :] == c).astype(int) # 热独码：208*208的img标签->208*208*2的seg_labels标签 40 | # (208,208,2)->(208*208,2) 41 | seg_labels = np.reshape(seg_labels, (-1, NCLASSES)) 42 | Y_train.append(seg_labels) 43 | # 读完一个周期后重新开始 44 | i = (i+1) % n 45 | yield (np.array(X_train), np.array(Y_train)) 46 | 47 | def loss_function(y_true, y_pred): 48 | loss = losses.categorical_crossentropy(y_true, y_pred) 49 | return loss 50 | 51 | if __name__ == "__main__": 52 | 53 | model = SegNet_Mobile(n_classes=NCLASSES, input_height=HEIGHT, input_width=WIDTH) 54 | weights_path = "/home/fmc/WX/Segmentation/SegNet-Mobile-tf2/pretrained/mobilenet_1_0_224_tf_no_top.h5" 55 | model.load_weights(weights_path, by_name=True) 56 | with open("/home/fmc/WX/Segmentation/SegNet-Mobile-tf2/dataset/train.txt", "r") as f: 57 | lines = f.readlines() 58 | 59 | # 打乱的数据更有利于训练 60 | np.random.seed(10101) 61 | np.random.shuffle(lines) 62 | np.random.seed(None) 63 | # 90%用于训练，10%用于估计。 64 | num_val = int(len(lines) * 0.1) 65 | num_train = len(lines) - num_val 66 | 67 | # 保存的方式，3代保存一次 68 | log_dir = "models/" 69 | checkpoint_period = ModelCheckpoint( 70 | log_dir + 'ep{epoch:03d}-loss{loss:.3f}-val_loss{val_loss:.3f}.h5', 71 | monitor='val_loss', 72 | save_weights_only=True, 73 | save_best_only=True, 74 | period=3 75 | # save_freq=3 76 | ) 77 | # 学习率下降的方式，val_loss3次不下降就下降学习率继续训练 78 | reduce_lr = ReduceLROnPlateau( 79 | monitor='val_loss', 80 | factor=0.5, 81 | patience=3, 82 | verbose=2 83 | ) 84 | # 是否需要早停，当val_loss连续10个epochs不下降的时候意味着模型基本训练完毕，可以停止 85 | early_stopping = EarlyStopping( 86 | monitor='val_loss', 87 | min_delta=0, 88 | patience=10, 89 | verbose=2 90 | ) 91 | 92 | model.compile(loss=loss_function, # 交叉熵损失函数 93 | optimizer=optimizers.Adam(lr=1e-3), # 优化器 94 | metrics=['accuracy']) # 评价标准 95 | batch_size = 4 96 | print('Train on {} samples, val on {} samples, with batch size {}.'.format(num_train, num_val, batch_size)) 97 | 98 | # 开始训练 99 | model.fit_generator(generate_arrays_from_file(lines[:num_train], batch_size), # 训练集 100 | steps_per_epoch=max(1, num_train // batch_size), # 每一个epos的steps数 101 | validation_data=generate_arrays_from_file(lines[num_train:], batch_size), # 验证集 102 | validation_steps=max(1, num_val // batch_size), 103 | epochs=50, 104 | initial_epoch=0, 105 | callbacks=[checkpoint_period, reduce_lr, early_stopping]) # 回调 106 | 107 | model.save_weights(log_dir + 'last1.h5') # 保存最后一次模型权重 108 | --------------------------------------------------------------------------------