Update app.py

app.py

@@ -37,16 +37,17 @@ class_labels = ["glass", "metal", "organic", "paper", "plastic"]
 
 
 def load_model():
-    """Load the trained TensorFlow/Keras model from model/ directory"""
+    """Load the trained TensorFlow/Keras model with version compatibility handling"""
     try:
-        # Try loading different formats from model/ directory
+        # Debug info
+        logger.info(f"TensorFlow version: {tf.__version__}")
+        logger.info("=== DEBUGGING MODEL LOADING ===")
+        
         model_files = [
-            './model/waste_model.keras',  # Keras format (recommended)
-            './model/waste_model.h5',     # H5 format
-            './model/best_model.keras',   # Checkpoint from training
             'model/waste_model.keras',
             'model/waste_model.h5',
-            'model/best_model.keras',
+            './model/waste_model.keras',
+            './model/waste_model.h5',
         ]
         
         model = None
@@ -55,28 +56,64 @@ def load_model():
         for model_file in model_files:
             if os.path.exists(model_file):
                 try:
-                    model = tf.keras.models.load_model(model_file)
-                    loaded_from = model_file
-                    logger.info(f"Model loaded successfully from {model_file}")
-                    break
+                    logger.info(f"Attempting to load: {model_file}")
+                    
+                    # Try different loading methods for compatibility
+                    try:
+                        # Method 1: Standard loading
+                        model = tf.keras.models.load_model(model_file, compile=False)
+                        logger.info(f"✅ Loaded with standard method: {model_file}")
+                        loaded_from = model_file
+                        break
+                    except Exception as e1:
+                        logger.warning(f"Standard loading failed: {e1}")
+                        
+                        # Method 2: Load with custom objects (for compatibility)
+                        try:
+                            custom_objects = {
+                                'InputLayer': tf.keras.layers.InputLayer,
+                                'Rescaling': tf.keras.layers.Rescaling,
+                            }
+                            model = tf.keras.models.load_model(
+                                model_file, 
+                                custom_objects=custom_objects,
+                                compile=False
+                            )
+                            logger.info(f"✅ Loaded with custom objects: {model_file}")
+                            loaded_from = model_file
+                            break
+                        except Exception as e2:
+                            logger.warning(f"Custom objects loading failed: {e2}")
+                            
+                            # Method 3: Try loading weights only
+                            try:
+                                # Create model architecture first, then load weights
+                                model = create_model_architecture()
+                                if model_file.endswith('.h5'):
+                                    model.load_weights(model_file)
+                                    logger.info(f"✅ Loaded weights only: {model_file}")
+                                    loaded_from = f"{model_file} (weights only)"
+                                    break
+                            except Exception as e3:
+                                logger.warning(f"Weights loading failed: {e3}")
+                                continue
+                
                 except Exception as e:
                     logger.warning(f"Failed to load {model_file}: {e}")
                     continue
         
         if model is None:
-            logger.error("No model file found in any location. Creating dummy model for testing.")
-            # Create dummy model with same architecture for testing
-            model = tf.keras.Sequential([
-                tf.keras.layers.Rescaling(1./255, input_shape=(224, 224, 3)),
-                tf.keras.applications.MobileNetV2(input_shape=(224, 224, 3), include_top=False, weights='imagenet'),
-                tf.keras.layers.GlobalAveragePooling2D(),
-                tf.keras.layers.Dense(128, activation='relu'),
-                tf.keras.layers.Dropout(0.2),
-                tf.keras.layers.Dense(5, activation='softmax')
-            ])
-            logger.warning("Using dummy model - predictions will be random!")
+            logger.warning("All loading methods failed. Creating model from architecture...")
+            model = create_model_architecture()
+            logger.warning("⚠️ Using untrained model - predictions will be random!")
         else:
-            logger.info(f"Successfully loaded model from: {loaded_from}")
+            logger.info(f"✅ Model loaded successfully from: {loaded_from}")
+            # Recompile the model
+            model.compile(
+                optimizer='adam',
+                loss='categorical_crossentropy',
+                metrics=['accuracy']
+            )
             
         return model
         
@@ -84,6 +121,37 @@ def load_model():
         logger.error(f"Critical error loading model: {e}")
         raise Exception(f"Model loading failed: {e}")
 
+def create_model_architecture():
+    """Create the model architecture matching your training setup"""
+    try:
+        # Create the same architecture as in your training notebook
+        base_model = tf.keras.applications.MobileNetV2(
+            weights='imagenet',
+            include_top=False,
+            input_shape=(224, 224, 3)
+        )
+        
+        # Freeze base model
+        base_model.trainable = False
+        
+        # Create complete model
+        model = tf.keras.Sequential([
+            tf.keras.layers.Rescaling(1./255, input_shape=(224, 224, 3)),
+            base_model,
+            tf.keras.layers.GlobalAveragePooling2D(),
+            tf.keras.layers.Dropout(0.2),
+            tf.keras.layers.Dense(128, activation='relu'),
+            tf.keras.layers.Dropout(0.5),
+            tf.keras.layers.Dense(5, activation='softmax')  # 5 classes
+        ])
+        
+        logger.info("Created model architecture successfully")
+        return model
+        
+    except Exception as e:
+        logger.error(f"Failed to create model architecture: {e}")
+        raise
+
 def preprocess_image(image_data):
     """
     Preprocess image to match training pipeline: