Update app.py

app.py

@@ -12,6 +12,8 @@ import time
 import os
 import tempfile
 from urllib.parse import urlparse
+from tensorflow import keras
+from tensorflow.keras import layers, models
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -21,6 +23,41 @@ logger = logging.getLogger(__name__)
 tf.get_logger().setLevel('ERROR')
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 
+# ========== CUSTOM LAYERS DEFINITIONS ==========
+@keras.saving.register_keras_serializable()
+class RepeatChannels(keras.layers.Layer):
+    """Converts single channel (depth) to 3 channels for RGB models"""
+    def __init__(self, **kwargs):
+        super(RepeatChannels, self).__init__(**kwargs)
+    
+    def call(self, inputs):
+        return tf.repeat(inputs, 3, axis=-1)
+    
+    def get_config(self):
+        config = super(RepeatChannels, self).get_config()
+        return config
+
+# Add any other custom layers your models might need
+@keras.saving.register_keras_serializable()
+class CustomLayer(keras.layers.Layer):
+    """Template for additional custom layers if needed"""
+    def __init__(self, **kwargs):
+        super(CustomLayer, self).__init__(**kwargs)
+    
+    def call(self, inputs):
+        return inputs
+    
+    def get_config(self):
+        config = super(CustomLayer, self).get_config()
+        return config
+
+# Custom objects dictionary for model loading
+CUSTOM_OBJECTS = {
+    'RepeatChannels': RepeatChannels,
+    'CustomLayer': CustomLayer,
+    # Add more custom layers here as needed
+}
+
 # Class mappings
 CLASS_NAMES = {
     0: 'AcrimSat', 1: 'Aquarius', 2: 'Aura', 3: 'Calipso', 4: 'Cloudsat',
@@ -118,7 +155,7 @@ def download_model_with_progress(url, timeout=120):
         return None, f"Download error: {str(e)}"
 
 def load_model(model_name):
-    """Load model from Hugging Face with enhanced error handling"""
+    """Load model from Hugging Face with enhanced error handling and custom objects"""
     
     # Check cache first
     if model_name in model_cache:
@@ -150,8 +187,12 @@ def load_model(model_name):
                 tmp_file.write(model_bytes.read())
                 tmp_file_path = tmp_file.name
             
-            # Load model from temporary file
-            model = tf.keras.models.load_model(tmp_file_path, compile=False)
+            # Load model from temporary file with custom objects
+            model = tf.keras.models.load_model(
+                tmp_file_path, 
+                custom_objects=CUSTOM_OBJECTS,
+                compile=False
+            )
             
             # Clean up temporary file
             try:
@@ -192,6 +233,31 @@ def preprocess_image(image, target_size=(224, 224)):
     except Exception as e:
         return None, f"Error preprocessing image: {str(e)}"
 
+def handle_multi_input_prediction(model, image, model_name):
+    """Handle models that expect multiple inputs (RGB + Depth + Tabular)"""
+    try:
+        # For multi-input models, we need to provide dummy inputs for missing modalities
+        rgb_input = image
+        
+        # Create dummy depth input (grayscale version of RGB)
+        depth_input = np.mean(image, axis=-1, keepdims=True)  # Convert RGB to grayscale
+        depth_input = np.repeat(depth_input, 3, axis=-1)  # Repeat to make it 3-channel
+        
+        # Create dummy tabular input
+        if model_name == "Custom CNN":
+            tabular_input = np.random.random((image.shape[0], 10))  # Adjust size as needed
+        else:
+            tabular_input = np.random.random((image.shape[0], 1))
+        
+        # Try multi-input prediction
+        predictions = model.predict([rgb_input, depth_input, tabular_input], verbose=0)
+        return predictions
+        
+    except Exception as e:
+        logger.warning(f"Multi-input prediction failed for {model_name}: {e}")
+        # Fallback to single input
+        return model.predict(image, verbose=0)
+
 def predict_with_model(model, image, model_name):
     """Make prediction with a specific model"""
     if model is None:
@@ -199,8 +265,14 @@ def predict_with_model(model, image, model_name):
     try:
         start_time = time.time()
         
-        # Make prediction
-        predictions = model.predict(image, verbose=0)
+        # Check if model expects multiple inputs
+        if len(model.input_shape) > 1 or (hasattr(model, 'input') and isinstance(model.input, list)):
+            # Multi-input model
+            predictions = handle_multi_input_prediction(model, image, model_name)
+        else:
+            # Single input model
+            predictions = model.predict(image, verbose=0)
+            
         inference_time = (time.time() - start_time) * 1000
         
         # Handle different output shapes