1 代码详解
得到模型权重后，对一张图片进行语义分割处理，其详细过程如下所示，结合代码注释即可。
import colorsys
import copy
import time
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
#-----------------------------------------------------------------------------------#
# DeepLab网络及代码
# 详细介绍可见 https://blog.csdn.net/weixin_45377629/article/details/124083978
#-----------------------------------------------------------------------------------#
from nets.deeplabv3_plus import DeepLab
#-----------------------------------------------------------------------------------#
# 三个函数代码下方给出
# cvtColor： 将图像转换成RGB图像，防止灰度图在预测时报错。
# preprocess_input： 归一化
# resize_image： 对输入图像进行resize，letterbox_image方式，不失真resize
#-----------------------------------------------------------------------------------#
from utils.utils import cvtColor, preprocess_input, resize_image
#-----------------------------------------------------------------------------------#
# 使用自己训练好的模型预测需要修改3个参数
# model_path、backbone和num_classes都需要修改！
# 如果出现shape不匹配，一定要注意训练时的model_path、backbone和num_classes的修改
#-----------------------------------------------------------------------------------#
class DeeplabV3(object):
_defaults = {
#-------------------------------------------------------------------#
# model_path指向logs文件夹下的权值文件
# 训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
# 验证集损失较低不代表miou较高，仅代表该权值在验证集上泛化性能较好。
# 链接：https://pan.baidu.com/s/1TrBlnZUd6xwxUvgFjbz7TQ 提取码：cj80
#-------------------------------------------------------------------#
"model_path" : 'model_data/deeplab_mobilenetv2.pth',
#----------------------------------------#
# 所需要区分的类的个数+1
#----------------------------------------#
"num_classes" : 21,
#----------------------------------------#
# 所使用的的主干网络：
# mobilenet
# xception
#----------------------------------------#
"backbone" : "mobilenet",
#----------------------------------------#
# 输入图片的大小
#----------------------------------------#
"input_shape" : [512, 512],
#----------------------------------------#
# 下采样的倍数，一般可选的为8和16
# 与训练时设置的一样即可
#----------------------------------------#
"downsample_factor" : 16,
#-------------------------------------------------#
# mix_type参数用于控制检测结果的可视化方式
#
# mix_type = 0的时候代表原图与生成的图进行混合
# mix_type = 1的时候代表仅保留生成的图
# mix_type = 2的时候代表仅扣去背景，仅保留原图中的目标
# 下方有给出三种可视化结果的区别
#-------------------------------------------------#
"mix_type" : 0,
#-------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
#-------------------------------#
"cuda" : False,
}
#---------------------------------------------------#
# 初始化Deeplab
#---------------------------------------------------#
def __init__(self, **kwargs):
#---------------------------------------------------#
# _defaults字典原来是这么用起来的
#---------------------------------------------------#
self.__dict__.update(self._defaults)
for name, value in kwargs.items():
#-----------------------------------------------#
# 设置属性 name 值，即self.name==value
#-----------------------------------------------#
setattr(self, name, value)
#---------------------------------------------------#
# 画框设置不同的颜色
#---------------------------------------------------#
if self.num_classes <= 21:
self.colors = [ (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128), (0, 128, 128),
(128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), (192, 128, 0), (64, 0, 128), (192, 0, 128),
(64, 128, 128), (192, 128, 128), (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128),
(128, 64, 12)]
else:
hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
#---------------------------------------------------#
# 获得模型
#---------------------------------------------------#
self.generate()
#---------------------------------------------------#
# 获得所有的分类
#---------------------------------------------------#
def generate(self):
#-------------------------------#
# 载入模型与权值
# 详细介绍可见 https://blog.csdn.net/weixin_45377629/article/details/124083978
#-------------------------------#
self.net = DeepLab(num_classes=self.num_classes, backbone=self.backbone, downsample_factor=self.downsample_factor, pretrained=False)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.net.load_state_dict(torch.load(self.model_path, map_location=device))
self.net = self.net.eval()
print('{} model, and classes loaded.'.format(self.model_path))
if self.cuda:
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
#---------------------------------------------------#
# 检测图片
#---------------------------------------------------#
def detect_image(self, image):
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
image = cvtColor(image)
#---------------------------------------------------#
# 对输入图像进行一个备份，后面用于绘图
#---------------------------------------------------#
old_img = copy.deepcopy(image)
orininal_h = np.array(image).shape[0]
orininal_w = np.array(image).shape[1]
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize，即letterbox_image
# 也可以直接resize进行识别，此处是不失真resize
#---------------------------------------------------------#
image_data, nw, nh = resize_image(image, (self.input_shape[1],self.input_shape[0]))
#---------------------------------------------------------#
# 添加上batch_size维度，缺少这个维度，网络没法预测，
# preprocess_input表示数据归一化 image /= 255.0
#---------------------------------------------------------#
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)
with torch.no_grad():
images = torch.from_numpy(image_data)
if self.cuda:
images = images.cuda()
#---------------------------------------------------#
# 图片传入网络进行预测
# VOC为例，self.net(images) shape：torch.size([1,21,512,512])
# pr :tensor, shape：torch.size([21,512,512])
#---------------------------------------------------#
pr = self.net(images)[0]
#---------------------------------------------------#
# 取出每一个像素点的种类
# pr.permute(1,2,0)：通道交换
# F.softmax(input, dim=-1)：在行上softmax
# pr :array, shape：(512,512,21)
#---------------------------------------------------#
pr = F.softmax(pr.permute(1,2,0), dim=-1).cpu().numpy()
#----------------------------------------------------#
# 将灰条部分截取掉
# letterbox_image一般会引入灰条
# pr :array, shape：(512,512,21)，有灰条w、h尺寸会变
#----------------------------------------------------#
pr = pr[int((self.input_shape[0] - nh) // 2) : int((self.input_shape[0] - nh) // 2 + nh), \
int((self.input_shape[1] - nw) // 2) : int((self.input_shape[1] - nw) // 2 + nw)]
#---------------------------------------------------#
# 进行图片的resize
# 灰条去掉后，resize回原图大小
# pr :array, shape：(orininal_w, orininal_h,21)
#---------------------------------------------------#
pr = cv2.resize(pr, (orininal_w, orininal_h), interpolation = cv2.INTER_LINEAR)
#---------------------------------------------------#
# 取出每一个像素点的种类
# pr :array, shape：(orininal_w, orininal_h)
#---------------------------------------------------#
pr = pr.argmax(axis=-1)
#-------------------------------------------------#
# mix_type参数用于控制检测结果的可视化方式
#
# mix_type = 0的时候代表原图与生成的图进行混合
# mix_type = 1的时候代表仅保留生成的图
# mix_type = 2的时候代表仅扣去背景，仅保留原图中的目标
# 下方有给出三种可视化结果
#-------------------------------------------------#
if self.mix_type == 0:
#------------------------------------------------#
# seg_img shape：(orininal_w, orininal_h,3)
#------------------------------------------------#
# seg_img = np.zeros((np.shape(pr)[0], np.shape(pr)[1], 3))
# for c in range(self.num_classes):
# seg_img[:, :, 0] += ((pr[:, :] == c ) * self.colors[c][0]).astype('uint8')
# seg_img[:, :, 1] += ((pr[:, :] == c ) * self.colors[c][1]).astype('uint8')
# seg_img[:, :, 2] += ((pr[:, :] == c ) * self.colors[c][2]).astype('uint8')
seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
#------------------------------------------------#
# 将新图片转换成Image的形式
#------------------------------------------------#
image = Image.fromarray(np.uint8(seg_img))
#------------------------------------------------#
# 将新图与原图及进行混合
# 0.7是加权融合系数
#------------------------------------------------#
image = Image.blend(old_img, image, 0.7)
elif self.mix_type == 1:
# seg_img = np.zeros((np.shape(pr)[0], np.shape(pr)[1], 3))
# for c in range(self.num_classes):
# seg_img[:, :, 0] += ((pr[:, :] == c ) * self.colors[c][0]).astype('uint8')
# seg_img[:, :, 1] += ((pr[:, :] == c ) * self.colors[c][1]).astype('uint8')
# seg_img[:, :, 2] += ((pr[:, :] == c ) * self.colors[c][2]).astype('uint8')
seg_img = np.reshape(np.array(self.colors, np.uint8)[np.reshape(pr, [-1])], [orininal_h, orininal_w, -1])
#------------------------------------------------#
# 将新图片转换成Image的形式
#------------------------------------------------#
image = Image.fromarray(np.uint8(seg_img))
elif self.mix_type == 2:
seg_img = (np.expand_dims(pr != 0, -1) * np.array(old_img, np.float32)).astype('uint8')
#------------------------------------------------#
# 将新图片转换成Image的形式
#------------------------------------------------#
image = Image.fromarray(np.uint8(seg_img))
return image
if __name__ == "__main__":
deeplab = DeeplabV3()
while True:
img = input('Input image filename:')
try:
image = Image.open(img)
except:
print('Open Error! Try again!')
continue
else:
r_image = deeplab.detect_image(image)
r_image.show()
2 单张图片输入时结果输出
mix_type = 0的时候代表原图与生成的图进行混合
mix_type = 1的时候代表仅保留生成的图
mix_type = 2的时候代表仅扣去背景，仅保留原图中的目标
3 图片预处理代码
import numpy as np
from PIL import Image
#---------------------------------------------------------#
# 将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
def cvtColor(image):
if len(np.shape(image)) == 3 and np.shape(image)[2] == 3:
return image
else:
image = image.convert('RGB')
return image
#---------------------------------------------------#
# 对输入图像进行resize
# 这是letterbox_image，不失真resize
#---------------------------------------------------#
def resize_image(image, size):
iw, ih = image.size
w, h = size
scale = min(w/iw, h/ih)
nw = int(iw*scale)
nh = int(ih*scale)
image = image.resize((nw,nh), Image.BICUBIC)
new_image = Image.new('RGB', size, (128,128,128))
new_image.paste(image, ((w-nw)//2, (h-nh)//2))
return new_image, nw, nh
#---------------------------------------------------#
# 归一化
#---------------------------------------------------#
def preprocess_input(image):
image /= 255.0
return image
4 视频分割时代码及结果
对于视频，只需要将上面的 main 函数修改如下：
if __name__ == "__main__":
deeplab = DeeplabV3()
#------------------------------------------------------------------------------------#
# video_path用于指定视频的路径，当video_path=0时表示检测摄像头
# 想要检测视频，则设置如video_path = "xxx.mp4"即可，代表读取出根目录下的xxx.mp4文件。
# video_save_path表示视频保存的路径，当video_save_path=""时表示不保存
# 想要保存视频，则设置如video_save_path = "yyy.mp4"即可，代表保存为根目录下的yyy.mp4文件。
# video_fps用于保存的视频的fps
#------------------------------------------------------------------------------------#
video_path = 0
video_save_path = ""
video_fps = 25.0
capture=cv2.VideoCapture(video_path)
if video_save_path != "":
fourcc = cv2.VideoWriter_fourcc(*'XVID')
size = (int(capture.get(cv2.CAP_PROP_FRAME_WIDTH)), int(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)))
out = cv2.VideoWriter(video_save_path, fourcc, video_fps, size)
ref, frame = capture.read()
if not ref:
raise ValueError("未能正确读取摄像头（视频），请注意是否正确安装摄像头（是否正确填写视频路径）。")
fps = 0.0
while(True):
t1 = time.time()
# 读取某一帧
ref, frame = capture.read()
if not ref:
break
# 格式转变，BGRtoRGB
frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
# 转变成Image
frame = Image.fromarray(np.uint8(frame))
# 进行检测
frame = np.array(deeplab.detect_image(frame))
# RGBtoBGR满足opencv显示格式
frame = cv2.cvtColor(frame,cv2.COLOR_RGB2BGR)
fps = ( fps + (1./(time.time()-t1)) ) / 2
print("fps= %.2f"%(fps))
frame = cv2.putText(frame, "fps= %.2f"%(fps), (0, 40), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
cv2.imshow("video",frame)
c= cv2.waitKey(1) & 0xff
if video_save_path != "":
out.write(frame)
if c==27:
capture.release()
break
print("Video Detection Done!")
capture.release()
if video_save_path != "":
print("Save processed video to the path :" + video_save_path)
out.release()
cv2.destroyAllWindows()
结果：
5 遍历文件夹下多图片分割实现
对于文件夹下多图片分割，只需要将上面的 main 函数修改如下：
if __name__ == "__main__":
deeplab = DeeplabV3()
#-------------------------------------------------------------------------#
# dir_origin_path指定了用于检测的图片的文件夹路径
# dir_save_path指定了检测完图片的保存路径
#-------------------------------------------------------------------------#
dir_origin_path = "img/"
dir_save_path = "img_out/"
import os
#-------------------------------------------------#
# tqdm：进度条的效果，可参考链接：
# https://www.jianshu.com/p/b171f5f127b5
#-------------------------------------------------#
from tqdm import tqdm
img_names = os.listdir(dir_origin_path)
for img_name in tqdm(img_names):
if img_name.lower().endswith(('.bmp', '.dib', '.png', '.jpg', '.jpeg', '.pbm', '.pgm', '.ppm', '.tif', '.tiff')):
image_path = os.path.join(dir_origin_path, img_name)
image = Image.open(image_path)
r_image = deeplab.detect_image(image)
if not os.path.exists(dir_save_path):
os.makedirs(dir_save_path)
r_image.save(os.path.join(dir_save_path, img_name))
结果：
6 想测模型FPS时
import colorsys
import copy
import time
import cv2
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch import nn
#-----------------------------------------------------------------------------------#
# DeepLab网络及代码
# 详细介绍可见 https://blog.csdn.net/weixin_45377629/article/details/124083978
#-----------------------------------------------------------------------------------#
from nets.deeplabv3_plus import DeepLab
#-----------------------------------------------------------------------------------#
# 三个函数代码下方给出
# cvtColor： 将图像转换成RGB图像，防止灰度图在预测时报错。
# preprocess_input： 归一化
# resize_image： 对输入图像进行resize，letterbox_image方式，不失真resize
#-----------------------------------------------------------------------------------#
from utils.utils import cvtColor, preprocess_input, resize_image
#-----------------------------------------------------------------------------------#
# 使用自己训练好的模型预测需要修改3个参数
# model_path、backbone和num_classes都需要修改！
# 如果出现shape不匹配，一定要注意训练时的model_path、backbone和num_classes的修改
#-----------------------------------------------------------------------------------#
class DeeplabV3(object):
_defaults = {
#-------------------------------------------------------------------#
# model_path指向logs文件夹下的权值文件
# 训练好后logs文件夹下存在多个权值文件，选择验证集损失较低的即可。
# 验证集损失较低不代表miou较高，仅代表该权值在验证集上泛化性能较好。
# 链接：https://pan.baidu.com/s/1TrBlnZUd6xwxUvgFjbz7TQ 提取码：cj80
#-------------------------------------------------------------------#
"model_path" : 'model_data/deeplab_mobilenetv2.pth',
#----------------------------------------#
# 所需要区分的类的个数+1
#----------------------------------------#
"num_classes" : 21,
#----------------------------------------#
# 所使用的的主干网络：
# mobilenet
# xception
#----------------------------------------#
"backbone" : "mobilenet",
#----------------------------------------#
# 输入图片的大小
#----------------------------------------#
"input_shape" : [512, 512],
#----------------------------------------#
# 下采样的倍数，一般可选的为8和16
# 与训练时设置的一样即可
#----------------------------------------#
"downsample_factor" : 16,
#-------------------------------------------------#
# mix_type参数用于控制检测结果的可视化方式
#
# mix_type = 0的时候代表原图与生成的图进行混合
# mix_type = 1的时候代表仅保留生成的图
# mix_type = 2的时候代表仅扣去背景，仅保留原图中的目标
# 下方有给出三种可视化结果的区别
#-------------------------------------------------#
"mix_type" : 0,
#-------------------------------#
# 是否使用Cuda
# 没有GPU可以设置成False
#-------------------------------#
"cuda" : False,
}
#---------------------------------------------------#
# 初始化Deeplab
#---------------------------------------------------#
def __init__(self, **kwargs):
#---------------------------------------------------#
# _defaults字典原来是这么用起来的
#---------------------------------------------------#
self.__dict__.update(self._defaults)
for name, value in kwargs.items():
#-----------------------------------------------#
# 设置属性 name 值，即self.name==value
#-----------------------------------------------#
setattr(self, name, value)
#---------------------------------------------------#
# 画框设置不同的颜色
#---------------------------------------------------#
if self.num_classes <= 21:
self.colors = [ (0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128), (128, 0, 128), (0, 128, 128),
(128, 128, 128), (64, 0, 0), (192, 0, 0), (64, 128, 0), (192, 128, 0), (64, 0, 128), (192, 0, 128),
(64, 128, 128), (192, 128, 128), (0, 64, 0), (128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128),
(128, 64, 12)]
else:
hsv_tuples = [(x / self.num_classes, 1., 1.) for x in range(self.num_classes)]
self.colors = list(map(lambda x: colorsys.hsv_to_rgb(*x), hsv_tuples))
self.colors = list(map(lambda x: (int(x[0] * 255), int(x[1] * 255), int(x[2] * 255)), self.colors))
#---------------------------------------------------#
# 获得模型
#---------------------------------------------------#
self.generate()
#---------------------------------------------------#
# 获得所有的分类
#---------------------------------------------------#
def generate(self):
#-------------------------------#
# 载入模型与权值
# 详细介绍可见 https://blog.csdn.net/weixin_45377629/article/details/124083978
#-------------------------------#
self.net = DeepLab(num_classes=self.num_classes, backbone=self.backbone, downsample_factor=self.downsample_factor, pretrained=False)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
self.net.load_state_dict(torch.load(self.model_path, map_location=device))
self.net = self.net.eval()
print('{} model, and classes loaded.'.format(self.model_path))
if self.cuda:
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
#---------------------------------------------------#
# 测模型FPS
#---------------------------------------------------#
def get_FPS(self, image, test_interval):
#---------------------------------------------------------#
# 在这里将图像转换成RGB图像，防止灰度图在预测时报错。
# 代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
#---------------------------------------------------------#
image = cvtColor(image)
#---------------------------------------------------------#
# 给图像增加灰条，实现不失真的resize
# 也可以直接resize进行识别
#---------------------------------------------------------#
image_data, nw, nh = resize_image(image, (self.input_shape[1],self.input_shape[0]))
#---------------------------------------------------------#
# 添加上batch_size维度
#---------------------------------------------------------#
image_data = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, np.float32)), (2, 0, 1)), 0)
with torch.no_grad():
images = torch.from_numpy(image_data)
if self.cuda:
images = images.cuda()
#---------------------------------------------------#
# 图片传入网络进行预测
#---------------------------------------------------#
pr = self.net(images)[0]
#---------------------------------------------------#
# 取出每一个像素点的种类
#---------------------------------------------------#
pr = F.softmax(pr.permute(1,2,0),dim = -1).cpu().numpy().argmax(axis=-1)
#--------------------------------------#
# 将灰条部分截取掉
#--------------------------------------#
pr = pr[int((self.input_shape[0] - nh) // 2) : int((self.input_shape[0] - nh) // 2 + nh), \
int((self.input_shape[1] - nw) // 2) : int((self.input_shape[1] - nw) // 2 + nw)]
t1 = time.time()
for _ in range(test_interval):
with torch.no_grad():
#---------------------------------------------------#
# 图片传入网络进行预测
#---------------------------------------------------#
pr = self.net(images)[0]
#---------------------------------------------------#
# 取出每一个像素点的种类
#---------------------------------------------------#
pr = F.softmax(pr.permute(1,2,0),dim = -1).cpu().numpy().argmax(axis=-1)
#--------------------------------------#
# 将灰条部分截取掉
#--------------------------------------#
pr = pr[int((self.input_shape[0] - nh) // 2) : int((self.input_shape[0] - nh) // 2 + nh), \
int((self.input_shape[1] - nw) // 2) : int((self.input_shape[1] - nw) // 2 + nw)]
t2 = time.time()
tact_time = (t2 - t1) / test_interval
return tact_time
if __name__ == "__main__":
deeplab = DeeplabV3()
#-------------------------------------------------------------------------#
# test_interval用于指定测量fps的时候，图片检测的次数
# 理论上test_interval越大，fps越准确。
#-------------------------------------------------------------------------#
test_interval = 10
img = Image.open('img/street.jpg')
tact_time = deeplab.get_FPS(img, test_interval)
print(str(tact_time) + ' seconds, ' + str(1/tact_time) + 'FPS, @batch_size 1')
结果输出：
model_data/deeplab_mobilenetv2.pth model, and classes loaded.
0.5868569374084472 seconds, 1.7039928068601988FPS, @batch_size 1
7 感谢链接
https://blog.csdn.net/weixin_44791964/article/details/120113686
1万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
