Commit 0b73611d authored by liucong's avatar liucong
Browse files

删除部分代码和文档

parent 40dd18f7
opencv-python
numpy
pillow
\ No newline at end of file
import numpy as np
from PIL import Image
import cv2
import migraphx
def Preprocess(pil_img, newW, newH):
assert newW > 0 and newH > 0, 'Scale is too small'
pil_img = pil_img.resize((newW, newH))
img_nd = np.array(pil_img)
if len(img_nd.shape) == 2:
img_nd = np.expand_dims(img_nd, axis=2)
# HWC to CHW
img_print = pil_img
img_trans = img_nd.transpose((2, 0, 1))
if img_trans.max() > 1:
img_trans = img_trans / 255
img_trans = np.expand_dims(img_trans, 0)
return img_trans, img_print
def Sigmoid(x):
return 1 / (1 + np.exp(-x))
if __name__ == '__main__':
# 加载模型
model = migraphx.parse_onnx("./model/unet_13_256.onnx")
inputName = model.get_parameter_names()
inputShape = model.get_parameter_shapes()
print("inputName:{0} \ninputShape:{1}".format(inputName, inputShape))
# 编译模型
model.compile(migraphx.get_target("gpu"), device_id=0) # device_id: 设置GPU设备,默认为0号设备
# 图像预处理
img = Image.open("./car.jpeg")
img, imPrint = Preprocess(img, 256, 256)
input_img = np.zeros((1,3,256,256),dtype='float32')
np.lib.stride_tricks.as_strided(input_img, shape=img.shape, strides=input_img.strides)[:] = img
# 模型推理
mask = model.run({'inputs':input_img})
output_mask = np.array(mask[0])[0] # 获取推理结果,shape为(1,256,256)
probs = Sigmoid(output_mask) # 计算sigmoid值
# 0/1像素值,当大于0.996时,值为255,小于等于0.996时,值为0
output_mask[probs > 0.996] = 255
output_mask[probs <= 0.996] = 0
output = output_mask.astype(np.uint8)[0] # 将浮点型转换为uint8整型,shape为(256,256)
cv2.imwrite("output.jpg", output) # 保存图像分割结果
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment