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()