commit bc14484576ed5db36465c2b501e9299bd7902593 Author: rudals252 Date: Mon Nov 17 18:00:06 2025 +0900 yolo, mmedet 라벨링 수정도구 diff --git a/.claude/settings.local.json b/.claude/settings.local.json new file mode 100644 index 0000000..b2e2fe8 --- /dev/null +++ b/.claude/settings.local.json @@ -0,0 +1,9 @@ +{ + "permissions": { + "allow": [ + "Bash(python yolo_annotation_modify.py:*)" + ], + "deny": [], + "ask": [] + } +} diff --git a/README.md b/README.md new file mode 100644 index 0000000..2717beb --- /dev/null +++ b/README.md @@ -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 프로젝트의 일부로 동일한 라이선스를 따릅니다. diff --git a/mmdet_annotation_modify.py b/mmdet_annotation_modify.py new file mode 100644 index 0000000..21d5e5a --- /dev/null +++ b/mmdet_annotation_modify.py @@ -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() diff --git a/yolo_annotation_modify.py b/yolo_annotation_modify.py new file mode 100644 index 0000000..3d2b689 --- /dev/null +++ b/yolo_annotation_modify.py @@ -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()