CVPR/lama-example

1#!/usr/bin/env python32 3import os4 5import pandas as pd6 7from saicinpainting.evaluation.data import PrecomputedInpaintingResultsDataset8from saicinpainting.evaluation.evaluator import InpaintingEvaluator, lpips_fid100_f19from saicinpainting.evaluation.losses.base_loss import SegmentationAwareSSIM, \10    SegmentationClassStats, SSIMScore, LPIPSScore, FIDScore, SegmentationAwareLPIPS, SegmentationAwareFID11from saicinpainting.evaluation.utils import load_yaml12 13 14def main(args):15    config = load_yaml(args.config)16 17    dataset = PrecomputedInpaintingResultsDataset(args.datadir, args.predictdir, **config.dataset_kwargs)18 19    metrics = {20        'ssim': SSIMScore(),21        'lpips': LPIPSScore(),22        'fid': FIDScore()23    }24    enable_segm = config.get('segmentation', dict(enable=False)).get('enable', False)25    if enable_segm:26        weights_path = os.path.expandvars(config.segmentation.weights_path)27        metrics.update(dict(28            segm_stats=SegmentationClassStats(weights_path=weights_path),29            segm_ssim=SegmentationAwareSSIM(weights_path=weights_path),30            segm_lpips=SegmentationAwareLPIPS(weights_path=weights_path),31            segm_fid=SegmentationAwareFID(weights_path=weights_path)32        ))33    evaluator = InpaintingEvaluator(dataset, scores=metrics,34                                    integral_title='lpips_fid100_f1', integral_func=lpips_fid100_f1,35                                    **config.evaluator_kwargs)36 37    os.makedirs(os.path.dirname(args.outpath), exist_ok=True)38 39    results = evaluator.evaluate()40 41    results = pd.DataFrame(results).stack(1).unstack(0)42    results.dropna(axis=1, how='all', inplace=True)43    results.to_csv(args.outpath, sep='\t', float_format='%.4f')44 45    if enable_segm:46        only_short_results = results[[c for c in results.columns if not c[0].startswith('segm_')]].dropna(axis=1, how='all')47        only_short_results.to_csv(args.outpath + '_short', sep='\t', float_format='%.4f')48 49        print(only_short_results)50 51        segm_metrics_results = results[['segm_ssim', 'segm_lpips', 'segm_fid']].dropna(axis=1, how='all').transpose().unstack(0).reorder_levels([1, 0], axis=1)52        segm_metrics_results.drop(['mean', 'std'], axis=0, inplace=True)53 54        segm_stats_results = results['segm_stats'].dropna(axis=1, how='all').transpose()55        segm_stats_results.index = pd.MultiIndex.from_tuples(n.split('/') for n in segm_stats_results.index)56        segm_stats_results = segm_stats_results.unstack(0).reorder_levels([1, 0], axis=1)57        segm_stats_results.sort_index(axis=1, inplace=True)58        segm_stats_results.dropna(axis=0, how='all', inplace=True)59 60        segm_results = pd.concat([segm_metrics_results, segm_stats_results], axis=1, sort=True)61        segm_results.sort_values(('mask_freq', 'total'), ascending=False, inplace=True)62 63        segm_results.to_csv(args.outpath + '_segm', sep='\t', float_format='%.4f')64    else:65        print(results)66 67 68if __name__ == '__main__':69    import argparse70 71    aparser = argparse.ArgumentParser()72    aparser.add_argument('config', type=str, help='Path to evaluation config')73    aparser.add_argument('datadir', type=str,74                         help='Path to folder with images and masks (output of gen_mask_dataset.py)')75    aparser.add_argument('predictdir', type=str,76                         help='Path to folder with predicts (e.g. predict_hifill_baseline.py)')77    aparser.add_argument('outpath', type=str, help='Where to put results')78 79    main(aparser.parse_args())80