code/src/train_detection.py · savastakan/hitit-cuneiform-ocr at main

hitit-cuneiform-ocr / code /src /train_detection.py

Initial upload: code + 5 record checkpoints + fuse

f211247 verified 6 months ago

6.87 kB

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

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