Efficient Neuron Segmentation in Electron Microscopy by Affinity-Guided Queries

Authors: Hang Chen, Chufeng Tang, Xiao Li, Xiaolin Hu
Primary area: applications to neuroscience & cognitive science
Keywords: Neuron Segmentation; Biomedical Image Segmentation; Electron Microscopy Image

Abstract:
Accurate segmentation of neurons in electron microscopy (EM) images plays a crucial role in understanding the intricate wiring patterns of the brain. Existing automatic neuron segmentation methods rely on traditional clustering algorithms, where affinities are predicted first, and then watershed and post-processing algorithms are applied to yield segmentation results. Due to the nature of watershed algorithm, this paradigm has deficiency in both prediction quality and speed. Inspired by recent advances in natural image segmentation, we propose to use query-based methods to address the problem because they do not necessitate watershed algorithms. However, we find that directly applying existing query-based methods faces great challenges due to the large memory requirement of the 3D data and considerably different morphology of neurons. To tackle these challenges, we introduce affinity-guided queries and integrate them into a lightweight query-based framework. Specifically, we first predict affinities with a lightweight branch, which provides coarse neuron structure information. The affinities are then used to construct affinity-guided queries, facilitating segmentation with bottom-up cues. These queries, along with additional learnable queries, interact with the image features to directly predict the final segmentation results. Experiments on benchmark datasets demonstrated that our method achieved better results over state-of-the-art methods with a 2$\sim$3$\times$ speedup in inference. Code is available at https://github.com/chenhang98/AGQ.

Decision: accept (poster)
Reviewer scores: 6, 5, 8, 6

Reviews and discussion:
<thread>
creation date:
1730018780363

thread_type:
official_review

<thread_object>
creator: reviewer_wcu2

<content>
summary:
The paper presents an innovative approach to neuron segmentation in electron microscopy (EM) images, a critical task in connectomics research that aims to reconstruct neural connections in the brain. The authors propose an affinity-guided query (AGQ) method, which introduces a lightweight query-based framework to address the limitations of traditional clustering algorithms.

strengths:
The authors introduce AGQ to incorporate coarse neuron structure information into the query generation process, reducing the learning difficulty and improving segmentation accuracy.

weaknesses:
While the work is commendable and addresses significant challenges in the field, there are areas where further development and clarification could enhance the contribution and impact of the research. Here are some specific points for consideration:

1. The paper could benefit from a more comprehensive comparison with the latest state-of-the-art (SOTA) methods, such as those referenced in [1][2]. Including these comparisons will provide a clearer picture of how the AGQ method stacks up against the most recent advancements in the field.
2. Figure 1 (b) suggests that the proposed method relies on agglomeration operations from 'waterz' post-processing. This seems to contradict the claim of directly predicting segmentation results. The authors should clarify whether the proposed method truly bypasses the need for watershed-based post-processing or if it still relies on some aspects of it.
3. The experiments presented in Table 2, while informative, may not be based on benchmarks that are widely recognized in the community. Utilizing more established benchmarks would strengthen the paper's findings and allow for better comparisons with other methods.
4. The motivation behind the proposed affinity-guided queries appears to have similarities with modules found in other works, such as [3]. The authors should elaborate on the novel aspects of their approach and how it differs from or improves upon existing techniques.
5. 'typically the neuron gap' in line168 is not unprofessional

[1]Sheridan A, Nguyen T M, Deb D, et al. Local shape descriptors for neuron segmentation[J]. Nature methods, 2023, 20(2): 295-303.
[2]Liu X, Cai M, Chen Y, et al. Cross-dimension affinity distillation for 3d em neuron segmentation[C]//2024 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR). IEEE Computer Society, 2024: 11104-11113.
[3]Sun R, Luo N, Pan Y, et al. Appearance Prompt Vision Transformer for Connectome Reconstruction[C]//IJCAI. 2023: 1423-1431.

questions:
Comparison with State-of-the-Art Baselines

Clarification on the Use of Agglomerations

Use of Well-Established Benchmarks


Originality of Affinity-Guided Queries
</content>
</thread_object>

</thread>
<thread>
creation date:
1730585135377

thread_type:
official_review

<thread_object>
creator: reviewer_pzeh

