RaspberryPI/sampleCaptureScript.py at main · InhaCapstoneProject/RaspberryPI · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
import cv2
import os
import time

# ---------------------------------------------------------
# [설정] YOLOv4-tiny 입력 크기에 맞춤
# ---------------------------------------------------------
SAVE_DIR = "dataset"
IMG_SIZE = 416  # 가로, 세로 416
# ---------------------------------------------------------

# 1. 저장 폴더 생성
if not os.path.exists(SAVE_DIR):
    os.makedirs(SAVE_DIR)
    print(f"...[정보] '{SAVE_DIR}' 폴더에 이미지가 저장됩니다.")

# 2. OpenCV VideoCapture 초기화
# libcamerify를 통해 실행하면 이 부분이 libcamera로 자동 연결됩니다.
print("...[정보] 카메라 초기화 중...")

# [수정됨] V4L2 백엔드 명시 및 YUYV 포맷 설정
# 'reshape' 에러(버퍼 크기 불일치)를 피하기 위해 Packed 포맷인 'YUYV'를 사용합니다.
# 기본 Planar 포맷들은 stride(여백) 계산 문제로 종종 충돌이 발생합니다.
cap = cv2.VideoCapture(0, cv2.CAP_V4L2)

# YUYV 포맷 설정 (중요: 이 설정이 reshape 에러를 방지합니다)
cap.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc('Y', 'U', 'Y', 'V'))
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

# 카메라 워밍업 시간
time.sleep(2)

if not cap.isOpened():
    print("❌ 카메라를 열 수 없습니다. 'libcamerify'로 실행했는지 확인해주세요.")
    exit()

# 기존 파일 수 확인
existing_files = os.listdir(SAVE_DIR)
count = len(existing_files)

print("--------------------------------------------------")
print(f"  목표 이미지 크기: {IMG_SIZE} x {IMG_SIZE}")
print(f"  현재 저장된 이미지 수: {count}장")
print("  [s] 키: 이미지 저장")
print("  [q] 키: 프로그램 종료")
print("--------------------------------------------------")

try:
    while True:
        # 3. 프레임 읽기
        ret, frame = cap.read()

        if not ret:
            print("❌ 프레임을 읽을 수 없습니다.")
            break

        # 4. 크기 강제 조정 (416x416)
        # 카메라가 어떤 해상도(640x480, 1920x1080 등)로 열리든 상관없이
        # 우리가 원하는 YOLO 입력 크기로 무조건 리사이징합니다.
        if frame.shape[0] != IMG_SIZE or frame.shape[1] != IMG_SIZE:
            frame = cv2.resize(frame, (IMG_SIZE, IMG_SIZE))

        # 5. 화면 표시용 (가이드라인 추가)
        display_frame = frame.copy()

        # 중앙 십자가 가이드라인
        cv2.line(display_frame, (IMG_SIZE//2, 0), (IMG_SIZE//2, IMG_SIZE), (0, 255, 0), 1)
        cv2.line(display_frame, (0, IMG_SIZE//2), (IMG_SIZE, IMG_SIZE//2), (0, 255, 0), 1)

        cv2.imshow("YOLO Data Collector (libcamerify)", display_frame)

        # 6. 키보드 입력 처리
        key = cv2.waitKey(1) & 0xFF

        # 's' 키: 저장
        if key == ord('s'):
            filename = os.path.join(SAVE_DIR, f"image_{count:03d}.jpg")
            cv2.imwrite(filename, frame) # 416x416 이미지 저장

            print(f"[저장됨] {filename} ({frame.shape[1]}x{frame.shape[0]})")

            # 화면 깜빡임 효과
            cv2.putText(display_frame, "SAVED!", (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
            cv2.imshow("YOLO Data Collector (libcamerify)", display_frame)
            cv2.waitKey(200)

            count += 1

        # 'q' 키: 종료
        elif key == ord('q'):
            break

finally:
    print("\n...[정보] 종료합니다.")
    cap.release()
    cv2.destroyAllWindows()