Update app.py

app.py

@@ -1,6 +1,7 @@
 import tensorflow as tf
 from PIL import Image
 import numpy as np
+import cv2
 import gradio as gr
 from tensorflow import keras
 from keras.models import load_model
@@ -9,32 +10,49 @@ from keras.models import load_model
 model = "Hugging_face_model_final.h5"
 model = tf.keras.models.load_model(model)
 
-# Load the mineral detection model mineral_detection_model_Final_4_18_2024.h5
+# Load the mineral detection model
 mineral_detection_model = tf.keras.models.load_model("mineral_detection_model_Final_4_18_2024.h5")
 
 # Define the class labels
 class_labels = ['biotite', 'granite', 'olivine', 'plagioclase', 'staurolite']
 
+# Function to preprocess the image for mineral detection
+def preprocess_image_detection(img_array):
+    if img_array is None:
+        return None
+    img = (img_array * 255).astype(np.uint8)  # Convert back to uint8
+    img_array = cv2.resize(img, (150, 150))  # Resize to 150x150
+    img_array = img_array.astype(np.uint8)
+    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
+    return img_array
+
+# Function to preprocess the image for classification
+def preprocess_image_classification(img_array):
+    if img_array is None:
+        return None
+    img = (img_array * 255).astype(np.uint8)  # Convert back to uint8
+    img_array = cv2.resize(img, (224, 224))  # Resize to 224x224
+    img_array = img_array.astype(np.uint8)
+    img_array = np.expand_dims(img_array, axis=0)  # Add batch dimension
+    return img_array
+
 # Define the function to detect if the input is a mineral
 def detect_mineral(image):
-    # Preprocess the image for mineral detection model
     if image is not None:
         image = Image.fromarray(np.array(image).astype(np.uint8), 'RGB')
-        # Rest of your code for detecting minerals
+        image = np.array(image)
+        image = Image.fromarray(image.astype(np.uint8), 'RGB')
+        image = image.resize((150, 150))  # Assuming the model expects 150x150 images
+        image = np.array(image) / 255.0
+        image = np.expand_dims(image, axis=0)
+
+        # Make prediction
+        prediction = mineral_detection_model.predict(image)
+        is_mineral = prediction[0][0] < 0.5  # Assuming binary classification
+        return is_mineral
     else:
         # Handle the case where no image is provided
         return "No image provided."
-    image = np.array(image)
-    image = Image.fromarray(image.astype(np.uint8), 'RGB')
-    image = image.resize((150, 150))  # Assuming the model expects 224x224 images
-    image = np.array(image) / 255.0
-    image = np.expand_dims(image, axis=0)
-
-    # Make prediction
-    prediction = mineral_detection_model.predict(image)
-    is_mineral = prediction[0][0] < 0.5  # Assuming binary classification
-
-    return is_mineral
 
 # Define the function to make predictions
 def classify_image(image):
@@ -44,11 +62,9 @@ def classify_image(image):
         return "Input is not a mineral.", ""
 
     # Preprocess the image for classification
-    image = np.array(image)
-    image = Image.fromarray(image.astype(np.uint8), 'RGB')
-    image = image.resize((224, 224))
-    image = np.array(image) / 255.0
-    image = np.expand_dims(image, axis=0)
+    image = preprocess_image_classification(np.array(image))
+    if image is None:
+        return "Error preprocessing image.", ""
 
     # Make prediction
     prediction = model.predict(image)
@@ -74,4 +90,5 @@ with gr.Blocks() as demo:
         ]
         image_button = gr.Button("Classify Image")
         image_button.click(classify_image, inputs=image_input, outputs=output_components)
-    demo.launch( share=True, auth_message="Welcome To Mineral Identification App, Please login to use the app.")
+
+demo.launch(share=True, auth_message="Welcome To Mineral Identification App, Please login to use the app.")