from __future__ import annotations
import copy
import json
import pickle
from collections.abc import Generator, Iterator
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path
from typing import TYPE_CHECKING, Any, TypedDict, TypeVar
import numpy as np
import numpy.typing as npt
from py3dtiles.exceptions import TilerException
from py3dtiles.points import Points
from py3dtiles.tilers.point.pnts import MIN_POINT_SIZE
from py3dtiles.tilers.point.pnts.pnts_writer import points_to_pnts_file
from py3dtiles.tileset.bounding_volume_box import BoundingVolumeBox
from py3dtiles.tileset.content import read_binary_tile_content
from py3dtiles.tileset.content.pnts_feature_table import SemanticPoint
from py3dtiles.tileset.tile import Tile
from py3dtiles.tileset.tileset import TileSet
from py3dtiles.typing import ExtraFieldsDescription
from py3dtiles.utils import (
SubdivisionType,
aabb_size_to_subdivision_type,
node_name_to_path,
split_aabb,
)
from .distance import xyz_to_child_index
from .points_grid import Grid
if TYPE_CHECKING:
from typing_extensions import NotRequired
from .node_catalog import NodeCatalog
_T = TypeVar("_T", bound=npt.NBitBase)
[docs]
def node_to_tileset(
args: tuple[Node, Path, npt.NDArray[np.float32], Node | None, int]
) -> Tile | None:
return args[0].to_tileset(args[1], args[2], args[3], args[4], None)
[docs]
class DummyNodeDictType(TypedDict):
children: NotRequired[list[bytes]]
grid: NotRequired[Grid]
points: NotRequired[list[Points]]
[docs]
class DummyNode:
def __init__(self, _bytes: DummyNodeDictType) -> None:
if "children" in _bytes:
self.children: list[bytes] | None = _bytes["children"]
self.grid = _bytes["grid"]
else:
self.children = None
self.points = _bytes["points"]
[docs]
def get_points(self) -> Points | None:
if self.children:
return self.grid.get_points()
else:
points = self.points
if len(points) == 0:
return None
xyz = np.concatenate(tuple([pt.positions for pt in points]))
if points[0].colors is None:
# assume we don't have color
rgb: npt.NDArray[np.uint8] | None = None
else:
# is not none only to make mypy happy
rgb = np.concatenate(
tuple([pt.colors for pt in points if pt.colors is not None])
)
extra_fields = {}
for f in self.points[0].extra_fields.keys():
extra_fields[f] = np.concatenate(
tuple([pt.extra_fields[f] for pt in points])
)
return Points(positions=xyz, colors=rgb, extra_fields=extra_fields)
[docs]
class Node:
"""docstring for Node"""
__slots__ = (
"name",
"aabb",
"aabb_size",
"inv_aabb_size",
"aabb_center",
"spacing",
"include_rgb",
"extra_fields",
"pending_points",
"children",
"grid",
"points",
"dirty",
)
def __init__(
self,
name: bytes,
aabb: npt.NDArray[np.float64 | np.float32],
spacing: float,
include_rgb: bool,
extra_fields: list[ExtraFieldsDescription],
) -> None:
super().__init__()
self.name = name
self.aabb = aabb.astype(
np.float32
) # TODO remove astype once the whole typing is done (and once data type issues on numpy arrays are fixed).
self.aabb_size = np.maximum(self.aabb[1] - self.aabb[0], MIN_POINT_SIZE)
self.inv_aabb_size = 1.0 / self.aabb_size
self.aabb_center = (self.aabb[0] + self.aabb[1]) * 0.5
self.spacing = spacing
self.include_rgb = include_rgb
self.extra_fields = extra_fields
self.pending_points: list[Points] = []
self.children: list[bytes] | None = None
self.grid = Grid(self.spacing, self.include_rgb, self.extra_fields)
self.points: list[Points] = []
self.dirty = False
[docs]
@staticmethod
def create_child_node_from_parent(
name: bytes,
parent_aabb: npt.NDArray[np.floating[_T]],
parent_spacing: float,
include_rgb: bool,
extra_fields: list[ExtraFieldsDescription],
) -> Node:
spacing = parent_spacing * 0.5
aabb = split_aabb(parent_aabb, int(name[-1])) if len(name) > 0 else parent_aabb
# let's build a new Node
return Node(name, aabb.astype(np.float64), spacing, include_rgb, extra_fields)
[docs]
def save_to_bytes(self) -> bytes:
sub_pickle: dict[str, Any] = {}
if self.children is not None:
sub_pickle["children"] = self.children
sub_pickle["grid"] = self.grid
else:
sub_pickle["points"] = self.points
return pickle.dumps(sub_pickle)
[docs]
def load_from_bytes(self, byt: bytes) -> None:
sub_pickle = pickle.loads(byt)
if "children" in sub_pickle:
self.children = sub_pickle["children"]
self.grid = sub_pickle["grid"]
else:
self.points = sub_pickle["points"]
[docs]
def insert(
self,
scale: float,
points: Points,
make_empty_node: bool = False,
) -> None:
if make_empty_node:
self.children = []
self.pending_points.append(points)
return
# fastpath
if self.children is None:
self.points.append(points)
count = sum([pt.positions.shape[0] for pt in self.points])
# stop subdividing if spacing is 1mm
if count >= 20000 and self.spacing > 0.001 * scale:
self._split(scale)
self.dirty = True
return
# grid based insertion
(
remainder_xyz,
remainder_rgb,
remainder_extra_fields,
needs_balance,
) = self.grid.insert(
self.aabb[0],
self.inv_aabb_size,
points.positions,
points.colors,
points.extra_fields,
)
if needs_balance:
self.grid.balance(self.aabb_size, self.aabb[0], self.inv_aabb_size)
self.dirty = True
self.dirty = self.dirty or (len(remainder_xyz) != len(points.positions))
if len(remainder_xyz) > 0:
self.pending_points.append(
Points(
positions=remainder_xyz,
colors=remainder_rgb,
extra_fields=remainder_extra_fields,
)
)
[docs]
def needs_balance(self) -> bool:
if self.children is not None:
return self.grid.needs_balance()
return False
[docs]
def flush_pending_points(self, catalog: NodeCatalog, scale: float) -> None:
for name, pt in self._get_pending_points():
catalog.get_node(name).insert(scale, pt)
self.pending_points = []
[docs]
def dump_pending_points(self) -> list[tuple[bytes, bytes, int]]:
result = []
for name, pt in self._get_pending_points():
points_dict = {
"xyz": pt.positions,
"rgb": pt.colors,
"extra_fields": pt.extra_fields,
}
result.append((name, pickle.dumps(points_dict), len(pt.positions)))
self.pending_points = []
return result
[docs]
def get_pending_points_count(self) -> int:
return sum([pt.positions.shape[0] for pt in self.pending_points])
def _get_pending_points(
self,
) -> Iterator[tuple[bytes, Points]]:
if not self.pending_points:
return
pending_xyz_arr = np.concatenate([pt.positions for pt in self.pending_points])
if self.include_rgb:
pending_rgb_arr = np.concatenate([pt.colors for pt in self.pending_points])
pending_extra_fields: dict[str, npt.NDArray[Any]] = {}
for pt in self.pending_points:
for field, arr in pt.extra_fields.items():
if pending_extra_fields.get(field) is None:
pending_extra_fields[field] = arr
else:
pending_extra_fields[field] = np.append(
pending_extra_fields[field], arr, axis=0
)
t = aabb_size_to_subdivision_type(self.aabb_size)
if t == SubdivisionType.QUADTREE:
indices = xyz_to_child_index(
pending_xyz_arr,
np.array(
[self.aabb_center[0], self.aabb_center[1], self.aabb[1][2]],
dtype=np.float32,
),
)
else:
indices = xyz_to_child_index(pending_xyz_arr, self.aabb_center)
# unique children list
childs = np.unique(indices)
# make sure all children nodes exist
for child in childs:
name = "{}{}".format(self.name.decode("ascii"), child).encode("ascii")
# create missing nodes, only for remembering they exist.
# We don't want to serialize them
# probably not needed...
if self.children is not None and name not in self.children:
self.children += [name]
self.dirty = True
# print('Added node {}'.format(name))
mask = np.where(indices - child == 0)
xyz = pending_xyz_arr[mask]
if len(xyz) > 0:
extra_fields = {}
for f, arr in pending_extra_fields.items():
extra_fields[f] = arr[mask]
yield name, Points(
positions=xyz,
colors=pending_rgb_arr[mask] if self.include_rgb else None,
extra_fields=extra_fields,
)
def _split(self, scale: float) -> None:
self.children = []
for pt in self.points:
self.insert(scale, pt)
self.points = []
[docs]
def get_point_count(
self, node_catalog: NodeCatalog, max_depth: int, depth: int = 0
) -> int:
if self.children is None:
return sum([pt.positions.shape[0] for pt in self.points])
else:
count = self.grid.get_point_count()
if depth < max_depth:
for n in self.children:
count += node_catalog.get_node(n).get_point_count(
node_catalog, max_depth, depth + 1
)
return count
[docs]
def get_points(self) -> Points | None:
if self.children is None:
if len(self.points) == 0:
return None
xyz = np.concatenate(tuple([pt.positions for pt in self.points]))
if self.points[0].colors is not None:
# the "or []" is just to make mypy happy. Normally it shouldn't be None here
rgb: npt.NDArray[np.uint8 | np.uint16] | None = np.concatenate(
tuple([pt.colors or [] for pt in self.points])
)
else:
rgb = None
extra_fields = {}
for f in self.extra_fields:
extra_fields[f.name] = np.concatenate(
tuple([pt.extra_fields[f.name] for pt in self.points])
)
return Points(positions=xyz, colors=rgb, extra_fields=extra_fields)
else:
return self.grid.get_points()
[docs]
def get_child_names(self) -> Generator[bytes, None, None]:
for number_child in range(8):
yield f"{self.name.decode('ascii')}{number_child}".encode("ascii")
[docs]
def to_tileset(
self,
folder: Path,
scale: npt.NDArray[np.float32],
parent_node: Node | None = None,
depth: int = 0,
pool_executor: ProcessPoolExecutor | None = None,
) -> Tile | None:
# create child tileset parts
# if their size is below of 100 points, they will be merged in this node.
children_tileset_parts: list[Tile] = []
parameter_to_compute: list[
tuple[Node, Path, npt.NDArray[np.float32], Node, int]
] = []
for child_name in self.get_child_names():
child_node = Node.create_child_node_from_parent(
child_name, self.aabb, self.spacing, self.include_rgb, self.extra_fields
)
child_pnts_path = node_name_to_path(folder, child_name, ".pnts")
if child_pnts_path.exists():
# multi thread is only allowed on nodes where there are no prune
# a simple rule is: only is there is not a parent node
if pool_executor and parent_node is None:
parameter_to_compute.append(
(child_node, folder, scale, self, depth + 1)
)
else:
children_tileset_part = child_node.to_tileset(
folder, scale, self, depth + 1
)
if (
children_tileset_part is not None
): # return None if the child has been merged
children_tileset_parts.append(children_tileset_part)
if pool_executor and parent_node is None:
children_tileset_parts = [
t
for t in pool_executor.map(node_to_tileset, parameter_to_compute)
if t is not None
]
pnts_path = node_name_to_path(folder, self.name, ".pnts")
tile_content = read_binary_tile_content(pnts_path)
fth = tile_content.body.feature_table.header
xyz = tile_content.body.feature_table.body.position.reshape((-1, 3))
# check if this node should be merged in the parent.
prune = False # prune only if the node is a leaf
# If this child is small enough, merge in the current tile
if parent_node is not None and depth > 1 and fth.points_length < 100:
parent_pnts_path = node_name_to_path(folder, parent_node.name, ".pnts")
parent_tile = read_binary_tile_content(parent_pnts_path)
parent_fth = parent_tile.body.feature_table.header
parent_xyz = parent_tile.body.feature_table.body.position.reshape(
(parent_fth.points_length, 3)
)
if (
parent_fth.colors != SemanticPoint.NONE
and parent_tile.body.feature_table.body.color is not None
):
parent_rgb = parent_tile.body.feature_table.body.color.reshape((-1, 3))
else:
parent_rgb = None
parent_extra_fields = {}
for field in parent_tile.body.batch_table.header.data:
parent_extra_fields[
field
] = parent_tile.body.batch_table.get_binary_property(field)
# update aabb based on real values
parent_bounding_volume = BoundingVolumeBox.from_points(parent_xyz)
parent_xyz = np.concatenate((parent_xyz, xyz))
if fth.colors != SemanticPoint.NONE:
if tile_content.body.feature_table.body.color is None:
raise TilerException(
"If the parent has color data, the children must also have color data."
)
parent_rgb = np.concatenate(
(
parent_rgb,
tile_content.body.feature_table.body.color.reshape((-1, 3)),
)
)
for field in tile_content.body.batch_table.header.data:
parent_extra_fields[field] = np.concatenate(
(
parent_extra_fields[field],
tile_content.body.batch_table.get_binary_property(field),
)
)
# update aabb
xyz_float = xyz.view(np.float32)
new_bounding_volume_box = BoundingVolumeBox.from_points(xyz_float)
parent_bounding_volume.add(new_bounding_volume_box)
parent_pnts_path.unlink()
points_to_pnts_file(
folder,
parent_node.name,
Points(
positions=parent_xyz,
colors=parent_rgb,
extra_fields=parent_extra_fields,
),
)
pnts_path.unlink()
prune = True
content_uri = None
if not prune:
content_uri = pnts_path.relative_to(folder)
xyz_float = xyz.view(np.float32)
# update aabb based on real values
bounding_box = BoundingVolumeBox.from_points(xyz_float)
else:
# if it is a leaf that should be pruned
if not children_tileset_parts:
return None
# recompute the aabb in function of children
bounding_box = BoundingVolumeBox()
for child_tileset_part in children_tileset_parts:
if child_tileset_part.bounding_volume is not None:
bounding_box.add(child_tileset_part.bounding_volume)
if bounding_box is None:
raise TilerException("bounding_box shouldn't be None")
tile: Tile = Tile(
geometric_error=10 * self.spacing / scale[0], bounding_volume=bounding_box
)
if content_uri is not None:
tile.content_uri = content_uri
if children_tileset_parts:
tile.children = children_tileset_parts
else:
tile.geometric_error = 0.0
if (
len(self.name) > 0
and children_tileset_parts
and len(json.dumps(tile.to_dict())) > 100000
):
tile = split_tileset(tile, self.name.decode(), folder)
return tile
[docs]
def split_tileset(tile: Tile, split_name: str, folder: Path) -> Tile:
tile.set_refine_mode("ADD")
tileset = TileSet(geometric_error=tile.geometric_error)
tileset.root_tile = copy.deepcopy(tile)
tileset_name = Path(f"tileset.{split_name}.json")
tileset.write_as_json(folder / tileset_name)
tile.content_uri = tileset_name
tile.children = []
return tile