Spaces:
Runtime error
Runtime error
File size: 2,550 Bytes
19677a1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
"""
command line example:
$ python -i -m jaxnerf.nerf.precompute --data_dir {path-to-data-dir} --split train \
--dataset blender --factor 4 --dtype float16
"""
import os
import argparse
from typing import Optional
import jax.numpy as np
from jaxnerf.nerf import utils
from jaxnerf.nerf import clip_utils
from jaxnerf.nerf import datasets
def precompute_image_features(data_dir: str, split: str, dataset: str, factor: int, dtype: str,
model_name: Optional[str], render_path: Optional[str]):
if dataset == "blender":
if render_path:
raise ValueError("render_path cannot be used for the blender dataset.")
# image in numpy.ndarray
_, images, _ = datasets.Blender.load_files(data_dir, split, factor)
clip_model = clip_utils.init_CLIP(dtype, model_name)
# CLIP output in jax.numpy.ndarray
images = np.stack(images).transpose(0, 3, 1, 2)
images = images[:, :3, :, :]
images = clip_utils.preprocess_for_CLIP(images)
embeddings = clip_model.get_image_features(pixel_values=images)
embeddings /= np.linalg.norm(embeddings, axis=-1, keepdims=True)
print(f'completed precomputing CLIP embeddings: ({embeddings.shape[0]} images)')
# write as pickle
write_path = os.path.join(data_dir, f'clip_cache_{split}_factor{factor}_{dtype}.pkl')
utils.write_pickle(embeddings, write_path)
print(f'precompute written as pickle: {write_path}')
elif dataset == "llff":
raise NotImplementedError
else:
raise ValueError(f"invalid dataset: {dataset}")
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument("--data_dir", type=str, required=True)
parser.add_argument("--split", type=str, required=True, help="train/val/test")
parser.add_argument("--dataset", type=str, required=True)
parser.add_argument("--factor", type=int, required=True,
help="downsampling factor: 0/2/4")
parser.add_argument("--dtype", type=str, required=True,
help="float32/float16 (float16 is used to save memory)")
parser.add_argument("--model_name", type=str, required=False, default=None)
parser.add_argument("--render_path", type=str, required=False, default=None)
args = parser.parse_args()
precompute_image_features(args.data_dir, args.split, args.dataset, args.factor,
args.dtype, args.model_name, args.render_path)
|