Spaces:
Runtime error
Runtime error
File size: 3,256 Bytes
07b6f54 |
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 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 |
import base64
import numpy as np
import cv2
import tensorflow as tf
from tensorflow.keras.utils import CustomObjectScope
def iou(y_true, y_pred):
def f(y_true, y_pred):
intersection = (y_true * y_pred).sum()
union = y_true.sum() + y_pred.sum() - intersection
x = (intersection + 1e-15) / (union + 1e-15)
x = x.astype(np.float32)
return x
return tf.numpy_function(f, [y_true, y_pred], tf.float32)
smooth = 1e-15
def dice_coef(y_true, y_pred):
y_true = tf.keras.layers.Flatten()(y_true)
y_pred = tf.keras.layers.Flatten()(y_pred)
intersection = tf.reduce_sum(y_true * y_pred)
return (2. * intersection + smooth) / (tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) + smooth)
def dice_loss(y_true, y_pred):
return 1.0 - dice_coef(y_true, y_pred)
def predict():
#Load Model
interpreter = tf.lite.Interpreter(model_path='model/segmentation_model.tflite')
interpreter.allocate_tensors()
# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# Global Parameters
H = 512
W = 384
# Set up your input data.
image = cv2.imread('model/test.jpg', cv2.IMREAD_COLOR)
resized_image = cv2.resize(image, (W, H))
x = resized_image/255.0
x = x.astype(np.float32)
x = np.expand_dims(x,0)
# Input Data to Model
interpreter.set_tensor(input_details[0]['index'], x)
interpreter.invoke()
#Geting Output
prediction = interpreter.get_tensor(output_details[0]['index'])[0]
prediction = cv2.resize(prediction, (W, H))
prediction = np.expand_dims(prediction, axis=-1)
prediction = prediction > 0.5
mask = np.uint8(prediction) * 255
prediction = cv2.merge([mask, mask, mask])
mask_inv = cv2.cvtColor(prediction, cv2.COLOR_BGR2GRAY)
output_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2BGRA)
output_image[..., 3] = mask_inv
# cv2.imshow(output_image)
cv2.imwrite("model/output.png", output_image)
return "done"
def predict_h5(file):
#Load Model
with CustomObjectScope({'iou': iou, 'dice_coef': dice_coef, 'dice_loss': dice_loss}):
model = tf.keras.models.load_model("model/model.h5")
# Global Parameters
H = 512
W = 384
# Set up your input data.
# Convert File to np array
file_bytes = np.fromstring(file, np.uint8)
# convert numpy array to image
image = cv2.imdecode(file_bytes, cv2.IMREAD_UNCHANGED)
resized_image = cv2.resize(image, (W, H))
x = resized_image/255.0
x = x.astype(np.float32)
x = np.expand_dims(x,0)
# Input Data to Model
#Geting Output
prediction = model.predict(x)[0]
prediction = cv2.resize(prediction, (W, H))
prediction = np.expand_dims(prediction, axis=-1)
prediction = prediction > 0.5
mask = np.uint8(prediction) * 255
prediction = cv2.merge([mask, mask, mask])
mask_inv = cv2.cvtColor(prediction, cv2.COLOR_BGR2GRAY)
output_image = cv2.cvtColor(resized_image, cv2.COLOR_BGR2BGRA)
output_image[..., 3] = mask_inv
retval, buffer = cv2.imencode('.png', output_image)
encoded_image = base64.b64encode(buffer).decode('utf-8')
return encoded_image |