|
@@ -456,6 +456,8 @@ std::shared_ptr<data::InstanceSegmentMap> YoloModelImpl::decode_segment(int imem
|
|
|
|
|
|
int start_x = slice_->slice_start_point_.cpu()[batch_index*2];
|
|
|
int start_y = slice_->slice_start_point_.cpu()[batch_index*2+1];
|
|
|
+
|
|
|
+ int mask_dim = segment_head_dims_[1];
|
|
|
|
|
|
float *mask_weights = bbox_output_device +
|
|
|
(batch_index * bbox_head_dims_[1] + row_index) * bbox_head_dims_[2] +
|
|
@@ -578,7 +580,7 @@ bool YoloModelImpl::load(const std::string &engine_file, ModelType model_type, c
|
|
|
num_classes_ = bbox_head_dims_[2] - 4 - KEY_POINT_NUM * 3;
|
|
|
// NUM_BOX_ELEMENT = 8 + KEY_POINT_NUM * 3;
|
|
|
}
|
|
|
- else if (this->yolo_type_ == ModelType::YOLO11SEG)
|
|
|
+ else if (this->model_type_ == ModelType::YOLO11SEG)
|
|
|
{
|
|
|
num_classes_ = bbox_head_dims_[2] - 4 - segment_head_dims_[1];
|
|
|
}
|
|
@@ -714,7 +716,7 @@ Result YoloModelImpl::forwards(void *stream)
|
|
|
{
|
|
|
decode_kernel_invoker_v11pose(image_based_bbox_output, bbox_head_dims_[1], num_classes_,
|
|
|
bbox_head_dims_[2], confidence_threshold_, nms_threshold_,
|
|
|
- affine_matrix_device, boxarray_device, box_count, MAX_IMAGE_BOXES, start_x, start_y, stream_);
|
|
|
+ affine_matrix_device, boxarray_device, box_count, MAX_IMAGE_BOXES, start_x, start_y, ib, stream_);
|
|
|
}
|
|
|
|
|
|
}
|