import streamlit as st
import numpy as np
from PIL import Image
import MNN

def load_deeplab_model():
    inter = MNN.Interpreter("saved_model/segment/seg_deeplab_v3_mnv2_pascal.mnn")
    #create the MNN session by using MNN interpreter
    session = inter.createSession()
    #get the input dict from the session
    input = inter.getSessionInput(session)
    #get the output dict from the session
    output = inter.getSessionOutput(session)
    return inter, session, input, output

inter, session, input, output = load_deeplab_model()

def create_pascal_label_colormap():
  """Creates a label colormap used in PASCAL VOC segmentation benchmark.

  Returns:
    A Colormap for visualizing segmentation results.
  """
  colormap = np.zeros((256, 3), dtype=int)
  ind = np.arange(256, dtype=int)

  for shift in reversed(range(8)):
    for channel in range(3):
      colormap[:, channel] |= ((ind >> channel) & 1) << shift
    ind >>= 3

  return colormap

def label_to_color_image(label):
  """Adds color defined by the dataset colormap to the label.

  Args:
    label: A 2D array with integer type, storing the segmentation label.

  Returns:
    result: A 2D array with floating type. The element of the array
      is the color indexed by the corresponding element in the input label
      to the PASCAL color map.

  Raises:
    ValueError: If label is not of rank 2 or its value is larger than color
      map maximum entry.
  """
  if label.ndim != 2:
    raise ValueError('Expect 2-D input label')

  colormap = create_pascal_label_colormap()

  if np.max(label) >= len(colormap):
    raise ValueError('label value too large.')

  return colormap[label]

def vis_segmentation(image, seg_map):
  """Visualizes input image, segmentation map and overlay view."""
  LABEL_NAMES = np.asarray([
    "背景", "飞机","自行车", "鸟", "船", "瓶子", "公共汽车",
    '汽车', '猫', '椅子', '牛', '餐桌', '狗', '马', '摩托车',
     "人", "盆栽", "羊", "沙发", "火车", "电视"
    ])

  FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
  FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)

  seg_image = label_to_color_image(seg_map).astype(np.uint8)
  seg_image = Image.fromarray(seg_image)
  alpha_image = image.convert("RGBA")
  seg_image = seg_image.convert("RGBA")
  masked_image = Image.blend(alpha_image, seg_image, 0.6)
  col1, col2 = st.columns(2)
  col1.image(image, "原始图")
  col2.image(masked_image, "分割图")
  
def preprocess_deeplab(image_file):
    image = Image.open(image_file)
    INPUT_SIZE = 513
    width, height = image.size
    resize_ratio = 1.0 * INPUT_SIZE / max(width, height)
    target_h = int(resize_ratio * height)
    target_w = int(resize_ratio * width)
    image = image.convert('RGB').resize((target_w, target_h),  Image.Resampling.LANCZOS)
    return image, target_h, target_w

def predict_deeplab(image_file):
    image, target_h, target_w = preprocess_deeplab(image_file)
    #resize APIs 
    inter.resizeTensor(input, (1, target_h, target_w, 3))
    inter.resizeSession(session)
    tmp_input = MNN.Tensor((1, target_h, target_w, 3), MNN.Halide_Type_Uint8, \
                           np.asarray(image).astype(np.uint8), MNN.Tensor_DimensionType_Tensorflow)
    tmp_output = MNN.Tensor((target_h, target_w), MNN.Halide_Type_Int, \
                            tuple(target_h * target_w *[1]), MNN.Tensor_DimensionType_Tensorflow)
    #copy tensors values as inputs
    input.copyFrom(tmp_input)
    #run session
    inter.runSession(session)
    #copy tensors value as outputs
    output.copyToHostTensor(tmp_output)
    seg_map = tmp_output.getData()
    seg_map = np.reshape(np.asarray(seg_map), [target_h, target_w])
    vis_segmentation(image, seg_map)


def main_loop():
    st.title("MNN segment demo")
    st.subheader("图像分割Demo，经过MNN量化后的model, Size=8MB")
    st.text("使用语义分割图片，图像大小默认处理为最大513 pixel")

    image_file = st.file_uploader("上传一个图片", ["jpg", "png", "jpeg"])
    if not image_file:
        return None
    predict_deeplab(image_file)

if __name__ == "__main__":
    main_loop()