Upload 17 files

.gitattributes

@@ -32,3 +32,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/attn_plot.png filter=lfs diff=lfs merge=lfs -text
+assets/examples.png filter=lfs diff=lfs merge=lfs -text
+checkpoints/cxr_validator_model.tf/variables/variables.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text
+checkpoints/RATCHET.tf/variables/variables.data-00000-of-00001 filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,209 @@
+import tqdm
+import datetime
+
+import matplotlib.pyplot as plt
+import numpy as np
+import streamlit as st
+import tensorflow as tf
+
+from skimage import io
+from transformer import Transformer
+from tokenizers import ByteLevelBPETokenizer
+
+
+@st.cache_resource
+def load_validator():
+    validator_model = tf.keras.models.load_model('checkpoints/cxr_validator_model.tf')
+    print('Validator Model Loaded!')
+    return validator_model
+
+
+@st.cache_resource
+def load_model():
+
+    # Load Tokenizer
+    tokenizer = ByteLevelBPETokenizer(
+        'mimic/mimic-vocab.json',
+        'mimic/mimic-merges.txt',
+    )
+
+    # Load Model
+    hparams = default_hparams()
+    transformer = Transformer(
+        num_layers=hparams['num_layers'],
+        d_model=hparams['d_model'],
+        num_heads=hparams['num_heads'],
+        dff=hparams['dff'],
+        target_vocab_size=tokenizer.get_vocab_size(),
+        dropout_rate=hparams['dropout_rate'])
+    transformer.load_weights('checkpoints/RATCHET.tf')
+    print(f'Model Loaded! Checkpoint file: checkpoints/RATCHET.tf')
+
+    return transformer, tokenizer
+
+
+def top_k_logits(logits, k):
+    if k == 0:
+        # no truncation
+        return logits
+
+    def _top_k():
+        values, _ = tf.nn.top_k(logits, k=k)
+        min_values = values[:, -1, tf.newaxis]
+        return tf.where(
+            logits < min_values,
+            tf.ones_like(logits, dtype=logits.dtype) * -1e10,
+            logits,
+        )
+    return tf.cond(
+       tf.equal(k, 0),
+       lambda: logits,
+       lambda: _top_k(),
+    )
+
+
+def top_p_logits(logits, p):
+    """Nucleus sampling"""
+    batch, _ = logits.shape.as_list()
+    sorted_logits = tf.sort(logits, direction='DESCENDING', axis=-1)
+    cumulative_probs = tf.cumsum(tf.nn.softmax(sorted_logits, axis=-1), axis=-1)
+    indices = tf.stack([
+        tf.range(0, batch),
+        # number of indices to include
+        tf.maximum(tf.reduce_sum(tf.cast(cumulative_probs <= p, tf.int32), axis=-1) - 1, 0),
+    ], axis=-1)
+    min_values = tf.gather_nd(sorted_logits, indices)
+    return tf.where(
+        logits < min_values,
+        tf.ones_like(logits) * -1e10,
+        logits,
+    )
+
+
+def evaluate(inp_img, tokenizer, transformer, temperature, top_k, top_p, options, seed, MAX_LENGTH=128):
+
+    # The first token to the transformer should be the start token
+    output = tf.convert_to_tensor([[tokenizer.token_to_id('<s>')]])
+
+    my_bar = st.progress(0)
+    for i in tqdm.tqdm(range(MAX_LENGTH)):
+        my_bar.progress(i/MAX_LENGTH)
+
+        # predictions.shape == (batch_size, seq_len, vocab_size)
+        predictions = transformer([inp_img, output], training=False)
+
+        # select the last word from the seq_len dimension
+        predictions = predictions[:, -1, :] / temperature  # (batch_size, vocab_size)
+        predictions = top_k_logits(predictions, k=top_k)
+        predictions = top_p_logits(predictions, p=top_p)
+
+        if options == 'Greedy':
+            predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)[:, tf.newaxis]
+        elif options == 'Sampling':
+            predicted_id = tf.random.categorical(predictions, num_samples=1, dtype=tf.int32, seed=seed)
+        else:
+            st.write('SHOULD NOT HAPPEN')
+
+        # return the result if the predicted_id is equal to the end token
+        if predicted_id == 2:  # stop token #tokenizer_en.vocab_size + 1:
+            my_bar.empty()
+            break
+
+        # concatentate the predicted_id to the output which is given to the decoder
+        # as its input.
+        output = tf.concat([output, predicted_id], axis=-1)
+
+    my_bar.empty()
+
+    # transformer([inp_img, output[:, :-1]], training=False)
+    return tf.squeeze(output, axis=0)[1:], transformer.decoder.last_attn_scores
+
+
+def main():
+
+    st.title('Chest X-ray AI Diagnosis Demo')
+    st.text('Made with Streamlit and Attention RNN')
+
+    transformer, tokenizer = load_model()
+    cxr_validator_model = load_validator()
+
+    st.sidebar.title('Configuration')
+    options = st.sidebar.selectbox('Generation Method', ('Greedy', 'Sampling'))
+    seed = st.sidebar.number_input('Sampling Seed:', value=42)
+    temperature = st.sidebar.number_input('Temperature', value=1.)
+    top_k = st.sidebar.slider('top_k', min_value=0, max_value=tokenizer.get_vocab_size(), value=6, step=1)
+    top_p = st.sidebar.slider('top_p', min_value=0., max_value=1., value=1., step=0.01)
+    attention_head = st.sidebar.slider('attention_head', min_value=-1, max_value=7, value=-1, step=1)
+
+    st.sidebar.info('PRIVACY POLICY: Uploaded images are never stored on disk.')
+
+    st.set_option('deprecation.showfileUploaderEncoding', False)
+    uploaded_file = st.file_uploader('Choose an image...', type=('png', 'jpg', 'jpeg'))
+
+    if uploaded_file:
+
+        # Read input image with size [1, H, W, 1] and range (0, 255)
+        img_array = io.imread(uploaded_file, as_gray=True)[None, ..., None]
+
+        # Convert image to float values in (0, 1)
+        img_array = tf.image.convert_image_dtype(img_array, tf.float32)
+
+        # Resize image with padding to [1, 224, 224, 1]
+        img_array = tf.image.resize_with_pad(img_array, 224, 224, method=tf.image.ResizeMethod.BILINEAR)
+
+        # Display input image
+        st.image(np.squeeze(img_array.numpy()), caption='Uploaded Image')
+
+        # Check image
+        valid = tf.nn.sigmoid(cxr_validator_model(img_array))
+        if valid < 0.1:
+            st.info('Image is not a Chest X-ray')
+            return
+
+        # Log datetime
+        print('[{}] Running Analysis...'
+              .format(datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")))
+
+        # Generate radiology report
+        with st.spinner('Generating report... Do not refresh or close window.'):
+            result, attention_weights = evaluate(img_array, tokenizer, transformer,
+                                                 temperature, top_k, top_p,
+                                                 options, seed)
+            predicted_sentence = tokenizer.decode(result)
+
+        # Display generated text
+        st.subheader('Generated Report:')
+        st.write(predicted_sentence)
+        # st.info(predicted_sentence)
+
+        st.subheader('Attention Plot:')
+
+        attn_map = attention_weights[0]  # squeeze
+        if attention_head == -1:  # average attention heads
+            attn_map = tf.reduce_mean(attn_map, axis=0)
+        else:  # select attention heads
+            attn_map = attn_map[attention_head]
+        attn_map = attn_map / attn_map.numpy().max() * 255
+
+        fig = plt.figure(figsize=(40, 80))
+
+        for i in range(attn_map.shape[0] - 1):
+            attn_token = attn_map[i, ...]
+            attn_token = tf.reshape(attn_token, [7, 7])
+
+            ax = fig.add_subplot(16, 8, i + 1)
+            ax.set_title(tokenizer.decode([result.numpy()[i]]))
+            img = ax.imshow(np.squeeze(img_array))
+            ax.imshow(attn_token, cmap='gray', alpha=0.6, extent=img.get_extent())
+
+        st.pyplot(plt)
+
+        # Run again?
+        st.button('Regenerate Report')
+
+
+if __name__ == '__main__':
+
+    tf.config.set_visible_devices([], 'GPU')
+
+    main()

checkpoints/RATCHET.tf/keras_metadata.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a8fa018ac83d10617e20e3f03de3718d9d3d6e1b89673707cb510318fd3198b3
+size 1065144

checkpoints/RATCHET.tf/saved_model.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:84e9d837b881c58edee113c7bbdc793159e6e57c2ddcf9d2a3e4da7c5104a7db
+size 26013311

checkpoints/RATCHET.tf/variables/variables.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae18face6fa821f8c6c62923ef5533fca681e01b6bb8ae511a9c94844f618c8e
+size 1669994429

checkpoints/cxr_validator_model.tf/fingerprint.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:21a31ac72a46d124de283ecbd75c35efc8ac0c5f597efd3040ed8dd00d071ef2
+size 53

checkpoints/cxr_validator_model.tf/keras_metadata.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:19106ee698a03e8b9ec11b0092fd65c32654380171a3c55a7976d56313e4438a
+size 2538679

checkpoints/cxr_validator_model.tf/saved_model.pb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:16e7434007981626733e6f925cd0b226e1f4130cfaec7e79ba81ffd16d7ab1cb
+size 14320368

checkpoints/cxr_validator_model.tf/variables/variables.data-00000-of-00001

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e2edd5cef46c1624f31464e13f3b5fb8c0ceb4ce8a1d834a6cde9c2e71dd509e
+size 224256098

requirements.txt

@@ -0,0 +1,6 @@
+matplotlib
+numpy
+scikit-image
+tensorflow
+tokenizers
+tqdm

transformer.py

@@ -0,0 +1,263 @@
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+from __future__ import unicode_literals
+
+import datetime
+
+import numpy as np
+import tensorflow as tf
+
+
+def default_hparams():
+    return {
+        'img_x': 224,
+        'img_y': 224,
+        'img_ch': 1,
+        'd_model': 512,
+        'dff': 2048,
+        'num_heads': 8,
+        'num_layers': 6,
+        'dropout_rate': 0.1
+    }
+
+
+def positional_encoding(length, depth):
+    depth = depth / 2
+
+    positions = np.arange(length)[:, np.newaxis]  # (seq, 1)
+    depths = np.arange(depth)[np.newaxis, :] / depth  # (1, depth)
+
+    angle_rates = 1 / (10000 ** depths)  # (1, depth)
+    angle_rads = positions * angle_rates  # (pos, depth)
+
+    pos_encoding = np.concatenate(
+        [np.sin(angle_rads), np.cos(angle_rads)],
+        axis=-1)
+
+    return tf.cast(pos_encoding, dtype=tf.float32)
+
+
+class PositionalEmbedding(tf.keras.layers.Layer):
+    def __init__(self, vocab_size, d_model):
+        super().__init__()
+        self.d_model = d_model
+        self.embedding = tf.keras.layers.Embedding(vocab_size, d_model, mask_zero=True)
+        self.pos_encoding = positional_encoding(length=2048, depth=d_model)
+
+    def compute_mask(self, *args, **kwargs):
+        return self.embedding.compute_mask(*args, **kwargs)
+
+    def call(self, x):
+        length = tf.shape(x)[1]
+        x = self.embedding(x)
+        # This factor sets the relative scale of the embedding and positonal_encoding.
+        x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
+        x = x + self.pos_encoding[tf.newaxis, :length, :]
+        return x
+
+
+class BaseAttention(tf.keras.layers.Layer):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.mha = tf.keras.layers.MultiHeadAttention(**kwargs)
+        self.layernorm = tf.keras.layers.LayerNormalization()
+        self.add = tf.keras.layers.Add()
+
+
+class CrossAttention(BaseAttention):
+    def call(self, x, context):
+        attn_output, attn_scores = self.mha(
+            query=x,
+            key=context,
+            value=context,
+            return_attention_scores=True)
+
+        # Cache the attention scores for plotting later.
+        self.last_attn_scores = attn_scores
+
+        x = self.add([x, attn_output])
+        x = self.layernorm(x)
+
+        return x
+
+
+class CausalSelfAttention(BaseAttention):
+    def call(self, x):
+        attn_output = self.mha(
+            query=x,
+            value=x,
+            key=x,
+            use_causal_mask=True)
+        x = self.add([x, attn_output])
+        x = self.layernorm(x)
+        return x
+
+
+class FeedForward(tf.keras.layers.Layer):
+    def __init__(self, d_model, dff, dropout_rate=0.1):
+        super().__init__()
+        self.seq = tf.keras.Sequential([
+            tf.keras.layers.Dense(dff, activation='relu'),
+            tf.keras.layers.Dense(d_model),
+            tf.keras.layers.Dropout(dropout_rate)
+        ])
+        self.add = tf.keras.layers.Add()
+        self.layer_norm = tf.keras.layers.LayerNormalization()
+
+    def call(self, x):
+        x = self.add([x, self.seq(x)])
+        x = self.layer_norm(x)
+        return x
+
+
+class DecoderLayer(tf.keras.layers.Layer):
+    def __init__(self,
+                 *,
+                 d_model,
+                 num_heads,
+                 dff,
+                 dropout_rate=0.1):
+        super(DecoderLayer, self).__init__()
+
+        self.causal_self_attention = CausalSelfAttention(
+            num_heads=num_heads,
+            key_dim=d_model,
+            dropout=dropout_rate)
+
+        self.cross_attention = CrossAttention(
+            num_heads=num_heads,
+            key_dim=d_model,
+            dropout=dropout_rate)
+
+        self.ffn = FeedForward(d_model, dff)
+
+    def call(self, x, context):
+        x = self.causal_self_attention(x=x)
+        x = self.cross_attention(x=x, context=context)
+
+        # Cache the last attention scores for plotting later
+        self.last_attn_scores = self.cross_attention.last_attn_scores
+
+        x = self.ffn(x)  # Shape `(batch_size, seq_len, d_model)`.
+        return x
+
+
+class Encoder(tf.keras.layers.Layer):
+    def __init__(self, embedding_dim, input_shape, pretrain_weights=None):
+        super(Encoder, self).__init__()
+
+        # shape after fc == (batch_size, nf * nf, embedding_dim)
+        self.fc = tf.keras.layers.Dense(embedding_dim, activation='relu')
+
+        # Use DenseNet-121 as feature extraction model
+        self.base_model = tf.keras.applications.DenseNet121(
+            include_top=False, weights=None, input_shape=input_shape)
+
+        # Load pre-trained weights if present
+        if pretrain_weights:
+            print(f'{datetime.datetime.now()}: I Loading Pretrained DenseNet-121 weights: {pretrain_weights}')
+            self.base_model.load_weights(pretrain_weights)
+        else:
+            print(f'{datetime.datetime.now()}: I No Pretrained DenseNet-121 weights specified')
+
+    def call(self, x, **kwargs):
+        x = self.base_model(x)
+        # DenseNet-121 output is (batch_size, ?, ?, 1024)
+        s = tf.shape(x)
+        x = tf.reshape(x, (s[0], s[1] * s[2], x.shape[3]))
+        x = self.fc(x)
+        return x
+
+
+class Decoder(tf.keras.layers.Layer):
+    def __init__(self, *, num_layers, d_model, num_heads, dff, vocab_size,
+                 dropout_rate=0.1):
+        super(Decoder, self).__init__()
+
+        self.d_model = d_model
+        self.num_layers = num_layers
+
+        self.pos_embedding = PositionalEmbedding(vocab_size=vocab_size,
+                                                 d_model=d_model)
+        self.dropout = tf.keras.layers.Dropout(dropout_rate)
+        self.dec_layers = [
+            DecoderLayer(d_model=d_model, num_heads=num_heads,
+                         dff=dff, dropout_rate=dropout_rate)
+            for _ in range(num_layers)]
+
+        self.last_attn_scores = None
+
+    def call(self, x, context):
+        # `x` is token-IDs shape (batch, target_seq_len)
+        x = self.pos_embedding(x)  # (batch_size, target_seq_len, d_model)
+
+        x = self.dropout(x)
+
+        for i in range(self.num_layers):
+            x = self.dec_layers[i](x, context)
+
+        self.last_attn_scores = self.dec_layers[-1].last_attn_scores
+
+        # The shape of x is (batch_size, target_seq_len, d_model).
+        return x
+
+
+class Transformer(tf.keras.Model):
+    def __init__(self, num_layers, d_model, num_heads, dff,
+                 target_vocab_size, dropout_rate=0.1, input_shape=(224, 224, 1),
+                 classifier_weights=None):
+        super(Transformer, self).__init__()
+
+        self.encoder = Encoder(d_model, input_shape,
+                               pretrain_weights=classifier_weights)
+
+        self.decoder = Decoder(num_layers=num_layers, d_model=d_model,
+                               num_heads=num_heads, dff=dff,
+                               vocab_size=target_vocab_size,
+                               dropout_rate=dropout_rate)
+
+        self.final_layer = tf.keras.layers.Dense(target_vocab_size)
+
+    def call(self, inputs):
+        # To use a Keras model with `.fit` you must pass all your inputs in the
+        # first argument.
+        context, x = inputs
+
+        context = self.encoder(context)  # (batch_size, context_len, d_model)
+
+        x = self.decoder(x, context)  # (batch_size, target_len, d_model)
+
+        # Final linear layer output.
+        logits = self.final_layer(x)  # (batch_size, target_len, target_vocab_size)
+
+        try:
+            # Drop the keras mask, so it doesn't scale the losses/metrics.
+            # b/250038731
+            del logits._keras_mask
+        except AttributeError:
+            pass
+
+        # Return the final output and the attention weights.
+        return logits
+
+
+if __name__ == "__main__":
+
+    hparams = default_hparams()
+
+    transformer = Transformer(
+        num_layers=hparams['num_layers'],
+        d_model=hparams['d_model'],
+        num_heads=hparams['num_heads'],
+        dff=hparams['dff'],
+        target_vocab_size=2048,
+        dropout_rate=hparams['dropout_rate'])
+
+    a=1
+
+
+    image = np.random.rand(1,224,224,1).astype('float32')
+    text = np.random.randint(0, 2048, size=(1, 27))
+
+    output = transformer((image, text))