Merge pull request #857 from roboflow/feat/model-block-id-output

Add model id output param
roboflow · Dec 11, 2024 · f5f4d6f · f5f4d6f
2 parents b727647 + ae6b983
commit f5f4d6f
Show file tree

Hide file tree

Showing 7 changed files with 130 additions and 45 deletions.
diff --git a/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v2.py b/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v2.py
@@ -156,6 +156,7 @@ def describe_outputs(cls) -> List[OutputDefinition]:
                 name="predictions",
                 kind=[INSTANCE_SEGMENTATION_PREDICTION_KIND],
             ),
+            OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]),
         ]
 
     @classmethod
@@ -281,6 +282,7 @@ def run_locally(
             images=images,
             predictions=predictions,
             class_filter=class_filter,
+            model_id=model_id,
         )
 
     def run_remotely(
@@ -336,13 +338,15 @@ def run_remotely(
             images=images,
             predictions=predictions,
             class_filter=class_filter,
+            model_id=model_id,
         )
 
     def _post_process_result(
         self,
         images: Batch[WorkflowImageData],
         predictions: List[dict],
         class_filter: Optional[List[str]],
+        model_id: str,
     ) -> BlockResult:
         inference_ids = [p.get(INFERENCE_ID_KEY, None) for p in predictions]
         predictions = convert_inference_detections_batch_to_sv_detections(predictions)
@@ -359,6 +363,10 @@ def _post_process_result(
             predictions=predictions,
         )
         return [
-            {"inference_id": inference_id, "predictions": prediction}
+            {
+                "inference_id": inference_id,
+                "predictions": prediction,
+                "model_id": model_id,
+            }
             for inference_id, prediction in zip(inference_ids, predictions)
         ]
diff --git a/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v2.py b/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v2.py
@@ -147,6 +147,7 @@ def describe_outputs(cls) -> List[OutputDefinition]:
             OutputDefinition(
                 name="predictions", kind=[KEYPOINT_DETECTION_PREDICTION_KIND]
             ),
+            OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]),
         ]
 
     @classmethod
@@ -267,6 +268,7 @@ def run_locally(
             images=images,
             predictions=predictions,
             class_filter=class_filter,
+            model_id=model_id,
         )
 
     def run_remotely(
@@ -320,13 +322,15 @@ def run_remotely(
             images=images,
             predictions=predictions,
             class_filter=class_filter,
+            model_id=model_id,
         )
 
     def _post_process_result(
         self,
         images: Batch[WorkflowImageData],
         predictions: List[dict],
         class_filter: Optional[List[str]],
+        model_id: str,
     ) -> BlockResult:
         inference_ids = [p.get(INFERENCE_ID_KEY, None) for p in predictions]
         detections = convert_inference_detections_batch_to_sv_detections(predictions)
@@ -348,6 +352,10 @@ def _post_process_result(
             predictions=detections,
         )
         return [
-            {"inference_id": inference_id, "predictions": image_detections}
+            {
+                "inference_id": inference_id,
+                "predictions": image_detections,
+                "model_id": model_id,
+            }
             for inference_id, image_detections in zip(inference_ids, detections)
         ]
diff --git a/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v2.py b/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v2.py
@@ -101,6 +101,7 @@ def describe_outputs(cls) -> List[OutputDefinition]:
         return [
             OutputDefinition(name="predictions", kind=[CLASSIFICATION_PREDICTION_KIND]),
             OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]),
+            OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]),
         ]
 
     @classmethod
@@ -191,6 +192,7 @@ def run_locally(
         return self._post_process_result(
             predictions=predictions,
             images=images,
+            model_id=model_id,
         )
 
     def run_remotely(
@@ -231,12 +233,14 @@ def run_remotely(
         return self._post_process_result(
             predictions=predictions,
             images=images,
+            model_id=model_id,
         )
 
     def _post_process_result(
         self,
         images: Batch[WorkflowImageData],
         predictions: List[dict],
+        model_id: str,
     ) -> BlockResult:
         predictions = attach_prediction_type_info(
             predictions=predictions,
@@ -251,6 +255,7 @@ def _post_process_result(
             {
                 "inference_id": prediction.get(INFERENCE_ID_KEY),
                 "predictions": prediction,
+                "model_id": model_id,
             }
             for prediction in predictions
         ]
diff --git a/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v2.py b/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v2.py
@@ -101,6 +101,7 @@ def describe_outputs(cls) -> List[OutputDefinition]:
         return [
             OutputDefinition(name="predictions", kind=[CLASSIFICATION_PREDICTION_KIND]),
             OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]),
+            OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]),
         ]
 
     @classmethod
@@ -191,6 +192,7 @@ def run_locally(
         return self._post_process_result(
             predictions=predictions,
             images=images,
+            model_id=model_id,
         )
 
     def run_remotely(
@@ -228,12 +230,15 @@ def run_remotely(
         )
         if not isinstance(predictions, list):
             predictions = [predictions]
-        return self._post_process_result(images=images, predictions=predictions)
+        return self._post_process_result(
+            images=images, predictions=predictions, model_id=model_id
+        )
 
     def _post_process_result(
         self,
         images: Batch[WorkflowImageData],
         predictions: List[dict],
+        model_id: str,
     ) -> List[dict]:
         predictions = attach_prediction_type_info(
             predictions=predictions,
@@ -248,6 +253,7 @@ def _post_process_result(
             {
                 "inference_id": prediction.get(INFERENCE_ID_KEY),
                 "predictions": prediction,
+                "model_id": model_id,
             }
             for prediction in predictions
         ]
diff --git a/inference/core/workflows/core_steps/models/roboflow/object_detection/v2.py b/inference/core/workflows/core_steps/models/roboflow/object_detection/v2.py
@@ -136,6 +136,7 @@ def describe_outputs(cls) -> List[OutputDefinition]:
             OutputDefinition(
                 name="predictions", kind=[OBJECT_DETECTION_PREDICTION_KIND]
             ),
+            OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]),
         ]
 
     @classmethod
@@ -251,6 +252,7 @@ def run_locally(
             images=images,
             predictions=predictions,
             class_filter=class_filter,
+            model_id=model_id,
         )
 
     def run_remotely(
@@ -302,13 +304,15 @@ def run_remotely(
             images=images,
             predictions=predictions,
             class_filter=class_filter,
+            model_id=model_id,
         )
 
     def _post_process_result(
         self,
         images: Batch[WorkflowImageData],
         predictions: List[dict],
         class_filter: Optional[List[str]],
+        model_id: str,
     ) -> BlockResult:
         inference_ids = [p.get(INFERENCE_ID_KEY, None) for p in predictions]
         predictions = convert_inference_detections_batch_to_sv_detections(predictions)
@@ -325,6 +329,10 @@ def _post_process_result(
             predictions=predictions,
         )
         return [
-            {"inference_id": inference_id, "predictions": prediction}
+            {
+                "inference_id": inference_id,
+                "predictions": prediction,
+                "model_id": model_id,
+            }
             for inference_id, prediction in zip(inference_ids, predictions)
         ]