diff --git a/.gitignore b/.gitignore index cf6286a..b104af5 100644 --- a/.gitignore +++ b/.gitignore @@ -49,3 +49,6 @@ labels_backup_*/ *.tmp temp/ tmp/ + +#claude +.claude/ \ No newline at end of file diff --git a/mmdet_annotation_modify.py b/mmdet_annotation_modify.py index 21d5e5a..281ddca 100644 --- a/mmdet_annotation_modify.py +++ b/mmdet_annotation_modify.py @@ -6,19 +6,19 @@ COCO format annotation file manipulation utility Usage examples: # Interactive mode (recommended) - python Utility_lableing_tool.py + python mmdet_annotation_modify.py # Remove classes - python Utility_lableing_tool.py remove --input annotations.json --output new.json --classes "person,car" + python mmdet_annotation_modify.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" + python mmdet_annotation_modify.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" + python mmdet_annotation_modify.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 + python mmdet_annotation_modify.py info --input annotations.json """ import json @@ -34,18 +34,59 @@ import os class COCOAnnotationEditor: """Class for editing COCO format annotations""" - def __init__(self, annotation_path: str): + 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""" - with open(self.annotation_path, 'r', encoding='utf-8') as f: - return json.load(f) + 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): """ @@ -57,16 +98,31 @@ class COCOAnnotationEditor: """ 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}") + if self.dry_run: + print(f"[DRY RUN] 파일이 저장될 예정입니다: {output_path}") + if backup and output_path.exists(): + print(f"[DRY RUN] 백업이 생성될 예정입니다") + return - # 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}") + 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""" @@ -75,19 +131,31 @@ class COCOAnnotationEditor: # 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 + 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: - 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) - } + 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 @@ -161,13 +229,22 @@ class COCOAnnotationEditor: # 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"[*] {original_count}개 어노테이션 검사 중...") - print(f"[+] {removed_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 @@ -232,11 +309,19 @@ class COCOAnnotationEditor: self.data['categories'].append(new_category) # Update annotations (change all source IDs to new ID) - for ann in self.data.get('annotations', []): + 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 - print(f"[+] '{target_name}' (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': @@ -273,7 +358,10 @@ class COCOAnnotationEditor: else: print(f"[!] 이름 '{old_key}'를 찾을 수 없습니다") - print(f"[+] {renamed_count}개 클래스 이름이 변경되었습니다") + 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': @@ -298,15 +386,32 @@ class COCOAnnotationEditor: cat['id'] = id_mapping[cat['id']] # Update annotation category IDs - for ann in self.data.get('annotations', []): + 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']] - print(f"[+] 클래스 ID가 재할당되었습니다 (시작: {start_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) @@ -329,7 +434,9 @@ def interactive_mode(): found_files.extend(glob(pattern, recursive=True)) input_file = None - while not input_file: + editor = None + + while not editor: if found_files: print("\n발견된 어노테이션 파일:") for idx, f in enumerate(found_files[:10], 1): @@ -337,28 +444,53 @@ def interactive_mode(): 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 - choice = input("\n파일 번호 선택 (직접 입력은 0): ").strip() if choice == '0': - input_file = input("어노테이션 파일 경로 입력: ").strip() + input_file = input("어노테이션 파일 경로 입력 (취소: q): ").strip() + if input_file.lower() == 'q': + print("[*] 종료합니다.") + return else: try: - input_file = found_files[int(choice) - 1] - except (ValueError, IndexError): - print("[!] 잘못된 선택입니다. 다시 선택해주세요.") + 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("어노테이션 파일 경로 입력: ").strip() + input_file = input("어노테이션 파일 경로 입력 (종료: q): ").strip() + if input_file.lower() == 'q': + print("[*] 종료합니다.") + return - if not Path(input_file).exists(): - print(f"[!] 파일을 찾을 수 없습니다: {input_file}") - print("[!] 다시 입력해주세요.") + 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 - - # 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', [])} @@ -376,13 +508,26 @@ def interactive_mode(): print(" 4. 클래스 ID 재할당") print(" 5. 현재 정보 표시") print(" 0. 완료 (저장 단계로)") + print(" h. 도움말") + print(" q. 저장하지 않고 종료") - choice = input("\n작업 선택 (0-5): ").strip() + 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("[!] 선택된 작업이 없습니다. 종료합니다...") - return + print("[!] 선택된 작업이 없습니다.") + confirm = input("그래도 종료하시겠습니까? (y/N): ").strip().lower() + if confirm == 'y': + print("[*] 종료합니다.") + return + continue break elif choice == '1': # Remove @@ -390,12 +535,27 @@ def interactive_mode(): 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)}") + 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현재 클래스:") @@ -451,6 +611,30 @@ def interactive_mode(): 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("[!] 잘못된 선택입니다") diff --git a/yolo_annotation_modify.py b/yolo_annotation_modify.py index 3d2b689..375d0d8 100644 --- a/yolo_annotation_modify.py +++ b/yolo_annotation_modify.py @@ -35,12 +35,25 @@ from collections import defaultdict class YOLOAnnotationEditor: """Class for editing YOLO format annotations""" - def __init__(self, dataset_path: str): + 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(): @@ -49,14 +62,36 @@ class YOLOAnnotationEditor: 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""" - with open(self.yaml_path, 'r', encoding='utf-8') as f: - return yaml.safe_load(f) + 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: """ @@ -68,11 +103,19 @@ class YOLOAnnotationEditor: 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}") + if self.dry_run: + print(f"[DRY RUN] labels 폴더 백업이 생성될 예정입니다: {backup_path}") + return 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): """ @@ -81,14 +124,26 @@ class YOLOAnnotationEditor: 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}") + if self.dry_run: + print(f"[DRY RUN] data.yaml이 저장될 예정입니다: {self.yaml_path}") + if backup: + print(f"[DRY RUN] 백업이 생성될 예정입니다") + return - 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 저장 완료") + 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""" @@ -100,13 +155,24 @@ class YOLOAnnotationEditor: for label_file in self.label_files: try: - with open(label_file, 'r') as f: - for line in f: + with open(label_file, 'r', encoding='utf-8') as f: + for line_num, line in enumerate(f, 1): line = line.strip() if line: - class_id = int(line.split()[0]) - class_counts[class_id] += 1 - total_annotations += 1 + 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}") @@ -160,6 +226,7 @@ class YOLOAnnotationEditor: 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 @@ -182,32 +249,72 @@ class YOLOAnnotationEditor: 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 - for label_file in self.label_files: + 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') as f: + with open(label_file, 'r', encoding='utf-8') 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: + 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) - else: - total_removed += 1 - with open(label_file, 'w') as f: - f.writelines(new_lines) + 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) - print(f"[+] {total_removed}개의 어노테이션이 제거되었습니다") + 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': @@ -246,6 +353,9 @@ class YOLOAnnotationEditor: # 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 @@ -253,31 +363,67 @@ class YOLOAnnotationEditor: # Update label files total_merged = 0 - for label_file in self.label_files: + 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') as f: + 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() - 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: + 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) - with open(label_file, 'w') as f: - f.writelines(new_lines) + 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) - print(f"[+] '{target_name}' (ID: {new_class_id})로 통합되었습니다") - print(f"[+] {total_merged}개의 어노테이션이 변경되었습니다") + 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': @@ -313,7 +459,10 @@ class YOLOAnnotationEditor: print(f"[!] 이름 '{old_key}'를 찾을 수 없습니다") self.config['names'] = names - print(f"[+] {renamed_count}개 클래스 이름이 변경되었습니다") + 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': @@ -340,35 +489,70 @@ class YOLOAnnotationEditor: # Update label files total_updated = 0 - for label_file in self.label_files: + 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') as f: + 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() - 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: + 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) - with open(label_file, 'w') as f: - f.writelines(new_lines) + 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) - print(f"[+] 클래스 ID가 재할당되었습니다 (시작: {start_id})") - print(f"[+] {total_updated}개의 어노테이션이 업데이트되었습니다") + 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) @@ -382,36 +566,62 @@ def interactive_mode(): yaml_files = glob("**/data.yaml", recursive=True) dataset_path = None - while not dataset_path: + 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 - choice = input("\n데이터셋 번호 선택 (직접 입력은 0): ").strip() if choice == '0': - dataset_path = input("데이터셋 디렉토리 경로 입력: ").strip() + dataset_path = input("데이터셋 디렉토리 경로 입력 (취소: q): ").strip() + if dataset_path.lower() == 'q': + print("[*] 종료합니다.") + return else: try: - dataset_path = str(Path(yaml_files[int(choice) - 1]).parent) - except (ValueError, IndexError): - print("[!] 잘못된 선택입니다. 다시 선택해주세요.") + 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("데이터셋 디렉토리 경로 입력: ").strip() + dataset_path = input("데이터셋 디렉토리 경로 입력 (종료: q): ").strip() + if dataset_path.lower() == 'q': + print("[*] 종료합니다.") + return - yaml_path = Path(dataset_path) / "data.yaml" - if not yaml_path.exists(): - print(f"[!] data.yaml을 찾을 수 없습니다: {yaml_path}") - print("[!] 다시 입력해주세요.") + 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 - - # Load and show info - print(f"\n[+] 로딩 중: {dataset_path}") - editor = YOLOAnnotationEditor(dataset_path) - editor.print_info() # 2. Select operations print("\n[2/3] 작업 선택") @@ -426,13 +636,26 @@ def interactive_mode(): print(" 4. 클래스 ID 재할당") print(" 5. 현재 정보 표시") print(" 0. 완료 (저장 단계로)") + print(" h. 도움말") + print(" q. 저장하지 않고 종료") - choice = input("\n작업 선택 (0-5): ").strip() + 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("[!] 선택된 작업이 없습니다. 종료합니다...") - return + print("[!] 선택된 작업이 없습니다.") + confirm = input("그래도 종료하시겠습니까? (y/N): ").strip().lower() + if confirm == 'y': + print("[*] 종료합니다.") + return + continue break elif choice == '1': # Remove @@ -440,26 +663,62 @@ def interactive_mode(): 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)}") + 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 또는 이름, 쉼표로 구분): ").strip() - if source: - target = input("통합 후 클래스 이름 입력: ").strip() + 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(',')] - print(f"\n[*] '{target}'로 통합 중: {source_list}") - editor.merge_classes(source_list, target) - operations.append(f"통합: {', '.join(source_list)} → {target}") + 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현재 클래스:") @@ -468,31 +727,75 @@ def interactive_mode(): print("\n매핑 입력 (형식: 기존:새이름,기존:새이름 - 쉼표로 구분)") print("예시: 0:vehicle,1:person 또는 truck:vehicle,person:human") - mapping_input = input("\n입력: ").strip() + mapping_input = input("\n입력 (취소: q): ").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_input.lower() == 'q': + print("[*] 작업을 취소했습니다.") + continue + if not mapping_input: + print("[!] 입력이 비어있습니다. 작업을 건너뜁니다.") + continue - if mapping: - print(f"\n[*] 클래스 이름 변경 중") - editor.rename_classes(mapping) - operations.append(f"이름 변경: {len(mapping)}개 클래스") + 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("\n시작 ID 입력 (기본값: 0): ").strip() - if not start_id: + 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) + start_id = int(start_id_input) + if start_id < 0: + print("[!] 0 이상의 숫자를 입력해주세요.") + continue break except ValueError: - print("[!] 숫자를 입력해주세요.") + 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) @@ -501,6 +804,30 @@ def interactive_mode(): 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("[!] 잘못된 선택입니다")