# YOLO Annotation Utility Tool

YOLO 형식의 어노테이션 파일을 자유롭게 조작할 수 있는 유틸리티 스크립트입니다.

## 주요 기능

- **대화형 모드** (interactive): 터미널에서 대화형으로 편집 (추천!)
- **클래스 제거** (remove): 특정 클래스 삭제 (ID 또는 이름으로 지정 가능)
- **클래스 통합** (merge): 여러 클래스를 하나로 병합 (ID 또는 이름으로 지정 가능)
- **클래스 이름 변경** (rename): 클래스 이름 수정 (ID 또는 이름으로 지정 가능)
- **클래스 ID 재할당** (reindex): ID를 순차적으로 재정렬
- **정보 확인** (info): 어노테이션 파일 정보 출력 (클래스별 개수 및 합계 표시)

## 설치

PyYAML 라이브러리가 필요합니다:

```bash
pip install pyyaml
```

## YOLO 데이터셋 구조

이 도구는 다음과 같은 YOLO 데이터셋 구조를 기준으로 동작합니다:

```
dataset/
├── data.yaml          # 클래스 정의 파일
├── images/
│   ├── train/
│   │   ├── image1.jpg
│   │   └── image2.jpg
│   └── val/
│       ├── image3.jpg
│       └── image4.jpg
└── labels/
    ├── train/
    │   ├── image1.txt
    │   └── image2.txt
    └── val/
        ├── image3.txt
        └── image4.txt
```

**data.yaml 예시:**
```yaml
names:
  0: person
  1: car
  2: truck
  3: bus
train: images/train
val: images/val
```

**label 파일(.txt) 예시:**
```
0 0.5 0.5 0.3 0.4
1 0.7 0.3 0.2 0.3
```
형식: `class_id center_x center_y width height` (모든 좌표는 정규화된 값 0-1)

## 빠른 시작 - 대화형 모드 (권장)

인수 없이 실행하면 대화형 모드로 진입합니다:

```bash
python yolo_annotation_modify.py
```

대화형 모드에서는:
1. **데이터셋 선택**: 자동으로 찾은 YOLO 데이터셋 목록에서 선택하거나 직접 경로 입력
2. **작업 선택**: 메뉴에서 원하는 작업을 선택하고 여러 작업을 순차적으로 수행 가능
3. **저장**: 모든 작업 완료 후 결과를 저장

### 대화형 모드 예시

```
============================================================
  YOLO 어노테이션 편집기 - 대화형 모드
============================================================

[1/3] 데이터셋 선택
------------------------------------------------------------

발견된 YOLO 데이터셋:
  1. dataset/yolo
  2. data/custom_dataset
  0. 직접 경로 입력

데이터셋 번호 선택 (직접 입력은 0): 1

[+] 로딩 중: dataset/yolo
[+] Found 150 label files

============================================================
Dataset: dataset/yolo
============================================================

전체 통계:
  - 라벨 파일 수: 150
  - 클래스 수: 3

클래스 상세:
ID     클래스명                        어노테이션 수
------------------------------------------------------------
0      person                        450
1      car                           320
2      truck                         180
------------------------------------------------------------
합계                                  950
============================================================

[2/3] 작업 선택
------------------------------------------------------------

사용 가능한 작업:
  1. 클래스 제거
  2. 클래스 통합 (병합)
  3. 클래스 이름 변경
  4. 클래스 ID 재할당
  5. 현재 정보 표시
  0. 완료 (저장 단계로)

작업 선택 (0-5): 2

현재 클래스:
  0: person
  1: car
  2: truck

통합할 클래스 입력 (ID 또는 이름, 쉼표로 구분): 1,2

통합 후 클래스 이름 입력: vehicle

[*] 'vehicle'로 통합 중: ['1', '2']
  - 통합 대상: car (ID: 1)
  - 통합 대상: truck (ID: 2)
[+] 'vehicle' (ID: 1)로 통합되었습니다
[+] 500개의 어노테이션이 변경되었습니다

작업 선택 (0-5): 0

[3/3] 결과 저장
------------------------------------------------------------

수행된 작업:
  1. 통합: 1, 2 → vehicle

[!] 주의: labels 폴더의 모든 .txt 파일이 수정됩니다.
labels 폴더 백업 생성? (Y/n): Y

[*] labels 폴더 백업 중...
[+] 백업 생성됨: dataset/yolo/labels_backup_20250117_143025

data.yaml 백업 생성? (Y/n): Y

[*] 저장 중...
[+] data.yaml 백업 생성됨: dataset/yolo/data_backup_20250117_143026.yaml
[+] data.yaml 저장 완료

============================================================
  작업이 성공적으로 완료되었습니다!
============================================================

데이터셋: dataset/yolo
설정 파일: dataset/yolo/data.yaml
라벨 파일: 150개 업데이트됨
```

---

## CLI 모드 사용법

명령행 인수를 사용하여 스크립트 방식으로도 실행할 수 있습니다.

### 1. 정보 확인 (info)

데이터셋의 전체 정보를 확인합니다.

```bash
python yolo_annotation_modify.py info --dataset dataset/yolo
```

출력 예시:
```
============================================================
Dataset: dataset/yolo
============================================================

전체 통계:
  - 라벨 파일 수: 150
  - 클래스 수: 3

클래스 상세:
ID     클래스명                        어노테이션 수
------------------------------------------------------------
0      person                        450
1      car                           320
2      truck                         180
------------------------------------------------------------
합계                                  950
============================================================
```

### 2. 클래스 제거 (remove)

특정 클래스와 해당 어노테이션을 삭제합니다. **클래스 이름 또는 ID로 지정 가능**합니다.

```bash
# 이름으로 클래스 제거
python yolo_annotation_modify.py remove \
  --dataset dataset/yolo \
  --classes "person,bicycle"

# ID로 클래스 제거
python yolo_annotation_modify.py remove \
  --dataset dataset/yolo \
  --classes "0,5"

# 이름과 ID 혼합 가능
python yolo_annotation_modify.py remove \
  --dataset dataset/yolo \
  --classes "0,bicycle,5"

# 백업 생성 안 함
python yolo_annotation_modify.py remove \
  --dataset dataset/yolo \
  --classes "0,1" \
  --no-backup
```

### 3. 클래스 통합 (merge)

여러 클래스를 하나로 병합합니다. **클래스 이름 또는 ID로 지정 가능**합니다.

```bash
# 이름으로 클래스 통합
python yolo_annotation_modify.py merge \
  --dataset dataset/yolo \
  --source "car,truck,bus,motorcycle" \
  --target "vehicle"

# ID로 클래스 통합
python yolo_annotation_modify.py merge \
  --dataset dataset/yolo \
  --source "1,2,3,4" \
  --target "vehicle"

# 이름과 ID 혼합 가능
python yolo_annotation_modify.py merge \
  --dataset dataset/yolo \
  --source "1,truck,3" \
  --target "vehicle"

# 백업 생성 안 함
python yolo_annotation_modify.py merge \
  --dataset dataset/yolo \
  --source "1,2,3" \
  --target "vehicle" \
  --no-backup
```

### 4. 클래스 이름 변경 (rename)

클래스 이름을 변경합니다. **기존 클래스를 이름 또는 ID로 지정 가능**합니다.

```bash
# 이름으로 클래스 이름 변경
python yolo_annotation_modify.py rename \
  --dataset dataset/yolo \
  --mapping "person:human,car:automobile"

# ID로 클래스 이름 변경
python yolo_annotation_modify.py rename \
  --dataset dataset/yolo \
  --mapping "0:human,1:automobile"

# 이름과 ID 혼합 가능
python yolo_annotation_modify.py rename \
  --dataset dataset/yolo \
  --mapping "0:human,car:automobile"
```

### 5. 클래스 ID 재할당 (reindex)

클래스 ID를 순차적으로 재할당합니다. 클래스를 삭제하거나 통합한 후 ID를 정리할 때 유용합니다.

```bash
# ID를 0부터 재할당 (YOLO 기본값)
python yolo_annotation_modify.py reindex \
  --dataset dataset/yolo

# ID를 1부터 재할당
python yolo_annotation_modify.py reindex \
  --dataset dataset/yolo \
  --start 1

# 백업 생성 안 함
python yolo_annotation_modify.py reindex \
  --dataset dataset/yolo \
  --no-backup
```

## 고급 사용 예시

### 예시 1: 데이터셋 정리 워크플로우

```bash
# 1. 현재 상태 확인
python yolo_annotation_modify.py info --dataset dataset/yolo

# 2. ID로 불필요한 클래스 제거
python yolo_annotation_modify.py remove \
  --dataset dataset/yolo \
  --classes "5,7,9"

# 3. 유사 클래스 통합
python yolo_annotation_modify.py merge \
  --dataset dataset/yolo \
  --source "1,2,3" \
  --target "vehicle"

# 4. 클래스 이름 정리
python yolo_annotation_modify.py rename \
  --dataset dataset/yolo \
  --mapping "0:person,vehicle:car"

# 5. ID 재할당
python yolo_annotation_modify.py reindex \
  --dataset dataset/yolo \
  --start 0

# 6. 최종 결과 확인
python yolo_annotation_modify.py info --dataset dataset/yolo
```

### 예시 2: Train/Valid 데이터셋 동시 처리

YOLO 데이터셋의 경우 train과 val이 같은 data.yaml을 공유하므로, 한 번의 명령으로 모든 라벨 파일이 함께 업데이트됩니다:

```bash
# 데이터셋 전체 클래스 이름 변경 (train/val 모두 자동 처리)
python yolo_annotation_modify.py rename \
  --dataset dataset/yolo \
  --mapping "0:person,1:vehicle"
```

### 예시 3: 대규모 클래스 통합

COCO 80개 클래스를 커스텀 클래스로 통합하는 예시:

```bash
# 모든 vehicle 클래스를 하나로 통합
python yolo_annotation_modify.py merge \
  --dataset dataset/yolo \
  --source "car,truck,bus,motorcycle,bicycle" \
  --target "vehicle"

# 모든 animal 클래스를 하나로 통합
python yolo_annotation_modify.py merge \
  --dataset dataset/yolo \
  --source "cat,dog,bird,horse,cow" \
  --target "animal"

# ID 재정렬
python yolo_annotation_modify.py reindex \
  --dataset dataset/yolo
```

