Add docstrings, typing and TODOs to DiffOSM/data_extraction/helpers

data_extraction/helpers/get_osm_data.py

@@ -1,44 +1,49 @@
+'''
+Utility for read Open Street Map (OSM) data and extracting the required geometries
+'''
+
 import open3d as o3d
 import osmnx as ox
 
-osm_file_path = '/home/donceykong/kitti_360/kitti360Scripts/data/KITTI-360/data_osm/map_0009.osm'
-
-# Filter features for buildings and sidewalks
-buildings = ox.geometries_from_xml(osm_file_path, tags={'building': True})
-sidewalks = ox.geometries_from_xml(osm_file_path, tags={'highway': 'footway', 'footway': 'sidewalk'})
-
-# Process Sidewalks
-sidewalk_lines = []
-for _, sidewalk in sidewalks.iterrows():
-    if sidewalk.geometry.type == 'LineString':
-        x, y = sidewalk.geometry.xy
-        for i in range(len(x) - 1):
-            sidewalk_lines.append([[x[i], y[i], 0], [x[i + 1], y[i + 1], 0]])
-
-sidewalk_line_set = o3d.geometry.LineSet()
-sidewalk_points = [point for line in sidewalk_lines for point in line]
-sidewalk_lines_idx = [[i, i + 1] for i in range(0, len(sidewalk_points), 2)]
-sidewalk_line_set.points = o3d.utility.Vector3dVector(sidewalk_points)
-sidewalk_line_set.lines = o3d.utility.Vector2iVector(sidewalk_lines_idx)
-sidewalk_line_set.paint_uniform_color([0, 1, 0])  # Green color for sidewalks
-
-# Process Buildings as LineSets
-building_lines = []
-for _, building in buildings.iterrows():
-    if building.geometry.type == 'Polygon':
-        exterior_coords = building.geometry.exterior.coords
-        for i in range(len(exterior_coords) - 1):
-            start_point = [exterior_coords[i][0], exterior_coords[i][1], 0]
-            end_point = [exterior_coords[i + 1][0], exterior_coords[i + 1][1], 0]
-            building_lines.append([start_point, end_point])
-
-building_line_set = o3d.geometry.LineSet()
-building_points = [point for line in building_lines for point in line]
-building_lines_idx = [[i, i + 1] for i in range(0, len(building_points), 2)]
-building_line_set.points = o3d.utility.Vector3dVector(building_points)
-building_line_set.lines = o3d.utility.Vector2iVector(building_lines_idx)
-building_line_set.paint_uniform_color([0, 0, 1])  # Blue color for buildings
-
-# Visualize
-o3d.visualization.draw_geometries([sidewalk_line_set, building_line_set])
+if __name__ == '__main__':
+    osm_file_path = '/home/donceykong/kitti_360/kitti360Scripts/data/KITTI-360/data_osm/map_0009.osm'
+
+    # Filter features for buildings and sidewalks
+    buildings = ox.geometries_from_xml(osm_file_path, tags={'building': True})
+    sidewalks = ox.geometries_from_xml(osm_file_path, tags={'highway': 'footway', 'footway': 'sidewalk'})
+
+    # Process Sidewalks
+    sidewalk_lines = []
+    for _, sidewalk in sidewalks.iterrows():
+        if sidewalk.geometry.type == 'LineString':
+            x, y = sidewalk.geometry.xy
+            for i in range(len(x) - 1):
+                sidewalk_lines.append([[x[i], y[i], 0], [x[i + 1], y[i + 1], 0]])
+
+    sidewalk_line_set = o3d.geometry.LineSet()
+    sidewalk_points = [point for line in sidewalk_lines for point in line]
+    sidewalk_lines_idx = [[i, i + 1] for i in range(0, len(sidewalk_points), 2)]
+    sidewalk_line_set.points = o3d.utility.Vector3dVector(sidewalk_points)
+    sidewalk_line_set.lines = o3d.utility.Vector2iVector(sidewalk_lines_idx)
+    sidewalk_line_set.paint_uniform_color([0, 1, 0])  # Green color for sidewalks
+
+    # Process Buildings as LineSets
+    building_lines = []
+    for _, building in buildings.iterrows():
+        if building.geometry.type == 'Polygon':
+            exterior_coords = building.geometry.exterior.coords
+            for i in range(len(exterior_coords) - 1):
+                start_point = [exterior_coords[i][0], exterior_coords[i][1], 0]
+                end_point = [exterior_coords[i + 1][0], exterior_coords[i + 1][1], 0]
+                building_lines.append([start_point, end_point])
+
+    building_line_set = o3d.geometry.LineSet()
+    building_points = [point for line in building_lines for point in line]
+    building_lines_idx = [[i, i + 1] for i in range(0, len(building_points), 2)]
+    building_line_set.points = o3d.utility.Vector3dVector(building_points)
+    building_line_set.lines = o3d.utility.Vector2iVector(building_lines_idx)
+    building_line_set.paint_uniform_color([0, 0, 1])  # Blue color for buildings
+
+    # Visualize
+    o3d.visualization.draw_geometries([sidewalk_line_set, building_line_set])
 

