UniSketch

This repository contains the official implementation and dataset for the paper:

"UniSketch: A Unified Multi-Task Framework for Parametric Sketch Generation and Constraint Prediction"

1. Installation

conda create -n unisketch python=3.10
conda activate unisketch
pip install -r requirements.txt

2. Dataset

We support two different ways to store and read the dataset.

1. Using a split.json file to define the train, validation, and test splits.
2. Using separate folders named train, validation, and test that directly contain the corresponding data.

Our experiments use three types of data formats, each stored in a separate folder (sequence, image_clean, and image_noisy) under the main data directory.

• sequence: sketch sequences
• image_clean: original (clean) sketch images
• image_noisy: noisy sketch images

You can use myHDF5 to view the .h5-formatted sequence data directly in your browser.

The full UniSketch dataset, which includes sketch sequences and their corresponding original images, can be downloaded from unisketch-full. Note that noisy images are not yet included in the full dataset.

A 1,000,000-sample random subset of the dataset used for experiments is available at unisketch-subset.

Additionally, the Vitruvion dataset, used for comparative experiments with the Vitruvion, SketchGraphs and UniSketch networks, can be downloaded from vitruvion.

3. Usage

3.1 Training

python train.py

The default configuration is as follows:

• data_loader = json
• train_type = seq_only
• use_noise is disabled (set to False)

These and other hyperparameters can be adjusted for different experimental setups. Please refer to training/config.py for the complete configuration options.

3.2 Testing

python test.py --ckpt path/to/checkpoint.pth --mode tf

The default configuration is as follows:

• test_type = seq_only

The --mode argument controls the generation strategy and supports the following options:

• tf: teacher forcing
• ar: autoregressive (constraint prediction)
• ug: unconditional generation
• sc: sketch completion

Please refer to training/config.py for more details.

This script generates and saves sequences for each sketch in an .h5 file:

• For tf and ar modes, both the ground truth (gt_vec) and predicted sequences (out_vec) are saved.
• For ug and sc modes, only the generated sequences (gen_vec) are saved.

3.2 Evaluation

python evaluate.py --test_path path/to/evaluation --mode tf

The --mode argument specifies the evaluation strategy:

• tf: teacher forcing, evaluates primitive-/constraint-/parameter-level accuracy

• ar: autoregressive, evaluates constraint-level precision, recall, and F1 score

For more details, please refer to evaluate.py.