yolo, mmedet 라벨링 수정도구
This commit is contained in:
9
.claude/settings.local.json
Normal file
9
.claude/settings.local.json
Normal file
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"permissions": {
|
||||
"allow": [
|
||||
"Bash(python yolo_annotation_modify.py:*)"
|
||||
],
|
||||
"deny": [],
|
||||
"ask": []
|
||||
}
|
||||
}
|
||||
395
README.md
Normal file
395
README.md
Normal file
@@ -0,0 +1,395 @@
|
||||
# COCO Annotation Utility Tool
|
||||
|
||||
COCO 형식의 어노테이션 파일을 자유롭게 조작할 수 있는 유틸리티 스크립트입니다.
|
||||
|
||||
## 주요 기능
|
||||
|
||||
- **대화형 모드** (interactive): 터미널에서 대화형으로 편집 (추천!)
|
||||
- **클래스 제거** (remove): 특정 클래스 삭제 (ID 또는 이름으로 지정 가능)
|
||||
- **클래스 통합** (merge): 여러 클래스를 하나로 병합 (ID 또는 이름으로 지정 가능)
|
||||
- **클래스 이름 변경** (rename): 클래스 이름 수정 (ID 또는 이름으로 지정 가능)
|
||||
- **클래스 ID 재할당** (reindex): ID를 순차적으로 재정렬
|
||||
- **정보 확인** (info): 어노테이션 파일 정보 출력 (클래스별 개수 및 합계 표시)
|
||||
|
||||
## 설치
|
||||
|
||||
별도 설치 불필요. Python 3 기본 라이브러리만 사용합니다.
|
||||
|
||||
## 빠른 시작 - 대화형 모드 (권장)
|
||||
|
||||
인수 없이 실행하면 대화형 모드로 진입합니다:
|
||||
|
||||
```bash
|
||||
python Utility_lableing_tool.py
|
||||
```
|
||||
|
||||
대화형 모드에서는:
|
||||
1. **파일 선택**: 자동으로 찾은 어노테이션 파일 목록에서 선택하거나 직접 경로 입력
|
||||
2. **작업 선택**: 메뉴에서 원하는 작업을 선택하고 여러 작업을 순차적으로 수행 가능
|
||||
3. **저장**: 모든 작업 완료 후 결과를 저장
|
||||
|
||||
### 대화형 모드 예시
|
||||
|
||||
```
|
||||
============================================================
|
||||
COCO 어노테이션 편집기 - 대화형 모드
|
||||
============================================================
|
||||
|
||||
[1/3] 어노테이션 파일 선택
|
||||
------------------------------------------------------------
|
||||
|
||||
발견된 어노테이션 파일:
|
||||
1. dataset/annotations/instances_train.json
|
||||
2. dataset/annotations/instances_valid.json
|
||||
0. 직접 경로 입력
|
||||
|
||||
파일 번호 선택 (직접 입력은 0): 1
|
||||
|
||||
[+] 로딩 중: dataset/annotations/instances_train.json
|
||||
|
||||
============================================================
|
||||
파일: dataset/annotations/instances_train.json
|
||||
============================================================
|
||||
|
||||
전체 통계:
|
||||
- 이미지 수: 10
|
||||
- 어노테이션 수: 29
|
||||
- 클래스 수: 2
|
||||
|
||||
클래스 상세:
|
||||
ID 클래스명 어노테이션 수
|
||||
-------------------------------------------------------
|
||||
1 tt 13
|
||||
2 ss 16
|
||||
-------------------------------------------------------
|
||||
합계 29
|
||||
============================================================
|
||||
|
||||
[2/3] 작업 선택
|
||||
------------------------------------------------------------
|
||||
|
||||
사용 가능한 작업:
|
||||
1. 클래스 제거
|
||||
2. 클래스 통합 (병합)
|
||||
3. 클래스 이름 변경
|
||||
4. 클래스 ID 재할당
|
||||
5. 현재 정보 표시
|
||||
0. 완료 (저장 단계로)
|
||||
|
||||
작업 선택 (0-5): 3
|
||||
|
||||
현재 클래스:
|
||||
- tt (ID: 1)
|
||||
- ss (ID: 2)
|
||||
|
||||
매핑 입력 (형식: 기존:새이름,기존:새이름 - 쉼표로 구분)
|
||||
예시: tt:transformer,ss:substation 또는 1:transformer,2:substation
|
||||
|
||||
입력: tt:transformer,ss:substation
|
||||
|
||||
[*] 클래스 이름 변경 중
|
||||
- 변경: 'tt' → 'transformer' (ID: 1)
|
||||
- 변경: 'ss' → 'substation' (ID: 2)
|
||||
[+] 2개 클래스 이름이 변경되었습니다
|
||||
|
||||
작업 선택 (0-5): 0
|
||||
|
||||
[3/3] 결과 저장
|
||||
------------------------------------------------------------
|
||||
|
||||
수행된 작업:
|
||||
1. 이름 변경: 2개 클래스
|
||||
|
||||
입력 파일: dataset/annotations/instances_train.json
|
||||
|
||||
출력 파일 경로 입력 (기본값: dataset/annotations/instances_train_edited.json):
|
||||
|
||||
출력 파일이 존재하면 백업 생성? (Y/n): Y
|
||||
|
||||
[*] 저장 중: dataset/annotations/instances_train_edited.json...
|
||||
[+] 저장 완료: dataset/annotations/instances_train_edited.json
|
||||
|
||||
============================================================
|
||||
작업이 성공적으로 완료되었습니다!
|
||||
============================================================
|
||||
|
||||
입력: dataset/annotations/instances_train.json
|
||||
출력: dataset/annotations/instances_train_edited.json
|
||||
(파일이 존재했다면 백업이 생성되었습니다)
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## CLI 모드 사용법
|
||||
|
||||
명령행 인수를 사용하여 스크립트 방식으로도 실행할 수 있습니다.
|
||||
|
||||
### 1. 정보 확인 (info)
|
||||
|
||||
어노테이션 파일의 전체 정보를 확인합니다.
|
||||
|
||||
```bash
|
||||
python Utility_lableing_tool.py info --input dataset/annotations/instances_train.json
|
||||
```
|
||||
|
||||
출력 예시:
|
||||
```
|
||||
============================================================
|
||||
파일: dataset/annotations/instances_train.json
|
||||
============================================================
|
||||
|
||||
전체 통계:
|
||||
- 이미지 수: 10
|
||||
- 어노테이션 수: 29
|
||||
- 클래스 수: 2
|
||||
|
||||
클래스 상세:
|
||||
ID 클래스명 어노테이션 수
|
||||
-------------------------------------------------------
|
||||
1 tt 13
|
||||
2 ss 16
|
||||
-------------------------------------------------------
|
||||
합계 29
|
||||
============================================================
|
||||
```
|
||||
|
||||
### 2. 클래스 제거 (remove)
|
||||
|
||||
특정 클래스와 해당 어노테이션을 삭제합니다. **클래스 이름 또는 ID로 지정 가능**합니다.
|
||||
|
||||
```bash
|
||||
# 이름으로 클래스 제거
|
||||
python Utility_lableing_tool.py remove \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_filtered.json \
|
||||
--classes "person,car,truck"
|
||||
|
||||
# ID로 클래스 제거
|
||||
python Utility_lableing_tool.py remove \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_filtered.json \
|
||||
--classes "1,2,5"
|
||||
|
||||
# 이름과 ID 혼합 가능
|
||||
python Utility_lableing_tool.py remove \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_filtered.json \
|
||||
--classes "1,person,5"
|
||||
```
|
||||
|
||||
### 3. 클래스 통합 (merge)
|
||||
|
||||
여러 클래스를 하나로 병합합니다. **클래스 이름 또는 ID로 지정 가능**합니다.
|
||||
|
||||
```bash
|
||||
# 이름으로 클래스 통합
|
||||
python Utility_lableing_tool.py merge \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_merged.json \
|
||||
--source "car,truck,bus,motorcycle" \
|
||||
--target "vehicle"
|
||||
|
||||
# ID로 클래스 통합
|
||||
python Utility_lableing_tool.py merge \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_merged.json \
|
||||
--source "1,2,3,4" \
|
||||
--target "vehicle"
|
||||
|
||||
# 이름과 ID 혼합 가능
|
||||
python Utility_lableing_tool.py merge \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_merged.json \
|
||||
--source "1,truck,3" \
|
||||
--target "vehicle"
|
||||
```
|
||||
|
||||
### 4. 클래스 이름 변경 (rename)
|
||||
|
||||
클래스 이름을 변경합니다. **기존 클래스를 이름 또는 ID로 지정 가능**합니다.
|
||||
|
||||
```bash
|
||||
# 이름으로 클래스 이름 변경
|
||||
python Utility_lableing_tool.py rename \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_renamed.json \
|
||||
--mapping "tt:transformer,ss:substation"
|
||||
|
||||
# ID로 클래스 이름 변경
|
||||
python Utility_lableing_tool.py rename \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_renamed.json \
|
||||
--mapping "1:transformer,2:substation"
|
||||
|
||||
# 이름과 ID 혼합 가능
|
||||
python Utility_lableing_tool.py rename \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_renamed.json \
|
||||
--mapping "1:transformer,ss:substation"
|
||||
```
|
||||
|
||||
### 5. 클래스 ID 재할당 (reindex)
|
||||
|
||||
클래스 ID를 순차적으로 재할당합니다. 클래스를 삭제하거나 통합한 후 ID를 정리할 때 유용합니다.
|
||||
|
||||
```bash
|
||||
# ID를 1부터 재할당
|
||||
python Utility_lableing_tool.py reindex \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_reindexed.json
|
||||
|
||||
# ID를 0부터 재할당 (COCO는 보통 1부터 시작)
|
||||
python Utility_lableing_tool.py reindex \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_reindexed.json \
|
||||
--start 0
|
||||
```
|
||||
|
||||
## 고급 사용 예시
|
||||
|
||||
### 예시 1: 데이터셋 정리 워크플로우
|
||||
|
||||
```bash
|
||||
# 1. 현재 상태 확인
|
||||
python Utility_lableing_tool.py info --input instances_train.json
|
||||
|
||||
# 2. ID로 불필요한 클래스 제거
|
||||
python Utility_lableing_tool.py remove \
|
||||
--input instances_train.json \
|
||||
--output instances_train_step1.json \
|
||||
--classes "5,7"
|
||||
|
||||
# 3. ID로 유사 클래스 통합
|
||||
python Utility_lableing_tool.py merge \
|
||||
--input instances_train_step1.json \
|
||||
--output instances_train_step2.json \
|
||||
--source "1,2" \
|
||||
--target "vehicle"
|
||||
|
||||
# 4. 클래스 이름 정리
|
||||
python Utility_lableing_tool.py rename \
|
||||
--input instances_train_step2.json \
|
||||
--output instances_train_step3.json \
|
||||
--mapping "tt:transformer,ss:substation"
|
||||
|
||||
# 5. ID 재할당
|
||||
python Utility_lableing_tool.py reindex \
|
||||
--input instances_train_step3.json \
|
||||
--output instances_train_final.json
|
||||
|
||||
# 6. 최종 결과 확인
|
||||
python Utility_lableing_tool.py info --input instances_train_final.json
|
||||
```
|
||||
|
||||
### 예시 2: Train/Valid 데이터셋 동시 처리
|
||||
|
||||
```bash
|
||||
# Train 데이터
|
||||
python Utility_lableing_tool.py rename \
|
||||
--input dataset/annotations/instances_train.json \
|
||||
--output dataset/annotations/instances_train_new.json \
|
||||
--mapping "1:transformer,2:substation"
|
||||
|
||||
# Valid 데이터 (동일한 변경 적용)
|
||||
python Utility_lableing_tool.py rename \
|
||||
--input dataset/annotations/instances_valid.json \
|
||||
--output dataset/annotations/instances_valid_new.json \
|
||||
--mapping "1:transformer,2:substation"
|
||||
```
|
||||
|
||||
## Python 스크립트에서 사용
|
||||
|
||||
유틸리티를 Python 코드에서 직접 사용할 수도 있습니다.
|
||||
|
||||
```python
|
||||
from Utility_lableing_tool import COCOAnnotationEditor
|
||||
|
||||
# 어노테이션 로드
|
||||
editor = COCOAnnotationEditor('dataset/annotations/instances_train.json')
|
||||
|
||||
# 정보 확인
|
||||
editor.print_info()
|
||||
|
||||
# 작업 수행 (이름 또는 ID로 지정 가능)
|
||||
editor.remove_categories(['1', '2']) # ID로 제거
|
||||
editor.remove_categories(['person', 'car']) # 이름으로 제거
|
||||
editor.rename_categories({'1': 'transformer', 'ss': 'substation'}) # ID와 이름 혼합 가능
|
||||
editor.reindex_categories(start_id=1)
|
||||
|
||||
# 저장
|
||||
editor.save('dataset/annotations/instances_train_modified.json')
|
||||
|
||||
# 체이닝도 가능
|
||||
editor = COCOAnnotationEditor('input.json')
|
||||
editor.remove_categories(['1', '2']) \
|
||||
.merge_categories(['3', '4'], 'merged_class') \
|
||||
.reindex_categories() \
|
||||
.save('output.json')
|
||||
```
|
||||
|
||||
## 옵션
|
||||
|
||||
### 백업 관련
|
||||
|
||||
기본적으로 출력 파일이 이미 존재하면 자동으로 백업이 생성됩니다.
|
||||
|
||||
```bash
|
||||
# 백업 생성 (기본값)
|
||||
python Utility_lableing_tool.py rename -i input.json -o output.json -m "old:new"
|
||||
|
||||
# 백업 생성 안 함
|
||||
python Utility_lableing_tool.py rename -i input.json -o output.json -m "old:new" --no-backup
|
||||
```
|
||||
|
||||
백업 파일명 형식: `{original_name}_backup_{YYYYMMDD_HHMMSS}.json`
|
||||
|
||||
## 주요 특징
|
||||
|
||||
### ID와 이름 모두 지원
|
||||
- 모든 작업(제거, 통합, 이름 변경)에서 클래스를 **ID 또는 이름**으로 지정 가능
|
||||
- 숫자로 입력하면 자동으로 ID로 인식, 문자로 입력하면 이름으로 인식
|
||||
- ID와 이름을 혼합하여 사용 가능
|
||||
|
||||
### 어노테이션 통계
|
||||
- 클래스별 어노테이션 개수 자동 계산
|
||||
- 전체 어노테이션 합계 표시
|
||||
- 실시간 정보 확인 가능
|
||||
|
||||
### 입력 오류 처리
|
||||
- 잘못된 파일 경로 입력 시 재입력 요청
|
||||
- 잘못된 선택 입력 시 재입력 요청
|
||||
- 존재하지 않는 클래스 ID/이름 입력 시 경고 표시
|
||||
|
||||
## 주의사항
|
||||
|
||||
1. **데이터 백업**: 중요한 데이터는 항상 백업 후 작업하세요.
|
||||
2. **Train/Valid 일치**: Train과 Valid 데이터셋의 클래스는 동일하게 유지해야 합니다.
|
||||
3. **ID 순서**: `reindex` 명령은 기존 ID 순서를 기준으로 재할당합니다.
|
||||
4. **병합 ID**: `merge` 명령은 가장 작은 ID를 새 카테고리 ID로 사용합니다.
|
||||
5. **원본 보호**: 모든 작업은 새 파일로 저장되며 원본 파일은 수정되지 않습니다.
|
||||
|
||||
## 문제 해결
|
||||
|
||||
### 파일을 찾을 수 없음
|
||||
|
||||
```bash
|
||||
# 절대 경로 사용
|
||||
python Utility_lableing_tool.py info --input /full/path/to/annotations.json
|
||||
|
||||
# 또는 현재 디렉토리에서 상대 경로 사용
|
||||
python Utility_lableing_tool.py info --input ./dataset/annotations/instances_train.json
|
||||
```
|
||||
|
||||
### JSON 형식 오류
|
||||
|
||||
COCO 형식 확인:
|
||||
```python
|
||||
import json
|
||||
with open('annotations.json') as f:
|
||||
data = json.load(f)
|
||||
print(data.keys()) # 'images', 'annotations', 'categories' 포함 확인
|
||||
```
|
||||
|
||||
## 라이선스
|
||||
|
||||
MMDetection 프로젝트의 일부로 동일한 라이선스를 따릅니다.
|
||||
578
mmdet_annotation_modify.py
Normal file
578
mmdet_annotation_modify.py
Normal file
@@ -0,0 +1,578 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
COCO Annotation Utility Tool
|
||||
COCO format annotation file manipulation utility
|
||||
|
||||
Usage examples:
|
||||
# Interactive mode (recommended)
|
||||
python Utility_lableing_tool.py
|
||||
|
||||
# Remove classes
|
||||
python Utility_lableing_tool.py remove --input annotations.json --output new.json --classes "person,car"
|
||||
|
||||
# Merge classes
|
||||
python Utility_lableing_tool.py merge --input annotations.json --output new.json --source "car,truck,bus" --target "vehicle"
|
||||
|
||||
# Rename classes
|
||||
python Utility_lableing_tool.py rename --input annotations.json --output new.json --mapping "old_name:new_name,another_old:another_new"
|
||||
|
||||
# Show info
|
||||
python Utility_lableing_tool.py info --input annotations.json
|
||||
"""
|
||||
|
||||
import json
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Set, Union
|
||||
import shutil
|
||||
from datetime import datetime
|
||||
import sys
|
||||
import os
|
||||
|
||||
|
||||
class COCOAnnotationEditor:
|
||||
"""Class for editing COCO format annotations"""
|
||||
|
||||
def __init__(self, annotation_path: str):
|
||||
"""
|
||||
Args:
|
||||
annotation_path: Path to COCO format annotation JSON file
|
||||
"""
|
||||
self.annotation_path = Path(annotation_path)
|
||||
self.data = self._load_annotation()
|
||||
|
||||
def _load_annotation(self) -> dict:
|
||||
"""Load annotation file"""
|
||||
with open(self.annotation_path, 'r', encoding='utf-8') as f:
|
||||
return json.load(f)
|
||||
|
||||
def save(self, output_path: str, backup: bool = True):
|
||||
"""
|
||||
Save modified annotation
|
||||
|
||||
Args:
|
||||
output_path: Output file path
|
||||
backup: Whether to backup original file
|
||||
"""
|
||||
output_path = Path(output_path)
|
||||
|
||||
# Create backup
|
||||
if backup and output_path.exists():
|
||||
backup_path = output_path.parent / f"{output_path.stem}_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}{output_path.suffix}"
|
||||
shutil.copy2(output_path, backup_path)
|
||||
print(f"[+] 백업 생성됨: {backup_path}")
|
||||
|
||||
# Save
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(self.data, f, ensure_ascii=False, indent=2)
|
||||
print(f"[+] 저장 완료: {output_path}")
|
||||
|
||||
def get_category_info(self) -> Dict:
|
||||
"""Get category information"""
|
||||
categories = self.data.get('categories', [])
|
||||
annotations = self.data.get('annotations', [])
|
||||
|
||||
# Count annotations per category
|
||||
category_counts = {}
|
||||
for ann in annotations:
|
||||
cat_id = ann['category_id']
|
||||
category_counts[cat_id] = category_counts.get(cat_id, 0) + 1
|
||||
|
||||
info = {}
|
||||
for cat in categories:
|
||||
cat_id = cat['id']
|
||||
info[cat_id] = {
|
||||
'id': cat_id,
|
||||
'name': cat['name'],
|
||||
'supercategory': cat.get('supercategory', ''),
|
||||
'annotation_count': category_counts.get(cat_id, 0)
|
||||
}
|
||||
|
||||
return info
|
||||
|
||||
def print_info(self):
|
||||
"""Print annotation information"""
|
||||
print("\n" + "="*60)
|
||||
print(f"파일: {self.annotation_path}")
|
||||
print("="*60)
|
||||
|
||||
print(f"\n전체 통계:")
|
||||
print(f" - 이미지 수: {len(self.data.get('images', []))}")
|
||||
print(f" - 어노테이션 수: {len(self.data.get('annotations', []))}")
|
||||
print(f" - 클래스 수: {len(self.data.get('categories', []))}")
|
||||
|
||||
print(f"\n클래스 상세:")
|
||||
print(f"{'ID':<6} {'클래스명':<30} {'어노테이션 수':<15}")
|
||||
print("-" * 55)
|
||||
|
||||
category_info = self.get_category_info()
|
||||
total_annotations = 0
|
||||
for cat_id, info in sorted(category_info.items()):
|
||||
ann_count = info['annotation_count']
|
||||
total_annotations += ann_count
|
||||
print(f"{info['id']:<6} {info['name']:<30} {ann_count:<15}")
|
||||
|
||||
print("-" * 55)
|
||||
print(f"{'합계':<37} {total_annotations:<15}")
|
||||
print("="*60 + "\n")
|
||||
|
||||
def remove_categories(self, category_names: List[str]) -> 'COCOAnnotationEditor':
|
||||
"""
|
||||
Remove specific categories
|
||||
|
||||
Args:
|
||||
category_names: List of category names or IDs to remove
|
||||
"""
|
||||
# Build ID to name mapping
|
||||
id_to_cat = {cat['id']: cat for cat in self.data.get('categories', [])}
|
||||
name_to_cat = {cat['name']: cat for cat in self.data.get('categories', [])}
|
||||
|
||||
# Find category IDs to remove
|
||||
categories_to_remove = set()
|
||||
|
||||
for item in category_names:
|
||||
# Try to parse as ID first
|
||||
try:
|
||||
cat_id = int(item)
|
||||
if cat_id in id_to_cat:
|
||||
cat = id_to_cat[cat_id]
|
||||
categories_to_remove.add(cat['id'])
|
||||
print(f" - 제거: {cat['name']} (ID: {cat['id']})")
|
||||
else:
|
||||
print(f"[!] ID {cat_id}를 찾을 수 없습니다")
|
||||
except ValueError:
|
||||
# Not a number, treat as name
|
||||
if item in name_to_cat:
|
||||
cat = name_to_cat[item]
|
||||
categories_to_remove.add(cat['id'])
|
||||
print(f" - 제거: {cat['name']} (ID: {cat['id']})")
|
||||
else:
|
||||
print(f"[!] 이름 '{item}'을 찾을 수 없습니다")
|
||||
|
||||
# Update categories
|
||||
remaining_categories = [
|
||||
cat for cat in self.data.get('categories', [])
|
||||
if cat['id'] not in categories_to_remove
|
||||
]
|
||||
|
||||
# Update categories
|
||||
self.data['categories'] = remaining_categories
|
||||
|
||||
# Remove annotations of removed categories
|
||||
original_count = len(self.data.get('annotations', []))
|
||||
self.data['annotations'] = [
|
||||
ann for ann in self.data.get('annotations', [])
|
||||
if ann['category_id'] not in categories_to_remove
|
||||
]
|
||||
removed_count = original_count - len(self.data['annotations'])
|
||||
|
||||
print(f"[+] {removed_count}개의 어노테이션이 제거되었습니다")
|
||||
return self
|
||||
|
||||
|
||||
def merge_categories(self, source_names: List[str], target_name: str,
|
||||
target_supercategory: str = '') -> 'COCOAnnotationEditor':
|
||||
"""
|
||||
Merge multiple categories into one
|
||||
|
||||
Args:
|
||||
source_names: List of source category names or IDs to merge
|
||||
target_name: Target category name after merge
|
||||
target_supercategory: Target supercategory (optional)
|
||||
"""
|
||||
# Find source category IDs
|
||||
source_ids = set()
|
||||
source_categories = []
|
||||
|
||||
# Build ID to name mapping
|
||||
id_to_cat = {cat['id']: cat for cat in self.data.get('categories', [])}
|
||||
name_to_cat = {cat['name']: cat for cat in self.data.get('categories', [])}
|
||||
|
||||
for item in source_names:
|
||||
# Try to parse as ID first
|
||||
try:
|
||||
cat_id = int(item)
|
||||
if cat_id in id_to_cat:
|
||||
cat = id_to_cat[cat_id]
|
||||
source_ids.add(cat['id'])
|
||||
source_categories.append(cat)
|
||||
print(f" - 통합 대상: {cat['name']} (ID: {cat['id']})")
|
||||
else:
|
||||
print(f"[!] ID {cat_id}를 찾을 수 없습니다")
|
||||
except ValueError:
|
||||
# Not a number, treat as name
|
||||
if item in name_to_cat:
|
||||
cat = name_to_cat[item]
|
||||
source_ids.add(cat['id'])
|
||||
source_categories.append(cat)
|
||||
print(f" - 통합 대상: {cat['name']} (ID: {cat['id']})")
|
||||
else:
|
||||
print(f"[!] 이름 '{item}'을 찾을 수 없습니다")
|
||||
|
||||
if not source_ids:
|
||||
print(f"[!] 통합할 클래스를 찾을 수 없습니다")
|
||||
return self
|
||||
|
||||
# New category ID (use first source ID)
|
||||
new_category_id = min(source_ids)
|
||||
|
||||
# Create new category
|
||||
new_category = {
|
||||
'id': new_category_id,
|
||||
'name': target_name,
|
||||
'supercategory': ''
|
||||
}
|
||||
|
||||
# Remove source categories and add new category
|
||||
self.data['categories'] = [
|
||||
cat for cat in self.data.get('categories', [])
|
||||
if cat['id'] not in source_ids
|
||||
]
|
||||
self.data['categories'].append(new_category)
|
||||
|
||||
# Update annotations (change all source IDs to new ID)
|
||||
for ann in self.data.get('annotations', []):
|
||||
if ann['category_id'] in source_ids:
|
||||
ann['category_id'] = new_category_id
|
||||
|
||||
print(f"[+] '{target_name}' (ID: {new_category_id})로 통합되었습니다")
|
||||
return self
|
||||
|
||||
def rename_categories(self, mapping: Dict[str, str]) -> 'COCOAnnotationEditor':
|
||||
"""
|
||||
Rename categories
|
||||
|
||||
Args:
|
||||
mapping: Dictionary of {old_name_or_id: new_name}
|
||||
"""
|
||||
# Build ID to cat mapping
|
||||
id_to_cat = {cat['id']: cat for cat in self.data.get('categories', [])}
|
||||
name_to_cat = {cat['name']: cat for cat in self.data.get('categories', [])}
|
||||
|
||||
renamed_count = 0
|
||||
for old_key, new_name in mapping.items():
|
||||
# Try to parse as ID first
|
||||
try:
|
||||
cat_id = int(old_key)
|
||||
if cat_id in id_to_cat:
|
||||
cat = id_to_cat[cat_id]
|
||||
old_name = cat['name']
|
||||
cat['name'] = new_name
|
||||
print(f" - 변경: '{old_name}' → '{new_name}' (ID: {cat['id']})")
|
||||
renamed_count += 1
|
||||
else:
|
||||
print(f"[!] ID {cat_id}를 찾을 수 없습니다")
|
||||
except ValueError:
|
||||
# Not a number, treat as name
|
||||
if old_key in name_to_cat:
|
||||
cat = name_to_cat[old_key]
|
||||
cat['name'] = new_name
|
||||
print(f" - 변경: '{old_key}' → '{new_name}' (ID: {cat['id']})")
|
||||
renamed_count += 1
|
||||
else:
|
||||
print(f"[!] 이름 '{old_key}'를 찾을 수 없습니다")
|
||||
|
||||
print(f"[+] {renamed_count}개 클래스 이름이 변경되었습니다")
|
||||
return self
|
||||
|
||||
def reindex_categories(self, start_id: int = 1) -> 'COCOAnnotationEditor':
|
||||
"""
|
||||
Reindex category IDs sequentially
|
||||
|
||||
Args:
|
||||
start_id: Starting ID (default: 1)
|
||||
"""
|
||||
# Create mapping: old ID -> new ID
|
||||
id_mapping = {}
|
||||
new_id = start_id
|
||||
|
||||
for cat in sorted(self.data.get('categories', []), key=lambda x: x['id']):
|
||||
old_id = cat['id']
|
||||
id_mapping[old_id] = new_id
|
||||
print(f" - ID 변경: {old_id} → {new_id} ({cat['name']})")
|
||||
new_id += 1
|
||||
|
||||
# Update category IDs
|
||||
for cat in self.data.get('categories', []):
|
||||
cat['id'] = id_mapping[cat['id']]
|
||||
|
||||
# Update annotation category IDs
|
||||
for ann in self.data.get('annotations', []):
|
||||
ann['category_id'] = id_mapping[ann['category_id']]
|
||||
|
||||
print(f"[+] 클래스 ID가 재할당되었습니다 (시작: {start_id})")
|
||||
return self
|
||||
|
||||
|
||||
def interactive_mode():
|
||||
"""Interactive mode for user-friendly editing"""
|
||||
print("\n" + "="*60)
|
||||
print(" COCO 어노테이션 편집기 - 대화형 모드")
|
||||
print("="*60)
|
||||
|
||||
# 1. Select annotation file
|
||||
print("\n[1/3] 어노테이션 파일 선택")
|
||||
print("-" * 60)
|
||||
|
||||
# Find annotation files in common locations
|
||||
common_paths = [
|
||||
"dataset/*/annotations/*.json",
|
||||
"data/*/annotations/*.json",
|
||||
"annotations/*.json",
|
||||
"*.json"
|
||||
]
|
||||
|
||||
from glob import glob
|
||||
found_files = []
|
||||
for pattern in common_paths:
|
||||
found_files.extend(glob(pattern, recursive=True))
|
||||
|
||||
input_file = None
|
||||
while not input_file:
|
||||
if found_files:
|
||||
print("\n발견된 어노테이션 파일:")
|
||||
for idx, f in enumerate(found_files[:10], 1):
|
||||
print(f" {idx}. {f}")
|
||||
if len(found_files) > 10:
|
||||
print(f" ... 외 {len(found_files) - 10}개 더")
|
||||
print(f" 0. 직접 경로 입력")
|
||||
|
||||
choice = input("\n파일 번호 선택 (직접 입력은 0): ").strip()
|
||||
if choice == '0':
|
||||
input_file = input("어노테이션 파일 경로 입력: ").strip()
|
||||
else:
|
||||
try:
|
||||
input_file = found_files[int(choice) - 1]
|
||||
except (ValueError, IndexError):
|
||||
print("[!] 잘못된 선택입니다. 다시 선택해주세요.")
|
||||
continue
|
||||
else:
|
||||
input_file = input("어노테이션 파일 경로 입력: ").strip()
|
||||
|
||||
if not Path(input_file).exists():
|
||||
print(f"[!] 파일을 찾을 수 없습니다: {input_file}")
|
||||
print("[!] 다시 입력해주세요.")
|
||||
input_file = None
|
||||
|
||||
# Load and show info
|
||||
print(f"\n[+] 로딩 중: {input_file}")
|
||||
editor = COCOAnnotationEditor(input_file)
|
||||
editor.print_info()
|
||||
|
||||
# Store original categories for reference
|
||||
original_categories = {cat['id']: cat['name'] for cat in editor.data.get('categories', [])}
|
||||
|
||||
# 2. Select operations
|
||||
print("\n[2/3] 작업 선택")
|
||||
print("-" * 60)
|
||||
operations = []
|
||||
|
||||
while True:
|
||||
print("\n사용 가능한 작업:")
|
||||
print(" 1. 클래스 제거")
|
||||
print(" 2. 클래스 통합 (병합)")
|
||||
print(" 3. 클래스 이름 변경")
|
||||
print(" 4. 클래스 ID 재할당")
|
||||
print(" 5. 현재 정보 표시")
|
||||
print(" 0. 완료 (저장 단계로)")
|
||||
|
||||
choice = input("\n작업 선택 (0-5): ").strip()
|
||||
|
||||
if choice == '0':
|
||||
if not operations:
|
||||
print("[!] 선택된 작업이 없습니다. 종료합니다...")
|
||||
return
|
||||
break
|
||||
|
||||
elif choice == '1': # Remove
|
||||
print("\n현재 클래스:")
|
||||
for cat in editor.data.get('categories', []):
|
||||
print(f" - {cat['name']} (ID: {cat['id']})")
|
||||
|
||||
classes = input("\n제거할 클래스 입력 (이름 또는 ID, 쉼표로 구분): ").strip()
|
||||
if classes:
|
||||
class_list = [c.strip() for c in classes.split(',')]
|
||||
print(f"\n[*] 클래스 제거 중: {class_list}")
|
||||
editor.remove_categories(class_list)
|
||||
operations.append(f"제거: {', '.join(class_list)}")
|
||||
|
||||
elif choice == '2': # Merge
|
||||
print("\n현재 클래스:")
|
||||
for cat in editor.data.get('categories', []):
|
||||
print(f" - {cat['name']} (ID: {cat['id']})")
|
||||
|
||||
source = input("\n통합할 클래스 입력 (이름 또는 ID, 쉼표로 구분): ").strip()
|
||||
if source:
|
||||
target = input("통합 후 클래스 이름 입력: ").strip()
|
||||
|
||||
source_list = [c.strip() for c in source.split(',')]
|
||||
print(f"\n[*] '{target}'로 통합 중: {source_list}")
|
||||
editor.merge_categories(source_list, target, '')
|
||||
operations.append(f"통합: {', '.join(source_list)} → {target}")
|
||||
|
||||
elif choice == '3': # Rename
|
||||
print("\n현재 클래스:")
|
||||
for cat in editor.data.get('categories', []):
|
||||
print(f" - {cat['name']} (ID: {cat['id']})")
|
||||
|
||||
print("\n매핑 입력 (형식: 기존:새이름,기존:새이름 - 쉼표로 구분)")
|
||||
print("예시: tt:transformer,ss:substation 또는 1:transformer,2:substation")
|
||||
mapping_input = input("\n입력: ").strip()
|
||||
|
||||
if mapping_input:
|
||||
mapping = {}
|
||||
for pair in mapping_input.split(','):
|
||||
if ':' in pair:
|
||||
old, new = pair.split(':', 1)
|
||||
mapping[old.strip()] = new.strip()
|
||||
|
||||
if mapping:
|
||||
print(f"\n[*] 클래스 이름 변경 중")
|
||||
editor.rename_categories(mapping)
|
||||
operations.append(f"이름 변경: {len(mapping)}개 클래스")
|
||||
|
||||
elif choice == '4': # Reindex
|
||||
while True:
|
||||
start_id = input("\n시작 ID 입력 (기본값: 1): ").strip()
|
||||
if not start_id:
|
||||
start_id = 1
|
||||
break
|
||||
try:
|
||||
start_id = int(start_id)
|
||||
break
|
||||
except ValueError:
|
||||
print("[!] 숫자를 입력해주세요.")
|
||||
|
||||
print(f"\n[*] 클래스 ID 재할당 중 (시작: {start_id})")
|
||||
editor.reindex_categories(start_id=start_id)
|
||||
operations.append(f"ID 재할당: {start_id}부터")
|
||||
|
||||
elif choice == '5': # Show info
|
||||
editor.print_info()
|
||||
|
||||
else:
|
||||
print("[!] 잘못된 선택입니다")
|
||||
|
||||
# 3. Save
|
||||
print("\n[3/3] 결과 저장")
|
||||
print("-" * 60)
|
||||
print("\n수행된 작업:")
|
||||
for idx, op in enumerate(operations, 1):
|
||||
print(f" {idx}. {op}")
|
||||
|
||||
print(f"\n입력 파일: {input_file}")
|
||||
default_output = str(Path(input_file).parent / f"{Path(input_file).stem}_edited{Path(input_file).suffix}")
|
||||
output_file = input(f"\n출력 파일 경로 입력 (기본값: {default_output}): ").strip()
|
||||
if not output_file:
|
||||
output_file = default_output
|
||||
|
||||
create_backup = input("출력 파일이 존재하면 백업 생성? (Y/n): ").strip().lower()
|
||||
backup = create_backup != 'n'
|
||||
|
||||
print(f"\n[*] 저장 중: {output_file}...")
|
||||
editor.save(output_file, backup=backup)
|
||||
|
||||
print("\n" + "="*60)
|
||||
print(" 작업이 성공적으로 완료되었습니다!")
|
||||
print("="*60)
|
||||
print(f"\n입력: {input_file}")
|
||||
print(f"출력: {output_file}")
|
||||
if backup and Path(output_file).exists():
|
||||
print(f"(파일이 존재했다면 백업이 생성되었습니다)")
|
||||
print()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='COCO Annotation Editing Utility',
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest='command', help='Available commands')
|
||||
|
||||
# info command
|
||||
info_parser = subparsers.add_parser('info', help='Show annotation info')
|
||||
info_parser.add_argument('-i', '--input', required=True, help='Input annotation file')
|
||||
|
||||
# remove command
|
||||
remove_parser = subparsers.add_parser('remove', help='Remove categories')
|
||||
remove_parser.add_argument('-i', '--input', required=True, help='Input annotation file')
|
||||
remove_parser.add_argument('-o', '--output', required=True, help='Output annotation file')
|
||||
remove_parser.add_argument('-c', '--classes', required=True, help='Classes to remove (comma-separated)')
|
||||
remove_parser.add_argument('--no-backup', action='store_true', help='Do not create backup')
|
||||
|
||||
# merge command
|
||||
merge_parser = subparsers.add_parser('merge', help='Merge categories')
|
||||
merge_parser.add_argument('-i', '--input', required=True, help='Input annotation file')
|
||||
merge_parser.add_argument('-o', '--output', required=True, help='Output annotation file')
|
||||
merge_parser.add_argument('-s', '--source', required=True, help='Source classes to merge (comma-separated)')
|
||||
merge_parser.add_argument('-t', '--target', required=True, help='Target class name after merge')
|
||||
merge_parser.add_argument('--supercategory', default='', help='Supercategory for merged class')
|
||||
merge_parser.add_argument('--no-backup', action='store_true', help='Do not create backup')
|
||||
|
||||
# rename command
|
||||
rename_parser = subparsers.add_parser('rename', help='Rename categories')
|
||||
rename_parser.add_argument('-i', '--input', required=True, help='Input annotation file')
|
||||
rename_parser.add_argument('-o', '--output', required=True, help='Output annotation file')
|
||||
rename_parser.add_argument('-m', '--mapping', required=True,
|
||||
help='Name mapping (old:new,old2:new2 format)')
|
||||
rename_parser.add_argument('--no-backup', action='store_true', help='Do not create backup')
|
||||
|
||||
# reindex command
|
||||
reindex_parser = subparsers.add_parser('reindex', help='Reindex category IDs')
|
||||
reindex_parser.add_argument('-i', '--input', required=True, help='Input annotation file')
|
||||
reindex_parser.add_argument('-o', '--output', required=True, help='Output annotation file')
|
||||
reindex_parser.add_argument('--start', type=int, default=1, help='Starting ID (default: 1)')
|
||||
reindex_parser.add_argument('--no-backup', action='store_true', help='Do not create backup')
|
||||
|
||||
# If no arguments provided, enter interactive mode
|
||||
if len(sys.argv) == 1:
|
||||
interactive_mode()
|
||||
return
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.command:
|
||||
parser.print_help()
|
||||
return
|
||||
|
||||
# Execute command
|
||||
if args.command == 'info':
|
||||
editor = COCOAnnotationEditor(args.input)
|
||||
editor.print_info()
|
||||
|
||||
elif args.command == 'remove':
|
||||
classes = [c.strip() for c in args.classes.split(',')]
|
||||
print(f"\n[*] 클래스 제거 중...")
|
||||
editor = COCOAnnotationEditor(args.input)
|
||||
editor.remove_categories(classes)
|
||||
editor.save(args.output, backup=not args.no_backup)
|
||||
|
||||
elif args.command == 'merge':
|
||||
source_classes = [c.strip() for c in args.source.split(',')]
|
||||
print(f"\n[*] 클래스 통합 중...")
|
||||
editor = COCOAnnotationEditor(args.input)
|
||||
editor.merge_categories(source_classes, args.target, '')
|
||||
editor.save(args.output, backup=not args.no_backup)
|
||||
|
||||
elif args.command == 'rename':
|
||||
mapping = {}
|
||||
for pair in args.mapping.split(','):
|
||||
old, new = pair.split(':')
|
||||
mapping[old.strip()] = new.strip()
|
||||
print(f"\n[*] 클래스 이름 변경 중...")
|
||||
editor = COCOAnnotationEditor(args.input)
|
||||
editor.rename_categories(mapping)
|
||||
editor.save(args.output, backup=not args.no_backup)
|
||||
|
||||
elif args.command == 'reindex':
|
||||
print(f"\n[*] 클래스 ID 재할당 중...")
|
||||
editor = COCOAnnotationEditor(args.input)
|
||||
editor.reindex_categories(start_id=args.start)
|
||||
editor.save(args.output, backup=not args.no_backup)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
646
yolo_annotation_modify.py
Normal file
646
yolo_annotation_modify.py
Normal file
@@ -0,0 +1,646 @@
|
||||
#!/usr/bin/env python3
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
YOLO Annotation Utility Tool
|
||||
YOLO format annotation file manipulation utility
|
||||
|
||||
Usage examples:
|
||||
# Interactive mode (recommended)
|
||||
python yolo_annotation_modify.py
|
||||
|
||||
# Remove classes
|
||||
python yolo_annotation_modify.py remove --dataset dataset/yolo --classes "0,1"
|
||||
|
||||
# Merge classes
|
||||
python yolo_annotation_modify.py merge --dataset dataset/yolo --source "0,1,2" --target "vehicle"
|
||||
|
||||
# Rename classes
|
||||
python yolo_annotation_modify.py rename --dataset dataset/yolo --mapping "0:vehicle,1:person"
|
||||
|
||||
# Show info
|
||||
python yolo_annotation_modify.py info --dataset dataset/yolo
|
||||
"""
|
||||
|
||||
import yaml
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Set, Union
|
||||
import shutil
|
||||
from datetime import datetime
|
||||
import sys
|
||||
import os
|
||||
from collections import defaultdict
|
||||
|
||||
|
||||
class YOLOAnnotationEditor:
|
||||
"""Class for editing YOLO format annotations"""
|
||||
|
||||
def __init__(self, dataset_path: str):
|
||||
"""
|
||||
Args:
|
||||
dataset_path: Path to YOLO dataset directory (containing data.yaml)
|
||||
"""
|
||||
self.dataset_path = Path(dataset_path)
|
||||
self.yaml_path = self.dataset_path / "data.yaml"
|
||||
|
||||
if not self.yaml_path.exists():
|
||||
raise FileNotFoundError(f"data.yaml not found in {dataset_path}")
|
||||
|
||||
self.config = self._load_yaml()
|
||||
self.labels_path = self.dataset_path / "labels"
|
||||
|
||||
# Find all label files
|
||||
self.label_files = list(self.labels_path.rglob("*.txt"))
|
||||
print(f"[+] Found {len(self.label_files)} label files")
|
||||
|
||||
def _load_yaml(self) -> dict:
|
||||
"""Load YOLO data.yaml file"""
|
||||
with open(self.yaml_path, 'r', encoding='utf-8') as f:
|
||||
return yaml.safe_load(f)
|
||||
|
||||
def backup_labels(self) -> Path:
|
||||
"""
|
||||
Backup entire labels folder
|
||||
|
||||
Returns:
|
||||
Path to backup folder
|
||||
"""
|
||||
timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
|
||||
backup_path = self.dataset_path / f"labels_backup_{timestamp}"
|
||||
|
||||
print(f"[*] labels 폴더 백업 중...")
|
||||
shutil.copytree(self.labels_path, backup_path)
|
||||
print(f"[+] 백업 생성됨: {backup_path}")
|
||||
|
||||
return backup_path
|
||||
|
||||
def save_yaml(self, backup: bool = True):
|
||||
"""
|
||||
Save modified data.yaml
|
||||
|
||||
Args:
|
||||
backup: Whether to backup original file
|
||||
"""
|
||||
if backup and self.yaml_path.exists():
|
||||
backup_path = self.yaml_path.parent / f"data_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.yaml"
|
||||
shutil.copy2(self.yaml_path, backup_path)
|
||||
print(f"[+] data.yaml 백업 생성됨: {backup_path}")
|
||||
|
||||
with open(self.yaml_path, 'w', encoding='utf-8') as f:
|
||||
yaml.dump(self.config, f, allow_unicode=True, sort_keys=False)
|
||||
print(f"[+] data.yaml 저장 완료")
|
||||
|
||||
def get_class_info(self) -> Dict:
|
||||
"""Get class information and count annotations"""
|
||||
names = self.config.get('names', {})
|
||||
|
||||
# Count annotations per class
|
||||
class_counts = defaultdict(int)
|
||||
total_annotations = 0
|
||||
|
||||
for label_file in self.label_files:
|
||||
try:
|
||||
with open(label_file, 'r') as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line:
|
||||
class_id = int(line.split()[0])
|
||||
class_counts[class_id] += 1
|
||||
total_annotations += 1
|
||||
except Exception as e:
|
||||
print(f"[!] Error reading {label_file}: {e}")
|
||||
|
||||
info = {}
|
||||
for class_id, class_name in names.items():
|
||||
info[class_id] = {
|
||||
'id': class_id,
|
||||
'name': class_name,
|
||||
'annotation_count': class_counts.get(class_id, 0)
|
||||
}
|
||||
|
||||
return info
|
||||
|
||||
def print_info(self):
|
||||
"""Print annotation information"""
|
||||
print("\n" + "="*60)
|
||||
print(f"Dataset: {self.dataset_path}")
|
||||
print("="*60)
|
||||
|
||||
names = self.config.get('names', {})
|
||||
print(f"\n전체 통계:")
|
||||
print(f" - 라벨 파일 수: {len(self.label_files)}")
|
||||
print(f" - 클래스 수: {len(names)}")
|
||||
|
||||
print(f"\n클래스 상세:")
|
||||
print(f"{'ID':<6} {'클래스명':<35} {'어노테이션 수':<15}")
|
||||
print("-" * 60)
|
||||
|
||||
class_info = self.get_class_info()
|
||||
total_annotations = 0
|
||||
for class_id in sorted(class_info.keys()):
|
||||
info = class_info[class_id]
|
||||
ann_count = info['annotation_count']
|
||||
total_annotations += ann_count
|
||||
print(f"{info['id']:<6} {info['name']:<35} {ann_count:<15}")
|
||||
|
||||
print("-" * 60)
|
||||
print(f"{'합계':<42} {total_annotations:<15}")
|
||||
print("="*60 + "\n")
|
||||
|
||||
def remove_classes(self, class_ids: List[Union[int, str]]) -> 'YOLOAnnotationEditor':
|
||||
"""
|
||||
Remove specific classes
|
||||
|
||||
Args:
|
||||
class_ids: List of class IDs or names to remove
|
||||
"""
|
||||
names = self.config.get('names', {})
|
||||
|
||||
# Convert names to IDs
|
||||
name_to_id = {name: cid for cid, name in names.items()}
|
||||
ids_to_remove = set()
|
||||
|
||||
for item in class_ids:
|
||||
try:
|
||||
# Try as ID first
|
||||
class_id = int(item)
|
||||
if class_id in names:
|
||||
ids_to_remove.add(class_id)
|
||||
print(f" - 제거: {names[class_id]} (ID: {class_id})")
|
||||
else:
|
||||
print(f"[!] ID {class_id}를 찾을 수 없습니다")
|
||||
except ValueError:
|
||||
# Try as name
|
||||
if item in name_to_id:
|
||||
class_id = name_to_id[item]
|
||||
ids_to_remove.add(class_id)
|
||||
print(f" - 제거: {item} (ID: {class_id})")
|
||||
else:
|
||||
print(f"[!] 이름 '{item}'을 찾을 수 없습니다")
|
||||
|
||||
if not ids_to_remove:
|
||||
print("[!] 제거할 클래스를 찾을 수 없습니다")
|
||||
return self
|
||||
|
||||
# Remove from data.yaml
|
||||
new_names = {cid: name for cid, name in names.items() if cid not in ids_to_remove}
|
||||
self.config['names'] = new_names
|
||||
|
||||
# Remove annotations from label files
|
||||
total_removed = 0
|
||||
for label_file in self.label_files:
|
||||
try:
|
||||
with open(label_file, 'r') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
new_lines = []
|
||||
for line in lines:
|
||||
if line.strip():
|
||||
class_id = int(line.split()[0])
|
||||
if class_id not in ids_to_remove:
|
||||
new_lines.append(line)
|
||||
else:
|
||||
total_removed += 1
|
||||
|
||||
with open(label_file, 'w') as f:
|
||||
f.writelines(new_lines)
|
||||
except Exception as e:
|
||||
print(f"[!] Error processing {label_file}: {e}")
|
||||
|
||||
print(f"[+] {total_removed}개의 어노테이션이 제거되었습니다")
|
||||
return self
|
||||
|
||||
def merge_classes(self, source_ids: List[Union[int, str]], target_name: str) -> 'YOLOAnnotationEditor':
|
||||
"""
|
||||
Merge multiple classes into one
|
||||
|
||||
Args:
|
||||
source_ids: List of source class IDs or names to merge
|
||||
target_name: Target class name after merge
|
||||
"""
|
||||
names = self.config.get('names', {})
|
||||
name_to_id = {name: cid for cid, name in names.items()}
|
||||
|
||||
# Find source class IDs
|
||||
source_class_ids = set()
|
||||
for item in source_ids:
|
||||
try:
|
||||
class_id = int(item)
|
||||
if class_id in names:
|
||||
source_class_ids.add(class_id)
|
||||
print(f" - 통합 대상: {names[class_id]} (ID: {class_id})")
|
||||
else:
|
||||
print(f"[!] ID {class_id}를 찾을 수 없습니다")
|
||||
except ValueError:
|
||||
if item in name_to_id:
|
||||
class_id = name_to_id[item]
|
||||
source_class_ids.add(class_id)
|
||||
print(f" - 통합 대상: {item} (ID: {class_id})")
|
||||
else:
|
||||
print(f"[!] 이름 '{item}'을 찾을 수 없습니다")
|
||||
|
||||
if not source_class_ids:
|
||||
print("[!] 통합할 클래스를 찾을 수 없습니다")
|
||||
return self
|
||||
|
||||
# Use the smallest ID as the new class ID
|
||||
new_class_id = min(source_class_ids)
|
||||
|
||||
# Update data.yaml
|
||||
new_names = {cid: name for cid, name in names.items() if cid not in source_class_ids}
|
||||
new_names[new_class_id] = target_name
|
||||
self.config['names'] = new_names
|
||||
|
||||
# Update label files
|
||||
total_merged = 0
|
||||
for label_file in self.label_files:
|
||||
try:
|
||||
with open(label_file, 'r') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
new_lines = []
|
||||
for line in lines:
|
||||
if line.strip():
|
||||
parts = line.split()
|
||||
class_id = int(parts[0])
|
||||
if class_id in source_class_ids:
|
||||
# Replace with new class ID
|
||||
parts[0] = str(new_class_id)
|
||||
new_lines.append(' '.join(parts) + '\n')
|
||||
total_merged += 1
|
||||
else:
|
||||
new_lines.append(line)
|
||||
|
||||
with open(label_file, 'w') as f:
|
||||
f.writelines(new_lines)
|
||||
except Exception as e:
|
||||
print(f"[!] Error processing {label_file}: {e}")
|
||||
|
||||
print(f"[+] '{target_name}' (ID: {new_class_id})로 통합되었습니다")
|
||||
print(f"[+] {total_merged}개의 어노테이션이 변경되었습니다")
|
||||
return self
|
||||
|
||||
def rename_classes(self, mapping: Dict[Union[int, str], str]) -> 'YOLOAnnotationEditor':
|
||||
"""
|
||||
Rename classes
|
||||
|
||||
Args:
|
||||
mapping: Dictionary of {old_id_or_name: new_name}
|
||||
"""
|
||||
names = self.config.get('names', {})
|
||||
name_to_id = {name: cid for cid, name in names.items()}
|
||||
|
||||
renamed_count = 0
|
||||
for old_key, new_name in mapping.items():
|
||||
try:
|
||||
# Try as ID first
|
||||
class_id = int(old_key)
|
||||
if class_id in names:
|
||||
old_name = names[class_id]
|
||||
names[class_id] = new_name
|
||||
print(f" - 변경: '{old_name}' → '{new_name}' (ID: {class_id})")
|
||||
renamed_count += 1
|
||||
else:
|
||||
print(f"[!] ID {class_id}를 찾을 수 없습니다")
|
||||
except ValueError:
|
||||
# Try as name
|
||||
if old_key in name_to_id:
|
||||
class_id = name_to_id[old_key]
|
||||
names[class_id] = new_name
|
||||
print(f" - 변경: '{old_key}' → '{new_name}' (ID: {class_id})")
|
||||
renamed_count += 1
|
||||
else:
|
||||
print(f"[!] 이름 '{old_key}'를 찾을 수 없습니다")
|
||||
|
||||
self.config['names'] = names
|
||||
print(f"[+] {renamed_count}개 클래스 이름이 변경되었습니다")
|
||||
return self
|
||||
|
||||
def reindex_classes(self, start_id: int = 0) -> 'YOLOAnnotationEditor':
|
||||
"""
|
||||
Reindex class IDs sequentially
|
||||
|
||||
Args:
|
||||
start_id: Starting ID (default: 0 for YOLO)
|
||||
"""
|
||||
names = self.config.get('names', {})
|
||||
|
||||
# Create ID mapping: old ID -> new ID
|
||||
id_mapping = {}
|
||||
new_id = start_id
|
||||
|
||||
for old_id in sorted(names.keys()):
|
||||
id_mapping[old_id] = new_id
|
||||
print(f" - ID 변경: {old_id} → {new_id} ({names[old_id]})")
|
||||
new_id += 1
|
||||
|
||||
# Update data.yaml
|
||||
new_names = {id_mapping[old_id]: name for old_id, name in names.items()}
|
||||
self.config['names'] = new_names
|
||||
|
||||
# Update label files
|
||||
total_updated = 0
|
||||
for label_file in self.label_files:
|
||||
try:
|
||||
with open(label_file, 'r') as f:
|
||||
lines = f.readlines()
|
||||
|
||||
new_lines = []
|
||||
for line in lines:
|
||||
if line.strip():
|
||||
parts = line.split()
|
||||
old_class_id = int(parts[0])
|
||||
if old_class_id in id_mapping:
|
||||
parts[0] = str(id_mapping[old_class_id])
|
||||
new_lines.append(' '.join(parts) + '\n')
|
||||
total_updated += 1
|
||||
else:
|
||||
new_lines.append(line)
|
||||
|
||||
with open(label_file, 'w') as f:
|
||||
f.writelines(new_lines)
|
||||
except Exception as e:
|
||||
print(f"[!] Error processing {label_file}: {e}")
|
||||
|
||||
print(f"[+] 클래스 ID가 재할당되었습니다 (시작: {start_id})")
|
||||
print(f"[+] {total_updated}개의 어노테이션이 업데이트되었습니다")
|
||||
return self
|
||||
|
||||
|
||||
def interactive_mode():
|
||||
"""Interactive mode for user-friendly editing"""
|
||||
print("\n" + "="*60)
|
||||
print(" YOLO 어노테이션 편집기 - 대화형 모드")
|
||||
print("="*60)
|
||||
|
||||
# 1. Select dataset
|
||||
print("\n[1/3] 데이터셋 선택")
|
||||
print("-" * 60)
|
||||
|
||||
# Find YOLO datasets
|
||||
from glob import glob
|
||||
yaml_files = glob("**/data.yaml", recursive=True)
|
||||
|
||||
dataset_path = None
|
||||
while not dataset_path:
|
||||
if yaml_files:
|
||||
print("\n발견된 YOLO 데이터셋:")
|
||||
for idx, f in enumerate(yaml_files, 1):
|
||||
dataset_dir = str(Path(f).parent)
|
||||
print(f" {idx}. {dataset_dir}")
|
||||
print(f" 0. 직접 경로 입력")
|
||||
|
||||
choice = input("\n데이터셋 번호 선택 (직접 입력은 0): ").strip()
|
||||
if choice == '0':
|
||||
dataset_path = input("데이터셋 디렉토리 경로 입력: ").strip()
|
||||
else:
|
||||
try:
|
||||
dataset_path = str(Path(yaml_files[int(choice) - 1]).parent)
|
||||
except (ValueError, IndexError):
|
||||
print("[!] 잘못된 선택입니다. 다시 선택해주세요.")
|
||||
continue
|
||||
else:
|
||||
dataset_path = input("데이터셋 디렉토리 경로 입력: ").strip()
|
||||
|
||||
yaml_path = Path(dataset_path) / "data.yaml"
|
||||
if not yaml_path.exists():
|
||||
print(f"[!] data.yaml을 찾을 수 없습니다: {yaml_path}")
|
||||
print("[!] 다시 입력해주세요.")
|
||||
dataset_path = None
|
||||
|
||||
# Load and show info
|
||||
print(f"\n[+] 로딩 중: {dataset_path}")
|
||||
editor = YOLOAnnotationEditor(dataset_path)
|
||||
editor.print_info()
|
||||
|
||||
# 2. Select operations
|
||||
print("\n[2/3] 작업 선택")
|
||||
print("-" * 60)
|
||||
operations = []
|
||||
|
||||
while True:
|
||||
print("\n사용 가능한 작업:")
|
||||
print(" 1. 클래스 제거")
|
||||
print(" 2. 클래스 통합 (병합)")
|
||||
print(" 3. 클래스 이름 변경")
|
||||
print(" 4. 클래스 ID 재할당")
|
||||
print(" 5. 현재 정보 표시")
|
||||
print(" 0. 완료 (저장 단계로)")
|
||||
|
||||
choice = input("\n작업 선택 (0-5): ").strip()
|
||||
|
||||
if choice == '0':
|
||||
if not operations:
|
||||
print("[!] 선택된 작업이 없습니다. 종료합니다...")
|
||||
return
|
||||
break
|
||||
|
||||
elif choice == '1': # Remove
|
||||
print("\n현재 클래스:")
|
||||
for cid, name in sorted(editor.config['names'].items()):
|
||||
print(f" {cid}: {name}")
|
||||
|
||||
classes = input("\n제거할 클래스 입력 (ID 또는 이름, 쉼표로 구분): ").strip()
|
||||
if classes:
|
||||
class_list = [c.strip() for c in classes.split(',')]
|
||||
print(f"\n[*] 클래스 제거 중: {class_list}")
|
||||
editor.remove_classes(class_list)
|
||||
operations.append(f"제거: {', '.join(class_list)}")
|
||||
|
||||
elif choice == '2': # Merge
|
||||
print("\n현재 클래스:")
|
||||
for cid, name in sorted(editor.config['names'].items()):
|
||||
print(f" {cid}: {name}")
|
||||
|
||||
source = input("\n통합할 클래스 입력 (ID 또는 이름, 쉼표로 구분): ").strip()
|
||||
if source:
|
||||
target = input("통합 후 클래스 이름 입력: ").strip()
|
||||
|
||||
source_list = [c.strip() for c in source.split(',')]
|
||||
print(f"\n[*] '{target}'로 통합 중: {source_list}")
|
||||
editor.merge_classes(source_list, target)
|
||||
operations.append(f"통합: {', '.join(source_list)} → {target}")
|
||||
|
||||
elif choice == '3': # Rename
|
||||
print("\n현재 클래스:")
|
||||
for cid, name in sorted(editor.config['names'].items()):
|
||||
print(f" {cid}: {name}")
|
||||
|
||||
print("\n매핑 입력 (형식: 기존:새이름,기존:새이름 - 쉼표로 구분)")
|
||||
print("예시: 0:vehicle,1:person 또는 truck:vehicle,person:human")
|
||||
mapping_input = input("\n입력: ").strip()
|
||||
|
||||
if mapping_input:
|
||||
mapping = {}
|
||||
for pair in mapping_input.split(','):
|
||||
if ':' in pair:
|
||||
old, new = pair.split(':', 1)
|
||||
mapping[old.strip()] = new.strip()
|
||||
|
||||
if mapping:
|
||||
print(f"\n[*] 클래스 이름 변경 중")
|
||||
editor.rename_classes(mapping)
|
||||
operations.append(f"이름 변경: {len(mapping)}개 클래스")
|
||||
|
||||
elif choice == '4': # Reindex
|
||||
while True:
|
||||
start_id = input("\n시작 ID 입력 (기본값: 0): ").strip()
|
||||
if not start_id:
|
||||
start_id = 0
|
||||
break
|
||||
try:
|
||||
start_id = int(start_id)
|
||||
break
|
||||
except ValueError:
|
||||
print("[!] 숫자를 입력해주세요.")
|
||||
|
||||
print(f"\n[*] 클래스 ID 재할당 중 (시작: {start_id})")
|
||||
editor.reindex_classes(start_id=start_id)
|
||||
operations.append(f"ID 재할당: {start_id}부터")
|
||||
|
||||
elif choice == '5': # Show info
|
||||
editor.print_info()
|
||||
|
||||
else:
|
||||
print("[!] 잘못된 선택입니다")
|
||||
|
||||
# 3. Save
|
||||
print("\n[3/3] 결과 저장")
|
||||
print("-" * 60)
|
||||
print("\n수행된 작업:")
|
||||
for idx, op in enumerate(operations, 1):
|
||||
print(f" {idx}. {op}")
|
||||
|
||||
# Ask for labels backup
|
||||
print("\n[!] 주의: labels 폴더의 모든 .txt 파일이 수정됩니다.")
|
||||
create_labels_backup = input("labels 폴더 백업 생성? (Y/n): ").strip().lower()
|
||||
if create_labels_backup != 'n':
|
||||
editor.backup_labels()
|
||||
|
||||
create_yaml_backup = input("data.yaml 백업 생성? (Y/n): ").strip().lower()
|
||||
yaml_backup = create_yaml_backup != 'n'
|
||||
|
||||
print(f"\n[*] 저장 중...")
|
||||
editor.save_yaml(backup=yaml_backup)
|
||||
|
||||
print("\n" + "="*60)
|
||||
print(" 작업이 성공적으로 완료되었습니다!")
|
||||
print("="*60)
|
||||
print(f"\n데이터셋: {dataset_path}")
|
||||
print(f"설정 파일: {editor.yaml_path}")
|
||||
print(f"라벨 파일: {len(editor.label_files)}개 업데이트됨")
|
||||
print()
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description='YOLO Annotation Editing Utility',
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__
|
||||
)
|
||||
|
||||
subparsers = parser.add_subparsers(dest='command', help='Available commands')
|
||||
|
||||
# info command
|
||||
info_parser = subparsers.add_parser('info', help='Show annotation info')
|
||||
info_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')
|
||||
|
||||
# remove command
|
||||
remove_parser = subparsers.add_parser('remove', help='Remove classes')
|
||||
remove_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')
|
||||
remove_parser.add_argument('-c', '--classes', required=True, help='Classes to remove (comma-separated IDs or names)')
|
||||
remove_parser.add_argument('--no-backup', action='store_true', help='Do not create backup for labels and yaml')
|
||||
remove_parser.add_argument('--no-labels-backup', action='store_true', help='Do not create backup for labels folder')
|
||||
|
||||
# merge command
|
||||
merge_parser = subparsers.add_parser('merge', help='Merge classes')
|
||||
merge_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')
|
||||
merge_parser.add_argument('-s', '--source', required=True, help='Source classes to merge (comma-separated)')
|
||||
merge_parser.add_argument('-t', '--target', required=True, help='Target class name after merge')
|
||||
merge_parser.add_argument('--no-backup', action='store_true', help='Do not create backup for labels and yaml')
|
||||
merge_parser.add_argument('--no-labels-backup', action='store_true', help='Do not create backup for labels folder')
|
||||
|
||||
# rename command
|
||||
rename_parser = subparsers.add_parser('rename', help='Rename classes')
|
||||
rename_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')
|
||||
rename_parser.add_argument('-m', '--mapping', required=True,
|
||||
help='Name mapping (old:new,old2:new2 format)')
|
||||
rename_parser.add_argument('--no-backup', action='store_true', help='Do not create backup for labels and yaml')
|
||||
rename_parser.add_argument('--no-labels-backup', action='store_true', help='Do not create backup for labels folder')
|
||||
|
||||
# reindex command
|
||||
reindex_parser = subparsers.add_parser('reindex', help='Reindex class IDs')
|
||||
reindex_parser.add_argument('-d', '--dataset', required=True, help='Dataset directory')
|
||||
reindex_parser.add_argument('--start', type=int, default=0, help='Starting ID (default: 0)')
|
||||
reindex_parser.add_argument('--no-backup', action='store_true', help='Do not create backup for labels and yaml')
|
||||
reindex_parser.add_argument('--no-labels-backup', action='store_true', help='Do not create backup for labels folder')
|
||||
|
||||
# If no arguments provided, enter interactive mode
|
||||
if len(sys.argv) == 1:
|
||||
interactive_mode()
|
||||
return
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if not args.command:
|
||||
parser.print_help()
|
||||
return
|
||||
|
||||
# Execute command
|
||||
if args.command == 'info':
|
||||
editor = YOLOAnnotationEditor(args.dataset)
|
||||
editor.print_info()
|
||||
|
||||
elif args.command == 'remove':
|
||||
classes = [c.strip() for c in args.classes.split(',')]
|
||||
editor = YOLOAnnotationEditor(args.dataset)
|
||||
|
||||
# Backup labels if requested
|
||||
if not args.no_backup and not args.no_labels_backup:
|
||||
editor.backup_labels()
|
||||
|
||||
print(f"\n[*] 클래스 제거 중...")
|
||||
editor.remove_classes(classes)
|
||||
editor.save_yaml(backup=not args.no_backup)
|
||||
|
||||
elif args.command == 'merge':
|
||||
source_classes = [c.strip() for c in args.source.split(',')]
|
||||
editor = YOLOAnnotationEditor(args.dataset)
|
||||
|
||||
# Backup labels if requested
|
||||
if not args.no_backup and not args.no_labels_backup:
|
||||
editor.backup_labels()
|
||||
|
||||
print(f"\n[*] 클래스 통합 중...")
|
||||
editor.merge_classes(source_classes, args.target)
|
||||
editor.save_yaml(backup=not args.no_backup)
|
||||
|
||||
elif args.command == 'rename':
|
||||
mapping = {}
|
||||
for pair in args.mapping.split(','):
|
||||
old, new = pair.split(':')
|
||||
mapping[old.strip()] = new.strip()
|
||||
editor = YOLOAnnotationEditor(args.dataset)
|
||||
|
||||
# Backup labels if requested
|
||||
if not args.no_backup and not args.no_labels_backup:
|
||||
editor.backup_labels()
|
||||
|
||||
print(f"\n[*] 클래스 이름 변경 중...")
|
||||
editor.rename_classes(mapping)
|
||||
editor.save_yaml(backup=not args.no_backup)
|
||||
|
||||
elif args.command == 'reindex':
|
||||
editor = YOLOAnnotationEditor(args.dataset)
|
||||
|
||||
# Backup labels if requested
|
||||
if not args.no_backup and not args.no_labels_backup:
|
||||
editor.backup_labels()
|
||||
|
||||
print(f"\n[*] 클래스 ID 재할당 중...")
|
||||
editor.reindex_classes(start_id=args.start)
|
||||
editor.save_yaml(backup=not args.no_backup)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user