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 +119 -0
- nps_cli-0.1.5/README.md +99 -0
- nps_cli-0.1.5/nps/blockwise/__init__.py +2 -0
- nps_cli-0.1.5/nps/blockwise/block_mask.py +141 -0
- nps_cli-0.1.5/nps/blockwise/sample_points.py +248 -0
- nps_cli-0.1.5/nps/main.py +117 -0
- nps_cli-0.1.5/nps_cli.egg-info/PKG-INFO +119 -0
- {nps_cli-0.1.3 → nps_cli-0.1.5}/nps_cli.egg-info/SOURCES.txt +1 -0
- {nps_cli-0.1.3 → nps_cli-0.1.5}/nps_cli.egg-info/requires.txt +2 -1
- {nps_cli-0.1.3 → nps_cli-0.1.5}/pyproject.toml +6 -2
- nps_cli-0.1.3/PKG-INFO +0 -100
- nps_cli-0.1.3/README.md +0 -81
- nps_cli-0.1.3/nps/blockwise/__init__.py +0 -1
- nps_cli-0.1.3/nps/blockwise/sample_points.py +0 -132
- nps_cli-0.1.3/nps/main.py +0 -79
- nps_cli-0.1.3/nps_cli.egg-info/PKG-INFO +0 -100
- {nps_cli-0.1.3 → nps_cli-0.1.5}/nps/__init__.py +0 -0
- {nps_cli-0.1.3 → nps_cli-0.1.5}/nps_cli.egg-info/dependency_links.txt +0 -0
- {nps_cli-0.1.3 → nps_cli-0.1.5}/nps_cli.egg-info/entry_points.txt +0 -0
- {nps_cli-0.1.3 → nps_cli-0.1.5}/nps_cli.egg-info/top_level.txt +0 -0
- {nps_cli-0.1.3 → nps_cli-0.1.5}/setup.cfg +0 -0
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
|
+
[](https://badge.fury.io/py/nps-cli)
|
|
22
|
+
[](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
|
+
```
|
nps_cli-0.1.5/README.md
ADDED
|
@@ -0,0 +1,99 @@
|
|
|
1
|
+
[](https://badge.fury.io/py/nps-cli)
|
|
2
|
+
[](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,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
|