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