[#51386] nodes: yolact_kenning.py: Limit predictions by score and amount

Signed-off-by: Illia Vysochyn <[email protected]>
antmicro · Apr 16, 2024 · a891354 · a891354
1 parent d8f4a19
commit a891354
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/cvnode_base/nodes/yolact_kenning.py b/cvnode_base/nodes/yolact_kenning.py
@@ -47,7 +47,7 @@ def prepare(self) -> bool:
         if not model_path.exists():
             self.node.get_logger().error(f"File {model_path} does not exist")
             return False
-        self.model = YOLACT(model_path, None)
+        self.model = YOLACT(model_path, None, top_k=100, score_threshold=0.3)
 
         device = self.node.get_parameter("device").value
         if not device:
@@ -108,7 +108,7 @@ def prepare(self) -> bool:
                 "Please specify device for TVM runtime"
             )
             return False
-        self.model = YOLACT(model_path, None)
+        self.model = YOLACT(model_path, None, top_k=100, score_threshold=0.3)
         self.runtime = TVMRuntime(
             model_path,
             contextname=device,
@@ -146,7 +146,7 @@ def prepare(self) -> bool:
         if not model_path.exists():
             self.node.get_logger().error(f"File {model_path} does not exist")
             return False
-        self.model = YOLACT(model_path, None)
+        self.model = YOLACT(model_path, None, top_k=100, score_threshold=0.3)
         self.runtime = TFLiteRuntime(
             model_path, disable_performance_measurements=True
         )