Health/Assets/Scripts/PoseCheck/EstimateModel/OpenPoseEsimater.cs

203 lines
7.1 KiB
C#

using OpenCVForUnity.CoreModule;
using OpenCVForUnity.DnnModule;
using OpenCVForUnity.ImgprocModule;
using OpenCVForUnity.UnityUtils;
using System;
using System.Collections.Generic;
using UnityEngine;
using UnityEngine.UIElements;
using Yoga;
public class OpenPoseEsimater : Estimator
{
KeypointsModel _openPoseModel;
private Net _net;
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;
private double threshold = 0.5;
public override void InitModel()
{
_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*/);
_net = null;
var modelFilePath = Utils.getFilePath(YogaConfig.MODEL_PATHS[ModelType.OpenPose]);
if (string.IsNullOrEmpty(modelFilePath))
{
Debug.LogError("modelFilePath is empty. Please copy from “OpenCVForUnity/StreamingAssets/” to “Assets/StreamingAssets/” folder. ");
return;
}
_net = Dnn.readNet(modelFilePath);
_openPoseModel = new KeypointsModel(_net);
_openPoseModel.setInputScale(new Scalar(YogaConfig.InScale));
_openPoseModel.setInputSize(new Size(YogaConfig.InWidth, YogaConfig.InHeight));
_openPoseModel.setInputMean(new Scalar(YogaConfig.InMean));
_openPoseModel.setInputSwapRB(false);
_openPoseModel.setInputCrop(false);
}
public override bool Esitmate(Mat bgrMat, Mat rgbaMat, out List<Point> points)
{
points = null;
if (_objectDetector == null)
{
Debug.LogWarning("ObjectDetector is not ready. ");
return false; //重新检测
}
//volo 人物检测
//获取数据
Mat results = _objectDetector.infer(bgrMat);
var voloResultBox = new List<float[]>();
bool hasValidObject = false;
List<List<float[]>> boxList = new();
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)) //不符合规则的,跳过
continue;
hasValidObject = true;
List<float[]> item= new List<float[]>();
item.Add(box);
item.Add(conf);
item.Add(cls);
boxList.Add(item);
}
float area = 0;
List<float[]> maxSizeBox = new List<float[]>();
//选择最大的一个
foreach(var box in boxList)
{
var rect = box[0];
var tmp = (rect[2] - rect[0]) * (rect[3] - rect[1]);
area = Math.Max(area, tmp);
if (area == tmp)
{
maxSizeBox = box;
}
}
voloResultBox.Clear();
voloResultBox = maxSizeBox;
if (!hasValidObject) //没有检测到人体
{
Debug.Log("No person block found. Re-Estimation.");
return false; //重新检测
}
//更新人物框检测结果
YogaManager.Instance.VoloResult = voloResultBox;
points = null;
if (_openPoseModel == null)
throw new ArgumentException("CVEstimator.Init: args[1] is not KeypointsModel")
;
OpenCVForUnity.CoreModule.Rect roiRect = new OpenCVForUnity.CoreModule.Rect(
(int)voloResultBox[0][0],
(int)voloResultBox[0][1],
Math.Abs((int)(voloResultBox[0][2] - voloResultBox[0][0])),
Math.Abs((int)(voloResultBox[0][3] - voloResultBox[0][1])));
if (roiRect.y < 0 || //0 <= _rowRange.start
(roiRect.y + roiRect.height) < roiRect.y || // _rowRange.start <= _rowRange.end
bgrMat.rows() < (roiRect.y + roiRect.height)) //_rowRange.end <= m.rows
return false; //重新检测
Mat personRectImg = new Mat(bgrMat, roiRect);//获取人体区域
points = _openPoseModel.estimate(personRectImg, (float)threshold).toList();
//将人体区域的坐标转换为原图坐标
for (int j = 0; j < points.Count; j++)
{
if (points[j] == null ||
(points[j].x == -1 && points[j].y == -1)) //没找到的点,跳过
continue;
points[j].x += roiRect.x;
points[j].y += roiRect.y;
}
if (!YogaManager.Instance.ActionCheckPoints(points))
{
Debug.Log("ActionCheckPoints failed. Re-Estimation.");
return false; //重新检测
}
return true;
}
public override void DisposeModel()
{
if (_net != null)
_net.Dispose();
}
public override void Check(ref Mat img)
{
if (_net == null)
{
Imgproc.putText(img, "model file is not loaded.", new Point(5, img.rows() - 30), Imgproc.FONT_HERSHEY_SIMPLEX, 0.7, new Scalar(255, 255, 255), 2, Imgproc.LINE_AA, false);
Imgproc.putText(img, "Please read console message.", new Point(5, img.rows() - 10), Imgproc.FONT_HERSHEY_SIMPLEX, 0.7, new Scalar(255, 255, 255), 2, Imgproc.LINE_AA, false);
}
}
public override void DebugPrint(Mat image, bool isRGB = false)
{
if (YogaManager.Instance.VoloResult.Count > 2)
return;
var box = YogaManager.Instance.VoloResult[0];
var conf = YogaManager.Instance.VoloResult[1];
var cls = YogaManager.Instance.VoloResult[2];
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);
}
}