| 36 | |
| 37 | |
| 38 | class ClipMatcher(SetCriterion): |
| 39 | def __init__(self, num_classes, |
| 40 | matcher, |
| 41 | weight_dict, |
| 42 | losses): |
| 43 | """ Create the criterion. |
| 44 | Parameters: |
| 45 | num_classes: number of object categories, omitting the special no-object category |
| 46 | matcher: module able to compute a matching between targets and proposals |
| 47 | weight_dict: dict containing as key the names of the losses and as values their relative weight. |
| 48 | eos_coef: relative classification weight applied to the no-object category |
| 49 | losses: list of all the losses to be applied. See get_loss for list of available losses. |
| 50 | """ |
| 51 | super().__init__(num_classes, matcher, weight_dict, losses) |
| 52 | self.num_classes = num_classes |
| 53 | self.matcher = matcher |
| 54 | self.weight_dict = weight_dict |
| 55 | self.losses = losses |
| 56 | self.focal_loss = True |
| 57 | self.losses_dict = {} |
| 58 | self._current_frame_idx = 0 |
| 59 | |
| 60 | def initialize_for_single_clip(self, gt_instances: List[Instances]): |
| 61 | self.gt_instances = gt_instances |
| 62 | self.num_samples = 0 |
| 63 | self.sample_device = None |
| 64 | self._current_frame_idx = 0 |
| 65 | self.losses_dict = {} |
| 66 | |
| 67 | def _step(self): |
| 68 | self._current_frame_idx += 1 |
| 69 | |
| 70 | def calc_loss_for_track_scores(self, track_instances: Instances): |
| 71 | frame_id = self._current_frame_idx - 1 |
| 72 | gt_instances = self.gt_instances[frame_id] |
| 73 | outputs = { |
| 74 | 'pred_logits': track_instances.track_scores[None], |
| 75 | } |
| 76 | device = track_instances.track_scores.device |
| 77 | |
| 78 | num_tracks = len(track_instances) |
| 79 | src_idx = torch.arange(num_tracks, dtype=torch.long, device=device) |
| 80 | tgt_idx = track_instances.matched_gt_idxes # -1 for FP tracks and disappeared tracks |
| 81 | |
| 82 | track_losses = self.get_loss('labels', |
| 83 | outputs=outputs, |
| 84 | gt_instances=[gt_instances], |
| 85 | indices=[(src_idx, tgt_idx)], |
| 86 | num_boxes=1) |
| 87 | self.losses_dict.update( |
| 88 | {'frame_{}_track_{}'.format(frame_id, key): value for key, value in |
| 89 | track_losses.items()}) |
| 90 | |
| 91 | def get_num_boxes(self, num_samples): |
| 92 | num_boxes = torch.as_tensor(num_samples, dtype=torch.float, device=self.sample_device) |
| 93 | if is_dist_avail_and_initialized(): |
| 94 | torch.distributed.all_reduce(num_boxes) |
| 95 | num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item() |