<content>
summary:
This paper presents an approach for neuron segmentation in electron microscopy images by a query-based segmentation framework. Traditional methods often rely on predicting affinities followed by watershed algorithms, which can introduce artifacts and inefficiencies. The proposed method predicts affinities using a lightweight branch to obtain coarse neuron structure information, which is then used to construct Affinity-Guided Queries (AGQs). These queries, along with learnable queries, interact with image features to directly predict the final segmentation results, eliminating the need for watershed algorithms. The method achieves better accuracy and a 2–3× speedup in inference time.

strengths:
1. Novel and intersting idea that uses a query-based framework to eliminate the need for watershed algorithms used in traditional methods. The proposed method greatly improves efficiency.
2. Very well written manuscript. The paper clearly articulates the shortcomings of existing methods and motivates the need for the query based approach.
3. Comprehensive ablation studies show the contributions of different modules.

weaknesses:
1. Despite the successful and interesting use of AGQs, I have difficulties to see why similar results couldn't be achieved by inputting image features, predicted affinities, and learnable queries directly into the neuron decoder (with modified architecture). While I understand that predicted affinities are used to create AGQ, I wonder if the LQ could "attend" to the predicted affinity map to construct the necessary queries themselves. It appears that the current neuron decoder is attempting to "rescue" the low-quality segmentation by AGQ in $M_{coarse}$. 
2. The choice of $a$ in ConnectedComponent algorithm in eq 8 is not reported and discussed.
3. The computational efficiency is evaluated in terms of inference latency. It would be great if the manuscript could educate the readers about the model's memory consumption, which is important for deployment.

questions:
1. There are, in fact, parallel/CUDA implementations of watershed algorithms. How would the proposed method compare to these optimized implementations in terms of efficiency?
2. For the Transformer Decoder implementation in Table 3, how do you input the image features and queries?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730655596480

thread_type:
official_review

<thread_object>
creator: reviewer_nq34

<content>
summary:
This paper introduces Affinity-Guided Queries (AGQ) as a way to apply query-based instance segmentation to the field of connectomics. AGQ predicts the segments directly, bypassing postprocessing steps such as watershed used in some other approaches. The method is evaluated on two small-scale datasets, showing good results in terms of efficiency and accuracy.

strengths:
- Evaluation on two datasets acquired with different EM techniques.
- Novel approach to segmentation of 3d EM data (query-based model).
- Good results in terms of efficiency and promising accuracy (but see comments about metrics below).
- Ablations showing the impact of the various modules.
- The paper is clearly written and well-organized.

weaknesses:
- Evaluation is performed on a very small scale (tens of um^3). Experience in the field of connectomics shows that such results are often not predictive of reconstruction quality over larger volumes. One of the datasets (ZEBRAFINCH) provides full volume (1M um^3 -scale) skeleton tracings that make it possible to compute topological metrics, which can then be compared to those reported https://doi.org/10.1038/s41592-022-01711-z and https://doi.org/10.1038/s41592-018-0049-4. Having these available, would make it much clearer how the proposed method compares to the state of the art in a practical setting.
- VOI advantage seems to be driven by VOI_split reduction at the cost of VOI_merge. If the mergers happen at the level of supervoxels, the results are in practice less useful for proofreading, where many systems do not make it possible to manually fix this type of error.
- The model is quite architecturally complex, with different losses, modules, and an explicit affinity modeling branch.