## Python 스크립트에서 사용

유틸리티를 Python 코드에서 직접 사용할 수도 있습니다.

```python
from yolo_annotation_modify import YOLOAnnotationEditor

# 데이터셋 로드
editor = YOLOAnnotationEditor('dataset/yolo')

# 정보 확인
editor.print_info()

# 작업 수행 (이름 또는 ID로 지정 가능)
editor.remove_classes(['0', '1'])  # ID로 제거
editor.remove_classes(['person', 'car'])  # 이름으로 제거
editor.rename_classes({'0': 'human', 'car': 'automobile'})  # ID와 이름 혼합 가능
editor.reindex_classes(start_id=0)

# 저장 (data.yaml만 저장, label 파일은 이미 수정됨)
editor.save_yaml(backup=True)

# 체이닝도 가능
editor = YOLOAnnotationEditor('dataset/yolo')
editor.remove_classes(['5', '7']) \
      .merge_classes(['1', '2'], 'vehicle') \
      .reindex_classes(start_id=0) \
      .save_yaml(backup=True)
```

## 옵션

### 백업 관련

기본적으로 labels 폴더와 data.yaml 파일 모두 백업이 생성됩니다.

```bash
# 백업 생성 (기본값)
python yolo_annotation_modify.py rename --dataset dataset/yolo --mapping "0:human"

# 모든 백업 생성 안 함
python yolo_annotation_modify.py rename --dataset dataset/yolo --mapping "0:human" --no-backup

# labels 폴더 백업만 생성 안 함 (data.yaml은 백업)
python yolo_annotation_modify.py rename --dataset dataset/yolo --mapping "0:human" --no-labels-backup
```

백업 파일명 형식:
- Labels 폴더: `labels_backup_{YYYYMMDD_HHMMSS}/`
- data.yaml: `data_backup_{YYYYMMDD_HHMMSS}.yaml`

## 주요 특징

### ID와 이름 모두 지원
- 모든 작업(제거, 통합, 이름 변경)에서 클래스를 **ID 또는 이름**으로 지정 가능
- 숫자로 입력하면 자동으로 ID로 인식, 문자로 입력하면 이름으로 인식
- ID와 이름을 혼합하여 사용 가능

### 어노테이션 통계
- 클래스별 어노테이션 개수 자동 계산
- 전체 어노테이션 합계 표시
- 실시간 정보 확인 가능

### 자동 백업
- labels 폴더 전체 백업
- data.yaml 파일 백업
- 타임스탬프로 백업 파일 구분

### 입력 오류 처리
- 잘못된 경로 입력 시 재입력 요청
- 잘못된 선택 입력 시 재입력 요청
- 존재하지 않는 클래스 ID/이름 입력 시 경고 표시

## YOLO vs COCO 차이점

| 항목 | YOLO | COCO (MMDetection) |
|------|------|---------------------|
| 설정 파일 | data.yaml | JSON |
| 클래스 ID | 0부터 시작 | 1부터 시작 |
| 어노테이션 파일 | 이미지별 .txt 파일 | 단일 JSON 파일 |
| 좌표 형식 | 정규화된 중심점+크기 | 픽셀 좌표 (bbox/polygon) |
| 백업 대상 | labels 폴더 전체 + data.yaml | JSON 파일 하나 |

## 주의사항

1. **데이터 백업**: 중요한 데이터는 항상 백업 후 작업하세요.
2. **Train/Val 일치**: Train과 Val은 같은 data.yaml을 공유하므로 한 번의 작업으로 모두 수정됩니다.
3. **ID 순서**: `reindex` 명령은 기존 ID 순서를 기준으로 재할당합니다.
4. **병합 ID**: `merge` 명령은 가장 작은 ID를 새 카테고리 ID로 사용합니다.
5. **좌표 유지**: 모든 작업은 bounding box 좌표를 그대로 유지하고 클래스 ID만 수정합니다.
6. **파일 직접 수정**: label 파일(.txt)은 원본이 직접 수정되므로 반드시 백업을 생성하세요.

## 문제 해결

### data.yaml을 찾을 수 없음

```bash
# 올바른 데이터셋 디렉토리 지정 (data.yaml이 있는 폴더)
python yolo_annotation_modify.py info --dataset dataset/yolo

# 절대 경로 사용
python yolo_annotation_modify.py info --dataset /full/path/to/dataset
```

### label 파일 형식 오류

YOLO 형식 확인:
```python
# 올바른 형식: class_id center_x center_y width height
0 0.5 0.5 0.3 0.4
1 0.7 0.3 0.2 0.3
```

모든 값은 0-1 사이의 정규화된 값이어야 합니다.

### data.yaml 형식 확인

```yaml
names:
  0: class_name_1
  1: class_name_2
  2: class_name_3
train: images/train
val: images/val
```

## 라이선스

이 도구는 YOLO 형식 어노테이션 작업을 위한 유틸리티입니다.