computer-visiondatasetcocosqlitedata-engineering

SQLite DB 기반으로 COCO annotation 서브셋 추출하기

2026-01-25 • 3 min read

소스코드

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 split 별 JSON 생성

전체 흐름

스크립트의 전체 동작 흐름은 다음과 같다.

SQLite DB에서 category 테이블 로드
category name → id 매핑 (선택)
images 테이블에서 이미지 폴더 기준으로 필터링
annotations 테이블에서 image_id 기준으로 chunk 단위 조회
(선택) category_id 기준 추가 필터링
COCO 형식 JSON으로 재구성 후 저장

이미지 기준 필터링

먼저 대상 이미지 폴더에 실제로 존재하는 파일명을 기준으로 DB의 images 테이블에서 해당 이미지들만 선별한다.

이런 방식으로 JSON 전체를 순회하는 대신 DB 쿼리 + 파일명 매칭으로 빠르게 필터링한다.

extract-subset-from-sqlite.py

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)이 있으며, 동시에 SQL 문장 자체의 길이에도 제한(SQLITE_MAX_SQL_LENGTH)이 존재한다. IN (...) 절에 수천 개 이상의 ID가 포함될 경우 변수 개수 제한 또는 SQL 텍스트 길이 제한 중 하나에 걸려 쿼리가 실패할 수 있다. 이를 피하기 위해 image_id 목록을 chunk 단위로 나누어 여러 번 조회하도록 구현하였다.

extract-subset-from-sqlite.py

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_names
target_category_ids

중 하나를 넘기면 내부적으로 category_id 기준으로 변환하여 쿼리에 반영된다.

extract-subset-from-sqlite.py

def load_category_mapping(cursor):
    mapping = {}
    cursor.execute("SELECT id, name FROM categories")
    for cat_id, name in cursor:
        mapping[name] = cat_id
    return mapping

COCO JSON 생성

최종적으로 다음 항목들을 모아 COCO JSON을 생성한다.

images
annotations
categories

필요한 category만 남기도록 정리한 다음 patch 단위로 JSON 파일을 생성한다.

extract-subset-from-sqlite.py

def build_coco_json(images, annotations, categories):
    return {
        "images": images,
        "annotations": annotations,
        "categories": categories,
        "info": {"description": "Filtered COCO dataset from DB"},
        "license": [],
    }

extract-subset-from-sqlite.py

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로 변환해 두고 쿼리 기반으로 서브셋을 추출하면:

속도
메모리 사용량
반복 작업 자동화

측면에서 훨씬 효율성이 개선된다.

Back to all posts