def __init__(self, input_file, features, n_neighbors, all_labels=False):
"Initialization"
data = ply2dict(input_file)
try:
all_features = ["x", "y", "z"] + features
X = np.vstack([data[f] for f in all_features]).T
except KeyError:
print(f"ERROR: Input features {features} not recognized")
return
labels = data["labels"]
self.index = np.arange(X.shape[0])
if not all_labels:
X, labels = self.filter_labels(X, labels)
self.X = torch.from_numpy(X)
self.labels = torch.from_numpy(labels)
self.n_samples = self.labels.shape[0]
tree = KDTree(self.X[:, :3])
_, self.neighbors_idx = tree.query(
self.X[:, :3], k=n_neighbors, sort_results=True
)
for path_ply in args.files:
path_ply = os.path.join(args.prefix_path, path_ply)
print(f"\nProcessing file: {path_ply}")
print("* Preparing dataloader..", end=" ", flush=True)
dataset = AerialPointDataset(path_ply, **config["data"])
loader = DataLoader(
dataset=dataset,
batch_size=args.batch_size,
num_workers=args.num_workers,
shuffle=False,
)
print("DONE")
# Create and fill point cloud field
data = ply2dict(path_ply)
true_labels = data["labels"]
names = NAMES_9
# in the 4-labels case
if not config["data"]["all_labels"]:
true_labels = convert_labels(true_labels).astype(np.int32)
names = NAMES_4
n = len(true_labels)
predictions = -np.ones(n, dtype=np.int32)
raw_predictions = predict(loader, len(dataset)).astype(np.int32)
predictions[dataset.index] = raw_predictions
errors = predictions != true_labels
data["predictions"] = predictions
data["errors"] = errors.astype(np.uint8)
def compute_features(path_ply, steps_params):
filename = os.path.split(path_ply)[-1]
data = ply2dict(path_ply)
coords = np.vstack((data["x"], data["y"], data["z"])).T
grid_ground_3d = None
for (step, params) in steps_params.items():
if step == "descriptors":
print("Computing local descriptors..")
all_descriptors = descriptors.compute_descriptors(coords, **params)
data.update(all_descriptors)
if step == "region_growing":
print("\nComputing regions..")
normals = np.vstack((data["nx"], data["ny"], data["nz"])).T
params_copy = params.copy()
descriptor_selected = params_copy.pop("descriptor")
print("* descriptor selected : "
f"{'min' if params['minimize'] else 'max'} "
f"{descriptor_selected}")
print(f"* thresholds : {params['thresholds']}")
print(f"* radius : {params['radius']}")
try:
descriptor_vals = data[descriptor_selected]
region_labels = region_growing.multi_region_growing(
coords, normals, descriptor_vals, **params_copy)
data["regions"] = region_labels
except KeyError:
print(
f"Descriptor '{descriptor_selected}' has not been computed"
", run 'python3 compute_features.py --descriptors "
f"{descriptor_selected}'")
sys.exit(-1)
if step == "ground_extraction":
print("\nExtracting ground from regions..")
region_labels = data["regions"]
ground_mask = ground_extraction.stitch_regions(
coords, region_labels, **params)
ground_only = {
field: data[field][ground_mask]
for field in list(data.keys())
}
data["ground"] = ground_mask.astype(np.uint8)
os.makedirs(PATH_GROUND_ONLY, exist_ok=True)
path_ground = os.path.join(PATH_GROUND_ONLY, filename)
if dict2ply(ground_only, path_ground):
print(f"* PLY ground file successfully saved to {path_ground}")
if step == "ground_rasterization":
print("\nComputing ground rasterization..")
ground_mask = data["ground"].astype(bool)
grid_ground_3d = ground_extraction.rasterize_ground(
coords, ground_mask, **params)
ground_rasterized = {
"x": grid_ground_3d[:, 0],
"y": grid_ground_3d[:, 1],
"z": grid_ground_3d[:, 2],
"ground_altitude": grid_ground_3d[:, 2],
}
path_rasterized = os.path.join(PATH_GROUND_RASTERIZED, filename)
if dict2ply(ground_rasterized, path_rasterized):
print("* PLY ground rasterized file successfully saved to "
f"{path_rasterized}")
if step == "height_above_ground":
print("\nComputing height above ground..")
if grid_ground_3d is None:
path_rasterized = os.path.join(PATH_GROUND_RASTERIZED,
filename)
print(f"* Loading rasterized ground : {path_rasterized}")
ground_rasterized = ply2dict(path_rasterized)
grid_ground_3d = np.vstack((
ground_rasterized["x"],
ground_rasterized["y"],
ground_rasterized["z"],
)).T
ground_mask = data["ground"].astype(bool)
heights = ground_extraction.height_above_ground(
coords, ground_mask, grid_ground_3d)
data["height_above_ground"] = heights
print("DONE")
# saving data
path_output = os.path.join(PATH_FEATURES, filename)
if dict2ply(data, path_output):
print(f"\nPLY features file successfully saved to {path_output}")
steps["height_above_ground"] = {}
if args.rasterize_ground or args.full_pipeline:
steps["rasterize_ground"] = {
"step": args.rasterize_step,
}
os.makedirs(PATH_GROUND_RASTERIZED, exist_ok=True)
if len(steps.keys()) == 0:
print("ERROR : No steps to compute")
sys.exit(-1)
for file in args.files:
print(f"\nComputing features of file {file}")
data = ply2dict(file)
data, ground_only, ground_rasterized = compute_features(data, steps)
# save PLY files
filename = os.path.split(file)[-1]
f_data = os.path.join(PATH_FEATURES, filename)
if dict2ply(data, f_data):
print(f"PLY file successfully saved to {f_data}")
if ground_only:
f_ground_only = os.path.join(PATH_GROUND_ONLY, filename)
if dict2ply(ground_only, f_ground_only):
print(f"PLY ground file successfully saved to {f_ground_only}")
if ground_rasterized:
f_ground_rasterized = os.path.join(PATH_GROUND_RASTERIZED,