车道线数据集例图
1、数据处理
import torchkeras
from argparse import Namespace
from pathlib import Path
from PIL import Image
import numpy as np
import torch
from torch import nn
from torch.utils.data import Dataset,DataLoader
import os
from torchkeras.data import resize_and_pad_image
from torchkeras.plots import joint_imgs_col
import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2
import glob
config = Namespace(
img_size = 128,
lr = 1e-4,
batch_size = 4,
)
class MyDataset(Dataset):
def __init__(self, img_files, img_size, transforms = None):
self.__dict__.update(locals())
def __len__(self) -> int:
return len(self.img_files)
def get(self, index):
img_path = self.img_files[index]
mask_path = img_path.replace('images','masks').replace('.jpg','.png')
image = Image.open(img_path).convert('RGB')
mask = Image.open(mask_path).convert('L')
return image, mask
def __getitem__(self, index):
image,mask = self.get(index)
image = resize_and_pad_image(image,self.img_size,self.img_size)
mask = resize_and_pad_image(mask,self.img_size,self.img_size)
image_arr = np.array(image, dtype=np.float32)/255.0
mask_arr = np.array(mask,dtype=np.float32)
mask_arr = np.where(mask_arr>100.0,1.0,0.0).astype(np.int64)
sample = {
"image": image_arr,
"mask": mask_arr
}
if self.transforms is not None:
sample = self.transforms(**sample)
sample['mask'] = sample['mask'][None,...]
return sample
def show_sample(self, index):
image, mask = self.get(index)
image_result = joint_imgs_col(image,mask)
return image_result
def get_train_transforms():
return A.Compose(
[
A.OneOf([A.HorizontalFlip(p=0.5),A.VerticalFlip(p=0.5)]),
ToTensorV2(p=1),
],
p=1.0
)
def get_val_transforms():
return A.Compose(
[
ToTensorV2(p=1),
],
p=1.0
)
# 训练和验证图像分别存放在两个不同的目录下
train_dir = 'datasets/lane_lines/train/images'
val_dir = 'datasets/lane_lines/val/images'
# 获取所有图像文件的路径
train_imgs = glob.glob(os.path.join(train_dir, '*.jpg')) # 图像格式为 .jpg
val_imgs = glob.glob(os.path.join(val_dir, '*.jpg')) # 图像格式为 .jpg
train_transforms=get_train_transforms()
val_transforms=get_val_transforms()
ds_train = MyDataset(train_imgs,img_size=config.img_size,transforms=train_transforms)
ds_val = MyDataset(val_imgs,img_size=config.img_size,transforms=val_transforms)
dl_train = DataLoader(ds_train,batch_size=config.batch_size)
dl_val = DataLoader(ds_val,batch_size=config.batch_size)
ds_train.show_sample(10)
2、模型定义与训练
import segmentation_models_pytorch as smp
import torch
from torchkeras import KerasModel
from torch.nn import functional as F
from torchkeras.metrics import IOU
from torchkeras.kerascallbacks import WandbCallback
num_classes = 1
net = smp.DeepLabV3Plus(
encoder_name="mobilenet_v2", # choose encoder, e.g. mobilenet_v2 or efficientnet-b7
encoder_weights='imagenet', # use `imagenet` pretrained weights for encoder initialization
in_channels=3, # model input channels (1 for grayscale images, 3 for RGB, etc.)
classes=num_classes, # model output channels (number of classes in your dataset)
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 由于输入数据batch结构差异,需要重写StepRunner并覆盖
class StepRunner:
def __init__(self, net, loss_fn, accelerator, stage = "train", metrics_dict = None,
optimizer = None, lr_scheduler = None
):
self.net,self.loss_fn,self.metrics_dict,self.stage = net,loss_fn,metrics_dict,stage
self.optimizer,self.lr_scheduler = optimizer,lr_scheduler
self.accelerator = accelerator
if self.stage=='train':
self.net.train()
else:
self.net.eval()
def __call__(self, batch):
features,labels = batch['image'],batch['mask']
#loss
preds = self.net(features)
loss = self.loss_fn(preds,labels)
#backward()
if self.optimizer is not None and self.stage=="train":
self.accelerator.backward(loss)
self.optimizer.step()
if self.lr_scheduler is not None:
self.lr_scheduler.step()
self.optimizer.zero_grad()
all_preds = self.accelerator.gather(preds)
all_labels = self.accelerator.gather(labels)
all_loss = self.accelerator.gather(loss).sum()
#losses
step_losses = {self.stage+"_loss":all_loss.item()}
#metrics
step_metrics = {self.stage+"_"+name:metric_fn(all_preds, all_labels).item()
for name,metric_fn in self.metrics_dict.items()}
if self.optimizer is not None and self.stage=="train":
step_metrics['lr'] = self.optimizer.state_dict()['param_groups'][0]['lr']
return step_losses,step_metrics
KerasModel.StepRunner = StepRunner
class DiceLoss(nn.Module):
def __init__(self,smooth=0.001,num_classes=1,weights = None):
super().__init__()
self.smooth = smooth
self.num_classes = num_classes
self.weights = weights
def forward(self, logits, targets):
# binary
if self.num_classes==1:
preds = logits.contiguous().view(logits.size()[0], -1).sigmoid()
targets = targets.contiguous().view(targets.size()[0], -1).float()
loss = self.compute_loss(preds,targets)
return loss
# multiclass
else:
preds = logits.softmax(axis=1).contiguous().view(
logits.size()[0],self.num_classes, -1)
t = targets.contiguous().view(
targets.size()[0], -1)
targets = torch.nn.functional.one_hot(t,self.num_classes).permute(0,2,1)
totalLoss = 0.0
for i in range(self.num_classes):
diceLoss = self.compute_loss(preds[:, i], targets[:, i])
if self.weights is not None:
diceLoss *= self.weights[i]
totalLoss += diceLoss
return totalLoss
def compute_loss(self,preds,targets):
a = torch.sum(preds * targets, 1) # |X⋂Y|
b = torch.sum(preds * preds, 1) + self.smooth # |X|
c = torch.sum(targets * targets, 1) + self.smooth # |Y|
score = (2 * a) / (b + c)
loss = torch.mean(1-score)
return loss
class MixedLoss(nn.Module):
def __init__(self,bce_ratio=0.5):
super().__init__()
self.bce = nn.BCEWithLogitsLoss()
self.dice = DiceLoss()
self.bce_ratio = bce_ratio
def forward(self,logits,targets):
bce_loss = self.bce(logits,targets.float())
dice_loss = self.dice(logits,targets)
total_loss = bce_loss*self.bce_ratio + dice_loss*(1-self.bce_ratio)
return total_loss
optimizer = torch.optim.AdamW(net.parameters(), lr=config.lr)
lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
optimizer = optimizer,
T_max=8,
eta_min=0
)
metrics_dict = {'iou': IOU(num_classes=1)}
model = KerasModel(net,
loss_fn=MixedLoss(bce_ratio=0.5),
metrics_dict=metrics_dict,
optimizer=optimizer,
lr_scheduler = lr_scheduler
)
dfhistory=model.fit(train_data=dl_train,
val_data=dl_val,
epochs=100,
ckpt_path='checkpoint.pt',
patience=10,
monitor="val_iou",
mode="max",
mixed_precision='no',
plot = True
)
3、模型评估与可视化
from matplotlib import pyplot as plt
metrics_dict = {'iou': IOU(num_classes=1,if_print=True)}
model = KerasModel(net,
loss_fn=MixedLoss(bce_ratio=0.5),
metrics_dict=metrics_dict,
optimizer=optimizer,
lr_scheduler = lr_scheduler
)
model.evaluate(dl_val)
batch = next(iter(dl_val))
with torch.no_grad():
model.eval()
logits = model(batch["image"].cuda())
pr_masks = logits.sigmoid()
for image, gt_mask, pr_mask in zip(batch["image"], batch["mask"], pr_masks):
plt.figure(figsize=(16, 10))
plt.subplot(1, 3, 1)
plt.imshow(image.numpy().transpose(1, 2, 0)) # convert CHW -> HWC
plt.title("Image")
plt.axis("off")
plt.subplot(1, 3, 2)
plt.imshow(gt_mask.numpy().squeeze())
plt.title("Ground truth")
plt.axis("off")
plt.subplot(1, 3, 3)
plt.imshow(pr_mask.cpu().numpy().squeeze())
plt.title("Prediction")
plt.axis("off")
plt.show()