135 lines
4.3 KiB
C#
135 lines
4.3 KiB
C#
using OpenCVForUnity.CoreModule;
|
||
using OpenCVForUnity.ImgprocModule;
|
||
using OpenCVForUnity.UnityUtils;
|
||
using OpenCVForUnity.UnityUtils.Helper;
|
||
using System;
|
||
using System.Collections.Generic;
|
||
using UnityEngine;
|
||
using Yoga;
|
||
|
||
public abstract class Estimator
|
||
{
|
||
private YOLOv7ObjectDetector _objectDetector;
|
||
public int inpWidth = 416;
|
||
public int inpHeight = 416;
|
||
public float confThreshold = 0.35f;
|
||
public float nmsThreshold = 0.6f;
|
||
public int topK = 1000;
|
||
|
||
//初始化
|
||
public void Init()
|
||
{
|
||
_objectDetector = new YOLOv7ObjectDetector(
|
||
Utils.getFilePath("OpenCVForUnity/dnn/yolov7-tiny.weights"),
|
||
Utils.getFilePath("OpenCVForUnity/dnn/yolov7-tiny.cfg"),
|
||
Utils.getFilePath("OpenCVForUnity/dnn/coco.names"),
|
||
new Size(inpWidth, inpHeight), confThreshold, nmsThreshold/*, topK*/);
|
||
|
||
InitModel();
|
||
}
|
||
|
||
public abstract void InitModel();
|
||
|
||
public void Dispose()
|
||
{
|
||
DisposeModel();
|
||
}
|
||
public abstract void DisposeModel();
|
||
|
||
//检测到的人体框 默认使用VOLO7
|
||
public virtual bool TryGetROIRegionMat(Mat rgbaMat, Mat bgrMat, out List<float[]> voloResult)
|
||
{
|
||
voloResult = null;
|
||
if (_objectDetector == null)
|
||
{
|
||
Debug.LogWarning("ObjectDetector is not ready. ");
|
||
return false; //重新检测
|
||
}
|
||
|
||
//获取数据
|
||
Mat results = _objectDetector.infer(bgrMat);
|
||
var voloResultBox = new List<float[]>();
|
||
bool hasValidObject = false;
|
||
for (int i = results.rows() - 1; i >= 0; --i)
|
||
{
|
||
float[] box = new float[4];
|
||
results.get(i, 0, box); //方框
|
||
float[] conf = new float[1];
|
||
results.get(i, 4, conf); //检测数据
|
||
float[] cls = new float[1];
|
||
results.get(i, 5, cls); //类别
|
||
|
||
if (!IsObjectValid(box, conf, cls, rgbaMat))
|
||
{
|
||
return false;
|
||
}
|
||
hasValidObject = true;
|
||
voloResultBox.Clear();
|
||
voloResultBox.Add(box);
|
||
voloResultBox.Add(conf);
|
||
voloResultBox.Add(cls);
|
||
break;
|
||
}
|
||
|
||
if (!hasValidObject) //没有检测到人体
|
||
{
|
||
Debug.Log("No person block found. Re-Estimation.");
|
||
return false; //重新检测
|
||
}
|
||
|
||
voloResult = voloResultBox;
|
||
return true;
|
||
}
|
||
|
||
|
||
protected bool IsObjectValid(float[] box, float[] confidence, float[] classID, Mat rgbaMat)
|
||
{
|
||
if ((int)classID[0] != 0 || confidence[0] < 0.8f) //只检测人体，且置信度大于80%
|
||
return false;
|
||
|
||
//是否满足主驾/副驾的位置条件
|
||
float width = rgbaMat.width();
|
||
//获取矩形的中点
|
||
float centerX = (box[0] + box[2]) / 2;
|
||
//左边副驾驶，右边主驾驶
|
||
|
||
return true;
|
||
}
|
||
|
||
public void DebugPrintObjectLayout(Mat image, float[] box, float[] conf, float[] cls, bool isRGB = false)
|
||
{
|
||
|
||
float left = box[0];
|
||
float top = box[1];
|
||
float right = box[2];
|
||
float bottom = box[3];
|
||
int classId = (int)cls[0];
|
||
|
||
Scalar c = _objectDetector.palette[classId % _objectDetector.palette.Count];
|
||
Scalar color = isRGB ? c : new Scalar(c.val[2], c.val[1], c.val[0], c.val[3]);
|
||
|
||
Imgproc.rectangle(image, new Point(left, top), new Point(right, bottom), color, 2);
|
||
|
||
string label = String.Format("{0:0.00}", conf[0]);
|
||
if (_objectDetector.classNames != null && _objectDetector.classNames.Count != 0)
|
||
{
|
||
if (classId < (int)_objectDetector.classNames.Count)
|
||
{
|
||
label = _objectDetector.classNames[classId] + " " + label;
|
||
}
|
||
}
|
||
|
||
int[] baseLine = new int[1];
|
||
Size labelSize = Imgproc.getTextSize(label, Imgproc.FONT_HERSHEY_SIMPLEX, 0.5, 1, baseLine);
|
||
|
||
top = Mathf.Max((float)top, (float)labelSize.height);
|
||
Imgproc.rectangle(image, new Point(left, top - labelSize.height),
|
||
new Point(left + labelSize.width, top + baseLine[0]), color, Core.FILLED);
|
||
Imgproc.putText(image, label, new Point(left, top), Imgproc.FONT_HERSHEY_SIMPLEX, 0.5, Scalar.all(255), 1, Imgproc.LINE_AA);
|
||
}
|
||
|
||
public abstract bool Esitmate(Mat bgrMat, List<float[]> voloResultBox, out List<Point> points);
|
||
|
||
public virtual void Check(ref Mat img) { }
|
||
}
|