devbox
秋刀鱼在做梦
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

【PyTorch】多对象分割项目

作者：秋刀鱼在做梦 | 2024-08-06 23:13:04

赞

踩

【PyTorch】多对象分割项目

 【PyTorch】单对象分割项目

对象分割任务的目标是找到图像中目标对象的边界。实际应用例如自动驾驶汽车和医学成像分析。这里将使用PyTorch开发一个深度学习模型来完成多对象分割任务。多对象分割的主要目标是自动勾勒出图像中多个目标对象的边界。

对象的边界通常由与图像大小相同的分割掩码定义，在分割掩码中属于目标对象的所有像素基于预定义的标记被标记为相同。

目录

创建数据集

创建数据加载器

创建模型

部署模型

定义损失函数和优化器

训练和验证模型

创建数据集

  1. from torchvision.datasets import VOCSegmentation
  2. from PIL import Image   
  3. from torchvision.transforms.functional import to_tensor, to_pil_image
  4.  
  5. class myVOCSegmentation(VOCSegmentation):
  6.     def __getitem__(self, index):
  7.         img = Image.open(self.images[index]).convert('RGB')
  8.         target = Image.open(self.masks[index])
  9.  
  10.         if self.transforms is not None:
  11.             augmented= self.transforms(image=np.array(img), mask=np.array(target))
  12.             img = augmented['image']
  13.             target = augmented['mask']                  
  14.             target[target>20]=0
  15.  
  16.         img= to_tensor(img)            
  17.         target= torch.from_numpy(target).type(torch.long)
  18.         return img, target
  19.  
  20. from albumentations import (
  21.     HorizontalFlip,
  22.     Compose,
  23.     Resize,
  24.     Normalize)
  25.  
  26. mean = [0.485, 0.456, 0.406] 
  27. std = [0.229, 0.224, 0.225]
  28. h,w=520,520
  29.  
  30. transform_train = Compose([ Resize(h,w),
  31.                 HorizontalFlip(p=0.5), 
  32.                 Normalize(mean=mean,std=std)])
  33.  
  34. transform_val = Compose( [ Resize(h,w),
  35.                           Normalize(mean=mean,std=std)])            
  36.  
  37. path2data="./data/"    
  38. train_ds=myVOCSegmentation(path2data, 
  39.                 year='2012', 
  40.                 image_set='train', 
  41.                 download=False, 
  42.                 transforms=transform_train) 
  43. print(len(train_ds)) 
  44.  
  45.  
  46. val_ds=myVOCSegmentation(path2data, 
  47.                 year='2012', 
  48.                 image_set='val', 
  49.                 download=False, 
  50.                 transforms=transform_val)
  51. print(len(val_ds)) 
  52.  
  1. import torch
  2. import numpy as np
  3. from skimage.segmentation import mark_boundaries
  4. import matplotlib.pylab as plt
  5. %matplotlib inline
  6. np.random.seed(0)
  7. num_classes=21
  8. COLORS = np.random.randint(0, 2, size=(num_classes+1, 3),dtype="uint8")
  9.  
  10. def show_img_target(img, target):
  11.     if torch.is_tensor(img):
  12.         img=to_pil_image(img)
  13.         target=target.numpy()
  14.     for ll in range(num_classes):
  15.         mask=(target==ll)
  16.         img=mark_boundaries(np.array(img) , 
  17.                             mask,
  18.                             outline_color=COLORS[ll],
  19.                             color=COLORS[ll])
  20.     plt.imshow(img)
  21.  
  22.  
  23. def re_normalize (x, mean = mean, std= std):
  24.     x_r= x.clone()
  25.     for c, (mean_c, std_c) in enumerate(zip(mean, std)):
  26.         x_r [c] *= std_c
  27.         x_r [c] += mean_c
  28.     return x_r

 展示训练数据集示例图像

  1. img, mask = train_ds[10]
  2. print(img.shape, img.type(),torch.max(img))
  3. print(mask.shape, mask.type(),torch.max(mask))
  4.  
  5. plt.figure(figsize=(20,20))
  6.  
  7. img_r= re_normalize(img)
  8. plt.subplot(1, 3, 1) 
  9. plt.imshow(to_pil_image(img_r))
  10.  
  11. plt.subplot(1, 3, 2) 
  12. plt.imshow(mask)
  13.  
  14. plt.subplot(1, 3, 3) 
  15. show_img_target(img_r, mask)

