[Python] OAK-D Lite를 이용하여 실시간 이미지 분류 앱 만들기 - 데이터셋 준비

먼저 OAK-D Lite를 이용해 MP4 파일을 생성하겠습니다. 이를 위해 다음과 같이 의존성 패키지를 설치합니다.

pip install opencv-python==4.11.0.86
pip install depthai==2.30.0.0

 

아래는 make_mp4.py의 전체 코드입니다.

# make_mp4.py
import cv2
import depthai as dai

pipeline = dai.Pipeline()

camRgb = pipeline.createColorCamera()
xoutRgb = pipeline.createXLinkOut()
xoutRgb.setStreamName("rgb")

camRgb.preview.link(xoutRgb.input)

with dai.Device(pipeline) as device:
    qRgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)
    video_writer = None
    while True:
        inRgb = qRgb.get()
        frame = inRgb.getCvFrame()
        cv2.imshow("rgb", frame)

        key = cv2.waitKey(1)
        
        if key == ord('s'):
            print("Start recording")
            video_writer = cv2.VideoWriter("temp.mp4", 
                                           cv2.VideoWriter_fourcc(*'mp4v'), # mp4 코덱
                                           30,                              # 프레임 수
                                           (300,300))                       # preview 디폴트 사이즈
        elif key == ord('e'):
            print("End recording")  
            video_writer.release()
            video_writer = None

        if video_writer is not None:
            video_writer.write(frame)

        if key == ord('q'):
            break

 

 

pipeline = dai.Pipeline()

camRgb = pipeline.createColorCamera()
xoutRgb = pipeline.createXLinkOut()
xoutRgb.setStreamName("rgb")

camRgb.preview.link(xoutRgb.input)

OAK-D Lite 카메라를 사용하려면 먼저 파이프라인을 통해 노드(Node)를 생성해야 합니다. XLinkOut 노드를 통해 RGB 스트림을 외부에서 읽을 수 있도록 스트림 이름 "rgb"를 설정했습니다. 고해상도 영상은 필요하지 않으므로 RGB preview 스트림을 사용했으며 기본 해상도는 300x300입니다.

 

qRgb = device.getOutputQueue(name="rgb", maxSize=4, blocking=False)

장치에서 전송되는 메세지가 호스트에서 읽는 속도보다 빠르면 메세지 큐가 쌓일 수 있습니다. 이때 maxSize와 blocking 옵션이 큐의 동작 방식을 결정합니다. 메모리(RAM)의 크기가 작을 때, maxSize를 크게 설정하면 메모리가 금방 부하에 걸릴 수 있습니다. blocking이 True라면 큐가 꽉 찼을 때 장치가 호스트에서 메세지를 처리할 때까지 기다립니다. 이와 반대로 blocking이 False라면 큐가 꽉 찼을 때 기존에 있던 메세지를 지우고 새 메세지를 덮어쓰기 합니다.

 

while True:
    inRgb = qRgb.get()
    frame = inRgb.getCvFrame()
    cv2.imshow("rgb", frame)

    key = cv2.waitKey(1)

    if key == ord('s'):
        print("Start recording")
        video_writer = cv2.VideoWriter("temp.mp4", 
                                       cv2.VideoWriter_fourcc(*'mp4v'), # mp4 코덱
                                       30,                              # 프레임 수
                                       (300,300))                       # preview 디폴트 사이즈
    elif key == ord('e'):
        print("End recording")  
        video_writer.release()
        video_writer = None

    if video_writer is not None:
        video_writer.write(frame)

    if key == ord('q'):
        break

실행 도중 실시간으로 프레임이 출력되는 상황에서 's'키를 누르면 video_writer를 설정하여 프레임을 녹화합니다. 이어 'e'키를 눌러 녹화를 종료합니다. 'q'키를 누르면 프로그램이 종료됩니다.

 

아래는 extract_frames.py의 전체 코드 입니다.

# extract_frames.py
import cv2
import os

bodypart="temp" # 파일 이름
output_dir = f"bodyparts/{bodypart}"

if not os.path.exists(output_dir):
    os.makedirs(output_dir)

cap = cv2.VideoCapture(f"{bodypart}.mp4")

if not cap.isOpened(): 
    print("ERROR")
    exit()

frame_count = 1

while True:
    ret, frame = cap.read()

    if not ret:
        break

    frame_filename = os.path.join(output_dir, f'{bodypart}_{frame_count:04d}.jpg')

    cv2.imwrite(frame_filename, frame)
    print(f'Saved {frame_filename}')

    frame_count += 1

cap.release()

 

bodypart="temp" # 파일 이름
output_dir = f"bodyparts/{bodypart}"

if not os.path.exists(output_dir):
    os.makedirs(output_dir)

cap = cv2.VideoCapture(f"{bodypart}.mp4")

if not cap.isOpened(): 
    print("ERROR")
    exit()

기계학습을 위해 출력 디렉토리는 "bodyparts/{bodypart}" 형태로 설정했습니다. 영상 파일이 존재하지 않더라도 cv2.VideoCapture는 객체를 반환하기 때문에, 이를 고려해 에러 체크를 추가했습니다.

 

frame_count = 1

while True:
    ret, frame = cap.read()

    if not ret:
        break

    frame_filename = os.path.join(output_dir, f'{bodypart}_{frame_count:04d}.jpg')

    cv2.imwrite(frame_filename, frame)
    print(f'Saved {frame_filename}')

    frame_count += 1

cap.release()

cv2.imwrite를 사용하여 영상의 각 프레임을 .jpg 형식으로 출력 폴더에 저장합니다. 파일 이름은 {bodypart}_{프레임번호}.jpg 형식으로 지정되며, 프레임 번호는 네 자리 숫자로 패딩됩니다. 

 

손의 모양에 따라 LAT , PA로 구분하여 데이터셋을 준비합니다. LAT는 엄지와 검지로 'O'자를 만들고 손의 측면을 위치시키고 PA는 손바닥을 물체를 놓는 테이블과 맞닿게 위치시키고 손등을 카메라를 향하여 위치시킵니다. 

<그림 1. 손의 위치에 따른 mp4 추출 모습>

 

extract_frames.py를 이용하여 생성된 mp4로부터 jpg 프레임을 추출하여 준비를 마칩니다.

<그림 2. 추출된 jpg 프레임 준비된 모습>