RTDETR update

docs/source/models/rt_detr.rst

@@ -1,6 +1,28 @@
 Real-Time Detection Transformer (RT-DETR)
 =========================================
 
+.. code-block:: python
+
+    from kornia.io import load_image
+    from kornia.models.detector.rtdetr import RTDETRDetectorBuilder
+
+    input_img = load_image(img_path)[None]  # Load image to BCHW
+
+    # NOTE: available models: 'rtdetr_r18vd', 'rtdetr_r34vd', 'rtdetr_r50vd_m', 'rtdetr_r50vd', 'rtdetr_r101vd'.
+    # NOTE: recommended image scales: [480, 512, 544, 576, 608, 640, 640, 640, 672, 704, 736, 768, 800]
+    detector = RTDETRDetectorBuilder.build("rtdetr_r18vd", image_size=640)
+
+    # get the output boxes
+    boxes = detector(input_img)
+
+    # draw the bounding boxes on the images directly.
+    output = detector.draw(input_img, output_type="pil")
+    output[0].save("Kornia-RTDETR-output.png")
+
+    # convert the whole model to ONNX directly
+    RTDETRDetectorBuilder.to_onnx("RTDETR-640.onnx", model_name="rtdetr_r18vd", image_size=640)
+
+
 .. card::
     :link: https://arxiv.org/abs/2304.08069
 

kornia/color/yuv.py

@@ -122,7 +122,7 @@ def yuv_to_rgb(image: Tensor) -> Tensor:
     if not isinstance(image, Tensor):
         raise TypeError(f"Input type is not a Tensor. Got {type(image)}")
 
-    if len(image.shape) < 3 or image.shape[-3] != 3:
+    if image.dim() < 3 or image.shape[-3] != 3:
         raise ValueError(f"Input size must have a shape of (*, 3, H, W). Got {image.shape}")
 
     y: Tensor = image[..., 0, :, :]

kornia/contrib/models/common.py

@@ -9,7 +9,16 @@
 
 class ConvNormAct(nn.Sequential):
     def __init__(
-        self, in_channels: int, out_channels: int, kernel_size: int, stride: int = 1, act: str = "relu", groups: int = 1
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        act: str = "relu",
+        groups: int = 1,
+        conv_naming: str = "conv",
+        norm_naming: str = "norm",
+        act_naming: str = "act",
     ) -> None:
         super().__init__()
         if kernel_size % 2 == 0:
@@ -23,9 +32,13 @@ def __init__(
             padding = 0
         else:
             padding = (kernel_size - 1) // 2
-        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, 1, groups, False)
-        self.norm = nn.BatchNorm2d(out_channels)
-        self.act = {"relu": nn.ReLU, "silu": nn.SiLU, "none": nn.Identity}[act](inplace=True)
+        conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding, 1, groups, False)
+        norm = nn.BatchNorm2d(out_channels)
+        activation = {"relu": nn.ReLU, "silu": nn.SiLU, "none": nn.Identity}[act](inplace=True)
+
+        self.__setattr__(conv_naming, conv)
+        self.__setattr__(norm_naming, norm)
+        self.__setattr__(act_naming, activation)
 
 
 # Lightly adapted from

kornia/contrib/models/rt_detr/architecture/hybrid_encoder.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import copy
 from typing import Optional
 
 import torch
@@ -80,15 +81,16 @@ class AIFI(Module):
     def __init__(self, embed_dim: int, num_heads: int, dim_feedforward: int, dropout: float = 0.0) -> None:
         super().__init__()
         self.self_attn = nn.MultiheadAttention(embed_dim, num_heads, dropout)  # NOTE: batch_first = False
-        self.dropout1 = nn.Dropout(dropout)
-        self.norm1 = nn.LayerNorm(embed_dim)
 
         self.linear1 = nn.Linear(embed_dim, dim_feedforward)
-        self.act = nn.GELU()
         self.dropout = nn.Dropout(dropout)
         self.linear2 = nn.Linear(dim_feedforward, embed_dim)
+
+        self.dropout1 = nn.Dropout(dropout)
         self.dropout2 = nn.Dropout(dropout)
+        self.norm1 = nn.LayerNorm(embed_dim)
         self.norm2 = nn.LayerNorm(embed_dim)
