"""EchoNet-Dynamic Dataset."""
import os
import collections
import pandas
import numpy as np
import skimage.draw
import torchvision
import echonet
class Echo(torchvision.datasets.VisionDataset):
"""EchoNet-Dynamic Dataset.
Args:
root (string): Root directory of dataset (defaults to `echonet.config.DATA_DIR`)
split (string): One of {``train'', ``val'', ``test'', ``all'', or ``external_test''}
target_type (string or list, optional): Type of target to use,
``Filename'', ``EF'', ``EDV'', ``ESV'', ``LargeIndex'',
``SmallIndex'', ``LargeFrame'', ``SmallFrame'', ``LargeTrace'',
or ``SmallTrace''
Can also be a list to output a tuple with all specified target types.
The targets represent:
``Filename'' (string): filename of video
``EF'' (float): ejection fraction
``EDV'' (float): end-diastolic volume
``ESV'' (float): end-systolic volume
``LargeIndex'' (int): index of large (diastolic) frame in video
``SmallIndex'' (int): index of small (systolic) frame in video
``LargeFrame'' (np.array shape=(3, height, width)): normalized large (diastolic) frame
``SmallFrame'' (np.array shape=(3, height, width)): normalized small (systolic) frame
``LargeTrace'' (np.array shape=(height, width)): left ventricle large (diastolic) segmentation
value of 0 indicates pixel is outside left ventricle
1 indicates pixel is inside left ventricle
``SmallTrace'' (np.array shape=(height, width)): left ventricle small (systolic) segmentation
value of 0 indicates pixel is outside left ventricle
1 indicates pixel is inside left ventricle
Defaults to ``EF''.
mean (int, float, or np.array shape=(3,), optional): means for all (if scalar) or each (if np.array) channel.
Used for normalizing the video. Defaults to 0 (video is not shifted).
std (int, float, or np.array shape=(3,), optional): standard deviation for all (if scalar) or each (if np.array) channel.
Used for normalizing the video. Defaults to 0 (video is not scaled).
length (int or None, optional): Number of frames to clip from video. If ``None'', longest possible clip is returned.
Defaults to 16.
period (int, optional): Sampling period for taking a clip from the video (i.e. every ``period''-th frame is taken)
Defaults to 2.
max_length (int or None, optional): Maximum number of frames to clip from video (main use is for shortening excessively
long videos when ``length'' is set to None). If ``None'', shortening is not applied to any video.
Defaults to 250.
clips (int, optional): Number of clips to sample. Main use is for test-time augmentation with random clips.
Defaults to 1.
pad (int or None, optional): Number of pixels to pad all frames on each side (used as augmentation).
and a window of the original size is taken. If ``None'', no padding occurs.
Defaults to ``None''.
noise (float or None, optional): Fraction of pixels to black out as simulated noise. If ``None'', no simulated noise is added.
Defaults to ``None''.
target_transform (callable, optional): A function/transform that takes in the target and transforms it.
external_test_location (string): Path to videos to use for external testing.
"""
def __init__(self, root=None,
split="train", target_type="EF",
mean=0., std=1.,
length=16, period=2,
max_length=250,
clips=1,
pad=None,
noise=None,
target_transform=None,
external_test_location=None):
if root is None:
root = echonet.config.DATA_DIR
super().__init__(root, target_transform=target_transform)
self.split = split.upper()
if not isinstance(target_type, list):
target_type = [target_type]
self.target_type = target_type
self.mean = mean
self.std = std
self.length = length
self.max_length = max_length
self.period = period
self.clips = clips
self.pad = pad
self.noise = noise
self.target_transform = target_transform
self.external_test_location = external_test_location
self.fnames, self.outcome = [], []
if self.split == "EXTERNAL_TEST":
self.fnames = sorted(os.listdir(self.external_test_location))
else:
# Load video-level labels
with open(os.path.join(self.root, "FileList.csv")) as f:
data = pandas.read_csv(f)
data["Split"].map(lambda x: x.upper())
if self.split != "ALL":
data = data[data["Split"] == self.split]
self.header = data.columns.tolist()
self.fnames = data["FileName"].tolist()
self.fnames = [fn + ".avi" for fn in self.fnames if os.path.splitext(fn)[1] == ""] # Assume avi if no suffix
self.outcome = data.values.tolist()
# Check that files are present
missing = set(self.fnames) - set(os.listdir(os.path.join(self.root, "Videos")))
if len(missing) != 0:
print("{} videos could not be found in {}:".format(len(missing), os.path.join(self.root, "Videos")))
for f in sorted(missing):
print("\t", f)
raise FileNotFoundError(os.path.join(self.root, "Videos", sorted(missing)[0]))
# Load traces
self.frames = collections.defaultdict(list)
self.trace = collections.defaultdict(_defaultdict_of_lists)
with open(os.path.join(self.root, "VolumeTracings.csv")) as f:
header = f.readline().strip().split(",")
assert header == ["FileName", "X1", "Y1", "X2", "Y2", "Frame"]
for line in f:
filename, x1, y1, x2, y2, frame = line.strip().split(',')
x1 = float(x1)
y1 = float(y1)
x2 = float(x2)
y2 = float(y2)
frame = int(frame)
if frame not in self.trace[filename]:
self.frames[filename].append(frame)
self.trace[filename][frame].append((x1, y1, x2, y2))
for filename in self.frames:
for frame in self.frames[filename]:
self.trace[filename][frame] = np.array(self.trace[filename][frame])
# A small number of videos are missing traces; remove these videos
keep = [len(self.frames[f]) >= 2 for f in self.fnames]
self.fnames = [f for (f, k) in zip(self.fnames, keep) if k]
self.outcome = [f for (f, k) in zip(self.outcome, keep) if k]
def __getitem__(self, index):
# Find filename of video
if self.split == "EXTERNAL_TEST":
video = os.path.join(self.external_test_location, self.fnames[index])
elif self.split == "CLINICAL_TEST":
video = os.path.join(self.root, "ProcessedStrainStudyA4c", self.fnames[index])
else:
video = os.path.join(self.root, "Videos", self.fnames[index])
# Load video into np.array
video = echonet.utils.loadvideo(video).astype(np.float32)
# Add simulated noise (black out random pixels)
# 0 represents black at this point (video has not been normalized yet)
if self.noise is not None:
n = video.shape[1] * video.shape[2] * video.shape[3]
ind = np.random.choice(n, round(self.noise * n), replace=False)
f = ind % video.shape[1]
ind //= video.shape[1]
i = ind % video.shape[2]
ind //= video.shape[2]
j = ind
video[:, f, i, j] = 0
# Apply normalization
if isinstance(self.mean, (float, int)):
video -= self.mean
else:
video -= self.mean.reshape(3, 1, 1, 1)
if isinstance(self.std, (float, int)):
video /= self.std
else:
video /= self.std.reshape(3, 1, 1, 1)
# Set number of frames
c, f, h, w = video.shape
if self.length is None:
# Take as many frames as possible
length = f // self.period
else:
# Take specified number of frames
length = self.length
if self.max_length is not None:
# Shorten videos to max_length
length = min(length, self.max_length)
if f < length * self.period:
# Pad video with frames filled with zeros if too short
# 0 represents the mean color (dark grey), since this is after normalization
video = np.concatenate((video, np.zeros((c, length * self.period - f, h, w), video.dtype)), axis=1)
c, f, h, w = video.shape # pylint: disable=E0633
if self.clips == "all":
# Take all possible clips of desired length
start = np.arange(f - (length - 1) * self.period)
else:
# Take random clips from video
start = np.random.choice(f - (length - 1) * self.period, self.clips)
# Gather targets
target = []
for t in self.target_type:
key = self.fnames[index]
if t == "Filename":
target.append(self.fnames[index])
elif t == "LargeIndex":
# Traces are sorted by cross-sectional area
# Largest (diastolic) frame is last
target.append(np.int(self.frames[key][-1]))
elif t == "SmallIndex":
# Largest (diastolic) frame is first
target.append(np.int(self.frames[key][0]))
elif t == "LargeFrame":
target.append(video[:, self.frames[key][-1], :, :])
elif t == "SmallFrame":
target.append(video[:, self.frames[key][0], :, :])
elif t in ["LargeTrace", "SmallTrace"]:
if t == "LargeTrace":
t = self.trace[key][self.frames[key][-1]]
else:
t = self.trace[key][self.frames[key][0]]
x1, y1, x2, y2 = t[:, 0], t[:, 1], t[:, 2], t[:, 3]
x = np.concatenate((x1[1:], np.flip(x2[1:])))
y = np.concatenate((y1[1:], np.flip(y2[1:])))
r, c = skimage.draw.polygon(np.rint(y).astype(np.int), np.rint(x).astype(np.int), (video.shape[2], video.shape[3]))
mask = np.zeros((video.shape[2], video.shape[3]), np.float32)
mask[r, c] = 1
target.append(mask)
else:
if self.split == "CLINICAL_TEST" or self.split == "EXTERNAL_TEST":
target.append(np.float32(0))
else:
target.append(np.float32(self.outcome[index][self.header.index(t)]))
if target != []:
target = tuple(target) if len(target) > 1 else target[0]
if self.target_transform is not None:
target = self.target_transform(target)
# Select clips from video
video = tuple(video[:, s + self.period * np.arange(length), :, :] for s in start)
if self.clips == 1:
video = video[0]
else:
video = np.stack(video)
if self.pad is not None:
# Add padding of zeros (mean color of videos)
# Crop of original size is taken out
# (Used as augmentation)
c, l, h, w = video.shape
temp = np.zeros((c, l, h + 2 * self.pad, w + 2 * self.pad), dtype=video.dtype)
temp[:, :, self.pad:-self.pad, self.pad:-self.pad] = video # pylint: disable=E1130
i, j = np.random.randint(0, 2 * self.pad, 2)
video = temp[:, :, i:(i + h), j:(j + w)]
return video, target
def __len__(self):
return len(self.fnames)
def extra_repr(self) -> str:
"""Additional information to add at end of __repr__."""
lines = ["Target type: {target_type}", "Split: {split}"]
return '\n'.join(lines).format(**self.__dict__)
def _defaultdict_of_lists():
"""Returns a defaultdict of lists.
This is used to avoid issues with Windows (if this function is anonymous,
the Echo dataset cannot be used in a dataloader).
"""
return collections.defaultdict(list)
Datasets
Models
