-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
86 lines (71 loc) · 3.14 KB
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import cv2
import numpy as np
import torch
# def normalize_depth(depth, min_val=250.0, max_val=2000.0):#min_val=250.0, max_val=1500.0):
def normalize_depth(depth, min_val=250.0, max_val=2000.0):
""" normalize the input depth (mm) and return depth image (0 ~ 255)
Args:
depth ([np.float]): depth array [H, W] (mm)
min_val (float, optional): [min depth]. Defaults to 250 mm
max_val (float, optional): [max depth]. Defaults to 1500 mm.
Returns:
[np.uint8]: normalized depth array [H, W, 3] (0 ~ 255)
"""
depth[depth < min_val] = min_val
depth[depth > max_val] = max_val
depth = (depth - min_val) / (max_val - min_val) * 255
depth = np.expand_dims(depth, -1)
depth = np.uint8(np.repeat(depth, 3, -1))
return depth
def unnormalize_depth(depth, min_val=250.0, max_val=2000.0):
""" unnormalize the input depth (0 ~ 255) and return depth image (mm)
Args:
depth([np.uint8]): normalized depth array [H, W, 3] (0 ~ 255)
min_val (float, optional): [min depth]. Defaults to 250 mm
max_val (float, optional): [max depth]. Defaults to 1500 mm.
Returns:
[np.float]: depth array [H, W] (mm)
"""
depth = np.float32(depth) / 255
depth = depth * (max_val - min_val) + min_val
return depth
def inpaint_depth(depth, factor=1, kernel_size=3, dilate=False):
""" inpaint the input depth where the value is equal to zero
Args:
depth ([np.uint8]): normalized depth array [H, W, 3] (0 ~ 255)
factor (int, optional): resize factor in depth inpainting. Defaults to 4.
kernel_size (int, optional): kernel size in depth inpainting. Defaults to 5.
Returns:
[np.uint8]: inpainted depth array [H, W, 3] (0 ~ 255)
"""
H, W, _ = depth.shape
resized_depth = cv2.resize(depth, (W//factor, H//factor))
mask = np.all(resized_depth == 0, axis=2).astype(np.uint8)
if dilate:
mask = cv2.dilate(mask, np.ones((kernel_size, kernel_size), np.uint8), iterations=1)
inpainted_data = cv2.inpaint(resized_depth, mask, kernel_size, cv2.INPAINT_TELEA)
inpainted_data = cv2.resize(inpainted_data, (W, H))
depth = np.where(depth == 0, inpainted_data, depth)
return depth
def array_to_tensor(array):
""" Converts a numpy.ndarray (N x H x W x C) to a torch.FloatTensor of shape (N x C x H x W)
OR
converts a nump.ndarray (H x W x C) to a torch.FloatTensor of shape (C x H x W)
"""
if array.ndim == 4: # NHWC
tensor = torch.from_numpy(array).permute(0,3,1,2).float()
elif array.ndim == 3: # HWC
tensor = torch.from_numpy(array).permute(2,0,1).float()
else: # everything else
tensor = torch.from_numpy(array).float()
return tensor
def standardize_image(image):
""" Convert a numpy.ndarray [H x W x 3] of images to [0,1] range, and then standardizes
@return: a [H x W x 3] numpy array of np.float32
"""
image_standardized = np.zeros_like(image).astype(np.float32)
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]
for i in range(3):
image_standardized[...,i] = (image[...,i]/255. - mean[i]) / std[i]
return image_standardized