drone-pathgen/examples/visualization.py at main · u-k-g/drone-pathgen · GitHub

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
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
#!/usr/bin/env python3
"""
trajectory visualization example

demonstrates 3d visualization of:
1. obstacle voxels
2. planned trajectory
3. safe flight corridors
4. start/goal markers

requires: uv add "git+https://github.com/u-k-g/drone-pathgen.git[viz]"
"""

import numpy as np
from basic_pathgen import generate_basic_trajectory # look at basic_pathgen.py before reading this

def visualize_trajectory():
    """run trajectory planning and visualize results with open3d"""

    try:
        import open3d as o3d
    except ImportError:
        print("❌ open3d not installed")
        print("   install with: uv add \"git+https://github.com/u-k-g/drone-pathgen.git[viz]\"")
        return

    print("🚁 trajectory visualization example")

    # reuse the basic trajectory generation
    api, success = generate_basic_trajectory()

    if not success:
        print("❌ trajectory optimization failed!")
        return

    print("✅ trajectory generated successfully!")

    # get visualization data
    result = api.get_visualization_data(show_initial_route=True)
    success, trajectory_points, voxel_data, voxel_size, start, goal, initial_route = result

    if not success:
        print("❌ failed to get visualization data")
        return

    print("🎨 building 3d visualization...")

    # --- build visualization geometries ---
    geometries = []

    # helper function for voxel meshes
    def build_voxel_mesh(indices, voxel_size_val, origin, color):
        cube = o3d.geometry.TriangleMesh.create_box(
            width=voxel_size_val, height=voxel_size_val, depth=voxel_size_val
        )
        cube.compute_vertex_normals()
        cube.paint_uniform_color(color)

        mesh = o3d.geometry.TriangleMesh()
        for x, y, z in indices:
            cube_copy = cube.translate(
                origin + np.array([x, y, z]) * voxel_size_val, relative=False
            )
            mesh += cube_copy
        return mesh

    # extract voxel indices by type
    occupied_indices, dilated_indices = [], []
    for z, layer in enumerate(voxel_data):
        for y, row in enumerate(layer):
            for x, val in enumerate(row):
                if val == 1:  # occupied
                    occupied_indices.append((x, y, z))
                elif val == 2:  # dilated (safety margin)
                    dilated_indices.append((x, y, z))

    # create obstacle mesh (red)
    if occupied_indices:
        occupied_mesh = build_voxel_mesh(
            occupied_indices, voxel_size, start, [0.8, 0.2, 0.2]
        )
        geometries.append(occupied_mesh)

    # create safety margin mesh (yellow/transparent)
    if dilated_indices:
        dilated_mesh = build_voxel_mesh(
            dilated_indices, voxel_size, start, [0.9, 0.7, 0.2]
        )
        geometries.append(dilated_mesh)

    # trajectory line (blue)
    if len(trajectory_points) > 1:
        trajectory_line = o3d.geometry.LineSet()
        trajectory_line.points = o3d.utility.Vector3dVector(trajectory_points)
        lines = [[i, i + 1] for i in range(len(trajectory_points) - 1)]
        trajectory_line.lines = o3d.utility.Vector2iVector(lines)
        trajectory_line.paint_uniform_color([0.2, 0.5, 0.9])
        geometries.append(trajectory_line)

    # initial route (gray, dashed-like)
    if initial_route and len(initial_route) > 1:
        initial_line = o3d.geometry.LineSet()
        initial_line.points = o3d.utility.Vector3dVector(initial_route)
        lines = [[i, i + 1] for i in range(len(initial_route) - 1)]
        initial_line.lines = o3d.utility.Vector2iVector(lines)
        initial_line.paint_uniform_color([0.5, 0.5, 0.5])
        geometries.append(initial_line)

    # start marker (green sphere)
    start_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.3)
    start_sphere.translate(start)
    start_sphere.paint_uniform_color([0.2, 0.8, 0.2])
    geometries.append(start_sphere)

    # goal marker (purple sphere)
    goal_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.3)
    goal_sphere.translate(goal)
    goal_sphere.paint_uniform_color([0.8, 0.2, 0.8])
    geometries.append(goal_sphere)

    # coordinate frame
    coord_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=1.0)
    coord_frame.translate(start)
    geometries.append(coord_frame)

    # launch visualization
    print("🚀 launching open3d viewer...")
    print("   🟢 green sphere = start")
    print("   🟣 purple sphere = goal")
    print("   🔴 red cubes = obstacles")
    print("   🟡 yellow cubes = safety margins")
    print("   🔵 blue line = optimized trajectory")
    print("   ⚪ gray line = initial route")

    o3d.visualization.draw_geometries(
        geometries,
        window_name="gcopter trajectory visualization",
        width=1200,
        height=800
    )

if __name__ == "__main__":
    visualize_trajectory()