diff --git a/application/backend/app/execution/common/geti_config_converter.py b/application/backend/app/execution/common/geti_config_converter.py index 29f6d24777b..80310096c87 100644 --- a/application/backend/app/execution/common/geti_config_converter.py +++ b/application/backend/app/execution/common/geti_config_converter.py @@ -777,6 +777,26 @@ def convert(config: dict) -> dict: "status": ModelStatus.ACTIVE, "default": False, }, + "object-detection-edgecrafter-s": { + "recipe_path": RECIPE_PATH / "detection" / "edgecrafter_s.yaml", + "status": ModelStatus.ACTIVE, + "default": False, + }, + "object-detection-edgecrafter-m": { + "recipe_path": RECIPE_PATH / "detection" / "edgecrafter_m.yaml", + "status": ModelStatus.ACTIVE, + "default": False, + }, + "object-detection-edgecrafter-l": { + "recipe_path": RECIPE_PATH / "detection" / "edgecrafter_l.yaml", + "status": ModelStatus.ACTIVE, + "default": False, + }, + "object-detection-edgecrafter-x": { + "recipe_path": RECIPE_PATH / "detection" / "edgecrafter_x.yaml", + "status": ModelStatus.ACTIVE, + "default": False, + }, "object-detection-dinov3-detr-s": { "recipe_path": RECIPE_PATH / "detection" / "deimv2_s.yaml", "status": ModelStatus.ACTIVE, diff --git a/application/backend/app/supported_models/manifests/detection/edgecrafter_l.yaml b/application/backend/app/supported_models/manifests/detection/edgecrafter_l.yaml new file mode 100644 index 00000000000..67a9b7bd5ee --- /dev/null +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_l.yaml @@ -0,0 +1,42 @@ +id: object-detection-edgecrafter-l +name: ECDet-L + +pretrained_weights: + url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecdet_l.pth + mirror_url: https://huggingface.co/Intellindust/ECDet_L/resolve/main/model.safetensors + sha_sum: 60e48c51b33d70acf8e3a3faddbdb70b0876c64509b74cb623106cc7c31e85a1 + +description: "EdgeCrafter detection model (ECDet-L) combines an efficient ECViT backbone with a hybrid encoder and transformer decoder for real-time, high-accuracy object detection with low latency." + +stats: + gigaflops: 101.0 + trainable_parameters: 31.0 + benchmark_metrics: + coco_map_50_95: 57.0 + +capabilities: + tiling: false + +hyperparameters: + training: + learning_rate: 0.0005 + weight_decay: 0.0001 + batch_size: 4 + accumulate_grad_batches: 8 + early_stopping: + patience: 10 + scheduler: + factor: 0.5 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 640 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + deim_framework: true + random_zoom_out: + enable: true diff --git a/application/backend/app/supported_models/manifests/detection/edgecrafter_m.yaml b/application/backend/app/supported_models/manifests/detection/edgecrafter_m.yaml new file mode 100644 index 00000000000..985486d6b16 --- /dev/null +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_m.yaml @@ -0,0 +1,42 @@ +id: object-detection-edgecrafter-m +name: ECDet-M + +pretrained_weights: + url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecdet_m.pth + mirror_url: https://huggingface.co/Intellindust/ECDet_M/resolve/main/model.safetensors + sha_sum: c4cdf8bcd3b27c7903e422acd03caf733b5e1bfd664550bce43e50e2a3bbdc6e + +description: "EdgeCrafter detection model (ECDet-M) combines an efficient ECViT backbone with a hybrid encoder and transformer decoder for real-time, high-accuracy object detection with low latency." + +stats: + gigaflops: 53.0 + trainable_parameters: 18.0 + benchmark_metrics: + coco_map_50_95: 54.3 + +capabilities: + tiling: false + +hyperparameters: + training: + learning_rate: 0.0005 + weight_decay: 0.0001 + batch_size: 8 + accumulate_grad_batches: 4 + early_stopping: + patience: 10 + scheduler: + factor: 0.5 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 640 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + deim_framework: true + random_zoom_out: + enable: true diff --git a/application/backend/app/supported_models/manifests/detection/edgecrafter_s.yaml b/application/backend/app/supported_models/manifests/detection/edgecrafter_s.yaml new file mode 100644 index 00000000000..9bbbe5cf229 --- /dev/null +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_s.yaml @@ -0,0 +1,43 @@ +id: object-detection-edgecrafter-s +name: ECDet-S + +pretrained_weights: + url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecdet_s.pth + mirror_url: https://huggingface.co/Intellindust/ECDet_S/resolve/main/model.safetensors + sha_sum: 35340a45101d24a094becf8776aeaa327c9e6b37453e5e81dfd2923c134c4619 + +description: "EdgeCrafter detection model (ECDet-S) combines an efficient ECViT backbone with a hybrid encoder and transformer decoder for real-time, high-accuracy object detection with low latency." + +stats: + gigaflops: 26.0 + trainable_parameters: 10.0 + benchmark_metrics: + coco_map_50_95: 51.7 + +capabilities: + tiling: false + +hyperparameters: + training: + learning_rate: 0.0005 + weight_decay: 0.0001 + batch_size: 8 + accumulate_grad_batches: 4 + early_stopping: + patience: 10 + scheduler: + factor: 0.5 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 640 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + deim_framework: true + random_zoom_out: + enable: true + fill: 0 diff --git a/application/backend/app/supported_models/manifests/detection/edgecrafter_x.yaml b/application/backend/app/supported_models/manifests/detection/edgecrafter_x.yaml new file mode 100644 index 00000000000..2f95a59e252 --- /dev/null +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_x.yaml @@ -0,0 +1,42 @@ +id: object-detection-edgecrafter-x +name: ECDet-X + +pretrained_weights: + url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecdet_x.pth + mirror_url: https://huggingface.co/Intellindust/ECDet_X/resolve/main/model.safetensors + sha_sum: c87394be021b573fd8e25969b85a8670127120c9d0209ac7807fbadae1a5cec6 + +description: "EdgeCrafter detection model (ECDet-X) combines an efficient ECViT backbone with a hybrid encoder and transformer decoder for real-time, high-accuracy object detection with low latency." + +stats: + gigaflops: 151.0 + trainable_parameters: 49.0 + benchmark_metrics: + coco_map_50_95: 57.9 + +capabilities: + tiling: false + +hyperparameters: + training: + learning_rate: 0.0005 + weight_decay: 0.0001 + batch_size: 4 + accumulate_grad_batches: 8 + early_stopping: + patience: 10 + scheduler: + factor: 0.5 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 640 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + deim_framework: true + random_zoom_out: + enable: true diff --git a/application/backend/tests/unit/api/routers/test_model_architectures.py b/application/backend/tests/unit/api/routers/test_model_architectures.py index e73e804bca4..8f01cc6137f 100644 --- a/application/backend/tests/unit/api/routers/test_model_architectures.py +++ b/application/backend/tests/unit/api/routers/test_model_architectures.py @@ -17,7 +17,7 @@ def test_get_all_model_architectures(self, fxt_client: TestClient): data = response.json() assert "model_architectures" in data - assert len(data["model_architectures"]) == 32 + assert len(data["model_architectures"]) == 36 # Verify structure of first detection model detection_model = next( @@ -57,7 +57,7 @@ def test_get_all_model_architectures(self, fxt_client: TestClient): @pytest.mark.parametrize( "task_filter, total_models", [ - ("detection", 32), + ("detection", 36), ("instance_segmentation", 20), ("classification", 6), ], diff --git a/library/benchmark_manifest.yaml b/library/benchmark_manifest.yaml index f6800dab39d..12798d89e92 100644 --- a/library/benchmark_manifest.yaml +++ b/library/benchmark_manifest.yaml @@ -166,6 +166,18 @@ experiments: - name: yolo26_m priority: extended recipe: detection/yolo26_m.yaml + - name: edgecrafter_s + priority: extended + recipe: detection/edgecrafter_s.yaml + - name: edgecrafter_m + priority: extended + recipe: detection/edgecrafter_m.yaml + - name: edgecrafter_l + priority: extended + recipe: detection/edgecrafter_l.yaml + - name: edgecrafter_x + priority: extended + recipe: detection/edgecrafter_x.yaml - name: yolo26_l priority: extended recipe: detection/yolo26_l.yaml diff --git a/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py new file mode 100644 index 00000000000..e8af83bfe54 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py @@ -0,0 +1,427 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EdgeCrafter Mixin providing shared logic for Detection and Instance Segmentation models. + +Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). +""" + +from __future__ import annotations + +import copy +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Protocol, cast + +import torch +from torch import Tensor, nn +from torchvision import tv_tensors +from torchvision.ops import box_convert +from torchvision.tv_tensors import BoundingBoxFormat + +from getitune.backend.lightning.models.detection.backbones.ecvit import ECViTAdapter +from getitune.backend.lightning.models.detection.detectors.edgecrafter import ECDETRDetector +from getitune.backend.lightning.models.detection.heads.ec_decoder import ECTransformer +from getitune.backend.lightning.models.detection.losses.ec_loss import ECCriterion +from getitune.backend.lightning.models.detection.necks.dfine_hybrid_encoder import HybridEncoder +from getitune.backend.lightning.models.utils.utils import load_checkpoint +from getitune.data.entity.base import BatchLoss +from getitune.data.entity.sample import PredictionBatch, SampleBatch + +if TYPE_CHECKING: + from lightning.pytorch.cli import OptimizerCallable + + from getitune.backend.lightning.models.base import DataInputParams + + +class _EdgeCrafterHost(Protocol): + """Attributes/behaviour :class:`EdgeCrafterMixin` expects from its host class. + + Concrete host classes (:class:`~getitune.backend.lightning.models.detection.edgecrafter.EdgeCrafter` + and :class:`~getitune.backend.lightning.models.instance_segmentation.edgecrafter_inst.EdgeCrafterInst`) + provide these either directly (as instance/class attributes) or by inheriting them from + :class:`~getitune.backend.lightning.models.base.LightningModel`. + """ + + model: nn.Module + model_name: str + data_input_params: DataInputParams + optimizer_callable: OptimizerCallable + # NOTE: typed permissively (rather than `LRSchedulerCallable | LRSchedulerListCallable`) + # because pyrefly infers a narrower concrete type for this read-write attribute on the + # host classes, and Protocol attribute matching is invariant; `configure_optimizers` + # below already handles both the single-scheduler and list-of-schedulers cases. + scheduler_callable: Callable[..., Any] + multi_scale: bool + training: bool + + @property + def explain_mode(self) -> bool: + """Whether the host model is in explain (XAI) mode; read-only from the Mixin's perspective.""" + ... + + @property + def num_classes(self) -> int: + """Number of target classes; read-only from the Mixin's perspective.""" + ... + + # Per-variant backbone/proj_dim/LR config; shared, read-only lookup table. + _EC_MODEL_CFGS: ClassVar[dict[str, dict[str, Any]]] + + # Per-task configuration, defined as ClassVars on concrete subclasses. + _pretrained_weights: ClassVar[dict[str, str]] + _loss_weights: ClassVar[dict[str, float]] + _matcher_cost_dict: ClassVar[dict[str, int | float] | None] + _mask_downsample_ratio: ClassVar[int | None] + _backbone_key: ClassVar[str] + + +class EdgeCrafterMixin: + """Mixin providing shared EdgeCrafter functionality for detection and instance segmentation. + + Concrete sub-classes must set: + + * ``_pretrained_weights`` — ``{model_name: url}`` for checkpoint downloading. + * ``_loss_weights`` — per-task loss weight dict passed to :class:`ECCriterion`. + * ``model_name`` — one of ``"edgecrafter_{s,m,l,x}"``. + * ``data_input_params`` — :class:`DataInputParams` with ``input_size``. + * ``multi_scale`` — whether multi-scale training is enabled. + * ``num_classes`` — number of target classes. + + Concrete sub-classes may override (defaults below suit plain detection): + + * ``_matcher_cost_dict`` — Hungarian matcher cost overrides. Defaults to ``None``. + * ``_mask_downsample_ratio`` — mask head output stride ratio. Defaults to ``None`` + (no segmentation head). + * ``_backbone_key`` — key into ``_EC_MODEL_CFGS`` selecting the backbone variant. + Defaults to ``"backbone_name"``. + + The mixin overrides ``_customize_inputs``, ``_customize_outputs``, and + ``configure_optimizers`` so that both :class:`EdgeCrafter` + (detection) and :class:`EdgeCrafterInst` (instance segmentation) can share + the same core logic. The mixin itself stays agnostic to which concrete task + it is serving — everything task-specific is read from the host class's + ClassVar configuration (see :class:`_EdgeCrafterHost`). + """ + + _pretrained_weights: ClassVar[dict[str, str]] + _matcher_cost_dict: ClassVar[dict[str, int | float] | None] = None + _mask_downsample_ratio: ClassVar[int | None] = None + _backbone_key: ClassVar[str] = "backbone_name" + + # Per-variant backbone, proj_dim, and per-layer LR config. + _EC_MODEL_CFGS: ClassVar[dict[str, dict[str, Any]]] = { + "edgecrafter_s": { + "backbone_name": "ecvitt", + "seg_backbone_name": "ecseg_vitt", + "proj_dim": None, + "backbone_lr": 0.000025, + }, + "edgecrafter_m": { + "backbone_name": "ecvittplus", + "seg_backbone_name": "ecseg_vittplus", + "proj_dim": None, + "backbone_lr": 0.000025, + }, + "edgecrafter_l": { + "backbone_name": "ecvits", + "seg_backbone_name": "ecseg_vits", + "proj_dim": 256, + "backbone_lr": 0.000005, + }, + "edgecrafter_x": { + "backbone_name": "ecvitsplus", + "seg_backbone_name": "ecseg_vitsplus", + "proj_dim": 256, + "backbone_lr": 0.0000025, + }, + } + + def _build_ec_model( + self: _EdgeCrafterHost, + num_classes: int, + *, + backbone_lr: float | None = None, + ) -> ECDETRDetector: + """Construct the full EdgeCrafter model for detection or instance segmentation. + + Steps: + 1. Build :class:`ECViTAdapter` backbone (det or seg weights variant, selected + via ``self._backbone_key``). + 2. Build :class:`HybridEncoder` neck. + 3. Build :class:`ECTransformer` decoder (with seg head when + ``self._mask_downsample_ratio`` is set). + 4. Build :class:`ECCriterion` from ``self._loss_weights`` / + ``self._matcher_cost_dict``. + 5. Wrap everything in :class:`ECDETRDetector`. + 6. Load pretrained checkpoint via :func:`load_checkpoint`. + + Args: + num_classes: Number of target classes. + backbone_lr: Optional override for the backbone learning rate. + Defaults to the per-variant value in ``_EC_MODEL_CFGS``. + + Returns: + Configured :class:`ECDETRDetector` instance. + """ + cfg = self._EC_MODEL_CFGS[self.model_name] + + if self.data_input_params.input_size is None: + msg = "input_size must not be None." + raise ValueError(msg) + input_size: tuple[int, int] = self.data_input_params.input_size + + backbone = ECViTAdapter(model_name=cfg[self._backbone_key], proj_dim=cfg["proj_dim"]) + encoder = HybridEncoder(model_name=self.model_name) + decoder = ECTransformer( + model_name=self.model_name, + num_classes=num_classes, + eval_spatial_size=input_size, + mask_downsample_ratio=self._mask_downsample_ratio, + ) + + criterion = ECCriterion( + weight_dict=self._loss_weights, + alpha=0.75, + gamma=1.5, + reg_max=32, + num_classes=num_classes, + matcher_cost_dict=self._matcher_cost_dict, + ) + + backbone_lr = backbone_lr if backbone_lr is not None else cfg["backbone_lr"] + optimizer_configuration = [ + {"params": r"^(?=.*backbone)(?!.*(?:norm|bn|bias)).*$", "lr": backbone_lr}, + {"params": r"^(?=.*backbone)(?=.*(?:norm|bn|bias)).*$", "lr": backbone_lr, "weight_decay": 0.0}, + {"params": r"^(?=.*(?:encoder|decoder))(?=.*(?:norm|bn|bias)).*$", "weight_decay": 0.0}, + ] + + model = ECDETRDetector( + backbone=backbone, + encoder=encoder, + decoder=decoder, + criterion=criterion, + num_classes=num_classes, + optimizer_configuration=optimizer_configuration, + multi_scale=self.multi_scale, + input_size=input_size[0], + ) + model.init_weights() + load_checkpoint(model, self._pretrained_weights[self.model_name], map_location="cpu") + return model + + def _customize_inputs( # pyrefly: ignore[bad-override] + self: _EdgeCrafterHost, + entity: SampleBatch, + ) -> dict[str, Any]: + """Convert getitune :class:`SampleBatch` to EdgeCrafter input format. + + Handles detection (boxes + labels) and instance segmentation + (boxes + labels + masks) depending on what the entity contains. + + Args: + entity: getitune data batch. + + Returns: + Dict with ``images`` (Tensor) and ``targets`` (list of per-image dicts). + """ + targets: list[dict[str, Any]] = [] + has_masks = entity.masks is not None + + if entity.bboxes is not None and entity.labels is not None: + iterables: tuple + if has_masks and entity.masks is not None: + iterables = (entity.bboxes, entity.labels, entity.masks) + else: + iterables = (entity.bboxes, entity.labels) + + for items in zip(*iterables): + bb, ll = items[0], items[1] + mm = items[2] if has_masks else None + + if len(bb) > 0 and getattr(bb, "canvas_size", None) is not None: + h, w = bb.canvas_size + converted = ( + box_convert(bb, in_fmt="xyxy", out_fmt="cxcywh") if bb.format == BoundingBoxFormat.XYXY else bb + ) + device = converted.device + scaled_bboxes = converted / torch.tensor([w, h, w, h], device=device, dtype=torch.float32) + else: + h, w = entity.images.shape[2:] # pyrefly: ignore[missing-attribute] + device = getattr(entity.images, "device", torch.device("cpu")) + scaled_bboxes = torch.zeros((0, 4), device=device, dtype=torch.float32) + + target: dict[str, Any] = { + "boxes": scaled_bboxes, + "labels": ll, + "size": torch.tensor([h, w], device=device), + "orig_size": torch.tensor([h, w], device=device), + } + if mm is not None: + target["masks"] = mm + + targets.append(target) + + if self.explain_mode: + return {"entity": entity} + + return { + "images": entity.images, + "targets": targets, + } + + def _customize_outputs( # pyrefly: ignore[bad-override] + self: _EdgeCrafterHost, + outputs: dict[str, Any] | tuple | list, + inputs: SampleBatch, + ) -> PredictionBatch | BatchLoss: + """Convert model outputs to getitune format. + + During training, wraps the loss dict in a :class:`BatchLoss`. + During inference, calls ``self.model.postprocess`` and formats predictions. + When masks are returned (instance segmentation), they are included in the + :class:`PredictionBatch`. + + Args: + outputs: Loss dict (training) or raw decoder output dict (inference). + inputs: Original getitune data batch. + + Returns: + :class:`BatchLoss` during training, :class:`PredictionBatch` during inference. + """ + if self.training: + if not isinstance(outputs, dict): + msg = f"Expected dict during training, got {type(outputs)}" + raise TypeError(msg) + + losses = BatchLoss() + for k, v in outputs.items(): + if isinstance(v, Tensor): + losses[k] = v + elif isinstance(v, list): + losses[k] = sum( + (_loss.mean() for _loss in v if isinstance(_loss, Tensor)), + torch.tensor(0.0), + ) # pyrefly: ignore[unsupported-operation] + # nested dicts (e.g. dn_meta) are silently skipped + return losses + + original_sizes = [img_info.ori_shape for img_info in inputs.imgs_info] # type: ignore[union-attr] + model = cast("ECDETRDetector", self.model) + result = model.postprocess(cast("dict[str, Tensor]", outputs), original_sizes) + + prediction_kwargs: dict[str, Any] = { + "images": inputs.images, + "imgs_info": inputs.imgs_info, + } + + if len(result) == 4: # detection + masks (instance seg) + scores_list, boxes_list, labels_list, masks_list = result + prediction_kwargs["scores"] = scores_list + prediction_kwargs["bboxes"] = boxes_list + prediction_kwargs["labels"] = labels_list + + formatted_masks = [] + for masks, img_info in zip(masks_list, inputs.imgs_info): # type: ignore[union-attr, arg-type] + if masks is not None and masks.numel() > 0: + formatted_masks.append(tv_tensors.Mask(masks, dtype=torch.uint8)) # type: ignore[call-overload] + else: + formatted_masks.append( + tv_tensors.Mask( # type: ignore[call-overload] + torch.zeros( + (0, img_info.img_shape[0], img_info.img_shape[1]), # type: ignore[union-attr] + dtype=torch.bool, + ), + ) + ) + prediction_kwargs["masks"] = formatted_masks + else: # detection only + scores_list, boxes_list, labels_list = result + prediction_kwargs["scores"] = scores_list + prediction_kwargs["bboxes"] = boxes_list + prediction_kwargs["labels"] = labels_list + + return PredictionBatch(**prediction_kwargs) + + def configure_optimizers( # pyrefly: ignore[bad-override] + self: _EdgeCrafterHost, + ) -> tuple[list[torch.optim.Optimizer], list[dict[str, Any]]]: + """Configure optimizer and learning-rate schedulers. + + Uses :meth:`RTDETR._get_optim_params` to build per-param-group + configurations from the regex patterns in + ``self.model.optimizer_configuration``. + + Returns: + Two lists: optimizer list and lr-scheduler config list. + """ + from getitune.backend.lightning.models.detection.rtdetr import RTDETR + + param_groups = RTDETR._get_optim_params( # noqa: SLF001 + cast("ECDETRDetector", self.model).optimizer_configuration, + self.model, + ) + optimizer = self.optimizer_callable(param_groups) + schedulers = self.scheduler_callable(optimizer) + + def _ensure_list(item: Any) -> list: # noqa: ANN401 + return item if isinstance(item, list) else [item] + + lr_scheduler_configs = [] + for scheduler in _ensure_list(schedulers): + cfg: dict[str, Any] = {"scheduler": scheduler} + if hasattr(scheduler, "interval"): + cfg["interval"] = scheduler.interval + if hasattr(scheduler, "monitor"): + cfg["monitor"] = scheduler.monitor + lr_scheduler_configs.append(cfg) + + return [optimizer], lr_scheduler_configs + + @property + def _optimization_config(self) -> dict[str, Any]: + """PTQ config for EdgeCrafter.""" + return {"model_type": "transformer"} + + # ------------------------------------------------------------------ + # Helpers used by sub-classes + # ------------------------------------------------------------------ + + @staticmethod + def _make_is_export_prediction( + _inputs: SampleBatch | Tensor, + result: dict[str, Tensor], + ) -> dict[str, Tensor]: + """Reformat deploy-mode dict for MaskRCNN-compatible ONNX export. + + Merges ``bboxes`` + ``scores`` into a ``[B, Q, 5]`` tensor, + scales bboxes from normalised to pixel coordinates, and returns + the three ONNX output tensors: ``boxes``, ``labels``, ``masks``. + + Args: + inputs: Original image batch (used for spatial dimensions). + result: Deploy-mode postprocess output dict. + + Returns: + Dict with ``boxes`` [B, Q, 5], ``labels`` [B, Q], ``masks`` [B, Q, H, W]. + """ + bboxes = result["bboxes"] # [B, Q, 4] - already pixel coords + scores = result["scores"] # [B, Q] + labels = result["labels"] # [B, Q] + masks = result.get("masks") # [B, Q, H, W] or None + + boxes_with_scores = torch.cat([bboxes, scores.unsqueeze(-1)], dim=-1) # [B, Q, 5] + out: dict[str, Tensor] = {"boxes": boxes_with_scores, "labels": labels} + if masks is not None: + out["masks"] = masks + return out + + @staticmethod + def _default_is_meta(inputs: Tensor) -> list[dict[str, Any]]: + """Build a minimal meta-info list for export tracing.""" + shape = (int(inputs.shape[2]), int(inputs.shape[3])) + meta = { + "pad_shape": shape, + "batch_input_shape": shape, + "img_shape": shape, + "scale_factor": (1.0, 1.0), + } + return [copy.copy(meta) for _ in range(inputs.shape[0])] diff --git a/library/src/getitune/backend/lightning/models/common/layers/vit_blocks.py b/library/src/getitune/backend/lightning/models/common/layers/vit_blocks.py new file mode 100644 index 00000000000..2520134cfe0 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/common/layers/vit_blocks.py @@ -0,0 +1,142 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""Shared Vision Transformer building blocks (Attention, Block). + +Unifies the near-identical ``Attention``/``Block`` implementations that previously +existed separately in ``detection/backbones/ecvit.py`` (EdgeCrafter) and +``detection/backbones/vit_tiny.py`` (DEIMv2). Each caller keeps its own MLP +implementation (``Mlp`` / ``MLP2L``) and injects it into :class:`Block` via the +``mlp_layer`` constructor parameter, so this unification does not change either +model family's MLP dropout-placement behavior. + +Modified from DINOv3 (https://github.com/facebookresearch/dinov3). +Modified from https://huggingface.co/spaces/Hila/RobustViT/blob/main/ViT/ViT_new.py +""" + +from __future__ import annotations + +from typing import Callable + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import Tensor, nn + +from getitune.backend.lightning.models.common.layers.transformer_layers import MLP2L +from getitune.backend.lightning.models.modules.drop import DropPath + +__all__ = ["Attention", "Block"] + + +def _rotate_half(x: Tensor) -> Tensor: + x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rope(x: Tensor, sin: Tensor, cos: Tensor) -> Tensor: + return (x * cos) + (_rotate_half(x) * sin) + + +class Attention(nn.Module): + """Multi-head self-attention with optional 2-D RoPE. + + Args: + dim: Input dimension. + num_heads: Number of attention heads. Defaults to 8. + qkv_bias: Whether to add bias to QKV projection. Defaults to False. + attn_drop: Attention dropout rate. Defaults to 0.0. + proj_drop: Output projection dropout rate. Defaults to 0.0. + """ + + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + ) -> None: + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + self.scale = head_dim**-0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.attn_drop = attn_drop + self.proj = nn.Linear(dim, dim) + self.proj_drop = nn.Dropout(proj_drop) + + def forward(self, x: Tensor, rope_sincos: tuple[Tensor, Tensor] | None = None) -> Tensor: + """Forward pass.""" + b, n, c = x.shape + head_dim = c // self.num_heads + qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, head_dim) + q, k, v = qkv.unbind(2) + q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) + + if rope_sincos is not None: + sin, cos = rope_sincos + # register/cls token is at index 0; apply RoPE only to patch tokens + q_cls, q_patch = q[:, :, :1, :], q[:, :, 1:, :] + k_cls, k_patch = k[:, :, :1, :], k[:, :, 1:, :] + q_patch = _apply_rope(q_patch, sin, cos) + k_patch = _apply_rope(k_patch, sin, cos) + q = torch.cat((q_cls, q_patch), dim=2) + k = torch.cat((k_cls, k_patch), dim=2) + + x = F.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop if self.training else 0.0) + x = x.transpose(1, 2).reshape(b, n, c) + return self.proj_drop(self.proj(x)) + + +class Block(nn.Module): + """ViT transformer block with pre-normalization. + + Args: + dim: Input dimension. + num_heads: Number of attention heads. + mlp_ratio: Ratio of MLP hidden dim to embedding dim. Defaults to 4.0. + qkv_bias: Whether to add bias to QKV projection. Defaults to False. + drop: Dropout rate for MLP. Defaults to 0.0. + attn_drop: Attention dropout rate. Defaults to 0.0. + drop_path: Drop path (stochastic depth) rate. Defaults to 0.0. + act_layer: Activation layer class. Defaults to nn.GELU. + norm_layer: Normalization layer class. Defaults to nn.LayerNorm. + mlp_layer: MLP implementation, constructed as + ``mlp_layer(in_features=dim, hidden_features=int(dim * mlp_ratio), + out_features=dim, act_layer=act_layer, drop=drop)``. Callers with + different MLP dropout-placement semantics (e.g. EdgeCrafter's + single-dropout ``Mlp`` vs. DEIMv2's double-dropout ``MLP2L``) should + pass their own class here rather than relying on the default. + Defaults to :class:`~getitune.backend.lightning.models.common.layers.transformer_layers.MLP2L`. + """ + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + drop: float = 0.0, + attn_drop: float = 0.0, + drop_path: float = 0.0, + act_layer: type[nn.Module] = nn.GELU, + norm_layer: type[nn.Module] | Callable[..., nn.Module] = nn.LayerNorm, + mlp_layer: Callable[..., nn.Module] = MLP2L, + ) -> None: + super().__init__() + self.norm1 = norm_layer(dim) + self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) + self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) + self.mlp = mlp_layer( + in_features=dim, + hidden_features=int(dim * mlp_ratio), + out_features=dim, + act_layer=act_layer, + drop=drop, + ) + + def forward(self, x: Tensor, rope_sincos: tuple[Tensor, Tensor] | None = None) -> Tensor: + """Forward pass.""" + x = x + self.drop_path(self.attn(self.norm1(x), rope_sincos=rope_sincos)) + return x + self.drop_path(self.mlp(self.norm2(x))) diff --git a/library/src/getitune/backend/lightning/models/common/utils/assigners/hungarian_matcher.py b/library/src/getitune/backend/lightning/models/common/utils/assigners/hungarian_matcher.py index 01d238b1a59..9f3552d1bc7 100644 --- a/library/src/getitune/backend/lightning/models/common/utils/assigners/hungarian_matcher.py +++ b/library/src/getitune/backend/lightning/models/common/utils/assigners/hungarian_matcher.py @@ -223,6 +223,11 @@ def mask_cost( out_mask = output["pred_masks"] target_mask = output["target_mask"] + if out_mask is None: + num_queries = output["pred_logits"].shape[0] + num_targets = output["target_boxes"].shape[0] + return torch.zeros(num_queries, num_targets, device=output["pred_logits"].device) + out_mask = out_mask[:, None] target_mask = target_mask[:, None] diff --git a/library/src/getitune/backend/lightning/models/detection/__init__.py b/library/src/getitune/backend/lightning/models/detection/__init__.py index 9eb5b849769..a060f96d08f 100644 --- a/library/src/getitune/backend/lightning/models/detection/__init__.py +++ b/library/src/getitune/backend/lightning/models/detection/__init__.py @@ -7,9 +7,10 @@ from .d_fine import DFine from .deim import DEIMDFine from .deimv2 import DEIMV2 +from .edgecrafter import EdgeCrafter from .rfdetr import RFDETR from .rtdetr import RTDETR from .ssd import SSD from .yolox import YOLOX -__all__ = ["ATSS", "DEIMV2", "RFDETR", "RTDETR", "SSD", "YOLOX", "DEIMDFine", "DFine"] +__all__ = ["ATSS", "DEIMV2", "RFDETR", "RTDETR", "SSD", "YOLOX", "DEIMDFine", "DFine", "EdgeCrafter"] diff --git a/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py b/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py new file mode 100644 index 00000000000..310e3867aa2 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py @@ -0,0 +1,332 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EC-ViT backbone (ViTAdapter) for EdgeCrafter models. + +Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). +Modified from DINOv3 (https://github.com/facebookresearch/dinov3). +Modified from https://huggingface.co/spaces/Hila/RobustViT/blob/main/ViT/ViT_new.py +""" + +from __future__ import annotations + +import math +from functools import partial +from typing import Callable, ClassVar + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import nn + +from getitune.backend.lightning.models.common.layers.position_embed import RopePositionEmbedding +from getitune.backend.lightning.models.common.layers.vit_blocks import Block +from getitune.backend.lightning.models.detection.necks.dfine_hybrid_encoder import ConvNormLayerFusable +from getitune.backend.lightning.models.utils.weight_init import trunc_normal_ + +__all__ = ["ECViTAdapter"] + + +# --------------------------------------------------------------------------- +# Patch embedding +# --------------------------------------------------------------------------- + + +class ConvPyramidPatchEmbed(nn.Module): + """Convolutional pyramid patch embedding (patch_size=16). + + Replaces a single strided convolution with a chain of 3 x 3 stride-2 convs + followed by a final 3 x 3 stride-2 projection. Supports Conv-BN fusion + via ``ConvNormLayerFusable``. + + Args: + embed_dim: Output embedding dimension. + patch_size: Patch size (only 16 supported). + act: Activation name forwarded to ``ConvNormLayerFusable``. + """ + + def __init__(self, embed_dim: int = 192, patch_size: int = 16, act: str = "relu") -> None: + super().__init__() + if patch_size != 16: + msg = "Only patch_size=16 is supported for ConvPyramidPatchEmbed" + raise ValueError(msg) + + num_stages = int(math.log2(patch_size)) - 1 # 3 + ratios = [2**i for i in range(num_stages, 0, -1)] # [8, 4, 2] + channels = [embed_dim // r for r in ratios] # [24, 48, 96] for embed_dim=192 + + in_ch_list = [3] + channels[:-1] + self.convs = nn.ModuleList( + [ + ConvNormLayerFusable(in_ch, out_ch, 3, 2, act=_build_act(act)) + for in_ch, out_ch in zip(in_ch_list, channels) + ] + ) + self.proj = nn.Conv2d(channels[-1], embed_dim, kernel_size=3, stride=2, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass.""" + for conv in self.convs: + x = conv(x) + return self.proj(x) + + +def _build_act(act: str) -> type[nn.Module] | None: + """Map activation string to nn.Module class.""" + _act_map: dict[str, type[nn.Module]] = { + "relu": nn.ReLU, + "silu": nn.SiLU, + "gelu": nn.GELU, + } + return _act_map.get(act.lower()) + + +# --------------------------------------------------------------------------- +# ViT components +# --------------------------------------------------------------------------- + + +class Mlp(nn.Module): + """MLP block (SiLU activation by default).""" + + def __init__( + self, + in_features: int, + hidden_features: int | None = None, + out_features: int | None = None, + act_layer: type[nn.Module] = nn.SiLU, + drop: float = 0.0, + ) -> None: + super().__init__() + hidden_features = hidden_features or in_features + out_features = out_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features) + self.act = act_layer() + self.fc2 = nn.Linear(hidden_features, out_features) + self.drop = nn.Dropout(drop) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass.""" + return self.fc2(self.drop(self.act(self.fc1(x)))) + + +class VisionTransformer(nn.Module): + """ViT backbone used by EC-ViT models. + + Args: + embed_dim: Embedding dimension. + depth: Number of transformer blocks. + num_heads: Number of attention heads. + ffn_ratio: MLP hidden-dim ratio. + qkv_bias: Whether to use bias in QKV projection. + drop_rate: Dropout rate. + attn_drop_rate: Attention dropout rate. + drop_path_rate: Stochastic depth drop path rate. + patch_size: Patch size (only 16 supported). + return_layers: Block indices whose output tokens are collected. + norm_layer: Norm layer class. + """ + + def __init__( + self, + embed_dim: int = 192, + depth: int = 12, + num_heads: int = 3, + ffn_ratio: float = 4.0, + qkv_bias: bool = True, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.0, + patch_size: int = 16, + return_layers: list[int] | None = None, + norm_layer: type[nn.Module] | Callable[..., nn.Module] | None = None, + ) -> None: + super().__init__() + if return_layers is None: + return_layers = [10, 11] + self.embed_dim = embed_dim + self.return_layers = return_layers + norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6) + + self.patch_embed = ConvPyramidPatchEmbed(embed_dim=embed_dim, patch_size=patch_size) + self.patch_size = patch_size + self.register_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + + dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] + self.blocks = nn.ModuleList( + [ + Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=ffn_ratio, + qkv_bias=qkv_bias, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path=dpr[i], + act_layer=nn.GELU, + norm_layer=norm_layer, # type: ignore[arg-type] + mlp_layer=Mlp, + ) + for i in range(depth) + ] + ) + + self.rope_embed = RopePositionEmbedding( + embed_dim=embed_dim, + num_heads=num_heads, + base=100.0, + normalize_coords="separate", + ) + self._init_weights() + + def _init_weights(self) -> None: + self.apply(self._init_vit_weights) + self.rope_embed._init_weights() # noqa: SLF001 + trunc_normal_(self.register_token, std=0.02) + + @staticmethod + def _init_vit_weights(m: nn.Module) -> None: + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.zeros_(m.bias) + elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)): + nn.init.zeros_(m.bias) # type: ignore[arg-type] + nn.init.ones_(m.weight) # type: ignore[arg-type] + + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: + """Forward pass, returns patch tokens for each ``return_layers`` block.""" + x_embed = self.patch_embed(x) + _, _, H, W = x_embed.shape # noqa: N806 + x_embed = x_embed.flatten(2).transpose(1, 2) # [B, N, C] + reg = self.register_token.expand(x_embed.shape[0], -1, -1) + x = torch.cat((reg, x_embed), dim=1) + rope = self.rope_embed(h=H, w=W) + + outs: list[torch.Tensor] = [] + for i, blk in enumerate(self.blocks): + x = blk(x, rope_sincos=rope) + if i in self.return_layers: + outs.append(x[:, 1:]) # exclude register token + return outs + + +# --------------------------------------------------------------------------- +# ViTAdapter (= ECViTAdapter) +# --------------------------------------------------------------------------- + +_BACKBONE_CONFIGS: dict[str, dict] = { + # ECDet backbones + "ecvitt": {"embed_dim": 192, "num_heads": 3, "depth": 12, "ffn_ratio": 4}, + "ecvittplus": {"embed_dim": 256, "num_heads": 4, "depth": 12, "ffn_ratio": 4}, + "ecvits": {"embed_dim": 384, "num_heads": 6, "depth": 12, "ffn_ratio": 4}, + "ecvitsplus": {"embed_dim": 384, "num_heads": 6, "depth": 12, "ffn_ratio": 6}, + # ECSeg backbones + "ecseg_vitt": {"embed_dim": 192, "num_heads": 3, "depth": 12, "ffn_ratio": 4}, + "ecseg_vittplus": {"embed_dim": 256, "num_heads": 4, "depth": 12, "ffn_ratio": 4}, + "ecseg_vits": {"embed_dim": 384, "num_heads": 6, "depth": 12, "ffn_ratio": 4}, + "ecseg_vitsplus": {"embed_dim": 384, "num_heads": 6, "depth": 12, "ffn_ratio": 6}, +} + + +class ECViTAdapter(nn.Module): + """ViT-based backbone for EdgeCrafter detection/segmentation models. + + Wraps a ``VisionTransformer`` and projects its multi-scale outputs to the + feature pyramid expected by ``HybridEncoder``. Three spatial levels are + produced by interpolating the fused ViT output at scales x2, x1, x0.5 + relative to the ViT feature-map resolution. + + Args: + model_name: One of the keys in ``_BACKBONE_CONFIGS`` (e.g. ``"ecvitt"``). + interaction_indexes: Indices of ViT blocks whose outputs are fused. + proj_dim: If not None, all levels are projected from ``embed_dim`` to + ``proj_dim`` (used for L/X variants with larger backbones). + num_levels: Number of output pyramid levels (always 3). + patch_size: Patch size of the ViT (only 16 supported). + """ + + _pretrained_urls: ClassVar[dict[str, str]] = { + "ecvitt": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecvitt.pth", + "ecvittplus": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecvittplus.pth", + "ecvits": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecvits.pth", + "ecvitsplus": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecvitsplus.pth", + "ecseg_vitt": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_vitt.pth", + "ecseg_vittplus": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_vittplus.pth", + "ecseg_vits": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_vits.pth", + "ecseg_vitsplus": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_vitsplus.pth", + } + + def __init__( + self, + model_name: str, + interaction_indexes: list[int] | None = None, + proj_dim: int | None = None, + num_levels: int = 3, + patch_size: int = 16, + ) -> None: + super().__init__() + if model_name not in _BACKBONE_CONFIGS: + msg = f"Unknown ECViT model name '{model_name}'. Available: {list(_BACKBONE_CONFIGS)}" + raise ValueError(msg) + + cfg = _BACKBONE_CONFIGS[model_name] + embed_dim: int = cfg["embed_dim"] + + if interaction_indexes is None: + interaction_indexes = [10, 11] + + self.backbone = VisionTransformer( + embed_dim=embed_dim, + depth=cfg["depth"], + num_heads=cfg["num_heads"], + ffn_ratio=cfg["ffn_ratio"], + return_layers=interaction_indexes, + patch_size=patch_size, + ) + + if num_levels != 3: + msg = "Only num_levels=3 is supported for ECViTAdapter" + raise ValueError(msg) + + self.num_levels = num_levels + self.patch_size = patch_size + self.embed_dim = embed_dim + + # When proj_dim is set: project all levels from embed_dim → proj_dim. + # When None: a single projector reduces only the coarsest level, keeping + # all levels at embed_dim (matching the S/M checkpoint key layout). + if proj_dim is not None: + self.projector = nn.ModuleList([ConvNormLayerFusable(embed_dim, proj_dim, 1, 1) for _ in range(num_levels)]) + self.out_channels = proj_dim + else: + self.projector = nn.ModuleList([ConvNormLayerFusable(embed_dim, embed_dim, 1, 1)]) + self.out_channels = embed_dim + + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: + """Return a 3-level feature pyramid from the ViT backbone.""" + H_c = x.shape[2] // self.patch_size # noqa: N806 + W_c = x.shape[3] // self.patch_size # noqa: N806 + bs = x.shape[0] + + return_layers = self.backbone(x) + fused = torch.mean(torch.stack(return_layers), dim=0) # [B, N, C] + fused = fused.transpose(1, 2).contiguous().view(bs, -1, H_c, W_c) # [B, C, H, W] + + proj_feats: list[torch.Tensor] = [] + for i in range(self.num_levels): + scale = 2 ** (1 - i) + feat = F.interpolate( + fused, + size=[int(H_c * scale), int(W_c * scale)], + mode="bilinear", + align_corners=False, + ) + proj_feats.append(feat) + + # Project only the last (coarsest) level when using a single projector + if len(self.projector) == 1: + proj_feats[-1] = self.projector[0](proj_feats[-1]) + else: + proj_feats = [layer(feat) for layer, feat in zip(self.projector, proj_feats)] + + return proj_feats diff --git a/library/src/getitune/backend/lightning/models/detection/backbones/vit_tiny.py b/library/src/getitune/backend/lightning/models/detection/backbones/vit_tiny.py index 22c221d69e3..d10f74530e6 100644 --- a/library/src/getitune/backend/lightning/models/detection/backbones/vit_tiny.py +++ b/library/src/getitune/backend/lightning/models/detection/backbones/vit_tiny.py @@ -17,7 +17,7 @@ from torch import Tensor, nn from getitune.backend.lightning.models.common.layers.position_embed import RopePositionEmbedding -from getitune.backend.lightning.models.common.layers.transformer_layers import MLP2L as MLP +from getitune.backend.lightning.models.common.layers.vit_blocks import Attention, Block # noqa: F401 from getitune.backend.lightning.models.utils.weight_init import trunc_normal_ @@ -135,122 +135,6 @@ def forward(self, x: Tensor) -> Tensor: return drop_path(x, self.drop_prob or 0.0, self.training) -class Attention(nn.Module): - """Multi-head self-attention module with optional RoPE support. - - Args: - dim: Input dimension. - num_heads: Number of attention heads. Defaults to 8. - qkv_bias: Whether to add bias to QKV projection. Defaults to False. - attn_drop: Attention dropout rate. Defaults to 0.0. - proj_drop: Output projection dropout rate. Defaults to 0.0. - """ - - def __init__( - self, - dim: int, - num_heads: int = 8, - qkv_bias: bool = False, - attn_drop: float = 0.0, - proj_drop: float = 0.0, - ) -> None: - super().__init__() - self.num_heads = num_heads - head_dim = dim // num_heads - self.scale = head_dim**-0.5 - self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) - self.attn_drop = attn_drop - self.proj = nn.Linear(dim, dim) - self.proj_drop = nn.Dropout(proj_drop) - - def forward( - self, - x: Tensor, - rope_sincos: tuple[Tensor, Tensor] | None = None, - ) -> Tensor: - """Forward pass for multi-head attention. - - Args: - x: Input tensor of shape (B, N, C). - rope_sincos: Optional tuple of (sin, cos) tensors for RoPE. - - Returns: - Output tensor of shape (B, N, C). - """ - B, N, C = x.shape # noqa: N806 - qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4) - q, k, v = qkv.unbind(0) - - if rope_sincos is not None: - sin, cos = rope_sincos - q_cls, q_patch = q[:, :, :1, :], q[:, :, 1:, :] - k_cls, k_patch = k[:, :, :1, :], k[:, :, 1:, :] - - q_patch = apply_rope(q_patch, sin, cos) - k_patch = apply_rope(k_patch, sin, cos) - - q = torch.cat((q_cls, q_patch), dim=2) - k = torch.cat((k_cls, k_patch), dim=2) - - x = torch.nn.functional.scaled_dot_product_attention(q, k, v, dropout_p=self.attn_drop) - x = x.transpose(1, 2).reshape([B, N, C]) - x = self.proj(x) - return self.proj_drop(x) - - -class Block(nn.Module): - """Transformer block with attention and MLP. - - Standard transformer encoder block with pre-normalization. - - Args: - dim: Input dimension. - num_heads: Number of attention heads. - mlp_ratio: Ratio of MLP hidden dim to embedding dim. Defaults to 4.0. - qkv_bias: Whether to add bias to QKV projection. Defaults to False. - drop: Dropout rate for MLP. Defaults to 0.0. - attn_drop: Attention dropout rate. Defaults to 0.0. - drop_path: Drop path rate. Defaults to 0.0. - act_layer: Activation layer class. Defaults to nn.GELU. - norm_layer: Normalization layer class. Defaults to nn.LayerNorm. - """ - - def __init__( - self, - dim: int, - num_heads: int, - mlp_ratio: float = 4.0, - qkv_bias: bool = False, - attn_drop: float = 0.0, - drop_path: float = 0.0, - drop: float = 0.0, - act_layer: type[nn.Module] = nn.GELU, - norm_layer: type[nn.Module] | Callable[..., nn.Module] = nn.LayerNorm, - ) -> None: - super().__init__() - self.norm1 = norm_layer(dim) - self.attn = Attention(dim, num_heads=num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) - self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() - self.norm2 = norm_layer(dim) - self.mlp = MLP( - in_features=dim, hidden_features=int(dim * mlp_ratio), out_features=dim, act_layer=act_layer, drop=drop - ) - - def forward(self, x: Tensor, rope_sincos: tuple[Tensor, Tensor] | None = None) -> Tensor: - """Forward pass through transformer block. - - Args: - x: Input tensor of shape (B, N, C). - rope_sincos: Optional tuple of (sin, cos) tensors for RoPE. - - Returns: - Output tensor of shape (B, N, C). - """ - attn_output = self.attn(self.norm1(x), rope_sincos=rope_sincos) - x = x + self.drop_path(attn_output) - return x + self.drop_path(self.mlp(self.norm2(x))) - - class VisionTransformer(nn.Module): """Vision Transformer (ViT) backbone for object detection. diff --git a/library/src/getitune/backend/lightning/models/detection/detectors/__init__.py b/library/src/getitune/backend/lightning/models/detection/detectors/__init__.py index 3ce58ef7f50..381e93e112c 100644 --- a/library/src/getitune/backend/lightning/models/detection/detectors/__init__.py +++ b/library/src/getitune/backend/lightning/models/detection/detectors/__init__.py @@ -4,7 +4,8 @@ """Base models classes implementations for detection task.""" from .detection_transformer import DETR +from .edgecrafter import ECDETRDetector from .rfdetr import RFDETRDetector from .single_stage_detector import SingleStageDetector -__all__ = ["DETR", "RFDETRDetector", "SingleStageDetector"] +__all__ = ["DETR", "ECDETRDetector", "RFDETRDetector", "SingleStageDetector"] diff --git a/library/src/getitune/backend/lightning/models/detection/detectors/edgecrafter.py b/library/src/getitune/backend/lightning/models/detection/detectors/edgecrafter.py new file mode 100644 index 00000000000..657437d45b1 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/detectors/edgecrafter.py @@ -0,0 +1,159 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EdgeCrafter DETR-style detector wrapper.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, cast + +if TYPE_CHECKING: + from getitune.backend.lightning.models.detection.heads.ec_decoder import ECTransformer + +import torch +from torch import Tensor +from torchvision.tv_tensors import BoundingBoxes + +from .detection_transformer import DETR + + +class ECDETRDetector(DETR): + """EdgeCrafter DETR-style detector wrapping :class:`ECTransformer`. + + Extends :class:`DETR` with: + + * ``_forward_features`` — passes the stride-8 backbone feature to the decoder + so that the optional segmentation head (ECSeg) has access to fine-grained + spatial context. + * ``postprocess`` — converts ``ECTransformer.postprocess``'s per-image + list-of-dicts output into the tuple format expected by the getitune + training/inference pipeline. When the model produces masks a 4-tuple + is returned; otherwise a 3-tuple. + * ``export`` — strips the ``explain_mode`` flag (not supported by + ECTransformer) and runs the inference path without teacher inputs. + + Args: + backbone: :class:`ECViTAdapter` feature extractor. + encoder: :class:`HybridEncoder` neck. + decoder: :class:`ECTransformer` head. + num_classes: Number of object classes. + criterion: Loss function (:class:`ECCriterion`). + optimizer_configuration: Per-param-group LR / WD overrides. + multi_scale: Whether to use multi-scale training. + num_top_queries: Number of top scoring queries to keep. Defaults to 300. + input_size: Training input resolution for multi-scale sampling. Defaults to 640. + """ + + def _forward_features( + self, + images: Tensor, + targets: object = None, + ) -> dict[str, Any]: + """Backbone → encoder → decoder forward with spatial feature propagation. + + The stride-8 backbone feature (``backbone_feats[0]``) is forwarded to + :class:`ECTransformer` as ``spatial_feat`` so that the optional + :class:`SegmentationHead` (ECSeg variants) can produce instance masks. + For detection-only models the decoder ignores this argument. + + Args: + images: Input image tensor [B, C, H, W]. + targets: Ground-truth target dicts (used during training for CDN denoising). + + Returns: + ECTransformer output dict. + """ + backbone_feats = self.backbone(images) + encoder_feats = self.encoder(backbone_feats) + return self.decoder(encoder_feats, cast("list[dict] | None", targets), spatial_feat=backbone_feats[0]) + + def postprocess( + self, + outputs: dict[str, Tensor], + original_sizes: list[tuple[int, int]], + deploy_mode: bool = False, + ) -> dict[str, Tensor] | tuple: + """Post-process :class:`ECTransformer` outputs to getitune format. + + Delegates actual decoding to ``self.decoder.postprocess`` which handles + box de-normalisation, top-k selection, and optional mask upsampling. + + Args: + outputs: Raw model output dict (``pred_logits``, ``pred_boxes``, + optionally ``pred_masks``). + original_sizes: Per-image ``(H, W)`` tuples for rescaling. + deploy_mode: If ``True``, returns a compact dict of stacked tensors + suitable for ONNX / OpenVINO export. + + Returns: + *deploy mode*: ``dict(bboxes, labels, scores[, masks])``. + + *inference mode*: ``(scores_list, boxes_list, labels_list)`` + or ``(scores_list, boxes_list, labels_list, masks_list)`` + when instance-segmentation masks are present. + """ + device = outputs["pred_logits"].device + sizes_tensor = torch.tensor(original_sizes, device=device, dtype=torch.float32) # [B, 2] as (H, W) + + results = cast("ECTransformer", self.decoder).postprocess(outputs, sizes_tensor, self.num_top_queries) + + has_masks = bool(results) and "masks" in results[0] + + if deploy_mode: + scores = torch.stack([r["scores"] for r in results]) + labels = torch.stack([r["labels"] for r in results]) + boxes = torch.stack([r["boxes"] for r in results]) + out: dict[str, Tensor] = {"bboxes": boxes, "labels": labels, "scores": scores} + if has_masks: + out["masks"] = torch.stack([r["masks"].float() for r in results]) + return out + + scores_list: list[Tensor] = [] + boxes_list: list[BoundingBoxes] = [] + labels_list: list[Tensor] = [] + masks_list: list[Tensor] = [] + + for res, orig_size in zip(results, original_sizes): + scores_list.append(res["scores"]) + boxes_list.append( + BoundingBoxes( # pyrefly: ignore[no-matching-overload] + res["boxes"], format="xyxy", canvas_size=orig_size + ) + ) + labels_list.append(res["labels"].long()) + if has_masks: + masks_list.append(res["masks"]) + + if has_masks: + return scores_list, boxes_list, labels_list, masks_list + return scores_list, boxes_list, labels_list + + def export( + self, + batch_inputs: Tensor, + batch_img_metas: list[dict], + explain_mode: bool = False, + ) -> dict[str, Tensor]: + """Export-mode forward pass (no teacher, no denoising, deploy postprocess). + + ``explain_mode`` is accepted for API compatibility with the base class and + the instance-segmentation base's ``forward_for_tracing``, but is not used + because ``ECTransformer`` does not implement XAI features. + + Args: + batch_inputs: Input image batch [B, C, H, W]. + batch_img_metas: List of per-image meta dicts; ``"img_shape"`` key + is used for rescaling. + explain_mode: Ignored. + + Returns: + Dict with ``bboxes``, ``labels``, ``scores`` (and ``masks`` for ECSeg). + """ + backbone_feats = self.backbone(batch_inputs) + encoder_feats = self.encoder(backbone_feats) + predictions = self.decoder(encoder_feats, spatial_feat=backbone_feats[0]) + return self.postprocess( # type: ignore[return-value] + predictions, + [meta["img_shape"] for meta in batch_img_metas], + deploy_mode=True, + ) diff --git a/library/src/getitune/backend/lightning/models/detection/edgecrafter.py b/library/src/getitune/backend/lightning/models/detection/edgecrafter.py new file mode 100644 index 00000000000..42fd22d6a64 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/edgecrafter.py @@ -0,0 +1,166 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EdgeCrafter model implementation for object detection.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, ClassVar, Literal, cast + +import torch +from torch import Tensor + +from getitune.backend.lightning.exporter.base import ModelExporter +from getitune.backend.lightning.exporter.native import LightningModelExporter +from getitune.backend.lightning.models.base import DataInputParams, DefaultOptimizerCallable, DefaultSchedulerCallable +from getitune.backend.lightning.models.common.edgecrafter_mixin import EdgeCrafterMixin +from getitune.backend.lightning.models.detection.base import LightningDetectionModel +from getitune.backend.lightning.models.detection.detectors.edgecrafter import ECDETRDetector +from getitune.config.data import TileConfig +from getitune.metrics.fmeasure import MeanAveragePrecisionFMeasureCallable + +if TYPE_CHECKING: + from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable + + from getitune.backend.lightning.schedulers import LRSchedulerListCallable + from getitune.metrics import MetricCallable + from getitune.types.label import LabelInfoTypes + + +class EdgeCrafter(EdgeCrafterMixin, LightningDetectionModel): # pyrefly: ignore[inconsistent-inheritance] + """getitune Detection model class for EdgeCrafter. + + EdgeCrafter is a family of compact ViT-based DETR models designed for edge + deployment via task-specialised distillation. Four sizes are available: + S (small), M (medium), L (large), and X (extra-large). + + Original repository: https://github.com/Intellindust-AI-Lab/EdgeCrafter + Paper: https://arxiv.org/abs/2603.18739 + + The model should be used with + :class:`~getitune.backend.lightning.callbacks.aug_scheduler.AdaptiveTrainScheduling` + and optional augmentation scheduling callbacks. + + Attributes: + _pretrained_weights: Mapping from model-name to pretrained weight URL. + input_size_multiplier: Patch-size aligned input divisor (16). + + Args: + label_info: Information about the labels. + data_input_params: Preprocessing parameters (input size, mean, std). + When ``None`` the ``_default_preprocessing_params`` dict is used. + model_name: One of ``"edgecrafter_{s,m,l,x}"``. + optimizer: Optimizer callable. Defaults to :data:`DefaultOptimizerCallable`. + scheduler: LR-scheduler callable. Defaults to :data:`DefaultSchedulerCallable`. + metric: Metric callable. Defaults to MAP F-measure. + multi_scale: Whether to use multi-scale training. Defaults to ``False``. + backbone_lr: Learning rate for the backbone. Defaults to the upstream + per-variant value in :class:`EdgeCrafterMixin`. + torch_compile: Whether to use ``torch.compile``. Defaults to ``False``. + tile_config: Tiling configuration. Defaults to disabled tiler. + """ + + _pretrained_weights: ClassVar[dict[str, str]] = { + "edgecrafter_s": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecdet_s.pth", + "edgecrafter_m": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecdet_m.pth", + "edgecrafter_l": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecdet_l.pth", + "edgecrafter_x": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecdet_x.pth", + } + + # Detection loss weights (matcher/backbone-key/mask-ratio use EdgeCrafterMixin defaults). + _loss_weights: ClassVar[dict[str, float]] = { + "loss_mal": 1.0, + "loss_bbox": 5.0, + "loss_giou": 2.0, + "loss_fgl": 0.15, + "loss_ddf": 1.5, + } + + input_size_multiplier = 16 + + def __init__( + self, + label_info: LabelInfoTypes, + data_input_params: DataInputParams | dict | None = None, + model_name: Literal[ + "edgecrafter_s", + "edgecrafter_m", + "edgecrafter_l", + "edgecrafter_x", + ] = "edgecrafter_s", + optimizer: OptimizerCallable = DefaultOptimizerCallable, + scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable, + metric: MetricCallable = MeanAveragePrecisionFMeasureCallable, + multi_scale: bool = False, + backbone_lr: float | None = None, + torch_compile: bool = False, + tile_config: TileConfig = TileConfig(enable_tiler=False), + ) -> None: + self.multi_scale = multi_scale + self.backbone_lr = backbone_lr + super().__init__( + label_info=label_info, + data_input_params=data_input_params, + model_name=model_name, + optimizer=optimizer, + scheduler=scheduler, + metric=metric, + torch_compile=torch_compile, + tile_config=tile_config, + ) + + def _create_model(self, num_classes: int | None = None) -> ECDETRDetector: + """Create EdgeCrafter detection model. + + Args: + num_classes: Number of classes. Defaults to ``self.num_classes``. + + Returns: + Configured :class:`ECDETRDetector`. + """ + num_classes = num_classes if num_classes is not None else self.num_classes + return self._build_ec_model(num_classes, backbone_lr=self.backbone_lr) + + def forward_for_tracing(self, inputs: Tensor) -> dict[str, Tensor]: + """Export-mode forward. + + The detector returns boxes in pixel coordinates, but ModelAPI expects + DETR-style normalized ``xyxy`` boxes (like DEIM/D-FINE). Normalize the + exported boxes by the input spatial size here. + """ + result = cast("dict[str, Tensor]", super().forward_for_tracing(inputs)) + _, _, height, width = inputs.shape + scale = torch.tensor([width, height, width, height], device=inputs.device, dtype=inputs.dtype) + result["bboxes"] = result["bboxes"] / scale + return result + + @property + def _exporter(self) -> ModelExporter: + """Creates :class:`ModelExporter` for EdgeCrafter detection.""" + return LightningModelExporter( + task_level_export_parameters=self._export_parameters, + data_input_params=self.data_input_params, + resize_mode="standard", + swap_rgb=False, + via_onnx=False, + onnx_export_configuration={ + "input_names": ["images"], + "output_names": ["bboxes", "labels", "scores"], + "dynamic_axes": {"images": {0: "batch"}}, + "dynamo": True, + "opset_version": 18, + }, + output_names=["bboxes", "labels", "scores"], + ) + + @property + def _default_preprocessing_params(self) -> dict[str, DataInputParams]: + """Default preprocessing parameters for each EdgeCrafter variant.""" + imagenet_mean = (0.485, 0.456, 0.406) + imagenet_std = (0.229, 0.224, 0.225) + return { + "edgecrafter_s": DataInputParams(input_size=(640, 640), mean=imagenet_mean, std=imagenet_std), + "edgecrafter_m": DataInputParams(input_size=(640, 640), mean=imagenet_mean, std=imagenet_std), + "edgecrafter_l": DataInputParams(input_size=(640, 640), mean=imagenet_mean, std=imagenet_std), + "edgecrafter_x": DataInputParams(input_size=(640, 640), mean=imagenet_mean, std=imagenet_std), + } diff --git a/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py b/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py new file mode 100644 index 00000000000..8c84724308d --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py @@ -0,0 +1,873 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EdgeCrafter Transformer Decoder (ECTransformer). + +Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). +Modified from D-FINE (https://github.com/Peterande/D-FINE). +Copyright (c) 2024 D-FINE Authors. All Rights Reserved. +Modified from RT-DETR (https://github.com/lyuwenyu/RT-DETR). +Copyright(c) 2023 lyuwenyu. All Rights Reserved. +""" + +from __future__ import annotations + +import copy +from collections import OrderedDict +from typing import TYPE_CHECKING, Any, Callable, cast + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import Tensor, nn +from torch.nn import init + +from getitune.backend.lightning.models.common.layers.transformer_layers import ( + LQE, + MLP, + Gate, + Integral, + MSDeformableAttentionV2, + get_contrastive_denoising_training_group, +) +from getitune.backend.lightning.models.common.utils.utils import inverse_sigmoid +from getitune.backend.lightning.models.detection.utils.utils import dfine_distance2bbox, dfine_weighting_function +from getitune.backend.lightning.models.utils.weight_init import bias_init_with_prob + +if TYPE_CHECKING: + # Deferred: importing the *package* `instance_segmentation.heads` eagerly executes + # `instance_segmentation/__init__.py`, which re-exports `EdgeCrafterInst`, which in + # turn imports `EdgeCrafterMixin` -> this module, completing a circular import. + # Safe here because `from __future__ import annotations` makes every annotation in + # this file lazy; the real class is imported locally where it is instantiated below. + from getitune.backend.lightning.models.instance_segmentation.heads import SegmentationHead + +__all__ = ["ECTransformer"] + +# --------------------------------------------------------------------------- +# Per-model hyper-parameters +# --------------------------------------------------------------------------- + +_MODEL_CFGS: dict[str, dict[str, Any]] = { + "edgecrafter_s": { + "hidden_dim": 192, + "num_heads": 8, + "dim_feedforward": 512, + "num_layers": 4, + "feat_channels": [192, 192, 192], + }, + "edgecrafter_m": { + "hidden_dim": 256, + "num_heads": 8, + "dim_feedforward": 1024, + "num_layers": 4, + "feat_channels": [256, 256, 256], + }, + "edgecrafter_l": { + "hidden_dim": 256, + "num_heads": 8, + "dim_feedforward": 1024, + "num_layers": 4, + "feat_channels": [256, 256, 256], + }, + "edgecrafter_x": { + "hidden_dim": 256, + "num_heads": 8, + "dim_feedforward": 2048, + "num_layers": 4, + "feat_channels": [256, 256, 256], + }, +} + +# --------------------------------------------------------------------------- +# Transformer Decoder Layer +# --------------------------------------------------------------------------- + + +class ECTransformerDecoderLayer(nn.Module): + """Single EC-Transformer decoder layer. + + Uses: + * ``nn.MultiheadAttention`` for self-attention (matches checkpoint keys). + * ``MSDeformableAttentionV2`` for cross-attention. + * ``Gate`` for gated cross-attention fusion. + * Standard Linear FFN with SiLU activation. + + Args: + d_model: Model dimension. + n_head: Number of attention heads. + dim_feedforward: FFN hidden dimension. + dropout: Dropout rate. + n_levels: Number of feature levels. + n_points: Number of deformable sampling points per level (int or list). + layer_scale: Optional scale factor for enlarged wide layers. + activation: Activation layer class. + """ + + def __init__( + self, + d_model: int = 256, + n_head: int = 8, + dim_feedforward: int = 1024, + dropout: float = 0.0, + n_levels: int = 3, + n_points: int | list[int] = 4, + layer_scale: float | None = None, + activation: Callable[..., nn.Module] = nn.SiLU, + ) -> None: + super().__init__() + + if layer_scale is not None: + dim_feedforward = round(layer_scale * dim_feedforward) + d_model = round(layer_scale * d_model) + + # self-attention (uses nn.MultiheadAttention to match checkpoint key layout) + self.self_attn = nn.MultiheadAttention(d_model, n_head, dropout=dropout, batch_first=True) + self.dropout1 = nn.Dropout(dropout) + self.norm1 = nn.LayerNorm(d_model) + + # cross-attention + n_points_list = [n_points] * n_levels if isinstance(n_points, int) else list(n_points) + self.cross_attn = MSDeformableAttentionV2(d_model, n_head, n_levels, n_points_list) + self.dropout2 = nn.Dropout(dropout) + self.gateway = Gate(d_model) + + # FFN + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.activation = activation() + self.dropout3 = nn.Dropout(dropout) + self.linear2 = nn.Linear(dim_feedforward, d_model) + self.dropout4 = nn.Dropout(dropout) + self.norm2 = nn.LayerNorm(d_model) + + @staticmethod + def _with_pos(tensor: Tensor, pos: Tensor | None) -> Tensor: + return tensor if pos is None else tensor + pos + + def forward( + self, + target: Tensor, + reference_points: Tensor, + value: list[Tensor], + spatial_shapes: list[list[int]], + attn_mask: Tensor | None = None, + query_pos_embed: Tensor | None = None, + ) -> Tensor: + """Forward pass.""" + # self-attention + q = k = self._with_pos(target, query_pos_embed) + target2, _ = self.self_attn(q, k, value=target, attn_mask=attn_mask) + target = self.norm1(target + self.dropout1(target2)) + + # cross-attention + target2 = self.cross_attn( + self._with_pos(target, query_pos_embed), + reference_points, + value, + spatial_shapes, + ) + target = self.gateway(target, self.dropout2(target2)) + + # FFN + target2 = self.linear2(self.dropout3(self.activation(self.linear1(target)))) + return self.norm2((target + self.dropout4(target2)).clamp(min=-65504, max=65504)) + + +# --------------------------------------------------------------------------- +# Transformer Decoder (iterative FDR) +# --------------------------------------------------------------------------- + + +class ECTransformerDecoder(nn.Module): + """Iterative Fine-grained Distribution Refinement (FDR) decoder.""" + + def __init__( + self, + hidden_dim: int, + decoder_layer: ECTransformerDecoderLayer, + decoder_layer_wide: ECTransformerDecoderLayer, + segmentation_head: SegmentationHead | None, + num_layers: int, + num_head: int, + reg_max: int, + reg_scale: Tensor, + up: Tensor, + eval_idx: int = -1, + layer_scale: int = 1, + act: Callable[..., nn.Module] = nn.SiLU, + ) -> None: + super().__init__() + self.hidden_dim = hidden_dim + self.num_layers = num_layers + self.layer_scale = layer_scale + self.num_head = num_head + self.eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx + self.up = up + self.reg_scale = reg_scale + self.reg_max = reg_max + + self.layers = nn.ModuleList( + [copy.deepcopy(decoder_layer) for _ in range(self.eval_idx + 1)] + + [copy.deepcopy(decoder_layer_wide) for _ in range(num_layers - self.eval_idx - 1)] + ) + self.segmentation_head = segmentation_head + self.lqe_layers = nn.ModuleList( + [copy.deepcopy(LQE(4, 64, 2, reg_max, activation=act)) for _ in range(num_layers)], + ) + + def _value_op( + self, + memory: Tensor, + memory_spatial_shapes: list[list[int]], + ) -> list[Tensor]: + """Pre-process encoder memory into per-level value tensors.""" + bs = memory.shape[0] + memory_reshaped = memory.reshape(bs, memory.shape[1], self.num_head, -1) + split_shape = [h * w for h, w in memory_spatial_shapes] + return list(memory_reshaped.permute(0, 2, 3, 1).split(split_shape, dim=-1)) + + def convert_to_deploy(self) -> None: + """Trim layers and freeze LQE for deployment.""" + project = dfine_weighting_function(self.reg_max, self.up, self.reg_scale) + self.register_buffer("project", project) + self.layers = self.layers[: self.eval_idx + 1] + self.lqe_layers = nn.ModuleList( + [nn.Identity()] * self.eval_idx + [self.lqe_layers[self.eval_idx]] # type: ignore[list-item] + ) + + def forward( + self, + spatial_feat: Tensor | None, + target: Tensor, + ref_points_unact: Tensor, + memory: Tensor, + spatial_shapes: list[list[int]], + bbox_head: nn.ModuleList, + score_head: nn.ModuleList, + query_pos_head: nn.Module, + pre_bbox_head: nn.Module, + integral: Integral, + up: Tensor, + reg_scale: Tensor, + attn_mask: Tensor | None = None, + dn_meta: dict | None = None, + ) -> tuple: + """Iterative decode pass.""" + output = target + output_detach = pred_corners_undetach = 0 + + value = self._value_op(memory, spatial_shapes) + + dec_out_bboxes: list[Tensor] = [] + dec_out_logits: list[Tensor] = [] + dec_out_pred_corners: list[Tensor] = [] + dec_out_refs: list[Tensor] = [] + dec_out_hs: list[Tensor] = [] + + if not hasattr(self, "project"): + project = dfine_weighting_function(self.reg_max, up, reg_scale) + else: + project = self.project # type: ignore[attr-defined] + + ref_points_detach = F.sigmoid(ref_points_unact) + query_pos_embed = query_pos_head(ref_points_detach).clamp(min=-10, max=10) + + ref_points_initial = ref_points_detach + pre_bboxes = ref_points_detach + pre_scores: Tensor = ref_points_detach + + for i, layer in enumerate(self.layers): + ref_points_input = ref_points_detach.unsqueeze(2) + + if i >= self.eval_idx + 1 and self.layer_scale > 1: + query_pos_embed = F.interpolate(query_pos_embed, scale_factor=float(self.layer_scale)) + value = self._value_op(memory, spatial_shapes) # recompute for new dim + output = F.interpolate(output, size=query_pos_embed.shape[-1]) + output_detach = output.detach() + + output = layer(output, ref_points_input, value, spatial_shapes, attn_mask, query_pos_embed) + + if i == 0: + pre_bboxes = F.sigmoid(pre_bbox_head(output) + inverse_sigmoid(ref_points_detach)) + pre_scores = score_head[0](output) + ref_points_initial = pre_bboxes.detach() + + pred_corners = bbox_head[i](output + output_detach) + pred_corners_undetach # type: ignore[operator] + inter_ref_bbox = dfine_distance2bbox(ref_points_initial, integral(pred_corners, project), reg_scale) + + if self.training or i == self.eval_idx: + scores = self.lqe_layers[i](score_head[i](output), pred_corners) + dec_out_logits.append(scores) + dec_out_bboxes.append(inter_ref_bbox) + dec_out_pred_corners.append(pred_corners) + dec_out_refs.append(ref_points_initial) + dec_out_hs.append(output) + + if not self.training: + break + + pred_corners_undetach = pred_corners + ref_points_detach = inter_ref_bbox.detach() + output_detach = output.detach() + + # Segmentation + if spatial_feat is not None and self.segmentation_head is not None: + dec_out_segs = self.segmentation_head( + spatial_features=spatial_feat, + query_features=dec_out_hs, + ) + return ( + torch.stack(dec_out_bboxes), + torch.stack(dec_out_logits), + torch.stack(dec_out_pred_corners), + torch.stack(dec_out_refs), + torch.stack(dec_out_segs), + pre_bboxes, + pre_scores, + dec_out_segs[-1], + ) + + return ( + torch.stack(dec_out_bboxes), + torch.stack(dec_out_logits), + torch.stack(dec_out_pred_corners), + torch.stack(dec_out_refs), + None, + pre_bboxes, + pre_scores, + None, + ) + + +# --------------------------------------------------------------------------- +# ECTransformer (top-level decoder module) +# --------------------------------------------------------------------------- + + +class ECTransformer(nn.Module): + """EdgeCrafter Transformer decoder. + + Handles both detection (``spatial_feat=None``) and instance segmentation + (``spatial_feat`` is the stride-8 backbone feature map). + + Args: + model_name: One of ``edgecrafter_{s,m,l,x}``. + num_classes: Number of object classes. + eval_spatial_size: ``(H, W)`` used to pre-compute static anchor cache. + num_queries: Number of object queries. + num_levels: Number of encoder feature levels. + num_points: Deformable attention points per level (int or list). + dropout: Dropout rate. + activation: FFN activation layer class. + num_denoising: Number of denoising groups. + label_noise_ratio: Label noise ratio for CDN. + box_noise_scale: Box noise scale for CDN. + reg_max: Bin count for distribution-based box regression. + reg_scale: Curvature parameter for weighting function. + layer_scale: Dimension scale factor for wide decoder layers (usually 1). + mask_downsample_ratio: When set, creates a ``SegmentationHead`` with this + output/input stride ratio (e.g. 4 for 160x160 masks from 640 input). + """ + + def __init__( + self, + model_name: str, + num_classes: int = 80, + eval_spatial_size: tuple[int, int] | list[int] = (640, 640), + num_queries: int = 300, + num_levels: int = 3, + num_points: int | list[int] = [3, 6, 3], # noqa: B006 + dropout: float = 0.0, + activation: Callable[..., nn.Module] = nn.SiLU, + num_denoising: int = 100, + label_noise_ratio: float = 0.5, + box_noise_scale: float = 1.0, + reg_max: int = 32, + reg_scale: float = 4.0, + layer_scale: int = 1, + mask_downsample_ratio: int | None = None, + ) -> None: + super().__init__() + if model_name not in _MODEL_CFGS: + msg = f"Unknown EdgeCrafter model name '{model_name}'. Available: {list(_MODEL_CFGS)}" + raise ValueError(msg) + + cfg = _MODEL_CFGS[model_name] + hidden_dim: int = cfg["hidden_dim"] + nhead: int = cfg["num_heads"] + dim_feedforward: int = cfg["dim_feedforward"] + num_layers: int = cfg["num_layers"] + feat_channels: list[int] = cfg["feat_channels"] + + self.hidden_dim = hidden_dim + self.nhead = nhead + self.num_levels = num_levels + self.num_classes = num_classes + self.num_queries = num_queries + self.num_layers = num_layers + self.eps = 1e-2 + self.aux_loss = True + self.reg_max = reg_max + self.eval_spatial_size = tuple(eval_spatial_size) + self.feat_strides = [8, 16, 32] + + # Learnable reg params (frozen after init) + self.up = nn.Parameter(torch.tensor([0.5]), requires_grad=False) + self.reg_scale = nn.Parameter(torch.tensor([reg_scale]), requires_grad=False) + + # Input projections + self._build_input_proj_layer(feat_channels) + + eval_idx = -1 # always use last layer at eval + + # Decoder layers + n_points_list = [num_points] * num_levels if isinstance(num_points, int) else list(num_points) + decoder_layer = ECTransformerDecoderLayer( + hidden_dim, nhead, dim_feedforward, dropout, num_levels, n_points_list, activation=activation + ) + decoder_layer_wide = ECTransformerDecoderLayer( + hidden_dim, + nhead, + dim_feedforward, + dropout, + num_levels, + n_points_list, + layer_scale=float(layer_scale) if layer_scale > 1 else None, + activation=activation, + ) + + seg_head: SegmentationHead | None = None + if mask_downsample_ratio is not None: + from getitune.backend.lightning.models.instance_segmentation.heads import ( + SegmentationHead as _SegmentationHead, + ) + + seg_head = _SegmentationHead( + in_dim=hidden_dim, + num_blocks=num_layers, + downsample_ratio=mask_downsample_ratio, + image_size=cast("tuple[int, int]", self.eval_spatial_size), + ) + + self.decoder = ECTransformerDecoder( + hidden_dim=hidden_dim, + decoder_layer=decoder_layer, + decoder_layer_wide=decoder_layer_wide, + segmentation_head=seg_head, + num_layers=num_layers, + num_head=nhead, + reg_max=reg_max, + reg_scale=self.reg_scale, + up=self.up, + eval_idx=eval_idx, + layer_scale=layer_scale, + act=activation, + ) + + # Denoising + self.num_denoising = num_denoising + self.label_noise_ratio = label_noise_ratio + self.box_noise_scale = box_noise_scale + if num_denoising > 0: + self.denoising_class_embed = nn.Embedding(num_classes + 1, hidden_dim, padding_idx=num_classes) + init.normal_(self.denoising_class_embed.weight[:-1]) + + # Encoder output heads + self.enc_score_head = nn.Linear(hidden_dim, num_classes) + self.enc_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, activation=activation) + self.query_pos_head = MLP(4, hidden_dim, hidden_dim, 3, activation=activation) + self.pre_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, activation=activation) + self.integral = Integral(reg_max) + + actual_eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx + self._eval_idx = actual_eval_idx + + dec_score_head = nn.Linear(hidden_dim, num_classes) + self.dec_score_head = nn.ModuleList( + [copy.deepcopy(dec_score_head) for _ in range(actual_eval_idx + 1)] + + [copy.deepcopy(dec_score_head) for _ in range(num_layers - actual_eval_idx - 1)] + ) + + scaled_dim = round(layer_scale * hidden_dim) + dec_bbox_head = MLP(hidden_dim, hidden_dim, 4 * (reg_max + 1), 3, activation=activation) + wide_bbox_head = MLP(scaled_dim, scaled_dim, 4 * (reg_max + 1), 3, activation=activation) + self.dec_bbox_head = nn.ModuleList( + [copy.deepcopy(dec_bbox_head) for _ in range(actual_eval_idx + 1)] + + [ + wide_bbox_head if i > actual_eval_idx else copy.deepcopy(dec_bbox_head) + for i in range(actual_eval_idx + 1, num_layers) + ] + ) + + # Static anchor cache + if self.eval_spatial_size: + anchors, valid_mask = self._generate_anchors() + self.register_buffer("anchors", anchors) + self.register_buffer("valid_mask", valid_mask) + + self._reset_parameters(feat_channels) + + # ------------------------------------------------------------------ + def _build_input_proj_layer(self, feat_channels: list[int]) -> None: + self.input_proj = nn.ModuleList() + for in_ch in feat_channels: + if in_ch == self.hidden_dim: + self.input_proj.append(nn.Identity()) + else: + _layers1: list[tuple[str, nn.Module]] = [ + ("conv", nn.Conv2d(in_ch, self.hidden_dim, 1, bias=False)), + ("norm", nn.BatchNorm2d(self.hidden_dim)), + ] + self.input_proj.append(nn.Sequential(OrderedDict(_layers1))) + for _ in range(self.num_levels - len(feat_channels)): + in_ch = feat_channels[-1] + _layers2: list[tuple[str, nn.Module]] = [ + ("conv", nn.Conv2d(in_ch, self.hidden_dim, 3, 2, padding=1, bias=False)), + ("norm", nn.BatchNorm2d(self.hidden_dim)), + ] + self.input_proj.append(nn.Sequential(OrderedDict(_layers2))) + + def _reset_parameters(self, feat_channels: list[int]) -> None: + bias = bias_init_with_prob(0.01) + init.constant_(self.enc_score_head.bias, bias) + _enc_last = cast("nn.Linear", self.enc_bbox_head.layers[-1]) + init.constant_(_enc_last.weight, 0) + init.constant_(_enc_last.bias, 0) + _pre_last = cast("nn.Linear", self.pre_bbox_head.layers[-1]) + init.constant_(_pre_last.weight, 0) + init.constant_(_pre_last.bias, 0) + for cls_, reg_ in zip(self.dec_score_head, self.dec_bbox_head): + init.constant_(cast("nn.Linear", cls_).bias, bias) + if hasattr(reg_, "layers"): + _reg = cast("MLP", reg_) + _reg_last = cast("nn.Linear", _reg.layers[-1]) + init.constant_(_reg_last.weight, 0) + init.constant_(_reg_last.bias, 0) + _q0 = cast("nn.Linear", self.query_pos_head.layers[0]) + init.xavier_uniform_(_q0.weight) + _q1 = cast("nn.Linear", self.query_pos_head.layers[1]) + init.xavier_uniform_(_q1.weight) + _qlast = cast("nn.Linear", self.query_pos_head.layers[-1]) + init.xavier_uniform_(_qlast.weight) + + def _generate_anchors( + self, + spatial_shapes: list[list[int]] | None = None, + grid_size: float = 0.05, + dtype: torch.dtype = torch.float32, + device: str | torch.device = "cpu", + ) -> tuple[Tensor, Tensor]: + if spatial_shapes is None: + spatial_shapes = [ + [self.eval_spatial_size[0] // s, self.eval_spatial_size[1] // s] for s in self.feat_strides + ] + anchors: list[Tensor] = [] + for lvl, (h, w) in enumerate(spatial_shapes): + gy, gx = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij") + grid_xy = torch.stack([gx, gy], dim=-1) + grid_xy = (grid_xy.unsqueeze(0) + 0.5) / torch.tensor([w, h], dtype=dtype) + wh = torch.ones_like(grid_xy) * grid_size * (2.0**lvl) + lvl_anchors = torch.cat([grid_xy, wh], dim=-1).reshape(-1, h * w, 4) + anchors.append(lvl_anchors) + anchors_cat = torch.cat(anchors, dim=1).to(device) + valid_mask = ((anchors_cat > self.eps) & (anchors_cat < 1 - self.eps)).all(-1, keepdim=True) + anchors_cat = torch.log(anchors_cat / (1 - anchors_cat)) + anchors_cat = torch.where(valid_mask, anchors_cat, torch.full_like(anchors_cat, float("inf"))) + return anchors_cat, valid_mask + + def _get_encoder_input(self, feats: list[Tensor]) -> tuple[Tensor, list[list[int]]]: + proj_feats = [self.input_proj[i](feat) for i, feat in enumerate(feats)] + feat_flatten: list[Tensor] = [] + spatial_shapes: list[list[int]] = [] + for feat in proj_feats: + _, _, h, w = feat.shape + feat_flatten.append(feat.flatten(2).permute(0, 2, 1)) + spatial_shapes.append([h, w]) + return torch.cat(feat_flatten, 1), spatial_shapes + + def _select_topk( + self, + memory: Tensor, + outputs_logits: Tensor, + outputs_anchors_unact: Tensor, + topk: int, + ) -> tuple[Tensor, Tensor | None, Tensor]: + _, topk_ind = torch.topk(outputs_logits.max(-1).values, topk, dim=-1) + topk_anchors = outputs_anchors_unact.gather( + dim=1, + index=topk_ind.unsqueeze(-1).expand(-1, -1, outputs_anchors_unact.shape[-1]), + ) + topk_logits = ( + outputs_logits.gather( + dim=1, + index=topk_ind.unsqueeze(-1).expand(-1, -1, outputs_logits.shape[-1]), + ) + if self.training + else None + ) + topk_memory = memory.gather( + dim=1, + index=topk_ind.unsqueeze(-1).expand(-1, -1, memory.shape[-1]), + ) + return topk_memory, topk_logits, topk_anchors + + def _get_decoder_input( + self, + memory: Tensor, + spatial_shapes: list[list[int]], + denoising_logits: Tensor | None = None, + denoising_bbox_unact: Tensor | None = None, + ) -> tuple[Tensor, Tensor, list[Tensor], list[Tensor | None]]: + if self.training or not hasattr(self, "anchors"): + anchors, valid_mask = self._generate_anchors(spatial_shapes, device=memory.device) + else: + anchors: Tensor = self.anchors # type: ignore[attr-defined] + valid_mask: Tensor = self.valid_mask # type: ignore[attr-defined] + + if memory.shape[0] > 1: + anchors = anchors.expand(memory.shape[0], -1, -1) + + memory = valid_mask.to(memory.dtype) * memory + enc_outputs_logits = self.enc_score_head(memory) + + topk_memory, topk_logits, topk_anchors = self._select_topk( + memory, enc_outputs_logits, anchors, self.num_queries + ) + enc_topk_bbox_unact = self.enc_bbox_head(topk_memory) + topk_anchors + + enc_bboxes_list: list[Tensor] = [] + enc_logits_list: list[Tensor | None] = [] + if self.training: + enc_bboxes_list.append(F.sigmoid(enc_topk_bbox_unact)) + enc_logits_list.append(topk_logits) + + content = topk_memory.detach() + enc_topk_bbox_unact = enc_topk_bbox_unact.detach() + + if denoising_bbox_unact is not None and denoising_logits is not None: + enc_topk_bbox_unact = torch.cat([denoising_bbox_unact, enc_topk_bbox_unact], dim=1) + content = torch.cat([denoising_logits, content], dim=1) + + return content, enc_topk_bbox_unact, enc_bboxes_list, enc_logits_list + + @staticmethod + def _split(x: Tensor | None, dim: int, s_idx: Tensor | list[int] | None) -> tuple[Tensor | None, Tensor | None]: + if x is None or s_idx is None: + return None, x + split_sizes: list[int] = s_idx.tolist() if isinstance(s_idx, Tensor) else s_idx + return torch.split(x, split_sizes, dim=dim) # type: ignore[return-value] + + def forward( + self, + feats: list[Tensor], + targets: list[dict] | None = None, + spatial_feat: Tensor | None = None, + ) -> dict[str, Any]: + """Forward pass. + + Args: + feats: List of multi-scale features from the encoder. + targets: Ground-truth targets (required during training). + spatial_feat: Stride-8 spatial feature for segmentation head (ECSeg only). + + Returns: + Dict with ``pred_logits``, ``pred_boxes``, optional ``pred_masks``, + and auxiliary losses during training. + """ + memory, spatial_shapes = self._get_encoder_input(feats) + + # Denoising + denoising_logits = denoising_bbox_unact = attn_mask = dn_meta = None + if self.training and self.num_denoising > 0 and targets is not None: + denoising_logits, denoising_bbox_unact, attn_mask, dn_meta = get_contrastive_denoising_training_group( + targets, + self.num_classes, + self.num_queries, + self.denoising_class_embed, + num_denoising=self.num_denoising, + label_noise_ratio=self.label_noise_ratio, + box_noise_scale=self.box_noise_scale, + ) + + init_ref_contents, init_ref_points_unact, enc_topk_bboxes_list, enc_topk_logits_list = self._get_decoder_input( + memory, spatial_shapes, denoising_logits, denoising_bbox_unact + ) + + out_bboxes, out_logits, out_corners, out_refs, out_masks, pre_bboxes, pre_logits, pre_segs = self.decoder( + spatial_feat, + init_ref_contents, + init_ref_points_unact, + memory, + spatial_shapes, + self.dec_bbox_head, + self.dec_score_head, + self.query_pos_head, + self.pre_bbox_head, + self.integral, + self.up, + self.reg_scale, + attn_mask=attn_mask, + dn_meta=dn_meta, + ) + + s_idx = dn_meta["dn_num_split"] if dn_meta is not None else None + + if self.training and dn_meta is not None: + dn_pre_logits, pre_logits = self._split(pre_logits, 1, s_idx) + dn_pre_bboxes, pre_bboxes = self._split(pre_bboxes, 1, s_idx) + dn_pre_segs, pred_segs = self._split(pre_segs, 1, s_idx) + dn_out_logits, out_logits = self._split(out_logits, 2, s_idx) + dn_out_bboxes, out_bboxes = self._split(out_bboxes, 2, s_idx) + dn_out_masks, out_masks = self._split(out_masks, 2, s_idx) + dn_out_corners, out_corners = self._split(out_corners, 2, s_idx) + dn_out_refs, out_refs = self._split(out_refs, 2, s_idx) + else: + pred_segs = pre_segs + dn_out_logits = dn_out_bboxes = dn_out_masks = dn_out_corners = dn_out_refs = None + dn_pre_logits = dn_pre_bboxes = dn_pre_segs = None + + last_masks = out_masks[-1] if out_masks is not None else None + if self.training: + out: dict[str, Any] = { + "pred_logits": out_logits[-1], + "pred_boxes": out_bboxes[-1], + "pred_corners": out_corners[-1], + "ref_points": out_refs[-1], + "up": self.up, + "reg_scale": self.reg_scale, + } + if last_masks is not None: + out["pred_masks"] = last_masks + else: + out = { + "pred_logits": out_logits[-1], + "pred_boxes": out_bboxes[-1], + } + if last_masks is not None: + out["pred_masks"] = last_masks + + if self.training and self.aux_loss: + aux_masks: list[Tensor | None] = ( + list(out_masks[:-1]) if out_masks is not None else [None] * (len(out_logits) - 1) + ) + out["aux_outputs"] = self._set_aux_loss2( + out_logits[:-1], + out_bboxes[:-1], + out_corners[:-1], + out_refs[:-1], + aux_masks, + out_corners[-1], + out_logits[-1], + ) + out["enc_aux_outputs"] = self._set_aux_loss(enc_topk_logits_list, enc_topk_bboxes_list) + pre_out: dict[str, Any] = { + "pred_logits": pre_logits, + "pred_boxes": pre_bboxes, + } + if pred_segs is not None: + pre_out["pred_masks"] = pred_segs + out["pre_outputs"] = pre_out + out["enc_meta"] = {"class_agnostic": False} + + if dn_meta is not None: + dn_masks: list[Tensor | None] = ( + list(dn_out_masks[:-1]) if dn_out_masks is not None else [None] * (self.num_layers - 1) + ) + out["dn_outputs"] = self._set_aux_loss2( + dn_out_logits, + dn_out_bboxes, + dn_out_corners, + dn_out_refs, + dn_masks, + dn_out_corners[-1] if dn_out_corners is not None else None, + dn_out_logits[-1] if dn_out_logits is not None else None, + ) + dn_pre_out: dict[str, Any] = { + "pred_logits": dn_pre_logits, + "pred_boxes": dn_pre_bboxes, + } + if dn_pre_segs is not None: + dn_pre_out["pred_masks"] = dn_pre_segs + out["dn_pre_outputs"] = dn_pre_out + out["dn_meta"] = dn_meta + + return out + + @torch.jit.unused + def _set_aux_loss(self, outputs_class: list[Tensor | None], outputs_coord: list[Tensor]) -> list[dict[str, Tensor]]: + return [{"pred_logits": a, "pred_boxes": b} for a, b in zip(outputs_class, outputs_coord) if a is not None] + + @torch.jit.unused + def _set_aux_loss2( + self, + outputs_class: Tensor | None, + outputs_coord: Tensor | None, + outputs_corners: Tensor | None, + outputs_ref: Tensor | None, + outputs_masks: list[Tensor | None] | None = None, + teacher_corners: Tensor | None = None, + teacher_logits: Tensor | None = None, + ) -> list[dict[str, Any]]: + if outputs_class is None or outputs_coord is None: + return [] + results = [] + n = outputs_class.shape[0] + for i in range(n): + result: dict[str, Any] = { + "pred_logits": outputs_class[i], + "pred_boxes": outputs_coord[i], + "pred_corners": outputs_corners[i] if outputs_corners is not None else None, + "ref_points": outputs_ref[i] if outputs_ref is not None else None, + "teacher_corners": teacher_corners, + "teacher_logits": teacher_logits, + } + if outputs_masks is not None and i < len(outputs_masks) and outputs_masks[i] is not None: + result["pred_masks"] = outputs_masks[i] + results.append(result) + return results + + # ------------------------------------------------------------------ + # Postprocessor (used by ECDETRDetector.postprocess) + # ------------------------------------------------------------------ + + @torch.no_grad() + def postprocess( + self, + outputs: dict[str, Tensor], + orig_target_sizes: Tensor, + num_top_queries: int = 300, + ) -> list[dict[str, Tensor]]: + """Decode raw model outputs into boxes/labels/scores (and masks if present).""" + import torchvision + + logits = outputs["pred_logits"] + boxes = outputs["pred_boxes"] + mask_pred = outputs.get("pred_masks") + + bbox_pred = torchvision.ops.box_convert(boxes, in_fmt="cxcywh", out_fmt="xyxy") + # orig_target_sizes is (B, 2) with (H, W); xyxy needs (W, H, W, H) scale. + bbox_pred = bbox_pred * orig_target_sizes.flip(-1).repeat(1, 2).unsqueeze(1) + + scores = F.sigmoid(logits) + scores, index = torch.topk(scores.flatten(1), num_top_queries, dim=-1) + labels = index % self.num_classes + index = index // self.num_classes + boxes_out = bbox_pred.gather(dim=1, index=index.unsqueeze(-1).expand(-1, -1, 4)) + + masks_out: Tensor | None = None + if mask_pred is not None: + masks_out = mask_pred.gather( + dim=1, + index=index.unsqueeze(-1).unsqueeze(-1).expand(-1, -1, mask_pred.shape[-2], mask_pred.shape[-1]), + ) + + results = [] + for i, (s, lb, b) in enumerate(zip(scores, labels, boxes_out)): + res: dict[str, Tensor] = {"scores": s, "labels": lb, "boxes": b} + if masks_out is not None: + h, w = orig_target_sizes[i].tolist() + m = F.interpolate( + masks_out[i].unsqueeze(1), size=(int(h), int(w)), mode="bilinear", align_corners=False + ) + res["masks"] = (m > 0.0).squeeze(1) + results.append(res) + return results diff --git a/library/src/getitune/backend/lightning/models/detection/losses/__init__.py b/library/src/getitune/backend/lightning/models/detection/losses/__init__.py index 07c877027d9..b41bea7b882 100644 --- a/library/src/getitune/backend/lightning/models/detection/losses/__init__.py +++ b/library/src/getitune/backend/lightning/models/detection/losses/__init__.py @@ -5,6 +5,7 @@ from .atss_loss import ATSSCriterion from .dfine_loss import DFINECriterion +from .ec_loss import ECCriterion from .rtdetr_loss import DetrCriterion from .rtmdet_loss import RTMDetCriterion from .ssd_loss import SSDCriterion @@ -14,6 +15,7 @@ "ATSSCriterion", "DFINECriterion", "DetrCriterion", + "ECCriterion", "RTMDetCriterion", "SSDCriterion", "YOLOXCriterion", diff --git a/library/src/getitune/backend/lightning/models/detection/losses/dfine_loss.py b/library/src/getitune/backend/lightning/models/detection/losses/dfine_loss.py index 61f5ee8e07d..ad844eb4c17 100644 --- a/library/src/getitune/backend/lightning/models/detection/losses/dfine_loss.py +++ b/library/src/getitune/backend/lightning/models/detection/losses/dfine_loss.py @@ -487,7 +487,7 @@ def fgl_loss( loss = loss_left + loss_right - if iou_weight is not None and iou_weight.sum() > 0: + if iou_weight is not None: iou_weight = iou_weight.float() loss = loss * iou_weight diff --git a/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py b/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py new file mode 100644 index 00000000000..3452a7dbc18 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py @@ -0,0 +1,127 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EdgeCrafter criterion implementation. + +Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). +Modified from D-FINE (https://github.com/Peterande/D-FINE). +Copyright (c) 2024 D-FINE Authors. All Rights Reserved. +""" + +from __future__ import annotations + +from typing import Callable + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import Tensor + +from getitune.backend.lightning.models.common.utils.assigners.hungarian_matcher import HungarianMatcher + +from .deim_loss import DEIMCriterion + + +class ECCriterion(DEIMCriterion): + """EdgeCrafter criterion extending DEIMCriterion with instance-segmentation mask losses. + + Adds sigmoid-BCE and Dice losses for matched prediction/GT mask pairs on top of the + VFL, MAL, box regression, FGL, and DDF losses inherited from :class:`DEIMCriterion`. + + Mask losses are computed only when ``pred_masks`` is present in ``outputs``, + so auxiliary encoder branches (which do not produce masks) are handled + transparently without extra bookkeeping. + + Args: + weight_dict: Loss weight dictionary; may include ``loss_mask_ce`` and + ``loss_mask_dice`` keys for mask losses. + alpha: VFL/MAL alpha parameter. Defaults to 0.2. + gamma: VFL/MAL gamma parameter. Defaults to 2.0. + num_classes: Number of object classes. Defaults to 80. + reg_max: Bin count for distribution-based box regression. Defaults to 32. + matcher_cost_dict: Optional custom cost dictionary for the Hungarian matcher. + When provided, overrides the default detection matcher. + """ + + def __init__( + self, + weight_dict: dict[str, int | float], + alpha: float = 0.2, + gamma: float = 2.0, + num_classes: int = 80, + reg_max: int = 32, + matcher_cost_dict: dict[str, int | float] | None = None, + ) -> None: + super().__init__(weight_dict, alpha, gamma, num_classes, reg_max) + if matcher_cost_dict is not None: + self.matcher = HungarianMatcher(cost_dict=matcher_cost_dict) + + def loss_masks( + self, + outputs: dict[str, Tensor], + targets: list[dict[str, Tensor]], + indices: list[tuple[Tensor, Tensor]], + num_boxes: int, + ) -> dict[str, Tensor]: + """Sigmoid-BCE and Dice losses on matched prediction/GT mask pairs. + + Silently returns an empty dict when ``pred_masks`` is absent in + ``outputs`` (e.g. for encoder auxiliary branches or pure detection). + + Args: + outputs: Model output dict; must contain ``pred_masks`` [B, Q, H, W] + when mask supervision is desired. + targets: Per-image target dicts; must contain ``masks`` [N, Ht, Wt]. + indices: Hungarian-matched (src_idx, tgt_idx) pairs per image. + num_boxes: Normalisation denominator (total matched boxes). + + Returns: + Dict with ``loss_mask_ce`` and ``loss_mask_dice``, or empty dict. + """ + pred_masks = outputs.get("pred_masks") + if pred_masks is None: + return {} + + idx = self._get_src_permutation_idx(indices) + src_masks = pred_masks[idx] # [N, H_m, W_m] + + if src_masks.numel() == 0: + return { + "loss_mask_ce": src_masks.sum(), + "loss_mask_dice": src_masks.sum(), + } + + # Gather matched GT masks and resize to prediction spatial size + target_masks = torch.cat( + [t["masks"][j] for t, (_, j) in zip(targets, indices)], + dim=0, + ).float() # [N, Ht, Wt] + + h_m, w_m = src_masks.shape[-2:] + if target_masks.shape[-2:] != (h_m, w_m): + target_masks = F.interpolate( + target_masks.unsqueeze(1), + size=(h_m, w_m), + mode="nearest", + ).squeeze(1) + + # Sigmoid BCE loss + loss_ce = F.binary_cross_entropy_with_logits(src_masks, target_masks, reduction="none") + loss_ce = loss_ce.flatten(1).mean(1).sum() / num_boxes + + # Dice loss + pred_sig = src_masks.sigmoid().flatten(1) + tgt_flat = target_masks.flatten(1) + numerator = 2 * (pred_sig * tgt_flat).sum(-1) + denominator = pred_sig.sum(-1) + tgt_flat.sum(-1) + loss_dice = (1 - (numerator + 1) / (denominator + 1)).sum() / num_boxes + + return {"loss_mask_ce": loss_ce, "loss_mask_dice": loss_dice} + + @property + def _available_losses(self) -> tuple[Callable]: # type: ignore[return-value] + return ( # pyrefly: ignore[bad-return] + self.loss_boxes, + self.loss_labels_mal, + self.loss_local, + self.loss_masks, + ) diff --git a/library/src/getitune/backend/lightning/models/detection/necks/dfine_hybrid_encoder.py b/library/src/getitune/backend/lightning/models/detection/necks/dfine_hybrid_encoder.py index 5279a86b32d..3992e219c32 100644 --- a/library/src/getitune/backend/lightning/models/detection/necks/dfine_hybrid_encoder.py +++ b/library/src/getitune/backend/lightning/models/detection/necks/dfine_hybrid_encoder.py @@ -905,6 +905,43 @@ class HybridEncoder: "fuse_op": "sum", "use_fusable_layers": True, }, + # EdgeCrafter models (use sum fusion, fusable layers) + "edgecrafter_s": { + "in_channels": [192, 192, 192], + "hidden_dim": 192, + "dim_feedforward": 512, + "depth_mult": 0.67, + "expansion": 0.34, + "fuse_op": "sum", + "use_fusable_layers": True, + }, + "edgecrafter_m": { + "in_channels": [256, 256, 256], + "hidden_dim": 256, + "dim_feedforward": 512, + "depth_mult": 0.67, + "expansion": 0.75, + "fuse_op": "sum", + "use_fusable_layers": True, + }, + "edgecrafter_l": { + "in_channels": [256, 256, 256], + "hidden_dim": 256, + "dim_feedforward": 1024, + "depth_mult": 1.0, + "expansion": 0.75, + "fuse_op": "sum", + "use_fusable_layers": True, + }, + "edgecrafter_x": { + "in_channels": [256, 256, 256], + "hidden_dim": 256, + "dim_feedforward": 2048, + "depth_mult": 1.0, + "expansion": 1.5, + "fuse_op": "sum", + "use_fusable_layers": True, + }, } def __new__(cls, model_name: str) -> HybridEncoderModule: diff --git a/library/src/getitune/backend/lightning/models/instance_segmentation/__init__.py b/library/src/getitune/backend/lightning/models/instance_segmentation/__init__.py index 303608cfff1..8205497116e 100644 --- a/library/src/getitune/backend/lightning/models/instance_segmentation/__init__.py +++ b/library/src/getitune/backend/lightning/models/instance_segmentation/__init__.py @@ -3,9 +3,10 @@ """Module for getitune instance segmentation models.""" +from .edgecrafter_inst import EdgeCrafterInst from .maskrcnn import MaskRCNN from .maskrcnn_tv import MaskRCNNTV from .rfdetr_inst import RFDETRInst from .rtmdet_inst import RTMDetInst -__all__ = ["MaskRCNN", "MaskRCNNTV", "RFDETRInst", "RTMDetInst"] +__all__ = ["EdgeCrafterInst", "MaskRCNN", "MaskRCNNTV", "RFDETRInst", "RTMDetInst"] diff --git a/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py b/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py new file mode 100644 index 00000000000..aae38f213ae --- /dev/null +++ b/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py @@ -0,0 +1,187 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EdgeCrafter model implementation for instance segmentation.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any, ClassVar, Literal + +from getitune.backend.lightning.exporter.base import ModelExporter +from getitune.backend.lightning.exporter.native import LightningModelExporter +from getitune.backend.lightning.models.base import DataInputParams, DefaultOptimizerCallable, DefaultSchedulerCallable +from getitune.backend.lightning.models.common.edgecrafter_mixin import EdgeCrafterMixin +from getitune.backend.lightning.models.detection.detectors.edgecrafter import ECDETRDetector +from getitune.backend.lightning.models.instance_segmentation.base import LightningInstanceSegModel +from getitune.config.data import TileConfig +from getitune.metrics.fmeasure import MaskRLEMeanAPFMeasureCallable +from getitune.types.export import TaskLevelExportParameters + +if TYPE_CHECKING: + import torch + from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable + + from getitune.backend.lightning.schedulers import LRSchedulerListCallable + from getitune.metrics import MetricCallable + from getitune.types.label import LabelInfoTypes + + +class EdgeCrafterInst(EdgeCrafterMixin, LightningInstanceSegModel): # pyrefly: ignore[inconsistent-inheritance] + """getitune Instance Segmentation model class for EdgeCrafter. + + Extends the EdgeCrafter detection backbone with a dedicated instance-segmentation + head (``SegmentationHead`` inside :class:`ECTransformer`). The backbone is + initialised from the ECSeg pretrained weights + (``ecseg_vitt/vittplus/vits/vitsplus``) which are trained jointly for + detection and segmentation. + + Original repository: https://github.com/Intellindust-AI-Lab/EdgeCrafter + Paper: https://arxiv.org/abs/2603.18739 + + Args: + label_info: Information about the labels. + data_input_params: Preprocessing parameters (input size, mean, std). + When ``None`` the ``_default_preprocessing_params`` dict is used. + model_name: One of ``"edgecrafter_s"``, ``"edgecrafter_m"``, + ``"edgecrafter_l"``, ``"edgecrafter_x"``. + optimizer: Optimizer callable. Defaults to :data:`DefaultOptimizerCallable`. + scheduler: LR-scheduler callable. Defaults to :data:`DefaultSchedulerCallable`. + metric: Metric callable. Defaults to mask RLE mean AP. + multi_scale: Whether to use multi-scale training. Defaults to ``False``. + backbone_lr: Learning rate for the backbone. Defaults to the upstream + per-variant value in :class:`EdgeCrafterMixin`. + torch_compile: Whether to use ``torch.compile``. Defaults to ``False``. + tile_config: Tiling configuration. Defaults to disabled tiler. + """ + + # ECSeg weights - loaded in place of ECDet weights. + _pretrained_weights: ClassVar[dict[str, str]] = { + "edgecrafter_s": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_s.pth", + "edgecrafter_m": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_m.pth", + "edgecrafter_l": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_l.pth", + "edgecrafter_x": "https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_x.pth", + } + + # Instance-segmentation loss weights + matcher costs (adds mask terms on top of detection). + _loss_weights: ClassVar[dict[str, float]] = { + "loss_mal": 2.0, + "loss_bbox": 1.0, + "loss_giou": 1.0, + "loss_fgl": 0.15, + "loss_ddf": 1.5, + "loss_mask_ce": 5.0, + "loss_mask_dice": 5.0, + } + _matcher_cost_dict: ClassVar[dict[str, int | float]] = { + "cost_class": 2, + "cost_bbox": 1, + "cost_giou": 1, + "cost_mask": 5, + "cost_dice": 5, + } + _mask_downsample_ratio: ClassVar[int] = 4 + _backbone_key: ClassVar[str] = "seg_backbone_name" + + input_size_multiplier = 16 + + def __init__( + self, + label_info: LabelInfoTypes, + data_input_params: DataInputParams | dict | None = None, + model_name: Literal[ + "edgecrafter_s", + "edgecrafter_m", + "edgecrafter_l", + "edgecrafter_x", + ] = "edgecrafter_s", + optimizer: OptimizerCallable = DefaultOptimizerCallable, + scheduler: LRSchedulerCallable | LRSchedulerListCallable = DefaultSchedulerCallable, + metric: MetricCallable = MaskRLEMeanAPFMeasureCallable, + multi_scale: bool = False, + backbone_lr: float | None = None, + torch_compile: bool = False, + tile_config: TileConfig = TileConfig(enable_tiler=False), + ) -> None: + self.multi_scale = multi_scale + self.backbone_lr = backbone_lr + super().__init__( + label_info=label_info, + data_input_params=data_input_params, + model_name=model_name, + optimizer=optimizer, + scheduler=scheduler, + metric=metric, + torch_compile=torch_compile, + tile_config=tile_config, + ) + + def _create_model(self, num_classes: int | None = None) -> ECDETRDetector: + """Create EdgeCrafter instance segmentation model. + + Args: + num_classes: Number of classes. Defaults to ``self.num_classes``. + + Returns: + Configured :class:`ECDETRDetector` with instance-segmentation head. + """ + num_classes = num_classes if num_classes is not None else self.num_classes + return self._build_ec_model(num_classes, backbone_lr=self.backbone_lr) + + @property + def _export_parameters(self) -> TaskLevelExportParameters: + """Export parameters for EdgeCrafter instance segmentation. + + Uses the ``DETRInstSeg`` model type which triggers full-image mask + postprocessing in ModelAPI. A conservative IoU threshold (0.8) suppresses + only near-duplicate predictions. + """ + return super()._export_parameters.wrap(model_type="DETRInstSeg", iou_threshold=0.8) + + @property + def _exporter(self) -> ModelExporter: + """Creates :class:`ModelExporter` for EdgeCrafter instance segmentation.""" + return LightningModelExporter( + task_level_export_parameters=self._export_parameters, + data_input_params=self.data_input_params, + resize_mode="standard", + swap_rgb=False, + via_onnx=False, + onnx_export_configuration={ + "input_names": ["images"], + "output_names": ["boxes", "labels", "masks"], + "dynamic_axes": {"images": {0: "batch"}}, + "dynamo": False, + "do_constant_folding": True, + "opset_version": 18, + }, + output_names=["boxes", "labels", "masks"], + ) + + def forward_for_tracing(self, inputs: torch.Tensor) -> dict[str, Any]: + """Forward pass used for ONNX / OpenVINO export. + + Reformats the deploy-mode output into the MaskRCNN-compatible format: + ``boxes`` [B, Q, 5] (x1,y1,x2,y2,score in pixel coords), + ``labels`` [B, Q], ``masks`` [B, Q, H, W]. + + Args: + inputs: Image batch [B, C, H, W]. + + Returns: + Dict with ``boxes``, ``labels``, ``masks``. + """ + meta_info_list = self._default_is_meta(inputs) + result = self.model.export(inputs, meta_info_list) # type: ignore[attr-defined] # pyrefly: ignore[not-callable] + return self._make_is_export_prediction(inputs, result) + + @property + def _default_preprocessing_params(self) -> dict[str, DataInputParams]: + """Default preprocessing parameters for each EdgeCrafter instance-seg variant.""" + imagenet_mean = (0.485, 0.456, 0.406) + imagenet_std = (0.229, 0.224, 0.225) + return { + "edgecrafter_s": DataInputParams(input_size=(640, 640), mean=imagenet_mean, std=imagenet_std), + "edgecrafter_m": DataInputParams(input_size=(640, 640), mean=imagenet_mean, std=imagenet_std), + "edgecrafter_l": DataInputParams(input_size=(640, 640), mean=imagenet_mean, std=imagenet_std), + "edgecrafter_x": DataInputParams(input_size=(640, 640), mean=imagenet_mean, std=imagenet_std), + } diff --git a/library/src/getitune/backend/lightning/models/instance_segmentation/heads/__init__.py b/library/src/getitune/backend/lightning/models/instance_segmentation/heads/__init__.py index cfebe90c385..7649662de81 100644 --- a/library/src/getitune/backend/lightning/models/instance_segmentation/heads/__init__.py +++ b/library/src/getitune/backend/lightning/models/instance_segmentation/heads/__init__.py @@ -4,6 +4,7 @@ """Custom head implementations for instance segmentation task.""" from .bbox_head import ConvFCBBoxHead +from .ec_segmentation_head import SegmentationHead from .fcn_mask_head import FCNMaskHead from .roi_head import RoIHead from .roi_head_tv import TVRoIHeads @@ -16,5 +17,6 @@ "RPNHead", "RTMDetInstSepBNHead", "RoIHead", + "SegmentationHead", "TVRoIHeads", ] diff --git a/library/src/getitune/backend/lightning/models/instance_segmentation/heads/ec_segmentation_head.py b/library/src/getitune/backend/lightning/models/instance_segmentation/heads/ec_segmentation_head.py new file mode 100644 index 00000000000..ee80c98f019 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/instance_segmentation/heads/ec_segmentation_head.py @@ -0,0 +1,119 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EdgeCrafter instance-segmentation head (query-conditioned dot-product mask head). + +Used by :class:`~getitune.backend.lightning.models.detection.heads.ec_decoder.ECTransformer` +when ``mask_downsample_ratio`` is set (i.e. for :class:`EdgeCrafterInst`). Kept as its own +module in ``instance_segmentation/heads/`` (rather than inlined in ``ec_decoder.py``) to +keep the primary decoder file focused on detection logic and to group this +instance-segmentation-specific head alongside its siblings (``fcn_mask_head.py``, +``rtmdet_inst_head.py``, etc.), following the "one head per file" convention used +throughout ``instance_segmentation/heads/``. + +Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). +""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F # noqa: N812 +from torch import Tensor, nn + +__all__ = ["SegmentationHead"] + + +class _DepthwiseConvBlock(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=3, padding=1, groups=dim) + self.norm = nn.LayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, dim) + self.act = nn.GELU() + + def forward(self, x: Tensor) -> Tensor: + """Forward pass.""" + res = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.act(self.pwconv1(self.norm(x))) + return x.permute(0, 3, 1, 2) + res + + +class _MLPBlock(nn.Module): + def __init__(self, dim: int) -> None: + super().__init__() + self.norm_in = nn.LayerNorm(dim) + # Named "layers" to match checkpoint key layout (layers.0, layers.2). + self.layers = nn.Sequential( + nn.Linear(dim, dim * 4), + nn.GELU(), + nn.Linear(dim * 4, dim), + ) + + def forward(self, x: Tensor) -> Tensor: + """Forward pass.""" + return x + self.layers(self.norm_in(x)) + + +class SegmentationHead(nn.Module): + """Lightweight dot-product segmentation head. + + Args: + in_dim: Feature dimension. + num_blocks: Number of DepthwiseConvBlocks applied to spatial features. + downsample_ratio: Spatial downsampling ratio for the mask output. + image_size: Reference image size ``(H, W)`` (used to compute mask target size). + """ + + def __init__( + self, + in_dim: int, + num_blocks: int, + downsample_ratio: int = 4, + image_size: tuple[int, int] | list[int] = (640, 640), + ) -> None: + super().__init__() + self.downsample_ratio = downsample_ratio + self.image_size = tuple(image_size) + self.blocks = nn.ModuleList([_DepthwiseConvBlock(in_dim) for _ in range(num_blocks)]) + # 1x1 conv projects spatial features channel-wise (matches checkpoint key layout). + self.spatial_features_proj = nn.Conv2d(in_dim, in_dim, 1) + self.query_features_block = _MLPBlock(in_dim) + # Linear projection on query features (matches checkpoint key layout). + self.query_features_proj = nn.Linear(in_dim, in_dim) + self.bias = nn.Parameter(torch.zeros(1)) + + def forward( + self, + spatial_features: Tensor, + query_features: list[Tensor], + ) -> list[Tensor]: + """Forward pass during training (one mask per decoder layer).""" + h_out = self.image_size[0] // self.downsample_ratio + w_out = self.image_size[1] // self.downsample_ratio + sf = F.interpolate(spatial_features, size=(h_out, w_out), mode="bilinear", align_corners=False) + + mask_logits = [] + for block, qf_in in zip(self.blocks, query_features): + sf = block(sf) + sf_proj = self.spatial_features_proj(sf) + qf = self.query_features_proj(self.query_features_block(qf_in)) + mask_logits.append(torch.einsum("bchw,bnc->bnhw", sf_proj, qf) + self.bias) + return mask_logits + + def forward_export( + self, + spatial_features: Tensor, + query_features: list[Tensor], + ) -> list[Tensor]: + """Forward at export time (single query feature, no dropout).""" + assert len(query_features) == 1 # noqa: S101 + h_out = self.image_size[0] // self.downsample_ratio + w_out = self.image_size[1] // self.downsample_ratio + sf = F.interpolate(spatial_features, size=(h_out, w_out), mode="bilinear", align_corners=False) + for block in self.blocks: + sf = block(sf) + sf = self.spatial_features_proj(sf) + qf = self.query_features_proj(self.query_features_block(query_features[0])) + return [torch.einsum("bchw,bnc->bnhw", sf, qf) + self.bias] diff --git a/library/src/getitune/recipe/detection/edgecrafter_l.yaml b/library/src/getitune/recipe/detection/edgecrafter_l.yaml new file mode 100644 index 00000000000..320afc32bc2 --- /dev/null +++ b/library/src/getitune/recipe/detection/edgecrafter_l.yaml @@ -0,0 +1,249 @@ +task: DETECTION +model: + class_path: getitune.backend.lightning.models.detection.EdgeCrafter + init_args: + model_name: edgecrafter_l + label_info: 80 + multi_scale: false + backbone_lr: 0.000005 + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + scheduler: + class_path: getitune.backend.lightning.schedulers.LinearWarmupSchedulerCallable + init_args: + num_warmup_steps: 30 + main_scheduler_callable: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + mode: max + factor: 0.5 + patience: 10 + monitor: val/map_50 + +engine: + device: auto + accumulate_grad_batches: 8 + +callback_monitor: val/map_50 + +callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 3 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + mode: max + patience: 10 + min_delta: 0.001 + warmup_iters: 50 + warmup_epochs: 10 + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + dirpath: "" + monitor: val/map_50 + mode: max + save_top_k: 1 + save_last: true + auto_insert_metric_name: false + filename: "checkpoints/epoch_{epoch:03d}" + - class_path: getitune.backend.lightning.callbacks.aug_scheduler.AugmentationSchedulerCallback + init_args: + data_aug_switch: + class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch + init_args: + policy_epochs: [4, 23] + input_size: [640, 640] + policies: + no_aug: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + light_aug: + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + strong_aug_1: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + random_pop: true + max_cached_images: 20 + img_scale: [640, 640] + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + alpha: 1.5 + p: 0.5 + random_pop: true + max_cached_images: 10 + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + strong_aug_2: + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + side_range: [1.0, 2.0] + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + alpha: 1.5 + p: 0.5 + random_pop: true + max_cached_images: 10 + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +data: ../_base_/data/detection.yaml +overrides: + gradient_clip_val: 0.1 + callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 3 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + warmup_iters: 100 + warmup_epochs: 7 + + data: + input_size: + - 640 + - 640 + task: DETECTION + train_subset: + batch_size: 4 + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + sampler: + class_path: getitune.data.samplers.balanced_sampler.BalancedSampler + + val_subset: + batch_size: 4 + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + resize_targets: false + augmentations_gpu: + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + test_subset: + batch_size: 4 + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + resize_targets: false + augmentations_gpu: + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] diff --git a/library/src/getitune/recipe/detection/edgecrafter_m.yaml b/library/src/getitune/recipe/detection/edgecrafter_m.yaml new file mode 100644 index 00000000000..c80544156f1 --- /dev/null +++ b/library/src/getitune/recipe/detection/edgecrafter_m.yaml @@ -0,0 +1,249 @@ +task: DETECTION +model: + class_path: getitune.backend.lightning.models.detection.EdgeCrafter + init_args: + model_name: edgecrafter_m + label_info: 80 + multi_scale: false + backbone_lr: 0.000025 + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + scheduler: + class_path: getitune.backend.lightning.schedulers.LinearWarmupSchedulerCallable + init_args: + num_warmup_steps: 30 + main_scheduler_callable: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + mode: max + factor: 0.5 + patience: 10 + monitor: val/map_50 + +engine: + device: auto + accumulate_grad_batches: 4 + +callback_monitor: val/map_50 + +callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 3 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + mode: max + patience: 10 + min_delta: 0.001 + warmup_iters: 50 + warmup_epochs: 10 + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + dirpath: "" + monitor: val/map_50 + mode: max + save_top_k: 1 + save_last: true + auto_insert_metric_name: false + filename: "checkpoints/epoch_{epoch:03d}" + - class_path: getitune.backend.lightning.callbacks.aug_scheduler.AugmentationSchedulerCallback + init_args: + data_aug_switch: + class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch + init_args: + policy_epochs: [4, 40] + input_size: [640, 640] + policies: + no_aug: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + light_aug: + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + strong_aug_1: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + random_pop: true + max_cached_images: 20 + img_scale: [640, 640] + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + alpha: 1.5 + p: 0.5 + random_pop: true + max_cached_images: 10 + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + strong_aug_2: + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + side_range: [1.0, 2.0] + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + alpha: 1.5 + p: 0.5 + random_pop: true + max_cached_images: 10 + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +data: ../_base_/data/detection.yaml +overrides: + gradient_clip_val: 0.1 + callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 3 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + warmup_iters: 100 + warmup_epochs: 7 + + data: + input_size: + - 640 + - 640 + task: DETECTION + train_subset: + batch_size: 8 + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + sampler: + class_path: getitune.data.samplers.balanced_sampler.BalancedSampler + + val_subset: + batch_size: 8 + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + resize_targets: false + augmentations_gpu: + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + test_subset: + batch_size: 8 + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + resize_targets: false + augmentations_gpu: + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] diff --git a/library/src/getitune/recipe/detection/edgecrafter_s.yaml b/library/src/getitune/recipe/detection/edgecrafter_s.yaml new file mode 100644 index 00000000000..6d75670f707 --- /dev/null +++ b/library/src/getitune/recipe/detection/edgecrafter_s.yaml @@ -0,0 +1,249 @@ +task: DETECTION +model: + class_path: getitune.backend.lightning.models.detection.EdgeCrafter + init_args: + model_name: edgecrafter_s + label_info: 80 + multi_scale: false + backbone_lr: 0.000025 + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + scheduler: + class_path: getitune.backend.lightning.schedulers.LinearWarmupSchedulerCallable + init_args: + num_warmup_steps: 30 + main_scheduler_callable: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + mode: max + factor: 0.5 + patience: 10 + monitor: val/map_50 + +engine: + device: auto + accumulate_grad_batches: 4 + +callback_monitor: val/map_50 + +callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 3 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + mode: max + patience: 10 + min_delta: 0.001 + warmup_iters: 50 + warmup_epochs: 10 + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + dirpath: "" + monitor: val/map_50 + mode: max + save_top_k: 1 + save_last: true + auto_insert_metric_name: false + filename: "checkpoints/epoch_{epoch:03d}" + - class_path: getitune.backend.lightning.callbacks.aug_scheduler.AugmentationSchedulerCallback + init_args: + data_aug_switch: + class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch + init_args: + policy_epochs: [4, 40] + input_size: [640, 640] + policies: + no_aug: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + light_aug: + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + strong_aug_1: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + random_pop: true + max_cached_images: 20 + img_scale: [640, 640] + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + alpha: 1.5 + p: 0.5 + random_pop: true + max_cached_images: 10 + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + strong_aug_2: + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + side_range: [1.0, 2.0] + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + alpha: 1.5 + p: 0.5 + random_pop: true + max_cached_images: 10 + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +data: ../_base_/data/detection.yaml +overrides: + gradient_clip_val: 0.1 + callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 3 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + warmup_iters: 100 + warmup_epochs: 7 + + data: + input_size: + - 640 + - 640 + task: DETECTION + train_subset: + batch_size: 8 + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + sampler: + class_path: getitune.data.samplers.balanced_sampler.BalancedSampler + + val_subset: + batch_size: 8 + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + resize_targets: false + augmentations_gpu: + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + test_subset: + batch_size: 8 + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + resize_targets: false + augmentations_gpu: + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] diff --git a/library/src/getitune/recipe/detection/edgecrafter_x.yaml b/library/src/getitune/recipe/detection/edgecrafter_x.yaml new file mode 100644 index 00000000000..929127c36c8 --- /dev/null +++ b/library/src/getitune/recipe/detection/edgecrafter_x.yaml @@ -0,0 +1,249 @@ +task: DETECTION +model: + class_path: getitune.backend.lightning.models.detection.EdgeCrafter + init_args: + model_name: edgecrafter_x + label_info: 80 + multi_scale: false + backbone_lr: 0.0000025 + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0005 + betas: [0.9, 0.999] + weight_decay: 0.0001 + + scheduler: + class_path: getitune.backend.lightning.schedulers.LinearWarmupSchedulerCallable + init_args: + num_warmup_steps: 30 + main_scheduler_callable: + class_path: lightning.pytorch.cli.ReduceLROnPlateau + init_args: + mode: max + factor: 0.5 + patience: 10 + monitor: val/map_50 + +engine: + device: auto + accumulate_grad_batches: 8 + +callback_monitor: val/map_50 + +callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 3 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + mode: max + patience: 10 + min_delta: 0.001 + warmup_iters: 50 + warmup_epochs: 10 + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + dirpath: "" + monitor: val/map_50 + mode: max + save_top_k: 1 + save_last: true + auto_insert_metric_name: false + filename: "checkpoints/epoch_{epoch:03d}" + - class_path: getitune.backend.lightning.callbacks.aug_scheduler.AugmentationSchedulerCallback + init_args: + data_aug_switch: + class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch + init_args: + policy_epochs: [4, 23] + input_size: [640, 640] + policies: + no_aug: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + light_aug: + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + strong_aug_1: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + random_pop: true + max_cached_images: 20 + img_scale: [640, 640] + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + alpha: 1.5 + p: 0.5 + random_pop: true + max_cached_images: 10 + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + strong_aug_2: + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + side_range: [1.0, 2.0] + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + alpha: 1.5 + p: 0.5 + random_pop: true + max_cached_images: 10 + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: [640, 640] + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.ColorJiggle + init_args: + brightness: 0.125 + contrast: 0.5 + saturation: 0.5 + hue: 0.05 + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + +data: ../_base_/data/detection.yaml +overrides: + gradient_clip_val: 0.1 + callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 3 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + warmup_iters: 100 + warmup_epochs: 7 + + data: + input_size: + - 640 + - 640 + task: DETECTION + train_subset: + batch_size: 4 + augmentations_cpu: + - class_path: torchvision.transforms.v2.RandomZoomOut + init_args: + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop + - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes + init_args: + min_size: 1 + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + augmentations_gpu: + - class_path: kornia.augmentation.RandomHorizontalFlip + init_args: + p: 0.5 + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + sampler: + class_path: getitune.data.samplers.balanced_sampler.BalancedSampler + + val_subset: + batch_size: 4 + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + resize_targets: false + augmentations_gpu: + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] + test_subset: + batch_size: 4 + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.Resize + init_args: + size: $(input_size) + keep_aspect_ratio: false + resize_targets: false + augmentations_gpu: + - class_path: kornia.augmentation.Normalize + init_args: + mean: [0.485, 0.456, 0.406] + std: [0.229, 0.224, 0.225] diff --git a/library/tests/unit/backend/lightning/models/detection/backbones/test_vit_tiny.py b/library/tests/unit/backend/lightning/models/detection/backbones/test_vit_tiny.py index 2c186f6c0ba..b59b95c4e40 100644 --- a/library/tests/unit/backend/lightning/models/detection/backbones/test_vit_tiny.py +++ b/library/tests/unit/backend/lightning/models/detection/backbones/test_vit_tiny.py @@ -16,6 +16,7 @@ drop_path, rotate_half, ) +from getitune.backend.lightning.models.modules.drop import DropPath as SharedDropPath class TestHelperFunctions: @@ -244,7 +245,11 @@ def test_init(self, block): def test_init_with_droppath(self, block_with_droppath): """Test block with drop path initialization.""" - assert isinstance(block_with_droppath.drop_path, DropPath) + # `Block` is now shared (common/layers/vit_blocks.py) and uses the canonical + # `modules.drop.DropPath` internally, not vit_tiny's own standalone `DropPath` + # (the two are behaviorally identical; `vit_tiny.DropPath` is kept for its own + # direct unit tests below and is no longer used internally by `Block`). + assert isinstance(block_with_droppath.drop_path, SharedDropPath) assert block_with_droppath.drop_path.drop_prob == 0.1 def test_forward_without_rope(self, block):