-
Notifications
You must be signed in to change notification settings - Fork 0
/
Main.py
79 lines (55 loc) · 2.13 KB
/
Main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import os
import open3d as o3d
from Picker import pick_points
import pandas as pd
import numpy as np
def get_sorted_ply_files(folder_path):
"""Returns a sorted list of .ply files in the given folder."""
ply_files = [f for f in os.listdir(folder_path) if f.endswith('.ply')]
ply_files.sort()
return ply_files
def process_ply(ply_file, file_name, output_dir="annotations"):
pcd = o3d.io.read_point_cloud(ply_file)
print(f"Processing file: {file_name}")
if not pcd.is_empty():
print("Point cloud loaded successfully.")
else:
print("Failed to load point cloud. Please check the file path.")
return
picked_points = pick_points(pcd)
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# Extract points based on picked_points indices
points = [pcd.points[i] for i in picked_points]
# Convert points to a NumPy array
points_array = np.array(points)
# Create a DataFrame from the NumPy array
df = pd.DataFrame(points_array, columns=['x', 'y', 'z'])
# Save DataFrame to CSV
csv_path = os.path.join(output_dir, file_name+'.csv')
df.to_csv(csv_path, index=False)
print(f"Points saved to {csv_path}")
def main():
folder_path = input("Enter the folder path containing .ply files: ")
if not os.path.isdir(folder_path):
print("Invalid folder path.")
return
ply_files = get_sorted_ply_files(folder_path)
if not ply_files:
print("No .ply files found in the folder.")
return
start_index_str = input(
f"Enter the start index (0 to {len(ply_files) - 1}), or press Enter to start from the beginning: ")
start_index = 0
if start_index_str.isdigit():
start_index = int(start_index_str)
if start_index < 0 or start_index >= len(ply_files):
print("Invalid start index. Starting from the beginning.")
start_index = 0
for i in range(start_index, len(ply_files)):
print(i)
ply_file = ply_files[i]
file_path = os.path.join(folder_path, ply_file)
process_ply(file_path, ply_file)
if __name__ == "__main__":
main()