Can 2D segmentation function be added to UNetR?

Project-MONAI / MONAI

AI Toolkit for Healthcare Imaging

https://monai.io/

Apache License 2.0

5.86k stars 1.08k forks source link

Can 2D segmentation function be added to UNetR? #5062

Closed lbwdruid closed 2 years ago

lbwdruid commented 2 years ago

Thanks for the wonderful work of UNetR. I was wondering if 2D segmentation function could be added to UNetR? Right now, based on my understanding of the codes, the UNetR only works with 3D datasets. I'm working on 2D ultrasound segmentation, so it would be great if 2D segmentation could work with UNetR.

wyli commented 2 years ago

I think you can set spatial_dims=2 https://github.com/Project-MONAI/MONAI/blob/fdb16acbf606a425c412fecaddcdd0c782925581/monai/networks/nets/unetr.py#L42

(converting this to a discussion.)