Initial commit: project structure, scripts and documentation

2026-08-22 13:17:29 +03:00
commit 5042305df0
54 changed files with 14166 additions and 0 deletions
@@ -0,0 +1,99 @@
+# scripts/train_yolo.py
+"""
+Обучение YOLOv8 на детекцию.
+Принимает параметры через командную строку.
+"""
+
+import sys
+import time
+import argparse
+from pathlib import Path
+
+import torch
+from ultralytics import YOLO
+
+
+def check_environment():
+    print("=" * 60)
+    print("ПРОВЕРКА ОКРУЖЕНИЯ (DETECT)")
+    print("=" * 60)
+    print(f"Python: {sys.version.split()[0]}")
+    print(f"PyTorch: {torch.__version__}")
+    print(f"CUDA доступна: {torch.cuda.is_available()}")
+    if torch.cuda.is_available():
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+        print(f"VRAM: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB")
+    print("=" * 60)
+
+
+def train(data_path: str, model: str, epochs: int, imgsz: int,
+          batch: int, device: str, project: str, name: str):
+    check_environment()
+
+    log_dir = Path(project) / name
+    log_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"\nЗагрузка модели: {model}")
+    yolo = YOLO(model)
+
+    args = {
+        'task': 'detect',
+        'data': data_path,
+        'epochs': epochs,
+        'imgsz': imgsz,
+        'batch': batch,
+        'device': device,
+        'project': project,
+        'name': name,
+        'exist_ok': True,
+        'pretrained': True,
+        'optimizer': 'auto',
+        'lr0': 0.01,
+        'lrf': 0.01,
+        'momentum': 0.937,
+        'weight_decay': 0.0005,
+        'warmup_epochs': 3.0,
+        'warmup_momentum': 0.8,
+        'warmup_bias_lr': 0.1,
+        'box': 7.5,
+        'cls': 0.5,
+        'dfl': 1.5,
+        'patience': 20,
+        'save': True,
+        'save_period': 10,
+        'plots': True,
+        'verbose': True,
+    }
+
+    print(f"\nЗапуск обучения DETECT...")
+    print(f"Датасет: {data_path}")
+    print(f"Модель: {model}")
+    print(f"Эпохи: {epochs}, Batch: {batch}, imgsz: {imgsz}")
+    print(f"Результаты: {log_dir}")
+    print("=" * 60)
+
+    t0 = time.time()
+    results = yolo.train(**args)
+    print(f"\nВремя обучения: {(time.time() - t0) / 60:.1f} мин")
+    print(f"Лучшая модель: {log_dir / 'weights' / 'best.pt'}")
+    return results
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Обучение YOLOv8 detect')
+    parser.add_argument('--data', type=str, required=True, help='Путь к data.yaml')
+    parser.add_argument('--model', type=str, default='yolov8s.pt')
+    parser.add_argument('--epochs', type=int, default=100)
+    parser.add_argument('--imgsz', type=int, default=512)
+    parser.add_argument('--batch', type=int, default=16)
+    parser.add_argument('--device', type=str, default='0')
+    parser.add_argument('--project', type=str, default='runs/detect')
+    parser.add_argument('--name', type=str, default='experiment')
+    args = parser.parse_args()
+
+    train(args.data, args.model, args.epochs, args.imgsz,
+          args.batch, args.device, args.project, args.name)
+
+
+if __name__ == '__main__':
+    main()