展示验证数据集示例图像

  1. img, mask = val_ds[10]
  2. print(img.shape, img.type(),torch.max(img))
  3. print(mask.shape, mask.type(),torch.max(mask))
  4.  
  5. plt.figure(figsize=(20,20))
  6.  
  7. img_r= re_normalize(img)
  8. plt.subplot(1, 3, 1) 
  9. plt.imshow(to_pil_image(img_r))
  10.  
  11. plt.subplot(1, 3, 2) 
  12. plt.imshow(mask)
  13.  
  14. plt.subplot(1, 3, 3) 
  15. show_img_target(img_r, mask)

创建数据加载器

 通过torch.utils.data针对训练和验证集分别创建Dataloader，打印示例观察效果

  1. from torch.utils.data import DataLoader
  2. train_dl = DataLoader(train_ds, batch_size=4, shuffle=True)
  3. val_dl = DataLoader(val_ds, batch_size=8, shuffle=False) 
  4.  
  5. for img_b, mask_b in train_dl:
  6.     print(img_b.shape,img_b.dtype)
  7.     print(mask_b.shape, mask_b.dtype)
  8.     break
  9.  
  10. for img_b, mask_b in val_dl:
  11.     print(img_b.shape,img_b.dtype)
  12.     print(mask_b.shape, mask_b.dtype)
  13.     break

创建模型

创建并打印deeplab_resnet模型结构，使用预训练权重

  1. from torchvision.models.segmentation import deeplabv3_resnet101
  2. import torch
  3.  
  4. model=deeplabv3_resnet101(pretrained=True, num_classes=21)
  5. device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
  6. model=model.to(device)
  7. print(model)

部署模型

在验证数据集的数据批次上部署模型观察效果 

  1. from torch import nn
  2.  
  3. model.eval()
  4. with torch.no_grad():
  5.     for xb, yb in val_dl:
  6.         yb_pred = model(xb.to(device))
  7.         yb_pred = yb_pred["out"].cpu()
  8.         print(yb_pred.shape)    
  9.         yb_pred = torch.argmax(yb_pred,axis=1)
  10.         break
  11. print(yb_pred.shape)
  12.  
  13. plt.figure(figsize=(20,20))
  14.  
  15. n=2
  16. img, mask= xb[n], yb_pred[n]
  17. img_r= re_normalize(img)
  18. plt.subplot(1, 3, 1) 
  19. plt.imshow(to_pil_image(img_r))
  20.  
  21. plt.subplot(1, 3, 2) 
  22. plt.imshow(mask)
  23.  
  24. plt.subplot(1, 3, 3) 
  25. show_img_target(img_r, mask)

可见勾勒对象方面效果很好 

定义损失函数和优化器

  1. from torch import nn
  2. criterion = nn.CrossEntropyLoss(reduction="sum")
  1. from torch import optim
  2. opt = optim.Adam(model.parameters(), lr=1e-6)
  3.  
  4. def loss_batch(loss_func, output, target, opt=None):   
  5.     loss = loss_func(output, target)
  6.     
  7.     if opt is not None:
  8.         opt.zero_grad()
  9.         loss.backward()
  10.         opt.step()
  11.  
  12.     return loss.item(), None
  13.  
  14. from torch.optim.lr_scheduler import ReduceLROnPlateau
  15. lr_scheduler = ReduceLROnPlateau(opt, mode='min',factor=0.5, patience=20,verbose=1)
  16.  
  17. def get_lr(opt):
  18.     for param_group in opt.param_groups:
  19.         return param_group['lr']
  20.  
  21. current_lr=get_lr(opt)
  22. print('current lr={}'.format(current_lr))

