Update app.py

app.py

@@ -43,6 +43,12 @@ else:
         st.image(rescaled)
         input = jnp.array(rescaled).reshape(1, 28, 28, 1) / 255.
         st.write("Model Prediction: " + cnn.apply({"params": params}, input).argmax(axis=1)[0])
+  input = jnp.array([jnp.array(Image.open(uploaded_file).convert("HSV").split()[2].resize((28, 28))).reshape(1, 28, 28, 1) / 255. for uploaded_file in uploaded_files])
+  prediction = cnn.apply({"params": params}, input)
+  for (index, image) in enumerate(uploaded_files):
+    st.image(Image.open(image))
+    st.write("Model Prediction: " + prediction[index].argmax(axis=0)[0])
+
 
 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,))