This commit is contained in:
student_2333 2024-10-30 01:33:03 +08:00
parent c118baade8
commit 3148fda51a
3 changed files with 31 additions and 27 deletions

View File

@ -176,6 +176,10 @@ plugins = [
## 📝 更新日志
### 2.1.2
- 重构部分代码,修复潜在 Bug
### 2.1.1
- 新增变量 `$checked_image`

View File

@ -9,7 +9,7 @@ require("nonebot_plugin_uninfo")
from . import handler as handler
from .config import Config
__version__ = "2.1.1.post2"
__version__ = "2.1.2"
__plugin_meta__ = PluginMetadata(
name="自动撤回奶龙",
description="一个基于图像分类模型的简单插件~",

View File

@ -1,6 +1,5 @@
import asyncio
import io
from asyncio import Semaphore
from typing import (
Any,
Awaitable,
@ -13,6 +12,7 @@ from typing import (
TypeVar,
cast,
)
from typing_extensions import TypeAlias
import cv2
import numpy as np
@ -115,35 +115,35 @@ async def nailong_rule(
)
async def check_frames(
frames: Iterator[np.ndarray],
) -> Tuple[bool, Optional[np.ndarray]]:
signal = asyncio.Future[Tuple[bool, Optional[np.ndarray]]]()
sem = Semaphore(config.nailong_concurrency)
CheckFrameResult: TypeAlias = Tuple[bool, Optional[np.ndarray]]
async def task(f: np.ndarray):
try:
async with sem:
ok = await check_image(f)
except Exception as e:
signal.set_exception(e)
else:
if isinstance(ok, bool):
ok = (ok, None)
signal.set_result(ok)
async def wait_result():
nonlocal signal
ok = await signal
signal = asyncio.Future()
return ok
async def check_frames(frames: Iterator[np.ndarray]) -> CheckFrameResult:
async def worker() -> CheckFrameResult:
while True:
try:
frame = next(frames)
except StopIteration:
return False, None
res = await check_image(frame)
if not isinstance(res, tuple):
res = res, None
if res[0]:
return res
for f in frames:
if (sem.locked() or signal.done()) and (r := await wait_result())[0]:
return r
asyncio.create_task(task(f))
tasks = [asyncio.create_task(worker()) for _ in range(config.nailong_concurrency)]
while True:
if not tasks:
break
done, pending = await asyncio.wait(tasks, return_when=asyncio.FIRST_COMPLETED)
for t in done:
if (res := t.result())[0]:
for pt in pending:
pt.cancel()
return res
tasks = pending
return await wait_result()
return False, None
nailong = on_message(rule=Rule(nailong_rule), priority=config.nailong_priority)