mirror of
https://github.com/Refound-445/nonebot-plugin-nailongremove.git
synced 2025-11-04 21:22:43 +08:00
up
This commit is contained in:
parent
3148fda51a
commit
3507754317
@ -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)
|
||||||
|
|||||||
@ -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)
|
||||||
|
|||||||
@ -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)
|
||||||
|
|||||||
@ -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)
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user