Merge pull request #10 from kaspervn/inference_extra_options

Add possibility to pass extra arguments to the inference function

Author: Koldim2001

patched_yolo_infer/elements/CropElement.py

@@ -32,10 +32,11 @@ def __init__(
         self.detected_masks_real = None # List of np arrays containing masks in case of yolo-seg with the size of source_image_resized or source_image
         self.detected_polygons_real = None # List of polygons points in case of using memory optimaze in values from source_image_resized or source_image
 
-    def calculate_inference(self, model, imgsz=640, conf=0.35, iou=0.7, segment=False, classes_list=None, memory_optimize=False):
+    def calculate_inference(self, model, imgsz=640, conf=0.35, iou=0.7, segment=False, classes_list=None, memory_optimize=False, extra_args=None):
 
         # Perform inference
-        predictions = model.predict(self.crop, imgsz=imgsz, conf=conf, iou=iou, classes=classes_list, verbose=False)
+        extra_args = {} if extra_args is None else extra_args
+        predictions = model.predict(self.crop, imgsz=imgsz, conf=conf, iou=iou, classes=classes_list, verbose=False, **extra_args)
 
         pred = predictions[0]
 

patched_yolo_infer/nodes/MakeCropsDetectThem.py

@@ -68,7 +68,8 @@ def __init__(
         show_crops=False,
         resize_initial_size=False,
         model=None,
-        memory_optimize=True
+        memory_optimize=True,
+        inference_extra_args=None,
     ) -> None:
         if model is None:
             self.model = YOLO(model_path)  # Load the model from the specified path
@@ -89,6 +90,7 @@ def __init__(
         self.resize_initial_size = resize_initial_size  # slow operation !
         self.memory_optimize = memory_optimize # memory opimization option for segmentation
         self.class_names_dict = self.model.names
+        self.inference_extra_args = inference_extra_args
 
         self.crops = self.get_crops_xy(
             self.image,
@@ -199,7 +201,8 @@ def _detect_objects(self):
                 iou=self.iou,
                 segment=self.segment,
                 classes_list=self.classes_list,
-                memory_optimize=self.memory_optimize
+                memory_optimize=self.memory_optimize,
+                extra_args=self.inference_extra_args
             )
             crop.calculate_real_values()
             if self.resize_initial_size: