File size: 1,201 Bytes
0f90f73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
import logging
import torch
from saicinpainting.evaluation.evaluator import InpaintingEvaluatorOnline, ssim_fid100_f1, lpips_fid100_f1
from saicinpainting.evaluation.losses.base_loss import SSIMScore, LPIPSScore, FIDScore
def make_evaluator(kind='default', ssim=True, lpips=True, fid=True, integral_kind=None, **kwargs):
logging.info(f'Make evaluator {kind}')
device = "cuda" if torch.cuda.is_available() else "cpu"
metrics = {}
if ssim:
metrics['ssim'] = SSIMScore()
if lpips:
metrics['lpips'] = LPIPSScore()
if fid:
metrics['fid'] = FIDScore().to(device)
if integral_kind is None:
integral_func = None
elif integral_kind == 'ssim_fid100_f1':
integral_func = ssim_fid100_f1
elif integral_kind == 'lpips_fid100_f1':
integral_func = lpips_fid100_f1
else:
raise ValueError(f'Unexpected integral_kind={integral_kind}')
if kind == 'default':
return InpaintingEvaluatorOnline(scores=metrics,
integral_func=integral_func,
integral_title=integral_kind,
**kwargs)
|