This commit is contained in:
student_2333 2024-11-06 01:55:49 +08:00
parent 3291728604
commit c9ef92349d
No known key found for this signature in database
GPG Key ID: 665F083BEC56F2A6
9 changed files with 20 additions and 20 deletions

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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