COMPMID-1582 Add YOLOLayer to graph API

Change-Id: I8eb7cf220f038b0ac70cb3df1f399a3b66bbca1b
Reviewed-on: https://eu-gerrit-1.euhpc.arm.com/148245
Reviewed-by: Georgios Pinitas <georgios.pinitas@arm.com>
Tested-by: bsgcomp <bsgcomp@arm.com>

1 files changed, 26 insertions, 0 deletions

diff --git a/arm_compute/graph/frontend/Layers.h b/arm_compute/graph/frontend/Layers.h
index 9afb1787f3..bbfda9ca82 100644
--- a/arm_compute/graph/frontend/Layers.h
+++ b/arm_compute/graph/frontend/Layers.h
@@ -741,6 +741,32 @@ public:
 private:
     float _beta;
 };
+
+/** YOLO Layer */
+class YOLOLayer final : public ILayer
+{
+public:
+    /** Construct a YOLO layer.
+     *
+     * @param[in] act_info    Activation info
+     * @param[in] num_classes Number of classes to activate
+     */
+    YOLOLayer(ActivationLayerInfo act_info, int32_t num_classes)
+        : _act_info(act_info), _num_classes(num_classes)
+    {
+    }
+
+    NodeID create_layer(IStream &s) override
+    {
+        NodeParams  common_params = { name(), s.hints().target_hint };
+        NodeIdxPair input         = { s.tail_node(), 0 };
+        return GraphBuilder::add_yolo_node(s.graph(), common_params, input, _act_info, _num_classes);
+    }
+
+private:
+    ActivationLayerInfo _act_info;
+    int32_t             _num_classes;
+};
 } // namespace frontend
 } // namespace graph
 } // namespace arm_compute