utils/data.py

import torch
import torch.nn as nn
import torch.utils.data as Data
import torchvision.transforms as transforms
from PIL import Image
import numpy as np
from PIL import Image, ImageOps, ImageFilter
import cv2
import os
import os.path as osp
import sys
import random
import scipy.io as scio
import numpy as np

__all__ = ['SirstDataset', 'SirstAugDataset', 'IRSTD1kDataset', 'NUDTDataset',
           'DriveDatasetTrain', 'DriveDatasetTest','CHASEDB1DatasetTrain', 'CHASEDB1DatasetTest','STAREDatasetTrain', 'STAREDatasetTest',
           'NEUDatasetTrain', 'NEUDatasetTest', 'SaliencyDatasetTrain', 'SaliencyDatasetTest']

class NUDTDataset(Data.Dataset):
    '''
    Return: Single channel
    '''

    def __init__(self, base_dir=r'D:/WFY/datasets/NUDT',
                 mode='train', base_size=256):
        assert mode in ['train', 'test']


        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'trainval')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError
        self.base_size = base_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('png'):
                self.names.append(filename)

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name)
        label_path = osp.join(self.data_dir, 'masks', name)

        img, mask = cv2.imread(img_path, 0), cv2.imread(label_path, 0)
        img = cv2.resize(img, [self.base_size, self.base_size], interpolation=cv2.INTER_LINEAR)
        mask = cv2.resize(mask, [self.base_size, self.base_size], interpolation=cv2.INTER_NEAREST)
        img = img.reshape(1, self.base_size, self.base_size) / 255.
        if np.max(mask) > 0:
            mask = mask.reshape(1, self.base_size, self.base_size) / np.max(mask)
        else:
            mask = mask.reshape(1, self.base_size, self.base_size)

        img = torch.from_numpy(img).type(torch.FloatTensor)
        mask = torch.from_numpy(mask).type(torch.FloatTensor)

        return img, mask

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

class IRSTD1kDataset(Data.Dataset):
    '''
    Return: Single channel
    '''

    def __init__(self, base_dir=r'D:/WFY/datasets/IRSTD-1k',
                 mode='train', base_size=256):
        assert mode in ['train', 'test']


        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'trainval')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError
        self.base_size = base_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('png'):
                self.names.append(filename)

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name)
        label_path = osp.join(self.data_dir, 'masks', name)

        img, mask = cv2.imread(img_path, 0), cv2.imread(label_path, 0)
        row, col = img.shape
        img = img.reshape(1, row, col) / 255.


        img, mask = cv2.imread(img_path, 0), cv2.imread(label_path, 0)

        img = cv2.resize(img, [self.base_size, self.base_size], interpolation=cv2.INTER_LINEAR)
        mask = cv2.resize(mask, [self.base_size, self.base_size], interpolation=cv2.INTER_NEAREST)
        img = img.reshape(1, self.base_size, self.base_size) / 255.
        if np.max(mask) > 0:
            mask = mask.reshape(1, self.base_size, self.base_size) / np.max(mask)
        else:
            mask = mask.reshape(1, self.base_size, self.base_size)

        img = torch.from_numpy(img).type(torch.FloatTensor)
        mask = torch.from_numpy(mask).type(torch.FloatTensor)

        return img, mask

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

