nonebot-plugin-nailongremove/nonebot_plugin_nailongremove/model/target_detection.py
student_2333 8ff3a8054d
up
2024-10-31 19:38:15 +08:00

82 lines
2.2 KiB
Python

from typing import TYPE_CHECKING
import numpy as np
import onnxruntime
from ..config import config
from .update import GitHubLatestReleaseModelUpdater
from .yolox_utils import demo_postprocess, multiclass_nms, preprocess, vis
if TYPE_CHECKING:
from . import CheckResult
model_path = GitHubLatestReleaseModelUpdater(
"nkxingxh",
"NailongDetection",
lambda x: x.endswith(f"_{config.nailong_model1_type}.onnx"),
).get()
labels_path = GitHubLatestReleaseModelUpdater(
"nkxingxh",
"NailongDetection",
lambda x: x == "labels.txt",
).get()
labels = labels_path.read_text("u8").splitlines()
session = onnxruntime.InferenceSession(
model_path,
providers=(
[
"TensorrtExecutionProvider",
"CUDAExecutionProvider",
"CPUExecutionProvider",
]
if config.nailong_model1_try_to_use_gpu
else ["CPUExecutionProvider"]
),
)
input_shape = config.nailong_model1_yolox_size
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)
predictions = demo_postprocess(output[0], input_shape)[0]
boxes = predictions[:, :4]
scores = predictions[:, 4:5] * predictions[:, 5:]
boxes_xyxy = np.ones_like(boxes)
boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2.0
boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2.0
boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2.0
boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2.0
boxes_xyxy /= ratio
dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1)
if dets is None:
return False
final_boxes, final_scores, final_cls_inds = (
dets[:, :4], # type: ignore
dets[:, 4], # type: ignore
dets[:, 5], # type: ignore
)
has = any(
True
for c, s in zip(final_cls_inds, final_scores)
if c == 1 and s >= config.nailong_model1_score
)
if has:
image = vis(
image,
final_boxes,
final_scores,
final_cls_inds,
conf=0.3,
class_names=labels,
)
return True, image
return False