Create app.py

app.py

@@ -0,0 +1,64 @@
+import streamlit as st
+from PIL import Image
+import jax
+import numpy as np
+import jax.numpy as jnp  # JAX NumPy
+from flax.training import train_state  # Useful dataclass to keep train state
+from flax import linen as nn  # Linen API
+from huggingface_hub import HfFileSystem
+from flax.serialization import msgpack_restore, from_state_dict
+import os
+import tensorflow as tf
+
+hf_key =  text_input = st.text_input("Access token")
+
+class CNN(nn.Module):
+  """A simple CNN model."""
+
+  @nn.compact
+  def __call__(self, x):
+    x = nn.Conv(features=32, kernel_size=(3, 3))(x)
+    x = nn.relu(x)
+    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
+    x = nn.Conv(features=64, kernel_size=(3, 3))(x)
+    x = nn.relu(x)
+    x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))
+    x = x.reshape((x.shape[0], -1)) # flatten
+    x = nn.Dense(features=256)(x)
+    x = nn.relu(x)
+    x = nn.Dense(features=16)(x)
+    x = nn.relu(x)
+    x = nn.Dense(features=2)(x)
+    return x
+
+cnn = CNN()
+params = cnn.init(jax.random.PRNGKey(0), jnp.ones([1, 50, 50, 3]))['params']
+
+fs = HfFileSystem(token=hf_key)
+with fs.open("PrakhAI/CatVsDog/checkpoint.msgpack", "rb") as f:
+    params = from_state_dict(params, msgpack_restore(f.read())["params"])
+
+uploaded_files = st.file_uploader("Input images of cats or dogs (examples in files)", type=['jpg','png','tif'], accept_multiple_files=True)
+
+if len(uploaded_files) == 0:
+    st.write("Please upload an image!")
+else:
+    for uploaded_file in uploaded_files:
+        img = Image.open(uploaded_file)
+        st.image(img)
+        input = tf.image_resize(tf.convert_to_tensor(img), [50, 50])
+        st.write("Model Prediction: " + cnn.apply({"params": params}, input))
+        st.write("Model Prediction type: " + type(cnn.apply({"params": params}, input)))
+        st.write("Model Prediction type dir: " + dir(cnn.apply({"params": params}, input)))
+
+def gridify(kernel, grid, kernel_size, scaling=5, padding=1):
+    scaled_and_padded = np.pad(np.repeat(np.repeat(kernel, repeats=scaling, axis=0), repeats=scaling, axis=1), ((padding,),(padding,),(0,),(0,)), 'constant', constant_values=(-1,))
+    grid = np.pad(np.array(scaled_and_padded.reshape((kernel_size[0]*scaling+2*padding, kernel_size[1]*scaling+2*padding, grid[0], grid[1])).transpose(2,0,3,1).reshape(grid[0]*(kernel_size[0]*scaling+2*padding), grid[1]*(kernel_size[1]*scaling+2*padding))+1)*127., (padding,), 'constant', constant_values=(0,))
+    st.image(Image.fromarray(np.repeat(np.expand_dims(grid, axis=0), repeats=3, axis=0).astype(np.uint8).transpose(1,2,0), mode="RGB"))
+
+with st.expander("See first convolutional layer"):
+    gridify(params["Conv_0"]["kernel"], grid=(4,8), kernel_size=(3,3))
+    
+with st.expander("See second convolutional layer"):
+    print(params["Conv_1"]["kernel"].shape)
+    gridify(params["Conv_1"]["kernel"], grid=(32,64), kernel_size=(3,3))