leon 4 tuần trước cách đây
mục cha
commit
4afa039aa9
1 tập tin đã thay đổi với 4 bổ sung2 xóa
  1. 4 2
      src/infer/trt/yolo/yolo.cu

+ 4 - 2
src/infer/trt/yolo/yolo.cu

@@ -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_);
         }
         
     }