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()
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
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
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
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