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

zea/data/convert/echonetlvh/precompute_crop.py

@@ -3,58 +3,27 @@ Script to precompute cone parameters for the EchoNet-LVH dataset.
 This script should be run separately before the main conversion process.
 """
 
-import argparse
 import csv
 import json
-import os
 from pathlib import Path
 
 from tqdm import tqdm
 
-# Set Keras backend to numpy for best CPU performance
-os.environ["KERAS_BACKEND"] = "numpy"
-
+from zea import log
 from zea.tools.fit_scan_cone import fit_and_crop_around_scan_cone
 
 
-def get_args():
-    """Parse command line arguments."""
-    parser = argparse.ArgumentParser(
-        description="Precompute cone parameters for EchoNet-LVH dataset"
-    )
-    parser.add_argument(
-        "--source",
-        type=str,
-        required=True,
-    )
-    parser.add_argument(
-        "--output",
-        type=str,
-        required=True,
-    )
-    parser.add_argument(
-        "--batch",
-        type=str,
-        help="Specify which BatchX directory to process, e.g. --batch=Batch2",
-    )
-    parser.add_argument(
-        "--max_files",
-        type=int,
-        default=None,
-        help="Maximum number of files to process (for testing)",
-    )
-    parser.add_argument(
-        "--force",
-        action="store_true",
-        help="Force recomputation even if parameters already exist",
-    )
-    return parser.parse_args()
-
-
 def load_splits(source_dir):
-    """Load splits from MeasurementsList.csv and return avi filenames"""
+    """
+    Load splits from MeasurementsList.csv and return avi filenames
+
+    Args:
+        source_dir: Source directory containing MeasurementsList.csv
+    Returns:
+        Dictionary with keys 'train', 'val', 'test', 'rejected' and values as lists of avi filenames
+    """
     csv_path = Path(source_dir) / "MeasurementsList.csv"
-    splits = {"train": [], "val": [], "test": []}
+    splits = {"train": [], "val": [], "test": [], "rejected": []}
     # Read CSV using built-in csv module
     with open(csv_path, newline="", encoding="utf-8") as csvfile:
         reader = csv.DictReader(csvfile)
@@ -71,7 +40,17 @@ def load_splits(source_dir):
 
 
 def find_avi_file(source_dir, hashed_filename, batch=None):
-    """Find AVI file in the specified batch directory or any batch if not specified."""
+    """
+    Find AVI file in the specified batch directory or any batch if not specified.
+
+    Args:
+        source_dir: Source directory containing BatchX subdirectories
+        hashed_filename: Hashed filename (with or without .avi extension)
+        batch: Specific batch directory to search in (e.g., "Batch2"), or None to search all batches
+
+    Returns:
+        Path to the AVI file if found, else None
+    """
     # If filename already has .avi extension, strip it
     if hashed_filename.endswith(".avi"):
         hashed_filename = hashed_filename[:-4]
@@ -98,7 +77,7 @@ def load_first_frame(avi_file):
         avi_file: Path to the video file
 
     Returns:
-        First frame as numpy array
+        First frame as numpy array of shape (H, W) and dtype np.uint8 (grayscale)
     """
     try:
         import cv2
@@ -129,9 +108,20 @@ def precompute_cone_parameters(args):
     This function loads the first frame from each AVI file, applies fit_scan_cone
     to determine cropping parameters, and saves these parameters to a CSV file
     for later use during the actual data conversion.
+
+    Args:
+        args: Argument parser namespace with the following attributes:
+            src: Source directory containing EchoNet-LVH data
+            dst: Destination directory to save cone parameters
+            batch: Specific batch to process (e.g., "Batch2") or None for all
+            max_files: Maximum number of files to process (or None for all)
+            force: Whether to recompute parameters if they already exist
+    Returns:
+        Path to the CSV file containing cone parameters
     """
-    source_path = Path(args.source)
-    output_path = Path(args.output)
+
+    source_path = Path(args.src)
+    output_path = Path(args.dst)
     output_path.mkdir(parents=True, exist_ok=True)
 
     # Output file for cone parameters
@@ -140,7 +130,7 @@ def precompute_cone_parameters(args):
 
     # Check if parameters already exist
     if cone_params_csv.exists() and not args.force:
-        print(f"Parameters already exist at {cone_params_csv}. Use --force to recompute.")
+        log.warning(f"Parameters already exist at {cone_params_csv}. Use --force to recompute.")
         return cone_params_csv
 
     # Get list of files to process
@@ -151,21 +141,21 @@ def precompute_cone_parameters(args):
         for avi_filename in split_files:
             # Strip .avi if present
             base_filename = avi_filename[:-4] if avi_filename.endswith(".avi") else avi_filename
-            avi_file = find_avi_file(args.source, base_filename, batch=args.batch)
+            avi_file = find_avi_file(args.src, base_filename, batch=args.batch)
             if avi_file:
                 files_to_process.append((avi_file, avi_filename))
             else:
-                print(
-                    f"Warning: Could not find AVI file for {base_filename} in batch "
+                log.warning(
+                    f"Could not find AVI file for {base_filename} in batch "
                     f"{args.batch if args.batch else 'any'}"
                 )
 
     # Limit files if max_files is specified
     if args.max_files is not None:
         files_to_process = files_to_process[: args.max_files]
-        print(f"Limited to processing {args.max_files} files due to max_files parameter")
+        log.info(f"Limited to processing {args.max_files} files due to max_files parameter")
 
-    print(f"Computing cone parameters for {len(files_to_process)} files")
+    log.info(f"Computing cone parameters for {len(files_to_process)} files")
 
     # Dictionary to store parameters for each file
     all_cone_params = {}
@@ -217,7 +207,7 @@ def precompute_cone_parameters(args):
                 all_cone_params[avi_filename] = essential_params
 
             except Exception as e:
-                print(f"Error processing {avi_file}: {str(e)}")
+                log.error(f"Error processing {avi_file}: {str(e)}")
 
                 # Write failure record
                 failure_record = {
@@ -236,16 +226,5 @@ def precompute_cone_parameters(args):
     with open(cone_params_json, "w", encoding="utf-8") as jsonfile:
         json.dump(all_cone_params, jsonfile)
 
-    print(f"Cone parameters saved to {cone_params_csv} and {cone_params_json}")
+    log.info(f"Cone parameters saved to {cone_params_csv} and {cone_params_json}")
     return cone_params_csv
-
-
-if __name__ == "__main__":
-    args = get_args()
-    print("Using Keras backend: numpy (forced for best performance)")
-
-    # Precompute cone parameters
-    cone_params_csv = precompute_cone_parameters(args)
-
-    print(f"Precomputation completed. Parameters saved to {cone_params_csv}")
-    print("You can now run the main conversion script.")

zea/data/convert/picmus.py

@@ -1,15 +1,11 @@
 """
 Script to convert the PICMUS database to the zea format.
 
-Example usage:
-```bash
-python zea/data/convert/picmus.py \
---src_dir /mnt/data/PICMUS \
---output_dir converted_PICMUS_dir
-```
+For more information about the dataset, resort to the following links:
+
+- The original dataset can be found at `this link <https://www.creatis.insa-lyon.fr/Challenge/IEEE_IUS_2016/download>`_.
 """
 
-import argparse
 import logging
 import os
 from pathlib import Path
@@ -19,11 +15,13 @@ import numpy as np
 
 from zea import log
 from zea.beamform.delays import compute_t0_delays_planewave
+from zea.data.convert.utils import unzip
 from zea.data.data_format import generate_zea_dataset
 
 
-def convert_picmus(source_path, output_path, overwrite=False):
-    """Converts the PICMUS database to the zea format.
+def convert(source_path, output_path, overwrite=False):
+    """
+    Converts and writes a single PICMUS file to the zea format.
 
     Args:
         source_path (str, pathlike): The path to the original PICMUS file.
@@ -112,37 +110,37 @@ def convert_picmus(source_path, output_path, overwrite=False):
     )
 
 
-def get_args():
-    """Parse command line arguments."""
-    parser = argparse.ArgumentParser(
-        description=(
-            "Converts the PICMUS database to the zea format. The "
-            "src_dir is scanned for hdf5 files ending in iq or rf. These files are"
-            "converted and stored in output_dir under the same relative path as "
-            "they came from in src_dir."
-        )
-    )
-    parser.add_argument(
-        "--src_dir",
-        type=str,
-        help="Source directory where the original PICMUS data is stored.",
-    )
-
-    parser.add_argument("--output_dir", type=str, help="Output directory of the converted database")
-    return parser.parse_args()
-
-
-if __name__ == "__main__":
-    # Parse the arguments
-    args = get_args()
+def convert_picmus(args):
+    """
+    Convert PICMUS HDF5 files under a source directory into the zea dataset format,
+    preserving relative paths in the destination.
 
