def postprocess_detections( self, logits: torch.Tensor, proposals: List[torch.Tensor]) -> Tuple[TensorList, TensorList]: keypoints_probs, kp_scores = keypointrcnn_inference(logits, proposals) proposals = [] scores = [] for keypoint_prob, kps in zip(keypoints_probs, kp_scores): proposals.append(keypoint_prob) scores.append(kps) return proposals, scores
def get_joints(self, features, keypoint_proposals, image_shapes): keypoint_features = self.roi_heads.keypoint_roi_pool( features, keypoint_proposals, image_shapes) keypoint_features = self.roi_heads.keypoint_head(keypoint_features) keypoint_logits = self.roi_heads.keypoint_predictor(keypoint_features) assert keypoint_logits is not None assert keypoint_proposals is not None keypoints_probs, kp_scores = keypointrcnn_inference( keypoint_logits, keypoint_proposals) return keypoints_probs, kp_scores
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, proposals, image_shapes, targets=None): """ Args: features: List proposals: List image_shapes: List targets: List (Default value = None) Returns: """ maskrcnn_loss_func = maskrcnn_loss fastrcnn_loss_func = fastrcnn_loss keypointrcnn_loss_func = keypointrcnn_loss eval_when_train = not self.training try: if self._eval_when_train: eval_when_train = True except AttributeError: pass if self.maskrcnn_loss_customized is not None: maskrcnn_loss_func = self.maskrcnn_loss_customized if self.fastrcnn_loss_customized is not None: fastrcnn_loss_func = self.fastrcnn_loss_customized if self.keypointrcnn_loss_customized is not None: keypointrcnn_loss_func = self.keypointrcnn_loss_customized 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_func( class_logits, box_regression, labels, regression_targets) losses = dict(loss_classifier=loss_classifier, loss_box_reg=loss_box_reg) if eval_when_train: 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: # 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_func(mask_logits, mask_proposals, gt_masks, gt_labels, pos_matched_idxs) loss_mask = dict(loss_mask=loss_mask) if eval_when_train: 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]) 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_func(keypoint_logits, keypoint_proposals, gt_keypoints, pos_matched_idxs) loss_keypoint = dict(loss_keypoint=loss_keypoint) if eval_when_train: 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): # 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 roi_heads_forward(self, features, proposals, image_shapes, targets=None): """Hack into the torchvision model to obtain features for training caption model; training is assumed to be false https://github.com/pytorch/vision/blob/master/ torchvision/models/detection/roi_heads.py""" 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]], []) boxes, scores, labels, box_features = roi_postprocess_detections( self, class_logits, box_regression, proposals, image_shapes, box_features) num_images = len(boxes) for i in range(num_images): result.append({ "boxes_features": box_features[i], "boxes": boxes[i], "labels": labels[i], "scores": scores[i] }) if self.has_mask(): mask_proposals = [p["boxes"] for p in result] 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") 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_features"] = mask_features r["masks"] = mask_prob # 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] 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) 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_features"] = keypoint_features r["keypoints"] = keypoint_prob r["keypoints_scores"] = kps return result, dict()