def generate(self):
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
self.net = pspnet(num_classes=self.num_classes, downsample_factor=self.downsample_factor, pretrained=False, backbone=self.backbone, aux_branch=False)
self.net = self.net.eval()
state_dict = torch.load(self.model_path)
self.net.load_state_dict(state_dict, strict=False)
if self.cuda:
self.net = nn.DataParallel(self.net)
self.net = self.net.cuda()
print('{} model, anchors, and classes loaded.'.format(self.model_path))
# 画框设置不同的颜色
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 / len(self.class_names), 1., 1.)
for x in range(len(self.class_names))]
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))
def generate(self):
self.model = pspnet(self.num_classes,
self.model_image_size,
downsample_factor=self.downsample_factor,
backbone=self.backbone,
aux_branch=False)
self.model.load_weights(self.model_path, by_name=True)
print('{} model loaded.'.format(self.model_path))
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 / len(self.class_names), 1., 1.)
for x in range(len(self.class_names))]
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))
def generate(self):
#-------------------------------#
# 载入模型与权值
#-------------------------------#
self.model = pspnet([self.input_shape[0], self.input_shape[1], 3],
self.num_classes,
backbone=self.backbone,
downsample_factor=self.downsample_factor,
aux_branch=False)
self.model.load_weights(self.model_path, by_name=True)
print('{} model loaded.'.format(self.model_path))
def generate(self):
#-------------------------------#
# 载入模型与权值
#-------------------------------#
self.net = pspnet(num_classes=self.num_classes,
downsample_factor=self.downsample_factor,
pretrained=False,
backbone=self.backbone,
aux_branch=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),
strict=False)
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()
# 是否使用辅助分支，会占用大量显存
#---------------------------------------------------------------------#
aux_branch = False
#-------------------------------------------------------------------#
# 用于设置是否使用多线程读取数据，1代表关闭多线程
# 开启后会加快数据读取速度，但是会占用更多内存
# 在IO为瓶颈的时候再开启多线程，即GPU运算速度远大于读取图片的速度。
#-------------------------------------------------------------------#
num_workers = 1
#------------------------------------------------------#
# 获取model
#------------------------------------------------------#
model = pspnet([input_shape[0], input_shape[1], 3],
num_classes,
downsample_factor=downsample_factor,
backbone=backbone,
aux_branch=aux_branch)
if model_path != '':
#------------------------------------------------------#
# 载入预训练权重
#------------------------------------------------------#
print('Load weights {}.'.format(model_path))
model.load_weights(model_path, by_name=True, skip_mismatch=True)
if focal_loss:
if dice_loss:
loss = dice_loss_with_Focal_Loss(cls_weights)
else:
loss = Focal_Loss(cls_weights)
else:
# 是否使用辅助分支
#------------------------------#
aux_branch = True
#------------------------------#
# 下采样的倍数
#------------------------------#
downsample_factor = 16
#------------------------------#
# 数据集路径
#------------------------------#
dataset_path = "dataset_processing/classfication/"
# 获取model
model = pspnet(num_classes,
inputs_size,
downsample_factor=downsample_factor,
backbone=backbone,
aux_branch=aux_branch)
#-------------------------------------------#
# 权值文件位置
#-------------------------------------------#
model_path = "model_data/pspnet_mobilenetv2.h5"
model.load_weights(model_path, by_name=True, skip_mismatch=True)
# 打开训练集
with open(os.path.join(dataset_path, "ImageSets/Segmentation/train.txt"),
"r") as f:
train_lines = f.readlines()
# 打开验证集
from nets.pspnet import pspnet
model = pspnet(21, [473, 473, 3],
downsample_factor=16,
backbone='resnet50',
aux_branch=False)
model.summary()
for i, layer in enumerate(model.layers):
print(i, layer.name)
#--------------------------------------------#
# 该部分代码只用于看网络结构，并非测试代码
#--------------------------------------------#
from nets.pspnet import pspnet
if __name__ == "__main__":
model = pspnet(21, [473, 473, 3],
downsample_factor=16,
backbone='mobilenet',
aux_branch=False)
model.summary()
for i, layer in enumerate(model.layers):
print(i, layer.name)
