Drone Equinox

drone-equinox.py

from pyparrot.Minidrone import Mambo
from pyparrot.DroneVision import DroneVision
from pynput.keyboard import Key, Listener
import cv2
import threading
import face_recognition
import os
import time

print("Identificando super villano...")
os.system("say Identificando super villano")
print("Preparando despegue...")
os.system("say preparando despegue")

mamboAddr = "e0:14:d0:63:3d:d0"

mambo = Mambo(mamboAddr, use_wifi=True)

pics = ["Eva", "Miguel"]

known_face_encodings = []
known_face_names = []

for pic in pics:
    image = face_recognition.load_image_file(f'{pic}.jpg')
    face_encoding = face_recognition.face_encodings(image)[0]

    known_face_encodings.append(face_encoding)
    known_face_names.append(pic)


class UserVision:
    def __init__(self, vision):
        self.vision = vision
        self.identified = []

    def process_pic(self, args):
        frame = self.vision.get_latest_valid_picture()

        if (frame is None):
            return

        small_frame = cv2.resize(frame, (0, 0), fx=0.25, fy=0.25)
        rgb_small_frame = small_frame[:, :, ::-1]

        face_locations = face_recognition.face_locations(rgb_small_frame)
        face_encodings = face_recognition.face_encodings(rgb_small_frame, face_locations)

        face_names = []
        for face_encoding in face_encodings:
            # See if the face is a match for the known face(s)
            matches = face_recognition.compare_faces(known_face_encodings, face_encoding)

            # If a match was found in known_face_encodings, just use the first one.
            if True in matches:
                first_match_index = matches.index(True)
                name = known_face_names[first_match_index]

                face_names.append(name)

                # Display the results
                for (top, right, bottom, left), name in zip(face_locations, face_names):
                    # Scale back up face locations since the frame we detected in was scaled to 1/4 size
                    top *= 4
                    right *= 4
                    bottom *= 4
                    left *= 4

                    # Draw a box around the face
                    cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255), 2)

                    # Draw a label with a name below the face
                    cv2.rectangle(frame, (left, bottom - 35),
                                  (right, bottom), (0, 0, 255), cv2.FILLED)
                    font = cv2.FONT_HERSHEY_DUPLEX
                    cv2.putText(frame, name, (left + 6, bottom - 6), font, 1.0, (255, 255, 255), 1)

                    cv2.imwrite(f'{name}_suspect.jpg', frame)

                    if name not in self.identified:
                      self.identified.append(name)

                      if name == 'Miguel':
                        cv2.imwrite(f'{name}_villain.jpg', frame)
                        os.system(f'say super villano {name} identificado')
                        os.system(f'say alerta')
                        os.system(f'say super villano {name} identificado')
                        os.system(f'say alerta')
                        os.system(f'say super villano {name} identificado')
                        os.system(f'say alerta')
                        os.system(f'say super villano {name} identificado')
                      else:
                        os.system(f'say Sospechoso {name} identificado')
                        os.system(f'Inocente')
                        
            else:
                print("no match!")


roll = 0
pitch = 0
yaw = 0
vertical_movement = 0
cont = True


def on_press(key):
    global pitch
    global yaw
    global roll

    if key == Key.up:
        pitch = 50
    elif key == Key.down:
        pitch = -50
    elif hasattr(key, 'char') and key.char == 'd':
        yaw = 70
    elif hasattr(key, 'char') and key.char == 'a':
        yaw = -70
    elif hasattr(key, 'char') and key.char == 'w':
        vertical_movement = 10
    elif key == Key.right:
        roll = 50
    elif key == Key.left:
        roll = -50


def on_release(key):
    global pitch
    global yaw
    global roll

    if key == Key.up or key == Key.down:
        pitch = 0
    elif key == Key.right or key == Key.left:
        roll = 0
    elif hasattr(key, 'char') and (key.char == 'd' or key.char == 'a'):
        yaw = 0
    elif hasattr(key, 'char') and key.char == 'w':
        vertical_movement = 0
    elif hasattr(key, 'char') and key.char == 's':
        vertical_movement = 0
    elif key == Key.esc:
        global cont
        cont = False

        print("landing")
        mambo.safe_land(5)
        mamboVision.close_video()

        print("disconnect")
        mambo.disconnect()

        return False


def fly():
    while(cont):
        global pitch
        global yaw
        global roll

        if yaw == 0 and pitch == 0 and roll == 0 and vertical_movement == 0:
            mambo.smart_sleep(0.1)
        else:
            mambo.fly_direct(roll=roll, pitch=pitch, yaw=yaw,
                             vertical_movement=vertical_movement, duration=0.1)

success = mambo.connect(num_retries=3)

if (success):
    mambo.smart_sleep(2)
    mambo.ask_for_state_update()
    mambo.smart_sleep(2)

    mamboVision = DroneVision(mambo, is_bebop=False, buffer_size=30)
    userVision = UserVision(mamboVision)
    mamboVision.set_user_callback_function(userVision.process_pic, user_callback_args=None)
    mamboVision.open_video()

    mambo.safe_takeoff(5)
    mambo.smart_sleep(3)

    mambo.fly_direct(roll=0, pitch=0, yaw=0, vertical_movement=10, duration=1)
    
    threading.Thread(target=fly).start()

    with Listener(on_press=on_press, on_release=on_release) as listener:
        listener.join()