|
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") |
|
|
|
session = inter.createSession() |
|
|
|
input = inter.getSessionInput(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) |
|
|
|
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) |
|
|
|
input.copyFrom(tmp_input) |
|
|
|
inter.runSession(session) |
|
|
|
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() |
|
|
