nick-localhost/Sign-language-detection

1import torch
2import torch.nn as nn
3import torch.nn.functional as F
4from scipy.optimize import linear_sum_assignment
5import sys
6from colorama import Fore 
7from utils.boxes import box_cxcywh_to_xyxy, generalized_box_iou
8
9class HungarianMatcher(nn.Module):
10    """
11    This class computes an assignment between y and predictions using 
12    the Hungarian algorithm (via scipy's linear_sum_assignment).
13    """
14    
15    def __init__(self, weight_dict:dict):
16        """
17        Args:
18            class_weighting: relative weight of the classification error in the matching cost
19            bbox_weighting: relative weight of the L1 error of the bounding box coordinates
20            giou_weighting: relative weight of the giou loss of the bounding box
21        """
22        super().__init__()
23        assert weight_dict.get('class_weighting') != None and weight_dict.get('bbox_weighting')  != None and weight_dict.get('giou_weighting')  != None, "Weight dict must contain weighting for all three losses, giou, class and bbox."
24        assert weight_dict.get('class_weighting') != 0 or weight_dict.get('bbox_weighting')  != 0 or weight_dict.get('giou_weighting')  != 0, "All loss weights cant be 0."
25
26        self.class_weighting = weight_dict.get('class_weighting') 
27        self.bbox_weighting = weight_dict.get('bbox_weighting') 
28        self.giou_weighting = weight_dict.get('giou_weighting') 
29
30    @torch.no_grad()
31    def forward(self, yhat, y):
32        """
33        Performs the matching
34        
35        Params:
36            yhat: dict containing:
37                "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with class logits
38                "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with predicted box coordinates
39                
40            y: list of y (len(y) = batch_size), where each target is a dict containing:
41                "labels": Tensor of dim [num_target_boxes] containing class indices
42                "boxes": Tensor of dim [num_target_boxes, 4] containing target box coordinates
43                
44        Returns:
45            A list of size batch_size, containing tuples of (index_i, index_j) where:
46                - index_i is the indices of selected predictions (in order)
47                - index_j is the indices of corresponding selected y (in order)
48        """
49        indices = []        
50        for batch_idx, target in enumerate(y):
51            # Get predictions for this batch
52            batch_logits = yhat["pred_logits"][batch_idx]  # [num_queries, num_classes]
53            batch_boxes = yhat["pred_boxes"][batch_idx]    # [num_queries, 4]
54            batch_prob = batch_logits.softmax(-1)          # [num_queries, num_classes]
55            
56            # Get targets for this batch
57            tgt_labels = target["labels"].to(torch.long)   # [num_targets]
58            tgt_boxes = target["boxes"].to(batch_boxes.dtype)  # [num_targets, 4]
59            
60            # Compute cost matrix for this batch only
61            # cost_class[i, j] = -probability that query i predicts class of target j
62            cost_class = -batch_prob[:, tgt_labels]  # [num_queries, num_targets]
63            cost_bbox = torch.cdist(batch_boxes, tgt_boxes, p=1)  # [num_queries, num_targets]
64            cost_giou = -generalized_box_iou(
65                box_cxcywh_to_xyxy(batch_boxes), 
66                box_cxcywh_to_xyxy(tgt_boxes)
67            )  # [num_queries, num_targets]
68            
69            # Final cost matrix for this batch
70            C_batch = (self.bbox_weighting * cost_bbox + 
71                      self.class_weighting * cost_class + 
72                      self.giou_weighting * cost_giou).cpu()
73            
74            # Solve assignment for this batch
75            ii, jj = linear_sum_assignment(C_batch)
76            indices.append(
77                (torch.as_tensor(ii, dtype=torch.int64), torch.as_tensor(jj, dtype=torch.int64))
78            )
79            
80        return indices
81
82class DETRLoss(nn.Module):
83    """
84    This class computes the loss for DETR.
85    The process happens in two steps:
86        1) Compute hungarian assignment between ground truth boxes and the yhat of the model
87        2) Supervise each pair of matched ground-truth / prediction (supervise class and box)
88    """
89    def __init__(self, num_classes, matcher, weight_dict, eos_coef):
90        """
91        Args:
92            num_classes: number of object categories, omitting the special no-object category
93            matcher: module able to compute a matching between y and proposals
94            weight_dict: dict containing as key the names of the losses and as values their relative weight.
95            eos_coef: relative classification weight applied to the no-object category
96        """
97        super().__init__()
98        self.num_classes = num_classes
99        self.matcher = matcher
100        self.weight_dict = weight_dict
101        self.eos_coef = eos_coef
102        empty_weight = torch.ones(self.num_classes + 1)
103        empty_weight[-1] = self.eos_coef
104        self.register_buffer('empty_weight', empty_weight)
105
106    def classification_loss(self, yhat, y, indices):
107        """Classification loss"""
108        assert 'pred_logits' in yhat
109        src_logits = yhat['pred_logits']
110        idx = self.get_matched_query_indices(indices)
111
112        target_classes_o = torch.cat([t["labels"][J] for t, (_, J) in zip(y, indices)])
113        target_classes = torch.full(src_logits.shape[:2], self.num_classes,
114                                  dtype=torch.int64, device=src_logits.device)
115        target_classes[idx] = target_classes_o
116        
117        loss_ce = F.cross_entropy(src_logits.transpose(1, 2), target_classes, self.empty_weight)
118        losses = {'loss_ce': loss_ce}
119
120        return losses
121
122    def box_loss(self, yhat, y, indices, num_boxes):
123        """Compute the losses related to the bounding boxes, the L1 regression loss and the GIoU loss
124           y dicts must contain the key "boxes" containing a tensor of dim [nb_target_boxes, 4]
125           The target boxes are expected in format (center_x, center_y, w, h), normalized by the image size.
126        """
127        assert 'pred_boxes' in yhat
128        idx = self.get_matched_query_indices(indices)
129        src_boxes = yhat['pred_boxes'][idx]
130        target_boxes = torch.cat([t['boxes'][i] for t, (_, i) in zip(y, indices)], dim=0)
131
132        loss_bbox = F.l1_loss(src_boxes, target_boxes, reduction='none')
133
134        losses = {}
135        losses['loss_bbox'] = loss_bbox.sum() / num_boxes
136
137        loss_giou = 1 - torch.diag(generalized_box_iou(
138            box_cxcywh_to_xyxy(src_boxes),
139            box_cxcywh_to_xyxy(target_boxes)))
140        losses['loss_giou'] = loss_giou.sum() / num_boxes
141        return losses
142
143    def get_matched_query_indices(self, indices):
144        # Takes in [[query_index], [which_tgt_pos]] x Per Batch -> Returns [batch_index] [query_index] 
145        batch_idx = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
146        src_idx = torch.cat([src for (src, _) in indices])
147        return batch_idx, src_idx
148
149    def forward(self, yhat, y):
150        """
151        This performs the loss computation.
152        
153        Args:
154            yhat: dict containing:
155                "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with class logits
156                "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with predicted box coordinates
157                
158            y: list of dicts, such that len(y) == batch_size.
159                     Each dict contains:
160                     "labels": Tensor of dim [num_target_boxes] containing class labels
161                     "boxes": Tensor of dim [num_target_boxes, 4] containing box coordinates
162        """
163
164        # Retrieve the matching between the yhat of the last layer and the y
165        indices = self.matcher(yhat, y)
166
167        # Compute the average number of target boxes across all nodes, for normalization purposes
168        device = next(iter(yhat.values())).device
169        # optional: coerce target dtypes defensively
170        y = [
171            {'labels': t['labels'].to(torch.long), 'boxes': t['boxes'].to(torch.float32)}
172            for t in y
173        ]
174
175        num_boxes = sum(len(t["labels"]) for t in y)
176        num_boxes = torch.as_tensor([num_boxes], dtype=torch.float, device=device).clamp(min=1)
177        
178        return {'labels':self.classification_loss(yhat, y, indices), 'boxes':self.box_loss(yhat, y, indices, num_boxes)}
179
180if __name__ == "__main__":
181    num_classes = 5
182    weight_dict = {'class_weighting': 1, 'bbox_weighting': 5, 'giou_weighting': 2}
183    matcher = HungarianMatcher(weight_dict)
184    criterion = DETRLoss(num_classes=num_classes, matcher=matcher, weight_dict=weight_dict, eos_coef=0.1)
185    
186    # # Example predictions and y
187    batch_size = 2
188    num_queries = 10
189    
190    # # Mock model yhat
191    yhat = {
192        'pred_logits': torch.randn(batch_size, num_queries, num_classes+1),
193        'pred_boxes': torch.rand(batch_size, num_queries, 4)
194    }
195    
196    # # Mock ground truth y
197    y = [
198        {
199            'labels': torch.tensor([1, 1, 2]),  
200            'boxes': torch.tensor([[0.5, 0.5, 0.2, 0.3],  
201                                 [0.3, 0.7, 0.1, 0.2],
202                                 [0.8, 0.2, 0.15, 0.25]])
203        },
204        {
205            'labels': torch.tensor([1]),  
206            'boxes': torch.tensor([[0.4, 0.6, 0.3, 0.4]])
207        }
208    ]
209    
210    # # Compute loss
211    loss_dict = criterion(yhat, y)
212    losses = loss_dict['labels']['loss_ce']*weight_dict['class_weighting'] + loss_dict['boxes']['loss_bbox']*weight_dict['bbox_weighting'] + loss_dict['boxes']['loss_giou']*weight_dict['giou_weighting']
213    print(losses)
214    
215    print("Loss components:")
216    for k, v in loss_dict.items():
217        print(f"{k}: {v}")
218    print(f"\nTotal weighted loss: {losses.item():.4f}")