Update app.py

app.py

@@ -41,14 +41,14 @@ uploaded_files = st.file_uploader("Input images of cats or dogs (examples in fil
 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 = jnp.array(tf.cast(tf.image.resize(tf.convert_to_tensor(img), [50, 50]), tf.float32) / 255.)
-    st.write("Input shape: " + str(input.shape))
-    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)))
+  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)))
+  for (index, image) in enumerate(uploaded_files):
+    st.image(Image.open(image))
 
 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,))