From 091656c24a77319a19ba64640de3831d95417858 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Tue, 30 Jun 2026 04:23:23 +0900 Subject: [PATCH 01/14] add edge crafter. First revision. --- .../execution/common/geti_config_converter.py | 40 + .../manifests/detection/edgecrafter_l.yaml | 51 + .../manifests/detection/edgecrafter_m.yaml | 51 + .../manifests/detection/edgecrafter_s.yaml | 51 + .../manifests/detection/edgecrafter_x.yaml | 51 + .../edgecrafter_inst_l.yaml | 56 + .../edgecrafter_inst_m.yaml | 56 + .../edgecrafter_inst_s.yaml | 56 + .../edgecrafter_inst_x.yaml | 56 + .../api/routers/test_model_architectures.py | 6 +- .../models/common/edgecrafter_mixin.py | 375 ++++++ .../lightning/models/detection/__init__.py | 3 +- .../models/detection/backbones/ecvit.py | 566 +++++++++ .../models/detection/detectors/__init__.py | 3 +- .../models/detection/detectors/edgecrafter.py | 152 +++ .../lightning/models/detection/edgecrafter.py | 145 +++ .../models/detection/heads/ec_decoder.py | 1022 +++++++++++++++++ .../models/detection/losses/__init__.py | 2 + .../models/detection/losses/ec_loss.py | 112 ++ .../detection/necks/dfine_hybrid_encoder.py | 37 + .../models/instance_segmentation/__init__.py | 3 +- .../instance_segmentation/edgecrafter_inst.py | 163 +++ .../recipe/detection/edgecrafter_l.yaml | 123 ++ .../recipe/detection/edgecrafter_m.yaml | 123 ++ .../recipe/detection/edgecrafter_s.yaml | 123 ++ .../recipe/detection/edgecrafter_x.yaml | 123 ++ .../edgecrafter_inst_l.yaml | 106 ++ .../edgecrafter_inst_m.yaml | 106 ++ .../edgecrafter_inst_s.yaml | 106 ++ .../edgecrafter_inst_x.yaml | 106 ++ 30 files changed, 3967 insertions(+), 6 deletions(-) create mode 100644 application/backend/app/supported_models/manifests/detection/edgecrafter_l.yaml create mode 100644 application/backend/app/supported_models/manifests/detection/edgecrafter_m.yaml create mode 100644 application/backend/app/supported_models/manifests/detection/edgecrafter_s.yaml create mode 100644 application/backend/app/supported_models/manifests/detection/edgecrafter_x.yaml create mode 100644 application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_l.yaml create mode 100644 application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_m.yaml create mode 100644 application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_s.yaml create mode 100644 application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_x.yaml create mode 100644 library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py create mode 100644 library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py create mode 100644 library/src/getitune/backend/lightning/models/detection/detectors/edgecrafter.py create mode 100644 library/src/getitune/backend/lightning/models/detection/edgecrafter.py create mode 100644 library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py create mode 100644 library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py create mode 100644 library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py create mode 100644 library/src/getitune/recipe/detection/edgecrafter_l.yaml create mode 100644 library/src/getitune/recipe/detection/edgecrafter_m.yaml create mode 100644 library/src/getitune/recipe/detection/edgecrafter_s.yaml create mode 100644 library/src/getitune/recipe/detection/edgecrafter_x.yaml create mode 100644 library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml create mode 100644 library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml create mode 100644 library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml create mode 100644 library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml diff --git a/application/backend/app/execution/common/geti_config_converter.py b/application/backend/app/execution/common/geti_config_converter.py index 2c56c5abf5e..ab7f64a2505 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, @@ -858,6 +878,26 @@ def convert(config: dict) -> dict: "status": ModelStatus.ACCURACY, "default": False, }, + "instance-segmentation-edgecrafter-s": { + "recipe_path": RECIPE_PATH / "instance_segmentation" / "edgecrafter_inst_s.yaml", + "status": ModelStatus.ACTIVE, + "default": False, + }, + "instance-segmentation-edgecrafter-m": { + "recipe_path": RECIPE_PATH / "instance_segmentation" / "edgecrafter_inst_m.yaml", + "status": ModelStatus.ACTIVE, + "default": False, + }, + "instance-segmentation-edgecrafter-l": { + "recipe_path": RECIPE_PATH / "instance_segmentation" / "edgecrafter_inst_l.yaml", + "status": ModelStatus.ACTIVE, + "default": False, + }, + "instance-segmentation-edgecrafter-x": { + "recipe_path": RECIPE_PATH / "instance_segmentation" / "edgecrafter_inst_x.yaml", + "status": ModelStatus.ACTIVE, + "default": False, + }, "instance-segmentation-yolo26-n": { "recipe_path": RECIPE_PATH / "instance_segmentation" / "yolo26_n_seg.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..cd4c03e7b0e --- /dev/null +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_l.yaml @@ -0,0 +1,51 @@ +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.0002 + weight_decay: 0.0001 + early_stopping: + patience: 15 + scheduler: + factor: 0.1 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 320 + - 448 + - 512 + - 576 + - 640 + - 704 + - 768 + - 896 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + random_zoom_out: + enable: true + fill: 0 + side_range: + - 1.0 + - 4.0 + probability: 0.5 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..481fe761530 --- /dev/null +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_m.yaml @@ -0,0 +1,51 @@ +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.0004 + weight_decay: 0.0001 + early_stopping: + patience: 15 + scheduler: + factor: 0.1 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 320 + - 448 + - 512 + - 576 + - 640 + - 704 + - 768 + - 896 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + random_zoom_out: + enable: true + fill: 0 + side_range: + - 1.0 + - 4.0 + probability: 0.5 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..da607737658 --- /dev/null +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_s.yaml @@ -0,0 +1,51 @@ +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.0004 + weight_decay: 0.0001 + early_stopping: + patience: 15 + scheduler: + factor: 0.1 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 320 + - 448 + - 512 + - 576 + - 640 + - 704 + - 768 + - 896 + input_size_width: 512 + input_size_height: 512 + dataset_preparation: + augmentation: + random_zoom_out: + enable: true + fill: 0 + side_range: + - 1.0 + - 4.0 + probability: 0.5 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..79508692adf --- /dev/null +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_x.yaml @@ -0,0 +1,51 @@ +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.0001 + weight_decay: 0.0001 + early_stopping: + patience: 15 + scheduler: + factor: 0.1 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 320 + - 448 + - 512 + - 576 + - 640 + - 704 + - 768 + - 896 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + random_zoom_out: + enable: true + fill: 0 + side_range: + - 1.0 + - 4.0 + probability: 0.5 diff --git a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_l.yaml b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_l.yaml new file mode 100644 index 00000000000..ff4054b33c1 --- /dev/null +++ b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_l.yaml @@ -0,0 +1,56 @@ +id: instance-segmentation-edgecrafter-l +name: ECSeg-L + +pretrained_weights: + url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_l.pth + mirror_url: https://huggingface.co/Intellindust/ECSeg_L/resolve/main/model.safetensors + sha_sum: 6fa7e8ffdd277feeef44bb1b8046315ae860023af9432bb75ecf400d74fe2396 + +description: "EdgeCrafter instance segmentation model (ECSeg-L) extends ECDet with a mask prediction head for high-accuracy real-time instance segmentation." + +stats: + gigaflops: 111.0 + trainable_parameters: 34.0 + benchmark_metrics: + coco_map_50_95: 47.1 + +capabilities: + tiling: false + +hyperparameters: + training: + learning_rate: 0.0002 + weight_decay: 0.0001 + early_stopping: + patience: 15 + scheduler: + factor: 0.1 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 320 + - 448 + - 512 + - 576 + - 640 + - 704 + - 768 + - 896 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + random_zoom_out: + enable: true + fill: 0 + side_range: + - 1.0 + - 4.0 + probability: 0.5 + iou_random_crop: + enable: true + probability: 1.0 + min_scale: 0.3 + max_scale: 1.0 diff --git a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_m.yaml b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_m.yaml new file mode 100644 index 00000000000..ebc74cbe6ee --- /dev/null +++ b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_m.yaml @@ -0,0 +1,56 @@ +id: instance-segmentation-edgecrafter-m +name: ECSeg-M + +pretrained_weights: + url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_m.pth + mirror_url: https://huggingface.co/Intellindust/ECSeg_M/resolve/main/model.safetensors + sha_sum: f1abcbd6ed747ed7364466a9fc4d3dc9ceb8d48cc123b353f717235c9c52069d + +description: "EdgeCrafter instance segmentation model (ECSeg-M) extends ECDet with a mask prediction head for high-accuracy real-time instance segmentation." + +stats: + gigaflops: 64.0 + trainable_parameters: 20.0 + benchmark_metrics: + coco_map_50_95: 45.2 + +capabilities: + tiling: false + +hyperparameters: + training: + learning_rate: 0.0004 + weight_decay: 0.0001 + early_stopping: + patience: 15 + scheduler: + factor: 0.1 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 320 + - 448 + - 512 + - 576 + - 640 + - 704 + - 768 + - 896 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + random_zoom_out: + enable: true + fill: 0 + side_range: + - 1.0 + - 4.0 + probability: 0.5 + iou_random_crop: + enable: true + probability: 1.0 + min_scale: 0.3 + max_scale: 1.0 diff --git a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_s.yaml b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_s.yaml new file mode 100644 index 00000000000..0d742763e31 --- /dev/null +++ b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_s.yaml @@ -0,0 +1,56 @@ +id: instance-segmentation-edgecrafter-s +name: ECSeg-S + +pretrained_weights: + url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_s.pth + mirror_url: https://huggingface.co/Intellindust/ECSeg_S/resolve/main/model.safetensors + sha_sum: b83b07d71b942c6255aaad1cf1f92eee075c5a548624d27464c987cd618d5a93 + +description: "EdgeCrafter instance segmentation model (ECSeg-S) extends ECDet with a mask prediction head for high-accuracy real-time instance segmentation." + +stats: + gigaflops: 33.0 + trainable_parameters: 10.0 + benchmark_metrics: + coco_map_50_95: 43.0 + +capabilities: + tiling: false + +hyperparameters: + training: + learning_rate: 0.0004 + weight_decay: 0.0001 + early_stopping: + patience: 15 + scheduler: + factor: 0.1 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 320 + - 448 + - 512 + - 576 + - 640 + - 704 + - 768 + - 896 + input_size_width: 512 + input_size_height: 512 + dataset_preparation: + augmentation: + random_zoom_out: + enable: true + fill: 0 + side_range: + - 1.0 + - 4.0 + probability: 0.5 + iou_random_crop: + enable: true + probability: 1.0 + min_scale: 0.3 + max_scale: 1.0 diff --git a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_x.yaml b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_x.yaml new file mode 100644 index 00000000000..4fa062f4915 --- /dev/null +++ b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_x.yaml @@ -0,0 +1,56 @@ +id: instance-segmentation-edgecrafter-x +name: ECSeg-X + +pretrained_weights: + url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_x.pth + mirror_url: https://huggingface.co/Intellindust/ECSeg_X/resolve/main/model.safetensors + sha_sum: 5e097fdc49cdf31dcc117f56b6504ee4afc9198378a02b3e7e2da69f1d4275a7 + +description: "EdgeCrafter instance segmentation model (ECSeg-X) extends ECDet with a mask prediction head for high-accuracy real-time instance segmentation." + +stats: + gigaflops: 168.0 + trainable_parameters: 50.0 + benchmark_metrics: + coco_map_50_95: 48.4 + +capabilities: + tiling: false + +hyperparameters: + training: + learning_rate: 0.0001 + weight_decay: 0.0001 + early_stopping: + patience: 15 + scheduler: + factor: 0.1 + patience: 10 + gradient_clip: + enable: true + max_grad_norm: 0.1 + allowed_values_input_size: + - 320 + - 448 + - 512 + - 576 + - 640 + - 704 + - 768 + - 896 + input_size_width: 640 + input_size_height: 640 + dataset_preparation: + augmentation: + random_zoom_out: + enable: true + fill: 0 + side_range: + - 1.0 + - 4.0 + probability: 0.5 + iou_random_crop: + enable: true + probability: 1.0 + min_scale: 0.3 + max_scale: 1.0 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 592d06a47e7..0f51ab427b6 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"]) == 20 + assert len(data["model_architectures"]) == 24 # Verify structure of first detection model detection_model = next( @@ -57,8 +57,8 @@ def test_get_all_model_architectures(self, fxt_client: TestClient): @pytest.mark.parametrize( "task_filter, total_models", [ - ("detection", 20), - ("instance_segmentation", 13), + ("detection", 24), + ("instance_segmentation", 17), ("classification", 6), ], ) 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..ddcd47c7563 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py @@ -0,0 +1,375 @@ +# 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). +Copyright (c) 2026 The EdgeCrafter Authors. All Rights Reserved. +""" + +from __future__ import annotations + +import copy +from typing import Any, ClassVar + +import torch +from torch import Tensor +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 + + +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. + * ``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. + + 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. + """ + + _pretrained_weights: ClassVar[dict[str, str]] + + # 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.0001, + }, + "edgecrafter_m": { + "backbone_name": "ecvittplus", + "seg_backbone_name": "ecseg_vittplus", + "proj_dim": None, + "backbone_lr": 0.0001, + }, + "edgecrafter_l": { + "backbone_name": "ecvits", + "seg_backbone_name": "ecseg_vits", + "proj_dim": 256, + "backbone_lr": 0.00005, + }, + "edgecrafter_x": { + "backbone_name": "ecvitsplus", + "seg_backbone_name": "ecseg_vitsplus", + "proj_dim": 256, + "backbone_lr": 0.00005, + }, + } + + def _build_ec_model( + self, + num_classes: int, + *, + with_seg: bool = False, + ) -> ECDETRDetector: + """Construct the full EdgeCrafter model for detection or instance segmentation. + + Steps: + 1. Build :class:`ECViTAdapter` backbone (det or seg weights variant). + 2. Build :class:`HybridEncoder` neck. + 3. Build :class:`ECTransformer` decoder (with seg head for ECSeg). + 4. Build :class:`ECCriterion` with mask losses added for ECSeg. + 5. Wrap everything in :class:`ECDETRDetector`. + 6. Load pretrained checkpoint via :func:`load_checkpoint`. + + Args: + num_classes: Number of target classes. + with_seg: When ``True``, builds the ECSeg variant (adds segmentation + head and mask losses). + + Returns: + Configured :class:`ECDETRDetector` instance. + """ + cfg = self._EC_MODEL_CFGS[self.model_name] # type: ignore[attr-defined] + backbone_key = "seg_backbone_name" if with_seg else "backbone_name" + + if self.data_input_params.input_size is None: # type: ignore[attr-defined] + msg = "input_size must not be None." + raise ValueError(msg) + input_size: tuple[int, int] = self.data_input_params.input_size # type: ignore[attr-defined] + + backbone = ECViTAdapter(model_name=cfg[backbone_key], proj_dim=cfg["proj_dim"]) + encoder = HybridEncoder(model_name=self.model_name) # type: ignore[attr-defined] + decoder = ECTransformer( + model_name=self.model_name, # type: ignore[attr-defined] + num_classes=num_classes, + eval_spatial_size=input_size, + mask_downsample_ratio=4 if with_seg else None, + ) + + weight_dict: dict[str, float] = { + "loss_vfl": 1.0, + "loss_bbox": 5.0, + "loss_giou": 2.0, + "loss_fgl": 0.15, + "loss_ddf": 1.5, + "loss_mal": 1.0, + } + if with_seg: + weight_dict["loss_mask_ce"] = 2.0 + weight_dict["loss_mask_dice"] = 5.0 + + criterion = ECCriterion( + weight_dict=weight_dict, + alpha=0.75, + gamma=1.5, + reg_max=32, + num_classes=num_classes, + ) + + backbone_lr: float = 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, # type: ignore[attr-defined] + input_size=input_size[0], + ) + model.init_weights() + load_checkpoint(model, self._pretrained_weights[self.model_name], map_location="cpu") # type: ignore[attr-defined] + return model + + def _customize_inputs( # pyrefly: ignore[bad-override] + self, + 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: # type: ignore[attr-defined] + return {"entity": entity} + + return { + "images": entity.images, + "targets": targets, + } + + def _customize_outputs( # pyrefly: ignore[bad-override] + self, + 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: # type: ignore[attr-defined] + 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] + result = self.model.postprocess(outputs, original_sizes) # type: ignore[attr-defined] + + 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, + ) -> 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 + self.model.optimizer_configuration, # type: ignore[attr-defined] + self.model, # type: ignore[attr-defined] + ) + optimizer = self.optimizer_callable(param_groups) # type: ignore[attr-defined] + schedulers = self.scheduler_callable(optimizer) # type: ignore[attr-defined] + + 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, + 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/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..0341ddc7f63 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py @@ -0,0 +1,566 @@ +# 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). +Copyright (c) 2026 The EdgeCrafter Authors. All Rights Reserved. +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 +import warnings +from functools import partial +from typing import ClassVar, Literal + +import numpy as np +import torch +import torch.nn.functional as F # noqa: N812 +from torch import nn + +from getitune.backend.lightning.models.detection.necks.dfine_hybrid_encoder import ConvNormLayerFusable + +__all__ = ["ECViTAdapter"] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _trunc_normal_(tensor: torch.Tensor, mean: float, std: float, a: float, b: float) -> torch.Tensor: + def _norm_cdf(x: float) -> float: + return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 + + if (mean < a - 2 * std) or (mean > b + 2 * std): + warnings.warn( + "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " + "The distribution of values may be incorrect.", + stacklevel=2, + ) + with torch.no_grad(): + lo = _norm_cdf((a - mean) / std) + hi = _norm_cdf((b - mean) / std) + tensor.uniform_(2 * lo - 1, 2 * hi - 1) + tensor.erfinv_() + tensor.mul_(std * math.sqrt(2.0)) + tensor.add_(mean) + tensor.clamp_(min=a, max=b) + return tensor + + +def trunc_normal_( + tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0 +) -> torch.Tensor: + """Fill tensor with truncated normal values.""" + return _trunc_normal_(tensor, mean, std, a, b) + + +def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor: + """Stochastic depth drop path.""" + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) + return x.div(keep_prob) * random_tensor.floor() + + +class DropPath(nn.Module): + """Drop paths (stochastic depth) per sample.""" + + def __init__(self, drop_prob: float = 0.0) -> None: + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass.""" + return drop_path(x, self.drop_prob, self.training) + + +# --------------------------------------------------------------------------- +# RoPE +# --------------------------------------------------------------------------- + + +class RopePositionEmbedding(nn.Module): + """2-D Rotary Position Embedding for ViT. + + Args: + embed_dim: Transformer embedding dimension. + num_heads: Number of attention heads. + base: RoPE base period (used when min/max_period are None). + min_period: Minimum period (used together with max_period). + max_period: Maximum period (used together with min_period). + normalize_coords: Coordinate normalisation strategy. + shift_coords: Optional coordinate shift jitter during training. + jitter_coords: Optional coordinate scale jitter during training. + rescale_coords: Optional global scale jitter during training. + """ + + def __init__( + self, + embed_dim: int, + *, + num_heads: int, + base: float | None = 100.0, + min_period: float | None = None, + max_period: float | None = None, + normalize_coords: Literal["min", "max", "separate"] = "separate", + shift_coords: float | None = None, + jitter_coords: float | None = None, + rescale_coords: float | None = None, + ) -> None: + super().__init__() + head_dim = embed_dim // num_heads + if head_dim % 4 != 0: + msg = "Head dimension must be divisible by 4 for 2D RoPE" + raise ValueError(msg) + both = min_period is not None and max_period is not None + if (base is None and not both) or (base is not None and both): + msg = "Either `base` or both `min_period`+`max_period` must be provided." + raise ValueError(msg) + + self.base = base + self.min_period = min_period + self.max_period = max_period + self.D_head = head_dim + self.normalize_coords = normalize_coords + self.shift_coords = shift_coords + self.jitter_coords = jitter_coords + self.rescale_coords = rescale_coords + self.register_buffer("periods", torch.empty(head_dim // 4), persistent=True) + self._init_weights() + + def forward(self, *, H: int, W: int) -> tuple[torch.Tensor, torch.Tensor]: # noqa: N803 + """Compute sin/cos RoPE tables for an HxW feature map.""" + device = self.periods.device # type: ignore[union-attr] + dtype = torch.get_default_dtype() + dd: dict = {"device": device, "dtype": dtype} + + if self.normalize_coords == "max": + m = max(H, W) + coords_h = torch.arange(0.5, H, **dd) / m + coords_w = torch.arange(0.5, W, **dd) / m + elif self.normalize_coords == "separate": + coords_h = torch.arange(0.5, H, **dd) / H + coords_w = torch.arange(0.5, W, **dd) / W + else: # min + m = min(H, W) + coords_h = torch.arange(0.5, H, **dd) / m + coords_w = torch.arange(0.5, W, **dd) / m + + coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1) + coords = coords.flatten(0, 1) + coords = 2.0 * coords - 1.0 + + if self.training and self.shift_coords is not None: + coords = coords + torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords)[None, :] + if self.training and self.jitter_coords is not None: + j = torch.empty(2, **dd).uniform_(-np.log(self.jitter_coords), np.log(self.jitter_coords)).exp() + coords = coords * j[None, :] + if self.training and self.rescale_coords is not None: + r = torch.empty(1, **dd).uniform_(-np.log(self.rescale_coords), np.log(self.rescale_coords)).exp() + coords = coords * r + + angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :] # type: ignore[index] + angles = angles.flatten(1, 2).repeat(1, 2) + return torch.sin(angles).unsqueeze(0).unsqueeze(0), torch.cos(angles).unsqueeze(0).unsqueeze(0) + + def _init_weights(self) -> None: + """Initialise period buffer.""" + device = self.periods.device # type: ignore[union-attr] + dtype = torch.get_default_dtype() + if self.base is not None: + periods = self.base ** (2 * torch.arange(self.D_head // 4, device=device, dtype=dtype) / (self.D_head // 2)) + else: + base_ratio = self.max_period / self.min_period # type: ignore[operator] + exponents = torch.linspace(0, 1, self.D_head // 4, device=device, dtype=dtype) + periods = self.max_period * (base_ratio ** (exponents - 1)) # type: ignore[operator] + self.periods.data.copy_(periods) # type: ignore[union-attr] + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def _apply_rope(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor: + return (x * cos) + (_rotate_half(x) * sin) + + +# --------------------------------------------------------------------------- +# 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 Attention(nn.Module): + """Multi-head self-attention with optional RoPE.""" + + 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 + 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: torch.Tensor, rope_sincos: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor: + """Forward pass.""" + B, N, C = x.shape # noqa: N806 + 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 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.""" + + def __init__( + self, + dim: int, + num_heads: int, + ffn_ratio: float = 4.0, + qkv_bias: bool = False, + drop: float = 0.0, + attn_drop: float = 0.0, + drop_path_rate: float = 0.0, + act_layer: type[nn.Module] = nn.GELU, + norm_layer: type[nn.Module] = nn.LayerNorm, + ) -> None: + super().__init__() + self.norm1 = norm_layer(dim) # type: ignore[call-arg] + self.attn = Attention(dim, num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) + self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() + self.norm2 = norm_layer(dim) # type: ignore[call-arg] + self.mlp = Mlp(dim, int(dim * ffn_ratio), act_layer=act_layer, drop=drop) + + def forward(self, x: torch.Tensor, rope_sincos: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.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))) + + +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] | 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, + ffn_ratio=ffn_ratio, + qkv_bias=qkv_bias, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path_rate=dpr[i], + act_layer=nn.GELU, + norm_layer=norm_layer, # type: ignore[arg-type] + ) + 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/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..7a2a1bc5581 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/detectors/edgecrafter.py @@ -0,0 +1,152 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EdgeCrafter DETR-style detector wrapper.""" + +from __future__ import annotations + +from typing import Any + +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: list[dict] | None = 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, 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 = 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(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..7cbccd79b55 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/edgecrafter.py @@ -0,0 +1,145 @@ +# 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 + +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 +from getitune.types.export import TaskLevelExportParameters + +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 paper / repository: + https://github.com/Intellindust-AI-Lab/EdgeCrafter + + 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``. + 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", + } + + 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, + torch_compile: bool = False, + tile_config: TileConfig = TileConfig(enable_tiler=False), + ) -> None: + self.multi_scale = multi_scale + 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, with_seg=False) + + @property + def _export_parameters(self) -> TaskLevelExportParameters: + """Export parameters for EdgeCrafter detection.""" + # Use a conservative NMS IoU threshold: DETR produces one-to-one predictions + # via Hungarian matching, so only suppress near-duplicate boxes. + return super()._export_parameters.wrap(iou_threshold=0.8, nms_execute=True) + + @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=(512, 512), 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..d841a5860f1 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py @@ -0,0 +1,1022 @@ +# 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). +Copyright (c) 2026 The EdgeCrafter Authors. All Rights Reserved. +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 Any, ClassVar + +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 ( + 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 + +__all__ = ["ECTransformer"] + +# --------------------------------------------------------------------------- +# Per-model hyper-parameters +# --------------------------------------------------------------------------- + +_MODEL_CFGS: ClassVar[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], + }, +} + +# --------------------------------------------------------------------------- +# Sub-modules +# --------------------------------------------------------------------------- + + +class MLP(nn.Module): + """Simple multi-layer perceptron.""" + + def __init__( + self, + input_dim: int, + hidden_dim: int, + output_dim: int, + num_layers: int = 3, + act: str = "silu", + ) -> None: + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim, *h], [*h, output_dim])) + self._act = nn.SiLU() if act == "silu" else nn.ReLU() + + def forward(self, x: Tensor) -> Tensor: + """Forward pass.""" + for i, layer in enumerate(self.layers): + x = self._act(layer(x)) if i < self.num_layers - 1 else layer(x) + return x + + +class Gate(nn.Module): + """Gated cross-attention fusion (LayerNorm variant).""" + + def __init__(self, d_model: int) -> None: + super().__init__() + self.gate = nn.Linear(2 * d_model, 2 * d_model) + bias = bias_init_with_prob(0.5) + init.constant_(self.gate.bias, bias) + init.constant_(self.gate.weight, 0) + self.norm = nn.LayerNorm(d_model) + + def forward(self, x1: Tensor, x2: Tensor) -> Tensor: + """Forward pass.""" + gates = torch.sigmoid(self.gate(torch.cat([x1, x2], dim=-1))) + g1, g2 = gates.chunk(2, dim=-1) + return self.norm(g1 * x1 + g2 * x2) + + +class Integral(nn.Module): + """Distribution-to-distance integral (non-uniform weighting).""" + + def __init__(self, reg_max: int = 32) -> None: + super().__init__() + self.reg_max = reg_max + + def forward(self, x: Tensor, project: Tensor) -> Tensor: + """Forward pass.""" + shape = x.shape + x = F.softmax(x.reshape(-1, self.reg_max + 1), dim=1) + x = F.linear(x, project.to(x.device)).reshape(-1, 4) + return x.reshape([*list(shape[:-1]), -1]) + + +class LQE(nn.Module): + """Localization Quality Estimator.""" + + def __init__(self, k: int, hidden_dim: int, num_layers: int, reg_max: int, act: str = "silu") -> None: + super().__init__() + self.k = k + self.reg_max = reg_max + self.reg_conf = MLP(4 * (k + 1), hidden_dim, 1, num_layers, act=act) + init.constant_(self.reg_conf.layers[-1].bias, 0) + init.constant_(self.reg_conf.layers[-1].weight, 0) + + def forward(self, scores: Tensor, pred_corners: Tensor) -> Tensor: + """Forward pass.""" + B, L, _ = pred_corners.size() # noqa: N806 + prob = F.softmax(pred_corners.reshape(B, L, 4, self.reg_max + 1), dim=-1) + prob_topk, _ = prob.topk(self.k, dim=-1) + stat = torch.cat([prob_topk, prob_topk.mean(dim=-1, keepdim=True)], dim=-1) + quality_score = self.reg_conf(stat.reshape(B, L, -1)) + return scores + quality_score + + +# --------------------------------------------------------------------------- +# Segmentation Head +# --------------------------------------------------------------------------- + + +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) + self.fc1 = nn.Linear(dim, dim * 4) + self.act = nn.GELU() + self.fc2 = nn.Linear(dim * 4, dim) + + def forward(self, x: Tensor) -> Tensor: + """Forward pass.""" + return x + self.fc2(self.act(self.fc1(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)]) + self.spatial_features_proj = nn.Identity() + self.query_features_block = _MLPBlock(in_dim) + self.query_features_proj = nn.Identity() + 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) + qf = self.query_features_proj(self.query_features_block(query_features[0])) + return [torch.einsum("bchw,bnc->bnhw", sf, qf) + self.bias] + + +# --------------------------------------------------------------------------- +# 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 name. + """ + + 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: str = "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 + _act = nn.SiLU if activation == "silu" else nn.ReLU + self.linear1 = nn.Linear(d_model, dim_feedforward) + self.activation = _act() + 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: str = "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, act=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) + + 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 (``"silu"`` or ``"relu"``). + 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] = 4, + dropout: float = 0.0, + activation: str = "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: + seg_head = SegmentationHead( + in_dim=hidden_dim, + num_blocks=num_layers, + downsample_ratio=mask_downsample_ratio, + image_size=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, act=activation) + self.query_pos_head = MLP(4, hidden_dim, hidden_dim, 3, act=activation) + self.pre_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, act=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, act=activation) + wide_bbox_head = MLP(scaled_dim, scaled_dim, 4 * (reg_max + 1), 3, act=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: + self.input_proj.append( + nn.Sequential( + OrderedDict( + [ + ("conv", nn.Conv2d(in_ch, self.hidden_dim, 1, bias=False)), + ("norm", nn.BatchNorm2d(self.hidden_dim)), + ] + ) + ) + ) + for _ in range(self.num_levels - len(feat_channels)): + in_ch = feat_channels[-1] + self.input_proj.append( + nn.Sequential( + OrderedDict( + [ + ("conv", nn.Conv2d(in_ch, self.hidden_dim, 3, 2, padding=1, bias=False)), + ("norm", nn.BatchNorm2d(self.hidden_dim)), + ] + ) + ) + ) + + def _reset_parameters(self, feat_channels: list[int]) -> None: + bias = bias_init_with_prob(0.01) + init.constant_(self.enc_score_head.bias, bias) + init.constant_(self.enc_bbox_head.layers[-1].weight, 0) + init.constant_(self.enc_bbox_head.layers[-1].bias, 0) + init.constant_(self.pre_bbox_head.layers[-1].weight, 0) + init.constant_(self.pre_bbox_head.layers[-1].bias, 0) + for cls_, reg_ in zip(self.dec_score_head, self.dec_bbox_head): + init.constant_(cls_.bias, bias) + if hasattr(reg_, "layers"): + init.constant_(reg_.layers[-1].weight, 0) + init.constant_(reg_.layers[-1].bias, 0) + init.xavier_uniform_(self.query_pos_head.layers[0].weight) + init.xavier_uniform_(self.query_pos_head.layers[1].weight) + init.xavier_uniform_(self.query_pos_head.layers[-1].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 = self.anchors # type: ignore[attr-defined] + valid_mask = 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: list[int] | None) -> tuple[Tensor | None, Tensor | None]: + if x is None or s_idx is None: + return None, x + return torch.split(x, s_idx, 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], + "pred_masks": last_masks, + "ref_points": out_refs[-1], + "up": self.up, + "reg_scale": self.reg_scale, + } + else: + out = { + "pred_logits": out_logits[-1], + "pred_boxes": out_bboxes[-1], + "pred_masks": last_masks, + } + + if self.training and self.aux_loss: + aux_masks = 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) + out["pre_outputs"] = { + "pred_logits": pre_logits, + "pred_boxes": pre_bboxes, + "pred_masks": pred_segs, + } + out["enc_meta"] = {"class_agnostic": False} + + if dn_meta is not None: + dn_masks = 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, + ) + out["dn_pre_outputs"] = { + "pred_logits": dn_pre_logits, + "pred_boxes": dn_pre_bboxes, + "pred_masks": dn_pre_segs, + } + 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") + bbox_pred = bbox_pred * orig_target_sizes.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/ec_loss.py b/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py new file mode 100644 index 00000000000..e571c5b7c37 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py @@ -0,0 +1,112 @@ +# Copyright (C) 2026 Intel Corporation +# SPDX-License-Identifier: Apache-2.0 + +"""EdgeCrafter criterion implementation. + +Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). +Copyright (c) 2026 The EdgeCrafter Authors. All Rights Reserved. +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 .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. + """ + + 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_vfl, + 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..5439adb9f39 --- /dev/null +++ b/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py @@ -0,0 +1,163 @@ +# 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 paper / repository: + https://github.com/Intellindust-AI-Lab/EdgeCrafter + + 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``. + 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", + } + + 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, + torch_compile: bool = False, + tile_config: TileConfig = TileConfig(enable_tiler=False), + ) -> None: + self.multi_scale = multi_scale + 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, with_seg=True) + + @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) # type: ignore[arg-type] + + @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=(512, 512), 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/recipe/detection/edgecrafter_l.yaml b/library/src/getitune/recipe/detection/edgecrafter_l.yaml new file mode 100644 index 00000000000..29f9d5017ec --- /dev/null +++ b/library/src/getitune/recipe/detection/edgecrafter_l.yaml @@ -0,0 +1,123 @@ +task: DETECTION +model: + class_path: getitune.backend.lightning.models.detection.EdgeCrafter + init_args: + model_name: edgecrafter_l + label_info: 80 + multi_scale: false + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0002 + 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 + +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}" + +data: ../_base_/data/detection.yaml +overrides: + 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..bd808271f2a --- /dev/null +++ b/library/src/getitune/recipe/detection/edgecrafter_m.yaml @@ -0,0 +1,123 @@ +task: DETECTION +model: + class_path: getitune.backend.lightning.models.detection.EdgeCrafter + init_args: + model_name: edgecrafter_m + label_info: 80 + multi_scale: false + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0004 + 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 + +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}" + +data: ../_base_/data/detection.yaml +overrides: + 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..770ac927ff0 --- /dev/null +++ b/library/src/getitune/recipe/detection/edgecrafter_s.yaml @@ -0,0 +1,123 @@ +task: DETECTION +model: + class_path: getitune.backend.lightning.models.detection.EdgeCrafter + init_args: + model_name: edgecrafter_s + label_info: 80 + multi_scale: false + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0004 + 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 + +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}" + +data: ../_base_/data/detection.yaml +overrides: + 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: + - 512 + - 512 + 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..d22ee1051c9 --- /dev/null +++ b/library/src/getitune/recipe/detection/edgecrafter_x.yaml @@ -0,0 +1,123 @@ +task: DETECTION +model: + class_path: getitune.backend.lightning.models.detection.EdgeCrafter + init_args: + model_name: edgecrafter_x + label_info: 80 + multi_scale: false + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0001 + 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 + +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}" + +data: ../_base_/data/detection.yaml +overrides: + 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/instance_segmentation/edgecrafter_inst_l.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml new file mode 100644 index 00000000000..c0b506120c0 --- /dev/null +++ b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml @@ -0,0 +1,106 @@ +task: INSTANCE_SEGMENTATION + +model: + class_path: getitune.backend.lightning.models.instance_segmentation.EdgeCrafterInst + init_args: + model_name: edgecrafter_l + label_info: 80 + multi_scale: false + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0002 + 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 + +data: ../_base_/data/instance_segmentation.yaml + +callback_monitor: val/map_50 + +callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 5 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + mode: max + patience: 15 + check_on_train_epoch_end: false + monitor: val/map_50 + min_delta: 0.001 + warmup_iters: 100 + 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: lightning.pytorch.callbacks.EMAWeightAveraging + init_args: + decay: 0.993 + +overrides: + max_epochs: 100 + gradient_clip_val: 0.1 + data: + task: INSTANCE_SEGMENTATION + input_size: + - 640 + - 640 + 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 + 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 + + 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 diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml new file mode 100644 index 00000000000..c5a50193bb6 --- /dev/null +++ b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml @@ -0,0 +1,106 @@ +task: INSTANCE_SEGMENTATION + +model: + class_path: getitune.backend.lightning.models.instance_segmentation.EdgeCrafterInst + init_args: + model_name: edgecrafter_m + label_info: 80 + multi_scale: false + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0004 + 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 + +data: ../_base_/data/instance_segmentation.yaml + +callback_monitor: val/map_50 + +callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 5 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + mode: max + patience: 15 + check_on_train_epoch_end: false + monitor: val/map_50 + min_delta: 0.001 + warmup_iters: 100 + 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: lightning.pytorch.callbacks.EMAWeightAveraging + init_args: + decay: 0.993 + +overrides: + max_epochs: 100 + gradient_clip_val: 0.1 + data: + task: INSTANCE_SEGMENTATION + input_size: + - 640 + - 640 + 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 + 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 + + 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 diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml new file mode 100644 index 00000000000..e9d792e353c --- /dev/null +++ b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml @@ -0,0 +1,106 @@ +task: INSTANCE_SEGMENTATION + +model: + class_path: getitune.backend.lightning.models.instance_segmentation.EdgeCrafterInst + init_args: + model_name: edgecrafter_s + label_info: 80 + multi_scale: false + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0004 + 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 + +data: ../_base_/data/instance_segmentation.yaml + +callback_monitor: val/map_50 + +callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 5 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + mode: max + patience: 15 + check_on_train_epoch_end: false + monitor: val/map_50 + min_delta: 0.001 + warmup_iters: 100 + 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: lightning.pytorch.callbacks.EMAWeightAveraging + init_args: + decay: 0.993 + +overrides: + max_epochs: 100 + gradient_clip_val: 0.1 + data: + task: INSTANCE_SEGMENTATION + input_size: + - 512 + - 512 + 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 + 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 + + 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 diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml new file mode 100644 index 00000000000..3640f6a24a4 --- /dev/null +++ b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml @@ -0,0 +1,106 @@ +task: INSTANCE_SEGMENTATION + +model: + class_path: getitune.backend.lightning.models.instance_segmentation.EdgeCrafterInst + init_args: + model_name: edgecrafter_x + label_info: 80 + multi_scale: false + + optimizer: + class_path: torch.optim.AdamW + init_args: + lr: 0.0001 + 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 + +data: ../_base_/data/instance_segmentation.yaml + +callback_monitor: val/map_50 + +callbacks: + - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling + init_args: + max_interval: 1 + min_lrschedule_patience: 5 + - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup + init_args: + mode: max + patience: 15 + check_on_train_epoch_end: false + monitor: val/map_50 + min_delta: 0.001 + warmup_iters: 100 + 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: lightning.pytorch.callbacks.EMAWeightAveraging + init_args: + decay: 0.993 + +overrides: + max_epochs: 100 + gradient_clip_val: 0.1 + data: + task: INSTANCE_SEGMENTATION + input_size: + - 640 + - 640 + 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 + 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 + + 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 From 231053e1d6b5f095e3c04e58a710f6d1d328b81f Mon Sep 17 00:00:00 2001 From: kprokofi Date: Tue, 30 Jun 2026 16:16:32 +0900 Subject: [PATCH 02/14] fix linter --- .../models/detection/backbones/ecvit.py | 6 +- .../models/detection/detectors/edgecrafter.py | 17 ++-- .../models/detection/heads/ec_decoder.py | 89 ++++++++++--------- 3 files changed, 63 insertions(+), 49 deletions(-) diff --git a/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py b/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py index 0341ddc7f63..cf4a58845c4 100644 --- a/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py +++ b/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py @@ -14,7 +14,7 @@ import math import warnings from functools import partial -from typing import ClassVar, Literal +from typing import Callable, ClassVar, Literal import numpy as np import torch @@ -172,7 +172,7 @@ def forward(self, *, H: int, W: int) -> tuple[torch.Tensor, torch.Tensor]: # no def _init_weights(self) -> None: """Initialise period buffer.""" - device = self.periods.device # type: ignore[union-attr] + device: torch.device = self.periods.device # type: ignore[union-attr] dtype = torch.get_default_dtype() if self.base is not None: periods = self.base ** (2 * torch.arange(self.D_head // 4, device=device, dtype=dtype) / (self.D_head // 2)) @@ -373,7 +373,7 @@ def __init__( drop_path_rate: float = 0.0, patch_size: int = 16, return_layers: list[int] | None = None, - norm_layer: type[nn.Module] | None = None, + norm_layer: type[nn.Module] | Callable[..., nn.Module] | None = None, ) -> None: super().__init__() if return_layers is None: diff --git a/library/src/getitune/backend/lightning/models/detection/detectors/edgecrafter.py b/library/src/getitune/backend/lightning/models/detection/detectors/edgecrafter.py index 7a2a1bc5581..657437d45b1 100644 --- a/library/src/getitune/backend/lightning/models/detection/detectors/edgecrafter.py +++ b/library/src/getitune/backend/lightning/models/detection/detectors/edgecrafter.py @@ -5,7 +5,10 @@ from __future__ import annotations -from typing import Any +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 @@ -44,7 +47,7 @@ class ECDETRDetector(DETR): def _forward_features( self, images: Tensor, - targets: list[dict] | None = None, + targets: object = None, ) -> dict[str, Any]: """Backbone → encoder → decoder forward with spatial feature propagation. @@ -62,7 +65,7 @@ def _forward_features( """ backbone_feats = self.backbone(images) encoder_feats = self.encoder(backbone_feats) - return self.decoder(encoder_feats, targets, spatial_feat=backbone_feats[0]) + return self.decoder(encoder_feats, cast("list[dict] | None", targets), spatial_feat=backbone_feats[0]) def postprocess( self, @@ -92,7 +95,7 @@ def postprocess( device = outputs["pred_logits"].device sizes_tensor = torch.tensor(original_sizes, device=device, dtype=torch.float32) # [B, 2] as (H, W) - results = self.decoder.postprocess(outputs, sizes_tensor, self.num_top_queries) + results = cast("ECTransformer", self.decoder).postprocess(outputs, sizes_tensor, self.num_top_queries) has_masks = bool(results) and "masks" in results[0] @@ -112,7 +115,11 @@ def postprocess( for res, orig_size in zip(results, original_sizes): scores_list.append(res["scores"]) - boxes_list.append(BoundingBoxes(res["boxes"], format="xyxy", canvas_size=orig_size)) + 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"]) 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 index d841a5860f1..8aaa5ddd8f0 100644 --- a/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py +++ b/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py @@ -15,7 +15,7 @@ import copy from collections import OrderedDict -from typing import Any, ClassVar +from typing import Any, cast import torch import torch.nn.functional as F # noqa: N812 @@ -36,7 +36,7 @@ # Per-model hyper-parameters # --------------------------------------------------------------------------- -_MODEL_CFGS: ClassVar[dict[str, dict[str, Any]]] = { +_MODEL_CFGS: dict[str, dict[str, Any]] = { "edgecrafter_s": { "hidden_dim": 192, "num_heads": 8, @@ -137,8 +137,9 @@ def __init__(self, k: int, hidden_dim: int, num_layers: int, reg_max: int, act: self.k = k self.reg_max = reg_max self.reg_conf = MLP(4 * (k + 1), hidden_dim, 1, num_layers, act=act) - init.constant_(self.reg_conf.layers[-1].bias, 0) - init.constant_(self.reg_conf.layers[-1].weight, 0) + _reg_last = cast("nn.Linear", self.reg_conf.layers[-1]) + init.constant_(_reg_last.bias, 0) + init.constant_(_reg_last.weight, 0) def forward(self, scores: Tensor, pred_corners: Tensor) -> Tensor: """Forward pass.""" @@ -437,6 +438,10 @@ def forward( 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) @@ -603,7 +608,7 @@ def __init__( in_dim=hidden_dim, num_blocks=num_layers, downsample_ratio=mask_downsample_ratio, - image_size=self.eval_spatial_size, + image_size=cast("tuple[int, int]", self.eval_spatial_size), ) self.decoder = ECTransformerDecoder( @@ -671,44 +676,41 @@ def _build_input_proj_layer(self, feat_channels: list[int]) -> None: if in_ch == self.hidden_dim: self.input_proj.append(nn.Identity()) else: - self.input_proj.append( - nn.Sequential( - OrderedDict( - [ - ("conv", nn.Conv2d(in_ch, self.hidden_dim, 1, bias=False)), - ("norm", nn.BatchNorm2d(self.hidden_dim)), - ] - ) - ) - ) + _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] - self.input_proj.append( - nn.Sequential( - OrderedDict( - [ - ("conv", nn.Conv2d(in_ch, self.hidden_dim, 3, 2, padding=1, bias=False)), - ("norm", nn.BatchNorm2d(self.hidden_dim)), - ] - ) - ) - ) + _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) - init.constant_(self.enc_bbox_head.layers[-1].weight, 0) - init.constant_(self.enc_bbox_head.layers[-1].bias, 0) - init.constant_(self.pre_bbox_head.layers[-1].weight, 0) - init.constant_(self.pre_bbox_head.layers[-1].bias, 0) + _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_(cls_.bias, bias) + init.constant_(cast("nn.Linear", cls_).bias, bias) if hasattr(reg_, "layers"): - init.constant_(reg_.layers[-1].weight, 0) - init.constant_(reg_.layers[-1].bias, 0) - init.xavier_uniform_(self.query_pos_head.layers[0].weight) - init.xavier_uniform_(self.query_pos_head.layers[1].weight) - init.xavier_uniform_(self.query_pos_head.layers[-1].weight) + _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, @@ -781,8 +783,8 @@ def _get_decoder_input( if self.training or not hasattr(self, "anchors"): anchors, valid_mask = self._generate_anchors(spatial_shapes, device=memory.device) else: - anchors = self.anchors # type: ignore[attr-defined] - valid_mask = self.valid_mask # type: ignore[attr-defined] + 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) @@ -811,10 +813,11 @@ def _get_decoder_input( return content, enc_topk_bbox_unact, enc_bboxes_list, enc_logits_list @staticmethod - def _split(x: Tensor | None, dim: int, s_idx: list[int] | None) -> tuple[Tensor | None, Tensor | None]: + 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 - return torch.split(x, s_idx, dim=dim) # type: ignore[return-value] + 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, @@ -904,7 +907,9 @@ def forward( } if self.training and self.aux_loss: - aux_masks = list(out_masks[:-1]) if out_masks is not None else [None] * (len(out_logits) - 1) + 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], @@ -923,7 +928,9 @@ def forward( out["enc_meta"] = {"class_agnostic": False} if dn_meta is not None: - dn_masks = list(dn_out_masks[:-1]) if dn_out_masks is not None else [None] * (self.num_layers - 1) + 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, From 25a91b1dffb467e34ad78a807fc53d534cb3a423 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Tue, 30 Jun 2026 22:57:42 +0900 Subject: [PATCH 03/14] fix weights and bbox scale --- .../models/detection/heads/ec_decoder.py | 38 ++++++++++++------- 1 file changed, 25 insertions(+), 13 deletions(-) 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 index 8aaa5ddd8f0..20b36929242 100644 --- a/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py +++ b/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py @@ -177,13 +177,16 @@ class _MLPBlock(nn.Module): def __init__(self, dim: int) -> None: super().__init__() self.norm_in = nn.LayerNorm(dim) - self.fc1 = nn.Linear(dim, dim * 4) - self.act = nn.GELU() - self.fc2 = nn.Linear(dim * 4, 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.fc2(self.act(self.fc1(self.norm_in(x)))) + return x + self.layers(self.norm_in(x)) class SegmentationHead(nn.Module): @@ -207,9 +210,11 @@ def __init__( self.downsample_ratio = downsample_ratio self.image_size = tuple(image_size) self.blocks = nn.ModuleList([_DepthwiseConvBlock(in_dim) for _ in range(num_blocks)]) - self.spatial_features_proj = nn.Identity() + # 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) - self.query_features_proj = nn.Identity() + # 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( @@ -894,17 +899,19 @@ def forward( "pred_logits": out_logits[-1], "pred_boxes": out_bboxes[-1], "pred_corners": out_corners[-1], - "pred_masks": last_masks, "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], - "pred_masks": last_masks, } + if last_masks is not None: + out["pred_masks"] = last_masks if self.training and self.aux_loss: aux_masks: list[Tensor | None] = ( @@ -920,11 +927,13 @@ def forward( out_logits[-1], ) out["enc_aux_outputs"] = self._set_aux_loss(enc_topk_logits_list, enc_topk_bboxes_list) - out["pre_outputs"] = { + pre_out: dict[str, Any] = { "pred_logits": pre_logits, "pred_boxes": pre_bboxes, - "pred_masks": pred_segs, } + 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: @@ -940,11 +949,13 @@ def forward( 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, ) - out["dn_pre_outputs"] = { + dn_pre_out: dict[str, Any] = { "pred_logits": dn_pre_logits, "pred_boxes": dn_pre_bboxes, - "pred_masks": dn_pre_segs, } + 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 @@ -1001,7 +1012,8 @@ def postprocess( mask_pred = outputs.get("pred_masks") bbox_pred = torchvision.ops.box_convert(boxes, in_fmt="cxcywh", out_fmt="xyxy") - bbox_pred = bbox_pred * orig_target_sizes.repeat(1, 2).unsqueeze(1) + # 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) From b86ce36daac181537588666ccadd805e2eb96854 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Tue, 30 Jun 2026 22:58:35 +0900 Subject: [PATCH 04/14] change header --- .../backend/lightning/models/common/edgecrafter_mixin.py | 1 - .../backend/lightning/models/detection/backbones/ecvit.py | 1 - .../backend/lightning/models/detection/heads/ec_decoder.py | 1 - .../backend/lightning/models/detection/losses/ec_loss.py | 1 - 4 files changed, 4 deletions(-) diff --git a/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py index ddcd47c7563..ce13608230a 100644 --- a/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py +++ b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py @@ -4,7 +4,6 @@ """EdgeCrafter Mixin providing shared logic for Detection and Instance Segmentation models. Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). -Copyright (c) 2026 The EdgeCrafter Authors. All Rights Reserved. """ from __future__ import annotations diff --git a/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py b/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py index cf4a58845c4..4ac93a6433f 100644 --- a/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py +++ b/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py @@ -4,7 +4,6 @@ """EC-ViT backbone (ViTAdapter) for EdgeCrafter models. Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). -Copyright (c) 2026 The EdgeCrafter Authors. All Rights Reserved. Modified from DINOv3 (https://github.com/facebookresearch/dinov3). Modified from https://huggingface.co/spaces/Hila/RobustViT/blob/main/ViT/ViT_new.py """ 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 index 8aaa5ddd8f0..a2b9e3c5cf2 100644 --- a/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py +++ b/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py @@ -4,7 +4,6 @@ """EdgeCrafter Transformer Decoder (ECTransformer). Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). -Copyright (c) 2026 The EdgeCrafter Authors. All Rights Reserved. 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). 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 index e571c5b7c37..f29421a262c 100644 --- a/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py +++ b/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py @@ -4,7 +4,6 @@ """EdgeCrafter criterion implementation. Modified from EdgeCrafter (https://github.com/Intellindust-AI-Lab/EdgeCrafter). -Copyright (c) 2026 The EdgeCrafter Authors. All Rights Reserved. Modified from D-FINE (https://github.com/Peterande/D-FINE). Copyright (c) 2024 D-FINE Authors. All Rights Reserved. """ From fb0f880e6642c2803526e7556b995bfbb3d048a3 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Wed, 1 Jul 2026 19:03:36 +0900 Subject: [PATCH 05/14] fix OV inference and update hyperparameters --- .../models/common/edgecrafter_mixin.py | 49 ++++-- .../lightning/models/detection/edgecrafter.py | 30 ++-- .../models/detection/losses/ec_loss.py | 17 +++ .../instance_segmentation/edgecrafter_inst.py | 6 +- .../recipe/detection/edgecrafter_l.yaml | 134 ++++++++++++++++- .../recipe/detection/edgecrafter_m.yaml | 134 ++++++++++++++++- .../recipe/detection/edgecrafter_s.yaml | 138 ++++++++++++++++- .../recipe/detection/edgecrafter_x.yaml | 134 ++++++++++++++++- .../edgecrafter_inst_l.yaml | 135 ++++++++++++++++- .../edgecrafter_inst_m.yaml | 135 ++++++++++++++++- .../edgecrafter_inst_s.yaml | 139 +++++++++++++++++- .../edgecrafter_inst_x.yaml | 135 ++++++++++++++++- 12 files changed, 1133 insertions(+), 53 deletions(-) diff --git a/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py index ce13608230a..993c88edc30 100644 --- a/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py +++ b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py @@ -52,25 +52,25 @@ class EdgeCrafterMixin: "backbone_name": "ecvitt", "seg_backbone_name": "ecseg_vitt", "proj_dim": None, - "backbone_lr": 0.0001, + "backbone_lr": 0.000025, }, "edgecrafter_m": { "backbone_name": "ecvittplus", "seg_backbone_name": "ecseg_vittplus", "proj_dim": None, - "backbone_lr": 0.0001, + "backbone_lr": 0.000025, }, "edgecrafter_l": { "backbone_name": "ecvits", "seg_backbone_name": "ecseg_vits", "proj_dim": 256, - "backbone_lr": 0.00005, + "backbone_lr": 0.000005, }, "edgecrafter_x": { "backbone_name": "ecvitsplus", "seg_backbone_name": "ecseg_vitsplus", "proj_dim": 256, - "backbone_lr": 0.00005, + "backbone_lr": 0.0000025, }, } @@ -79,6 +79,7 @@ def _build_ec_model( num_classes: int, *, with_seg: bool = False, + backbone_lr: float | None = None, ) -> ECDETRDetector: """Construct the full EdgeCrafter model for detection or instance segmentation. @@ -94,6 +95,8 @@ def _build_ec_model( num_classes: Number of target classes. with_seg: When ``True``, builds the ECSeg variant (adds segmentation head and mask losses). + backbone_lr: Optional override for the backbone learning rate. + Defaults to the per-variant value in ``_EC_MODEL_CFGS``. Returns: Configured :class:`ECDETRDetector` instance. @@ -115,17 +118,32 @@ def _build_ec_model( mask_downsample_ratio=4 if with_seg else None, ) - weight_dict: dict[str, float] = { - "loss_vfl": 1.0, - "loss_bbox": 5.0, - "loss_giou": 2.0, - "loss_fgl": 0.15, - "loss_ddf": 1.5, - "loss_mal": 1.0, - } if with_seg: - weight_dict["loss_mask_ce"] = 2.0 - weight_dict["loss_mask_dice"] = 5.0 + weight_dict: 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: dict[str, int | float] | None = { + "cost_class": 2, + "cost_bbox": 1, + "cost_giou": 1, + "cost_mask": 5, + "cost_dice": 5, + } + else: + weight_dict = { + "loss_mal": 1.0, + "loss_bbox": 5.0, + "loss_giou": 2.0, + "loss_fgl": 0.15, + "loss_ddf": 1.5, + } + matcher_cost_dict = None criterion = ECCriterion( weight_dict=weight_dict, @@ -133,9 +151,10 @@ def _build_ec_model( gamma=1.5, reg_max=32, num_classes=num_classes, + matcher_cost_dict=matcher_cost_dict, ) - backbone_lr: float = cfg["backbone_lr"] + 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}, diff --git a/library/src/getitune/backend/lightning/models/detection/edgecrafter.py b/library/src/getitune/backend/lightning/models/detection/edgecrafter.py index 7cbccd79b55..c95f98dd2c7 100644 --- a/library/src/getitune/backend/lightning/models/detection/edgecrafter.py +++ b/library/src/getitune/backend/lightning/models/detection/edgecrafter.py @@ -5,7 +5,10 @@ from __future__ import annotations -from typing import TYPE_CHECKING, ClassVar, Literal +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 @@ -15,7 +18,6 @@ from getitune.backend.lightning.models.detection.detectors.edgecrafter import ECDETRDetector from getitune.config.data import TileConfig from getitune.metrics.fmeasure import MeanAveragePrecisionFMeasureCallable -from getitune.types.export import TaskLevelExportParameters if TYPE_CHECKING: from lightning.pytorch.cli import LRSchedulerCallable, OptimizerCallable @@ -52,6 +54,8 @@ class EdgeCrafter(EdgeCrafterMixin, LightningDetectionModel): # pyrefly: ignore 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. """ @@ -79,10 +83,12 @@ def __init__( 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, @@ -104,14 +110,20 @@ def _create_model(self, num_classes: int | None = None) -> ECDETRDetector: Configured :class:`ECDETRDetector`. """ num_classes = num_classes if num_classes is not None else self.num_classes - return self._build_ec_model(num_classes, with_seg=False) + return self._build_ec_model(num_classes, with_seg=False, backbone_lr=self.backbone_lr) - @property - def _export_parameters(self) -> TaskLevelExportParameters: - """Export parameters for EdgeCrafter detection.""" - # Use a conservative NMS IoU threshold: DETR produces one-to-one predictions - # via Hungarian matching, so only suppress near-duplicate boxes. - return super()._export_parameters.wrap(iou_threshold=0.8, nms_execute=True) + 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: 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 index f29421a262c..445dd7b5da3 100644 --- a/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py +++ b/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py @@ -16,6 +16,8 @@ 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 @@ -36,8 +38,23 @@ class ECCriterion(DEIMCriterion): 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], 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 index 5439adb9f39..38325caa971 100644 --- a/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py +++ b/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py @@ -48,6 +48,8 @@ class EdgeCrafterInst(EdgeCrafterMixin, LightningInstanceSegModel): # pyrefly: 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. """ @@ -76,10 +78,12 @@ def __init__( 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, @@ -101,7 +105,7 @@ def _create_model(self, num_classes: int | None = None) -> ECDETRDetector: 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, with_seg=True) + return self._build_ec_model(num_classes, with_seg=True, backbone_lr=self.backbone_lr) @property def _export_parameters(self) -> TaskLevelExportParameters: diff --git a/library/src/getitune/recipe/detection/edgecrafter_l.yaml b/library/src/getitune/recipe/detection/edgecrafter_l.yaml index 29f9d5017ec..94b53d2375f 100644 --- a/library/src/getitune/recipe/detection/edgecrafter_l.yaml +++ b/library/src/getitune/recipe/detection/edgecrafter_l.yaml @@ -5,11 +5,12 @@ model: model_name: edgecrafter_l label_info: 80 multi_scale: false + backbone_lr: 0.000005 optimizer: class_path: torch.optim.AdamW init_args: - lr: 0.0002 + lr: 0.0005 betas: [0.9, 0.999] weight_decay: 0.0001 @@ -27,6 +28,7 @@ model: engine: device: auto + accumulate_grad_batches: 8 callback_monitor: val/map_50 @@ -51,9 +53,139 @@ callbacks: 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: [2, 48] + 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.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: + img_scale: [640, 640] + p: 1.0 + max_cached_images: 50 + random_pop: true + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 1.0 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + strong_aug_2: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + img_scale: [640, 640] + p: 1.0 + max_cached_images: 50 + random_pop: true + - 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: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 1.0 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + light_aug: + 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.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: diff --git a/library/src/getitune/recipe/detection/edgecrafter_m.yaml b/library/src/getitune/recipe/detection/edgecrafter_m.yaml index bd808271f2a..d647b16fa98 100644 --- a/library/src/getitune/recipe/detection/edgecrafter_m.yaml +++ b/library/src/getitune/recipe/detection/edgecrafter_m.yaml @@ -5,11 +5,12 @@ model: model_name: edgecrafter_m label_info: 80 multi_scale: false + backbone_lr: 0.000025 optimizer: class_path: torch.optim.AdamW init_args: - lr: 0.0004 + lr: 0.0005 betas: [0.9, 0.999] weight_decay: 0.0001 @@ -27,6 +28,7 @@ model: engine: device: auto + accumulate_grad_batches: 4 callback_monitor: val/map_50 @@ -51,9 +53,139 @@ callbacks: 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: [2, 60] + 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.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: + img_scale: [640, 640] + p: 0.75 + max_cached_images: 50 + random_pop: true + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 0.75 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + strong_aug_2: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + img_scale: [640, 640] + p: 0.75 + max_cached_images: 50 + random_pop: true + - 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: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 0.75 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + light_aug: + 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.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: diff --git a/library/src/getitune/recipe/detection/edgecrafter_s.yaml b/library/src/getitune/recipe/detection/edgecrafter_s.yaml index 770ac927ff0..3686320a39f 100644 --- a/library/src/getitune/recipe/detection/edgecrafter_s.yaml +++ b/library/src/getitune/recipe/detection/edgecrafter_s.yaml @@ -5,11 +5,12 @@ model: model_name: edgecrafter_s label_info: 80 multi_scale: false + backbone_lr: 0.000025 optimizer: class_path: torch.optim.AdamW init_args: - lr: 0.0004 + lr: 0.0005 betas: [0.9, 0.999] weight_decay: 0.0001 @@ -27,6 +28,7 @@ model: engine: device: auto + accumulate_grad_batches: 4 callback_monitor: val/map_50 @@ -51,9 +53,139 @@ callbacks: 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: [2, 72] + 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.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: + img_scale: [640, 640] + p: 0.75 + max_cached_images: 50 + random_pop: true + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 0.75 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + strong_aug_2: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + img_scale: [640, 640] + p: 0.75 + max_cached_images: 50 + random_pop: true + - 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: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 0.75 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + light_aug: + 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.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: @@ -66,8 +198,8 @@ overrides: data: input_size: - - 512 - - 512 + - 640 + - 640 task: DETECTION train_subset: batch_size: 8 diff --git a/library/src/getitune/recipe/detection/edgecrafter_x.yaml b/library/src/getitune/recipe/detection/edgecrafter_x.yaml index d22ee1051c9..7247acec950 100644 --- a/library/src/getitune/recipe/detection/edgecrafter_x.yaml +++ b/library/src/getitune/recipe/detection/edgecrafter_x.yaml @@ -5,11 +5,12 @@ model: model_name: edgecrafter_x label_info: 80 multi_scale: false + backbone_lr: 0.0000025 optimizer: class_path: torch.optim.AdamW init_args: - lr: 0.0001 + lr: 0.0005 betas: [0.9, 0.999] weight_decay: 0.0001 @@ -27,6 +28,7 @@ model: engine: device: auto + accumulate_grad_batches: 8 callback_monitor: val/map_50 @@ -51,9 +53,139 @@ callbacks: 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: [2, 48] + 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.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: + img_scale: [640, 640] + p: 1.0 + max_cached_images: 50 + random_pop: true + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 1.0 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + strong_aug_2: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + img_scale: [640, 640] + p: 1.0 + max_cached_images: 50 + random_pop: true + - 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: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 1.0 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + light_aug: + 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.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: diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml index c0b506120c0..724c724a090 100644 --- a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml +++ b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml @@ -6,11 +6,12 @@ model: model_name: edgecrafter_l label_info: 80 multi_scale: false + backbone_lr: 0.000005 optimizer: class_path: torch.optim.AdamW init_args: - lr: 0.0002 + lr: 0.0005 betas: [0.9, 0.999] weight_decay: 0.0001 @@ -57,9 +58,135 @@ callbacks: save_last: true auto_insert_metric_name: false filename: "checkpoints/epoch_{epoch:03d}" - - class_path: lightning.pytorch.callbacks.EMAWeightAveraging + - class_path: getitune.backend.lightning.callbacks.aug_scheduler.AugmentationSchedulerCallback init_args: - decay: 0.993 + data_aug_switch: + class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch + init_args: + policy_epochs: [2, 48] + 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.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: + img_scale: [640, 640] + p: 1.0 + max_cached_images: 50 + random_pop: true + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 1.0 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + strong_aug_2: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + img_scale: [640, 640] + p: 1.0 + max_cached_images: 50 + random_pop: true + - 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: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 1.0 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + light_aug: + 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.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] overrides: max_epochs: 100 @@ -84,8 +211,6 @@ overrides: init_args: size: $(input_size) keep_aspect_ratio: false - sampler: - class_path: getitune.data.samplers.balanced_sampler.BalancedSampler val_subset: batch_size: 4 diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml index c5a50193bb6..f0e91a1db6d 100644 --- a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml +++ b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml @@ -6,11 +6,12 @@ model: model_name: edgecrafter_m label_info: 80 multi_scale: false + backbone_lr: 0.000025 optimizer: class_path: torch.optim.AdamW init_args: - lr: 0.0004 + lr: 0.0005 betas: [0.9, 0.999] weight_decay: 0.0001 @@ -57,9 +58,135 @@ callbacks: save_last: true auto_insert_metric_name: false filename: "checkpoints/epoch_{epoch:03d}" - - class_path: lightning.pytorch.callbacks.EMAWeightAveraging + - class_path: getitune.backend.lightning.callbacks.aug_scheduler.AugmentationSchedulerCallback init_args: - decay: 0.993 + data_aug_switch: + class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch + init_args: + policy_epochs: [2, 60] + 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.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: + img_scale: [640, 640] + p: 0.75 + max_cached_images: 50 + random_pop: true + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 0.75 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + strong_aug_2: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + img_scale: [640, 640] + p: 0.75 + max_cached_images: 50 + random_pop: true + - 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: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 0.75 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + light_aug: + 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.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] overrides: max_epochs: 100 @@ -84,8 +211,6 @@ overrides: init_args: size: $(input_size) keep_aspect_ratio: false - sampler: - class_path: getitune.data.samplers.balanced_sampler.BalancedSampler val_subset: batch_size: 4 diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml index e9d792e353c..afc65f170d1 100644 --- a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml +++ b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml @@ -6,11 +6,12 @@ model: model_name: edgecrafter_s label_info: 80 multi_scale: false + backbone_lr: 0.000025 optimizer: class_path: torch.optim.AdamW init_args: - lr: 0.0004 + lr: 0.0005 betas: [0.9, 0.999] weight_decay: 0.0001 @@ -57,9 +58,135 @@ callbacks: save_last: true auto_insert_metric_name: false filename: "checkpoints/epoch_{epoch:03d}" - - class_path: lightning.pytorch.callbacks.EMAWeightAveraging + - class_path: getitune.backend.lightning.callbacks.aug_scheduler.AugmentationSchedulerCallback init_args: - decay: 0.993 + data_aug_switch: + class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch + init_args: + policy_epochs: [2, 72] + 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.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: + img_scale: [640, 640] + p: 0.75 + max_cached_images: 50 + random_pop: true + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 0.75 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + strong_aug_2: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + img_scale: [640, 640] + p: 0.75 + max_cached_images: 50 + random_pop: true + - 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: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 0.75 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + light_aug: + 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.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] overrides: max_epochs: 100 @@ -67,8 +194,8 @@ overrides: data: task: INSTANCE_SEGMENTATION input_size: - - 512 - - 512 + - 640 + - 640 train_subset: batch_size: 4 augmentations_cpu: @@ -84,8 +211,6 @@ overrides: init_args: size: $(input_size) keep_aspect_ratio: false - sampler: - class_path: getitune.data.samplers.balanced_sampler.BalancedSampler val_subset: batch_size: 4 diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml index 3640f6a24a4..4b2ed92822b 100644 --- a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml +++ b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml @@ -6,11 +6,12 @@ model: model_name: edgecrafter_x label_info: 80 multi_scale: false + backbone_lr: 0.0000025 optimizer: class_path: torch.optim.AdamW init_args: - lr: 0.0001 + lr: 0.0005 betas: [0.9, 0.999] weight_decay: 0.0001 @@ -57,9 +58,135 @@ callbacks: save_last: true auto_insert_metric_name: false filename: "checkpoints/epoch_{epoch:03d}" - - class_path: lightning.pytorch.callbacks.EMAWeightAveraging + - class_path: getitune.backend.lightning.callbacks.aug_scheduler.AugmentationSchedulerCallback init_args: - decay: 0.993 + data_aug_switch: + class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch + init_args: + policy_epochs: [2, 48] + 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.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: + img_scale: [640, 640] + p: 1.0 + max_cached_images: 50 + random_pop: true + - class_path: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 1.0 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + strong_aug_2: + augmentations_cpu: + - class_path: getitune.data.augmentation.transforms.CachedMosaic + init_args: + img_scale: [640, 640] + p: 1.0 + max_cached_images: 50 + random_pop: true + - 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: getitune.data.augmentation.transforms.CachedMixUp + init_args: + img_scale: [640, 640] + p: 1.0 + alpha: 1.5 + max_cached_images: 20 + random_pop: true + - 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] + light_aug: + 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.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] overrides: max_epochs: 100 @@ -84,8 +211,6 @@ overrides: init_args: size: $(input_size) keep_aspect_ratio: false - sampler: - class_path: getitune.data.samplers.balanced_sampler.BalancedSampler val_subset: batch_size: 4 From 05413e98e8778fca7212dd7b05e56bb8585eb960 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Fri, 3 Jul 2026 20:28:27 +0900 Subject: [PATCH 06/14] Guard Hungarian matcher without masks --- .../edgecrafter_inst_l.yaml | 56 ----- .../edgecrafter_inst_m.yaml | 56 ----- .../edgecrafter_inst_s.yaml | 56 ----- .../edgecrafter_inst_x.yaml | 56 ----- .../utils/assigners/hungarian_matcher.py | 5 + .../edgecrafter_inst_l.yaml | 231 ------------------ .../edgecrafter_inst_m.yaml | 231 ------------------ .../edgecrafter_inst_s.yaml | 231 ------------------ .../edgecrafter_inst_x.yaml | 231 ------------------ 9 files changed, 5 insertions(+), 1148 deletions(-) delete mode 100644 application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_l.yaml delete mode 100644 application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_m.yaml delete mode 100644 application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_s.yaml delete mode 100644 application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_x.yaml delete mode 100644 library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml delete mode 100644 library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml delete mode 100644 library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml delete mode 100644 library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml diff --git a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_l.yaml b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_l.yaml deleted file mode 100644 index ff4054b33c1..00000000000 --- a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_l.yaml +++ /dev/null @@ -1,56 +0,0 @@ -id: instance-segmentation-edgecrafter-l -name: ECSeg-L - -pretrained_weights: - url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_l.pth - mirror_url: https://huggingface.co/Intellindust/ECSeg_L/resolve/main/model.safetensors - sha_sum: 6fa7e8ffdd277feeef44bb1b8046315ae860023af9432bb75ecf400d74fe2396 - -description: "EdgeCrafter instance segmentation model (ECSeg-L) extends ECDet with a mask prediction head for high-accuracy real-time instance segmentation." - -stats: - gigaflops: 111.0 - trainable_parameters: 34.0 - benchmark_metrics: - coco_map_50_95: 47.1 - -capabilities: - tiling: false - -hyperparameters: - training: - learning_rate: 0.0002 - weight_decay: 0.0001 - early_stopping: - patience: 15 - scheduler: - factor: 0.1 - patience: 10 - gradient_clip: - enable: true - max_grad_norm: 0.1 - allowed_values_input_size: - - 320 - - 448 - - 512 - - 576 - - 640 - - 704 - - 768 - - 896 - input_size_width: 640 - input_size_height: 640 - dataset_preparation: - augmentation: - random_zoom_out: - enable: true - fill: 0 - side_range: - - 1.0 - - 4.0 - probability: 0.5 - iou_random_crop: - enable: true - probability: 1.0 - min_scale: 0.3 - max_scale: 1.0 diff --git a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_m.yaml b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_m.yaml deleted file mode 100644 index ebc74cbe6ee..00000000000 --- a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_m.yaml +++ /dev/null @@ -1,56 +0,0 @@ -id: instance-segmentation-edgecrafter-m -name: ECSeg-M - -pretrained_weights: - url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_m.pth - mirror_url: https://huggingface.co/Intellindust/ECSeg_M/resolve/main/model.safetensors - sha_sum: f1abcbd6ed747ed7364466a9fc4d3dc9ceb8d48cc123b353f717235c9c52069d - -description: "EdgeCrafter instance segmentation model (ECSeg-M) extends ECDet with a mask prediction head for high-accuracy real-time instance segmentation." - -stats: - gigaflops: 64.0 - trainable_parameters: 20.0 - benchmark_metrics: - coco_map_50_95: 45.2 - -capabilities: - tiling: false - -hyperparameters: - training: - learning_rate: 0.0004 - weight_decay: 0.0001 - early_stopping: - patience: 15 - scheduler: - factor: 0.1 - patience: 10 - gradient_clip: - enable: true - max_grad_norm: 0.1 - allowed_values_input_size: - - 320 - - 448 - - 512 - - 576 - - 640 - - 704 - - 768 - - 896 - input_size_width: 640 - input_size_height: 640 - dataset_preparation: - augmentation: - random_zoom_out: - enable: true - fill: 0 - side_range: - - 1.0 - - 4.0 - probability: 0.5 - iou_random_crop: - enable: true - probability: 1.0 - min_scale: 0.3 - max_scale: 1.0 diff --git a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_s.yaml b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_s.yaml deleted file mode 100644 index 0d742763e31..00000000000 --- a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_s.yaml +++ /dev/null @@ -1,56 +0,0 @@ -id: instance-segmentation-edgecrafter-s -name: ECSeg-S - -pretrained_weights: - url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_s.pth - mirror_url: https://huggingface.co/Intellindust/ECSeg_S/resolve/main/model.safetensors - sha_sum: b83b07d71b942c6255aaad1cf1f92eee075c5a548624d27464c987cd618d5a93 - -description: "EdgeCrafter instance segmentation model (ECSeg-S) extends ECDet with a mask prediction head for high-accuracy real-time instance segmentation." - -stats: - gigaflops: 33.0 - trainable_parameters: 10.0 - benchmark_metrics: - coco_map_50_95: 43.0 - -capabilities: - tiling: false - -hyperparameters: - training: - learning_rate: 0.0004 - weight_decay: 0.0001 - early_stopping: - patience: 15 - scheduler: - factor: 0.1 - patience: 10 - gradient_clip: - enable: true - max_grad_norm: 0.1 - allowed_values_input_size: - - 320 - - 448 - - 512 - - 576 - - 640 - - 704 - - 768 - - 896 - input_size_width: 512 - input_size_height: 512 - dataset_preparation: - augmentation: - random_zoom_out: - enable: true - fill: 0 - side_range: - - 1.0 - - 4.0 - probability: 0.5 - iou_random_crop: - enable: true - probability: 1.0 - min_scale: 0.3 - max_scale: 1.0 diff --git a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_x.yaml b/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_x.yaml deleted file mode 100644 index 4fa062f4915..00000000000 --- a/application/backend/app/supported_models/manifests/instance_segmentation/edgecrafter_inst_x.yaml +++ /dev/null @@ -1,56 +0,0 @@ -id: instance-segmentation-edgecrafter-x -name: ECSeg-X - -pretrained_weights: - url: https://github.com/capsule2077/edgecrafter/releases/download/edgecrafterv1/ecseg_x.pth - mirror_url: https://huggingface.co/Intellindust/ECSeg_X/resolve/main/model.safetensors - sha_sum: 5e097fdc49cdf31dcc117f56b6504ee4afc9198378a02b3e7e2da69f1d4275a7 - -description: "EdgeCrafter instance segmentation model (ECSeg-X) extends ECDet with a mask prediction head for high-accuracy real-time instance segmentation." - -stats: - gigaflops: 168.0 - trainable_parameters: 50.0 - benchmark_metrics: - coco_map_50_95: 48.4 - -capabilities: - tiling: false - -hyperparameters: - training: - learning_rate: 0.0001 - weight_decay: 0.0001 - early_stopping: - patience: 15 - scheduler: - factor: 0.1 - patience: 10 - gradient_clip: - enable: true - max_grad_norm: 0.1 - allowed_values_input_size: - - 320 - - 448 - - 512 - - 576 - - 640 - - 704 - - 768 - - 896 - input_size_width: 640 - input_size_height: 640 - dataset_preparation: - augmentation: - random_zoom_out: - enable: true - fill: 0 - side_range: - - 1.0 - - 4.0 - probability: 0.5 - iou_random_crop: - enable: true - probability: 1.0 - min_scale: 0.3 - max_scale: 1.0 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/recipe/instance_segmentation/edgecrafter_inst_l.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml deleted file mode 100644 index 724c724a090..00000000000 --- a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_l.yaml +++ /dev/null @@ -1,231 +0,0 @@ -task: INSTANCE_SEGMENTATION - -model: - class_path: getitune.backend.lightning.models.instance_segmentation.EdgeCrafterInst - 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: 4 - -data: ../_base_/data/instance_segmentation.yaml - -callback_monitor: val/map_50 - -callbacks: - - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling - init_args: - max_interval: 1 - min_lrschedule_patience: 5 - - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup - init_args: - mode: max - patience: 15 - check_on_train_epoch_end: false - monitor: val/map_50 - min_delta: 0.001 - warmup_iters: 100 - 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: [2, 48] - 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.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: - img_scale: [640, 640] - p: 1.0 - max_cached_images: 50 - random_pop: true - - class_path: getitune.data.augmentation.transforms.CachedMixUp - init_args: - img_scale: [640, 640] - p: 1.0 - alpha: 1.5 - max_cached_images: 20 - random_pop: true - - 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] - strong_aug_2: - augmentations_cpu: - - class_path: getitune.data.augmentation.transforms.CachedMosaic - init_args: - img_scale: [640, 640] - p: 1.0 - max_cached_images: 50 - random_pop: true - - 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: getitune.data.augmentation.transforms.CachedMixUp - init_args: - img_scale: [640, 640] - p: 1.0 - alpha: 1.5 - max_cached_images: 20 - random_pop: true - - 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] - light_aug: - 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.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] - -overrides: - max_epochs: 100 - gradient_clip_val: 0.1 - data: - task: INSTANCE_SEGMENTATION - input_size: - - 640 - - 640 - 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 - - 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 - - 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 diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml deleted file mode 100644 index f0e91a1db6d..00000000000 --- a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_m.yaml +++ /dev/null @@ -1,231 +0,0 @@ -task: INSTANCE_SEGMENTATION - -model: - class_path: getitune.backend.lightning.models.instance_segmentation.EdgeCrafterInst - 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 - -data: ../_base_/data/instance_segmentation.yaml - -callback_monitor: val/map_50 - -callbacks: - - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling - init_args: - max_interval: 1 - min_lrschedule_patience: 5 - - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup - init_args: - mode: max - patience: 15 - check_on_train_epoch_end: false - monitor: val/map_50 - min_delta: 0.001 - warmup_iters: 100 - 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: [2, 60] - 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.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: - img_scale: [640, 640] - p: 0.75 - max_cached_images: 50 - random_pop: true - - class_path: getitune.data.augmentation.transforms.CachedMixUp - init_args: - img_scale: [640, 640] - p: 0.75 - alpha: 1.5 - max_cached_images: 20 - random_pop: true - - 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] - strong_aug_2: - augmentations_cpu: - - class_path: getitune.data.augmentation.transforms.CachedMosaic - init_args: - img_scale: [640, 640] - p: 0.75 - max_cached_images: 50 - random_pop: true - - 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: getitune.data.augmentation.transforms.CachedMixUp - init_args: - img_scale: [640, 640] - p: 0.75 - alpha: 1.5 - max_cached_images: 20 - random_pop: true - - 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] - light_aug: - 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.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] - -overrides: - max_epochs: 100 - gradient_clip_val: 0.1 - data: - task: INSTANCE_SEGMENTATION - input_size: - - 640 - - 640 - 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 - - 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 - - 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 diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml deleted file mode 100644 index afc65f170d1..00000000000 --- a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_s.yaml +++ /dev/null @@ -1,231 +0,0 @@ -task: INSTANCE_SEGMENTATION - -model: - class_path: getitune.backend.lightning.models.instance_segmentation.EdgeCrafterInst - 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 - -data: ../_base_/data/instance_segmentation.yaml - -callback_monitor: val/map_50 - -callbacks: - - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling - init_args: - max_interval: 1 - min_lrschedule_patience: 5 - - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup - init_args: - mode: max - patience: 15 - check_on_train_epoch_end: false - monitor: val/map_50 - min_delta: 0.001 - warmup_iters: 100 - 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: [2, 72] - 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.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: - img_scale: [640, 640] - p: 0.75 - max_cached_images: 50 - random_pop: true - - class_path: getitune.data.augmentation.transforms.CachedMixUp - init_args: - img_scale: [640, 640] - p: 0.75 - alpha: 1.5 - max_cached_images: 20 - random_pop: true - - 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] - strong_aug_2: - augmentations_cpu: - - class_path: getitune.data.augmentation.transforms.CachedMosaic - init_args: - img_scale: [640, 640] - p: 0.75 - max_cached_images: 50 - random_pop: true - - 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: getitune.data.augmentation.transforms.CachedMixUp - init_args: - img_scale: [640, 640] - p: 0.75 - alpha: 1.5 - max_cached_images: 20 - random_pop: true - - 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] - light_aug: - 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.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] - -overrides: - max_epochs: 100 - gradient_clip_val: 0.1 - data: - task: INSTANCE_SEGMENTATION - input_size: - - 640 - - 640 - 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 - - 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 - - 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 diff --git a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml b/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml deleted file mode 100644 index 4b2ed92822b..00000000000 --- a/library/src/getitune/recipe/instance_segmentation/edgecrafter_inst_x.yaml +++ /dev/null @@ -1,231 +0,0 @@ -task: INSTANCE_SEGMENTATION - -model: - class_path: getitune.backend.lightning.models.instance_segmentation.EdgeCrafterInst - 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: 4 - -data: ../_base_/data/instance_segmentation.yaml - -callback_monitor: val/map_50 - -callbacks: - - class_path: getitune.backend.lightning.callbacks.adaptive_train_scheduling.AdaptiveTrainScheduling - init_args: - max_interval: 1 - min_lrschedule_patience: 5 - - class_path: getitune.backend.lightning.callbacks.adaptive_early_stopping.EarlyStoppingWithWarmup - init_args: - mode: max - patience: 15 - check_on_train_epoch_end: false - monitor: val/map_50 - min_delta: 0.001 - warmup_iters: 100 - 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: [2, 48] - 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.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: - img_scale: [640, 640] - p: 1.0 - max_cached_images: 50 - random_pop: true - - class_path: getitune.data.augmentation.transforms.CachedMixUp - init_args: - img_scale: [640, 640] - p: 1.0 - alpha: 1.5 - max_cached_images: 20 - random_pop: true - - 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] - strong_aug_2: - augmentations_cpu: - - class_path: getitune.data.augmentation.transforms.CachedMosaic - init_args: - img_scale: [640, 640] - p: 1.0 - max_cached_images: 50 - random_pop: true - - 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: getitune.data.augmentation.transforms.CachedMixUp - init_args: - img_scale: [640, 640] - p: 1.0 - alpha: 1.5 - max_cached_images: 20 - random_pop: true - - 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] - light_aug: - 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.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] - -overrides: - max_epochs: 100 - gradient_clip_val: 0.1 - data: - task: INSTANCE_SEGMENTATION - input_size: - - 640 - - 640 - 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 - - 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 - - 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 From 18506836bc889faa12db3ad84df671ffdc534bdf Mon Sep 17 00:00:00 2001 From: kprokofi Date: Fri, 3 Jul 2026 20:56:09 +0900 Subject: [PATCH 07/14] align hyperparameters --- .../manifests/detection/edgecrafter_l.yaml | 21 ++----- .../manifests/detection/edgecrafter_m.yaml | 21 ++----- .../manifests/detection/edgecrafter_s.yaml | 24 +++---- .../manifests/detection/edgecrafter_x.yaml | 21 ++----- .../recipe/detection/edgecrafter_l.yaml | 62 +++++++++---------- .../recipe/detection/edgecrafter_m.yaml | 62 +++++++++---------- .../recipe/detection/edgecrafter_s.yaml | 62 +++++++++---------- .../recipe/detection/edgecrafter_x.yaml | 62 +++++++++---------- 8 files changed, 138 insertions(+), 197 deletions(-) diff --git a/application/backend/app/supported_models/manifests/detection/edgecrafter_l.yaml b/application/backend/app/supported_models/manifests/detection/edgecrafter_l.yaml index cd4c03e7b0e..67a9b7bd5ee 100644 --- a/application/backend/app/supported_models/manifests/detection/edgecrafter_l.yaml +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_l.yaml @@ -19,33 +19,24 @@ capabilities: hyperparameters: training: - learning_rate: 0.0002 + learning_rate: 0.0005 weight_decay: 0.0001 + batch_size: 4 + accumulate_grad_batches: 8 early_stopping: - patience: 15 + patience: 10 scheduler: - factor: 0.1 + factor: 0.5 patience: 10 gradient_clip: enable: true max_grad_norm: 0.1 allowed_values_input_size: - - 320 - - 448 - - 512 - - 576 - 640 - - 704 - - 768 - - 896 input_size_width: 640 input_size_height: 640 dataset_preparation: augmentation: + deim_framework: true random_zoom_out: enable: true - fill: 0 - side_range: - - 1.0 - - 4.0 - probability: 0.5 diff --git a/application/backend/app/supported_models/manifests/detection/edgecrafter_m.yaml b/application/backend/app/supported_models/manifests/detection/edgecrafter_m.yaml index 481fe761530..985486d6b16 100644 --- a/application/backend/app/supported_models/manifests/detection/edgecrafter_m.yaml +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_m.yaml @@ -19,33 +19,24 @@ capabilities: hyperparameters: training: - learning_rate: 0.0004 + learning_rate: 0.0005 weight_decay: 0.0001 + batch_size: 8 + accumulate_grad_batches: 4 early_stopping: - patience: 15 + patience: 10 scheduler: - factor: 0.1 + factor: 0.5 patience: 10 gradient_clip: enable: true max_grad_norm: 0.1 allowed_values_input_size: - - 320 - - 448 - - 512 - - 576 - 640 - - 704 - - 768 - - 896 input_size_width: 640 input_size_height: 640 dataset_preparation: augmentation: + deim_framework: true random_zoom_out: enable: true - fill: 0 - side_range: - - 1.0 - - 4.0 - probability: 0.5 diff --git a/application/backend/app/supported_models/manifests/detection/edgecrafter_s.yaml b/application/backend/app/supported_models/manifests/detection/edgecrafter_s.yaml index da607737658..9bbbe5cf229 100644 --- a/application/backend/app/supported_models/manifests/detection/edgecrafter_s.yaml +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_s.yaml @@ -19,33 +19,25 @@ capabilities: hyperparameters: training: - learning_rate: 0.0004 + learning_rate: 0.0005 weight_decay: 0.0001 + batch_size: 8 + accumulate_grad_batches: 4 early_stopping: - patience: 15 + patience: 10 scheduler: - factor: 0.1 + factor: 0.5 patience: 10 gradient_clip: enable: true max_grad_norm: 0.1 allowed_values_input_size: - - 320 - - 448 - - 512 - - 576 - 640 - - 704 - - 768 - - 896 - input_size_width: 512 - input_size_height: 512 + input_size_width: 640 + input_size_height: 640 dataset_preparation: augmentation: + deim_framework: true random_zoom_out: enable: true fill: 0 - side_range: - - 1.0 - - 4.0 - probability: 0.5 diff --git a/application/backend/app/supported_models/manifests/detection/edgecrafter_x.yaml b/application/backend/app/supported_models/manifests/detection/edgecrafter_x.yaml index 79508692adf..2f95a59e252 100644 --- a/application/backend/app/supported_models/manifests/detection/edgecrafter_x.yaml +++ b/application/backend/app/supported_models/manifests/detection/edgecrafter_x.yaml @@ -19,33 +19,24 @@ capabilities: hyperparameters: training: - learning_rate: 0.0001 + learning_rate: 0.0005 weight_decay: 0.0001 + batch_size: 4 + accumulate_grad_batches: 8 early_stopping: - patience: 15 + patience: 10 scheduler: - factor: 0.1 + factor: 0.5 patience: 10 gradient_clip: enable: true max_grad_norm: 0.1 allowed_values_input_size: - - 320 - - 448 - - 512 - - 576 - 640 - - 704 - - 768 - - 896 input_size_width: 640 input_size_height: 640 dataset_preparation: augmentation: + deim_framework: true random_zoom_out: enable: true - fill: 0 - side_range: - - 1.0 - - 4.0 - probability: 0.5 diff --git a/library/src/getitune/recipe/detection/edgecrafter_l.yaml b/library/src/getitune/recipe/detection/edgecrafter_l.yaml index 94b53d2375f..320afc32bc2 100644 --- a/library/src/getitune/recipe/detection/edgecrafter_l.yaml +++ b/library/src/getitune/recipe/detection/edgecrafter_l.yaml @@ -58,7 +58,7 @@ callbacks: data_aug_switch: class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch init_args: - policy_epochs: [2, 48] + policy_epochs: [4, 23] input_size: [640, 640] policies: no_aug: @@ -68,25 +68,20 @@ callbacks: 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] - strong_aug_1: + light_aug: augmentations_cpu: - - class_path: getitune.data.augmentation.transforms.CachedMosaic - init_args: - img_scale: [640, 640] - p: 1.0 - max_cached_images: 50 - random_pop: true - - class_path: getitune.data.augmentation.transforms.CachedMixUp + - class_path: torchvision.transforms.v2.RandomZoomOut init_args: - img_scale: [640, 640] - p: 1.0 - alpha: 1.5 - max_cached_images: 20 - random_pop: true + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes init_args: min_size: 1 @@ -95,9 +90,6 @@ callbacks: 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 @@ -105,38 +97,30 @@ callbacks: 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_2: + strong_aug_1: augmentations_cpu: - class_path: getitune.data.augmentation.transforms.CachedMosaic init_args: - img_scale: [640, 640] - p: 1.0 - max_cached_images: 50 random_pop: true - - 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 + max_cached_images: 20 + img_scale: [640, 640] - class_path: getitune.data.augmentation.transforms.CachedMixUp init_args: img_scale: [640, 640] - p: 1.0 alpha: 1.5 - max_cached_images: 20 + 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: @@ -152,7 +136,7 @@ callbacks: init_args: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] - light_aug: + strong_aug_2: augmentations_cpu: - class_path: torchvision.transforms.v2.RandomZoomOut init_args: @@ -160,6 +144,16 @@ callbacks: 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 diff --git a/library/src/getitune/recipe/detection/edgecrafter_m.yaml b/library/src/getitune/recipe/detection/edgecrafter_m.yaml index d647b16fa98..c80544156f1 100644 --- a/library/src/getitune/recipe/detection/edgecrafter_m.yaml +++ b/library/src/getitune/recipe/detection/edgecrafter_m.yaml @@ -58,7 +58,7 @@ callbacks: data_aug_switch: class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch init_args: - policy_epochs: [2, 60] + policy_epochs: [4, 40] input_size: [640, 640] policies: no_aug: @@ -68,25 +68,20 @@ callbacks: 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] - strong_aug_1: + light_aug: augmentations_cpu: - - class_path: getitune.data.augmentation.transforms.CachedMosaic - init_args: - img_scale: [640, 640] - p: 0.75 - max_cached_images: 50 - random_pop: true - - class_path: getitune.data.augmentation.transforms.CachedMixUp + - class_path: torchvision.transforms.v2.RandomZoomOut init_args: - img_scale: [640, 640] - p: 0.75 - alpha: 1.5 - max_cached_images: 20 - random_pop: true + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes init_args: min_size: 1 @@ -95,9 +90,6 @@ callbacks: 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 @@ -105,38 +97,30 @@ callbacks: 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_2: + strong_aug_1: augmentations_cpu: - class_path: getitune.data.augmentation.transforms.CachedMosaic init_args: - img_scale: [640, 640] - p: 0.75 - max_cached_images: 50 random_pop: true - - 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 + max_cached_images: 20 + img_scale: [640, 640] - class_path: getitune.data.augmentation.transforms.CachedMixUp init_args: img_scale: [640, 640] - p: 0.75 alpha: 1.5 - max_cached_images: 20 + 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: @@ -152,7 +136,7 @@ callbacks: init_args: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] - light_aug: + strong_aug_2: augmentations_cpu: - class_path: torchvision.transforms.v2.RandomZoomOut init_args: @@ -160,6 +144,16 @@ callbacks: 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 diff --git a/library/src/getitune/recipe/detection/edgecrafter_s.yaml b/library/src/getitune/recipe/detection/edgecrafter_s.yaml index 3686320a39f..6d75670f707 100644 --- a/library/src/getitune/recipe/detection/edgecrafter_s.yaml +++ b/library/src/getitune/recipe/detection/edgecrafter_s.yaml @@ -58,7 +58,7 @@ callbacks: data_aug_switch: class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch init_args: - policy_epochs: [2, 72] + policy_epochs: [4, 40] input_size: [640, 640] policies: no_aug: @@ -68,25 +68,20 @@ callbacks: 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] - strong_aug_1: + light_aug: augmentations_cpu: - - class_path: getitune.data.augmentation.transforms.CachedMosaic - init_args: - img_scale: [640, 640] - p: 0.75 - max_cached_images: 50 - random_pop: true - - class_path: getitune.data.augmentation.transforms.CachedMixUp + - class_path: torchvision.transforms.v2.RandomZoomOut init_args: - img_scale: [640, 640] - p: 0.75 - alpha: 1.5 - max_cached_images: 20 - random_pop: true + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes init_args: min_size: 1 @@ -95,9 +90,6 @@ callbacks: 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 @@ -105,38 +97,30 @@ callbacks: 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_2: + strong_aug_1: augmentations_cpu: - class_path: getitune.data.augmentation.transforms.CachedMosaic init_args: - img_scale: [640, 640] - p: 0.75 - max_cached_images: 50 random_pop: true - - 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 + max_cached_images: 20 + img_scale: [640, 640] - class_path: getitune.data.augmentation.transforms.CachedMixUp init_args: img_scale: [640, 640] - p: 0.75 alpha: 1.5 - max_cached_images: 20 + 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: @@ -152,7 +136,7 @@ callbacks: init_args: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] - light_aug: + strong_aug_2: augmentations_cpu: - class_path: torchvision.transforms.v2.RandomZoomOut init_args: @@ -160,6 +144,16 @@ callbacks: 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 diff --git a/library/src/getitune/recipe/detection/edgecrafter_x.yaml b/library/src/getitune/recipe/detection/edgecrafter_x.yaml index 7247acec950..929127c36c8 100644 --- a/library/src/getitune/recipe/detection/edgecrafter_x.yaml +++ b/library/src/getitune/recipe/detection/edgecrafter_x.yaml @@ -58,7 +58,7 @@ callbacks: data_aug_switch: class_path: getitune.backend.lightning.callbacks.aug_scheduler.DataAugSwitch init_args: - policy_epochs: [2, 48] + policy_epochs: [4, 23] input_size: [640, 640] policies: no_aug: @@ -68,25 +68,20 @@ callbacks: 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] - strong_aug_1: + light_aug: augmentations_cpu: - - class_path: getitune.data.augmentation.transforms.CachedMosaic - init_args: - img_scale: [640, 640] - p: 1.0 - max_cached_images: 50 - random_pop: true - - class_path: getitune.data.augmentation.transforms.CachedMixUp + - class_path: torchvision.transforms.v2.RandomZoomOut init_args: - img_scale: [640, 640] - p: 1.0 - alpha: 1.5 - max_cached_images: 20 - random_pop: true + fill: 0 + p: 0.5 + - class_path: getitune.data.augmentation.transforms.RandomIoUCrop - class_path: torchvision.transforms.v2.SanitizeBoundingBoxes init_args: min_size: 1 @@ -95,9 +90,6 @@ callbacks: 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 @@ -105,38 +97,30 @@ callbacks: 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_2: + strong_aug_1: augmentations_cpu: - class_path: getitune.data.augmentation.transforms.CachedMosaic init_args: - img_scale: [640, 640] - p: 1.0 - max_cached_images: 50 random_pop: true - - 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 + max_cached_images: 20 + img_scale: [640, 640] - class_path: getitune.data.augmentation.transforms.CachedMixUp init_args: img_scale: [640, 640] - p: 1.0 alpha: 1.5 - max_cached_images: 20 + 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: @@ -152,7 +136,7 @@ callbacks: init_args: mean: [0.485, 0.456, 0.406] std: [0.229, 0.224, 0.225] - light_aug: + strong_aug_2: augmentations_cpu: - class_path: torchvision.transforms.v2.RandomZoomOut init_args: @@ -160,6 +144,16 @@ callbacks: 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 From d36f6db071e4d50ce53c3f07bb8f4f52c5108e04 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Fri, 3 Jul 2026 22:27:07 +0900 Subject: [PATCH 08/14] fix unit tests --- .../tests/unit/api/routers/test_model_architectures.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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..2618444ad28 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,8 +57,8 @@ def test_get_all_model_architectures(self, fxt_client: TestClient): @pytest.mark.parametrize( "task_filter, total_models", [ - ("detection", 32), - ("instance_segmentation", 20), + ("detection", 36), + ("instance_segmentation", 24), ("classification", 6), ], ) From df2dce1a139d7f9b3253e65d5a7f0aea2b5dd5b1 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Fri, 3 Jul 2026 23:56:00 +0900 Subject: [PATCH 09/14] Fix model architecture count test --- .../backend/tests/unit/api/routers/test_model_architectures.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 2618444ad28..8f01cc6137f 100644 --- a/application/backend/tests/unit/api/routers/test_model_architectures.py +++ b/application/backend/tests/unit/api/routers/test_model_architectures.py @@ -58,7 +58,7 @@ def test_get_all_model_architectures(self, fxt_client: TestClient): "task_filter, total_models", [ ("detection", 36), - ("instance_segmentation", 24), + ("instance_segmentation", 20), ("classification", 6), ], ) From 3aa7f2094c91a78b4ec86e8f30eafd892bf05f44 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Mon, 6 Jul 2026 22:19:29 +0900 Subject: [PATCH 10/14] Fix EdgeCrafter PR review comments --- .../execution/common/geti_config_converter.py | 20 ------------------- .../api/routers/test_model_architectures.py | 2 +- library/benchmark_manifest.yaml | 3 +++ .../lightning/models/detection/edgecrafter.py | 2 +- .../instance_segmentation/edgecrafter_inst.py | 2 +- 5 files changed, 6 insertions(+), 23 deletions(-) diff --git a/application/backend/app/execution/common/geti_config_converter.py b/application/backend/app/execution/common/geti_config_converter.py index eb23dd9ed48..80310096c87 100644 --- a/application/backend/app/execution/common/geti_config_converter.py +++ b/application/backend/app/execution/common/geti_config_converter.py @@ -938,26 +938,6 @@ def convert(config: dict) -> dict: "status": ModelStatus.ACCURACY, "default": False, }, - "instance-segmentation-edgecrafter-s": { - "recipe_path": RECIPE_PATH / "instance_segmentation" / "edgecrafter_inst_s.yaml", - "status": ModelStatus.ACTIVE, - "default": False, - }, - "instance-segmentation-edgecrafter-m": { - "recipe_path": RECIPE_PATH / "instance_segmentation" / "edgecrafter_inst_m.yaml", - "status": ModelStatus.ACTIVE, - "default": False, - }, - "instance-segmentation-edgecrafter-l": { - "recipe_path": RECIPE_PATH / "instance_segmentation" / "edgecrafter_inst_l.yaml", - "status": ModelStatus.ACTIVE, - "default": False, - }, - "instance-segmentation-edgecrafter-x": { - "recipe_path": RECIPE_PATH / "instance_segmentation" / "edgecrafter_inst_x.yaml", - "status": ModelStatus.ACTIVE, - "default": False, - }, "instance-segmentation-yolo26-n": { "recipe_path": RECIPE_PATH / "instance_segmentation" / "yolo26_n_seg.yaml", "status": ModelStatus.ACTIVE, 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 8f01cc6137f..2618444ad28 100644 --- a/application/backend/tests/unit/api/routers/test_model_architectures.py +++ b/application/backend/tests/unit/api/routers/test_model_architectures.py @@ -58,7 +58,7 @@ def test_get_all_model_architectures(self, fxt_client: TestClient): "task_filter, total_models", [ ("detection", 36), - ("instance_segmentation", 20), + ("instance_segmentation", 24), ("classification", 6), ], ) diff --git a/library/benchmark_manifest.yaml b/library/benchmark_manifest.yaml index 75fd961fb34..8298be9c1cc 100644 --- a/library/benchmark_manifest.yaml +++ b/library/benchmark_manifest.yaml @@ -139,6 +139,9 @@ experiments: - name: yolo26_m priority: extended recipe: detection/yolo26_m.yaml + - name: edgecrafter_m + priority: extended + recipe: detection/edgecrafter_m.yaml datasets: - bccd criteria: diff --git a/library/src/getitune/backend/lightning/models/detection/edgecrafter.py b/library/src/getitune/backend/lightning/models/detection/edgecrafter.py index c95f98dd2c7..3a15d796b1d 100644 --- a/library/src/getitune/backend/lightning/models/detection/edgecrafter.py +++ b/library/src/getitune/backend/lightning/models/detection/edgecrafter.py @@ -150,7 +150,7 @@ def _default_preprocessing_params(self) -> dict[str, DataInputParams]: imagenet_mean = (0.485, 0.456, 0.406) imagenet_std = (0.229, 0.224, 0.225) return { - "edgecrafter_s": DataInputParams(input_size=(512, 512), mean=imagenet_mean, std=imagenet_std), + "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/edgecrafter_inst.py b/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py index 38325caa971..735e9bb8e5c 100644 --- a/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py +++ b/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py @@ -160,7 +160,7 @@ def _default_preprocessing_params(self) -> dict[str, DataInputParams]: imagenet_mean = (0.485, 0.456, 0.406) imagenet_std = (0.229, 0.224, 0.225) return { - "edgecrafter_s": DataInputParams(input_size=(512, 512), mean=imagenet_mean, std=imagenet_std), + "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), From ad0e60851ce21f83ddd13983132651b07d7dc433 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Mon, 6 Jul 2026 22:36:18 +0900 Subject: [PATCH 11/14] fix unit test --- .../backend/tests/unit/api/routers/test_model_architectures.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 2618444ad28..8f01cc6137f 100644 --- a/application/backend/tests/unit/api/routers/test_model_architectures.py +++ b/application/backend/tests/unit/api/routers/test_model_architectures.py @@ -58,7 +58,7 @@ def test_get_all_model_architectures(self, fxt_client: TestClient): "task_filter, total_models", [ ("detection", 36), - ("instance_segmentation", 24), + ("instance_segmentation", 20), ("classification", 6), ], ) From d2abda16899d1fd5a9c9fbb08d7e4925b6ccb701 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Mon, 6 Jul 2026 22:46:21 +0900 Subject: [PATCH 12/14] fix spatial features projection --- .../backend/lightning/models/common/edgecrafter_mixin.py | 2 +- .../backend/lightning/models/detection/heads/ec_decoder.py | 1 + .../lightning/models/instance_segmentation/edgecrafter_inst.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py index 993c88edc30..84c3bc1d99f 100644 --- a/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py +++ b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py @@ -353,7 +353,7 @@ def _optimization_config(self) -> dict[str, Any]: @staticmethod def _make_is_export_prediction( - _inputs: SampleBatch, + _inputs: SampleBatch | Tensor, result: dict[str, Tensor], ) -> dict[str, Tensor]: """Reformat deploy-mode dict for MaskRCNN-compatible ONNX export. 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 index a667aca1344..5019b66d364 100644 --- a/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py +++ b/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py @@ -246,6 +246,7 @@ def forward_export( 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/backend/lightning/models/instance_segmentation/edgecrafter_inst.py b/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py index 735e9bb8e5c..fa63f2d70be 100644 --- a/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py +++ b/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py @@ -152,7 +152,7 @@ def forward_for_tracing(self, inputs: torch.Tensor) -> dict[str, Any]: """ 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) # type: ignore[arg-type] + return self._make_is_export_prediction(inputs, result) @property def _default_preprocessing_params(self) -> dict[str, DataInputParams]: From 8ba934c8d367d4fb4ed15ad263b7d0ae4e1a8c1f Mon Sep 17 00:00:00 2001 From: kprokofi Date: Tue, 7 Jul 2026 18:20:56 +0900 Subject: [PATCH 13/14] fixes of losses and num_points --- .../backend/lightning/models/detection/heads/ec_decoder.py | 2 +- .../backend/lightning/models/detection/losses/dfine_loss.py | 2 +- .../backend/lightning/models/detection/losses/ec_loss.py | 1 - 3 files changed, 2 insertions(+), 3 deletions(-) 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 index 5019b66d364..df6d27fa3df 100644 --- a/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py +++ b/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py @@ -547,7 +547,7 @@ def __init__( eval_spatial_size: tuple[int, int] | list[int] = (640, 640), num_queries: int = 300, num_levels: int = 3, - num_points: int | list[int] = 4, + num_points: int | list[int] = (3, 6, 3), dropout: float = 0.0, activation: str = "silu", num_denoising: int = 100, 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 index 445dd7b5da3..3452a7dbc18 100644 --- a/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py +++ b/library/src/getitune/backend/lightning/models/detection/losses/ec_loss.py @@ -121,7 +121,6 @@ def loss_masks( def _available_losses(self) -> tuple[Callable]: # type: ignore[return-value] return ( # pyrefly: ignore[bad-return] self.loss_boxes, - self.loss_labels_vfl, self.loss_labels_mal, self.loss_local, self.loss_masks, From bbb56d1060397c722264d5f403acd53c7e9c39a6 Mon Sep 17 00:00:00 2001 From: kprokofi Date: Tue, 7 Jul 2026 21:09:01 +0900 Subject: [PATCH 14/14] do clean up --- .../models/common/edgecrafter_mixin.py | 152 ++++++----- .../models/common/layers/vit_blocks.py | 142 ++++++++++ .../models/detection/backbones/ecvit.py | 249 +----------------- .../models/detection/backbones/vit_tiny.py | 118 +-------- .../lightning/models/detection/edgecrafter.py | 15 +- .../models/detection/heads/ec_decoder.py | 234 +++------------- .../instance_segmentation/edgecrafter_inst.py | 26 +- .../instance_segmentation/heads/__init__.py | 2 + .../heads/ec_segmentation_head.py | 119 +++++++++ .../detection/backbones/test_vit_tiny.py | 7 +- 10 files changed, 439 insertions(+), 625 deletions(-) create mode 100644 library/src/getitune/backend/lightning/models/common/layers/vit_blocks.py create mode 100644 library/src/getitune/backend/lightning/models/instance_segmentation/heads/ec_segmentation_head.py diff --git a/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py index 84c3bc1d99f..e8af83bfe54 100644 --- a/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py +++ b/library/src/getitune/backend/lightning/models/common/edgecrafter_mixin.py @@ -9,10 +9,10 @@ from __future__ import annotations import copy -from typing import Any, ClassVar +from typing import TYPE_CHECKING, Any, Callable, ClassVar, Protocol, cast import torch -from torch import Tensor +from torch import Tensor, nn from torchvision import tv_tensors from torchvision.ops import box_convert from torchvision.tv_tensors import BoundingBoxFormat @@ -26,6 +26,53 @@ 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. @@ -33,18 +80,32 @@ class EdgeCrafterMixin: 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 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]]] = { @@ -75,83 +136,55 @@ class EdgeCrafterMixin: } def _build_ec_model( - self, + self: _EdgeCrafterHost, num_classes: int, *, - with_seg: bool = False, 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). + 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 for ECSeg). - 4. Build :class:`ECCriterion` with mask losses added for ECSeg. + 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. - with_seg: When ``True``, builds the ECSeg variant (adds segmentation - head and mask losses). 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] # type: ignore[attr-defined] - backbone_key = "seg_backbone_name" if with_seg else "backbone_name" + cfg = self._EC_MODEL_CFGS[self.model_name] - if self.data_input_params.input_size is None: # type: ignore[attr-defined] + 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 # type: ignore[attr-defined] + input_size: tuple[int, int] = self.data_input_params.input_size - backbone = ECViTAdapter(model_name=cfg[backbone_key], proj_dim=cfg["proj_dim"]) - encoder = HybridEncoder(model_name=self.model_name) # type: ignore[attr-defined] + 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, # type: ignore[attr-defined] + model_name=self.model_name, num_classes=num_classes, eval_spatial_size=input_size, - mask_downsample_ratio=4 if with_seg else None, + mask_downsample_ratio=self._mask_downsample_ratio, ) - if with_seg: - weight_dict: 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: dict[str, int | float] | None = { - "cost_class": 2, - "cost_bbox": 1, - "cost_giou": 1, - "cost_mask": 5, - "cost_dice": 5, - } - else: - weight_dict = { - "loss_mal": 1.0, - "loss_bbox": 5.0, - "loss_giou": 2.0, - "loss_fgl": 0.15, - "loss_ddf": 1.5, - } - matcher_cost_dict = None - criterion = ECCriterion( - weight_dict=weight_dict, + weight_dict=self._loss_weights, alpha=0.75, gamma=1.5, reg_max=32, num_classes=num_classes, - matcher_cost_dict=matcher_cost_dict, + matcher_cost_dict=self._matcher_cost_dict, ) backbone_lr = backbone_lr if backbone_lr is not None else cfg["backbone_lr"] @@ -168,15 +201,15 @@ def _build_ec_model( criterion=criterion, num_classes=num_classes, optimizer_configuration=optimizer_configuration, - multi_scale=self.multi_scale, # type: ignore[attr-defined] + 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") # type: ignore[attr-defined] + load_checkpoint(model, self._pretrained_weights[self.model_name], map_location="cpu") return model def _customize_inputs( # pyrefly: ignore[bad-override] - self, + self: _EdgeCrafterHost, entity: SampleBatch, ) -> dict[str, Any]: """Convert getitune :class:`SampleBatch` to EdgeCrafter input format. @@ -227,7 +260,7 @@ def _customize_inputs( # pyrefly: ignore[bad-override] targets.append(target) - if self.explain_mode: # type: ignore[attr-defined] + if self.explain_mode: return {"entity": entity} return { @@ -236,7 +269,7 @@ def _customize_inputs( # pyrefly: ignore[bad-override] } def _customize_outputs( # pyrefly: ignore[bad-override] - self, + self: _EdgeCrafterHost, outputs: dict[str, Any] | tuple | list, inputs: SampleBatch, ) -> PredictionBatch | BatchLoss: @@ -254,7 +287,7 @@ def _customize_outputs( # pyrefly: ignore[bad-override] Returns: :class:`BatchLoss` during training, :class:`PredictionBatch` during inference. """ - if self.training: # type: ignore[attr-defined] + if self.training: if not isinstance(outputs, dict): msg = f"Expected dict during training, got {type(outputs)}" raise TypeError(msg) @@ -272,7 +305,8 @@ def _customize_outputs( # pyrefly: ignore[bad-override] return losses original_sizes = [img_info.ori_shape for img_info in inputs.imgs_info] # type: ignore[union-attr] - result = self.model.postprocess(outputs, original_sizes) # type: ignore[attr-defined] + model = cast("ECDETRDetector", self.model) + result = model.postprocess(cast("dict[str, Tensor]", outputs), original_sizes) prediction_kwargs: dict[str, Any] = { "images": inputs.images, @@ -308,7 +342,7 @@ def _customize_outputs( # pyrefly: ignore[bad-override] return PredictionBatch(**prediction_kwargs) def configure_optimizers( # pyrefly: ignore[bad-override] - self, + self: _EdgeCrafterHost, ) -> tuple[list[torch.optim.Optimizer], list[dict[str, Any]]]: """Configure optimizer and learning-rate schedulers. @@ -322,11 +356,11 @@ def configure_optimizers( # pyrefly: ignore[bad-override] from getitune.backend.lightning.models.detection.rtdetr import RTDETR param_groups = RTDETR._get_optim_params( # noqa: SLF001 - self.model.optimizer_configuration, # type: ignore[attr-defined] - self.model, # type: ignore[attr-defined] + cast("ECDETRDetector", self.model).optimizer_configuration, + self.model, ) - optimizer = self.optimizer_callable(param_groups) # type: ignore[attr-defined] - schedulers = self.scheduler_callable(optimizer) # type: ignore[attr-defined] + 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] 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/detection/backbones/ecvit.py b/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py index 4ac93a6433f..310e3867aa2 100644 --- a/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py +++ b/library/src/getitune/backend/lightning/models/detection/backbones/ecvit.py @@ -11,186 +11,21 @@ from __future__ import annotations import math -import warnings from functools import partial -from typing import Callable, ClassVar, Literal +from typing import Callable, ClassVar -import numpy as np 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"] -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _trunc_normal_(tensor: torch.Tensor, mean: float, std: float, a: float, b: float) -> torch.Tensor: - def _norm_cdf(x: float) -> float: - return (1.0 + math.erf(x / math.sqrt(2.0))) / 2.0 - - if (mean < a - 2 * std) or (mean > b + 2 * std): - warnings.warn( - "mean is more than 2 std from [a, b] in nn.init.trunc_normal_. " - "The distribution of values may be incorrect.", - stacklevel=2, - ) - with torch.no_grad(): - lo = _norm_cdf((a - mean) / std) - hi = _norm_cdf((b - mean) / std) - tensor.uniform_(2 * lo - 1, 2 * hi - 1) - tensor.erfinv_() - tensor.mul_(std * math.sqrt(2.0)) - tensor.add_(mean) - tensor.clamp_(min=a, max=b) - return tensor - - -def trunc_normal_( - tensor: torch.Tensor, mean: float = 0.0, std: float = 1.0, a: float = -2.0, b: float = 2.0 -) -> torch.Tensor: - """Fill tensor with truncated normal values.""" - return _trunc_normal_(tensor, mean, std, a, b) - - -def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor: - """Stochastic depth drop path.""" - if drop_prob == 0.0 or not training: - return x - keep_prob = 1 - drop_prob - shape = (x.shape[0],) + (1,) * (x.ndim - 1) - random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) - return x.div(keep_prob) * random_tensor.floor() - - -class DropPath(nn.Module): - """Drop paths (stochastic depth) per sample.""" - - def __init__(self, drop_prob: float = 0.0) -> None: - super().__init__() - self.drop_prob = drop_prob - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass.""" - return drop_path(x, self.drop_prob, self.training) - - -# --------------------------------------------------------------------------- -# RoPE -# --------------------------------------------------------------------------- - - -class RopePositionEmbedding(nn.Module): - """2-D Rotary Position Embedding for ViT. - - Args: - embed_dim: Transformer embedding dimension. - num_heads: Number of attention heads. - base: RoPE base period (used when min/max_period are None). - min_period: Minimum period (used together with max_period). - max_period: Maximum period (used together with min_period). - normalize_coords: Coordinate normalisation strategy. - shift_coords: Optional coordinate shift jitter during training. - jitter_coords: Optional coordinate scale jitter during training. - rescale_coords: Optional global scale jitter during training. - """ - - def __init__( - self, - embed_dim: int, - *, - num_heads: int, - base: float | None = 100.0, - min_period: float | None = None, - max_period: float | None = None, - normalize_coords: Literal["min", "max", "separate"] = "separate", - shift_coords: float | None = None, - jitter_coords: float | None = None, - rescale_coords: float | None = None, - ) -> None: - super().__init__() - head_dim = embed_dim // num_heads - if head_dim % 4 != 0: - msg = "Head dimension must be divisible by 4 for 2D RoPE" - raise ValueError(msg) - both = min_period is not None and max_period is not None - if (base is None and not both) or (base is not None and both): - msg = "Either `base` or both `min_period`+`max_period` must be provided." - raise ValueError(msg) - - self.base = base - self.min_period = min_period - self.max_period = max_period - self.D_head = head_dim - self.normalize_coords = normalize_coords - self.shift_coords = shift_coords - self.jitter_coords = jitter_coords - self.rescale_coords = rescale_coords - self.register_buffer("periods", torch.empty(head_dim // 4), persistent=True) - self._init_weights() - - def forward(self, *, H: int, W: int) -> tuple[torch.Tensor, torch.Tensor]: # noqa: N803 - """Compute sin/cos RoPE tables for an HxW feature map.""" - device = self.periods.device # type: ignore[union-attr] - dtype = torch.get_default_dtype() - dd: dict = {"device": device, "dtype": dtype} - - if self.normalize_coords == "max": - m = max(H, W) - coords_h = torch.arange(0.5, H, **dd) / m - coords_w = torch.arange(0.5, W, **dd) / m - elif self.normalize_coords == "separate": - coords_h = torch.arange(0.5, H, **dd) / H - coords_w = torch.arange(0.5, W, **dd) / W - else: # min - m = min(H, W) - coords_h = torch.arange(0.5, H, **dd) / m - coords_w = torch.arange(0.5, W, **dd) / m - - coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1) - coords = coords.flatten(0, 1) - coords = 2.0 * coords - 1.0 - - if self.training and self.shift_coords is not None: - coords = coords + torch.empty(2, **dd).uniform_(-self.shift_coords, self.shift_coords)[None, :] - if self.training and self.jitter_coords is not None: - j = torch.empty(2, **dd).uniform_(-np.log(self.jitter_coords), np.log(self.jitter_coords)).exp() - coords = coords * j[None, :] - if self.training and self.rescale_coords is not None: - r = torch.empty(1, **dd).uniform_(-np.log(self.rescale_coords), np.log(self.rescale_coords)).exp() - coords = coords * r - - angles = 2 * math.pi * coords[:, :, None] / self.periods[None, None, :] # type: ignore[index] - angles = angles.flatten(1, 2).repeat(1, 2) - return torch.sin(angles).unsqueeze(0).unsqueeze(0), torch.cos(angles).unsqueeze(0).unsqueeze(0) - - def _init_weights(self) -> None: - """Initialise period buffer.""" - device: torch.device = self.periods.device # type: ignore[union-attr] - dtype = torch.get_default_dtype() - if self.base is not None: - periods = self.base ** (2 * torch.arange(self.D_head // 4, device=device, dtype=dtype) / (self.D_head // 2)) - else: - base_ratio = self.max_period / self.min_period # type: ignore[operator] - exponents = torch.linspace(0, 1, self.D_head // 4, device=device, dtype=dtype) - periods = self.max_period * (base_ratio ** (exponents - 1)) # type: ignore[operator] - self.periods.data.copy_(periods) # type: ignore[union-attr] - - -def _rotate_half(x: torch.Tensor) -> torch.Tensor: - x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -def _apply_rope(x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor: - return (x * cos) + (_rotate_half(x) * sin) - - # --------------------------------------------------------------------------- # Patch embedding # --------------------------------------------------------------------------- @@ -274,75 +109,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.fc2(self.drop(self.act(self.fc1(x)))) -class Attention(nn.Module): - """Multi-head self-attention with optional RoPE.""" - - 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 - 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: torch.Tensor, rope_sincos: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.Tensor: - """Forward pass.""" - B, N, C = x.shape # noqa: N806 - 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 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.""" - - def __init__( - self, - dim: int, - num_heads: int, - ffn_ratio: float = 4.0, - qkv_bias: bool = False, - drop: float = 0.0, - attn_drop: float = 0.0, - drop_path_rate: float = 0.0, - act_layer: type[nn.Module] = nn.GELU, - norm_layer: type[nn.Module] = nn.LayerNorm, - ) -> None: - super().__init__() - self.norm1 = norm_layer(dim) # type: ignore[call-arg] - self.attn = Attention(dim, num_heads, qkv_bias=qkv_bias, attn_drop=attn_drop, proj_drop=drop) - self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity() - self.norm2 = norm_layer(dim) # type: ignore[call-arg] - self.mlp = Mlp(dim, int(dim * ffn_ratio), act_layer=act_layer, drop=drop) - - def forward(self, x: torch.Tensor, rope_sincos: tuple[torch.Tensor, torch.Tensor] | None = None) -> torch.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))) - - class VisionTransformer(nn.Module): """ViT backbone used by EC-ViT models. @@ -391,13 +157,14 @@ def __init__( Block( dim=embed_dim, num_heads=num_heads, - ffn_ratio=ffn_ratio, + mlp_ratio=ffn_ratio, qkv_bias=qkv_bias, drop=drop_rate, attn_drop=attn_drop_rate, - drop_path_rate=dpr[i], + drop_path=dpr[i], act_layer=nn.GELU, norm_layer=norm_layer, # type: ignore[arg-type] + mlp_layer=Mlp, ) for i in range(depth) ] @@ -433,7 +200,7 @@ def forward(self, x: torch.Tensor) -> list[torch.Tensor]: 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) + rope = self.rope_embed(h=H, w=W) outs: list[torch.Tensor] = [] for i, blk in enumerate(self.blocks): 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/edgecrafter.py b/library/src/getitune/backend/lightning/models/detection/edgecrafter.py index 3a15d796b1d..42fd22d6a64 100644 --- a/library/src/getitune/backend/lightning/models/detection/edgecrafter.py +++ b/library/src/getitune/backend/lightning/models/detection/edgecrafter.py @@ -34,8 +34,8 @@ class EdgeCrafter(EdgeCrafterMixin, LightningDetectionModel): # pyrefly: ignore deployment via task-specialised distillation. Four sizes are available: S (small), M (medium), L (large), and X (extra-large). - Original paper / repository: - https://github.com/Intellindust-AI-Lab/EdgeCrafter + 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` @@ -67,6 +67,15 @@ class EdgeCrafter(EdgeCrafterMixin, LightningDetectionModel): # pyrefly: ignore "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__( @@ -110,7 +119,7 @@ def _create_model(self, num_classes: int | None = None) -> ECDETRDetector: Configured :class:`ECDETRDetector`. """ num_classes = num_classes if num_classes is not None else self.num_classes - return self._build_ec_model(num_classes, with_seg=False, backbone_lr=self.backbone_lr) + 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. 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 index df6d27fa3df..8c84724308d 100644 --- a/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py +++ b/library/src/getitune/backend/lightning/models/detection/heads/ec_decoder.py @@ -14,7 +14,7 @@ import copy from collections import OrderedDict -from typing import Any, cast +from typing import TYPE_CHECKING, Any, Callable, cast import torch import torch.nn.functional as F # noqa: N812 @@ -22,6 +22,10 @@ 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, ) @@ -29,6 +33,14 @@ 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"] # --------------------------------------------------------------------------- @@ -66,191 +78,6 @@ }, } -# --------------------------------------------------------------------------- -# Sub-modules -# --------------------------------------------------------------------------- - - -class MLP(nn.Module): - """Simple multi-layer perceptron.""" - - def __init__( - self, - input_dim: int, - hidden_dim: int, - output_dim: int, - num_layers: int = 3, - act: str = "silu", - ) -> None: - super().__init__() - self.num_layers = num_layers - h = [hidden_dim] * (num_layers - 1) - self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim, *h], [*h, output_dim])) - self._act = nn.SiLU() if act == "silu" else nn.ReLU() - - def forward(self, x: Tensor) -> Tensor: - """Forward pass.""" - for i, layer in enumerate(self.layers): - x = self._act(layer(x)) if i < self.num_layers - 1 else layer(x) - return x - - -class Gate(nn.Module): - """Gated cross-attention fusion (LayerNorm variant).""" - - def __init__(self, d_model: int) -> None: - super().__init__() - self.gate = nn.Linear(2 * d_model, 2 * d_model) - bias = bias_init_with_prob(0.5) - init.constant_(self.gate.bias, bias) - init.constant_(self.gate.weight, 0) - self.norm = nn.LayerNorm(d_model) - - def forward(self, x1: Tensor, x2: Tensor) -> Tensor: - """Forward pass.""" - gates = torch.sigmoid(self.gate(torch.cat([x1, x2], dim=-1))) - g1, g2 = gates.chunk(2, dim=-1) - return self.norm(g1 * x1 + g2 * x2) - - -class Integral(nn.Module): - """Distribution-to-distance integral (non-uniform weighting).""" - - def __init__(self, reg_max: int = 32) -> None: - super().__init__() - self.reg_max = reg_max - - def forward(self, x: Tensor, project: Tensor) -> Tensor: - """Forward pass.""" - shape = x.shape - x = F.softmax(x.reshape(-1, self.reg_max + 1), dim=1) - x = F.linear(x, project.to(x.device)).reshape(-1, 4) - return x.reshape([*list(shape[:-1]), -1]) - - -class LQE(nn.Module): - """Localization Quality Estimator.""" - - def __init__(self, k: int, hidden_dim: int, num_layers: int, reg_max: int, act: str = "silu") -> None: - super().__init__() - self.k = k - self.reg_max = reg_max - self.reg_conf = MLP(4 * (k + 1), hidden_dim, 1, num_layers, act=act) - _reg_last = cast("nn.Linear", self.reg_conf.layers[-1]) - init.constant_(_reg_last.bias, 0) - init.constant_(_reg_last.weight, 0) - - def forward(self, scores: Tensor, pred_corners: Tensor) -> Tensor: - """Forward pass.""" - B, L, _ = pred_corners.size() # noqa: N806 - prob = F.softmax(pred_corners.reshape(B, L, 4, self.reg_max + 1), dim=-1) - prob_topk, _ = prob.topk(self.k, dim=-1) - stat = torch.cat([prob_topk, prob_topk.mean(dim=-1, keepdim=True)], dim=-1) - quality_score = self.reg_conf(stat.reshape(B, L, -1)) - return scores + quality_score - - -# --------------------------------------------------------------------------- -# Segmentation Head -# --------------------------------------------------------------------------- - - -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] - - # --------------------------------------------------------------------------- # Transformer Decoder Layer # --------------------------------------------------------------------------- @@ -273,7 +100,7 @@ class ECTransformerDecoderLayer(nn.Module): 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 name. + activation: Activation layer class. """ def __init__( @@ -285,7 +112,7 @@ def __init__( n_levels: int = 3, n_points: int | list[int] = 4, layer_scale: float | None = None, - activation: str = "silu", + activation: Callable[..., nn.Module] = nn.SiLU, ) -> None: super().__init__() @@ -305,9 +132,8 @@ def __init__( self.gateway = Gate(d_model) # FFN - _act = nn.SiLU if activation == "silu" else nn.ReLU self.linear1 = nn.Linear(d_model, dim_feedforward) - self.activation = _act() + self.activation = activation() self.dropout3 = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model) self.dropout4 = nn.Dropout(dropout) @@ -367,7 +193,7 @@ def __init__( up: Tensor, eval_idx: int = -1, layer_scale: int = 1, - act: str = "silu", + act: Callable[..., nn.Module] = nn.SiLU, ) -> None: super().__init__() self.hidden_dim = hidden_dim @@ -384,7 +210,9 @@ def __init__( + [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, act=act)) for _ in range(num_layers)]) + self.lqe_layers = nn.ModuleList( + [copy.deepcopy(LQE(4, 64, 2, reg_max, activation=act)) for _ in range(num_layers)], + ) def _value_op( self, @@ -529,7 +357,7 @@ class ECTransformer(nn.Module): num_levels: Number of encoder feature levels. num_points: Deformable attention points per level (int or list). dropout: Dropout rate. - activation: FFN activation (``"silu"`` or ``"relu"``). + 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. @@ -547,9 +375,9 @@ def __init__( 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), + num_points: int | list[int] = [3, 6, 3], # noqa: B006 dropout: float = 0.0, - activation: str = "silu", + activation: Callable[..., nn.Module] = nn.SiLU, num_denoising: int = 100, label_noise_ratio: float = 0.5, box_noise_scale: float = 1.0, @@ -609,7 +437,11 @@ def __init__( seg_head: SegmentationHead | None = None if mask_downsample_ratio is not None: - seg_head = SegmentationHead( + 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, @@ -641,9 +473,9 @@ def __init__( # Encoder output heads self.enc_score_head = nn.Linear(hidden_dim, num_classes) - self.enc_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, act=activation) - self.query_pos_head = MLP(4, hidden_dim, hidden_dim, 3, act=activation) - self.pre_bbox_head = MLP(hidden_dim, hidden_dim, 4, 3, act=activation) + 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 @@ -656,8 +488,8 @@ def __init__( ) scaled_dim = round(layer_scale * hidden_dim) - dec_bbox_head = MLP(hidden_dim, hidden_dim, 4 * (reg_max + 1), 3, act=activation) - wide_bbox_head = MLP(scaled_dim, scaled_dim, 4 * (reg_max + 1), 3, act=activation) + 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)] + [ 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 index fa63f2d70be..aae38f213ae 100644 --- a/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py +++ b/library/src/getitune/backend/lightning/models/instance_segmentation/edgecrafter_inst.py @@ -35,8 +35,8 @@ class EdgeCrafterInst(EdgeCrafterMixin, LightningInstanceSegModel): # pyrefly: (``ecseg_vitt/vittplus/vits/vitsplus``) which are trained jointly for detection and segmentation. - Original paper / repository: - https://github.com/Intellindust-AI-Lab/EdgeCrafter + Original repository: https://github.com/Intellindust-AI-Lab/EdgeCrafter + Paper: https://arxiv.org/abs/2603.18739 Args: label_info: Information about the labels. @@ -62,6 +62,26 @@ class EdgeCrafterInst(EdgeCrafterMixin, LightningInstanceSegModel): # pyrefly: "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__( @@ -105,7 +125,7 @@ def _create_model(self, num_classes: int | None = None) -> ECDETRDetector: 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, with_seg=True, backbone_lr=self.backbone_lr) + return self._build_ec_model(num_classes, backbone_lr=self.backbone_lr) @property def _export_parameters(self) -> TaskLevelExportParameters: 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/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):