+    Args:
+        args (argparse.Namespace): An object with the following attributes.
+
+            - src (str or Path): Path to the PICMUS source directory or archive.
+            - dst (str or Path): Path to the output directory where converted .hdf5 files
+              will be written.
+
+    Note:
+        - Scans `src` (after unzipping if needed) for `.hdf5` files containing IQ/RF data and
+          converts each to the zea format.
+        - Preserves the relative directory structure under `dst` and places each converted
+          file in its own subdirectory named after the file stem.
+        - Fails fast if `src` does not exist or if `dst` already exists.
+    """
     # Get the source and output directories
-    base_dir = Path(args.src_dir)
-    output_dir = Path(args.output_dir)
+    base_dir = Path(args.src)
+    dst = Path(args.dst)
 
     # Check if the source directory exists and create the output directory
     assert base_dir.exists(), f"Source directory {base_dir} does not exist."
-    output_dir.mkdir(parents=True, exist_ok=False)
+
+    assert not dst.exists(), f"Destination directory {dst} already exists, Exiting."
+
+    # Unzip the PICMUS dataset if necessary
+    base_dir = unzip(base_dir, "picmus")
+    dst.mkdir(parents=True, exist_ok=False)
 
     # Traverse the source directory and convert all files
     for file in base_dir.rglob("*.hdf5"):
@@ -151,9 +149,8 @@ if __name__ == "__main__":
         # Select only the data files that actually contain rf or iq data
         # (There are also files containing the geometry of the phantoms or
         # images)
-        if (
-            not str_file.endswith("iq.hdf5") or not str_file.endswith("rf.hdf5")
-        ) and "img" in str_file:
+        is_data_file = str_file.endswith("iq.hdf5") or str_file.endswith("rf.hdf5")
+        if not is_data_file or "img" in str_file:
             log.info("Skipping %s", file.name)
             continue
 
@@ -161,7 +158,7 @@ if __name__ == "__main__":
 
         # Find the folder relative to the base directory to retain the
         # folder structure in the output directory
-        output_file = output_dir / file.relative_to(base_dir)
+        output_file = dst / file.relative_to(base_dir)
 
         # Define the output path
         # NOTE: I added output_file.stem to put each file in its own
@@ -175,7 +172,7 @@ if __name__ == "__main__":
             # Create the output directory if it does not exist already
             output_file.parent.mkdir(parents=True, exist_ok=True)
 
-            convert_picmus(file, output_file, overwrite=True)
+            convert(file, output_file, overwrite=True)
         except Exception:
             output_file.parent.rmdir()
             log.error("Failed to convert %s", str_file)

zea/data/convert/utils.py

@@ -0,0 +1,86 @@
+import zipfile
+from pathlib import Path
+
+import imageio
+import numpy as np
+from PIL import Image
+
+from zea import log
+
+
+def load_avi(file_path, mode="L"):
+    """Load a .avi file and return a numpy array of frames.
+
+    Args:
+        filename (str): The path to the video file.
+        mode (str, optional): Color mode: "L" (grayscale) or "RGB".
+            Defaults to "L".
+
+    Returns:
+        numpy.ndarray: Array of frames (num_frames, H, W) or (num_frames, H, W, C)
+    """
+    frames = []
+    with imageio.get_reader(file_path) as reader:
+        for frame in reader:
+            img = Image.fromarray(frame)
+            img = img.convert(mode)
+            img = np.array(img)
+            frames.append(img)
+    return np.stack(frames)
+
+
+def unzip(src: str | Path, dataset: str) -> Path:
+    """
+    Checks if data folder exist in src.
+    Otherwise, unzip dataset.zip in src.
+
+    Args:
+        src (str | Path): The source directory containing the zip file or unzipped folder.
+        dataset (str): The name of the dataset to unzip.
+            Options are "picmus", "camus", "echonet", "echonetlvh".
+
+    Returns:
+        Path: The path to the unzipped dataset directory.
+    """
+    src = Path(src)
+    if dataset == "picmus":
+        zip_name = "picmus.zip"
+        folder_name = "archive_to_download"
+        unzip_dir = src / folder_name
+    elif dataset == "camus":
+        zip_name = "CAMUS_public.zip"
+        folder_name = "CAMUS_public"
+        unzip_dir = src / folder_name
+    elif dataset == "echonet":
+        zip_name = "EchoNet-Dynamic.zip"
+        folder_name = "EchoNet-Dynamic"
+        unzip_dir = src / folder_name / "Videos"
+    elif dataset == "echonetlvh":
+        zip_name = "EchoNet-LVH.zip"
+        folder_name = "Batch1"
+        unzip_dir = src
+    else:
+        raise ValueError(f"Dataset {dataset} not recognized for unzip.")
+
+    if (src / folder_name).exists():
+        if dataset == "echonetlvh":
+            # EchoNetLVH dataset unzips into four folders. Check they all exist.
+            assert (src / "Batch2").exists(), f"Missing Batch2 folder in {src}."
+            assert (src / "Batch3").exists(), f"Missing Batch3 folder in {src}."
+            assert (src / "Batch4").exists(), f"Missing Batch4 folder in {src}."
+            assert (src / "MeasurementsList.csv").exists(), (
+                f"Missing MeasurementsList.csv in {src}."
+            )
+            log.info(f"Found Batch1, Batch2, Batch3, Batch4 and MeasurementsList.csv in {src}.")
+        return unzip_dir
+
+    zip_path = src / zip_name
+    if not zip_path.exists():
+        raise FileNotFoundError(f"Could not find {zip_name} or {folder_name} folder in {src}.")
+
+    log.info(f"Unzipping {zip_path} to {src}...")
+    with zipfile.ZipFile(zip_path, "r") as zip_ref:
+        zip_ref.extractall(src)
+    log.info("Unzipping completed.")
+    log.info(f"Starting conversion from {src / folder_name}.")
+    return unzip_dir

zea/data/convert/{matlab.py → verasonics.py}

@@ -1,6 +1,6 @@
-"""Functionality to convert Verasonics matlab raw files to the zea format.
+"""Functionality to convert Verasonics MATLAB workspace to the zea format.
 
-Example (MATLAB):
+Example of saving the entire workspace to a .mat file (MATLAB):
 
     .. code-block:: matlab
 
@@ -8,19 +8,19 @@ Example (MATLAB):
         >> VSX;
         >> save_raw('C:/path/to/raw_data.mat');
 
-Then in python:
+Then convert the saved `raw_data.mat` file to zea format using the following code (Python):
 
     .. code-block:: python
 
-        from zea.data_format.zea_from_matlab_raw import zea_from_matlab_raw
+        from zea.data.convert.verasonics import zea_from_verasonics_workspace
 
-        zea_from_matlab_raw("C:/path/to/raw_data.mat", "C:/path/to/output.hdf5")
+        zea_from_verasonics_workspace("C:/path/to/raw_data.mat", "C:/path/to/output.hdf5")
 
 Or alternatively, use the script below to convert all .mat files in a directory:
 
     .. code-block:: bash
 
-        python zea/data/convert/matlab.py "C:/path/to/directory"
+        python zea/data/convert/verasonics.py "C:/path/to/directory"
 
 or without the directory argument, the script will prompt you to select a directory
 using a file dialog.
@@ -72,7 +72,7 @@ Adding additional elements
 
 You can add additional elements to the dataset by defining a function that reads the
 data from the file and returns a ``DatasetElement``. Then pass the function to the
-``zea_from_matlab_raw`` function as a list.
+``zea_from_verasonics_workspace`` function as a list.
 
 .. code-block:: python
 
@@ -91,14 +91,13 @@ data from the file and returns a ``DatasetElement``. Then pass the function to t
         )
 
 
-    zea_from_matlab_raw(
+    zea_from_verasonics_workspace(
         "C:/path/to/raw_data.mat",
         "C:/path/to/output.hdf5",
         [read_high_voltage_func],
     )
 """  # noqa: E501
 
