DHEIVER's picture
Update app.py
59c942f
raw
history blame contribute delete
No virus
2.08 kB
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
model = tf.keras.models.load_model('stomach.h5')
# Define the class names
class_names = {
0: 'Esophagitis',
1: 'Dyed lifted polyps'
}
# 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
# 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])
if confidence < 0.6:
result = "Unable to detect"
details = "The model was unable to confidently classify the image."
else:
class_name = class_names[predicted_class.numpy()]
result = class_name
details = f"The image is classified as {class_name} with a confidence of {confidence:.2f}."
# 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=color_gradient_segmentation,
inputs=components.Image(shape=(256, 256)),
outputs=[
components.Textbox(label="Result"),
components.Number(label="Confidence"),
components.Textbox(label="Details"),
components.Image(label="Segmented Image")
],
examples=[
['examples/0.jpg'],
['examples/1.jpg'],
['examples/2.jpg'],
['examples/3.jpg']
]
)
iface.launch()