import cv2
import numpy as np
import yaml
import random
import os
import codecs
from paddle.inference import Config
from paddle.inference import PrecisionType
from paddle.inference import create_predictor
import paddleseg.transforms as T
from paddleseg.cvlibs import manager
from paddleseg.utils import get_sys_env, logger, get_image_list
from paddleseg.utils.visualize import get_pseudo_color_map
class DeployYmlConfig:
def __init__(self, path):
yml_path = os.path.join(path, "deploy.yaml")
with codecs.open(yml_path, 'r', 'utf-8') as file:
self.dic = yaml.load(file)
self._transforms = self._load_transforms(
self.dic['Deploy']['transforms'])
self._dir = path
@property
def transforms(self):
print("transforms")
return self._transforms
@property
def model_file(self):
return os.path.join(self._dir, self.dic['Deploy']['model'])
@property
def params_file(self):
return os.path.join(self._dir, self.dic['Deploy']['params'])
def _load_transforms(self, t_list):
com = manager.TRANSFORMS
transforms = []
for t in t_list:
ctype = t.pop('type')
transforms.append(com[ctype](**t))
return T.Compose(transforms)
class Paddle_Seg:
def __init__(self, model_folder_dir, infer_img_size=224, use_gpu=False,
gpu_memory=500, use_tensorrt=False, precision_mode="fp32"):
self.model_folder_dir = model_folder_dir
self.infer_img_size = infer_img_size # 模型预测的输入图像尺寸
self.use_gpu = use_gpu # 是否使用GPU，默认False
self.gpu_memory = gpu_memory # GPU的显存，默认500
self.use_tensorrt = use_tensorrt # 是否使用TensorRT，默认False
self.precision = precision_mode # TensorRT的precision_mode为"fp16"、"fp32"、"int8"
def init(self,camera_width=640,camera_height=480):
img = np.zeros(shape=(int(camera_height), int(camera_width),3),dtype="float32")
# 从deploy.yml中读出模型相关信息
self.cfg = DeployYmlConfig(self.model_folder_dir)
# 初始化预测模型
self.predictor = self.predict_config()
self.color_list = [(random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)) for i in range(10)]
def predict_config(self):
# ——————————————模型配置、预测相关函数————————————————— #
# 根据预测部署的实际情况，设置Config
config = Config()
# 读取模型文件
config.set_prog_file(self.cfg.model_file)
config.set_params_file(self.cfg.params_file)
precision_map = {
"int8": PrecisionType.Int8,
"fp16": PrecisionType.Half,
"fp32": PrecisionType.Float32}
if self.use_gpu == True:
config.enable_use_gpu(self.gpu_memory, 0)
if self.use_tensorrt == True:
if self.precision == "int8":
use_calib_mode = True
use_static = True
if self.precision == "fp16":
use_calib_mode = False
use_static = True
else:
use_calib_mode = False
use_static = True
config.enable_tensorrt_engine(workspace_size=1 << 30, precision_mode=precision_map[self.precision],
max_batch_size=1, min_subgraph_size=50,
use_static=use_static, use_calib_mode=use_calib_mode)
print("----------------------------------------------")
print(" RUNNING CONFIG ")
print("----------------------------------------------")
print("Model input size: {}".format([self.infer_img_size, self.infer_img_size, 3])) # 0
print("Use GPU is: {}".format(config.use_gpu())) # True
print("GPU device id is: {}".format(config.gpu_device_id())) # 0
print("Init mem size is: {}".format(config.memory_pool_init_size_mb())) # 100
print("Use TensorRT: {}".format(self.use_tensorrt)) # 0
print("Precision mode: {}".format(precision_map[self.precision])) # 0
print("----------------------------------------------")
# 可以设置开启IR优化、开启内存优化
config.switch_ir_optim()
config.enable_memory_optim()
predictor = create_predictor(config)
return predictor
def preprocess(self, img):
data = np.array([self.cfg.transforms(img)[0]])
return data
def predict(self, predictor, data):
input_names = predictor.get_input_names()
input_handle = predictor.get_input_handle(input_names[0])
output_names = predictor.get_output_names()
output_handle = predictor.get_output_handle(output_names[0])
input_handle.reshape(data.shape)
input_handle.copy_from_cpu(data)
# 执行Predictor
predictor.run()
# 获取输出
results = []
results = output_handle.copy_to_cpu()
# 获取输出
return results
def infer(self, img):
data = self.preprocess(img)
result = self.predict(self.predictor, data)
return result
def draw_img(self, res, img):
res = res[0].astype("int8")
mask = cv2.merge([res, res*255, res])
img = cv2.addWeighted(img, 1, mask, 0.5, 0)
return img
def save_img(self, result, save_name='res.png'):
result_img = get_pseudo_color_map(result[0])
result_img.save(save_name)
