Update app.py

app.py

@@ -3,12 +3,6 @@ from tensorflow.keras.models import load_model
 import cv2
 import numpy as np
 
-# Load the trained model
-model = load_model('FightOS_CNN_Models.h5')
-
-# Define the class labels
-class_labels = ['positive', 'negative']
-
 # Function to preprocess the image (resize, normalize, etc.)
 def preprocess_image(image):
     # Resize the image to 224x224
@@ -22,7 +16,7 @@ def preprocess_image(image):
     
     return preprocessed_image
 
-# Function to segment the image using Otsu's thresholding
+# Function to segment the image using Otsu's thresholding and post-processing
 def segment_image(image):
     # Convert the image to grayscale
     grayscale_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
@@ -30,31 +24,38 @@ def segment_image(image):
     # Apply Otsu's thresholding
     _, thresholded_image = cv2.threshold(grayscale_image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
     
-    return thresholded_image
+    # Perform morphological operations for post-processing
+    kernel = np.ones((5, 5), np.uint8)
+    segmented_image = cv2.morphologyEx(thresholded_image, cv2.MORPH_OPEN, kernel, iterations=2)
+    segmented_image = cv2.morphologyEx(segmented_image, cv2.MORPH_CLOSE, kernel, iterations=2)
+    
+    return segmented_image
 
 # Function to make the classification prediction and return the segmented image
 def classify_pcos(image):
-    try:
-        # Preprocess the image (resize, normalize, etc.)
-        preprocessed_image = preprocess_image(image)
-        
-        # Make the prediction using the model
-        prediction = model.predict(preprocessed_image)
-        
-        # Get the predicted class
-        class_index = np.argmax(prediction)
-        class_label = class_labels[class_index]
-        
-        # Get the probability of each class
-        probabilities = {label: round(float(prediction[0][i]), 2) for i, label in enumerate(class_labels)}
-        
-        # Segment the image
-        segmented_image = segment_image(image)
-        
-        return class_label, probabilities, segmented_image
-    except Exception as e:
-        error_message = f"Error occurred: {str(e)}"
-        return error_message, {}, None
+    # Preprocess the image (resize, normalize, etc.)
+    preprocessed_image = preprocess_image(image)
+    
+    # Make the prediction using the model
+    prediction = model.predict(preprocessed_image)
+    
+    # Get the predicted class
+    class_index = np.argmax(prediction)
+    class_label = class_labels[class_index]
+    
+    # Get the probability of each class
+    probabilities = {label: round(float(prediction[0][i]), 2) for i, label in enumerate(class_labels)}
+    
+    # Segment the image
+    segmented_image = segment_image(image)
+    
+    return class_label, probabilities, segmented_image
+
+# Load the trained model
+model = load_model('FightOS_CNN_Models.h5')
+
+# Define the class labels
+class_labels = ['positive', 'negative']
 
 # Gradio app interface with multiple outputs
 iface = gr.Interface(