Modern Object Detection: DETR and Transformer-Based Approaches

import torch
import torch.nn as nn
from torchvision.models import resnet50

class SimpleDETR(nn.Module):
    def __init__(self, num_classes=91, num_queries=100, hidden_dim=256):
        super().__init__()

        # CNN Backbone (ResNet-50, remove final layers)
        backbone = resnet50(pretrained=True)
        self.backbone = nn.Sequential(*list(backbone.children())[:-2])
        self.conv = nn.Conv2d(2048, hidden_dim, 1)  # Reduce channels

        # Transformer
        self.transformer = nn.Transformer(
            d_model=hidden_dim,
            nhead=8,
            num_encoder_layers=6,
            num_decoder_layers=6,
            dim_feedforward=2048,
            batch_first=True
        )

        # Object queries (learned embeddings)
        self.query_embed = nn.Embedding(num_queries, hidden_dim)

        # Prediction heads
        self.class_head = nn.Linear(hidden_dim, num_classes + 1)  # +1 for "no object"
        self.bbox_head = nn.Sequential(
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, 4),
            nn.Sigmoid()  # Normalized coordinates
        )

        # Positional encoding
        self.pos_encoder = PositionalEncoding2D(hidden_dim)

    def forward(self, x):
        # Extract features: (B, 3, H, W) -> (B, 2048, H/32, W/32)
        features = self.backbone(x)
        features = self.conv(features)  # (B, 256, H/32, W/32)

        B, C, H, W = features.shape

        # Add positional encoding
        pos = self.pos_encoder(H, W).to(x.device)
        features = features + pos

        # Flatten spatial dimensions: (B, 256, H*W) -> (B, H*W, 256)
        src = features.flatten(2).permute(0, 2, 1)

        # Object queries: (num_queries, 256) -> (B, num_queries, 256)
        queries = self.query_embed.weight.unsqueeze(0).repeat(B, 1, 1)

        # Transformer: encode image, decode with queries
        output = self.transformer(src, queries)  # (B, num_queries, 256)

        # Predictions
        class_logits = self.class_head(output)  # (B, 100, num_classes+1)
        bbox_pred = self.bbox_head(output)       # (B, 100, 4)

        return {'pred_logits': class_logits, 'pred_boxes': bbox_pred}

class PositionalEncoding2D(nn.Module):
    def __init__(self, hidden_dim, temperature=10000):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.temperature = temperature

    def forward(self, h, w):
        y_embed = torch.arange(h).unsqueeze(1).repeat(1, w)
        x_embed = torch.arange(w).unsqueeze(0).repeat(h, 1)

        dim_t = torch.arange(self.hidden_dim // 2)
        dim_t = self.temperature ** (2 * dim_t / self.hidden_dim)

        pos_x = x_embed.unsqueeze(-1) / dim_t
        pos_y = y_embed.unsqueeze(-1) / dim_t

        pos_x = torch.stack([pos_x.sin(), pos_x.cos()], dim=-1).flatten(-2)
        pos_y = torch.stack([pos_y.sin(), pos_y.cos()], dim=-1).flatten(-2)

        pos = torch.cat([pos_y, pos_x], dim=-1)
        return pos.permute(2, 0, 1).unsqueeze(0)  # (1, hidden_dim, H, W)

# Usage
model = SimpleDETR(num_classes=91)
image = torch.randn(2, 3, 800, 1066)
output = model(image)
print(f"Class predictions: {output['pred_logits'].shape}")  # (2, 100, 92)
print(f"Box predictions: {output['pred_boxes'].shape}")    # (2, 100, 4)