Update app.py

app.py

@@ -1,16 +1,49 @@
 import gradio as gr
 import tensorflow as tf
 import numpy as np
-from huggingface_hub import hf_hub_download
+from huggingface_hub import hf_hub_download, list_repo_files
 
-# Function to load model from H5 file
-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 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 []
 
-# Load models from Hugging Face Hub
-model1 = load_model_from_hub("arsath-sm/face_classification_model1", "model.h5")
-model2 = load_model_from_hub("arsath-sm/face_classification_model2", "model.h5")
+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
 
 def preprocess_image(image):
     img = tf.convert_to_tensor(image)
@@ -19,22 +52,31 @@ def preprocess_image(image):
     return tf.expand_dims(img, 0)
 
 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 = {}
     
-    # Make predictions using both models
-    pred1 = model1.predict(preprocessed_image)[0][0]
-    pred2 = model2.predict(preprocessed_image)[0][0]
+    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"
     
-    # 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
+    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"
     
-    return {
-        "Model 1 Prediction": f"{result1} (Confidence: {confidence1:.2f})",
-        "Model 2 Prediction": f"{result2} (Confidence: {confidence2:.2f})"
-    }
+    return results
 
 # Create the Gradio interface
 iface = gr.Interface(
@@ -42,7 +84,8 @@ iface = gr.Interface(
     inputs=gr.Image(),
     outputs={
         "Model 1 Prediction": gr.Textbox(),
-        "Model 2 Prediction": gr.Textbox()
+        "Model 2 Prediction": gr.Textbox(),
+        "Error": 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."