Update app.py

app.py

@@ -34,29 +34,29 @@ 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"])
+  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!")
+  st.write("Please upload an image!")
 else:
-    for uploaded_file in uploaded_files:
-        img = Image.open(uploaded_file)
-        st.image(img)
-        input = tf.cast(tf.image.resize(tf.convert_to_tensor(img), [50, 50]), tf.float32) / 255.
-        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)))
+  for uploaded_file in uploaded_files:
+    img = Image.open(uploaded_file)
+    st.image(img)
+    input = jnp.array(tf.cast(tf.image.resize(tf.convert_to_tensor(img), [50, 50]), tf.float32) / 255.)
+    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, 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"))
+  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"))
 
 with st.expander("See first convolutional layer"):
-    gridify(params["Conv_0"]["kernel"], grid=(4,8), kernel_size=(3,3))
+  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))
+  print(params["Conv_1"]["kernel"].shape)
+  gridify(params["Conv_1"]["kernel"], grid=(64,96), kernel_size=(3,3))