From 1f21034b4c2a6c1ab3e32676cac295d5b8ffc48b Mon Sep 17 00:00:00 2001 From: student_2333 Date: Wed, 6 Nov 2024 02:15:43 +0800 Subject: [PATCH] up --- README.md | 24 ++++++++++++------- nonebot_plugin_nailongremove/config.py | 21 ++++++++++++---- .../model/target_detection.py | 6 ++--- 3 files changed, 35 insertions(+), 16 deletions(-) diff --git a/README.md b/README.md index 00679ed..a7b868c 100644 --- a/README.md +++ b/README.md @@ -145,6 +145,7 @@ plugins = [ | `NAILONG_MUTE_SECONDS` | 否 | `0` | 设置禁言时间,默认为 0 即不禁言
单位:秒 | | `NAILONG_TIP` | 否 | `本群禁止发奶龙!` | 发送的提示,使用 [Alconna 的消息模板](https://nonebot.dev/docs/best-practice/alconna/uniseg#%E4%BD%BF%E7%94%A8%E6%B6%88%E6%81%AF%E6%A8%A1%E6%9D%BF),可用变量见下 | | `NAILONG_FAILED_TIP` | 否 | `{:Reply($message_id)}呜,不要发奶龙了嘛 🥺 👉👈` | 撤回失败或禁用撤回时发送的提示,同上 | +| `NAILONG_CHECK_ALL_FRAMES` | 否 | `False` | 使用模型 1 时是否检查图片中的所有帧,启用该项后消息模板中的 `$checked_result` 变量当原图为动图时会变为动图 | | **模型通用配置** | | | | | `NAILONG_MODEL_DIR` | 否 | `./data/nailongremove` | 模型的下载位置 | | `NAILONG_MODEL` | 否 | `0` | 选择需要加载的模型,可用模型见下 | @@ -152,8 +153,8 @@ plugins = [ | `NAILONG_CONCURRENCY` | 否 | `1` | 当图片为动图时,针对该图片并发识别图片帧的最大并发数 | | `NAILONG_ONNX_TRY_TO_USE_GPU` | 否 | `True` | 加载 onnx 模型时是否尝试使用 GPU,如果失败则会显示一串警告但是对插件并无影响,如果不想看见警告关闭此配置项即可 | | **模型 1 特定配置** | | | | -| `NAILONG_MODEL1_YOLOX_SIZE` | 否 | `[416,416]` | 针对模型 1,自定义模型输入可能会有尺寸更改 | | `NAILONG_MODEL1_TYPE` | 否 | `tiny` | 模型 1 使用的模型类型,可用 `tiny` / `m` | +| `NAILONG_MODEL1_YOLOX_SIZE` | 否 | `None` | 针对模型 1,自定义模型输入可能会有尺寸更改 | | `NAILONG_MODEL1_SCORE` | 否 | `0.5` | 模型 1 置信度阈值,范围 `0` ~ `1` | | **杂项配置** | | | | | `NAILONG_GITHUB_TOKEN` | 否 | `None` | GitHub Access Token,遇到模型下载或更新问题时可尝试填写 | @@ -165,14 +166,14 @@ plugins = [ ### 消息模板可用变量 -| 变量名 | 类型 | 说明 | -| ---------------- | ---------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------- | -| `$event` | [`Event`](https://nonebot.dev/docs/api/adapters/#Event) | 当前事件 | -| `$target` | [`Target`](https://nonebot.dev/docs/best-practice/alconna/uniseg#%E6%B6%88%E6%81%AF%E5%8F%91%E9%80%81%E5%AF%B9%E8%B1%A1) | 事件目标 | -| `$message_id` | `str` | 消息 ID | -| `$msg` | [`UniMessage`](https://nonebot.dev/docs/best-practice/alconna/uniseg#%E9%80%9A%E7%94%A8%E6%B6%88%E6%81%AF%E5%BA%8F%E5%88%97) | 当前消息 | -| `$ss` | [`Session`](https://github.com/RF-Tar-Railt/nonebot-plugin-uninfo?tab=readme-ov-file#session) | 当前会话 | -| `$checked_image` | `bytes` | 框选出对应目标后的图片,仅在模型配置为 `1` 时存在 | +| 变量名 | 类型 | 说明 | +| ----------------- | ---------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------- | +| `$event` | [`Event`](https://nonebot.dev/docs/api/adapters/#Event) | 当前事件 | +| `$target` | [`Target`](https://nonebot.dev/docs/best-practice/alconna/uniseg#%E6%B6%88%E6%81%AF%E5%8F%91%E9%80%81%E5%AF%B9%E8%B1%A1) | 事件目标 | +| `$message_id` | `str` | 消息 ID | +| `$msg` | [`UniMessage`](https://nonebot.dev/docs/best-practice/alconna/uniseg#%E9%80%9A%E7%94%A8%E6%B6%88%E6%81%AF%E5%BA%8F%E5%88%97) | 当前消息 | +| `$ss` | [`Session`](https://github.com/RF-Tar-Railt/nonebot-plugin-uninfo?tab=readme-ov-file#session) | 当前会话 | +| `$checked_result` | [`Image`](https://nonebot.dev/docs/best-practice/alconna/uniseg#%E9%80%9A%E7%94%A8%E6%B6%88%E6%81%AF%E6%AE%B5) | 框选出对应目标后的图片,仅在模型配置为 `1` 时存在 | ## 🎉 使用 @@ -188,6 +189,11 @@ plugins = [ ## 📝 更新日志 +### dev + +- 支持了检查 GIF 中的所有帧并将结果重新封成 GIF,默认禁用,同时弃用 `$checked_image` 变量,新增 `$checked_result` 变量 +- 现在模型 1 的输入大小可以根据模型类型自动配置了,但是如果配置项指定了那么会优先使用 + ### 2.2.1 - 优化模型自动更新(可能是反向优化) diff --git a/nonebot_plugin_nailongremove/config.py b/nonebot_plugin_nailongremove/config.py index 716513e..2aa7efb 100644 --- a/nonebot_plugin_nailongremove/config.py +++ b/nonebot_plugin_nailongremove/config.py @@ -1,7 +1,8 @@ -from enum import Enum +from enum import Enum, auto from pathlib import Path from typing import List, Optional, Tuple +from cookit import StrEnum from nonebot import get_plugin_config from pydantic import BaseModel, Field @@ -11,6 +12,18 @@ class ModelType(int, Enum): TARGET_DETECTION = 1 +class Model1Type(StrEnum): + TINY = auto() + M = auto() + + @property + def yolox_size(self) -> Tuple[int, int]: + return { + Model1Type.TINY: (416, 416), + Model1Type.M: (640, 640), + }[self] + + class Config(BaseModel): proxy: Optional[str] = None @@ -25,7 +38,7 @@ class Config(BaseModel): nailong_mute_seconds: int = 0 nailong_tip: str = "本群禁止发送奶龙!" nailong_failed_tip: str = "{:Reply($message_id)}呜,不要发奶龙了嘛 🥺 👉👈" - nailong_checked_result_all: bool = False + nailong_check_all_frames: bool = False nailong_model_dir: Path = Field( default_factory=lambda: Path.cwd() / "data" / "nailongremove", @@ -35,8 +48,8 @@ class Config(BaseModel): nailong_concurrency: int = 1 nailong_onnx_try_to_use_gpu: bool = True - nailong_model1_yolox_size: Tuple[int, int] = (416, 416) - nailong_model1_type: str = "tiny" + nailong_model1_type: Model1Type = Model1Type.TINY + nailong_model1_yolox_size: Optional[Tuple[int, int]] = None nailong_model1_score: float = 0.5 nailong_github_token: Optional[str] = None diff --git a/nonebot_plugin_nailongremove/model/target_detection.py b/nonebot_plugin_nailongremove/model/target_detection.py index 60ac2e3..8e51cdb 100644 --- a/nonebot_plugin_nailongremove/model/target_detection.py +++ b/nonebot_plugin_nailongremove/model/target_detection.py @@ -14,7 +14,7 @@ from .utils.common import CheckResult, CheckSingleResult, race_check from .utils.update import GitHubLatestReleaseModelUpdater, ModelInfo, UpdaterGroup from .utils.yolox import demo_postprocess, multiclass_nms, preprocess, vis -model_filename_sfx = f"_{config.nailong_model1_type}.onnx" +model_filename_sfx = f"_{config.nailong_model1_type.value}.onnx" class ModelUpdater(GitHubLatestReleaseModelUpdater): @@ -54,7 +54,7 @@ session = onnxruntime.InferenceSession( else ["CPUExecutionProvider"] ), ) -input_shape = config.nailong_model1_yolox_size +input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yolox_size @dataclass @@ -129,7 +129,7 @@ async def check_single(frame: np.ndarray) -> CheckSingleResult[FrameInfo]: async def check(source: FrameSource) -> CheckResult: extra_vars = {} - if config.nailong_checked_result_all: + if config.nailong_check_all_frames: sem = asyncio.Semaphore(config.nailong_concurrency) results = asyncio.gather( *(with_semaphore(sem)(check_single)(frame) for frame in source),