mirror of
https://github.com/Refound-445/nonebot-plugin-nailongremove.git
synced 2025-11-04 21:22:43 +08:00
up
This commit is contained in:
parent
3b45bab8ff
commit
26389a66ee
@ -9,7 +9,7 @@ require("nonebot_plugin_uninfo")
|
|||||||
from . import handler as handler
|
from . import handler as handler
|
||||||
from .config import Config
|
from .config import Config
|
||||||
|
|
||||||
__version__ = "2.3.2"
|
__version__ = "2.3.2.post1"
|
||||||
__plugin_meta__ = PluginMetadata(
|
__plugin_meta__ = PluginMetadata(
|
||||||
name="自动撤回奶龙",
|
name="自动撤回奶龙",
|
||||||
description="一个基于图像分类模型的简单插件~",
|
description="一个基于图像分类模型的简单插件~",
|
||||||
|
|||||||
@ -55,7 +55,7 @@ class Config(BaseModel):
|
|||||||
nailong_model: ModelType = ModelType.TARGET_DETECTION
|
nailong_model: ModelType = ModelType.TARGET_DETECTION
|
||||||
nailong_auto_update_model: bool = True
|
nailong_auto_update_model: bool = True
|
||||||
nailong_concurrency: int = 1
|
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_type: Model1Type = Model1Type.TINY
|
||||||
nailong_model1_yolox_size: Optional[Tuple[int, int]] = None
|
nailong_model1_yolox_size: Optional[Tuple[int, int]] = None
|
||||||
@ -77,9 +77,7 @@ class Config(BaseModel):
|
|||||||
mode="before",
|
mode="before",
|
||||||
)
|
)
|
||||||
def transform_to_dict(cls, v: Any): # noqa: N805
|
def transform_to_dict(cls, v: Any): # noqa: N805
|
||||||
if not isinstance(v, dict):
|
return v if isinstance(v, dict) else {DEFAULT_LABEL: v}
|
||||||
return {DEFAULT_LABEL: v}
|
|
||||||
return v
|
|
||||||
|
|
||||||
@field_validator(
|
@field_validator(
|
||||||
"nailong_tip",
|
"nailong_tip",
|
||||||
@ -92,5 +90,20 @@ class Config(BaseModel):
|
|||||||
raise ValueError(f"Please ensure default label {DEFAULT_LABEL} in dict")
|
raise ValueError(f"Please ensure default label {DEFAULT_LABEL} in dict")
|
||||||
return v
|
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)
|
config = get_plugin_config(Config)
|
||||||
|
|||||||
@ -13,6 +13,7 @@ from .model import check
|
|||||||
from .uniapi import mute, recall
|
from .uniapi import mute, recall
|
||||||
from .model.utils.common import process_gif_and_save_jpgs
|
from .model.utils.common import process_gif_and_save_jpgs
|
||||||
|
|
||||||
|
|
||||||
T = TypeVar("T")
|
T = TypeVar("T")
|
||||||
|
|
||||||
|
|
||||||
@ -83,7 +84,6 @@ async def nailong_rule(
|
|||||||
nailong = on_message(rule=Rule(nailong_rule), priority=config.nailong_priority)
|
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
|
input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yolox_size
|
||||||
|
|
||||||
|
|
||||||
@nailong.handle()
|
@nailong.handle()
|
||||||
async def handle_function(bot: BaseBot, ev: BaseEvent, msg: UniMsg, session: Uninfo):
|
async def handle_function(bot: BaseBot, ev: BaseEvent, msg: UniMsg, session: Uninfo):
|
||||||
save_img = False
|
save_img = False
|
||||||
|
|||||||
@ -22,11 +22,20 @@ if config.nailong_model is ModelType.CLASSIFICATION:
|
|||||||
raise_extra_import_error(e, "model0")
|
raise_extra_import_error(e, "model0")
|
||||||
|
|
||||||
elif config.nailong_model is ModelType.TARGET_DETECTION:
|
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:
|
elif config.nailong_model is ModelType.HF_DETECTION:
|
||||||
from .hf_detection import check as check
|
from .hf_detection import check as check
|
||||||
|
|
||||||
|
|
||||||
else:
|
else:
|
||||||
raise ValueError("Invalid model type")
|
raise NotImplementedError # never reach here
|
||||||
|
|||||||
@ -19,7 +19,6 @@ if config.nailong_model2_online:
|
|||||||
import base64
|
import base64
|
||||||
import io
|
import io
|
||||||
import shutil
|
import shutil
|
||||||
|
|
||||||
FILENAME = "nailong_yolo11.pt"
|
FILENAME = "nailong_yolo11.pt"
|
||||||
client = Client("Hakureirm/NailongKiller")
|
client = Client("Hakureirm/NailongKiller")
|
||||||
logger.info(f"Using model {FILENAME} online")
|
logger.info(f"Using model {FILENAME} online")
|
||||||
@ -35,8 +34,6 @@ else:
|
|||||||
api = hf_api.HfApi()
|
api = hf_api.HfApi()
|
||||||
file_path = os.path.join(str(config.nailong_model_dir), FILENAME)
|
file_path = os.path.join(str(config.nailong_model_dir), FILENAME)
|
||||||
model_info = api.model_info(REPO_ID)
|
model_info = api.model_info(REPO_ID)
|
||||||
|
|
||||||
|
|
||||||
def get_file_last_modified_time(file_path):
|
def get_file_last_modified_time(file_path):
|
||||||
try:
|
try:
|
||||||
timestamp = os.path.getmtime(file_path)
|
timestamp = os.path.getmtime(file_path)
|
||||||
@ -44,8 +41,6 @@ else:
|
|||||||
return last_modified_time
|
return last_modified_time
|
||||||
except FileNotFoundError:
|
except FileNotFoundError:
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
local_time = get_file_last_modified_time(file_path)
|
local_time = get_file_last_modified_time(file_path)
|
||||||
if local_time is None or model_info.last_modified >= local_time:
|
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)
|
hf_hub_download(repo_id=REPO_ID, filename=FILENAME, local_dir=config.nailong_model_dir)
|
||||||
@ -57,7 +52,6 @@ else:
|
|||||||
|
|
||||||
input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yolox_size
|
input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yolox_size
|
||||||
|
|
||||||
|
|
||||||
@run_sync
|
@run_sync
|
||||||
def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult:
|
def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult:
|
||||||
if is_gif:
|
if is_gif:
|
||||||
@ -71,8 +65,7 @@ def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult:
|
|||||||
|
|
||||||
if not os.path.exists(os.path.join(str(config.nailong_model_dir),"online_temp")):
|
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"))
|
os.makedirs(os.path.join(str(config.nailong_model_dir),"online_temp"))
|
||||||
image_path = 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")))
|
||||||
"temp_{}.jpg".format(datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")))
|
|
||||||
while os.path.exists(image_path):
|
while os.path.exists(image_path):
|
||||||
basename=os.path.basename(image_path)
|
basename=os.path.basename(image_path)
|
||||||
image_path=os.path.join(str(config.nailong_model_dir),"online_temp",f"exist-{basename}")
|
image_path=os.path.join(str(config.nailong_model_dir),"online_temp",f"exist-{basename}")
|
||||||
@ -82,8 +75,7 @@ def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult:
|
|||||||
api_name="/predict"
|
api_name="/predict"
|
||||||
)
|
)
|
||||||
os.remove(image_path)
|
os.remove(image_path)
|
||||||
if "检测到的目标数量: " in result_info and int(
|
if "检测到的目标数量: " in result_info and int(result_info.split("检测到的目标数量: ")[1].split("\n")[0])<1:
|
||||||
result_info.split("检测到的目标数量: ")[1].split("\n")[0]) < 1:
|
|
||||||
return CheckSingleResult(ok=False,label=None,extra=frame)
|
return CheckSingleResult(ok=False,label=None,extra=frame)
|
||||||
if isinstance(result_image, str):
|
if isinstance(result_image, str):
|
||||||
if result_image.startswith('data:image'):
|
if result_image.startswith('data:image'):
|
||||||
@ -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 = Image.new('RGB', (max_size, max_size), (114, 114, 114))
|
||||||
padded_img.paste(input_image, (pad_w // 2, pad_h // 2))
|
padded_img.paste(input_image, (pad_w // 2, pad_h // 2))
|
||||||
|
|
||||||
|
|
||||||
img_array = np.array(padded_img)
|
img_array = np.array(padded_img)
|
||||||
|
|
||||||
|
|
||||||
results = model.predict(
|
results = model.predict(
|
||||||
img_array,
|
img_array,
|
||||||
conf=config.nailong_model1_score['nailong'],
|
conf=config.nailong_model1_score['nailong'],
|
||||||
|
|||||||
@ -4,17 +4,17 @@ from typing import Optional
|
|||||||
from typing_extensions import override
|
from typing_extensions import override
|
||||||
|
|
||||||
import numpy as np
|
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 cookit import with_semaphore
|
||||||
from nonebot.utils import run_sync
|
from nonebot.utils import run_sync
|
||||||
|
|
||||||
from plugins.nonebot_plugin_nailongremove.config import config
|
from ..config import config
|
||||||
from plugins.nonebot_plugin_nailongremove.frame_source import FrameSource, repack_save
|
from ..frame_source import FrameSource, repack_save
|
||||||
from plugins.nonebot_plugin_nailongremove.model.utils.common import CheckResult, CheckSingleResult, race_check, \
|
from .utils.common import CheckResult, CheckSingleResult, race_check, similarity_process
|
||||||
similarity_process
|
from .utils.update import GitHubLatestReleaseModelUpdater, ModelInfo, UpdaterGroup
|
||||||
from plugins.nonebot_plugin_nailongremove.model.utils.update import GitHubLatestReleaseModelUpdater, ModelInfo, \
|
from .utils.yolox import demo_postprocess, multiclass_nms, preprocess, vis
|
||||||
UpdaterGroup
|
|
||||||
from plugins.nonebot_plugin_nailongremove.model.utils.yolox import demo_postprocess, multiclass_nms, preprocess, vis
|
|
||||||
import itertools
|
import itertools
|
||||||
|
|
||||||
model_filename_sfx = f"_{config.nailong_model1_type.value}.onnx"
|
model_filename_sfx = f"_{config.nailong_model1_type.value}.onnx"
|
||||||
@ -47,15 +47,7 @@ labels = labels_path.read_text("u8").splitlines()
|
|||||||
|
|
||||||
session = onnxruntime.InferenceSession(
|
session = onnxruntime.InferenceSession(
|
||||||
model_path,
|
model_path,
|
||||||
providers=(
|
providers=config.nailong_onnx_providers,
|
||||||
[
|
|
||||||
"TensorrtExecutionProvider",
|
|
||||||
"CUDAExecutionProvider",
|
|
||||||
"CPUExecutionProvider",
|
|
||||||
]
|
|
||||||
if config.nailong_onnx_try_to_use_gpu
|
|
||||||
else ["CPUExecutionProvider"]
|
|
||||||
),
|
|
||||||
)
|
)
|
||||||
input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yolox_size
|
input_shape = config.nailong_model1_yolox_size or config.nailong_model1_type.yolox_size
|
||||||
|
|
||||||
|
|||||||
@ -115,7 +115,6 @@ async def race_check(
|
|||||||
return res
|
return res
|
||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
def similarity_process(image1: np.ndarray, dsize) -> Optional[CheckSingleResult]:
|
def similarity_process(image1: np.ndarray, dsize) -> Optional[CheckSingleResult]:
|
||||||
path = list(glob.glob(os.path.join(config.nailong_model_dir, 'records/*/*.jpg')))
|
path = list(glob.glob(os.path.join(config.nailong_model_dir, 'records/*/*.jpg')))
|
||||||
if len(path) == 0:
|
if len(path) == 0:
|
||||||
|
|||||||
@ -99,12 +99,10 @@ class ModelInfo(Generic[T]):
|
|||||||
|
|
||||||
class ModelUpdater(ABC):
|
class ModelUpdater(ABC):
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def find_from_local(self) -> Optional[Path]:
|
def find_from_local(self) -> Optional[Path]: ...
|
||||||
...
|
|
||||||
|
|
||||||
@abstractmethod
|
@abstractmethod
|
||||||
def get_info(self) -> ModelInfo:
|
def get_info(self) -> ModelInfo: ...
|
||||||
...
|
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def root_dir(self) -> Path:
|
def root_dir(self) -> Path:
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user