974 lines
38 KiB
Python
974 lines
38 KiB
Python
#!/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, dry_run: bool = False):
|
|
"""
|
|
Args:
|
|
dataset_path: Path to YOLO dataset directory (containing data.yaml)
|
|
dry_run: If True, only simulate changes without writing files
|
|
"""
|
|
self.dataset_path = Path(dataset_path)
|
|
self.dry_run = dry_run
|
|
|
|
if self.dry_run:
|
|
print("[DRY RUN] 시뮬레이션 모드 - 파일은 수정되지 않습니다\n")
|
|
|
|
# Validate dataset path exists
|
|
if not self.dataset_path.exists():
|
|
raise FileNotFoundError(f"Dataset directory not found: {dataset_path}")
|
|
|
|
if not self.dataset_path.is_dir():
|
|
raise NotADirectoryError(f"Path is not a directory: {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"
|
|
|
|
# Validate labels directory exists
|
|
if not self.labels_path.exists():
|
|
raise FileNotFoundError(f"labels directory not found in {dataset_path}")
|
|
|
|
# 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"""
|
|
try:
|
|
with open(self.yaml_path, 'r', encoding='utf-8') as f:
|
|
data = yaml.safe_load(f)
|
|
|
|
# Validate YAML structure
|
|
if not isinstance(data, dict):
|
|
raise ValueError(f"Invalid YAML format: Expected dictionary, got {type(data)}")
|
|
|
|
if 'names' not in data:
|
|
raise ValueError("data.yaml must contain 'names' field")
|
|
|
|
if not isinstance(data['names'], dict):
|
|
raise ValueError("'names' field must be a dictionary")
|
|
|
|
return data
|
|
|
|
except yaml.YAMLError as e:
|
|
raise ValueError(f"Failed to parse YAML file: {e}")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Error loading data.yaml: {e}")
|
|
|
|
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}"
|
|
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] labels 폴더 백업이 생성될 예정입니다: {backup_path}")
|
|
return backup_path
|
|
|
|
try:
|
|
print(f"[*] labels 폴더 백업 중...")
|
|
shutil.copytree(self.labels_path, backup_path)
|
|
print(f"[+] 백업 생성됨: {backup_path}")
|
|
return backup_path
|
|
except PermissionError:
|
|
raise PermissionError(f"Permission denied creating backup at {backup_path}")
|
|
except OSError as e:
|
|
raise OSError(f"Failed to create backup: {e}")
|
|
|
|
def save_yaml(self, backup: bool = True):
|
|
"""
|
|
Save modified data.yaml
|
|
|
|
Args:
|
|
backup: Whether to backup original file
|
|
"""
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] data.yaml이 저장될 예정입니다: {self.yaml_path}")
|
|
if backup:
|
|
print(f"[DRY RUN] 백업이 생성될 예정입니다")
|
|
return
|
|
|
|
try:
|
|
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 저장 완료")
|
|
|
|
except PermissionError:
|
|
raise PermissionError(f"Permission denied writing to {self.yaml_path}")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Failed to save data.yaml: {e}")
|
|
|
|
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', encoding='utf-8') as f:
|
|
for line_num, line in enumerate(f, 1):
|
|
line = line.strip()
|
|
if line:
|
|
parts = line.split()
|
|
if len(parts) < 5:
|
|
print(f"[!] Warning: Invalid annotation format in {label_file}:{line_num} (expected 5 values, got {len(parts)})")
|
|
continue
|
|
try:
|
|
class_id = int(parts[0])
|
|
class_counts[class_id] += 1
|
|
total_annotations += 1
|
|
except ValueError:
|
|
print(f"[!] Warning: Invalid class ID in {label_file}:{line_num}")
|
|
except UnicodeDecodeError:
|
|
print(f"[!] Error: Cannot decode file {label_file} (invalid encoding)")
|
|
except PermissionError:
|
|
print(f"[!] Error: Permission denied reading {label_file}")
|
|
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()
|
|
|
|
print("\n[변경 예정]")
|
|
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
|
|
|
|
# Store statistics before change
|
|
class_count_before = len(names)
|
|
|
|
# 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
|
|
failed_files = []
|
|
total_files = len(self.label_files)
|
|
|
|
print(f"[*] {total_files}개 파일 처리 중...")
|
|
for idx, label_file in enumerate(self.label_files, 1):
|
|
if idx % 100 == 0 or idx == total_files:
|
|
print(f"[*] 진행: {idx}/{total_files} ({idx*100//total_files}%)")
|
|
try:
|
|
with open(label_file, 'r', encoding='utf-8') as f:
|
|
lines = f.readlines()
|
|
|
|
new_lines = []
|
|
for line in lines:
|
|
if line.strip():
|
|
parts = line.split()
|
|
if len(parts) < 5:
|
|
print(f"[!] Warning: Skipping invalid line in {label_file}")
|
|
continue
|
|
try:
|
|
class_id = int(parts[0])
|
|
if class_id not in ids_to_remove:
|
|
new_lines.append(line)
|
|
else:
|
|
total_removed += 1
|
|
except ValueError:
|
|
print(f"[!] Warning: Skipping line with invalid class ID in {label_file}")
|
|
new_lines.append(line)
|
|
|
|
if not self.dry_run:
|
|
with open(label_file, 'w', encoding='utf-8') as f:
|
|
f.writelines(new_lines)
|
|
|
|
except PermissionError:
|
|
print(f"[!] Error: Permission denied writing to {label_file}")
|
|
failed_files.append(label_file)
|
|
except Exception as e:
|
|
print(f"[!] Error processing {label_file}: {e}")
|
|
failed_files.append(label_file)
|
|
|
|
if failed_files:
|
|
print(f"[!] Warning: {len(failed_files)} files failed to process")
|
|
|
|
# Print summary
|
|
print(f"\n{'='*60}")
|
|
print(f"작업 요약 - 클래스 제거")
|
|
print(f"{'='*60}")
|
|
print(f"제거된 클래스 수: {len(ids_to_remove)}")
|
|
print(f"클래스 수 변화: {class_count_before} → {len(self.config['names'])}")
|
|
print(f"제거된 어노테이션 수: {total_removed}")
|
|
if failed_files:
|
|
print(f"실패한 파일 수: {len(failed_files)}")
|
|
print(f"{'='*60}\n")
|
|
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] 위 변경사항이 적용될 예정입니다")
|
|
else:
|
|
print(f"[+] 작업이 완료되었습니다")
|
|
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)
|
|
|
|
# Store statistics before change
|
|
class_count_before = len(names)
|
|
|
|
# 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
|
|
failed_files = []
|
|
total_files = len(self.label_files)
|
|
|
|
print(f"[*] {total_files}개 파일 처리 중...")
|
|
for idx, label_file in enumerate(self.label_files, 1):
|
|
if idx % 100 == 0 or idx == total_files:
|
|
print(f"[*] 진행: {idx}/{total_files} ({idx*100//total_files}%)")
|
|
try:
|
|
with open(label_file, 'r', encoding='utf-8') as f:
|
|
lines = f.readlines()
|
|
|
|
new_lines = []
|
|
for line in lines:
|
|
if line.strip():
|
|
parts = line.split()
|
|
if len(parts) < 5:
|
|
print(f"[!] Warning: Skipping invalid line in {label_file}")
|
|
continue
|
|
try:
|
|
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)
|
|
except ValueError:
|
|
print(f"[!] Warning: Skipping line with invalid class ID in {label_file}")
|
|
new_lines.append(line)
|
|
|
|
if not self.dry_run:
|
|
with open(label_file, 'w', encoding='utf-8') as f:
|
|
f.writelines(new_lines)
|
|
|
|
except PermissionError:
|
|
print(f"[!] Error: Permission denied writing to {label_file}")
|
|
failed_files.append(label_file)
|
|
except Exception as e:
|
|
print(f"[!] Error processing {label_file}: {e}")
|
|
failed_files.append(label_file)
|
|
|
|
if failed_files:
|
|
print(f"[!] Warning: {len(failed_files)} files failed to process")
|
|
|
|
# Print summary
|
|
print(f"\n{'='*60}")
|
|
print(f"작업 요약 - 클래스 통합")
|
|
print(f"{'='*60}")
|
|
print(f"통합된 클래스 수: {len(source_class_ids)}")
|
|
print(f"클래스 수 변화: {class_count_before} → {len(self.config['names'])}")
|
|
print(f"새 클래스: {target_name} (ID: {new_class_id})")
|
|
print(f"변경된 어노테이션 수: {total_merged}")
|
|
if failed_files:
|
|
print(f"실패한 파일 수: {len(failed_files)}")
|
|
print(f"{'='*60}\n")
|
|
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] 위 변경사항이 적용될 예정입니다")
|
|
else:
|
|
print(f"[+] 작업이 완료되었습니다")
|
|
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
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] {renamed_count}개 클래스 이름이 변경될 예정입니다")
|
|
else:
|
|
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
|
|
failed_files = []
|
|
total_files = len(self.label_files)
|
|
|
|
print(f"[*] {total_files}개 파일 처리 중...")
|
|
for idx, label_file in enumerate(self.label_files, 1):
|
|
if idx % 100 == 0 or idx == total_files:
|
|
print(f"[*] 진행: {idx}/{total_files} ({idx*100//total_files}%)")
|
|
try:
|
|
with open(label_file, 'r', encoding='utf-8') as f:
|
|
lines = f.readlines()
|
|
|
|
new_lines = []
|
|
for line in lines:
|
|
if line.strip():
|
|
parts = line.split()
|
|
if len(parts) < 5:
|
|
print(f"[!] Warning: Skipping invalid line in {label_file}")
|
|
continue
|
|
try:
|
|
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)
|
|
except ValueError:
|
|
print(f"[!] Warning: Skipping line with invalid class ID in {label_file}")
|
|
new_lines.append(line)
|
|
|
|
if not self.dry_run:
|
|
with open(label_file, 'w', encoding='utf-8') as f:
|
|
f.writelines(new_lines)
|
|
|
|
except PermissionError:
|
|
print(f"[!] Error: Permission denied writing to {label_file}")
|
|
failed_files.append(label_file)
|
|
except Exception as e:
|
|
print(f"[!] Error processing {label_file}: {e}")
|
|
failed_files.append(label_file)
|
|
|
|
if failed_files:
|
|
print(f"[!] Warning: {len(failed_files)} files failed to process")
|
|
|
|
if self.dry_run:
|
|
print(f"[DRY RUN] 클래스 ID가 재할당될 예정입니다 (시작: {start_id})")
|
|
print(f"[DRY RUN] {total_updated}개의 어노테이션이 업데이트될 예정입니다")
|
|
else:
|
|
print(f"[+] 클래스 ID가 재할당되었습니다 (시작: {start_id})")
|
|
print(f"[+] {total_updated}개의 어노테이션이 업데이트되었습니다")
|
|
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(" 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
|
|
editor = None
|
|
|
|
while not editor:
|
|
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. 직접 경로 입력")
|
|
print(f" q. 종료")
|
|
|
|
choice = input("\n데이터셋 번호 선택 (직접 입력은 0, 종료는 q): ").strip().lower()
|
|
|
|
if choice == 'q':
|
|
print("[*] 종료합니다.")
|
|
return
|
|
|
|
if choice == '0':
|
|
dataset_path = input("데이터셋 디렉토리 경로 입력 (취소: q): ").strip()
|
|
if dataset_path.lower() == 'q':
|
|
print("[*] 종료합니다.")
|
|
return
|
|
else:
|
|
try:
|
|
idx = int(choice)
|
|
if 1 <= idx <= len(yaml_files):
|
|
dataset_path = str(Path(yaml_files[idx - 1]).parent)
|
|
else:
|
|
print(f"[!] 1에서 {len(yaml_files)} 사이의 숫자를 입력해주세요.")
|
|
continue
|
|
except ValueError:
|
|
print("[!] 잘못된 입력입니다. 숫자를 입력해주세요.")
|
|
continue
|
|
else:
|
|
dataset_path = input("데이터셋 디렉토리 경로 입력 (종료: q): ").strip()
|
|
if dataset_path.lower() == 'q':
|
|
print("[*] 종료합니다.")
|
|
return
|
|
|
|
if not dataset_path:
|
|
print("[!] 경로를 입력해주세요.")
|
|
continue
|
|
|
|
# Try to load dataset
|
|
try:
|
|
print(f"\n[+] 로딩 중: {dataset_path}")
|
|
editor = YOLOAnnotationEditor(dataset_path)
|
|
editor.print_info()
|
|
except FileNotFoundError as e:
|
|
print(f"[!] 오류: {e}")
|
|
print("[!] 다시 시도해주세요.")
|
|
dataset_path = None
|
|
except Exception as e:
|
|
print(f"[!] 데이터셋 로딩 실패: {e}")
|
|
print("[!] 다시 시도해주세요.")
|
|
dataset_path = None
|
|
|
|
# 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 cid, name in sorted(editor.config['names'].items()):
|
|
print(f" {cid}: {name}")
|
|
|
|
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_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 또는 이름, 쉼표로 구분, 취소: q): ").strip()
|
|
if source.lower() == 'q':
|
|
print("[*] 작업을 취소했습니다.")
|
|
continue
|
|
if not source:
|
|
print("[!] 입력이 비어있습니다. 작업을 건너뜁니다.")
|
|
continue
|
|
|
|
source_list = [c.strip() for c in source.split(',') if c.strip()]
|
|
if len(source_list) < 2:
|
|
print("[!] 최소 2개 이상의 클래스를 입력해야 합니다.")
|
|
continue
|
|
|
|
target = input("통합 후 클래스 이름 입력 (취소: q): ").strip()
|
|
if target.lower() == 'q':
|
|
print("[*] 작업을 취소했습니다.")
|
|
continue
|
|
if not target:
|
|
print("[!] 클래스 이름을 입력해주세요.")
|
|
continue
|
|
|
|
confirm = input(f"\n{len(source_list)}개의 클래스를 '{target}'로 통합하시겠습니까? (y/N): ").strip().lower()
|
|
if confirm != 'y':
|
|
print("[*] 작업을 취소했습니다.")
|
|
continue
|
|
|
|
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입력 (취소: q): ").strip()
|
|
|
|
if mapping_input.lower() == 'q':
|
|
print("[*] 작업을 취소했습니다.")
|
|
continue
|
|
if not mapping_input:
|
|
print("[!] 입력이 비어있습니다. 작업을 건너뜁니다.")
|
|
continue
|
|
|
|
mapping = {}
|
|
invalid_pairs = []
|
|
for pair in mapping_input.split(','):
|
|
pair = pair.strip()
|
|
if ':' in pair:
|
|
parts = pair.split(':', 1)
|
|
old, new = parts[0].strip(), parts[1].strip()
|
|
if old and new:
|
|
mapping[old] = new
|
|
else:
|
|
invalid_pairs.append(pair)
|
|
else:
|
|
invalid_pairs.append(pair)
|
|
|
|
if invalid_pairs:
|
|
print(f"[!] 잘못된 형식의 매핑: {', '.join(invalid_pairs)}")
|
|
|
|
if not mapping:
|
|
print("[!] 유효한 매핑이 없습니다.")
|
|
continue
|
|
|
|
print(f"\n변경 예정:")
|
|
for old, new in mapping.items():
|
|
print(f" {old} → {new}")
|
|
|
|
confirm = input(f"\n{len(mapping)}개의 클래스 이름을 변경하시겠습니까? (y/N): ").strip().lower()
|
|
if confirm != 'y':
|
|
print("[*] 작업을 취소했습니다.")
|
|
continue
|
|
|
|
print(f"\n[*] 클래스 이름 변경 중")
|
|
editor.rename_classes(mapping)
|
|
operations.append(f"이름 변경: {len(mapping)}개 클래스")
|
|
|
|
elif choice == '4': # Reindex
|
|
while True:
|
|
start_id_input = input("\n시작 ID 입력 (기본값: 0, 취소: q): ").strip().lower()
|
|
if start_id_input == 'q':
|
|
print("[*] 작업을 취소했습니다.")
|
|
start_id = None
|
|
break
|
|
if not start_id_input:
|
|
start_id = 0
|
|
break
|
|
try:
|
|
start_id = int(start_id_input)
|
|
if start_id < 0:
|
|
print("[!] 0 이상의 숫자를 입력해주세요.")
|
|
continue
|
|
break
|
|
except ValueError:
|
|
print("[!] 유효한 숫자를 입력해주세요.")
|
|
|
|
if start_id is None:
|
|
continue
|
|
|
|
confirm = input(f"\n클래스 ID를 {start_id}부터 재할당하시겠습니까? (y/N): ").strip().lower()
|
|
if confirm != 'y':
|
|
print("[*] 작업을 취소했습니다.")
|
|
continue
|
|
|
|
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()
|
|
|
|
elif choice == 'h': # Help
|
|
print("\n" + "="*60)
|
|
print(" 도움말")
|
|
print("="*60)
|
|
print("\n1. 클래스 제거")
|
|
print(" - 지정한 클래스와 해당 어노테이션을 완전히 제거합니다")
|
|
print(" - 예: '0,1' 또는 'person,car'")
|
|
print("\n2. 클래스 통합 (병합)")
|
|
print(" - 여러 클래스를 하나로 합칩니다")
|
|
print(" - 예: car,truck,bus → vehicle")
|
|
print("\n3. 클래스 이름 변경")
|
|
print(" - 클래스 이름만 변경합니다 (ID는 유지)")
|
|
print(" - 예: 0:vehicle,1:person")
|
|
print("\n4. 클래스 ID 재할당")
|
|
print(" - 모든 클래스 ID를 순차적으로 재할당합니다")
|
|
print(" - 예: 시작 ID 0 → 0,1,2,3...")
|
|
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}")
|
|
|
|
# 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()
|