[ENH] Finally, batched bundle recognition#187
Conversation
There was a problem hiding this comment.
Pull request overview
This PR refactors pyAFQ’s bundle recognition pipeline to support batched (chunked) recognition—intended to bound peak memory while segmenting very large tractograms (e.g., TRX) by running “chunk-local” criteria per chunk and then merging survivors for a “global” phase.
Changes:
- Introduces a chunked recognition flow (
recognize_bundles) with per-chunk filtering and a merged global phase. - Replaces the prior
immlib-based preprocessing plan with a lightweightPreprocPlanusing cached properties. - Adjusts TRX handling to sometimes pass a TRX path through segmentation/recognition, plus a few error/logging message updates.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
AFQ/tasks/tractography.py |
Updates the “no streamlines generated” error message text. |
AFQ/tasks/segmentation.py |
Adjusts TRX loading behavior to optionally pass a TRX path through to recognition. |
AFQ/recognition/utils.py |
Updates recognition utilities (logging → tqdm.write) and adds chunk-survivor export/merge helpers. |
AFQ/recognition/recognize.py |
Routes recognition through the new batched recognition entrypoint and adds chunk_size. |
AFQ/recognition/preprocess.py |
Replaces the previous preprocessing-plan machinery with PreprocPlan. |
AFQ/recognition/criteria.py |
Major refactor: splits criteria into chunk-local/global phases and implements recognize_bundles chunking/merge logic. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
36000
changed the title
|
@arokem this is ready for review/merge |
arokem
left a comment
arokem
left a comment
There was a problem hiding this comment.
This is great. I had some small not-so-important comments.
I wonder whether this could be leveraged to parallelize across chunks when memory is available. Could make this even faster?
| return abu.tolerance_mm_to_vox(img, dist_to_waypoint, input_dist_to_atlas) | ||
| @cached_property | ||
| def fgarray(self): | ||
| return np.asarray(abu.resample_tg(self.tg, 20), dtype=np.float32) |
There was a problem hiding this comment.
Should we really hard-code 20 here?
There was a problem hiding this comment.
It is used for cross_midline, prob_map, and start and end points. For start and end points, the resampling number does not matter; DIPY always puts a point at the exact start and end. For prob map and cross_midline, it could potentially matter; however, we have had this hard-coded in the past, and I have never found a reason to tweak it.
Expect to segment ~100M streamlines in ~1 hour on your PC