data_agu.py

import cv2
import numpy as np
import os
from PIL import Image, ImageEnhance
from torch.utils.data import Dataset
import pandas as pd
import random


def resize(image,gt,insize,outsize):
    x = np.random.randint(-512, 512)
    y = np.random.randint(-512, 512)
    if x < 0:
        if y < 0:
            image = image[0:x, 0:y, :]
            gt = gt[0:x, 0:y]
        else:
            image = image[0:x, y:insize, :]
            gt = gt[0:x, y:insize]
    else:
        if y < 0:
            image = image[x:insize, 0:y, :]
            gt = gt[x:insize, 0:y]
        else:
            image = image[x:insize, y:insize, :]
            gt = gt[x:insize, y:insize]
    image = cv2.resize(image, (outsize, outsize), interpolation=cv2.INTER_LINEAR)
    gt = cv2.resize(gt, (outsize, outsize), interpolation=cv2.INTER_NEAREST)
    return image,gt


def randomHueSaturationValue(image, hue_shift_limit=(-180, 180),
                             sat_shift_limit=(-255, 255),
                             val_shift_limit=(-255, 255), u=0.5):
    if np.random.random() < u:
        image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
        h, s, v = cv2.split(image)
        hue_shift = np.random.randint(hue_shift_limit[0], hue_shift_limit[1] + 1)
        hue_shift = np.uint8(hue_shift)
        h += hue_shift
        sat_shift = np.random.uniform(sat_shift_limit[0], sat_shift_limit[1])
        s = cv2.add(s, sat_shift)
        val_shift = np.random.uniform(val_shift_limit[0], val_shift_limit[1])
        v = cv2.add(v, val_shift)
        image = cv2.merge((h, s, v))
        # image = cv2.merge((s, v))
        image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
    return image


def randomShiftScaleRotate(image, mask,
                           shift_limit=(-0.0, 0.0),
                           scale_limit=(-0.0, 0.0),
                           rotate_limit=(-0.0, 0.0),
                           aspect_limit=(-0.0, 0.0),
                           borderMode=cv2.BORDER_CONSTANT, u=0.5):
    if np.random.random() < u:
        height, width, channel = image.shape

        angle = np.random.uniform(rotate_limit[0], rotate_limit[1])
        scale = np.random.uniform(1 + scale_limit[0], 1 + scale_limit[1])
        aspect = np.random.uniform(1 + aspect_limit[0], 1 + aspect_limit[1])
        sx = scale * aspect / (aspect ** 0.5)
        sy = scale / (aspect ** 0.5)
        dx = round(np.random.uniform(shift_limit[0], shift_limit[1]) * width)
        dy = round(np.random.uniform(shift_limit[0], shift_limit[1]) * height)

        cc = np.math.cos(angle / 180 * np.math.pi) * sx
        ss = np.math.sin(angle / 180 * np.math.pi) * sy
        rotate_matrix = np.array([[cc, -ss], [ss, cc]])

        box0 = np.array([[0, 0], [width, 0], [width, height], [0, height], ])
        box1 = box0 - np.array([width / 2, height / 2])
        box1 = np.dot(box1, rotate_matrix.T) + np.array([width / 2 + dx, height / 2 + dy])

        box0 = box0.astype(np.float32)
        box1 = box1.astype(np.float32)
        mat = cv2.getPerspectiveTransform(box0, box1)
        image = cv2.warpPerspective(image, mat, (width, height), flags=cv2.INTER_LINEAR, borderMode=borderMode,
                                    borderValue=(
                                        0, 0,
                                        0,))
        mask = cv2.warpPerspective(mask, mat, (width, height), flags=cv2.INTER_LINEAR, borderMode=borderMode,
                                   borderValue=(
                                       0, 0,
                                       0,))

    return image, mask


def randomHorizontalFlip(image, mask, u=0.5):
    if np.random.random() < u:
        image = cv2.flip(image, 1)
        mask = cv2.flip(mask, 1)

    return image, mask


def randomVerticleFlip(image, mask, u=0.5):
    if np.random.random() < u:
        image = cv2.flip(image, 0)
        mask = cv2.flip(mask, 0)

    return image, mask


def randomRotate90(image, mask, u=0.5):
    if np.random.random() < u:
        angle = np.random.randint(1,4)
        for i in range(angle):
            image = np.rot90(image)
            mask = np.rot90(mask)
        return image, mask
    return image, mask
# def randomRotate90(image, mask, u=0.5):
#     if np.random.random() < u:
#         image = np.rot90(image)
#         mask = np.rot90(mask)
#
#     return image, mask


def grade(img):
    x = cv2.Sobel(img, cv2.CV_32F, 1, 0, ksize=1)
    y = cv2.Sobel(img, cv2.CV_32F, 0, 1, ksize=1)
    absX = cv2.convertScaleAbs(x)
    absY = cv2.convertScaleAbs(y)
    dst = cv2.addWeighted(absX, 0.5, absY, 0.5, 0)
    mi = np.min(dst)
    ma = np.max(dst)
    res = (dst - mi) / (0.000000001 + (ma - mi))
    res[np.isnan(res)] = 0
    return res


class Mydataset(Dataset):
    def __init__(self, path, augment=False, transform=None, target_transform=None):

        self.aug = augment
        self.file_path = os.path.dirname(path)
        data = pd.read_csv(path)  # 获取csv表中的数据
        imgs = []
        for i in range(len(data)):
            imgs.append((data.iloc[i, 0], data.iloc[i, 1]))
        self.imgs = imgs
        self.transform = transform
        self.target_transform = target_transform

    def __getitem__(self, item):
        if self.aug == False:
            fn, lab = self.imgs[item]
            # fn = os.path.join(self.file_path, "image_A/" + fn)
            # label = os.path.join(self.file_path, "image_A/" + lab)
            fn = os.path.join(self.file_path, "images/" + fn)
            label = os.path.join(self.file_path, "labels/" + lab)

            bgr_img = cv2.imread(fn, -1)
            rgb_img = bgr_img[..., ::-1]  # bgr2rgb
            img = Image.fromarray(rgb_img)

            # gray = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2GRAY)
            # grad = (255 * grade(gray)).astype(np.uint8)
            #
            # # img = Image.open(fn).convert('RGB')
            # img = cv2.merge([rgb_img, grad])
            # img = Image.fromarray(img, mode="CMYK")
            if self.transform is not None:
                img = self.transform(img)

            gt = cv2.imread(label, -1)/255
            return img, gt, lab


        else:
            # 进行数据增强
            fn, lab = self.imgs[item]
            # train with data.cvs
            fn = os.path.join(self.file_path, "images/" + fn)
            label = os.path.join(self.file_path, "labels/" + lab)

            gt = cv2.imread(label, -1)/255
            image = cv2.imread(fn, -1)

            # image,gt = resize(image, gt, 1500, 512)
            image = randomHueSaturationValue(image,
                                             hue_shift_limit=(-30, 30),
                                             sat_shift_limit=(-25, 25),
                                             val_shift_limit=(-15, 15))

            image, gt = randomShiftScaleRotate(image, gt,
                                               shift_limit=(-0.1, 0.1),
                                               scale_limit=(-0.1, 0.1),
                                               aspect_limit=(-0.1, 0.1),
                                               rotate_limit=(-5, 5))

            image, gt = randomHorizontalFlip(image, gt)
            image, gt = randomVerticleFlip(image, gt)
            image, gt = randomRotate90(image, gt)

            # rgb_img = image[..., ::-1]  # bgr2rgb
            batch = [[0, 1, 2], [0, 2, 1], [1, 0, 2], [1, 2, 0], [2, 0, 1], [2, 1, 0]]
            sort = batch[np.random.randint(0, 6)]
            rgb_img = cv2.merge([image[:, :, sort[0]], image[:, :, sort[1]], image[:, :, sort[2]]])

            img = Image.fromarray(rgb_img)

            # gray = cv2.cvtColor(rgb_img, cv2.COLOR_RGB2GRAY)
            # grad = (255 * grade(gray)).astype(np.uint8)
            # img = cv2.merge([rgb_img, grad])
            # img = Image.fromarray(img, mode="CMYK")
            if self.transform is not None:
                img = self.transform(img.copy())
            return img, gt.copy(), lab

    def __len__(self):
        return len(self.imgs)