fix bugs

app.py

@@ -3,8 +3,11 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 import tensorflow as tf
+from tensorflow.keras import layers, models
 
 print("Building model architecture...")
+
+# 1. Rebuild the exact backbone used in training
 base_model = tf.keras.applications.MobileNetV2(
     input_shape=(224, 224, 3),
     include_top=False,
@@ -12,16 +15,28 @@ base_model = tf.keras.applications.MobileNetV2(
 )
 base_model.trainable = False
 
-model = tf.keras.models.Sequential([
-    base_model,
-    tf.keras.layers.GlobalAveragePooling2D(),
-    tf.keras.layers.Dense(128, activation="relu"),
-    tf.keras.layers.Dropout(0.3),
-    tf.keras.layers.Dense(2, activation="softmax")
-])
+# 2. Rebuild the exact head used in training (Functional API)
+inputs = tf.keras.Input(shape=(224, 224, 3))
+x = base_model(inputs, training=False)
+x = layers.GlobalAveragePooling2D()(x)
+
+# --- The Missing Layers ---
+x = layers.BatchNormalization()(x)
+x = layers.Dense(256, activation="relu")(x) # This matches your trained (1280, 256) weights
+x = layers.BatchNormalization()(x)
+x = layers.Dropout(0.5)(x)
+x = layers.Dense(128, activation="relu")(x)
+x = layers.BatchNormalization()(x)
+x = layers.Dropout(0.4)(x)
+# --------------------------
+
+outputs = layers.Dense(2, activation="softmax")(x)
+
+model = models.Model(inputs, outputs)
 
 print("Loading weights...")
-model.load_weights("model.keras")
+# Now the shapes match perfectly
+model.load_weights("model.keras") 
 print("Model ready!")
 
 LABELS = ["preserved", "looted"]
@@ -34,6 +49,8 @@ def preprocess(img: Image.Image):
     return arr
 
 def predict(image: Image.Image):
+    if image is None:
+        return None, "No image uploaded."
     x = preprocess(image)
     probs = model.predict(x)[0]
     
@@ -48,8 +65,8 @@ iface = gr.Interface(
     inputs=gr.Image(type="pil", label="Upload photo"),
     outputs=[gr.Label(num_top_classes=2), gr.Textbox(label="Summary")],
     title="Preserved vs Looted Classifier",
-    description="Upload a photo to classify."
+    description="Upload a photo to classify archeological sites."
 )
 
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+    iface.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))