99 lines
3.3 KiB
C#
99 lines
3.3 KiB
C#
// Copyright (c) 2021 homuler
|
|
//
|
|
// Use of this source code is governed by an MIT-style
|
|
// license that can be found in the LICENSE file or at
|
|
// https://opensource.org/licenses/MIT.
|
|
|
|
using Mediapipe.Unity.CoordinateSystem;
|
|
using UnityEngine;
|
|
|
|
namespace Mediapipe.Unity
|
|
{
|
|
#pragma warning disable IDE0065
|
|
using Color = UnityEngine.Color;
|
|
#pragma warning restore IDE0065
|
|
|
|
public sealed class DetectionAnnotation : HierarchicalAnnotation
|
|
{
|
|
[SerializeField] private RectangleAnnotation _locationDataAnnotation;
|
|
[SerializeField] private PointListAnnotation _relativeKeypointsAnnotation;
|
|
[SerializeField] private LabelAnnotation _labelAnnotation;
|
|
|
|
public override bool isMirrored
|
|
{
|
|
set
|
|
{
|
|
_locationDataAnnotation.isMirrored = value;
|
|
_relativeKeypointsAnnotation.isMirrored = value;
|
|
_labelAnnotation.isMirrored = value;
|
|
base.isMirrored = value;
|
|
}
|
|
}
|
|
|
|
public override RotationAngle rotationAngle
|
|
{
|
|
set
|
|
{
|
|
_locationDataAnnotation.rotationAngle = value;
|
|
_relativeKeypointsAnnotation.rotationAngle = value;
|
|
_labelAnnotation.rotationAngle = value;
|
|
base.rotationAngle = value;
|
|
}
|
|
}
|
|
|
|
public void SetLineWidth(float lineWidth)
|
|
{
|
|
_locationDataAnnotation.SetLineWidth(lineWidth);
|
|
}
|
|
|
|
public void SetKeypointRadius(float radius)
|
|
{
|
|
_relativeKeypointsAnnotation.SetRadius(radius);
|
|
}
|
|
|
|
/// <param name="threshold">
|
|
/// Score threshold. This value must be between 0 and 1.
|
|
/// This will affect the rectangle's color. For example, if the score is below the threshold, the rectangle will be transparent.
|
|
/// The default value is 0.
|
|
/// </param>
|
|
public void Draw(Detection target, float threshold = 0.0f)
|
|
{
|
|
if (ActivateFor(target))
|
|
{
|
|
var score = target.Score.Count > 0 ? target.Score[0] : 1.0f;
|
|
var color = GetColor(score, Mathf.Clamp(threshold, 0.0f, 1.0f));
|
|
|
|
// Assume that location data's format is always RelativeBoundingBox
|
|
// TODO: fix if there are cases where this assumption is not correct.
|
|
var rectVertices = GetScreenRect().GetRectVertices(target.LocationData.RelativeBoundingBox, rotationAngle, isMirrored);
|
|
_locationDataAnnotation.SetColor(GetColor(score, Mathf.Clamp(threshold, 0.0f, 1.0f)));
|
|
_locationDataAnnotation.Draw(rectVertices);
|
|
|
|
var width = rectVertices[2].x - rectVertices[0].x;
|
|
var height = rectVertices[2].y - rectVertices[0].y;
|
|
var labelText = target.Label.Count > 0 ? target.Label[0] : null;
|
|
var vertexId = (((int)rotationAngle / 90) + 1) % 4;
|
|
var isInverted = ImageCoordinate.IsInverted(rotationAngle);
|
|
var (maxWidth, maxHeight) = isInverted ? (height, width) : (width, height);
|
|
_labelAnnotation.Draw(labelText, rectVertices[vertexId], color, maxWidth, maxHeight);
|
|
|
|
_relativeKeypointsAnnotation.Draw(target.LocationData.RelativeKeypoints);
|
|
}
|
|
}
|
|
|
|
private Color GetColor(float score, float threshold)
|
|
{
|
|
var t = (score - threshold) / (1 - threshold);
|
|
var h = Mathf.Lerp(90, 0, t) / 360; // from yellow-green to red
|
|
var color = Color.HSVToRGB(h, 1, 1);
|
|
|
|
if (t < 0)
|
|
{
|
|
// below the threshold
|
|
color.a = 0.5f;
|
|
}
|
|
return color;
|
|
}
|
|
}
|
|
}
|