Skip to content
BranchesTags

Repository files navigation

FlashImgs

Fast 2D Gaussian Splatting for image representation. Fit any image as a set of 2D Gaussians — train in seconds, render at 12,000+ FPS.

Installation

Requires Python 3.12+, CUDA 12.x, and uv.

git clone https://github.com/OpsiClear/flashimgs-clean.git
cd flashimgs-clean

# Create virtual environment
uv venv --python 3.12

# 1. Install PyTorch with CUDA (must come first)
uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu128

# 2. Install Python dependencies
uv pip install -e .

# 3. Build in-tree CUDA extensions from source (requires torch in the env)
uv pip install ./fussim/ --no-build-isolation
uv pip install ./simple_sum_backend/ --no-build-isolation

Note: If nvcc is not at /usr/local/cuda/bin/nvcc, set CUDA_HOME before step 3:

export CUDA_HOME=/usr/local/cuda-12.8  # adjust to your CUDA path

Verify:

uv run flashimgs info kodim01

Kodak test images are downloaded automatically on first use.

After uv pip install -e ., uv run flashimgs ... is the supported console entry point for image fitting.

Windows: Install Visual Studio Build Tools with "Desktop development with C++". Make sure cl.exe is on PATH before Git's link.exe.

Python 3.12 note: if a CUDA extension fails to build with AttributeError: module 'pkgutil' has no attribute 'ImpImporter', your environment has an old setuptools. This repo now requires setuptools>=70; upgrade it in the venv before rebuilding the CUDA extensions.

Quick Start

# Train on a Kodak image (auto-selects ~50K Gaussians)
uv run flashimgs train kodim01

# Train on your own image
uv run flashimgs train path/to/photo.png

# Best quality: overcomplete train -> merge -> finetune (~40 dB)
uv run flashimgs merge kodim01

# Benchmark on 6 test images
uv run flashimgs bench

# Launch live training UI in browser
uv run flashimgs viz

Training Features

  • Auto Gaussians — count scales with calibrated scene complexity (gradients, edges, Laplacian texture, entropy, and color variance). Typical outputs stay around ~3-13% of pixels, with a size-dependent floor for small images.
  • Density-weighted init — positions sampled proportional to gradient magnitude + local color variance. More Gaussians in detailed regions.
  • Adaptive density control (ADC) — every 300 steps from step 150, dead/out-of-bounds/bloated Gaussians are relocated to high-error pixels.
  • Asymmetric LR decay — position/color decay fast (0.16^x), scale/rotation decay slow (0.50^x). Scale/rotation have flatter loss landscapes.
  • Fixed simple-sum LR defaults — train/bench/export default to an 11x LR multiplier tuned for the simple-sum branch.
  • Color resampling — 10% GT color blend every 20 steps in first half. Fixes stale colors when Gaussians migrate.
  • Divergence recovery — parameters checkpointed every 400 steps. On 10x loss spike, rollback to last good state with halved LR instead of continuing from corrupted parameters.

Commands

train — Fit Gaussians to an image

uv run flashimgs train [IMAGE] [-n GAUSSIANS] [-s STEPS] [--lr LR] [--scale S] [--loss LOSS]
Option Default Description
IMAGE kodim01 Image name (e.g. kodim07) or file path
-n auto 0 = auto (~3-13% of pixels based on calibrated complexity, with a small size-dependent floor)
-s 2800 Training steps
--lr 11 LR multiplier
--scale 1.5 Initial Gaussian size in pixels
--loss l2 Loss function: l2, l1, l2+ssim, l1+ssim
--init adaptive Init: adaptive (distortion-ranked quadtree)

merge — Best quality pipeline

Trains more Gaussians than needed, merges similar ones down, then finetunes. Three clean phases:

Phase 1: Train 80K Gaussians (auto scene-aware steps, 11x LR)
Phase 2: Merge 80K -> 50K via moment-matching
Phase 3: Finetune 50K Gaussians (auto scene-aware steps, 6x LR)

uv run flashimgs merge [IMAGE] [--target N] [--ratio R] [--steps-pre N] [--steps-post N]
Option Default Description
--target auto Final Gaussian count
--ratio 1.4 Overcomplete ratio (start with 1.4x target)
--steps-pre auto 0 = auto scene-aware pre-merge schedule
--steps-post auto 0 = auto scene-aware post-merge schedule
--lr-pre 11 LR multiplier for overcomplete phase
--lr-post 6 LR multiplier for finetune phase

bench — Benchmark speed and quality

uv run flashimgs bench              # all 6 test images
uv run flashimgs bench kodim07      # single image

Output:

Image           N  Steps    Time     PSNR    SSIM
-------------------------------------------------------
kodim01     51,000   2000    2.1s  37.42dB  0.9821
kodim02     51,000   2000    2.0s  39.15dB  0.9887
...

info — Image info

uv run flashimgs info kodim01

viz — Live training visualizer

uv run flashimgs viz

Opens a browser UI at http://localhost:8767 with 4-panel view (render, ground truth, Gaussians, error map), PSNR chart, and 9 Gaussian visualization modes. Works over SSH — access via http://<server-ip>:8767.

The k-means clustering visualization modes additionally require flash-kmeans (vendored as a git submodule). The viz server runs without it — those modes simply report an error if used. To enable them: git submodule update --init then uv pip install ./third_party/flash-kmeans.

render — Render from checkpoint

uv run flashimgs render path/to/checkpoint.pt [--height 2048]

Python API

After uv pip install -e ., FlashImgs is importable as a package — other repos can drop import flashimgs and skip the CLI entirely. Two layers are exposed:

One-shot recipes (mirror the CLI)

import flashimgs

# Fit Gaussians to an image on disk (Kodak name or path)
result = flashimgs.train("kodim01", steps=2000)
print(result.psnr, result.ssim, result.elapsed_s)
print(result.params["xy"].shape)        # (N, 2) normalized coords

# Same recipe but from an in-memory CHW tensor (no disk I/O)
result = flashimgs.train_tensor(image_chw, gaussians=8000, steps=1500)

# Best-quality overcomplete -> merge -> finetune
result = flashimgs.merge("photo.png", target=8000)

Result carries the four learnable parameter tensors (xy, scale, rot, feat) with the gate already folded into feat, plus PSNR / SSIM / elapsed seconds and the source image's resolution. Pass return_session=True to also get back the live Session for further training or rendering.

Lifecycle / building-block API

flashimgs.Session exposes each phase of training individually — useful for custom schedules, NeRF priors, or any setting where the canned recipes are the wrong shape.

import flashimgs

s = flashimgs.Session.from_image("kodim01", gaussians=8000)
# or:  s = flashimgs.Session.from_tensor(img_chw, gaussians=8000)

# Custom schedule
s.optimize_steps(800)                  # fixed-LR-schedule training
moved = s.densify()                    # one ADC pass; returns relocations
s.add_gaussians(1500)                  # error-guided addition
s.optimize_steps(600)
s.merge_to(6000)                       # merge down to a target count
s.reset_optimizer(lr_mult=0.1)
s.optimize_steps(1000)                 # finetune at 0.1x base LR

# Eval / render / export
print(s.evaluate())                    # (psnr, ssim)
print(s.evaluate_full())               # + lpips / flip / msssim
img = s.render(height=2048)            # CHW tensor, any resolution
times = s.benchmark_render(num_reps=50)
s.materialize_gate()
s.save_checkpoint("ckpt.pt")
s.export_splat2d("out.splat2d")        # for the web viewer

Helpers

flashimgs.auto_gaussians("kodim01")          # complexity-aware budget
flashimgs.image_info("kodim01")              # height / width / suggested count
flashimgs.parse_loss("l2+ssim")              # → (l1, l2, ssim) ratios
cfg = flashimgs.default_config(gaussians=10_000, log_dir=None)
s = flashimgs.Session.from_config(cfg)       # escape hatch for niche knobs

For low-level operations the CUDA library still exposes merge_gaussians, project_gaussians_2d_scale_rot, and rasterize_gaussians_simple_sum directly; see simple_sum_backend/ for the tensor-level entry points.

Web viewer

Export a trained model for browser rendering and standard 3DGS viewers:

