251 lines
9.4 KiB
Python
251 lines
9.4 KiB
Python
from __future__ import annotations
|
||
|
||
import random
|
||
from pathlib import Path
|
||
from typing import Any
|
||
from uuid import uuid4
|
||
|
||
import yaml
|
||
|
||
|
||
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():
|
||
if isinstance(raw_key, bool):
|
||
raise ValueError(f"Некорректный ID класса в '{source}': {raw_key!r}.")
|
||
try:
|
||
class_id = int(raw_key)
|
||
except (TypeError, ValueError) as exc:
|
||
raise ValueError(
|
||
f"Некорректный ID класса в '{source}': {raw_key!r}."
|
||
) from exc
|
||
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_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 labels_dir.rglob("*.txt"):
|
||
if txt_file.name == "classes.txt":
|
||
continue
|
||
try:
|
||
with txt_file.open("r", encoding="utf-8") as label_file:
|
||
for line in label_file:
|
||
parts = line.strip().split()
|
||
if parts:
|
||
class_ids.add(int(parts[0]))
|
||
except (OSError, ValueError):
|
||
continue
|
||
|
||
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 split_dataset(
|
||
dataset_dir: str, train_ratio: float, classes_path: str
|
||
) -> tuple[int, int, str]:
|
||
if not dataset_dir.strip():
|
||
raise ValueError("Укажите каталог датасета.")
|
||
if not 0.1 <= train_ratio <= 0.95:
|
||
raise ValueError("Доля обучающей выборки должна быть от 0.1 до 0.95.")
|
||
|
||
base_dir = Path(dataset_dir.strip()).expanduser().absolute()
|
||
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) * train_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)
|
||
|
||
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)
|