up

nonebot_plugin_nailongremove/handler.py

@@ -7,12 +7,9 @@
     Iterable,
     Iterator,
     List,
-    Optional,
-    Tuple,
     TypeVar,
     cast,
 )
-from typing_extensions import TypeAlias
 
 import cv2
 import numpy as np
@@ -28,7 +25,7 @@
 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)

nonebot_plugin_nailongremove/model/__init__.py

@@ -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 @@ def raise_extra_import_error(e: BaseException, group: str) -> NoReturn:
     raise ValueError("Invalid model type")
 
 
-check_image = run_sync(check_image_sync)
+check_image: Callable[[np.ndarray], Awaitable[CheckResult]] = run_sync(check_image_sync)

nonebot_plugin_nailongremove/model/classification.py

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

nonebot_plugin_nailongremove/model/target_detection.py

@@ -1,10 +1,15 @@
+from typing import TYPE_CHECKING
+
 import numpy as np
 import onnxruntime
 
 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 @@
 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)