+        self.act = nn.GELU()
 
     def forward(self, x: Tensor) -> Tensor:
         # using post-norm
@@ -149,6 +151,20 @@ def build_2d_sincos_pos_emb(
         return pos_emb.unsqueeze(1)  # (H * W, 1, C)
 
 
+class TransformerEncoder(nn.Module):
+    def __init__(self, encoder_layer: nn.Module, num_layers: int) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([copy.deepcopy(encoder_layer) for _ in range(num_layers)])
+        self.num_layers = num_layers
+
+    def forward(self, src: Tensor) -> Tensor:  # NOTE: Missing src_mask: Tensor = None, pos_embed: Tensor = None
+        output = src
+        for layer in self.layers:
+            output = layer(output)
+
+        return output
+
+
 class CCFM(Module):
     def __init__(self, num_fmaps: int, hidden_dim: int, expansion: float = 1.0) -> None:
         super().__init__()
@@ -192,12 +208,20 @@ def forward(self, fmaps: list[Tensor]) -> list[Tensor]:
 class HybridEncoder(Module):
     def __init__(self, in_channels: list[int], hidden_dim: int, dim_feedforward: int, expansion: float = 1.0) -> None:
         super().__init__()
-        self.input_proj = nn.ModuleList([ConvNormAct(in_ch, hidden_dim, 1, act="none") for in_ch in in_channels])
-        self.aifi = AIFI(hidden_dim, 8, dim_feedforward)
+        self.input_proj = nn.ModuleList(
+            [
+                ConvNormAct(  # To align the naming strategy for the official weights
+                    in_ch, hidden_dim, 1, act="none", conv_naming="0", norm_naming="1", act_naming="2"
+                )
+                for in_ch in in_channels
+            ]
+        )
+        encoder_layer = AIFI(hidden_dim, 8, dim_feedforward)
+        self.encoder = nn.Sequential(TransformerEncoder(encoder_layer, 1))
         self.ccfm = CCFM(len(in_channels), hidden_dim, expansion)
 
     def forward(self, fmaps: list[Tensor]) -> list[Tensor]:
         projected_maps = [proj(fmap) for proj, fmap in zip(self.input_proj, fmaps)]
-        projected_maps[-1] = self.aifi(projected_maps[-1])
+        projected_maps[-1] = self.encoder(projected_maps[-1])
         new_fmaps = self.ccfm(projected_maps)
         return new_fmaps

kornia/contrib/models/rt_detr/architecture/resnet_d.py

@@ -6,6 +6,8 @@
 
 from __future__ import annotations
 
+from collections import OrderedDict
+
 from torch import nn
 
 from kornia.contrib.models.common import ConvNormAct
@@ -15,7 +17,9 @@
 
 def _make_shortcut(in_channels: int, out_channels: int, stride: int) -> Module:
     return (
-        nn.Sequential(nn.AvgPool2d(2, 2), ConvNormAct(in_channels, out_channels, 1, act="none"))
+        nn.Sequential(
+            OrderedDict([("pool", nn.AvgPool2d(2, 2)), ("conv", ConvNormAct(in_channels, out_channels, 1, act="none"))])
+        )
         if stride == 2
         else ConvNormAct(in_channels, out_channels, 1, act="none")
     )
@@ -28,14 +32,18 @@ def __init__(self, in_channels: int, out_channels: int, stride: int, shortcut: b
         KORNIA_CHECK(stride in {1, 2})
         super().__init__()
         self.convs = nn.Sequential(
-            ConvNormAct(in_channels, out_channels, 3, stride=stride),
-            ConvNormAct(out_channels, out_channels, 3, act="none"),
+            OrderedDict(
+                [
+                    ("branch2a", ConvNormAct(in_channels, out_channels, 3, stride=stride)),
+                    ("branch2b", ConvNormAct(out_channels, out_channels, 3, act="none")),
+                ]
+            )
         )
-        self.shortcut = nn.Identity() if shortcut else _make_shortcut(in_channels, out_channels, stride)
+        self.short = nn.Identity() if shortcut else _make_shortcut(in_channels, out_channels, stride)
         self.relu = nn.ReLU(inplace=True)
 
     def forward(self, x: Tensor) -> Tensor:
-        return self.relu(self.convs(x) + self.shortcut(x))
+        return self.relu(self.convs(x) + self.short(x))
 
 
 class BottleneckD(Module):
@@ -46,15 +54,25 @@ def __init__(self, in_channels: int, out_channels: int, stride: int, shortcut: b
         super().__init__()
         expanded_out_channels = out_channels * self.expansion
         self.convs = nn.Sequential(
-            ConvNormAct(in_channels, out_channels, 1),
-            ConvNormAct(out_channels, out_channels, 3, stride=stride),
-            ConvNormAct(out_channels, expanded_out_channels, 1, act="none"),
+            OrderedDict(
+                [
+                    ("branch2a", ConvNormAct(in_channels, out_channels, 1)),
+                    ("branch2b", ConvNormAct(out_channels, out_channels, 3, stride=stride)),
+                    ("branch2c", ConvNormAct(out_channels, expanded_out_channels, 1, act="none")),
+                ]
+            )
         )
-        self.shortcut = nn.Identity() if shortcut else _make_shortcut(in_channels, expanded_out_channels, stride)
+        self.short = nn.Identity() if shortcut else _make_shortcut(in_channels, expanded_out_channels, stride)
         self.relu = nn.ReLU(inplace=True)
 
     def forward(self, x: Tensor) -> Tensor:
-        return self.relu(self.convs(x) + self.shortcut(x))
+        return self.relu(self.convs(x) + self.short(x))
+
+
+class Block(nn.Sequential):
+    def __init__(self, blocks: Module) -> None:
+        super().__init__()
+        self.blocks = blocks
 
 
 class ResNetD(Module):
@@ -63,35 +81,43 @@ def __init__(self, n_blocks: list[int], block: type[BasicBlockD | BottleneckD])
         super().__init__()
         in_channels = 64
         self.conv1 = nn.Sequential(
-            ConvNormAct(3, in_channels // 2, 3, stride=2),
-            ConvNormAct(in_channels // 2, in_channels // 2, 3),
-            ConvNormAct(in_channels // 2, in_channels, 3),
-            nn.MaxPool2d(3, stride=2, padding=1),
+            OrderedDict(
+                [
+                    ("conv1_1", ConvNormAct(3, in_channels // 2, 3, stride=2)),
+                    ("conv1_2", ConvNormAct(in_channels // 2, in_channels // 2, 3)),
+                    ("conv1_3", ConvNormAct(in_channels // 2, in_channels, 3)),
+                    ("pool", nn.MaxPool2d(3, stride=2, padding=1)),
+                ]
+            )
         )
 
-        self.res2, in_channels = self.make_stage(in_channels, 64, 1, n_blocks[0], block)
-        self.res3, in_channels = self.make_stage(in_channels, 128, 2, n_blocks[1], block)
-        self.res4, in_channels = self.make_stage(in_channels, 256, 2, n_blocks[2], block)
-        self.res5, in_channels = self.make_stage(in_channels, 512, 2, n_blocks[3], block)
+        res2, in_channels = self.make_stage(in_channels, 64, 1, n_blocks[0], block)
+        res3, in_channels = self.make_stage(in_channels, 128, 2, n_blocks[1], block)
+        res4, in_channels = self.make_stage(in_channels, 256, 2, n_blocks[2], block)
+        res5, in_channels = self.make_stage(in_channels, 512, 2, n_blocks[3], block)
+
+        self.res_layers = nn.ModuleList([res2, res3, res4, res5])
 
         self.out_channels = [ch * block.expansion for ch in [128, 256, 512]]
 
     @staticmethod
     def make_stage(
         in_channels: int, out_channels: int, stride: int, n_blocks: int, block: type[BasicBlockD | BottleneckD]
     ) -> tuple[Module, int]:
-        stage = nn.Sequential(
-            block(in_channels, out_channels, stride, False),
-            *[block(out_channels * block.expansion, out_channels, 1, True) for _ in range(n_blocks - 1)],
+        stage = Block(
+            nn.Sequential(
+                block(in_channels, out_channels, stride, False),
+                *[block(out_channels * block.expansion, out_channels, 1, True) for _ in range(n_blocks - 1)],
+            )
         )
         return stage, out_channels * block.expansion
 
     def forward(self, x: Tensor) -> list[Tensor]:
         x = self.conv1(x)
-        res2 = self.res2(x)
-        res3 = self.res3(res2)
-        res4 = self.res4(res3)
-        res5 = self.res5(res4)
+        res2 = self.res_layers[0](x)
+        res3 = self.res_layers[1](res2)
+        res4 = self.res_layers[2](res3)
+        res5 = self.res_layers[3](res4)
         return [res3, res4, res5]
 
     @staticmethod

kornia/contrib/models/rt_detr/architecture/rtdetr_head.py

@@ -3,6 +3,7 @@
 
 from __future__ import annotations
 
+import copy
 from typing import Optional
 
 import torch
@@ -192,14 +193,10 @@ def forward(
         return out
 
 
-class TransformerDecoder:
-    def __init__(self, hidden_dim: int, decoder_layers: nn.ModuleList, num_layers: int, eval_idx: int = -1) -> None:
+class TransformerDecoder(Module):
+    def __init__(self, hidden_dim: int, decoder_layer: nn.Module, num_layers: int, eval_idx: int = -1) -> None:
         super().__init__()
-        self.layers = decoder_layers
-        # TODO: come back to this later
-        # self.layers = nn.ModuleList([
-        #    copy.deepcopy(decoder_layer) for _ in range(num_layers)
-        # ])
+        self.layers = nn.ModuleList([copy.deepcopy(decoder_layer) for _ in range(num_layers)])
         self.hidden_dim = hidden_dim
         self.num_layers = num_layers
         self.eval_idx = eval_idx if eval_idx >= 0 else num_layers + eval_idx
@@ -272,13 +269,16 @@ def __init__(
         num_decoder_layers: int,
         num_heads: int = 8,
         num_decoder_points: int = 4,
-        # num_levels: int = 3,
+        num_levels: int = 3,
         dropout: float = 0.0,
+        num_denoising: int = 100,
     ) -> None:
         super().__init__()
         self.num_queries = num_queries
         # TODO: verify this is correct
-        self.num_levels = len(in_channels)
+        if len(in_channels) > num_levels:
+            raise ValueError(f"`num_levels` cannot be greater than {len(in_channels)}. Got {num_levels}.")
+        self.num_levels = num_levels
 
         # build the input projection layers
         self.input_proj = nn.ModuleList()
@@ -288,25 +288,23 @@ def __init__(
         # https://github.com/lyuwenyu/RT-DETR/blob/main/rtdetr_pytorch/src/zoo/rtdetr/rtdetr_decoder.py#L403-L410
 
         # NOTE: need to be integrated with the TransformerDecoderLayer
-        self.decoder_layers = nn.ModuleList(
-            [
-                TransformerDecoderLayer(
-                    embed_dim=hidden_dim,
-                    num_heads=num_heads,
-                    dropout=dropout,
-                    num_levels=len(in_channels),
-                    num_points=num_decoder_points,
-                )
-                for _ in range(num_decoder_layers)
-            ]
+        decoder_layer = TransformerDecoderLayer(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            dropout=dropout,
+            num_levels=self.num_levels,
+            num_points=num_decoder_points,
         )
 
         self.decoder = TransformerDecoder(
-            hidden_dim=hidden_dim, decoder_layers=self.decoder_layers, num_layers=num_decoder_layers
+            hidden_dim=hidden_dim, decoder_layer=decoder_layer, num_layers=num_decoder_layers
         )
 
         # denoising part
-        self.denoising_class_embed = nn.Embedding(num_classes, hidden_dim)  # not used in evaluation
+        if num_denoising > 0:
+            self.denoising_class_embed = nn.Embedding(
+                num_classes + 1, hidden_dim, padding_idx=num_classes
+            )  # not used in evaluation
 
         # decoder embedding
         self.query_pos_head = MLP(4, 2 * hidden_dim, hidden_dim, num_layers=2)
@@ -334,7 +332,7 @@ def forward(self, feats: Tensor) -> tuple[Tensor, Tensor]:
         )
 
         # decoder
-        out_bboxes, out_logits = self.decoder.forward(
+        out_bboxes, out_logits = self.decoder(
             target,
             init_ref_points_unact,
             memory,