(self, current_iter, img_name)
| 138 | |
| 139 | @master_only |
| 140 | def save_image(self, current_iter, img_name): |
| 141 | visuals = self.get_current_visuals() |
| 142 | sr_img = tensor2img([visuals['result']]) |
| 143 | out_img = [sr_img] |
| 144 | |
| 145 | if 'mask_img' in visuals: |
| 146 | mask_img = tensor2img([visuals['mask_img']]) |
| 147 | out_img.append(mask_img) |
| 148 | if 'gt' in visuals: |
| 149 | gt_img = tensor2img([visuals['gt']]) |
| 150 | out_img.append(gt_img) |
| 151 | |
| 152 | sr_img = np.hstack(out_img) |
| 153 | |
| 154 | # tentative for out of GPU memory |
| 155 | del self.lq |
| 156 | del self.output |
| 157 | torch.cuda.empty_cache() |
| 158 | |
| 159 | if self.opt['is_train']: |
| 160 | save_img_path = osp.join(self.opt['path']['visualization'], |
| 161 | f'{current_iter}_{img_name}.png') |
| 162 | else: |
| 163 | dataset_name = self.opt['datasets']['test']['name'] |
| 164 | if self.opt['val']['suffix']: |
| 165 | save_img_path = osp.join(self.opt['path']['visualization'], dataset_name, |
| 166 | f'{img_name}_{self.opt["val"]["suffix"]}.png') |
| 167 | else: |
| 168 | save_img_path = osp.join(self.opt['path']['visualization'], dataset_name, |
| 169 | f'{img_name}_{self.opt["name"]}.png') |
| 170 | |
| 171 | imwrite(sr_img, save_img_path) |
| 172 | def save(self, epoch, current_iter): |
| 173 | if hasattr(self, 'net_g_ema'): |
| 174 | self.save_network([self.net_g, self.net_g_ema], 'net_g', current_iter, param_key=['params', 'params_ema']) |
no test coverage detected