Security_AI_integrated/inference/engine.py

import os

os.environ["TENSORRT_DISABLE_MYELIN"] = "1"

import time
from typing import Any, Dict, List, Optional, Tuple

import cv2
import numpy as np
import tensorrt as trt
import torch
import onnxruntime as ort
from ultralytics import YOLO
from ultralytics.engine.results import Results, Boxes as UltralyticsBoxes

from config import get_config


class ONNXEngine:
    def __init__(self, onnx_path: Optional[str] = None, device: int = 0):
        config = get_config()
        self.onnx_path = onnx_path or config.model.onnx_path
        self.device = device
        self.imgsz = tuple(config.model.imgsz)
        self.conf_thresh = config.model.conf_threshold
        self.iou_thresh = config.model.iou_threshold

        self.session = None
        self.input_names = None
        self.output_names = None
        self.load_model()

    def load_model(self):
        if not os.path.exists(self.onnx_path):
            raise FileNotFoundError(f"ONNX模型文件不存在: {self.onnx_path}")

        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider'] if self.device >= 0 else ['CPUExecutionProvider']
        self.session = ort.InferenceSession(self.onnx_path, providers=providers)

        self.input_names = [inp.name for inp in self.session.get_inputs()]
        self.output_names = [out.name for out in self.session.get_outputs()]

    def preprocess(self, frame: np.ndarray) -> np.ndarray:
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, self.imgsz)

        img = img.transpose(2, 0, 1).astype(np.float32) / 255.0

        return img

    def postprocess(self, output: np.ndarray, orig_img: np.ndarray) -> List[Results]:
        c, n = output.shape
        output = output.T

        boxes = output[:, :4]
        scores = output[:, 4]
        classes = output[:, 5:].argmax(axis=1) if output.shape[1] > 5 else np.zeros(len(output), dtype=np.int32)

        mask = scores > self.conf_thresh
        boxes = boxes[mask]
        scores = scores[mask]
        classes = classes[mask]

        if len(boxes) == 0:
            return [Results(orig_img=orig_img, path="", names={0: "person"})]

        indices = cv2.dnn.NMSBoxes(
            boxes.tolist(),
            scores.tolist(),
            self.conf_thresh,
            self.iou_thresh,
        )

        orig_h, orig_w = orig_img.shape[:2]
        scale_x, scale_y = orig_w / self.imgsz[1], orig_h / self.imgsz[0]

        filtered_boxes = []
        for idx in indices:
            if idx >= len(boxes):
                continue
            box = boxes[idx]
            x1, y1, x2, y2 = box
            w, h = x2 - x1, y2 - y1
            filtered_boxes.append([
                int(x1 * scale_x),
                int(y1 * scale_y),
                int(w * scale_x),
                int(h * scale_y),
                float(scores[idx]),
                int(classes[idx])
            ])

        from ultralytics.engine.results import Boxes as BoxesObj
        if filtered_boxes:
            box_tensor = torch.tensor(filtered_boxes)
            boxes_obj = BoxesObj(
                box_tensor,
                orig_shape=(orig_h, orig_w)
            )
            result = Results(
                orig_img=orig_img,
                path="",
                names={0: "person"},
                boxes=boxes_obj
            )
            return [result]

        return [Results(orig_img=orig_img, path="", names={0: "person"})]

    def inference(self, images: List[np.ndarray]) -> List[Results]:
        if not images:
            return []

        batch_imgs = []
        for frame in images:
            img = self.preprocess(frame)
            batch_imgs.append(img)

        batch = np.stack(batch_imgs, axis=0)

        inputs = {self.input_names[0]: batch}
        outputs = self.session.run(self.output_names, inputs)

        results = []
        output = outputs[0]
        if output.shape[0] == 1:
            result = self.postprocess(output[0], images[0])
            results.extend(result)
        else:
            for i in range(output.shape[0]):
                result = self.postprocess(output[i], images[i])
                results.extend(result)

        return results

    def inference_single(self, frame: np.ndarray) -> List[Results]:
        return self.inference([frame])

    def warmup(self, num_warmup: int = 10):
        dummy_frame = np.zeros((640, 640, 3), dtype=np.uint8)
        for _ in range(num_warmup):
            self.inference_single(dummy_frame)

    def __del__(self):
        if self.session:
            try:
                self.session.end_profiling()
            except Exception:
                pass


class TensorRTEngine:
    def __init__(self, engine_path: Optional[str] = None, device: int = 0):
        config = get_config()
        self.engine_path = engine_path or config.model.engine_path
        self.device = device
        self.imgsz = tuple(config.model.imgsz)
        self.conf_thresh = config.model.conf_threshold
        self.iou_thresh = config.model.iou_threshold
        self.half = config.model.half

        self.logger = trt.Logger(trt.Logger.INFO)
        self.engine = None
        self.context = None
        self.stream = torch.cuda.Stream(device=self.device)
        self.input_buffer = None
        self.output_buffers = []
        self.input_name = None
        self.output_name = None

        self._load_engine()

    def _load_engine(self):
        if not os.path.exists(self.engine_path):
            raise FileNotFoundError(f"TensorRT引擎文件不存在: {self.engine_path}")

        with open(self.engine_path, "rb") as f:
            serialized_engine = f.read()

        runtime = trt.Runtime(self.logger)
        self.engine = runtime.deserialize_cuda_engine(serialized_engine)

        self.context = self.engine.create_execution_context()

        self.stream = torch.cuda.Stream(device=self.device)
        self.batch_size = 1

        for i in range(self.engine.num_io_tensors):
            name = self.engine.get_tensor_name(i)
            dtype = self.engine.get_tensor_dtype(name)
            shape = list(self.engine.get_tensor_shape(name))

            if self.engine.get_tensor_mode(name) == trt.TensorIOMode.INPUT:
                if -1 in shape:
                    shape = [self.batch_size if d == -1 else d for d in shape]
                if dtype == trt.float16:
                    buffer = torch.zeros(shape, dtype=torch.float16, device=self.device)
                else:
                    buffer = torch.zeros(shape, dtype=torch.float32, device=self.device)
                self.input_buffer = buffer
                self.input_name = name
            else:
                if -1 in shape:
                    shape = [self.batch_size if d == -1 else d for d in shape]
                if dtype == trt.float16:
                    buffer = torch.zeros(shape, dtype=torch.float16, device=self.device)
                else:
                    buffer = torch.zeros(shape, dtype=torch.float32, device=self.device)
                self.output_buffers.append(buffer)
                if self.output_name is None:
                    self.output_name = name

            self.context.set_tensor_address(name, buffer.data_ptr())

        stream_handle = torch.cuda.current_stream(self.device).cuda_stream
        self.context.set_optimization_profile_async(0, stream_handle)

        self.batch_size = 1

    def preprocess(self, frame: np.ndarray) -> torch.Tensor:
        img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, self.imgsz)

        img = img.transpose(2, 0, 1).astype(np.float32) / 255.0

        if self.half:
            img = img.astype(np.float16)

        tensor = torch.from_numpy(img).unsqueeze(0).to(self.device)

        return tensor

    def inference(self, images: List[np.ndarray]) -> List[Results]:
        batch_size = len(images)
        if batch_size == 0:
            return []

        input_tensor = self.preprocess(images[0])

        if batch_size > 1:
            for i in range(1, batch_size):
                input_tensor = torch.cat(
                    [input_tensor, self.preprocess(images[i])], dim=0
                )

        self.context.set_tensor_address(
            self.input_name, input_tensor.contiguous().data_ptr()
        )

        input_shape = list(input_tensor.shape)
        self.context.set_input_shape(self.input_name, input_shape)

        torch.cuda.synchronize(self.stream)
        self.context.execute_async_v3(self.stream.cuda_stream)
        torch.cuda.synchronize(self.stream)

        results = []
        for i in range(batch_size):
            pred = self.output_buffers[0][i].cpu().numpy()
            pred = pred.T  # 转置: (8400, 84)
            boxes = pred[:, :4]
            scores = pred[:, 4]
            classes = pred[:, 5].astype(np.int32)

            mask = scores > self.conf_thresh
            boxes = boxes[mask]
            scores = scores[mask]
            classes = classes[mask]

            indices = cv2.dnn.NMSBoxes(
                boxes.tolist(),
                scores.tolist(),
                self.conf_thresh,
                self.iou_thresh,
            )

            if len(indices) > 0:
                for idx in indices:
                    box = boxes[idx]
                    x1, y1, x2, y2 = box
                    w, h = x2 - x1, y2 - y1
                    conf = scores[idx]
                    cls = classes[idx]

                    orig_h, orig_w = images[i].shape[:2]
                    scale_x, scale_y = orig_w / self.imgsz[1], orig_h / self.imgsz[0]
                    box_orig = [
                        int(x1 * scale_x),
                        int(y1 * scale_y),
                        int(w * scale_x),
                        int(h * scale_y),
                    ]

                    result = Results(
                        orig_img=images[i],
                        path="",
                        names={0: "person"},
                        boxes=UltralyticsBoxes(
                            torch.tensor([box_orig + [conf, cls]]),
                            orig_shape=(orig_h, orig_w),
                        ),
                    )
                    results.append(result)

        return results

    def inference_single(self, frame: np.ndarray) -> List[Results]:
        return self.inference([frame])

    def warmup(self, num_warmup: int = 10):
        dummy_frame = np.zeros((480, 640, 3), dtype=np.uint8)
        for _ in range(num_warmup):
            self.inference_single(dummy_frame)

    def __del__(self):
        if self.context:
            try:
                self.context.synchronize()
            except Exception:
                pass
        if self.stream:
            try:
                self.stream.synchronize()
            except Exception:
                pass


