import cv2
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import gradio as gr

# Load the DeepLabV3 model from TensorFlow Hub
model = hub.load("https://tfhub.dev/tensorflow/deeplabv3/1")

# Function to run inference on a single image
def run_inference(model, image):
    image = tf.convert_to_tensor(image)
    image = tf.image.convert_image_dtype(image, tf.uint8)
    image = tf.image.resize(image, [513, 513])
    image = tf.expand_dims(image, axis=0)
    result = model(image)
    return result['semantic_pred'][0]

# Function to apply background removal
def remove_background(image, mask):
    mask = cv2.resize(mask.numpy(), (image.shape[1], image.shape[0]), interpolation=cv2.INTER_NEAREST)
    mask = mask == 15  # Assuming class 15 is the foreground (this might change based on the model and dataset)
    result = np.where(mask[..., None], image, 0)
    return result

# Gradio function to process the input and return the result
def process_image(image):
    # Convert the image to RGB
    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
    
    # Run inference to get the segmentation mask
    mask = run_inference(model, image)
    
    # Remove the background
    result = remove_background(image, mask)
    
    # Convert the result back to BGR for display
    result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
    
    return result

# Create a Gradio interface
iface = gr.Interface(fn=process_image, 
                     inputs=gr.inputs.Image(source="webcam", tool="editor"), 
                     outputs=gr.outputs.Image(type="auto"),
                     live=True)

# Launch the interface
iface.launch()