sample_following_drone.py

# -*- coding: utf-8 -*-

import tellopy
import av
import numpy as np
import cv2
import time


# drone settings
drone = tellopy.Tello()
drone.connect()
drone.wait_for_connection(30.0)

# stream settings
stream = av.open(drone.get_video_stream())

# face detector
cascade = cv2.CascadeClassifier("haarcascade_frontalface_alt.xml")

# AR marker detector
dic = aruco.getPredefinedDictionary(aruco.DICT_4X4_50)

# hyperparameters
GOAL_AREA = 20000


# initial state
def init():
	drone.takeoff()
	return state1


# state1: searching
def state1():
	while True:
		for frame in stream.devode(video=0):
			image = cv2.cvtColor(np.array(frame.to_image()), cv2.COLOR_RGB2BGR)
			
			# face detection
			image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
			faces = cascade.detectMultiScale(image_gray)
			
			# goto state2 if faces are detected
			if len(faces) > 0:
				return state2
			
			# window
			cv2.imshow("window", image)
			key = cv2.waitKey(1)
			
			# goto fin() if q is pushed
			if key & 0xFF == ord("q"):
				return fin

			
# state2: following
def state2():
	while True:
		for frame in stream.devode(video=0):
			image = cv2.cvtColor(np.array(frame.to_image()), cv2.COLOR_RGB2BGR)
			
			# face detection
			image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
			faces = cascade.detectMultiScale(image_gray)
			
			# goto state1 if face is not detected
			if len(faces) == 0:
				return state1
			
			# compute area & center point
			x, y, w, h = face[0]
			area = w * h
			center = [y + h // 2, x + w // 2]
			
			# forward or backward
			if area < GOAL_AREA:
				drone.forward(50)
			else:
				drone.backward(25)
				
			# up or down
			if center[0] < image.shape[0] // 2:
				drone.up(25)
			else:
				drone.down(25)
				
			# left or right
			if center[1] < image.shape[1] // 2:
				drone.left(25)
			else:
				drone.right(25)
				
			# AR marker detection
			corns, ids, _ = aruco.detectMarkers(image, dic)
			
			# goto fin if AR marker 0 is detected
			if 0 in ids:
				return fin
			
			# window
			cv2.imshow("window", image)
			key = cv2.waitKey(1)
			
			# goto fin() if q is pushed
			if key & 0xFF == ord("q"):
				return fin


# final state
def fin():
	drone.land()
	drone.quit()
	return None


# state machine
state = init
while True:
	state = state()
	if state == None:
		break