-import argparse
 import os
 import sys
 import traceback
@@ -150,11 +149,6 @@ def dereference_index(file, dataset, index, event=None, subindex=None):
         else:
             return file[reference][subindex]
     else:
-        if index > 0:
-            log.warning(
-                f"index {index} is not a reference. You are probably "
-                "incorrectly indexing a dataset."
-            )
         return dataset
 
 
@@ -935,7 +929,7 @@ def get_frame_indices(file, frames):
         frame_indices (np.ndarray): The frame indices.
     """
     # Read the number of frames from the file
-    n_frames = int(file["Resource"]["RcvBuffer"]["numFrames"][0][0])
+    n_frames = int(dereference_index(file, file["Resource"]["RcvBuffer"]["numFrames"], 0)[0][0])
 
     if isinstance(frames, str) and frames == "all":
         # Create an array of all frame-indices
@@ -956,8 +950,8 @@ def get_frame_indices(file, frames):
     return frame_indices
 
 
-def zea_from_matlab_raw(input_path, output_path, additional_functions=None, frames="all"):
-    """Converts a Verasonics matlab raw file to the zea format. The MATLAB file
+def zea_from_verasonics_workspace(input_path, output_path, additional_functions=None, frames="all"):
+    """Converts a Verasonics MATLAB workspace file (.mat) to the zea format. The .mat file
     should be created using the `save_raw` function and be stored in "v7.3" format.
 
     Args:
@@ -970,6 +964,7 @@ def zea_from_matlab_raw(input_path, output_path, additional_functions=None, fram
             a list of integers, a range of integers (e.g. 4-8), or 'all'. Defaults to
             'all'.
     """
+
     # Create the output directory if it does not exist
     input_path = Path(input_path)
     output_path = Path(output_path)
