Update app.py

app.py

@@ -3,33 +3,27 @@ import gradio as gr
 import numpy as np
 import tensorflow as tf
 import cv2
+from mrcnn import model as modellib
+from mrcnn.config import Config
 
 # Set environment variable to avoid floating-point errors
 os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0'
 
-# 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
+# Define Mask R-CNN configuration
+class InferenceConfig(Config):
+    NAME = "mask_rcnn"
+    NUM_CLASSES = 1 + 80  # Update according to your dataset
+    GPU_COUNT = 1
+    IMAGES_PER_GPU = 1
+
+config = InferenceConfig()
 
-# Build the model and load weights
-model = build_mask_rcnn_model()
+# Rebuild the Mask R-CNN model
+model = modellib.MaskRCNN(mode="inference", config=config, model_dir=os.getcwd())
 
 # 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)
+model_path = os.path.join('toolkit', 'condmodel_100.h5')
+model.load_weights(model_path, by_name=True)
 print("Mask R-CNN model loaded successfully with weights.")
 
 # Function to apply Mask R-CNN for image segmentation
@@ -39,15 +33,17 @@ def apply_mask_rcnn(image):
         if image.shape[2] == 4:  # Convert RGBA to RGB
             image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
 
-        # 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
+        # Resize the image to match the model input size (for inference)
+        resized_image = cv2.resize(image, (1024, 1024))  # 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 = model.predict(input_image)
+        # Use Mask R-CNN to predict
+        result = model.detect(input_image, verbose=0)
+        r = result[0]
 
-        # Extract mask (assumed to be the first output)
-        mask = np.squeeze(prediction[0])
+        # Create a mask for the detected objects
+        mask = r['masks']
+        mask = np.sum(mask, axis=-1)  # Combine masks for all objects
 
         # Resize mask back to the original image size
         mask = cv2.resize(mask, (image.shape[1], image.shape[0]))