clean-dataset 0.2.0

데이터셋 정리 및 시각화 라이브러리

clean-dataset

데이터셋 정리 및 시각화 라이브러리

train/val/test로 구성된 비디오/이미지 데이터셋을 정리하고 시각화하는 도구입니다.

설치

pip install clean-dataset

또는 소스에서 설치:

git clone https://github.com/haechan/clean-dataset.git
cd clean-dataset
pip install -e .

주요 기능

1. 데이터 정리

이미지 재구성

from clean_dataset import reorganize_images

# train/{video_id}/frame.jpg -> train/image/{video_id}@{frame}.jpg
reorganize_images(
    input_dir="./data",
    output_dir="./output",
    train_json="train.json",
    val_json="val.json",
    test_json="test.json"  # 선택
)

JSON 어노테이션 생성

from clean_dataset import create_annotations, create_annotations_from_masks

# 원본 JSON에서 단순화된 bbox JSON 생성
create_annotations(
    output_dir="./output",
    train_json="train.json",
    val_json="val.json",
    use_thing_only=True  # stuff 제외
)

# 마스크에서 bbox JSON 생성
create_annotations_from_masks(
    mask_dir="./mask_semantic",
    output_dir="./json",
    category_json="train.json",  # 선택
    use_thing_only=True
)

출력 JSON 구조:

[
    {
        "id": 1,
        "tag": "person",
        "xmin": 701,
        "ymin": 0,
        "xmax": 1280,
        "ymax": 709
    }
]

마스크 생성

from clean_dataset import generate_masks

generate_masks(
    annotation_json="train.json",
    output_dir="./output",
    mask_type="instance",  # "instance" | "semantic" | "panoptic"
    split="train"  # 선택
)

카테고리 JSON 생성

from clean_dataset import generate_category_json, generate_category_json_from_mask

# 어노테이션 JSON에서 카테고리 JSON 자동 생성
generate_category_json(
    annotation_json="train.json",
    output_path="categories.json"
)

# 마스크 폴더에서 사용된 클래스만 추출하여 카테고리 JSON 생성
generate_category_json_from_mask(
    mask_dir="./masks",
    output_path="categories.json",
    category_json="train.json"  # 선택: 클래스명 매핑용
)

출력 형식 (isthing 정보가 있는 경우):

{
    "names": {
        "wall": "stuff",
        "person": "thing"
    },
    "ignore": 255
}

출력 형식 (isthing 정보가 없는 경우):

{
    "names": ["wall", "person"],
    "ignore": 255
}

2. 데이터셋 분할

비율 기반 분할

from clean_dataset import split_dataset

# JSON 또는 TXT 파일을 train/val/test로 분할
split_dataset(
    input_file="data.json",  # 또는 data.txt
    output_dir="./output",
    train_ratio=0.8,
    val_ratio=0.1,
    test_ratio=0.1,
    shuffle=True,
    seed=42
)

리스트 기반 분할

from clean_dataset import split_by_list

# 미리 정의된 리스트로 분할
split_by_list(
    input_file="annotations.json",
    train_list="train_ids.txt",
    val_list="val_ids.txt",
    test_list="test_ids.txt",
    output_dir="./output"
)

폴더 분할

from clean_dataset import split_folder

# 폴더 내 파일을 train/val/test로 분할
split_folder(
    input_dir="./images",
    output_dir="./output",
    train_ratio=0.8,
    val_ratio=0.1,
    test_ratio=0.1,
    extensions=['.jpg', '.png'],
    copy_files=True  # True: 파일 복사, False: 리스트만 저장
)

리스트 기반 파일 복사

from clean_dataset import copy_files_by_list

# txt 리스트를 기반으로 파일 복사
copy_files_by_list(
    source_dir="./images",
    output_dir="./output",
    train_list="train.txt",
    val_list="val.txt",
    test_list="test.txt"
)
# 결과: output/train/, output/val/, output/test/ 폴더에 파일 복사

K-Fold 분할

from clean_dataset import create_kfold_splits

# K-Fold 교차 검증용 분할 생성
create_kfold_splits(
    input_file="data.json",
    output_dir="./kfold",
    k=5,
    shuffle=True
)

지원 JSON 형식:

• YouTube-VOS 형식: {"videos": {...}}
• COCO 형식: {"images": [...], "annotations": [...]}
• Annotations 형식: {"annotations": [...]}
• 리스트 형식: ["item1", "item2"] 또는 [{"id": 1, ...}]

3. 시각화

팔레트 이미지 생성

from clean_dataset import generate_palette_image

generate_palette_image(
    output_path="palette.jpg",
    annotation_json="train.json"
)

Detection 시각화

from clean_dataset import visualize_detection

visualize_detection(
    image_path="image.jpg",
    json_path="annotation.json",
    output_dir="./viz",
    show_class_name=True,
    show_legend=True  # 오른쪽에 팔레트 범례 표시
)

Segmentation 시각화

from clean_dataset import visualize_segmentation

visualize_segmentation(
    image_dir="./images",
    mask_dir="./masks",
    json_dir="./json",
    output_dir="./viz",
    is_video=True,  # True: strip으로 연결, False: 단일 이미지
    alpha=0.5
)

카테고리 관리

from clean_dataset import load_categories, category_manager

# JSON에서 카테고리 로드
load_categories("train.json")

# 카테고리 정보 조회
print(category_manager.get_class_name(1))  # "person"
print(category_manager.is_thing(1))  # True
print(len(category_manager))  # 카테고리 수

지원 형식

• 입력: YouTube-VOS, VIPSeg, COCO 형식의 JSON
• 카테고리: 딕셔너리 형식 ({"wall": "stuff"}) 또는 리스트 형식 ([{"id": 1, "name": "person"}])
• 마스크: RLE 인코딩, RGB panoptic 마스크

의존성

• numpy >= 1.20.0
• Pillow >= 8.0.0
• opencv-python >= 4.5.0
• sanghyunjo

라이선스

MIT License