calligraphy-repair-cgan

🖋️ Calligraphy Repair cGAN

A two-stage system for repairing damaged handwriting and calligraphy using combined pathfinding + conditional GAN approach.

Stage 1 uses a deterministic A*/Bezier pathfinding algorithm to repair large structural gaps in strokes. Stage 2 uses a conditional GAN to re-apply the artistic style, producing visually coherent results.

Architecture Overview

┌─────────────────────────────────────────────────────────────┐
│                    INPUT: Damaged Calligraphy                │
│                       (RGB, H×W×3)                          │
└────────────────────────┬────────────────────────────────────┘
                         │
         ┌───────────────▼───────────────┐
         │      STAGE 1: PATHFINDING     │
         │   (Deterministic Algorithm)    │
         │                               │
         │  1. Binarize → ink mask       │
         │  2. Skeletonize (Zhang-Suen)  │
         │  3. Find gap endpoints        │
         │  4. A*/Bezier gap bridging    │
         │  5. Variable-width rendering  │
         └───────┬───────────┬───────────┘
                 │           │
         stroke_mask    gap_mask
           (1,H,W)     (1,H,W)
                 │           │
         ┌───────▼───────────▼───────────┐
         │      STAGE 2: cGAN REFINEMENT │
         │    (Learned Style Transfer)    │
         │                               │
         │  Generator Input (5ch):       │
         │  [damaged_RGB(3) +            │
         │   stroke_mask(1) +            │
         │   gap_mask(1)]                │
         │                               │
         │  ┌─────────────────────────┐  │
         │  │   U-Net Generator       │  │
         │  │   + Gated Convolutions  │  │
         │  │   + 8 Dilated ResBlocks │  │
         │  │   + Skip Connections    │  │
         │  └─────────────────────────┘  │
         │                               │
         │  ┌─────────────────────────┐  │
         │  │  70×70 SN-PatchGAN      │  │
         │  │  Discriminator          │  │
         │  │  (Spectral Normalized)  │  │
         │  └─────────────────────────┘  │
         └───────────────┬───────────────┘
                         │
         ┌───────────────▼───────────────┐
         │    OUTPUT: Repaired Image      │
         │        (RGB, H×W×3)           │
         └───────────────────────────────┘

Key Features

• 🔍 Pathfinding Stage: A*/Bezier curve gap bridging with direction-aware endpoint matching, variable stroke thickness estimation, and smooth calligraphic curve generation
• 🎨 cGAN Stage: EdgeConnect-inspired architecture with gated convolutions, dilated residual blocks, and multi-loss training (adversarial + L1 + perceptual VGG + style Gram + feature matching)
• 📊 Synthetic Data Pipeline: Generates training pairs from clean calligraphy with realistic damage (erosion, gaps, fading, bleeding, stains, scratches)
• ⚡ Flexible: Works with any writing system (Latin, Chinese, Arabic, Japanese, etc.)

Installation

pip install -r requirements.txt

Quick Start

1. Generate Training Data

# Generate 5000 synthetic training pairs
python damage_generator.py \
    --output_dir data \
    --num_train 5000 \
    --num_val 500 \
    --image_size 256

# Or use your own clean calligraphy images:
python damage_generator.py \
    --output_dir data \
    --source_dir /path/to/clean/calligraphy/ \
    --num_train 5000

This creates:

data/
├── train/
│   ├── clean/      # Ground truth images
│   ├── damaged/    # Synthetically damaged images
│   └── mask/       # Damage location masks
└── val/
    ├── clean/
    ├── damaged/
    └── mask/

2. Train the Model

# Full training (recommended)
python train.py \
    --data_dir data \
    --epochs 200 \
    --batch_size 8 \
    --image_size 256 \
    --generator_type unet \
    --gan_type lsgan \
    --lr_g 1e-4 \
    --lr_d 1e-4

# Quick test run
python train.py \
    --data_dir data \
    --epochs 5 \
    --batch_size 4 \
    --generate_data \
    --num_train 100 \
    --num_val 20

# With on-the-fly damage generation (only needs clean/ directory)
python train.py \
    --data_dir data \
    --on_the_fly_damage \
    --epochs 200

# Resume from checkpoint
python train.py \
    --data_dir data \
    --resume checkpoints/best_model.pth \
    --epochs 300

3. Repair Damaged Images

# With trained GAN (best quality)
python inference.py \
    --input damaged_calligraphy.png \
    --output repaired.png \
    --checkpoint checkpoints/generator_best.pth

# Pathfinding only (no GAN needed, instant)
python inference.py \
    --input damaged_calligraphy.png \
    --output repaired.png \
    --pathfinding_only

# Batch repair a directory
python inference.py \
    --input_dir damaged_images/ \
    --output_dir repaired_images/ \
    --checkpoint checkpoints/generator_best.pth

# Save all intermediate stages for visualization
python inference.py \
    --input damaged_calligraphy.png \
    --output repaired.png \
    --checkpoint checkpoints/generator_best.pth \
    --save_stages

