File size: 4,382 Bytes
f211247 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | #!/usr/bin/env python3
"""Train-time augmentation pipeline (Albumentations-based).
SOTA recipe: elastic + grid distortion + morph + CLAHE + color jitter.
Reference: Albumentations OCR recipe, arXiv:2508.11499 historical HTR.
"""
from pathlib import Path
import yaml
ROOT = Path("/arf/scratch/stakan/hitit-proje")
def build_train_transform(config=None):
"""Returns Albumentations.Compose — import edilir training'de."""
import albumentations as A
from albumentations.pytorch import ToTensorV2
if config is None:
config = yaml.safe_load(open(ROOT / "hitit_ocr" / "configs" / "preprocessing.yaml"))
aug = config.get('augment_train', {})
norm = config.get('normalization', {})
# Pipeline adımları
steps = []
# Geometric distortion (elastic + grid)
geo_aug = []
if aug.get('elastic_transform', {}).get('enabled'):
et = aug['elastic_transform']
geo_aug.append(A.ElasticTransform(
alpha=et.get('alpha', 40.0),
sigma=et.get('sigma', 6.0),
alpha_affine=et.get('alpha_affine', 8.0),
p=1.0,
border_mode=0
))
if aug.get('grid_distortion', {}).get('enabled'):
gd = aug['grid_distortion']
geo_aug.append(A.GridDistortion(
num_steps=gd.get('num_steps', 5),
distort_limit=gd.get('distort_limit', 0.15),
p=1.0,
border_mode=0
))
if geo_aug:
steps.append(A.OneOf(geo_aug, p=0.5))
# Morphological (wedge thickness variation)
if aug.get('morphological', {}).get('enabled'):
m = aug['morphological']
# Custom implementation (Albumentations'da tam yok)
# Kullanım: training script'te manuel apply
pass
# Color jitter
if aug.get('color_jitter', {}).get('enabled'):
cj = aug['color_jitter']
steps.append(A.ColorJitter(
brightness=cj.get('brightness', 0.2),
contrast=cj.get('contrast', 0.2),
saturation=cj.get('saturation', 0.1),
hue=0.0,
p=cj.get('p', 0.5)
))
# Random horizontal flip — cuneiform için DEĞİL (sign orientation önemli)
# if aug.get('horizontal_flip', {}).get('enabled'):
# steps.append(A.HorizontalFlip(p=0.5))
# Normalization
if norm.get('strategy') == 'dataset_specific':
# norm_stats'tan yükle
import json
stats_path = ROOT / "datasets" / "processed" / "normalization_stats.json"
if stats_path.exists():
ns = json.load(open(stats_path))
mean = ns['global_rgb']['mean']
std = ns['global_rgb']['std']
else:
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
elif norm.get('strategy') == 'imagenet':
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
else:
mean = [0.0, 0.0, 0.0]
std = [1.0, 1.0, 1.0]
steps.append(A.Normalize(mean=mean, std=std))
steps.append(ToTensorV2())
return A.Compose(steps)
def build_val_transform(config=None):
"""Validation/test: sadece normalize."""
import albumentations as A
from albumentations.pytorch import ToTensorV2
import json
if config is None:
config = yaml.safe_load(open(ROOT / "hitit_ocr" / "configs" / "preprocessing.yaml"))
norm = config.get('normalization', {})
if norm.get('strategy') == 'dataset_specific':
stats_path = ROOT / "datasets" / "processed" / "normalization_stats.json"
if stats_path.exists():
ns = json.load(open(stats_path))
mean, std = ns['global_rgb']['mean'], ns['global_rgb']['std']
else:
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
else:
mean, std = [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
return A.Compose([
A.Normalize(mean=mean, std=std),
ToTensorV2(),
])
if __name__ == '__main__':
# Test/print pipeline
try:
import albumentations
t = build_train_transform()
print(f"Train pipeline: {len(t.transforms)} steps")
for step in t.transforms:
print(f" - {type(step).__name__}: p={getattr(step, 'p', 'n/a')}")
except ImportError:
print("albumentations kurulu değil. pip install albumentations")
|