inference-models 0.18.3__py3-none-any.whl

inference_models/models/rfdetr/backbone_builder.py

@@ -0,0 +1,101 @@
+from typing import Callable
+
+import torch
+from torch import nn
+
+from inference_models.models.rfdetr.misc import NestedTensor
+from inference_models.models.rfdetr.position_encoding import build_position_encoding
+from inference_models.models.rfdetr.rfdetr_backbone_pytorch import Backbone
+
+
+class Joiner(nn.Sequential):
+    def __init__(self, backbone, position_embedding):
+        super().__init__(backbone, position_embedding)
+        self._export = False
+
+    def forward(self, tensor_list: NestedTensor):
+        """ """
+        x = self[0](tensor_list)
+        pos = []
+        for x_ in x:
+            pos.append(self[1](x_, align_dim_orders=False).to(x_.tensors.dtype))
+        return x, pos
+
+    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_export(self, inputs: torch.Tensor):
+        feats, masks = self[0](inputs)
+        poss = []
+        for feat, mask in zip(feats, masks):
+            poss.append(self[1](mask, align_dim_orders=False).to(feat.dtype))
+        return feats, None, poss
+
+
+def build_backbone(
+    encoder,
+    vit_encoder_num_layers,
+    pretrained_encoder,
+    window_block_indexes,
+    drop_path,
+    out_channels,
+    out_feature_indexes,
+    projector_scale,
+    use_cls_token,
+    hidden_dim,
+    position_embedding,
+    freeze_encoder,
+    layer_norm,
+    target_shape,
+    rms_norm,
+    backbone_lora,
+    force_no_pretrain,
+    gradient_checkpointing,
+    load_dinov2_weights,
+    patch_size,
+    num_windows,
+    positional_encoding_size,
+):
+    """
+    Useful args:
+        - encoder: encoder name
+        - lr_encoder:
+        - dilation
+        - use_checkpoint: for swin only for now
+
+    """
+    position_embedding = build_position_encoding(hidden_dim, position_embedding)
+
+    backbone = Backbone(
+        encoder,
+        pretrained_encoder,
+        window_block_indexes=window_block_indexes,
+        drop_path=drop_path,
+        out_channels=out_channels,
+        out_feature_indexes=out_feature_indexes,
+        projector_scale=projector_scale,
+        use_cls_token=use_cls_token,
+        layer_norm=layer_norm,
+        freeze_encoder=freeze_encoder,
+        target_shape=target_shape,
+        rms_norm=rms_norm,
+        backbone_lora=backbone_lora,
+        gradient_checkpointing=gradient_checkpointing,
+        load_dinov2_weights=load_dinov2_weights,
+        patch_size=patch_size,
+        num_windows=num_windows,
+        positional_encoding_size=positional_encoding_size,
+    )
+
+    model = Joiner(backbone, position_embedding)
+    return model

inference_models/models/rfdetr/class_remapping.py

@@ -0,0 +1,41 @@
+from collections import namedtuple
+from typing import List, Tuple, Union
+
+import torch
+
+from inference_models.models.common.roboflow.model_packages import ClassNameRemoval
+
+ClassesReMapping = namedtuple(
+    "ClassesReMapping", ["remaining_class_ids", "class_mapping"]
+)
+
+
+def prepare_class_remapping(
+    class_names: List[str],
+    class_names_operations: List[
+        Union[ClassNameRemoval]
+    ],  # be ready for different elements of union type
+    device: torch.device,
+) -> Tuple[List[str], ClassesReMapping]:
+    removed_classes = {
+        o.class_name for o in class_names_operations if isinstance(o, ClassNameRemoval)
+    }
+    removed_class_ids = set()
+    remaining_class_ids = []
+    result_classes = []
+    class_mapping = []
+    for class_id, class_name in enumerate(class_names):
+        if class_name in removed_classes:
+            removed_class_ids.add(class_id)
+            class_mapping.append(-1)
+            continue
+        remaining_class_ids.append(class_id)
+        class_mapping.append(class_id - len(removed_class_ids))
+        result_classes.append(class_name)
+    classes_re_mapping = ClassesReMapping(
+        remaining_class_ids=torch.tensor(
+            remaining_class_ids, dtype=torch.int64, device=device
+        ),
+        class_mapping=torch.tensor(class_mapping, dtype=torch.int64, device=device),
+    )
+    return result_classes, classes_re_mapping

inference_models/models/rfdetr/common.py

