Update app.py

app.py

@@ -12,6 +12,9 @@ import tensorflow as tf
 
 hf_key =  text_input = st.text_input("Access token")
 
+def sigmoid(x):
+  return 1/(1+e**-x)
+
 class CNN(nn.Module):
   @nn.compact
   def __call__(self, x):
@@ -42,13 +45,14 @@ if len(uploaded_files) == 0:
   st.write("Please upload an image!")
 else:
   input = jnp.array([tf.cast(tf.image.resize(tf.convert_to_tensor(Image.open(uploaded_file)), [50, 50]), tf.float32) / 255. for uploaded_file in uploaded_files])
-  st.write("Input shape: " + str(input.shape))
-  pred = cnn.apply({"params": params}, input)
-  st.write("Model Prediction: " + str(pred))
-  st.write("Model Prediction type: " + str(type(pred)))
-  st.write("Model Prediction type dir: " + str(dir(pred)))
+  prediction = cnn.apply({"params": params}, input)
   for (index, image) in enumerate(uploaded_files):
     st.image(Image.open(image))
+    [cat_prob, dog_prob] = jax.nn.softmax(prediction[index])
+    if cat_prob > dog_prob:
+      st.write(f"Model Prediction - Cat ({100*cat_prob}%), Dog ({100*dog_prob}%)")
+    else:
+      st.write(f"Model Prediction - Dog ({100*dog_prob}%), Cat ({100*cat_prob}%)")
 
 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,))
@@ -60,4 +64,4 @@ with st.expander("See first convolutional layer"):
     
 with st.expander("See second convolutional layer"):
   print(params["Conv_1"]["kernel"].shape)
-  gridify(params["Conv_1"]["kernel"], grid=(64,96), kernel_size=(3,3))
+  gridify(params["Conv_1"]["kernel"], grid=(32,64), kernel_size=(3,3))