Spaces:
Sleeping
Sleeping
File size: 2,419 Bytes
fa811d4 a2e8183 fa811d4 edf8500 fa811d4 edf8500 fa811d4 172db72 edf8500 a2e8183 7bc6fbb a2e8183 7bc6fbb a2e8183 fa811d4 edf8500 7bc6fbb edf8500 7bc6fbb edf8500 7bc6fbb 2fca5ba edf8500 2fca5ba fa811d4 2fca5ba fa811d4 a2e8183 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 |
import gradio as gr
import tensorflow as tf
import numpy as np
import cv2
from keras.utils import normalize
from PIL import Image
def dice_coef(y_true, y_pred):
smooth = 1e-5
intersection = K.sum(y_true * y_pred, axis=[1, 2, 3])
union = K.sum(y_true, axis=[1, 2, 3]) + K.sum(y_pred, axis=[1, 2, 3])
return K.mean((2.0 * intersection + smooth) / (union + smooth), axis=0)
def predict_segmentation(image):
original_size = (image.shape[1], image.shape[0]) # (width, height)
# Resize to the model's input size
SIZE_X = 128
SIZE_Y = 128
img = cv2.resize(image, (SIZE_Y, SIZE_X))
if len(img.shape) == 3 and img.shape[2] == 3: # If the image is RGB
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) # Convert to grayscale
img = np.expand_dims(img, axis=2) # Add the channel dimension
img = normalize(img, axis=1)
X_test = np.expand_dims(img, axis=0) # Add the batch dimension
custom_objects = {'dice_coef': dice_coef}
with tf.keras.utils.custom_object_scope(custom_objects):
model = tf.keras.models.load_model("model100.h5")
# Get the prediction
prediction = model.predict(X_test)
predicted_img = np.argmax(prediction, axis=3)[0, :, :]
# Resize prediction back to original image size
predicted_img_resized = cv2.resize(predicted_img, original_size, interpolation=cv2.INTER_NEAREST)
# Create an RGBA image with a transparent background
rgba_img = np.zeros((predicted_img_resized.shape[0], predicted_img_resized.shape[1], 4), dtype=np.uint8)
# Define the color for the segmented area (e.g., red)
segmented_color = [255, 0, 0] # Red color in RGB
# Set the segmented area to the desired color
for i in range(3):
rgba_img[:, :, i] = np.where(predicted_img_resized > 0, segmented_color[i], 0)
# Create an alpha channel: 255 where there is segmentation, 0 otherwise
rgba_img[:, :, 3] = np.where(predicted_img_resized > 0, 255, 0)
# Convert the numpy array to an image
output_image = Image.fromarray(rgba_img)
# Save the image as PNG to return it
output_image_path = "/tmp/segmented_output.png"
output_image.save(output_image_path)
return output_image_path
# Gradio Interface
iface = gr.Interface(
fn=predict_segmentation,
inputs="image",
outputs="file", # Return the file path to download the PNG
live=False
)
iface.launch(share=True)
|