语义分割功能代码

图像数据扩充

# import tensorflow as tf
import Augmentor

# 读入原始图像存储路径   也要是png格式的,路径需要是英文的。
p = Augmentor.Pipeline("C:\\Users\\ZWSZJC\\Desktop\\leafimage\\timagepng")
# 读取掩码文件存储路径
p.ground_truth("C:\\Users\\ZWSZJC\\Desktop\\leafimage\\tlabel")
# 默认会生成输出的路径，在Augmentor.Pipeline后的路径下自动生成output文件夹，不需要自己新建。
# output文件夹中包含图像和标注图扩充后的图片，都是png格式的

# 图像旋转：按照概率0.8执行，最大左旋角度15，最大右旋角度15
p.rotate(probability=0.8, max_left_rotation=15, max_right_rotation=15)

# 图像左右互换： 按照概率1执行
p.flip_left_right(probability=1)

# 图像放大缩小： 按照概率0.8执行，面积为原始图0.85倍
p.zoom_random(probability=0.8, percentage_area=0.85)

# 按概率随即上下翻转
p.flip_top_bottom(probability=1)

# 逆时针随机旋转90度（随机概率可自行设定），在指定所有操作之后再运行，否则只会进行旋转90度这个操作
p.rotate90(probability=1)

# 图片顺时针随机旋转90度（随机概率可自行设定）
p.rotate270(probability=1)

# 透视形变-垂直方向形变：magnitude取（0,1），指的是形变程度
# p.skew_tilt(probability=1,magnitude=1)

# 透视形变-斜四角形变形变：magnitude取（0,1），指的是形变程度
# p.skew_corner(probability=1,magnitude=1)

# 弹性扭曲，类似区域扭曲的感觉
# p.random_distortion(probability=1,grid_height=5,grid_width=16,magnitude=8)

# 错切变换
# p.shear(probability=1,max_shear_left=15,max_shear_right=15)

# 随机区域擦除
# p.random_erasing(probability=1,rectangle_area=0.5)

# 最终扩充的数据样本数
p.sample(800)

移除图片背景

通过对比图片和标注，将背景的像素点置为0

import os
from PIL import Image
import numpy as np

def removeBG(imagePath, maskPath, savePath):
	img = np.array(Image.open(imagePath))
	mask = np.array(Image.open(maskPath))

	for x in range(img.shape[0]):   # 图片的高
		for y in range(img.shape[1]):   # 图片的宽
			if mask[x,y] == 0:
				img[x,y] = [0, 0, 0]

	image = Image.fromarray(img)
	image.save(savePath)


split = ['train', 'val', 'test']
_base_dir = "/Users/trinasolar/Desktop/深度学习资源/个人数据集/Disease"
_image_dir = os.path.join(_base_dir, 'Images')
_cat_dir = os.path.join(_base_dir, 'SegmentationClass')
_splits_dir = os.path.join(_base_dir, 'ImageSets')
_new_img_dir = os.path.join(_base_dir, 'ImagesNoBG')
if not os.path.exists(_new_img_dir):
	os.makedirs(_new_img_dir)

for splt in split:
	with open(os.path.join(os.path.join(_splits_dir, splt + '.txt')), "r") as f:
	    lines = f.read().splitlines()

	for ii, line in enumerate(lines):
	    _image = os.path.join(_image_dir, line)
	    _cat = os.path.join(_cat_dir, line)
	    _newImg = os.path.join(_new_img_dir, line)
	    removeBG(_image, _cat, _newImg)

获取图片所有像素值

from PIL import Image 
import numpy as np

image = Image.open("/Users/trinasolar/Downloads/SegmentationClass/001.png")
image_arr = np.array(image) # 转化成numpy数组
print(np.unique(image_arr))

预测与原图合并

import cv2
from PIL import Image
from matplotlib import gridspec
import os
import numpy as np
import matplotlib.pyplot as plt

