Update app.py

app.py

@@ -1,87 +1,75 @@
 import gradio as gr
-import tensorflow as tf
+from PIL import Image
 import numpy as np
-from tensorflow.keras.preprocessing import image
+from tensorflow.keras.preprocessing import image as keras_image
+from tensorflow.keras.applications.resnet50 import preprocess_input as resnet_preprocess_input
+from tensorflow.keras.models import load_model
 
 # Klassenlabels (ersetzen Sie durch Ihre spezifischen Klassen)
 class_labels = ['bench_press', 'deadlift', 'hip_thrust', 'lat_pulldown', 'pull_up', 'squat', 'tricep_dips']
 
-def preprocess(img):
+# Load your trained models
+resnet_model = load_model('exercise_classification_model_resnet50.h5')
+cnn_model = load_model('exercise_classification_model.h5')
+
+def predict_exercise(img, model_type):
     try:
-        if img is None:
-            raise ValueError("No image uploaded")
-        print("Image received:", type(img), img)
+        # Convert the image to RGB format
+        img = Image.fromarray(img.astype('uint8'), 'RGB')
+        # Resize the image to the required input size of the model
         img = img.resize((150, 150))
-        img_array = image.img_to_array(img)
-        print("Image array shape:", img_array.shape)
+        # Convert the image to an array
+        img_array = keras_image.img_to_array(img)
+        # Expand dimensions to match the model's input shape
         img_array = np.expand_dims(img_array, axis=0)
-        img_array /= 255.0
-        print("Processed image array shape:", img_array.shape)
-        return img_array
+        
+        if model_type == 'ResNet50':
+            # Preprocess the input as expected by ResNet50
+            img_array = resnet_preprocess_input(img_array)
+            # Predict using the ResNet50 model
+            prediction = resnet_model.predict(img_array)
+        elif model_type == 'CNN':
+            # Preprocess the input as expected by the CNN
+            img_array /= 255.0  # Normalisierung
+            # Predict using the CNN model
+            prediction = cnn_model.predict(img_array)
+        else:
+            return {"error": "Invalid model type selected"}
+        
+        # Debugging: print the prediction shape
+        print(f"Prediction shape: {prediction.shape}")
+        
+        # Check prediction shape
+        if prediction.shape[1] == len(class_labels):  # Ensure the prediction matches the number of classes
+            # Return the prediction as a dictionary with class probabilities
+            return {class_labels[i]: float(prediction[0][i]) for i in range(len(class_labels))}
+        else:
+            return {"error": f"Unexpected prediction shape: {prediction.shape}"}
     except Exception as e:
-        print(f"Error in preprocessing: {e}")
-        return None
-
-# Laden der Modelle
-try:
-    model_cnn = tf.keras.models.load_model('exercise_classification_model.h5')
-    print("CNN model loaded successfully")
-except Exception as e:
-    print(f"Error loading CNN model: {e}")
-
-try:
-    model_resnet50 = tf.keras.models.load_model('exercise_classification_model_resnet50.h5')
-    print("ResNet50 model loaded successfully")
-except Exception as e:
-    print(f"Error loading ResNet50 model: {e}")
+        return {"error": str(e)}
 
-def predict_cnn(img):
-    try:
-        img_array = preprocess(img)
-        if img_array is None:
-            return {label: 0.0 for label in class_labels}
-        prediction = model_cnn.predict(img_array)
-        print("Prediction:", prediction)
-        return {class_labels[i]: float(prediction[0][i]) for i in range(len(class_labels))}
-    except Exception as e:
-        print(f"Error in CNN prediction: {e}")
-        return {label: 0.0 for label in class_labels}
+# Custom CSS for Dark Mode and Elegant Design
+custom_css = """
+body {background-color: #121212; color: #e0e0e0; font-family: 'Arial', sans-serif;}
+h1 {color: #ff5722;}
+label {color: #ff9800;}
+input[type=radio] {accent-color: #ff5722;}
+button:hover {background-color: #e64a19;}
+.footer {display: none !important;}
+"""
 
-def predict_resnet50(img):
-    try:
-        img_array = preprocess(img)
-        if img_array is None:
-            return {label: 0.0 for label in class_labels}
-        print(f"Processed image array: {img_array}")  # Debugging-Information
-        prediction = model_resnet50.predict(img_array)
-        print("Prediction:", prediction)
-        return {class_labels[i]: float(prediction[0][i]) for i in range(len(class_labels))}
-    except Exception as e:
-        print(f"Error in ResNet50 prediction: {e}")
-        return {label: 0.0 for label in class_labels}
-
-# Gradio Interface für das CNN Modell
-interface_cnn = gr.Interface(
-    fn=predict_cnn,
-    inputs=gr.Image(type="pil"),  # Verwenden Sie gr.Image
-    outputs=gr.Label(num_top_classes=3),
-    live=True,
-    title="Exercise Classification with CNN",
-    description="Upload an image to classify the exercise using CNN model."
+# Define the Gradio interface
+interface = gr.Interface(
+    fn=predict_exercise,
+    inputs=[
+        gr.Image(type="numpy", label="Upload an image of the exercise"),
+        gr.Radio(['ResNet50', 'CNN'], label="Choose Model")
+    ],
+    outputs=gr.Label(num_top_classes=3, label="Prediction"),
+    title="Exercise Classifier",
+    description="Upload an image of an exercise and the classifier will predict the exercise.",
+    css=custom_css  # Apply the custom CSS
 )
 
-# Gradio Interface für das ResNet50 Modell
-interface_resnet50 = gr.Interface(
-    fn=predict_resnet50,
-    inputs=gr.Image(type="pil"),  # Verwenden Sie gr.Image
-    outputs=gr.Label(num_top_classes=3),
-    live=True,
-    title="Exercise Classification with ResNet50",
-    description="Upload an image to classify the exercise using ResNet50 model."
-)
-
-# Kombiniere die Interfaces in eine Gradio App
-demo = gr.TabbedInterface([interface_cnn, interface_resnet50], ["CNN Model", "ResNet50 Model"])
-
-if __name__ == "__main__":
-    demo.launch()
+# Launch the interface
+interface.launch()