zea 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

zea/data/convert/camus.py

@@ -1,14 +1,19 @@
 """Functionality to convert the camus dataset to the zea format.
-Requires SimpleITK to be installed: pip install SimpleITK.
+
+.. note::
+    Requires SimpleITK to be installed: `pip install SimpleITK`.
+
+For more information about the dataset, resort to the following links:
+
+- The original dataset can be found at `this link <https://humanheart-project.creatis.insa-lyon.fr/database/#collection/6373703d73e9f0047faa1bc8>`_.
+
 """
 
 from __future__ import annotations
 
-import argparse
-import importlib.util
 import logging
 import os
-import sys
+from concurrent.futures import ProcessPoolExecutor
 from pathlib import Path
 from typing import Any, Dict, Tuple
 
@@ -17,10 +22,11 @@ import scipy
 from skimage.transform import resize
 from tqdm import tqdm
 
-# from zea.display import transform_sc_image_to_polar
 from zea import log
+from zea.data.convert.utils import unzip
 from zea.data.data_format import generate_zea_dataset
-from zea.utils import find_first_nonzero_index, translate
+from zea.internal.utils import find_first_nonzero_index
+from zea.tensor_ops import translate
 
 
 def transform_sc_image_to_polar(image_sc, output_size=None, fit_outline=True):
@@ -123,10 +129,12 @@ def sitk_load(filepath: str | Path) -> Tuple[np.ndarray, Dict[str, Any]]:
         filepath: Path to the image.
 
     Returns:
-        - ([N], H, W), Image array.
+        - Image array of shape (num_frames, height, width).
         - Collection of metadata.
     """
     # Load image and save info
+    import SimpleITK as sitk
+
     image = sitk.ReadImage(str(filepath))
 
     all_metadata = {}
@@ -147,7 +155,7 @@ def sitk_load(filepath: str | Path) -> Tuple[np.ndarray, Dict[str, Any]]:
     return im_array, metadata
 
 
-def convert_camus(source_path, output_path, overwrite=False):
+def process_camus(source_path, output_path, overwrite=False):
     """Converts the camus database to the zea format.
 
     Args:
@@ -187,29 +195,22 @@ def convert_camus(source_path, output_path, overwrite=False):
     )
 
 
-def get_args():
-    """Parse command line arguments."""
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--source",
-        type=str,
-        # path to CAMUS_public/database_nifti
-        required=True,
-    )
-    parser.add_argument(
-        "--output",
-        type=str,
-        required=True,
-    )
-    args = parser.parse_args()
-    return args
-
-
 splits = {"train": [1, 401], "val": [401, 451], "test": [451, 501]}
 
 
 def get_split(patient_id: int) -> str:
-    """Determine the dataset split for a given patient ID."""
+    """
+    Determine which dataset split a patient ID belongs to.
+
+    Args:
+        patient_id: Integer ID of the patient.
+
+    Returns:
+        The split name: "train", "val", or "test".
+
+    Raises:
+        ValueError: If the patient_id does not fall into any defined split range.
+    """
     if splits["train"][0] <= patient_id < splits["train"][1]:
         return "train"
     elif splits["val"][0] <= patient_id < splits["val"][1]:
@@ -220,16 +221,58 @@ def get_split(patient_id: int) -> str:
         raise ValueError(f"Did not find split for patient: {patient_id}")
 
 
-if __name__ == "__main__":
-    if importlib.util.find_spec("SimpleITK") is None:
-        log.error("SimpleITK not installed. Please install SimpleITK: `pip install SimpleITK`")
-        sys.exit()
-    import SimpleITK as sitk
+def _process_task(task):
+    """
+    Unpack a task tuple and invoke process_camus in a worker process.
+
+    Creates parent directories for the target outputs, calls process_camus
+    with the unpacked paths, and logs then re-raises any exception raised by processing.
+
+    Args:
+        task (tuple): (source_file_str, output_file_str)
+            - source_file_str: filesystem path to the source CAMUS file as a string.
+            - output_file_str: filesystem path for the ZEA output file as a string.
+    """
+    source_file_str, output_file_str = task
+    source_file = Path(source_file_str)
+    output_file = Path(output_file_str)
 
