add scripts

Author: tzole1155

data splits/scenes_test.txt

@@ -0,0 +1,18 @@
+2t7WUuJeko7
+5ZKStnWn8Zo
+ARNzJeq3xxb
+fzynW3qQPVF
+jtcxE69GiFV
+pa4otMbVnkk
+q9vSo1VnCiC
+rqfALeAoiTq
+UwV83HsGsw3
+wc2JMjhGNzB
+WYY7iVyf5p8
+YFuZgdQ5vWj
+yqstnuAEVhm
+YVUC4YcDtcY
+gxdoqLR6rwA
+gYvKGZ5eRqb
+RPmz2sHmrrY
+Vt2qJdWjCF2

data splits/scenes_train.txt

@@ -0,0 +1,61 @@
+17DRP5sb8fy
+1LXtFkjw3qL
+1pXnuDYAj8r
+29hnd4uzFmX
+5LpN3gDmAk7
+5q7pvUzZiYa
+759xd9YjKW5
+7y3sRwLe3Va
+82sE5b5pLXE
+8WUmhLawc2A
+aayBHfsNo7d
+ac26ZMwG7aT
+B6ByNegPMKs
+b8cTxDM8gDG
+cV4RVeZvu5T
+D7N2EKCX4Sj
+e9zR4mvMWw7
+EDJbREhghzL
+GdvgFV5R1Z5
+gTV8FGcVJC9
+HxpKQynjfin
+i5noydFURQK
+JeFG25nYj2p
+JF19kD82Mey
+jh4fc5c5qoQ
+kEZ7cmS4wCh
+mJXqzFtmKg4
+p5wJjkQkbXX
+Pm6F8kyY3z2
+pRbA3pwrgk9
+PuKPg4mmafe
+PX4nDJXEHrG
+qoiz87JEwZ2
+rPc6DW4iMge
+s8pcmisQ38h
+S9hNv5qa7GM
+sKLMLpTHeUy
+SN83YJsR3w2
+sT4fr6TAbpF
+ULsKaCPVFJR
+uNb9QFRL6hY
+Uxmj2M2itWa
+V2XKFyX4ASd
+VFuaQ6m2Qom
+VVfe2KiqLaN
+Vvot9Ly1tCj
+vyrNrziPKCB
+VzqfbhrpDEA
+XcA2TqTSSAj
+2n8kARJN3HM
+D7G3Y4RVNrH
+dhjEzFoUFzH
+E9uDoFAP3SH
+gZ6f7yhEvPG
+JmbYfDe2QKZ
+r1Q1Z4BcV1o
+r47D5H71a5s
+ur6pFq6Qu1A
+VLzqgDo317F
+YmJkqBEsHnH
+ZMojNkEp431

data splits/scenes_val.txt

@@ -0,0 +1,11 @@
+2azQ1b91cZZ
+8194nk5LbLH
+EU6Fwq7SyZv
+oLBMNvg9in8
+QUCTc6BB5sX
+TbHJrupSAjP
+X7HyMhZNoso
+pLe4wQe7qrG
+x8F5xyUWy9e
+Z6MFQCViBuw
+zsNo4HB9uLZ

dataset/__init__.py

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

dataset/dataloader.py

