nick-localhost/Sign-language-detection

1import torch
2from torch import nn
3from torchvision.models import resnet50, ResNet50_Weights
4import sys 
5from colorama import Fore 
6from utils.logger import get_logger
7from utils.rich_handlers import ModelHandler
8from torchinfo import summary
9import sys 
10import math
11
12
13def _get_1d_sincos_pos_embed(length: int, dim: int, temperature: float = 10000.0, device=None):
14    assert dim % 2 == 0
15    position = torch.arange(length, device=device, dtype=torch.float32).unsqueeze(1)  # (L,1)
16    div_term = torch.exp(
17        torch.arange(0, dim, 2, device=device, dtype=torch.float32) * (-math.log(temperature) / dim)
18    )  # (dim/2)
19    pe = torch.zeros(length, dim, device=device, dtype=torch.float32)
20    pe[:, 0::2] = torch.sin(position * div_term)
21    pe[:, 1::2] = torch.cos(position * div_term)
22    return pe  # (L, dim)
23
24
25def build_2d_sincos_position_embedding(height: int, width: int, dim: int, device=None):
26    """Create 2D sine-cos positional encoding of shape (1, H*W, dim).
27    Half dims for Y, half for X.
28    """
29    assert dim % 2 == 0, "positional dim must be even"
30    dim_half = dim // 2
31    pe_y = _get_1d_sincos_pos_embed(height, dim_half, device=device)  # (H, dim/2)
32    pe_x = _get_1d_sincos_pos_embed(width, dim_half, device=device)   # (W, dim/2)
33    # Combine to (H, W, dim)
34    pos = torch.zeros(height, width, dim, device=device, dtype=torch.float32)
35    pos[:, :, :dim_half] = pe_y[:, None, :].expand(-1, width, -1)
36    pos[:, :, dim_half:] = pe_x[None, :, :].expand(height, -1, -1)
37    pos = pos.view(1, height * width, dim)  # (1, H*W, dim)
38    return pos
39
40
41class DETR(nn.Module):
42    def __init__(self, num_classes, hidden_dim=256, nheads=8,
43                 num_encoder_layers=1, num_decoder_layers=1, num_queries=25):
44        super().__init__()
45        
46        # Initialize logger and model handler
47        self.logger = get_logger("model")
48        self.model_handler = ModelHandler()
49        
50        # Log model configuration
51        model_config = {
52            "Model Type": "DETR (Detection Transformer)",
53            "Number of Classes": num_classes,
54            "Hidden Dimension": hidden_dim,
55            "Attention Heads": nheads,
56            "Encoder Layers": num_encoder_layers,
57            "Decoder Layers": num_decoder_layers,
58            "Object Queries": num_queries,
59            "Backbone": "ResNet-50 (ImageNet pretrained)"
60        }
61        self.model_handler.log_model_architecture(model_config)
62
63        # create ResNet-50 backbone
64        self.backbone = resnet50(weights=ResNet50_Weights.IMAGENET1K_V1)
65        self.backbone.fc = nn.Identity()
66
67        # create conversion layer
68        self.conv = nn.Conv2d(2048, hidden_dim, 1)
69
70        # create a default PyTorch transformer
71        self.transformer = nn.Transformer(
72            hidden_dim, nheads, num_encoder_layers, num_decoder_layers, batch_first=True, dropout=0.1)
73
74        # prediction heads, one extra class for predicting non-empty slots
75        # note that in baseline DETR linear_bbox layer is 3-layer MLP
76        self.linear_class = nn.Linear(hidden_dim, num_classes + 1)
77        self.linear_bbox = nn.Linear(hidden_dim, 4)
78
79        # number of object queries
80        self.num_queries = num_queries
81        # learned query positional encodings
82        self.query_pos = nn.Parameter(torch.randn(self.num_queries, hidden_dim))
83
84        # normalizations
85        self.norm_src = nn.LayerNorm(hidden_dim)
86        self.norm_tgt = nn.LayerNorm(hidden_dim)
87
88    def forward(self, inputs):
89        # propagate inputs through ResNet-50 up to avg-pool layer
90        x = self.backbone.conv1(inputs)
91        x = self.backbone.bn1(x)
92        x = self.backbone.relu(x)
93        x = self.backbone.maxpool(x)
94
95        x = self.backbone.layer1(x)
96        x = self.backbone.layer2(x)
97        x = self.backbone.layer3(x)
98        x = self.backbone.layer4(x)
99
100        # convert from 2048 to hidden_dim feature planes for the transformer
101        feat = self.conv(x)  # (b, d, Hf, Wf)
102        bsz, d_model, Hf, Wf = feat.shape
103        src = feat.flatten(2).permute(0, 2, 1)  # (b, Hf*Wf, d)
104
105        # dynamic 2D sine-cos positional encoding
106        pos = build_2d_sincos_position_embedding(Hf, Wf, d_model, device=feat.device)  # (1, Hf*Wf, d)
107        src = self.norm_src(src + pos)
108
109        # decoder target: zero content + learned query positional encodings
110        tgt = torch.zeros(bsz, self.num_queries, d_model, device=feat.device)
111        query_pos = self.query_pos.unsqueeze(0).expand(bsz, -1, -1)
112        tgt = self.norm_tgt(tgt + query_pos)
113
114        # propagate through the transformer
115        hs = self.transformer(src=src, tgt=tgt)  # (b, num_queries, d)
116
117        # finally project transformer outputs to class labels and bounding boxes
118        return {
119            'pred_logits': self.linear_class(hs),
120            'pred_boxes': self.linear_bbox(hs).sigmoid()
121        }
122    
123    def log_model_info(self):
124        """Log model parameter information."""
125        total_params = sum(p.numel() for p in self.parameters())
126        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
127        self.model_handler.log_parameters_count(total_params, trainable_params)
128        
129    def load_pretrained(self, checkpoint_path: str):
130        """Load pretrained weights with logging."""
131        try:
132            self.load_state_dict(torch.load(checkpoint_path))
133            self.model_handler.log_model_loading(checkpoint_path, success=True)
134        except Exception as e:
135            self.logger.error(f"Failed to load checkpoint: {str(e)}")
136            self.model_handler.log_model_loading(checkpoint_path, success=False)
137
138
139if __name__ == '__main__': 
140    model = DETR(num_classes=3)
141    summary(model, (5,3,224,224))