nps-cli 0.1.3__tar.gz → 0.1.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
nps_cli-0.1.5/PKG-INFO ADDED
@@ -0,0 +1,119 @@
1
+ Metadata-Version: 2.4
2
+ Name: nps-cli
3
+ Version: 0.1.5
4
+ Summary: A simple CLI for sampling point clouds from large volumetric datasets
5
+ Author: Jakob Troidl
6
+ License-Expression: MIT
7
+ Classifier: Programming Language :: Python :: 3.11
8
+ Requires-Python: >=3.11
9
+ Description-Content-Type: text/markdown
10
+ Requires-Dist: click
11
+ Requires-Dist: cloud-volume!=12.13.*
12
+ Requires-Dist: volara>=2.0.1
13
+ Requires-Dist: pocaduck
14
+ Provides-Extra: dev
15
+ Requires-Dist: pytest; extra == "dev"
16
+ Requires-Dist: black; extra == "dev"
17
+ Requires-Dist: build; extra == "dev"
18
+ Requires-Dist: twine; extra == "dev"
19
+ Requires-Dist: ruff; extra == "dev"
20
+
21
+ [![pypi](https://badge.fury.io/py/nps-cli.svg)](https://badge.fury.io/py/nps-cli)
22
+ [![Python Version](https://img.shields.io/pypi/pyversions/nps-cli.svg)](https://pypi.org/project/nps-cli/)
23
+
24
+ # nps: distributed point cloud sampling
25
+
26
+ `nps` samples point clouds of every segment in a large segmentation volumes
27
+ * Compatible with [CloudVolume](https://github.com/seung-lab/cloud-volume) (precomputed or graphene).
28
+ * It processes the volume blockwise on local, LSF or Slurm workers and writes the points in [pocaduck](https://github.com/JaneliaSciComp/pocaduck) format.
29
+
30
+ ```bash
31
+ pip install nps-cli
32
+ nps --help
33
+ ```
34
+
35
+ ## Usage
36
+
37
+ ```bash
38
+ # Whole volume
39
+ nps --cv-path precomputed://gs://neuroglancer-janelia-flyem-hemibrain/v1.0/segmentation
40
+
41
+ # Subvolume (begin_x begin_y begin_z end_x end_y end_z, in voxels), 1% of each segment
42
+ nps --cv-path precomputed://gs://neuroglancer-janelia-flyem-hemibrain/v1.0/segmentation \
43
+ --bbox 15347 19712 18606 15859 20224 19118 --fraction 0.01
44
+
45
+ # Also store the supervoxel ID of every point, looked up in a separate volume
46
+ # that shares the voxel grid of --cv-path
47
+ nps --cv-path precomputed://gs://<bucket>/agglomerated \
48
+ --svid-cv-path precomputed://gs://<bucket>/supervoxels
49
+
50
+ # Graphene: read supervoxels from the same volume (agglomerate=False)
51
+ nps --cv-path graphene://https://<server>/segmentation/table/<table> --sample_svids
52
+ ```
53
+
54
+ Key options (run `nps --help` for all of them):
55
+
56
+ | Option | Default | Meaning |
57
+ | --- | --- | --- |
58
+ | `--fraction` | `0.001` | Fraction of each segment's voxels to sample (at least 1 per segment and block). |
59
+ | `--mip` | `0` | Resolution level to sample at. Coordinates are voxels of this mip. |
60
+ | `--bbox` | whole volume | ROI in voxels. |
61
+ | `--block-size` | `128 128 128` | Block size per task, in voxels. |
62
+ | `--worker-type` / `--num-workers` / `--cpus-per-worker` / `--queue` | `LocalWorker` / `8` / `4` / `local` | Where and how the blocks are processed. |
63
+ | `--fill-missing` | off | Treat missing chunks as empty instead of raising an error. |
64
+ | `--overwrite` | off | Delete points of an earlier run in `--output-dir` instead of refusing to start. |
65
+ | `--skip-empty` / `--block-mask` / `--mask-mip` | on | Skipping empty blocks, see below. |
66
+
67
+ ### Output
68
+
69
+ ```
70
+ <output-dir>/
71
+ block_mask.npy # occupancy mask from the empty-block pre-scan
72
+ points/
73
+ unified_index.db # pocaduck index over all workers
74
+ worker_<id>/... # parquet files + per-worker index
75
+ ```
76
+
77
+ Each point is `[x, y, z]`, or `[x, y, z, svid]` when SVIDs are sampled, stored under its segment label.
78
+
79
+ `nps` refuses to write into an `--output-dir` that already holds points, since every
80
+ run re-samples all blocks and old points would show up as duplicates. Pass
81
+ `--overwrite` to delete them first.
82
+
83
+ ### Skipping empty blocks
84
+
85
+ Before sampling, `nps` scans the ROI at a coarse mip (`--mask-mip`, default the
86
+ coarsest one) and never dispatches blocks without labels. The mask is dilated by
87
+ one block, so thin structures lost in downsampling are still sampled. Use
88
+ `--no-skip-empty` to disable the scan, or `--block-mask <path>` to reuse a
89
+ `block_mask.npy` from an earlier run with the same `--bbox`, `--block-size` and `--mip`.
90
+
91
+ ## Reading point clouds
92
+
93
+ ```python
94
+ from pocaduck import Query, StorageConfig
95
+
96
+ query = Query(storage_config=StorageConfig(base_path="<output-dir>/points"))
97
+ labels = query.get_labels()
98
+ points = query.get_points(label=12345) # all points of one segment
99
+ query.close()
100
+ ```
101
+
102
+ See the [pocaduck docs](https://github.com/JaneliaSciComp/pocaduck?tab=readme-ov-file#querying-point-clouds)
103
+ for more, including [optimizing](https://github.com/JaneliaSciComp/pocaduck?tab=readme-ov-file#running-the-optimization-pipeline) the storage for fast reads.
104
+
105
+ ## Example scripts
106
+
107
+ | Script | What it does |
108
+ | --- | --- |
109
+ | [`liconn_local.sh`](scripts/liconn_local.sh) | Points only, local workers |
110
+ | [`liconn_svids_local.sh`](scripts/liconn_svids_local.sh) | Points + SVIDs from a separate precomputed volume |
111
+ | [`flywire_cluster.sh`](scripts/flywire_cluster.sh) | Points + SVIDs for FlyWire on LSF (`LABEL_NAME=flywire_v141_m783` for the 2023 snapshot) |
112
+
113
+ ## Release
114
+
115
+ ```bash
116
+ pip install -e ".[dev]"
117
+ python -m build
118
+ twine upload dist/*
119
+ ```
@@ -0,0 +1,99 @@
1
+ [![pypi](https://badge.fury.io/py/nps-cli.svg)](https://badge.fury.io/py/nps-cli)
2
+ [![Python Version](https://img.shields.io/pypi/pyversions/nps-cli.svg)](https://pypi.org/project/nps-cli/)
3
+
4
+ # nps: distributed point cloud sampling
5
+
6
+ `nps` samples point clouds of every segment in a large segmentation volumes
7
+ * Compatible with [CloudVolume](https://github.com/seung-lab/cloud-volume) (precomputed or graphene).
8
+ * It processes the volume blockwise on local, LSF or Slurm workers and writes the points in [pocaduck](https://github.com/JaneliaSciComp/pocaduck) format.
9
+
10
+ ```bash
11
+ pip install nps-cli
12
+ nps --help
13
+ ```
14
+
15
+ ## Usage
16
+
17
+ ```bash
18
+ # Whole volume
19
+ nps --cv-path precomputed://gs://neuroglancer-janelia-flyem-hemibrain/v1.0/segmentation
20
+
21
+ # Subvolume (begin_x begin_y begin_z end_x end_y end_z, in voxels), 1% of each segment
22
+ nps --cv-path precomputed://gs://neuroglancer-janelia-flyem-hemibrain/v1.0/segmentation \
23
+ --bbox 15347 19712 18606 15859 20224 19118 --fraction 0.01
24
+
25
+ # Also store the supervoxel ID of every point, looked up in a separate volume
26
+ # that shares the voxel grid of --cv-path
27
+ nps --cv-path precomputed://gs://<bucket>/agglomerated \
28
+ --svid-cv-path precomputed://gs://<bucket>/supervoxels
29
+
30
+ # Graphene: read supervoxels from the same volume (agglomerate=False)
31
+ nps --cv-path graphene://https://<server>/segmentation/table/<table> --sample_svids
32
+ ```
33
+
34
+ Key options (run `nps --help` for all of them):
35
+
36
+ | Option | Default | Meaning |
37
+ | --- | --- | --- |
38
+ | `--fraction` | `0.001` | Fraction of each segment's voxels to sample (at least 1 per segment and block). |
39
+ | `--mip` | `0` | Resolution level to sample at. Coordinates are voxels of this mip. |
40
+ | `--bbox` | whole volume | ROI in voxels. |
41
+ | `--block-size` | `128 128 128` | Block size per task, in voxels. |
42
+ | `--worker-type` / `--num-workers` / `--cpus-per-worker` / `--queue` | `LocalWorker` / `8` / `4` / `local` | Where and how the blocks are processed. |
43
+ | `--fill-missing` | off | Treat missing chunks as empty instead of raising an error. |
44
+ | `--overwrite` | off | Delete points of an earlier run in `--output-dir` instead of refusing to start. |
45
+ | `--skip-empty` / `--block-mask` / `--mask-mip` | on | Skipping empty blocks, see below. |
46
+
47
+ ### Output
48
+
49
+ ```
50
+ <output-dir>/
51
+ block_mask.npy # occupancy mask from the empty-block pre-scan
52
+ points/
53
+ unified_index.db # pocaduck index over all workers
54
+ worker_<id>/... # parquet files + per-worker index
55
+ ```
56
+
57
+ Each point is `[x, y, z]`, or `[x, y, z, svid]` when SVIDs are sampled, stored under its segment label.
58
+
59
+ `nps` refuses to write into an `--output-dir` that already holds points, since every
60
+ run re-samples all blocks and old points would show up as duplicates. Pass
61
+ `--overwrite` to delete them first.
62
+
63
+ ### Skipping empty blocks
64
+
65
+ Before sampling, `nps` scans the ROI at a coarse mip (`--mask-mip`, default the
66
+ coarsest one) and never dispatches blocks without labels. The mask is dilated by
67
+ one block, so thin structures lost in downsampling are still sampled. Use
68
+ `--no-skip-empty` to disable the scan, or `--block-mask <path>` to reuse a
69
+ `block_mask.npy` from an earlier run with the same `--bbox`, `--block-size` and `--mip`.
70
+
71
+ ## Reading point clouds
72
+
73
+ ```python
74
+ from pocaduck import Query, StorageConfig
75
+
76
+ query = Query(storage_config=StorageConfig(base_path="<output-dir>/points"))
77
+ labels = query.get_labels()
78
+ points = query.get_points(label=12345) # all points of one segment
79
+ query.close()
80
+ ```
81
+
82
+ See the [pocaduck docs](https://github.com/JaneliaSciComp/pocaduck?tab=readme-ov-file#querying-point-clouds)
83
+ for more, including [optimizing](https://github.com/JaneliaSciComp/pocaduck?tab=readme-ov-file#running-the-optimization-pipeline) the storage for fast reads.
84
+
85
+ ## Example scripts
86
+
87
+ | Script | What it does |
88
+ | --- | --- |
89
+ | [`liconn_local.sh`](scripts/liconn_local.sh) | Points only, local workers |
90
+ | [`liconn_svids_local.sh`](scripts/liconn_svids_local.sh) | Points + SVIDs from a separate precomputed volume |
91
+ | [`flywire_cluster.sh`](scripts/flywire_cluster.sh) | Points + SVIDs for FlyWire on LSF (`LABEL_NAME=flywire_v141_m783` for the 2023 snapshot) |
92
+
93
+ ## Release
94
+
95
+ ```bash
96
+ pip install -e ".[dev]"
97
+ python -m build
98
+ twine upload dist/*
99
+ ```
@@ -0,0 +1,2 @@
1
+ from .sample_points import SamplePoints as SamplePoints
2
+ from .sample_points import consolidate_indexes as consolidate_indexes
@@ -0,0 +1,141 @@
1
+ """
2
+ Per-block occupancy mask computed from a coarse mip level.
3
+
4
+ Reading a chunk from a (sharded) precomputed volume costs roughly the same
5
+ regardless of how many voxels are requested, so checking each block for
6
+ emptiness at sampling time does not pay off. Instead we make one pass over the
7
+ ROI at the coarsest available mip, where a single chunk covers hundreds of
8
+ sampling blocks, and record for every block whether it contains any label.
9
+ The mask is then used by :class:`SamplePoints` to skip empty blocks inside the
10
+ daisy scheduler, so workers never touch them.
11
+ """
12
+
13
+ from concurrent.futures import ThreadPoolExecutor
14
+ from itertools import product
15
+
16
+ import numpy as np
17
+ from cloudvolume import CloudVolume
18
+ from funlib.geometry import Coordinate, Roi
19
+ from volara.datasets import CloudVolumeWrapper
20
+
21
+ # Approximate edge length, in coarse voxels, of the region read per job.
22
+ _TARGET_TILE_VOXELS = 128
23
+
24
+
25
+ def _open_coarse(store: str, mip: int, timestamp: int, agglomerate: bool) -> CloudVolume:
26
+ return CloudVolume(
27
+ store,
28
+ mip=mip,
29
+ use_https=True,
30
+ agglomerate=agglomerate,
31
+ timestamp=timestamp,
32
+ fill_missing=True, # missing chunks are, by definition, empty
33
+ bounded=False, # tiles at the ROI edge may poke out of the volume
34
+ )
35
+
36
+
37
+ def _tile_job(args):
38
+ """Read one coarse tile and return which of its blocks are non-empty."""
39
+ vol, tile_index, block_lo, block_hi, roi_begin, roi_end, block_size, factor = args
40
+ block_lo, block_hi = np.array(block_lo), np.array(block_hi)
41
+ roi_begin, roi_end = np.array(roi_begin), np.array(roi_end)
42
+ block_size, factor = np.array(block_size), np.array(factor)
43
+
44
+ # Coarse voxel range covering all blocks in this tile (clipped to the ROI).
45
+ fine_lo = roi_begin + block_lo * block_size
46
+ fine_hi = np.minimum(roi_begin + block_hi * block_size, roi_end)
47
+ coarse_lo = fine_lo // factor
48
+ coarse_hi = -(-fine_hi // factor) # ceil
49
+ data = np.asarray(
50
+ vol[tuple(slice(int(a), int(b)) for a, b in zip(coarse_lo, coarse_hi))]
51
+ ).squeeze(axis=-1)
52
+
53
+ out = np.zeros(block_hi - block_lo, dtype=bool)
54
+ for idx in product(*(range(n) for n in out.shape)):
55
+ b = block_lo + np.array(idx)
56
+ lo = (roi_begin + b * block_size) // factor - coarse_lo
57
+ hi = -(-(np.minimum(roi_begin + (b + 1) * block_size, roi_end)) // factor) - coarse_lo
58
+ out[idx] = data[tuple(slice(int(a), int(c)) for a, c in zip(lo, hi))].any()
59
+ return tile_index, out
60
+
61
+
62
+ def _dilate(mask: np.ndarray, radius: int) -> np.ndarray:
63
+ """Binary dilation with a cube of the given radius (no scipy dependency)."""
64
+ if radius <= 0:
65
+ return mask
66
+ padded = np.pad(mask, radius)
67
+ out = np.zeros_like(mask)
68
+ for shift in product(range(2 * radius + 1), repeat=mask.ndim):
69
+ sl = tuple(slice(s, s + n) for s, n in zip(shift, mask.shape))
70
+ out |= padded[sl]
71
+ return out
72
+
73
+
74
+ def compute_block_mask(
75
+ labels: CloudVolumeWrapper,
76
+ roi: Roi,
77
+ block_size: Coordinate,
78
+ mask_mip: int | None = None,
79
+ num_workers: int = 8,
80
+ dilate: int = 1,
81
+ ) -> np.ndarray:
82
+ """
83
+ Compute a boolean array, indexed by block index within ``roi``, that is
84
+ True for blocks containing at least one non-zero label.
85
+
86
+ Args:
87
+ labels: The label volume that will be sampled.
88
+ roi: ROI to be processed, in voxel coordinates of ``labels.mip``
89
+ (i.e. the frame daisy blocks are expressed in).
90
+ block_size: Sampling block size in the same frame.
91
+ mask_mip: Mip level to read for the occupancy check. Defaults to the
92
+ coarsest mip available.
93
+ num_workers: Number of threads reading coarse tiles. Reads are I/O and
94
+ decompression bound and release the GIL, so threads scale as well
95
+ as processes here while avoiding multiprocessing entirely (forked
96
+ and spawned pools both died with BrokenProcessPool on LSF nodes).
97
+ dilate: Grow the mask by this many blocks in every direction. Guards
98
+ against thin structures that vanish when downsampled.
99
+ """
100
+ base = CloudVolume(str(labels.store), mip=labels.mip, use_https=True)
101
+ if mask_mip is None:
102
+ mask_mip = max(base.available_mips)
103
+ factor = np.asarray(base.mip_resolution(mask_mip)) / np.asarray(
104
+ base.mip_resolution(labels.mip)
105
+ )
106
+ assert np.allclose(factor, np.round(factor)), (
107
+ f"mip {mask_mip} resolution is not an integer multiple of mip {labels.mip}"
108
+ )
109
+ factor = np.round(factor).astype(int)
110
+
111
+ roi_begin = np.array(roi.begin)
112
+ roi_end = np.array(roi.end)
113
+ block_size = np.array(block_size)
114
+ n_blocks = -(-(roi_end - roi_begin) // block_size)
115
+
116
+ # Blocks per tile per axis: enough that a tile is ~_TARGET_TILE_VOXELS wide.
117
+ blocks_per_tile = np.maximum(
118
+ 1, _TARGET_TILE_VOXELS // np.maximum(1, block_size // factor)
119
+ )
120
+ n_tiles = -(-n_blocks // blocks_per_tile)
121
+
122
+ coarse = _open_coarse(
123
+ str(labels.store), mask_mip, labels.timestamp, labels.agglomerate
124
+ )
125
+ jobs = []
126
+ for t in product(*(range(n) for n in n_tiles)):
127
+ t = np.array(t)
128
+ lo = t * blocks_per_tile
129
+ hi = np.minimum(lo + blocks_per_tile, n_blocks)
130
+ jobs.append(
131
+ (coarse, tuple(t), tuple(lo), tuple(hi), tuple(roi_begin),
132
+ tuple(roi_end), tuple(block_size), tuple(factor))
133
+ )
134
+
135
+ mask = np.zeros(n_blocks, dtype=bool)
136
+ with ThreadPoolExecutor(max_workers=num_workers) as pool:
137
+ for tile_index, sub in pool.map(_tile_job, jobs):
138
+ lo = np.array(tile_index) * blocks_per_tile
139
+ mask[tuple(slice(a, a + n) for a, n in zip(lo, sub.shape))] = sub
140
+
141
+ return _dilate(mask, dilate)
@@ -0,0 +1,248 @@
1
+ import glob
2
+ import logging
3
+ import os
4
+ import shutil
5
+ import uuid
6
+ from contextlib import contextmanager
7
+ from typing import Literal
8
+
9
+ import daisy
10
+ import duckdb
11
+ import numpy as np
12
+ import pandas as pd
13
+ import pyarrow as pa
14
+ import pyarrow.parquet as pq
15
+ from daisy import Block
16
+ from funlib.geometry import Coordinate, Roi
17
+ from volara.blockwise import BlockwiseTask
18
+ from volara.datasets import CloudVolumeWrapper, Dataset
19
+ from volara.utils import PydanticCoordinate
20
+
21
+ _INDEX_SCHEMA = """
22
+ CREATE TABLE IF NOT EXISTS {name} (
23
+ label UBIGINT,
24
+ block_id VARCHAR,
25
+ file_path VARCHAR,
26
+ point_count UBIGINT
27
+ )
28
+ """
29
+
30
+
31
+ class SamplePoints(BlockwiseTask):
32
+ task_type: Literal["sample_pc"] = "sample_pc"
33
+ out_dir: str
34
+ labels: CloudVolumeWrapper
35
+ svids: CloudVolumeWrapper | None = None
36
+ block_mask: str | None = None
37
+ block_size: PydanticCoordinate
38
+ fraction: float
39
+ fit: Literal["shrink"] = "shrink"
40
+ read_write_conflict: Literal[False] = False
41
+
42
+ @property
43
+ def task_name(self) -> str:
44
+ return f"{self.labels.name}-{self.task_type}"
45
+
46
+ @property
47
+ def write_roi(self) -> Roi:
48
+ total_roi = self.labels.array("r").roi
49
+ if self.roi is not None:
50
+ total_roi = total_roi.intersect(self.roi)
51
+ return total_roi
52
+
53
+ @property
54
+ def voxel_size(self) -> Coordinate:
55
+ return self.labels.voxel_size
56
+
57
+ @property
58
+ def write_size(self) -> Coordinate:
59
+ return self.block_size
60
+
61
+ @property
62
+ def context_size(self) -> Coordinate:
63
+ return Coordinate((0,) * self.write_size.dims)
64
+
65
+ def drop_artifacts(self):
66
+ # init() recreates the directory when the task is run.
67
+ shutil.rmtree(self.out_dir, ignore_errors=True)
68
+
69
+ def check_block_func(self):
70
+ """
71
+ Daisy calls this in the scheduler before dispatching a block; returning
72
+ True means "already done". If a block mask is given, report empty blocks
73
+ as done so they are never sent to a worker.
74
+ """
75
+ base = super().check_block_func()
76
+ if self.block_mask is None:
77
+ return base
78
+
79
+ mask = np.load(self.block_mask)
80
+ roi_begin = np.array(self.write_roi.begin)
81
+ block_size = np.array(self.block_size)
82
+
83
+ def check_block(block: Block) -> bool:
84
+ index = (np.array(block.write_roi.begin) - roi_begin) // block_size
85
+ if not mask[tuple(index)]:
86
+ return True
87
+ return base(block)
88
+
89
+ return check_block
90
+
91
+ @property
92
+ def output_datasets(self) -> list[Dataset]:
93
+ return []
94
+
95
+ def sample_pc_in_block(self, block: Block, labels: np.ndarray, svids: np.ndarray | None, offset: Coordinate, writer: "_PointWriter"):
96
+ block_id = block.block_id[1]
97
+
98
+ seg_ids, pts = self.sample_segment_points(labels, self.fraction)
99
+ logging.info(f"sampled {len(pts)} points in block {block_id}")
100
+
101
+ packed_pts = (pts + np.array(offset)).astype(np.uint64)
102
+ if self.svids is not None:
103
+ sampled_svids = svids[pts[:, 0], pts[:, 1], pts[:, 2]].reshape(-1, 1)
104
+ packed_pts = np.concatenate((packed_pts, sampled_svids.astype(np.uint64)), axis=1)
105
+
106
+ writer.write(block_id, seg_ids, packed_pts)
107
+
108
+ def init(self):
109
+ os.makedirs(self.out_dir, exist_ok=True)
110
+
111
+ @contextmanager
112
+ def process_block_func(self):
113
+ data = self.labels.array("r")
114
+ if self.svids is not None:
115
+ supervoxels = self.svids.array("r")
116
+
117
+ try:
118
+ worker_id = daisy.Context.from_env()["worker_id"]
119
+ except KeyError:
120
+ worker_id = 0
121
+ # One writer per worker process instead of a pocaduck Ingestor per
122
+ # block, which re-read and re-wrote the worker's whole parquet file
123
+ # and ran two DuckDB queries per segment on every block.
124
+ writer = _PointWriter(self.out_dir, worker_id)
125
+
126
+ def process_block(block: Block):
127
+ # Index with the Roi (world/absolute voxel coords), not with
128
+ # to_slices(): numpy-style keys on a funlib Array are zero-based and
129
+ # would read from the wrong place on volumes with a voxel offset.
130
+ # Only squeeze the channel axis; edge blocks can be 1 voxel thick.
131
+ labels = np.asarray(data[block.write_roi]).squeeze(axis=-1)
132
+
133
+ offset = block.write_roi.get_begin()
134
+
135
+ if self.svids is not None:
136
+ s = np.asarray(supervoxels[block.write_roi]).squeeze(axis=-1)
137
+ assert s.shape == labels.shape, (
138
+ f"label block {labels.shape} and svid block {s.shape} differ in "
139
+ f"{block.write_roi}; the label and supervoxel volumes must share "
140
+ "resolution and offset at the chosen mip"
141
+ )
142
+ else:
143
+ s = None
144
+
145
+ self.sample_pc_in_block(block, labels, s, offset, writer)
146
+
147
+ # Index consolidation happens once in the driver after all blocks are
148
+ # done (see consolidate_indexes), not concurrently in every worker.
149
+ yield process_block
150
+
151
+ def sample_segment_points(self, labels, fraction, background=0):
152
+ """
153
+ Sample max(1, int(n * fraction)) voxels without replacement from every
154
+ segment with n voxels. Returns the segment ID of every sampled point
155
+ and the (N, 3) voxel coordinates within ``labels``.
156
+
157
+ Vectorized: one sort over the foreground voxels instead of one full
158
+ ``labels == seg`` scan per segment.
159
+ """
160
+ flat = labels.ravel()
161
+ idx = np.flatnonzero(flat != background) if background is not None else np.arange(flat.size)
162
+ vals = flat[idx]
163
+
164
+ # Shuffle, then stable-sort by segment: groups segments together with a
165
+ # random order inside each group, so the first k of a group are a
166
+ # uniform sample without replacement.
167
+ perm = np.random.permutation(len(idx))
168
+ perm = perm[np.argsort(vals[perm], kind="stable")]
169
+ idx, vals = idx[perm], vals[perm]
170
+
171
+ _, starts, counts = np.unique(vals, return_index=True, return_counts=True)
172
+ k = np.maximum(1, (counts * fraction).astype(np.int64))
173
+ rank = np.arange(len(vals)) - np.repeat(starts, counts)
174
+ keep = rank < np.repeat(k, counts)
175
+
176
+ coords = np.column_stack(np.unravel_index(idx[keep], labels.shape))
177
+ return vals[keep], coords
178
+
179
+
180
+ class _PointWriter:
181
+ """
182
+ Writes sampled points of one worker in the layout of pocaduck's Ingestor
183
+ (``worker_<id>/data/*.parquet`` plus ``worker_<id>/index_<id>.db``), so
184
+ :func:`consolidate_indexes` and ``pocaduck.Query`` work unchanged.
185
+
186
+ Every block is written (parquet first, then its index rows) before the
187
+ block returns: daisy marks a block done as soon as process_block returns
188
+ and does not wait for worker teardown, so buffering across blocks would
189
+ lose data.
190
+ """
191
+
192
+ def __init__(self, out_dir: str, worker_id):
193
+ worker_dir = os.path.join(out_dir, f"worker_{worker_id}")
194
+ self.data_dir = os.path.join(worker_dir, "data")
195
+ self.db_path = os.path.join(worker_dir, f"index_{worker_id}.db")
196
+ os.makedirs(self.data_dir, exist_ok=True)
197
+ # Unique prefix: a restarted worker can reuse a worker_id and must not
198
+ # overwrite files of a previous run.
199
+ self.prefix = f"{worker_id}-{uuid.uuid4().hex[:8]}"
200
+
201
+ def write(self, block_id, seg_ids: np.ndarray, points: np.ndarray):
202
+ if len(points) == 0:
203
+ return
204
+ file_path = os.path.join(self.data_dir, f"{self.prefix}-{block_id}.parquet")
205
+ data = pa.FixedSizeListArray.from_arrays(pa.array(points.ravel()), points.shape[1])
206
+ table = pa.table({
207
+ "label": seg_ids.astype(np.int64),
208
+ "block_id": np.full(len(points), block_id, dtype=np.int64),
209
+ # Same list<uint64> column the pocaduck Ingestor writes.
210
+ "data": data.cast(pa.list_(pa.uint64())),
211
+ })
212
+ # Write-then-rename so a killed worker never leaves a truncated file.
213
+ pq.write_table(table, file_path + ".tmp")
214
+ os.replace(file_path + ".tmp", file_path)
215
+
216
+ segs, counts = np.unique(seg_ids, return_counts=True)
217
+ index = pd.DataFrame({
218
+ "label": segs.astype(np.uint64),
219
+ "block_id": str(block_id),
220
+ "file_path": file_path,
221
+ "point_count": counts.astype(np.uint64),
222
+ })
223
+ # Short-lived connection: the driver attaches these files for
224
+ # consolidation while worker processes may still be alive.
225
+ con = duckdb.connect(self.db_path)
226
+ con.execute(_INDEX_SCHEMA.format(name="point_cloud_index"))
227
+ con.register("new_rows", index)
228
+ con.execute("INSERT INTO point_cloud_index SELECT * FROM new_rows")
229
+ con.close()
230
+
231
+
232
+ def consolidate_indexes(out_dir: str) -> str:
233
+ """
234
+ Merge all worker indexes into ``unified_index.db`` (what pocaduck's Query
235
+ reads). Replaces ``Ingestor.consolidate_indexes``, which pulls every row
236
+ through Python and was run concurrently by every worker.
237
+ """
238
+ output_path = os.path.join(out_dir, "unified_index.db")
239
+ if os.path.exists(output_path):
240
+ os.remove(output_path)
241
+ con = duckdb.connect(output_path)
242
+ con.execute(_INDEX_SCHEMA.format(name="point_cloud_index"))
243
+ for path in sorted(glob.glob(os.path.join(out_dir, "worker_*", "index_*.db"))):
244
+ con.execute(f"ATTACH '{path}' AS w (READ_ONLY)")
245
+ con.execute("INSERT INTO point_cloud_index SELECT * FROM w.point_cloud_index")
246
+ con.execute("DETACH w")
247
+ con.close()
248
+ return output_path