针对mnn模型处理,输出增加shape
This commit is contained in:
@ -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>;
|
||||
|
@ -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";
|
@ -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);
|
||||
}))
|
||||
}
|
||||
|
@ -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);
|
||||
}));
|
||||
}
|
||||
|
@ -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; }
|
||||
|
@ -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) {
|
||||
|
@ -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) {
|
||||
|
@ -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",
|
||||
|
@ -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);
|
||||
|
@ -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);
|
||||
}
|
||||
|
@ -44,7 +44,7 @@ export class Insightface extends FaceRecognition {
|
||||
}
|
||||
}).then(res => res[output.name]);
|
||||
|
||||
return new Array(...embedding);
|
||||
return new Array(...embedding.data);
|
||||
}
|
||||
|
||||
}
|
||||
|
11
src/test.ts
11
src/test.ts
@ -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();
|
||||
}
|
||||
|
||||
|
Reference in New Issue
Block a user