Add progress bars to single checks runs (#1236)

deepchecks/vision/base_checks.py

@@ -22,6 +22,7 @@
     ModelOnlyBaseCheck,
     DatasetKind
 )
+from deepchecks.core.display_suite import ProgressBar
 from deepchecks.vision.context import Context
 from deepchecks.vision.vision_data import VisionData
 from deepchecks.vision.batch_wrapper import Batch
@@ -52,6 +53,8 @@ def run(
     ) -> CheckResult:
         """Run check."""
         assert self.context_type is not None
+        p_bar = ProgressBar('Validating Input', 1, unit='')
+
         # Context is copying the data object, then not using the original after the init
         context: Context = self.context_type(dataset,
                                              model=model,
@@ -60,20 +63,29 @@ def run(
                                              n_samples=n_samples)
 
         self.initialize_run(context, DatasetKind.TRAIN)
+        p_bar.inc_progress()
+        p_bar.close()
 
+        p_bar = ProgressBar('Ingesting Batches', len(context.train), unit='Batch')
         context.train.init_cache()
         batch_start_index = 0
         for batch in context.train:
             batch = Batch(batch, context, DatasetKind.TRAIN, batch_start_index)
             context.train.update_cache(batch)
             self.update(context, batch, DatasetKind.TRAIN)
             batch_start_index += len(batch)
+            p_bar.inc_progress()
+        p_bar.close()
 
+        p_bar = ProgressBar('Computing Check', 1, unit='Check')
         result = self.compute(context, DatasetKind.TRAIN)
         footnote = context.get_is_sampled_footnote(DatasetKind.TRAIN)
         if footnote:
             result.display.append(footnote)
-        return self.finalize_check_result(result)
+        result = self.finalize_check_result(result)
+        p_bar.inc_progress()
+        p_bar.close()
+        return result
 
     def initialize_run(self, context: Context, dataset_kind: DatasetKind):
         """Initialize run before starting updating on batches. Optional."""
@@ -107,6 +119,7 @@ def run(
     ) -> CheckResult:
         """Run check."""
         assert self.context_type is not None
+        p_bar = ProgressBar('Validating Input', 1, unit='')
         # Context is copying the data object, then not using the original after the init
         context: Context = self.context_type(train_dataset,
                                              test_dataset,
@@ -116,28 +129,40 @@ def run(
                                              n_samples=n_samples)
 
         self.initialize_run(context)
+        p_bar.inc_progress()
+        p_bar.close()
 
+        p_bar = ProgressBar('Ingesting Batches - Train Dataset', len(context.train), unit='Batch')
         context.train.init_cache()
         batch_start_index = 0
         for batch in context.train:
             batch = Batch(batch, context, DatasetKind.TRAIN, batch_start_index)
             context.train.update_cache(batch)
             self.update(context, batch, DatasetKind.TRAIN)
             batch_start_index += len(batch)
+            p_bar.inc_progress()
+        p_bar.close()
 
+        p_bar = ProgressBar('Ingesting Batches - Test Dataset', len(context.train), unit='Batch')
         context.test.init_cache()
         batch_start_index = 0
         for batch in context.test:
             batch = Batch(batch, context, DatasetKind.TEST, batch_start_index)
             context.test.update_cache(batch)
             self.update(context, batch, DatasetKind.TEST)
             batch_start_index += len(batch)
+            p_bar.inc_progress()
+        p_bar.close()
 
+        p_bar = ProgressBar('Computing Check', 1, unit='Check')
         result = self.compute(context)
         footnote = context.get_is_sampled_footnote()
         if footnote:
             result.display.append(footnote)
-        return self.finalize_check_result(result)
+        result = self.finalize_check_result(result)
+        p_bar.inc_progress()
+        p_bar.close()
+        return result
 
     def initialize_run(self, context: Context):
         """Initialize run before starting updating on batches. Optional."""
@@ -165,10 +190,19 @@ def run(
     ) -> CheckResult:
         """Run check."""
         assert self.context_type is not None
+        p_bar = ProgressBar('Validating Input', 1, unit='')
         context: Context = self.context_type(model=model, device=device, random_state=random_state)
 
         self.initialize_run(context)
-        return self.finalize_check_result(self.compute(context))
+
+        p_bar.inc_progress()
+        p_bar.close()
+
+        p_bar = ProgressBar('Computing Check', 1, unit='Check')
+        result = self.finalize_check_result(self.compute(context))
+        p_bar.inc_progress()
+        p_bar.close()
+        return result
 
     def initialize_run(self, context: Context):
         """Initialize run before starting updating on batches. Optional."""