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
3291728604
commit
c9ef92349d
@ -180,7 +180,7 @@ async def extract_source(seg: Segment) -> FrameSource:
|
|||||||
return await source_extractors[k](seg)
|
return await source_extractors[k](seg)
|
||||||
|
|
||||||
|
|
||||||
async def iter_frames_in_message(
|
async def iter_sources_in_message(
|
||||||
message: UniMessage,
|
message: UniMessage,
|
||||||
) -> AsyncIterator[Tuple[FrameSource, Segment]]:
|
) -> AsyncIterator[Tuple[FrameSource, Segment]]:
|
||||||
for seg in message:
|
for seg in message:
|
||||||
|
|||||||
@ -8,7 +8,7 @@ from nonebot_plugin_alconna.builtins.uniseg.market_face import MarketFace
|
|||||||
from nonebot_plugin_alconna.uniseg import Image, UniMessage, UniMsg
|
from nonebot_plugin_alconna.uniseg import Image, UniMessage, UniMsg
|
||||||
from nonebot_plugin_uninfo import QryItrface, Uninfo
|
from nonebot_plugin_uninfo import QryItrface, Uninfo
|
||||||
|
|
||||||
from nonebot_plugin_nailongremove.frame_source import iter_frames_in_message
|
from nonebot_plugin_nailongremove.frame_source import iter_sources_in_message
|
||||||
|
|
||||||
from .config import config
|
from .config import config
|
||||||
from .model import check
|
from .model import check
|
||||||
@ -84,9 +84,9 @@ nailong = on_message(rule=Rule(nailong_rule), priority=config.nailong_priority)
|
|||||||
|
|
||||||
@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):
|
||||||
async for frames, seg in iter_frames_in_message(msg):
|
async for source, seg in iter_sources_in_message(msg):
|
||||||
try:
|
try:
|
||||||
check_res = await check(frames)
|
check_res = await check(source)
|
||||||
except Exception:
|
except Exception:
|
||||||
logger.exception(f"Failed to check {seg!r}")
|
logger.exception(f"Failed to check {seg!r}")
|
||||||
continue
|
continue
|
||||||
|
|||||||
@ -2,7 +2,7 @@ from typing import Awaitable, Callable, NoReturn
|
|||||||
|
|
||||||
from ..config import ModelType, config
|
from ..config import ModelType, config
|
||||||
from ..frame_source import FrameSource
|
from ..frame_source import FrameSource
|
||||||
from .common import CheckResult as CheckResult
|
from .utils.common import CheckResult as CheckResult
|
||||||
|
|
||||||
|
|
||||||
def raise_extra_import_error(e: BaseException, group: str) -> NoReturn:
|
def raise_extra_import_error(e: BaseException, group: str) -> NoReturn:
|
||||||
|
|||||||
@ -8,8 +8,8 @@ from torch import nn
|
|||||||
from torchvision import transforms
|
from torchvision import transforms
|
||||||
|
|
||||||
from ..frame_source import FrameSource
|
from ..frame_source import FrameSource
|
||||||
from .common import CheckResult, CheckSingleResult, race_check
|
from .utils.common import CheckResult, CheckSingleResult, race_check
|
||||||
from .update import GitHubRepoModelUpdater
|
from .utils.update import GitHubRepoModelUpdater
|
||||||
|
|
||||||
model_path = GitHubRepoModelUpdater(
|
model_path = GitHubRepoModelUpdater(
|
||||||
"spawner1145",
|
"spawner1145",
|
||||||
@ -47,6 +47,6 @@ def check_single(image: np.ndarray) -> CheckSingleResult[None]:
|
|||||||
return CheckSingleResult(ok=pred.item() == 1, extra=None)
|
return CheckSingleResult(ok=pred.item() == 1, extra=None)
|
||||||
|
|
||||||
|
|
||||||
async def check(frames: FrameSource):
|
async def check(source: FrameSource):
|
||||||
res = await race_check(check_single, frames)
|
res = await race_check(check_single, source)
|
||||||
return CheckResult(ok=bool(res))
|
return CheckResult(ok=bool(res))
|
||||||
|
|||||||
@ -10,9 +10,9 @@ from nonebot.utils import run_sync
|
|||||||
|
|
||||||
from ..config import config
|
from ..config import config
|
||||||
from ..frame_source import FrameSource, repack_save
|
from ..frame_source import FrameSource, repack_save
|
||||||
from .common import CheckResult, CheckSingleResult, race_check
|
from .utils.common import CheckResult, CheckSingleResult, race_check
|
||||||
from .update import GitHubLatestReleaseModelUpdater, ModelInfo, UpdaterGroup
|
from .utils.update import GitHubLatestReleaseModelUpdater, ModelInfo, UpdaterGroup
|
||||||
from .yolox_utils import demo_postprocess, multiclass_nms, preprocess, vis
|
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}.onnx"
|
||||||
|
|
||||||
@ -127,25 +127,25 @@ async def check_single(frame: np.ndarray) -> CheckSingleResult[FrameInfo]:
|
|||||||
return CheckSingleResult(ok=res.ok, extra=FrameInfo(frame, res.extra))
|
return CheckSingleResult(ok=res.ok, extra=FrameInfo(frame, res.extra))
|
||||||
|
|
||||||
|
|
||||||
async def check(frames: FrameSource) -> CheckResult:
|
async def check(source: FrameSource) -> CheckResult:
|
||||||
extra_vars = {}
|
extra_vars = {}
|
||||||
if config.nailong_checked_result_all:
|
if config.nailong_checked_result_all:
|
||||||
sem = asyncio.Semaphore(config.nailong_concurrency)
|
sem = asyncio.Semaphore(config.nailong_concurrency)
|
||||||
results = asyncio.gather(
|
results = asyncio.gather(
|
||||||
*(with_semaphore(sem)(check_single)(frame) for frame in frames),
|
*(with_semaphore(sem)(check_single)(frame) for frame in source),
|
||||||
)
|
)
|
||||||
ok = any(r.ok for r in results)
|
ok = any(r.ok for r in results)
|
||||||
if ok:
|
if ok:
|
||||||
extra_vars["$checked_result"] = await repack_save(
|
extra_vars["$checked_result"] = await repack_save(
|
||||||
frames,
|
source,
|
||||||
(r.extra.vis() for r in results),
|
(r.extra.vis() for r in results),
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
res = await race_check(check_single, frames)
|
res = await race_check(check_single, source)
|
||||||
ok = bool(res)
|
ok = bool(res)
|
||||||
if res:
|
if res:
|
||||||
extra_vars["$checked_result"] = await repack_save(
|
extra_vars["$checked_result"] = await repack_save(
|
||||||
frames,
|
source,
|
||||||
iter((res.extra.vis(),)),
|
iter((res.extra.vis(),)),
|
||||||
)
|
)
|
||||||
return CheckResult(ok, extra_vars)
|
return CheckResult(ok, extra_vars)
|
||||||
|
|||||||
@ -5,8 +5,8 @@ from typing_extensions import TypeAlias
|
|||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
from ..config import config
|
from ...config import config
|
||||||
from ..frame_source import FrameSource
|
from ...frame_source import FrameSource
|
||||||
|
|
||||||
T = TypeVar("T")
|
T = TypeVar("T")
|
||||||
|
|
||||||
@ -10,7 +10,7 @@ from githubkit import GitHub
|
|||||||
from nonebot import logger
|
from nonebot import logger
|
||||||
from tqdm import tqdm
|
from tqdm import tqdm
|
||||||
|
|
||||||
from ..config import config
|
from ...config import config
|
||||||
|
|
||||||
T = TypeVar("T")
|
T = TypeVar("T")
|
||||||
|
|
||||||
Loading…
Reference in New Issue
Block a user