You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
237 lines
7.7 KiB
Python
237 lines
7.7 KiB
Python
"""Functions and classes for image preprocessing"""
|
|
|
|
import math
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
def apply_min_size(sample, size, image_interpolation_method=cv2.INTER_AREA):
|
|
"""
|
|
Rezise the sample to ensure the given size. Keeps aspect ratio.
|
|
|
|
Args:
|
|
sample (dict): sample
|
|
size (tuple): image size
|
|
|
|
Returns:
|
|
tuple: new size
|
|
"""
|
|
shape = list(sample["disparity"].shape)
|
|
|
|
if shape[0] >= size[0] and shape[1] >= size[1]:
|
|
return sample
|
|
|
|
scale = [0, 0]
|
|
scale[0] = size[0] / shape[0]
|
|
scale[1] = size[1] / shape[1]
|
|
|
|
scale = max(scale)
|
|
|
|
shape[0] = math.ceil(scale * shape[0])
|
|
shape[1] = math.ceil(scale * shape[1])
|
|
|
|
# resize
|
|
sample["image"] = cv2.resize(
|
|
sample["image"], tuple(shape[::-1]), interpolation=image_interpolation_method
|
|
)
|
|
|
|
sample["disparity"] = cv2.resize(
|
|
sample["disparity"], tuple(shape[::-1]), interpolation=cv2.INTER_NEAREST
|
|
)
|
|
sample["mask"] = cv2.resize(
|
|
sample["mask"].astype(np.float32),
|
|
tuple(shape[::-1]),
|
|
interpolation=cv2.INTER_NEAREST,
|
|
)
|
|
sample["mask"] = sample["mask"].astype(bool)
|
|
|
|
return tuple(shape)
|
|
|
|
|
|
class Resize:
|
|
"""Resize sample to given size (width, height)."""
|
|
|
|
def __init__(
|
|
self,
|
|
width,
|
|
height,
|
|
resize_target=True,
|
|
keep_aspect_ratio=False,
|
|
ensure_multiple_of=1,
|
|
resize_method="lower_bound",
|
|
image_interpolation_method=cv2.INTER_AREA,
|
|
):
|
|
"""
|
|
Init.
|
|
|
|
Args:
|
|
width (int): desired output width
|
|
height (int): desired output height
|
|
resize_target (bool, optional):
|
|
True: Resize the full sample (image, mask, target).
|
|
False: Resize image only.
|
|
Defaults to True.
|
|
keep_aspect_ratio (bool, optional):
|
|
True: Keep the aspect ratio of the input sample.
|
|
Output sample might not have the given width and height, and
|
|
resize behaviour depends on the parameter 'resize_method'.
|
|
Defaults to False.
|
|
ensure_multiple_of (int, optional):
|
|
Output width and height is constrained to be multiple of this parameter.
|
|
Defaults to 1.
|
|
resize_method (str, optional):
|
|
"lower_bound": Output will be at least as large as the given size.
|
|
"upper_bound": Output will be at max as large as the given size. (Output size might be smaller than given size.)
|
|
"minimal": Scale as least as possible. (Output size might be smaller than given size.)
|
|
Defaults to "lower_bound".
|
|
"""
|
|
self.__width = width
|
|
self.__height = height
|
|
|
|
self.__resize_target = resize_target
|
|
self.__keep_aspect_ratio = keep_aspect_ratio
|
|
self.__multiple_of = ensure_multiple_of
|
|
self.__resize_method = resize_method
|
|
self.__image_interpolation_method = image_interpolation_method
|
|
|
|
def constrain_to_multiple_of(self, x, min_val=0, max_val=None):
|
|
y = (np.round(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
|
|
|
if max_val is not None and y > max_val:
|
|
y = (np.floor(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
|
|
|
if y < min_val:
|
|
y = (np.ceil(x / self.__multiple_of) * self.__multiple_of).astype(int)
|
|
|
|
return y
|
|
|
|
def get_size(self, width, height):
|
|
# determine new height and width
|
|
scale_height = self.__height / height
|
|
scale_width = self.__width / width
|
|
|
|
if self.__keep_aspect_ratio:
|
|
if self.__resize_method == "lower_bound":
|
|
# scale such that output size is lower bound
|
|
if scale_width > scale_height:
|
|
# fit width
|
|
scale_height = scale_width
|
|
else:
|
|
# fit height
|
|
scale_width = scale_height
|
|
elif self.__resize_method == "upper_bound":
|
|
# scale such that output size is upper bound
|
|
if scale_width < scale_height:
|
|
# fit width
|
|
scale_height = scale_width
|
|
else:
|
|
# fit height
|
|
scale_width = scale_height
|
|
elif self.__resize_method == "minimal":
|
|
# scale as least as possbile
|
|
if abs(1 - scale_width) < abs(1 - scale_height):
|
|
# fit width
|
|
scale_height = scale_width
|
|
else:
|
|
# fit height
|
|
scale_width = scale_height
|
|
else:
|
|
msg = f"resize_method {self.__resize_method} not implemented"
|
|
raise ValueError(msg)
|
|
|
|
if self.__resize_method == "lower_bound":
|
|
new_height = self.constrain_to_multiple_of(
|
|
scale_height * height, min_val=self.__height
|
|
)
|
|
new_width = self.constrain_to_multiple_of(
|
|
scale_width * width, min_val=self.__width
|
|
)
|
|
elif self.__resize_method == "upper_bound":
|
|
new_height = self.constrain_to_multiple_of(
|
|
scale_height * height, max_val=self.__height
|
|
)
|
|
new_width = self.constrain_to_multiple_of(
|
|
scale_width * width, max_val=self.__width
|
|
)
|
|
elif self.__resize_method == "minimal":
|
|
new_height = self.constrain_to_multiple_of(scale_height * height)
|
|
new_width = self.constrain_to_multiple_of(scale_width * width)
|
|
else:
|
|
msg = f"resize_method {self.__resize_method} not implemented"
|
|
raise ValueError(msg)
|
|
|
|
return (new_width, new_height)
|
|
|
|
def __call__(self, sample):
|
|
width, height = self.get_size(
|
|
sample["image"].shape[1], sample["image"].shape[0]
|
|
)
|
|
|
|
# resize sample
|
|
sample["image"] = cv2.resize(
|
|
sample["image"],
|
|
(width, height),
|
|
interpolation=self.__image_interpolation_method,
|
|
)
|
|
|
|
if self.__resize_target:
|
|
if "disparity" in sample:
|
|
sample["disparity"] = cv2.resize(
|
|
sample["disparity"],
|
|
(width, height),
|
|
interpolation=cv2.INTER_NEAREST,
|
|
)
|
|
|
|
if "depth" in sample:
|
|
sample["depth"] = cv2.resize(
|
|
sample["depth"], (width, height), interpolation=cv2.INTER_NEAREST
|
|
)
|
|
|
|
sample["mask"] = cv2.resize(
|
|
sample["mask"].astype(np.float32),
|
|
(width, height),
|
|
interpolation=cv2.INTER_NEAREST,
|
|
)
|
|
sample["mask"] = sample["mask"].astype(bool)
|
|
|
|
return sample
|
|
|
|
|
|
class NormalizeImage:
|
|
"""Normlize image by given mean and std."""
|
|
|
|
def __init__(self, mean, std):
|
|
self.__mean = mean
|
|
self.__std = std
|
|
|
|
def __call__(self, sample):
|
|
sample["image"] = (sample["image"] - self.__mean) / self.__std
|
|
|
|
return sample
|
|
|
|
|
|
class PrepareForNet:
|
|
"""Prepare sample for usage as network input."""
|
|
|
|
def __init__(self):
|
|
pass
|
|
|
|
def __call__(self, sample):
|
|
image = np.transpose(sample["image"], (2, 0, 1))
|
|
sample["image"] = np.ascontiguousarray(image).astype(np.float32)
|
|
|
|
if "mask" in sample:
|
|
sample["mask"] = sample["mask"].astype(np.float32)
|
|
sample["mask"] = np.ascontiguousarray(sample["mask"])
|
|
|
|
if "disparity" in sample:
|
|
disparity = sample["disparity"].astype(np.float32)
|
|
sample["disparity"] = np.ascontiguousarray(disparity)
|
|
|
|
if "depth" in sample:
|
|
depth = sample["depth"].astype(np.float32)
|
|
sample["depth"] = np.ascontiguousarray(depth)
|
|
|
|
return sample
|