benchmark/metrics/arniqa.py

from typing import Dict

import numpy as np
import torch
from PIL import Image
from torchmetrics.image.arniqa import ARNIQA


class ARNIQAMetric:
    def __init__(self):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.metric = ARNIQA(
            regressor_dataset="koniq10k",
            reduction="mean",
            normalize=True,
            autocast=False
        )
        self.metric.to(self.device)
    @property
    def name(self) -> str:
        return "arniqa"
    
    def compute_score(self, image: Image.Image, prompt: str) -> Dict[str, float]:
        image_tensor = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0
        image_tensor = image_tensor.unsqueeze(0).to(self.device)
        score = self.metric(image_tensor)
        return {"arniqa": score.item()}