@@ -1004,6 +999,7 @@ def zea_from_matlab_raw(input_path, output_path, additional_functions=None, fram
             # convert dict of events to dict of lists
             data = {key: [data[event][key] for event in data] for key in data[0]}
             description = ["Verasonics data with multiple events"] * num_events
+
             # Generate the zea dataset
             generate_zea_dataset(
                 path=output_path,
@@ -1024,40 +1020,6 @@ def zea_from_matlab_raw(input_path, output_path, additional_functions=None, fram
     log.success(f"Converted {log.yellow(input_path)} to {log.yellow(output_path)}")
 
 
-def parse_args():
-    """Parse command line arguments."""
-    parser = argparse.ArgumentParser(
-        description="Convert Verasonics matlab raw files to the zea format."
-        "Example usage: python zea/data/convert/matlab.py raw_file.mat output.hdf5 --frames 1-5 7"
-    )
-    parser.add_argument(
-        "input_path",
-        default=None,
-        type=str,
-        nargs="?",
-        help="The path to a file or directory containing raw Verasonics data.",
-    )
-
-    parser.add_argument(
-        "output_path",
-        default=None,
-        type=str,
-        nargs="?",
-        help="The path to the output file or directory.",
-    )
-
-    parser.add_argument(
-        "--frames",
-        default=["all"],
-        type=str,
-        nargs="+",
-        help="The frames to add to the file. This can be a list of integers, a range "
-        "of integers (e.g. 4-8), or 'all'.",
-    )
-
-    return parser.parse_args()
-
-
 def get_answer(prompt, additional_options=None):
     """Get a yes or no answer from the user. There is also the option to provide
     additional options. In case yes or no is selected, the function returns a boolean.
@@ -1083,15 +1045,25 @@ def get_answer(prompt, additional_options=None):
         log.warning("Invalid input.")
 
 
-if __name__ == "__main__":
-    args = parse_args()
+def convert_verasonics(args):
+    """
+    Converts a Verasonics MATLAB workspace file (.mat) or a directory containing multiple
+    such files to the zea format.
+
+    Args:
+        args (argparse.Namespace): An object with attributes:
+
+            - src (str): Source folder path.
+            - dst (str): Destination folder path.
+            - frames (list[str]): MATLAB frames spec (e.g., ["all"], integers, or ranges like "4-8")
+    """
 
     # Variable to indicate what to do with existing files.
     # Is set by the user in case these are found.
     existing_file_policy = None
 
-    if args.input_path is None:
-        log.info("Select a directory containing Verasonics matlab raw files.")
+    if args.src is None:
+        log.info("Select a directory containing Verasonics MATLAB workspace files.")
         # Create a Tkinter root window
         try:
             import tkinter as tk
@@ -1115,7 +1087,7 @@ if __name__ == "__main__":
                 )
             ) from e
     else:
-        selected_path = args.input_path
+        selected_path = args.src
 
     # Exit when no path is selected
     if not selected_path:
@@ -1128,14 +1100,14 @@ if __name__ == "__main__":
 
     # Set the output path to be next to the input directory with _zea appended
     # to the name
-    if args.output_path is None:
+    if args.dst is None:
         if selected_path_is_directory:
             output_path = selected_path.parent / (Path(selected_path).name + "_zea")
         else:
             output_path = str(selected_path.with_suffix("")) + "_zea.hdf5"
             output_path = Path(output_path)
     else:
-        output_path = Path(args.output_path)
+        output_path = Path(args.dst)
         if selected_path.is_file() and output_path.suffix not in (".hdf5", ".h5"):
             log.error(
                 "When converting a single file, the output path should have the .hdf5 "
@@ -1145,7 +1117,7 @@ if __name__ == "__main__":
         elif selected_path.is_dir() and output_path.is_file():
             log.error("When converting a directory, the output path should be a directory.")
             sys.exit()
-        #
+
         if output_path.is_dir() and not selected_path_is_directory:
             output_path = output_path / (selected_path.name + "_zea.hdf5")
 
@@ -1180,7 +1152,7 @@ if __name__ == "__main__":
             else:
                 log.info("Aborting...")
                 sys.exit()
-        zea_from_matlab_raw(selected_path, output_path, frames=frames)
+        zea_from_verasonics_workspace(selected_path, output_path, frames=frames)
     else:
         # Continue with the rest of your code...
         for root, dirs, files in os.walk(selected_path):
@@ -1227,7 +1199,7 @@ if __name__ == "__main__":
                         file_output_path.unlink(missing_ok=False)
 
                 try:
-                    zea_from_matlab_raw(full_path, file_output_path, frames=frames)
+                    zea_from_verasonics_workspace(full_path, file_output_path, frames=frames)
                 except Exception:
                     # Print error message without raising it
                     log.error(f"Failed to convert {mat_file}")