def get_instance(self, frame_id): if frame_id not in self.frame_ids: return None else: i = self.frame_ids.index(frame_id) if self.proposal_instances[i] is not None: return self.proposal_instances[i] else: # This has been a skipped frame ... interpolate box from the neighboring instances index_before = index_after = i while self.proposal_instances[ index_before] is None and index_before > 0: index_before -= 1 while self.proposal_instances[ index_after] is None and index_after < len( self.proposal_instances): index_after += 1 instance_before = self.proposal_instances[index_before] instance_after = self.proposal_instances[index_after] interpolation_factor = (i - index_before) / (index_after - index_before) interpolated_instance = Instances(instance_before.image_size) interpolated_instance.pred_boxes = Boxes( instance_before.pred_boxes.tensor + interpolation_factor * (instance_after.pred_boxes.tensor - instance_before.pred_boxes.tensor)) interpolated_instance.scores = torch.tensor([0]) interpolated_instance.pred_classes = instance_before.pred_classes interpolated_instance.class_distributions = instance_before.class_distributions interpolated_instance.generation_process = ["I"] return interpolated_instance
def generate_poposals(images, model, score_threshold=0): inputs = [{ "image": torch.as_tensor(image.astype("float32").transpose(2, 0, 1)), "height": image.shape[0], "width": image.shape[1] } for image in images] with torch.no_grad(): images = model.preprocess_image(inputs) features = model.backbone(images.tensor) proposals, _ = model.proposal_generator(images, features, None) features_ = [features[f] for f in model.roi_heads.box_in_features] box_features = model.roi_heads.box_pooler( features_, [x.proposal_boxes for x in proposals]) box_features = model.roi_heads.box_head(box_features) proposals_scores, proposals_deltas = model.roi_heads.box_predictor( box_features) boxes_tensors = model.roi_heads.box_predictor.predict_boxes( (proposals_scores, proposals_deltas), proposals) scores = model.roi_heads.box_predictor.predict_probs( (proposals_scores, proposals_deltas), proposals) result = [] for i in range(len(inputs)): image_size = proposals[i].image_size num_bbox_reg_classes = boxes_tensors[i].shape[1] // 4 boxes = Boxes(boxes_tensors[i].reshape(-1, 4)) boxes.clip(image_size) boxes = boxes.tensor.view(-1, num_bbox_reg_classes, 4) img_scores = scores[i][:, :-1] max_scores, pred_classes = torch.max(img_scores, dim=1) keep_mask = max_scores > score_threshold filtered_scores = img_scores[keep_mask, :] filtered_max_scores = max_scores[keep_mask] filtered_pred_classes = pred_classes[keep_mask] boxes = boxes[keep_mask, filtered_pred_classes, :] result_instance = Instances(image_size) result_instance.pred_boxes = Boxes(boxes) result_instance.scores = filtered_max_scores result_instance.pred_classes = filtered_pred_classes result_instance.class_distributions = filtered_scores result.append(result_instance) return result
def project_proposal_instance(self, frame_index): if len(self.proposal_instances) == 1: # We consider the first frame, there is nothing to project here return self.proposal_instances[0] instance_index_current = self.last_key_instance_index[-1] instance_index_before = self.last_key_instance_index[-2] frame_index_current = self.frame_ids[instance_index_current] assert frame_index > frame_index_current if frame_index == frame_index_current: # We are replacing the most recent proposal instance if len(self.proposal_instances) < 3: instance_index_current = self.last_key_instance_index[-2] return self.proposal_instances[instance_index_current] else: instance_index_current = self.last_key_instance_index[-2] instance_index_before = self.last_key_instance_index[-3] frame_index_current = self.frame_ids[instance_index_current] instance_current = self.proposal_instances[instance_index_current] instance_before = self.proposal_instances[instance_index_before] centers_current = instance_current.pred_boxes.get_centers() centers_before = instance_before.pred_boxes.get_centers() centers_delta = (centers_current - centers_before) / ( instance_index_current - instance_index_before) * (frame_index - frame_index_current) projected_instance = Instances(instance_current.image_size) projected_instance.pred_boxes = Boxes( instance_current.pred_boxes.tensor + centers_delta.repeat(1, 2)) projected_instance.scores = instance_current.scores projected_instance.pred_classes = instance_current.pred_classes projected_instance.class_distributions = instance_current.class_distributions projected_instance.generation_process = ["P"] return projected_instance