Update main.py

main.py

@@ -9,6 +9,11 @@ from tensorflow.keras.preprocessing.image import img_to_array
 from tensorflow.keras.applications import resnet50
 from tensorflow.keras.applications.resnet50 import preprocess_input
 import uvicorn
+import tensorflow_addons as tfa
+# Load the h5 model
+custom_objects = {
+    "Addons>CohenKappa": tfa.metrics.CohenKappa,
+}
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -59,7 +64,7 @@ class_labels = {
 
 # Load the model if it exists
 if os.path.exists(model_path):
-    model = tf.keras.models.load_model(model_path)
+    model = tf.keras.models.load_model('model.h5', custom_objects=custom_objects)
     print("Model loaded successfully.")
 else:
     print(f"Model file not found at {model_path}. Please upload the model.")