waifuc/action/ccip.py

import logging
from enum import IntEnum
from typing import Iterator, Optional, List, Tuple

import numpy as np
from hbutils.string import plural_word
from hbutils.testing import disable_output
from imgutils.metrics import ccip_extract_feature, ccip_default_threshold, ccip_clustering, ccip_batch_differences

from .base import BaseAction
from ..model import ImageItem


class CCIPStatus(IntEnum):
    INIT = 0x1
    APPROACH = 0x2
    EVAL = 0x3
    INIT_WITH_SOURCE = 0x4


class CCIPAction(BaseAction):
    def __init__(self, init_source=None, *, min_val_count: int = 15, step: int = 5,
                 ratio_threshold: float = 0.6, min_clu_dump_ratio: float = 0.3, cmp_threshold: float = 0.5,
                 eps: Optional[float] = None, min_samples: Optional[int] = None,
                 model='ccip-caformer-24-randaug-pruned', threshold: Optional[float] = None):
        self.init_source = init_source

        self.min_val_count = min_val_count
        self.step = step
        self.ratio_threshold = ratio_threshold
        self.min_clu_dump_ratio = min_clu_dump_ratio
        self.cmp_threshold = cmp_threshold
        self.eps, self.min_samples = eps, min_samples
        self.model = model
        self.threshold = threshold or ccip_default_threshold(self.model)

        self.items = []
        self.item_released = []
        self.feats = []
        if self.init_source is not None:
            self.status = CCIPStatus.INIT_WITH_SOURCE
        else:
            self.status = CCIPStatus.INIT

    def _extract_feature(self, item: ImageItem):
        if 'ccip_feature' in item.meta:
            return item.meta['ccip_feature']
        else:
            return ccip_extract_feature(item.image, model=self.model)

    def _try_cluster(self) -> bool:
        with disable_output():
            clu_ids = ccip_clustering(self.feats, method='optics', model=self.model,
                                      eps=self.eps, min_samples=self.min_samples)
        clu_counts = {}
        for id_ in clu_ids:
            if id_ != -1:
                clu_counts[id_] = clu_counts.get(id_, 0) + 1

        clu_total = sum(clu_counts.values()) if clu_counts else 0
        chosen_id = None
        for id_, count in clu_counts.items():
            if count >= clu_total * self.ratio_threshold:
                chosen_id = id_
                break

        if chosen_id is not None:
            feats = [feat for i, feat in enumerate(self.feats) if clu_ids[i] == chosen_id]
            clu_dump_ratio = np.array([
                self._compare_to_exists(feat, base_set=feats)
                for feat in feats
            ]).astype(float).mean()

            if clu_dump_ratio >= self.min_clu_dump_ratio:
                self.items = [item for i, item in enumerate(self.items) if clu_ids[i] == chosen_id]
                self.item_released = [False] * len(self.items)
                self.feats = [feat for i, feat in enumerate(self.feats) if clu_ids[i] == chosen_id]
                return True
            else:
                return False
        else:
            return False

    def _compare_to_exists(self, feat, base_set=None) -> Tuple[bool, List[int]]:
        diffs = ccip_batch_differences([feat, *(base_set or self.feats)], model=self.model)[0, 1:]
        matches = diffs <= self.threshold
        return matches.astype(float).mean() >= self.cmp_threshold

    def _dump_items(self) -> Iterator[ImageItem]:
        for i in range(len(self.items)):
            if not self.item_released[i]:
                if self._compare_to_exists(self.feats[i]):
                    self.item_released[i] = True
                    yield self.items[i]

    def _eval_iter(self, item: ImageItem) -> Iterator[ImageItem]:
        feat = self._extract_feature(item)
        if self._compare_to_exists(feat):
            self.feats.append(feat)
            yield item

            if (len(self.feats) - len(self.items)) % self.step == 0:
                yield from self._dump_items()

    def iter(self, item: ImageItem) -> Iterator[ImageItem]:
        if self.status == CCIPStatus.INIT_WITH_SOURCE:
            cnt = 0
            logging.info('Existing anchor detected.')
            for item_ in self.init_source:
                self.feats.append(self._extract_feature(item_))
                yield item_
                cnt += 1
            logging.info(f'{plural_word(cnt, "items")} loaded from anchor.')

            self.status = CCIPStatus.EVAL
            yield from self._eval_iter(item)

        elif self.status == CCIPStatus.INIT:
            self.items.append(item)
            self.feats.append(self._extract_feature(item))

            if len(self.items) >= self.min_val_count:
                if self._try_cluster():
                    self.status = CCIPStatus.EVAL
                    yield from self._dump_items()
                else:
                    self.status = CCIPStatus.APPROACH

        elif self.status == CCIPStatus.APPROACH:
            self.items.append(item)
            self.feats.append(self._extract_feature(item))

            if (len(self.items) - self.min_val_count) % self.step == 0:
                if self._try_cluster():
                    self.status = CCIPStatus.EVAL
                    yield from self._dump_items()

        elif self.status == CCIPStatus.EVAL:
            yield from self._eval_iter(item)

        else:
            raise ValueError(f'Unknown status for {self.__class__.__name__} - {self.status!r}.')

    def reset(self):
        self.items.clear()
        self.item_released.clear()
        self.feats.clear()
        if self.init_source:
            self.status = CCIPStatus.INIT_WITH_SOURCE
        else:
            self.status = CCIPStatus.INIT