Fixing tests (WIP)

darwin/dataset/local_dataset.py

@@ -8,6 +8,7 @@
 from darwin.dataset.utils import get_classes, get_release_path, load_pil_image
 from darwin.utils import (
     SUPPORTED_IMAGE_EXTENSIONS,
+    get_darwin_json_version,
     get_image_path_from_stream,
     parse_darwin_json,
     stream_darwin_json,
@@ -131,12 +132,23 @@ def _setup_annotations_and_images(
         partition,
         split_type,
     ):
+        # Determine if the release is V1 or V2 JSON
+        json_version = get_darwin_json_version(annotations_dir)
+
+        #
+        annotation_files = list(annotations_dir.glob("**/*.json"))
+
+        for annotation_file in annotation_files:
+            with open(annotation_file, "r") as file:
+                data_str = file.read()
+                print(data_str)
+
         # Find all the annotations and their corresponding images
         with_folders = any([item.is_dir() for item in images_dir.iterdir()])
         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, with_folders
+                darwin_json, images_dir, with_folders, json_version
             )
             if image_path.exists():
                 self.images_path.append(image_path)

darwin/dataset/utils.py

@@ -17,6 +17,7 @@
     SUPPORTED_EXTENSIONS,
     SUPPORTED_VIDEO_EXTENSIONS,
     attempt_decode,
+    get_darwin_json_version,
     get_image_path_from_stream,
     is_unix_like_os,
     parse_darwin_json,
@@ -568,13 +569,20 @@ def _map_annotations_to_images(
     Raises:
         ValueError: If there are inconsistencies with the annotations and images.
     """
+
     images_paths = []
     annotations_paths = []
     invalid_annotation_paths = []
+
+    # Determine if the release is V1 or V2 JSON
+    json_version = get_darwin_json_version(annotations_dir)
+
     with_folders = any([item.is_dir() for item in images_dir.iterdir()])
     for annotation_path in annotations_dir.glob("**/*.json"):
         darwin_json = stream_darwin_json(annotation_path)
-        image_path = get_image_path_from_stream(darwin_json, images_dir, with_folders)
+        image_path = get_image_path_from_stream(
+            darwin_json, images_dir, with_folders, json_version
+        )
         if image_path.exists():
             images_paths.append(image_path)
             annotations_paths.append(annotation_path)

darwin/utils/utils.py

@@ -495,45 +495,74 @@ def stream_darwin_json(path: Path) -> PersistentStreamingJSONObject:
 
 
 def get_image_path_from_stream(
-    darwin_json: PersistentStreamingJSONObject, images_dir: Path, with_folders: bool
+    darwin_json: PersistentStreamingJSONObject,
+    images_dir: Path,
+    with_folders: bool,
+    json_version: str,
 ) -> Path:
     """
-    Returns the path to the image file associated with the given darwin json file (V1 or V2).
+    Returns the path to the image file associated with the given darwin json file.
+    Compatible with V1 & V2 Darwin JSON, as well as releases in folders and flat structures.
 
     Parameters
     ----------
     darwin_json : PersistentStreamingJSONObject
         A stream of the JSON file.
     images_dir : Path
         Path to the directory containing the images.
-    with_folders: Bool
+    with_folders: bool
         Flag to determine if the release was pulled with or without folders.
+    json_version: str
+        String representing the version of the Darwin JSON
 
     Returns
     -------
     Path
         Path to the image file.
     """
-    if not with_folders:
-        try:
+    if json_version == "2.0":
+        if not with_folders:
             return images_dir / Path(darwin_json["item"]["name"])
-        except KeyError:
-            return images_dir / Path(darwin_json["image"]["filename"])
-    else:
-        try:
+        else:
             return (
                 images_dir
                 / (Path(darwin_json["item"]["path"].lstrip("/\\")))
                 / Path(darwin_json["item"]["name"])
             )
-        except KeyError:
+    else:
+        if not with_folders:
+            try:
+                return images_dir / Path(darwin_json["image"]["filename"])
+            except Exception:
+                pass
+        else:
             return (
                 images_dir
                 / (Path(darwin_json["image"]["path"].lstrip("/\\")))
                 / Path(darwin_json["image"]["filename"])
             )
 
 
+def get_darwin_json_version(annotations_dir: Path) -> str:
+    """
+    Returns true is the input Darwin JSON file is 2.0, and False if 1.0.
+
+    Parameters
+    ----------
+    annotations_dir : Path
+        Path to the directory containing the annotation files.
+
+    Returns
+    -------
+    str
+        A str representing the Darwin JSON version.
+    """
+    with open(next(annotations_dir.glob("*.json")), "r") as file:
+        data_str = file.read()
+        data = json.loads(data_str)
+        return "2.0" if "version" in data and data["version"] == "2.0" else "1.0"
+
+
 def _parse_darwin_v2(path: Path, data: Dict[str, Any]) -> dt.AnnotationFile:
     item = data["item"]
     item_source = item.get("source_info", {})