+ training, inference, dataset, exporters, importers, models, supervi…

…sion and utils scripts added

dataset/__init__.py

@@ -0,0 +1 @@
+from .dataloader import *

dataset/dataloader.py

@@ -0,0 +1,137 @@
+import os
+import sys
+import torch
+from torch.utils.data.dataset import Dataset
+from torch.utils.data import DataLoader
+import importers
+import warnings
+
+#testing
+# sys.path.append('E:\\Projects\\vsc\\deep_depth_denoising\\denoise')
+# import importers
+
+
+'''
+Dataset importer. We assume that data follows the below structure.
+root_path
+device_repository.json
+	|
+	|-----recording_i
+	|		|-----Data
+	|		|-----Calibration
+	|
+	|-----recording_i+1
+	|		|-----Data
+	|		|-----Calibration
+	|
+'''
+
+class DataLoaderParams:
+	def __init__(self
+	,root_path 
+	,device_list
+	,decimation_scale = 2
+	,device_repository_path = "."
+	,depth_scale = 0.001
+	,depth_threshold = 5):
+		self.root_path = root_path
+		self.device_list = device_list
+		self.device_repository_path = device_repository_path
+		self.depth_scale = depth_scale
+		self.decimation_scale = decimation_scale
+		self.depth_threshold = depth_threshold
+
+class DataLoad(Dataset):
+	def __init__(self, params):
+		super(DataLoad,self).__init__()
+		self.params = params
+
+		device_repo_path = os.path.join(self.params.device_repository_path,"device_repository.json")
+		if not os.path.exists(device_repo_path):
+			raise ValueError("{} does not exist, exiting.".format(device_repo_path))
+		self.device_repository = importers.intrinsics.load_intrinsics_repository(device_repo_path)
+
+		root_path = self.params.root_path
+
+
+		if not os.path.exists(root_path):
+			raise ValueError("{} does not exist, exiting.".format(root_path))
+
+		self.data = {}
+
+		# iterate over each recorded folder
+		for recording in os.listdir(root_path):
+			abs_recording_path = os.path.join(root_path,recording)
+			if not os.path.isdir(abs_recording_path):
+				continue
+			# path where data supposed to be stored
+			data_path = os.path.join(abs_recording_path,"Data")
+
+			if not os.path.exists(data_path):
+				warnings.warn("Folder {} does not containt \"Data\" folder".format(abs_recording_path))
+				continue
+
+			# path to the calibration of that particular recording
+			calibration_path = os.path.join(abs_recording_path,"Calibration")
+			if not os.path.exists(calibration_path):
+				warnings.warn("Folder {} does not containt \"Calibration\" folder".format(calibration_path))
+				continue
+
+			# data iteration
+			for file in os.listdir(data_path):
+				full_filename = os.path.join(data_path,file)
+
+				_, ext = os.path.splitext(full_filename)
+				if ext != ".png":
+					continue
+
+				_id,_name,_type,_ = file.split("_")
+				unique_name = recording + "-" + str(_id)
+
+				# skip names that we do not want to load
+				if _name not in self.params.device_list:
+					continue
+
+				if unique_name not in self.data:
+					self.data[unique_name] = {}
+					self.data[unique_name]["calibration"] = calibration_path
+
+				if _name not in self.data[unique_name]:
+					self.data[unique_name][_name] = {}
+
+				self.data[unique_name][_name][_type] = full_filename
+
+
+	def __len__(self):
+		return len(self.data)
+
+	def __getitem__(self, idx):
+		#get an entry
+		key = list(self.data.keys())[idx]
+		datum = self.data[key]
+
+		datum_out = {}
+
+		for device in self.params.device_list:
+			color_img = importers.image.load_image(datum[device]["color"])
+			depth_range_mask = (importers.image.load_depth(datum[device]["depth"]) < self.params.depth_threshold).float()
+			depth_img = importers.image.load_depth(datum[device]["depth"]) * depth_range_mask
+			intrinsics, intrinsics_inv = importers.intrinsics.get_intrinsics(\
+				device, self.device_repository, self.params.decimation_scale)
+			extrinsics, extrinsics_inv = importers.extrinsics.load_extrinsics(\
+				os.path.join(datum["calibration"], device + ".extrinsics"))
+
+			datum_out.update({
+				device : {
+				"color" : color_img.squeeze(0),
+				"depth" : depth_img.squeeze(0), 
+				"intrinsics" : intrinsics,
+				"intrinsics_inv" : intrinsics_inv,
+				"extrinsics" : extrinsics,
+				"extrinsics_inv" : extrinsics_inv
+				}})
+
+		return datum_out
+
+	def get_data(self):
+		return self.data

exporters/__init__.py

