Merge branch 'feature/damian/openpifpaf_field_output' into feature/da…

…mian/openpifpaf_val

src/deepsparse/yolo/pipelines.py

@@ -210,6 +210,7 @@ def process_inputs(
             multi_label=inputs.multi_label,
             original_image_shapes=original_image_shapes,
             return_masks=inputs.return_masks,
+            return_intermediate_outputs=inputs.return_intermediate_outputs,
         )
         return [image_batch], postprocessing_kwargs
 
@@ -309,6 +310,9 @@ def process_engine_outputs(
             boxes=batch_boxes,
             scores=batch_scores,
             labels=batch_labels,
+            intermediate_outputs=engine_outputs[0]
+            if kwargs.get("return_intermediate_outputs")
+            else None,
         )
 
     def _make_batch(self, image_batch: List[numpy.ndarray]) -> numpy.ndarray:

src/deepsparse/yolo/schemas.py

@@ -18,7 +18,7 @@
 """
 
 from collections import namedtuple
-from typing import Iterable, List, TextIO
+from typing import Any, Iterable, List, Optional, TextIO
 
 import numpy
 from PIL import Image
@@ -61,6 +61,11 @@ class YOLOInput(ComputerVisionSchema):
         description="Controls whether the pipeline should additionally "
         "return segmentation masks (if running a segmentation model)",
     )
+    return_intermediate_outputs: bool = Field(
+        default=False,
+        description="Controls whether the pipeline should additionally "
+        "return intermediate outputs from the model",
+    )
 
     @classmethod
     def from_files(
@@ -105,6 +110,10 @@ class YOLOOutput(BaseModel):
     labels: List[List[str]] = Field(
         description="List of labels, one for each prediction"
     )
+    intermediate_outputs: Optional[Any] = Field(
+        default=None,
+        description="Intermediate outputs from the YOLOv8 segmentation model.",
+    )
 
     def __getitem__(self, index):
         if index >= len(self.boxes):

src/deepsparse/yolov8/pipelines.py

@@ -131,7 +131,7 @@ def process_engine_outputs_seg(
             nm=nm,
             multi_label=kwargs.get("multi_label", False),
         )
-        detections_output = torch.stack(detections_output)
+
         mask_protos = numpy.stack(mask_protos)
         original_image_shapes = kwargs.get("original_image_shapes")
         batch_boxes, batch_scores, batch_labels, batch_masks = [], [], [], []
@@ -145,6 +145,15 @@ def process_engine_outputs_seg(
             )
 
             bboxes = detection_output[:, :4]
+
+            # check if empty detection
+            if bboxes.shape[0] == 0:
+                batch_boxes.append([None])
+                batch_scores.append([None])
+                batch_labels.append([None])
+                batch_masks.append([None])
+                continue
+
             bboxes = self._scale_boxes(bboxes, original_image_shape)
             scores = detection_output[:, 4]
             labels = detection_output[:, 5]
@@ -155,6 +164,7 @@ def process_engine_outputs_seg(
             batch_boxes.append(bboxes.tolist())
             batch_scores.append(scores.tolist())
             batch_labels.append(labels.tolist())
+
             batch_masks.append(
                 process_mask_upsample(
                     protos=protos,
@@ -179,4 +189,7 @@ def process_engine_outputs_seg(
             scores=batch_scores,
             classes=batch_labels,
             masks=batch_masks if kwargs.get("return_masks") else None,
+            intermediate_outputs=(detections, mask_protos)
+            if kwargs.get("return_intermediate_outputs")
+            else None,
         )

src/deepsparse/yolov8/schemas.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, List
+from typing import Any, List, Optional, Tuple
 
 from pydantic import BaseModel, Field
 
@@ -25,13 +25,23 @@ class YOLOSegOutput(BaseModel):
     Output model for YOLOv8 Segmentation model
     """
 
-    boxes: List[List[List[float]]] = Field(
+    boxes: List[List[Optional[List[float]]]] = Field(
         description="List of bounding boxes, one for each prediction"
     )
-    scores: List[List[float]] = Field(
+    scores: List[List[Optional[float]]] = Field(
         description="List of scores, one for each prediction"
     )
-    classes: List[List[str]] = Field(
+    classes: List[List[Optional[str]]] = Field(
         description="List of labels, one for each prediction"
     )
-    masks: List[Any] = Field(description="List of masks, one for each prediction")
+    masks: Optional[List[Any]] = Field(
+        description="List of masks, one for each prediction"
+    )
+
+    intermediate_outputs: Optional[Tuple[Any, Any]] = Field(
+        default=None,
+        description="A tuple that contains of intermediate outputs "
+        "from the YOLOv8 segmentation model. The tuple"
+        "contains two items: predictions from the model"
+        "and mask prototypes",
+    )