from __future__ import annotations import math import random from pathlib import Path from typing import Any from uuid import uuid4 import yaml def _normalize_class_key(raw_key: Any, source: Path) -> int: if isinstance(raw_key, bool): raise ValueError(f"Некорректный ID класса в '{source}': {raw_key!r}.") if isinstance(raw_key, int): class_id = raw_key elif isinstance(raw_key, str): try: class_id = int(raw_key.strip()) except ValueError as exc: raise ValueError( f"Некорректный ID класса в '{source}': {raw_key!r}." ) from exc else: raise ValueError(f"Некорректный ID класса в '{source}': {raw_key!r}.") if class_id < 0: raise ValueError(f"ID класса в '{source}' не может быть отрицательным.") return class_id def _normalize_names(names: Any, source: Path) -> dict[int, str]: if isinstance(names, list): items = enumerate(names) elif isinstance(names, dict): normalized_items: list[tuple[int, Any]] = [] seen: set[int] = set() for raw_key, value in names.items(): class_id = _normalize_class_key(raw_key, source) if class_id in seen: raise ValueError(f"Повторяющийся ID класса {class_id} в '{source}'.") seen.add(class_id) normalized_items.append((class_id, value)) items = normalized_items else: raise ValueError( f"Поле names в '{source}' должно быть списком или словарём." ) result: dict[int, str] = {} for class_id, raw_name in items: if isinstance(raw_name, (dict, list)): name = "" else: name = str(raw_name).strip() if not name: raise ValueError(f"Пустое имя класса {class_id} в '{source}'.") result[class_id] = name if not result: raise ValueError(f"Список классов в '{source}' пуст.") expected_ids = list(range(len(result))) if sorted(result) != expected_ids: raise ValueError( f"ID классов в '{source}' должны идти подряд, начиная с 0." ) return dict(sorted(result.items())) def _parse_text_classes(path: Path) -> dict[int, str]: with path.open("r", encoding="utf-8") as classes_file: lines = [line.strip() for line in classes_file if line.strip()] return _normalize_names(lines, path) def _load_yaml(path: Path) -> Any: try: with path.open("r", encoding="utf-8") as yaml_file: return yaml.safe_load(yaml_file) except yaml.YAMLError as exc: raise ValueError(f"Не удалось разобрать YAML классов '{path}': {exc}.") from exc def _parse_yaml_classes(path: Path) -> dict[int, str]: data = _load_yaml(path) if not isinstance(data, dict) or "names" not in data: raise ValueError(f"В YAML-файле '{path}' отсутствует поле names.") return _normalize_names(data["names"], path) def _parse_classes_file(path: Path) -> dict[int, str]: suffix = path.suffix.lower() if suffix == ".txt": return _parse_text_classes(path) if suffix in {".yaml", ".yml"}: return _parse_yaml_classes(path) raise ValueError( f"Неподдерживаемый формат файла классов '{path}'. " "Используйте .txt, .yaml или .yml." ) def _read_label_class_ids(path: Path) -> set[int]: class_ids: set[int] = set() try: with path.open("r", encoding="utf-8") as label_file: for line_number, line in enumerate(label_file, start=1): parts = line.strip().split() if not parts: continue try: numeric_id = float(parts[0]) except ValueError as exc: raise ValueError( f"Некорректный ID класса в '{path}', строка {line_number}: " f"{parts[0]!r}." ) from exc if ( not math.isfinite(numeric_id) or not numeric_id.is_integer() or numeric_id < 0 ): raise ValueError( f"Некорректный ID класса в '{path}', строка {line_number}: " f"{parts[0]!r}." ) class_ids.add(int(numeric_id)) except (OSError, UnicodeError) as exc: raise ValueError(f"Не удалось прочитать файл разметки '{path}': {exc}.") from exc return class_ids def read_classes(dataset_dir: Path, custom_classes_path: str) -> dict[int, str]: # An explicit path is authoritative: typos and malformed files must not fall back. if custom_classes_path.strip(): path = Path(custom_classes_path.strip()).expanduser() if not path.exists(): raise ValueError(f"Указанный файл классов '{path}' не существует.") if not path.is_file(): raise ValueError(f"Путь к классам '{path}' не является файлом.") try: return _parse_classes_file(path) except OSError as exc: raise ValueError(f"Не удалось прочитать файл классов '{path}': {exc}.") from exc for path in (dataset_dir / "classes.txt", dataset_dir / "labels" / "classes.txt"): if path.is_file(): return _parse_text_classes(path) yaml_files = sorted( (*dataset_dir.glob("*.yaml"), *dataset_dir.glob("*.yml")), key=lambda item: item.name, ) for path in yaml_files: try: data = _load_yaml(path) except (OSError, ValueError): continue if isinstance(data, dict) and "names" in data: return _normalize_names(data["names"], path) # Infer the range from every nested label file when no class list is available. class_ids: set[int] = set() labels_dir = dataset_dir / "labels" if labels_dir.exists(): for txt_file in sorted(labels_dir.rglob("*.txt")): if txt_file.name == "classes.txt": continue class_ids.update(_read_label_class_ids(txt_file)) if class_ids: max_id = max(class_ids) return {class_id: f"class_{class_id}" for class_id in range(max_id + 1)} raise ValueError( "Не удалось найти список классов. Создайте classes.txt или YAML с полем " "names в корневой папке датасета либо укажите путь к нему." ) def _write_new(path: Path, content: str) -> None: with path.open("x", encoding="utf-8") as output_file: output_file.write(content) def _validate_image_labels( images_dir: Path, labels_dir: Path, image_files: list[Path], classes: dict[int, str], ) -> None: checked_paths: set[Path] = set() for image_path in image_files: label_path = labels_dir / image_path.relative_to(images_dir).with_suffix(".txt") if label_path in checked_paths: continue checked_paths.add(label_path) if not label_path.exists(): # YOLO treats an image without a label file as a background image. continue if not label_path.is_file(): raise ValueError(f"Путь к разметке '{label_path}' не является файлом.") for class_id in _read_label_class_ids(label_path): if class_id not in classes: raise ValueError( f"В файле '{label_path}' указан класс {class_id}, " f"но список классов содержит ID от 0 до {len(classes) - 1}." ) def split_dataset( dataset_dir: str, train_ratio: float, classes_path: str ) -> tuple[int, int, str]: if not isinstance(dataset_dir, str): raise ValueError("Каталог датасета должен быть строкой.") if not dataset_dir.strip(): raise ValueError("Укажите каталог датасета.") if not isinstance(classes_path, str): raise ValueError("Путь к файлу классов должен быть строкой.") if isinstance(train_ratio, bool) or not isinstance(train_ratio, (int, float)): raise ValueError("Доля обучающей выборки должна быть числом.") try: ratio = float(train_ratio) except OverflowError as exc: raise ValueError("Доля обучающей выборки должна быть конечным числом.") from exc if not math.isfinite(ratio): raise ValueError("Доля обучающей выборки должна быть конечным числом.") if not 0.1 <= ratio <= 0.95: raise ValueError("Доля обучающей выборки должна быть от 0.1 до 0.95.") base_dir = Path(dataset_dir.strip()).expanduser().resolve(strict=False) images_dir = base_dir / "images" labels_dir = base_dir / "labels" if not base_dir.is_dir(): raise ValueError(f"Каталог датасета '{base_dir}' не существует.") if not images_dir.is_dir(): raise ValueError(f"Папка с изображениями '{images_dir}' не найдена.") if not labels_dir.is_dir(): raise ValueError(f"Папка с разметкой '{labels_dir}' не найдена.") valid_extensions = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".tif", ".tiff"} image_files = sorted( ( path for path in images_dir.rglob("*") if path.is_file() and path.suffix.lower() in valid_extensions ), key=lambda item: item.as_posix(), ) if not image_files: raise ValueError(f"В папке '{images_dir}' не найдено изображений.") if len(image_files) == 1: raise ValueError( "Для разделения нужно минимум 2 изображения; найдено: 1." ) rng = random.Random(42) rng.shuffle(image_files) split_idx = int(len(image_files) * ratio) split_idx = max(1, min(split_idx, len(image_files) - 1)) train_images = image_files[:split_idx] val_images = image_files[split_idx:] # Resolve classes before creating output so invalid input leaves no partial split. classes = read_classes(base_dir, classes_path) _validate_image_labels(images_dir, labels_dir, image_files, classes) relative_split_dir = Path(".yolo-webui") / "splits" / uuid4().hex split_dir = base_dir / relative_split_dir split_dir.mkdir(parents=True, exist_ok=False) train_txt_path = split_dir / "train.txt" val_txt_path = split_dir / "val.txt" dataset_yaml_path = split_dir / "dataset.yaml" _write_new(train_txt_path, "".join(f"{image}\n" for image in train_images)) _write_new(val_txt_path, "".join(f"{image}\n" for image in val_images)) # Read existing dataset YAML if available to preserve custom tags (e.g., kpt_shape) existing_data = {} if classes_path.strip(): cp = Path(classes_path.strip()).expanduser() if cp.is_file() and cp.suffix.lower() in (".yaml", ".yml"): try: existing_data = _load_yaml(cp) or {} except Exception: pass if not existing_data: yaml_files = sorted( (*base_dir.glob("*.yaml"), *base_dir.glob("*.yml")), key=lambda item: item.name, ) for yf in yaml_files: try: data = _load_yaml(yf) if isinstance(data, dict): existing_data = data break except Exception: continue # Build dataset metadata, merging existing keys dataset_data = {} if isinstance(existing_data, dict): dataset_data.update(existing_data) dataset_data.update({ "path": str(base_dir), "train": (relative_split_dir / train_txt_path.name).as_posix(), "val": (relative_split_dir / val_txt_path.name).as_posix(), }) # `read_classes()` has already applied the explicit-path precedence and validated # the IDs. A different root YAML must never replace that authoritative result. dataset_data["names"] = classes _write_new( dataset_yaml_path, yaml.safe_dump(dataset_data, allow_unicode=True, sort_keys=False), ) return len(train_images), len(val_images), str(dataset_yaml_path)