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)