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():
    res = []
    for idx, imgname in enumerate(os.listdir(args.root_path)):
        img_path = os.path.join(args.root_path, imgname)
        img = cv2.imread(img_path)         
        with warnings.catch_warnings():
            warnings.simplefilter('ignore', category=RuntimeWarning)
            niqe_result = calculate_niqe(img, 0, input_order='HWC', convert_to='y')
            res.append(niqe_result)
        print('{} img: {} niqe: {}'.format(idx, img_path, niqe_result))
    avg = float(sum(res)/len(res))
    print('avg niqe', avg)

if __name__ == '__main__':
    main()