data/custom_transforms.py

from email.mime import base
import numpy as np
from PIL import Image
import os
import sys
import torch
import torch.nn.functional as F
from PIL import Image, ImageOps, ImageFilter, ImageEnhance
from typing import Optional

filepath = os.path.split(__file__)[0]
repopath = os.path.split(filepath)[0]
sys.path.append(repopath)

from utils.misc import *
        
class static_resize:
    # Resize for training
    # size: h x w
    def __init__(self, size=[384, 384], base_size=None):
        self.size = size[::-1]
        self.base_size = base_size[::-1] if base_size is not None else None
            
    def __call__(self, sample):
        sample['image'] = sample['image'].resize(self.size, Image.BILINEAR)
        if 'gt' in sample.keys():
            sample['gt'] = sample['gt'].resize(self.size, Image.NEAREST)
            
        if self.base_size is not None:
            sample['image_resized'] = sample['image'].resize(self.size, Image.BILINEAR)
            if 'gt' in sample.keys():
                sample['gt_resized'] = sample['gt'].resize(self.size, Image.NEAREST)
                
        return sample

class dynamic_resize:
    # base_size: h x w
    def __init__(self, L=1280, base_size=[384, 384]): 
        self.L = L
        self.base_size = base_size[::-1]
                    
    def __call__(self, sample):
        size = list(sample['image'].size)
        if (size[0] >= size[1]) and size[1] > self.L: 
            size[0] = size[0] / (size[1] / self.L)
            size[1] = self.L
        elif (size[1] > size[0]) and size[0] > self.L:
            size[1] = size[1] / (size[0] / self.L)
            size[0] = self.L
        size = (int(round(size[0] / 32)) * 32, int(round(size[1] / 32)) * 32)
    
        if 'image' in sample.keys():
            sample['image_resized'] = sample['image'].resize(self.base_size, Image.BILINEAR)
            sample['image'] = sample['image'].resize(size, Image.BILINEAR)
            
        if 'gt' in sample.keys():
            sample['gt_resized'] = sample['gt'].resize(self.base_size, Image.NEAREST)
            sample['gt'] = sample['gt'].resize(size, Image.NEAREST)
        
        return sample
        
class random_scale_crop:
    def __init__(self, range=[0.75, 1.25]):
        self.range = range

    def __call__(self, sample):
        scale = np.random.random() * (self.range[1] - self.range[0]) + self.range[0]
        if np.random.random() < 0.5:
            for key in sample.keys():
                if key in ['image', 'gt']:
                    base_size = sample[key].size

                    scale_size = tuple((np.array(base_size) * scale).round().astype(int))
                    sample[key] = sample[key].resize(scale_size)

                    lf = (sample[key].size[0] - base_size[0]) // 2
                    up = (sample[key].size[1] - base_size[1]) // 2
                    rg = (sample[key].size[0] + base_size[0]) // 2
                    lw = (sample[key].size[1] + base_size[1]) // 2

                    border = -min(0, min(lf, up))
                    sample[key] = ImageOps.expand(sample[key], border=border)
                    sample[key] = sample[key].crop((lf + border, up + border, rg + border, lw + border))
        return sample

class random_flip:
    def __init__(self, lr=True, ud=True):
        self.lr = lr
        self.ud = ud

    def __call__(self, sample):
        lr = np.random.random() < 0.5 and self.lr is True
        ud = np.random.random() < 0.5 and self.ud is True

        for key in sample.keys():
            if key in ['image', 'gt']:
                sample[key] = np.array(sample[key])
                if lr:
                    sample[key] = np.fliplr(sample[key])
                if ud:
                    sample[key] = np.flipud(sample[key])
                sample[key] = Image.fromarray(sample[key])

        return sample

class random_rotate:
    def __init__(self, range=[0, 360], interval=1):
        self.range = range
        self.interval = interval

    def __call__(self, sample):
        rot = (np.random.randint(*self.range) // self.interval) * self.interval
        rot = rot + 360 if rot < 0 else rot

        if np.random.random() < 0.5:
            for key in sample.keys():
                if key in ['image', 'gt']:
                    base_size = sample[key].size
                    sample[key] = sample[key].rotate(rot, expand=True, fillcolor=255 if key == 'depth' else None)

                    sample[key] = sample[key].crop(((sample[key].size[0] - base_size[0]) // 2,
                                                    (sample[key].size[1] - base_size[1]) // 2,
                                                    (sample[key].size[0] + base_size[0]) // 2,
                                                    (sample[key].size[1] + base_size[1]) // 2))

        return sample

class random_image_enhance:
    def __init__(self, methods=['contrast', 'brightness', 'sharpness']):
        self.enhance_method = []
        if 'contrast' in methods:
            self.enhance_method.append(ImageEnhance.Contrast)
        if 'brightness' in methods:
            self.enhance_method.append(ImageEnhance.Brightness)
        if 'sharpness' in methods:
            self.enhance_method.append(ImageEnhance.Sharpness)

    def __call__(self, sample):
        if 'image' in sample.keys():
            np.random.shuffle(self.enhance_method)

            for method in self.enhance_method:
                if np.random.random() > 0.5:
                    enhancer = method(sample['image'])
                    factor = float(1 + np.random.random() / 10)
                    sample['image'] = enhancer.enhance(factor)

        return sample

class tonumpy:
    def __init__(self):
        pass

    def __call__(self, sample):
        for key in sample.keys():
            if key in ['image', 'image_resized', 'gt', 'gt_resized']:
                sample[key] = np.array(sample[key], dtype=np.float32)

        return sample

class normalize:
    def __init__(self, mean: Optional[list]=None, std: Optional[list]=None, div=255):
        self.mean = mean if mean is not None else 0.0
        self.std = std if std is not None else 1.0
        self.div = div
        
    def __call__(self, sample):
        if 'image' in sample.keys():
            sample['image'] /= self.div
            sample['image'] -= self.mean
            sample['image'] /= self.std
            
        if 'image_resized' in sample.keys():
            sample['image_resized'] /= self.div
            sample['image_resized'] -= self.mean
            sample['image_resized'] /= self.std

        if 'gt' in sample.keys():
            sample['gt'] /= self.div
            
        if 'gt_resized' in sample.keys():
            sample['gt_resized'] /= self.div

        return sample
class totensor:
    def __init__(self):
        pass

    def __call__(self, sample):
        if 'image' in sample.keys():
            sample['image'] = sample['image'].transpose((2, 0, 1))
            sample['image'] = torch.from_numpy(sample['image']).float()
        
        if 'image_resized' in sample.keys():
            sample['image_resized'] = sample['image_resized'].transpose((2, 0, 1))
            sample['image_resized'] = torch.from_numpy(sample['image_resized']).float()
            
        if 'gt' in sample.keys():
            sample['gt'] = torch.from_numpy(sample['gt'])
            sample['gt'] = sample['gt'].unsqueeze(dim=0)
            
        if 'gt_resized' in sample.keys():
            sample['gt_resized'] = torch.from_numpy(sample['gt_resized'])
            sample['gt_resized'] = sample['gt_resized'].unsqueeze(dim=0)

        return sample