class Boxes:
    def __init__(
        self,
        data: torch.Tensor,
        orig_shape: Tuple[int, int],
        is_track: bool = False,
    ):
        self.data = data
        self.orig_shape = orig_shape
        self.is_track = is_track

    @property
    def xyxy(self):
        if self.is_track:
            return self.data[:, :4]
        return self.data[:, :4]

    @property
    def conf(self):
        if self.is_track:
            return self.data[:, 4]
        return self.data[:, 4]

    @property
    def cls(self):
        if self.is_track:
            return self.data[:, 5]
        return self.data[:, 5]


def _check_pt_file_valid(pt_path: str) -> bool:
    try:
        with open(pt_path, 'rb') as f:
            header = f.read(10)
        return len(header) == 10
    except Exception:
        return False


class YOLOEngine:
    def __init__(
        self,
        model_path: Optional[str] = None,
        device: int = 0,
        use_trt: bool = True,
    ):
        self.use_trt = False
        self.onnx_engine = None
        self.trt_engine = None
        self.device = device
        config = get_config()

        if use_trt:
            try:
                self.trt_engine = TensorRTEngine(device=device)
                self.trt_engine.warmup()
                self.use_trt = True
                print("TensorRT引擎加载成功")
                return
            except Exception as e:
                print(f"TensorRT加载失败: {e}")

        try:
            onnx_path = config.model.onnx_path
            if os.path.exists(onnx_path):
                self.onnx_engine = ONNXEngine(device=device)
                self.onnx_engine.warmup()
                print("ONNX引擎加载成功")
                return
            else:
                print(f"ONNX模型不存在: {onnx_path}")
        except Exception as e:
            print(f"ONNX加载失败: {e}")

        try:
            pt_path = model_path or config.model.pt_model_path
            if os.path.exists(pt_path) and _check_pt_file_valid(pt_path):
                self.model = YOLO(pt_path)
                self.model.to(device)
                print(f"PyTorch模型加载成功: {pt_path}")
            else:
                raise FileNotFoundError(f"PT文件无效或不存在: {pt_path}")
        except Exception as e:
            print(f"PyTorch加载失败: {e}")
            raise RuntimeError("所有模型加载方式均失败")

    def __call__(self, frame: np.ndarray, **kwargs) -> List[Results]:
        if self.use_trt and self.trt_engine:
            try:
                return self.trt_engine.inference_single(frame)
            except Exception as e:
                print(f"TensorRT推理失败，切换到ONNX: {e}")
                self.use_trt = False
                if self.onnx_engine:
                    return self.onnx_engine.inference_single(frame)
                elif self.model:
                    return self.model(frame, imgsz=get_config().model.imgsz, **kwargs)
                else:
                    return []
        elif self.onnx_engine:
            return self.onnx_engine.inference_single(frame)
        else:
            results = self.model(frame, imgsz=get_config().model.imgsz, **kwargs)
            return results

    def __del__(self):
        if self.trt_engine:
            del self.trt_engine
        if self.onnx_engine:
            del self.onnx_engine