SQLite DB 기반으로 COCO annotation 서브셋 추출하기
COCO 형식 대용량 annotation JSON을 다루다 보면 특정 디렉토리의 이미지 라벨이나 특정 클래스 라벨만 추출해야 하는 상황이 발생한다.
Objects365처럼 용량이 큰 데이터셋의 한계:
- JSON 전체 로딩/파싱 비용 증가
- 단순 스크립트로는 메모리 사용량 급증
- 반복적인 필터링 작업의 비효율성
이에 annotation 데이터를 SQLite DB로 변환한 뒤, 필요할 때 DB 쿼리로 서브셋을 추출하는 방식을 적용했다.
본 문서에서는 다음 사항을 한 번에 처리하는 스크립트를 정리한다:
- 특정 이미지 폴더 기준 이미지 및 annotation 필터링
- 필요 시 특정 category 추가 필터링
- 대용량 처리를 위한 chunk 단위 쿼리
사용 사례
주요 사용 목적은 다음과 같다.
- 특정 폴더에 위치한 이미지 기준으로 COCO JSON 생성
- patch 단위로 JSON을 분할하여 관리 (patch_0, patch_1, …)
- 필요 시 category name이나 id 기준으로 클래스 필터링
예시:
- Objects365 데이터셋에서 특정 이미지 세트만 추출
- 클래스 일부만 포함한 경량 서브셋 생성
- 학습용 train/val 분할 JSON 생성
전체 흐름
스크립트의 전체 동작 흐름은 다음과 같다.
- SQLite DB에서 category 테이블 로드
- category name → id 매핑 (선택)
- images 테이블에서 이미지 폴더 기준으로 데이터 필터링
- annotations 테이블에서 image_id 기준으로 chunk 단위 조회
- (선택) category_id 기준 추가 필터링
- COCO 형식 JSON으로 재구성 후 저장
이미지 기준 필터링
우선 대상 이미지 폴더에 존재하는 파일명을 기준으로 DB images 테이블에서 해당 이미지들만 선별한다.
JSON 전체를 순회하는 대신 DB 쿼리와 파일명 매칭을 활용해 속도를 개선했다.
def fetch_filtered_images(cursor, image_path: Path):
target_images = set(os.listdir(image_path))
images = []
image_ids = []
cursor.execute(
"SELECT id, file_name, width, height, license, flickr_url, coco_url, date_captured FROM images"
)
for row in tqdm(cursor, desc="Filter DB Images"):
(
img_id,
file_path,
width,
height,
license,
flickr_url,
coco_url,
date_captured,
) = row
file_name = Path(file_path).name
if file_name in target_images:
image_ids.append(img_id)
images.append(
{
"id": img_id,
"file_name": file_name,
"width": width,
"height": height,
"license": license,
"flickr_url": flickr_url,
"coco_url": coco_url,
"date_captured": date_captured,
}
)
return images, image_ids대용량 annotations chunk 처리
SQLite는 한 SQL 문장에 바인딩할 수 있는 변수 개수 제한(SQLITE_MAX_VARIABLE_NUMBER)과 문장 길이 제한(SQLITE_MAX_SQL_LENGTH)이 존재한다. IN (...) 절에 ID가 수천 개 이상 포함되면 제한으로 인해 쿼리가 실패할 수 있다. 이 문제를 해결하기 위해 image_id 목록을 chunk 단위로 분할하여 조회하도록 구현했다.
def fetch_filtered_annotations(
cursor, image_ids, target_category_ids=None, chunk_size=1000
):
annotations = []
for i in tqdm(
range(0, len(image_ids), chunk_size),
desc="Filter DB annotations",
):
chunk_ids = image_ids[i : i + chunk_size]
image_placeholders = ",".join("?" * len(chunk_ids))
params = list(chunk_ids)
sql = f"""
SELECT id, image_id, category_id, bbox, area, iscrowd, segmentation
FROM annotations
WHERE image_id IN ({image_placeholders})
"""
if target_category_ids:
cat_placeholders = ",".join("?" * len(target_category_ids))
sql += f" AND category_id IN ({cat_placeholders})"
params.extend(target_category_ids)
cursor.execute(sql, params)적용 효과는 다음과 같다:
- 대용량 이미지 세트에서도 안정적인 처리 가능
- 메모리 사용량 제어 용이
category 기준 추가 필터링
필요 시 category name이나 category id 기준으로 annotation을 추가로 필터링할 수 있다.
target_category_namestarget_category_ids
중 하나를 인자로 입력하면 내부에서 category_id 기준으로 변환해 쿼리에 반영한다.
def load_category_mapping(cursor):
mapping = {}
cursor.execute("SELECT id, name FROM categories")
for cat_id, name in cursor:
mapping[name] = cat_id
return mappingCOCO JSON 생성
최종적으로 다음 데이터를 모아 COCO JSON을 구성한다.
- images
- annotations
- categories
필요한 category만 남긴 뒤 patch 단위로 JSON 파일을 생성한다.
def build_coco_json(images, annotations, categories):
return {
"images": images,
"annotations": annotations,
"categories": categories,
"info": {"description": "Filtered COCO dataset from DB"},
"license": [],
}def save_coco_json(coco_data, output_dir: Path, patch_number, split_name):
output_dir = output_dir / split_name
output_dir.mkdir(parents=True, exist_ok=True)
output_json = output_dir / f"patch_{patch_number}.json"
with open(output_json, "w", encoding="utf-8") as f:
json.dump(coco_data, f, indent=4, ensure_ascii=False)
print(f"[OK] Saved: {output_json}")정리
COCO 형식 대용량 annotation JSON을 통째로 메모리에 로드하는 대신 SQLite DB로 변환해 서브셋을 쿼리하면:
- 속도
- 메모리 사용량
- 반복 작업 자동화
측면에서 처리 효율을 개선할 수 있다.