def compute_loss(self, model, net_output, sample): # get rpn proposals and box detections rpn_proposals = model.get_rpn_proposals(net_output) box_detections = model.get_box_detections(net_output) # compute proposal losses anchors = model.rpn.get_anchors(rpn_proposals) objectness = model.rpn.get_objectness(rpn_proposals) pred_bbox_deltas = model.rpn.get_pred_bbox_deltas(rpn_proposals) labels, matched_gt_boxes = model.rpn.assign_targets_to_anchors(anchors, sample['target']) regression_targets = model.rpn.box_coder.encode(matched_gt_boxes, anchors) loss_objectness, loss_rpn_box_reg = model.rpn.compute_loss( objectness, pred_bbox_deltas, labels, regression_targets, ) # compute detector loss labels = model.roi_heads.get_labels(box_detections) regression_targets = model.roi_heads.get_regression_targets(box_detections) class_logits = model.roi_heads.get_class_logits(box_detections) box_regression = model.roi_heads.get_box_regression(box_detections) loss_classifier, loss_box_reg = fastrcnn_loss( class_logits, box_regression, labels, regression_targets, ) losses = { # proposal losses 'loss_objectness': loss_objectness, 'loss_rpn_box_reg': loss_rpn_box_reg, # detector loss 'loss_classifier': loss_classifier, 'loss_box_reg': loss_box_reg, } return losses
def forward(self, features, proposals, image_shapes, targets=None): """ Arguments: 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: assert t["boxes"].dtype.is_floating_point, \ 'target boxes must of float type' assert t["labels"].dtype == torch.int64, \ 'target labels must of int64 type' if self.has_keypoint(): assert t["keypoints"].dtype == torch.float32, \ 'target keypoints must of float type' if self.training: proposals, matched_idxs, labels, regression_targets = \ self.select_training_samples(proposals, targets) # from IPython import embed # embed() roi_pooled_features = \ self.box_roi_pool(features, proposals, image_shapes) rcnn_features = self.box_head(roi_pooled_features) cls_logits, box_regression = self.box_predictor(rcnn_features) person_feat = self.reid_head(rcnn_features) result, losses = [], {} if self.training: det_labels = [y.clamp(0, 1) for y in labels] #just 0 or 1 loss_detection, loss_box_reg = \ fastrcnn_loss(cls_logits, box_regression, det_labels, regression_targets) loss_reid = self.reid_loss(person_feat, labels) losses = dict(loss_detection=loss_detection, loss_box_reg=loss_box_reg, loss_reid=loss_reid) else: boxes, scores, person_feat, labels = \ self.oim_postprocess_detections(cls_logits, box_regression,person_feat, proposals, image_shapes) num_images = len(boxes) for i in range(num_images): result.append( dict( boxes=boxes[i], labels=labels[i], scores=scores[i], person_feat=person_feat[i], )) # Mask and Keypoint losses are deleted return result, losses
def forward(self, features, proposals, image_shapes, targets=None): # type: (Dict[str, Tensor], List[Tensor], List[Tuple[int, int]], Optional[List[Dict[str, Tensor]]]) """ Arguments: 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) assert t[ "boxes"].dtype in floating_point_types, 'target boxes must of float type' assert t[ "labels"].dtype == torch.int64, 'target labels must of int64 type' if self.has_keypoint(): assert t[ "keypoints"].dtype == torch.float32, 'target keypoints must of float type' 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 = torch.jit.annotate(List[Dict[str, torch.Tensor]], []) losses = {} if self.training: assert labels is not None and regression_targets is not 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, props, prob_max, scores_cls = 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], "props": props[i], "prob_max": prob_max[i], "scores_cls": scores_cls[i], "features": features }) return result, losses
def get_detections_losses(self, model, class_logits, box_regression, labels, regression_targets, proposals, image_sizes, original_image_sizes, name=''): """ Arguments: model (torchvision FasterRCNN model) class_logits, box_regression: output of model.roi_heads.box_predictor(box_features) labels (List[Tensor[N']]): ground truth class labels for training regression_targets (List[Tensor[N', 4]]): ground truth targets for box regression for training proposals (List[Tensor[N, 4]]): box proposals image_sizes (list[tuple[int, int]]): original shapes of all images original_image_sizes (List[Tuple[int, int]]): to postprocess detections name (str): append this name to keys in losses Returns: detections (List[Dict[str, torch.Tensor]]): predicted boxes, labels, and scores (in eval mode) detector_losses (Dict): classifier and box_regression losses for final detection (in train mode) """ detections = torch.jit.annotate(List[Dict[str, torch.Tensor]], []) detector_losses = {} if self.training: assert labels is not None and regression_targets is not None loss_classifier, loss_box_reg = fastrcnn_loss( class_logits, box_regression, labels, regression_targets) detector_losses = { 'loss_classifier' + name: loss_classifier, 'loss_box_reg' + name: loss_box_reg } else: boxes, scores, labels = model.roi_heads.postprocess_detections( class_logits, box_regression, proposals, image_sizes) num_images = len(boxes) for i in range(num_images): detections.append({ 'boxes': boxes[i], 'labels': labels[i], 'scores': scores[i] }) detections = self.transform.postprocess(detections, image_sizes, original_image_sizes) return detections, detector_losses
def forward(self, features, proposals, image_shapes, targets=None): """ Arguments: 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: assert t["boxes"].dtype.is_floating_point, 'target boxes must of float type' assert t["labels"].dtype == torch.int64, 'target labels must of int64 type' if self.has_keypoint: assert t["keypoints"].dtype == torch.float32, 'target keypoints must of float type' if self.training: proposals, matched_idxs, labels, regression_targets = self.select_training_samples(proposals, targets) 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, losses = [], {} if self.training: loss_classifier, loss_box_reg = fastrcnn_loss( class_logits, box_regression, labels, regression_targets) losses = dict(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( dict( boxes=boxes[i], labels=labels[i], scores=scores[i], ) ) return result, losses
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]]] ): # modified to return box features # 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) assert ( t["boxes"].dtype in floating_point_types ), "target boxes must of float type" assert ( t["labels"].dtype == torch.int64 ), "target labels must of int64 type" if self.has_keypoint(): assert ( t["keypoints"].dtype == torch.float32 ), "target keypoints must of float type" 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_features0 = self.box_roi_pool(features, proposals, image_shapes) ## box_featuers0 is a 256x7x7 (the 7x7 is a parameter from the box_roi_palign). # it is then flattened to a single 1024 vector by the box_head box_features = self.box_head(box_features0) # there is a final box for each class (it is category dependent). class_logits, box_regression = self.box_predictor(box_features) result: List[Dict[str, torch.Tensor]] = [] losses = {} if self.training: assert labels is not None and regression_targets is not 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], } ) return {"result": result, "losses": losses, "box_features": box_features}
def forward(self, features, proposals, image_shapes, targets=None): """ Arguments: 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: assert t[ "boxes"].dtype.is_floating_point, 'target boxes must of float type' assert t[ "labels"].dtype == torch.int64, 'target labels must of int64 type' assert t["poses"].dtype.is_floating_point if self.has_keypoint: assert t[ "keypoints"].dtype == torch.float32, 'target keypoints must of float type' if self.training: proposals, matched_idxs, labels, regression_targets = self.select_training_samples( proposals, targets) box_features = self.box_roi_pool( features, proposals, image_shapes) # torch.Size([bs*1000, 256, 7, 7]) box_features = self.box_head( box_features) # torch.Size([bs*1000, 1024]) class_logits, box_regression = self.box_predictor( box_features) # torch.Size([bs*1000, 2]) torch.Size([bs*1000, 8]) result, losses = [], {} # result 是一个字典的列表, 每一个字典存着每张图片的预测值 if self.training: loss_classifier, loss_box_reg = fastrcnn_loss( class_logits, box_regression, labels, regression_targets) losses = dict(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( dict( boxes=boxes[i], labels=labels[i], scores=scores[i], )) # 如果是Mask R-CNN if self.has_mask: mask_proposals = [p["boxes"] for p in result] if self.training: # 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.nonzero(labels[img_id] > 0).squeeze(1) mask_proposals.append(proposals[img_id][pos]) pos_matched_idxs.append(matched_idxs[img_id][pos]) 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) loss_mask = {} if self.training: gt_masks = [t["masks"] for t in targets] gt_labels = [t["labels"] for t in targets] loss_mask = maskrcnn_loss(mask_logits, mask_proposals, gt_masks, gt_labels, pos_matched_idxs) loss_mask = dict(loss_mask=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) if self.has_keypoint: 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 = [] for img_id in range(num_images): pos = torch.nonzero(labels[img_id] > 0).squeeze(1) keypoint_proposals.append( proposals[img_id][pos]) # shape=(num_pos, 4) pos_matched_idxs.append(matched_idxs[img_id][pos]) 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: gt_keypoints = [t["keypoints"] for t in targets] loss_keypoint = keypointrcnn_loss(keypoint_logits, keypoint_proposals, gt_keypoints, pos_matched_idxs) loss_keypoint = dict(loss_keypoint=loss_keypoint) else: 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) if self.has_pose: pose_proposals = [p["boxes"] for p in result] if self.training: # during training, only focus on positive boxes num_images = len(proposals) pose_proposals = [] pos_matched_idxs = [] for img_id in range(num_images): pos = torch.nonzero(labels[img_id] > 0).squeeze( 1) # 所有被分配为正样本的proposal的下标 pose_proposals.append( proposals[img_id] [pos]) # proposal的box(xmin, ymin, xmax, ymax) pos_matched_idxs.append( matched_idxs[img_id][pos]) # 每个proposal对应哪个target pose pose_features = self.pose_roi_pool(features, pose_proposals, image_shapes) pose_features = self.pose_head(pose_features) # pose_regression = self.pose_predictor(pose_features) loss_pose = {} if self.training: gt_poses = [t["poses"] for t in targets] # a list of (rx, ry, rz, tz) loss_pose = posercnn_loss(pose_regression, gt_poses, labels, pos_matched_idxs) loss_pose = dict(loss_pose=loss_pose) else: pred_poses = postprocess_poses(pose_regression, pose_proposals) for poses, r in zip(pred_poses, result): r['poses'] = poses losses.update(loss_pose) if self.has_trans: trans_proposals = [p["boxes"] for p in result] if self.training: # during training, only focus on positive boxes num_images = len(proposals) trans_proposals = [] pos_matched_idxs = [] for img_id in range(num_images): # keep_only_positive_boxes pos = torch.nonzero(labels[img_id] > 0).squeeze( 1) # 所有被分配为正样本的proposal的下标 trans_proposals.append( proposals[img_id] [pos]) # proposal的box(xmin, ymin, xmax, ymax) pos_matched_idxs.append( matched_idxs[img_id] [pos]) # 每个proposal对应哪个target pose box_features = torch.cat( trans_proposals, dim=0) # [N, 4] N=batch_size*num_proposal_per_image trans_features = self.translation_head(box_features) trans_pred = self.translation_predictor( trans_features, pose_features) loss_trans = {} if self.training: gt_trans = [t["translations"] for t in targets] # 6DVNET中平移损失的权重是0.05 loss_trans = 0.05 * trans_loss(trans_pred, gt_trans, labels, pos_matched_idxs) loss_trans = dict(loss_trans=loss_trans) else: pred_trans = postprocess_trans(trans_pred, trans_proposals) for translations, r in zip(pred_trans, result): r['translations'] = translations losses.update(loss_trans) return result, losses
def forward(self, features, proposals, image_shapes, targets=None, return_loss=False): """ Arguments: features (List[Tensor]) proposals (List[Tensor[N, 4]]) image_shapes (List[Tuple[H, W]]) targets (List[Dict]) return_loss (Bool): return the loss (even if we are in eval mode) """ if targets is not None: for t in targets: assert t["boxes"].dtype.is_floating_point, 'target boxes must of float type' assert t["labels"].dtype == torch.int64, 'target labels must of int64 type' # if self.has_keypoint: # assert t["keypoints"].dtype == torch.float32, 'target keypoints must of float type' if self.training or return_loss: proposals, matched_idxs, labels, regression_targets = self.select_training_samples(proposals, targets) 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, losses = [], {} if self.training or return_loss: loss_classifier, loss_box_reg = fastrcnn_loss( class_logits, box_regression, labels, regression_targets) losses = dict(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( dict( boxes=boxes[i], labels=labels[i], scores=scores[i], ) ) if self.has_mask: mask_proposals = [p["boxes"] for p in result] if self.training or return_loss: # 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.nonzero(labels[img_id] > 0).squeeze(1) mask_proposals.append(proposals[img_id][pos]) pos_matched_idxs.append(matched_idxs[img_id][pos]) 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) loss_mask = {} if self.training or return_loss: gt_masks = [t["masks"] for t in targets] gt_labels = [t["labels"] for t in targets] loss_mask = maskrcnn_loss( mask_logits, mask_proposals, gt_masks, gt_labels, pos_matched_idxs) loss_mask = dict(loss_mask=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) if self.has_keypoint: keypoint_proposals = [p["boxes"] for p in result] if self.training or return_loss: # during training, only focus on positive boxes num_images = len(proposals) keypoint_proposals = [] pos_matched_idxs = [] for img_id in range(num_images): pos = torch.nonzero(labels[img_id] > 0).squeeze(1) keypoint_proposals.append(proposals[img_id][pos]) pos_matched_idxs.append(matched_idxs[img_id][pos]) 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 or return_loss: gt_keypoints = [t["keypoints"] for t in targets] loss_keypoint = keypointrcnn_loss( keypoint_logits, keypoint_proposals, gt_keypoints, pos_matched_idxs) loss_keypoint = dict(loss_keypoint=loss_keypoint) else: 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
def forward( self, features: Dict[str, torch.Tensor], proposals: List[torch.Tensor], image_shapes: List[Tuple[int, int]], targets: Optional[List[Dict[str, torch.Tensor]]] = None ) -> Tuple[List[Dict[str, torch.Tensor]], Dict[str, torch.Tensor]]: if targets is not None: for t in targets: floating_point_types = (torch.float, torch.double, torch.half) assert t[ "boxes"].dtype in floating_point_types, 'target boxes must of float type' assert t[ "ellipse_matrices"].dtype in floating_point_types, 'target ellipse_offsets must of float type' assert t[ "labels"].dtype == torch.int64, 'target labels must of int64 type' 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: assert labels is not None and regression_targets is not 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_ellipse_reg(): ellipse_proposals = [p["boxes"] for p in result] if self.training: assert matched_idxs is not None # during training, only focus on positive boxes num_images = len(proposals) ellipse_proposals = [] pos_matched_idxs = [] for img_id in range(num_images): pos = torch.where(labels[img_id] > 0)[0] ellipse_proposals.append(proposals[img_id][pos]) pos_matched_idxs.append(matched_idxs[img_id][pos]) else: pos_matched_idxs = None if self.ellipse_roi_pool is not None: ellipse_features = self.ellipse_roi_pool( features, ellipse_proposals, image_shapes) ellipse_features = self.ellipse_head(ellipse_features) ellipse_shapes_normalised = self.ellipse_predictor( ellipse_features) else: raise Exception("Expected ellipse_roi_pool to be not None") loss_ellipse_regressor = {} if self.training: assert targets is not None assert pos_matched_idxs is not None assert ellipse_shapes_normalised is not None ellipse_matrix_targets = [ t["ellipse_matrices"] for t in targets ] rcnn_loss_ellipse = self.ellipse_loss_fn( ellipse_shapes_normalised, ellipse_matrix_targets, pos_matched_idxs, ellipse_proposals) loss_ellipse_regressor = {"loss_ellipse": rcnn_loss_ellipse} else: ellipses_per_image = [l.shape[0] for l in labels] for e_l, r, box in zip( ellipse_shapes_normalised.split(ellipses_per_image, dim=0), result, ellipse_proposals): d_a = e_l[:, 0] d_b = e_l[:, 1] d_angle = e_l[:, 2] r["ellipse_matrices"] = postprocess_ellipse_predictor( d_a, d_b, d_angle, box) losses.update(loss_ellipse_regressor) return result, losses
def forward(self, features, proposals, image_shapes, targets=None, task_heads=None): """ Arguments: 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) assert t["boxes"].dtype in floating_point_types, 'target boxes must of float type' assert t["labels"].dtype == torch.int64, 'target labels must of int64 type' 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) task_results = {task_head.name: task_head(box_features) for task_head in task_heads} result = torch.jit.annotate(List[Dict[str, torch.Tensor]], []) losses = {} if self.training: assert labels is not None and regression_targets is not 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, } # Calculate losses for all the tasks for name, preds in task_results.items(): actuals = [] for idxs, target in zip(matched_idxs, targets): actuals.append(target[name][idxs]) actuals = torch.cat(actuals, dim=0) # Discount the loss for task so that model still prioritises learning # the box / mask properly losses[f"loss_{name}"] = F.cross_entropy(preds, actuals) else: boxes, scores, labels = self.postprocess_detections(class_logits, box_regression, proposals, image_shapes) boxes_per_image = [len(boxes_in_image) for boxes_in_image in proposals] if len(boxes_per_image) == 1: task_results_list = task_results else: task_results_list = {name: preds.split(boxes_per_image, 0) for name, preds in task_results.items()} num_images = len(boxes) for i in range(num_images): result.append( { "boxes": boxes[i], "labels": labels[i], "scores": scores[i], "task_results": {name: preds[i] for name, preds in task_results_list.items()} } ) if self.has_mask(): mask_proposals = [p["boxes"] for p in result] if self.training: assert matched_idxs is not 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.nonzero(labels[img_id] > 0).squeeze(1) 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: mask_logits = torch.tensor(0) raise Exception("Expected mask_roi_pool to be not None") loss_mask = {} if self.training: assert targets is not None assert pos_matched_idxs is not None assert mask_logits is not None 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) return result, losses
def forward(self, features, proposals, image_shapes, targets=None): # type: (Dict[str, Tensor], List[Tensor], List[Tuple[int, int]], Optional[List[Dict[str, Tensor]]]) """ Arguments: 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) assert t[ "boxes"].dtype in floating_point_types, 'target boxes must of float type' assert t[ "labels"].dtype == torch.int64, 'target labels must of int64 type' if self.has_keypoint(): assert t[ "keypoints"].dtype == torch.float32, 'target keypoints must of float type' if self.training: proposals, matched_idxs, labels, regression_targets = self.select_training_samples( proposals, targets) else: labels = None regression_targets = None matched_idxs = None # compute bbox embedding position_matrix = extract_position_matrix(proposals, Nongt_dim) position_embedding = extract_position_embedding(position_matrix, feat_dim=64) box_features = self.box_roi_pool(features, proposals, image_shapes) box_features = self.box_head(box_features, position_embedding) class_logits, box_regression = self.box_predictor(box_features) result = torch.jit.annotate(List[Dict[str, torch.Tensor]], []) losses = {} if self.training: assert labels is not None and regression_targets is not 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_mask(): mask_proposals = [p["boxes"] for p in result] if self.training: assert matched_idxs is not 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.nonzero(labels[img_id] > 0).squeeze(1) 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: mask_logits = torch.tensor(0) raise Exception("Expected mask_roi_pool to be not None") loss_mask = {} if self.training: assert targets is not None assert pos_matched_idxs is not None assert mask_logits is not None 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 = [] assert matched_idxs is not None for img_id in range(num_images): pos = torch.nonzero(labels[img_id] > 0).squeeze(1) 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: assert targets is not None assert pos_matched_idxs is not None 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: assert keypoint_logits is not None assert keypoint_proposals is not None 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
def forward(self, features, proposals, image_shapes, targets=None): """ Arguments: features (List[Tensor]) proposals (List[Tensor[N, 4]]) image_shapes (List[Tuple[H, W]]) targets (List[Dict]) """ feature_dims = np.array( [features[layer].shape[1] for layer in features]) if np.all(feature_dims == self.in_channels): # RGB only features_rgb = features elif np.all(feature_dims == 2 * self.in_channels): # RGB-depth 6 channel, two backbones from collections import OrderedDict features_rgb = OrderedDict() for key in features.keys(): features_rgb[key] = features[key][:, :self.in_channels] else: # RGB-D 4 channel features_rgb = features # Detection if self.training: proposals, matched_idxs, labels, regression_targets = self.select_training_samples( proposals, targets) box_features = self.box_roi_pool(features_rgb, proposals, image_shapes) box_features = self.box_head(box_features) class_logits, box_regression = self.box_predictor(box_features) result, losses = [], {} if self.training: loss_classifier, loss_box_reg = fastrcnn_loss( class_logits, box_regression, labels, regression_targets) losses = dict(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): if boxes[i].shape[0] == 0: return result, losses result.append( dict( boxes=boxes[i], labels=labels[i], scores=scores[i], )) # Proposals selected by detection stage is shared by all other branches box_proposals = [p["boxes"] for p in result] if self.training: # during training, only focus on positive boxes num_images = len(proposals) box_proposals = [] pos_matched_idxs = [] for img_id in range(num_images): pos = torch.nonzero(labels[img_id] > 0).squeeze(1) box_proposals.append(proposals[img_id][pos]) pos_matched_idxs.append(matched_idxs[img_id][pos]) # MultiStage RoIAlign is shared by all other branches shared_features_rgb = self.shared_roi_pool(features_rgb, box_proposals, image_shapes) # Segmentation mask_features = self.mask_head(shared_features_rgb) mask_logits = self.mask_predictor(mask_features) loss_mask = {} masks_on_features = None if self.training: gt_masks = [t["masks"] for t in targets] gt_labels = [t["labels"] for t in targets] gt_is = [t["instance_ids"] for t in targets] loss_mask, masks_for_paf, masks_for_vote = maskrcnn_loss_updated( mask_logits, box_proposals, gt_masks, gt_labels, gt_is, pos_matched_idxs) loss_mask = dict(loss_mask=loss_mask) else: ref_labels = [r["labels"] for r in result] masks_probs = maskrcnn_inference(mask_logits, ref_labels) for mask_prob, r in zip(masks_probs, result): r["masks"] = mask_prob losses.update(loss_mask) if self.with_paf_branch: paf_features = self.paf_head(shared_features_rgb) paf_logits = self.paf_predictor(paf_features) loss_paf = {} if self.training: gt_pafs = [t["target_pafs"] for t in targets] loss_paf = paf_loss_updated(paf_logits, masks_for_paf, pos_matched_idxs, gt_pafs, gt_labels) if torch.isnan(loss_paf): print('error') loss_paf = dict(loss_paf=loss_paf) else: paf_ref_labels = torch.cat(ref_labels) - 1 N, _, H, W = paf_logits.shape paf_logits = paf_logits.view(N, -1, 2, H, W)[torch.arange(N), paf_ref_labels] paf_probs = [paf_logits] for paf_prob, r in zip(paf_probs, result): r["pafs"] = F.normalize(paf_prob, dim=1) losses.update(loss_paf) if self.input_mode == config.INPUT_RGBD: shared_features = self.attention_block( shared_features_rgb ) # shared_features_rgb actually has 4-channel RGBD input bs, c, _, _ = shared_features.shape # shared_features = shared_features.view(bs, c, -1) # for conv1d if self.with_3d_keypoints: keypoint_features = self.vote_keypoint_head(shared_features) # keypoint_features = keypoint_features.view(bs, self.keypoint_dim_reduced, 14, 14) keypoint_offsets = self.vote_keypoint_predictor( keypoint_features) loss_keypoint = {} if self.training: gt_3d_keypoints = [t["frame"][:, :3] for t in targets] ori_depth = [t["ori_image_depth"] for t in targets] gt_labels = [t["labels"] for t in targets] loss_keypoint = vote_keypoint_loss( keypoint_offsets, box_proposals, ori_depth, gt_3d_keypoints, pos_matched_idxs, masks_for_vote, gt_labels) loss_keypoint = dict(loss_keypoint=loss_keypoint) else: ref_labels = torch.cat(ref_labels) - 1 N, _, H, W = keypoint_offsets.shape keypoint_offsets = keypoint_offsets.view( N, -1, 3, H, W)[torch.arange(N), ref_labels] keypoints = [keypoint_offsets] for kps, r in zip(keypoints, result): r["keypoints_offset"] = kps losses.update(loss_keypoint) if self.with_axis_keypoints: keypoint_features = self.orientation_keypoint_head( shared_features) # keypoint_features = keypoint_features.view(bs, self.keypoint_dim_reduced, 14, 14) axis_keypoint_offsets = self.orientation_keypoint_predictor( keypoint_features) N, _, H, W = axis_keypoint_offsets.shape axis_keypoint_offsets = axis_keypoint_offsets.view( N, -1, 2, 3, H, W) loss_orientation = {} if self.training: gt_3d_keypoints = [t["axis_keypoints"] for t in targets] ori_depth = [t["ori_image_depth"] for t in targets] loss_orientation = vote_orientation_loss( axis_keypoint_offsets, box_proposals, ori_depth, gt_3d_keypoints, pos_matched_idxs, masks_for_vote, gt_labels) loss_orientation = dict(loss_orientation=loss_orientation) else: axis_keypoint_offsets = axis_keypoint_offsets[ torch.arange(N), ref_labels] axis_keypoints = [axis_keypoint_offsets] for kps, r in zip(axis_keypoints, result): r["axis_keypoint_offsets"] = kps losses.update(loss_orientation) if self.regress_axis: keypoint_features = self.axis_head(shared_features) keypoint_features = keypoint_features.view( bs, self.keypoint_dim_reduced, 14, 14) axis_keypoint_offsets = self.axis_predictor(keypoint_features) N, _, H, W = axis_keypoint_offsets.shape axis_keypoint_offsets = axis_keypoint_offsets.view( N, -1, 4, H, W) loss_axis = {} if self.training: gt_3d_keypoints = [t["axis_keypoints"] for t in targets] ori_depth = [t["ori_image_depth"] for t in targets] loss_axis = vote_axis_loss(axis_keypoint_offsets, box_proposals, ori_depth, gt_3d_keypoints, pos_matched_idxs, masks_for_vote, gt_labels) loss_axis = dict(loss_axis=loss_axis) else: axis_keypoint_offsets = axis_keypoint_offsets[ torch.arange(N), ref_labels] axis_keypoints = [axis_keypoint_offsets] for kps, r in zip(axis_keypoints, result): r["axis_offsets"] = kps losses.update(loss_axis) if self.estimate_norm_vector: keypoint_features = self.norm_vector_head(shared_features) keypoint_features = keypoint_features.view( bs, self.keypoint_dim_reduced, 14, 14) norm_vectors = self.norm_vector_predictor(keypoint_features) N, _, H, W = norm_vectors.shape norm_vectors = norm_vectors.view(N, -1, 3, H, W) loss_norm_vector = {} if self.training: gt_3d_keypoints = [t["axis_keypoints"] for t in targets] loss_norm_vector = calculate_norm_vectors( norm_vectors, gt_3d_keypoints, pos_matched_idxs, masks_for_vote, gt_labels) loss_norm_vector = dict(loss_norm_vector=loss_norm_vector) else: norm_vectors = F.normalize(norm_vectors, dim=2) norm_vectors = norm_vectors[torch.arange(N), ref_labels] estimate_norm_vectors = [norm_vectors] for norm_v, r in zip(estimate_norm_vectors, result): r["norm_vector"] = norm_v losses.update(loss_norm_vector) return result, losses