Unet網路實作葉子病蟲害影像分割-有解無憂

作者|李秋鍵

出品|AI科技大本營(ID:rgznai100)

引言

智能化農業作為人工智能應用的重要領域，對較高的影像處理能力要求較高，其中影像分割作為影像處理方法在其中起著重要作用，影像分割是影像分析的關鍵步驟, 在復雜的自然背景下進行影像分割, 難度較大，

在傳統的病蟲害影像分割演算法中應用最多的是基于統計模式識別、K-均值聚類、模糊C-均值聚類、Otsu、levelset、EM等演算法，Gui等人提出了一種基于復雜背景的彩色葉片影像的大豆葉病檢測系統，該系統根據K均值演算法和結合經驗閾值，基于顯著性映射從大豆葉片病害影像中分割出顯著區域.Kaur等人提出了一種基于K均值聚類演算法的葡萄病害葉片影像分割方法.Chaudhary等人提出了一種基于影像處理技術的植物葉片病斑分割演算法，比較了CIELAB、HSI和YCbCr顏色空間在病害檢測程序中的效果.Mohammad等人比較了閾值法、分水嶺法、邊緣檢測法、模糊C-均值聚類法和K-均值聚類法在植物病害檢測中的應用，指出它們適合于任務重疊聚類，

而Unet網路作為影像分割的經典深度學習演算法，在植物葉片病蟲害區域分割中也起著重要作用，故本專案通過使用python語言搭建Unet影像分割技術實作對植物病蟲害區域的分割，其簡單訓練幾輪后的模型效果如下圖可見：

1.基本介紹

1.1 環境要求

本次環境使用的是python3.6.5+windows平臺，主要用的庫有：

opencv模塊，在計算機視覺專案的開發中，opencv作為較大眾的開源庫，擁有了豐富的常用影像處理函式庫，采用C/C++語言撰寫，可以運行在Linux/Windows/Mac等作業系統上，能夠快速的實作一些影像處理和識別的任務，

numpy模塊，numpy系統是Python的一種開源的數值計算擴展，這種工具可用來存盤和處理大型矩陣，比Python自身的嵌套串列結構(nested list structure)要高效得多(該結構也可以用來表示矩陣(matrix))，

pytorch模塊，pytorch是一個基于Torch的Python開源機器學習庫，用于自然語言處理等應用程式，它主要由Facebookd的人工智能小組開發，不僅能夠實作強大的GPU加速，同時還支持動態神經網路，這一點是現在很多主流框架如TensorFlow都不支持的，PyTorch提供了兩個高級功能：1.具有強大的GPU加速的張量計算（如Numpy） 2.包含自動求導系統的深度神經網路除了Facebook之外，Twitter、GMU和Salesforce等機構都采用了pytorch，

visdom模塊，visdom一個靈活的可視化工具，可用來對于實時，富資料的創建，組織和共享，支持Torch和numpy還有pytorch，visdom 可以實作遠程資料的可視化，對科學實驗有很大幫助，我們可以遠程的發送圖片和資料，并進行在ui界面顯示出來，檢查實驗結果，或者debug，

1.2 Unet模型介紹

U-Net 網路模型是在 2015 年由 Ronneberger 等人提出的，U-Net 基于全卷積網路 FCN，其網路結構與 FCN 相似，都采用了編碼器和解碼器，以及跳躍連接的拓撲結構，能夠實作在少量訓練影像上進行更加精準的分割，但是 U-Net 與 FCN 不同之處在于 U-Net 網路是左右對稱的，其左側是用于捕獲背景關系資訊的收縮路徑，網路右側是用于精確定位的擴張路徑，與收縮路徑對稱，以恢復影像大小，編碼器對應層的輸出特征圖經復制、裁減后與反卷積后的特征圖通過跳躍連接進行特征融合，然后作為下一層的輸入，繼續上采樣，U-Net 網路在上采樣程序中有大量的特征通道，這使得網路能夠將背景關系資訊傳到更高解析度的層，

2.模型搭建

2.1資料集準備

首先我們使用labelme工具對需要準備好的資料集進行處理標注，生成對應的json檔案放置同一目錄即可，其中labelme可以實作對影像進行多邊形，矩形，圓形，多段線，線段，點形式的標注（可用于目標檢測，影像分割，等任務），處理格式如下：

2.2模型創建

U-Net 網路模型結構主要包括編碼器、解碼器和跳躍連接部分，編碼器用于抽象特征、提取資訊，解碼器部分使得影像逐步恢復原始尺寸，而跳躍連接則將不同層次的特征進行了融合，在這里我們使用segmentation_models_pytorch庫實作對unet的直接呼叫

其中UNet編解碼器初始化代碼如下：

def __init__(
        self,
        encoder_name: str = "resnet34",
        encoder_depth: int = 5,
        encoder_weights: Optional[str] = "imagenet",
        decoder_use_batchnorm: bool = True,
        decoder_channels: List[int] = (256, 128, 64, 32, 16),
        decoder_attention_type: Optional[str] = None,
        in_channels: int = 3,
        classes: int = 1,
        activation: Optional[Union[str, callable]] = None,
        aux_params: Optional[dict] = None,
    ):
        super().__init__()
        self.encoder = get_encoder(
            encoder_name,
            in_channels=in_channels,
            depth=encoder_depth,
            weights=encoder_weights,
        )
        self.decoder = UnetDecoder(
            encoder_channels=self.encoder.out_channels,
            decoder_channels=decoder_channels,
            n_blocks=encoder_depth,
            use_batchnorm=decoder_use_batchnorm,
            center=True if encoder_name.startswith("vgg") else False,
            attention_type=decoder_attention_type,
        )
        self.segmentation_head = SegmentationHead(
            in_channels=decoder_channels[-1],
            out_channels=classes,
            activation=activation,
            kernel_size=3,
        )

2.3 模型訓練

設定模型基本引數，其中包括學習率，batchsize，迭代次數，損失值等初始化，UNet 網路及大部分使用 Relu 函式進行激活，能夠有效避免和糾正梯度消失問題，

def __init__(self, model):
    self.num_workers = 0
    self.batch_size = {"train": 8, "val":1}
    self.accumulation_steps = 32 // self.batch_size['train']
    self.lr = 1e-3
    self.num_epochs = 100
    self.best_loss = float("inf")
    self.best_dice = float(0)
    self.phases = ["train", "val"]
    self.device = torch.device("cuda:0")
    torch.set_default_tensor_type("torch.cuda.FloatTensor")
    self.net = model
    self.criterion = nn.BCEWithLogitsLoss()
    self.optimizer = optim.Adam(self.net.parameters(), lr=self.lr)
    self.scheduler = ReduceLROnPlateau(self.optimizer, mode="min", patience=4, verbose=True)
    self.net = self.net.to(self.device)
    cudnn.benchmark = True
    self.dataloaders = {
        phase: provider(
            image_path=image_path,
            phase=phase,
            mean=(0.485, 0.456, 0.406),
            std=(0.229, 0.224, 0.225),
            batch_size=self.batch_size[phase],
            num_workers=self.num_workers,
        )
        for phase in self.phases
    }
    self.losses = {phase: [] for phase in self.phases}
    self.iou_scores = {phase: [] for phase in self.phases}
    self.dice_scores = {phase: [] for phase in self.phases}

2.4模型評估

損失函式是神經網路尋找最優權重引數的指標，常用的損失函式有均方誤差、交叉熵損失函式等，U-Net 網路中使用了交叉熵損失函式在最后的特征圖上通過像素級的 soft-max進行計算，Loss可以通過標準二值交叉熵（binary cross-entropy）和 Dice 損失計算，這是評估影像分割成功與否的常用性能標準，交并比（IoU）是基于像素的評價標準，通常用于評估分割性能，這里考慮了目標矩陣與結果矩陣之間的不同像素比，這個度量與Dice計算也有關，

def __init__(self, phase, epoch):
    self.base_threshold = 0.5 # <<<<<<<<<<< here's the threshold
    self.base_dice_scores = []
    self.dice_neg_scores = []
    self.dice_pos_scores = []
    self.iou_scores = []
def update(self, targets, outputs):
    probs = torch.sigmoid(outputs)
    dice, dice_neg, dice_pos, _, _ = metric(probs, targets, self.base_threshold)
    self.base_dice_scores.append(dice)
    self.dice_pos_scores.append(dice_pos)
    self.dice_neg_scores.append(dice_neg)
    preds = predict(probs, self.base_threshold)
    iou = compute_iou_batch(preds, targets, classes=[1])
    self.iou_scores.append(iou)
def get_metrics(self):
    dice = np.mean(self.base_dice_scores)
    dice_neg = np.mean(self.dice_neg_scores)
    dice_pos = np.mean(self.dice_pos_scores)
    dices = [dice, dice_neg, dice_pos]
    iou = np.nanmean(self.iou_scores)
    return dices, iou

2.5可視化

通過設定visdom模塊中的provider建立訓練程序中的可視化，然后使用命令“python -m visdom.server”實作瀏覽器訪問訓練程序，

def provider(image_path, phase, mean=None, std=None, batch_size=4, num_workers=0,
):
    image_list = glob(os.path.join(image_path, "*"))
    train_idx, val_idx = train_test_split(range(len(image_list)), random_state=4396, test_size=0.1)
    index = train_idx if phase == "train" else val_idx
    dataset = CatDataset(index, image_list, phase=phase)
    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers, pin_memory=True, shuffle=True,)
    return dataloader