File size: 4,178 Bytes
7463eed 18f531b 7463eed |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
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()
|