dataset.py

import os
import cv2
import torch
import numpy as np
from torch.utils.data import Dataset
from PIL import Image
import random

########################### Data Augmentation ###########################
class Normalize(object):
    def __init__(self, mean, std):
        self.mean = mean 
        self.std  = std
    
    def __call__(self, image, mask=None, body=None, detail=None):
        image = (image - self.mean)/self.std
        if mask is None:
            return image
        return image, mask/255

class RandomCrop(object):
    def __call__(self, image, mask=None, body=None, detail=None):
        H,W,_   = image.shape
        randw   = np.random.randint(W/8)
        randh   = np.random.randint(H/8)
        offseth = 0 if randh == 0 else np.random.randint(randh)
        offsetw = 0 if randw == 0 else np.random.randint(randw)
        p0, p1, p2, p3 = offseth, H+offseth-randh, offsetw, W+offsetw-randw
        if mask is None:
            return image[p0:p1,p2:p3, :]
        return image[p0:p1,p2:p3, :], mask[p0:p1,p2:p3] 

class RandomFlip(object):
    def __call__(self, image, mask=None, body=None, detail=None):
        if np.random.randint(2)==0:
            if mask is None:
                return image[:,::-1,:].copy()
            return image[:,::-1,:].copy(), mask[:, ::-1].copy() 
        else:
            if mask is None:
                return image
            return image, mask 

class Resize(object):
    def __init__(self, H, W):
        self.H = H
        self.W = W

    def __call__(self, image, mask=None, body=None, detail=None):
        image = cv2.resize(image, dsize=(self.W, self.H), interpolation=cv2.INTER_LINEAR)
        if mask is None:
            return image
        mask  = cv2.resize( mask, dsize=(self.W, self.H), interpolation=cv2.INTER_LINEAR)
        body  = cv2.resize( body, dsize=(self.W, self.H), interpolation=cv2.INTER_LINEAR)
        detail= cv2.resize( detail, dsize=(self.W, self.H), interpolation=cv2.INTER_LINEAR)
        return image, mask 

class RandomRotate(object):
    def rotate(self, x, random_angle, mode='image'):

        if mode == 'image':
            H, W, _ = x.shape
        else:
            H, W = x.shape

        random_angle %= 360
        image_change = cv2.getRotationMatrix2D((W/2, H/2), random_angle, 1)
        image_rotated = cv2.warpAffine(x, image_change, (W, H))
    
        angle_crop = random_angle % 180
        if random_angle > 90:
            angle_crop = 180 - angle_crop
        theta = angle_crop * np.pi / 180
        hw_ratio = float(H) / float(W)
        tan_theta = np.tan(theta)
        numerator = np.cos(theta) + np.sin(theta) * np.tan(theta)
        r = hw_ratio if H > W else 1 / hw_ratio
        denominator = r * tan_theta + 1
        crop_mult = numerator / denominator

        w_crop = int(crop_mult * W)
        h_crop = int(crop_mult * H)
        x0 = int((W - w_crop) / 2)
        y0 = int((H - h_crop) / 2)
        crop_image = lambda img, x0, y0, W, H: img[y0:y0+h_crop, x0:x0+w_crop ]
        output = crop_image(image_rotated, x0, y0, w_crop, h_crop)

        return output

    def __call__(self, image, mask=None, body=None, detail=None):

        do_seed = np.random.randint(0,3)
        if do_seed != 2:
            if mask is None:
                return image
            return image, mask
        
        random_angle = np.random.randint(-10, 10)
        image = self.rotate(image, random_angle, 'image')

        if mask is None:
            return image
        mask = self.rotate(mask, random_angle, 'mask')

        return image, mask 


class ColorEnhance(object):
    def __init__(self):

        #A:0.5~1.5, G: 5-15
        self.A = np.random.randint(7, 13, 1)[0]/10
        self.G = np.random.randint(7, 13, 1)[0]
        

    def __call__(self, image, mask=None, body=None, detail=None):

        do_seed = np.random.randint(0,3)
        if do_seed > 1:#1: # 1/3
            H, W, _   = image.shape
            dark_matrix = np.zeros([H, W, _], image.dtype)
            image = cv2.addWeighted(image, self.A, dark_matrix, 1-self.A, self.G) 
        else:
            pass
            
        if mask is None:
            return image
        return image, mask 

class GaussNoise(object):
    def __init__(self):
        self.Mean = 0
        self.Var = 0.001

    def __call__(self, image, mask=None, body=None, detail=None):
        H, W, _   = image.shape
        do_seed = np.random.randint(0,3)


        if do_seed == 0: #1: # 1/3
            factor = np.random.randint(0,10)
            noise = np.random.normal(self.Mean, self.Var ** 0.5, image.shape) * factor
            noise = noise.astype(image.dtype)
            image = cv2.add(image, noise)
        else:
            pass

        if mask is None:
            return image
        return image, mask 



class ToTensor(object):
    def __call__(self, image, mask=None, body=None, detail=None):
        image = torch.from_numpy(image)
        image = image.permute(2, 0, 1)
        if mask is None:
            return image
        mask  = torch.from_numpy(mask)
        return image, mask 


########################### Config File ###########################
class Config(object):
    def __init__(self, **kwargs):
        self.kwargs = kwargs
        self.mean   = np.array([[[124.55, 118.90, 102.94]]])
        self.std    = np.array([[[ 56.77,  55.97,  57.50]]])
        print('\nParameters...')
        for k, v in self.kwargs.items():
            print('%-10s: %s'%(k, v))

    def __getattr__(self, name):
        if name in self.kwargs:
            return self.kwargs[name]
        else:
            return None


########################### Dataset Class ###########################
class Data(Dataset):
    def __init__(self, cfg, model_name):
        self.cfg        = cfg
        self.model_name = model_name
        self.normalize  = Normalize(mean=cfg.mean, std=cfg.std)
        self.randomcrop = RandomCrop()
        self.randomflip = RandomFlip()
    
        self.resize     = Resize(384, 384)
        
        self.randomrotate = RandomRotate()
        self.colorenhance = ColorEnhance()
        self.gaussnoise = GaussNoise()
        self.totensor   = ToTensor()

        self.samples=os.listdir(cfg.datapath+'/Image')

    def __getitem__(self, idx):
        name  = self.samples[idx]
        try:
            image = cv2.imread(self.cfg.datapath+'/Image/'+name.replace('.jpg','')+'.jpg') 
        except:
            print(str(name)+' not found!')


        if self.cfg.mode=='train':
            try:
                mask  = cv2.imread(self.cfg.datapath+'/GT/' +name.replace('.jpg','')+'.png', 0).astype(np.float32)
            except:
                print(str(name)+' not found!')

            image, mask = self.normalize(image, mask)
            image, mask = self.randomcrop(image, mask)
            image, mask = self.randomflip(image, mask)
 
            return image, mask 
        else:
            shape = image.shape[:2]
            image = self.normalize(image)
            image = self.resize(image)
            image = self.totensor(image)

            return image, shape, name

    def __len__(self):
        return len(self.samples)
    
    def collate(self, batch):
        size = 384
        image, mask = [list(item) for item in zip(*batch)]
        for i in range(len(batch)):
            image[i] = cv2.resize(image[i], dsize=(size, size), interpolation=cv2.INTER_LINEAR)
            mask[i]  = cv2.resize(mask[i],  dsize=(size, size), interpolation=cv2.INTER_LINEAR)

        image  = torch.from_numpy(np.stack(image, axis=0)).permute(0,3,1,2)
        mask   = torch.from_numpy(np.stack(mask, axis=0)).unsqueeze(1)
        return image, mask