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

zea/data/convert/echonetlvh/{convert_raw_to_usbmd.py → __init__.py}

@@ -4,24 +4,18 @@ Script to convert the EchoNet-LVH database to zea format.
 Each video is cropped so that the scan cone is centered
 without padding, such that it can be converted to polar domain.
 
-This cropping requires first computing scan cone parameters
-using `data/convert/echonetlvh/precompute_crop.py`, which
-are then passed to this script.
-"""
-
-import os
-
-os.environ["KERAS_BACKEND"] = "jax"
+.. note::
+    This cropping requires first computing scan cone parameters
+    using :mod:`zea.data.convert.echonetlvh.precompute_crop`, which
+    are then passed to this script.
 
+For more information about the dataset, resort to the following links:
 
-if __name__ == "__main__":
-    from zea import init_device
-
-    init_device("auto:1")
+- The original dataset can be found at `this link <https://stanfordaimi.azurewebsites.net/datasets/5b7fcc28-579c-4285-8b72-e4238eac7bd1>`_.
+"""
 
-import argparse
 import csv
-import sys
+import os
 from concurrent.futures import ProcessPoolExecutor, as_completed
 from pathlib import Path
 
@@ -30,50 +24,71 @@ import numpy as np
 from jax import jit, vmap
 from tqdm import tqdm
 
+from zea import log
 from zea.data import generate_zea_dataset
 from zea.data.convert.echonet import H5Processor
+from zea.data.convert.echonetlvh.precompute_crop import precompute_cone_parameters
+from zea.data.convert.utils import load_avi, unzip
 from zea.display import cartesian_to_polar_matrix
-from zea.io_lib import load_video
 from zea.tensor_ops import translate
 
 
-def get_args():
-    """Parse command line arguments."""
-    parser = argparse.ArgumentParser(description="Convert EchoNet-LVH to zea format")
-    parser.add_argument(
-        "--source",
-        type=str,
-        required=True,
-    )
-    parser.add_argument("--output", type=str, required=True)
-    parser.add_argument("--output_numpy", type=str, default=None)
-    parser.add_argument("--file_list", type=str, help="Optional path to list of files")
-    parser.add_argument("--use_hyperthreading", action="store_true", help="Enable hyperthreading")
-    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)",
-    )
-    # if neither is specified, both will be converted
-    parser.add_argument(
-        "--convert_measurements",
-        action="store_true",
-        help="Only convert measurements CSV file",
-    )
-    parser.add_argument("--convert_images", action="store_true", help="Only convert image files")
-    return parser.parse_args()
+def overwrite_splits(source_dir):
+    """
+    Overwrite MeasurementsList.csv splits based on manual_rejections.txt
+
+    Args:
+        source_dir: Source directory containing MeasurementsList.csv and manual_rejections.txt
+    Returns:
+        None
+    """
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    rejection_path = os.path.join(current_dir, "manual_rejections.txt")
+    try:
+        with open(rejection_path) as f:
+            rejected_hashes = [line.strip() for line in f]
+    except FileNotFoundError:
+        log.warning(f"{rejection_path} not found, skipping rejections.")
+        return
+
+    csv_path = Path(source_dir) / "MeasurementsList.csv"
+    temp_path = Path(source_dir) / "MeasurementsList_temp.csv"
+    try:
+        rejection_counter = 0
+        with (
+            csv_path.open("r", newline="", encoding="utf-8") as infile,
+            temp_path.open("w", encoding="utf-8", newline="") as outfile,
+        ):
+            reader = csv.DictReader(infile)
+            writer = csv.DictWriter(outfile, fieldnames=reader.fieldnames)
+            writer.writeheader()
+            for row in reader:
+                if row["HashedFileName"] in rejected_hashes:
+                    row["split"] = "rejected"
+                    rejection_counter += 1
+                writer.writerow(row)
+            assert rejection_counter == 278, (
+                f"Expected 278 rejections, but applied only {rejection_counter}."
+            )
+    except FileNotFoundError:
+        log.warning(f"{csv_path} not found, skipping rejections.")
+        return
+    temp_path.replace(csv_path)
+    log.info(f"Overwritten {rejection_counter}/278 rejections to {csv_path}")
+    return
 
 
 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": []}
     with open(csv_path, newline="", encoding="utf-8") as csvfile:
         reader = csv.DictReader(csvfile)
         file_split_map = {}
