zea 0.0.7__py3-none-any.whl → 0.0.9__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