Skip to content

feat: add LiteMobileUnet5ch and model selection through env - #79

Merged
pFornagiel merged 6 commits into
mainfrom
@pFornagiel/alt-arch
Oct 25, 2026
Merged

pFornagiel merged 6 commits into
mainfrom
@pFornagiel/alt-arch

Conversation

@pFornagiel

Copy link
Copy Markdown
Owner

This PR adds a second 2.5D segmentation network based on architecture suggested by nnUnet plans with mobile blocks and makes it the default. The existing 3-channel MobileNetV3 U-Net stays available, and which model runs is chosen with an environment variable.

Summary

  • Renamed the previous architecture to MobilenetV3Unet3Ch
  • Added new default model: LiteMobileUNet5Ch, 5-channel U-Net with a MobileNet-style encoder, modelled after nnUnet plans based on the dataset. Seven encoder stages of plain inverted residuals (blocks mobilenetv3 uses) go down to 1/64 resolution, and the decoder uses depthwise-separable blocks. Each slice is stacked with two pairs of neighbours, at ±9 and ±27 slices (SliceContext25D5Channel, edges clamped). Z-score normalisation runs on the volume before that stacking, so all five channels share one set of statistics.
  • The rest of the pipeline is unchanged: resample to 0.3 mm, reorient to RPI, paddubg, sliding-window 2D prediction, then reverse pad, drop labels above 32, remove components under 40 mm^3, and resample back to the input grid.

Selecting a model

Model can now be configured through the environmental variables:
TOOTH_MODEL picks lite_mobile_unet_5ch (default) or mobilenetv3_unet_3ch.
TOOTH_MODEL_PATH still points at an ONNX file and overrides the above.

@pFornagiel
pFornagiel merged commit 753e2b7 into main Oct 25, 2026
2 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant