KHPR / Scripts /generate_paragraphs.py
Karez's picture
Upload folder using huggingface_hub
449d173 verified
Raw
History Blame Contribute Delete
16.8 kB
"""
Synthetic Handwritten Paragraph Generator
Single-Writer Consistency | Cross-Source Mixing | Zero Duplicate Orders
Generates synthetic paragraph images from handwritten line sources for
pre-training paragraph recognition models. Supports RTL scripts.
Guarantees:
1. Single-writer consistency: all lines in each paragraph from one writer
2. Cross-source mixing: multi-line paragraphs use lines from 2+ sources
3. Zero duplicate text orderings across the entire dataset
4. Source-level isolation: training and validation use separate line pools
5. Configurable reuse caps per source to control line repetition
Usage:
python generate_paragraphs.py \
--unique_train_dir ./data/UniqueLines/Training \
--fixed_train_dir ./data/FixedLines/Training \
--synthetic_train_dir ./data/SyntheticLines/Training \
--unique_val_dir ./data/UniqueLines/Validation \
--fixed_val_dir ./data/FixedLines/Validation \
--synthetic_val_dir ./data/SyntheticLines/Validation \
--output_dir ./SyntheticParagraphs_12000 \
--dataset_size 12000
"""
import os, glob, random, argparse, gc
import numpy as np
from PIL import Image
from tqdm import tqdm
from datetime import datetime
from collections import defaultdict
def parse_args():
p = argparse.ArgumentParser(description="Synthetic Paragraph Generator")
p.add_argument("--unique_train_dir", type=str, required=True)
p.add_argument("--fixed_train_dir", type=str, default=None)
p.add_argument("--synthetic_train_dir", type=str, default=None)
p.add_argument("--unique_val_dir", type=str, required=True)
p.add_argument("--fixed_val_dir", type=str, default=None)
p.add_argument("--synthetic_val_dir", type=str, default=None)
p.add_argument("--output_dir", type=str, required=True)
p.add_argument("--output_format", type=str, default="TIFF", choices=["TIFF","PNG","JPEG"])
p.add_argument("--dataset_size", type=int, default=12000)
p.add_argument("--train_ratio", type=float, default=0.85)
p.add_argument("--min_lines", type=int, default=1)
p.add_argument("--max_lines", type=int, default=7)
p.add_argument("--spacing_min", type=int, default=15)
p.add_argument("--spacing_max", type=int, default=35)
p.add_argument("--canvas_width", type=int, default=2470)
p.add_argument("--canvas_height", type=int, default=1200)
p.add_argument("--padding", type=int, default=40)
p.add_argument("--train_fixed_cap", type=float, default=1.5)
p.add_argument("--train_synthetic_cap", type=float, default=2.5)
p.add_argument("--val_fixed_cap", type=float, default=1.0)
p.add_argument("--val_synthetic_cap", type=float, default=2.0)
p.add_argument("--crop_whitespace", action="store_true", default=True)
p.add_argument("--no_crop_whitespace", action="store_true")
p.add_argument("--clean_left_edge", action="store_true", default=True)
p.add_argument("--no_clean_left_edge", action="store_true")
p.add_argument("--whitespace_threshold", type=int, default=250)
p.add_argument("--edge_pixels", type=int, default=8)
p.add_argument("--max_attempts", type=int, default=500)
p.add_argument("--seed", type=int, default=42)
p.add_argument("--gc_interval", type=int, default=100)
return p.parse_args()
def extract_writer_id(filename):
basename = os.path.splitext(os.path.basename(filename))[0]
parts = basename.split('_')
return parts[0] if parts else basename
def load_line_dataset(directory, name="Dataset"):
if not directory or not os.path.exists(directory):
return []
files = []
for ext in ["*.tif","*.tiff","*.png","*.jpg","*.jpeg","*.bmp"]:
files.extend(glob.glob(os.path.join(directory, ext)))
files.extend(glob.glob(os.path.join(directory, ext.upper())))
files = sorted(list(set(files)))
data, skipped = [], 0
for p in files:
lp = os.path.splitext(p)[0] + ".txt"
if not os.path.exists(lp):
skipped += 1; continue
try:
with open(lp, "r", encoding="utf-8") as f: label = f.readline().strip()
except:
try:
with open(lp, "r", encoding="utf-8-sig") as f: label = f.readline().strip()
except: skipped += 1; continue
if label: data.append((p, label))
print(f" {name}: {len(data)} lines, {skipped} skipped")
return data
def merge_lines_by_writer(source_pairs):
merged = defaultdict(list)
counts = defaultdict(lambda: defaultdict(int))
for src, lines in source_pairs:
for path, label in lines:
wid = extract_writer_id(path)
merged[wid].append((path, label, src))
counts[wid][src] += 1
return dict(merged), {k: dict(v) for k, v in counts.items()}
class SingleWriterParagraphGenerator:
def __init__(self, merged, totals, fixed_cap, synth_cap, min_l, max_l, max_att):
self.merged = merged
self.used = set()
self.totals = totals
self.usage = defaultdict(int)
self.fcap, self.scap = fixed_cap, synth_cap
self.min_l, self.max_l, self.max_att = min_l, max_l, max_att
self.line_usage = defaultdict(int)
self.line_src = {}
for lines in merged.values():
for p, _, s in lines: self.line_src[p] = s
self.valid = {w: l for w, l in merged.items() if len(l) >= min_l}
self.wlist = list(self.valid.keys())
self.writers_used = set()
self.src_para = defaultdict(int)
self.n_paras = 0
self.n_lines = 0
self.n_dups = 0
def _capped(self, s):
if s == "unique": return False
c = self.fcap if s == "fixed" else self.scap
return self.usage[s] >= int(c * self.totals.get(s, 0))
def _both_capped(self):
return self._capped("fixed") and self._capped("synthetic")
def _avail(self, wlines):
return [x for x in wlines if x[2] == "unique" or not self._capped(x[2])]
def get_paragraph_lines(self):
if not self.wlist: return None
bc = self._both_capped()
for _ in range(self.max_att):
wid = random.choice(self.wlist)
av = self._avail(self.valid[wid])
if len(av) < self.min_l: continue
nl = random.randint(self.min_l, min(self.max_l, len(av)))
sel = None
if nl >= 2 and not bc:
srcs = set(s for _, _, s in av)
if len(srcs) < 2:
if self.min_l <= 1: nl = 1; sel = random.sample(av, 1)
else: continue
else:
for _ in range(30):
c = random.sample(av, nl)
if len(set(s for _, _, s in c)) >= 2: sel = c; break
if sel is None: continue
else:
sel = random.sample(av, nl)
if sel is None: continue
key = tuple(l for _, l, _ in sel)
if key in self.used: self.n_dups += 1; continue
tmp = defaultdict(int)
for _, _, s in sel:
if s in ("fixed", "synthetic"): tmp[s] += 1
ok = True
for s in ("fixed", "synthetic"):
if tmp[s] > 0:
c = self.fcap if s == "fixed" else self.scap
if self.usage[s] + tmp[s] > int(c * self.totals.get(s, 0)): ok = False; break
if not ok: continue
self.used.add(key); self.n_paras += 1; self.n_lines += nl
self.writers_used.add(wid)
si = set()
for p, _, s in sel:
self.usage[s] += 1; self.line_usage[p] += 1; si.add(s)
for s in si: self.src_para[s] += 1
return sel, wid
return None
def get_stats(self):
st = {}
for s in ["unique", "fixed", "synthetic"]:
t = self.totals.get(s, 0); u = self.usage.get(s, 0)
uu = sum(1 for p, c in self.line_usage.items() if c > 0 and self.line_src.get(p) == s)
st[s] = {'available': t, 'used': u, 'unique_used': uu,
'ratio': u / max(t, 1), 'utilisation': uu / max(t, 1) * 100}
return st
def load_image(path):
try: return Image.open(path).convert("RGB")
except: return None
def crop_whitespace(image, threshold=250, margin=5):
g = np.array(image.convert('L'))
m = g < threshold
r, c = np.any(m, axis=1), np.any(m, axis=0)
if not np.any(r) or not np.any(c): return image
ri, ci = np.where(r)[0], np.where(c)[0]
return image.crop((max(0, ci[0]-margin), max(0, ri[0]-margin),
min(image.width, ci[-1]+margin+1), min(image.height, ri[-1]+margin+1)))
def clean_left_edge(image, edge_px=8, wt=240, rs=10, vt=500):
a = np.array(image, dtype=np.float32)
_, w = a.shape[:2]
if w <= edge_px: return image
for col in range(min(edge_px, w)):
cd = a[:, col, :]
rm, gm, bm = np.mean(cd[:,0]), np.mean(cd[:,1]), np.mean(cd[:,2])
ov = (rm+gm+bm)/3; v = np.var(cd)
if (ov > wt or (rm > gm+rs and rm > bm+rs) or
(v < vt and ov > 180) or (rm > 200 and rm > gm and rm > bm and ov > 180)):
a[:, col, :] = 255.0
return Image.fromarray(a.astype(np.uint8))
def process_line(img, cw, do_crop, do_clean, wst, epx):
if do_crop: img = crop_whitespace(img, threshold=wst, margin=3)
if do_clean: img = clean_left_edge(img, edge_px=epx)
if img.width > cw:
s = cw / img.width
img = img.resize((cw, max(int(img.height * s), 20)), Image.Resampling.LANCZOS)
return img
def create_paragraph(imgs, sp, cw, ch, pad, content_w, do_crop, do_clean, wst, epx):
proc = [process_line(i, content_w, do_crop, do_clean, wst, epx)
for i in imgs if i.width > 0 and i.height > 0]
if not proc: return None, 0
th = pad*2 + sum(p.height for p in proc) + sp*(len(proc)-1)
ah = min(th, ch)
canvas = Image.new('RGB', (cw + pad*2, ah), (255, 255, 255))
y, used = pad, 0
for p in proc:
if y + p.height > ah - pad: break
x = max(cw + pad - p.width, pad)
canvas.paste(p, (x, y)); y += p.height + sp; used += 1
return canvas, used
def generate_split(gen, n, out_dir, name, cw, ch, pad, smin, smax,
do_crop, do_clean, wst, epx, fmt, gc_int):
os.makedirs(out_dir, exist_ok=True)
content_w = cw - pad*2
wc = defaultdict(int); ld = defaultdict(int)
cnt, err = 0, 0
pbar = tqdm(range(n), desc=f"Generating {name}")
for i in pbar:
try:
r = gen.get_paragraph_lines()
if r is None: err += 1; continue
sel, wid = r
imgs = [(load_image(p), l) for p, l, _ in sel]
imgs = [(im, l) for im, l in imgs if im is not None]
if not imgs: err += 1; continue
sp = random.randint(smin, smax)
pi, lu = create_paragraph([im for im, _ in imgs], sp, cw, ch, pad,
content_w, do_crop, do_clean, wst, epx)
if pi is None or lu == 0: err += 1; continue
cnt += 1; wc[wid] += 1; ld[lu] += 1
ext = {"TIFF": "tif", "PNG": "png", "JPEG": "jpg"}[fmt]
pi.save(os.path.join(out_dir, f"{wid}_para_{wc[wid]:04d}.{ext}"), fmt)
with open(os.path.join(out_dir, f"{wid}_para_{wc[wid]:04d}.txt"), "w", encoding="utf-8") as f:
f.write("\n".join(l for _, l in imgs[:lu]))
del pi
if (i+1) % gc_int == 0: gc.collect()
pbar.set_postfix({"saved": cnt, "err": err})
except Exception as e:
err += 1
if err < 10: print(f"\nError: {e}")
gc.collect()
gc.collect()
print(f" {name}: {cnt:,} saved, {err:,} errors")
return cnt, err, dict(ld)
def print_stats(name, gen, count, ld):
st = gen.get_stats()
tl = sum(k*v for k, v in ld.items())
tp = sum(ld.values())
print(f"\n {name}:")
print(f" Paragraphs: {count:,}, Writers: {len(gen.writers_used)}")
if tp > 0: print(f" Avg lines/para: {tl/tp:.2f}")
for s in ["unique", "fixed", "synthetic"]:
d = st[s]
if d['available'] > 0:
print(f" {s.capitalize():12s}: {d['ratio']:.2f}x reuse, "
f"{d['unique_used']:,}/{d['available']:,} ({d['utilisation']:.1f}%)")
print(f" Duplicates rejected: {gen.n_dups:,}")
def main():
args = parse_args()
random.seed(args.seed); np.random.seed(args.seed)
do_crop = args.crop_whitespace and not args.no_crop_whitespace
do_clean = args.clean_left_edge and not args.no_clean_left_edge
ts = int(args.dataset_size * args.train_ratio); vs = args.dataset_size - ts
print("\n" + "="*70)
print("SYNTHETIC PARAGRAPH GENERATOR")
print("="*70)
print(f"Size: {args.dataset_size:,} (train={ts:,}, val={vs:,})")
print("\n[1] Loading lines...")
ut = load_line_dataset(args.unique_train_dir, "Unique Train")
ft = load_line_dataset(args.fixed_train_dir, "Fixed Train")
st = load_line_dataset(args.synthetic_train_dir, "Synth Train")
uv = load_line_dataset(args.unique_val_dir, "Unique Val")
fv = load_line_dataset(args.fixed_val_dir, "Fixed Val")
sv = load_line_dataset(args.synthetic_val_dir, "Synth Val")
tt = {"unique": len(ut), "fixed": len(ft), "synthetic": len(st)}
vt = {"unique": len(uv), "fixed": len(fv), "synthetic": len(sv)}
print("\n[2] Verifying isolation...")
tp = set(os.path.abspath(p) for p, _ in ut+ft+st)
vp = set(os.path.abspath(p) for p, _ in uv+fv+sv)
ov = tp & vp
print(f" {'WARNING: '+str(len(ov))+' overlap!' if ov else 'Zero overlap confirmed'}")
del tp, vp
print("\n[3] Merging by writer...")
src_t = [("unique", ut)] + ([("fixed", ft)] if ft else []) + ([("synthetic", st)] if st else [])
src_v = [("unique", uv)] + ([("fixed", fv)] if fv else []) + ([("synthetic", sv)] if sv else [])
tm, _ = merge_lines_by_writer(src_t)
vm, _ = merge_lines_by_writer(src_v)
print(f" Train: {len(tm)} writers | Val: {len(vm)} writers")
tg = SingleWriterParagraphGenerator(tm, tt, args.train_fixed_cap, args.train_synthetic_cap,
args.min_lines, args.max_lines, args.max_attempts)
vg = SingleWriterParagraphGenerator(vm, vt, args.val_fixed_cap, args.val_synthetic_cap,
args.min_lines, args.max_lines, args.max_attempts)
td = os.path.join(args.output_dir, "Training")
vd = os.path.join(args.output_dir, "Validation")
print(f"\n[4] Generating training ({ts:,})...")
tc, te, tld = generate_split(tg, ts, td, "Training", args.canvas_width, args.canvas_height,
args.padding, args.spacing_min, args.spacing_max,
do_crop, do_clean, args.whitespace_threshold, args.edge_pixels,
args.output_format, args.gc_interval)
print(f"\n[5] Generating validation ({vs:,})...")
vc, ve, vld = generate_split(vg, vs, vd, "Validation", args.canvas_width, args.canvas_height,
args.padding, args.spacing_min, args.spacing_max,
do_crop, do_clean, args.whitespace_threshold, args.edge_pixels,
args.output_format, args.gc_interval)
print("\n" + "="*70)
print("COMPLETE")
print("="*70)
print(f" Total: {tc+vc:,} (train={tc:,}, val={vc:,}, errors={te+ve:,})")
print_stats("Training", tg, tc, tld)
print_stats("Validation", vg, vc, vld)
print(f"\n Output: {args.output_dir}")
print(f" Finished: {datetime.now():%Y-%m-%d %H:%M:%S}")
info = os.path.join(args.output_dir, "generation_info.txt")
with open(info, "w", encoding="utf-8") as f:
f.write(f"Generated: {datetime.now():%Y-%m-%d %H:%M:%S}\n")
f.write(f"Size: {args.dataset_size}, Train: {tc}, Val: {vc}\n")
f.write(f"Config: {vars(args)}\n")
for nm, g in [("Training", tg), ("Validation", vg)]:
s = g.get_stats(); f.write(f"\n{nm}:\n")
for src in ["unique","fixed","synthetic"]:
d = s[src]
if d['available'] > 0:
f.write(f" {src}: {d['used']:,}/{d['available']:,} ({d['ratio']:.2f}x)\n")
print(f" Info: {info}")
gc.collect()
if __name__ == "__main__":
main()