-    args = get_args()
+    # Ensure destination directories exist (safe to call from multiple processes)
+    output_file.parent.mkdir(parents=True, exist_ok=True)
+
+    # Call the real processing function (must be importable in the worker)
+    # If process_camus lives in another module, import it there instead.
+    try:
+        process_camus(source_file, output_file, overwrite=False)
+    except Exception:
+        # Log and re-raise so the main process can handle it
+        log.error("Error processing %s", source_file)
+        raise
+
+
+def convert_camus(args):
+    """
+    Converts the CAMUS dataset into ZEA HDF5 files across dataset splits.
 
-    camus_source_folder = Path(args.source)
-    camus_output_folder = Path(args.output)
+    Processes files found under the CAMUS source folder (after unzipping if needed),
+    assigns each patient to a train/val/test split, creates matching output paths,
+    and executes per-file conversion tasks either serially or in parallel.
+    Ensures output directories do not pre-exist, and logs progress and failures.
+
+    Args:
+        args (argparse.Namespace): An object with attributes:
+
+            - src (str | Path): Path to the CAMUS archive or extracted folder.
+            - dst (str | Path): Root destination folder for ZEA HDF5 outputs;
+              split subfolders will be created.
+            - no_hyperthreading (bool, optional): If True, run tasks serially instead
+              of using a process pool.
+    """
+    camus_source_folder = Path(args.src)
+    camus_output_folder = Path(args.dst)
+
+    # Look for either CAMUS_public.zip or folders database_nifti, database_split
+    camus_source_folder = unzip(camus_source_folder, "camus")
 
     # check if output folders already exist
     for split in splits:
@@ -238,9 +281,9 @@ if __name__ == "__main__":
         )
 
     # clone folder structure of source to output using pathlib
-    # and run convert_camus() for every hdf5 found in there
     files = list(camus_source_folder.glob("**/*_half_sequence.nii.gz"))
-    for source_file in tqdm(files):
+    tasks = []
+    for source_file in files:
         # check if source file in camus database (ignore other files)
         if "database_nifti" not in source_file.parts:
             continue
@@ -250,10 +293,28 @@ if __name__ == "__main__":
         split = get_split(patient_id)
 
         output_file = camus_output_folder / split / source_file.relative_to(camus_source_folder)
-
         # Replace .nii.gz with .hdf5
         output_file = output_file.with_suffix("").with_suffix(".hdf5")
-
         # make sure folder exists
         output_file.parent.mkdir(parents=True, exist_ok=True)
-        convert_camus(source_file, output_file, overwrite=False)
+
+        tasks.append((str(source_file), str(output_file)))
+    if not tasks:
+        log.info("No files found to process.")
+        return
+
+    if getattr(args, "no_hyperthreading", False):
+        log.info("no_hyperthreading is True — running tasks serially (no ProcessPoolExecutor)")
+        for t in tqdm(tasks, desc="Processing files (serial)"):
+            try:
+                _process_task(t)
+            except Exception as e:
+                log.error("Task processing failed: %s", e)
+        log.info("Processing finished for %d files (serial)", len(tasks))
+        return
+
+    # Submit tasks to the process pool and track progress
+    with ProcessPoolExecutor() as exe:
+        for _ in tqdm(exe.map(_process_task, tasks), total=len(tasks), desc="Processing files"):
+            pass
+    log.info("Processing finished for %d files", len(tasks))

zea/data/convert/echonet.py

@@ -1,49 +1,30 @@
 """
-Script to convert the EchoNet database to .npy and zea formats.
-Will segment the images and convert them to polar coordinates.
-"""
+Script to convert the EchoNet database to zea format.
 
-import os
+.. note::
+    Will segment the images and convert them to polar coordinates.
+
+For more information about the dataset, resort to the following links:
 
