YOLO 객체 검출 데이터 SQLite 저장
• 1 min read
YOLO 객체 검출 모델은 결과를 프레임별 txt 파일로 분할 저장하는 단점이 있다. 파일 수량이 많아질 뿐만 아니라, 파일 내부 클래스가 인덱스 번호로 기록되어 직관적인 확인이 어렵다.
과거 다수 영상 내 객체 검출 업무를 진행하며, 폴더 내 영상을 일괄 처리하고 YOLO로 타겟 클래스를 검출한 뒤 FFmpeg로 영상에 Bounding Box를 그리는 파이프라인을 구축했다. 본 문서에서는 해당 스크립트를 변경해 검출 결과를 SQLite3 데이터베이스에 통합 저장하는 구조를 정리한다.
- DB 스키마
데이터 중복을 방지하기 위해 영상 메타데이터(videos)와 검출 결과(detections) 테이블을 분리했다. video_id 컬럼에 인덱스를 추가해 특정 영상의 데이터 추출 속도를 개선했다.
conn.executescript("""
CREATE TABLE IF NOT EXISTS videos (
id INTEGER PRIMARY KEY AUTOINCREMENT,
filename TEXT UNIQUE
);
CREATE TABLE IF NOT EXISTS detections (
id INTEGER PRIMARY KEY AUTOINCREMENT,
video_id INTEGER,
frame_no INTEGER,
class_name TEXT,
confidence REAL,
x_center_rel REAL, y_center_rel REAL,
width_rel REAL, height_rel REAL,
FOREIGN KEY (video_id) REFERENCES videos (id)
);
CREATE INDEX IF NOT EXISTS idx_video_id ON detections(video_id);
""")- ROI 필터링 및 데이터 추출
모든 객체를 저장하는 대신, 박스 중심점(cx, cy) 좌표가 설정한 ROI 범위에 들어왔을 때만 데이터베이스에 기록하도록 필터링을 적용했다.
for frame_idx, res in enumerate(results):
batch = []
if res.boxes is not None:
for box in res.boxes:
x1, y1, x2, y2 = box.xyxy[0].tolist()
cx, cy = (x1 + x2) / 2, (y1 + y2) / 2
if ROI[0] <= cx <= ROI[2] and ROI[1] <= cy <= ROI[3]:
conf = float(box.conf[0])
cls = int(box.cls[0])
xywhn = box.xywhn[0].tolist()
batch.append(
(video_id, frame_idx, names[cls], conf, *xywhn)
)- 배치 처리
프레임 단위로 DB 커밋을 실행하면 I/O 병목으로 성능이 저하된다. 따라서 데이터 INSERT는 프레임 단위로 실행하고, COMMIT 연산은 단일 영상 처리가 끝나는 시점에 일괄 처리하도록 변경했다.