|
|
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: |
|
|
|
|
|
|
|
|
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: |
|
|
|
|
|
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 |
|
|
|
