306 lines
8.6 KiB
Python
306 lines
8.6 KiB
Python
import mindspore.dataset as ms
|
||
import mindspore.dataset.vision.c_transforms as c_vision
|
||
import mindspore.dataset.vision.py_transforms as py_vision
|
||
import mindspore.dataset.vision.py_transforms_util as py_util
|
||
import numpy as np
|
||
from PIL import Image, ImageOps, ImageEnhance, __version__
|
||
|
||
__all__ = [
|
||
'CentralCrop', 'HsvToRgb', 'AdjustBrightness', 'AdjustContrast', 'AdjustHue', 'Crop', 'FlipHorizontal',
|
||
'FlipVertical', 'GrayToRgb', 'RgbToGray', 'PadToBoundingBox'
|
||
]
|
||
|
||
augment_error_message = 'img should be PIL image. Got {}. Use Decode() for encoded data or ToPIL() for decoded data.'
|
||
|
||
|
||
def CentralCrop(image, central_fraction=None, size=None):
|
||
'''
|
||
|
||
Parameters
|
||
----------
|
||
image :
|
||
input Either a 3-D float Tensor of shape [height, width, depth],
|
||
or a 4-D Tensor of shape [batch_size, height, width, depth].
|
||
central_fraction :
|
||
float (0, 1], fraction of size to crop
|
||
size:
|
||
size (Union[int, sequence]) – The output size of the cropped image. If size is an integer, a square crop of size (size, size) is returned.
|
||
If size is a sequence of length 2, it should be (height, width).
|
||
Returns :
|
||
3-D / 4-D float Tensor, as per the input.
|
||
-------
|
||
'''
|
||
if size is None and central_fraction is None:
|
||
raise ValueError('central_fraction and size can not be both None')
|
||
|
||
if size is None:
|
||
outshape = np.shape(image)
|
||
if len(outshape) == 3:
|
||
h_axis = 0
|
||
w_axis = 1
|
||
elif len(outshape) == 4:
|
||
h_axis = 1
|
||
w_axis = 2
|
||
|
||
height = outshape[h_axis]
|
||
width = outshape[w_axis]
|
||
|
||
target_height = height * central_fraction
|
||
target_width = width * central_fraction
|
||
|
||
size = (target_height, target_width)
|
||
|
||
return py_util.center_crop(image, size)
|
||
|
||
|
||
def HsvToRgb(image, is_hwc=True):
|
||
|
||
image = np.asarray(image)
|
||
|
||
return py_util.hsv_to_rgbs(image, is_hwc=is_hwc)
|
||
|
||
|
||
def AdjustBrightness(image, factor):
|
||
'''
|
||
|
||
Parameters
|
||
----------
|
||
image:
|
||
input NumPy image array or PIL image
|
||
factor:
|
||
factor should be in the range (-1,1)
|
||
Returns:
|
||
-------
|
||
np darray image
|
||
'''
|
||
|
||
image = np.asarray(image)
|
||
image = image / 255
|
||
image = image + factor
|
||
index = np.where(image > 1)
|
||
image[index] = 1
|
||
index = np.where(image < 0)
|
||
image[index] = 0
|
||
image = image * 255
|
||
|
||
return image
|
||
|
||
|
||
def AdjustContrast(image, factor):
|
||
|
||
if isinstance(image, np.ndarray):
|
||
image = Image.fromarray(image)
|
||
if not isinstance(image, Image.Image):
|
||
raise TypeError(augment_error_message.format(type(image)))
|
||
|
||
image = ImageEnhance.Contrast(image).enhance(factor)
|
||
|
||
image = np.array(image)
|
||
|
||
return image
|
||
|
||
|
||
def AdjustHue(image, factor):
|
||
|
||
if isinstance(image, np.ndarray):
|
||
image = Image.fromarray(image)
|
||
if not isinstance(image, Image.Image):
|
||
raise TypeError(augment_error_message.format(type(image)))
|
||
|
||
image_hue_factor = factor
|
||
if not -1 <= image_hue_factor <= 1:
|
||
raise ValueError('image_hue_factor {} is not in [-1, 1].'.format(image_hue_factor))
|
||
|
||
if not isinstance(image, Image.Image):
|
||
raise TypeError(augment_error_message.format(type(image)))
|
||
|
||
mode = image.mode
|
||
if mode in {'L', '1', 'I', 'F'}:
|
||
return image
|
||
|
||
hue, saturation, value = image.convert('HSV').split()
|
||
|
||
np_hue = np.array(hue, dtype=np.uint8)
|
||
|
||
with np.errstate(over='ignore'):
|
||
np_hue += np.uint8(image_hue_factor * 255)
|
||
hue = Image.fromarray(np_hue, 'L')
|
||
|
||
image = Image.merge('HSV', (hue, saturation, value)).convert(mode)
|
||
return image
|
||
|
||
|
||
def AdjustSaturation(image, factor):
|
||
|
||
if isinstance(image, np.ndarray):
|
||
image = Image.fromarray(image)
|
||
if not isinstance(image, Image.Image):
|
||
raise TypeError(augment_error_message.format(type(image)))
|
||
|
||
enhancer = ImageEnhance.Color(image)
|
||
image = enhancer.enhance(factor)
|
||
return image
|
||
|
||
|
||
def Crop(image, offset_height, offset_width, target_height, target_width):
|
||
|
||
if isinstance(image, np.ndarray):
|
||
image = Image.fromarray(image)
|
||
if not isinstance(image, Image.Image):
|
||
raise TypeError(augment_error_message.format(type(image)))
|
||
image = np.array(
|
||
image.crop((offset_width, offset_height, offset_width + target_width, offset_width + target_height))
|
||
)
|
||
return image
|
||
|
||
|
||
def FlipHorizontal(image):
|
||
|
||
if isinstance(image, np.ndarray):
|
||
image = Image.fromarray(image)
|
||
if not isinstance(image, Image.Image):
|
||
raise TypeError(augment_error_message.format(type(image)))
|
||
|
||
image = np.array(image.transpose(Image.FLIP_LEFT_RIGHT))
|
||
|
||
return image
|
||
|
||
|
||
def FlipVertical(image):
|
||
|
||
if isinstance(image, np.ndarray):
|
||
image = Image.fromarray(image)
|
||
if not isinstance(image, Image.Image):
|
||
raise TypeError(augment_error_message.format(type(image)))
|
||
|
||
image = np.array(image.transpose(Image.FLIP_TOP_BOTTOM))
|
||
|
||
return image
|
||
|
||
|
||
def GrayToRgb(image):
|
||
|
||
image = np.asarray(image)
|
||
shape = image.shape
|
||
output_image = np.zeros((shape[0], shape[1], 3), dtype=np.uint8)
|
||
if len(shape) == 3:
|
||
for i in range(3):
|
||
output_image[:, :, i] = image[:, :, 1]
|
||
elif len(shape) == 2:
|
||
for i in range(3):
|
||
output_image[:, :, i] = image
|
||
|
||
return output_image
|
||
|
||
|
||
def RgbToGray(image):
|
||
|
||
if isinstance(image, np.ndarray):
|
||
image = Image.fromarray(image)
|
||
if not isinstance(image, Image.Image):
|
||
raise TypeError(augment_error_message.format(type(image)))
|
||
'''
|
||
将彩色图像转换为灰度（模式“L”）时，库使用ITU-R 601-2 Luma转换：
|
||
L = R * 299/1000 + G * 587/1000 + B * 114/1000
|
||
'''
|
||
image = image.convert('L')
|
||
image = np.asarray(image)
|
||
|
||
return image
|
||
|
||
|
||
def PadToBoundingBox(image, offset_height, offset_width, target_height, target_width):
|
||
'''
|
||
|
||
Parameters
|
||
----------
|
||
image:
|
||
A PIL image
|
||
offset_height:
|
||
Number of rows of zeros to add on top.
|
||
offset_width:
|
||
Number of columns of zeros to add on the left.
|
||
target_height:
|
||
Height of output image.
|
||
target_width
|
||
Width of output image.
|
||
Returns
|
||
A numpy ndarray image
|
||
-------
|
||
'''
|
||
|
||
if offset_height < 0:
|
||
raise ValueError("offset_height must be >= 0")
|
||
if offset_width < 0:
|
||
raise ValueError("offset_width must be >= 0")
|
||
image = np.array(image)
|
||
shape = image.shape
|
||
top = offset_height
|
||
bottom = target_height - shape[0] - top
|
||
left = offset_width
|
||
right = target_width - shape[1] - left
|
||
|
||
if bottom < 0:
|
||
raise ValueError("target_height must be >= offset_height + height")
|
||
|
||
if right < 0:
|
||
raise ValueError("target_width must be >= offset_width + width")
|
||
|
||
return np.pad(image, ((top, bottom), (left, right), (0, 0)), mode='constant')
|
||
|
||
|
||
def Standardization(image, mean=None, std=None, channel_mode=False):
|
||
'''
|
||
|
||
Parameters
|
||
----------
|
||
image:
|
||
An n-D Tensor with at least 3 dimensions, the last 3 of which are the dimensions of each image.
|
||
mean:
|
||
List or tuple of mean values for each channel, with respect to channel order.
|
||
std:
|
||
List or tuple of standard deviations for each channel.
|
||
channel_mode:
|
||
Decide to implement standardization on whole image or each channel of image.
|
||
Returns:
|
||
A Tensor with the same shape and dtype as image.
|
||
-------
|
||
'''
|
||
image = np.array(image, dtype=np.float32)
|
||
num_shape = image.shape
|
||
if mean is not None and std is not None:
|
||
if len(mean) != len(std):
|
||
raise ValueError("Length of mean and std must be equal")
|
||
if len(mean) == 1:
|
||
mean = [mean[0]] * num_shape[2]
|
||
std = [std[0]] * num_shape[2]
|
||
mean = np.array(mean, dtype=image.dtype)
|
||
std = np.array(std, dtype=image.dtype)
|
||
return (image - mean[:, None, None]) / std[:, None, None]
|
||
elif mean is None and std is None:
|
||
if channel_mode:
|
||
num_pixels = num_shape[0] * num_shape[1]
|
||
image_mean = np.mean(image, axis=(0, 1))
|
||
stddev = np.std(image, axis=(0, 1))
|
||
min_sttdev = 1 / np.sqrt(num_pixels)
|
||
min_sttdev = [min_sttdev] * num_shape[2]
|
||
adjusted_sttdev = np.maximum(stddev, min_sttdev)
|
||
|
||
image -= image_mean
|
||
image = np.divide(image, adjusted_sttdev)
|
||
return image
|
||
else:
|
||
num_pixels = num_shape[0] * num_shape[1] * num_shape[2]
|
||
image_mean = np.mean(image, axis=(0, 1, 2))
|
||
image_mean = [image_mean] * 3
|
||
stddev = np.std(image, axis=(0, 1, 2))
|
||
min_sttdev = 1 / np.sqrt(num_pixels)
|
||
adjusted_sttdev = np.maximum(stddev, min_sttdev)
|
||
adjusted_sttdev = [adjusted_sttdev] * 3
|
||
|
||
image -= image_mean
|
||
image = np.divide(image, adjusted_sttdev)
|
||
return image
|
||
else:
|
||
raise ValueError('std and mean must both be None or not None')
|