训练和验证模型

  1. def loss_epoch(model,loss_func,dataset_dl,sanity_check=False,opt=None):
  2.     running_loss=0.0
  3.     len_data=len(dataset_dl.dataset)
  4.  
  5.     for xb, yb in dataset_dl:
  6.         xb=xb.to(device)
  7.         yb=yb.to(device)
  8.         
  9.         output=model(xb)["out"]
  10.         loss_b, _ = loss_batch(loss_func, output, yb, opt)
  11.         running_loss += loss_b
  12.         
  13.         if sanity_check is True:
  14.             break
  15.     
  16.     loss=running_loss/float(len_data)
  17.     return loss, None
  18.  
  19. import copy
  20. def train_val(model, params):
  21.     num_epochs=params["num_epochs"]
  22.     loss_func=params["loss_func"]
  23.     opt=params["optimizer"]
  24.     train_dl=params["train_dl"]
  25.     val_dl=params["val_dl"]
  26.     sanity_check=params["sanity_check"]
  27.     lr_scheduler=params["lr_scheduler"]
  28.     path2weights=params["path2weights"]
  29.     
  30.     loss_history={
  31.         "train": [],
  32.         "val": []}
  33.     
  34.     metric_history={
  35.         "train": [],
  36.         "val": []}    
  37.     
  38.     
  39.     best_model_wts = copy.deepcopy(model.state_dict())
  40.     best_loss=float('inf')    
  41.     
  42.     for epoch in range(num_epochs):
  43.         current_lr=get_lr(opt)
  44.         print('Epoch {}/{}, current lr={}'.format(epoch, num_epochs - 1, current_lr))   
  45.  
  46.         model.train()
  47.         train_loss, train_metric=loss_epoch(model,loss_func,train_dl,sanity_check,opt)
  48.  
  49.         loss_history["train"].append(train_loss)
  50.         metric_history["train"].append(train_metric)
  51.         
  52.         model.eval()
  53.         with torch.no_grad():
  54.             val_loss, val_metric=loss_epoch(model,loss_func,val_dl,sanity_check)
  55.        
  56.         loss_history["val"].append(val_loss)
  57.         metric_history["val"].append(val_metric)   
  58.         
  59.         if val_loss < best_loss:
  60.             best_loss = val_loss
  61.             best_model_wts = copy.deepcopy(model.state_dict())
  62.             
  63.             torch.save(model.state_dict(), path2weights)
  64.             print("Copied best model weights!")
  65.             
  66.         lr_scheduler.step(val_loss)
  67.         if current_lr != get_lr(opt):
  68.             print("Loading best model weights!")
  69.             model.load_state_dict(best_model_wts) 
  70.             
  71.         print("train loss: %.6f" %(train_loss))
  72.         print("val loss: %.6f" %(val_loss))
  73.         print("-"*10) 
  74.     model.load_state_dict(best_model_wts)
  75.     return model, loss_history, metric_history        
  1. import os
  2. opt = optim.Adam(model.parameters(), lr=1e-6)
  3. lr_scheduler = ReduceLROnPlateau(opt, mode='min',factor=0.5, patience=20,verbose=1)
  4.  
  5. path2models= "./models/"
  6. if not os.path.exists(path2models):
  7.         os.mkdir(path2models)
  8.  
  9. params_train={
  10.     "num_epochs": 10,
  11.     "optimizer": opt,
  12.     "loss_func": criterion,
  13.     "train_dl": train_dl,
  14.     "val_dl": val_dl,
  15.     "sanity_check": True,
  16.     "lr_scheduler": lr_scheduler,
  17.     "path2weights": path2models+"sanity_weights.pt",
  18. }
  19.  
  20. model, loss_hist, _ = train_val(model, params_train)

绘制了训练和验证损失曲线 

  1. num_epochs=params_train["num_epochs"]
  2.  
  3. plt.title("Train-Val Loss")
  4. plt.plot(range(1,num_epochs+1),loss_hist["train"],label="train")
  5. plt.plot(range(1,num_epochs+1),loss_hist["val"],label="val")
  6. plt.ylabel("Loss")
  7. plt.xlabel("Training Epochs")
  8. plt.legend()
  9. plt.show()

声明：本文内容由网友自发贡献，转载请注明出处：【wpsshop】
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号