230 lines
7.0 KiB
Python
230 lines
7.0 KiB
Python
import warnings
|
|
|
|
import cv2
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
from imageio import imread
|
|
from skimage.transform import rotate, resize
|
|
|
|
warnings.filterwarnings("ignore", category=UserWarning)
|
|
|
|
|
|
class Image:
|
|
"""
|
|
Wrapper around an image with some convenient functions.
|
|
"""
|
|
|
|
def __init__(self, img):
|
|
self.img = img
|
|
|
|
def __getattr__(self, attr):
|
|
# Pass along any other methods to the underlying ndarray
|
|
return getattr(self.img, attr)
|
|
|
|
@classmethod
|
|
def from_file(cls, fname):
|
|
return cls(imread(fname))
|
|
|
|
def copy(self):
|
|
"""
|
|
:return: Copy of self.
|
|
"""
|
|
return self.__class__(self.img.copy())
|
|
|
|
def crop(self, top_left, bottom_right, resize=None):
|
|
"""
|
|
Crop the image to a bounding box given by top left and bottom right pixels.
|
|
:param top_left: tuple, top left pixel.
|
|
:param bottom_right: tuple, bottom right pixel
|
|
:param resize: If specified, resize the cropped image to this size
|
|
"""
|
|
self.img = self.img[top_left[0]:bottom_right[0], top_left[1]:bottom_right[1]]
|
|
if resize is not None:
|
|
self.resize(resize)
|
|
|
|
def cropped(self, *args, **kwargs):
|
|
"""
|
|
:return: Cropped copy of the image.
|
|
"""
|
|
i = self.copy()
|
|
i.crop(*args, **kwargs)
|
|
return i
|
|
|
|
def normalise(self):
|
|
"""
|
|
Normalise the image by converting to float [0,1] and zero-centering
|
|
"""
|
|
self.img = self.img.astype(np.float32) / 255.0
|
|
self.img -= self.img.mean()
|
|
|
|
def resize(self, shape):
|
|
"""
|
|
Resize image to shape.
|
|
:param shape: New shape.
|
|
"""
|
|
if self.img.shape == shape:
|
|
return
|
|
self.img = resize(self.img, shape, preserve_range=True).astype(self.img.dtype)
|
|
|
|
def resized(self, *args, **kwargs):
|
|
"""
|
|
:return: Resized copy of the image.
|
|
"""
|
|
i = self.copy()
|
|
i.resize(*args, **kwargs)
|
|
return i
|
|
|
|
def rotate(self, angle, center=None):
|
|
"""
|
|
Rotate the image.
|
|
:param angle: Angle (in radians) to rotate by.
|
|
:param center: Center pixel to rotate if specified, otherwise image center is used.
|
|
"""
|
|
if center is not None:
|
|
center = (center[1], center[0])
|
|
self.img = rotate(self.img, angle / np.pi * 180, center=center, mode='symmetric', preserve_range=True).astype(
|
|
self.img.dtype)
|
|
|
|
def rotated(self, *args, **kwargs):
|
|
"""
|
|
:return: Rotated copy of image.
|
|
"""
|
|
i = self.copy()
|
|
i.rotate(*args, **kwargs)
|
|
return i
|
|
|
|
def show(self, ax=None, **kwargs):
|
|
"""
|
|
Plot the image
|
|
:param ax: Existing matplotlib axis (optional)
|
|
:param kwargs: kwargs to imshow
|
|
"""
|
|
if ax:
|
|
ax.imshow(self.img, **kwargs)
|
|
else:
|
|
plt.imshow(self.img, **kwargs)
|
|
plt.show()
|
|
|
|
def zoom(self, factor):
|
|
"""
|
|
"Zoom" the image by cropping and resizing.
|
|
:param factor: Factor to zoom by. e.g. 0.5 will keep the center 50% of the image.
|
|
"""
|
|
sr = int(self.img.shape[0] * (1 - factor)) // 2
|
|
sc = int(self.img.shape[1] * (1 - factor)) // 2
|
|
orig_shape = self.img.shape
|
|
self.img = self.img[sr:self.img.shape[0] - sr, sc: self.img.shape[1] - sc].copy()
|
|
self.img = resize(self.img, orig_shape, mode='symmetric', preserve_range=True).astype(self.img.dtype)
|
|
|
|
def zoomed(self, *args, **kwargs):
|
|
"""
|
|
:return: Zoomed copy of the image.
|
|
"""
|
|
i = self.copy()
|
|
i.zoom(*args, **kwargs)
|
|
return i
|
|
|
|
|
|
class DepthImage(Image):
|
|
def __init__(self, img):
|
|
super().__init__(img)
|
|
|
|
@classmethod
|
|
def from_pcd(cls, pcd_filename, shape, default_filler=0, index=None):
|
|
"""
|
|
Create a depth image from an unstructured PCD file.
|
|
If index isn't specified, use euclidean distance, otherwise choose x/y/z=0/1/2
|
|
"""
|
|
img = np.zeros(shape)
|
|
if default_filler != 0:
|
|
img += default_filler
|
|
|
|
with open(pcd_filename) as f:
|
|
for l in f.readlines():
|
|
ls = l.split()
|
|
|
|
if len(ls) != 5:
|
|
# Not a point line in the file.
|
|
continue
|
|
try:
|
|
# Not a number, carry on.
|
|
float(ls[0])
|
|
except ValueError:
|
|
continue
|
|
|
|
i = int(ls[4])
|
|
r = i // shape[1]
|
|
c = i % shape[1]
|
|
|
|
if index is None:
|
|
x = float(ls[0])
|
|
y = float(ls[1])
|
|
z = float(ls[2])
|
|
|
|
img[r, c] = np.sqrt(x ** 2 + y ** 2 + z ** 2)
|
|
|
|
else:
|
|
img[r, c] = float(ls[index])
|
|
|
|
return cls(img / 1000.0)
|
|
|
|
@classmethod
|
|
def from_tiff(cls, fname):
|
|
return cls(imread(fname))
|
|
|
|
def inpaint(self, missing_value=0):
|
|
"""
|
|
Inpaint missing values in depth image.
|
|
:param missing_value: Value to fill in teh depth image.
|
|
"""
|
|
# cv2 inpainting doesn't handle the border properly
|
|
# https://stackoverflow.com/questions/25974033/inpainting-depth-map-still-a-black-image-border
|
|
self.img = cv2.copyMakeBorder(self.img, 1, 1, 1, 1, cv2.BORDER_DEFAULT)
|
|
mask = (self.img == missing_value).astype(np.uint8)
|
|
|
|
# Scale to keep as float, but has to be in bounds -1:1 to keep opencv happy.
|
|
scale = np.abs(self.img).max()
|
|
self.img = self.img.astype(np.float32) / scale # Has to be float32, 64 not supported.
|
|
self.img = cv2.inpaint(self.img, mask, 1, cv2.INPAINT_NS)
|
|
|
|
# Back to original size and value range.
|
|
self.img = self.img[1:-1, 1:-1]
|
|
self.img = self.img * scale
|
|
|
|
def gradients(self):
|
|
"""
|
|
Compute gradients of the depth image using Sobel filtesr.
|
|
:return: Gradients in X direction, Gradients in Y diretion, Magnitude of XY gradients.
|
|
"""
|
|
grad_x = cv2.Sobel(self.img, cv2.CV_64F, 1, 0, borderType=cv2.BORDER_DEFAULT)
|
|
grad_y = cv2.Sobel(self.img, cv2.CV_64F, 0, 1, borderType=cv2.BORDER_DEFAULT)
|
|
grad = np.sqrt(grad_x ** 2 + grad_y ** 2)
|
|
|
|
return DepthImage(grad_x), DepthImage(grad_y), DepthImage(grad)
|
|
|
|
def normalise(self):
|
|
"""
|
|
Normalise by subtracting the mean and clippint [-1, 1]
|
|
"""
|
|
self.img = np.clip((self.img - self.img.mean()), -1, 1)
|
|
|
|
|
|
class WidthImage(Image):
|
|
"""
|
|
A width image is one that describes the desired gripper width at each pixel.
|
|
"""
|
|
|
|
def zoom(self, factor):
|
|
"""
|
|
"Zoom" the image by cropping and resizing. Also scales the width accordingly.
|
|
:param factor: Factor to zoom by. e.g. 0.5 will keep the center 50% of the image.
|
|
"""
|
|
super().zoom(factor)
|
|
self.img = self.img / factor
|
|
|
|
def normalise(self):
|
|
"""
|
|
Normalise by mapping [0, 150] -> [0, 1]
|
|
"""
|
|
self.img = np.clip(self.img, 0, 150.0) / 150.0
|