class Transformer(object):

    def __init__(self, rootPath, std=[0.229, 0.224, 0.225], mean=[0.485, 0.456, 0.406]):
        self.rootPath = rootPath
        if not os.path.exists(self.rootPath):
            os.makedirs(self.rootPath)
        self.std = std
        self.mean = mean

    # 图片逆归一化
    def UnNormalize(self, x):
        x[0] = x[0] * self.std[0] + self.mean[0]
        x[1] = x[1] * self.std[1] + self.mean[1]
        x[2] = x[2].mul(self.std[2]) + self.mean[2]
        img = x.mul(255).byte()
        img = img.numpy().transpose((1, 2, 0))
        return img

    # 合并图片
    def maskToImage(self, pred, image, index, isPred=True):
        if isPred:
            resultPath = os.path.join(self.rootPath, 'pred')
        else:
            resultPath = os.path.join(self.rootPath, 'mask')
        if not os.path.exists(resultPath):
            os.makedirs(resultPath)

        img_path = os.path.join(resultPath, 'result_%d.png' % (index))

        # 预测
        predResult = np.reshape(pred[0], (224, 224))
        # result = result.astype(np.float32) * 255.
        predResult = np.clip(predResult, 0, 255).astype('uint8')
        predMask = Image.fromarray(predResult)

        # 图片反归一化
        x = image[0]
        img = Image.fromarray(self.UnNormalize(x))

        self.vis_segmentation(img, predMask, img_path)

    def label_to_color_image(self, label):
        colormap = np.array([[0, 0, 0], [0, 128, 0], [255, 0, 0]])

        if np.max(label) >= len(colormap):
            raise ValueError('label value too large.')

        return colormap[label]

    def vis_segmentation(self, image, seg_map, imagePath):
        """
        输入图片和分割 mask 的可视化.
        """
        plt.figure(figsize=(15, 5))
        grid_spec = gridspec.GridSpec(1, 4, width_ratios=[6, 6, 6, 1])

        LABEL_NAMES = np.asarray(['background', 'leaf', 'disease'])
        FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
        FULL_COLOR_MAP = self.label_to_color_image(FULL_LABEL_MAP)

        plt.subplot(grid_spec[0])
        plt.imshow(image)
        plt.axis('off')
        plt.title('input image')

        plt.subplot(grid_spec[1])
        seg_image = self.label_to_color_image(seg_map).astype(np.uint8)
        plt.imshow(seg_image)
        plt.axis('off')
        plt.title('segmentation map')

        plt.subplot(grid_spec[2])
        plt.imshow(image)
        plt.imshow(seg_image, alpha=0.7)
        plt.axis('off')
        plt.title('segmentation overlay')

        unique_labels = np.unique(seg_map)
        ax = plt.subplot(grid_spec[3])
        plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation='nearest')
        ax.yaxis.tick_right()
        plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
        plt.xticks([], [])
        ax.tick_params(width=0.0)
        plt.grid('off')
        plt.savefig(imagePath)
        # plt.show()
        plt.close('all')

    def vis_segmentation2(self, image, seg_map, imagePath):
        """
        输入图片和分割 mask 的统一可视化.
        """
        seg_image = self.label_to_color_image(seg_map).astype(np.uint8)
        plt.figure()
        plt.imshow(image)
        plt.imshow(seg_image, alpha=0.6)
        plt.axis('off')
        plt.savefig(imagePath)
        # plt.show()
        plt.close('all')

图像预处理函数(pytorch)

import torch
import random
import numpy as np

from PIL import Image, ImageOps, ImageFilter

class Normalize(object):
    """Normalize a tensor image with mean and standard deviation.
    Args:
        mean (tuple): means for each channel.
        std (tuple): standard deviations for each channel.
    """
    def __init__(self, mean=(0., 0., 0.), std=(1., 1., 1.)):
        self.mean = mean
        self.std = std

    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        img = np.array(img).astype(np.float32)
        mask = np.array(mask).astype(np.float32)
        img /= 255.0
        img -= self.mean
        img /= self.std

        return {'image': img,
                'label': mask}


class ToTensor(object):
    """Convert ndarrays in sample to Tensors."""

    def __call__(self, sample):
        # swap color axis because
        # numpy image: H x W x C
        # torch image: C X H X W
        img = sample['image']
        mask = sample['label']
        img = np.array(img).astype(np.float32).transpose((2, 0, 1))
        mask = np.array(mask).astype(np.float32)

        img = torch.from_numpy(img).float()
        mask = torch.from_numpy(mask).float()

        return {'image': img,
                'label': mask}


class RandomHorizontalFlip(object):
    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        if random.random() < 0.5:
            img = img.transpose(Image.FLIP_LEFT_RIGHT)
            mask = mask.transpose(Image.FLIP_LEFT_RIGHT)

        return {'image': img,
                'label': mask}


class RandomRotate(object):
    def __init__(self, degree):
        self.degree = degree

    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        rotate_degree = random.uniform(-1*self.degree, self.degree)
        img = img.rotate(rotate_degree, Image.BILINEAR)
        mask = mask.rotate(rotate_degree, Image.NEAREST)

        return {'image': img,
                'label': mask}


class RandomGaussianBlur(object):
    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        if random.random() < 0.5:
            img = img.filter(ImageFilter.GaussianBlur(
                radius=random.random()))

        return {'image': img,
                'label': mask}


class RandomScaleCrop(object):
    def __init__(self, base_size, crop_size, fill=0):
        self.base_size = base_size
        self.crop_size = crop_size
        self.fill = fill

    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        # random scale (short edge)
        short_size = random.randint(int(self.base_size * 0.5), int(self.base_size * 2.0))
        w, h = img.size
        if h > w:
            ow = short_size
            oh = int(1.0 * h * ow / w)
        else:
            oh = short_size
            ow = int(1.0 * w * oh / h)
        img = img.resize((ow, oh), Image.BILINEAR)
        mask = mask.resize((ow, oh), Image.NEAREST)
        # pad crop
        if short_size < self.crop_size:
            padh = self.crop_size - oh if oh < self.crop_size else 0
            padw = self.crop_size - ow if ow < self.crop_size else 0
            img = ImageOps.expand(img, border=(0, 0, padw, padh), fill=0)
            mask = ImageOps.expand(mask, border=(0, 0, padw, padh), fill=self.fill)
        # random crop crop_size
        w, h = img.size
        x1 = random.randint(0, w - self.crop_size)
        y1 = random.randint(0, h - self.crop_size)
        img = img.crop((x1, y1, x1 + self.crop_size, y1 + self.crop_size))
        mask = mask.crop((x1, y1, x1 + self.crop_size, y1 + self.crop_size))

        return {'image': img,
                'label': mask}


class FixScaleCrop(object):
    def __init__(self, crop_size):
        self.crop_size = crop_size

    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']
        w, h = img.size
        if w > h:
            oh = self.crop_size
            ow = int(1.0 * w * oh / h)
        else:
            ow = self.crop_size
            oh = int(1.0 * h * ow / w)
        img = img.resize((ow, oh), Image.BILINEAR)
        mask = mask.resize((ow, oh), Image.NEAREST)
        # center crop
        w, h = img.size
        x1 = int(round((w - self.crop_size) / 2.))
        y1 = int(round((h - self.crop_size) / 2.))
        img = img.crop((x1, y1, x1 + self.crop_size, y1 + self.crop_size))
        mask = mask.crop((x1, y1, x1 + self.crop_size, y1 + self.crop_size))

        return {'image': img,
                'label': mask}

class FixedResize(object):
    def __init__(self, size):
        self.size = (size, size)  # size: (h, w)

    def __call__(self, sample):
        img = sample['image']
        mask = sample['label']

        assert img.size == mask.size

        img = img.resize(self.size, Image.BILINEAR)
        mask = mask.resize(self.size, Image.NEAREST)

        return {'image': img,
                'label': mask}

计算类别权重

import os
from tqdm import tqdm
import numpy as np
from dataloaders.dataPath import Path

def calculate_weigths_labels(dataset, dataloader, num_classes):
    # Create an instance from the data loader
    z = np.zeros((num_classes,))
    # Initialize tqdm
    tqdm_batch = tqdm(dataloader)
    print('Calculating classes weights')
    for sample in tqdm_batch:
        y = sample['label']
        y = y.detach().cpu().numpy()
        mask = (y >= 0) & (y < num_classes)
        labels = y[mask].astype(np.uint8)
        count_l = np.bincount(labels, minlength=num_classes)
        z += count_l
    tqdm_batch.close()
    total_frequency = np.sum(z)
    class_weights = []
    for frequency in z:
        class_weight = 1 / (np.log(1.02 + (frequency / total_frequency)))
        class_weights.append(class_weight)
    ret = np.array(class_weights)
    classes_weights_path = os.path.join(Path.db_root_dir(dataset), dataset+'_classes_weights.npy')
    np.save(classes_weights_path, ret)

    return ret

生成数据集文档(txt)

def getAllImageNames():
    face_path = "E:\\myself\\papercode\\1.数据集\\1.个人数据集\\disease-nobg\\Images"

    fileTrain = open("E:\\myself\\papercode\\1.数据集\\1.个人数据集\\disease-nobg\\ImageSets\\train.txt", 'w')
    fileVal = open("E:\\myself\\papercode\\1.数据集\\1.个人数据集\\disease-nobg\\ImageSets\\val.txt", 'w')
    for dirname, dirnames, filenames in os.walk(face_path):
        list = filenames
        random.shuffle(list)
        total = len(list)
        trainNum = int(total * 0.9)
        train = list[:trainNum]
        val = list[trainNum:]

        print(len(train))
        print(len(val))

        for filename in train:
            fileTrain.write(filename)
            fileTrain.write("\n")

        for filename in val:
            fileVal.write(filename)
            fileVal.write("\n")

    fileVal.close()
    fileTrain.close()

检查标注文件是否出错

#!/usr/bin/python
# -*- coding: UTF-8 -*-
import argparse
import time
import os
from PIL import Image
from matplotlib import gridspec
import numpy as np
import matplotlib.pyplot as plt

def main():

    # 加载图片
    imageDir = '/Users/trinasolar/Desktop/深度学习资源/个人数据集/Disease/Images'
    maskDir = '/Users/trinasolar/Desktop/深度学习资源/个人数据集/Disease/SegmentationClass'
    resultDir = '/Users/trinasolar/Desktop/深度学习资源/个人数据集/Disease/mergeResult'
    if not os.path.exists(resultDir):
        os.makedirs(resultDir)

    onlyPredMaskPath = os.path.join(resultDir, 'onlymask')
    if not os.path.exists(onlyPredMaskPath):
        os.makedirs(onlyPredMaskPath)

    # assert 1==2, '1 不等于 2'
    for root, dirs, files in os.walk(maskDir):
        for f in files:
            imagePath = os.path.join(imageDir, f)
            maskPath = os.path.join(maskDir, f)
            assert os.path.exists(imagePath), '%s 图片不存在' % imagePath
            _img = Image.open(imagePath).convert('RGB')
            _target = Image.open(maskPath)

            resultPath = os.path.join(resultDir, f)

            # 将mask显示出来
            onlyMaskPath = os.path.join(onlyPredMaskPath, f)
            seg_mask = label_to_color_image(_target).astype(np.uint8)
            seg_image = Image.fromarray(seg_mask)
            seg_image.save(onlyMaskPath)

            vis_segmentation(_img, _target, resultPath)


def label_to_color_image(label):
    colormap = np.array([[0, 0, 0], [0, 128, 0], [255, 0, 0]])

    if np.max(label) >= len(colormap):
        raise ValueError('label value too large.')

    return colormap[label]

def vis_segmentation(image, seg_map, imagePath):
    """
    输入图片和分割 mask 的可视化.
    """
    plt.figure(figsize=(15, 5))
    grid_spec = gridspec.GridSpec(1, 4, width_ratios=[6, 6, 6, 1])

    LABEL_NAMES = np.asarray(['background', 'leaf', 'disease'])
    FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
    FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)

    plt.subplot(grid_spec[0])
    plt.imshow(image)
    plt.axis('off')
    plt.title('input image')

    plt.subplot(grid_spec[1])
    seg_image = label_to_color_image(seg_map).astype(np.uint8)
    plt.imshow(seg_image)
    plt.axis('off')
    plt.title('segmentation map')

    plt.subplot(grid_spec[2])
    plt.imshow(image)
    plt.imshow(seg_image, alpha=0.3)
    plt.axis('off')
    plt.title('segmentation overlay')

    unique_labels = np.unique(seg_map)
    ax = plt.subplot(grid_spec[3])
    plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation='nearest')
    ax.yaxis.tick_right()
    plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
    plt.xticks([], [])
    ax.tick_params(width=0.0)
    plt.grid('off')
    plt.savefig(imagePath)
    # 手动清除内存中的图像
    plt.close('all')
    # plt.show()


if __name__ == "__main__":
    main()