@@ -0,0 +1,115 @@
+from typing import List, Optional, Tuple
+
+import torch
+from torchvision.transforms import functional
+
+from inference_models import InstanceDetections
+from inference_models.entities import ImageDimensions
+from inference_models.errors import CorruptedModelPackageError
+from inference_models.models.common.roboflow.model_packages import (
+    PreProcessingMetadata,
+    StaticCropOffset,
+)
+from inference_models.models.common.roboflow.post_processing import (
+    align_instance_segmentation_results,
+)
+from inference_models.models.rfdetr.class_remapping import ClassesReMapping
+from inference_models.utils.file_system import read_json
+
+
+def parse_model_type(config_path: str) -> str:
+    try:
+        parsed_config = read_json(path=config_path)
+        if not isinstance(parsed_config, dict):
+            raise ValueError(
+                f"decoded value is {type(parsed_config)}, but dictionary expected"
+            )
+        if "model_type" not in parsed_config or not isinstance(
+            parsed_config["model_type"], str
+        ):
+            raise ValueError(
+                "could not find required entries in config - either "
+                "'model_type' field is missing or not a string"
+            )
+        return parsed_config["model_type"]
+    except (IOError, OSError, ValueError) as error:
+        raise CorruptedModelPackageError(
+            message=f"Model type config file is malformed: "
+            f"{error}. In case that the package is "
+            f"hosted on the Roboflow platform - contact support. If you created model package manually, please "
+            f"verify its consistency in docs.",
+            help_url="https://todo",
+        ) from error
+
+
+def post_process_instance_segmentation_results(
+    bboxes: torch.Tensor,
+    logits: torch.Tensor,
+    masks: torch.Tensor,
+    pre_processing_meta: List[PreProcessingMetadata],
+    threshold: float,
+    classes_re_mapping: Optional[ClassesReMapping],
+) -> List[InstanceDetections]:
+    logits_sigmoid = torch.nn.functional.sigmoid(logits)
+    results = []
+    device = bboxes.device
+    for image_bboxes, image_logits, image_masks, image_meta in zip(
+        bboxes, logits_sigmoid, masks, 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]
+        selected_masks = image_masks[confidence_mask]
+        confidence, sorted_indices = torch.sort(confidence, descending=True)
+        top_classes = top_classes[sorted_indices]
+        selected_boxes = selected_boxes[sorted_indices]
+        selected_masks = selected_masks[sorted_indices]
+        if classes_re_mapping is not None:
+            remapping_mask = torch.isin(
+                top_classes, classes_re_mapping.remaining_class_ids
+            )
+            top_classes = 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=device,
+        )
+        padding = (
+            image_meta.pad_left,
+            image_meta.pad_top,
+            image_meta.pad_right,
+            image_meta.pad_bottom,
+        )
+        selected_boxes_xyxy = selected_boxes_xyxy_pct * inference_size_hwhw
+        aligned_boxes, aligned_masks = align_instance_segmentation_results(
+            image_bboxes=selected_boxes_xyxy,
+            masks=selected_masks,
+            padding=padding,
+            scale_height=image_meta.scale_height,
+            scale_width=image_meta.scale_width,
+            original_size=image_meta.original_size,
+            size_after_pre_processing=image_meta.size_after_pre_processing,
+            inference_size=image_meta.inference_size,
+            static_crop_offset=image_meta.static_crop_offset,
+        )
+        detections = InstanceDetections(
+            xyxy=aligned_boxes.round().int(),
+            confidence=confidence,
+            class_id=top_classes.int(),
+            mask=aligned_masks,
+        )
+        results.append(detections)
+    return results

inference_models/models/rfdetr/default_labels.py

@@ -0,0 +1,108 @@
+from copy import deepcopy
+from typing import List
+
+from inference_models.errors import ModelLoadingError
+
+COCO_LABELS = [
+    "background",
+    "person",
+    "bicycle",
+    "car",
+    "motorcycle",
+    "airplane",
+    "bus",
+    "train",
+    "truck",
+    "boat",
+    "traffic light",
+    "fire hydrant",
+    "background",
+    "stop sign",
+    "parking meter",
+    "bench",
+    "bird",
+    "cat",
+    "dog",
+    "horse",
+    "sheep",
+    "cow",
+    "elephant",
+    "bear",
+    "zebra",
+    "giraffe",
+    "background",
+    "backpack",
+    "umbrella",
+    "background",
+    "background",
+    "handbag",
+    "tie",
+    "suitcase",
+    "frisbee",
+    "skis",
+    "snowboard",
+    "sports ball",
+    "kite",
+    "baseball bat",
+    "baseball glove",
+    "skateboard",
+    "surfboard",
+    "tennis racket",
+    "bottle",
+    "background",
+    "wine glass",
+    "cup",
+    "fork",
+    "knife",
+    "spoon",
+    "bowl",
+    "banana",
+    "apple",
+    "sandwich",
+    "orange",
+    "broccoli",
+    "carrot",
+    "hot dog",
+    "pizza",
+    "donut",
+    "cake",
+    "chair",
+    "couch",
+    "potted plant",
+    "bed",
+    "background",
+    "dining table",
+    "background",
+    "background",
+    "toilet",
+    "background",
+    "tv",
+    "laptop",
+    "mouse",
+    "remote",
+    "keyboard",
+    "cell phone",
+    "microwave",
+    "oven",
+    "toaster",
+    "sink",
+    "refrigerator",
+    "background",
+    "book",
+    "clock",
+    "vase",
+    "scissors",
+    "teddy bear",
+    "hair drier",
+    "toothbrush",
+]
+
+
+def resolve_labels(labels: str) -> List[str]:
+    if labels != "coco":
+        raise ModelLoadingError(
+            message=f"While loading RFDetr model, `labels` parameter was set to `{labels}` which is invalid. "
+            f"Supported set of labels: `coco`.",
+            help_url="https://todo",
+        )
+    return deepcopy(COCO_LABELS)