uv run flashimgs export kodim01
# → web/kodim01.splat2d (exact FlashImgs web format)
# → web/kodim01.ply     (planar standard 3DGS compatibility export)

Masked fits treat mask edges that touch the image frame as real mask boundaries by default, adding virtual outside-mask padding during training so splats are penalized for bleeding past the border. Use --no-mask-border-padding to restore the old flush-to-image-edge behavior.

The PLY export places the image on a flat XY plane with image height normalized to 1 world unit, stores 3DGS-compatible SH DC color / opacity / log-scale / quaternion fields, and includes fi_* attributes for exact 2D parameter recovery. The default --ply-y-axis down preserves image-space vertical orientation in common 3DGS viewers; use --ply-y-axis up if your viewer expects positive world Y to point toward the image top. The PLY export uses factor_rgb, factoring FlashImgs' additive RGB into per-splat alpha and color so alpha * color matches the 2D amplitude before standard 3DGS compositing. Use .splat2d for exact FlashImgs rendering; use .ply for standard 3DGS viewers. Tune --ply-opacity-scale, --ply-max-opacity, --ply-scale-mult, and --ply-thickness for a specific viewer.

Render in any web page:

<script type="module" src="flashimgs-viewer.js"></script>
<flashimgs-viewer src="kodim01.splat2d"></flashimgs-viewer>

The <flashimgs-viewer> component is a single self-contained JS module (web/flashimgs-viewer.js) with two backends:

Backend Render time Notes
WebGPU ~4 ms (250 FPS) Tile-binned compute, no CPU readback
CPU ~240 ms Web Workers + fast exp table, fallback

Serve with any HTTP server (python -m http.server 8080).

Configuration

Defaults in cfgs/default.yaml:

Parameter Default Description
num_gaussians 0 (auto) 0 = auto, typically ~3-13% of pixels depending on calibrated scene complexity
init_scale 1.5 Initial Gaussian size in pixels
max_steps 2800 Training iterations
l2_loss_ratio 1.0 L2/MSE loss weight
init_mode gradient Legacy model-side init (CLI uses adaptive): gradient, structure_tensor, gradient_rotinit, sublinear_gradient, saliency (else random)
rasterizer simple_sum Fixed simple-sum renderer in this branch

Performance

RTX 4090, Kodak images (768x512):

Mode Gaussians Steps Time PSNR
train auto (30-80K) 2800 5-10s 35-48 dB
merge overcomplete -> target 3000 6-12s 40+ dB
Cached rendering — — — 12,000+ FPS

PSNR varies by image complexity. Simple scenes reach 45+ dB, complex textures 35-40 dB.

Project Structure

flashimgs/           Importable Python package (programmatic API)
  __init__.py          Public re-exports (train / Session / helpers)
  api.py               train / train_tensor / merge / Result / auto_gaussians / image_info
  session.py           Session lifecycle (optimize / densify / add / merge / render / checkpoint)
main.py              CLI (train / merge / bench / export / info / viz / render)
model.py             GaussianSplatting2D (training, evaluation, rendering)
viz_server.py        Live browser-based training visualizer (single-port HTTP+WS)
web/
  flashimgs-viewer.js  <flashimgs-viewer> web component (WebGPU + CPU backends)
  index.html           Demo viewer page
cfgs/default.yaml    Default configuration
utils/               Image I/O, PSNR, loss functions, config loading
simple_sum_backend/  Simple-sum rasterizer CUDA extension (default renderer on this branch)
  rasterize_simple_sum.py        Inference-time rasterization
  rasterize_simple_sum_train.py  Training-time rasterization (fused fwd/bwd)
  merge_fast.py                  merge_gaussians() + fast_merge() for this backend
  cuda/csrc/                     CUDA kernels (forward_fused, simple_sum_train, ...)
tools/               Auxiliary scripts (batch_export, eval_merge_curve, process_gaussians)

License

GNU Affero General Public License v3.0 (AGPL-3.0). Third-party attributions are listed in NOTICE.

About

FlashImgs: fast 2D Gaussian Splatting for image representation

Resources

Readme
AGPL-3.0 license
Activity
Custom properties

Stars

1 star

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages