lama

1
import logging
2

3
import torch
4

5
from saicinpainting.evaluation.evaluator import InpaintingEvaluatorOnline, ssim_fid100_f1, lpips_fid100_f1
6
from saicinpainting.evaluation.losses.base_loss import SSIMScore, LPIPSScore, FIDScore
7

8

9
def make_evaluator(kind='default', ssim=True, lpips=True, fid=True, integral_kind=None, **kwargs):
10
    logging.info(f'Make evaluator {kind}')
11
    device = "cuda" if torch.cuda.is_available() else "cpu"
12
    metrics = {}
13
    if ssim:
14
        metrics['ssim'] = SSIMScore()
15
    if lpips:
16
        metrics['lpips'] = LPIPSScore()
17
    if fid:
18
        metrics['fid'] = FIDScore().to(device)
19
        
20
    if integral_kind is None:
21
        integral_func = None
22
    elif integral_kind == 'ssim_fid100_f1':
23
        integral_func = ssim_fid100_f1
24
    elif integral_kind == 'lpips_fid100_f1':
25
        integral_func = lpips_fid100_f1
26
    else:
27
        raise ValueError(f'Unexpected integral_kind={integral_kind}')
28

29
    if kind == 'default':
30
        return InpaintingEvaluatorOnline(scores=metrics,
31
                                         integral_func=integral_func,
32
                                         integral_title=integral_kind,
33
                                         **kwargs)
34