HAWP (Holistically-Attracted Wireframe Parsing)

repository·main·Indexed 18 days ago

https://github.com/cherubicxn/hawp

A framework for wireframe parsing that provides robust line segment representations for complex structures in images. It includes HAWPv2, a fully-supervised parser, and HAWPv3, a self-supervised parser designed for out-of-distribution images. The library provides modules for training (hawp.fsl.train, hawp.ssl.train), inference (hawp.ssl.predict), and benchmarking (hawp.fsl.benchmark), along with a multi-step training pipeline for HAWPv3 involving synthetic learning and homographic adaptation.

Tokens
3.3K
Snippets
13
Records
15
Agent score
61%

What's inside HAWP

  1. Overview of HAWP versions

    main

    The project provides several versions of the Holistically-Attracted Wireframe Parsing model:

    • HAWPv2: A state-of-the-art fully-supervised wireframe parser. See docs/HAWPv2.md for details.
    • HAWPv3: A state-of-the-art self-supervised wireframe parser. It is highly effective for out-of-distribution images. See docs/HAWPv3.md for details and docs/HAWPv3.train.md for the training recipe.
  2. Locate generated pseudo wireframe data

    main

    After completing the Homographic Adaptation step, the generated wireframe and auxiliary files are stored in a directory named after the training round: data-ssl/{name} (where {name} is the --name provided during the synthetic training phase).

    File Naming Convention:

    • Individual wireframe files follow the pattern: {hash}-model-{epoch:05d}.
    • The generated dataset configuration file is located at: data-ssl/{name}/{hash}-model-{epoch:05d}.yaml.
  3. Download HAWPv2 training and testing datasets

    main

    The training and testing data for HAWPv2 (including the Wireframe and YorkUrban datasets) can be downloaded via Google Drive.

    You can use gdown to download and unzip the data directly in your terminal:

    gdown 134L-u9pgGtnzw0auPv8ykHqMjjZ2claO
    unzip data.zip
  4. Perform Step 2: Learning from Real-World images

    main

    Once pseudo wireframe labels are generated, you can train HAWPv3 on real-world images using the generated YAML configuration file as the --datacfg.

    Workflow:

    1. Train on real-world images using the pseudo labels.
    2. (Optional) Run homographic adaptation again using the new model checkpoints to generate improved labels.
    3. (Optional) Fine-tune the model further to improve repeatability.
    python -m hawp.ssl.train  --datacfg data-ssl/export_datasets/{name}/{hash}-model-00010.yaml --modelcfg hawp/ssl/config/hawpv3.yaml --base-lr 0.0004 --epochs 30  --name hawpv3-round1 --batch-size 6
  5. Run inference with HAWPv3 to obtain wireframes

    main

    Use the hawp.ssl.predict module to extract wireframes from images using a HAWPv3 checkpoint. You can use either the model trained on the Wireframe dataset or the one trained on ImageNet.

    Available Checkpoints:

    • hawpv3-fdc5487a.pth: Trained on Wireframe dataset images.
    • hawpv3-imagenet-03a84.pth: Trained on 100k ImageNet images.

    Arguments:

    • --ckpt: Path to the model checkpoint file.
    • --threshold: Confidence threshold for wireframe detection (e.g., 0.05).
    • --img: Path to the input image or a pattern of images (e.g., *.png).
    • --saveto: Directory to save the output results.
    • --ext: File extension for the saved output (e.g., png).
    # Example: Batch processing images from a dataset
    python -m hawp.ssl.predict --ckpt checkpoints/hawpv3-imagenet-03a84.pth  \\
      --threshold 0.05  \\
      --img ~/datasets/DTU/scan24/image/*.png \\
      --saveto docs/figures/dtu-24 --ext png \
  6. Perform Step 1: Homographic Adaptation for Pseudo Wireframe Generation

    main

    This step generates pseudo wireframe labels using the model trained in the previous phase. You can choose between single-image mode (lower GPU memory footprint) or batch mode (faster processing).

    Single-image mode

    Use this mode to minimize GPU memory usage.

    Batch mode

    Use this mode for faster processing of large datasets. Note that batch processing requires significant GPU memory (e.g., a batch size of 16 may require ~40GB of VRAM on an NVIDIA A6000).

    Arguments:

    • --metarch: The architecture method (e.g., HAWP-heatmap).
    • --datacfg: Path to the export configuration YAML.
    • --workdir: The directory where results are stored.
    • --epoch: The specific epoch of the model to use.
    • --modelcfg: Path to the model configuration file.
    • --min_score / --min-score: Minimum score threshold for label generation.
    • --batch-size: (Batch mode only) Number of images per batch.
    # Single-image mode
    python -m hawp.ssl.homoadp --metarch HAWP-heatmap \
        --datacfg hawp/ssl/config/export/wireframe-10iters.yaml \
        --workdir exp-ssl/hawpv3-round0 \
        --epoch 10 \
        --modelcfg exp-ssl/hawpv3-round0/model.yaml \
        --min_score 0.75 
    
    # Batch mode
    python -m hawp.ssl.homoadp-bm --metarch HAWP-heatmap \
        --datacfg hawp/ssl/config/exports/wireframe-10iters.yaml \
        --workdir exp-ssl/hawpv3-round0 \
        --epoch 10 \
        --modelcfg exp-ssl/hawpv3-round0/model.yaml \
        --min-score 0.75   --batch-size=16
  7. Train HAWPv2 using hawp.fsl.train

    main

    To train HAWPv2 on a dataset (e.g., Wireframe), run the hawp.fsl.train module. You must provide a path to a configuration file. You can also specify a directory for logs using --logdir.

    python -m hawp.fsl.train configs/hawpv2.yaml --logdir outputs
  8. Perform Step 0: Synthetic learning

    main

    The first phase of HAWPv3 training involves synthetic learning. Use the hawp.ssl.train module with a synthetic dataset configuration and the HAWPv3 model configuration.

    Arguments:

    • --datacfg: Path to the synthetic dataset YAML configuration.
    • --modelcfg: Path to the HAWPv3 model YAML configuration.
    • --base-lr: Base learning rate.
    • --epochs: Number of training epochs.
    • --batch-size: Number of samples per batch.
    • --name: The name for this training round (used for organizing output).
    python -m hawp.ssl.train \
        --datacfg hawp/ssl/config/synthetic_dataset.yaml \
        --modelcfg hawp/ssl/config/hawpv3.yaml \
        --base-lr 0.0004 \
        --epochs 10  \
        --batch-size 6 \
        --name hawpv3-round0
  9. Install HAWP via Anaconda

    main

    Follow these steps to set up the HAWP environment using Conda. Note that you must install ninja-build on your system first.

    1. Clone the repository.
    2. Install ninja-build via sudo apt install ninja-build.
    3. Create and activate a Python 3.9 environment.
    4. Install the package in editable mode.
    5. Install PyTorch (ensure the CUDA version matches your machine) and other dependencies.
    6. Verify CUDA support and download official checkpoints.
    # Clone the code repo
    git clone https://github.com/cherubicXN/hawp.git
    
    # Install ninja-build
    sudo apt install ninja-build
    
    # Create conda environment
    conda create -n hawp python==3.9
    conda activate hawp
    pip install -e .
    
    # Install pytorch (adjust CUDA version if necessary)
    pip install torch==1.12.0+cu116 torchvision==0.13.0+cu116 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu116 
    
    # Install other dependencies
    pip install -r requirement.txt
    
    # Verify installation
    python -c "import torch; print(torch.cuda.is_available())"
    
    # Download checkpoints
    sh downloads.sh
  10. Evaluate HAWPv2 models on Wireframe or YorkUrban datasets

    main

    To evaluate the official HAWPv2 model, use the hawp.fsl.benchmark module. You must provide a configuration file and a checkpoint file. You can specify the dataset using the --dataset flag.

    Note: The jhm (Joint Heatmap) parameter can be adjusted to change evaluation sensitivity. The default setting is 0.008.

    # Evaluation on the Wireframe dataset
    python -m hawp.fsl.benchmark configs/hawpv2.yaml \
      --ckpt checkpoints/hawpv2-edb9b23f.pth \
      --dataset wireframe
    
    # Evaluation on the YorkUrban dataset
    python -m hawp.fsl.benchmark configs/hawpv2.yaml \
      --ckpt checkpoints/hawpv2-edb9b23f.pth \
      --dataset york
  11. Run HAWPv3 inference on images

    main

    You can perform wireframe parsing on a set of images using the hawp.ssl.predict module. This is particularly useful for HAWPv3, which is designed to handle out-of-distribution images effectively.

    To run inference, use the following command structure, ensuring you provide a valid checkpoint path and the directory of your input images.

    python -m hawp.ssl.predict --ckpt checkpoints/hawpv3-imagenet-03a84.pth  \    --threshold 0.05  \    --img ~/datasets/DTU/scan24/image/*.png \    --saveto docs/figures/dtu-24 --ext png \