|
|
@@ -0,0 +1,903 @@
|
|
|
+from typing import Dict, List, Optional, Tuple
|
|
|
+
|
|
|
+import torch
|
|
|
+import torch.nn.functional as F
|
|
|
+import torchvision
|
|
|
+from torch import nn, Tensor
|
|
|
+from torchvision.ops import boxes as box_ops, roi_align
|
|
|
+
|
|
|
+from libs.vision_libs.models.detection import _utils as det_utils
|
|
|
+
|
|
|
+###计算多头损失
|
|
|
+def line_loss():
|
|
|
+ pass
|
|
|
+
|
|
|
+def fastrcnn_loss(class_logits, box_regression, labels, regression_targets):
|
|
|
+ # type: (Tensor, Tensor, List[Tensor], List[Tensor]) -> Tuple[Tensor, Tensor]
|
|
|
+ """
|
|
|
+ Computes the loss for Faster R-CNN.
|
|
|
+
|
|
|
+ Args:
|
|
|
+ class_logits (Tensor)
|
|
|
+ box_regression (Tensor)
|
|
|
+ labels (list[BoxList])
|
|
|
+ regression_targets (Tensor)
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ classification_loss (Tensor)
|
|
|
+ box_loss (Tensor)
|
|
|
+ """
|
|
|
+
|
|
|
+ labels = torch.cat(labels, dim=0)
|
|
|
+ regression_targets = torch.cat(regression_targets, dim=0)
|
|
|
+
|
|
|
+ classification_loss = F.cross_entropy(class_logits, labels)
|
|
|
+
|
|
|
+ # get indices that correspond to the regression targets for
|
|
|
+ # the corresponding ground truth labels, to be used with
|
|
|
+ # advanced indexing
|
|
|
+ sampled_pos_inds_subset = torch.where(labels > 0)[0]
|
|
|
+ labels_pos = labels[sampled_pos_inds_subset]
|
|
|
+ N, num_classes = class_logits.shape
|
|
|
+ box_regression = box_regression.reshape(N, box_regression.size(-1) // 4, 4)
|
|
|
+
|
|
|
+ box_loss = F.smooth_l1_loss(
|
|
|
+ box_regression[sampled_pos_inds_subset, labels_pos],
|
|
|
+ regression_targets[sampled_pos_inds_subset],
|
|
|
+ beta=1 / 9,
|
|
|
+ reduction="sum",
|
|
|
+ )
|
|
|
+ box_loss = box_loss / labels.numel()
|
|
|
+
|
|
|
+ return classification_loss, box_loss
|
|
|
+
|
|
|
+
|
|
|
+def maskrcnn_inference(x, labels):
|
|
|
+ # type: (Tensor, List[Tensor]) -> List[Tensor]
|
|
|
+ """
|
|
|
+ From the results of the CNN, post process the masks
|
|
|
+ by taking the mask corresponding to the class with max
|
|
|
+ probability (which are of fixed size and directly output
|
|
|
+ by the CNN) and return the masks in the mask field of the BoxList.
|
|
|
+
|
|
|
+ Args:
|
|
|
+ x (Tensor): the mask logits
|
|
|
+ labels (list[BoxList]): bounding boxes that are used as
|
|
|
+ reference, one for ech image
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ results (list[BoxList]): one BoxList for each image, containing
|
|
|
+ the extra field mask
|
|
|
+ """
|
|
|
+ mask_prob = x.sigmoid()
|
|
|
+
|
|
|
+ # select masks corresponding to the predicted classes
|
|
|
+ num_masks = x.shape[0]
|
|
|
+ boxes_per_image = [label.shape[0] for label in labels]
|
|
|
+ labels = torch.cat(labels)
|
|
|
+ index = torch.arange(num_masks, device=labels.device)
|
|
|
+ mask_prob = mask_prob[index, labels][:, None]
|
|
|
+ mask_prob = mask_prob.split(boxes_per_image, dim=0)
|
|
|
+
|
|
|
+ return mask_prob
|
|
|
+
|
|
|
+
|
|
|
+def project_masks_on_boxes(gt_masks, boxes, matched_idxs, M):
|
|
|
+ # type: (Tensor, Tensor, Tensor, int) -> Tensor
|
|
|
+ """
|
|
|
+ Given segmentation masks and the bounding boxes corresponding
|
|
|
+ to the location of the masks in the image, this function
|
|
|
+ crops and resizes the masks in the position defined by the
|
|
|
+ boxes. This prepares the masks for them to be fed to the
|
|
|
+ loss computation as the targets.
|
|
|
+ """
|
|
|
+ matched_idxs = matched_idxs.to(boxes)
|
|
|
+ rois = torch.cat([matched_idxs[:, None], boxes], dim=1)
|
|
|
+ gt_masks = gt_masks[:, None].to(rois)
|
|
|
+ return roi_align(gt_masks, rois, (M, M), 1.0)[:, 0]
|
|
|
+
|
|
|
+
|
|
|
+def maskrcnn_loss(mask_logits, proposals, gt_masks, gt_labels, mask_matched_idxs):
|
|
|
+ # type: (Tensor, List[Tensor], List[Tensor], List[Tensor], List[Tensor]) -> Tensor
|
|
|
+ """
|
|
|
+ Args:
|
|
|
+ proposals (list[BoxList])
|
|
|
+ mask_logits (Tensor)
|
|
|
+ targets (list[BoxList])
|
|
|
+
|
|
|
+ Return:
|
|
|
+ mask_loss (Tensor): scalar tensor containing the loss
|
|
|
+ """
|
|
|
+
|
|
|
+ discretization_size = mask_logits.shape[-1]
|
|
|
+ labels = [gt_label[idxs] for gt_label, idxs in zip(gt_labels, mask_matched_idxs)]
|
|
|
+ mask_targets = [
|
|
|
+ project_masks_on_boxes(m, p, i, discretization_size) for m, p, i in zip(gt_masks, proposals, mask_matched_idxs)
|
|
|
+ ]
|
|
|
+
|
|
|
+ labels = torch.cat(labels, dim=0)
|
|
|
+ mask_targets = torch.cat(mask_targets, dim=0)
|
|
|
+
|
|
|
+ # torch.mean (in binary_cross_entropy_with_logits) doesn't
|
|
|
+ # accept empty tensors, so handle it separately
|
|
|
+ if mask_targets.numel() == 0:
|
|
|
+ return mask_logits.sum() * 0
|
|
|
+
|
|
|
+ mask_loss = F.binary_cross_entropy_with_logits(
|
|
|
+ mask_logits[torch.arange(labels.shape[0], device=labels.device), labels], mask_targets
|
|
|
+ )
|
|
|
+ return mask_loss
|
|
|
+
|
|
|
+
|
|
|
+def keypoints_to_heatmap(keypoints, rois, heatmap_size):
|
|
|
+ # type: (Tensor, Tensor, int) -> Tuple[Tensor, Tensor]
|
|
|
+ offset_x = rois[:, 0]
|
|
|
+ offset_y = rois[:, 1]
|
|
|
+ scale_x = heatmap_size / (rois[:, 2] - rois[:, 0])
|
|
|
+ scale_y = heatmap_size / (rois[:, 3] - rois[:, 1])
|
|
|
+
|
|
|
+ offset_x = offset_x[:, None]
|
|
|
+ offset_y = offset_y[:, None]
|
|
|
+ scale_x = scale_x[:, None]
|
|
|
+ scale_y = scale_y[:, None]
|
|
|
+
|
|
|
+ x = keypoints[..., 0]
|
|
|
+ y = keypoints[..., 1]
|
|
|
+
|
|
|
+ x_boundary_inds = x == rois[:, 2][:, None]
|
|
|
+ y_boundary_inds = y == rois[:, 3][:, None]
|
|
|
+
|
|
|
+ x = (x - offset_x) * scale_x
|
|
|
+ x = x.floor().long()
|
|
|
+ y = (y - offset_y) * scale_y
|
|
|
+ y = y.floor().long()
|
|
|
+
|
|
|
+ x[x_boundary_inds] = heatmap_size - 1
|
|
|
+ y[y_boundary_inds] = heatmap_size - 1
|
|
|
+
|
|
|
+ valid_loc = (x >= 0) & (y >= 0) & (x < heatmap_size) & (y < heatmap_size)
|
|
|
+ vis = keypoints[..., 2] > 0
|
|
|
+ valid = (valid_loc & vis).long()
|
|
|
+
|
|
|
+ lin_ind = y * heatmap_size + x
|
|
|
+ heatmaps = lin_ind * valid
|
|
|
+
|
|
|
+ return heatmaps, valid
|
|
|
+
|
|
|
+
|
|
|
+def _onnx_heatmaps_to_keypoints(
|
|
|
+ maps, maps_i, roi_map_width, roi_map_height, widths_i, heights_i, offset_x_i, offset_y_i
|
|
|
+):
|
|
|
+ num_keypoints = torch.scalar_tensor(maps.size(1), dtype=torch.int64)
|
|
|
+
|
|
|
+ width_correction = widths_i / roi_map_width
|
|
|
+ height_correction = heights_i / roi_map_height
|
|
|
+
|
|
|
+ roi_map = F.interpolate(
|
|
|
+ maps_i[:, None], size=(int(roi_map_height), int(roi_map_width)), mode="bicubic", align_corners=False
|
|
|
+ )[:, 0]
|
|
|
+
|
|
|
+ w = torch.scalar_tensor(roi_map.size(2), dtype=torch.int64)
|
|
|
+ pos = roi_map.reshape(num_keypoints, -1).argmax(dim=1)
|
|
|
+
|
|
|
+ x_int = pos % w
|
|
|
+ y_int = (pos - x_int) // w
|
|
|
+
|
|
|
+ x = (torch.tensor(0.5, dtype=torch.float32) + x_int.to(dtype=torch.float32)) * width_correction.to(
|
|
|
+ dtype=torch.float32
|
|
|
+ )
|
|
|
+ y = (torch.tensor(0.5, dtype=torch.float32) + y_int.to(dtype=torch.float32)) * height_correction.to(
|
|
|
+ dtype=torch.float32
|
|
|
+ )
|
|
|
+
|
|
|
+ xy_preds_i_0 = x + offset_x_i.to(dtype=torch.float32)
|
|
|
+ xy_preds_i_1 = y + offset_y_i.to(dtype=torch.float32)
|
|
|
+ xy_preds_i_2 = torch.ones(xy_preds_i_1.shape, dtype=torch.float32)
|
|
|
+ xy_preds_i = torch.stack(
|
|
|
+ [
|
|
|
+ xy_preds_i_0.to(dtype=torch.float32),
|
|
|
+ xy_preds_i_1.to(dtype=torch.float32),
|
|
|
+ xy_preds_i_2.to(dtype=torch.float32),
|
|
|
+ ],
|
|
|
+ 0,
|
|
|
+ )
|
|
|
+
|
|
|
+ # TODO: simplify when indexing without rank will be supported by ONNX
|
|
|
+ base = num_keypoints * num_keypoints + num_keypoints + 1
|
|
|
+ ind = torch.arange(num_keypoints)
|
|
|
+ ind = ind.to(dtype=torch.int64) * base
|
|
|
+ end_scores_i = (
|
|
|
+ roi_map.index_select(1, y_int.to(dtype=torch.int64))
|
|
|
+ .index_select(2, x_int.to(dtype=torch.int64))
|
|
|
+ .view(-1)
|
|
|
+ .index_select(0, ind.to(dtype=torch.int64))
|
|
|
+ )
|
|
|
+
|
|
|
+ return xy_preds_i, end_scores_i
|
|
|
+
|
|
|
+
|
|
|
+@torch.jit._script_if_tracing
|
|
|
+def _onnx_heatmaps_to_keypoints_loop(
|
|
|
+ maps, rois, widths_ceil, heights_ceil, widths, heights, offset_x, offset_y, num_keypoints
|
|
|
+):
|
|
|
+ xy_preds = torch.zeros((0, 3, int(num_keypoints)), dtype=torch.float32, device=maps.device)
|
|
|
+ end_scores = torch.zeros((0, int(num_keypoints)), dtype=torch.float32, device=maps.device)
|
|
|
+
|
|
|
+ for i in range(int(rois.size(0))):
|
|
|
+ xy_preds_i, end_scores_i = _onnx_heatmaps_to_keypoints(
|
|
|
+ maps, maps[i], widths_ceil[i], heights_ceil[i], widths[i], heights[i], offset_x[i], offset_y[i]
|
|
|
+ )
|
|
|
+ xy_preds = torch.cat((xy_preds.to(dtype=torch.float32), xy_preds_i.unsqueeze(0).to(dtype=torch.float32)), 0)
|
|
|
+ end_scores = torch.cat(
|
|
|
+ (end_scores.to(dtype=torch.float32), end_scores_i.to(dtype=torch.float32).unsqueeze(0)), 0
|
|
|
+ )
|
|
|
+ return xy_preds, end_scores
|
|
|
+
|
|
|
+
|
|
|
+def heatmaps_to_keypoints(maps, rois):
|
|
|
+ """Extract predicted keypoint locations from heatmaps. Output has shape
|
|
|
+ (#rois, 4, #keypoints) with the 4 rows corresponding to (x, y, logit, prob)
|
|
|
+ for each keypoint.
|
|
|
+ """
|
|
|
+ # This function converts a discrete image coordinate in a HEATMAP_SIZE x
|
|
|
+ # HEATMAP_SIZE image to a continuous keypoint coordinate. We maintain
|
|
|
+ # consistency with keypoints_to_heatmap_labels by using the conversion from
|
|
|
+ # Heckbert 1990: c = d + 0.5, where d is a discrete coordinate and c is a
|
|
|
+ # continuous coordinate.
|
|
|
+ offset_x = rois[:, 0]
|
|
|
+ offset_y = rois[:, 1]
|
|
|
+
|
|
|
+ widths = rois[:, 2] - rois[:, 0]
|
|
|
+ heights = rois[:, 3] - rois[:, 1]
|
|
|
+ widths = widths.clamp(min=1)
|
|
|
+ heights = heights.clamp(min=1)
|
|
|
+ widths_ceil = widths.ceil()
|
|
|
+ heights_ceil = heights.ceil()
|
|
|
+
|
|
|
+ num_keypoints = maps.shape[1]
|
|
|
+
|
|
|
+ if torchvision._is_tracing():
|
|
|
+ xy_preds, end_scores = _onnx_heatmaps_to_keypoints_loop(
|
|
|
+ maps,
|
|
|
+ rois,
|
|
|
+ widths_ceil,
|
|
|
+ heights_ceil,
|
|
|
+ widths,
|
|
|
+ heights,
|
|
|
+ offset_x,
|
|
|
+ offset_y,
|
|
|
+ torch.scalar_tensor(num_keypoints, dtype=torch.int64),
|
|
|
+ )
|
|
|
+ return xy_preds.permute(0, 2, 1), end_scores
|
|
|
+
|
|
|
+ xy_preds = torch.zeros((len(rois), 3, num_keypoints), dtype=torch.float32, device=maps.device)
|
|
|
+ end_scores = torch.zeros((len(rois), num_keypoints), dtype=torch.float32, device=maps.device)
|
|
|
+ for i in range(len(rois)):
|
|
|
+ roi_map_width = int(widths_ceil[i].item())
|
|
|
+ roi_map_height = int(heights_ceil[i].item())
|
|
|
+ width_correction = widths[i] / roi_map_width
|
|
|
+ height_correction = heights[i] / roi_map_height
|
|
|
+ roi_map = F.interpolate(
|
|
|
+ maps[i][:, None], size=(roi_map_height, roi_map_width), mode="bicubic", align_corners=False
|
|
|
+ )[:, 0]
|
|
|
+ # roi_map_probs = scores_to_probs(roi_map.copy())
|
|
|
+ w = roi_map.shape[2]
|
|
|
+ pos = roi_map.reshape(num_keypoints, -1).argmax(dim=1)
|
|
|
+
|
|
|
+ x_int = pos % w
|
|
|
+ y_int = torch.div(pos - x_int, w, rounding_mode="floor")
|
|
|
+ # assert (roi_map_probs[k, y_int, x_int] ==
|
|
|
+ # roi_map_probs[k, :, :].max())
|
|
|
+ x = (x_int.float() + 0.5) * width_correction
|
|
|
+ y = (y_int.float() + 0.5) * height_correction
|
|
|
+ xy_preds[i, 0, :] = x + offset_x[i]
|
|
|
+ xy_preds[i, 1, :] = y + offset_y[i]
|
|
|
+ xy_preds[i, 2, :] = 1
|
|
|
+ end_scores[i, :] = roi_map[torch.arange(num_keypoints, device=roi_map.device), y_int, x_int]
|
|
|
+
|
|
|
+ return xy_preds.permute(0, 2, 1), end_scores
|
|
|
+
|
|
|
+
|
|
|
+def keypointrcnn_loss(keypoint_logits, proposals, gt_keypoints, keypoint_matched_idxs):
|
|
|
+ # type: (Tensor, List[Tensor], List[Tensor], List[Tensor]) -> Tensor
|
|
|
+ N, K, H, W = keypoint_logits.shape
|
|
|
+ if H != W:
|
|
|
+ raise ValueError(
|
|
|
+ f"keypoint_logits height and width (last two elements of shape) should be equal. Instead got H = {H} and W = {W}"
|
|
|
+ )
|
|
|
+ discretization_size = H
|
|
|
+ heatmaps = []
|
|
|
+ valid = []
|
|
|
+ for proposals_per_image, gt_kp_in_image, midx in zip(proposals, gt_keypoints, keypoint_matched_idxs):
|
|
|
+ kp = gt_kp_in_image[midx]
|
|
|
+ heatmaps_per_image, valid_per_image = keypoints_to_heatmap(kp, proposals_per_image, discretization_size)
|
|
|
+ heatmaps.append(heatmaps_per_image.view(-1))
|
|
|
+ valid.append(valid_per_image.view(-1))
|
|
|
+
|
|
|
+ keypoint_targets = torch.cat(heatmaps, dim=0)
|
|
|
+ valid = torch.cat(valid, dim=0).to(dtype=torch.uint8)
|
|
|
+ valid = torch.where(valid)[0]
|
|
|
+
|
|
|
+ # torch.mean (in binary_cross_entropy_with_logits) doesn't
|
|
|
+ # accept empty tensors, so handle it sepaartely
|
|
|
+ if keypoint_targets.numel() == 0 or len(valid) == 0:
|
|
|
+ return keypoint_logits.sum() * 0
|
|
|
+
|
|
|
+ keypoint_logits = keypoint_logits.view(N * K, H * W)
|
|
|
+
|
|
|
+ keypoint_loss = F.cross_entropy(keypoint_logits[valid], keypoint_targets[valid])
|
|
|
+ return keypoint_loss
|
|
|
+
|
|
|
+
|
|
|
+def keypointrcnn_inference(x, boxes):
|
|
|
+ # type: (Tensor, List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
|
|
|
+ kp_probs = []
|
|
|
+ kp_scores = []
|
|
|
+
|
|
|
+ boxes_per_image = [box.size(0) for box in boxes]
|
|
|
+ x2 = x.split(boxes_per_image, dim=0)
|
|
|
+
|
|
|
+ for xx, bb in zip(x2, boxes):
|
|
|
+ kp_prob, scores = heatmaps_to_keypoints(xx, bb)
|
|
|
+ kp_probs.append(kp_prob)
|
|
|
+ kp_scores.append(scores)
|
|
|
+
|
|
|
+ return kp_probs, kp_scores
|
|
|
+
|
|
|
+
|
|
|
+def _onnx_expand_boxes(boxes, scale):
|
|
|
+ # type: (Tensor, float) -> Tensor
|
|
|
+ w_half = (boxes[:, 2] - boxes[:, 0]) * 0.5
|
|
|
+ h_half = (boxes[:, 3] - boxes[:, 1]) * 0.5
|
|
|
+ x_c = (boxes[:, 2] + boxes[:, 0]) * 0.5
|
|
|
+ y_c = (boxes[:, 3] + boxes[:, 1]) * 0.5
|
|
|
+
|
|
|
+ w_half = w_half.to(dtype=torch.float32) * scale
|
|
|
+ h_half = h_half.to(dtype=torch.float32) * scale
|
|
|
+
|
|
|
+ boxes_exp0 = x_c - w_half
|
|
|
+ boxes_exp1 = y_c - h_half
|
|
|
+ boxes_exp2 = x_c + w_half
|
|
|
+ boxes_exp3 = y_c + h_half
|
|
|
+ boxes_exp = torch.stack((boxes_exp0, boxes_exp1, boxes_exp2, boxes_exp3), 1)
|
|
|
+ return boxes_exp
|
|
|
+
|
|
|
+
|
|
|
+# the next two functions should be merged inside Masker
|
|
|
+# but are kept here for the moment while we need them
|
|
|
+# temporarily for paste_mask_in_image
|
|
|
+def expand_boxes(boxes, scale):
|
|
|
+ # type: (Tensor, float) -> Tensor
|
|
|
+ if torchvision._is_tracing():
|
|
|
+ return _onnx_expand_boxes(boxes, scale)
|
|
|
+ w_half = (boxes[:, 2] - boxes[:, 0]) * 0.5
|
|
|
+ h_half = (boxes[:, 3] - boxes[:, 1]) * 0.5
|
|
|
+ x_c = (boxes[:, 2] + boxes[:, 0]) * 0.5
|
|
|
+ y_c = (boxes[:, 3] + boxes[:, 1]) * 0.5
|
|
|
+
|
|
|
+ w_half *= scale
|
|
|
+ h_half *= scale
|
|
|
+
|
|
|
+ boxes_exp = torch.zeros_like(boxes)
|
|
|
+ boxes_exp[:, 0] = x_c - w_half
|
|
|
+ boxes_exp[:, 2] = x_c + w_half
|
|
|
+ boxes_exp[:, 1] = y_c - h_half
|
|
|
+ boxes_exp[:, 3] = y_c + h_half
|
|
|
+ return boxes_exp
|
|
|
+
|
|
|
+
|
|
|
+@torch.jit.unused
|
|
|
+def expand_masks_tracing_scale(M, padding):
|
|
|
+ # type: (int, int) -> float
|
|
|
+ return torch.tensor(M + 2 * padding).to(torch.float32) / torch.tensor(M).to(torch.float32)
|
|
|
+
|
|
|
+
|
|
|
+def expand_masks(mask, padding):
|
|
|
+ # type: (Tensor, int) -> Tuple[Tensor, float]
|
|
|
+ M = mask.shape[-1]
|
|
|
+ if torch._C._get_tracing_state(): # could not import is_tracing(), not sure why
|
|
|
+ scale = expand_masks_tracing_scale(M, padding)
|
|
|
+ else:
|
|
|
+ scale = float(M + 2 * padding) / M
|
|
|
+ padded_mask = F.pad(mask, (padding,) * 4)
|
|
|
+ return padded_mask, scale
|
|
|
+
|
|
|
+
|
|
|
+def paste_mask_in_image(mask, box, im_h, im_w):
|
|
|
+ # type: (Tensor, Tensor, int, int) -> Tensor
|
|
|
+ TO_REMOVE = 1
|
|
|
+ w = int(box[2] - box[0] + TO_REMOVE)
|
|
|
+ h = int(box[3] - box[1] + TO_REMOVE)
|
|
|
+ w = max(w, 1)
|
|
|
+ h = max(h, 1)
|
|
|
+
|
|
|
+ # Set shape to [batchxCxHxW]
|
|
|
+ mask = mask.expand((1, 1, -1, -1))
|
|
|
+
|
|
|
+ # Resize mask
|
|
|
+ mask = F.interpolate(mask, size=(h, w), mode="bilinear", align_corners=False)
|
|
|
+ mask = mask[0][0]
|
|
|
+
|
|
|
+ im_mask = torch.zeros((im_h, im_w), dtype=mask.dtype, device=mask.device)
|
|
|
+ x_0 = max(box[0], 0)
|
|
|
+ x_1 = min(box[2] + 1, im_w)
|
|
|
+ y_0 = max(box[1], 0)
|
|
|
+ y_1 = min(box[3] + 1, im_h)
|
|
|
+
|
|
|
+ im_mask[y_0:y_1, x_0:x_1] = mask[(y_0 - box[1]) : (y_1 - box[1]), (x_0 - box[0]) : (x_1 - box[0])]
|
|
|
+ return im_mask
|
|
|
+
|
|
|
+
|
|
|
+def _onnx_paste_mask_in_image(mask, box, im_h, im_w):
|
|
|
+ one = torch.ones(1, dtype=torch.int64)
|
|
|
+ zero = torch.zeros(1, dtype=torch.int64)
|
|
|
+
|
|
|
+ w = box[2] - box[0] + one
|
|
|
+ h = box[3] - box[1] + one
|
|
|
+ w = torch.max(torch.cat((w, one)))
|
|
|
+ h = torch.max(torch.cat((h, one)))
|
|
|
+
|
|
|
+ # Set shape to [batchxCxHxW]
|
|
|
+ mask = mask.expand((1, 1, mask.size(0), mask.size(1)))
|
|
|
+
|
|
|
+ # Resize mask
|
|
|
+ mask = F.interpolate(mask, size=(int(h), int(w)), mode="bilinear", align_corners=False)
|
|
|
+ mask = mask[0][0]
|
|
|
+
|
|
|
+ x_0 = torch.max(torch.cat((box[0].unsqueeze(0), zero)))
|
|
|
+ x_1 = torch.min(torch.cat((box[2].unsqueeze(0) + one, im_w.unsqueeze(0))))
|
|
|
+ y_0 = torch.max(torch.cat((box[1].unsqueeze(0), zero)))
|
|
|
+ y_1 = torch.min(torch.cat((box[3].unsqueeze(0) + one, im_h.unsqueeze(0))))
|
|
|
+
|
|
|
+ unpaded_im_mask = mask[(y_0 - box[1]) : (y_1 - box[1]), (x_0 - box[0]) : (x_1 - box[0])]
|
|
|
+
|
|
|
+ # TODO : replace below with a dynamic padding when support is added in ONNX
|
|
|
+
|
|
|
+ # pad y
|
|
|
+ zeros_y0 = torch.zeros(y_0, unpaded_im_mask.size(1))
|
|
|
+ zeros_y1 = torch.zeros(im_h - y_1, unpaded_im_mask.size(1))
|
|
|
+ concat_0 = torch.cat((zeros_y0, unpaded_im_mask.to(dtype=torch.float32), zeros_y1), 0)[0:im_h, :]
|
|
|
+ # pad x
|
|
|
+ zeros_x0 = torch.zeros(concat_0.size(0), x_0)
|
|
|
+ zeros_x1 = torch.zeros(concat_0.size(0), im_w - x_1)
|
|
|
+ im_mask = torch.cat((zeros_x0, concat_0, zeros_x1), 1)[:, :im_w]
|
|
|
+ return im_mask
|
|
|
+
|
|
|
+
|
|
|
+@torch.jit._script_if_tracing
|
|
|
+def _onnx_paste_masks_in_image_loop(masks, boxes, im_h, im_w):
|
|
|
+ res_append = torch.zeros(0, im_h, im_w)
|
|
|
+ for i in range(masks.size(0)):
|
|
|
+ mask_res = _onnx_paste_mask_in_image(masks[i][0], boxes[i], im_h, im_w)
|
|
|
+ mask_res = mask_res.unsqueeze(0)
|
|
|
+ res_append = torch.cat((res_append, mask_res))
|
|
|
+ return res_append
|
|
|
+
|
|
|
+
|
|
|
+def paste_masks_in_image(masks, boxes, img_shape, padding=1):
|
|
|
+ # type: (Tensor, Tensor, Tuple[int, int], int) -> Tensor
|
|
|
+ masks, scale = expand_masks(masks, padding=padding)
|
|
|
+ boxes = expand_boxes(boxes, scale).to(dtype=torch.int64)
|
|
|
+ im_h, im_w = img_shape
|
|
|
+
|
|
|
+ if torchvision._is_tracing():
|
|
|
+ return _onnx_paste_masks_in_image_loop(
|
|
|
+ masks, boxes, torch.scalar_tensor(im_h, dtype=torch.int64), torch.scalar_tensor(im_w, dtype=torch.int64)
|
|
|
+ )[:, None]
|
|
|
+ res = [paste_mask_in_image(m[0], b, im_h, im_w) for m, b in zip(masks, boxes)]
|
|
|
+ if len(res) > 0:
|
|
|
+ ret = torch.stack(res, dim=0)[:, None]
|
|
|
+ else:
|
|
|
+ ret = masks.new_empty((0, 1, im_h, im_w))
|
|
|
+ return ret
|
|
|
+
|
|
|
+
|
|
|
+class RoIHeads(nn.Module):
|
|
|
+ __annotations__ = {
|
|
|
+ "box_coder": det_utils.BoxCoder,
|
|
|
+ "proposal_matcher": det_utils.Matcher,
|
|
|
+ "fg_bg_sampler": det_utils.BalancedPositiveNegativeSampler,
|
|
|
+ }
|
|
|
+
|
|
|
+ def __init__(
|
|
|
+ self,
|
|
|
+ box_roi_pool,
|
|
|
+ box_head,
|
|
|
+ box_predictor,
|
|
|
+ line_head,
|
|
|
+ line_predictor,
|
|
|
+ # Faster R-CNN training
|
|
|
+ fg_iou_thresh,
|
|
|
+ bg_iou_thresh,
|
|
|
+ batch_size_per_image,
|
|
|
+ positive_fraction,
|
|
|
+ bbox_reg_weights,
|
|
|
+ # Faster R-CNN inference
|
|
|
+ score_thresh,
|
|
|
+ nms_thresh,
|
|
|
+ detections_per_img,
|
|
|
+ # Mask
|
|
|
+ mask_roi_pool=None,
|
|
|
+ mask_head=None,
|
|
|
+ mask_predictor=None,
|
|
|
+ keypoint_roi_pool=None,
|
|
|
+ keypoint_head=None,
|
|
|
+ keypoint_predictor=None,
|
|
|
+ ):
|
|
|
+ super().__init__()
|
|
|
+
|
|
|
+ self.box_similarity = box_ops.box_iou
|
|
|
+ # assign ground-truth boxes for each proposal
|
|
|
+ self.proposal_matcher = det_utils.Matcher(fg_iou_thresh, bg_iou_thresh, allow_low_quality_matches=False)
|
|
|
+
|
|
|
+ self.fg_bg_sampler = det_utils.BalancedPositiveNegativeSampler(batch_size_per_image, positive_fraction)
|
|
|
+
|
|
|
+ if bbox_reg_weights is None:
|
|
|
+ bbox_reg_weights = (10.0, 10.0, 5.0, 5.0)
|
|
|
+ self.box_coder = det_utils.BoxCoder(bbox_reg_weights)
|
|
|
+
|
|
|
+ self.box_roi_pool = box_roi_pool
|
|
|
+ self.box_head = box_head
|
|
|
+ self.box_predictor = box_predictor
|
|
|
+
|
|
|
+ self.line_head=line_head
|
|
|
+ self.line_predictor=line_predictor
|
|
|
+
|
|
|
+ self.score_thresh = score_thresh
|
|
|
+ self.nms_thresh = nms_thresh
|
|
|
+ self.detections_per_img = detections_per_img
|
|
|
+
|
|
|
+ self.mask_roi_pool = mask_roi_pool
|
|
|
+ self.mask_head = mask_head
|
|
|
+ self.mask_predictor = mask_predictor
|
|
|
+
|
|
|
+ self.keypoint_roi_pool = keypoint_roi_pool
|
|
|
+ self.keypoint_head = keypoint_head
|
|
|
+ self.keypoint_predictor = keypoint_predictor
|
|
|
+
|
|
|
+ def has_line(self):
|
|
|
+ pass
|
|
|
+ def has_mask(self):
|
|
|
+ if self.mask_roi_pool is None:
|
|
|
+ return False
|
|
|
+ if self.mask_head is None:
|
|
|
+ return False
|
|
|
+ if self.mask_predictor is None:
|
|
|
+ return False
|
|
|
+ return True
|
|
|
+
|
|
|
+ def has_keypoint(self):
|
|
|
+ if self.keypoint_roi_pool is None:
|
|
|
+ return False
|
|
|
+ if self.keypoint_head is None:
|
|
|
+ return False
|
|
|
+ if self.keypoint_predictor is None:
|
|
|
+ return False
|
|
|
+ return True
|
|
|
+
|
|
|
+ def assign_targets_to_proposals(self, proposals, gt_boxes, gt_labels):
|
|
|
+ # type: (List[Tensor], List[Tensor], List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
|
|
|
+ matched_idxs = []
|
|
|
+ labels = []
|
|
|
+ for proposals_in_image, gt_boxes_in_image, gt_labels_in_image in zip(proposals, gt_boxes, gt_labels):
|
|
|
+
|
|
|
+ if gt_boxes_in_image.numel() == 0:
|
|
|
+ # Background image
|
|
|
+ device = proposals_in_image.device
|
|
|
+ clamped_matched_idxs_in_image = torch.zeros(
|
|
|
+ (proposals_in_image.shape[0],), dtype=torch.int64, device=device
|
|
|
+ )
|
|
|
+ labels_in_image = torch.zeros((proposals_in_image.shape[0],), dtype=torch.int64, device=device)
|
|
|
+ else:
|
|
|
+ # set to self.box_similarity when https://github.com/pytorch/pytorch/issues/27495 lands
|
|
|
+ match_quality_matrix = box_ops.box_iou(gt_boxes_in_image, proposals_in_image)
|
|
|
+ matched_idxs_in_image = self.proposal_matcher(match_quality_matrix)
|
|
|
+
|
|
|
+ clamped_matched_idxs_in_image = matched_idxs_in_image.clamp(min=0)
|
|
|
+
|
|
|
+ labels_in_image = gt_labels_in_image[clamped_matched_idxs_in_image]
|
|
|
+ labels_in_image = labels_in_image.to(dtype=torch.int64)
|
|
|
+
|
|
|
+ # Label background (below the low threshold)
|
|
|
+ bg_inds = matched_idxs_in_image == self.proposal_matcher.BELOW_LOW_THRESHOLD
|
|
|
+ labels_in_image[bg_inds] = 0
|
|
|
+
|
|
|
+ # Label ignore proposals (between low and high thresholds)
|
|
|
+ ignore_inds = matched_idxs_in_image == self.proposal_matcher.BETWEEN_THRESHOLDS
|
|
|
+ labels_in_image[ignore_inds] = -1 # -1 is ignored by sampler
|
|
|
+
|
|
|
+ matched_idxs.append(clamped_matched_idxs_in_image)
|
|
|
+ labels.append(labels_in_image)
|
|
|
+ return matched_idxs, labels
|
|
|
+
|
|
|
+ def subsample(self, labels):
|
|
|
+ # type: (List[Tensor]) -> List[Tensor]
|
|
|
+ sampled_pos_inds, sampled_neg_inds = self.fg_bg_sampler(labels)
|
|
|
+ sampled_inds = []
|
|
|
+ for img_idx, (pos_inds_img, neg_inds_img) in enumerate(zip(sampled_pos_inds, sampled_neg_inds)):
|
|
|
+ img_sampled_inds = torch.where(pos_inds_img | neg_inds_img)[0]
|
|
|
+ sampled_inds.append(img_sampled_inds)
|
|
|
+ return sampled_inds
|
|
|
+
|
|
|
+ def add_gt_proposals(self, proposals, gt_boxes):
|
|
|
+ # type: (List[Tensor], List[Tensor]) -> List[Tensor]
|
|
|
+ proposals = [torch.cat((proposal, gt_box)) for proposal, gt_box in zip(proposals, gt_boxes)]
|
|
|
+
|
|
|
+ return proposals
|
|
|
+
|
|
|
+ def check_targets(self, targets):
|
|
|
+ # type: (Optional[List[Dict[str, Tensor]]]) -> None
|
|
|
+ if targets is None:
|
|
|
+ raise ValueError("targets should not be None")
|
|
|
+ if not all(["boxes" in t for t in targets]):
|
|
|
+ raise ValueError("Every element of targets should have a boxes key")
|
|
|
+ if not all(["labels" in t for t in targets]):
|
|
|
+ raise ValueError("Every element of targets should have a labels key")
|
|
|
+ if self.has_mask():
|
|
|
+ if not all(["masks" in t for t in targets]):
|
|
|
+ raise ValueError("Every element of targets should have a masks key")
|
|
|
+
|
|
|
+ def select_training_samples(
|
|
|
+ self,
|
|
|
+ proposals, # type: List[Tensor]
|
|
|
+ targets, # type: Optional[List[Dict[str, Tensor]]]
|
|
|
+ ):
|
|
|
+ # type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor], List[Tensor]]
|
|
|
+ self.check_targets(targets)
|
|
|
+ if targets is None:
|
|
|
+ raise ValueError("targets should not be None")
|
|
|
+ dtype = proposals[0].dtype
|
|
|
+ device = proposals[0].device
|
|
|
+
|
|
|
+ gt_boxes = [t["boxes"].to(dtype) for t in targets]
|
|
|
+ gt_labels = [t["labels"] for t in targets]
|
|
|
+
|
|
|
+ # append ground-truth bboxes to propos
|
|
|
+ proposals = self.add_gt_proposals(proposals, gt_boxes)
|
|
|
+
|
|
|
+ # get matching gt indices for each proposal
|
|
|
+ matched_idxs, labels = self.assign_targets_to_proposals(proposals, gt_boxes, gt_labels)
|
|
|
+ # sample a fixed proportion of positive-negative proposals
|
|
|
+ sampled_inds = self.subsample(labels)
|
|
|
+ matched_gt_boxes = []
|
|
|
+ num_images = len(proposals)
|
|
|
+ for img_id in range(num_images):
|
|
|
+ img_sampled_inds = sampled_inds[img_id]
|
|
|
+ proposals[img_id] = proposals[img_id][img_sampled_inds]
|
|
|
+ labels[img_id] = labels[img_id][img_sampled_inds]
|
|
|
+ matched_idxs[img_id] = matched_idxs[img_id][img_sampled_inds]
|
|
|
+
|
|
|
+ gt_boxes_in_image = gt_boxes[img_id]
|
|
|
+ if gt_boxes_in_image.numel() == 0:
|
|
|
+ gt_boxes_in_image = torch.zeros((1, 4), dtype=dtype, device=device)
|
|
|
+ matched_gt_boxes.append(gt_boxes_in_image[matched_idxs[img_id]])
|
|
|
+
|
|
|
+ regression_targets = self.box_coder.encode(matched_gt_boxes, proposals)
|
|
|
+ return proposals, matched_idxs, labels, regression_targets
|
|
|
+
|
|
|
+ def postprocess_detections(
|
|
|
+ self,
|
|
|
+ class_logits, # type: Tensor
|
|
|
+ box_regression, # type: Tensor
|
|
|
+ proposals, # type: List[Tensor]
|
|
|
+ image_shapes, # type: List[Tuple[int, int]]
|
|
|
+ ):
|
|
|
+ # type: (...) -> Tuple[List[Tensor], List[Tensor], List[Tensor]]
|
|
|
+ device = class_logits.device
|
|
|
+ num_classes = class_logits.shape[-1]
|
|
|
+
|
|
|
+ boxes_per_image = [boxes_in_image.shape[0] for boxes_in_image in proposals]
|
|
|
+ pred_boxes = self.box_coder.decode(box_regression, proposals)
|
|
|
+
|
|
|
+ pred_scores = F.softmax(class_logits, -1)
|
|
|
+
|
|
|
+ pred_boxes_list = pred_boxes.split(boxes_per_image, 0)
|
|
|
+ pred_scores_list = pred_scores.split(boxes_per_image, 0)
|
|
|
+
|
|
|
+ all_boxes = []
|
|
|
+ all_scores = []
|
|
|
+ all_labels = []
|
|
|
+ for boxes, scores, image_shape in zip(pred_boxes_list, pred_scores_list, image_shapes):
|
|
|
+ boxes = box_ops.clip_boxes_to_image(boxes, image_shape)
|
|
|
+
|
|
|
+ # create labels for each prediction
|
|
|
+ labels = torch.arange(num_classes, device=device)
|
|
|
+ labels = labels.view(1, -1).expand_as(scores)
|
|
|
+
|
|
|
+ # remove predictions with the background label
|
|
|
+ boxes = boxes[:, 1:]
|
|
|
+ scores = scores[:, 1:]
|
|
|
+ labels = labels[:, 1:]
|
|
|
+
|
|
|
+ # batch everything, by making every class prediction be a separate instance
|
|
|
+ boxes = boxes.reshape(-1, 4)
|
|
|
+ scores = scores.reshape(-1)
|
|
|
+ labels = labels.reshape(-1)
|
|
|
+
|
|
|
+ # remove low scoring boxes
|
|
|
+ inds = torch.where(scores > self.score_thresh)[0]
|
|
|
+ boxes, scores, labels = boxes[inds], scores[inds], labels[inds]
|
|
|
+
|
|
|
+ # remove empty boxes
|
|
|
+ keep = box_ops.remove_small_boxes(boxes, min_size=1e-2)
|
|
|
+ boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
|
|
|
+
|
|
|
+ # non-maximum suppression, independently done per class
|
|
|
+ keep = box_ops.batched_nms(boxes, scores, labels, self.nms_thresh)
|
|
|
+ # keep only topk scoring predictions
|
|
|
+ keep = keep[: self.detections_per_img]
|
|
|
+ boxes, scores, labels = boxes[keep], scores[keep], labels[keep]
|
|
|
+
|
|
|
+ all_boxes.append(boxes)
|
|
|
+ all_scores.append(scores)
|
|
|
+ all_labels.append(labels)
|
|
|
+
|
|
|
+ return all_boxes, all_scores, all_labels
|
|
|
+
|
|
|
+
|
|
|
+ def forward(
|
|
|
+ self,
|
|
|
+ features, # type: Dict[str, Tensor]
|
|
|
+ proposals, # type: List[Tensor]
|
|
|
+ image_shapes, # type: List[Tuple[int, int]]
|
|
|
+ targets=None, # type: Optional[List[Dict[str, Tensor]]]
|
|
|
+ ):
|
|
|
+ # type: (...) -> Tuple[List[Dict[str, Tensor]], Dict[str, Tensor]]
|
|
|
+ """
|
|
|
+ Args:
|
|
|
+ features (List[Tensor])
|
|
|
+ proposals (List[Tensor[N, 4]])
|
|
|
+ image_shapes (List[Tuple[H, W]])
|
|
|
+ targets (List[Dict])
|
|
|
+ """
|
|
|
+ if targets is not None:
|
|
|
+ for t in targets:
|
|
|
+ # TODO: https://github.com/pytorch/pytorch/issues/26731
|
|
|
+ floating_point_types = (torch.float, torch.double, torch.half)
|
|
|
+ if not t["boxes"].dtype in floating_point_types:
|
|
|
+ raise TypeError(f"target boxes must of float type, instead got {t['boxes'].dtype}")
|
|
|
+ if not t["labels"].dtype == torch.int64:
|
|
|
+ raise TypeError(f"target labels must of int64 type, instead got {t['labels'].dtype}")
|
|
|
+ if self.has_keypoint():
|
|
|
+ if not t["keypoints"].dtype == torch.float32:
|
|
|
+ raise TypeError(f"target keypoints must of float type, instead got {t['keypoints'].dtype}")
|
|
|
+
|
|
|
+ if self.training:
|
|
|
+ proposals, matched_idxs, labels, regression_targets = self.select_training_samples(proposals, targets)
|
|
|
+ else:
|
|
|
+ labels = None
|
|
|
+ regression_targets = None
|
|
|
+ matched_idxs = None
|
|
|
+
|
|
|
+ box_features = self.box_roi_pool(features, proposals, image_shapes)
|
|
|
+ box_features = self.box_head(box_features)
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+
|
|
|
+ class_logits, box_regression = self.box_predictor(box_features)
|
|
|
+
|
|
|
+ result: List[Dict[str, torch.Tensor]] = []
|
|
|
+ losses = {}
|
|
|
+ if self.training:
|
|
|
+ if labels is None:
|
|
|
+ raise ValueError("labels cannot be None")
|
|
|
+ if regression_targets is None:
|
|
|
+ raise ValueError("regression_targets cannot be None")
|
|
|
+ loss_classifier, loss_box_reg = fastrcnn_loss(class_logits, box_regression, labels, regression_targets)
|
|
|
+
|
|
|
+
|
|
|
+ losses = {"loss_classifier": loss_classifier, "loss_box_reg": loss_box_reg}
|
|
|
+ else:
|
|
|
+ boxes, scores, labels = self.postprocess_detections(class_logits, box_regression, proposals, image_shapes)
|
|
|
+ num_images = len(boxes)
|
|
|
+ for i in range(num_images):
|
|
|
+ result.append(
|
|
|
+ {
|
|
|
+ "boxes": boxes[i],
|
|
|
+ "labels": labels[i],
|
|
|
+ "scores": scores[i],
|
|
|
+ }
|
|
|
+ )
|
|
|
+ if self.has_line():
|
|
|
+
|
|
|
+ line_features = self.line_head(features)
|
|
|
+ _ = self.line_predictor(line_features)
|
|
|
+ ### line_loss(multitasklearner)
|
|
|
+
|
|
|
+
|
|
|
+ ### infer
|
|
|
+
|
|
|
+
|
|
|
+ pass
|
|
|
+ if self.has_mask():
|
|
|
+ mask_proposals = [p["boxes"] for p in result]
|
|
|
+ if self.training:
|
|
|
+ if matched_idxs is None:
|
|
|
+ raise ValueError("if in training, matched_idxs should not be None")
|
|
|
+
|
|
|
+ # during training, only focus on positive boxes
|
|
|
+ num_images = len(proposals)
|
|
|
+ mask_proposals = []
|
|
|
+ pos_matched_idxs = []
|
|
|
+ for img_id in range(num_images):
|
|
|
+ pos = torch.where(labels[img_id] > 0)[0]
|
|
|
+ mask_proposals.append(proposals[img_id][pos])
|
|
|
+ pos_matched_idxs.append(matched_idxs[img_id][pos])
|
|
|
+ else:
|
|
|
+ pos_matched_idxs = None
|
|
|
+
|
|
|
+ if self.mask_roi_pool is not None:
|
|
|
+ mask_features = self.mask_roi_pool(features, mask_proposals, image_shapes)
|
|
|
+ mask_features = self.mask_head(mask_features)
|
|
|
+ mask_logits = self.mask_predictor(mask_features)
|
|
|
+ else:
|
|
|
+ raise Exception("Expected mask_roi_pool to be not None")
|
|
|
+
|
|
|
+ loss_mask = {}
|
|
|
+ if self.training:
|
|
|
+ if targets is None or pos_matched_idxs is None or mask_logits is None:
|
|
|
+ raise ValueError("targets, pos_matched_idxs, mask_logits cannot be None when training")
|
|
|
+
|
|
|
+ gt_masks = [t["masks"] for t in targets]
|
|
|
+ gt_labels = [t["labels"] for t in targets]
|
|
|
+ rcnn_loss_mask = maskrcnn_loss(mask_logits, mask_proposals, gt_masks, gt_labels, pos_matched_idxs)
|
|
|
+ loss_mask = {"loss_mask": rcnn_loss_mask}
|
|
|
+ else:
|
|
|
+ labels = [r["labels"] for r in result]
|
|
|
+ masks_probs = maskrcnn_inference(mask_logits, labels)
|
|
|
+ for mask_prob, r in zip(masks_probs, result):
|
|
|
+ r["masks"] = mask_prob
|
|
|
+
|
|
|
+ losses.update(loss_mask)
|
|
|
+
|
|
|
+ # keep none checks in if conditional so torchscript will conditionally
|
|
|
+ # compile each branch
|
|
|
+ if (
|
|
|
+ self.keypoint_roi_pool is not None
|
|
|
+ and self.keypoint_head is not None
|
|
|
+ and self.keypoint_predictor is not None
|
|
|
+ ):
|
|
|
+ keypoint_proposals = [p["boxes"] for p in result]
|
|
|
+ if self.training:
|
|
|
+ # during training, only focus on positive boxes
|
|
|
+ num_images = len(proposals)
|
|
|
+ keypoint_proposals = []
|
|
|
+ pos_matched_idxs = []
|
|
|
+ if matched_idxs is None:
|
|
|
+ raise ValueError("if in trainning, matched_idxs should not be None")
|
|
|
+
|
|
|
+ for img_id in range(num_images):
|
|
|
+ pos = torch.where(labels[img_id] > 0)[0]
|
|
|
+ keypoint_proposals.append(proposals[img_id][pos])
|
|
|
+ pos_matched_idxs.append(matched_idxs[img_id][pos])
|
|
|
+ else:
|
|
|
+ pos_matched_idxs = None
|
|
|
+
|
|
|
+ keypoint_features = self.keypoint_roi_pool(features, keypoint_proposals, image_shapes)
|
|
|
+ keypoint_features = self.keypoint_head(keypoint_features)
|
|
|
+ keypoint_logits = self.keypoint_predictor(keypoint_features)
|
|
|
+
|
|
|
+ loss_keypoint = {}
|
|
|
+ if self.training:
|
|
|
+ if targets is None or pos_matched_idxs is None:
|
|
|
+ raise ValueError("both targets and pos_matched_idxs should not be None when in training mode")
|
|
|
+
|
|
|
+ gt_keypoints = [t["keypoints"] for t in targets]
|
|
|
+ rcnn_loss_keypoint = keypointrcnn_loss(
|
|
|
+ keypoint_logits, keypoint_proposals, gt_keypoints, pos_matched_idxs
|
|
|
+ )
|
|
|
+ loss_keypoint = {"loss_keypoint": rcnn_loss_keypoint}
|
|
|
+ else:
|
|
|
+ if keypoint_logits is None or keypoint_proposals is None:
|
|
|
+ raise ValueError(
|
|
|
+ "both keypoint_logits and keypoint_proposals should not be None when not in training mode"
|
|
|
+ )
|
|
|
+
|
|
|
+ keypoints_probs, kp_scores = keypointrcnn_inference(keypoint_logits, keypoint_proposals)
|
|
|
+ for keypoint_prob, kps, r in zip(keypoints_probs, kp_scores, result):
|
|
|
+ r["keypoints"] = keypoint_prob
|
|
|
+ r["keypoints_scores"] = kps
|
|
|
+ losses.update(loss_keypoint)
|
|
|
+
|
|
|
+ return result, losses
|
