Update app.py

app.py

@@ -1,102 +1,113 @@
-
 import os
 import typing as t
 
 import gradio as gr
 import numpy as np
 import tensorflow as tf
+from tensorflow.keras.layers import (Input, Conv2D, BatchNormalization,
+                                     MaxPooling2D, Flatten, Dropout, Dense, LeakyReLU)
+from tensorflow.keras.models import Model
 from PIL import Image
 
-# Try to load a user-provided Keras model if available; else fallback to a pretrained MobileNetV2.
-CUSTOM_MODEL_PATH = "model.h5"
-LABELS_PATH = "labels.txt"
+# Paths
+CUSTOM_MODEL_PATH = "model.h5"               # optional: full Keras model
+MESO_WEIGHTS_PATH = "weights/Meso4_DF"       # your weights-only file
+LABELS = ["real", "fake"]                    # index 0..1 (we'll compute both scores)
+
+# Globals
+MODEL: t.Optional[tf.keras.Model] = None
+IS_MESO = False
+TARGET_SIZE = (256, 256)     # your notebook used 256×256
+THRESHOLD = 0.5              # sigmoid > 0.5 => fake
+
+def build_meso4() -> tf.keras.Model:
+    x = Input(shape=(TARGET_SIZE[0], TARGET_SIZE[1], 3))
+    x1 = Conv2D(8, (3, 3), padding='same', activation='relu')(x)
+    x1 = BatchNormalization()(x1)
+    x1 = MaxPooling2D(pool_size=(2, 2), padding='same')(x1)
+
+    x2 = Conv2D(8, (5, 5), padding='same', activation='relu')(x1)
+    x2 = BatchNormalization()(x2)
+    x2 = MaxPooling2D(pool_size=(2, 2), padding='same')(x2)
 
-# Global objects
-MODEL = None
-USE_IMAGENET_DECODE = False
-CLASS_NAMES: t.Optional[t.List[str]] = None
-TARGET_SIZE = (224, 224)
+    x3 = Conv2D(16, (5, 5), padding='same', activation='relu')(x2)
+    x3 = BatchNormalization()(x3)
+    x3 = MaxPooling2D(pool_size=(2, 2), padding='same')(x3)
 
-def _maybe_load_labels(path: str) -> t.Optional[t.List[str]]:
-    if os.path.exists(path):
-        with open(path, "r", encoding="utf-8") as f:
-            lines = [x.strip() for x in f.readlines() if x.strip()]
-            return lines
-    return None
+    x4 = Conv2D(16, (5, 5), padding='same', activation='relu')(x3)
+    x4 = BatchNormalization()(x4)
+    x4 = MaxPooling2D(pool_size=(4, 4), padding='same')(x4)
+
+    y = Flatten()(x4)
+    y = Dropout(0.5)(y)
+    y = Dense(16)(y)
+    y = LeakyReLU(alpha=0.1)(y)
+    y = Dropout(0.5)(y)
+    y = Dense(1, activation='sigmoid')(y)
+
+    model = Model(inputs=x, outputs=y)
+    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
+    return model
 
 def _load_model():
-    global MODEL, USE_IMAGENET_DECODE, CLASS_NAMES
+    """Load a full Keras model if present; otherwise build Meso4 and load weights."""
+    global MODEL, IS_MESO
+    # 1) Full model (optional)
     if os.path.exists(CUSTOM_MODEL_PATH):
         try:
             MODEL = tf.keras.models.load_model(CUSTOM_MODEL_PATH, compile=False)
-            CLASS_NAMES = _maybe_load_labels(LABELS_PATH)
-            USE_IMAGENET_DECODE = False
+            IS_MESO = False
             print("Loaded custom model from model.h5")
             return
         except Exception as e:
-            print("Failed to load custom model:", e)
+            print("Failed to load model.h5:", e)
+
+    # 2) Meso4 + weights (your case)
+    if os.path.exists(MESO_WEIGHTS_PATH):
+        MODEL = build_meso4()
+        MODEL.load_weights(MESO_WEIGHTS_PATH)
+        IS_MESO = True
+        print("Loaded Meso4 with weights:", MESO_WEIGHTS_PATH)
+        return
 