@@ -0,0 +1,15 @@
+from .image import *
+from .point_cloud import *
+
+def save_splatted_info(inout, path, super_list):
+    def generator(name, type, ext):
+        return os.path.join(path, name + "_" + type + "_#" + ext)
+    for key in inout.keys():
+        save_image(generator(key, "mask",".png"), inout[key]["depth"]["mask"])
+        save_image(generator(key, "masked_color",".png"), inout[key]["color"]["masked"])
+        save_depth(generator(key, "depth", ".png"), inout[key]["depth"]["original"])
+        save_data(generator(key, "prediction", ".exr"), inout[key]["depth"]["prediction"])
+        others = [other for other in inout.keys() if other != key]
+        for other in others:
+            save_image(generator(key, "splatted",".png"), inout[key]["color"]["splatted"])
+

exporters/image.py

@@ -0,0 +1,64 @@
+import torch
+import cv2
+import numpy
+
+def save_image(filename, tensor, scale=255.0):
+    b, _, __, ___ = tensor.size()
+    for n in range(b):
+        array = tensor[n, :, :, :].detach().cpu().numpy()
+        array = array.transpose(1, 2, 0) * scale
+        cv2.imwrite(filename.replace("#", str(n)), array)
+
+def save_depth(filename, tensor, scale=1000.0):
+    b, _, __, ___ = tensor.size()
+    for n in range(b):
+        array = tensor[n, :, :, :].detach().cpu().numpy()
+        array = array.transpose(1, 2, 0) * scale
+        array = numpy.uint16(array)
+        cv2.imwrite(filename.replace("#", str(n)), array)
+
+def save_data(filename, tensor, scale=1000.0):
+    b, _, __, ___ = tensor.size()
+    for n in range(b):
+        array = tensor[n, :, :, :].detach().cpu().numpy()
+        array = array.transpose(1, 2, 0) * scale
+        array = numpy.float32(array)
+        cv2.imwrite(filename.replace("#", str(n)), array)
+
+def save_depth_from_3d(filename, tensor, scale=1000.0):
+    b, _, __, ___ = tensor.size()
+    for n in range(b):
+        array = tensor[n, :, :, :].detach().cpu().numpy()
+        depth_channel = numpy.zeros((1, array.shape[1], array.shape[2]))
+        depth_channel[0,:,:] = array[2, :, :]
+        depth_channel = depth_channel.transpose(1, 2, 0) * scale
+        depth_channel = numpy.uint16(depth_channel)
+        cv2.imwrite(filename.replace("#", str(n)), depth_channel)
+
+def save_normals(filename, tensor, scale=255.0):
+    b, _, __, ___ = tensor.size()    
+    for n in range(b):   
+        normals = tensor[n, :, :, :].detach().cpu().numpy()
+        # transpose for consistent rendering, multiplied by scale
+        normals = (normals.transpose(1, 2, 0) + 1) * scale // 2 + 1
+        # image write
+        cv2.imwrite(filename.replace("#", str(n)), normals)
+
+def save_phong_normals(filename, tensor):
+    b, _, __, ___ = tensor.size()
+    for n in range(b):            
+        # the z-component data of each normal vector is retrieved 
+        z_comp = tensor[n, 2, :, :].detach().cpu().numpy()
+        # phong image (1, z_comp.shape[0], z_comp.shape[1]) is initialized
+        phong = numpy.zeros((1, z_comp.shape[0], z_comp.shape[1]))
+        # z value is inversed to paint properly based on misalignment from camera's FOV direction 
+        phong[0,:,:] = 1 - z_comp
+        # get max uint16 value
+        iui16 = numpy.iinfo(numpy.uint16)
+        scale = iui16.max
+        # transpose for consistent rendering
+        phong = phong.transpose(1, 2, 0) * scale
+        # to unsigned int16 
+        phong = numpy.uint16(phong)
+        # image write
+        cv2.imwrite(filename.replace("#", str(n)), phong)

exporters/point_cloud.py

