TensorRTBBoxDetector.py

import os
import logging
from typing import List, Tuple

import torch
import numpy as np
from ultralytics import YOLO

# Impact Pack (for SEG and SEGS helpers)
import impact.core as core
from impact.core import SEG

# Optional: TensorRT sanity check
try:
    import tensorrt as trt  # type: ignore
except Exception:
    trt = None

# Local helpers
try:
    # If this folder is a package (has __init__.py), use relative import
    from .utils_salia import (
        NODE_DIR,
        IMGSZ,
        list_local_pt_files,
        tensor_to_pil,
        make_crop_region,
        crop_image,
        crop_ndarray2,
        dilate_mask,
    )
except ImportError:
    # Fallback: direct import if utils_salia is on sys.path
    from utils_salia import (
        NODE_DIR,
        IMGSZ,
        list_local_pt_files,
        tensor_to_pil,
        make_crop_region,
        crop_image,
        crop_ndarray2,
        dilate_mask,
    )

logger = logging.getLogger(__name__)


# -------------------------------------------------------------------------
# YOLO TensorRT-based BBOX_DETECTOR implementation
# -------------------------------------------------------------------------


class TRTYOLOBBoxDetector:
    """
    BBOX_DETECTOR interface compatible with Impact Pack / FaceDetailer.

    Methods required:
      - setAux(x)
      - detect(image, threshold, dilation, crop_factor, drop_size=1, detailer_hook=None)
      - detect_combined(image, threshold, dilation)
    """

    def __init__(self, yolo_model: YOLO, device: str = "0"):
        self.bbox_model = yolo_model
        # Ultralytics accepts "0" or "cuda:0"; normalize a bit:
        if device in ("0", "cuda", "cuda:0"):
            self.device = "0"
        else:
            self.device = str(device)

        # FaceDetailer calls setAux('face'); we keep it for compatibility.
        self.aux = None

    def setAux(self, x):
        # Kept for API compatibility. You could use this
        # later to filter by specific labels/classes.
        self.aux = x

    def detect(
        self,
        image: torch.Tensor,
        threshold: float,
        dilation: int,
        crop_factor: float,
        drop_size: int = 1,
        detailer_hook=None,
    ) -> Tuple[Tuple[int, int], List[SEG]]:
        """
        Main detection method used by FaceDetailer.

        Args:
            image: ComfyUI IMAGE tensor [B, H, W, C] in 0..1
            threshold: confidence threshold
            dilation: mask dilation in pixels
            crop_factor: expansion factor for bbox when computing crop_region
            drop_size: minimum bbox width/height to keep
            detailer_hook: optional hook with post_crop_region / post_detection

        Returns:
            SEGS tuple: ( (H, W), [SEG, SEG, ...] )
        """

        if image.dim() != 4:
            raise ValueError(
                "[TRTYOLOBBoxDetector] Expected IMAGE tensor with 4 dims [B, H, W, C]."
            )

        if image.shape[0] != 1:
            logger.warning(
                "[TRTYOLOBBoxDetector] Batch > 1 detected, using only the first image for detection."
            )
            image = image[:1]

        h, w = int(image.shape[1]), int(image.shape[2])
        shape = (h, w)

        # -----------------------------------------------------------------
        # Run YOLO TensorRT via Ultralytics wrapper
        # -----------------------------------------------------------------
        pil_img = tensor_to_pil(image)  # should return a single PIL image for B=1

        # Ultralytics chooses TensorRT backend automatically when you pass an .engine
        # model to YOLO(). Here we only set device & threshold.
        pred_list = self.bbox_model(
            pil_img,
            conf=threshold,
            device=self.device,
            verbose=False,
        )

        if len(pred_list) == 0:
            return (shape, [])

        pred = pred_list[0]
        boxes = pred.boxes

        if boxes is None or boxes.xyxy is None or boxes.xyxy.shape[0] == 0:
            return (shape, [])

        xyxy = boxes.xyxy.cpu().numpy()  # [N, 4] (x1, y1, x2, y2)
        confs = boxes.conf.cpu().numpy()
        clses = boxes.cls.cpu().numpy().astype(int)
        names = pred.names  # dict: class_index -> class_name

        seg_items: List[SEG] = []

        for i in range(xyxy.shape[0]):
            x1, y1, x2, y2 = xyxy[i]
            score = float(confs[i])
            cls_id = int(clses[i])
            label = names.get(cls_id, str(cls_id))

            box_w = x2 - x1
            box_h = y2 - y1
            if box_w <= drop_size or box_h <= drop_size:
                continue

            # Clamp bbox to image bounds
            x1_i = max(int(np.floor(x1)), 0)
            y1_i = max(int(np.floor(y1)), 0)
            x2_i = min(int(np.ceil(x2)), w)
            y2_i = min(int(np.ceil(y2)), h)
            if x2_i <= x1_i or y2_i <= y1_i:
                continue

            # Rectangular mask from bbox, uint8 0..255
            mask = np.zeros((h, w), dtype=np.uint8)
            mask[y1_i:y2_i, x1_i:x2_i] = 255

            # Optional dilation
            if dilation != 0:
                mask = dilate_mask(mask, dilation)

            # Impact bbox order here is (x1, y1, x2, y2)
            item_bbox = [float(x1), float(y1), float(x2), float(y2)]

            # Compute crop region from bbox and crop_factor
            crop_region = make_crop_region(w, h, item_bbox, crop_factor)
            if detailer_hook is not None and hasattr(detailer_hook, "post_crop_region"):
                crop_region = detailer_hook.post_crop_region(
                    w, h, item_bbox, crop_region
                )

            # Crop image + mask
            cropped_image = crop_image(image, crop_region)  # torch [1, h', w', C]
            cropped_mask = crop_ndarray2(mask, crop_region)  # np.uint8 [h', w']

            seg = SEG(
                cropped_image,
                cropped_mask,
                score,
                crop_region,
                item_bbox,
                label,
                None,  # control_net_wrapper
            )
            seg_items.append(seg)

        segs = (shape, seg_items)

        if detailer_hook is not None and hasattr(detailer_hook, "post_detection"):
            segs = detailer_hook.post_detection(segs)

        return segs

    def detect_combined(
        self,
        image: torch.Tensor,
        threshold: float,
        dilation: int,
    ) -> torch.Tensor:
        """
        Optional helper API: returns a combined MASK of all detections.
        """
        segs = self.detect(
            image=image,
            threshold=threshold,
            dilation=dilation,
            crop_factor=1.0,
            drop_size=1,
            detailer_hook=None,
        )
        return core.segs_to_combined_mask(segs)