data_extraction/helpers/read_data.py

@@ -11,9 +11,12 @@
 '''
 
 import os
+from typing import Tuple
 import numpy as np
+import numpy.typing as npt
 import open3d as o3d
 
+
 try:
     if 'KITTI360_DATASET' in os.environ:
         kitti360Path = os.environ['KITTI360_DATASET']
@@ -24,17 +27,52 @@
     print("Set absolute path to where the directory for KITTI-360 dataset is. \nSet using \"export KITTI360_DATASET = <path/to/dataset>\"")
 
 
-def read_building_pc_file(file_path):
+def read_building_pc_file(file_path: str) -> npt.NDArray:
+    '''
+    Load point cloud from filepath and return properly shaped pointcloud
+
+    params:
+        file_path (str): Filepath to pointcloud file to read
+
+    returns:
+        point_cloud (npt.NDArray): Loaded and reshaped point cloud
+    '''
     point_cloud = np.fromfile(file_path)
     return point_cloud.reshape(-1, 3)
 
-def read_building_edges_file(building_edges_file):
+def read_building_edges_file(building_edges_file: str) -> npt.NDArray:
+    '''
+    Loads building edges from file and returns an arrya of edges
+
+    params:
+        building_edges_file (str): File path to load edges
+
+    returns:
+        edges (npt.NDArray): Array of edges
+    '''
     with open(building_edges_file, 'rb') as bin_file:
         edges_array = np.fromfile(bin_file, dtype=float).reshape(-1, 2, 3)  # Reshape to 3D array
     return edges_array
 
-def get_building_scan_pcd(seq, building_index):
-    per_building_dir = os.path.join(kitti360Path, 'data_3d_extracted', '2013_05_28_drive_%04d_sync' % seq, 'buildings/per_building')
+def get_building_scan_pcd(
+        seq: int,
+        building_index: int
+    ) -> Tuple[npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray, npt.NDArray]: ## TODO I think this should probably be broken into smaller functions.
+    '''
+    Get all data (at the current timestep) for a given building
+
+    params:
+        seq (int): sequence number to pull data for
+        building_index (int): index for building to pull data for
+
+    returns:
+        build_points_pcd (npt.NDArray): pointcloud associated wiht building (for this scan)
+        build_edge_pcd (npt.NDArray): Array of edges
+        build_points_accum_pcd (npt.NDArray): Accumulate point cloud associated with building
+        build_points_diff_pcd (npt.NDArray): ## TODO Not sure what this is suppose to be
+        build_edges_span (npt.NDArray): Some sort of distance to each edge? ## TODO fill this in
+    '''
+    per_building_dir = os.path.join(kitti360Path, 'data_3d_extracted', '2013_05_28_drive_%04d_sync' % seq, 'buildings/per_building') ## TODO use f-strings instead of % operator
     per_build_edges_file = os.path.join(per_building_dir, f'build_{building_index}_edges.bin', )
     per_build_file = os.path.join(per_building_dir, f'build_{building_index}_scan_{scan_num}.bin', )
     per_build_accum_file = os.path.join(per_building_dir, f'build_{building_index}_accum.bin' )
@@ -76,7 +114,17 @@ def get_building_scan_pcd(seq, building_index):
 
     return build_points_pcd, build_edges_pcd, build_points_accum_pcd, build_points_diff_pcd, build_edges_span
 
-def get_max_scan_for_build(seq, building_index):
+def get_max_scan_for_build(seq: int, building_index: int) -> int: ## TODO THis coul probably be implemented using os.listdirs and parse the number instead of a loop.
+    '''
+    Parses through all scans within a sequences to find the last scan corresponding to a particular building
+
+    params:
+        seq (int): sequence number to pull data for
+        building_index (int): index for building to pull data for
+
+    returns:
+        max_scan (int): index of maximum scan 
+    '''
     per_building_dir = os.path.join(kitti360Path, 'data_3d_extracted', '2013_05_28_drive_%04d_sync' % seq, 'buildings/per_building')
     max_scan = 1
     while True:
@@ -88,7 +136,17 @@ def get_max_scan_for_build(seq, building_index):
 
     return max_scan
 
-def change_frame(vis, key_code):
+def change_frame(vis, key_code): ## TODO Fix typing
+    '''
+    Changes the data currently visualized to the next frame
+
+    params:
+        vis (): o3d visualization object
+        key_code (): key code? ## TODO not sure about this
+
+    returns:
+        bool ## TODO might be okay to get rid of return here. If there is no meaning behind this turth value
+    '''
     global scan_num
     global building_index
     global seq

data_extraction/helpers/read_gps.py

@@ -1,12 +1,30 @@
+'''
+Utitlity for reading GPS data
+provides two functions
+
+Example of usuage provided in main thread
+'''
+
 import os
+from typing import List, Tuple
+
 import osmnx as ox
 import matplotlib.pyplot as plt
 import glob
 
 # ox.config(log_console=True)
 
 # Function to parse OXTS file
-def parse_oxts_file(file_path):
+def parse_oxts_file(file_path: str) -> List[Tuple[float, float]]:
+    '''
+    Parses given OXTS file
+
+    params:
+        file_path (str): Path to OXTS file
+
+    returns:
+        OXTS_data (List[Tuple[float, float]]): the OXTS data parsed in (lat, lon)
+    '''
     gps_coords = []
     with open(file_path, 'r') as file:
         for line in file:
@@ -16,7 +34,17 @@ def parse_oxts_file(file_path):
     return gps_coords
 
 # Function to load all GPS data
-def load_gps_data(directory):
+def load_gps_data(directory: str) -> List[Tuple[float, float]]:
+    '''
+    Load all gps files form a given directory.
+    Data is in OXTS format
+
+    params:
+        directory (str): directory to load GPS data from
+
+    returns:
+        all_gps_coords (List[Tuple[float, float]]): all gps data in (lat, lon)
+    '''
     all_gps_coords = []
     # for i in range(10)
     for sub_dir in glob.glob(os.path.join(directory, '2013_05_28_drive_0003_sync/oxts/data/')):
@@ -25,28 +53,30 @@ def load_gps_data(directory):
             all_gps_coords.extend(gps_coords)
     return all_gps_coords
 
-# Directory containing your .txt files
-directory = '/home/donceykong/kitti_360/kitti360Scripts/data/KITTI-360/data_poses'
-gps_data = load_gps_data(directory)
-print("GPS data read!")
 
-# Load your local OSM file
-# osm_file = 'karlsruhe-regbez-latest.osm'  # Replace with your OSM file path
-# graph = ox.graph_from_xml(osm_file)
-graph = ox.graph_from_place("Karlsruhe, Germany", network_type="drive")
-print("OSM graph made!")
+if __name__ == '__main__':
+    # Directory containing your .txt files
+    directory = '/home/donceykong/kitti_360/kitti360Scripts/data/KITTI-360/data_poses'
+    gps_data = load_gps_data(directory)
+    print("GPS data read!")
+
+    # Load your local OSM file
+    # osm_file = 'karlsruhe-regbez-latest.osm'  # Replace with your OSM file path
+    # graph = ox.graph_from_xml(osm_file)
+    graph = ox.graph_from_place("Karlsruhe, Germany", network_type="drive")
+    print("OSM graph made!")
 
-# Plot the graph
-fig, ax = ox.plot_graph(graph, show=False, close=False)
+    # Plot the graph
+    fig, ax = ox.plot_graph(graph, show=False, close=False)
 
-# # Plot the GPS data
-for lat, lon in gps_data:
-    ax.scatter(lon, lat, c='red', s=1)  # longitude, latitude
+    # # Plot the GPS data
+    for lat, lon in gps_data:
+        ax.scatter(lon, lat, c='red', s=1)  # longitude, latitude
 
-# Save the figure as a PNG file
-# output_filename = '0002_gps_data.png'
-# plt.savefig(output_filename, bbox_inches='tight')
+    # Save the figure as a PNG file
+    # output_filename = '0002_gps_data.png'
+    # plt.savefig(output_filename, bbox_inches='tight')
 
-# Display the plot
-plt.show()
+    # Display the plot
+    plt.show()
 

data_extraction/helpers/trans_pose_data.py

@@ -1,44 +1,49 @@
 '''
 Doncey Albin
 
