103 lines
3.1 KiB
TypeScript
103 lines
3.1 KiB
TypeScript
import { cv } from "../../cv"
|
||
import { ImageSource, Model } from "../common/model"
|
||
|
||
interface IFaceBoxConstructorOption {
|
||
x1: number
|
||
y1: number
|
||
x2: number
|
||
y2: number
|
||
score: number
|
||
imw: number
|
||
imh: number
|
||
}
|
||
|
||
export class FaceBox {
|
||
#option: IFaceBoxConstructorOption;
|
||
constructor(option: IFaceBoxConstructorOption) { this.#option = option; }
|
||
|
||
public get x1() { return this.#option.x1; }
|
||
public get y1() { return this.#option.y1; }
|
||
public get x2() { return this.#option.x2; }
|
||
public get y2() { return this.#option.y2; }
|
||
public get centerX() { return this.x1 + this.width / 2; }
|
||
public get centerY() { return this.y1 + this.height / 2; }
|
||
public get score() { return this.#option.score; }
|
||
public get left() { return this.x1; }
|
||
public get top() { return this.y1; }
|
||
public get width() { return this.x2 - this.x1; }
|
||
public get height() { return this.y2 - this.y1; }
|
||
|
||
/** 转换成整数 */
|
||
public toInt() {
|
||
return new FaceBox({
|
||
...this.#option,
|
||
x1: parseInt(this.#option.x1 as any), y1: parseInt(this.#option.y1 as any),
|
||
x2: parseInt(this.#option.x2 as any), y2: parseInt(this.#option.y2 as any),
|
||
});
|
||
}
|
||
|
||
/** 转换成正方形 */
|
||
public toSquare() {
|
||
const { imw, imh } = this.#option;
|
||
let size = Math.max(this.width, this.height) / 2;
|
||
const cx = this.centerX, cy = this.centerY;
|
||
|
||
if (cx - size < 0) size = cx;
|
||
if (cx + size > imw) size = imw - cx;
|
||
if (cy - size < 0) size = cy;
|
||
if (cy + size > imh) size = imh - cy;
|
||
|
||
return new FaceBox({
|
||
...this.#option,
|
||
x1: this.centerX - size, y1: this.centerY - size,
|
||
x2: this.centerX + size, y2: this.centerY + size,
|
||
});
|
||
}
|
||
}
|
||
|
||
export interface FaceDetectOption {
|
||
/** 阈值，默认0.5 */
|
||
threshold?: number
|
||
/** MNS阈值，默认0.3 */
|
||
mnsThreshold?: number
|
||
}
|
||
|
||
export abstract class FaceDetector extends Model {
|
||
|
||
protected abstract doPredict(image: cv.Mat, option?: FaceDetectOption): Promise<FaceBox[]>;
|
||
|
||
public async predict(image: ImageSource, option?: FaceDetectOption) { return Model.resolveImage(image, im => this.doPredict(im, option)); }
|
||
}
|
||
|
||
|
||
export function nms(input: FaceBox[], threshold: number) {
|
||
if (!input.length) return [];
|
||
input = input.sort((a, b) => b.score - a.score);
|
||
const merged = input.map(() => 0);
|
||
const output: FaceBox[] = [];
|
||
|
||
for (let i = 0; i < input.length; i++) {
|
||
if (merged[i]) continue;
|
||
output.push(input[i]);
|
||
|
||
for (let j = i + 1; j < input.length; j++) {
|
||
if (merged[j]) continue;
|
||
const inner_x0 = input[i].x1 > input[j].x1 ? input[i].x1 : input[j].x1;
|
||
const inner_y0 = input[i].y1 > input[j].y1 ? input[i].y1 : input[j].y1;
|
||
const inner_x1 = input[i].x2 < input[j].x2 ? input[i].x2 : input[j].x2;
|
||
const inner_y1 = input[i].y2 < input[j].y2 ? input[i].y2 : input[j].y2;
|
||
const inner_h = inner_y1 - inner_y0 + 1;
|
||
const inner_w = inner_x1 - inner_x0 + 1;
|
||
if (inner_h <= 0 || inner_w <= 0) continue;
|
||
const inner_area = inner_h * inner_w;
|
||
const h1 = input[j].y2 - input[j].y1 + 1;
|
||
const w1 = input[j].x2 - input[j].x1 + 1;
|
||
const area1 = h1 * w1;
|
||
const score = inner_area / area1;
|
||
if (score > threshold) merged[j] = 1;
|
||
}
|
||
}
|
||
|
||
return output;
|
||
}
|