This commit is contained in:
student_2333 2024-10-30 01:43:30 +08:00
parent 3148fda51a
commit 3507754317
4 changed files with 24 additions and 15 deletions

View File

@ -7,12 +7,9 @@ from typing import (
Iterable,
Iterator,
List,
Optional,
Tuple,
TypeVar,
cast,
)
from typing_extensions import TypeAlias
import cv2
import numpy as np
@ -28,7 +25,7 @@ from nonebot_plugin_uninfo import QryItrface, Uninfo
from PIL import Image as PilImage
from .config import config
from .model import check_image
from .model import CheckResultTuple, check_image
from .uniapi import mute, recall
T = TypeVar("T")
@ -115,11 +112,8 @@ async def nailong_rule(
)
CheckFrameResult: TypeAlias = Tuple[bool, Optional[np.ndarray]]
async def check_frames(frames: Iterator[np.ndarray]) -> CheckFrameResult:
async def worker() -> CheckFrameResult:
async def check_frames(frames: Iterator[np.ndarray]) -> CheckResultTuple:
async def worker() -> CheckResultTuple:
while True:
try:
frame = next(frames)

View File

@ -1,4 +1,5 @@
from typing import Callable, NoReturn, Tuple, Union
from typing import Awaitable, Callable, Literal, NoReturn, Tuple, Union
from typing_extensions import TypeAlias
import numpy as np
from nonebot.utils import run_sync
@ -14,7 +15,13 @@ def raise_extra_import_error(e: BaseException, group: str) -> NoReturn:
) from e
check_image_sync: Callable[[np.ndarray], Union[bool, Tuple[bool, np.ndarray]]]
CheckResultTuple: TypeAlias = Union[
Tuple[bool, None],
Tuple[Literal[True], np.ndarray],
]
CheckResult: TypeAlias = Union[bool, CheckResultTuple]
check_image_sync: Callable[[np.ndarray], CheckResult]
if config.nailong_model is ModelType.CLASSIFICATION:
from .classification import check_image as check_image_sync
@ -29,4 +36,4 @@ else:
raise ValueError("Invalid model type")
check_image = run_sync(check_image_sync)
check_image: Callable[[np.ndarray], Awaitable[CheckResult]] = run_sync(check_image_sync)

View File

@ -1,4 +1,4 @@
from typing import Any
from typing import TYPE_CHECKING, Any
import cv2
import numpy as np
@ -8,6 +8,9 @@ from torchvision import transforms
from ..utils import ensure_model_from_github_repo
if TYPE_CHECKING:
from . import CheckResult
cuda_available = torch.cuda.is_available()
device = torch.device("cuda" if cuda_available else "cpu")
transform = transforms.Compose([transforms.ToTensor()])
@ -30,7 +33,7 @@ if cuda_available:
model.cuda()
def check_image(image: np.ndarray):
def check_image(image: np.ndarray) -> "CheckResult":
if image.shape[0] < 224 or image.shape[1] < 224:
return False
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

View File

@ -1,3 +1,5 @@
from typing import TYPE_CHECKING
import numpy as np
import onnxruntime
@ -5,6 +7,9 @@ from ..config import config
from ..utils import ensure_model_from_github_release
from .yolox_utils import demo_postprocess, multiclass_nms, preprocess, vis
if TYPE_CHECKING:
from . import CheckResult
COCO_CLASSES = ("_background_", "nailong", "anime", "human", "emoji", "long", "other")
model_path = ensure_model_from_github_release(
@ -18,7 +23,7 @@ session = onnxruntime.InferenceSession(model_path)
input_shape = config.nailong_yolox_size
def check_image(image: np.ndarray):
def check_image(image: np.ndarray) -> "CheckResult":
img, ratio = preprocess(image, input_shape)
ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]}
output = session.run(None, ort_inputs)