353 lines
13 KiB
Python
353 lines
13 KiB
Python
# scripts/ttv_freq2.py
|
||
"""
|
||
Версия 6: Читает predefined_classes.txt для сохранения полной схемы классов.
|
||
Это позволяет использовать общую схему 12 классов для всех экспериментов,
|
||
даже если в конкретном датасете присутствуют не все классы.
|
||
"""
|
||
|
||
import os
|
||
import shutil
|
||
import random
|
||
import argparse
|
||
from pathlib import Path
|
||
|
||
try:
|
||
import yaml
|
||
except ImportError:
|
||
raise SystemExit("Установите PyYAML: pip install pyyaml")
|
||
|
||
|
||
def load_predefined_classes(path: Path) -> dict:
|
||
"""
|
||
Читает predefined_classes.txt и возвращает маппинг:
|
||
{имя_класса: (id_1based, yolo_index_0based)}
|
||
|
||
Формат файла (одна строка = один класс):
|
||
right_kidney
|
||
right_kidney_upper_segment
|
||
...
|
||
"""
|
||
class_map = {}
|
||
with open(path, 'r', encoding='utf-8') as f:
|
||
for idx, line in enumerate(f):
|
||
name = line.strip()
|
||
if name:
|
||
class_map[name] = {
|
||
'id_1based': idx + 1,
|
||
'yolo_idx': idx,
|
||
}
|
||
return class_map
|
||
|
||
|
||
def clean_caches(output_dir: Path):
|
||
"""Удаляет старые .cache файлы."""
|
||
for cache in output_dir.glob('*.cache'):
|
||
try:
|
||
cache.unlink()
|
||
print(f" Удалён кэш: {cache.name}")
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
def remap_labels(src_dir: Path, dst_dir: Path, id_to_yolo_idx: dict):
|
||
"""
|
||
Копирует .txt файлы с перенумерацией классов из 1-based в 0-based.
|
||
"""
|
||
if dst_dir.exists():
|
||
shutil.rmtree(dst_dir)
|
||
dst_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
copied = 0
|
||
skipped = 0
|
||
for txt_file in src_dir.glob('*.txt'):
|
||
with open(txt_file, 'r', encoding='utf-8') as f:
|
||
lines = f.readlines()
|
||
|
||
new_lines = []
|
||
for line in lines:
|
||
parts = line.strip().split()
|
||
if not parts:
|
||
continue
|
||
try:
|
||
old_class = int(parts[0])
|
||
if old_class in id_to_yolo_idx:
|
||
parts[0] = str(id_to_yolo_idx[old_class])
|
||
new_lines.append(' '.join(parts))
|
||
else:
|
||
skipped += 1
|
||
except ValueError:
|
||
continue
|
||
|
||
if new_lines:
|
||
with open(dst_dir / txt_file.name, 'w', encoding='utf-8') as f:
|
||
f.write('\n'.join(new_lines) + '\n')
|
||
copied += 1
|
||
|
||
return copied, skipped
|
||
|
||
|
||
def create_split(
|
||
images_dir: Path,
|
||
labels_dir: Path,
|
||
output_dir: Path,
|
||
task: str,
|
||
predefined_classes: dict,
|
||
train_ratio: float = 0.8,
|
||
val_ratio: float = 0.1,
|
||
seed: int = 42,
|
||
):
|
||
random.seed(seed)
|
||
|
||
images_dir = images_dir.resolve()
|
||
labels_dir = labels_dir.resolve()
|
||
output_dir = output_dir.resolve()
|
||
|
||
if not images_dir.exists():
|
||
raise FileNotFoundError(f"Папка с изображениями не найдена: {images_dir}")
|
||
if not labels_dir.exists():
|
||
raise FileNotFoundError(f"Папка с лейблами не найдена: {labels_dir}")
|
||
|
||
print("\nОчистка кэша...")
|
||
clean_caches(output_dir)
|
||
|
||
image_files = [f for f in os.listdir(images_dir)
|
||
if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
|
||
if not image_files:
|
||
raise ValueError(f"В {images_dir} не найдено изображений.")
|
||
|
||
print(f"\n{'='*60}")
|
||
print(f"Задача: {task.upper()}")
|
||
print(f"{'='*60}")
|
||
print(f"Изображения: {images_dir}")
|
||
print(f"Лейблы: {labels_dir}")
|
||
print(f"Найдено изображений: {len(image_files)}")
|
||
|
||
# Собираем статистику по классам
|
||
class_to_images = {}
|
||
classes_found = set()
|
||
|
||
for image_file in image_files:
|
||
label_file = os.path.splitext(image_file)[0] + '.txt'
|
||
label_path = labels_dir / label_file
|
||
|
||
if label_path.exists():
|
||
with open(label_path, 'r', encoding='utf-8') as f:
|
||
for line in f:
|
||
parts = line.strip().split()
|
||
if not parts:
|
||
continue
|
||
try:
|
||
class_label = int(parts[0])
|
||
classes_found.add(class_label)
|
||
class_to_images.setdefault(class_label, set()).add(image_file)
|
||
except ValueError:
|
||
continue
|
||
|
||
# Маппинг 1-based ID -> 0-based YOLO index
|
||
id_to_yolo_idx = {info['id_1based']: info['yolo_idx']
|
||
for info in predefined_classes.values()}
|
||
|
||
print(f"\n📋 Схема классов (из predefined_classes.txt):")
|
||
print(f" Всего классов: {len(predefined_classes)}")
|
||
present = []
|
||
absent = []
|
||
for name, info in predefined_classes.items():
|
||
if info['id_1based'] in classes_found:
|
||
present.append((name, info, len(class_to_images[info['id_1based']])))
|
||
else:
|
||
absent.append((name, info))
|
||
|
||
print(f"\n ✓ Присутствуют в датасете ({len(present)}):")
|
||
for name, info, count in present:
|
||
print(f" [{info['id_1based']:2d}] → YOLO idx {info['yolo_idx']:2d} : {name:<40s} ({count} изобр.)")
|
||
|
||
if absent:
|
||
print(f"\n ○ Отсутствуют ({len(absent)}):")
|
||
for name, info in absent:
|
||
print(f" [{info['id_1based']:2d}] → YOLO idx {info['yolo_idx']:2d} : {name}")
|
||
|
||
# Разбиение по наборам (стратифицированное)
|
||
sorted_classes = sorted(class_to_images, key=lambda k: len(class_to_images[k]))
|
||
|
||
train_files, val_files, test_files = set(), set(), set()
|
||
|
||
print(f"\n--- Распределение по наборам ---")
|
||
for class_label in sorted_classes:
|
||
images = set(class_to_images[class_label])
|
||
total = len(images)
|
||
|
||
train_count = max(1, int(total * train_ratio))
|
||
val_count = max(1, int(total * val_ratio))
|
||
|
||
if train_count + val_count > total:
|
||
val_count = max(0, total - train_count)
|
||
|
||
train_selected = set(random.sample(list(images), min(train_count, len(images))))
|
||
images -= train_selected
|
||
val_selected = set(random.sample(list(images), min(val_count, len(images))))
|
||
test_selected = images - val_selected
|
||
|
||
train_files |= train_selected
|
||
val_files |= val_selected
|
||
test_files |= test_selected
|
||
|
||
# Имя класса из predefined
|
||
class_name = next((n for n, info in predefined_classes.items()
|
||
if info['id_1based'] == class_label), f"class_{class_label}")
|
||
|
||
print(f" [{class_label:2d}] {class_name[:30]:<30s}: "
|
||
f"Train={len(train_selected):3d}, Val={len(val_selected):3d}, Test={len(test_selected):3d}")
|
||
|
||
for img in train_selected | val_selected | test_selected:
|
||
for cls_set in class_to_images.values():
|
||
cls_set.discard(img)
|
||
|
||
print(f"\nИТОГО: Train={len(train_files)}, Val={len(val_files)}, Test={len(test_files)}")
|
||
|
||
# Подготовка labels с перенумерацией
|
||
labels_dir_dst = output_dir / 'labels'
|
||
|
||
# Временная папка
|
||
raw_labels = output_dir / '_raw_labels'
|
||
if raw_labels.exists():
|
||
shutil.rmtree(raw_labels)
|
||
raw_labels.mkdir(parents=True, exist_ok=True)
|
||
|
||
print(f"\nКопирование лейблов...")
|
||
for img in train_files | val_files | test_files:
|
||
label_file = labels_dir / (Path(img).stem + '.txt')
|
||
if label_file.exists():
|
||
dst = raw_labels / (Path(img).stem + '.txt')
|
||
if not dst.exists():
|
||
shutil.copy2(label_file, dst)
|
||
|
||
print(f" Копирование с перенумерацией в YOLO 0-based формат:")
|
||
copied, skipped = remap_labels(raw_labels, labels_dir_dst, id_to_yolo_idx)
|
||
print(f" ✓ Обработано {copied} файлов, пропущено строк: {skipped}")
|
||
|
||
shutil.rmtree(raw_labels)
|
||
|
||
# Запись txt-файлов с абсолютными путями
|
||
def write_split(file_path: Path, file_set: set):
|
||
with open(file_path, 'w', encoding='utf-8') as f:
|
||
for img in sorted(file_set):
|
||
abs_path = str(images_dir / img).replace('\\', '/')
|
||
f.write(abs_path + '\n')
|
||
|
||
train_txt = output_dir / f'train_{task}.txt'
|
||
val_txt = output_dir / f'val_{task}.txt'
|
||
test_txt = output_dir / f'test_{task}.txt'
|
||
|
||
write_split(train_txt, train_files)
|
||
write_split(val_txt, val_files)
|
||
write_split(test_txt, test_files)
|
||
|
||
# data.yaml — ВСЕГДА с полной схемой из predefined_classes.txt
|
||
nc = len(predefined_classes)
|
||
class_names = list(predefined_classes.keys())
|
||
|
||
data_yaml = {
|
||
'path': str(output_dir).replace('\\', '/'),
|
||
'train': str(train_txt).replace('\\', '/'),
|
||
'val': str(val_txt).replace('\\', '/'),
|
||
'test': str(test_txt).replace('\\', '/'),
|
||
'nc': nc,
|
||
'names': class_names,
|
||
}
|
||
|
||
yaml_name = f'data_{task}.yaml'
|
||
yaml_path = output_dir / yaml_name
|
||
|
||
with open(yaml_path, 'w', encoding='utf-8') as f:
|
||
yaml.safe_dump(data_yaml, f, sort_keys=False, allow_unicode=True)
|
||
|
||
print(f"\n✓ Создан {yaml_name}: {yaml_path}")
|
||
print(f"✓ nc={nc}, классов: {nc}")
|
||
print(f"✓ Классы (0..{nc-1}): {class_names}")
|
||
|
||
return yaml_path
|
||
|
||
|
||
def main():
|
||
parser = argparse.ArgumentParser(
|
||
description='Создание train/val/test сплита с учётом predefined_classes.txt'
|
||
)
|
||
|
||
parser.add_argument('--task', type=str, required=True,
|
||
choices=['detect', 'segment', 'both'])
|
||
parser.add_argument('--images-dir', type=str, required=True)
|
||
parser.add_argument('--labels-dir', type=str, required=True)
|
||
parser.add_argument('--polygons-dir', type=str, default=None)
|
||
parser.add_argument('--output-dir', type=str, required=True)
|
||
parser.add_argument('--predefined-classes', type=str, default='predefined_classes.txt',
|
||
help='Путь к файлу с полным списком классов (по умолчанию: predefined_classes.txt)')
|
||
parser.add_argument('--train-ratio', type=float, default=0.8)
|
||
parser.add_argument('--val-ratio', type=float, default=0.1)
|
||
parser.add_argument('--seed', type=int, default=42)
|
||
|
||
args = parser.parse_args()
|
||
|
||
# Загружаем полную схему классов
|
||
predefined_path = Path(args.predefined_classes).resolve()
|
||
if not predefined_path.exists():
|
||
raise FileNotFoundError(f"Файл классов не найден: {predefined_path}\n"
|
||
f"Создайте его (по одному классу на строку) или укажите путь через --predefined-classes")
|
||
|
||
predefined_classes = load_predefined_classes(predefined_path)
|
||
print(f"📖 Загружена схема классов: {predefined_path}")
|
||
print(f" Всего классов: {len(predefined_classes)}")
|
||
|
||
if args.task == 'both':
|
||
if not args.polygons_dir:
|
||
raise ValueError("Для task=both необходимо указать --polygons-dir")
|
||
|
||
create_split(
|
||
images_dir=Path(args.images_dir),
|
||
labels_dir=Path(args.labels_dir),
|
||
output_dir=Path(args.output_dir),
|
||
task='detect',
|
||
predefined_classes=predefined_classes,
|
||
train_ratio=args.train_ratio,
|
||
val_ratio=args.val_ratio,
|
||
seed=args.seed,
|
||
)
|
||
|
||
create_split(
|
||
images_dir=Path(args.images_dir),
|
||
labels_dir=Path(args.polygons_dir),
|
||
output_dir=Path(args.output_dir),
|
||
task='segment',
|
||
predefined_classes=predefined_classes,
|
||
train_ratio=args.train_ratio,
|
||
val_ratio=args.val_ratio,
|
||
seed=args.seed,
|
||
)
|
||
|
||
print(f"\n{'='*60}")
|
||
print("✅ ОБА SPLITS СОЗДАНЫ!")
|
||
print(f"{'='*60}")
|
||
|
||
else:
|
||
create_split(
|
||
images_dir=Path(args.images_dir),
|
||
labels_dir=Path(args.labels_dir),
|
||
output_dir=Path(args.output_dir),
|
||
task=args.task,
|
||
predefined_classes=predefined_classes,
|
||
train_ratio=args.train_ratio,
|
||
val_ratio=args.val_ratio,
|
||
seed=args.seed,
|
||
)
|
||
|
||
print(f"\n{'='*60}")
|
||
print("✅ SPLT СОЗДАН!")
|
||
print(f"{'='*60}")
|
||
yaml_name = f'data_{args.task}.yaml'
|
||
if args.task == 'detect':
|
||
print(f"Запуск: python scripts/train_yolo.py --data {Path(args.output_dir) / yaml_name} --workers 0")
|
||
else:
|
||
print(f"Запуск: python scripts/train_yolo_seg.py --data {Path(args.output_dir) / yaml_name} --workers 0")
|
||
|
||
|
||
if __name__ == '__main__':
|
||
main() |