-    # Fallback: MobileNetV2 pretrained on ImageNet
-    MODEL = tf.keras.applications.MobileNetV2(weights="imagenet")
-    USE_IMAGENET_DECODE = True
-    CLASS_NAMES = None
-    print("Loaded MobileNetV2 (ImageNet) fallback.")
+    # 3) Hard fail (don’t silently switch to ImageNet; this is a deepfake app)
+    raise RuntimeError(
+        "No model found. Upload either model.h5 or weights/Meso4_DF to the Space."
+    )
 
 def _preprocess(img: Image.Image) -> np.ndarray:
-    # Convert to RGB and resize
     img = img.convert("RGB").resize(TARGET_SIZE)
-    arr = np.array(img).astype("float32")
-    # If it's the MobileNetV2 fallback, apply its preprocess; otherwise just scale 0..1
-    if USE_IMAGENET_DECODE:
-        arr = tf.keras.applications.mobilenet_v2.preprocess_input(arr)
-    else:
-        arr = arr / 255.0
-    arr = np.expand_dims(arr, axis=0)
-    return arr
-
-def _decode_predictions(preds: np.ndarray, top: int = 3):
-    # preds: (1, num_classes)
-    preds = preds[0]
-    if USE_IMAGENET_DECODE:
-        decoded = tf.keras.applications.imagenet_utils.decode_predictions(preds[np.newaxis, :], top=top)[0]
-        # decoded is list of tuples: (class_id, class_name, score)
-        return [(name, float(score)) for (_, name, score) in decoded]
-    else:
-        # For custom model: if CLASS_NAMES provided, map; else show class indices
-        top_indices = preds.argsort()[-top:][::-1]
-        out = []
-        for idx in top_indices:
-            label = CLASS_NAMES[idx] if (CLASS_NAMES is not None and idx < len(CLASS_NAMES)) else f"class_{idx}"
-            out.append((label, float(preds[idx])))
-        return out
+    arr = np.array(img).astype("float32") / 255.0
+    return np.expand_dims(arr, axis=0)
 
 def predict(image: Image.Image):
     if image is None:
-        return [], None
+        return {"real": 0.0, "fake": 0.0}, None, "Upload an image."
     x = _preprocess(image)
-    preds = MODEL.predict(x)
-    top3 = _decode_predictions(preds, top=3)
-    # Also return a bar plot-friendly structure for Gradio's Label component
-    scores = {label: score for (label, score) in top3}
-    return scores, image
+    prob_fake = float(MODEL.predict(x, verbose=0)[0][0])
+    prob_real = 1.0 - prob_fake
+    label = "fake" if prob_fake >= THRESHOLD else "real"
+    msg = f"Prediction: {label.upper()}  |  fake={prob_fake:.2f}, real={prob_real:.2f}"
+    # Return both scores for the Label component
+    return {"real": prob_real, "fake": prob_fake}, image, msg
 
-# Initialize
+# Init
 _load_model()
 
-with gr.Blocks(title="Image Classifier (Keras/TF)") as demo:
-    gr.Markdown("# Image Classifier\nUpload an image to classify using a Keras model.\n\n"
-                "- Drop in your own `model.h5` (and optional `labels.txt`) to switch from ImageNet to your custom model.\n"
-                "- For custom models, ensure input size is 224x224x3 or adjust code.\n")
+with gr.Blocks(title="Deepfake Detector (Meso4)") as demo:
+    gr.Markdown("# Deepfake Detector (Meso4)\n"
+                "Upload a face image (or a frame from a video). The model outputs real vs fake.")
 
     with gr.Row():
         with gr.Column(scale=1):
             inp = gr.Image(type="pil", label="Upload image")
             btn = gr.Button("Predict")
         with gr.Column(scale=1):
-            out_label = gr.Label(num_top_classes=3, label="Top Predictions")
+            out_label = gr.Label(num_top_classes=2, label="Scores")
             out_img = gr.Image(type="pil", label="Preview")
+            out_text = gr.Markdown()
 
-    btn.click(fn=predict, inputs=inp, outputs=[out_label, out_img])
+    btn.click(fn=predict, inputs=inp, outputs=[out_label, out_img, out_text])
 
 if __name__ == "__main__":
     demo.launch()