1. 定义数据集
https://pytorch.org/tutorials/intermediate/torchvision_tutorial.html
数据集和相关的代码都在官网都可以下载到,一切准备工作完成之后,第一步就是构建数据集:
import os
import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision.io import read_image
from torchvision.ops import masks_to_boxes
import torchvision
import os
import numpy as np
import torch
from PIL import Image
class PennFudanDataset(Dataset):
def __init__(self, root, transforms=None):
self.root = root # 图片的根目录
self.transforms = transforms # 对图片做转换的函数
# load all image files, sorting them to
# ensure that they are aligned
self.imgs = list(sorted(os.listdir(os.path.join(root, "PNGImages")))) # 存放图片文件名的列表
self.masks = list(sorted(os.listdir(os.path.join(root, "PedMasks")))) # 存放mask文件的列表
def __getitem__(self, idx):
# load images and masks
img_path = os.path.join(self.root, "PNGImages", self.imgs[idx])
mask_path = os.path.join(self.root, "PedMasks", self.masks[idx])
img = Image.open(img_path).convert("RGB") # 读取图片并转换为RGB模式
mask = Image.open(mask_path) # mask不需要转换为RGB模式,因为mask的每个像素值都代表了一种实例
mask = np.array(mask) # 将mask转为array
obj_ids = np.unique(mask) # 去重之后的每个像素值代表了一种实例
obj_ids = obj_ids[1:] # 0代表背景,因此去除
masks = mask == obj_ids[:, None, None] # [n_masks, h, w], 构建mask的列表,每个mask的高宽和图片一致,每个mask只有实例部分值为True
# 构建锚框
num_objs = len(obj_ids)
boxes = []
for i in range(num_objs):
pos = np.where(masks[i])
xmin = np.min(pos[1])
xmax = np.max(pos[1])
ymin = np.min(pos[0])
ymax = np.max(pos[0])
boxes.append([xmin, ymin, xmax, ymax])
boxes = torch.as_tensor(boxes, dtype=torch.float32)
# 本例中只有一种类别
labels = torch.ones((num_objs,), dtype=torch.int64)
masks = torch.as_tensor(masks, dtype=torch.uint8)
image_id = torch.tensor([idx])
area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]) # 计算锚框面积
iscrowd = torch.zeros((num_objs,), dtype=torch.int64) # iscrowd为True时,在测试阶段将会被跳过,这里假定都是False
target = {
}
target["boxes"] = boxes
target["labels"] = labels
target["masks"] = masks
target["image_id"] = image_id
target["area"] = area
target["iscrowd"] = iscrowd
if self.transforms is not None:
img, target = self.transforms(img, target)
return img, target
def __len__(self):
return len(self.imgs)
数据转换函数如下所示,这里需要对照官网下载正确版本的包,下错了会报错:
from engine import train_one_epoch, evaluate
import utils
import transforms as T
def get_transform(train):
transforms = []
# converts the image, a PIL image, into a PyTorch Tensor
transforms.append(T.ToTensor())
if train:
# during training, randomly flip the training images
# and ground-truth for data augmentation
transforms.append(T.RandomHorizontalFlip(0.5))
return T.Compose(transforms)
2. 定义模型
加载预训练模型,并重新定义分类器
import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
def get_instance_segmentation_model(num_classes):
# 加载预训练模型
model = torchvision.models.detection.maskrcnn_resnet50_fpn(pretrained=True)
# 获取输入至分类器的特征数
in_features = model.roi_heads.box_predictor.cls_score.in_features
# 替换成自己重新定义的分类器
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
# 获取mask分类器的输入特征数
in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
hidden_layer = 256
# 替换mask分类器
model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask,
hidden_layer,
num_classes)
return model
3. 训练
# use our dataset and defined transformations
dataset = PennFudanDataset('PennFudanPed', get_transform(train=True))
dataset_test = PennFudanDataset('PennFudanPed', get_transform(train=False))
# split the dataset in train and test set
torch.manual_seed(1)
indices = torch.randperm(len(dataset)).tolist()
dataset = torch.utils.data.Subset(dataset, indices[:-50])
dataset_test = torch.utils.data.Subset(dataset_test, indices[-50:])
# define training and validation data loaders
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=2, shuffle=True, num_workers=0,
collate_fn=utils.collate_fn)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=1, shuffle=False, num_workers=0,
collate_fn=utils.collate_fn)
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
# our dataset has two classes only - background and person
num_classes = 2
# get the model using our helper function
model = get_instance_segmentation_model(num_classes)
# move model to the right device
model.to(device)
# construct an optimizer
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
momentum=0.9, weight_decay=0.0005)
# and a learning rate scheduler which decreases the learning rate by
# 10x every 3 epochs
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=3,
gamma=0.1)
from torch.optim.lr_scheduler import StepLR
num_epochs = 5
for epoch in range(num_epochs):
# train for one epoch, printing every 10 iterations
train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq=10)
# update the learning rate
lr_scheduler.step()
# evaluate on the test dataset
evaluate(model, data_loader_test, device=device)