produce up to 6 prediction results for each click
(self, image_ori, image, point=None)
| 24 | self.point = None |
| 25 | |
| 26 | def predict(self, image_ori, image, point=None): |
| 27 | """ |
| 28 | produce up to 6 prediction results for each click |
| 29 | """ |
| 30 | width = image_ori.shape[1] |
| 31 | height = image_ori.shape[0] |
| 32 | |
| 33 | data = {"image": image, "height": height, "width": width} |
| 34 | # import ipdb; ipdb.set_trace() |
| 35 | if point is None: |
| 36 | point = torch.tensor([[0.5, 0.5, 0.006, 0.006]]).cuda() |
| 37 | else: |
| 38 | point = torch.tensor(point).cuda() |
| 39 | point_ = point |
| 40 | point = point_.clone() |
| 41 | point[0, 0] = point_[0, 0] |
| 42 | point[0, 1] = point_[0, 1] |
| 43 | # point = point[:, [1, 0]] |
| 44 | point = torch.cat([point, point.new_tensor([[0.005, 0.005]])], dim=-1) |
| 45 | |
| 46 | self.point = point[:, :2].clone()*(torch.tensor([width, height]).to(point)) |
| 47 | |
| 48 | data['targets'] = [dict()] |
| 49 | data['targets'][0]['points'] = point |
| 50 | data['targets'][0]['pb'] = point.new_tensor([0.]) |
| 51 | |
| 52 | batch_inputs = [data] |
| 53 | masks, ious = self.model.model.evaluate_demo(batch_inputs) |
| 54 | |
| 55 | return masks, ious |
| 56 | |
| 57 | def process_multi_mask(self, masks, ious, image_ori): |
| 58 | pred_masks_poses = masks |
no test coverage detected