Update app.py

app.py

@@ -1,35 +1,53 @@
+import os
 import gradio as gr
 import numpy as np
-import cv2
-import os
 import tensorflow as tf
-from tensorflow.keras.models import load_model
+import cv2
+
+# Set environment variable to avoid floating-point errors
+os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
 
-# Load the Mask R-CNN model
-model_path = os.path.join('toolkit', 'condmodel_100.h5')  # Path to your model
-mask_rcnn_model = load_model(model_path)
+# Define the Mask R-CNN model architecture
+def build_mask_rcnn_model():
+    input_layer = tf.keras.layers.Input(shape=(224, 224, 3))  # Adjust input shape to match your model
+    # Example architecture, you should modify it to match your actual Mask R-CNN model
+    x = tf.keras.layers.Conv2D(64, (3, 3), activation='relu', padding='same')(input_layer)
+    x = tf.keras.layers.MaxPooling2D((2, 2))(x)
+    x = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same')(x)
+    x = tf.keras.layers.MaxPooling2D((2, 2))(x)
+    x = tf.keras.layers.Conv2D(256, (3, 3), activation='relu', padding='same')(x)
+    x = tf.keras.layers.UpSampling2D((2, 2))(x)
+    x = tf.keras.layers.Conv2D(128, (3, 3), activation='relu', padding='same')(x)
+    x = tf.keras.layers.UpSampling2D((2, 2))(x)
+    output_layer = tf.keras.layers.Conv2D(1, (1, 1), activation='sigmoid')(x)
+    
+    model = tf.keras.models.Model(inputs=input_layer, outputs=output_layer)
+    return model
+
+# Build the model and load weights
+model = build_mask_rcnn_model()
+
+# Load the Mask R-CNN model weights
+model_path = os.path.join('toolkit', 'condmodel_100.h5')  # Update with correct path
+model.load_weights(model_path)
+print("Mask R-CNN model loaded successfully with weights.")
 
+# Function to apply Mask R-CNN for image segmentation
 def apply_mask_rcnn(image):
-    """
-    Function to apply the Mask R-CNN model and return the segmented image.
-    :param image: Input image in numpy array format
-    :return: Image with segmentation mask overlaid
-    """
     try:
         # Convert image to RGB (in case of RGBA or grayscale)
         if image.shape[2] == 4:  # Convert RGBA to RGB
             image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
 
-        # Resize the image to the input size of the model
-        resized_image = cv2.resize(image, (224, 224))  # Adjust according to model input size
+        # Resize the image to match the model input size
+        resized_image = cv2.resize(image, (224, 224))  # Adjust based on the input shape of your model
         input_image = np.expand_dims(resized_image, axis=0)
 
         # Use Mask R-CNN to predict the mask
-        prediction = mask_rcnn_model.predict(input_image)
+        prediction = model.predict(input_image)
 
-        # Assuming the first output is the mask, you may need to adjust based on your model's structure
-        mask = prediction[0]
-        mask = np.squeeze(mask)  # Remove any unnecessary dimensions
+        # Extract mask (assumed to be the first output)
+        mask = np.squeeze(prediction[0])
 
         # Resize mask back to the original image size
         mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
@@ -47,16 +65,16 @@ def apply_mask_rcnn(image):
         print(f"Error in segmentation: {e}")
         return image  # Return original image if segmentation fails
 
-# Update Gradio interface for image input/output
+# Gradio interface definition
 inputs = gr.Image(source="upload", tool="editor", type="numpy", label="Upload an image")
 outputs = gr.Image(type="numpy", label="Segmented Image")
 
-# Gradio interface definition
+# Gradio app layout
 with gr.Blocks() as demo:
     gr.Markdown("<h1 style='text-align: center;'>Image Segmentation with Mask R-CNN</h1>")
     gr.Markdown("Upload an image to see segmentation results using the Mask R-CNN model.")
     
-    # Input and output components
+    # Input and output layout
     with gr.Row():
         with gr.Column():
             gr.Markdown("### Upload an Image")