47 lines
1.7 KiB
TypeScript
47 lines
1.7 KiB
TypeScript
import cv from "@yizhi/cv";
|
|
import { ImageCropOption, ImageSource, Model } from "../common/model";
|
|
import { convertImage } from "../common/processors";
|
|
|
|
interface GenderAgePredictOption extends ImageCropOption {
|
|
}
|
|
|
|
export interface GenderAgePredictResult {
|
|
gender: "M" | "F"
|
|
age: number
|
|
}
|
|
|
|
const MODEL_URL_CONFIG = {
|
|
INSIGHT_GENDER_AGE_ONNX: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceattr/insight_gender_age.onnx`,
|
|
INSIGHT_GENDER_AGE_MNN: `https://www.modelscope.cn/models/luyizhi/basic_cv/resolve/master/faceattr/insight_gender_age.mnn`,
|
|
};
|
|
|
|
export class GenderAge extends Model {
|
|
public static async load(type?: keyof typeof MODEL_URL_CONFIG) {
|
|
return this.cacheModel(MODEL_URL_CONFIG[type ?? "INSIGHT_GENDER_AGE_ONNX"], { createModel: true }).then(r => r.model);
|
|
}
|
|
|
|
public predict(image: ImageSource, option?: GenderAgePredictOption) { return Model.resolveImage(image, im => this.doPredict(im, option)); }
|
|
|
|
private async doPredict(image: cv.Mat, option?: GenderAgePredictOption): Promise<GenderAgePredictResult> {
|
|
const input = this.input;
|
|
const output = this.output;
|
|
if (option?.crop) image = cv.crop(image, option.crop);
|
|
image = cv.resize(image, input.shape[3], input.shape[2]);
|
|
|
|
const nchwImage = convertImage(image.data, { sourceImageFormat: "bgr", targetColorFormat: "rgb", targetShapeFormat: "nchw", targetNormalize: { mean: [0], std: [1] } });
|
|
|
|
const result = await this.session.run({
|
|
[input.name]: {
|
|
shape: [1, 3, input.shape[2], input.shape[3]],
|
|
data: nchwImage,
|
|
type: "float32",
|
|
}
|
|
}).then(res => res[output.name].data);
|
|
|
|
return {
|
|
gender: result[0] > result[1] ? "F" : "M",
|
|
age: parseInt(result[2] * 100 as any),
|
|
}
|
|
}
|
|
}
|