763 lines
30 KiB
Python
763 lines
30 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
"""
|
|
COCO Annotation Utility Tool
|
|
COCO format annotation file manipulation utility
|
|
|
|
Usage examples:
|
|
# Interactive mode (recommended)
|
|
python mmdet_annotation_modify.py
|
|
|
|
# Remove classes
|
|
python mmdet_annotation_modify.py remove --input annotations.json --output new.json --classes "person,car"
|
|
|
|
# Merge classes
|
|
python mmdet_annotation_modify.py merge --input annotations.json --output new.json --source "car,truck,bus" --target "vehicle"
|
|
|
|
# Rename classes
|
|
python mmdet_annotation_modify.py rename --input annotations.json --output new.json --mapping "old_name:new_name,another_old:another_new"
|
|
|
|
# Show info
|
|
python mmdet_annotation_modify.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, dry_run: bool = False):
|
|
"""
|
|
Args:
|
|
annotation_path: Path to COCO format annotation JSON file
|
|
dry_run: If True, only simulate changes without writing files
|
|
"""
|
|
self.annotation_path = Path(annotation_path)
|
|
self.dry_run = dry_run
|
|
|
|
if self.dry_run:
|
|
print("[DRY RUN] 시뮬레이션 모드 - 파일은 수정되지 않습니다\n")
|
|
|
|
# Validate file exists
|
|
if not self.annotation_path.exists():
|
|
raise FileNotFoundError(f"Annotation file not found: {annotation_path}")
|
|
|
|
if not self.annotation_path.is_file():
|
|
raise ValueError(f"Path is not a file: {annotation_path}")
|
|
|
|
self.data = self._load_annotation()
|
|
|
|
def _load_annotation(self) -> dict:
|
|
"""Load annotation file"""
|
|
try:
|
|
with open(self.annotation_path, 'r', encoding='utf-8') as f:
|
|
data = json.load(f)
|
|
|
|
# Validate COCO structure
|
|
if not isinstance(data, dict):
|
|
raise ValueError(f"Invalid JSON format: Expected dictionary, got {type(data)}")
|
|
|
|
required_fields = ['images', 'annotations', 'categories']
|
|
for field in required_fields:
|
|
if field not in data:
|
|
raise ValueError(f"Missing required field: '{field}'")
|
|
|
|
if not isinstance(data['categories'], list):
|
|
raise ValueError("'categories' field must be a list")
|
|
|
|
if not isinstance(data['annotations'], list):
|
|
raise ValueError("'annotations' field must be a list")
|
|
|
|
if not isinstance(data['images'], list):
|
|
raise ValueError("'images' field must be a list")
|
|
|
|
return data
|
|
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(f"Failed to parse JSON file: {e}")
|
|
except UnicodeDecodeError:
|
|
raise ValueError(f"Cannot decode file (invalid encoding)")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Error loading annotation file: {e}")
|
|
|
|
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)
|
|
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] 파일이 저장될 예정입니다: {output_path}")
|
|
if backup and output_path.exists():
|
|
print(f"[DRY RUN] 백업이 생성될 예정입니다")
|
|
return
|
|
|
|
try:
|
|
# 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}")
|
|
|
|
# Ensure parent directory exists
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# 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}")
|
|
|
|
except PermissionError:
|
|
raise PermissionError(f"Permission denied writing to {output_path}")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to save annotation file: {e}")
|
|
|
|
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 idx, ann in enumerate(annotations):
|
|
if 'category_id' not in ann:
|
|
print(f"[!] Warning: Annotation {idx} missing 'category_id' field")
|
|
continue
|
|
try:
|
|
cat_id = ann['category_id']
|
|
category_counts[cat_id] = category_counts.get(cat_id, 0) + 1
|
|
except Exception as e:
|
|
print(f"[!] Warning: Error processing annotation {idx}: {e}")
|
|
|
|
info = {}
|
|
for cat in categories:
|
|
if 'id' not in cat or 'name' not in cat:
|
|
print(f"[!] Warning: Category missing required fields: {cat}")
|
|
continue
|
|
try:
|
|
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)
|
|
}
|
|
except Exception as e:
|
|
print(f"[!] Warning: Error processing category: {e}")
|
|
|
|
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', []))
|
|
print(f"[*] {original_count}개 어노테이션 검사 중...")
|
|
|
|
new_annotations = []
|
|
for idx, ann in enumerate(self.data.get('annotations', []), 1):
|
|
if idx % 1000 == 0:
|
|
print(f"[*] 진행: {idx}/{original_count} ({idx*100//original_count}%)")
|
|
if ann['category_id'] not in categories_to_remove:
|
|
new_annotations.append(ann)
|
|
|
|
self.data['annotations'] = new_annotations
|
|
removed_count = original_count - len(new_annotations)
|
|
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] {removed_count}개의 어노테이션이 제거될 예정입니다")
|
|
else:
|
|
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)
|
|
total_anns = len(self.data.get('annotations', []))
|
|
print(f"[*] {total_anns}개 어노테이션 업데이트 중...")
|
|
|
|
for idx, ann in enumerate(self.data.get('annotations', []), 1):
|
|
if idx % 1000 == 0:
|
|
print(f"[*] 진행: {idx}/{total_anns} ({idx*100//total_anns}%)")
|
|
if ann['category_id'] in source_ids:
|
|
ann['category_id'] = new_category_id
|
|
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] '{target_name}' (ID: {new_category_id})로 통합될 예정입니다")
|
|
else:
|
|
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}'를 찾을 수 없습니다")
|
|
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] {renamed_count}개 클래스 이름이 변경될 예정입니다")
|
|
else:
|
|
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
|
|
total_anns = len(self.data.get('annotations', []))
|
|
print(f"[*] {total_anns}개 어노테이션 업데이트 중...")
|
|
|
|
for idx, ann in enumerate(self.data.get('annotations', []), 1):
|
|
if idx % 1000 == 0:
|
|
print(f"[*] 진행: {idx}/{total_anns} ({idx*100//total_anns}%)")
|
|
ann['category_id'] = id_mapping[ann['category_id']]
|
|
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] 클래스 ID가 재할당될 예정입니다 (시작: {start_id})")
|
|
else:
|
|
print(f"[+] 클래스 ID가 재할당되었습니다 (시작: {start_id})")
|
|
return self
|
|
|
|
|
|
def interactive_mode():
|
|
"""Interactive mode for user-friendly editing"""
|
|
try:
|
|
_interactive_mode_impl()
|
|
except KeyboardInterrupt:
|
|
print("\n\n[*] 사용자에 의해 중단되었습니다.")
|
|
print("[*] 변경사항이 저장되지 않았습니다.")
|
|
|
|
|
|
def _interactive_mode_impl():
|
|
"""Internal implementation of interactive mode"""
|
|
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
|
|
editor = None
|
|
|
|
while not editor:
|
|
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. 직접 경로 입력")
|
|
print(f" q. 종료")
|
|
|
|
choice = input("\n파일 번호 선택 (직접 입력은 0, 종료는 q): ").strip().lower()
|
|
|
|
if choice == 'q':
|
|
print("[*] 종료합니다.")
|
|
return
|
|
|
|
if choice == '0':
|
|
input_file = input("어노테이션 파일 경로 입력 (취소: q): ").strip()
|
|
if input_file.lower() == 'q':
|
|
print("[*] 종료합니다.")
|
|
return
|
|
else:
|
|
try:
|
|
idx = int(choice)
|
|
if 1 <= idx <= len(found_files):
|
|
input_file = found_files[idx - 1]
|
|
else:
|
|
print(f"[!] 1에서 {len(found_files)} 사이의 숫자를 입력해주세요.")
|
|
continue
|
|
except ValueError:
|
|
print("[!] 잘못된 입력입니다. 숫자를 입력해주세요.")
|
|
continue
|
|
else:
|
|
input_file = input("어노테이션 파일 경로 입력 (종료: q): ").strip()
|
|
if input_file.lower() == 'q':
|
|
print("[*] 종료합니다.")
|
|
return
|
|
|
|
if not input_file:
|
|
print("[!] 경로를 입력해주세요.")
|
|
continue
|
|
|
|
# Try to load annotation
|
|
try:
|
|
print(f"\n[+] 로딩 중: {input_file}")
|
|
editor = COCOAnnotationEditor(input_file)
|
|
editor.print_info()
|
|
except FileNotFoundError as e:
|
|
print(f"[!] 오류: {e}")
|
|
print("[!] 다시 시도해주세요.")
|
|
input_file = None
|
|
except Exception as e:
|
|
print(f"[!] 어노테이션 파일 로딩 실패: {e}")
|
|
print("[!] 다시 시도해주세요.")
|
|
input_file = None
|
|
|
|
# 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. 완료 (저장 단계로)")
|
|
print(" h. 도움말")
|
|
print(" q. 저장하지 않고 종료")
|
|
|
|
choice = input("\n작업 선택 (0-5, h, q): ").strip().lower()
|
|
|
|
if choice == 'q':
|
|
confirm = input("저장하지 않고 종료하시겠습니까? (y/N): ").strip().lower()
|
|
if confirm == 'y':
|
|
print("[*] 종료합니다.")
|
|
return
|
|
continue
|
|
|
|
if choice == '0':
|
|
if not operations:
|
|
print("[!] 선택된 작업이 없습니다.")
|
|
confirm = input("그래도 종료하시겠습니까? (y/N): ").strip().lower()
|
|
if confirm == 'y':
|
|
print("[*] 종료합니다.")
|
|
return
|
|
continue
|
|
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, 쉼표로 구분, 취소: q): ").strip()
|
|
if classes.lower() == 'q':
|
|
print("[*] 작업을 취소했습니다.")
|
|
continue
|
|
if not classes:
|
|
print("[!] 입력이 비어있습니다. 작업을 건너뜁니다.")
|
|
continue
|
|
|
|
class_list = [c.strip() for c in classes.split(',') if c.strip()]
|
|
if not class_list:
|
|
print("[!] 유효한 클래스가 없습니다.")
|
|
continue
|
|
|
|
confirm = input(f"\n정말로 {len(class_list)}개의 클래스를 제거하시겠습니까? (y/N): ").strip().lower()
|
|
if confirm != 'y':
|
|
print("[*] 작업을 취소했습니다.")
|
|
continue
|
|
|
|
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()
|
|
|
|
elif choice == 'h': # Help
|
|
print("\n" + "="*60)
|
|
print(" 도움말")
|
|
print("="*60)
|
|
print("\n1. 클래스 제거")
|
|
print(" - 지정한 클래스와 해당 어노테이션을 완전히 제거합니다")
|
|
print(" - 예: 'person,car' 또는 '1,2'")
|
|
print("\n2. 클래스 통합 (병합)")
|
|
print(" - 여러 클래스를 하나로 합칩니다")
|
|
print(" - 예: car,truck,bus → vehicle")
|
|
print("\n3. 클래스 이름 변경")
|
|
print(" - 클래스 이름만 변경합니다 (ID는 유지)")
|
|
print(" - 예: transformer:tt,substation:ss")
|
|
print("\n4. 클래스 ID 재할당")
|
|
print(" - 모든 클래스 ID를 순차적으로 재할당합니다")
|
|
print(" - COCO 형식은 보통 1부터 시작합니다")
|
|
print("\n5. 현재 정보 표시")
|
|
print(" - 어노테이션 파일의 현재 상태를 확인합니다")
|
|
print("\n팁:")
|
|
print(" - 모든 작업은 적용 전 확인 메시지가 표시됩니다")
|
|
print(" - 'q'를 입력하여 언제든지 작업을 취소할 수 있습니다")
|
|
print(" - 저장 시 백업이 자동으로 생성됩니다")
|
|
print("="*60)
|
|
|
|
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()
|