diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/__init__.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/__init__.py index 7b48775..ca2b454 100644 --- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/__init__.py +++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/__init__.py @@ -9,7 +9,7 @@ require("nonebot_plugin_uninfo") from . import handler as handler from .config import Config -__version__ = "2.3.2" +__version__ = "2.3.2.post1" __plugin_meta__ = PluginMetadata( name="自动撤回奶龙", description="一个基于图像分类模型的简单插件~", diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/config.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/config.py index 785d39e..e97cc92 100644 --- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/config.py +++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/config.py @@ -55,14 +55,14 @@ class Config(BaseModel): nailong_model: ModelType = ModelType.TARGET_DETECTION nailong_auto_update_model: bool = True nailong_concurrency: int = 1 - nailong_onnx_try_to_use_gpu: bool = True + nailong_onnx_providers: List[str] = ["CPUExecutionProvider"] nailong_model1_type: Model1Type = Model1Type.TINY nailong_model1_yolox_size: Optional[Tuple[int, int]] = None nailong_model1_score: Dict[str, Optional[float]] = { DEFAULT_LABEL: 0.5, } - nailong_model2_online: bool = False + nailong_model2_online: bool=False nailong_check_mode: int = 0 nailong_similarity_on: bool = False nailong_similarity_max_storage: int = 10 @@ -77,9 +77,7 @@ class Config(BaseModel): mode="before", ) def transform_to_dict(cls, v: Any): # noqa: N805 - if not isinstance(v, dict): - return {DEFAULT_LABEL: v} - return v + return v if isinstance(v, dict) else {DEFAULT_LABEL: v} @field_validator( "nailong_tip", @@ -92,5 +90,20 @@ class Config(BaseModel): raise ValueError(f"Please ensure default label {DEFAULT_LABEL} in dict") return v + @field_validator("nailong_onnx_providers", mode="before") + def transform_to_list(cls, v: Any): # noqa: N805 + return v if isinstance(v, list) else [v] + + @field_validator("nailong_onnx_providers", mode="after") + def validate_provider_available(cls, v: Any): # noqa: N805 + try: + from onnxruntime.capi import _pybind_state as c + except ImportError: + pass + else: + available_providers: List[str] = c.get_available_providers() # type: ignore + if any(p not in available_providers for p in v): + raise ValueError(f"Provider {v} not available in onnxruntime") + return v config = get_plugin_config(Config) diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/handler.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/handler.py index 66357f6..109dd7d 100644 --- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/handler.py +++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/handler.py @@ -13,6 +13,7 @@ from .model import check from .uniapi import mute, recall from .model.utils.common import process_gif_and_save_jpgs + T = TypeVar("T") @@ -83,7 +84,6 @@ async def nailong_rule( nailong = on_message(rule=Rule(nailong_rule), priority=config.nailong_priority) input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yolox_size - @nailong.handle() async def handle_function(bot: BaseBot, ev: BaseEvent, msg: UniMsg, session: Uninfo): save_img = False diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/__init__.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/__init__.py index 6ef8585..fd08fee 100644 --- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/__init__.py +++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/__init__.py @@ -22,11 +22,20 @@ if config.nailong_model is ModelType.CLASSIFICATION: raise_extra_import_error(e, "model0") elif config.nailong_model is ModelType.TARGET_DETECTION: - pass + try: + from .target_detection import check as check + except ImportError as e: + raise ImportError( + "To avoid dependency issues, please install onnxruntime manually.\n" + "If you have a compatible GPU, " + "please run `pip install onnxruntime-gpu` in your project's environment, " + "then edit plugin's `NAILONG_ONNX_PROVIDERS` config to use it;\n" + "Otherwise run `pip install onnxruntime` in your project's environment " + "and use CPU to compute.", + ) from e elif config.nailong_model is ModelType.HF_DETECTION: from .hf_detection import check as check - else: - raise ValueError("Invalid model type") + raise NotImplementedError # never reach here diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/hf_detection.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/hf_detection.py index 33cce2c..edcbe8d 100644 --- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/hf_detection.py +++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/hf_detection.py @@ -19,7 +19,6 @@ if config.nailong_model2_online: import base64 import io import shutil - FILENAME = "nailong_yolo11.pt" client = Client("Hakureirm/NailongKiller") logger.info(f"Using model {FILENAME} online") @@ -30,13 +29,11 @@ else: REPO_ID = "Hakureirm/NailongKiller" FILENAME = "nailong_yolo11.pt" - model_path = os.path.join(str(config.nailong_model_dir), FILENAME) - if config.nailong_auto_update_model or not os.path.exists(model_path): + model_path=os.path.join(str(config.nailong_model_dir),FILENAME) + if config.nailong_auto_update_model or not os.path.exists(model_path): api = hf_api.HfApi() file_path = os.path.join(str(config.nailong_model_dir), FILENAME) model_info = api.model_info(REPO_ID) - - def get_file_last_modified_time(file_path): try: timestamp = os.path.getmtime(file_path) @@ -44,47 +41,42 @@ else: return last_modified_time except FileNotFoundError: return None - - local_time = get_file_last_modified_time(file_path) if local_time is None or model_info.last_modified >= local_time: hf_hub_download(repo_id=REPO_ID, filename=FILENAME, local_dir=config.nailong_model_dir) logger.info(f"Update model {FILENAME} successfully!") - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device=torch.device("cuda" if torch.cuda.is_available() else "cpu") model = YOLO(model_path).to(device) logger.info(f"Using model {FILENAME}") input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yolox_size - @run_sync def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult: if is_gif: res = similarity_process(frame, dsize=input_shape) if res is not None: - return CheckSingleResult(ok=res.ok, label=res.label, extra=frame) - return CheckSingleResult(ok=False, label=None, extra=frame) + return CheckSingleResult(ok=res.ok,label=res.label,extra=frame) + return CheckSingleResult(ok=False,label=None,extra=frame) else: if config.nailong_model2_online: input_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) - if not os.path.exists(os.path.join(str(config.nailong_model_dir), "online_temp")): - os.makedirs(os.path.join(str(config.nailong_model_dir), "online_temp")) - image_path = os.path.join(str(config.nailong_model_dir), "online_temp", - "temp_{}.jpg".format(datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))) + if not os.path.exists(os.path.join(str(config.nailong_model_dir),"online_temp")): + os.makedirs(os.path.join(str(config.nailong_model_dir),"online_temp")) + image_path=os.path.join(str(config.nailong_model_dir),"online_temp","temp_{}.jpg".format(datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S"))) while os.path.exists(image_path): - basename = os.path.basename(image_path) - image_path = os.path.join(str(config.nailong_model_dir), "online_temp", f"exist-{basename}") - input_image.save(image_path, format='JPEG') + basename=os.path.basename(image_path) + image_path=os.path.join(str(config.nailong_model_dir),"online_temp",f"exist-{basename}") + input_image.save(image_path,format='JPEG') result_image, result_info = client.predict( img=handle_file(image_path), api_name="/predict" ) os.remove(image_path) - if "检测到的目标数量: " in result_info and int( - result_info.split("检测到的目标数量: ")[1].split("\n")[0]) < 1: - return CheckSingleResult(ok=False, label=None, extra=frame) + if "检测到的目标数量: " in result_info and int(result_info.split("检测到的目标数量: ")[1].split("\n")[0])<1: + return CheckSingleResult(ok=False,label=None,extra=frame) if isinstance(result_image, str): if result_image.startswith('data:image'): img_data = base64.b64decode(result_image.split(',')[1]) @@ -95,8 +87,8 @@ def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult: img = Image.open(img_data) result_image = np.array(img) shutil.rmtree(os.path.dirname(os.path.dirname(img_data))) - result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB) - return CheckSingleResult(ok=True, label="nailong", extra=result_image) + result_image=cv2.cvtColor(result_image,cv2.COLOR_BGR2RGB) + return CheckSingleResult(ok=True,label="nailong",extra=result_image) else: input_image = Image.fromarray(frame) original_size = input_image.size @@ -108,8 +100,10 @@ def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult: padded_img = Image.new('RGB', (max_size, max_size), (114, 114, 114)) padded_img.paste(input_image, (pad_w // 2, pad_h // 2)) + img_array = np.array(padded_img) + results = model.predict( img_array, conf=config.nailong_model1_score['nailong'], @@ -119,7 +113,7 @@ def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult: ) cls = results[0].boxes.cls if len(cls) < 1: - return CheckSingleResult(ok=False, label=None, extra=frame) + return CheckSingleResult(ok=False,label=None,extra=frame) result_img = results[0].plot() if pad_w > 0 or pad_h > 0: diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/target_detection.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/target_detection.py index b95e93e..0754e23 100644 --- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/target_detection.py +++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/target_detection.py @@ -4,17 +4,17 @@ from typing import Optional from typing_extensions import override import numpy as np -import onnxruntime +# import torch before onnxruntime +import torch as torch # isort: skip +import onnxruntime # isort: skip from cookit import with_semaphore from nonebot.utils import run_sync -from plugins.nonebot_plugin_nailongremove.config import config -from plugins.nonebot_plugin_nailongremove.frame_source import FrameSource, repack_save -from plugins.nonebot_plugin_nailongremove.model.utils.common import CheckResult, CheckSingleResult, race_check, \ - similarity_process -from plugins.nonebot_plugin_nailongremove.model.utils.update import GitHubLatestReleaseModelUpdater, ModelInfo, \ - UpdaterGroup -from plugins.nonebot_plugin_nailongremove.model.utils.yolox import demo_postprocess, multiclass_nms, preprocess, vis +from ..config import config +from ..frame_source import FrameSource, repack_save +from .utils.common import CheckResult, CheckSingleResult, race_check, similarity_process +from .utils.update import GitHubLatestReleaseModelUpdater, ModelInfo, UpdaterGroup +from .utils.yolox import demo_postprocess, multiclass_nms, preprocess, vis import itertools model_filename_sfx = f"_{config.nailong_model1_type.value}.onnx" @@ -47,15 +47,7 @@ labels = labels_path.read_text("u8").splitlines() session = onnxruntime.InferenceSession( model_path, - providers=( - [ - "TensorrtExecutionProvider", - "CUDAExecutionProvider", - "CPUExecutionProvider", - ] - if config.nailong_onnx_try_to_use_gpu - else ["CPUExecutionProvider"] - ), + providers=config.nailong_onnx_providers, ) input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yolox_size diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/common.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/common.py index d38f104..b5758fc 100644 --- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/common.py +++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/common.py @@ -115,7 +115,6 @@ async def race_check( return res return None - def similarity_process(image1: np.ndarray, dsize) -> Optional[CheckSingleResult]: path = list(glob.glob(os.path.join(config.nailong_model_dir, 'records/*/*.jpg'))) if len(path) == 0: @@ -187,7 +186,7 @@ def process_gif_and_save_jpgs(frames, label, dsize, similarity_threshold=0.85): indices = torch.nonzero(similarities > similarity_threshold) index = indices.squeeze().tolist() if indices.numel() > 0 else None if type(index) is int: - index = [index] + index=[index] if index is not None: indexs.extend([frame2_num[i] for i in index]) frame_count = [i for i in frame_count if i not in indexs] diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/update.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/update.py index d27b8f1..f221f89 100644 --- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/update.py +++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/update.py @@ -56,10 +56,10 @@ def create_parent_dir(path: Path, create: bool = True): def find_file( - path: Path, - checker: Union[Callable[[Path], bool], str, None] = None, - recursive: bool = False, - last_modified: bool = True, + path: Path, + checker: Union[Callable[[Path], bool], str, None] = None, + recursive: bool = False, + last_modified: bool = True, ) -> Optional[Path]: if isinstance(checker, str) and checker: if (p := path / checker).exists(): @@ -99,12 +99,10 @@ class ModelInfo(Generic[T]): class ModelUpdater(ABC): @abstractmethod - def find_from_local(self) -> Optional[Path]: - ... + def find_from_local(self) -> Optional[Path]: ... @abstractmethod - def get_info(self) -> ModelInfo: - ... + def get_info(self) -> ModelInfo: ... @property def root_dir(self) -> Path: @@ -121,8 +119,8 @@ class ModelUpdater(ABC): def check_local_ver(self, info: ModelInfo) -> Optional[str]: if ( - self.get_path(info.filename).exists() - and (ver_path := self.get_ver_path(info.filename)).exists() + self.get_path(info.filename).exists() + and (ver_path := self.get_ver_path(info.filename)).exists() ): return ver_path.read_text(encoding="u8").strip() return None @@ -164,10 +162,10 @@ class ModelUpdater(ABC): return def validate_with_unlink( - self, - path: Path, - info: ModelInfo, - clear_ver: bool = True, + self, + path: Path, + info: ModelInfo, + clear_ver: bool = True, ) -> Any: try: return self.validate(path, info) @@ -179,9 +177,9 @@ class ModelUpdater(ABC): def _get(self, force_update: bool = False) -> Path: if ( - (not force_update) - and (not config.nailong_auto_update_model) - and (local := self.find_from_local()) + (not force_update) + and (not config.nailong_auto_update_model) + and (local := self.find_from_local()) ): logger.info("Update skipped") return local @@ -301,10 +299,10 @@ class GitHubRepoModelUpdater(GitHubModelUpdater): class GitHubLatestReleaseModelUpdater(GitHubModelUpdater): def __init__( - self, - owner: str, - repo: str, - local_filename_checker: Optional[Callable[[str], bool]] = None, + self, + owner: str, + repo: str, + local_filename_checker: Optional[Callable[[str], bool]] = None, ) -> None: super().__init__() self.owner = owner diff --git a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/yolox.py b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/yolox.py index a19d5ca..2aa1df3 100644 --- a/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/yolox.py +++ b/packages/nonebot-plugin-nailongremove-base/nonebot_plugin_nailongremove/model/utils/yolox.py @@ -92,7 +92,7 @@ _COLORS = ( ], ) .astype(np.float32) - .reshape(-1, 3) + .reshape(-1, 3) ) # fmt: skip