Update app.py

app.py

@@ -50,18 +50,22 @@ else:
     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}%)")
+      st.write(f"Model Prediction - Cat ({100*cat_prob:.2f}%), Dog ({100*dog_prob:.2f}%)")
     else:
-      st.write(f"Model Prediction - Dog ({100*dog_prob}%), Cat ({100*cat_prob}%)")
+      st.write(f"Model Prediction - Dog ({100*dog_prob:.2f}%), Cat ({100*cat_prob:.2f}%)")
 
-def gridify(kernel, grid, kernel_size, scaling=5, padding=1):
+def gridify_rgb(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, 3, grid[0], grid[1])).transpose(3,0,4,1,2).reshape(grid[0]*(kernel_size[0]*scaling+2*padding), grid[1]*(kernel_size[1]*scaling+2*padding), 3)+1)*127., ((padding,),(padding,),(0,)), 'constant', constant_values=(0,))
   st.image(Image.fromarray(grid.astype(np.uint8), mode="RGB"))
 
+def gridify_grayscale(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))
+  gridify_rgb(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))
+  gridify_grayscale(params["Conv_1"]["kernel"], grid=(32,64), kernel_size=(3,3))