Refactor prediction function: update data handling to reshape composite image for prediction

app.py

@@ -28,8 +28,11 @@ model = tf.keras.models.load_model("./sketch_recognition_numbers_model.h5")
 
 # Prediction function for sketch recognition
 def predict(data):
+    print(data)
+    # ValueError: cannot reshape array of size 1 into shape (1,28,28,1)
+
     # Reshape image to 28x28
-    img = np.reshape(data, (1, img_size, img_size, 1))
+    img = np.reshape(data['composite'], (1, img_size, img_size, 1))
     # Make prediction
     pred = model.predict(img)
     # Get top class