From b155703f74994f55ec0139250a4c49d4c039d2c0 Mon Sep 17 00:00:00 2001 From: Neucrack Date: Thu, 12 Dec 2024 15:21:43 +0800 Subject: [PATCH] optimize nn MUD parse labels strip space --- docs/doc/en/pro/customize_model.md | 61 ++++++++++++++++--- docs/doc/zh/pro/customize_model.md | 9 ++- .../vision/ai_vision/nn_custom_classifier.py | 9 ++- examples/vision/ai_vision/nn_yolo11_seg.py | 2 +- 4 files changed, 64 insertions(+), 17 deletions(-) diff --git a/docs/doc/en/pro/customize_model.md b/docs/doc/en/pro/customize_model.md index d3e50eca..473093c9 100644 --- a/docs/doc/en/pro/customize_model.md +++ b/docs/doc/en/pro/customize_model.md @@ -71,24 +71,65 @@ import numpy as np def parse_str_values(value: str) -> list[float]: return [float(v) for v in value.split(",")] -def load_labels(model_path, path_or_labels: str): - path = os.path.join(os.path.dirname(model_path), path_or_labels) - labels0 = open(path, encoding="utf-8").readlines() if os.path.exists(path) else path_or_labels.split(",") - return [label.strip() for label in labels0] +def load_labels(model_path, path_or_labels : str): + path = "" + if not ("," in path_or_labels or " " in path_or_labels or "\n" in path_or_labels): + path = os.path.join(os.path.dirname(model_path), path_or_labels) + if path and os.path.exists(path): + with open(path, encoding = "utf-8") as f: + labels0 = f.readlines() + else: + labels0 = path_or_labels.split(",") + labels = [] + for label in labels0: + labels.append(label.strip()) + return labels class My_Classifier: - def __init__(self, model: str): - self.model = nn.NN(model, dual_buff=False) + def __init__(self, model : str): + self.model = nn.NN(model, dual_buff = False) + self.extra_info = self.model.extra_info() + self.mean = parse_str_values(self.extra_info["mean"]) + self.scale = parse_str_values(self.extra_info["scale"]) + self.labels = self.model.extra_info_labels() + # self.labels = load_labels(model, self.extra_info["labels"]) # same as self.model.extra_info_labels() + + def classify(self, img : image.Image): + outs = self.model.forward_image(img, self.mean, self.scale, copy_result = False) + # 后处理， 以分类模型为例 + for k in outs.keys(): + out = nn.F.softmax(outs[k], replace=True) + out = tensor.tensor_to_numpy_float32(out, copy = False).flatten() + max_idx = out.argmax() + return self.labels[max_idx], out[max_idx]def load_labels(model_path, path_or_labels : str): + path = "" + if not ("," in path_or_labels or " " in path_or_labels or "\n" in path_or_labels): + path = os.path.join(os.path.dirname(model_path), path_or_labels) + if path and os.path.exists(path): + with open(path, encoding = "utf-8") as f: + labels0 = f.readlines() + else: + labels0 = path_or_labels.split(",") + labels = [] + for label in labels0: + labels.append(label.strip()) + return labels + +class My_Classifier: + def __init__(self, model : str): + self.model = nn.NN(model, dual_buff = False) self.extra_info = self.model.extra_info() self.mean = parse_str_values(self.extra_info["mean"]) self.scale = parse_str_values(self.extra_info["scale"]) - self.labels = load_labels(model, self.extra_info["labels"]) + self.labels = self.model.extra_info_labels() + # self.labels = load_labels(model, self.extra_info["labels"]) # same as self.model.extra_info_labels() - def classify(self, img: image.Image): - outs = self.model.forward_image(img, self.mean, self.scale, copy_result=False) + def classify(self, img : image.Image): + outs = self.model.forward_image(img, self.mean, self.scale, copy_result = False) + # 后处理， 以分类模型为例 for k in outs.keys(): out = nn.F.softmax(outs[k], replace=True) - out = tensor.tensor_to_numpy_float32(out, copy=False).flatten() + out = tensor.tensor_to_numpy_float32(out, copy = False).flatten() max_idx = out.argmax() return self.labels[max_idx], out[max_idx] diff --git a/docs/doc/zh/pro/customize_model.md b/docs/doc/zh/pro/customize_model.md index 5d9f81ee..a4bee93f 100644 --- a/docs/doc/zh/pro/customize_model.md +++ b/docs/doc/zh/pro/customize_model.md @@ -80,8 +80,10 @@ def parse_str_values(value : str) -> list[float]: return [float(value)] def load_labels(model_path, path_or_labels : str): - path = os.path.join(os.path.dirname(model_path), path_or_labels) - if os.path.exists(path): + path = "" + if not ("," in path_or_labels or " " in path_or_labels or "\n" in path_or_labels): + path = os.path.join(os.path.dirname(model_path), path_or_labels) + if path and os.path.exists(path): with open(path, encoding = "utf-8") as f: labels0 = f.readlines() else: @@ -97,7 +99,8 @@ class My_Classifier: self.extra_info = self.model.extra_info() self.mean = parse_str_values(self.extra_info["mean"]) self.scale = parse_str_values(self.extra_info["scale"]) - self.labels = load_labels(model, self.extra_info["labels"]) + self.labels = self.model.extra_info_labels() + # self.labels = load_labels(model, self.extra_info["labels"]) # same as self.model.extra_info_labels() def classify(self, img : image.Image): outs = self.model.forward_image(img, self.mean, self.scale, copy_result = False) diff --git a/examples/vision/ai_vision/nn_custom_classifier.py b/examples/vision/ai_vision/nn_custom_classifier.py index 8a71e083..cc401ba2 100644 --- a/examples/vision/ai_vision/nn_custom_classifier.py +++ b/examples/vision/ai_vision/nn_custom_classifier.py @@ -13,8 +13,10 @@ def parse_str_values(value : str) -> list[float]: return [float(value)] def load_labels(model_path, path_or_labels : str): - path = os.path.join(os.path.dirname(model_path), path_or_labels) - if os.path.exists(path): + path = "" + if not ("," in path_or_labels or " " in path_or_labels or "\n" in path_or_labels): + path = os.path.join(os.path.dirname(model_path), path_or_labels) + if path and os.path.exists(path): with open(path, encoding = "utf-8") as f: labels0 = f.readlines() else: @@ -30,7 +32,8 @@ def __init__(self, model : str): self.extra_info = self.model.extra_info() self.mean = parse_str_values(self.extra_info["mean"]) self.scale = parse_str_values(self.extra_info["scale"]) - self.labels = load_labels(model, self.extra_info["labels"]) + self.labels = self.model.extra_info_labels() + # self.labels = load_labels(model, self.extra_info["labels"]) # same as self.model.extra_info_labels() def classify(self, img : image.Image): outs = self.model.forward_image(img, self.mean, self.scale, copy_result = False) diff --git a/examples/vision/ai_vision/nn_yolo11_seg.py b/examples/vision/ai_vision/nn_yolo11_seg.py index dc616f55..17f06644 100644 --- a/examples/vision/ai_vision/nn_yolo11_seg.py +++ b/examples/vision/ai_vision/nn_yolo11_seg.py @@ -1,6 +1,6 @@ from maix import camera, display, image, nn, app, time -detector = nn.YOLO11(model="/root/models/yolov11n_seg.mud", dual_buff = True) +detector = nn.YOLO11(model="/root/models/yolo11n_seg.mud", dual_buff = True) cam = camera.Camera(detector.input_width(), detector.input_height(), detector.input_format()) disp = display.Display()