optimized_trt_predictor.py

import os
import time
import cv2
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from PIL import Image
from utils.utils import cvtColor, preprocess_input, resize_image
from threading import Lock

class OptimizedTensorRTPredictor:
    _instance = None
    _lock = Lock()
    _initialized = False
    
    @classmethod
    def get_instance(cls, engine_path='weights/model_complex.trt', input_shape=[512, 512], use_fp16=True):
        """单例模式获取预测器实例，避免重复加载模型"""
        with cls._lock:
            if cls._instance is None:
                cls._instance = cls(engine_path, input_shape, use_fp16)
            return cls._instance
    
    def __init__(self, engine_path='weights/model_complex.trt', input_shape=[512, 512], use_fp16=True):
        self.input_shape = input_shape
        self.engine_path = engine_path
        self.colors = [(0, 0, 0), (128, 0, 0), (0, 128, 0), (128, 128, 0)]
        self.num_classes = 4
        self.use_fp16 = use_fp16
        self.initialize_engine()
        
        # 预热模型，减少第一次推理的延迟
        if not OptimizedTensorRTPredictor._initialized:
            self._warmup()
            OptimizedTensorRTPredictor._initialized = True

    def initialize_engine(self):
        """初始化TensorRT引擎和内存"""
        # 加载TensorRT引擎
        logger = trt.Logger(trt.Logger.WARNING)
        trt.init_libnvinfer_plugins(logger, '')
        with open(self.engine_path, 'rb') as f, trt.Runtime(logger) as runtime:
            if self.use_fp16:
                print("Note: Using FP16 precision")
            self.engine = runtime.deserialize_cuda_engine(f.read())
        
        self.context = self.engine.create_execution_context()
        
        # 分配输入输出内存
        self.inputs = []
        self.outputs = []
        self.bindings = []
        self.stream = cuda.Stream()

        num_io = self.engine.num_io_tensors
        for i in range(num_io):
            name = self.engine.get_tensor_name(i)
            shape = self.engine.get_tensor_shape(name)
            size = trt.volume(shape)
            dtype = trt.nptype(self.engine.get_tensor_dtype(name))
            host_mem = cuda.pagelocked_empty(size, dtype)
            device_mem = cuda.mem_alloc(host_mem.nbytes)
            self.bindings.append(int(device_mem))
            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                self.inputs.append({'host': host_mem, 'device': device_mem, 'name': name})
            else:
                self.outputs.append({'host': host_mem, 'device': device_mem, 'name': name})
    
    def _warmup(self):
        """预热模型，减少第一次推理的延迟"""
        print("Warming up TensorRT engine...")
        # 创建一个空白图像进行预热
        dummy_input = np.zeros((self.input_shape[0], self.input_shape[1], 3), dtype=np.uint8)
        dummy_image = Image.fromarray(dummy_input)
        # 执行一次推理预热
        self.predict(dummy_image)
        print("TensorRT engine warmed up successfully")

    def preprocess_gpu(self, image):
        """图像预处理"""
        image = cvtColor(image)
        old_img = image.copy()
        # 将PIL Image转换为numpy数组
        image = np.array(image)
        original_h = image.shape[0]
        original_w = image.shape[1]
        
        # 使用OpenCV进行resize和预处理
        image_data = cv2.resize(image, (self.input_shape[1], self.input_shape[0]))
        image_data = cv2.cvtColor(image_data, cv2.COLOR_BGR2RGB)
        image_data = np.float32(image_data) / 255.0
        image_data = np.transpose(image_data, (2, 0, 1))
        image_data = np.expand_dims(image_data, 0)
        
        return image_data, old_img, original_h, original_w

    def predict(self, image, batch_size=1):
        """单张图像预测"""
        # 图像预处理
        image_data, old_img, original_h, original_w = self.preprocess_gpu(image)

        # 使用CUDA流水线优化内存拷贝和推理
        start_time = time.time()
        
        # 异步复制输入数据到GPU
        np.copyto(self.inputs[0]['host'], image_data.ravel())
        cuda.memcpy_htod_async(self.inputs[0]['device'], self.inputs[0]['host'], self.stream)
        
        # 设置输入输出地址并执行推理
        for inp in self.inputs:
            self.context.set_tensor_address(inp['name'], inp['device'])
        for out in self.outputs:
            self.context.set_tensor_address(out['name'], out['device'])
        
        # 使用CUDA事件优化同步
        start_event = cuda.Event()
        end_event = cuda.Event()
        
        start_event.record(self.stream)
        self.context.execute_async_v3(stream_handle=self.stream.handle)
        end_event.record(self.stream)
        
        # 异步复制输出数据回CPU
        for out in self.outputs:
            cuda.memcpy_dtoh_async(out['host'], out['device'], self.stream)
            
        # 等待所有操作完成
        end_event.synchronize()
        inference_time = time.time() - start_time

        # 后处理
        pred = self.outputs[0]['host'].reshape(1, self.num_classes, self.input_shape[0], self.input_shape[1])
        pred = np.transpose(pred[0], (1, 2, 0))
        
        # 使用OpenCV进行resize
        pred = cv2.resize(pred, (original_w, original_h))
        pred = pred.argmax(axis=-1)

        # 生成分割图像
        seg_img = np.zeros((original_h, original_w, 3), dtype=np.uint8)
        for i, color in enumerate(self.colors):
            mask = pred == i
            seg_img[mask] = color
        
        # 图像混合
        result = cv2.addWeighted(np.array(old_img), 0.3, seg_img, 0.7, 0)
        image = Image.fromarray(result)

        return image, inference_time
    
    def predict_batch(self, images, batch_size=4):
        """批量图像预测"""
        if not isinstance(images, list):
            images = [images]
        
        results = []
        total_time = 0
        
        # 分批处理图像
        for i in range(0, len(images), batch_size):
            batch = images[i:i+batch_size]
            batch_results = []
            
            # 处理每个批次
            for image in batch:
                result_image, inference_time = self.predict(image)
                batch_results.append((result_image, inference_time))
                total_time += inference_time
            
            results.extend(batch_results)
        
        return results, total_time


def save_batch_results(results, output_dir="predict_results"):
    """保存批处理结果到指定目录"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    for i, (image, _) in enumerate(results):
        output_path = os.path.join(output_dir, f"result_{i}.jpg")
        image.save(output_path)
    
    print(f"保存了 {len(results)} 张结果图像到 {output_dir} 目录")


if __name__ == "__main__":
    # 初始化预测器，启用FP16精度
    predictor = OptimizedTensorRTPredictor.get_instance(use_fp16=True)
    
    # 单张图像测试
    image_path = "E:\\one\\Downloads\\VOC2007\\JPEGImages\\4771_42.jpg"
    if os.path.exists(image_path):
        image = Image.open(image_path)
        
        # 执行预测
        result_image, inference_time = predictor.predict(image)
        
        # 打印推理时间
        print(f'单张图像推理时间: {inference_time*1000:.2f}ms')
        
        # 保存结果
        result_image.save('result_trt.jpg')
    else:
        print(f"测试图像 {image_path} 不存在，请修改路径")
        
    # 批处理测试 - 如果有多张图像
    test_dir = "E:\\one\\Downloads\\VOC2007\\JPEGImages"
    if os.path.exists(test_dir):
        image_files = [os.path.join(test_dir, f) for f in os.listdir(test_dir)[:5] if f.endswith(('.jpg', '.png', '.jpeg'))]
        if image_files:
            images = [Image.open(f) for f in image_files]
            
            # 批量预测
            batch_results, total_time = predictor.predict_batch(images, batch_size=4)
            
            # 打印批处理结果
            print(f'批处理 {len(images)} 张图像总时间: {total_time*1000:.2f}ms')
            print(f'平均每张图像时间: {(total_time/len(images))*1000:.2f}ms')
            
            # 保存批处理结果
            save_batch_results(batch_results)