Object filling method speed up, using cupy for GPU boost.

Author: Meli-0xFF

create_dataset.py

@@ -10,6 +10,6 @@
                                 def_maps=True,
                                 fill=True,
                                 canny=False,
-                                fill_texture=True)
+                                fill_tex=True)
 
 create_dataset_norm_data(dataset_name)

dataset.py

@@ -7,7 +7,7 @@
 from data_preparation.normalization import Normalization
 from data_preparation.filling import fill_depth_map
 from metrics import get_canny_mask
-from object_filling import fill_depth_map, fill_hr_texture
+from object_filling import fill_depth_map, fill_texture
 
 
 class DepthMapSRDataset(Dataset):
@@ -85,7 +85,7 @@ def __getitem__(self, idx):
     return sample[0], sample[1], sample[2], sample[3], sample[4]#, sample[5]
 
 
-def create_dataset(name, hr_dir, lr_dir, textures_dir, scale_lr=True, fill=True, def_maps=False, canny=False, fill_texture=False):
+def create_dataset(name, hr_dir, lr_dir, textures_dir, scale_lr=True, fill=True, def_maps=False, canny=False, fill_tex=False):
   print("--- CREATE DATASET: " + name + " ---")
   data = dict()
   hr_depth_maps = None
@@ -164,10 +164,10 @@ def create_dataset(name, hr_dir, lr_dir, textures_dir, scale_lr=True, fill=True,
     idx += 1
   data['tx'] = textures
 
-  if fill_texture:
+  if fill_tex:
     for i in range(len(data["hr"])):
       print("> Augmenting HR texture " + str(i + 1) + "/" + str(len(data["hr"])))
-      data['tx'][i] = fill_hr_texture(data['tx'][i], def_maps[i])
+      data['tx'][i] = fill_texture(data['tx'][i], def_maps[i])
 
   if canny:
     canny_masks = np.empty((len(data["hr"]), data["hr"][0].shape[0], data["hr"][0].shape[1]), dtype=float)

filling_inference.py

@@ -0,0 +1,48 @@
+from dataset import DepthMapSRDataset
+from torch.utils.data import DataLoader
+from evaluation.pointcloud import *
+import torch
+import matplotlib.pyplot as plt
+import matplotlib as mpl
+from object_filling import fill_depth_map, fill_texture
+import cupy as cp
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+dataset_name = 'NEAREST-LED-WARIOR-scale_4-ACTUAL'
+norm_file_path = 'dataset/' + dataset_name + '_norm.npy'
+
+assert os.path.isfile(norm_file_path), "Normalization file for dataset '" + dataset_name + "' does not exist"
+norm_data = np.load(norm_file_path, allow_pickle=True).tolist()
+depth_max = norm_data["depth"][1]
+
+dataset = DepthMapSRDataset(dataset_name, train=False, task='depth_map_sr', norm=False)
+dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
+
+lr_depth_map, texture, hr_depth_map, def_map, object_mask = next(iter(dataloader))
+unfilled_hr_depth_map = (hr_depth_map * def_map).clone()
+
+img = np.array((hr_depth_map * def_map)[0][0].float().numpy())
+tex = np.array((texture / torch.max(texture))[0][0].float().numpy())
+filled, object_map = fill_depth_map(img, depth_max)
+filled_texture = fill_texture(tex, def_map[0][0].float().numpy())
+
+cmap = mpl.cm.get_cmap("winter").copy()
+cmap.set_under(color='black')
+
+hr_pcl_unfilled = PointCloud(unfilled_hr_depth_map[0][0].numpy())
+hr_pcl_unfilled.create_ply("UNFILLED-PLANE-hr-ptcloud-actual")
+
+hr_pcl = PointCloud(cp.asnumpy(filled))
+hr_pcl.create_ply("FILLED-PLANE-hr-ptcloud-actual")
+
+plt.figure(plt.figure('HR Depth map UNFILLED'))
+plt.imshow(unfilled_hr_depth_map[0][0], cmap=cmap, vmin=0.0001)
+
+plt.figure(plt.figure('HR Depth map'))
+plt.imshow(cp.asnumpy(filled), cmap=cmap, vmin=0.0001)
+
+plt.figure(plt.figure('HR Filled Texture'))
+plt.imshow(cp.asnumpy(filled_texture), cmap='gray')
+
+plt.show()

object_filling.py

@@ -5,31 +5,31 @@
 import skimage.io, skimage.transform
 import matplotlib.pyplot as plt
 import matplotlib as mpl
-
+import cupy as cp
 
 def get_plane(img):
   img_loader_x = skimage.io.imread('scanner/normalized_vectors_x_plus_half.tif')
   img_loader_y = skimage.io.imread('scanner/normalized_vectors_y_plus_half.tif')
-  nv_x = skimage.img_as_float(img_loader_x) - 0.5
-  nv_y = skimage.img_as_float(img_loader_y) - 0.5
+  nv_x = cp.asarray(skimage.img_as_float(img_loader_x) - 0.5)
+  nv_y = cp.asarray(skimage.img_as_float(img_loader_y) - 0.5)
 
-  img_mean = np.mean(img[img != 0])
-  mean_error_map = np.power((img - img_mean), 2)
-  mean_error_map = np.where(img > 0, mean_error_map, np.infty)
+  img_mean = cp.mean(img[img != 0])
+  mean_error_map = cp.power((img - img_mean), 2)
+  mean_error_map = cp.where(img > 0, mean_error_map, cp.infty)
 
-  mean_sorted = np.dstack(np.unravel_index(np.argsort(mean_error_map.ravel()), img.shape))[0]
-  depth_sorted = np.dstack(np.unravel_index(np.argsort(img.ravel()), img.shape))[0]
+  mean_sorted = cp.dstack(cp.unravel_index(cp.argsort(mean_error_map.ravel()), img.shape))[0]
+  depth_sorted = cp.dstack(cp.unravel_index(cp.argsort(img.ravel()), img.shape))[0]
   a = mean_sorted[:20000]
   b = depth_sorted
   acceptable_mean = mean_error_map[a[19999, 0], a[19999, 1]]
   acceptable_depth = img[b[440000, 0], b[440000, 1]]
 
-  img1 = np.where(img == 0, 1.0, img)
+  img1 = cp.where(img == 0, 1.0, img)
   x = nv_x * img1
   y = nv_y * img1 * (-1.0)
   z = img1.copy()
 
-  mesh_points = np.empty((int(img.shape[0] / 20) * int(img.shape[1] / 20), 2))
+  mesh_points = cp.empty((int(img.shape[0] / 20) * int(img.shape[1] / 20), 2))
   for i in range(int(img.shape[0] / 20)):
     for j in range(int(img.shape[1] / 20)):
       mesh_points[int(i * img.shape[1] / 20) + j, 0] = i * 20
@@ -40,15 +40,15 @@ def get_plane(img):
   for i in range(mesh_points.shape[0]):
     if mean_error_map[int(mesh_points[i, 0]), int(mesh_points[i, 1])] < acceptable_mean:
       if mean_points is None:
-        mean_points = np.array([mesh_points[i]])
+        mean_points = cp.array([mesh_points[i]])
       else:
-        mean_points = np.vstack([mean_points, mesh_points[i]])
+        mean_points = cp.vstack([mean_points, mesh_points[i]])
 
     if img[int(mesh_points[i, 0]), int(mesh_points[i, 1])] > acceptable_depth:
       if depth_points is None:
-        depth_points = np.array([mesh_points[i]])
+        depth_points = cp.array([mesh_points[i]])
       else:
-        depth_points = np.vstack([depth_points, mesh_points[i]])
+        depth_points = cp.vstack([depth_points, mesh_points[i]])
 
   mean_points = mean_points[mean_points[:, 1].argsort()]
   depth_points = depth_points[depth_points[:, 1].argsort()]
@@ -68,27 +68,27 @@ def get_plane(img):
   plt.show()
   '''
 
-  mean_X = np.empty((mean_points.shape[0]))
-  mean_Y = np.empty((mean_points.shape[0]))
-  mean_Z = np.empty((mean_points.shape[0]))
+  mean_X = cp.empty((mean_points.shape[0]))
+  mean_Y = cp.empty((mean_points.shape[0]))
+  mean_Z = cp.empty((mean_points.shape[0]))
 
   for i in range(mean_points.shape[0]):
     mean_X[i] = x[int(mean_points[i, 0]), int(mean_points[i, 1])]
     mean_Y[i] = y[int(mean_points[i, 0]), int(mean_points[i, 1])]
     mean_Z[i] = z[int(mean_points[i, 0]), int(mean_points[i, 1])]
 
-  real_mean_points = np.c_[mean_X, mean_Y, mean_Z]
+  real_mean_points = cp.c_[mean_X, mean_Y, mean_Z]
 
-  depth_X = np.empty((depth_points.shape[0]))
-  depth_Y = np.empty((depth_points.shape[0]))
-  depth_Z = np.empty((depth_points.shape[0]))
+  depth_X = cp.empty((depth_points.shape[0]))
+  depth_Y = cp.empty((depth_points.shape[0]))
+  depth_Z = cp.empty((depth_points.shape[0]))
 
   for i in range(depth_points.shape[0]):
     depth_X[i] = x[int(depth_points[i, 0]), int(depth_points[i, 1])]
     depth_Y[i] = y[int(depth_points[i, 0]), int(depth_points[i, 1])]
     depth_Z[i] = z[int(depth_points[i, 0]), int(depth_points[i, 1])]
 
-  real_depth_points = np.c_[depth_X, depth_Y, depth_Z]
+  real_depth_points = cp.c_[depth_X, depth_Y, depth_Z]
 
   A = real_mean_points[0]
   B = real_depth_points[0]
@@ -107,61 +107,56 @@ def get_plane(img):
   plt.show()
   '''
 
-  normal = np.cross(v1, v2)
+  normal = cp.cross(v1, v2)
 
-  d = -np.sum(normal * A)
+  d = -cp.sum(normal * A)
   #plane = -(normal[0] * x + normal[1] * y + d) / normal[2]
   full_plane = -d / (normal[0] * nv_x + normal[1] * nv_y * (-1.0) + normal[2])
-  img_dis = abs(normal[0] * x + + normal[1] * y + normal[2] * img + d) / np.sqrt(normal[0]**2 + normal[1]**2 + normal[2]**2)
-  img_dis = np.where(img == 0, 0, img_dis)
+  img_dis = abs(normal[0] * x + + normal[1] * y + normal[2] * img + d) / cp.sqrt(normal[0]**2 + normal[1]**2 + normal[2]**2)
+  img_dis = cp.where(img == 0, 0, img_dis)
 
   return full_plane, img_dis
 
 
 def fill_depth_map(img, background_value):
   mask = np.where(img > 0, 0, 1).astype(np.uint8)
-  uint_mask = np.where(img > 0, 0, 255).astype(np.uint8)
-  rgb_mask = np.dstack((uint_mask, uint_mask, uint_mask))
 
   labels, holes = measure.label(mask, background=0, return_num=True)
-
+  img = cp.asarray(img)
+  labels = cp.asarray(labels)
   print("Filling " + str(holes) + " holes")
   with alive_bar(holes) as bar:
     for i in range(1, holes + 1):
       label = labels == i
-      hole_border = segmentation.mark_boundaries(np.zeros(rgb_mask.shape),
-                                                 label.astype(np.int32),
-                                                 (1, 0, 0),
-                                                 None,
-                                                 'outer')[:, :, 0].astype(float)
-
-      #hole_border1 = np.logical_or((np.cumsum(label.astype(np.int32), axis=0) == 1), (np.cumsum(label.astype(np.int32), axis=1) == 1))
-      #hole_border2 = np.logical_or(np.flipud(np.flipud(label.astype(np.int32).cumsum(axis=1)) == 1), np.fliplr(np.fliplr(label.astype(np.int32).cumsum(axis=0)) == 1))
-      #hole_border = np.logical_or(hole_border1, hole_border2).astype(float)
-
-      border_values = img * hole_border
-      label_img_border_pixels = np.sum(label.astype(np.int32)[0, :]) + \
-                                np.sum(label.astype(np.int32)[label.shape[0] - 1, :]) + \
-                                np.sum(label.astype(np.int32)[1:(label.shape[1] - 1), 0]) + \
-                                np.sum(label.astype(np.int32)[1:(label.shape[1] - 1), label.shape[1] - 1])
+
+      label_img_border_pixels = cp.sum(label.astype(cp.int32)[0, :]) + \
+                                cp.sum(label.astype(cp.int32)[label.shape[0] - 1, :]) + \
+                                cp.sum(label.astype(cp.int32)[1:(label.shape[1] - 1), 0]) + \
+                                cp.sum(label.astype(cp.int32)[1:(label.shape[1] - 1), label.shape[1] - 1])
+
+      if label_img_border_pixels == 0:
+        top = cp.roll(label, 1, axis=0)
+        bottom = cp.roll(label, -1, axis=0)
+        right = cp.roll(label, 1, axis=1)
+        left = cp.roll(label, -1, axis=1)
+        hole_border = cp.logical_or(cp.logical_or(top, bottom), cp.logical_or(right, left))
+        border_values = img * hole_border
 
       background = label_img_border_pixels > 0
 
       if not background:
-        mx = np.ma.masked_array(border_values, mask=border_values == 0)
-        line_values = mx.max(1)
-        hole = (labels == i).astype(float) * line_values[:, np.newaxis]
+        line_values = border_values.max(1)
+        hole = (labels == i).astype(float) * line_values[:, cp.newaxis]
       else:
-
         hole = (labels == i).astype(float) * 0
 
       bar()
 
-      img = np.where(hole > 0, hole, img)
+      img = cp.where(hole > 0, hole, img)
 
   plane, dis = get_plane(img)
-  object_map = np.where(dis < 10, 0, 1.0)
-  img = np.where(img == 0, background_value, img)
+  object_map = cp.where(dis < 10, 0, 1.0)
+  img = cp.where(img == 0, background_value, img)
 
   '''
   cmap = mpl.cm.get_cmap("winter").copy()
@@ -175,39 +170,39 @@ def fill_depth_map(img, background_value):
 
   return img, object_map
 
-def fill_hr_texture(tex, def_map):
+def fill_texture(tex, def_map):
   mask = np.where(def_map > 0, 0, 1).astype(np.uint8)
-  uint_mask = np.where(def_map > 0, 0, 255).astype(np.uint8)
-  rgb_mask = np.dstack((uint_mask, uint_mask, uint_mask))
 
   labels, holes = measure.label(mask, background=0, return_num=True)
-
+  tex = cp.asarray(tex)
+  labels = cp.asarray(labels)
   print("Filling " + str(holes) + " holes")
   with alive_bar(holes) as bar:
     for i in range(1, holes + 1):
       label = labels == i
 
-      hole_border = segmentation.mark_boundaries(np.zeros(rgb_mask.shape),
-                                                 label.astype(np.int32),
-                                                 (1, 0, 0),
-                                                 None,
-                                                 'outer')[:, :, 0].astype(float)
-      border_values = tex * hole_border
-      label_img_border_pixels = np.sum(label.astype(np.int32)[0, :]) + \
-                                np.sum(label.astype(np.int32)[label.shape[0] - 1, :]) + \
-                                np.sum(label.astype(np.int32)[1:(label.shape[1] - 1), 0]) + \
-                                np.sum(label.astype(np.int32)[1:(label.shape[1] - 1), label.shape[1] - 1])
+      label_img_border_pixels = cp.sum(label.astype(cp.int32)[0, :]) + \
+                                cp.sum(label.astype(cp.int32)[label.shape[0] - 1, :]) + \
+                                cp.sum(label.astype(cp.int32)[1:(label.shape[1] - 1), 0]) + \
+                                cp.sum(label.astype(cp.int32)[1:(label.shape[1] - 1), label.shape[1] - 1])
+
+      if label_img_border_pixels == 0:
+        top = cp.roll(label, 1, axis=0)
+        bottom = cp.roll(label, -1, axis=0)
+        right = cp.roll(label, 1, axis=1)
+        left = cp.roll(label, -1, axis=1)
+        hole_border = cp.logical_or(cp.logical_or(top, bottom), cp.logical_or(right, left))
+        border_values = tex * hole_border
 
-      background = label_img_border_pixels > 0 or np.sum((labels == i).astype(float)) > 1000
+      background = label_img_border_pixels > 0
 
       if not background:
-        mx = np.ma.masked_array(border_values, mask=border_values == 0)
-        line_values = mx.max(1)
-        hole = (labels == i).astype(float) * line_values[:, np.newaxis]
+        line_values = border_values.max(1)
+        hole = (labels == i).astype(float) * line_values[:, cp.newaxis]
       else:
-        hole = (labels == i).astype(float) * 1
+        hole = (labels == i).astype(float) * 0
 
       bar()
 
-      tex = np.where(hole > 0, hole, tex)
+      tex = cp.where(hole > 0, hole, tex)
     return tex