class SirstDataset(Data.Dataset):
    def __init__(self, base_dir=r'D:\WFY\datasets\sirst',
                 mode='train', base_size=256):
        if mode == 'train':
            txtfile = 'trainval.txt'
        elif mode == 'val':
            txtfile = 'test1.txt'

        self.list_dir = osp.join(base_dir, 'img_idx', txtfile)
        self.imgs_dir = osp.join(base_dir, 'images')
        self.label_dir = osp.join(base_dir, 'masks')

        self.names = []
        with open(self.list_dir, 'r') as f:
            self.names += [line.strip() for line in f.readlines()]

        self.mode = mode
        self.base_size = base_size
        self.tranform = augmentation()

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.imgs_dir, name+'.png')
        label_path = osp.join(self.label_dir, name+'.png')

        img, mask = cv2.imread(img_path, 0), cv2.imread(label_path, 0)


        if self.mode == 'train':
            img, mask = self.tranform(img, mask)
            img = cv2.resize(img, [self.base_size, self.base_size], interpolation=cv2.INTER_LINEAR)
            mask = cv2.resize(mask, [self.base_size, self.base_size], interpolation=cv2.INTER_NEAREST)
            img = img.reshape(1, self.base_size, self.base_size) / 255.
            mask = mask.reshape(1, self.base_size, self.base_size) / np.max(mask)
            img = torch.from_numpy(img).type(torch.FloatTensor)
            mask = torch.from_numpy(mask).type(torch.FloatTensor)
            return img, mask

        elif self.mode == 'val':
            img = cv2.resize(img, [self.base_size, self.base_size], interpolation=cv2.INTER_LINEAR)
            mask = cv2.resize(mask, [self.base_size, self.base_size], interpolation=cv2.INTER_NEAREST)
            img = img.reshape(1, self.base_size, self.base_size) / 255.
            mask = mask.reshape(1, self.base_size, self.base_size) / np.max(mask)
            _, h, w = img.shape
            img = PadImg(img)
            mask = PadImg(mask)
            img = torch.from_numpy(img).type(torch.FloatTensor)
            mask = torch.from_numpy(mask).type(torch.FloatTensor)
            return img, mask



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

class SirstAugDataset(Data.Dataset):
    '''
    Return: Single channel
    '''
    def __init__(self, base_dir=r'/Users/tianfangzhang/Program/DATASETS/sirst_aug',
                 mode='train', base_size=256):
        assert mode in ['train', 'test']

        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'trainval')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError

        self.base_size = base_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('png'):
                self.names.append(filename)
        self.tranform = augmentation()

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name)
        label_path = osp.join(self.data_dir, 'masks', name)

        img, mask = cv2.imread(img_path, 0), cv2.imread(label_path, 0)
        img, mask = self.tranform(img, mask)
        img = img.reshape(1, self.base_size, self.base_size) / 255.
        if np.max(mask) > 0:
            mask = mask.reshape(1, self.base_size, self.base_size) / np.max(mask)
        else:
            mask = mask.reshape(1, self.base_size, self.base_size)

        img = torch.from_numpy(img).type(torch.FloatTensor)
        mask = torch.from_numpy(mask).type(torch.FloatTensor)
        return img, mask

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

class DriveDatasetTrain(Data.Dataset):
    '''
    Return: Single channel
    '''
    def __init__(self, base_dir=r'../datasets/DRIVE',
                 mode='train', base_size=512, patch_size = 256, img_norm_cfg = None):
        assert mode in ['train', 'test']

        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError


        if img_norm_cfg == None:
            self.img_norm_cfg = get_img_norm_cfg('DRIVE', '../datasets/DRIVE')
        else:
            self.img_norm_cfg = img_norm_cfg

        self.base_size = base_size
        self.patch_size = patch_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('png'):
                self.names.append(filename)
        self.tranform = augmentation()

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name)
        label_path = osp.join(self.data_dir, 'masks', name)

        img = Image.open(img_path).convert('I')
        mask = Image.open(label_path).convert('I')  # Convert to grayscale
        img = Normalized(np.array(img, dtype=np.float32), self.img_norm_cfg)
        mask = np.array(mask, dtype=np.float32) / 255.0

        if len(mask.shape) > 2:
            mask = mask[:, :, 0]

        img_patch, mask_patch = random_crop(img, mask, self.patch_size, pos_prob=0.5)
        img_patch, mask_patch = self.tranform(img_patch, mask_patch)
        img_patch, mask_patch = img_patch[np.newaxis, :], mask_patch[np.newaxis, :]
        img_patch = torch.from_numpy(np.ascontiguousarray(img_patch))
        mask_patch = torch.from_numpy(np.ascontiguousarray(mask_patch))
        return img_patch, mask_patch

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