-os.environ["KERAS_BACKEND"] = "numpy"
+- The original dataset can be found at `this link <https://stanfordaimi.azurewebsites.net/datasets/834e1cd1-92f7-4268-9daa-d359198b310a>`_.
+- The project page is available `here <https://echonet.github.io/>`_.
+
+"""
 
-import argparse
+import os
 from concurrent.futures import ProcessPoolExecutor, as_completed
+from multiprocessing import Value
 from pathlib import Path
 
 import numpy as np
+import yaml
 from scipy.interpolate import griddata
 from tqdm import tqdm
 
-from zea.config import Config
+from zea import log
 from zea.data import generate_zea_dataset
-from zea.io_lib import load_video
-from zea.utils import translate
-
-
-def get_args():
-    """Parse command line arguments."""
-    parser = argparse.ArgumentParser(description="Convert EchoNet to zea format")
-    parser.add_argument(
-        "--source",
-        type=str,
-        # path to EchoNet-Dynamic/Videos
-        required=True,
-    )
-    parser.add_argument(
-        "--output",
-        type=str,
-        required=True,
-    )
-    parser.add_argument(
-        "--splits",
-        type=str,
-        default=None,
-    )
-    parser.add_argument("--output_numpy", type=str, default=None)
-    parser.add_argument("--no_hyperthreading", action="store_true")
-    args = parser.parse_args()
-    return args
+from zea.data.convert.utils import load_avi, unzip
+from zea.tensor_ops import translate
 
 
 def segment(tensor, number_erasing=0, min_clip=0):
@@ -52,6 +33,8 @@ def segment(tensor, number_erasing=0, min_clip=0):
     Args:
         tensor (ndarray): Input image (sc) with 3 dimensions. (N, 112, 112)
         number_erasing (float, optional): number to fill the background with.
+        min_clip (float, optional): If > 0, values on the computed cone edge will be clipped
+            to be at least this value. Defaults to 0.
     Returns:
         tensor (ndarray): Segmented matrix of same dimensions as input
 
@@ -264,13 +247,37 @@ def cartesian_to_polar_matrix(
 
 
 def find_split_for_file(file_dict, target_file):
-    """Function that finds the split for a given file in a dictionary."""
+    """
+    Locate which split contains a given filename.
+
+    Parameters:
+        file_dict (dict): Mapping from split name (e.g., "train", "val", "test", "rejected")
+            to an iterable of filenames.
+        target_file (str): Filename to search for within the split lists.
+
+    Returns:
+        str: The split name that contains `target_file`, or `"rejected"` if the file is not found.
+    """
     for split, files in file_dict.items():
         if target_file in files:
             return split
+    log.warning(f"File {target_file} not found in any split, defaulting to rejected.")
     return "rejected"
 
 
+def count_init(shared_counter):
+    """
+    Initialize the module-level shared counter used by worker processes.
+
+    Parameters:
+        shared_counter (multiprocessing.Value): A process-shared integer Value that
+            will be assigned to the module-global COUNTER for coordinated counting
+            across processes.
+    """
+    global COUNTER
+    COUNTER = shared_counter
+
+
 class H5Processor:
     """
     Stores a few variables and paths to allow for hyperthreading.
@@ -279,7 +286,6 @@ class H5Processor:
     def __init__(
         self,
         path_out_h5,
-        path_out=None,
         num_val=500,
         num_test=500,
         range_from=(0, 255),
@@ -287,7 +293,6 @@ class H5Processor:
         splits=None,
     ):
         self.path_out_h5 = Path(path_out_h5)
-        self.path_out = Path(path_out) if path_out else None
         self.num_val = num_val
         self.num_test = num_test
         self.range_from = range_from
@@ -297,20 +302,33 @@ class H5Processor:
 
         # Ensure train, val, test, rejected paths exist
         for folder in ["train", "val", "test", "rejected"]:
-            if self._to_numpy:
-                (self.path_out / folder).mkdir(parents=True, exist_ok=True)
             (self.path_out_h5 / folder).mkdir(parents=True, exist_ok=True)
 
-    @property
-    def _to_numpy(self):
-        return self.path_out is not None
-
-    def translate(self, data):
+    def _translate(self, data):
         """Translate the data from the processing range to final range."""
         return translate(data, self._process_range, self.range_to)
 
     def get_split(self, hdf5_file: str, sequence):
-        """Determine the split for a given file."""
+        """
+        Determine the dataset split label for a given file and its image sequence.
+
+        This method checks acceptance based on the first frame of `sequence`.
+        If explicit splits were provided to the processor, it returns the split
+        found for `hdf5_file` (and asserts that the acceptance result matches the split).
+        If no explicit splits are provided, rejected sequences are labeled `"rejected"`.
+        Accepted sequences increment a shared counter and are assigned
+        `"val"`, `"test"`, or `"train"` according to the processor's
+        `num_val` and `num_test` quotas.
+
+        Args:
+            hdf5_file (str): Filename or identifier used to look up an existing split
+                when splits are provided.
+            sequence (array-like): Time-ordered sequence of images; the first frame is
+                used for acceptance checking.
+
+        Returns:
+            str: One of `"train"`, `"val"`, `"test"`, or `"rejected"` indicating the assigned split.
+        """
         # Always check acceptance
         accepted = accept_shape(sequence[0])
 
@@ -324,25 +342,73 @@ class H5Processor:
         if not accepted:
             return "rejected"
 
-        # This inefficient counter works with hyperthreading
-        # TODO: but it is not reproducible!
-        val_counter = len(list((self.path_out_h5 / "val").iterdir()))
-        test_counter = len(list((self.path_out_h5 / "test").iterdir()))
+        # Increment the hyperthreading counter
+        # Note that some threads will start on subsequent splits
+        # while others are still processing
+        with COUNTER.get_lock():
+            COUNTER.value += 1
+            n = COUNTER.value
 
         # Determine the split
