Scripts/generate_paragraphs.py

"""

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()