Skip to content

Commit

Permalink
Merge branch 'PEDRA' of https://github.com/p-ai-org/P-Agent into PEDRA
Browse files Browse the repository at this point in the history
  • Loading branch information
@jaredmejia
jaredmejia committed May 10, 2021
2 parents 0769563 + 709b7c7 commit 5854d01
Show file tree
Hide file tree
Showing 2 changed files with 119 additions and 28 deletions.
31 changes: 25 additions & 6 deletions main.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,21 +8,40 @@

from Network.HardCode_Controller import HardCode_Controller

# Choose a random spawn location
random_choice = True

# Check if Connection with Airsim is Good
client = airsim.MultirotorClient()
client.confirmConnection()
client.enableApiControl(True)
client.armDisarm(True)

# Load in random position
# choice = np.random.randint(0,4)
# if random_choice:
# choice = 1
# if choice == 0:
# pass
# if choice == 1:
# client.simSetVehiclePose(airsim.Pose(airsim.Vector3r(2, float("nan"), float("nan")), airsim.to_quaternion(float("nan"),float("nan"),float("nan"))), True)
# if choice == 2:
# client.simSetVehiclePose
# if choice == 3:
# client.simSetVehiclePose

# Load CNN Model
loaded_model = torch.load(os.path.join(os.path.abspath(__name__),"..","Network","object_detection","mymodel3.pt"), map_location="cuda:0") # Assumes model is stored in \P-agent\Network\Object_detection
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')

#Takeoff and Prepare for instructions
client.takeoffAsync().join()

done = False
mem = 0
docking = False
found = False
done = False

while(done == False):
# Grab image and convert to RGB tensor
img = client.simGetImage("0", airsim.ImageType.Scene)
Expand All @@ -43,12 +62,12 @@
box_pred = []

#TODO: Train it to better detect literal edge cases
Controller = HardCode_Controller(client, box_pred, mem)
done, mem, collided = Controller.policy()
Controller = HardCode_Controller(client, box_pred, mem, docking, found)
done, mem, collided, docking, found = Controller.policy()

if collided:
print("I hit a wall, dumdum")
if not collided:
print("I got to the Package!")
print("I found the Package")
else:
print("I'm a dumdum. Something went wrong")


116 changes: 94 additions & 22 deletions network/HardCode_Controller.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,11 +19,13 @@ class HardCode_Controller:
Takes action in Unreal Environment
"""

def __init__(self, client, boundingBox_pts, mem):
def __init__(self, client, boundingBox_pts, mem, docking, found):
self.client = client
self.pts = boundingBox_pts
self.done = False
self.mem = mem
self.docking = docking
self.found = found
self.image_Coords = np.array([[0,144],[256,0]])

def center(self, in_frame, z, dir = 1, duration = 0.2):
Expand Down Expand Up @@ -52,16 +54,70 @@ def center(self, in_frame, z, dir = 1, duration = 0.2):
#TODO: If the distace of the max is less towards one side, move in that direction
self.client.moveByVelocityZAsync(0, 0, z, duration, airsim.DrivetrainType.MaxDegreeOfFreedom, airsim.YawMode(True, dir * 20 )).join()

# def move():
#TODO: move function will just move the drone and correct for that movement in the orientation of the camera and then just use the camera orientation to direct itself.
def centerBelow(self, rect, thresh_coord, z):
# Split up the image into quadrants and move accordingly
# right = thresh_coord
# up = thresh_coord

right_x = 0.5 * thresh_coord[1,0]
right_x = np.append(right_x, thresh_coord[1,0])
right_y = thresh_coord[:,1]
right_x = np.reshape(right_x, (2,1))
right_y = np.reshape(right_y, (2,1))
right = np.hstack((right_x, right_y))

up_x = thresh_coord[0,0]
up_x = np.append(up_x, thresh_coord[1,0])
up_y = thresh_coord[0,1]
up_y = np.append(up_y, 0.5 * thresh_coord[1,1])
up_x = np.reshape(up_x, (2,1))
up_y = np.reshape(up_y, (2,1))
up = np.hstack((up_x, up_y))

# right[0,0] = 0.5 * right[1,0]
# up[1,1] = 0.5 * up[1,1]

right = np.array([ [right[0,0], right[0,1]], [right[0,0], right[1,1]], [right[1,0], right[1,1]], [right[1,0], right[0,1]] ])
up = np.array([ [up[0,0], up[0,1]], [up[0,0], up[1,1]], [up[1,0], up[1,1]], [up[1,0], up[0,1]] ])
thresh = np.array([ [thresh_coord[0,0], thresh_coord[0,1]], [thresh_coord[0,0], thresh_coord[1,1]], [thresh_coord[1,0], thresh_coord[1,1]], [thresh_coord[1,0], thresh_coord[0,1]] ])
bb_pts = np.array([ [rect[0,0], rect[0,1]], [rect[0,0], rect[1,1]], [rect[1,0], rect[1,1]], [rect[1,0], rect[0,1]] ])

# Create polygons of image regions and bounding box
right_poly = Polygon(right) # Right side
up_poly = Polygon(up) # Upper section
thresh_poly = Polygon(thresh) # Threshold polygon (entire)
bb_poly = Polygon(bb_pts) # Bounding box

# Calculate Intersections
bb_thresh_int = thresh_poly.intersection(bb_poly).area
bb_right_int = right_poly.intersection(bb_poly).area
bb_up_int = right_poly.intersection(bb_poly).area

x = 0
y = 0
# if bb_thresh_int >= bb_poly.area:
if bb_thresh_int >= 0.1 * thresh_poly.area:
self.client.landAsync().join()
return True # Yay! We've found it!
elif bb_right_int == 0 and bb_up_int == 0:
# It's not directly underneath
x = 0.5
else:
if bb_right_int / bb_thresh_int > 0.5:
# If more of the box is on the right side, move right
x = 0.5
else:
x = -0.5

if bb_up_int / bb_thresh_int > 0.5:
# If more of the box is on the upper part of the threshold
y = 0.5
else:
y = -0.5

self.move(z, x, y)
return False

def area(self, coords):
"""
Calculate the area of the bounding box given a (2,2) Numpy array
"""
area = (coords[1,0] - coords[0,0]) * (coords[1,1] - coords[0,1])
return area

def overlap(self, rect, thresh_rect, ratio = .1):
"""
rect: Bounding box points (upper left, lower right)
Expand All @@ -73,7 +129,7 @@ def overlap(self, rect, thresh_rect, ratio = .1):
full_pts = np.array([ [rect[0,0], rect[0,1]], [rect[0,0], rect[1,1]], [rect[1,0], rect[1,1]], [rect[1,0], rect[0,1]] ])
full_pts = list(map(tuple, full_pts))

# Create an "inner box" as a threshold (move the box down?)
# Create an "inner box" as a threshold
thresh_x = np.add(thresh_rect[0,0], ratio * thresh_rect[1,0])
thresh_x = np.append(thresh_x, np.subtract([thresh_rect[1,0]], ratio * thresh_rect[1,0]))
thresh_x = np.reshape(thresh_x, (2,1))
Expand All @@ -97,26 +153,26 @@ def transformToEarthFrame(self, vector, q_):
q = Quaternion(q_)
return q.rotate(vector)

def move(self, z, duration = 0.2, desired_velocity = 1):
def move(self, z, x, y, duration = 0.2, desired_velocity = 1):
"""
Basic move command which just moves the drone forward
"""
# q = self.client.getCameraInfo(0).pose.orientation #this is different from the getOrientation, taking the latter can create problem due to the fact that the drone move slower than the cam?
# my_quaternion = Quaternion(w_val=q.w_val,x_val=q.x_val,y_val= q.y_val,z_val=q.z_val)
# mvm = my_quaternion.rotate(action)
# self.client.moveByVelocityZAsync(2, 0, z, duration, drivetrain = airsim.DrivetrainType.ForwardOnly, yaw_mode=YawMode(False, 0))
vx, vy, _ = self.transformToEarthFrame([desired_velocity, 0, 0], [self.client.simGetVehiclePose().orientation.w_val,\
vx, vy, _ = self.transformToEarthFrame([desired_velocity * x, desired_velocity * y, 0], [self.client.simGetVehiclePose().orientation.w_val,\
self.client.simGetVehiclePose().orientation.x_val,\
self.client.simGetVehiclePose().orientation.y_val,self.client.simGetVehiclePose().orientation.z_val])
self.client.moveByVelocityZAsync(vx=vx, vy=vy, z=z, yaw_mode=airsim.YawMode(True, 0), drivetrain=airsim.DrivetrainType.MaxDegreeOfFreedom, duration=duration).join()
self.client.hoverAsync().join()

def policy(self, center_thresh = .05, velocity = 3):

# Grab State Information
state = self.client.simGetVehiclePose()

# Centering
if self.pts == []:
if self.found == True:
centered_box = True
pass # If package has already been found then pass through
elif self.pts == []:
centered_box = False
self.center(False, state.position.z_val) # If you can't find the package at all
else:
Expand All @@ -129,6 +185,7 @@ def policy(self, center_thresh = .05, velocity = 3):

if self.overlap(self.pts, thresh_coord) > center_thresh:
centered_box = True
self.found = True
else:
centered_box = False
# Determine the direction to turn to fine tune.
Expand All @@ -150,7 +207,24 @@ def policy(self, center_thresh = .05, velocity = 3):

# Move
if centered_box:
self.move(state.position.z_val)
if self.docking == True:
# build a threshold that scales all points inward
thresh_coord = self.image_Coords
thresh_coord[1,1] = thresh_coord[0,1]
thresh_coord[0,0] = np.add(thresh_coord[0,0], 0.15 * thresh_coord[1,0])
thresh_coord[0,1] = np.subtract(thresh_coord[0,1], 0.75 * thresh_coord[0,1])
thresh_coord[1,:] = np.subtract(thresh_coord[1,:], 0.15 * thresh_coord[1,:])
self.done = self.centerBelow(self.pts, thresh_coord, state.position.z_val)
# Check if package is no longer found on camera dead ahead
elif self.pts == []:
x = self.client.simGetCameraInfo("0").pose.position.x_val
y = self.client.simGetCameraInfo("0").pose.position.y_val
z = self.client.simGetCameraInfo("0").pose.position.z_val
# Change orientation of camera to straight down and center
self.client.simSetCameraPose("0", airsim.Pose(airsim.Vector3r(x, y, z), airsim.to_quaternion(-90,0,0)))
self.docking = True
else:
self.move(state.position.z_val, x = 1, y = 0)

# Check if the drone collided or landed
collided = self.client.simGetCollisionInfo()
Expand All @@ -159,6 +233,4 @@ def policy(self, center_thresh = .05, velocity = 3):
if collided.has_collided:
self.done = True

return self.done, self.mem, collided

print('hello')
return self.done, self.mem, collided, self.docking, self.found

0 comments on commit 5854d01

Please sign in to comment.