def __call__(
        self, preds: BoxMaps, gt_box_batch: List[YoloBoxes], anchormap: BoxMap
    ) -> Tensor:
        device = preds.device
        _, _, h, w = preds.shape
        box_losses: List[Tensor] = []
        anchors = boxmap_to_boxes(anchormap)
        for diff_map, gt_boxes in zip(preds, gt_box_batch):
            if len(gt_boxes) == 0:
                continue

            pred_boxes = boxmap_to_boxes(BoxMap(diff_map))
            match_indices, positive_indices = self.matcher(anchors, gt_boxes, (w, h))
            num_pos = positive_indices.sum()
            if num_pos == 0:
                continue
            matched_gt_boxes = YoloBoxes(gt_boxes[match_indices][positive_indices])
            matched_pred_boxes = YoloBoxes(pred_boxes[positive_indices])
            if self.use_diff:
                matched_pred_boxes = YoloBoxes(
                    anchors[positive_indices] + matched_pred_boxes
                )

            box_losses.append(
                self.loss(
                    yolo_to_pascal(matched_pred_boxes, (1, 1)),
                    yolo_to_pascal(matched_gt_boxes, (1, 1)),
                )
            )
        if len(box_losses) == 0:
            return torch.tensor(0.0).to(device)
        return torch.stack(box_losses).mean()
Beispiel #2
0
    def __call__(self, pred_boxes: YoloBoxes, gt_boxes: YoloBoxes) -> float:
        if len(gt_boxes) == 0:
            return 1.0 if len(pred_boxes) == 0 else 0.0
        if len(pred_boxes) == 0:
            return 0.0

        iou_matrix = box_iou(
            yolo_to_pascal(pred_boxes, (1, 1)), yolo_to_pascal(gt_boxes, (1, 1)),
        )
        res = np.mean([precition(iou_matrix, t) for t in self.iou_thresholds])
        return res
Beispiel #3
0
 def __call__(
         self,
         inputs: NetOutput) -> Tuple[List[YoloBoxes], List[Confidences]]:
     anchormap, box_diffs, heatmap = inputs
     device = heatmap.device
     if self.use_peak:
         kpmap = (self.max_pool(heatmap)
                  == heatmap) & (heatmap > self.threshold)
     else:
         kpmap = heatmap > self.threshold
     batch_size, _, height, width = heatmap.shape
     original_wh = torch.tensor([width, height],
                                dtype=torch.float32).to(device)
     rows: List[Tuple[YoloBoxes, Confidences]] = []
     box_batch = []
     conf_batch = []
     for hm, km, box_diff in zip(heatmap.squeeze(1), kpmap.squeeze(1),
                                 box_diffs):
         kp = torch.nonzero(km, as_tuple=False)
         confidences = hm[kp[:, 0], kp[:, 1]]
         anchor = anchormap[:, kp[:, 0], kp[:, 1]].t()
         box_diff = box_diff[:, kp[:, 0], kp[:, 1]].t()
         boxes = anchor + box_diff
         boxes = yolo_clamp(YoloBoxes(boxes))
         sort_idx = nms(yolo_to_pascal(boxes, (1, 1)), confidences,
                        self.nms_thresold)
         box_batch.append(YoloBoxes(boxes[sort_idx]))
         conf_batch.append(Confidences(confidences[sort_idx]))
     return box_batch, conf_batch
Beispiel #4
0
 def __call__(
         self, net_output: NetOutput
 ) -> t.Tuple[List[YoloBoxes], List[Confidences]]:
     anchor_levels, box_diff_levels, labels_levels = net_output
     box_batch = []
     confidence_batch = []
     anchors = torch.cat(anchor_levels, dim=0)  # type: ignore
     box_diffs = torch.cat(box_diff_levels, dim=1)  # type:ignore
     labels_batch = torch.cat(labels_levels, dim=1)  # type:ignore
     for box_diff, confidences in zip(box_diffs, labels_batch):
         boxes = anchors + box_diff
         confidences, c_index = confidences.max(dim=1)
         filter_idx = confidences > self.confidence_threshold
         confidences = confidences[filter_idx][:self.limit]
         boxes = boxes[filter_idx][:self.limit]
         sort_idx = nms(
             yolo_to_pascal(YoloBoxes(boxes), (1, 1)),
             confidences,
             self.iou_threshold,
         )
         confidences.argsort(descending=True)
         boxes = YoloBoxes(boxes[sort_idx])
         boxes = yolo_clamp(boxes)
         confidences = confidences[sort_idx]
         box_batch.append(boxes)
         confidence_batch.append(Confidences(confidences))
     return box_batch, confidence_batch
Beispiel #5
0
    def __call__(
        self,
        images: ImageBatch,
        net_output: NetOutput,
        gt_boxes_list: List[YoloBoxes],
        gt_classes_list: List[Labels],
    ) -> Tuple[Tensor, Tensor, Tensor]:
        anchor_levels, box_diff_levels, classification_levels = net_output
        device = anchor_levels[0].device
        batch_size = classification_levels[0].shape[0]
        box_losses: List[Tensor] = []
        label_losses: List[Tensor] = []
        _, _, h, w = images.shape

        for anchors, box_diffs, pred_labels_level in zip(
                anchor_levels, box_diff_levels, classification_levels):
            for gt_boxes, gt_lables, box_diff, pred_labels in zip(
                    gt_boxes_list, gt_classes_list, box_diffs,
                    pred_labels_level):
                if len(gt_boxes) == 0:
                    continue
                iou_matrix = self.diou(
                    yolo_to_pascal(anchors, wh=(w, h)),
                    yolo_to_pascal(gt_boxes, wh=(w, h)),
                )
                match_score, match_indices = torch.min(iou_matrix, dim=1)
                label_losses.append(
                    self.label_loss(
                        match_score=match_score,
                        match_indices=match_indices,
                        pred_classes=pred_labels,
                        gt_classes=gt_lables,
                    ))
                box_losses.append(
                    self.box_loss(
                        match_score=match_score,
                        match_indices=match_indices,
                        anchors=anchors,
                        box_diff=BoxDiff(box_diff),
                        gt_boxes=gt_boxes,
                    ))
        all_box_loss = torch.stack(box_losses).mean() * self.box_weight
        all_label_loss = torch.stack(label_losses).mean() * self.label_weight
        loss = all_box_loss + all_label_loss
        return loss, all_box_loss, all_label_loss