class DriveDatasetTest(Data.Dataset):
    '''
    Return: Single channel
    '''
    def __init__(self, base_dir=r'../datasets/DRIVE',
                 mode='train', base_size=584, img_norm_cfg = None):
        assert mode in ['train', 'test']

        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError

        if img_norm_cfg == None:
            self.img_norm_cfg = get_img_norm_cfg('DRIVE', '../datasets/DRIVE')
        else:
            self.img_norm_cfg = img_norm_cfg

        self.base_size = base_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('png'):
                self.names.append(filename)
        self.tranform = augmentation()

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name)
        label_path = osp.join(self.data_dir, 'masks', name)

        img = Image.open(img_path).convert('I')
        mask = Image.open(label_path).convert('I')  # Convert to grayscale
        img = Normalized(np.array(img, dtype=np.float32), self.img_norm_cfg)
        mask = np.array(mask, dtype=np.float32) / 255.0
        if len(mask.shape) > 2:
            mask = mask[:, :, 0]

        h, w = img.shape
        img = PadImg(img)
        mask = PadImg(mask)

        img, mask = img[np.newaxis, :], mask[np.newaxis, :]

        img = torch.from_numpy(np.ascontiguousarray(img))
        mask = torch.from_numpy(np.ascontiguousarray(mask))
        return img, mask

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

class STAREDatasetTrain(Data.Dataset):

    def __init__(self, base_dir=r'../datasets/STARE',
                 mode='train', base_size=700, patch_size=256, img_norm_cfg=None):
        assert mode in ['train', 'test']

        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError

        if img_norm_cfg == None:
            self.img_norm_cfg = get_img_norm_cfg('STARE', r'../datasets/STARE')
        else:
            self.img_norm_cfg = img_norm_cfg
        self.base_size = base_size
        self.patch_size = patch_size
        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('ppm'):
                self.names.append(filename)
        self.tranform = augmentation()

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name)
        label_path = osp.join(self.data_dir, 'masks', name)

        img = Image.open(img_path).convert('I')  # It also can take use of .ppm file
        mask = Image.open(label_path).convert('I')  # Convert to grayscale
        img = np.array(img, dtype=np.float32)
        img = Normalized(img, self.img_norm_cfg)
        mask = np.array(mask, dtype=np.float32) / 255.0

        if len(mask.shape) > 2:  # len(x.shape) mean the channel of that x
            mask = mask[:, :, 0]
        img_patch, mask_patch = random_crop(img, mask, self.patch_size, pos_prob=0.5)
        img_patch, mask_patch = self.tranform(img_patch, mask_patch)
        img_patch, mask_patch = img_patch[np.newaxis, :], mask_patch[np.newaxis,
                                                          :]
        img_patch = torch.from_numpy(np.ascontiguousarray(img_patch))
        mask_patch = torch.from_numpy(np.ascontiguousarray(mask_patch))
        return img_patch, mask_patch

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

class STAREDatasetTest(Data.Dataset):

    def __init__(self, base_dir=r'../datasets/STARE',
                 mode='train', base_size=512, img_norm_cfg=None):
        assert mode in ['train', 'test']

        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError
        if img_norm_cfg == None:
            self.img_norm_cfg = get_img_norm_cfg('STARE', r'../datasets/STARE')
        else:
            self.img_norm_cfg = img_norm_cfg

        self.base_size = base_size
        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('ppm'):
                self.names.append(filename)
        self.tranform = augmentation()

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name)
        label_path = osp.join(self.data_dir, 'masks', name)

        img = Image.open(img_path).convert('I')
        mask = Image.open(label_path).convert('I')  # Convert to grayscale

        img = np.array(img, dtype=np.float32)  # Convert the data into numpy数组s
        img = Normalized(img, self.img_norm_cfg)
        mask = np.array(mask, dtype=np.float32) / 255.0  # Convert the data into numpy 数组 and normalization
        if len(mask.shape) > 2:
            mask = mask[:, :, 0]  # Only take the first channel
        h, w = img.shape
        img = PadImg(img)
        mask = PadImg(mask)
        img, mask = img[np.newaxis, :], mask[np.newaxis, :]
        img = torch.from_numpy(np.ascontiguousarray(img))
        mask = torch.from_numpy(np.ascontiguousarray(mask))
        return img, mask

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

