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

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

View File

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

View File

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

View File

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

View File

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

View File

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