import gradio as gr 
import numpy as np
from PIL import Image
import onnxruntime

def preprocess_image(directory):
    inputs = Image.open(directory)
    new_img = np.array(inputs)
    new_img = np.expand_dims(new_img.transpose((2, 0, 1)), 0)
    new_img = ((new_img - 127.5) / 127.5).astype(np.float32)
    return new_img

def color_map(N=256):
    def bitget(byteval, idx):
        return ((byteval & (1 << idx)) != 0)

    cmap = np.zeros((N, 3), dtype=np.uint8)
    for i in range(N):
        r = g = b = 0
        c = i
        for j in range(8):
            r = r | (bitget(c, 0) << 7-j)
            g = g | (bitget(c, 1) << 7-j)
            b = b | (bitget(c, 2) << 7-j)
            c = c >> 3

        cmap[i] = np.array([r, g, b])
    return cmap

def decode_segmap(temp):
    cmap = color_map()
    r = temp.copy()
    g = temp.copy()
    b = temp.copy()
    for l in range(0, 21):
        r[temp == l] = cmap[l][0]
        g[temp == l] = cmap[l][1]
        b[temp == l] = cmap[l][2]

    rgb = np.zeros((temp.shape[0], temp.shape[1], 3)).astype(np.uint8)
    rgb[:, :, 0] = r
    rgb[:, :, 1] = g 
    rgb[:, :, 2] = b 
    return rgb

def inference(image_dir):
    ort_session = onnxruntime.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
    image = preprocess_image(image_dir)
    outputs = ort_session.run(
        None,
        {"pixel_values": image})
    outputs = outputs[0].squeeze().argmax(axis=0)
    cmap = decode_segmap(outputs)
    return cmap

inputs_images = gr.components.Image(type="filepath", label="Input Image")
outputs_images = gr.components.Image(type="numpy", label="Output Image")

app = gr.Interface(
    fn=inference,
    inputs=inputs_images,
    outputs=outputs_images,
    title="Semantic segmentation",
    cache_examples=False,)

app.launch()