Skip to content

Commit

Permalink
Merge branch 'master' into pre-commit-ci-update-config
Browse files Browse the repository at this point in the history
  • Loading branch information
hvgazula authored Mar 22, 2024
2 parents 8409c83 + 902a124 commit 14f5c70
Show file tree
Hide file tree
Showing 9 changed files with 904 additions and 13 deletions.
24 changes: 24 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,27 @@
# 1.2.0 (Fri Mar 22 2024)

#### 🚀 Enhancement

- Dev [#295](https://github.com/neuronets/nobrainer/pull/295) ([@hvgazula](https://github.com/hvgazula) [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]) [@satra](https://github.com/satra))
- Update setup.cfg [#299](https://github.com/neuronets/nobrainer/pull/299) ([@satra](https://github.com/satra))

#### 🐛 Bug Fix

- Update release.yml ([@satra](https://github.com/satra))
- change hyphenation [#275](https://github.com/neuronets/nobrainer/pull/275) ([@satra](https://github.com/satra))
- update precommit checks [#275](https://github.com/neuronets/nobrainer/pull/275) ([@satra](https://github.com/satra))
- fix docker syntax [#275](https://github.com/neuronets/nobrainer/pull/275) ([@satra](https://github.com/satra))
- remove trained models [#275](https://github.com/neuronets/nobrainer/pull/275) ([@satra](https://github.com/satra))
- [pre-commit.ci] pre-commit autoupdate [#269](https://github.com/neuronets/nobrainer/pull/269) ([@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot]))

#### Authors: 3

- [@pre-commit-ci[bot]](https://github.com/pre-commit-ci[bot])
- H Gazula ([@hvgazula](https://github.com/hvgazula))
- Satrajit Ghosh ([@satra](https://github.com/satra))

---

# 1.1.1 (Sat Oct 07 2023)

#### 🐛 Bug Fix
Expand Down
3 changes: 3 additions & 0 deletions nobrainer/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,6 +121,9 @@ def from_tfrecords(
)
block_length = len([0 for _ in first_shard])

if not n_volumes:
n_volumes = block_length * len(files)

dataset = dataset.interleave(
map_func=lambda x: tf.data.TFRecordDataset(
x, compression_type=compression_type
Expand Down
File renamed without changes.
44 changes: 32 additions & 12 deletions nobrainer/models/__init__.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,32 @@
from pprint import pprint

from .attention_unet import attention_unet
from .attention_unet_with_inception import attention_unet_with_inception
from .autoencoder import autoencoder
from .bayesian_meshnet import variational_meshnet
from .dcgan import dcgan
from .highresnet import highresnet
from .meshnet import meshnet
from .progressiveae import progressiveae
from .progressivegan import progressivegan
from .unet import unet
from .unetr import unetr

__all__ = ["get", "list_available_models"]

_models = {
"highresnet": highresnet,
"meshnet": meshnet,
"unet": unet,
"autoencoder": autoencoder,
"progressivegan": progressivegan,
"progressiveae": progressiveae,
"dcgan": dcgan,
"attention_unet": attention_unet,
"attention_unet_with_inception": attention_unet_with_inception,
"unetr": unetr,
"variational_meshnet": variational_meshnet,
}


def get(name):
Expand All @@ -21,20 +43,18 @@ def get(name):
if not isinstance(name, str):
raise ValueError("Model name must be a string.")

models = {
"highresnet": highresnet,
"meshnet": meshnet,
"unet": unet,
"autoencoder": autoencoder,
"progressivegan": progressivegan,
"progressiveae": progressiveae,
"dcgan": dcgan,
}

try:
return models[name.lower()]
return _models[name.lower()]
except KeyError:
avail = ", ".join(models.keys())
avail = ", ".join(_models.keys())
raise ValueError(
"Unknown model: '{}'. Available models are {}.".format(name, avail)
)


def available_models():
return list(_models)


def list_available_models():
pprint(available_models())
80 changes: 80 additions & 0 deletions nobrainer/models/attention_unet.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
"""Model definition for Attention U-Net.
Adapted from https://github.com/nikhilroxtomar/Semantic-Segmentation-Architecture/blob/main/TensorFlow/attention-unet.py
""" # noqa: E501

from tensorflow.keras import layers
import tensorflow.keras.layers as L
from tensorflow.keras.models import Model


def conv_block(x, num_filters):
x = L.Conv3D(num_filters, 3, padding="same")(x)
x = L.BatchNormalization()(x)
x = L.Activation("relu")(x)

x = L.Conv3D(num_filters, 3, padding="same")(x)
x = L.BatchNormalization()(x)
x = L.Activation("relu")(x)

return x


def encoder_block(x, num_filters):
x = conv_block(x, num_filters)
p = L.MaxPool3D()(x)
return x, p


def attention_gate(g, s, num_filters):
Wg = L.Conv3D(num_filters, 1, padding="same")(g)
Wg = L.BatchNormalization()(Wg)

Ws = L.Conv3D(num_filters, 1, padding="same")(s)
Ws = L.BatchNormalization()(Ws)

out = L.Activation("relu")(Wg + Ws)
out = L.Conv3D(num_filters, 1, padding="same")(out)
out = L.Activation("sigmoid")(out)

return out * s


def decoder_block(x, s, num_filters):
x = L.UpSampling3D()(x)
s = attention_gate(x, s, num_filters)
x = L.Concatenate()([x, s])
x = conv_block(x, num_filters)
return x


def attention_unet(n_classes, input_shape):
"""Inputs"""
inputs = L.Input(input_shape)

""" Encoder """
s1, p1 = encoder_block(inputs, 64)
s2, p2 = encoder_block(p1, 128)
s3, p3 = encoder_block(p2, 256)

b1 = conv_block(p3, 512)

""" Decoder """
d1 = decoder_block(b1, s3, 256)
d2 = decoder_block(d1, s2, 128)
d3 = decoder_block(d2, s1, 64)

""" Outputs """
outputs = L.Conv3D(n_classes, 1, padding="same")(d3)

final_activation = "sigmoid" if n_classes == 1 else "softmax"
outputs = layers.Activation(final_activation)(outputs)

""" Model """
return Model(inputs=inputs, outputs=outputs, name="Attention_U-Net")


if __name__ == "__main__":
n_classes = 50
input_shape = (256, 256, 256, 3)
model = attention_unet(n_classes, input_shape)
model.summary()
Loading

0 comments on commit 14f5c70

Please sign in to comment.