def testMap(self):
grid = openvdb.BoolGrid()
grid.fill((-4, -4, -4), (5, 5, 5), grid.zeroValue) # make active
grid.mapOn(lambda x: not x) # replace active False values with True
n = sum(item.value for item in grid.iterOnValues())
self.assertEqual(n, 10 * 10 * 10)
grid = openvdb.FloatGrid()
grid.fill((-4, -4, -4), (5, 5, 5), grid.oneValue)
grid.mapOn(lambda x: x * 2)
n = sum(item.value for item in grid.iterOnValues())
self.assertEqual(n, 10 * 10 * 10 * 2)
grid = openvdb.Vec3SGrid()
grid.fill((-4, -4, -4), (5, 5, 5), grid.zeroValue)
grid.mapOn(lambda x: (0, 1, 0))
n = sum(item.value[1] for item in grid.iterOnValues())
self.assertEqual(n, 10 * 10 * 10)
def testMeshConversion(self):
import time
# Skip this test if NumPy is not available.
try:
import numpy as np
except ImportError:
return
# Test mesh to volume conversion.
# Generate the vertices of a cube.
cubeVertices = [(x, y, z) for x in (0, 100) for y in (0, 100) for z in (0, 100)]
cubePoints = np.array(cubeVertices, float)
# Generate the faces of a cube.
cubeQuads = np.array([
(0, 1, 3, 2), # left
(0, 2, 6, 4), # front
(4, 6, 7, 5), # right
(5, 7, 3, 1), # back
(2, 3, 7, 6), # top
(0, 4, 5, 1), # bottom
], float)
voxelSize = 2.0
halfWidth = 3.0
xform = openvdb.createLinearTransform(voxelSize)
# Only scalar, floating-point grids support createLevelSetFromPolygons()
# (and the OpenVDB module might have been compiled without DoubleGrid support).
grids = []
for gridType in [n for n in openvdb.GridTypes
if n.__name__ in ('FloatGrid', 'DoubleGrid')]:
# Skip this test if the OpenVDB module was built without NumPy support.
try:
grid = gridType.createLevelSetFromPolygons(
cubePoints, quads=cubeQuads, transform=xform, halfWidth=halfWidth)
except NotImplementedError:
return
#openvdb.write('/tmp/testMeshConversion.vdb', grid)
self.assertEqual(grid.transform, xform)
self.assertEqual(grid.background, halfWidth * voxelSize)
dim = grid.evalActiveVoxelDim()
self.assertTrue(50 < dim[0] < 58)
self.assertTrue(50 < dim[1] < 58)
self.assertTrue(50 < dim[2] < 58)
grids.append(grid)
# Boolean-valued grids can't be used to store level sets.
self.assertRaises(TypeError, lambda: openvdb.BoolGrid.createLevelSetFromPolygons(
cubePoints, quads=cubeQuads, transform=xform, halfWidth=halfWidth))
# Vector-valued grids can't be used to store level sets.
self.assertRaises(TypeError, lambda: openvdb.Vec3SGrid.createLevelSetFromPolygons(
cubePoints, quads=cubeQuads, transform=xform, halfWidth=halfWidth))
# The "points" argument to createLevelSetFromPolygons() must be a NumPy array.
self.assertRaises(TypeError, lambda: openvdb.FloatGrid.createLevelSetFromPolygons(
cubeVertices, quads=cubeQuads, transform=xform, halfWidth=halfWidth))
# The "points" argument to createLevelSetFromPolygons() must be a NumPy float or int array.
self.assertRaises(TypeError, lambda: openvdb.FloatGrid.createLevelSetFromPolygons(
np.array(cubeVertices, bool), quads=cubeQuads, transform=xform, halfWidth=halfWidth))
# The "triangles" argument to createLevelSetFromPolygons() must be an N x 3 NumPy array.
self.assertRaises(TypeError, lambda: openvdb.FloatGrid.createLevelSetFromPolygons(
cubePoints, triangles=cubeQuads, transform=xform, halfWidth=halfWidth))
# Test volume to mesh conversion.
# Vector-valued grids can't be meshed.
self.assertRaises(TypeError, lambda: openvdb.Vec3SGrid().convertToQuads())
for grid in grids:
points, quads = grid.convertToQuads()
# These checks are intended mainly to test the Python/C++ bindings,
# not the OpenVDB volume to mesh converter.
self.assertTrue(len(points) > 8)
self.assertTrue(len(quads) > 6)
pmin, pmax = points.min(0), points.max(0)
self.assertTrue(-2 < pmin[0] < 2)
self.assertTrue(-2 < pmin[1] < 2)
self.assertTrue(-2 < pmin[2] < 2)
self.assertTrue(98 < pmax[0] < 102)
self.assertTrue(98 < pmax[1] < 102)
self.assertTrue(98 < pmax[2] < 102)
points, triangles, quads = grid.convertToPolygons(adaptivity=1)
self.assertTrue(len(points) > 8)
pmin, pmax = points.min(0), points.max(0)
self.assertTrue(-2 < pmin[0] < 2)
self.assertTrue(-2 < pmin[1] < 2)
self.assertTrue(-2 < pmin[2] < 2)
self.assertTrue(98 < pmax[0] < 102)
self.assertTrue(98 < pmax[1] < 102)
self.assertTrue(98 < pmax[2] < 102)
def main():
args = get_args()
dataset_base_dir = os.path.abspath(os.path.split(args.dataset)[0])
dataset_name = os.path.splitext(os.path.split(args.dataset)[1])[0]
with open(args.dataset) as f:
dataset_data = yaml.full_load(f)
model_folder = os.path.normpath(args.model_folder)
with open(os.path.join(model_folder, args.model_config), "r") as json_file:
model_config = json.load(json_file)
_, nn_model_name = os.path.split(model_folder)
results_output_directory = os.path.join(args.base_output_directory,
dataset_name, nn_model_name)
model_data = {
"name": nn_model_name,
"config": model_config,
}
model_weights_filepath = os.path.join(model_folder, args.model_weights)
model_weights_file_md5 = md5sum_path(model_weights_filepath)
output_yaml_file = os.path.join(results_output_directory, "summary.yml")
if os.path.exists(output_yaml_file):
with open(output_yaml_file, "r") as yaml_file:
previous_prediction_data = yaml.full_load(yaml_file)
else:
previous_prediction_data = None
results_data = {}
dataset = dataset_data.get("items", {})
predictions = {}
prediction_masks = {}
gt_renderings = {}
to_predict_dataset = {}
predicted_items = set()
for item_key, dataset_item in dataset.items():
dataset_item["is_2D"] = is_2D_file(
get_dataset_item_volume_path(dataset_base_dir, dataset_item))
output_prediction_path = os.path.join(
results_output_directory,
"{}.{}".format(item_key,
"exr" if dataset_item["is_2D"] else "vdb"),
)
dataset_item["output_prediction_path"] = output_prediction_path
previous_prediction_md5 = ((previous_prediction_data.get(
"items", {}).get(item_key, {}).get("model_weights_file_md5"))
if previous_prediction_data else None)
md5_recalculate = (not args.skip_md5_check and previous_prediction_data
and previous_prediction_md5 is not None and
previous_prediction_md5 != model_weights_file_md5)
if not os.path.exists(
output_prediction_path) or args.recalculate or md5_recalculate:
to_predict_dataset[item_key] = dataset_item
else:
render_filename = get_dataset_item_render_path(
dataset_base_dir, dataset_item)
if dataset_item["is_2D"]:
predictions[item_key] = read_image(output_prediction_path)
gt_renderings[item_key] = read_image(render_filename)
else:
predictions[item_key] = pyopenvdb.read(output_prediction_path,
"prediction")
gt_renderings[item_key] = pyopenvdb.read(
render_filename, "rendering")
prediction_timings = {}
if len(to_predict_dataset.keys()):
predictions_new, gt_renderings_new, prediction_masks_new = predict_volume_appearance(
dataset=to_predict_dataset,
dataset_base_dir=dataset_base_dir,
materials_filename=dataset_data["materials_file"],
model_config=model_config,
model_filename=model_weights_filepath,
stencils_only=False,
timings=prediction_timings,
verbose_logging=args.verbose,
)
for item_key, prediction in predictions_new.items():
dataset_item = dataset[item_key]
folder_path, file_name = os.path.split(
dataset_item["output_prediction_path"])
ensure_dir(folder_path)
if dataset_item["is_2D"]:
dump_image(image=prediction,
filepath=dataset_item["output_prediction_path"])
else:
pyopenvdb.write(dataset_item["output_prediction_path"],
[prediction])
predicted_items.add(item_key)
predictions.update(predictions_new)
if prediction_masks_new:
prediction_masks.update(prediction_masks_new)
gt_renderings.update(gt_renderings_new)
if args.verbose:
print("Prediction timings:")
pprint(prediction_timings)
# That's temporary for the deadline - normalize diff images
max_diff_pos, max_diff_neg = 0.0, 0.0
for item_key, dataset_item in dataset.items():
prediction = predictions[item_key]
gt_rendering = gt_renderings[item_key]
if not dataset_item["is_2D"]:
prediction_grid = prediction
prediction_acc = prediction_grid.getConstAccessor()
try:
prediction_mask_accessor = prediction_masks[
item_key].getConstAccessor()
except KeyError:
prediction_mask_accessor = None
coords_list = []
diff_list = []
for item in gt_rendering.citerOnValues():
if item.count == 1:
target = item.value
prediction, prediction_active = prediction_acc.probeValue(
item.min)
if prediction_mask_accessor:
mask_value = prediction_mask_accessor.getValue(
item.min)
else:
mask_value = True
if prediction_active and mask_value:
coords_list.append(item.min)
diff_list.append([target, prediction])
if diff_list:
array = numpy.array(diff_list)
else:
array = numpy.zeros((1, 2, 3))
gt_rendering = array[:, 0, :]
prediction = array[:, 1, :]
diff = prediction - gt_rendering
diff_vis_scaling = 4.0
diff_vis_array = plt.get_cmap(get_colormap("error"))(
((diff_vis_scaling * diff) / 2.0 + 0.5).mean(axis=1))
diff_positive = numpy.maximum(diff, 0.0)
diff_negative = numpy.minimum(diff, 0.0)
max_diff_pos = max(max_diff_pos, diff_positive.mean(axis=1).max())
max_diff_neg = max(max_diff_neg,
(-1.0 * diff_negative.mean(axis=1)).max())
diff_grid = pyopenvdb.Vec3SGrid((-1, -1, -1))
diff_grid.transform = prediction_grid.transform
diff_grid.name = "diff_red-green"
diff_grid_accessor = diff_grid.getAccessor()
for coord, diff_vis_value in zip(coords_list, diff_vis_array):
diff_grid_accessor.setValueOn(
coord,
(diff_vis_value[0], diff_vis_value[1], diff_vis_value[2]))
diff_dE = get_difference_metric("ciede2000")(gt_rendering,
prediction)
diff_dE_vis_scaling = 20.0
diff_dE_vis_array = plt.get_cmap(get_colormap("ciede2000"))(
diff_dE / diff_dE_vis_scaling)
diff_dE_grid = pyopenvdb.Vec3SGrid((-1, -1, -1))
diff_dE_grid.transform = prediction_grid.transform
diff_dE_grid.name = "diff_dE2000_20max"
diff_dE_grid_accessor = diff_dE_grid.getAccessor()
for coord, diff_vis_value in zip(coords_list,
diff_dE_vis_array[0]):
diff_dE_grid_accessor.setValueOn(
coord,
(diff_vis_value[0], diff_vis_value[1], diff_vis_value[2]))
pyopenvdb.write(dataset_item["output_prediction_path"],
[prediction_grid, diff_grid, diff_dE_grid])
rmse_linear = float(
get_difference_metric("rms")(gt_rendering, prediction))
rmse_srgb = float(
get_difference_metric("rms")(linear_to_sRGB(gt_rendering),
linear_to_sRGB(prediction)))
if dataset_item["is_2D"]:
ssim_linear = float(
get_difference_metric("ssim")(gt_rendering, prediction))
ssim_srgb = float(
get_difference_metric("ssim")(linear_to_sRGB(gt_rendering),
linear_to_sRGB(prediction)))
else:
ssim_linear = 0.0
ssim_srgb = 0.0
print(item_key, "RMSE:", rmse_linear)
print(item_key, "RMSE SRGB:", rmse_srgb)
print(item_key, "SSIM:", ssim_linear)
print(item_key, "SSIM SRGB:", ssim_srgb)
volume_filename = get_dataset_item_volume_path(dataset_base_dir,
dataset_item)
render_filename = get_dataset_item_render_path(dataset_base_dir,
dataset_item)
results_data[item_key] = {
"volume_filename": os.path.abspath(volume_filename),
"render_filename": os.path.abspath(render_filename),
"prediction_filename": os.path.abspath(output_prediction_path),
"base_image_name": dataset_item.get("base_image_name", ""),
"rmse_linear": rmse_linear,
"rmse_srgb": rmse_srgb,
"ssim_linear": ssim_linear,
"ssim_srgb": ssim_srgb,
"model_weights_file_md5": model_weights_file_md5,
}
print("max_diff_pos=", max_diff_pos, "max_diff_neg=", max_diff_neg)
with open(output_yaml_file, "w") as f:
yaml.dump({"model": model_data, "items": results_data}, f)
print("Arguments: <num_frames> <output_dir> [-f(lip y and z)]")
exit(0)
num_frames = int(argv[0])
out_path = str(argv[1])
flip_y_z = len(argv) > 2 and argv[2] == "-f"
# Will look for ordered grid.***** files in dir_path
dir_path = os.path.dirname(os.path.realpath(__file__))
dir_path += "/grids"
for frame in range(0, num_frames):
density_grid = vdb.FloatGrid()
density_grid.name = "density"
v_grid = vdb.Vec3SGrid()
v_grid.name = "v"
density_accessor = density_grid.getAccessor()
v_accessor = v_grid.getAccessor()
filepath = dir_path + "/grid." + str(frame).zfill(5)
with open(filepath) as f:
for line in f:
if (" " in line) and ("[" in line) and ("]" in line) and ("density:" in line):
index = line.split(" ")[0].split("[")[1].split("]")[0]
i = int(index.split(",")[0])
j = int(index.split(",")[1])
k = int(index.split(",")[2])
if (flip_y_z):
def predict_volume_appearance(
dataset_base_dir: str,
dataset: dict,
model_filename,
model_config,
materials_filename,
stencils_only,
timings: Dict = None,
ignore_md5_checks: bool = False,
verbose_logging: bool = False,
):
model_params = model_config["model_params"]
if not os.path.isabs(materials_filename):
materials_filename = os.path.join(dataset_base_dir, materials_filename)
# backward compatibility
if type(model_params["stencil_channels"]) == int:
model_params["stencil_channels"] = [
"scattering", "absorption", "mask"
][:model_params["stencil_channels"]]
model_arch_name = model_params["model_arch_name"]
patch_size = model_params["patch_size"]
scale_levels = model_params["scale_levels"]
stencil_channels = model_params["stencil_channels"]
is_2D_dataset = classify_dataset_class(dataset_base_dir, dataset)
data_class = DataPlanar if is_2D_dataset else Data3D
alignment_z_centered = model_params.get("alignment_z_centered",
data_class == Data3D)
data = data_class(
alignment_z_centered=alignment_z_centered,
data_items=dataset,
dataset_base_dir=dataset_base_dir,
find_inner_material_voxels=model_params["find_inner_material_voxels"],
ignore_md5_checks=ignore_md5_checks,
materials_file=materials_filename,
mode=MODE_PREDICT,
patch_size=patch_size,
sat_object_class_name="TreeSAT",
scale_levels=scale_levels,
shuffle_patches=False,
sliding_window_length=1,
stencil_channels=stencil_channels,
stencils_only=stencils_only,
timings=timings,
verbose_logging=verbose_logging,
)
batch_size = int(os.getenv("BATCH_SIZE", 10000))
model_make_function = models_collection[model_arch_name]
model = model_make_function(params=model_params)
model.load_weights(model_filename)
make_batch_function = (make_batch_swap_axes if model_arch_name
in ("planar_first",
"first_baseline") else make_batch)
locations = []
predictions = model.predict_generator(
generator=make_batches_gen(data, batch_size, locations,
make_batch_function),
steps=math.ceil(len(data) / batch_size),
verbose=1,
)
predicted_images = {}
predicted_images_accessors = {}
gt_renderings = {}
predicted_masks = None
predicted_masks_accessors = {}
materials = populate_materials(materials_filename)
material_channels = [m.channels for m in materials.values()]
assert max(material_channels) == min(
material_channels
), "number of channels in materials file mismatch" # all elements equal
materials_channel_count = material_channels[0]
if materials_channel_count != 3:
# spectral mode
material_wavelengths = [m.wavelengths for m in materials.values()]
for i in range(1, len(material_wavelengths)):
assert numpy.array_equal(
material_wavelengths[i - 1], material_wavelengths[i]
), "wavelength definition mismatch in materials file"
# spectral prediction
X, Y, Z = CIEXYZ_primaries(material_wavelengths[0] *
10) # convert nm to Angstrom
if is_2D_dataset:
for (datafile_idx, channel,
pidx), pixel_prediction in zip(locations, predictions):
dataset_item_key, metadata = data.volume_files[datafile_idx]
if dataset_item_key not in predicted_images.keys():
if not stencils_only:
gt_renderings[dataset_item_key] = data.get_render(
dataset_item_key=dataset_item_key)
predicted_images[dataset_item_key] = numpy.empty(
shape=(metadata["height"], metadata["width"], 3),
dtype=numpy.float32)
x, y, z = data.convert_patch_index(datafile_idx, channel, pidx)
predicted_images[dataset_item_key][y, x,
channel] = pixel_prediction
else: # is 3D dataset
material_voxel_found_index_masks = {
dataset_item_key: data.material_voxel_found_index[datafile_idx]
for datafile_idx, (dataset_item_key,
_) in enumerate(data.volume_files)
}
predicted_masks = {}
locations_predictions = (pd.DataFrame(
numpy.concatenate(
[numpy.array(locations),
numpy.array(predictions)], axis=1),
columns=["datafile_idx", "channel", "pidx", "value"],
).astype({
"datafile_idx": "uint16",
"channel": "uint8",
"pidx": "uint32",
}).sort_values(by=["datafile_idx", "pidx", "channel"]))
num_channels = int(locations_predictions["channel"].max()) + 1
for lp_idx in range(0, len(locations_predictions), num_channels):
datafile_idx = int(locations_predictions.at[lp_idx,
"datafile_idx"])
pidx = int(locations_predictions.at[lp_idx, "pidx"])
dataset_item_key, metadata = data.volume_files[int(datafile_idx)]
if dataset_item_key not in predicted_images.keys():
predicted_images[dataset_item_key] = pyopenvdb.Vec3SGrid(
(-1, -1, -1))
predicted_images_accessors[
dataset_item_key] = predicted_images[
dataset_item_key].getAccessor()
predicted_masks[dataset_item_key] = pyopenvdb.BoolGrid(False)
predicted_masks_accessors[dataset_item_key] = predicted_masks[
dataset_item_key].getAccessor()
render, _, _, _, _ = data.get_render(
dataset_item_key=dataset_item_key)
gt_renderings[dataset_item_key] = render
if render is not None:
predicted_images[
dataset_item_key].transform = render.transform
else:
dataset_item = data.data_items[dataset_item_key]
filled_path = dataset_item["proxy_object_filled_path"]
if filled_path is not None:
normal_grid = pyopenvdb.read(filled_path, "normalGrid")
predicted_images[
dataset_item_key].transform = normal_grid.transform
else:
raise RuntimeError(
"One of: `render_filename`, `proxy_object_filled_path` "
"should be set to extract transform for the predicted grid"
)
predicted_images[dataset_item_key].name = "prediction"
x, y, z = data.convert_patch_index_to_render_coords(
datafile_idx, 0, pidx)
if num_channels == 3:
value = (
locations_predictions.at[lp_idx + 0, "value"],
locations_predictions.at[lp_idx + 1, "value"],
locations_predictions.at[lp_idx + 2, "value"],
)
else:
# convolve the renderings with the XYZ color matching functions
value = [0.0, 0.0, 0.0]
for i, wavelength in enumerate(material_wavelengths[0]):
value[0] += X[i] * locations_predictions.at[lp_idx + i,
"value"]
value[1] += Y[i] * locations_predictions.at[lp_idx + i,
"value"]
value[2] += Z[i] * locations_predictions.at[lp_idx + i,
"value"]
value = tuple(xyz2linear_rgb([[value]])[0][0])
predicted_images_accessors[dataset_item_key].setValueOn((x, y, z),
value)
mask_value = material_voxel_found_index_masks[dataset_item_key][
pidx]
predicted_masks_accessors[dataset_item_key].setValueOn((x, y, z),
mask_value)
return predicted_images, gt_renderings, predicted_masks