@@ -87,7 +102,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]
@@ -206,10 +231,9 @@ class LVHProcessor(H5Processor):
 
     def __init__(self, *args, cone_params=None, **kwargs):
         super().__init__(*args, **kwargs)
-        # self.cart2pol_jit = jit(cartesian_to_polar_matrix_jax)
+        # Store the pre-computed cone parameters
         self.cart2pol_jit = jit(cartesian_to_polar_matrix)
         self.cart2pol_batched = vmap(self.cart2pol_jit)
-        # Store the pre-computed cone parameters
         self.cone_parameters = cone_params or {}
 
     def get_split(self, avi_file: str, sequence):
@@ -232,22 +256,26 @@ class LVHProcessor(H5Processor):
         raise UserWarning("Unknown split for file: " + filename)
 
     def __call__(self, avi_file):
-        print(avi_file)
+        """Takes a single avi_file and generates a zea dataset
+
+        Args:
+            avi_file: String or path to avi_file to be processed
+
+        Returns:
+            zea dataset
+        """
+
         avi_filename = Path(avi_file).stem + ".avi"
-        sequence = jnp.array(load_video(avi_file))
+        sequence = np.array(load_avi(avi_file))
 
         sequence = translate(sequence, self.range_from, self._process_range)
-
         # Get pre-computed cone parameters for this file
         cone_params = self.cone_parameters.get(avi_filename)
-
         if cone_params is not None:
             # Apply pre-computed cropping parameters
             sequence = crop_sequence_with_params(sequence, cone_params)
         else:
-            print(f"Warning: No cone parameters for {avi_filename}, using original sequence")
-
-        # Convert to JAX array for polar conversion
+            log.warning(f"No cone parameters for {avi_filename}, using original sequence")
         sequence = jnp.array(sequence)
 
         split = self.get_split(avi_file, sequence)
