clean-dataset 0.1.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"  # 선택
)

2. 시각화

팔레트 이미지 생성

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 형식의 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