Skip to content

README.md

Latest commit

 

History

History
327 lines (241 loc) · 11.7 KB

README.md

File metadata and controls

327 lines (241 loc) · 11.7 KB
Raw

Fast 2DGS: Efficient Image Representation with Deep Gaussian Prior

As generative models become increasingly capable of producing high-fidelity visual content, the demand for efficient, interpretable, and editable image representations has grown substantially. Recent advances in 2D Gaussian Splatting (2DGS) have emerged as a promising solution, offering explicit control, high interpretability, and real-time rendering capabilities (>1000 FPS). However, high-quality 2DGS typically requires post-optimization. Existing methods adopt random or heuristics (e.g., gradient maps), which are often insensitive to image complexity and lead to slow convergence (>10s). More recent approaches introduce learnable networks to predict initial Gaussian configurations, but at the cost of increased computational and architectural complexity.

To bridge this gap, we present Fast-2DGS, a lightweight framework for efficient Gaussian image representation. Specifically, we introduce Deep Gaussian Prior, implemented as a conditional network to capture the spatial distribution of Gaussian primitives under different complexities. In addition, we propose an attribute regression network to predict dense Gaussian properties. Experiments demonstrate that this disentangled architecture achieves high-quality reconstruction in a single forward pass, followed by minimal fine-tuning. More importantly, our approach significantly reduces computational cost without compromising visual quality, bringing 2DGS closer to industry-ready deployment.

News

Thank you very much for following up on this work. We sincerely apologize for the delay due to our job shift.

This project is now managed by Aztech Labs. We will provide long-term maintenance for this project and will update industry-ready support in the next few months:

  1. Higher resolution support;
  2. End-to-end encoding architecture
  3. Efficient image file compression

Logs

[June 2026] We are working on

  • Entropy-aware encoding
  • File compression

[May 2026] Repo refresh

  • Inference API — Core logic moved to engine.py.
  • Training — Stage 1 & 2 training scripts shipped.
  • Benchmark — Paper-aligned evaluation.
  • Weights — Default heatmap checkpoint updated.

[March 2026] Our paper is presented in WACV 2026

Requirements

  • Linux / Windows, NVIDIA GPU
  • CUDA 12.6+ (matches torch==2.7.1+cu126)
  • Python 3.10+
  • C++ compiler + CUDA toolkit (to build gmod — Image-GS 2D rasterizer, renamed to avoid gsplat conflicts)

Installation

git clone https://github.com/Aztech-Lab/Fast-2DGS.git
cd Fast-2DGS

conda create -n 2dgs python=3.12 -y
conda activate 2dgs
pip install -r requirements.txt

# Install CUDA renderer (one-time)
# Windows:
powershell -ExecutionPolicy Bypass -File scripts/setup_gmod.ps1
# Linux:
bash scripts/setup_gmod.sh

python scripts/check_env.py

NOTE: We use the CUDA rasterizer from Image-GS, which is heavily optimized for 2DGS and is significantly different than classical gsplat (e.g., GaussianImage). To avoid conflicts, we rename it as gmod.

If install fails, ensure nvcc / CUDA path matches your PyTorch CUDA version (ours: torch==2.7.1+cu126).

Project Layout

Fast-2DGS/
├── engine.py              # Core API: encode / sample / rasterize / tune
├── inference.py           # CLI inference + batch runs
├── main_demo.py           # Inference quick start (recommended first run)
├── main_train_heat.py     # Stage 1 training quick start
├── main_train_feat.py     # Stage 2 training quick start
├── main_benchmark.py      # Paper-aligned dataset benchmark
├── benchmark.py           # Benchmark core (metrics + aggregation)
├── train_heatmap.py       # Stage 1: train HeatmapUNet (full script)
├── train_feature.py       # Stage 2: train GaussianUNet_Plus (full script)
├── tools.py               # Logging, heatmap, training helpers
├── dataset.py             # Image loader (center-crop by default)
├── models/GS_UNet.py      # HeatmapUNet, GaussianUNet, GaussianUNet_Plus
├── weights/               # Pretrained checkpoints
├── scripts/               # setup_gmod, check_env
└── old/                   # Archived legacy scripts

Quick Start

1. Python API demo (recommended)

python main_demo.py

Prints the API snippet, runs forward + fine-tune on assets/anime-1_2k.png, saves to outputs/demo/.

2. CLI inference

python inference.py --input assets/anime-1_2k.png --progress

Outputs go to outputs/infer/ (*_gt.png, *_pred_init.png, *_pred_tune.png, summary.json).

3. Batch inference (lightweight)

python inference.py --input_dir 2DGS_dataset/dataset/Kodak --progress

4. Benchmark

python main_benchmark.py
python main_benchmark.py --data_path 2DGS_dataset/dataset/Kodak --save_grid --progress

Reports init/tune PSNR, 1s/2s/5s PSNR, MS-SSIM, inference/tune/batch time, FPS → outputs/benchmark_kodak/summary.json.

Python API

from engine import Fast2DGEngine, load_image

engine = Fast2DGEngine(K=50_000, tune_steps=3000)
imgs = load_image("assets/anime-1_2k.png", image_size=512, crop=True, device=engine.device)

# Option A: step-by-step
pred, heatmap, params = engine.forward(imgs)       # encode -> sample -> rasterize
tune = engine.tune(imgs, params, tune_steps=3000)    # fine-tune