-
+Utility for loading pose data from KITTI360 dataset
 '''
 
+# Standard Library
 import os
+
+# Third party
+import numpy as np
+import osmnx as ox
 import open3d as o3d
 from open3d.visualization import gui
-import numpy as np
 from collections import namedtuple
-import osmnx as ox
 
 # Internal
 from tools.labels import labels
-from tools.utils import *
-from tools.convert_oxts_pose import *
+from tools.utils import * ## TODO remove * imports
+from tools.convert_oxts_pose import * ## TODO Remove * imports
+
 
-# Define the translation vector from IMU to LiDAR
-translation_vector = np.array([0.81, 0.32, -0.83])
+if __name__ == '__main__'
+    # Define the translation vector from IMU to LiDAR
+    translation_vector = np.array([0.81, 0.32, -0.83])
 
-# Define the rotation matrix for a 180-degree rotation about the X-axis (imu -> lidar)
-rotation_matrix = np.array([[1, 0, 0],
-                            [0, -1, 0],
-                            [0, 0, -1]])
+    # Define the rotation matrix for a 180-degree rotation about the X-axis (imu -> lidar)
+    rotation_matrix = np.array([[1, 0, 0],
+                                [0, -1, 0],
+                                [0, 0, -1]])
 
-imu_poses_file = 'kitti360Scripts/data/KITTI360/data_poses/2013_05_28_drive_0005_sync/poses.txt'
-vel_poses_file = 'kitti360Scripts/data/KITTI360/data_poses/2013_05_28_drive_0005_sync/vel_poses.txt'
+    imu_poses_file = 'kitti360Scripts/data/KITTI360/data_poses/2013_05_28_drive_0005_sync/poses.txt'
+    vel_poses_file = 'kitti360Scripts/data/KITTI360/data_poses/2013_05_28_drive_0005_sync/vel_poses.txt'
 
-# Read the IMU poses
-transformation_matrices = read_poses(imu_poses_file)
+    # Read the IMU poses
+    transformation_matrices = read_poses(imu_poses_file)
 
-# Extract frame indices (assuming they are the first element in each line)
-frame_indices = []
-with open(imu_poses_file, 'r') as file:
-    lines = file.readlines()
-    for line in lines:
-        frame_indices.append(line.split()[0])
+    # Extract frame indices (assuming they are the first element in each line)
+    frame_indices = []
+    with open(imu_poses_file, 'r') as file:
+        lines = file.readlines()
+        for line in lines:
+            frame_indices.append(line.split()[0])
 
-# Transform IMU poses to LiDAR poses
-lidar_poses = transform_imu_to_lidar(transformation_matrices, translation_vector, rotation_matrix)
+    # Transform IMU poses to LiDAR poses
+    lidar_poses = transform_imu_to_lidar(transformation_matrices, translation_vector, rotation_matrix)
 
-# Write the LiDAR poses to file
-write_poses(vel_poses_file, lidar_poses, frame_indices)
+    # Write the LiDAR poses to file
+    write_poses(vel_poses_file, lidar_poses, frame_indices)