Skip to content

Commit

Permalink
Improve reproducibility
Browse files Browse the repository at this point in the history
  • Loading branch information
@LoadingByte
LoadingByte committed Sep 12, 2024
1 parent 1a61b65 commit 624461b
Show file tree
Hide file tree
Showing 3 changed files with 71 additions and 34 deletions.
49 changes: 39 additions & 10 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,13 @@ install our package:
pip install -e .

Next, download the [dataset](https://neural-gaussian-scale-space-fields.mpi-inf.mpg.de/data.zip) and extract it into the
repository such that, e.g., the folder `/data/picture` exists.
repository such that the new `data/` folder and the `scripts/` folder are siblings.

Steps 1 and 2 take a long time. To skip them, download the
[pretrained models](https://neural-gaussian-scale-space-fields.mpi-inf.mpg.de/models.zip) and extract them into the
repository such that the new `results/` folder and the `scripts/` folder are siblings.

The outputs of any experiment or visualizer will be written to the `results/` folder.

### Step 1: Train

Expand All @@ -64,6 +70,12 @@ third argument. Look in the `data/` folder for all available names. Notice that
python scripts/train.py neural picture bbq
python scripts/train.py neural mesh armadillo

Training a field takes about 4 hours because it does not halt upon convergence, but stoically continues for the maximum
number of iterations. To speed up training, you can significantly lower the `n_iters` variables in
[`train.py`](scripts/train.py) without notably sacrificing quality.

Each field takes 24MB of space, summing to 1.3GB for all fields.

### Step 2: Calibrate

Use the [`calibrate.py`](scripts/calibrate.py) script to run our post-training calibration:
Expand All @@ -85,6 +97,9 @@ Also use this script to generate smoothed ground truths via Monte Carlo convolut
python scripts/benchmark.py gauss picture [name]
python scripts/benchmark.py gauss mesh [name]

Benchmarking a picture/mesh field takes about 5min/15min. Generating picture/mesh ground truths takes about 30min/4h.
Per picture/mesh, the smoothed versions take 5GB/6.5GB, summing to about 1TB for all results.

### Step 4: Metrics

Use the [`metrics.py`](scripts/metrics.py) script to compare the output of our field with the ground truth:
Expand All @@ -105,20 +120,20 @@ Numbers very similar to those found in Tables 1-4 in our paper should now be ava

Use the [`visualize.py`](scripts/visualize.py) script to get the images, videos, and meshes from our paper and website:

python scripts/visualize.py picture_isotropic
python scripts/visualize.py picture_anisotropic
python scripts/visualize.py picture_foveation
python scripts/visualize.py mesh_isotropic
python scripts/visualize.py mesh_anisotropic
python scripts/visualize.py picture_isotropic [name]
python scripts/visualize.py picture_anisotropic [name]
python scripts/visualize.py picture_foveation [name]
python scripts/visualize.py mesh_isotropic [name]
python scripts/visualize.py mesh_anisotropic [name]
python scripts/visualize.py lightstage
python scripts/visualize.py picture_video
python scripts/visualize.py mesh_video_objects
python scripts/visualize.py mesh_video_ellipsoids
python scripts/visualize.py picture_video [name]
python scripts/visualize.py mesh_video_objects [name]
python scripts/visualize.py mesh_video_ellipsoids [name]
python scripts/visualize.py lightstage_video

### Ablations

To reproduce our ablations, perform the above four above steps with `neural` replaced by one of the following:
To reproduce our ablations, perform steps 1-4 with `neural` replaced by one of the following:

| Configuration as in the paper | Script equivalent |
|-------------------------------|-----------------------------------------|
Expand All @@ -136,6 +151,11 @@ First install the required additional dependencies:

pip install -e .[neural-texture]

If you neither downloaded the pretrained models nor ran steps 1-2 yet, run these commands to prepare the neural texture:

python scripts/train.py neural textured
python scripts/calibrate.py neural textured

Open the 3D renderer window using the [`textured_render_uv.py`](scripts/textured_render_uv.py) script. Press ESC to
quit, W/A/S/D/SPACE/SHIFT to move, and use the mouse to look around. You will see the mesh of a fish, but instead of a
texture, it is covered in UV coordinates. If you feel dizzy, set `gizmo_and_grid` in the script to true to render a
Expand All @@ -149,3 +169,12 @@ factor here and increased the window size in `textured_render_uv.py` to still ge
demonstrate our method, you can replace `16` with `1`:

python scripts/textured_apply_neural_texture.py moving 16

### Performance

Use the [`performance.py`](scripts/performance.py) script to reproduce our timing experiment:

python scripts/performance.py vanilla picture 30
python scripts/performance.py vanilla mesh 0.004
python scripts/performance.py neural picture 30
python scripts/performance.py neural mesh 0.004
11 changes: 3 additions & 8 deletions scripts/performance.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,6 @@
def main(field_type, category, stop_metric):
stop_metric = float(stop_metric)

out_dir = ngssf.results.performance_dir()
pred_dir = out_dir / f"predictions_{field_type}_{category}_{stop_metric}"
pred_dir.mkdir(parents=True, exist_ok=True)

# Warm-up
time_method(field_type, category, perf_names(category)[0], 0, stop_metric)

Expand All @@ -33,7 +29,7 @@ def main(field_type, category, stop_metric):
its.append(it)
rows.append((f"{var_bench_idx}", field_type, f"{np.mean(tts):.4f}", f"{np.mean(its):.4f}"))

with open(out_dir / f"timings_{field_type}_{category}_{stop_metric}.csv", "w") as f:
with open(ngssf.results.performance_dir() / f"timings_{field_type}_{category}_{stop_metric}.csv", "w") as f:
f.write("\n".join(",".join(row) for row in rows))


Expand Down Expand Up @@ -70,8 +66,7 @@ def start_fn():
start_time = cur_time()

def loop_fn(itr, pred_fn):
# Bacon produces extremely complex meshes initially, which slow down our Chamfer code significantly.
if (itr + 1) % 50 != 0 or field_type == "bacon" and category == "mesh" and itr < 15_000:
if (itr + 1) % 50 != 0:
return False
inference_start_time = cur_time()
with torch.no_grad():
Expand Down Expand Up @@ -117,7 +112,7 @@ def train(field_type, category, true_grid, true_mesh, res, start_fn, loop_fn):
sampler = ngssf.MinibatchSampler(2 ** 24, ngssf.SDFSampler(true_mesh)).cuda()
enc_kw["length_distribution_param"] = 100
n_samples = 200_000
if field_type.startswith("neural"):
if field_type == "neural":
scaler = ngssf.MinibatchScaler(10_000_000, ngssf.RandomScaler(True, sig.coords)).cuda()
field = ngssf.nn.prefab.Smoothable4x1024NeuralField(sig.coords, sig.channels, True, None, enc_kw)
elif field_type == "vanilla":
Expand Down
45 changes: 29 additions & 16 deletions scripts/visualize.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,10 @@ def picture_anisotropic():


def _picture(base_dir, scale_set, indices):
for name in tqdm([
for name in tqdm(args_or([
"bbq", "cliffs", "colosseo", "crystals", "firenze", "firewood", "mutter", "peak", "portal", "rue",
"schaumbrunnen", "steepshore", "toomuchbleach", "tunnelrampe", "zebras"
], desc="name", leave=False):
]), desc="name", leave=False):
for index in tqdm(indices, desc="scale", leave=False):
gauss_img = _prepare_picture_image(ngssf.results.load_benchmark("gauss", "picture", name, scale_set, index))
pred_img = _prepare_picture_image(ngssf.results.load_benchmark("neural", "picture", name, scale_set, index))
Expand All @@ -45,12 +45,14 @@ def _picture(base_dir, scale_set, indices):


def picture_foveation():
field = ngssf.results.load_neural_field("neural", "picture", "squirrel").cuda()
name = arg_or("squirrel")
field = ngssf.results.load_neural_field("neural", "picture", name).cuda()
res = 512
X = ngssf.util.grid_coords(res, 2, device="cuda")
with torch.no_grad():
pred_img = field(X, (X.norm(dim=1) - 0.35).clamp(0) ** 3 / 200).T.reshape(3, res, res)
_write_image(ngssf.results.visualizations_dir() / "picture_foveation.jpg", _prepare_picture_image(pred_img))
img = _prepare_picture_image(pred_img)
_write_image(ngssf.results.visualizations_dir() / "picture_foveation" / f"{name}.jpg", img)


def _prepare_picture_image(img):
Expand All @@ -59,7 +61,7 @@ def _prepare_picture_image(img):

def mesh_isotropic():
scale_set = "variance_benchmark"
for name in tqdm(ngssf.data.names("mesh"), desc="name", leave=False):
for name in tqdm(args_or(ngssf.data.names("mesh")), desc="name", leave=False):
d = ngssf.results.visualizations_dir() / "mesh_isotropic"
dg = d / name / "gauss"
dn = d / name / "neural"
Expand All @@ -73,9 +75,10 @@ def mesh_isotropic():


def mesh_anisotropic():
d = ngssf.results.visualizations_dir() / "mesh_anisotropic"
name = arg_or("thai")
d = ngssf.results.visualizations_dir() / "mesh_anisotropic" / name
d.mkdir(parents=True, exist_ok=True)
field = ngssf.results.load_neural_field("neural", "mesh", "thai").cuda()
field = ngssf.results.load_neural_field("neural", "mesh", name).cuda()
for label, variances in [
("isotropic", [1e-2, 1e-2, 1e-2]),
("anisotropic_horizontal", [1e-2, 1e-8, 1e-2]),
Expand All @@ -84,12 +87,13 @@ def mesh_anisotropic():
scale = torch.diag(torch.tensor(variances))
with torch.no_grad():
grid = ngssf.util.eval_grid(256, field, scale.cuda(), batch_size=2 ** 18).cpu()
ngssf.util.mesh_from_grid(grid).export(d / f"thai_{label}.ply")
ngssf.util.mesh_from_grid(grid).export(d / f"{label}.ply")


def lightstage():
name = arg_or("cute")
light_positions = ngssf.data.lightstage_light_positions()
field = ngssf.results.load_neural_field("neural", "lightstage", "cute").cuda()
field = ngssf.results.load_neural_field("neural", "lightstage", name).cuda()
w, h = 512, 384
X = torch.cat([
torch.cartesian_prod(torch.linspace(-0.75, 0.75, h), torch.linspace(-1, 1, w)).flip(1),
Expand All @@ -99,7 +103,7 @@ def lightstage():
with torch.no_grad():
Y = field(X.cuda(), scale)
img = _prepare_image(Y.T.reshape(3, h, w))
_write_image(ngssf.results.visualizations_dir() / "lightstage" / f"{i}.jpg", img)
_write_image(ngssf.results.visualizations_dir() / "lightstage" / name / f"{i}.jpg", img)


def picture_video():
Expand All @@ -114,7 +118,7 @@ def picture_video():
spectrum_label = _label("Spectrum", 1300, (128, 32), overlay=True)
cov_label = _label("Covariance", 1300, (128, 32))

for name in tqdm(["bbq", "firewood", "schaumbrunnen", "tunnelrampe"], leave=False):
for name in tqdm(args_or(["bbq", "firewood", "schaumbrunnen", "tunnelrampe"]), leave=False):
orig_picture = torch.as_tensor(resize(ngssf.data.load("picture", name).numpy(), (3, 512, 512)), device="cuda")
neural_field = ngssf.results.load_neural_field("neural", "picture", name).cuda()
gauss_field = ngssf.GaussianMonteCarloSmoothableField(ngssf.GridField(orig_picture, padding_mode="reflection"))
Expand Down Expand Up @@ -146,7 +150,7 @@ def _spectrum(picture):
def mesh_video_objects():
cov_mats = torch.tensor(_mesh_video_covariance_matrices(), dtype=torch.float32, device="cuda")

for name in tqdm(ngssf.data.names("mesh"), leave=False):
for name in tqdm(args_or(ngssf.data.names("mesh")), leave=False):
orig_mesh = ngssf.data.load("mesh", name)
neural_field = ngssf.results.load_neural_field("neural", "mesh", name).cuda()
gauss_field = ngssf.GaussianMonteCarloSmoothableField(
Expand Down Expand Up @@ -191,12 +195,13 @@ def _mesh_video_covariance_matrices():


def lightstage_video():
d = ngssf.results.visualizations_dir()
name = arg_or("cute")
d = ngssf.results.visualizations_dir() / "lightstage_video"
d.mkdir(parents=True, exist_ok=True)

light_shots = ngssf.data.load("lightstage", "cute")
light_shots = ngssf.data.load("lightstage", name)
light_pos = ngssf.data.lightstage_light_positions()
neural_field = ngssf.results.load_neural_field("neural", "lightstage", "cute").cuda()
neural_field = ngssf.results.load_neural_field("neural", "lightstage", name).cuda()
neural_field.calibration_factors[1] = 500
gauss_field = ngssf.GaussianMonteCarloSmoothableField(ngssf.LightStageField(light_shots, light_pos), {2, 3}).cuda()

Expand All @@ -218,7 +223,7 @@ def lightstage_video():
cov_label = _label("Covariance", 1300, (128, 32))
plotted_pos = light_pos[(light_pos - (light_pos[2] + light_pos[22]) / 2).norm(dim=1) < 0.15]

video = VideoWriter(str(d / "lightstage_video.mp4"), VideoWriter.fourcc('a', 'v', 'c', '1'), 60, (1152, 384))
video = VideoWriter(str(d / f"{name}.mp4"), VideoWriter.fourcc('a', 'v', 'c', '1'), 60, (1152, 384))
for x_light, cov_mat in tqdm(list(zip(xs_light, cov_mats)), leave=False):
X = torch.cat([X_pixel, x_light.tile(w * h, 1)], dim=1).cuda()
with torch.no_grad():
Expand Down Expand Up @@ -304,5 +309,13 @@ def _write_image(file, img):
iio.imwrite(file, (img * 255).astype(np.uint8), quality=90)


def arg_or(default):
return sys.argv[2] if len(sys.argv) > 2 else default


def args_or(default):
return sys.argv[2:] if len(sys.argv) > 2 else default


if __name__ == "__main__":
globals()[sys.argv[1]]()

0 comments on commit 624461b

Please sign in to comment.