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