针对mnn模型处理,输出增加shape

This commit is contained in:
2025-03-10 17:49:53 +08:00
parent e362890c96
commit a966b82963
13 changed files with 51 additions and 29 deletions

View File

@ -2,7 +2,7 @@
export interface SessionNodeInfo {
name: string
type: number
type: DataType
shape: number[]
}
@ -30,8 +30,13 @@ export interface SessionRunInputOption {
shape?: number[]
}
export interface SessionRunOutput {
shape: number[]
data: Float32Array
}
export abstract class CommonSession {
public abstract run(inputs: Record<string, SessionNodeData | SessionRunInputOption>): Promise<Record<string, Float32Array>>
public abstract run(inputs: Record<string, SessionNodeData | SessionRunInputOption>): Promise<Record<string, SessionRunOutput>>
public abstract get inputs(): Record<string, SessionNodeInfo>;
public abstract get outputs(): Record<string, SessionNodeInfo>;

View File

@ -1,3 +1,3 @@
export * as common from "./common";
export { SessionNodeInfo, DataTypeString, DataType, SessionNodeData, SessionRunInputOption, SessionRunOutput, CommonSession } from "./common";
export * as ort from "./ort";
export * as mnn from "./mnn";

View File

@ -1,4 +1,4 @@
import { CommonSession, dataTypeFrom, isTypedArray, SessionNodeData, SessionNodeInfo, SessionRunInputOption } from "../common";
import { CommonSession, dataTypeFrom, isTypedArray, SessionNodeData, SessionNodeInfo, SessionRunInputOption, SessionRunOutput } from "../common";
export class MNNSession extends CommonSession {
#session: any
@ -11,16 +11,19 @@ export class MNNSession extends CommonSession {
this.#session = new addon.MNNSession(modelData);
}
public run(inputs: Record<string, SessionNodeData | SessionRunInputOption>): Promise<Record<string, Float32Array>> {
public run(inputs: Record<string, SessionNodeData | SessionRunInputOption>): Promise<Record<string, SessionRunOutput>> {
const inputArgs: Record<string, any> = {};
for (const [name, option] of Object.entries(inputs)) {
if (isTypedArray(option)) inputArgs[name] = { data: option }
else inputArgs[name] = { ...option, type: option.type ? dataTypeFrom(option.type) : undefined };
}
return new Promise((resolve, reject) => this.#session.Run(inputArgs, (err: any, res: any) => {
return new Promise((resolve, reject) => this.#session.Run(inputArgs, (err: any, res: Record<string, { data: ArrayBuffer, shape: number[] }>) => {
if (err) return reject(err);
const result: Record<string, Float32Array> = {};
for (const [name, val] of Object.entries(res)) result[name] = new Float32Array(val as ArrayBuffer);
const result: Record<string, SessionRunOutput> = {};
for (const [name, val] of Object.entries(res)) result[name] = {
shape: val.shape,
data: new Float32Array(val.data),
}
resolve(result);
}))
}

View File

@ -1,4 +1,4 @@
import { CommonSession, dataTypeFrom, isTypedArray, SessionNodeData, SessionNodeInfo, SessionRunInputOption } from "../common";
import { CommonSession, dataTypeFrom, isTypedArray, SessionNodeData, SessionNodeInfo, SessionRunInputOption, SessionRunOutput } from "../common";
export class OrtSession extends CommonSession {
#session: any;
@ -22,10 +22,10 @@ export class OrtSession extends CommonSession {
else inputArgs[name] = { ...option, type: option.type ? dataTypeFrom(option.type) : undefined };
}
return new Promise<Record<string, Float32Array>>((resolve, reject) => this.#session.Run(inputArgs, (err: any, res: any) => {
return new Promise<Record<string, SessionRunOutput>>((resolve, reject) => this.#session.Run(inputArgs, (err: any, res: Record<string, ArrayBuffer>) => {
if (err) return reject(err);
const result: Record<string, Float32Array> = {};
for (const [name, val] of Object.entries(res)) result[name] = new Float32Array(val as ArrayBuffer);
const result: Record<string, SessionRunOutput> = {};
for (const [name, val] of Object.entries(res)) result[name] = { data: new Float32Array(val), shape: this.outputs[name].shape };
resolve(result);
}));
}

View File

@ -1,7 +1,7 @@
import { backend } from "../../backend";
import { cv } from "../../cv";
export type ModelConstructor<T> = new (session: backend.common.CommonSession) => T;
export type ModelConstructor<T> = new (session: backend.CommonSession) => T;
export type ImageSource = cv.Mat | Uint8Array | string;
@ -27,7 +27,7 @@ export interface ModelCacheResult<T, Create extends boolean> {
const cacheDownloadTasks: Array<{ url: string, cacheDir: string, modelPath: Promise<string> }> = [];
export abstract class Model {
protected session: backend.common.CommonSession;
protected session: backend.CommonSession;
protected static async resolveImage<R>(image: ImageSource, resolver: (image: cv.Mat) => R | Promise<R>): Promise<R> {
if (typeof image === "string") {
@ -148,7 +148,7 @@ export abstract class Model {
return { modelPath, modelType, model: model as any }
}
public constructor(session: backend.common.CommonSession) { this.session = session; }
public constructor(session: backend.CommonSession) { this.session = session; }
public get inputs() { return this.session.inputs; }
public get outputs() { return this.session.outputs; }

View File

@ -46,7 +46,7 @@ export class FaceLandmark1000 extends Model {
data: nchwImageData,
type: "float32",
}
}).then(res => res[this.output.name]);
}).then(res => res[this.output.name].data);
const points: FacePoint[] = [];
for (let i = 0; i < res.length; i += 2) {

View File

@ -51,7 +51,7 @@ export class PFLD extends Model {
shape: [1, 3, input.shape[2], input.shape[3]],
}
});
const pointsBuffer = res[pointsOutput.name];
const pointsBuffer = res[pointsOutput.name].data;
const points: FacePoint[] = [];
for (let i = 0; i < pointsBuffer.length; i += 2) {

View File

@ -36,7 +36,7 @@ export class GenderAge extends Model {
data: nchwImage,
type: "float32",
}
}).then(res => res[output.name]);
}).then(res => res[output.name].data);
return {
gender: result[0] > result[1] ? "F" : "M",

View File

@ -27,7 +27,7 @@ export class Yolov5Face extends FaceDetector {
const threshold = option?.threshold ?? 0.5;
for (let i = 0; i < outShape[1]; i++) {
const beg = i * outShape[2];
const rectData = outputData.slice(beg, beg + outShape[2]);
const rectData = outputData.data.slice(beg, beg + outShape[2]);
const x = parseInt(rectData[0] * ratioWidth as any);
const y = parseInt(rectData[1] * ratioHeight as any);
const w = parseInt(rectData[2] * ratioWidth as any);

View File

@ -28,7 +28,7 @@ export class AdaFace extends FaceRecognition {
data: nchwImageData,
shape: [1, 3, input.shape[2], input.shape[3]],
}
}).then(res => res[output.name]);
}).then(res => res[output.name].data);
return new Array(...embedding);
}

