This commit is contained in:
Refound-445 2024-11-16 20:41:23 +08:00
parent 3b45bab8ff
commit 26389a66ee
9 changed files with 80 additions and 75 deletions

View File

@ -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="一个基于图像分类模型的简单插件~",

View File

@ -55,14 +55,14 @@ 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
nailong_model1_score: Dict[str, Optional[float]] = { nailong_model1_score: Dict[str, Optional[float]] = {
DEFAULT_LABEL: 0.5, DEFAULT_LABEL: 0.5,
} }
nailong_model2_online: bool = False nailong_model2_online: bool=False
nailong_check_mode: int = 0 nailong_check_mode: int = 0
nailong_similarity_on: bool = False nailong_similarity_on: bool = False
nailong_similarity_max_storage: int = 10 nailong_similarity_max_storage: int = 10
@ -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)

View File

@ -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

View File

@ -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

View File

@ -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")
@ -30,13 +29,11 @@ else:
REPO_ID = "Hakureirm/NailongKiller" REPO_ID = "Hakureirm/NailongKiller"
FILENAME = "nailong_yolo11.pt" FILENAME = "nailong_yolo11.pt"
model_path = os.path.join(str(config.nailong_model_dir), FILENAME) model_path=os.path.join(str(config.nailong_model_dir),FILENAME)
if config.nailong_auto_update_model or not os.path.exists(model_path): if config.nailong_auto_update_model or not os.path.exists(model_path):
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,47 +41,42 @@ 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)
logger.info(f"Update model {FILENAME} successfully!") 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) model = YOLO(model_path).to(device)
logger.info(f"Using model {FILENAME}") logger.info(f"Using model {FILENAME}")
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:
res = similarity_process(frame, dsize=input_shape) res = similarity_process(frame, dsize=input_shape)
if res is not None: if res is not None:
return CheckSingleResult(ok=res.ok, label=res.label, extra=frame) return CheckSingleResult(ok=res.ok,label=res.label,extra=frame)
return CheckSingleResult(ok=False, label=None, extra=frame) return CheckSingleResult(ok=False,label=None,extra=frame)
else: else:
if config.nailong_model2_online: if config.nailong_model2_online:
input_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) 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")): 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}")
input_image.save(image_path, format='JPEG') input_image.save(image_path,format='JPEG')
result_image, result_info = client.predict( result_image, result_info = client.predict(
img=handle_file(image_path), img=handle_file(image_path),
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'):
img_data = base64.b64decode(result_image.split(',')[1]) 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) img = Image.open(img_data)
result_image = np.array(img) result_image = np.array(img)
shutil.rmtree(os.path.dirname(os.path.dirname(img_data))) shutil.rmtree(os.path.dirname(os.path.dirname(img_data)))
result_image = cv2.cvtColor(result_image, cv2.COLOR_BGR2RGB) result_image=cv2.cvtColor(result_image,cv2.COLOR_BGR2RGB)
return CheckSingleResult(ok=True, label="nailong", extra=result_image) return CheckSingleResult(ok=True,label="nailong",extra=result_image)
else: else:
input_image = Image.fromarray(frame) input_image = Image.fromarray(frame)
original_size = input_image.size 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 = 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'],
@ -119,7 +113,7 @@ def _check_single(frame: np.ndarray, is_gif: bool = False) -> CheckSingleResult:
) )
cls = results[0].boxes.cls cls = results[0].boxes.cls
if len(cls) < 1: 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() result_img = results[0].plot()
if pad_w > 0 or pad_h > 0: if pad_w > 0 or pad_h > 0:

View File

@ -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

View File

@ -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:
@ -187,7 +186,7 @@ def process_gif_and_save_jpgs(frames, label, dsize, similarity_threshold=0.85):
indices = torch.nonzero(similarities > similarity_threshold) indices = torch.nonzero(similarities > similarity_threshold)
index = indices.squeeze().tolist() if indices.numel() > 0 else None index = indices.squeeze().tolist() if indices.numel() > 0 else None
if type(index) is int: if type(index) is int:
index = [index] index=[index]
if index is not None: if index is not None:
indexs.extend([frame2_num[i] for i in index]) indexs.extend([frame2_num[i] for i in index])
frame_count = [i for i in frame_count if i not in indexs] frame_count = [i for i in frame_count if i not in indexs]

View File

@ -56,10 +56,10 @@ def create_parent_dir(path: Path, create: bool = True):
def find_file( def find_file(
path: Path, path: Path,
checker: Union[Callable[[Path], bool], str, None] = None, checker: Union[Callable[[Path], bool], str, None] = None,
recursive: bool = False, recursive: bool = False,
last_modified: bool = True, last_modified: bool = True,
) -> Optional[Path]: ) -> Optional[Path]:
if isinstance(checker, str) and checker: if isinstance(checker, str) and checker:
if (p := path / checker).exists(): if (p := path / checker).exists():
@ -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:
@ -121,8 +119,8 @@ class ModelUpdater(ABC):
def check_local_ver(self, info: ModelInfo) -> Optional[str]: def check_local_ver(self, info: ModelInfo) -> Optional[str]:
if ( if (
self.get_path(info.filename).exists() self.get_path(info.filename).exists()
and (ver_path := self.get_ver_path(info.filename)).exists() and (ver_path := self.get_ver_path(info.filename)).exists()
): ):
return ver_path.read_text(encoding="u8").strip() return ver_path.read_text(encoding="u8").strip()
return None return None
@ -164,10 +162,10 @@ class ModelUpdater(ABC):
return return
def validate_with_unlink( def validate_with_unlink(
self, self,
path: Path, path: Path,
info: ModelInfo, info: ModelInfo,
clear_ver: bool = True, clear_ver: bool = True,
) -> Any: ) -> Any:
try: try:
return self.validate(path, info) return self.validate(path, info)
@ -179,9 +177,9 @@ class ModelUpdater(ABC):
def _get(self, force_update: bool = False) -> Path: def _get(self, force_update: bool = False) -> Path:
if ( if (
(not force_update) (not force_update)
and (not config.nailong_auto_update_model) and (not config.nailong_auto_update_model)
and (local := self.find_from_local()) and (local := self.find_from_local())
): ):
logger.info("Update skipped") logger.info("Update skipped")
return local return local
@ -301,10 +299,10 @@ class GitHubRepoModelUpdater(GitHubModelUpdater):
class GitHubLatestReleaseModelUpdater(GitHubModelUpdater): class GitHubLatestReleaseModelUpdater(GitHubModelUpdater):
def __init__( def __init__(
self, self,
owner: str, owner: str,
repo: str, repo: str,
local_filename_checker: Optional[Callable[[str], bool]] = None, local_filename_checker: Optional[Callable[[str], bool]] = None,
) -> None: ) -> None:
super().__init__() super().__init__()
self.owner = owner self.owner = owner