MetaScenario/map_visualization.py at main · THU-changc17/MetaScenario · 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
144
145
146
147
148
149
150
151
152
153
154
import matplotlib
import matplotlib.axes
import matplotlib.pyplot as plt
import DBtools.utils as utils
from DBtools.init_db import init_DB
import argparse
# from InsertDataBase.Interaction_Intersection_EP0_InsertParticipant import ProcessPolygon

def set_visible_area(x_list, y_list, axes):
    min_x = 10e9
    min_y = 10e9
    max_x = -10e9
    max_y = -10e9

    for i in range(len(x_list)):
        min_x = min(float(x_list[i]), min_x)
        max_x = max(float(x_list[i]), max_x)
        min_y = min(float(y_list[i]), min_y)
        max_y = max(float(y_list[i]), max_y)

    axes.set_aspect('equal', adjustable='box')
    axes.set_xlim([min_x - 10, max_x + 10])
    axes.set_ylim([min_y - 10, max_y + 10])
    print(min_x, max_x, min_y, max_y)


def draw_map(cursor, axes):
    assert isinstance(axes, matplotlib.axes.Axes)

    axes.set_aspect('equal', adjustable='box')
    axes.patch.set_facecolor('white')
    type_dict = dict()
    _, Xlist, Ylist = utils.SearchNodeIDFromDB(cursor)
    set_visible_area(Xlist, Ylist, axes)
    way_list = utils.SearchWayFromDB(cursor)
    for way_id in way_list:
        node_list, x_list, y_list = utils.SearchNodeIDOnWayFromDB(cursor, way_id = way_id)
        way_type, way_subtype = utils.SearchWayTypeFromDB(cursor, way_id = way_id)

        if way_type == "curbstone":
            type_dict = dict(color="black", linewidth=2) #zorder=10)
        elif way_type == "line_thin":
            if way_subtype == "dashed":
                type_dict = dict(color="darkgreen", linewidth=1,  dashes=[10, 10])
            else:
                type_dict = dict(color="darkgreen", linewidth=1)
        elif way_type == "virtual":
            type_dict = dict(color="black", linewidth=1, dashes=[2, 5])
        elif way_type == "wall" or way_type == "fence":
            continue
        else:
            type_dict = dict(color="darkgreen", linewidth=1, dashes=[10, 10])

        channelization = utils.SearchChannelizationOnWayFromDB(cursor, way_id = way_id)
        for key in channelization:
            if key == "pedestrian_marking":
                type_dict = dict(color="blue", linewidth=1, dashes=[5, 10])
            elif key == "stop_line":
                type_dict = dict(color="gray", linewidth=3)
            elif key == "road_border":
                type_dict = dict(color="black", linewidth=1)
            elif key == "lane_change":
                type_dict = dict(color="black", linewidth=1)
            elif key == "guard_rail":
                type_dict = dict(color="black", linewidth=2)
            elif key == "turn_direction":
                type_dict = dict(color="blue", linewidth=2)
                break
            elif key == "is_intersection":
                type_dict = dict(color="black", linewidth=2)
                break
        plt.plot(x_list, y_list, **type_dict)
        for idx in range(len(x_list)):
            plt.plot(x_list[idx], y_list[idx], '.y', markersize=2.)
            # examples for adding ID text:
            # axes.text(x_list[idx], y_list[idx], way_id, fontsize=8, color="r")

    # Only for Intersection Scenario plotting Polygon:
    # polygon_list, _ = ProcessPolygon(cursor)
    # for polygon in polygon_list:
    #     pylogon_x = list(polygon[:, 1])
    #     pylogon_y = list(polygon[:, 2])
    #     plt.fill(pylogon_x, pylogon_y, color="#DCDCDC")


def draw_map_part(cursor, axes, Xlist, Ylist):
    assert isinstance(axes, matplotlib.axes.Axes)

    axes.set_aspect('equal', adjustable='box')
    axes.patch.set_facecolor('white')
    type_dict = dict()
    set_visible_area(Xlist, Ylist, axes)

    way_list = utils.SearchWayFromDB(cursor)
    for way_id in way_list:
        node_list, x_list, y_list = utils.SearchNodeIDOnWayFromDB(cursor, way_id=way_id, x_range=Xlist, y_range=Ylist)
        way_type, way_subtype = utils.SearchWayTypeFromDB(cursor, way_id=way_id)

        if way_type == "curbstone":
            type_dict = dict(color="black", linewidth=2)
        elif way_type == "line_thin":
            if way_subtype == "dashed":
                type_dict = dict(color="darkgreen", linewidth=2, dashes=[10, 10])
            else:
                type_dict = dict(color="darkgreen", linewidth=2)
        elif way_type == "virtual":
            type_dict = dict(color="black", linewidth=1, dashes=[2, 5])
        elif way_type == "wall" or way_type == "fence":
            continue
        else:
            type_dict = dict(color="darkgreen", linewidth=1, dashes=[10, 10])

        channelization = utils.SearchChannelizationOnWayFromDB(cursor, way_id=way_id)

        for key in channelization:
            if key == "pedestrian_marking":
                type_dict = dict(color="blue", linewidth=1, dashes=[5, 10])
            elif key == "stop_line":
                type_dict = dict(color="gray", linewidth=2)
            elif key == "road_border":
                type_dict = dict(color="black", linewidth=2)
            elif key == "guard_rail":
                type_dict = dict(color="black", linewidth=2)
            elif key == "turn_direction":
                type_dict = dict(color="blue", linewidth=2)
                break
            elif key == "is_intersection":
                type_dict = dict(color="blue", linewidth=2)
                break
        plt.plot(x_list, y_list, **type_dict)
        for idx in range(len(x_list)):
            plt.plot(x_list[idx], y_list[idx], '.y', markersize=2.)
            # axes.text(x_list[idx], y_list[idx], node_list[idx], fontsize=5, color="b")

    # polygon_list, _ = ProcessPolygon(cursor)
    # for polygon in polygon_list:
    #     pylogon_x = list(polygon[:, 1])
    #     pylogon_y = list(polygon[:, 2])
    #     plt.fill(pylogon_x, pylogon_y, color="#DCDCDC")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--DB', type=str, default=None)
    args = parser.parse_args()
    conn, cursor = init_DB(args.DB)
    fig, axes = plt.subplots(1, 1)
    fig.canvas.set_window_title("DatasetMap Visualization")
    draw_map(cursor, axes)
    # examples for plot local range:
    # draw_map_part(cursor, axes, Xlist=[1000, 1040], Ylist=[950, 1000])
    fig.set_size_inches(19.2, 10.8)
    plt.show()