depth_prediction.py
Go to the documentation of this file.
1 from __future__ import division
2 
3 import os
4 import os.path as osp
5 
6 import chainer
7 import cv2
8 import numpy as np
9 import PIL.Image
10 import skimage.io
11 
12 
13 class DepthPredictionDataset(chainer.dataset.DatasetMixin):
14 
15  class_names = np.array([
16  '_background_',
17  'mirror',
18  ], dtype=np.str)
19  class_names.setflags(write=0)
20 
21  _files = set([
22  'depth.npz',
23  'depth_gt.npz',
24  'image.png',
25  'label.png'
26  ])
27 
28  mean_bgr = np.array([104.00698793, 116.66876762, 122.67891434])
29  min_value = 0.5
30  max_value = 5.0
31 
32  def __init__(self, root_dir, split, aug=False):
33  self.root_dir = root_dir
34  assert split in ['train', 'test']
35  self.split = split
36  self.aug = aug
37 
38  self._files_dirs = []
39  self._scenes = []
40  for date_dir in sorted(os.listdir(self.root_dir)):
41  date_dir = osp.join(self.root_dir, date_dir)
42  split_dir = osp.join(date_dir, split)
43  for files_dir in sorted(os.listdir(split_dir)):
44  files_dir = osp.join(split_dir, files_dir)
45  files = set(os.listdir(files_dir))
46  assert self._files.issubset(files), (
47  'In {}: File set does not match.\n'
48  'Expected: {}\nActual: {}'
49  .format(files_dir, self._files, files))
50  self._files_dirs.append(files_dir)
51 
52  def __len__(self):
53  return len(self._files_dirs)
54 
55  def get_example(self, i):
56  return self._get_example(self._files_dirs[i], i)
57 
58  def _get_example(self, files_dir, idx):
59  image_file = osp.join(files_dir, 'image.png')
60  image = skimage.io.imread(image_file)
61  assert image.dtype == np.uint8
62  assert image.ndim == 3
63 
64  depth_file = osp.join(files_dir, 'depth.npz')
65  depth = np.load(depth_file)['arr_0']
66  depth[depth == 0] = np.nan
67  if depth.dtype == np.uint16:
68  depth = depth.astype(np.float32)
69  depth /= 1000
70  depth_keep = ~np.isnan(depth)
71  depth[depth_keep] = np.maximum(depth[depth_keep], self.min_value)
72  depth[depth_keep] = np.minimum(depth[depth_keep], self.max_value)
73  assert depth.dtype == np.float32
74  assert depth.ndim == 2
75 
76  label_file = osp.join(files_dir, 'label.png')
77  with open(label_file, 'r') as f:
78  label = np.asarray(PIL.Image.open(f)).astype(np.int32)
79  assert label.dtype == np.int32
80  assert label.ndim == 2
81 
82  depth_gt_file = osp.join(files_dir, 'depth_gt.npz')
83  depth_gt = np.load(depth_gt_file)['arr_0']
84  depth_gt[depth_gt == 0] = np.nan
85  if depth_gt.dtype == np.uint16:
86  depth_gt = depth_gt.astype(np.float32)
87  depth /= 1000
88  depth_gt_keep = ~np.isnan(depth_gt)
89  depth_gt[depth_gt_keep] = np.maximum(
90  depth_gt[depth_gt_keep], self.min_value)
91  depth_gt[depth_gt_keep] = np.minimum(
92  depth_gt[depth_gt_keep], self.max_value)
93  assert depth_gt.dtype == np.float32
94  assert depth_gt.ndim == 2
95 
96  # Data augmentation
97  if self.aug:
98  # 1. Color augmentation
99  np.random.seed()
100  proba_color = 0.3
101  # 1-1. Add [-50, 50) to R, G, B channel each
102  if np.random.uniform() < proba_color:
103  image = image.astype(np.float64)
104  image[:, :, 0] += np.random.uniform() * 100 - 50
105  image[:, :, 1] += np.random.uniform() * 100 - 50
106  image[:, :, 2] += np.random.uniform() * 100 - 50
107  image = np.clip(image, 0, 255)
108  image = image.astype(np.uint8)
109  # 1-2. Multiply [0.5, 2) for S, V channel each
110  if np.random.uniform() < proba_color:
111  image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
112  image = image.astype(np.float64)
113  image[:, :, 1] *= np.random.uniform() * 1.5 + 0.5
114  image[:, :, 2] *= np.random.uniform() * 1.5 + 0.5
115  image = np.clip(image, 0, 255)
116  image = image.astype(np.uint8)
117  image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
118  # 1-3. Gaussian blur
119  if np.random.uniform() < proba_color:
120  image = cv2.GaussianBlur(image, (5, 5), np.random.uniform())
121  # 1-4. Add gaussian noise
122  if np.random.uniform() < proba_color:
123  image = image.astype(np.float64)
124  h = image.shape[0]
125  w = image.shape[1]
126  image[:, :, 0] += np.random.normal(
127  0, np.random.uniform() * 0.1 * 255, (h, w))
128  image[:, :, 1] += np.random.normal(
129  0, np.random.uniform() * 0.1 * 255, (h, w))
130  image[:, :, 2] += np.random.normal(
131  0, np.random.uniform() * 0.1 * 255, (h, w))
132  image = np.clip(image, 0, 255)
133  image = image.astype(np.uint8)
134 
135  # 2. Depth noise
136  if np.random.uniform() < 0.3:
137  noise_rate = np.random.uniform() * 0.25 + 0.05
138  depth[
139  np.random.rand(depth.shape[0], depth.shape[1]) < noise_rate
140  ] = np.nan
141 
142  # 3. Geometric augmentation
143  if np.random.uniform() < 0.5:
144  image = np.fliplr(image)
145  depth = np.fliplr(depth)
146  label = np.fliplr(label)
147  depth_gt = np.fliplr(depth_gt)
148  if np.random.uniform() < 0.5:
149  image = np.flipud(image)
150  depth = np.flipud(depth)
151  label = np.flipud(label)
152  depth_gt = np.flipud(depth_gt)
153  if np.random.uniform() < 0.5:
154  angle = (np.random.uniform() * 180) - 90
155  image = self.rotate_image(image, angle, cv2.INTER_LINEAR)
156  depth = self.rotate_depth_image(depth, angle, cv2.INTER_LINEAR)
157  label = self.rotate_image(label, angle, cv2.INTER_NEAREST)
158  depth_gt = self.rotate_depth_image(
159  depth_gt, angle, cv2.INTER_LINEAR)
160 
161  return image, depth, label, depth_gt, idx
162 
163  def rotate_image(self, in_img, angle, flags=cv2.INTER_LINEAR):
164  rot_mat = cv2.getRotationMatrix2D(
165  center=(in_img.shape[1] / 2, in_img.shape[0] / 2),
166  angle=angle, scale=1)
167  rot_img = cv2.warpAffine(
168  src=in_img, M=rot_mat,
169  dsize=(in_img.shape[1], in_img.shape[0]), flags=flags)
170  return rot_img
171 
172  def rotate_depth_image(self, in_img, angle, flags=cv2.INTER_LINEAR):
173  rot_mat = cv2.getRotationMatrix2D(
174  center=(in_img.shape[1] / 2, in_img.shape[0] / 2),
175  angle=angle, scale=1)
176  ones = np.ones(in_img.shape, dtype=np.int32)
177  rot_keep = cv2.warpAffine(
178  src=ones, M=rot_mat,
179  dsize=(in_img.shape[1], in_img.shape[0]),
180  flags=cv2.INTER_NEAREST)
181  rot_keep = rot_keep.astype(np.bool)
182  rot_img = cv2.warpAffine(
183  src=in_img, M=rot_mat,
184  dsize=(in_img.shape[1], in_img.shape[0]), flags=flags)
185  rot_img[rot_keep == False] = np.nan # NOQA
186  return rot_img
def rotate_image(self, in_img, angle, flags=cv2.INTER_LINEAR)
def rotate_depth_image(self, in_img, angle, flags=cv2.INTER_LINEAR)


jsk_recognition_utils
Author(s):
autogenerated on Wed Mar 3 2021 04:00:53