Training Recipe

Based on literature from EdgeConnect, pix2pix, DE-GAN, and pix2pixHD:

Loss Functions

The combined generator loss:

L_G = λ_adv · L_adversarial          (fool the discriminator)
    + λ_L1 · L_L1                     (pixel-level reconstruction)
    + λ_FM · L_feature_matching        (stabilize training via D features)
    + λ_perc · L_perceptual            (VGG relu1_2, relu2_2, relu3_3, relu4_3)
    + λ_style · L_style                (Gram matrices for texture matching)
    + λ_mask · L_masked_L1             (focus on damaged regions)

The discriminator loss:

L_D = L_adversarial_D + λ_R1 · L_R1_regularization

Architecture Details

Generator: U-Net with Dilated Residual Bottleneck

Encoder:
  [GatedConv 5→64, 7×7, stride 1]          # Level 0
  [GatedConv 64→128, 4×4, stride 2]         # Level 1 (↓2×)
  [GatedConv 128→256, 4×4, stride 2]        # Level 2 (↓2×)

Bottleneck:
  8× [DilatedResBlock 256→256, dilation=2]   # ~200px receptive field

Decoder:
  [ConvTranspose 512→128, 4×4, stride 2]    # Skip from Level 2 (↑2×)
  [ConvTranspose 256→64, 4×4, stride 2]     # Skip from Level 1 (↑2×)
  [Conv 128→3, 7×7, Tanh]                   # Skip from Level 0

All layers: InstanceNorm + ReLU (encoder: LeakyReLU)
Input: 5 channels (damaged_RGB + stroke_mask + gap_mask)
Output: 3 channels (repaired RGB in [-1, 1])

Discriminator: 70×70 SN-PatchGAN

C64(no norm) → C128 → C256 → C512 → Conv→1
All convs: 4×4, spectral normalized
InstanceNorm on layers 2-4
LeakyReLU(0.2) throughout
Input: 8 channels (damaged + output + masks)
Output: H/16 × W/16 patch predictions

Pathfinding Algorithm

The deterministic Stage 1 pipeline:

1. Binarization: Otsu's adaptive thresholding + morphological cleanup
2. Skeletonization: Zhang-Suen thinning to 1-pixel-wide skeleton
3. Thickness Estimation: Distance transform to measure local stroke width
4. Endpoint Detection: Find degree-1 nodes (gap openings) via 3×3 convolution kernel
5. Direction Analysis: Trace back along skeleton to compute tangent direction at each endpoint
6. Endpoint Matching: Score pairs by distance + direction alignment + collinearity; greedy matching
7. Gap Bridging:
   • A* mode: Cost = distance + direction_continuity + curvature_penalty + ink_proximity
   • Bezier mode: Cubic Bezier with control points guided by endpoint tangents
8. Stroke Rendering: Variable-width circular brush matching estimated local thickness

Project Structure

├── pathfinding.py          # Stage 1: Deterministic gap repair
├── damage_generator.py     # Synthetic training data generation
├── models.py               # cGAN architecture (Generator + Discriminator)
├── losses.py               # Loss functions + metrics
├── dataset.py              # Dataset loader with pathfinding integration
├── train.py                # Training pipeline
├── inference.py            # Inference / repair script
├── requirements.txt        # Dependencies
└── README.md               # This file

Using Your Own Data

Option A: Clean calligraphy images only (recommended)

Place clean calligraphy images in a directory and the system will synthetically damage them:

# Generate paired data from your clean images
python damage_generator.py \
    --source_dir /path/to/your/clean/calligraphy/ \
    --output_dir data \
    --num_train 5000

# Or train with on-the-fly damage
python train.py \
    --data_dir data \
    --on_the_fly_damage

Option B: Pre-paired damaged/clean images

Organize your data as:

data/
├── train/
│   ├── clean/001.png, 002.png, ...
│   └── damaged/001.png, 002.png, ...
└── val/
    ├── clean/001.png, 002.png, ...
    └── damaged/001.png, 002.png, ...

Filenames must match between clean/ and damaged/.

Tips for Best Results

1. More data = better results: Aim for 5000+ training pairs minimum
2. Use your own calligraphy: The model learns styles from training data; train on the style you want to repair
3. Pathfinding as baseline: Even without the GAN, Stage 1 gives usable structural repairs
4. Monitor training: Check checkpoints/samples/ for visual progress
5. Adjust gap distance: --max_gap_distance controls how large a gap the pathfinder will attempt to bridge
6. GPU recommended: Training takes ~4-8 hours on a single GPU (RTX 3080 or better)

References

• EdgeConnect — Two-stage edge + completion GAN
• pix2pix — Conditional adversarial image-to-image translation
• DE-GAN — Document enhancement with conditional GAN
• pix2pixHD — High-resolution synthesis with feature matching
• DeepFillv2 — Gated convolutions for inpainting
• DiffHDR — Historical document repair with masked perceptual loss

License

MIT