| 436 | return f'PSNR = {self.measure():.6f}' |
| 437 | |
| 438 | class LPIPSMeter: |
| 439 | def __init__(self, net='alex', device=None): |
| 440 | self.V = 0 |
| 441 | self.N = 0 |
| 442 | self.net = net |
| 443 | |
| 444 | self.device = device if device is not None else torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| 445 | self.fn = lpips.LPIPS(net=net).eval().to(self.device) |
| 446 | |
| 447 | def clear(self): |
| 448 | self.V = 0 |
| 449 | self.N = 0 |
| 450 | |
| 451 | def prepare_inputs(self, *inputs): |
| 452 | outputs = [] |
| 453 | for i, inp in enumerate(inputs): |
| 454 | inp = inp.permute(0, 3, 1, 2).contiguous() # [B, 3, H, W] |
| 455 | inp = inp.to(self.device) |
| 456 | outputs.append(inp) |
| 457 | return outputs |
| 458 | |
| 459 | def update(self, preds, truths): |
| 460 | preds, truths = self.prepare_inputs(preds, truths) # [B, H, W, 3] --> [B, 3, H, W], range in [0, 1] |
| 461 | v = self.fn(truths, preds, normalize=True).item() # normalize=True: [0, 1] to [-1, 1] |
| 462 | self.V += v |
| 463 | self.N += 1 |
| 464 | |
| 465 | def measure(self): |
| 466 | return self.V / self.N |
| 467 | |
| 468 | def write(self, writer, global_step, prefix=""): |
| 469 | writer.add_scalar(os.path.join(prefix, f"LPIPS ({self.net})"), self.measure(), global_step) |
| 470 | |
| 471 | def report(self): |
| 472 | return f'LPIPS ({self.net}) = {self.measure():.6f}' |
| 473 | |
| 474 | |
| 475 | class LMDMeter: |