import spaces
import pyiqa
import torch

class IQA:
    def __init__(self, model_name="nima"):
        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
        self.model = pyiqa.create_metric(model_name, device=device)
        print(self.model)
    def __call__(self, image_path):
        return self.model(image_path)

if __name__ == "__main__":
    import requests
    from PIL import Image
    import glob
    image_files = glob.glob("samples/*")
    iqa_metric = IQA(model_name="nima-vgg16-ava")
    for image_file in image_files:
        print(image_file)
        image = Image.open(image_file)
        score = iqa_metric(image)
        print(score)