Support for timm_3d models

Author: ZFTurbo

README.md

@@ -198,6 +198,20 @@ Note: In the official github repo the s0 variant has additional num_conv_branche
 </div>
 </details>
 
+### Timm 3D encoders
+
+We now support encoders from [timm_3d](https://github.com/ZFTurbo/timm_3d) library. Full list available [here](https://github.com/ZFTurbo/timm_3d/blob/main/docs/models_list.md). To use them add `tu-` before encoder name.
+Example:
+
+```python
+encoder_name = 'tu-maxvit_base_tf_224.in21k'
+model = smp.Unet(
+    encoder_name=encoder_name,
+    encoder_weights=None,
+    in_channels=3,
+    classes=1,
+)
+```
 
 ## Notes for 3D version
 

requirements.txt

@@ -3,6 +3,7 @@ torchvision>=0.5.0
 pretrainedmodels==0.7.4
 efficientnet-pytorch==0.7.1
 timm==0.9.7
+timm_3d==1.0.1
 tqdm
 pillow
 six

segmentation_models_pytorch_3d/encoders/timm_universal.py

@@ -1,4 +1,4 @@
-import timm
+import timm_3d
 import torch.nn as nn
 
 
@@ -17,7 +17,7 @@ def __init__(self, name, pretrained=True, in_channels=3, depth=5, output_stride=
         if output_stride == 32:
             kwargs.pop("output_stride")
 
-        self.model = timm.create_model(name, **kwargs)
+        self.model = timm_3d.create_model(name, **kwargs)
 
         self._in_channels = in_channels
         self._out_channels = [

test.py

@@ -244,4 +244,17 @@
                 strides=((1, 2, 2), (2, 2, 2), (2, 2, 2), (2, 2, 2), (2, 2, 2)),
             )
             o = model(torch.randn(2, 3, 32, 64, 128))
-            print(f'Result shape: {o.shape}')
+            print(f'Result shape: {o.shape}')
+
+    if 1:
+        encoder_name = 'tu-maxvit_base_tf_224.in21k'
+        print('Test Timm 3d model: {}...'.format(encoder_name))
+        print('Go for {}'.format(encoder_name))
+        model = smp.Unet(
+            encoder_name=encoder_name,
+            encoder_weights=None,
+            in_channels=3,
+            classes=1,
+        )
+        o = model(torch.randn(2, 3, 128, 64, 64))
+        print(f'Result shape: {o.shape}')