1import pychrono.core as chrono
2import pychrono.sensor as sens
3
4import numpy as np
5import time
6import random
7import cv2
8import math
9
10class simulation:
11
12    def __init__(self) -> None:
13        self.system = chrono.ChSystemNSC()
14        self.system.Set_G_acc(chrono.ChVectorD(0,0,0))
15
16        green = self.init_vis_mat(chrono.ChVectorF(0,1,0))
17        black = self.init_vis_mat(chrono.ChVectorF(1,1,1))
18        yellow = self.init_vis_mat(chrono.ChVectorF(1,1,0))
19
20        ground = chrono.ChBodyEasyBox(1000,40,1,1000,True,False)
21        ground.SetPos(chrono.ChVectorD(0,0,-1))
22        ground.SetBodyFixed(True)
23        asset = ground.GetAssets()[0]
24        visual_asset = chrono.CastToChVisualization(asset)
25        visual_asset.material_list.append(green)
26        self.system.Add(ground)
27
28        egocar = chrono.ChBodyEasyBox(5,2,2,1000,True,False)
29        egocar.SetPos(chrono.ChVectorD(0,1,1))
30        car_asset = egocar.GetAssets()[0]
31        car_visual_asset = chrono.CastToChVisualization(car_asset)
32        car_visual_asset.material_list.append(yellow)
33        self.system.Add(egocar)
34
35        frontcar = chrono.ChBodyEasyBox(5,2,2,1000,True,False)
36        frontcar.SetPos(chrono.ChVectorD(20,1,1))
37        frontcar.SetPos_dt(chrono.ChVectorD(5,0,0))
38        frontcar_asset = frontcar.GetAssets()[0]
39        frontcar_visual_asset = chrono.CastToChVisualization(frontcar_asset)
40        frontcar_visual_asset.material_list.append(yellow)
41        self.system.Add(frontcar)
42
43        # incoming cars on the left lane
44        for i in range(10):
45            leftcar = chrono.ChBodyEasyBox(5,2,2,1000,True,False)
46            leftcar.SetPos(chrono.ChVectorD(10 + i * 15 ,20,1))
47            leftcar.SetPos_dt(chrono.ChVectorD(-5,0,0))
48            leftcar_asset = frontcar.GetAssets()[0]
49            leftcar_visual_asset = chrono.CastToChVisualization(leftcar_asset)
50            leftcar_visual_asset.material_list.append(yellow)
51            self.system.Add(leftcar)
52
53
54        # cars in the right lane
55        for i in range(10):
56            rightcar = chrono.ChBodyEasyBox(5,2,2,1000,True,False)
57            rightcar.SetPos(chrono.ChVectorD(10 + i * 15 ,-20,1))
58            rightcar.SetPos_dt(chrono.ChVectorD(15,0,0))
59            rightcar_asset = rightcar.GetAssets()[0]
60            rightcar_visual_asset = chrono.CastToChVisualization(rightcar_asset)
61            rightcar_visual_asset.material_list.append(black)
62            self.system.Add(rightcar)
63
64        offset_pose = chrono.ChFrameD(chrono.ChVectorD(3,0,0), chrono.Q_from_AngZ(0))
65        self.adding_sensors(egocar, offset_pose)
66
67
68    # color should be a chrono.ChVectorF(float,float,float)
69    def init_vis_mat(self, color):
70        vis_mat = chrono.ChVisualMaterial()
71        vis_mat.SetDiffuseColor(color)
72        vis_mat.SetSpecularColor(chrono.ChVectorF(1,1,1))
73        return vis_mat
74
75    def adding_sensors(self, body, offset_pose):
76        self.manager = sens.ChSensorManager(self.system)
77        intensity = 1.0
78        self.manager.scene.AddPointLight(chrono.ChVectorF(
79            2, 2.5, 100), chrono.ChVectorF(intensity, intensity, intensity), 500.0)
80        self.manager.scene.AddPointLight(chrono.ChVectorF(
81            9, 2.5, 100), chrono.ChVectorF(intensity, intensity, intensity), 500.0)
82        self.manager.scene.AddPointLight(chrono.ChVectorF(
83            16, 2.5, 100), chrono.ChVectorF(intensity, intensity, intensity), 500.0)
84        self.manager.scene.AddPointLight(chrono.ChVectorF(
85            23, 2.5, 100), chrono.ChVectorF(intensity, intensity, intensity), 500.0)
86
87        update_rate = 30
88        lag = 0
89        exposure_time = 0
90
91        self.hfov = math.pi / 3
92        self.vfov = math.pi / 10
93        hfov = self.hfov
94        vfov = self.vfov
95        self.image_width = 1280
96        image_width = self.image_width
97        self.image_height = int(image_width * vfov / hfov)
98        image_height = self.image_height
99
100        self.cam = sens.ChCameraSensor(body, update_rate,offset_pose,image_width,image_height,hfov)
101        self.cam.SetName("Camera Sensor")
102        self.cam.SetLag(lag)
103        self.cam.SetCollectionWindow(exposure_time)
104        self.cam.PushFilter(sens.ChFilterRGBA8Access())
105        self.manager.AddSensor(self.cam)
106
107        h_samples = 200
108        v_samples = 200
109
110        self.radar = sens.ChRadarSensor(body,update_rate,offset_pose,h_samples,v_samples,hfov, vfov,100.0)
111        self.radar.PushFilter(sens.ChFilterRadarXYZReturn())
112        self.radar.PushFilter(sens.ChFilterRadarXYZAccess())
113        self.manager.AddSensor(self.radar)
114
115
116    def sim_advance(self):
117        step_size = 1e-3
118        self.manager.Update()
119        self.system.DoStepDynamics(step_size)
120        self.display_image()
121
122    def display_image(self):
123        rgba8_buffer = self.cam.GetMostRecentRGBA8Buffer()
124        if rgba8_buffer.HasData():
125            print("camera")
126            rgba8_data = rgba8_buffer.GetRGBA8Data()
127            np.flip(rgba8_data)
128            bgr = cv2.cvtColor(rgba8_data[::-1], cv2.COLOR_RGB2BGR)
129        radar_buffer = self.radar.GetMostRecentRadarXYZBuffer()
130        if radar_buffer.HasData():
131            radar_data = radar_buffer.GetRadarXYZData()[0]
132
133        if rgba8_buffer.HasData():
134            if radar_buffer.HasData():
135                print("radar")
136                for i in radar_data:
137                    box_x = self.image_width - (int(self.image_width / self.hfov * math.atan2(i[1],i[0])) + int(self.image_width / 2))
138                    box_y = self.image_height - (int(self.image_height / self.vfov * math.atan2(i[2],i[0])) + int(self.image_height / 2))
139                    if abs(i[3]) < 1e-2:
140                        # positive relative velocity -> blue
141                        cv2.rectangle(bgr,(box_x - 1, box_y - 1), (box_x+1,box_y+1),(0,0,0),int(2))
142                    elif i[3] < 0:
143                        # negative relative velocity -> red
144                        cv2.rectangle(bgr,(box_x - 1, box_y - 1), (box_x+1,box_y+1),(255,0,0),int(2))
145                        pass
146                    else:
147                        # neutral -> white
148                        cv2.rectangle(bgr,(box_x - 1, box_y + 1), (box_x+1,box_y-1),color=(0,0,255),thickness=int(2))
149
150            cv2.imshow("window", bgr)
151            if cv2.waitKey(1):
152                return
153
154
155def main():
156    sim = simulation()
157    while True:
158        sim.sim_advance()
159
160if __name__ == "__main__":
161    main()
162
163