View File

@ -44,7 +44,7 @@ export class Insightface extends FaceRecognition {
}
}).then(res => res[output.name]);
return new Array(...embedding);
return new Array(...embedding.data);
}
}

View File

@ -115,10 +115,13 @@ async function testFaceAlign() {
const fd = await deploy.facedet.Yolov5Face.load("YOLOV5S_MNN");
const fa = await deploy.facealign.PFLD.load("PFLD_106_LITE_MNN");
// const fa = await deploy.facealign.FaceLandmark1000.load("FACELANDMARK1000_ONNX");
let image = await cv.Mat.load("https://bkimg.cdn.bcebos.com/pic/d52a2834349b033b5bb5f183119c21d3d539b6001712");
let image = await cv.Mat.load("https://i0.hdslb.com/bfs/archive/64e47ec9fdac9e24bc2b49b5aaad5560da1bfe3e.jpg");
image = image.rotate(image.width / 2, image.height / 2, 0);
const face = await fd.predict(image).then(res => res[0].toSquare());
const face = await fd.predict(image).then(res => {
console.log(res);
return res[0].toSquare()
});
const points = await fa.predict(image, { crop: { sx: face.left, sy: face.top, sw: face.width, sh: face.height } });
points.points.forEach((point, idx) => {
@ -137,8 +140,8 @@ async function testFaceAlign() {
}
async function test() {
await testGenderTest();
await testFaceID();
// await testGenderTest();
// await testFaceID();
await testFaceAlign();
}