Mean Shift Clustering as a Loss Function for Accurate and Segmentation-aware Localization of Macromolecules in Cryo-electron Tomography

Cryo-electron tomography allows us to visualize and analyze the native cellular environment on a molecular level in 3D. To reliably study structures and interactions of proteins, they need to be accurately localized. Recent detection methods train a segmentation network and use post-processing to determine protein locations, often leading to inaccurate and inconsistent locations. We present an end-to-end learning approach for more accurate protein center identification by introducing a differentiable, scoremap-guided Mean Shift clustering implementation. To make training computationally feasible, we sample random cluster points instead of processing the entire image. We show that our Mean Shift loss leads to more accurate cluster center positions compared to the classical Dice loss. When combining these loss functions, we can enhance 3D protein shape preservation and improve clustering with more accurate, localization-focused score maps. In addition to improved protein localization, our method provides more efficient training with sparse ground truth annotations, due to our point sampling strategy.