default tqdm progress callback function

Author: Jordan-Pierce

patched_yolo_infer/nodes/MakeCropsDetectThem.py

@@ -27,14 +27,15 @@ class MakeCropsDetectThem:
         overlap_x (int): Percentage of overlap along the x-axis.
         overlap_y (int): Percentage of overlap along the y-axis.
         show_crops (bool): Whether to visualize the cropping.
+        show_processing_status (bool): Whether to show the processing status using tqdm.
         resize_initial_size (bool): Whether to resize the results to the original 
                                     image size (ps: slow operation).
         model: Pre-initialized model object. If provided, the model will be used directly 
                    instead of loading from model_path.
         memory_optimize (bool): Memory optimization option for segmentation (less accurate results)
         batch_inference (bool): Batch inference of image crops through a neural network instead of 
                     sequential passes of crops (ps: Faster inference, higher memory use)
-        progress_callback (function): Optional callback function, (task: str, current: int, total: int)
+        progress_callback (function): Optional custom callback function, (task: str, current: int, total: int)
         inference_extra_args (dict): Dictionary with extra ultralytics inference parameters
 
     Attributes:
@@ -52,13 +53,14 @@ class MakeCropsDetectThem:
         overlap_y (int): Percentage of overlap along the y-axis.
         crops (list): List to store the CropElement objects.
         show_crops (bool): Whether to visualize the cropping.
+        show_processing_status (bool): Whether to show the processing status using tqdm.
         resize_initial_size (bool): Whether to resize the results to the original  
                                     image size (ps: slow operation).
         class_names_dict (dict): Dictionary containing class names of the YOLO model.
         memory_optimize (bool): Memory optimization option for segmentation (less accurate results)
         batch_inference (bool): Batch inference of image crops through a neural network instead of 
                                     sequential passes of crops (ps: Faster inference, higher memory use)
-        progress_callback (function): Optional callback function, (task: str, current: int, total: int)
+        progress_callback (function): Optional custom callback function, (task: str, current: int, total: int)
         inference_extra_args (dict): Dictionary with extra ultralytics inference parameters
     """
     def __init__(
@@ -75,13 +77,27 @@ def __init__(
         overlap_x=25,
         overlap_y=25,
         show_crops=False,
+        show_processing_status=True,
         resize_initial_size=True,
         model=None,
         memory_optimize=True,
         inference_extra_args=None,
         batch_inference=False,
         progress_callback=None,
     ) -> None:
+        
+        # Add show_process_status parameter and initialize progress bars dict
+        self.show_process_status = show_processing_status
+        self._progress_bars = {}
+        
+        # Set up the progress callback based on parameters
+        if progress_callback is not None:
+            self.progress_callback = progress_callback
+        elif show_processing_status:
+            self.progress_callback = self._tqdm_callback
+        else:
+            self.progress_callback = None
+            
         if model is None:
             self.model = YOLO(model_path)  # Load the model from the specified path
         else:
@@ -103,7 +119,6 @@ def __init__(
         self.class_names_dict = self.model.names # dict with human-readable class names
         self.inference_extra_args = inference_extra_args # dict with extra ultralytics inference parameters
         self.batch_inference = batch_inference # batch inference of image crops through a neural network
-        self.progress_callback = progress_callback # callback function to report progress of the inference
 
         self.crops = self.get_crops_xy(
             self.image,
@@ -117,6 +132,31 @@ def __init__(
             self._detect_objects_batch() 
         else:
             self._detect_objects()
+            
+    def _tqdm_callback(self, task, current, total):
+        """Internal callback function that uses tqdm for progress tracking
+        
+        Args:
+            task (str): The name of the task being tracked
+            current (int): The current progress value
+            total (int): The total number of steps in the task
+            
+        """
+        if task not in self._progress_bars:
+            self._progress_bars[task] = tqdm(
+                total=total,
+                desc=task,
+                unit='items'
+            )
+        
+        # Update progress
+        self._progress_bars[task].n = current
+        self._progress_bars[task].refresh()
+        
+        # Close and cleanup if task is complete
+        if current >= total:
+            self._progress_bars[task].close()
+            del self._progress_bars[task]
 
     def get_crops_xy(
         self,
@@ -342,4 +382,10 @@ def patches_info(self):
                 # Append the formatted string to the patch_info list
                 output += f"\nOn patch № {i}, nothing was detected"
         print(output)
+        
+    def __del__(self):
+        """Cleanup method to ensure all progress bars are closed"""
+        for pbar in self._progress_bars.values():
+            pbar.close()
+        self._progress_bars.clear()