class CHASEDB1DatasetTrain(Data.Dataset):
    def __init__(self, base_dir=r'../datasets/CHASEDB1',
                 mode='train', base_size=960, patch_size = 256, img_norm_cfg = None):
        assert mode in ['train', 'test']

        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError


        if img_norm_cfg == None:
            self.img_norm_cfg = get_img_norm_cfg('CHASEDB1', '../datasets/CHASEDB1')
        else:
            self.img_norm_cfg = img_norm_cfg

        self.base_size = base_size
        self.patch_size = patch_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('jpg'):
                self.names.append(osp.splitext(filename)[0])  # Store filename without extension
        self.tranform = augmentation()

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name + '.jpg') # Construct the mask path
        label_path = osp.join(self.data_dir, 'masks', name + '.png')

        img = Image.open(img_path).convert('L')
        mask = Image.open(label_path).convert('L')  # Convert to grayscale
        img = Normalized(np.array(img, dtype=np.float32), self.img_norm_cfg)
        mask = np.array(mask, dtype=np.float32) / 255.0

        if len(mask.shape) > 2:
            mask = mask[:, :, 0]

        img_patch, mask_patch = random_crop(img, mask, self.patch_size, pos_prob=0.5)
        img_patch, mask_patch = self.tranform(img_patch, mask_patch)
        img_patch, mask_patch = img_patch[np.newaxis, :], mask_patch[np.newaxis, :]
        img_patch = torch.from_numpy(np.ascontiguousarray(img_patch))
        mask_patch = torch.from_numpy(np.ascontiguousarray(mask_patch))
        return img_patch, mask_patch

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

class CHASEDB1DatasetTest(Data.Dataset):
    def __init__(self, base_dir=r'../datasets/CHASEDB1',
                 mode='train', base_size=960, img_norm_cfg = None):
        assert mode in ['train', 'test']

        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError

        if img_norm_cfg == None:
            self.img_norm_cfg = get_img_norm_cfg('CHASEDB1', '../datasets/CHASEDB1')
        else:
            self.img_norm_cfg = img_norm_cfg

        self.base_size = base_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('jpg'):
                self.names.append(osp.splitext(filename)[0])  # Store filename without extension
        self.tranform = augmentation()

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name + '.jpg') # Construct the mask path
        label_path = osp.join(self.data_dir, 'masks', name + '.png')

        img = Image.open(img_path).convert('L')
        mask = Image.open(label_path).convert('L')  # Convert to grayscale
        img = Normalized(np.array(img, dtype=np.float32), self.img_norm_cfg)
        mask = np.array(mask, dtype=np.float32) / 255.0
        if len(mask.shape) > 2:
            mask = mask[:, :, 0]

        h, w = img.shape
        img = PadImg(img)
        mask = PadImg(mask)

        img, mask = img[np.newaxis, :], mask[np.newaxis, :]

        img = torch.from_numpy(np.ascontiguousarray(img))
        mask = torch.from_numpy(np.ascontiguousarray(mask))
        return img, mask

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


class NEUDatasetTrain(Data.Dataset):
    '''
    Return: Single channel
    '''

    def __init__(self, base_dir=r'../datasets/NEU',
                 mode='train', base_size=200, patch_size = 200):
        assert mode in ['train', 'test']


        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError
        self.base_size = base_size
        self.patch_size = patch_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('jpg'):
                self.names.append(osp.splitext(filename)[0])

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name + '.jpg')
        label_path = osp.join(self.data_dir, 'masks', name + '.png')

        img, mask = cv2.imread(img_path, 0), cv2.imread(label_path, 0)

        img = np.array(img, dtype=np.float32) / 255.0
        mask = np.array(mask, dtype=np.float32) / 255.0
        if len(mask.shape) > 2:
            mask = mask[:, :, 0]

        img_patch, mask_patch = random_crop(img, mask, self.patch_size, pos_prob=0.5)
        img_patch, mask_patch = img_patch[np.newaxis, :], mask_patch[np.newaxis, :]

        img = torch.from_numpy(img_patch).type(torch.FloatTensor)
        mask = torch.from_numpy(mask_patch).type(torch.FloatTensor)

        return img, mask

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