@@ -0,0 +1,48 @@
+import torch
+import os
+
+def save_ply(filename, tensor, scale, color='black' , normals = None):    
+    b, _, h, w = tensor.size()
+    for n in range(b):
+        coords = tensor[n, :, :, :].detach().cpu().numpy()
+        x_coords = coords[0, :] * scale
+        y_coords = coords[1, :] * scale
+        z_coords = coords[2, :] * scale
+        if normals is not None:
+            norms = normals[n, : , : , :].detach().cpu().numpy()
+            nx_coords = norms[0, :]
+            ny_coords = norms[1, :]
+            nz_coords = norms[2, :]
+        with open(filename.replace("#", str(n)), "w") as ply_file:        
+            ply_file.write("ply\n")
+            ply_file.write("format ascii 1.0\n")
+            ply_file.write("element vertex {}\n".format(w * h))
+            ply_file.write("property float x\n")
+            ply_file.write("property float y\n")
+            ply_file.write("property float z\n")
+            if normals is not None:
+                ply_file.write('property float nx\n')
+                ply_file.write('property float ny\n')
+                ply_file.write('property float nz\n')
+            ply_file.write("property uchar red\n")
+            ply_file.write("property uchar green\n")
+            ply_file.write("property uchar blue\n")
+            ply_file.write("end_header\n")
+
+            if normals is None:
+                for x in torch.arange(w):
+                    for y in torch.arange(h):
+                        ply_file.write("{} {} {} {} {} {}\n".format(\
+                            x_coords[y, x], y_coords[y, x], z_coords[y, x],\
+                            "255" if color=='red' else "0",
+                            "255" if color=='green' else "0",
+                            "255" if color=='blue' else "0"))
+            else:
+                for x in torch.arange(w):
+                    for y in torch.arange(h):
+                        ply_file.write("{} {} {} {} {} {} {} {} {}\n".format(\
+                            x_coords[y, x], y_coords[y, x], z_coords[y, x],\
+                            nx_coords[y, x], ny_coords[y, x], nz_coords[y, x],\
+                            "255" if color=='red' else "0",
+                            "255" if color=='green' else "0",
+                            "255" if color=='blue' else "0"))

importers/__init__.py

@@ -0,0 +1,3 @@
+from .extrinsics import *
+from .intrinsics import *
+from .image import *

importers/extrinsics.py

@@ -0,0 +1,11 @@
+import numpy
+import torch
+
+def load_extrinsics(filename, data_type=torch.float32):
+    data = numpy.loadtxt(filename)
+    pose = numpy.zeros([4, 4])
+    pose[3, 3] = 1
+    pose[:3, :3] = data[:3, :3]
+    pose[:3, 3] = data[3, :]
+    extrinsics = torch.tensor(pose, dtype=data_type)
+    return extrinsics , extrinsics.inverse()

importers/image.py

@@ -0,0 +1,28 @@
+import cv2
+import torch
+import numpy
+
+def load_image(filename, data_type=torch.float32):
+    color_img = numpy.array(cv2.imread(filename, cv2.IMREAD_ANYCOLOR))
+    h, w, c = color_img.shape
+    color_data = color_img.astype(numpy.float32).transpose(2, 0, 1)
+    return torch.from_numpy(
+        color_data.reshape(1, c, h, w)        
+    ).type(data_type) / 255.0
+
+def load_depth(filename, data_type=torch.float32, scale=0.001):
+    depth_img = numpy.array(cv2.imread(filename, cv2.IMREAD_ANYDEPTH))
+    h, w = depth_img.shape
+    depth_data = depth_img.astype(numpy.float32) * scale
+    return torch.from_numpy(
+        depth_data.reshape(1, 1, h, w)        
+    ).type(data_type)
+
+def crop_depth(filename, data_type=torch.float32, scale=0.001):
+    depth_img = numpy.array(cv2.imread(filename, cv2.IMREAD_ANYDEPTH))
+    center_cropped_depth_img = depth_img[60:420, 0:640]
+    h, w = center_cropped_depth_img.shape
+    depth_data = center_cropped_depth_img.astype(numpy.float32) * scale
+    return torch.from_numpy(
+        depth_data.reshape(1, 1, h, w)        
+    ).type(data_type)

importers/intrinsics.py

@@ -0,0 +1,27 @@
+import json
+import numpy
+import torch
+
+
+def load_intrinsics_repository(filename):    
+    #global intrinsics_dict
+    with open(filename, 'r') as json_file:
+        intrinsics_repository = json.load(json_file)
+        intrinsics_dict = dict((intrinsics['Device'], \
+            intrinsics['Depth Intrinsics'][0]['1280x720'])\
+                for intrinsics in intrinsics_repository)
+    return intrinsics_dict
+
+def get_intrinsics(name, intrinsics_dict, scale=1, data_type=torch.float32):
+    #global intrinsics_dict
+    if intrinsics_dict is not None:
+        intrinsics_data = numpy.array(intrinsics_dict[name])
+        intrinsics = torch.tensor(intrinsics_data).reshape(3, 3).type(data_type)    
+        intrinsics[0, 0] = intrinsics[0, 0] / scale
+        intrinsics[0, 2] = intrinsics[0, 2] / scale
+        intrinsics[1, 1] = intrinsics[1, 1] / scale
+        intrinsics[1, 2] = intrinsics[1, 2] / scale
+        intrinsics_inv = intrinsics.inverse()
+        return intrinsics, intrinsics_inv
+    raise ValueError("Intrinsics repository is empty")
+