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, Iterable,
Iterator, Iterator,
List, List,
Optional,
Tuple,
TypeVar, TypeVar,
cast, cast,
) )
from typing_extensions import TypeAlias
import cv2 import cv2
import numpy as np import numpy as np
@ -28,7 +25,7 @@ from nonebot_plugin_uninfo import QryItrface, Uninfo
from PIL import Image as PilImage from PIL import Image as PilImage
from .config import config from .config import config
from .model import check_image from .model import CheckResultTuple, check_image
from .uniapi import mute, recall from .uniapi import mute, recall
T = TypeVar("T") 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]) -> CheckResultTuple:
async def worker() -> CheckResultTuple:
async def check_frames(frames: Iterator[np.ndarray]) -> CheckFrameResult:
async def worker() -> CheckFrameResult:
while True: while True:
try: try:
frame = next(frames) 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 import numpy as np
from nonebot.utils import run_sync from nonebot.utils import run_sync
@ -14,7 +15,13 @@ def raise_extra_import_error(e: BaseException, group: str) -> NoReturn:
) from e ) 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: if config.nailong_model is ModelType.CLASSIFICATION:
from .classification import check_image as check_image_sync from .classification import check_image as check_image_sync
@ -29,4 +36,4 @@ else:
raise ValueError("Invalid model type") 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 cv2
import numpy as np import numpy as np
@ -8,6 +8,9 @@ from torchvision import transforms
from ..utils import ensure_model_from_github_repo from ..utils import ensure_model_from_github_repo
if TYPE_CHECKING:
from . import CheckResult
cuda_available = torch.cuda.is_available() cuda_available = torch.cuda.is_available()
device = torch.device("cuda" if cuda_available else "cpu") device = torch.device("cuda" if cuda_available else "cpu")
transform = transforms.Compose([transforms.ToTensor()]) transform = transforms.Compose([transforms.ToTensor()])
@ -30,7 +33,7 @@ if cuda_available:
model.cuda() 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: if image.shape[0] < 224 or image.shape[1] < 224:
return False return False
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

View File

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