Upload 6 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+model.rknn filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,139 @@
----
-license: agpl-3.0
----
+---
+license: agpl-3.0
+base_model: SmilingWolf/wd-convnext-tagger-v3
+tags:
+- rknn
+---
+
+# WD ConvNext Tagger v3 RKNN2
+
+## (English README see below)
+
+在RK3588上运行WaifuDiffusion图像标签模型!
+
+- 推理速度(RK3588):
+  - 单NPU核: 320ms
+  
+- 内存占用(RK3588): 
+  - 0.45GB
+
+## 使用方法
+
+1. 克隆或者下载此仓库到本地
+   
+2. 安装依赖
+
+```bash
+pip install numpy<2 pandas opencv-python rknn-toolkit-lite2
+```
+
+3. 运行
+   
+```bash
+python run_rknn.py input.jpg
+```
+
+输出结果示例:
+```log
+       tag_id           name     probs
+0     9999999        general  0.521484
+5      212816           solo  0.929199
+12      15080     short_hair  0.520508
+25     540830           1boy  0.947754
+40      16613        jewelry  0.577148
+72    1300281     male_focus  0.907227
+130     10926          pants  0.803223
+346   1094664   colored_skin  0.570312
+373      4009     turtleneck  0.552246
+1532  1314823  black_sweater  0.514160
+```
+
+## 模型转换
+
+1. 安装依赖
+
+```bash
+pip install numpy<2 onnxruntime rknn-toolkit2
+```
+
+2. 下载原始onnx模型
+
+3. 转换onnx模型到rknn模型:
+
+```bash
+python convert_rknn.py
+```
+
+## 已知问题
+
+- int8量化后精度损失极大, 基本不可用. 不建议使用量化推理.
+
+## 参考
+
+- [SmilingWolf/wd-convnext-tagger-v3](https://huggingface.co/SmilingWolf/wd-convnext-tagger-v3)
+
+## English README
+
+Run efficient image tagging model on RK3588!
+
+- Inference Speed (RK3588):
+  - Single NPU Core: 320ms
+
+- Memory Usage (RK3588):
+  - 0.45GB
+
+## Usage
+
+1. Clone or download this repository
+
+2. Install dependencies
+
+```bash
+pip install numpy<2 pandas opencv-python rknn-toolkit-lite2
+```
+
+3. Run
+
+```bash
+python run_rknn.py input.jpg
+```
+
+Output example:
+```log
+       tag_id           name     probs
+0     9999999        general  0.521484
+5      212816           solo  0.929199
+12      15080     short_hair  0.520508
+25     540830           1boy  0.947754
+40      16613        jewelry  0.577148
+72    1300281     male_focus  0.907227
+130     10926          pants  0.803223
+346   1094664   colored_skin  0.570312
+373      4009     turtleneck  0.552246
+1532  1314823  black_sweater  0.514160
+```
+
+
+## Model Conversion
+
+1. Install dependencies
+
+```bash
+pip install numpy<2 onnxruntime rknn-toolkit2
+```
+
+2. Download original onnx model
+
+3. Convert onnx model to rknn model:
+
+```bash
+python convert_rknn.py
+```
+
+## Known Issues
+
+- Huge precision loss after int8 quantization, not recommended to use quantized inference.
+
+## References
+
+- [SmilingWolf/wd-convnext-tagger-v3](https://huggingface.co/SmilingWolf/wd-convnext-tagger-v3)

convert_rknn.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+# coding: utf-8
+
+import datetime
+from rknn.api import RKNN
+from sys import exit
+import argparse
+
+
+ONNX_MODEL="model.onnx"
+detailed_performance_log = True
+
+# 添加命令行参数解析
+parser = argparse.ArgumentParser(description='将ONNX模型转换为RKNN模型')
+parser.add_argument('--quantize', '-q', action='store_true',
+                   help='是否对模型进行量化')
+parser.add_argument('--dataset', '-d', 
+                   default="/home/zt/rk3588-nn/rknn_model_zoo/datasets/COCO/coco_subset_20.txt",
+                   help='量化校准数据集的路径')
+args = parser.parse_args()
+
+DATASET=args.dataset  # 使用命令行参数指定数据集路径
+QUANTIZE=args.quantize
+
+# 根据是否量化决定输出文件名
+if args.quantize:
+    RKNN_MODEL=ONNX_MODEL.replace(".onnx","_int8.rknn")
+else:
+    RKNN_MODEL=ONNX_MODEL.replace(".onnx",".rknn")
+
+timedate_iso = datetime.datetime.now().isoformat()
+
+rknn = RKNN(verbose=True)
+rknn.config(
+    # mean_values=[x * 255 for x in [0.485, 0.456, 0.406]],
+    # std_values=[x * 255 for x in [0.229, 0.224, 0.225]],
+    quantized_dtype='w8a8',
+    quantized_algorithm='kl_divergence',
+    quantized_method='channel',
+    quantized_hybrid_level=0,
+    target_platform='rk3588',
+    quant_img_RGB2BGR = False,
+    float_dtype='float16',
+    optimization_level=3,
+    custom_string=f"converted by: qq: 232004040, email: 2302004040@qq.com at {timedate_iso}",
+    remove_weight=False,
+    compress_weight=False,
+    inputs_yuv_fmt=None,
+    single_core_mode=False,
+    dynamic_input=None,
+    model_pruning=False,
+    op_target=None,
+    quantize_weight=False,
+    remove_reshape=False,
+    sparse_infer=False,
+    enable_flash_attention=False,
+    )
+
+ret = rknn.load_onnx(model=ONNX_MODEL, inputs=["/Transpose_output_0"], input_size_list=[[1,3,448,448]])
+ret = rknn.build(do_quantization=QUANTIZE, dataset=DATASET, rknn_batch_size=None)
+ret = rknn.export_rknn(RKNN_MODEL)
+
+# ret = rknn.init_runtime(target='rk3588',device_id='cbb956772bf5dac9',core_mask=RKNN.NPU_CORE_0,perf_debug=detailed_performance_log)
+# rknn.eval_perf()
+ret = rknn.accuracy_analysis(inputs=['img.npy'], target='rk3588')
+

model.rknn

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:563126adc90241146f3d736593c2f86f203d1f4b539c2f7b44f937828b69cabc
+size 204725760

run_rknn.py

@@ -0,0 +1,88 @@
+import os
+import sys
+import time
+import cv2
+import numpy as np
+import pandas as pd
+from rknnlite.api import RKNNLite
+
+pd.set_option("display.max_rows", 1000)
+
+dim = 448
+thresh = 0.5
+
+# 初始化RKNNLite
+rknn_lite = RKNNLite(verbose=False)
+
+# 加载RKNN模型
+ret = rknn_lite.load_rknn("model.rknn")
+if ret != 0:
+    print('加载RKNN模型失败')
+    exit(ret)
+
+# 初始化运行时环境
+ret = rknn_lite.init_runtime()
+if ret != 0:
+    print('初始化运行时环境失败')
+    exit(ret)
+
+label_names = pd.read_csv("selected_tags.csv")
+
+target_img = "input.jpg" if len(sys.argv) < 2 else sys.argv[1]
+
+try:
+    # 图像预处理
+    # 1. 读取图像并转RGB
+    img = cv2.imread(target_img)
+    if img is None:
+        print(f"无法读取图像: {target_img}")
+        exit(1)
+    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    
+    # 2. 处理透明通道
+    if img.shape[2] == 4:  # RGBA图像
+        img = img[:, :, :3]  # 只保留RGB通道
+    elif len(img.shape) == 2:  # 灰度图
+        img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
+    
+    # 3. 填充为正方形
+    h, w = img.shape[:2]
+    if h > w:
+        diff = h - w
+        pad_left = diff // 2
+        pad_right = diff - pad_left
+        img = cv2.copyMakeBorder(img, 0, 0, pad_left, pad_right, 
+                                cv2.BORDER_CONSTANT, value=(255, 255, 255))
+    elif w > h:
+        diff = w - h
+        pad_top = diff // 2
+        pad_bottom = diff - pad_top
+        img = cv2.copyMakeBorder(img, pad_top, pad_bottom, 0, 0, 
+                                cv2.BORDER_CONSTANT, value=(255, 255, 255))
+    
+    # 4. resize到目标尺寸
+    img = cv2.resize(img, (dim, dim), interpolation=cv2.INTER_AREA)
+    
+    # 5. 转换为float32并添加batch维度
+    img = img.astype(np.float32)
+    img = np.expand_dims(img, 0)
+    print(img.shape) # 1, 448, 448, 3
+    # NHWC -> NWCH
+    # img = np.transpose(img, (0, 2, 3, 1))
+    # 执行推理
+    # np.save("img.npy",img)
+    start = time.time()
+    outputs = rknn_lite.inference(inputs=[img], data_format="nhwc")
+    end = time.time()
+    print(f"Inference time: {end - start} seconds")
+    probs = outputs[0]  # 获取第一个输出
+    print(probs.shape)
+
+    # 后处理
+    label_names["probs"] = probs[0]
+    found_tags = label_names[label_names["probs"] > thresh][["tag_id", "name", "probs"]]
+    print(found_tags)
+
+finally:
+    # 释放资源
+    rknn_lite.release()