Update app.py

app.py

@@ -27,8 +27,10 @@ model = tf.keras.models.load_model('my_model_1.h5', compile=False)
 
 
 def classify_image(image):
-    image = np.array(image['composite'])[:, :, 3] * 255
-    image = image[..., np.newaxis]
+    if len(np.array(image).shape) == 3:
+      image = tf.image.rgb_to_grayscale(image)
+    # image = np.array(image['composite'])[:, :, 3] * 255
+    # image = image[..., np.newaxis]
     image_tensor = tf.convert_to_tensor(image)
     image_tensor = tf.image.resize(image_tensor, (28, 28)),
     image_tensor = tf.cast(image_tensor, tf.float32)