__init__.py

"""
Implementation/adapter for "ViT-FIQA" based on the corresponding repository, i.e.:
<https://github.com/atzoriandrea/ViT-FIQA-Assessing-Face-Image-Quality-using-Vision-Transformers/tree/master>
"""

# - # Standard imports:
from pathlib import Path

# - # External imports:
# None

# - # Toolkit imports:
import fiqat

# - # Local imports:
from .QualityModel import QualityModel


class VitFiqaAdapter:

  def __init__(self, model_path: Path, color_channel: str = 'RGB'):
    self.model_type = model_path.stem.replace(' ', '')
    assert self.model_type in {'ViT-FIQA(T)', 'ViT-FIQA(C)'}, self.model_type
    if self.model_type == 'ViT-FIQA(T)':
      self.model_mode = 'token'
    else:
      assert self.model_type == 'ViT-FIQA(C)', self.model_type
      self.model_mode = 'crfiqa'
    # - #
    self.model_path = model_path
    self.model_epoch = None
    self.gpu_id = None  # 0
    self.backbone = 'vit_FC'
    self.color_channel = color_channel
    assert color_channel in {'BGR', 'RGB'}, color_channel
    self.batch_size = 16
    self.process_batch_size = self.batch_size * 4
    # - #
    self.quality_model = None

  def load_model(self):
    if self.quality_model is not None:
      return
    self.quality_model = QualityModel(
        model_prefix=self.model_path,
        model_epoch=self.model_epoch,
        gpu_id=self.gpu_id,
        backbone=self.backbone,
        model_mode=self.model_mode,
    )

  def _process(self, image_path_list: list):
    self.load_model()
    embedding_batch, quality_batch = self.quality_model.get_batch_feature(
        image_path_list=image_path_list,
        batch_size=self.batch_size,
        color=self.color_channel,
    )
    return embedding_batch, quality_batch

  def process(self, image_path_list: list, output_with_embedding: bool = False):
    for sub_list in fiqat.iterate_as_batches(image_path_list, self.process_batch_size):
      embedding_batch, quality_batch = self._process(sub_list)
      for embedding, quality in zip(embedding_batch, quality_batch):
        assert quality.shape == (1,), quality.shape
        quality_score = quality[0]
        if output_with_embedding:
          yield embedding, quality_score
        else:
          yield quality_score