Update app.py

app.py

@@ -15,10 +15,8 @@ import tensorflow as tf
 
 # Load your model (handle any custom layers if needed)
 # If you have custom layers like RandomHeight, include them in custom_objects:
-# model = tf.keras.models.load_model("stone_artifact_classifier.h5", compile=False, 
-#                                    custom_objects={"RandomHeight": tf.keras.layers.RandomHeight, 
-#                                                    "RandomWidth": tf.keras.layers.RandomWidth})
-model = tf.keras.models.load_model("stone_artifact_classifier.h5", compile=False)
+model = tf.keras.models.load_model("stone_artifact_classifier.h5", compile=False, custom_objects={"RandomHeight": tf.keras.layers.RandomHeight, "RandomWidth": tf.keras.layers.RandomWidth})
+#model = tf.keras.models.load_model("stone_artifact_classifier.h5", compile=False)
 
 # Replace with your classes
 class_names = ["arrowheads", "axehammerstones", "scrapers"]