import subprocess
import importlib

# List of required packages
packages = ['cv2', 'numpy', 'tensorflow', 'gr']

# Mapping of package names to their PyPI names
package_mapping = {'cv2': 'opencv-python-headless'}

# Check and install packages if needed
for package in packages:
    try:
        importlib.import_module(package)
    except ImportError:
        print(f"{package} not found. Installing...")
        
        # Get the correct package name from the mapping
        install_package = package_mapping.get(package, package)
        
        install_command = ['pip', 'install', install_package]
        subprocess.run(install_command, check=True)

# Now that all packages are installed, import them
import cv2
import numpy as np
import tensorflow as tf
import gradio as gr

# Your code using these packages can follow here
def remove_background_deeplab(image):
    # 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()

    # Convert the Gradio Image object to a NumPy array
    image_np = np.array(image.data)

    # Resize the image to match the expected input size of the model
    input_size = (257, 257)
    image_resized = cv2.resize(image_np, 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_np.shape[1], image_np.shape[0]))

    # Multiply the image with the binary mask to get the result
    result = image_np * binary_mask[:, :, np.newaxis]

    # Convert the result back to Gradio Image object
    result_image = gr.Image(result)

    return result_image

# Example usage
gr.Interface(
    fn=remove_background_deeplab,
    inputs=gr.Image(label='Drop an Image or Open Camera to Classify'),
    outputs=gr.Image()
).launch()