#!/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()