# -------------------------------------------------------------------------
# NODE 1: TRTYOLOEngineBuilder
# - Builds TensorRT engine from a .pt file sitting next to this .py
#   imgsz  = IMGSZ (H, W) from utils_salia
#   batch  = 1
#   half   = True  (FP16)
#   device = "0"
#   overwrite (exist_ok) = True
# -------------------------------------------------------------------------


class TRTYOLOEngineBuilder:
    @classmethod
    def INPUT_TYPES(cls):
        pt_files = list_local_pt_files()
        default_name = pt_files[0] if pt_files else "face.pt"

        return {
            "required": {
                "pt_model_name": (
                    pt_files if pt_files else ["face.pt"],
                    {
                        "default": default_name,
                        "tooltip": (
                            "Select a YOLO .pt file that lives in the SAME folder as this node file.\n"
                            "Example: 'face.pt' next to TensorRTBBoxDetector.py"
                        ),
                    },
                ),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("engine_path",)
    FUNCTION = "build"
    CATEGORY = "ImpactPack/TensorRT"

    def _check_tensorrt_available(self):
        """
        Optional: preflight check to give a clearer error message if TensorRT
        cannot initialize (instead of a raw pybind11::init() error deep inside Ultralytics).
        """
        if trt is None:
            raise RuntimeError(
                "[TRTYOLOEngineBuilder] TensorRT Python package is not available. "
                "Install it via pip (cu12 build) or use an image with TensorRT preinstalled."
            )
        try:
            logger_trt = trt.Logger(trt.Logger.ERROR)
            builder = trt.Builder(logger_trt)
            # If TensorRT has trouble with CUDA/driver, this is where it will fail.
            del builder
        except Exception as e:
            raise RuntimeError(
                "[TRTYOLOEngineBuilder] TensorRT failed to initialize. "
                "Check that your CUDA driver, CUDA runtime, and TensorRT versions match. "
                f"Original error: {e}"
            ) from e

    def build(self, pt_model_name: str):
        # Resolve .pt path relative to this .py
        pt_path = os.path.join(NODE_DIR, pt_model_name)
        if not os.path.isfile(pt_path):
            raise FileNotFoundError(
                f"[TRTYOLOEngineBuilder] .pt model not found next to this node: {pt_path}"
            )

        # Optional: sanity check TensorRT/driver before asking Ultralytics to export.
        self._check_tensorrt_available()

        logger.info(
            f"[TRTYOLOEngineBuilder] Exporting TensorRT engine from '{pt_path}' "
            f"with imgsz={IMGSZ} (H,W), batch=1, half=True, device='0', exist_ok=True"
        )

        try:
            # Ultralytics API: export TensorRT engine directly from .pt
            engine_path = YOLO(pt_path).export(
                format="engine",
                imgsz=IMGSZ,
                half=True,
                device="0",
                exist_ok=True,  # overwrite or reuse if same settings
            )
        except TypeError:
            # Fallback for older Ultralytics versions that might not recognize some args.
            engine_path = YOLO(pt_path).export(
                format="engine",
                imgsz=IMGSZ,
                half=True,
                device="0",
            )

        engine_path = str(engine_path)

        # If Ultralytics returned a relative path, treat it as relative to NODE_DIR
        if not os.path.isabs(engine_path):
            candidate = os.path.join(NODE_DIR, engine_path)
            if os.path.isfile(candidate):
                engine_path = candidate

        if not os.path.isfile(engine_path):
            raise FileNotFoundError(
                f"[TRTYOLOEngineBuilder] Export completed but engine file not found at: {engine_path}"
            )

        logger.info(f"[TRTYOLOEngineBuilder] Export done. Engine path: {engine_path}")

        return (engine_path,)


# -------------------------------------------------------------------------
# NODE 2: TRTYOLOBBoxDetectorProvider
# - Loads the TensorRT engine and wraps it as BBOX_DETECTOR
# - engine_path can be:
#     * Absolute path
#     * Relative to this Python file's folder
# -------------------------------------------------------------------------


class TRTYOLOBBoxDetectorProvider:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "engine_path": (
                    "STRING",
                    {
                        "default": "",
                        "multiline": False,
                        "tooltip": (
                            "Path to .engine file.\n"
                            "Can be an absolute path OR a path RELATIVE to this node's folder.\n"
                            "Typically, you connect this from TRTYOLOEngineBuilder."
                        ),
                    },
                ),
            }
        }

    RETURN_TYPES = ("BBOX_DETECTOR",)
    RETURN_NAMES = ("bbox_detector",)
    FUNCTION = "load"
    CATEGORY = "ImpactPack/TensorRT"

    def load(self, engine_path: str):
        if not engine_path:
            raise ValueError(
                "[TRTYOLOBBoxDetectorProvider] engine_path is empty. "
                "Connect the output from TRTYOLOEngineBuilder or type a path."
            )

        engine_path = engine_path.strip()

        # If relative, treat as relative to this file's folder
        if not os.path.isabs(engine_path):
            engine_path = os.path.join(NODE_DIR, engine_path)

        if not os.path.isfile(engine_path):
            raise FileNotFoundError(
                f"[TRTYOLOBBoxDetectorProvider] Engine file not found: {engine_path}"
            )

        logger.info(
            f"[TRTYOLOBBoxDetectorProvider] Loading YOLO TensorRT engine from '{engine_path}' on device '0'"
        )

        # Ultralytics will detect it's a TensorRT engine and use TRT backend internally.
        yolo_model = YOLO(engine_path)
        detector = TRTYOLOBBoxDetector(yolo_model, device="0")

        return (detector,)


# -------------------------------------------------------------------------
# ComfyUI registration
# -------------------------------------------------------------------------

NODE_CLASS_MAPPINGS = {
    "TRTYOLOEngineBuilder": TRTYOLOEngineBuilder,
    "TRTYOLOBBoxDetectorProvider": TRTYOLOBBoxDetectorProvider,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "TRTYOLOEngineBuilder": "TensorRT YOLO Engine Builder (1344x768, local .pt)",
    "TRTYOLOBBoxDetectorProvider": "TensorRT YOLO BBox Detector",
}