class NEUDatasetTest(Data.Dataset):
    def __init__(self, base_dir=r'../datasets/NEU',
                 mode='test', base_size=200, patch_size=200):
        assert mode in ['train', 'test']


        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError
        self.base_size = base_size
        self.patch_size = patch_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('jpg'):
                self.names.append(osp.splitext(filename)[0])

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name + '.jpg')
        label_path = osp.join(self.data_dir, 'masks', name + '.png')

        img, mask = cv2.imread(img_path, 0), cv2.imread(label_path, 0)
        img = np.array(img, dtype=np.float32) / 255.0
        mask = np.array(mask, dtype=np.float32) / 255.0
        if len(mask.shape) > 2:
            mask = mask[:, :, 0]

        h, w = img.shape
        img = PadImg(img)
        mask = PadImg(mask)
        img, mask = img[np.newaxis, :], mask[np.newaxis, :]

        img = torch.from_numpy(img).type(torch.FloatTensor)
        mask = torch.from_numpy(mask).type(torch.FloatTensor)

        return img, mask, [h, w]

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

class SaliencyDatasetTrain(Data.Dataset):
    def __init__(self, base_dir=r'../datasets/Saliency',
                 mode='train', base_size=200, patch_size=200):
        assert mode in ['train', 'test']


        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError
        self.base_size = base_size
        self.patch_size = patch_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('bmp'):
                self.names.append(osp.splitext(filename)[0])

    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name + '.bmp')
        label_path = osp.join(self.data_dir, 'masks', name + '.png')

        img, mask = cv2.imread(img_path, 0), cv2.imread(label_path, 0)
        img = np.array(img, dtype=np.float32) / 255.0
        mask = np.array(mask, dtype=np.float32) / 255.0
        if len(mask.shape) > 2:
            mask = mask[:, :, 0]

        img_patch, mask_patch = random_crop(img, mask, self.patch_size, pos_prob=0.5)
        img_patch, mask_patch = img_patch[np.newaxis, :], mask_patch[np.newaxis, :]

        img = torch.from_numpy(img_patch).type(torch.FloatTensor)
        mask = torch.from_numpy(mask_patch).type(torch.FloatTensor)

        return img, mask

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

class SaliencyDatasetTest(Data.Dataset):

    def __init__(self, base_dir=r'../datasets/Saliency900',
                 mode='test', base_size=200, patch_size =200):
        assert mode in ['train', 'test']


        if mode == 'train':
            self.data_dir = osp.join(base_dir, 'train')
        elif mode == 'test':
            self.data_dir = osp.join(base_dir, 'test')
        else:
            raise NotImplementedError
        self.base_size = base_size
        self.patch_size = patch_size

        self.names = []
        for filename in os.listdir(osp.join(self.data_dir, 'images')):
            if filename.endswith('bmp'):
                self.names.append(osp.splitext(filename)[0])
    def __getitem__(self, i):
        name = self.names[i]
        img_path = osp.join(self.data_dir, 'images', name + '.bmp')
        label_path = osp.join(self.data_dir, 'masks', name + '.png')

        img, mask = cv2.imread(img_path, 0), cv2.imread(label_path, 0)

        img = np.array(img, dtype=np.float32) / 255.0
        mask = np.array(mask, dtype=np.float32) / 255.0
        if len(mask.shape) > 2:
            mask = mask[:, :, 0]

        h, w = img.shape
        img = PadImg(img)
        mask = PadImg(mask)
        img, mask = img[np.newaxis, :], mask[np.newaxis, :]

        img = torch.from_numpy(img).type(torch.FloatTensor)
        mask = torch.from_numpy(mask).type(torch.FloatTensor)

        return img, mask, [h,w]

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




class augmentation(object):
    def __call__(self, input, target):
        if random.random()<0.5:
            input = input[::-1, :]
            target = target[::-1, :]
        if random.random()<0.5:
            input = input[:, ::-1]
            target = target[:, ::-1]
        if random.random()<0.5:
            input = input.transpose(1, 0)
            target = target.transpose(1, 0)
        return input.copy(), target.copy()