@@ -0,0 +1,150 @@
+import os
+import sys
+import torch
+from torch.utils.data.dataset import Dataset
+from torch.utils.data import DataLoader
+import importers
+import warnings
+from PIL import Image
+import utils
+
+class DataLoaderParams:
+    def __init__(self,
+        root_path,
+        trajectories_dir,
+        data_split,
+        drone_list,
+        view_list,
+        frame_list,
+        types_list,
+        transform):
+        
+        self.root_path = root_path
+        self.trajectories_dir = trajectories_dir
+        self.data_split = data_split
+        self.drone_list = drone_list
+        self.view_list = view_list
+        self.frame_list = frame_list
+        self.types_list = types_list
+        self.transform = transform
+
+class DataLoad(Dataset):
+    def __init__(self,params):
+        super(DataLoad,self).__init__()
+        self.params = params
+        root_path = self.params.root_path
+        trajectories_dir = self.params.trajectories_dir
+        data_split  = self.params.data_split
+        
+        self.data = {}
+
+        buildings = self.get_splits(data_split)
+
+        for drone in os.listdir(root_path):
+            #ignore not included drone models
+            if drone not in self.params.drone_list:
+                continue
+            root_path_ = os.path.join(root_path,drone)
+            for view in os.listdir(root_path_):
+                view_path = os.path.join(root_path_,view)
+                #get included types
+                for type_ in os.listdir(view_path):
+                    if type_ not in self.params.types_list:
+                            continue
+                    #get images for each building
+                    for building in buildings:
+                        img_path = os.path.join(view_path,type_,building)
+                        for img in os.listdir(img_path):
+                            if (len(img.split("_")) <= 6):
+                                row, date , frame , __ , view , _ = img.split("_")
+                            else:
+                                row, date , frame , __ ,_, view , _ = img.split("_")
+                            #get only included view
+                            if view not in self.params.view_list:
+                                continue
+                            #ignore not listed frame
+                            if int(frame) not in self.params.frame_list:
+                                continue
+                            #NOTE:set the path to the extrinsics file
+                            full_img_name = os.path.join(img_path,img)
+                            #set a unique name for each row
+                            unique_name = building + "_" + str(row) + "_" + str(date)
+                            if unique_name not in self.data:
+                                self.data[unique_name] = {}
+
+                            if view not in self.data[unique_name]:
+                                self.data[unique_name][view] = {}
+                            
+                            if frame not in self.data[unique_name][view]:
+                                self.data[unique_name][view][frame] = {}
+                            
+                            self.data[unique_name][view][frame][type_] = full_img_name
+    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 view in self.params.view_list:
+            datum_view = {}
+            for frame in self.params.frame_list:
+                frame_dict = {}
+                for type_ in self.params.types_list:
+                    if type_ == "colour":
+                        value = importers.image.load_image(datum[view][str(frame)][type_])
+                        #apply transform
+                        if self.params.transform is not None:
+                            value = self.params.transform(value)
+                    elif type_ == "silhouette" and view == "exocentric":
+                        value = importers.image.load_image(datum[view][str(frame)][type_])
+                    elif type_ == "colour" or type_ == "silhouette":
+                          continue
+                    elif type_ == "depth":
+                        value = importers.image.load_depth(datum[view][str(frame)][type_])
+                    elif type_ == "normal":
+                        value = importers.image.load_depth(datum[view][str(frame)][type_])
+                    else:
+                        vsalue = importers.image.load_image(datum[view][str(frame)][type_])
+                    frame_dict.update({str(type_) : value.squeeze(0)})
+                    #load pose for each frame
+                    if view == "exocentric":
+                        pose , pose_inv  = importers.load_pose(key,self.params.trajectories_dir,frame)
+                        frame_dict.update({"pose":pose,"pose_inv" : pose_inv, 'key':key})
+                    else:
+                        #NOTE:get source to target pose (t0,t1)
+                        if int(frame) == 1:
+                            #set source to 0
+                            source_ = 0
+                        else:
+                            #source is 1
+                            source_ = 1
+                        ego_pose = importers.load_egoPose(key,self.params.trajectories_dir,source_,int(frame))
+                        frame_dict.update({'key':key , 'source_to_target' : ego_pose})
+                datum_view.update({
+                    frame:frame_dict
+                })
+            datum_out.update({
+                view:datum_view
+            })
+        return datum_out
+    
+    def get_data(self):
+        return self.data
+
+    def get_splits(self,split):
+        
+        if split == "train":
+            txt_path = os.path.join(os.getcwd(),"data splits","scenes_train.txt")
+        
+        elif split == "test":
+            txt_path = os.path.join(os.getcwd(),"data splits","scenes_test.txt")
+        
+        elif split == "val":
+            txt_path = os.path.join(os.getcwd(),"data splits","scenes_val.txt")
+        
+        with open(txt_path,'r') as txt_file:
+            splits = [line.strip() for line in txt_file]
+        
+        return splits

importers/__init__.py

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

importers/extrinsics.py

@@ -0,0 +1,98 @@
+import numpy as np
+import torch
+import pandas as pd
+import os
+from utils import geometry
+
+
+def qmul(q, r):
+    """
+    Multiply quaternion(s) q with quaternion(s) r.
+    Expects two equally-sized tensors of shape (*, 4), where * denotes any number of dimensions.
+    Returns q*r as a tensor of shape (*, 4).
+    """
+    assert q.shape[-1] == 4
+    assert r.shape[-1] == 4
+    
+    original_shape = q.shape
+    
+    # Compute outer product
+    terms = torch.bmm(r.view(-1, 4, 1), q.view(-1, 1, 4))
+
+    w = terms[:, 0, 0] - terms[:, 1, 1] - terms[:, 2, 2] - terms[:, 3, 3]
+    x = terms[:, 0, 1] + terms[:, 1, 0] - terms[:, 2, 3] + terms[:, 3, 2]
+    y = terms[:, 0, 2] + terms[:, 1, 3] + terms[:, 2, 0] - terms[:, 3, 1]
+    z = terms[:, 0, 3] - terms[:, 1, 2] + terms[:, 2, 1] + terms[:, 3, 0]
+    return torch.stack((w, x, y, z), dim=1).view(original_shape)
+
+
+def load_extrinsics(filename, data_type=torch.float32):
+    data = np.loadtxt(filename)
+    pose = np.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()
+
+def load_pose(key,trajectories_dir,frame):
+    building , row , date = key.split("_")
+    i = int(row)
+    csv_file = os.listdir(os.path.join(trajectories_dir,building,date))
+    for file in csv_file:
+        if "blender" in file:
+            csv_name = file
+    
+    csv = pd.read_csv(os.path.join(trajectories_dir,building,date,csv_name),sep='\t',decimal='.')
+
+    if frame == 0:
+        Rd = np.array([csv['Orientation(x)_t1'][i],csv['Orientation(y)_t1'][i],csv['Orientation(z)_t1'][i],csv['Orientation(w)_t1'][i]])
+        Td = np.array([csv['Position(x)_t1'][i],csv['Position(y)_t1'][i],csv['Position(z)_t1'][i]])
+    else:
+        Rd = np.array([csv['Orientation(x)_t2'][i],csv['Orientation(y)_t2'][i],csv['Orientation(z)_t2'][i],csv['Orientation(w)_t2'][i]])
+        Td = np.array([csv['Position(x)_t2'][i],csv['Position(y)_t2'][i],csv['Position(z)_t2'][i]])
+    #read camera's pose
+    Rc = np.array([csv['PilotRot(x)'][i],csv['PilotRot(y)'][i],csv['PilotRot(z)'][i],csv['PilotRot(w)'][i]])
+    Tc = np.array([ csv['Pilot(x)'][i],csv['Pilot(y)'][i],csv['Pilot(z)'][i]])
+
+    #get drone matrix in world space
+    Pdw = geometry.objectMatrix(Rd,Td)
+    #get camera matrix in world space
+    Pcw = geometry.objectMatrix(Rc,Tc)
+    
+    PdroneToCamera = geometry.getRelativeMatrix(Pdw[:3,:3],Pcw[:3,:3],Pdw[:3,3],Pcw[:3,3])
+
+    
+    return PdroneToCamera , PdroneToCamera.inverse()
+
+def load_egoPose(key , trajectories_dir ,target_,source_):
+    building , row , date = key.split("_")
+    i = int(row)
+    csv_file = os.listdir(os.path.join(trajectories_dir,building,date))
+    for file in csv_file:
+        if "blender" in file:
+            csv_name = file
+    
+    csv = pd.read_csv(os.path.join(trajectories_dir,building,date,csv_name),sep='\t',decimal='.')
+     #read source pose
+    if target_ == 1:
+        Rsource = np.array([csv['Orientation(x)_t1'][i],csv['Orientation(y)_t1'][i],csv['Orientation(z)_t1'][i],csv['Orientation(w)_t1'][i]])
+        Tsource = np.array([csv['Position(x)_t1'][i],csv['Position(y)_t1'][i],csv['Position(z)_t1'][i]])
+        #read target pose
+        Rtarget = np.array([csv['Orientation(x)_t2'][i],csv['Orientation(y)_t2'][i],csv['Orientation(z)_t2'][i],csv['Orientation(w)_t2'][i]])
+        Ttarget = np.array([csv['Position(x)_t2'][i],csv['Position(y)_t2'][i],csv['Position(z)_t2'][i]])
+    else:
+        Rtarget = np.array([csv['Orientation(x)_t1'][i],csv['Orientation(y)_t1'][i],csv['Orientation(z)_t1'][i],csv['Orientation(w)_t1'][i]])
+        Ttarget = np.array([csv['Position(x)_t1'][i],csv['Position(y)_t1'][i],csv['Position(z)_t1'][i]])
+        #read source pose
+        Tsource = np.array([csv['Position(x)_t2'][i],csv['Position(y)_t2'][i],csv['Position(z)_t2'][i]])
+        Rsource = np.array([csv['Orientation(x)_t2'][i],csv['Orientation(y)_t2'][i],csv['Orientation(z)_t2'][i],csv['Orientation(w)_t2'][i]])
+
+    #get source pose matrix in world space
+    source_extrinsics = geometry.objectMatrix(Rsource,Tsource)
+    #get target pose matrix in world space
+    target_extrinsics = geometry.objectMatrix(Rtarget,Ttarget)
+
+    source_to_target_pose = target_extrinsics.inverse() @ source_extrinsics
+
+    return source_to_target_pose

importers/image.py

@@ -0,0 +1,33 @@
+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=1.0):
+    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)
+
+def load_normal(filename, data_type=torch.float32):
+    normal_img = numpy.array(cv2.imread(filename, cv2.IMREAD_UNCHANGED)).transpose(2, 0, 1)
+    c, h, w = normal_img.shape
+    return torch.from_numpy(normal_img).reshape(1,c, h, w).type(data_type)