File size: 1,665 Bytes
bc7f524 |
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 |
import cv2
import numpy as np
import tensorflow as tf
import gradio as gr
def remove_background_deeplab(image_path):
# Load the TensorFlow Lite model
interpreter = tf.lite.Interpreter(model_path="2.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# Load and preprocess the image
image = cv2.imread(image_path)
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# Resize the image to match the expected input size of the model
input_size = (257, 257)
image_resized = cv2.resize(image_rgb, input_size)
# Normalize the input image
input_tensor = (image_resized / 127.5 - 1.0).astype(np.float32)
# Set the input tensor to the model
interpreter.set_tensor(input_details[0]['index'], np.expand_dims(input_tensor, axis=0))
# Run inference
interpreter.invoke()
# Get the segmented mask
predictions = interpreter.get_tensor(output_details[0]['index'])
mask = np.argmax(predictions, axis=-1)[0]
# Resize the binary mask to match the shape of the image
binary_mask = cv2.resize(np.where(mask == 15, 1, 0).astype(np.uint8), (image.shape[1], image.shape[0]))
# Multiply the image with the binary mask to get the result
result = image * binary_mask[:, :, np.newaxis]
# Display the result or save it
cv2.imshow('Result', result)
cv2.waitKey(0)
cv2.destroyAllWindows()
# Example usage
# remove_background_deeplab('images (2).jpeg')
gr.Interface(fn=remove_background_deeplab,inputs=gr.Image(label='Drop an Image or Open Camera to Classify'),outputs=gr.Image()).launch() |