code/src/train_detection.py

#!/usr/bin/env python3
"""YOLO11-P2 detection training (Ultralytics wrapper).

5-fold CV support via tablet_view_fold.
BF16 + compile.
"""
import os, sys, json, argparse, tempfile, shutil, time
from pathlib import Path
import yaml
import torch

ROOT = Path("/arf/scratch/stakan/hitit-proje")

def _ensure_cuda():
    """Work around transient NVML / CUDA_VISIBLE_DEVICES init races.

    Retry torch.cuda probe up to 3×; on persistent failure, clear any stale
    CUDA_VISIBLE_DEVICES and re-probe."""
    for attempt in range(3):
        try:
            n = torch.cuda.device_count()
            ok = torch.cuda.is_available()
            if ok and n > 0: return n
        except Exception as e:
            print(f"CUDA probe {attempt}: {e}", flush=True)
        time.sleep(2)
    # Final attempt: unset CUDA_VISIBLE_DEVICES if SLURM set it without corresponding devices
    vis = os.environ.get('CUDA_VISIBLE_DEVICES', '')
    print(f"[warn] CUDA init failed; CUDA_VISIBLE_DEVICES={vis!r}. Attempting reset.", flush=True)
    os.environ.pop('CUDA_VISIBLE_DEVICES', None)
    try:
        import importlib
        importlib.reload(torch.cuda)
        n = torch.cuda.device_count()
        if n > 0: return n
    except Exception: pass
    return 0

def build_fold_data_yaml(fold, output_dir, n_folds=5):
    """
    Her fold için ultralytics data.yaml oluştur.
    tablet_view_fold != fold → train, == fold → val.
    Manifest'te tablet_view_fold yoksa tablet_id (yoksa path) hash'inden
    deterministik fold üret — aynı tabletin farklı view'leri aynı fold'a düşer.
    """
    manifest_path = ROOT / 'datasets/unified/detection/manifest.jsonl'

    train_imgs, val_imgs = [], []
    train_lbls, val_lbls = [], []

    def derive_fold(record):
        tvf = record.get('tablet_view_fold')
        if tvf is not None:
            return int(tvf)
        key = record.get('tablet_id') or record.get('path', '')
        return (hash(key) & 0x7fffffff) % n_folds

    with open(manifest_path) as f:
        for line in f:
            r = json.loads(line)
            if r.get('storage') != 'fs' or not r.get('path'): continue
            img_path = r['path']
            lbl_path = r.get('label_path')
            if not lbl_path or not Path(lbl_path).exists(): continue
            if not Path(img_path).exists(): continue

            tvf = derive_fold(r)
            if tvf == fold:
                val_imgs.append(img_path)
                val_lbls.append(lbl_path)
            else:
                train_imgs.append(img_path)
                train_lbls.append(lbl_path)

    print(f"Fold {fold}: train={len(train_imgs)}, val={len(val_imgs)}")

    # Ultralytics: write image lists
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    train_list = output_dir / 'train.txt'
    val_list = output_dir / 'val.txt'
    with open(train_list, 'w') as f:
        for p in train_imgs: f.write(f"{p}\n")
    with open(val_list, 'w') as f:
        for p in val_imgs: f.write(f"{p}\n")

    # Klas isimleri için label_dict
    label_dict_path = ROOT / 'datasets/sources/yeni_veri/label_dict.txt'
    if label_dict_path.exists():
        names = {}
        with open(label_dict_path) as f:
            for line in f:
                if ':' in line:
                    idx, name = line.split(':', 1)
                    names[int(idx.strip())] = name.strip()
        # Convert to list
        max_idx = max(names.keys()) + 1
        names_list = [names.get(i, f'class_{i}') for i in range(max_idx)]
    else:
        names_list = [f'class_{i}' for i in range(400)]

    data_yaml = {
        'path': str(output_dir),
        'train': 'train.txt',
        'val': 'val.txt',
        'names': names_list,
        'nc': len(names_list),
    }
    yaml_path = output_dir / 'data.yaml'
    with open(yaml_path, 'w') as f:
        yaml.dump(data_yaml, f)
    return str(yaml_path)

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--config', default=str(ROOT / 'hitit_ocr/configs/detection_hitit_v1.yaml'))
    ap.add_argument('--fold', type=int, default=0)
    ap.add_argument('--split-field', default='tablet_view_fold')
    ap.add_argument('--output', default=None)
    args = ap.parse_args()

    cfg = yaml.safe_load(open(args.config))
    output = Path(args.output or f'hitit_ocr/runs/detection_fold{args.fold}/')
    output.mkdir(parents=True, exist_ok=True)

    # 1) v2 (tablet + bg mining + pseudo-labels)
    # 2) tablet-only (>=3 box, +5% bg)
    # 3) tam ready/detection
    # 4) manifest fallback
    candidates = [
        ROOT / f'datasets/ready/detection_tablets_v2/fold_{args.fold}/data.yaml',
        ROOT / f'datasets/ready/detection_tablets/fold_{args.fold}/data.yaml',
        ROOT / f'datasets/ready/detection/fold_{args.fold}/data.yaml',
    ]
    data_yaml = None
    for c in candidates:
        if c.exists():
            data_yaml = str(c); print(f"Using fold data.yaml: {data_yaml}"); break
    if data_yaml is None:
        data_yaml = build_fold_data_yaml(args.fold, output / 'data')
        print(f"Built fold data.yaml from manifest: {data_yaml}")

    # Import ultralytics
    try:
        from ultralytics import YOLO
    except ImportError:
        print("Ultralytics not installed — pip install ultralytics")
        sys.exit(1)

    # Model
    model_name = cfg.get('model', 'yolo11m-p2.yaml')
    model = YOLO(model_name)

    # Train
    train_args = {
        'data': data_yaml,
        'epochs': cfg.get('epochs', 150),
        'imgsz': cfg.get('imgsz', 1280),
        'batch': cfg.get('batch', 4),
        'device': (lambda n: list(range(n)) if n > 0 else 'cpu')(_ensure_cuda()),
        'project': str(output.parent),
        'name': output.name,
        'optimizer': cfg.get('optimizer', 'AdamW'),
        'lr0': cfg.get('lr0', 0.001),
        'mosaic': cfg.get('mosaic', 0.8),
        'mixup': cfg.get('mixup', 0.15),
        'copy_paste': cfg.get('copy_paste', 0.15),
        'fliplr': cfg.get('fliplr', 0.0),
        'flipud': cfg.get('flipud', 0.0),
        'degrees': cfg.get('degrees', 5),
        'close_mosaic': cfg.get('close_mosaic', 10),
        'amp': True,
        'exist_ok': True,
        'cache': cfg.get('cache', 'ram'),
        'workers': cfg.get('workers', 16),
        'patience': cfg.get('patience', 20),
        'save_period': cfg.get('save_period', 5),  # SWA için son ckpt'leri sakla
        'iou': cfg.get('iou', 0.7),  # NMS IoU
    }

    # torch.compile
    if cfg.get('efficiency', {}).get('torch_compile'):
        train_args['compile'] = True

    print(f"Training with args: {train_args}")
    results = model.train(**train_args)

    # Save best
    best_src = output / 'weights' / 'best.pt'
    if best_src.exists():
        shutil.copy(best_src, output / 'best.pt')
    print(f"DONE: {output}")

if __name__ == '__main__':
    main()