第一部分:图像预处理
import numpy as np
import cv2
import acl
def imgProcess(imgPath):
image = cv2.imread(imgPath)
image = cv2.resize(image,(352, 352), interpolation = cv2.INTER_AREA)
yuvImg = cv2.cvtColor(image, cv2.COLOR_BGR2YUV)
y, u, v = cv2.split(yuvImg)
height, width = yuvImg.shape[:2]
u = cv2.resize(u, (width // 2, height // 2), interpolation=cv2.INTER_LINEAR)
v = cv2.resize(v, (width // 2, height // 2), interpolation=cv2.INTER_LINEAR)
uv = np.zeros((height // 2, width), dtype=yuvImg.dtype)
uv[:, 0::2] = u
uv[:, 1::2] = v
image = np.concatenate((y, uv), axis = 0)
source = ""
processData = imgProcess("source")
print(processData.shape)
第二部分:模型定义和推理
ACL_MEM_MALLOC_HUGE_FIRST = 0
ACL_MEMCPY_HOST_TO_DEVICE = 1
ACL_MEMCPY_DEVICE_TO_HOST = 2
class OmModel:
def __init__(self, modelPath):
self.deviceId = 5
ret = acl.init()
ret = acl.rt.set_device(self.deviceId)
self.modelId, ret = acl.mdl.load_from_file(modelPath)
self.modelDesc = acl.mdl.create_desc()
ret = acl.mdl.get_desc(self.modelDesc, self.modelId)
self.inputDatasets, self.inputData = self.prepareDatasets("input")
self.outputDatasets, self.outputData = self.prepareDatasets("output")
def prepareDataset(self, ioType):
if ioType == "input":
ioNum = acl.mdl.get_num_inputs(self.modelDesc)
aclMdlGetSizeByIndex = acl.mdl.get_input_size_by_index
else:
ioNum = acl.mdl.get_num_outputs(self.modelDesc)
aclMdlGetSizeByIndex = acl.mdl.get_output_size_by_index
dataset = acl.mdl.create_dataset()
datas = []
for i in range(ioNum):
bufferSize = aclMdlGetSizeByIndex(self.modelDesc, i)
buffer, ret = acl.rt.malloc(bufferSize, ACL_MEM_MALLOC_HUGE_FIRST)
dataBuffer = acl.create_data_buffer(buffer, bufferSize)
ret = acl.mdl.add_dataset_buffer(dataset, dataBuffer)
datas.append({"buffer": buffer, "data": dataBuffer, "size": bufferSize})
return dataset, datas
def forward(self, inputs):
inputNum = len(inputs)
for i in range(inputNum):
bytesData = inputs[i].tobytes()
bytesPtr = acl.util.bytes_to_ptr(bytesData)
ret = acl.rt.memcpy(self.inputData[i]["buffer"],
self.inputData[i]["size"],
bytesPtr,
len(bytesData),
ACL_MEMCPY_HOST_TO_DEVICE)
ret = acl.mdl.execute(self.modelId, self.inputDatasets, self.outputDatasets)
inferenceResults = []
for i, item in enumerate(self.outputData):
bufferHost, ret = acl.rt.malloc_host(self.outputData[i]["size"])
ret = acl.rt.memcpy(bufferHost,
self.outputData[i]["size"],
self.outputData[i]["buffer"],
self.outputData[i]["size"],
ACL_MEMCPY_DEVICE_TO_HOST
)
bytesOut = acl.util.ptr_to_bytes(bufferHost, self.outputData[i]["size"])
data = np.frombuffer(bytesOut, dtype = np.float32)
inferenceResults.append(data)
return inferenceResults
def __del__(self):
for dataset in [self.inputDatasets, self.outputDatasets]:
while dataset:
item = dataset.pop()
ret = acl.destroy_data_buffer(item["data"])
ret = acl.rt.free(item["buffer"])
ret = acl.mdl.destroy_dataset(self.inputDataset)
ret = acl.mdl.destroy_dataset(self.outputDataset)
ret = acl.destroy_desc(self.modelDesc)
ret = acl.mdl.unload(self.modelId)
ret = acl.rt.reset_device(self.deviceId)
ret = acl.finalize()
第三部分:主程序流程
pranetOmModelPath = "./PraNet-19_bs1.om"
pranetOmModel = OmModel(pranetOmModelPath)
outputRes = pranetOmModel.forward([processData])
tensorRes = outputRes[0].reshape(1, 1, 352, 352)
print(tensorRes.shape)
第四部分:后处理函数
THRESHOLD = 0.5
def decodeSegMap(labelMask, savePath):
segmentResult = labelMask[0][0]
segMap = np.zeros((segmentResult.shape[0], segmentResult.shape[1]), dtype=np.float32)
segMap[segmentResult > THRESHOLD] = 255
cv2.imwrite(savePath, segMap)
return
def enableContrastOutput(arr):
img1 = Image.open(arr[0])
img2 = Image.open(arr[1])
img2 = img2.resize(img1.size, Image.NEAREST)
toImage = Image.new('RGB', (img1.width + img2.width + 35, img2.height), 'white')
toImage.paste(img1, (0, 0))
toImage.paste(img2, (img1.width + 35, 0))
toImage.save(arr[2], quality=100)