Unet網絡實現葉子病蟲害圖像分割

引言

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

在傳統的病蟲害圖像分割算法中應用最多的是基于統計模式識別、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

完整代碼：

鏈接：https://pan.baidu.com/s/11quQOLw7uIP-JJ8LPBIIuQ

提取碼：dyrt

作者簡介：

李秋鍵，CSDN博客專家，CSDN達人課作者。碩士在讀于中國礦業大學，開發有taptap競賽獲獎等。