Repository: lightly-ai/lightly

Issue:
BUG: Active Learning agent not working with Scorer.model_output being a generator
We allow the `Scorer.model_output` to be a generator, as this is need for the huge memory sizes of segmentation:
https://github.com/lightly-ai/lightly/blob/ca83601cd72fde73ade02a156c213f0b0c07d034/lightly/active_learning/scorers/semantic_segmentation.py#L95-L97
In the `ActiveLearningAgent` it is checked if the length of the model output is as expected:
https://github.com/lightly-ai/lightly/blob/ca83601cd72fde73ade02a156c213f0b0c07d034/lightly/active_learning/agents/agent.py#L194-L196

However, this fails for generators: `TypeError: object of type 'generator' has no len()`

## Tasks

- [x] reproduce the error by creating a corresponding unittest
- [x] fix the error, e.g. by making the check on the values of the scores_dict, not on the model_output




Discussion:

Resolving patch:
diff --git a/lightly/active_learning/agents/agent.py b/lightly/active_learning/agents/agent.py
--- a/lightly/active_learning/agents/agent.py
+++ b/lightly/active_learning/agents/agent.py
@@ -190,15 +190,19 @@ def query(self,
         # calculate active learning scores
         scores_dict = None
         if al_scorer is not None:
-            no_query_samples = len(self.query_set)
-            no_query_samples_with_scores = len(al_scorer.model_output)
-            if no_query_samples != no_query_samples_with_scores:
-                raise ValueError(
-                    f'Number of query samples ({no_query_samples}) must match '
-                    f'the number of predictions ({no_query_samples_with_scores})!'
-                )
             scores_dict = al_scorer.calculate_scores()
 
+            # Check if the length of the query_set and each of the scores are the same
+            no_query_samples = len(self.query_set)
+            for score in scores_dict.values():
+                no_query_samples_with_scores = len(score)
+                if no_query_samples != no_query_samples_with_scores:
+                    raise ValueError(
+                        f'Number of query samples ({no_query_samples}) must match '
+                        f'the number of predictions ({no_query_samples_with_scores})!'
+                    )
+
+
         # perform the sampling
         new_tag_data = self.api_workflow_client.sampling(
             sampler_config=sampler_config,


Test patch:
diff --git a/tests/active_learning/test_active_learning_agent.py b/tests/active_learning/test_active_learning_agent.py
--- a/tests/active_learning/test_active_learning_agent.py
+++ b/tests/active_learning/test_active_learning_agent.py
@@ -2,6 +2,7 @@
 
 from lightly.active_learning.agents.agent import ActiveLearningAgent
 from lightly.active_learning.config.sampler_config import SamplerConfig
+from lightly.active_learning.scorers import ScorerSemanticSegmentation
 from lightly.active_learning.scorers.classification import ScorerClassification
 from lightly.openapi_generated.swagger_client import SamplingMethod
 from tests.api_workflow.mocked_api_workflow_client import MockedApiWorkflowSetup
@@ -52,7 +53,7 @@ def test_agent_wrong_number_of_scores(self):
         method = SamplingMethod.CORAL
         n_samples = len(agent.labeled_set) + 2
 
-        n_predictions = len(agent.query_set) - 3 # the -3 should cause en error
+        n_predictions = len(agent.query_set) - 3  # the -3 should cause an error
         predictions = np.random.rand(n_predictions, 10).astype(np.float32)
         predictions_normalized = predictions / np.sum(predictions, axis=1)[:, np.newaxis]
         al_scorer = ScorerClassification(predictions_normalized)
@@ -61,6 +62,25 @@ def test_agent_wrong_number_of_scores(self):
         with self.assertRaises(ValueError):
             agent.query(sampler_config=sampler_config, al_scorer=al_scorer)
 
+    def test_agent_with_generator(self):
+        self.api_workflow_client.embedding_id = "embedding_id_xyz"
+        width = 32
+        height = 32
+        no_classes = 13
+
+        agent = ActiveLearningAgent(self.api_workflow_client, preselected_tag_name="preselected_tag_name_xyz")
+        method = SamplingMethod.CORAL
+        n_samples = len(agent.labeled_set) + 2
+
+        n_predictions = len(agent.query_set)
+        predictions = np.random.rand(n_predictions, no_classes, width, height).astype(np.float32)
+        predictions_normalized = predictions / np.sum(predictions, axis=1)[:, np.newaxis]
+        predictions_generator = (predictions_normalized[i] for i in range(n_predictions))
+        al_scorer = ScorerSemanticSegmentation(predictions_generator)
+
+        sampler_config = SamplerConfig(n_samples=n_samples, method=method)
+        agent.query(sampler_config=sampler_config, al_scorer=al_scorer)
+
     def test_agent_added_set_before_query(self):
 
         self.api_workflow_client.embedding_id = "embedding_id_xyz"