-        if val_counter < self.num_val:
+        if n <= self.num_val:
             return "val"
-        elif test_counter < self.num_test:
+        elif n <= self.num_val + self.num_test:
             return "test"
         else:
             return "train"
 
+    def validate_split_copy(self, split_file):
+        """
+        Validate that a generated split YAML matches the original splits provided to the processor.
+
+        Reads the YAML at `split_file` and compares its `train`, `val`, `test`, and `rejected` lists
+        (or other split keys present in `self.splits`) against `self.splits`; logs confirmation
+        when a split matches and logs which entries are missing or extra when they differ. If the
+        processor was not initialized with `splits`, validation is skipped and a message is logged.
+
+        Args:
+            split_file (str or os.PathLike): Path to the YAML file containing the
+                generated dataset splits.
+        """
+        if self.splits is not None:
+            # Read the split_file and ensure contents of the train, val and split match
+            with open(split_file, "r") as f:
+                new_splits = yaml.safe_load(f)
+            for split in self.splits.keys():
+                if set(new_splits[split]) == set(self.splits[split]):
+                    log.info(f"Split {split} copied correctly.")
+                else:
+                    # Log which entry is missing or extra in the split_file
+                    missing = set(self.splits[split]) - set(new_splits[split])
+                    extra = set(new_splits[split]) - set(self.splits[split])
+                    if missing:
+                        log.warning(f"New dataset split {split} is missing entries: {missing}")
+                    if extra:
+                        log.warning(f"New dataset split {split} has extra entries: {extra}")
+        else:
+            log.info(
+                "Processor not initialized with a split, not validating if the split was copied."
+            )
+
     def __call__(self, avi_file):
         """
-        Processes a single h5 file using the class variables and the filename given.
+        Convert a single AVI file into a zea dataset entry.
+        Loads the AVI, validates and rescales pixel ranges, applies segmentation,
+        assigns a data split (train/val/test/rejected), converts accepted frames
+        to polar coordinates.
+        Constructs and returns the zea dataset descriptor used by
+        generate_zea_dataset; the descriptor always includes `path`, `image_sc`,
+        `probe_name`, and `description`, and includes `image` when the file is accepted.
+
+        Args:
+            avi_file (pathlib.Path): Path to the source .avi file to process.
+
+        Returns:
+            dict: The value returned by generate_zea_dataset containing the dataset
+                entry for the processed file.
         """
         hdf5_file = avi_file.stem + ".hdf5"
-        sequence = load_video(avi_file)
+        sequence = load_avi(avi_file)
 
         assert sequence.min() >= self.range_from[0], f"{sequence.min()} < {self.range_from[0]}"
         assert sequence.max() <= self.range_from[1], f"{sequence.max()} > {self.range_from[1]}"
@@ -356,25 +422,14 @@ class H5Processor:
         accepted = split != "rejected"
 
         out_h5 = self.path_out_h5 / split / hdf5_file
-        if self._to_numpy:
-            out_dir = self.path_out / split / avi_file.stem
-            out_dir.mkdir(parents=True, exist_ok=True)
 
         polar_im_set = []
-        for i, im in enumerate(sequence):
-            if self._to_numpy:
-                np.save(out_dir / f"sc{str(i).zfill(3)}.npy", im)
-
+        for _, im in enumerate(sequence):
             if not accepted:
                 continue
 
             polar_im = cartesian_to_polar_matrix(im, interpolation="cubic")
             polar_im = np.clip(polar_im, *self._process_range)
-            if self._to_numpy:
-                np.save(
-                    out_dir / f"polar{str(i).zfill(3)}.npy",
-                    polar_im,
-                )
             polar_im_set.append(polar_im)
 
         if accepted:
@@ -386,53 +441,99 @@ class H5Processor:
 
         zea_dataset = {
             "path": out_h5,
-            "image_sc": self.translate(sequence),
+            "image_sc": self._translate(sequence),
             "probe_name": "generic",
             "description": "EchoNet dataset converted to zea format",
         }
         if accepted:
-            zea_dataset["image"] = self.translate(polar_im_set)
+            zea_dataset["image"] = self._translate(polar_im_set)
         return generate_zea_dataset(**zea_dataset)
 
 
