Update app.py

app.py

@@ -1,6 +1,8 @@
 import gradio as gr
 import tensorflow as tf
 import numpy as np
+import cv2
+import matplotlib.cm as cm
 from gradio import components
 
 # Load the model
@@ -12,13 +14,17 @@ class_names = {
     1: 'Dyed lifted polyps'
 }
 
-
-def classify_image(image):
+# Define color gradient segmentation function
+def color_gradient_segmentation(image):
     # Preprocess the image
     img_array = tf.image.resize(image, [256, 256])
     img_array = tf.expand_dims(img_array, 0) / 255.0
 
-    # Make a prediction
+    # Apply color gradient segmentation
+    img_gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    _, img_binary = cv2.threshold(img_gray, 0, 255, cv2.THRESH_BINARY | cv2.THRESH_OTSU)
+
+    # Make a prediction on the segmented image
     prediction = model.predict(img_array)
     predicted_class = tf.argmax(prediction[0], axis=-1)
     confidence = np.max(prediction[0])
@@ -31,16 +37,23 @@ def classify_image(image):
         result = class_name
         details = f"The image is classified as {class_name} with a confidence of {confidence:.2f}."
 
-    return result, confidence, details
+    # Apply color overlay on the segmented region
+    img_overlay = cv2.cvtColor(img_binary, cv2.COLOR_GRAY2RGB)
+    img_overlay[np.where((img_overlay == [255, 255, 255]).all(axis=2))] = [255, 0, 0]  # Red overlay color
+
+    # Blend the original image with the color overlay
+    img_segmented = cv2.addWeighted(image, 0.8, img_overlay, 0.2, 0)
 
+    return result, confidence, details, img_segmented
 
 iface = gr.Interface(
-    fn=classify_image,
+    fn=color_gradient_segmentation,
     inputs=components.Image(shape=(256, 256)),
     outputs=[
         components.Textbox(label="Result"),
         components.Number(label="Confidence"),
-        components.Textbox(label="Details")
+        components.Textbox(label="Details"),
+        components.Image(label="Segmented Image")
     ],
     examples=[
         ['examples/0.jpg'],