에러처리 개선, 입력검증 강화, 진행률 상태 표시, 전후비교 표시 추가

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("[!] 잘못된 선택입니다")