Fix classification listing

worker_generic_training_dataset/__init__.py

@@ -156,14 +156,14 @@ class Extractor(DatasetWorker):
         """
         logger.info(f"Processing element ({element})")
 
+        polygon = element.polygon
+
         if isinstance(element, Element):
             image = element.image
-            polygon = element.polygon
             wk_version = get_object_id(element.worker_version)
             wk_run = get_object_id(element.worker_run)
         else:
             image = element.zone.image
-            polygon = element.zone.polygon
             wk_version = (
                 element.worker_version
                 if hasattr(element, "worker_version")
@@ -229,7 +229,7 @@ class Extractor(DatasetWorker):
                 f"Linking element {cached_element.id} to dataset ({self.cached_dataset.id})"
             )
             with cache_database.atomic():
-                cached_element: CachedDatasetElement = CachedDatasetElement.create(
+                CachedDatasetElement.create(
                     id=uuid.uuid4(),
                     element=cached_element,
                     dataset=self.cached_dataset,

worker_generic_training_dataset/from_api.py

@@ -8,7 +8,7 @@ from arkindex_worker.cache import (
     CachedTranscription,
     CachedTranscriptionEntity,
 )
-from arkindex_worker.models import Element, Transcription
+from arkindex_worker.models import Element, Set, Transcription
 from arkindex_worker.worker.classification import ClassificationMixin
 from arkindex_worker.worker.element import ElementMixin
 from arkindex_worker.worker.entity import EntityMixin
@@ -26,6 +26,16 @@ class DatasetExtractorFromAPI(
     TranscriptionMixin,
     MetaDataMixin,
 ):
+    def list_set_elements(self, dataset_set: Set) -> Iterator[Element]:
+        for element in super().list_set_elements(dataset_set):
+            # Classifications are not serialized in ListDatasetElements
+            yield Element(
+                **self.request(
+                    "RetrieveElement",
+                    id=element.id,
+                )
+            )
+
     def list_element_children(self, *args, **kwargs) -> Iterator[Element]:
         return map(
             Element, super().list_element_children(*args, **kwargs, with_classes=True)