Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 30 additions & 13 deletions app.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
from PyQt5.QtGui import QPixmap
from PyQt5.QtCore import QThread, pyqtSignal
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', write_through=True)

lang_id = 0
lang = {
Expand Down Expand Up @@ -82,9 +82,15 @@
"yolov8n-cls": ["yolov8n-cls", "object-classification-classifier"],
"yolov8n-seg": ["yolov8n-seg", "object-segmentation-segment"],
"yolov8n": ["yolov8n-det", "object-detection-detector"],
"yolo11n-cls": ["yolo11n-cls", "object-classification-classifier"],
"yolo11n-seg": ["yolo11n-seg", "object-segmentation-segment"],
"yolo11n": ["yolo11n-det", "object-detection-detector"],

"yolo11n-cls": ["yolov8n-cls", "object-classification-classifier"],
"yolov11n-cls": ["yolov8n-cls", "object-classification-classifier"],

"yolo11n-seg": ["yolov8n-seg", "object-segmentation-segment"],
"yolov11n-seg": ["yolov8n-seg", "object-segmentation-segment"],

"yolo11n": ["yolov8n-det", "object-detection-detector"],
"yolov11n": ["yolov8n-det", "object-detection-detector"],
}

def clean_name(name):
Expand Down Expand Up @@ -292,14 +298,19 @@ def __init__(self, onnx_path, kmodel_path, dataset_path, conf_path, output_zip_f
self.output_zip_file = output_zip_file

def run(self):
import convertor
# 耗时操作放在这里
shapes = get_input_shape(self.onnx_path)
if 1 == len(shapes):
shape = list(shapes.values())[0]
convertor.make(self.onnx_path, self.kmodel_path, self.dataset_path, self.conf_path,shape)
file_path = zip_with_md5(base_name=self.output_zip_file)
self.finished.emit(file_path) # 发射信号通知主线程
try:
import convertor
# 耗时操作放在这里
shapes = get_input_shape(self.onnx_path)
if 1 == len(shapes):
shape = list(shapes.values())[0]
convertor.make(self.onnx_path, self.kmodel_path, self.dataset_path, self.conf_path, shape)
file_path = zip_with_md5(base_name=self.output_zip_file)
self.finished.emit(file_path) # 发射信号通知主线程
except Exception as e:
import traceback
traceback.print_exc()
print(f"!!! 子线程发生严重错误: {e} !!!", flush=True)


class ModelExportApp(QWidget):
Expand Down Expand Up @@ -705,7 +716,7 @@ def export_model(self):

icon_file = self._conf["comm"]["icon_file"]
if os.path.exists(icon_file):
shutil.copy(icon_file, os.path.join("model_output/icon.png"))
shutil.copy(icon_file, os.path.join("model_output", "icon.png"))

# 创建空文件
open(f"model_output/app.{conf_data['conf']['application']}", "w").close()
Expand Down Expand Up @@ -736,6 +747,12 @@ def pack(self):
print("转换完成!")

if __name__ == "__main__":
import traceback
# 拦截 PyQt5 的主线程异常,防止 0xC0000409 闪退
def exception_hook(type, value, tb):
traceback.print_exception(type, value, tb)
print(f"\n[程序内部报错] {value}", flush=True)
sys.excepthook = exception_hook
language_code = locale.getdefaultlocale()[0]
print(f"默认语言环境: {language_code}")
if language_code.startswith("zh_CN"):
Expand Down
50 changes: 25 additions & 25 deletions convertor.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,35 +2,27 @@

import subprocess
import argparse
import nncase
import tomlkit
import sys
import os
import cv2
import numpy as np
from pathlib import Path

swapRB = False
preprocess = False

input_type = np.uint8

templs_shape = 640


# setup env
result = subprocess.run(["pip", "show", "nncase"], capture_output=True)
result = subprocess.run([sys.executable, "-m", "pip", "show", "nncase"], capture_output=True)
line_break = "\n"
if sys.platform == "win32":
line_break = "\r\n"
location_s = [i for i in result.stdout.decode().split(
line_break) if i.startswith("Location:")]
location = location_s[0].split(": ")[1]
if "PATH" in os.environ:
os.environ["PATH"] += os.pathsep + location
else:
os.environ["PATH"] = location
os.environ["NNCASE_PLUGIN_PATH"] = location
location_s = [i for i in result.stdout.decode().split(line_break) if i.startswith("Location:")]
if location_s:
location = location_s[0].split(": ")[1]
if "PATH" in os.environ:
os.environ["PATH"] += os.pathsep + location
else:
os.environ["PATH"] = location
os.environ["NNCASE_PLUGIN_PATH"] = location

import nncase
import tomlkit


class Convertor(nncase.Compiler):
Expand Down Expand Up @@ -134,30 +126,38 @@ def gen(_dir):
'''

def gen(_dir):
path = Path(_dir)

path = Path(_dir).resolve()
count = 0
for f in path.rglob('*'):
if not f.is_file():
continue

img_path = str(f) # �ؼ�����Ҫ f.name
img_path = str(f)

data = np.fromfile(img_path, dtype=np.uint8)
templ = cv2.imdecode(data, cv2.IMREAD_COLOR)

if templ is None:
print(f"[WARN] read failed: {img_path}")
continue

templ = process_img(templ)
count += 1
print(f" -> 成功载入校准图 [{count}/100]: {f.name}", flush=True)
yield templ

def make(onnx_file, kmodel_file, dataset, toml_file, input_shape):
global templs_shape
templs_shape = input_shape[2]
calib = []
max_calib_images = 100
for t in gen(dataset):
calib.append(t)
if len(calib) >= max_calib_images:
break

if len(calib) == 0:
print(f"\n[错误] 严重警告:在目录 '{dataset}' 下未找到任何有效的图片!", flush=True)
print("请检查该文件夹是否存在,且里面是否确实含有 .jpg 或 .png 图片。\n", flush=True)
raise FileNotFoundError(f"未能在 {dataset} 中找到校准图片")

npcalib = np.array(calib).astype(np.uint8)

Expand Down