with open(os.path.join(txt_path), “r”) as f:
file_names = [x.strip() for x in f.readlines() if len(x.strip()) > 0]
self.images = [os.path.join(image_dir, x + “.jpg”) for x in file_names]
self.masks = [os.path.join(mask_dir, x + “.png”) for x in file_names]
assert (len(self.images) == len(self.masks))
self.transforms = transforms
导入需要的包。
定义VOC数据集读取类VOCSegmentation。在init方法中，核心是读取image列表和mask列表。
def getitem(self, index):
img = Image.open(self.images[index]).convert(‘RGB’)
target = Image.open(self.masks[index])
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
__getitem__方法是获取单张图片和图片对应的mask，然后对其做数据增强。
def collate_fn(batch):
images, targets = list(zip(*batch))
batched_imgs = cat_list(images, fill_value=0)
batched_targets = cat_list(targets, fill_value=255)
return batched_imgs, batched_targets
collate_fn方法是对一个batch中数据调用cat_list做数据对齐。
在train.py中torch.utils.data.DataLoader调用
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=True,
pin_memory=True,
collate_fn=train_dataset.collate_fn)
val_loader = torch.utils.data.DataLoader(val_dataset,
batch_size=1,
num_workers=num_workers,
pin_memory=True,
collate_fn=val_dataset.collate_fn)
=============================================================
打开train.py,我们先认识一下重要的参数：
def parse_args():
import argparse
parser = argparse.ArgumentParser(description=“pytorch fcn training”)
数据集的根目录(VOCdevkit)所在的文件夹
parser.add_argument(“–data-path”, default=“data/”, help=“VOCdevkit root”)
parser.add_argument(“–num-classes”, default=20, type=int)
parser.add_argument(“–aux”, default=True, type=bool, help=“auxilier loss”)
parser.add_argument(“–device”, default=“cuda”, help=“training device”)
parser.add_argument(“-b”, “–batch-size”, default=32, type=int)
parser.add_argument(“–epochs”, default=30, type=int, metavar=“N”,
help=“number of total epochs to train”)
parser.add_argument(‘–lr’, default=0.0001, type=float, help=‘initial learning rate’)
parser.add_argument(‘–momentum’, default=0.9, type=float, metavar=‘M’,
help=‘momentum’)
parser.add_argument(‘–wd’, ‘–weight-decay’, default=1e-4, type=float,
metavar=‘W’, help=‘weight decay (default: 1e-4)’,
dest=‘weight_decay’)
parser.add_argument(‘–print-freq’, default=10, type=int, help=‘print frequency’)
parser.add_argument(‘–resume’, default=‘’, help=‘resume from checkpoint’)
parser.add_argument(‘–start-epoch’, default=0, type=int, metavar=‘N’,
help=‘start epoch’)
是否使用混合精度训练
parser.add_argument(“–amp”, default=False, type=bool,
help=“Use torch.cuda.amp for mixed precision training”)
args = parser.parse_args()
return args
data-path：定义数据集的根目录(VOCdevkit)所在的文件夹
num-classes：检测目标类别数(不包含背景)。
aux：是否使用aux_classifier。
device：使用cpu还是gpu训练，默认是cuda。
batch-size：BatchSize设置。
epochs：epoch的个数。
lr：学习率。
resume：继续训练时候，选择用的模型。
start-epoch：起始的epoch，针对再次训练时，可以不需要从0开始。
amp：是否使用torch的自动混合精度训练。
增强调用transforms.py中的方法。
训练集的增强如下：
class SegmentationPresetTrain:
def init(self, base_size, crop_size, hflip_prob=0.5, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
随机Resize的最小尺寸
min_size = int(0.5 * base_size)
随机Resize的最大尺寸
max_size = int(2.0 * base_size)
随机Resize增强。
trans = [T.RandomResize(min_size, max_size)]
if hflip_prob > 0:
#随机水平翻转
trans.append(T.RandomHorizontalFlip(hflip_prob))
trans.extend([
#随机裁剪
T.RandomCrop(crop_size),
T.ToTensor(),
T.Normalize(mean=mean, std=std),
])
self.transforms = T.Compose(trans)
def call(self, img, target):
return self.transforms(img, target)
训练集增强，包括随机Resize、随机水平翻转、随即裁剪。
验证集增强：
class SegmentationPresetEval:
def init(self, base_size, mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)):
self.transforms = T.Compose([
T.RandomResize(base_size, base_size),
T.ToTensor(),
T.Normalize(mean=mean, std=std),
])
def call(self, img, target):
return self.transforms(img, target)
验证集的增强比较简单，只有随机Resize。
对Main方法，我做了一些修改，修改的代码如下：
#定义模型，并加载预训练
model = fcn_resnet50(pretrained=True)
默认classes是21,如果不是21，则要修改类别。
if num_classes != 21:
model.classifier[4] = torch.nn.Conv2d(512, num_classes, kernel_size=(1, 1), stride=(1, 1))
model.aux_classifier[4] = torch.nn.Conv2d(256, num_classes, kernel_size=(1, 1), stride=(1, 1))
print(model)
model.to(device)
如果有多张显卡，则使用多张显卡
if torch.cuda.device_count() > 1:
print(“Let’s use”, torch.cuda.device_count(), “GPUs!”)
model = torch.nn.DataParallel(model)
模型，我改为pytorch官方的模型了，如果能使用官方的模型尽量使用官方的模型。
默认类别是21，如果不是21，则要修改类别。
检测系统中是否有多张卡，如果有多张卡则使用多张卡不能浪费资源。
如果不想使用所有的卡，而是指定其中的几张卡，可以使用：
os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘0,1’
也可以在DataParallel方法中设定：
model = torch.nn.DataParallel(model,device_ids=[0,1])
如果使用了多显卡，再使用模型的参数就需要改为model.module.xxx，例如：
params = [p for p in model.module.aux_classifier.parameters() if p.requires_grad]
params_to_optimize.append({“params”: params, “lr”: args.lr * 10})
上面的都完成了就可以开始训练了，如下图：
测试=====================================================================================================================================================================
在开始测试之前，我们还要获取到调色板，新建脚本get_palette.py，代码如下：
import json
import numpy as np
from PIL import Image
读取mask标签
target = Image.open(“./2007_001288.png”)
获取调色板
palette = target.getpalette()
palette = np.reshape(palette, (-1, 3)).tolist()
print(palette)
转换成字典子形式
pd = dict((i, color) for i, color in enumerate(palette))
json_str = json.dumps(pd)
with open(“palette.json”, “w”) as f:
f.write(json_str)
选取一张mask，然后使用getpalette方法获取，然后将其转为字典的格式保存。
接下来，开始预测部分，新建predict.py,插入以下代码：
import os
import time
import json
import torch
from torchvision import transforms
import numpy as np
from PIL import Image
from torchvision.models.segmentation import fcn_resnet50
导入程序需要的包文件，然在mian方法中：
def main():
aux = False # inference time not need aux_classifier
classes = 20
weights_path = “./save_weights/model_5.pth”
img_path = “./2007_000123.jpg”
palette_path = “./palette.json”
assert os.path.exists(weights_path), f"weights {weights_path} not found."
assert os.path.exists(img_path), f"image {img_path} not found."
assert os.path.exists(palette_path), f"palette {palette_path} not found."
with open(palette_path, “rb”) as f:
pallette_dict = json.load(f)
pallette = []
for v in pallette_dict.values():
pallette += v
-
定义是否需要aux_classifier，预测不需要aux_classifier，所以设置为False。
-
设置类别为20，不包括背景。
-
定义权重的路径。
-
定义调色板的路径。
-
读去调色板。
接下来，是加载模型，单显卡训练出来的模型和多显卡训练出来的模型加载有区别，我们先看单显卡训练出来的模型如何加载。
model = fcn_resnet50(num_classes=classes+1)
print(model)
单显卡训练出来的模型，加载
delete weights about aux_classifier
weights_dict = torch.load(weights_path, map_location=‘cpu’)[‘model’]
for k in list(weights_dict.keys()):
if “aux_classifier” in k:
del weights_dict[k]
load weights
model.load_state_dict(weights_dict)
model.to(device)
定义模型fcn_resnet50，num_classes设置为类别+1（背景）
加载训练好的模型，并将aux_classifier删除。
然后加载权重。
再看多显卡的模型如何加载
create model
model = fcn_resnet50(num_classes=classes+1)
model = torch.nn.DataParallel(model)
delete weights about aux_classifier
weights_dict = torch.load(weights_path, map_location=‘cpu’)[‘model’]
print(weights_dict)
for k in list(weights_dict.keys()):
if “aux_classifier” in k:
del weights_dict[k]
load weights
model.load_state_dict(weights_dict)
model=model.module
model.to(device)
定义模型fcn_resnet50，num_classes设置为类别+1（背景），将模型放入DataParallel类中。
加载训练好的模型，并将aux_classifier删除。
加载权重。
执行torch.nn.DataParallel(model)时，model被放在了model.module，所以model.module才真正需要的模型。所以我们在这里将model.module赋值给model。
接下来是图像数据的处理
load image
original_img = Image.open(img_path)
from pil image to tensor and normalize
data_transform = transforms.Compose([transforms.Resize(520),
transforms.ToTensor(),
transforms.Normalize(mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225))])
img = data_transform(original_img)
expand batch dimension
img = torch.unsqueeze(img, dim=0)
加载图像。
对图像做Resize、标准化、归一化处理。
使用torch.unsqueeze增加一个维度。
完成图像的处理后，就可以开始预测了。
model.eval() # 进入验证模式
with torch.no_grad():
init model
img_height, img_width = img.shape[-2:]
init_img = torch.zeros((1, 3, img_height, img_width), device=device)
model(init_img)
t_start = time_synchronized()
output = model(img.to(device))
t_end = time_synchronized()
print(“inference+NMS time: {}”.format(t_end - t_start))
prediction = output[‘out’].argmax(1).squeeze(0)
prediction = prediction.to(“cpu”).numpy().astype(np.uint8)
np.set_printoptions(threshold=sys.maxsize)
print(prediction.shape)
mask = Image.fromarray(prediction)
mask.putpalette(pallette)
mask.save(“test_result.png”)
将预测后的结果保存到test_result.png中。查看运行结果：
原图：
结果：
打印出来的数据：
类别列表：
{
“aeroplane”: 1,
“bicycle”: 2,
“bird”: 3,
“boat”: 4,
“bottle”: 5,
“bus”: 6,
“car”: 7,
“cat”: 8,
“chair”: 9,
“cow”: 10,
“diningtable”: 11,
“dog”: 12,
最后
Python崛起并且风靡，因为优点多、应用领域广、被大牛们认可。学习 Python 门槛很低，但它的晋级路线很多，通过它你能进入机器学习、数据挖掘、大数据，CS等更加高级的领域。Python可以做网络应用，可以做科学计算，数据分析，可以做网络爬虫，可以做机器学习、自然语言处理、可以写游戏、可以做桌面应用…Python可以做的很多，你需要学好基础，再选择明确的方向。这里给大家分享一份全套的 Python 学习资料，给那些想学习 Python 的小伙伴们一点帮助！
👉Python所有方向的学习路线👈
Python所有方向的技术点做的整理，形成各个领域的知识点汇总，它的用处就在于，你可以按照上面的知识点去找对应的学习资源，保证自己学得较为全面。
👉Python必备开发工具👈
工欲善其事必先利其器。学习Python常用的开发软件都在这里了，给大家节省了很多时间。
👉Python全套学习视频👈
我们在看视频学习的时候，不能光动眼动脑不动手，比较科学的学习方法是在理解之后运用它们，这时候练手项目就很适合了。
👉实战案例👈
学python就与学数学一样，是不能只看书不做题的，直接看步骤和答案会让人误以为自己全都掌握了，但是碰到生题的时候还是会一筹莫展。
因此在学习python的过程中一定要记得多动手写代码，教程只需要看一两遍即可。
👉大厂面试真题👈
我们学习Python必然是为了找到高薪的工作，下面这些面试题是来自阿里、腾讯、字节等一线互联网大厂最新的面试资料，并且有阿里大佬给出了权威的解答，刷完这一套面试资料相信大家都能找到满意的工作。
扫一扫
874
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
