221 lines
6.9 KiB
TypeScript
221 lines
6.9 KiB
TypeScript
import cv from '@techstark/opencv-js'
|
||
|
||
export interface TrackerParams {
|
||
lowerYellow: any;
|
||
upperYellow: any;
|
||
blurSize: number;
|
||
morphKernel: number;
|
||
minContourArea: number;
|
||
useRoi: boolean;
|
||
roiScaleX: number;
|
||
roiScaleY: number;
|
||
}
|
||
|
||
export interface Position {
|
||
x: number;
|
||
y: number;
|
||
}
|
||
|
||
export class YellowBallTracker {
|
||
public srcMat: cv.Mat | null = null
|
||
private hsvMat: cv.Mat | null = null
|
||
private blurredMat: cv.Mat | null = null
|
||
private maskMat: cv.Mat | null = null
|
||
private hierarchy: cv.Mat | null = null
|
||
private contours: cv.MatVector | null = null
|
||
private morphKernelMat: cv.Mat | null = null
|
||
private lowerBoundMat: cv.Mat | null = null
|
||
private upperBoundMat: cv.Mat | null = null
|
||
private params: TrackerParams
|
||
|
||
constructor(params: TrackerParams) {
|
||
this.params = params
|
||
}
|
||
|
||
public updateKernel() {
|
||
if (this.morphKernelMat) this.morphKernelMat.delete()
|
||
this.morphKernelMat = cv.Mat.ones(this.params.morphKernel, this.params.morphKernel, cv.CV_8U)
|
||
}
|
||
|
||
public initMats(w: number, h: number) {
|
||
this.cleanupMats()
|
||
|
||
this.srcMat = new cv.Mat(h, w, cv.CV_8UC4)
|
||
this.hsvMat = new cv.Mat()
|
||
this.blurredMat = new cv.Mat()
|
||
this.maskMat = new cv.Mat()
|
||
this.hierarchy = new cv.Mat()
|
||
this.contours = new cv.MatVector()
|
||
this.updateKernel()
|
||
// 创建阈值 Mat
|
||
this.lowerBoundMat = new cv.Mat(h, w, cv.CV_8UC3, this.params.lowerYellow)
|
||
this.upperBoundMat = new cv.Mat(h, w, cv.CV_8UC3, this.params.upperYellow)
|
||
}
|
||
|
||
private calculateROI(lastPos: Position): { x: number, y: number, width: number, height: number } {
|
||
const width = this.srcMat!.cols * this.params.roiScaleX
|
||
const height = this.srcMat!.rows * this.params.roiScaleY
|
||
const x = Math.max(0, Math.min(this.srcMat!.cols - width, lastPos.x - width / 2))
|
||
const y = Math.max(0, Math.min(this.srcMat!.rows - height, lastPos.y - height / 2))
|
||
|
||
return {
|
||
x: Math.round(x),
|
||
y: Math.round(y),
|
||
width: Math.round(width),
|
||
height: Math.round(height)
|
||
}
|
||
}
|
||
|
||
private drawROIDebugInfo(roi: { x: number, y: number, width: number, height: number }) {
|
||
cv.rectangle(this.srcMat!,
|
||
new cv.Point(roi.x, roi.y),
|
||
new cv.Point(roi.x + roi.width, roi.y + roi.height),
|
||
new cv.Scalar(255, 0, 0, 255), 2)
|
||
}
|
||
|
||
private processImage(inputMat: cv.Mat, lowerBound: cv.Mat, upperBound: cv.Mat) {
|
||
// 颜色空间转换和模糊处理
|
||
cv.cvtColor(inputMat, this.hsvMat!, cv.COLOR_RGBA2RGB)
|
||
cv.cvtColor(this.hsvMat!, this.hsvMat!, cv.COLOR_RGB2HSV)
|
||
cv.blur(this.hsvMat!, this.blurredMat!, new cv.Size(this.params.blurSize, this.params.blurSize))
|
||
|
||
// 颜色范围过滤
|
||
cv.inRange(this.blurredMat!, lowerBound, upperBound, this.maskMat!)
|
||
|
||
// 形态学操作
|
||
cv.morphologyEx(this.maskMat!, this.maskMat!, cv.MORPH_OPEN, this.morphKernelMat!)
|
||
cv.morphologyEx(this.maskMat!, this.maskMat!, cv.MORPH_CLOSE, this.morphKernelMat!)
|
||
cv.dilate(this.maskMat!, this.maskMat!, this.morphKernelMat!)
|
||
}
|
||
|
||
private findBestContour(offsetX: number = 0, offsetY: number = 0): Position | null {
|
||
// 查找轮廓
|
||
this.contours!.delete()
|
||
this.contours = new cv.MatVector()
|
||
cv.findContours(this.maskMat!, this.contours, this.hierarchy!, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)
|
||
|
||
// 找到最大面积的轮廓
|
||
let maxArea = 0
|
||
let best: Position | null = null
|
||
|
||
for (let i = 0; i < this.contours.size(); i++) {
|
||
const cnt = this.contours.get(i)
|
||
const area = cv.contourArea(cnt)
|
||
|
||
if (area > this.params.minContourArea && area > maxArea) {
|
||
maxArea = area
|
||
const circle = cv.minEnclosingCircle(cnt)
|
||
best = {
|
||
x: Math.round(circle.center.x + offsetX),
|
||
y: Math.round(circle.center.y + offsetY)
|
||
}
|
||
}
|
||
cnt.delete()
|
||
}
|
||
|
||
return best
|
||
}
|
||
|
||
private drawDetectionResult(position: Position) {
|
||
cv.circle(this.srcMat!, new cv.Point(position.x, position.y), 10, new cv.Scalar(0, 255, 0, 255), 2)
|
||
}
|
||
|
||
public detectYellowBall(imgData: ImageData, lastPos?: Position | null): Position | null {
|
||
if (!this.srcMat || !this.hsvMat || !this.blurredMat || !this.maskMat || !this.hierarchy ||
|
||
!this.contours || !this.morphKernelMat || !this.lowerBoundMat || !this.upperBoundMat) {
|
||
console.error('Mats not initialized')
|
||
return null
|
||
}
|
||
|
||
// Set image data
|
||
this.srcMat.data.set(imgData.data)
|
||
|
||
let best: Position | null = null
|
||
let srcRoiMat = new cv.Mat()
|
||
|
||
try {
|
||
// 尝试ROI检测
|
||
if (this.params.useRoi && lastPos) {
|
||
const roi = this.calculateROI(lastPos)
|
||
this.drawROIDebugInfo(roi)
|
||
|
||
// 提取ROI区域
|
||
const roiRect = new cv.Rect(roi.x, roi.y, roi.width, roi.height)
|
||
srcRoiMat = this.srcMat.roi(roiRect)
|
||
|
||
// 为ROI区域创建合适大小的阈值Mat
|
||
const roiLowerBound = new cv.Mat(roi.height, roi.width, cv.CV_8UC3, this.params.lowerYellow)
|
||
const roiUpperBound = new cv.Mat(roi.height, roi.width, cv.CV_8UC3, this.params.upperYellow)
|
||
|
||
try {
|
||
this.processImage(srcRoiMat, roiLowerBound, roiUpperBound)
|
||
best = this.findBestContour(roi.x, roi.y)
|
||
} finally {
|
||
roiLowerBound.delete()
|
||
roiUpperBound.delete()
|
||
}
|
||
}
|
||
|
||
// 如果ROI中没找到或者没启用ROI，就在全图搜索
|
||
if (!best) {
|
||
if (this.params.useRoi && lastPos) {
|
||
console.log("ROI失败，全域搜索")
|
||
}
|
||
|
||
this.processImage(this.srcMat, this.lowerBoundMat, this.upperBoundMat)
|
||
best = this.findBestContour()
|
||
}
|
||
|
||
// 在原图上绘制找到的点
|
||
if (best) {
|
||
this.drawDetectionResult(best)
|
||
}
|
||
|
||
} finally {
|
||
// 清理资源
|
||
srcRoiMat.delete()
|
||
}
|
||
|
||
return best
|
||
}
|
||
|
||
public getOutputCanvas(targetCanvas: HTMLCanvasElement) {
|
||
if (!this.srcMat) return null
|
||
cv.imshow(targetCanvas, this.srcMat)
|
||
return targetCanvas
|
||
}
|
||
|
||
public updateParams(params: Partial<TrackerParams>) {
|
||
this.params = { ...this.params, ...params }
|
||
|
||
// 如果更新了morphKernel，需要重新创建kernel
|
||
if ('morphKernel' in params) {
|
||
this.updateKernel()
|
||
}
|
||
}
|
||
|
||
public cleanupMats() {
|
||
// 清理资源
|
||
if (this.srcMat) this.srcMat.delete()
|
||
if (this.hsvMat) this.hsvMat.delete()
|
||
if (this.blurredMat) this.blurredMat.delete()
|
||
if (this.maskMat) this.maskMat.delete()
|
||
if (this.hierarchy) this.hierarchy.delete()
|
||
if (this.contours) this.contours.delete()
|
||
if (this.morphKernelMat) this.morphKernelMat.delete()
|
||
if (this.lowerBoundMat) this.lowerBoundMat.delete()
|
||
if (this.upperBoundMat) this.upperBoundMat.delete()
|
||
|
||
// 重置引用
|
||
this.srcMat = null
|
||
this.hsvMat = null
|
||
this.blurredMat = null
|
||
this.maskMat = null
|
||
this.hierarchy = null
|
||
this.contours = null
|
||
this.morphKernelMat = null
|
||
this.lowerBoundMat = null
|
||
this.upperBoundMat = null
|
||
}
|
||
}
|