@@ -154,12 +154,14 @@ void ModelYolo::prepareInputsOutputs(std::shared_ptr<ov::Model>& model) {
154
154
for (auto & out : outputs) {
155
155
ppp.output (out.get_any_name ()).tensor ().set_element_type (ov::element::f32 );
156
156
if (out.get_shape ().size () == 4 ) {
157
- if (out.get_shape ()[ov::layout::height_idx (yoloRegionLayout)] !=
158
- out.get_shape ()[ov::layout::width_idx (yoloRegionLayout)] &&
159
- out.get_shape ()[ov::layout::height_idx ({" NHWC" })] == out.get_shape ()[ov::layout::width_idx ({" NHWC" })]) {
160
- yoloRegionLayout = {" NHWC" };
157
+ if (out.get_shape ()[ov::layout::height_idx (" NCHW" )] != out.get_shape ()[ov::layout::width_idx (" NCHW" )] &&
158
+ out.get_shape ()[ov::layout::height_idx (" NHWC" )] == out.get_shape ()[ov::layout::width_idx (" NHWC" )]) {
159
+ ppp.output (out.get_any_name ()).model ().set_layout (" NHWC" );
160
+ // outShapes are saved before ppp.build() thus set yoloRegionLayout as it is in model before ppp.build()
161
+ yoloRegionLayout = " NHWC" ;
161
162
}
162
- ppp.output (out.get_any_name ()).tensor ().set_layout (yoloRegionLayout);
163
+ // yolo-v1-tiny-tf out shape is [1, 21125] thus set layout only for 4 dim tensors
164
+ ppp.output (out.get_any_name ()).tensor ().set_layout (" NCHW" );
163
165
}
164
166
outputsNames.push_back (out.get_any_name ());
165
167
outShapes[out.get_any_name ()] = out.get_shape ();
@@ -223,7 +225,7 @@ void ModelYolo::prepareInputsOutputs(std::shared_ptr<ov::Model>& model) {
223
225
for (const auto & name : outputsNames) {
224
226
const auto & shape = outShapes[name];
225
227
if (shape[ov::layout::channels_idx (yoloRegionLayout)] % num != 0 ) {
226
- throw std::logic_error (std::string (" Output tenosor " ) + name + " has wrong 2nd dimension" );
228
+ throw std::logic_error (std::string (" Output tensor " ) + name + " has wrong channel dimension" );
227
229
}
228
230
regions.emplace (
229
231
name,
@@ -328,8 +330,8 @@ void ModelYolo::parseYOLOOutput(const std::string& output_name,
328
330
case YOLO_V4:
329
331
case YOLO_V4_TINY:
330
332
case YOLOF:
331
- sideH = static_cast <int >(tensor.get_shape ()[ov::layout::height_idx (yoloRegionLayout )]);
332
- sideW = static_cast <int >(tensor.get_shape ()[ov::layout::width_idx (yoloRegionLayout )]);
333
+ sideH = static_cast <int >(tensor.get_shape ()[ov::layout::height_idx (" NCHW " )]);
334
+ sideW = static_cast <int >(tensor.get_shape ()[ov::layout::width_idx (" NCHW " )]);
333
335
scaleW = resized_im_w;
334
336
scaleH = resized_im_h;
335
337
break ;
0 commit comments