Update app.py

app.py

@@ -1,94 +1,50 @@
 import gradio as gr
 import tensorflow as tf
 import numpy as np
-from huggingface_hub import hf_hub_download, list_repo_files
+from huggingface_hub import hf_hub_download
 
-def list_files_in_repo(repo_id):
-    try:
-        files = list_repo_files(repo_id)
-        print(f"Files in {repo_id}:")
-        for file in files:
-            print(file)
-        return files
-    except Exception as e:
-        print(f"Error listing files in {repo_id}: {str(e)}")
-        return []
+# Function to load model from Hugging Face Hub
+def load_model_from_hub(repo_id, filename):
+    model_path = hf_hub_download(repo_id=repo_id, filename=filename)
+    return tf.keras.models.load_model(model_path)
 
-def load_model_from_hub(repo_id):
-    files = list_files_in_repo(repo_id)
-    model_file = next((f for f in files if f.endswith('.h5') or f.endswith('.keras')), None)
-    
-    if model_file is None:
-        raise ValueError(f"No .h5 or .keras file found in {repo_id}")
-    
-    try:
-        model_path = hf_hub_download(repo_id=repo_id, filename=model_file)
-        return tf.keras.models.load_model(model_path)
-    except Exception as e:
-        print(f"Error loading model from {repo_id}: {str(e)}")
-        raise
-
-# Try to load models
-try:
-    print("Attempting to load Model 1...")
-    model1 = load_model_from_hub("arsath-sm/face_classification_model1")
-    print("Model 1 loaded successfully.")
-except Exception as e:
-    print(f"Failed to load Model 1: {str(e)}")
-    model1 = None
-
-try:
-    print("Attempting to load Model 2...")
-    model2 = load_model_from_hub("arsath-sm/face_classification_model2")
-    print("Model 2 loaded successfully.")
-except Exception as e:
-    print(f"Failed to load Model 2: {str(e)}")
-    model2 = None
+# Load models from Hugging Face Hub
+model1 = load_model_from_hub("arsath-sm/face_classification_model1", "face_classification_model1.h5")
+model2 = load_model_from_hub("arsath-sm/face_classification_model2", "face_classification_model2.h5")
 
 def preprocess_image(image):
-    img = tf.convert_to_tensor(image)
-    img = tf.image.resize(img, (150, 150))
-    img = img / 255.0
-    return tf.expand_dims(img, 0)
+    img = tf.image.resize(image, (224, 224))  # Resize to match the input size of your models
+    img = tf.cast(img, tf.float32) / 255.0    # Normalize pixel values
+    return tf.expand_dims(img, 0)             # Add batch dimension
 
 def predict_image(image):
-    if model1 is None and model2 is None:
-        return {
-            "Error": "Both models failed to load. Please check the model repositories and try again."
-        }
-    
     preprocessed_image = preprocess_image(image)
-    results = {}
     
-    if model1 is not None:
-        pred1 = model1.predict(preprocessed_image)[0][0]
-        result1 = "Real" if pred1 > 0.5 else "Fake"
-        confidence1 = pred1 if pred1 > 0.5 else 1 - pred1
-        results["Model 1 Prediction"] = f"{result1} (Confidence: {confidence1:.2f})"
-    else:
-        results["Model 1 Prediction"] = "Model failed to load"
+    # Make predictions using both models
+    pred1 = model1.predict(preprocessed_image)[0][0]
+    pred2 = model2.predict(preprocessed_image)[0][0]
     
-    if model2 is not None:
-        pred2 = model2.predict(preprocessed_image)[0][0]
-        result2 = "Real" if pred2 > 0.5 else "Fake"
-        confidence2 = pred2 if pred2 > 0.5 else 1 - pred2
-        results["Model 2 Prediction"] = f"{result2} (Confidence: {confidence2:.2f})"
-    else:
-        results["Model 2 Prediction"] = "Model failed to load"
+    # Prepare results for each model
+    result1 = "Real" if pred1 > 0.5 else "Fake"
+    confidence1 = pred1 if pred1 > 0.5 else 1 - pred1
+    result2 = "Real" if pred2 > 0.5 else "Fake"
+    confidence2 = pred2 if pred2 > 0.5 else 1 - pred2
     
-    return results
+    return {
+        "Model 1 (ResNet) Prediction": f"{result1} (Confidence: {confidence1:.2f})",
+        "Model 2 (Inception) Prediction": f"{result2} (Confidence: {confidence2:.2f})"
+    }
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict_image,
     inputs=gr.Image(),
     outputs={
-        "Model 1 Prediction": gr.Textbox(),
-        "Model 2 Prediction": gr.Textbox(),
-        "Error": gr.Textbox()
+        "Model 1 (ResNet) Prediction": gr.Textbox(),
+        "Model 2 (Inception) Prediction": gr.Textbox()
     },
-    title="Real vs AI Face Classification",
-    description="Upload an image to classify whether it's a real face or an AI-generated face using two different models."
+    title="Real vs AI-Generated Face Classification",
+    description="Upload an image to classify whether it's a real face or an AI-generated face using two different models: ResNet-style and Inception-style."
 )
 
 # Launch the app