MSEmbGAN: Multi-Stitch Embroidery Synthesis

PyTorch implementation of "MSEmbGAN: Multi-Stitch Embroidery Synthesis via Region-Aware Texture Generation" (IEEE TVCG 2024, PubMed 39167500).

What This Does

Takes a flat design image (e.g., a flower drawing) and synthesizes a photorealistic multi-stitch embroidery rendering with automatic stitch type assignment:

• Satin stitch — narrow regions (needle drops on both sides, snakelike advance)
• Tatami stitch — wide regions (brick-like offset rows, resembling tatami weaving)
• Flat stitch — very small regions (uniform fill without directional variation)

Architecture

Input Content Image (RGB, 256x256)
    ↓
Stitch Classifier — Quick-shift segmentation + morphological erosion
    ↓
Region-Aware Texture Generation Network (WGAN, Lab L channel)
    ↓
Colorization Network (Lab ab channels)
    ↓
Final Colorized Embroidery Image

Files

• msembgan/modules.py — Custom layers (ReLUConv2d, INSResBlock, MisINSResBlock, etc.)
• msembgan/networks.py — All encoders, generators, discriminators from Supplemental Table 1
• msembgan/stitch_classifier.py — Quick-shift + erosion stitch classification + latent generator
• msembgan/model.py — Full MSEmbGAN assembly (RTGN + CN)
• msembgan/losses.py — All losses: KL, L1, WGAN, color consistency
• msembgan/dataset.py — Dataset loaders for aligned/unaligned data
• msembgan/train.py — 2-step training pipeline
• msembgan/inference.py — Inference script

Installation

pip install torch torchvision numpy pillow scikit-image scipy tqdm

Training

python -m msembgan.train \
    --data_dir ./data \
    --epochs 800 --batch_size 2 --lr 0.0002 \
    --step1_epochs 200 --step2_epochs 500

Inference

python -m msembgan.inference \
    --model ./checkpoints/final.pth \
    --input ./image.jpg --output ./embroidery.png

Dataset

Multi-Stitch Embroidery Dataset: https://drive.google.com/drive/folders/1VEg01D0vmO82oAYHr9D7KzDUCRza-DmU (Email fangfei@wtu.edu.cn if unavailable)

Paper

• PDF: https://github.com/CGVSL-NCKU/TVCG_MSEmbGAN/raw/HEAD/MSEmbGAN_paper.pdf
• DOI: 10.1109/TVCG.2024.3447351

License

MIT — clean-room implementation from the paper's architecture tables and equations.