"""
Script to convert the EchoNet database to zea format.
.. note::
Will segment the images and convert them to polar coordinates.
For more information about the dataset, resort to the following links:
- 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/dynamic/>`_.
"""
import os
from concurrent.futures import ProcessPoolExecutor, as_completed
from multiprocessing import Value
from pathlib import Path
from typing import TYPE_CHECKING
import numpy as np
import yaml
from scipy.interpolate import griddata
from tqdm import tqdm
from zea import log
from zea.data.convert.utils import load_avi, unzip
from zea.data.file import File
from zea.func.tensor import translate
if TYPE_CHECKING:
from multiprocessing.sharedctypes import Synchronized
# Process-shared counter, bound per worker process by ``count_init`` (see below)
# and read by ``H5Processor``. Annotation-only on purpose: the name genuinely
# does not exist until an initializer runs, but declaring it here gives the
# ``global COUNTER`` statement a module-scope declaration to rebind.
COUNTER: "Synchronized[int]"
[docs]
def segment(tensor, number_erasing=0, min_clip=0):
"""Segments the background of the echonet images by setting it to 0 and creating a hard edge.
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
"""
# Start with the upper part
# Height of the diagonal lines for the columns [0, 112]
rows_left = np.linspace(67, 7, 61)
rows_right = np.linspace(7, 57, 51)
rows = np.concatenate([rows_left, rows_right], axis=0)
for idx, row in enumerate(rows.astype(np.int32)):
# Set everything above the edge to the number_erasing value.
# Rows count up from 0 to 112 so row-1 is above.
tensor[:, 0 : row - 1, idx] = number_erasing
# Set minimum values for the edge
if min_clip > 0:
tensor[:, row, idx] = np.clip(tensor[:, row, idx], min_clip, 1)
# Bottom left curve (manual fit)
cols_left = np.linspace(0, 20, 21).astype(np.int32)
rows_left = np.array(
[
102,
103,
103,
104,
104,
105,
105,
106,
106,
107,
107,
107,
108,
108,
109,
109,
109,
110,
110,
111,
111,
]
)
# Bottom right curve (manual fit)
cols_right = np.linspace(89, 111, 23).astype(np.int32)
rows_right = np.array(
[
111,
111,
111,
110,
110,
110,
109,
109,
109,
108,
108,
107,
107,
107,
106,
106,
105,
105,
104,
104,
103,
103,
102,
]
)
rows = np.concatenate([rows_left, rows_right], axis=0)
cols = np.concatenate([cols_left, cols_right], axis=0)
for row, col in zip(rows, cols):
# Set everything under the edge to the number_erasing value.
# Rows count up from 0 to 112 so row-1 is above.
tensor[:, row:, col] = number_erasing
# Set minimum values for the edge
if min_clip > 0:
tensor[:, row - 1, col] = np.clip(tensor[:, row - 1, col], min_clip, 1)
return tensor
[docs]
def accept_shape(tensor):
"""Acceptance algorithm that determines whether to reject an image
based on left and right corner data.
Args:
tensor (ndarray): Input image (sc) with 2 dimensions. (112, 112)
Returns:
decision (bool): Whether or not the tensor should be rejected.
"""
decision = True
# Test one, check if left bottom corner is populated with values
rows_lower = np.linspace(78, 47, 21).astype(np.int32)
rows_upper = np.linspace(67, 47, 21).astype(np.int32)
counter = 0
for idx, row in enumerate(rows_lower):
counter += np.sum(tensor[rows_upper[idx] : row, idx])
# If it is not populated, reject the image
if counter < 0.1:
decision = False
# Test two, check if the bottom right cornered with values (that are not artifacts)
cols = np.linspace(70, 111, 42).astype(np.int32)
rows_bot = np.linspace(17, 57, 42).astype(np.int32)
rows_top = np.linspace(17, 80, 42).astype(np.int32)
# List all the values
counter = []
for i, col in enumerate(cols):
counter += [tensor[rows_bot[i] : rows_top[i], col]]
flattened_counter = [float(item) for sublist in counter for item in sublist]
# Sort and exclude the first 50 (likely artifacts)
flattened_counter.sort(reverse=True)
value = sum(flattened_counter[100:])
# Reject if the baseline is too low
if value < 5:
decision = False
return decision
[docs]
def rotate_coordinates(data_points, degrees):
"""Function that rotates the datapoints by a certain degree.
Args:
data_points (ndarray): tensor containing [N,2] (x and y) datapoints.
degrees (int): angle to rotate the datapoints with
Returns:
rotated_points (ndarray): the rotated data_points.
"""
angle_radians = np.radians(degrees)
cos_angle = np.cos(angle_radians)
sin_angle = np.sin(angle_radians)
rotation_matrix = np.array([[cos_angle, -sin_angle], [sin_angle, cos_angle]])
rotated_points = rotation_matrix @ data_points.T
return rotated_points.T
[docs]
def cartesian_to_polar_matrix(
cartesian_matrix, tip=(61, 7), r_max=107, angle=0.79, interpolation="nearest"
):
"""
Function that converts a timeseries of a cartesian cone to a polar representation
that is more compatible with CNN's/action selection.
Args:
- cartesian_matrix (2d array): (rows, cols) matrix containing time sequence
of scan-converted (Cartesian) image data.
- tip (tuple, optional): coordinates (in indices) of the tip of the cone.
Defaults to (61, 7).
- r_max (int, optional): expected radius of the cone. Defaults to 107.
- angle (float, optional): expected angle of the cone, will be used as (-angle, angle).
Defaults to 0.79.
- interpolation (str, optional): can be [nearest, linear, cubic]. Defaults to 'nearest'.
Returns:
polar_matrix (2d array): polar conversion of the input.
"""
rows, cols = cartesian_matrix.shape
center_x, center_y = tip
# Create cartesian coordinates of the image data
x = np.linspace(-center_x, cols - center_x - 1, cols)
y = np.linspace(-center_y, rows - center_y - 1, rows)
x, y = np.meshgrid(x, y)
# Flatten the grid and values
data_points = np.column_stack((x.ravel(), y.ravel()))
data_points = rotate_coordinates(data_points, -90)
data_values = cartesian_matrix.ravel()
# Define new points to sample from in the region of the data.
# R_max and Theta are found manually. R_max differs from the number of rows in EchoNet!
r = np.linspace(0, r_max, rows)
theta = np.linspace(-angle, angle, cols)
r, theta = np.meshgrid(r, theta)
x_polar = r * np.cos(theta)
y_polar = r * np.sin(theta)
new_points = np.column_stack((x_polar.ravel(), y_polar.ravel()))
# Interpolate and reshape to 2D matrix
polar_values = griddata(
data_points, data_values, new_points, method=interpolation, fill_value=0
)
polar_matrix = np.rot90(polar_values.reshape(cols, rows), k=-1)
return polar_matrix
[docs]
def find_split_for_file(file_dict, target_file):
"""
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"
[docs]
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
[docs]
class H5Processor:
"""
Stores a few variables and paths to allow for hyperthreading.
"""
def __init__(
self,
path_out_h5,
num_val=500,
num_test=500,
range_from=(0, 255),
range_to=(-60, 0),
splits=None,
):
self.path_out_h5 = Path(path_out_h5)
self.num_val = num_val
self.num_test = num_test
self.range_from = range_from
self.range_to = range_to
self.splits = splits
self._process_range = (0, 1)
# Ensure train, val, test, rejected paths exist
for folder in ["train", "val", "test", "rejected"]:
(self.path_out_h5 / folder).mkdir(parents=True, exist_ok=True)
def _translate(self, data):
"""Translate the data from the processing range to final range."""
return translate(data, self._process_range, self.range_to)
[docs]
def get_split(self, hdf5_file: str, sequence):
"""
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])
# Previous split
if self.splits is not None:
split = find_split_for_file(self.splits, hdf5_file)
assert accepted == (split != "rejected"), "Rejection mismatch"
return split
# New split
if not accepted:
return "rejected"
# 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 n <= self.num_val:
return "val"
elif n <= self.num_val + self.num_test:
return "test"
else:
return "train"
[docs]
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."
)
[docs]
def __call__(self, avi_file):
"""
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 and saves as a zea HDF5 file via File.create.
Args:
avi_file (pathlib.Path): Path to the source .avi file to process.
"""
hdf5_file = avi_file.stem + ".hdf5"
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]}"
# Translate to [0, 1]
sequence = translate(sequence, self.range_from, self._process_range)
sequence = segment(sequence, number_erasing=0, min_clip=0)
split = self.get_split(hdf5_file, sequence)
accepted = split != "rejected"
out_h5 = self.path_out_h5 / split / hdf5_file
polar_im_set = []
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)
polar_im_set.append(polar_im)
if accepted:
polar_im_set = np.stack(polar_im_set, axis=0)
# Check the ranges
assert sequence.min() >= self._process_range[0], sequence.min()
assert sequence.max() <= self._process_range[1], sequence.max()
if accepted:
# Store the polar (pre-scan-conversion) representation as the image.
polar_db = self._translate(polar_im_set)
polar_float32 = polar_db.astype(np.float32)
polar_float32 = np.expand_dims(polar_float32, axis=-1) # add y dim
data = {"image": {"values": polar_float32}}
else:
# Rejected sequences have no polar representation; store the original
# scan-converted (Cartesian) frames as the image instead.
cartesian_db = self._translate(sequence).astype(np.float32)
data = {"image": {"values": cartesian_db}}
File.create(
path=out_h5,
data=data,
probe={"name": "generic"},
description="EchoNet dataset converted to zea format",
)
def _resolve_path(src: str | Path) -> Path:
src = Path(src)
zip_name = "EchoNet-Dynamic.zip"
folder_name = "EchoNet-Dynamic"
unzip_dir = src / folder_name / "Videos"
if (src / folder_name).exists():
return unzip_dir
unzipped_path = unzip(src / zip_name, src)
return unzipped_path / folder_name / "Videos"
[docs]
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 = _resolve_path(args.src)
if args.split_path is not None:
# Reproduce a previous split...
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.dst):
for filename in filenames:
files_done.append(filename.replace(".hdf5", ""))
# List all files of echonet and exclude those already processed
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]
log.info(f"Files left to process: {len(h5_files)}")
# Run the processor
processor = H5Processor(path_out_h5=args.dst, splits=splits)
log.info("Starting the conversion process.")
if not args.no_hyperthreading:
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)
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")