def PadImg(img, times=32):

    if len(img.shape) == 2:
        h, w = img.shape
    elif len(img.shape) == 3:
        _, h, w = img.shape  # Unpack height, width, and channels
    else:
        raise ValueError("Unexpected number of dimensions in image")
    # h, w = img.shape
    if not h % times == 0:
        img = np.pad(img, ((0, (h//times+1)*times-h),(0, 0)), mode='constant')
    if not w % times == 0:
        img = np.pad(img, ((0, 0),(0, (w//times+1)*times-w)), mode='constant')
    return img


def random_crop(img, mask, patch_size, pos_prob=None):
    h, w = img.shape
    if min(h, w) < patch_size:
        img = np.pad(img, ((0, max(h, patch_size) - h), (0, max(w, patch_size) - w)), mode='constant')
        mask = np.pad(mask, ((0, max(h, patch_size) - h), (0, max(w, patch_size) - w)), mode='constant')
        h, w = img.shape

    cur_prob = random.random()
    if pos_prob == None or cur_prob > pos_prob or mask.max() == 0:
        h_start = random.randint(0, h - patch_size)
        w_start = random.randint(0, w - patch_size)
    else:
        loc = np.where(mask > 0)
        if len(loc[0]) <= 1:
            idx = 0
        else:
            idx = random.randint(0, len(loc[0]) - 1)
        h_start = random.randint(max(0, loc[0][idx] - patch_size), min(loc[0][idx], h - patch_size))
        w_start = random.randint(max(0, loc[1][idx] - patch_size), min(loc[1][idx], w - patch_size))

    h_end = h_start + patch_size
    w_end = w_start + patch_size
    img_patch = img[h_start:h_end, w_start:w_end]
    mask_patch = mask[h_start:h_end, w_start:w_end]

    return img_patch, mask_patch


def Normalized(img, img_norm_cfg):
    return (img - img_norm_cfg['mean']) / img_norm_cfg['std']


def Denormalization(img, img_norm_cfg):
    return img * img_norm_cfg['std'] + img_norm_cfg['mean']


def get_img_norm_cfg(dataset_name, dataset_dir):
    if dataset_name == 'CHASEDB1':
        img_norm_cfg = {'mean': 59.87125015258789, 'std': 46.84417724609375}
    elif dataset_name == 'DRIVE':
        img_norm_cfg = {'mean': 83.59488677978516, 'std': 54.175140380859375}
    elif dataset_name == 'STARE':
        img_norm_cfg = {'mean': 98.16558837890625, 'std': 52.33002853393555}

    else:
        with open(dataset_dir + '/' + '/img_idx/train_' + dataset_name + '.txt', 'r') as f:
            train_list = f.read().splitlines()
        if os.path.exists(dataset_dir + '/' +  '/img_idx/test_' + dataset_name + '.txt'):
            with open(dataset_dir + '/' +  '/img_idx/test_' + dataset_name + '.txt', 'r') as f:
                test_list = f.read().splitlines()
        else:
            test_list = []
        img_list = train_list + test_list
        img_dir = dataset_dir + '/' +  '/images/'
        mean_list = []
        std_list = []
        for img_pth in img_list:
            try:
                img = Image.open((img_dir + img_pth).replace('//', '/') + '.png').convert('I')
            except:
                try:
                    img = Image.open((img_dir + img_pth).replace('//', '/') + '.jpg').convert('I')
                except:
                    try:
                        img = Image.open((img_dir + img_pth).replace('//', '/') + '.pgm').convert('I')
                    except:
                        img = Image.open((img_dir + img_pth).replace('//', '/') + '.ppm').convert('I')
            img = np.array(img, dtype=np.float32)
            mean_list.append(img.mean())
            std_list.append(img.std())
        img_norm_cfg = dict(mean=float(np.array(mean_list).mean()), std=float(np.array(std_list).mean()))
        print(dataset_name + '\t' + str(img_norm_cfg))
    return img_norm_cfg