# Option B: one call
imgs, result = engine.run_image("assets/anime-1_2k.png")
print(result.psnr_init, result.psnr_tuned)

Key inference parameters

Parameter Default Description
--K 50000 Number of Gaussians
--tune_steps 3000 Fine-tuning steps (0 / --no_tune to skip)
--tune_lr 2e-3 AdamW learning rate for fit
--tune_weight_decay 0.05 AdamW weight decay
--no_xy_retain False Random init instead of predicted Gaussians
--image_size 512 Resize shorter side, then center crop
--crop / --no-crop crop Center crop (default, keeps aspect ratio)
--feat_plus True Use offset-refined GaussianUNet_Plus
--sampling multinomial multinomial or topk heatmap sampling
--heat_weight weights/smp_heat_div2k.pth Heatmap prior checkpoint
--feat_weight auto Attribute network checkpoint

Datasets

For training or batch benchmarks, clone 2DGS_dataset and download DIV2K HR splits:

git clone https://github.com/Aztech-Lab/2DGS_dataset.git
2DGS_dataset/
└── dataset/
    ├── Kodak/
    ├── DIV2K/
    │   ├── DIV2K_train_HR/    # 800 images (training)
    │   └── DIV2K_valid_HR/    # 100 images (validation)
    ├── ImageGS_anime/
    └── ImageGS_textures/

Point --data_path to any image folder. Training and inference use center crop by default (resize shorter side to 512, then crop — no aspect-ratio distortion).

Training

Two-stage pipeline:

Stage 1 (heatmap)          Stage 2 (feature)
─────────────────           ─────────────────
HeatmapUNet                 GaussianUNet_Plus (frozen heat)
  ↓ sample K xy               ↓ predict scale/color/rot/offset
  ↓ fit GT Gaussians          ↓ rasterize
  ↓ MSE vs blurred heatmap    ↓ L1+L2 recon loss

Stage 1 — Heatmap prior

# Quick start (prints API + runs 1-epoch Kodak demo)
python main_train_heat.py

# DIV2K full training
python train_heatmap.py \
  --data_path path/to/DIV2K_train_HR \
  --save_dir exp/heatmap_div2k \
  --init_weight weights/smp_heat_div2k.pth
from train_heatmap import build_parser, train

args = build_parser().parse_args([
    "--data_path", "path/to/DIV2K_train_HR",
    "--save_dir", "exp/heatmap_div2k",
    "--num_epochs", "500", "--lr", "5e-4",
])
train(args)
Parameter Default Description
--lr 5e-4 AdamW learning rate
--num_epochs 500 Training epochs
--batch_size 8 Batch size
--k_min / --k_max 10000 / 100000 Random K per batch
--gt_steps 200 Inner-loop GT Gaussian fit steps
--init_weight (none) Optional checkpoint to fine-tune from
--crop True Center crop preprocessing

Outputs in exp/heatmap*/: heat_best.pth, heat_last.pth, plot_train.png, results/epoch_*.png, train_log.csv.

Stage 2 — Attribute network

# Quick start (prints API + runs 1-epoch Kodak demo)
python main_train_feat.py

# Full training
python train_feature.py \
  --data_path path/to/DIV2K_train_HR \
  --heat_weight weights/smp_heat_div2k.pth \
  --save_dir exp/feature_div2k
from train_feature import build_parser, train

args = build_parser().parse_args([
    "--data_path", "path/to/DIV2K_train_HR",
    "--heat_weight", "weights/smp_heat_div2k.pth",
    "--save_dir", "exp/feature_div2k",
    "--num_epochs", "500", "--lr", "1e-3",
])
train(args)
Parameter Default Description
--heat_weight weights/smp_heat_div2k.pth Frozen heatmap checkpoint
--lr 1e-3 AdamW learning rate
--num_epochs 500 Training epochs
--loss_rec l1+l2 Reconstruction loss
--feat_plus True Use GaussianUNet_Plus with xy offset

Outputs in exp/feature*/: feat_best.pth (best PSNR), plot_train.png, results/epoch_*.png, train_log.csv.

Pretrained weights

File Description
weights/smp_heat_div2k.pth Heatmap prior (DIV2K fine-tuned)
weights/smp_feat_best_psnr_26_plus.pth Feature net with xy offset (default)
weights/smp_feat_best_psnr_26.pth Feature net without offset

Framework

Results

Gaussian Initialization Comparison

Impact of Gaussians

Additional Results

Acknowledgements

We sincerely appreciate the Image-GS team for providing the 2DGS rendering core and for sharing their high-quality datasets, and we thank Instant-GI team for their great work and deep inspiration. Moreover, we thank the GaussianImage team for their foundation work at this domain.

Citation

If you find this project helpful to your research, please consider citing:

@InProceedings{Wang_2026_WACV,
    author    = {Wang, Hao and Bastola, Ashish and Zhou, Chaoyi and Zhu, Wenhui and Chen, Xiwen and Dong, Xuanzhao and Huang, Siyu and Razi, Abolfazl},
    title     = {Fast 2DGS: Efficient Image Representation with Deep Gaussian Prior},
    booktitle = {Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision (WACV) Workshops},
    month     = {March},
    year      = {2026},
    pages     = {1184-1193}
}