Skip to content

Commit

Permalink
infer num_classes for yolov8-seg pipeline (#1261)
Browse files Browse the repository at this point in the history
non default num_classes not currently supported. this PR adds a quick patch to infer based on the scores dimension
  • Loading branch information
bfineran committed Sep 25, 2023
1 parent 098d1db commit 9a4c5cc
Showing 1 changed file with 7 additions and 4 deletions.
11 changes: 7 additions & 4 deletions src/deepsparse/yolov8/pipelines.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,24 +108,27 @@ def process_engine_outputs(
)

def process_engine_outputs_seg(
self, engine_outputs: List[numpy.ndarray], nc: int = 80, **kwargs
self,
engine_outputs: List[numpy.ndarray],
**kwargs,
) -> YOLOSegOutput:
"""
The pathway for processing the outputs of the engine for YOLOv8 segmentation.
:param engine_outputs: list of numpy arrays that are the output of the engine
forward pass
: params nc: number of classes. If not provided, calculated as
detection.shape[1] - 4
:return: outputs of engine post-processed into an object in the `output_schema`
format of this pipeline
"""

detections, mask_protos = engine_outputs

# defined per ultralytics documentation
num_classes = detections.shape[1] - 4

# NMS
detections_output = self.nms_function(
outputs=detections,
nc=nc,
nc=num_classes,
iou_thres=kwargs.get("iou_thres", 0.25),
conf_thres=kwargs.get("conf_thres", 0.45),
multi_label=kwargs.get("multi_label", False),
Expand Down

0 comments on commit 9a4c5cc

Please sign in to comment.