From 7bd6c15ce9e64e68fda282602540ed7239153351 Mon Sep 17 00:00:00 2001
From: rudals252 <rudals252@a2tec.co.kr>
Date: Thu, 27 Nov 2025 17:44:13 +0900
Subject: [PATCH] =?UTF-8?q?edit=20:=20=EB=94=94=EB=A0=89=ED=86=A0=EB=A6=AC?=
 =?UTF-8?q?=20=EA=B5=AC=EC=A1=B0=20=EB=B3=80=EA=B2=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 yolo_annotation_modify.py | 65 +++++++++++++++++++++++++++++++++------
 1 file changed, 55 insertions(+), 10 deletions(-)

diff --git a/yolo_annotation_modify.py b/yolo_annotation_modify.py
index 3d2b689..e3050d3 100644
--- a/yolo_annotation_modify.py
+++ b/yolo_annotation_modify.py
@@ -47,32 +47,77 @@ class YOLOAnnotationEditor:
             raise FileNotFoundError(f"data.yaml not found in {dataset_path}")
 
         self.config = self._load_yaml()
-        self.labels_path = self.dataset_path / "labels"
 
-        # Find all label files
-        self.label_files = list(self.labels_path.rglob("*.txt"))
-        print(f"[+] Found {len(self.label_files)} label files")
+        # Find label files - support both structures:
+        # 1. Old structure: dataset/labels/
+        # 2. Standard YOLO structure: dataset/train/labels/, dataset/valid/labels/
+        self.label_files = []
+
+        # Try old structure first
+        old_labels_path = self.dataset_path / "labels"
+        if old_labels_path.exists():
+            self.label_files.extend(old_labels_path.rglob("*.txt"))
+            print(f"[+] Found labels in: {old_labels_path}")
+
+        # Try standard YOLO structure
+        for split in ['train', 'valid', 'test']:
+            split_labels_path = self.dataset_path / split / "labels"
+            if split_labels_path.exists():
+                split_files = list(split_labels_path.rglob("*.txt"))
+                self.label_files.extend(split_files)
+                print(f"[+] Found {len(split_files)} labels in: {split_labels_path}")
+
+        if not self.label_files:
+            print("[!] Warning: No label files found")
+        else:
+            print(f"[+] Total: {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)
+            config = yaml.safe_load(f)
+
+        # Normalize names to dictionary format
+        # Support both list format ['class1', 'class2'] and dict format {0: 'class1', 1: 'class2'}
+        if 'names' in config:
+            if isinstance(config['names'], list):
+                # Convert list to dict with index as key
+                config['names'] = {i: name for i, name in enumerate(config['names'])}
+                print("[+] Converted names from list to dict format")
+
+        return config
 
     def backup_labels(self) -> Path:
         """
-        Backup entire labels folder
+        Backup entire labels folder(s)
 
         Returns:
             Path to backup folder
         """
         timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-        backup_path = self.dataset_path / f"labels_backup_{timestamp}"
+        backup_root = self.dataset_path / f"labels_backup_{timestamp}"
+        backup_root.mkdir(exist_ok=True)
 
         print(f"[*] labels 폴더 백업 중...")
-        shutil.copytree(self.labels_path, backup_path)
-        print(f"[+] 백업 생성됨: {backup_path}")
 
-        return backup_path
+        # Backup old structure
+        old_labels_path = self.dataset_path / "labels"
+        if old_labels_path.exists():
+            backup_dest = backup_root / "labels"
+            shutil.copytree(old_labels_path, backup_dest)
+            print(f"[+] 백업됨: labels/ -> {backup_dest}")
+
+        # Backup standard YOLO structure
+        for split in ['train', 'valid', 'test']:
+            split_labels_path = self.dataset_path / split / "labels"
+            if split_labels_path.exists():
+                backup_dest = backup_root / split / "labels"
+                backup_dest.parent.mkdir(parents=True, exist_ok=True)
+                shutil.copytree(split_labels_path, backup_dest)
+                print(f"[+] 백업됨: {split}/labels/ -> {backup_dest}")
+
+        print(f"[+] 백업 생성됨: {backup_root}")
+        return backup_root
 
     def save_yaml(self, backup: bool = True):
         """