AiDA_Python/app/service/design_fast/pipeline/keypoint.py

import logging

import numpy as np
# from pymilvus import MilvusClient

from app.core.config import KEYPOINT_RESULT_TABLE_FIELD_SET, MILVUS_TABLE_KEYPOINT, settings
from app.service.design_fast.utils.design_ensemble import get_keypoint_result
from app.service.utils.decorator import ClassCallRunTime, RunTime

logger = logging.getLogger(__name__)


class KeyPoint:
    name = "KeyPoint"

    @classmethod
    def get_name(cls):
        return cls.name

    @ClassCallRunTime
    def __call__(self, result):
        if result['name'] in ['blouse', 'skirt', 'dress', 'outwear', 'trousers', 'tops', 'bottoms']:  # 查询是否有数据 且类别相同 相同则直接读 不同则推理后更新
            # result['clothes_keypoint'] = self.infer_keypoint_result(result)
            # 'up' if result['name'] in ['blouse', 'outwear', 'dress', 'tops'] else 'down'
            # keypoint_cache = search_keypoint_cache(result["image_id"], site)
            # keypoint_cache = self.keypoint_cache(result, site)
            keypoint_cache = False
            # 取消向量查询 直接过模型推理
            if not keypoint_cache:
                keypoint_infer_result, site = self.infer_keypoint_result(result)
                result['clothes_keypoint'] = self.save_keypoint_cache(result["image_id"], keypoint_infer_result, site)
            else:
                result['clothes_keypoint'] = keypoint_cache
        return result

    @staticmethod
    def infer_keypoint_result(result):
        site = 'up' if result['name'] in ['blouse', 'outwear', 'dress', 'tops'] else 'down'
        keypoint_infer_result = get_keypoint_result(result["image"], site)  # 推理结果
        return keypoint_infer_result, site

    @staticmethod
    def save_keypoint_cache(keypoint_id, cache, site):
        if site == "down":
            zeros = np.zeros(20, dtype=int)
            result = np.concatenate([zeros, cache.flatten()])
        else:
            zeros = np.zeros(4, dtype=int)
            result = np.concatenate([cache.flatten(), zeros])
        # 取消向量保存 直接拿结果
        data = [
            {"keypoint_id": keypoint_id,
             "keypoint_site": site,
             "keypoint_vector": result.tolist()
             }
        ]
        return dict(zip(KEYPOINT_RESULT_TABLE_FIELD_SET, result.reshape(12, 2).astype(int).tolist()))

        # try:
        #     client = MilvusClient(uri=settings.MILVUS_URL, token=settings.MILVUS_TOKEN, db_name=settings.MILVUS_ALIAS)
        #     client.upsert(collection_name=MILVUS_TABLE_KEYPOINT, data=data)
        #     client.close()
        # except Exception as e:
        #     logger.info(f"save keypoint cache milvus error : {e}")
        #     return dict(zip(KEYPOINT_RESULT_TABLE_FIELD_SET, result.reshape(12, 2).astype(int).tolist()))

    # @staticmethod
    # def update_keypoint_cache(keypoint_id, infer_result, search_result, site):
    #     if site == "up":
    #         # 需要的是up 即推理出来的是up 那么查询的就是down
    #         result = np.concatenate([infer_result.flatten(), search_result[-4:]])
    #     else:
    #         # 需要的是down 即推理出来的是down 那么查询的就是up
    #         result = np.concatenate([search_result[:20], infer_result.flatten()])
    #     data = [
    #         {"keypoint_id": keypoint_id,
    #          "keypoint_site": "all",
    #          "keypoint_vector": result.tolist()
    #          }
    #     ]
    #
    #     try:
    #         client = MilvusClient(uri=settings.MILVUS_URL, token=settings.MILVUS_TOKEN, db_name=settings.MILVUS_ALIAS)
    #         client.upsert(
    #             collection_name=MILVUS_TABLE_KEYPOINT,
    #             data=data
    #         )
    #         return dict(zip(KEYPOINT_RESULT_TABLE_FIELD_SET, result.reshape(12, 2).astype(int).tolist()))
    #     except Exception as e:
    #         logger.info(f"save keypoint cache milvus error : {e}")
    #         return dict(zip(KEYPOINT_RESULT_TABLE_FIELD_SET, result.reshape(12, 2).astype(int).tolist()))

    # @RunTime
    # def keypoint_cache(self, result, site):
    #     try:
    #         client = MilvusClient(uri=settings.MILVUS_URL, token=settings.MILVUS_TOKEN, db_name=settings.MILVUS_ALIAS)
    #         keypoint_id = result['image_id']
    #         res = client.query(
    #             collection_name=MILVUS_TABLE_KEYPOINT,
    #             # ids=[keypoint_id],
    #             filter=f"keypoint_id == {keypoint_id}",
    #             output_fields=['keypoint_vector', 'keypoint_site']
    #         )
    #         if len(res) == 0:
    #             # 没有结果 直接推理拿结果 并保存
    #             keypoint_infer_result, site = self.infer_keypoint_result(result)
    #             return self.save_keypoint_cache(result['image_id'], keypoint_infer_result, site)
    #         elif res[0]["keypoint_site"] == "all" or res[0]["keypoint_site"] == site:
    #             # 需要的类型和查询的类型一致，或者查询的类型为all 则直接返回查询的结果
    #             return dict(zip(KEYPOINT_RESULT_TABLE_FIELD_SET, np.array(res[0]['keypoint_vector']).astype(int).reshape(12, 2).tolist()))
    #         elif res[0]["keypoint_site"] != site:
    #             # 需要的类型和查询到的不一致，则更新类型为all
    #             keypoint_infer_result, site = self.infer_keypoint_result(result)
    #             return self.update_keypoint_cache(result["image_id"], keypoint_infer_result, res[0]['keypoint_vector'], site)
    #     except Exception as e:
    #         logger.info(f"search keypoint cache milvus error {e}")
    #         return False