add test for slice indexing for builtin datasets

test/test_databases/test_afdb.py

@@ -30,7 +30,8 @@
 
 with pytest.warns(RuntimeWarning):
     reader = AFDB(_CWD, verbose=3)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestAFDB:

test/test_databases/test_apnea_ecg.py

@@ -29,7 +29,8 @@
 
 with pytest.warns(RuntimeWarning):
     reader = ApneaECG(_CWD)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestApneaECG:

test/test_databases/test_cinc2017.py

@@ -26,7 +26,8 @@
 
 with pytest.warns(RuntimeWarning):
     reader = CINC2017(_CWD)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestCINC2017:

test/test_databases/test_cinc2020.py

@@ -276,6 +276,11 @@ def test_getitem(self):
             )
             assert target.ndim == 1 and target.shape == (len(config.classes),)
 
+        # test slice indexing
+        data, target = ds[:2]
+        assert data.shape == (2, len(config.leads), config.input_len)
+        assert target.shape == (2, len(config.classes))
+
     def test_load_one_record(self):
         for rec in ds.records:
             data, target = ds._load_one_record(rec)

test/test_databases/test_cinc2021.py

@@ -299,6 +299,11 @@ def test_getitem(self):
             )
             assert target.ndim == 1 and target.shape == (len(config.classes),)
 
+        # test slice indexing
+        data, target = ds[:2]
+        assert data.shape == (2, len(config.leads), config.input_len)
+        assert target.shape == (2, len(config.classes))
+
     def test_load_one_record(self):
         for rec in ds.records:
             data, target = ds._load_one_record(rec)

test/test_databases/test_cpsc2019.py

@@ -28,7 +28,8 @@
 
 
 reader = CPSC2019(_CWD)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestCPSC2019:
@@ -144,5 +145,14 @@ def test_getitem(self):
             assert data.ndim == 2 and data.shape == (1, config_1.input_len)
             assert bin_mask.ndim == 2 and bin_mask.shape == (config_1.input_len, 1)
 
+        # test slice indexing
+        data, bin_mask = ds[:2]
+        assert data.ndim == 3 and data.shape == (2, 1, config.input_len)
+        assert bin_mask.ndim == 3 and bin_mask.shape == (
+            2,
+            config.input_len // config.reduction,
+            1,
+        )
+
     def test_properties(self):
         assert str(ds) == repr(ds)

test/test_databases/test_cpsc2020.py

@@ -27,7 +27,8 @@
 
 
 reader = CPSC2020(_CWD)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestCPSC2020:

test/test_databases/test_cpsc2021.py

@@ -389,6 +389,12 @@ def test_getitem(self):
             assert data.shape == (config.n_leads, input_len)
             assert qrs_mask.shape == (input_len, 1)
 
+        # test slice indexing
+        data, af_mask, weight_mask = ds[:2]
+        assert data.ndim == 3 and data.shape == (2, config.n_leads, input_len)
+        assert af_mask.ndim == 3 and af_mask.shape == (2, input_len, 1)
+        assert weight_mask.ndim == 3 and weight_mask.shape == (2, input_len, 1)
+
     def test_properties(self):
         assert ds.task == "main"
         assert ds_1.task == "rr_lstm"

test/test_databases/test_ltafdb.py

@@ -28,7 +28,8 @@
 
 with pytest.warns(RuntimeWarning):
     reader = LTAFDB(_CWD)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestLTAFDB:

test/test_databases/test_ludb.py

@@ -32,7 +32,8 @@
 
 
 reader = LUDB(_CWD)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestLUDB:
@@ -140,6 +141,14 @@ def test_getitem(self):
             assert signals.shape == (config.n_leads, config.input_len)
             assert labels.shape == (config.input_len, len(config.classes))
 
+        # test slice indexing
+        signals, labels = ds[:2]
+        assert signals.shape == (2, config.n_leads, config.input_len)
+        # NOTE that the (segmentation) labels have collapsed lead dimension
+        # so the shape is (n_samples, signal_len, n_classes)
+        # instead of (n_samples, n_leads, signal_len, n_classes)
+        assert labels.shape == (2, config.input_len, len(config.classes))
+
     def test_properties(self):
         signals_shape = ds.signals.shape  # (n_samples, n_leads, signal_len)
         labels_shape = ds.labels.shape  # (n_samples, n_leads, signal_len, n_classes)

test/test_databases/test_mitdb.py

@@ -30,7 +30,8 @@
 
 
 reader = MITDB(_CWD)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestMITDB:
@@ -164,6 +165,11 @@ def test_getitem(self):
 
         # `ds_rhythm` and `ds_af` have bugs now
 
+        # test slice indexing
+        data, ann = ds[:2]
+        assert data.shape == (2, config.n_leads, config[TASK].input_len)
+        assert ann.shape == (2, config[TASK].input_len, 1)
+
     def test_load_seg_data(self):
         seg = ds.all_segments[list(ds.all_segments)[0]][0]
         data = ds._load_seg_data(seg)

test/test_databases/test_ptb_xl.py

@@ -27,7 +27,8 @@
 
 
 reader = PTBXL(_CWD)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestPTBXL:

test/test_databases/test_qtdb.py

@@ -26,7 +26,8 @@
 
 
 reader = QTDB(_CWD)
-reader.download()
+if len(reader) == 0:
+    reader.download()
 
 
 class TestQTDB:

torch_ecg/databases/datasets/ludb/ludb_dataset.py

@@ -85,6 +85,8 @@ def __len__(self) -> int:
         return len(self.records)
 
     def __getitem__(self, index: Union[int, slice]) -> Tuple[np.ndarray, np.ndarray]:
+        if isinstance(index, slice):
+            return collate_fn([self[i] for i in range(*index.indices(len(self)))])
         if self.config.use_single_lead:
             rec_idx, lead_idx = divmod(index, len(self.leads))
         else:
@@ -100,7 +102,8 @@ def __getitem__(self, index: Union[int, slice]) -> Tuple[np.ndarray, np.ndarray]
             labels = labels[lead_idx, ...]
         else:
             # merge labels in all leads to one
-            # TODO: map via self.waveform_priority
+            # TODO: map via self.waveform_priority,
+            # or make it configurable whether to collapse the lead dimension
             labels = np.max(labels, axis=0)
         sampfrom = randint(self.config.start_from, signals.shape[1] - self.config.end_at - self.siglen)
         sampto = sampfrom + self.siglen