Merge branch 'main' into feature/try-to-beat-the-limitation-of-ee-in-terms-of-singular-elements-pushed-into-batch-inputs

Author: PawelPeczek-Roboflow

inference/enterprise/workflows/enterprise_blocks/sinks/PLC_modbus/v1.py

@@ -46,6 +46,13 @@ class ModbusTCPBlockManifest(WorkflowBlockManifest):
             "long_description": LONG_DESCRIPTION,
             "license": "Apache-2.0",
             "block_type": "analytics",
+            "ui_manifest": {
+                "section": "industrial",
+                "icon": "fal fa-network-wired",
+                "blockPriority": 14,
+                "enterprise_only": True,
+                "local_only": True,
+            },
         }
     )
 

inference/enterprise/workflows/enterprise_blocks/sinks/PLCethernetIP/v1.py

@@ -56,6 +56,13 @@ class PLCBlockManifest(WorkflowBlockManifest):
             "long_description": LONG_DESCRIPTION,
             "license": "Roboflow Enterprise License",
             "block_type": "sinks",
+            "ui_manifest": {
+                "section": "industrial",
+                "icon": "fal fa-microchip",
+                "blockPriority": 13,
+                "enterprise_only": True,
+                "local_only": True,
+            },
         }
     )
 

inference/enterprise/workflows/enterprise_blocks/sinks/microsoft_sql_server/v1.py

@@ -130,6 +130,8 @@ class BlockManifest(WorkflowBlockManifest):
                 "icon": "fal fa-database",
                 "blockPriority": 3,
                 "popular": True,
+                "enterprise_only": True,
+                "local_only": True,
             },
         }
     )

inference/enterprise/workflows/enterprise_blocks/sinks/mqtt_writer/v1.py

@@ -42,6 +42,13 @@ class BlockManifest(WorkflowBlockManifest):
             "long_description": LONG_DESCRIPTION,
             "license": "Roboflow Enterprise License",
             "block_type": "sink",
+            "ui_manifest": {
+                "section": "industrial",
+                "icon": "fal fa-network-wired",
+                "blockPriority": 10,
+                "enterprise_only": True,
+                "local_only": True,
+            },
         }
     )
     type: Literal["mqtt_writer_sink@v1"]

inference/enterprise/workflows/enterprise_blocks/sinks/opc_writer/v1.py

@@ -106,6 +106,13 @@ class BlockManifest(WorkflowBlockManifest):
             "long_description": LONG_DESCRIPTION,
             "license": "Roboflow Enterprise License",
             "block_type": "sink",
+            "ui_manifest": {
+                "section": "industrial",
+                "icon": "fal fa-industry",
+                "blockPriority": 11,
+                "enterprise_only": True,
+                "local_only": True,
+            },
         }
     )
     type: Literal[BLOCK_TYPE]

inference/models/owlv2/owlv2.py

@@ -564,6 +564,7 @@ def infer_from_embed(
         query_embeddings: Dict[str, PosNegDictType],
         confidence: float,
         iou_threshold: float,
+        max_detections: int = MAX_DETECTIONS,
     ) -> List[Dict]:
         image_embeds = self.get_image_embeds(image_hash)
         if image_embeds is None:
@@ -601,6 +602,11 @@ def infer_from_embed(
         all_predicted_classes = all_predicted_classes[survival_indices]
         all_predicted_scores = all_predicted_scores[survival_indices]
 
+        if len(all_predicted_boxes) > max_detections:
+            all_predicted_boxes = all_predicted_boxes[:max_detections]
+            all_predicted_classes = all_predicted_classes[:max_detections]
+            all_predicted_scores = all_predicted_scores[:max_detections]
+
         # move tensors to numpy before returning
         all_predicted_boxes = all_predicted_boxes.cpu().numpy()
         all_predicted_classes = all_predicted_classes.cpu().numpy()
@@ -626,13 +632,18 @@ def infer(
         training_data: Dict,
         confidence: float = 0.99,
         iou_threshold: float = 0.3,
+        max_detections: int = MAX_DETECTIONS,
         **kwargs,
     ):
         class_embeddings_dict = self.make_class_embeddings_dict(
             training_data, iou_threshold
         )
         return self.infer_from_embedding_dict(
-            image, class_embeddings_dict, confidence, iou_threshold
+            image,
+            class_embeddings_dict,
+            confidence,
+            iou_threshold,
+            max_detections=max_detections,
         )
 
     def infer_from_embedding_dict(
@@ -641,6 +652,7 @@ def infer_from_embedding_dict(
         class_embeddings_dict: Dict[str, PosNegDictType],
         confidence: float,
         iou_threshold: float,
+        max_detections: int = MAX_DETECTIONS,
         **kwargs,
     ):
         if not isinstance(image, list):
@@ -660,7 +672,11 @@ def infer_from_embedding_dict(
             image_hash = self.embed_image(image_wrapper)
             image_wrapper.unload_numpy_image()
             result = self.infer_from_embed(
-                image_hash, class_embeddings_dict, confidence, iou_threshold
+                image_hash,
+                class_embeddings_dict,
+                confidence,
+                iou_threshold,
+                max_detections=max_detections,
             )
             results.append(result)
         return self.make_response(
@@ -944,14 +960,20 @@ def weights_file(self):
         return self.weights_file_path
 
     def infer(
-        self, image, confidence: float = 0.99, iou_threshold: float = 0.3, **kwargs
+        self,
+        image,
+        confidence: float = 0.99,
+        iou_threshold: float = 0.3,
+        max_detections: int = MAX_DETECTIONS,
+        **kwargs,
     ):
         logger.info(f"Inferring OWLv2 model")
         result = self.owlv2.infer_from_embedding_dict(
             image,
             self.train_data_dict,
             confidence=confidence,
             iou_threshold=iou_threshold,
+            max_detections=max_detections,
             **kwargs,
         )
         logger.info(f"OWLv2 model inference complete")

tests/inference/unit_tests/models/test_owlv2_max_detections.py

@@ -0,0 +1,29 @@
+import torch
+from unittest.mock import MagicMock
+
+from inference.models.owlv2 import owlv2
+
+
+def test_infer_from_embed_respects_max_detections(monkeypatch):
+    model = owlv2.OwlV2.__new__(owlv2.OwlV2)
+    image_boxes = torch.tensor(
+        [[0, 0, 1, 1], [0, 0, 2, 2], [0, 0, 3, 3], [0, 0, 4, 4]],
+        dtype=torch.float32,
+    )
+    image_class_embeds = torch.zeros((4, 2))
+    model.get_image_embeds = MagicMock(return_value=(None, image_boxes, image_class_embeds, None, None))
+
+    def fake_get_class_preds_from_embeds(*args, **kwargs):
+        boxes = image_boxes
+        classes = torch.zeros(4, dtype=torch.int64)
+        scores = torch.tensor([0.9, 0.8, 0.7, 0.6])
+        return boxes, classes, scores
+
+    monkeypatch.setattr(owlv2, "get_class_preds_from_embeds", fake_get_class_preds_from_embeds)
+    monkeypatch.setattr(owlv2.torchvision.ops, "nms", lambda boxes, scores, iou: torch.arange(boxes.shape[0]))
+
+    query_embeddings = {"a": {"positive": torch.zeros((1, 2)), "negative": None}}
+    predictions = model.infer_from_embed(
+        "hash", query_embeddings, confidence=0.5, iou_threshold=0.5, max_detections=2
+    )
+    assert len(predictions) == 2