inference-models 0.18.3__py3-none-any.whl

inference_models/models/rfdetr/rfdetr_object_detection_trt.py

@@ -0,0 +1,270 @@
+import threading
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+
+from inference_models import Detections, ObjectDetectionModel
+from inference_models.configuration import DEFAULT_DEVICE
+from inference_models.entities import ColorFormat
+from inference_models.errors import (
+    CorruptedModelPackageError,
+    MissingDependencyError,
+    ModelRuntimeError,
+)
+from inference_models.models.common.cuda import (
+    use_cuda_context,
+    use_primary_cuda_context,
+)
+from inference_models.models.common.model_packages import get_model_package_contents
+from inference_models.models.common.roboflow.model_packages import (
+    InferenceConfig,
+    PreProcessingMetadata,
+    ResizeMode,
+    TRTConfig,
+    parse_class_names_file,
+    parse_inference_config,
+    parse_trt_config,
+)
+from inference_models.models.common.roboflow.post_processing import (
+    rescale_image_detections,
+)
+from inference_models.models.common.roboflow.pre_processing import (
+    pre_process_network_input,
+)
+from inference_models.models.common.trt import (
+    get_engine_inputs_and_outputs,
+    infer_from_trt_engine,
+    load_model,
+)
+from inference_models.models.rfdetr.class_remapping import (
+    ClassesReMapping,
+    prepare_class_remapping,
+)
+
+try:
+    import tensorrt as trt
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message=f"Could not import RFDetr model with TRT backend - this error means that some additional dependencies "
+        f"are not installed in the environment. If you run the `inference-models` library directly in your Python "
+        f"program, make sure the following extras of the package are installed: `trt10` - installation can only "
+        f"succeed for Linux and Windows machines with Cuda 12 installed. Jetson devices, should have TRT 10.x "
+        f"installed for all builds with Jetpack 6. "
+        f"If you see this error using Roboflow infrastructure, make sure the service you use does support the model. "
+        f"You can also contact Roboflow to get support.",
+        help_url="https://todo",
+    ) from import_error
+
+try:
+    import pycuda.driver as cuda
+except ImportError as import_error:
+    raise MissingDependencyError(
+        message="TODO", help_url="https://todo"
+    ) from import_error
+
+
+class RFDetrForObjectDetectionTRT(
+    (
+        ObjectDetectionModel[
+            torch.Tensor, PreProcessingMetadata, Tuple[torch.Tensor, torch.Tensor]
+        ]
+    )
+):
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_name_or_path: str,
+        device: torch.device = DEFAULT_DEVICE,
+        engine_host_code_allowed: bool = False,
+        **kwargs,
+    ) -> "RFDetrForObjectDetectionTRT":
+        if device.type != "cuda":
+            raise ModelRuntimeError(
+                message=f"TRT engine only runs on CUDA device - {device} device detected.",
+                help_url="https://todo",
+            )
+        model_package_content = get_model_package_contents(
+            model_package_dir=model_name_or_path,
+            elements=[
+                "class_names.txt",
+                "inference_config.json",
+                "trt_config.json",
+                "engine.plan",
+            ],
+        )
+        class_names = parse_class_names_file(
+            class_names_path=model_package_content["class_names.txt"]
+        )
+        inference_config = parse_inference_config(
+            config_path=model_package_content["inference_config.json"],
+            allowed_resize_modes={
+                ResizeMode.STRETCH_TO,
+                ResizeMode.LETTERBOX,
+                ResizeMode.CENTER_CROP,
+                ResizeMode.LETTERBOX_REFLECT_EDGES,
+            },
+        )
+        classes_re_mapping = None
+        if inference_config.class_names_operations:
+            class_names, classes_re_mapping = prepare_class_remapping(
+                class_names=class_names,
+                class_names_operations=inference_config.class_names_operations,
+                device=device,
+            )
+        trt_config = parse_trt_config(
+            config_path=model_package_content["trt_config.json"]
+        )
+        cuda.init()
+        cuda_device = cuda.Device(device.index or 0)
+        with use_primary_cuda_context(cuda_device=cuda_device) as cuda_context:
+            engine = load_model(
+                model_path=model_package_content["engine.plan"],
+                engine_host_code_allowed=engine_host_code_allowed,
+            )
+            execution_context = engine.create_execution_context()
+        inputs, outputs = get_engine_inputs_and_outputs(engine=engine)
+        if len(inputs) != 1:
+            raise CorruptedModelPackageError(
+                message=f"Implementation assume single model input, found: {len(inputs)}.",
+                help_url="https://todo",
+            )
+        if len(outputs) != 2:
+            raise CorruptedModelPackageError(
+                message=f"Implementation assume 2 model outputs, found: {len(outputs)}.",
+                help_url="https://todo",
+            )
+        if "dets" not in outputs or "labels" not in outputs:
+            raise CorruptedModelPackageError(
+                message=f"Expected model outputs to be named `output0` and `output1`, but found: {outputs}.",
+                help_url="https://todo",
+            )
+        return cls(
+            engine=engine,
+            input_name=inputs[0],
+            output_names=["dets", "labels"],
+            class_names=class_names,
+            classes_re_mapping=classes_re_mapping,
+            inference_config=inference_config,
+            trt_config=trt_config,
+            device=device,
+            cuda_context=cuda_context,
+            execution_context=execution_context,
+        )
+
+    def __init__(
+        self,
+        engine: trt.ICudaEngine,
+        input_name: str,
+        output_names: List[str],
+        class_names: List[str],
+        classes_re_mapping: Optional[ClassesReMapping],
+        inference_config: InferenceConfig,
+        trt_config: TRTConfig,
+        device: torch.device,
+        cuda_context: cuda.Context,
+        execution_context: trt.IExecutionContext,
+    ):
+        self._engine = engine
+        self._input_name = input_name
+        self._output_names = output_names
+        self._inference_config = inference_config
+        self._class_names = class_names
+        self._classes_re_mapping = classes_re_mapping
+        self._device = device
+        self._cuda_context = cuda_context
+        self._execution_context = execution_context
+        self._trt_config = trt_config
+        self._lock = threading.Lock()
+
+    @property
+    def class_names(self) -> List[str]:
+        return self._class_names
+
+    def pre_process(
+        self,
+        images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]],
+        input_color_format: Optional[ColorFormat] = None,
+        **kwargs,
+    ) -> Tuple[torch.Tensor, List[PreProcessingMetadata]]:
+        return pre_process_network_input(
+            images=images,
+            image_pre_processing=self._inference_config.image_pre_processing,
+            network_input=self._inference_config.network_input,
+            target_device=self._device,
+            input_color_format=input_color_format,
+        )
+
+    def forward(
+        self, pre_processed_images: torch.Tensor, **kwargs
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        with self._lock:
+            with use_cuda_context(context=self._cuda_context):
+                detections, labels = infer_from_trt_engine(
+                    pre_processed_images=pre_processed_images,
+                    trt_config=self._trt_config,
+                    engine=self._engine,
+                    context=self._execution_context,
+                    device=self._device,
+                    input_name=self._input_name,
+                    outputs=self._output_names,
+                )
+                return detections, labels
+
+    def post_process(
+        self,
+        model_results: Tuple[torch.Tensor, torch.Tensor],
+        pre_processing_meta: List[PreProcessingMetadata],
+        threshold: float = 0.5,
+        **kwargs,
+    ) -> List[Detections]:
+        bboxes, logits = model_results
+        logits_sigmoid = torch.nn.functional.sigmoid(logits)
+        results = []
+        for image_bboxes, image_logits, image_meta in zip(
+            bboxes, logits_sigmoid, pre_processing_meta
+        ):
+            confidence, top_classes = image_logits.max(dim=1)
+            confidence_mask = confidence > threshold
+            confidence = confidence[confidence_mask]
+            top_classes = top_classes[confidence_mask]
+            selected_boxes = image_bboxes[confidence_mask]
+            confidence, sorted_indices = torch.sort(confidence, descending=True)
+            top_classes = top_classes[sorted_indices]
+            selected_boxes = selected_boxes[sorted_indices]
+            if self._classes_re_mapping is not None:
+                remapping_mask = torch.isin(
+                    top_classes, self._classes_re_mapping.remaining_class_ids
+                )
+                top_classes = self._classes_re_mapping.class_mapping[
+                    top_classes[remapping_mask]
+                ]
+                selected_boxes = selected_boxes[remapping_mask]
+                confidence = confidence[remapping_mask]
+            cxcy = selected_boxes[:, :2]
+            wh = selected_boxes[:, 2:]
+            xy_min = cxcy - 0.5 * wh
+            xy_max = cxcy + 0.5 * wh
+            selected_boxes_xyxy_pct = torch.cat([xy_min, xy_max], dim=-1)
+            inference_size_hwhw = torch.tensor(
+                [
+                    image_meta.inference_size.height,
+                    image_meta.inference_size.width,
+                    image_meta.inference_size.height,
+                    image_meta.inference_size.width,
+                ],
+                device=self._device,
+            )
+            selected_boxes_xyxy = selected_boxes_xyxy_pct * inference_size_hwhw
+            selected_boxes_xyxy = rescale_image_detections(
+                image_detections=selected_boxes_xyxy,
+                image_metadata=image_meta,
+            )
+            detections = Detections(
+                xyxy=selected_boxes_xyxy.round().int(),
+                confidence=confidence,
+                class_id=top_classes.int(),
+            )
+            results.append(detections)
+        return results

inference_models/models/rfdetr/segmentation_head.py

@@ -0,0 +1,273 @@
+# ------------------------------------------------------------------------
+# RF-DETR
+# Copyright (c) 2025 Roboflow. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
+# ------------------------------------------------------------------------
+
+
+from typing import Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class DepthwiseConvBlock(nn.Module):
+    r"""Simplified ConvNeXt block without the MLP subnet"""
+
+    def __init__(self, dim, layer_scale_init_value=0):
+        super().__init__()
+        self.dwconv = nn.Conv2d(
+            dim, dim, kernel_size=3, padding=1, groups=dim
+        )  # depthwise conv
+        self.norm = nn.LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(
+            dim, dim
+        )  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        return x + input
+
+
+class MLPBlock(nn.Module):
+    def __init__(self, dim, layer_scale_init_value=0):
+        super().__init__()
+        self.norm_in = nn.LayerNorm(dim)
+        self.layers = nn.ModuleList(
+            [
+                nn.Linear(dim, dim * 4),
+                nn.GELU(),
+                nn.Linear(dim * 4, dim),
+            ]
+        )
+        self.gamma = (
+            nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True)
+            if layer_scale_init_value > 0
+            else None
+        )
+
+    def forward(self, x):
+        input = x
+        x = self.norm_in(x)
+        for layer in self.layers:
+            x = layer(x)
+        if self.gamma is not None:
+            x = self.gamma * x
+        return x + input
+
+
+class SegmentationHead(nn.Module):
+    def __init__(
+        self,
+        in_dim,
+        num_blocks: int,
+        bottleneck_ratio: int = 1,
+        downsample_ratio: int = 4,
+    ):
+        super().__init__()
+
+        self.downsample_ratio = downsample_ratio
+        self.interaction_dim = (
+            in_dim // bottleneck_ratio if bottleneck_ratio is not None else in_dim
+        )
+        self.blocks = nn.ModuleList(
+            [DepthwiseConvBlock(in_dim) for _ in range(num_blocks)]
+        )
+        self.spatial_features_proj = (
+            nn.Identity()
+            if bottleneck_ratio is None
+            else nn.Conv2d(in_dim, self.interaction_dim, kernel_size=1)
+        )
+
+        self.query_features_block = MLPBlock(in_dim)
+        self.query_features_proj = (
+            nn.Identity()
+            if bottleneck_ratio is None
+            else nn.Linear(in_dim, self.interaction_dim)
+        )
+
+        self.bias = nn.Parameter(torch.zeros(1), requires_grad=True)
+
+        self._export = False
+
+    def export(self):
+        self._export = True
+        self._forward_origin = self.forward
+        self.forward = self.forward_export
+        for name, m in self.named_modules():
+            if (
+                hasattr(m, "export")
+                and isinstance(m.export, Callable)
+                and hasattr(m, "_export")
+                and not m._export
+            ):
+                m.export()
+
+    def forward(
+        self,
+        spatial_features: torch.Tensor,
+        query_features: list[torch.Tensor],
+        image_size: tuple[int, int],
+        skip_blocks: bool = False,
+    ) -> list[torch.Tensor]:
+        # spatial features: (B, C, H, W)
+        # query features: [(B, N, C)] for each decoder layer
+        # output: (B, N, H*r, W*r)
+        target_size = (
+            image_size[0] // self.downsample_ratio,
+            image_size[1] // self.downsample_ratio,
+        )
+        spatial_features = F.interpolate(
+            spatial_features, size=target_size, mode="bilinear", align_corners=False
+        )
+
+        mask_logits = []
+        if not skip_blocks:
+            for block, qf in zip(self.blocks, query_features):
+                spatial_features = block(spatial_features)
+                spatial_features_proj = self.spatial_features_proj(spatial_features)
+                qf = self.query_features_proj(self.query_features_block(qf))
+                mask_logits.append(
+                    torch.einsum("bchw,bnc->bnhw", spatial_features_proj, qf)
+                    + self.bias
+                )
+        else:
+            assert (
+                len(query_features) == 1
+            ), "skip_blocks is only supported for length 1 query features"
+            qf = self.query_features_proj(self.query_features_block(query_features[0]))
+            mask_logits.append(
+                torch.einsum("bchw,bnc->bnhw", spatial_features, qf) + self.bias
+            )
+
+        return mask_logits
+
+    def forward_export(
+        self,
+        spatial_features: torch.Tensor,
+        query_features: list[torch.Tensor],
+        image_size: tuple[int, int],
+        skip_blocks: bool = False,
+    ) -> list[torch.Tensor]:
+        assert (
+            len(query_features) == 1
+        ), "at export time, segmentation head expects exactly one query feature"
+
+        target_size = (
+            image_size[0] // self.downsample_ratio,
+            image_size[1] // self.downsample_ratio,
+        )
+        spatial_features = F.interpolate(
+            spatial_features, size=target_size, mode="bilinear", align_corners=False
+        )
+
+        if not skip_blocks:
+            for block in self.blocks:
+                spatial_features = block(spatial_features)
+
+        spatial_features_proj = self.spatial_features_proj(spatial_features)
+
+        qf = self.query_features_proj(self.query_features_block(query_features[0]))
+        return [torch.einsum("bchw,bnc->bnhw", spatial_features_proj, qf) + self.bias]
+
+
+def point_sample(input, point_coords, **kwargs):
+    """
+    A wrapper around :function:`torch.nn.functional.grid_sample` to support 3D point_coords tensors.
+    Unlike :function:`torch.nn.functional.grid_sample` it assumes `point_coords` to lie inside
+    [0, 1] x [0, 1] square.
+    Args:
+        input (Tensor): A tensor of shape (N, C, H, W) that contains features map on a H x W grid.
+        point_coords (Tensor): A tensor of shape (N, P, 2) or (N, Hgrid, Wgrid, 2) that contains
+        [0, 1] x [0, 1] normalized point coordinates.
+    Returns:
+        output (Tensor): A tensor of shape (N, C, P) or (N, C, Hgrid, Wgrid) that contains
+            features for points in `point_coords`. The features are obtained via bilinear
+            interplation from `input` the same way as :function:`torch.nn.functional.grid_sample`.
+    """
+    add_dim = False
+    if point_coords.dim() == 3:
+        add_dim = True
+        point_coords = point_coords.unsqueeze(2)
+    output = F.grid_sample(input, 2.0 * point_coords - 1.0, **kwargs)
+    if add_dim:
+        output = output.squeeze(3)
+    return output
+
+
+def get_uncertain_point_coords_with_randomness(
+    coarse_logits,
+    uncertainty_func,
+    num_points,
+    oversample_ratio=3,
+    importance_sample_ratio=0.75,
+):
+    """
+    Sample points in [0, 1] x [0, 1] coordinate space based on their uncertainty. The unceratinties
+        are calculated for each point using 'uncertainty_func' function that takes point's logit
+        prediction as input.
+    See PointRend paper for details.
+    Args:
+        coarse_logits (Tensor): A tensor of shape (N, C, Hmask, Wmask) or (N, 1, Hmask, Wmask) for
+            class-specific or class-agnostic prediction.
+        uncertainty_func: A function that takes a Tensor of shape (N, C, P) or (N, 1, P) that
+            contains logit predictions for P points and returns their uncertainties as a Tensor of
+            shape (N, 1, P).
+        num_points (int): The number of points P to sample.
+        oversample_ratio (int): Oversampling parameter.
+        importance_sample_ratio (float): Ratio of points that are sampled via importnace sampling.
+    Returns:
+        point_coords (Tensor): A tensor of shape (N, P, 2) that contains the coordinates of P
+            sampled points.
+    """
+    assert oversample_ratio >= 1
+    assert importance_sample_ratio <= 1 and importance_sample_ratio >= 0
+    num_boxes = coarse_logits.shape[0]
+    num_sampled = int(num_points * oversample_ratio)
+    point_coords = torch.rand(num_boxes, num_sampled, 2, device=coarse_logits.device)
+    point_logits = point_sample(coarse_logits, point_coords, align_corners=False)
+    # It is crucial to calculate uncertainty based on the sampled prediction value for the points.
+    # Calculating uncertainties of the coarse predictions first and sampling them for points leads
+    # to incorrect results.
+    # To illustrate this: assume uncertainty_func(logits)=-abs(logits), a sampled point between
+    # two coarse predictions with -1 and 1 logits has 0 logits, and therefore 0 uncertainty value.
+    # However, if we calculate uncertainties for the coarse predictions first,
+    # both will have -1 uncertainty, and the sampled point will get -1 uncertainty.
+    point_uncertainties = uncertainty_func(point_logits)
+    num_uncertain_points = int(importance_sample_ratio * num_points)
+    num_random_points = num_points - num_uncertain_points
+    idx = torch.topk(point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1]
+    shift = num_sampled * torch.arange(
+        num_boxes, dtype=torch.long, device=coarse_logits.device
+    )
+    idx += shift[:, None]
+    point_coords = point_coords.view(-1, 2)[idx.view(-1), :].view(
+        num_boxes, num_uncertain_points, 2
+    )
+    if num_random_points > 0:
+        point_coords = torch.cat(
+            [
+                point_coords,
+                torch.rand(
+                    num_boxes, num_random_points, 2, device=coarse_logits.device
+                ),
+            ],
+            dim=1,
+        )
+    return point_coords