Initial commit: project structure, scripts and documentation
This commit is contained in:
@@ -0,0 +1,353 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user