Args: input (dict): one dataset dict with "image" field being a CHW tensor Returns: dict: one output dict
(self, input)
| 124 | return processed_results |
| 125 | |
| 126 | def _inference_one_image(self, input): |
| 127 | """ |
| 128 | Args: |
| 129 | input (dict): one dataset dict with "image" field being a CHW tensor |
| 130 | Returns: |
| 131 | dict: one output dict |
| 132 | """ |
| 133 | augmented_inputs, tfms = self._get_augmented_inputs(input) |
| 134 | |
| 135 | final_predictions = None |
| 136 | count_predictions = 0 |
| 137 | for input, tfm in zip(augmented_inputs, tfms): |
| 138 | count_predictions += 1 |
| 139 | with torch.no_grad(): |
| 140 | if final_predictions is None: |
| 141 | if any(isinstance(t, HFlipTransform) for t in tfm.transforms): |
| 142 | final_predictions = self.model(input, self.current_step)[0].pop("sem_seg").flip(dims=[2]) # should be [input] originally |
| 143 | else: |
| 144 | final_predictions = self.model(input, self.current_step)[0].pop("sem_seg") |
| 145 | else: |
| 146 | if any(isinstance(t, HFlipTransform) for t in tfm.transforms): |
| 147 | final_predictions += self.model(input, self.current_step)[0].pop("sem_seg").flip(dims=[2]) |
| 148 | else: |
| 149 | final_predictions += self.model(input, self.current_step)[0].pop("sem_seg") |
| 150 | |
| 151 | final_predictions = final_predictions / count_predictions |
| 152 | return {"sem_seg": final_predictions} |
| 153 | |
| 154 | def _get_augmented_inputs(self, input): |
| 155 | augmented_inputs = self.tta_mapper(input) |
no test coverage detected