[Python] 드론을 활용한 실시간 객체 인식 및 이미지 캡처 예제코드

드론을 활용한 실시간 객체 인식 및 이미지 캡처

이번 포스팅에서는 드론을 활용하여 실시간 객체 인식 및 이미지 캡처를 수행하는 방법을 소개합니다. 드론 제어와 비디오 스트림 처리를 통해 YOLO를 사용하여 사람 객체를 인식하고, OpenCV를 사용하여 이미지를 저장하는 과정을 다룹니다.

준비물

1. 드론: Tello 드론
2. 개발 환경: Python 3, DJITelloPy, OpenCV, YOLOv3
3. 라이브러리:
   • DJITelloPy: Tello 드론 제어 라이브러리
   • OpenCV: 컴퓨터 비전 라이브러리
   • YOLOv3: 실시간 객체 인식 모델

드론 제어 및 비디오 스트림 설정

먼저 드론을 제어하고 비디오 스트림을 시작하는 코드를 작성합니다.

드론 제어 및 비디오 스트림 시작 코드

 

from djitellopy import Tello
import time

# 드론 초기화 및 연결
drone = Tello()
drone.connect()

# 비디오 스트림 시작
drone.streamon()

# 드론 이륙
drone.takeoff()

# 드론 제어 명령 (예: 앞쪽으로 이동)
drone.move_forward(50)  # 50cm 앞으로 이동

# 드론 제자리 유지
time.sleep(2)

# 드론 착륙
drone.land()

# 비디오 스트림 종료
drone.streamoff()

비디오 스트림 처리 및 객체 인식

이제 드론으로부터 비디오 스트림을 수신하고, YOLO를 사용하여 객체 인식을 수행하는 코드를 작성합니다.

비디오 스트림 처리 및 객체 인식 코드

import cv2
import numpy as np
from djitellopy import Tello

# YOLO 모델 설정
YOLO_WEIGHTS = "yolov3.weights"
YOLO_CONFIG = "yolov3.cfg"
YOLO_CLASSES = "coco.names"

# YOLO 네트워크 로드
net = cv2.dnn.readNet(YOLO_WEIGHTS, YOLO_CONFIG)
with open(YOLO_CLASSES, "r") as f:
    classes = [line.strip() for line in f.readlines()]
layer_names = net.getLayerNames()
output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]

# 드론 초기화 및 연결
drone = Tello()
drone.connect()
drone.streamon()

def capture_and_detect():
    frame_read = drone.get_frame_read()
    frame = frame_read.frame

    height, width, channels = frame.shape
    blob = cv2.dnn.blobFromImage(frame, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
    net.setInput(blob)
    outs = net.forward(output_layers)

    class_ids = []
    confidences = []
    boxes = []

    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5 and classes[class_id] == "person":
                center_x = int(detection[0] * width)
                center_y = int(detection[1] * height)
                w = int(detection[2] * width)
                h = int(detection[3] * height)
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)
                boxes.append([x, y, w, h])
                confidences.append(float(confidence))
                class_ids.append(class_id)

    indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)

    for i in range(len(boxes)):
        if i in indexes:
            x, y, w, h = boxes[i]
            label = str(classes[class_ids[i]])
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            cv2.putText(frame, label, (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            person_img = frame[y:y+h, x:x+w]
            cv2.imwrite(f"captured_person_{i}.jpg", person_img)

    cv2.imshow("Frame", frame)
    cv2.waitKey(1)

if __name__ == "__main__":
    try:
        drone.takeoff()

        while True:
            capture_and_detect()

            key = cv2.waitKey(1) & 0xFF
            if key == ord('q'):
                break

        drone.land()
    finally:
        drone.streamoff()
        cv2.destroyAllWindows()

코드 설명

1. YOLO 모델 설정:
   • YOLO 가중치, 구성 파일, 클래스 이름 파일을 로드합니다.
   • YOLO 네트워크를 설정하고, 레이어 이름과 출력 레이어를 구성합니다.
2. 드론 초기화 및 연결:
   • Tello 드론을 초기화하고 연결합니다.
   • 비디오 스트림을 시작합니다.
3. capture_and_detect 함수:
   • 드론으로부터 프레임을 읽고 YOLO 네트워크로 객체 인식을 수행합니다.
   • 인식된 사람 객체를 바운딩 박스로 표시하고, 이미지를 저장합니다.
   • 결과 프레임을 화면에 표시합니다.
4. 메인 루프:
   • 드론을 이륙시키고, 객체 인식을 반복 수행합니다.
   • 'q' 키를 누르면 루프를 종료하고 드론을 착륙시킵니다.

결론

이 포스팅에서는 드론을 제어하고 비디오 스트림을 수신하여 실시간 객체 인식을 수행하는 방법을 소개했습니다. DJITelloPy를 사용하여 드론을 제어하고, OpenCV와 YOLO를 사용하여 객체 인식을 수행하며, 인식된 이미지를 저장하는 과정을 다뤘습니다.

이 예제를 통해 드론을 활용한 다양한 응용 프로그램을 개발할 수 있는 기초를 다질 수 있을 것입니다. 추가로 궁금한 점이 있으면 언제든지 댓글로 남겨주세요!