| 157 | return ax, fig, img.shape[0], img.shape[1] |
| 158 | |
| 159 | def inference(self): |
| 160 | self.model.eval() |
| 161 | |
| 162 | res_file = os.path.join(self.output, f'results.txt.rank{self.dist.rank}') |
| 163 | writer = open(res_file, 'w') |
| 164 | for batch_idx, batch in enumerate(self.val_data['loader']): |
| 165 | input = batch['image'] |
| 166 | input = input.cuda().half() if self.fp16 else input.cuda() |
| 167 | |
| 168 | # compute output |
| 169 | logits = self.model(input) |
| 170 | |
| 171 | scores = F.softmax(logits, dim=1) |
| 172 | # compute prediction |
| 173 | _, preds = logits.data.topk(k=1, dim=1) |
| 174 | preds = preds.view(-1) |
| 175 | # update batch information |
| 176 | batch.update({'prediction': preds.detach()}) |
| 177 | batch.update({'score': scores.detach()}) |
| 178 | # save prediction information |
| 179 | if self.cam: |
| 180 | heatmap = self.gradCam(input) |
| 181 | for idx in range(len(heatmap)): |
| 182 | basename = os.path.basename(batch["filename"][idx]) |
| 183 | ext = basename.split(".")[-1] |
| 184 | basename = basename.replace("." + ext, "_cam" + "." + ext) |
| 185 | |
| 186 | heatmap[idx].save(os.path.join(self.output, basename)) |
| 187 | |
| 188 | if self.visualize: |
| 189 | self.paint(batch["filename"], scores, preds, self.output) |
| 190 | self.val_data['loader'].dataset.dump(writer, batch) |
| 191 | |
| 192 | writer.close() |
| 193 | link.barrier() |
| 194 | |
| 195 | return |
| 196 | |
| 197 | def save_feature(self, module, input, output): |
| 198 | self.feature = output |