[Ai-1477]-darwin-v2-export-fix

darwin/dataset/local_dataset.py

@@ -9,6 +9,7 @@
 from darwin.utils import (
     SUPPORTED_IMAGE_EXTENSIONS,
     get_image_path_from_stream,
+    is_stream_list_empty,
     parse_darwin_json,
     stream_darwin_json,
 )
@@ -67,6 +68,7 @@ def __init__(
         split: str = "default",
         split_type: str = "random",
         release_name: Optional[str] = None,
+        keep_empty_annotations: bool = False,
     ):
         self.dataset_path = dataset_path
         self.annotation_type = annotation_type
@@ -75,10 +77,9 @@ def __init__(
         self.original_classes = None
         self.original_images_path: Optional[List[Path]] = None
         self.original_annotations_path: Optional[List[Path]] = None
+        self.keep_empty_annotations = keep_empty_annotations
 
-        release_path, annotations_dir, images_dir = self._initial_setup(
-            dataset_path, release_name
-        )
+        release_path, annotations_dir, images_dir = self._initial_setup(dataset_path, release_name)
         self._validate_inputs(partition, split_type, annotation_type)
         # Get the list of classes
 
@@ -101,6 +102,7 @@ def __init__(
             split,
             partition,
             split_type,
+            keep_empty_annotations,
         )
 
         if len(self.images_path) == 0:
@@ -117,9 +119,7 @@ def _validate_inputs(self, partition, split_type, annotation_type):
         if split_type not in ["random", "stratified"]:
             raise ValueError("split_type should be either 'random', 'stratified'")
         if annotation_type not in ["tag", "polygon", "bounding_box"]:
-            raise ValueError(
-                "annotation_type should be either 'tag', 'bounding_box', or 'polygon'"
-            )
+            raise ValueError("annotation_type should be either 'tag', 'bounding_box', or 'polygon'")
 
     def _setup_annotations_and_images(
         self,
@@ -130,19 +130,21 @@ def _setup_annotations_and_images(
         split,
         partition,
         split_type,
+        keep_empty_annotations: bool = False,
     ):
         # Find all the annotations and their corresponding images
         for annotation_path in sorted(annotations_dir.glob("**/*.json")):
             darwin_json = stream_darwin_json(annotation_path)
+
             image_path = get_image_path_from_stream(darwin_json, images_dir)
             if image_path.exists():
+                if not keep_empty_annotations and is_stream_list_empty(darwin_json["annotations"]):
+                    continue
                 self.images_path.append(image_path)
                 self.annotations_path.append(annotation_path)
                 continue
             else:
-                raise ValueError(
-                    f"Annotation ({annotation_path}) does not have a corresponding image"
-                )
+                raise ValueError(f"Annotation ({annotation_path}) does not have a corresponding image {image_path}")
 
     def _initial_setup(self, dataset_path, release_name):
         assert dataset_path is not None
@@ -201,9 +203,7 @@ def get_height_and_width(self, index: int) -> Tuple[float, float]:
         parsed = parse_darwin_json(self.annotations_path[index], index)
         return parsed.image_height, parsed.image_width
 
-    def extend(
-        self, dataset: "LocalDataset", extend_classes: bool = False
-    ) -> "LocalDataset":
+    def extend(self, dataset: "LocalDataset", extend_classes: bool = False) -> "LocalDataset":
         """
         Extends the current dataset with another one.
 
@@ -298,10 +298,7 @@ def parse_json(self, index: int) -> Dict[str, Any]:
         # Filter out unused classes and annotations of a different type
         if self.classes is not None:
             annotations = [
-                a
-                for a in annotations
-                if a.annotation_class.name in self.classes
-                and self.annotation_type_supported(a)
+                a for a in annotations if a.annotation_class.name in self.classes and self.annotation_type_supported(a)
             ]
         return {
             "image_id": index,
@@ -318,20 +315,15 @@ def annotation_type_supported(self, annotation) -> bool:
         elif self.annotation_type == "bounding_box":
             is_bounding_box = annotation_type == "bounding_box"
             is_supported_polygon = (
-                annotation_type in ["polygon", "complex_polygon"]
-                and "bounding_box" in annotation.data
+                annotation_type in ["polygon", "complex_polygon"] and "bounding_box" in annotation.data
             )
             return is_bounding_box or is_supported_polygon
         elif self.annotation_type == "polygon":
             return annotation_type in ["polygon", "complex_polygon"]
         else:
-            raise ValueError(
-                "annotation_type should be either 'tag', 'bounding_box', or 'polygon'"
-            )
+            raise ValueError("annotation_type should be either 'tag', 'bounding_box', or 'polygon'")
 
-    def measure_mean_std(
-        self, multi_threaded: bool = True
-    ) -> Tuple[np.ndarray, np.ndarray]:
+    def measure_mean_std(self, multi_threaded: bool = True) -> Tuple[np.ndarray, np.ndarray]:
         """
         Computes mean and std of trained images, given the train loader.
 
@@ -354,9 +346,7 @@ def measure_mean_std(
                 results = pool.map(self._return_mean, self.images_path)
                 mean = np.sum(np.array(results), axis=0) / len(self.images_path)
                 # Online image_classification deviation
-                results = pool.starmap(
-                    self._return_std, [[item, mean] for item in self.images_path]
-                )
+                results = pool.starmap(self._return_std, [[item, mean] for item in self.images_path])
                 std_sum = np.sum(np.array([item[0] for item in results]), axis=0)
                 total_pixel_count = np.sum(np.array([item[1] for item in results]))
                 std = np.sqrt(std_sum / total_pixel_count)
@@ -402,20 +392,14 @@ def _compute_weights(labels: List[int]) -> np.ndarray:
     @staticmethod
     def _return_mean(image_path: Path) -> np.ndarray:
         img = np.array(load_pil_image(image_path))
-        mean = np.array(
-            [np.mean(img[:, :, 0]), np.mean(img[:, :, 1]), np.mean(img[:, :, 2])]
-        )
+        mean = np.array([np.mean(img[:, :, 0]), np.mean(img[:, :, 1]), np.mean(img[:, :, 2])])
         return mean / 255.0
 
     # Loads an image with OpenCV and returns the channel wise std of the image.
     @staticmethod
     def _return_std(image_path: Path, mean: np.ndarray) -> Tuple[np.ndarray, float]:
         img = np.array(load_pil_image(image_path)) / 255.0
-        m2 = np.square(
-            np.array(
-                [img[:, :, 0] - mean[0], img[:, :, 1] - mean[1], img[:, :, 2] - mean[2]]
-            )
-        )
+        m2 = np.square(np.array([img[:, :, 0] - mean[0], img[:, :, 1] - mean[1], img[:, :, 2] - mean[2]]))
         return np.sum(np.sum(m2, axis=1), 1), m2.size / 3.0
 
     def __getitem__(self, index: int):
@@ -485,10 +469,7 @@ def build_stems(
     """
 
     if partition is None:
-        return (
-            str(e.relative_to(annotations_dir).parent / e.stem)
-            for e in sorted(annotations_dir.glob("**/*.json"))
-        )
+        return (str(e.relative_to(annotations_dir).parent / e.stem) for e in sorted(annotations_dir.glob("**/*.json")))
 
     if split_type == "random":
         split_filename = f"{split_type}_{partition}.txt"

darwin/dataset/remote_dataset.py

@@ -161,7 +161,7 @@ def split_video_annotations(self, release_name: str = "latest") -> None:
 
             frame_annotations = split_video_annotation(darwin_annotation)
             for frame_annotation in frame_annotations:
-                annotation = build_image_annotation(frame_annotation)
+                annotation = self._build_image_annotation(frame_annotation)
 
                 video_frame_annotations_path = annotations_path / annotation_file.stem
                 video_frame_annotations_path.mkdir(exist_ok=True, parents=True)
@@ -947,3 +947,8 @@ def local_images_path(self) -> Path:
     def identifier(self) -> DatasetIdentifier:
         """The ``DatasetIdentifier`` of this ``RemoteDataset``."""
         return DatasetIdentifier(team_slug=self.team, dataset_slug=self.slug)
+
+    def _build_image_annotation(
+        self, annotation_file: AnnotationFile
+    ) -> Dict[str, Any]:
+        return build_image_annotation(annotation_file)

darwin/dataset/remote_dataset_v1.py

@@ -12,8 +12,9 @@
     UploadHandlerV1,
 )
 from darwin.dataset.utils import is_relative_to
-from darwin.datatypes import ItemId, PathLike
+from darwin.datatypes import AnnotationFile, ItemId, PathLike
 from darwin.exceptions import NotFound, ValidationError
+from darwin.exporter.formats.darwin_1_0 import build_image_annotation
 from darwin.item import DatasetItem
 from darwin.item_sorter import ItemSorter
 from darwin.utils import find_files, urljoin
@@ -512,3 +513,8 @@ def import_annotation(self, item_id: ItemId, payload: Dict[str, Any]) -> None:
         """
 
         self.client.import_annotation(item_id, payload=payload)
+
+    def _build_image_annotation(
+        self, annotation_file: AnnotationFile
+    ) -> Dict[str, Any]:
+        return build_image_annotation(annotation_file)

darwin/dataset/remote_dataset_v2.py

@@ -22,8 +22,9 @@
     UploadHandlerV2,
 )
 from darwin.dataset.utils import is_relative_to
-from darwin.datatypes import ItemId, PathLike
+from darwin.datatypes import AnnotationFile, ItemId, PathLike
 from darwin.exceptions import NotFound, UnknownExportVersion
+from darwin.exporter.formats.darwin import build_image_annotation
 from darwin.item import DatasetItem
 from darwin.item_sorter import ItemSorter
 from darwin.utils import find_files, urljoin
@@ -543,3 +544,8 @@ def _fetch_stages(self, stage_type):
             workflow_id,
             [stage for stage in workflow["stages"] if stage["type"] == stage_type],
         )
+
+    def _build_image_annotation(
+        self, annotation_file: AnnotationFile
+    ) -> Dict[str, Any]:
+        return build_image_annotation(annotation_file)