-
Notifications
You must be signed in to change notification settings - Fork 0
/
DataCollectionPipeline.py
81 lines (61 loc) · 2.78 KB
/
DataCollectionPipeline.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
80
81
import airsim
import pprint
import os
import time
import math
import tempfile
from pynput import keyboard
import numpy as np
client = airsim.VehicleClient()
client.confirmConnection()
client.simSetVehiclePose(airsim.Pose(airsim.Vector3r(0, 0, -5), airsim.to_quaternion(0, 0, 0)), True)
i = 0
image_index = 0
def save_view(pressed_key):
global image_index
responses = client.simGetImages([
airsim.ImageRequest("0", airsim.ImageType.Scene, False, False),
])
response = responses[0]
# the original image has 4 channels. So the below code will make it RGB (3 channels)
img1d = np.fromstring(response.image_data_uint8, dtype=np.uint8)
# reshape array to 3 channel image array H X W X 3
img_rgb = img1d.reshape(response.height, response.width, 3)
# write to png
airsim.write_png(os.path.normpath('data\py' + str(image_index)+ '_'+ str(pressed_key)+ '.png'), img_rgb)
image_index+= 1
def on_key_press(key):
global i
try:
# Retrieve the character representation of the key
key_char = key.char
if key_char == 'w':
save_view(key_char)
client.simSetVehiclePose(airsim.Pose(airsim.Vector3r( client.simGetVehiclePose().position.x_val + math.cos(math.radians(i*15)), client.simGetVehiclePose().position.y_val + math.sin(math.radians(i*15)), -5), client.simGetVehiclePose().orientation), True)
if key_char == 's':
save_view(key_char)
client.simSetVehiclePose(airsim.Pose(airsim.Vector3r( client.simGetVehiclePose().position.x_val - math.cos(math.radians(i*15)), client.simGetVehiclePose().position.y_val - math.sin(math.radians(i*15)), -5), client.simGetVehiclePose().orientation), True)
if key_char == 'd':
save_view(key_char)
i= (i+1)%24
client.simSetVehiclePose(airsim.Pose(airsim.Vector3r(client.simGetVehiclePose().position.x_val, client.simGetVehiclePose().position.y_val, -5), airsim.to_quaternion(0, 0, math.radians(i*15))), True)
if key_char == 'a':
save_view(key_char)
i= (i-1)%24
client.simSetVehiclePose(airsim.Pose(airsim.Vector3r(client.simGetVehiclePose().position.x_val, client.simGetVehiclePose().position.y_val, -5), airsim.to_quaternion(0, 0, math.radians(i*15))), True)
print(client.simGetVehiclePose())
except AttributeError:
# Special keys like 'esc', 'enter', 'shift', etc.
key_name = key.name
print("Key pressed:", key_name)
def on_key_release(key):
# Handle key release event if needed
pass
listener = keyboard.Listener(
on_press=on_key_press,
on_release=on_key_release
)
# Start listening for keyboard events
listener.start()
# Keep the program running
listener.join()