evalution.py 816 Bytes
Newer Older
Rayyyyy's avatar
Rayyyyy committed
1
2
3
4
5
6
7
8
9
10
11
12
import os
import cv2
import warnings
import argparse

from basicsr.metrics import calculate_niqe

parser = argparse.ArgumentParser()
parser.add_argument('--root_path', type=str, default='results')
args = parser.parse_args()

def main():
13
14
    res = []
    for idx, imgname in enumerate(os.listdir(args.root_path)):
Rayyyyy's avatar
Rayyyyy committed
15
        img_path = os.path.join(args.root_path, imgname)
16
        img = cv2.imread(img_path)         
Rayyyyy's avatar
Rayyyyy committed
17
18
19
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', category=RuntimeWarning)
            niqe_result = calculate_niqe(img, 0, input_order='HWC', convert_to='y')
20
21
22
23
            res.append(niqe_result)
        print('{} img: {} niqe: {}'.format(idx, img_path, niqe_result))
    avg = float(sum(res)/len(res))
    print('avg niqe', avg)
Rayyyyy's avatar
Rayyyyy committed
24
25
26

if __name__ == '__main__':
    main()