-if __name__ == "__main__":
-    args = get_args()
+def convert_echonet(args):
+    """
+    Convert an EchoNet dataset into zea files, organizing results
+    into train/val/test/rejected splits.
+
+    Args:
+        args (argparse.Namespace): An object with the following attributes.
+
+            - src (str|Path): Path to the source archive or directory containing .avi files.
+                Will be unzipped if needed.
+            - dst (str|Path): Destination directory for generated zea files
+                per-split subdirectories (train, val, test, rejected) and a split.yaml
+                are created or updated.
+            - split_path (str|Path|None): If provided, must contain a split.yaml to reproduce
+                an existing split; function asserts the file exists.
+            - no_hyperthreading (bool): When false, processing uses a ProcessPoolExecutor
+                with a shared counter; when true, processing runs sequentially.
+
+    Note:
+        - May unzip the source into a working directory.
+        - Writes zea files into dst.
+        - Writes a split.yaml into dst summarizing produced files per split.
+        - Logs progress and validation results.
+        - Asserts that split.yaml exists at split_path when split reproduction is requested.
+    """
+    # Check if unzip is needed
+    src = unzip(args.src, "echonet")
 
-    if args.splits is not None:
+    if args.split_path is not None:
         # Reproduce a previous split...
-        split_yaml_dir = Path(args.splits)
-        splits = {"train": None, "val": None, "test": None}
-        for split in splits:
-            yaml_file = split_yaml_dir / (split + ".yaml")
-            assert yaml_file.exists(), f"File {yaml_file} does not exist."
-            splits[split] = Config.from_yaml(yaml_file)["file_paths"]
+        yaml_file = Path(args.split_path) / "split.yaml"
+        assert yaml_file.exists(), f"File {yaml_file} does not exist."
+        splits = {"train": None, "val": None, "test": None, "rejected": None}
+        with open(yaml_file, "r") as f:
+            splits = yaml.safe_load(f)
+        log.info(f"Processor initialized with train-val-test split from {yaml_file}.")
     else:
         splits = None
 
     # List the files that have an entry in path_out_h5 already
     files_done = []
-    for _, _, filenames in os.walk(args.output):
+    for _, _, filenames in os.walk(args.dst):
         for filename in filenames:
             files_done.append(filename.replace(".hdf5", ""))
 
     # List all files of echonet and exclude those already processed
-    path_in = Path(args.source)
+    path_in = Path(src)
     h5_files = path_in.glob("*.avi")
     h5_files = [file for file in h5_files if file.stem not in files_done]
-    print(f"Files left to process: {len(h5_files)}")
+    log.info(f"Files left to process: {len(h5_files)}")
 
     # Run the processor
-    processor = H5Processor(path_out_h5=args.output, path_out=args.output_numpy, splits=splits)
+    processor = H5Processor(path_out_h5=args.dst, splits=splits)
 
-    print("Starting the conversion process.")
+    log.info("Starting the conversion process.")
 
     if not args.no_hyperthreading:
-        with ProcessPoolExecutor() as executor:
-            futures = {executor.submit(processor, file): file for file in h5_files}
-            for future in tqdm(as_completed(futures), total=len(h5_files)):
-                future.result()
+        shared_counter = Value("i", 0)
+        with ProcessPoolExecutor(initializer=count_init, initargs=(shared_counter,)) as executor:
+            futures = [executor.submit(processor, file) for file in h5_files]
+            for future in tqdm(as_completed(futures), total=len(futures)):
+                try:
+                    future.result()
+                except Exception:
+                    log.warning("Task raised an exception")
     else:
+        # Initialize global variable for counting
+        count_init(Value("i", 0))
         for file in tqdm(h5_files):
             processor(file)
 
-    print("All tasks are completed.")
+    log.info("All tasks are completed.")
+
+    # Write to yaml split files
+    full_list = {}
+    for split in ["train", "val", "test", "rejected"]:
+        split_dir = Path(args.dst) / split
+
+        # Get only files (skip directories)
+        file_list = [f.name for f in split_dir.iterdir() if f.is_file()]
+        full_list[split] = file_list
+
+    with open(Path(args.dst) / "split.yaml", "w") as f:
+        yaml.dump(full_list, f)
+
+    # Validate that the split was copied correctly
+    processor.validate_split_copy(Path(args.dst) / "split.yaml")

zea/data/convert/echonetlvh/README.md

@@ -3,7 +3,6 @@
   vary from file to file. This is done by the `precompute_crop.py` script, which will
   produce as output a .csv file `cone_parameters.csv`. The cone parameters will indicate
   how the video should be cropped in order to bound the cone and remove margins.
-- Next, `convert_raw_to_zea.py` can be run to convert the dataset to zea format,
-  with cropping and scan conversion. The measurement locations stored in `MeasurementsList.csv`
+- Next, `__init__.py` converts the dataset to zea format,
+  with cropping and scan conversion. The original measurement locations stored in `MeasurementsList.csv`
   are also updated to match the new cropping / padding coordinates.
-- You can save the video and measurement plots for a converted video using `examples/echonetlvh/plot_sample.py`.