questions:
- Please include VOI_split and VOI_merge breakdown in the main text instead of only in the appendices.
- The intro, abstract and Fig. 1 make it sound as if all currently used neuron reconstruction methods rely on watershed. This is not true (the authors are aware of at least https://doi.org/10.1038/s41592-018-0049-4, which they cite later; there also exist other approaches such as https://openaccess.thecvf.com/content_CVPR_2019/html/Meirovitch_Cross-Classification_Clustering_An_Efficient_Multi-Object_Tracking_Technique_for_3-D_Instance_CVPR_2019_paper.html) and should be more clearly discussed earlier in the paper. The proposed model is far from from the first approach that directly converts 3d EM images into a segmentation.
- 0.08 mm^3 is a fairly specific number to quote as a "typical size of studied volume", and quotes a 6y old paper. It would be best to cite more works here, or just refrain from making statements about what is typical altogether. I would also suggest reformulating the comment about processing requiring "130 days of runtime", as this is highly specific to your specific cluster/hardware configuration.
- How often does your algorithm generate disconnected components for the same predicted object? (Eq. 3 does not guarantee connectivity)
- What's the motivation for predicting affinities and then immediately averaging them in Eq. 7? Why not predict a boundary map?
- Does N_a in Eq. 8 vary between passes through the network? If so, how does that interact the number of learned queries -- is N_l or N kept constant? How is N_l chosen in your work?
- Your method does not require block overlaps to generate consistent results. But wouldn't the additional spatial context help in making the predictions better? Have you explored this?
- What are the error bars for your results in Tab. 1?
- How does your network capacity (parameter count, FLOPS) compare with your baselines?
- Your examples in Fig. 5 focus on dendrites with broken off spines. Could you please also provide a representative figure for axons? (perhaps in the supplement)
</content>
</thread_object>

</thread>
<thread>
creation date:
1730716943714

thread_type:
official_review

<thread_object>
creator: reviewer_ua9h

<content>
summary:
The paper presents a novel method for neuron segmentation in electron microscopy (EM) images using affinity-guided queries within a query-based framework. The approach aims to overcome the limitations of traditional methods that rely on watershed algorithms by directly predicting segmentation results, achieving improved accuracy and efficiency. The method is evaluated on benchmark datasets, showing competitive performance.

strengths:
1. The use of affinity-guided queries in a query-based framework is a novel contribution to the field of neuron segmentation.
2. The method demonstrates a significant speedup (2-3x) over traditional methods, a notable advantage for large-scale data processing.
3. The query-based framework directly predicts the final segmentation results, avoiding the inaccuracy and inefficiency associated with the watershed algorithm used in traditional methods.

weaknesses:
1. The comparison with state-of-the-art methods is limited to a few selected approaches. A more extensive comparison with a broader range of recent methods would strengthen the paper's contributions, such as [1]. 
2. The paper does not provide a strong theoretical foundation for why affinity-guided queries outperform other query-based methods, especially in terms of handling neuron morphology.  It is better to give a more detailed explanation or analysis of the mechanisms by which affinity-guided queries better capture neuron morphology compared to other query-based approaches.

[1] Zhang, Yanchao, et al. "SegNeuron: 3D Neuron Instance Segmentation in Any EM Volume with a Generalist Model." International Conference on Medical Image Computing and Computer-Assisted Intervention. Cham: Springer Nature Switzerland, 2024.

questions:
1. How does the method handle variations in image quality or noise levels in EM datasets?
2. What are the specific computational requirements for training and inference, and how do these scale with larger volumes?
3. How do affinity-guided queries compare with other forms of guided queries in terms of performance and efficiency?
4. The paper mentions that the proposed method does not require overlap prediction to address border effects. However, it is unclear how the method handles the continuity and consistency of segmentation results across block boundaries. Can the authors elaborate on this aspect and provide visual examples?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732335581261

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
The authors sincerely appreciate the time and effort the reviewers put into this manuscript, which helps us a lot to improve the quality of the revised paper. It is pretty encouraging that the reviewers found the proposed method is *novel* [uA9h, nQ34, pZEH] and *interesting* [pZEH], and is *commendable and addresses significant challenges in the field* [wCU2], the experiments showed *a notable advantage* [uA9h] and *good results in terms of efficiency and promising accuracy* [nQ34], with *comprehensive ablation studies* [pZEH], and the written is *clearly* and *well-organized* [nQ34]. We carefully address the concerns below, and the paper has been revised accordingly (highlighted in blue).
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposes a network architecture for instance segmentation of electron microscopy (EM) images. Unlike previous methods with two stages of affinity prediction and watershed-based segmentation, the method here is end-to-end with several losses minimized. Unlike standard query-based segmentation, the network here includes a new module of affinity-based queries. The proposed pipeline is technically sound and is efficient during inference compared to standard query-based segmentation networks.
Three reviewers gave positive ratings with one reviewer giving a negative rating. For the review with a negative rating, the major concerns are about the complexity of the pipeline with several losses and the lack of experiments on large-scale 3D volume. However, other reviewers didn't find the proposed pipeline too complex or impractical. Also, the results on medium-scale 3D datasets demonstrated the effectiveness and efficiency of the proposed method.
Overall, the strengths overweight the weaknesses by a large margin, and the AC recommends acceptance of this paper.