PaddleX/paddlex/inference/pipelines/components/faisser.py

# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pickle
from pathlib import Path

import numpy as np

from ....utils import logging
from ....utils.deps import class_requires_deps, is_dep_available
from ...utils.io import YAMLReader, YAMLWriter

if is_dep_available("faiss-cpu"):
    import faiss


@class_requires_deps("faiss-cpu")
class IndexData:
    VECTOR_FN = "vector"
    VECTOR_SUFFIX = ".index"
    IDMAP_FN = "id_map"
    IDMAP_SUFFIX = ".yaml"

    def __init__(self, index, index_info):
        self._index = index
        self._index_info = index_info
        self._id_map = index_info["id_map"]
        self._metric_type = index_info["metric_type"]
        self._index_type = index_info["index_type"]

    @property
    def index(self):
        return self._index

    @property
    def index_bytes(self):
        return faiss.serialize_index(self._index)

    @property
    def id_map(self):
        return self._id_map

    @property
    def metric_type(self):
        return self._metric_type

    @property
    def index_type(self):
        return self._index_type

    @property
    def index_info(self):
        return {
            "index_type": self.index_type,
            "metric_type": self.metric_type,
            "id_map": self._convert_int(self.id_map),
        }

    @classmethod
    def from_bytes(cls, bytes):
        tup = pickle.loads(bytes)
        index = faiss.deserialize_index(tup[0])
        return cls(index, tup[1])

    def to_bytes(self):
        tup = (faiss.serialize_index(self._index), self.index_info)
        return pickle.dumps(tup)

    def _convert_int(self, id_map):
        return {int(k): str(v) for k, v in id_map.items()}

    @staticmethod
    def _convert_int64(id_map):
        return {np.int64(k): str(v) for k, v in id_map.items()}

    def save(self, save_dir):
        save_dir = Path(save_dir)
        save_dir.mkdir(parents=True, exist_ok=True)
        vector_path = (save_dir / f"{self.VECTOR_FN}{self.VECTOR_SUFFIX}").as_posix()
        index_info_path = (save_dir / f"{self.IDMAP_FN}{self.IDMAP_SUFFIX}").as_posix()

        if self.metric_type in FaissBuilder.BINARY_METRIC_TYPE:
            faiss.write_index_binary(self.index, vector_path)
        else:
            faiss.write_index(self.index, vector_path)

        yaml_writer = YAMLWriter()
        yaml_writer.write(
            index_info_path,
            self.index_info,
            default_flow_style=False,
            allow_unicode=True,
        )

    @classmethod
    def load(cls, index):
        if isinstance(index, str):
            index_root = Path(index)
            vector_path = index_root / f"{cls.VECTOR_FN}{cls.VECTOR_SUFFIX}"
            index_info_path = index_root / f"{cls.IDMAP_FN}{cls.IDMAP_SUFFIX}"

            assert (
                vector_path.exists()
            ), f"Not found the {cls.VECTOR_FN}{cls.VECTOR_SUFFIX} file in {index}!"
            assert (
                index_info_path.exists()
            ), f"Not found the {cls.IDMAP_FN}{cls.IDMAP_SUFFIX} file in {index}!"

            yaml_reader = YAMLReader()
            index_info = yaml_reader.read(index_info_path)
            assert (
                "id_map" in index_info
                and "metric_type" in index_info
                and "index_type" in index_info
            ), f"The index_info file({index_info_path}) may have been damaged, `id_map` or `metric_type` or `index_type` not found in `index_info`."
            id_map = IndexData._convert_int64(index_info["id_map"])

            if index_info["metric_type"] in FaissBuilder.BINARY_METRIC_TYPE:
                index = faiss.read_index_binary(vector_path.as_posix())
            else:
                index = faiss.read_index(vector_path.as_posix())
            assert index.ntotal == len(
                id_map
            ), "data number in index is not equal in in id_map"

            return index, id_map, index_info["metric_type"], index_info["index_type"]
        else:
            assert isinstance(index, IndexData)
            return index.index, index.id_map, index.metric_type, index.index_type


class FaissIndexer:

    def __init__(
        self,
        index,
    ):
        super().__init__()
        self._indexer, self.id_map, self.metric_type, index_type = IndexData.load(index)

    def __call__(self, feature, score_thres, hamming_radius, topk):
        scores_list, ids_list = self._indexer.search(np.array(feature), topk)
        preds = []
        for scores, ids in zip(scores_list, ids_list):
            preds.append({"score": [], "label": []})
            for score, id in zip(scores, ids):
                if id >= 0:
                    preds[-1]["score"].append(score)
                    preds[-1]["label"].append(self.id_map[id])

        if self.metric_type in FaissBuilder.BINARY_METRIC_TYPE:
            idxs = np.where(scores_list[:, 0] > hamming_radius)[0]
        else:
            idxs = np.where(scores_list[:, 0] < score_thres)[0]
        for idx in idxs:
            preds[idx] = {"score": None, "label": None}
        return preds