@@ -282,7 +310,7 @@ def transform_measurement_coordinates_with_cone_params(row, cone_params):
         A new row with transformed coordinates, or None if cone_params is None
     """
     if cone_params is None:
-        print(f"Warning: No cone parameters for file {row['HashedFileName']}")
+        log.warning(f"No cone parameters for file {row['HashedFileName']}")
         return None
 
     new_row = dict(row)
@@ -324,7 +352,7 @@ def transform_measurement_coordinates_with_cone_params(row, cone_params):
     )
 
     if is_out_of_bounds:
-        print(f"Warning: Transformed coordinates out of bounds for file {row['HashedFileName']}")
+        log.warning(f"Transformed coordinates out of bounds for file {row['HashedFileName']}")
 
     # Convert back to string if original was string
     for k in ["X1", "X2", "Y1", "Y2"]:
@@ -353,7 +381,7 @@ def convert_measurements_csv(source_csv, output_csv, cone_params_csv=None):
         if cone_params_csv and Path(cone_params_csv).exists():
             cone_parameters = load_cone_parameters(cone_params_csv)
         else:
-            print("Warning: No cone parameters file found. Measurements will not be transformed.")
+            log.warning("No cone parameters file found. Measurements will not be transformed.")
 
         # Apply coordinate transformation and track skipped rows
         transformed_rows = []
@@ -371,7 +399,7 @@ def convert_measurements_csv(source_csv, output_csv, cone_params_csv=None):
                 else:
                     skipped_files.add(row["HashedFileName"])
             except Exception as e:
-                print(f"Error processing row for file {row['HashedFileName']}: {str(e)}")
+                log.error(f"Error processing row for file {row['HashedFileName']}: {str(e)}")
                 skipped_files.add(row["HashedFileName"])
 
         # Save to new CSV file
@@ -389,30 +417,65 @@ def convert_measurements_csv(source_csv, output_csv, cone_params_csv=None):
                 writer.writeheader()
 
         # Print summary
-        print("\nConversion Summary:")
-        print(f"Total rows processed: {len(rows)}")
-        print(f"Rows successfully converted: {len(transformed_rows)}")
-        print(f"Rows skipped: {len(rows) - len(transformed_rows)}")
+        log.info("Conversion Summary:")
+        log.info(f"Total rows processed: {len(rows)}")
+        log.info(f"Rows successfully converted: {len(transformed_rows)}")
+        log.info(f"Rows skipped: {len(rows) - len(transformed_rows)}")
         if skipped_files:
-            print("\nSkipped files:")
+            log.info("Skipped files:")
             for filename in sorted(skipped_files):
-                print(f"  - {filename}")
-        print(f"\nConverted measurements saved to {output_csv}")
+                log.info(f"  - {filename}")
+        log.info(f"Converted measurements saved to {output_csv}")
 
     except Exception as e:
-        print(f"Error processing CSV file: {str(e)}")
+        log.error(f"Error processing CSV file: {str(e)}")
         raise
 
 
-if __name__ == "__main__":
-    args = get_args()
+def _process_file_worker(avi_file, dst, splits, cone_parameters, range_from, process_range):
+    """
+    Function for a hyperthreading worker to process a single file.
+
+    Args:
+        avi_file: Path to the AVI file to process
+        dst: Destination directory for output
+        splits: Dictionary of splits
+        cone_parameters: Dictionary of cone parameters
+        range_from: Range from value for processing
+        process_range: Process range value for processing
+    Returns:
+        Result of processing the file
+    """
+
+    # create a fresh processor inside the worker process
+    proc = LVHProcessor(path_out_h5=dst, splits=splits, cone_params=cone_parameters)
+    # if LVHProcessor needs range_from/_process_range set, set them here
+    proc.range_from = range_from
+    proc._process_range = process_range
+    return proc(avi_file)
+
+
+def convert_echonetlvh(args):
+    """
+    Conversion script for the EchoNet-LVH dataset.
+    Unzips, overwrites splits if needed, precomputes cone parameters,
+    and converts images and/or measurements to zea format and saves dataset.
+    Is called with argparse arguments through zea/zea/data/convert/__main__.py
+
+    Args:
+        args (argparse.Namespace): Command-line arguments
+    """
+    # Check if unzip is needed
+    src = unzip(args.src, "echonetlvh")
+
+    # Overwrite the splits if manual rejections are provided
+    if not args.no_rejection:
+        overwrite_splits(args.src)
 
     # Check that cone parameters exist
-    cone_params_csv = Path(args.output) / "cone_parameters.csv"
+    cone_params_csv = Path(args.dst) / "cone_parameters.csv"
     if not cone_params_csv.exists():
-        print(f"Error: Cone parameters not found at {cone_params_csv}")
-        print("Please run precompute_crop.py first to generate the parameters.")
-        sys.exit(1)
+        precompute_cone_parameters(args)
 
     # If no specific conversion is requested, convert both
     if not (args.convert_measurements or args.convert_images):
@@ -421,30 +484,30 @@ if __name__ == "__main__":
 
     # Convert images if requested
     if args.convert_images:
-        source_path = Path(args.source)
+        source_path = Path(src)
         splits = load_splits(source_path)
 
         # Load precomputed cone parameters
         cone_parameters = load_cone_parameters(cone_params_csv)
-        print(f"Loaded cone parameters for {len(cone_parameters)} files")
+        log.info(f"Loaded cone parameters for {len(cone_parameters)} files")
 
         files_to_process = []
         for split_files in splits.values():
             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(src, base_filename, batch=args.batch)
                 if avi_file:
                     files_to_process.append(avi_file)
                 else:
-                    print(
+                    log.warning(
                         f"Warning: Could not find AVI file for {base_filename} in batch "
                         f"{args.batch if args.batch else 'any'}"
                     )
 
         # List files that have already been processed
         files_done = []
-        for _, _, filenames in os.walk(args.output):
+        for _, _, filenames in os.walk(args.dst):
             for filename in filenames:
                 if filename.endswith(".hdf5"):
                     files_done.append(filename.replace(".hdf5", ""))
@@ -455,45 +518,53 @@ if __name__ == "__main__":
         # 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"Files left to process: {len(files_to_process)}")
+        log.info(f"Files left to process: {len(files_to_process)}")
 
         # Initialize processor with splits and cone parameters
-        processor = LVHProcessor(
-            path_out_h5=args.output,
-            path_out=args.output_numpy,
-            splits=splits,
-            cone_params=cone_parameters,
-        )
-
-        print("Starting the conversion process.")
-
-        if args.use_hyperthreading:
-            with ProcessPoolExecutor() as executor:
-                futures = {executor.submit(processor, file): file for file in files_to_process}
+        processor = LVHProcessor(path_out_h5=args.dst, splits=splits, cone_params=cone_parameters)
+
+        log.info("Starting the conversion process.")
+
+        if not args.no_hyperthreading:
+            # DO NOT create a processor here for submission
+            with ProcessPoolExecutor(max_workers=min(64, os.cpu_count())) as executor:
+                futures = {
+                    executor.submit(
+                        _process_file_worker,
+                        str(file),  # avi_file
+                        args.dst,  # dst (Path or str)
+                        splits,  # splits (picklable dict of lists)
+                        cone_parameters,  # cone params dict (picklable)
+                        processor.range_from,  # only if needed; better pass primitives
+                        processor._process_range,
+                    ): file
+                    for file in files_to_process
+                }
                 for future in tqdm(as_completed(futures), total=len(files_to_process)):
                     try:
                         future.result()
                     except Exception as e:
-                        print(f"Error processing file: {str(e)}")
+                        log.error(f"Error processing file: {str(e)}")
         else:
+            log.info("Converting without hyperthreading")
             for file in tqdm(files_to_process):
                 try:
                     processor(file)
                 except Exception as e:
-                    print(f"Error processing {file}: {str(e)}")
+                    log.error(f"Error processing {file}: {str(e)}")
 
-        print("All image conversion tasks are completed.")
+        log.info("All image conversion tasks are completed.")
 
     # Convert measurements if requested
     if args.convert_measurements:
-        source_path = Path(args.source)
+        source_path = Path(src)
         measurements_csv = source_path / "MeasurementsList.csv"
         if measurements_csv.exists():
-            output_csv = Path(args.output) / "MeasurementsList.csv"
+            output_csv = Path(args.dst) / "MeasurementsList.csv"
             convert_measurements_csv(measurements_csv, output_csv, cone_params_csv)
         else:
-            print("Warning: MeasurementsList.csv not found in source directory")
+            log.warning("MeasurementsList.csv not found in source directory")
 
-    print("All tasks are completed.")
+    log.info("All tasks are completed.")

zea/data/convert/echonetlvh/manual_rejections.txt

@@ -0,0 +1,73 @@
+0X11181B7E91E930CF
+0X111C213706B3A146
+0X13CC23E3AA9AB751
+0X1C6CAFF5CB0299C0
+0X1D7074550AF6F017
+0X227355B3DA6C4B7A
+0X23869862AA2B7A6C
+0X25BD4F8A91A524ED
+0X2E241BF70EED0CFA
+0X310DD4EE1DA13B80
+0X3FA4417F5B45F121
+0X488149AC9B0C76E2
+0X4885907B96B1BAB4
+0X4DAD29CD33A5EF93
+0X51E5F62690062FDC
+0X52ADC74C8CC6B8BF
+0X53E2CDA115A9A9F0
+0X56E51030B25F8380
+0X5AA2CEBB492A7CBD
+0X5C012F7585F8D394
+0X5F7DF2135352442B
+0X61B21A4616889011
+0X67AE0A2349838A2F
+0X687C7044FE93C8EE
+0X6ADC75E01F62D32A
+0X6CDFB4E3FBDC001B
+0X6E149B6ED5DB1AE9
+0X73DF29FFB2B683F0
+0X7E8F8D3A80E492FE
+0X80B0B9C83894F2F3
+0X80F8DC9F3491FC58
+0X83FC5AE3242FD375
+0X8E3726C5DC9819E3
+0X9B3A721521F6F71D
+0X9D011879D04756AD
+0XA685D6BE72086604
+0XA8DD5136FEA6B60E
+0XAD914D49A9577C03
+0XAF4C4BD0D33C9CE8
+0XB27AF4716B649D0E
+0XB5E17F4A0497BD92
+0XB616C5D99F1DD73
+0XBF449F4AAC0FA294
+0XC64829FD975EFBAD
+0XC6BFC34800F8397D
+0XCF0F583953B37563
+0XD0537A6E187120FF
+0XD1357C20D9D00A17
+0XD7D8BE5D94D25347
+0XDAD3FA61A927124B
+0XDDC98ED577E19D03
+0XEC65FFB6F7DCDA70
+0XED0CDC255553A242
+0XF49D9E0C9D9B266B
+0XF5465778B6D691DB
+0XF7B8A81B72C467DB
+0XFA0D127A284CCA1A
+0X1C9440E0E259157E
+0X36A0E840E9CB0CE1
+0X4547546741BB73E8
+0X4FE26154CFCECDD7
+0X5D1ADF0993864E58
+0X5F2928BCB0F4B94C
+0X6935644F26AE8EA
+0X6A4C7F69D0046092
+0X716C88F4E946105A
+0XC0D63F1F9088CA81
+0XEBF79B6D25606A2D
+0XFFB05BDB3A97322
+0X7D8208DCF7643708
+0XA836692EB71E582B
+0XB40B08D1642CEEE
+0XF8D61AF7DE6AC700