@class_requires_deps("faiss-cpu")
class FaissBuilder:

    SUPPORT_METRIC_TYPE = ("hamming", "IP", "L2")
    SUPPORT_INDEX_TYPE = ("Flat", "IVF", "HNSW32")
    BINARY_METRIC_TYPE = ("hamming",)
    BINARY_SUPPORT_INDEX_TYPE = ("Flat", "IVF", "BinaryHash")

    @classmethod
    def _get_index_type(cls, metric_type, index_type, num=None):
        # if IVF method, cal ivf number automatically
        if index_type == "IVF":
            index_type = index_type + str(min(int(num // 8), 65536))
            if metric_type in cls.BINARY_METRIC_TYPE:
                index_type += ",BFlat"
            else:
                index_type += ",Flat"

        # for binary index, add B at head of index_type
        if metric_type in cls.BINARY_METRIC_TYPE:
            assert (
                index_type in cls.BINARY_SUPPORT_INDEX_TYPE
            ), f"The metric type({metric_type}) only support {cls.BINARY_SUPPORT_INDEX_TYPE} index types!"
            index_type = "B" + index_type

        if index_type == "HNSW32":
            logging.warning("The HNSW32 method dose not support 'remove' operation")
            index_type = "HNSW32"

        if index_type == "Flat":
            index_type = "Flat"

        return index_type

    @classmethod
    def _get_metric_type(cls, metric_type):
        if metric_type == "hamming":
            return faiss.METRIC_Hamming
        elif metric_type == "jaccard":
            return faiss.METRIC_Jaccard
        elif metric_type == "IP":
            return faiss.METRIC_INNER_PRODUCT
        elif metric_type == "L2":
            return faiss.METRIC_L2

    @classmethod
    def build(
        cls,
        gallery_imgs,
        gallery_label,
        predict_func,
        metric_type="IP",
        index_type="HNSW32",
    ):
        assert (
            index_type in cls.SUPPORT_INDEX_TYPE
        ), f"Supported index types only: {cls.SUPPORT_INDEX_TYPE}!"

        assert (
            metric_type in cls.SUPPORT_METRIC_TYPE
        ), f"Supported metric types only: {cls.SUPPORT_METRIC_TYPE}!"

        if isinstance(gallery_label, str):
            gallery_docs, gallery_list = cls.load_gallery(gallery_label, gallery_imgs)
        else:
            gallery_docs, gallery_list = gallery_label, gallery_imgs

        features = [res["feature"] for res in predict_func(gallery_list)]
        dtype = np.uint8 if metric_type in cls.BINARY_METRIC_TYPE else np.float32
        features = np.array(features).astype(dtype)
        vector_num, vector_dim = features.shape

        if metric_type in cls.BINARY_METRIC_TYPE:
            index = faiss.index_binary_factory(
                vector_dim,
                cls._get_index_type(metric_type, index_type, vector_num),
                cls._get_metric_type(metric_type),
            )
        else:
            index = faiss.index_factory(
                vector_dim,
                cls._get_index_type(metric_type, index_type, vector_num),
                cls._get_metric_type(metric_type),
            )
            index = faiss.IndexIDMap2(index)
        ids = {}

        # calculate id for new data
        index, ids = cls._add_gallery(
            metric_type, index, ids, features, gallery_docs, mode="new"
        )
        return IndexData(
            index, {"id_map": ids, "metric_type": metric_type, "index_type": index_type}
        )

    @classmethod
    def remove(
        cls,
        remove_ids,
        index,
    ):
        index, ids, metric_type, index_type = IndexData.load(index)
        if index_type == "HNSW32":
            raise RuntimeError(
                "The index_type: HNSW32 dose not support 'remove' operation"
            )
        if isinstance(remove_ids, str):
            lines = []
            with open(remove_ids) as f:
                lines = f.readlines()
            remove_ids = []
            for line in lines:
                id_ = int(line.strip().split(" ")[0])
                remove_ids.append(id_)
            remove_ids = np.asarray(remove_ids)
        else:
            remove_ids = np.asarray(remove_ids)

        # remove ids in id_map, remove index data in faiss index
        index.remove_ids(remove_ids)
        ids = {k: v for k, v in ids.items() if k not in remove_ids}
        return IndexData(
            index, {"id_map": ids, "metric_type": metric_type, "index_type": index_type}
        )

    @classmethod
    def append(cls, gallery_imgs, gallery_label, predict_func, index):
        index, ids, metric_type, index_type = IndexData.load(index)
        assert (
            metric_type in cls.SUPPORT_METRIC_TYPE
        ), f"Supported metric types only: {cls.SUPPORT_METRIC_TYPE}!"

        if isinstance(gallery_label, str):
            gallery_docs, gallery_list = cls.load_gallery(gallery_label, gallery_imgs)
        else:
            gallery_docs, gallery_list = gallery_label, gallery_imgs

        features = [res["feature"] for res in predict_func(gallery_list)]
        dtype = np.uint8 if metric_type in cls.BINARY_METRIC_TYPE else np.float32
        features = np.array(features).astype(dtype)

        # calculate id for new data
        index, ids = cls._add_gallery(
            metric_type, index, ids, features, gallery_docs, mode="append"
        )
        return IndexData(
            index, {"id_map": ids, "metric_type": metric_type, "index_type": index_type}
        )

    @classmethod
    def _add_gallery(
        cls, metric_type, index, ids, gallery_features, gallery_docs, mode
    ):
        start_id = max(ids.keys()) + 1 if ids else 0
        ids_now = (np.arange(0, len(gallery_docs)) + start_id).astype(np.int64)

        # only train when new index file
        if mode == "new":
            if metric_type in cls.BINARY_METRIC_TYPE:
                index.add(gallery_features)
            else:
                index.train(gallery_features)

        if metric_type not in cls.BINARY_METRIC_TYPE:
            index.add_with_ids(gallery_features, ids_now)
        # TODO(gaotingquan): how append when using hamming metric type
        # else:
        #   pass

        for i, d in zip(list(ids_now), gallery_docs):
            ids[i] = d
        return index, ids

    @classmethod
    def load_gallery(cls, gallery_label_path, gallery_imgs_root="", delimiter=" "):
        lines = []
        files = []
        labels = []
        root = Path(gallery_imgs_root)
        with open(gallery_label_path, "r", encoding="utf-8") as f:
            lines = f.readlines()
        for line in lines:
            path, label = line.strip().split(delimiter)
            file_path = root / path
            files.append(file_path.as_posix())
            labels.append(label)
        return labels, files