refactor for new dataset design

Author: ryanccarelli

pathml/datasets/peso.py

@@ -3,17 +3,16 @@
 import shutil
 import sys
 import numpy as np
-import torch
-import torch.utils.data as data
+from torch.utils.data import Dataset
 from warnings import warn
 from pathlib import Path
 import cv2
 
-from pathml.datasets.base import BaseSlideDataset, BaseDataModule
+from pathml.datasets.base_data_module import BaseDataModule
 from pathml.datasets.utils import download_from_url
 from pathml.preprocessing.transforms import TissueDetectionHE
 from pathml.preprocessing.pipeline import Pipeline
-from pathml.core.slide_classes import HESlide
+from pathml.core.slide_data import HESlide
 from pathml.core.masks import Masks
 
 import cProfile
@@ -52,8 +51,8 @@ def _download_peso(self, download_dir):
             url = f'https://zenodo.org/record/1485967/files/'
             for file in files:
                 print(f"downloading {file}")
-                download_from_url(f"{url}{file}", f"{download_dir}") 
-            for root, _, files in os.walk(download_dir): 
+                download_from_url(f"{url}{file}", f"{download_dir}")
+            for root, _, files in os.walk(download_dir):
                 for file in files:
                     print(f"unzipping {file}")
                     if zipfile.is_zipfile(f"{root}/{file}"):
@@ -75,7 +74,7 @@ def _download_peso(self, download_dir):
                     profile.enable()
 
                     name = '_'.join(file.split('_')[:-1])
-                    maskpath = Path(name+'_HE_training_mask.tif') 
+                    maskpath = Path(name+'_HE_training_mask.tif')
                     mask = HESlide(filepath = str(Path(download_dir)/Path('peso_training_masks')/maskpath), name = name)
                     shape1, shape2 = mask.slide.get_image_shape()
                     shape = (shape2, shape1)
@@ -93,7 +92,7 @@ def _download_peso(self, download_dir):
                         ])
                     # TODO: choose tile size
                     wsi.run(pipeline, tile_size=256)
-                    
+
                     profile.disable()
                     ps = pstats.Stats(profile)
                     ps.print_stats()
@@ -111,7 +110,7 @@ def train_dataloader(self):
         """
         Dataloader for training set.
         """
-        return data.DataLoader(
+        return DataLoader(
             dataset = self._get_dataset(fold_ix = self.split),
             batch_size = self.batch_size,
             shuffle = self.shuffle
@@ -126,7 +125,7 @@ def valid_dataloader(self):
             fold_ix = 2
         else:
             fold_ix = 1
-        return data.DataLoader(
+        return DataLoader(
             self._get_dataset(fold_ix = fold_ix),
             batch_size = self.batch_size,
             shuffle = self.shuffle
@@ -141,20 +140,20 @@ def test_dataloader(self):
             fold_ix = 3
         else:
             fold_ix = 1
-        return data.DataLoader(
+        return DataLoader(
             self._get_dataset(fold_ix = fold_ix),
             batch_size = self.batch_size,
             shuffle = self.shuffle
         )
 
-class PesoDataset(BaseSlideDataset):
+class PesoDataset(Dataset):
     """
     Dataset object for Peso dataset.
     Raw data downloads:
     IHC color deconvolution (n=62) with p63, ck8/18 stainings
-    training masks (n=62) generated by segmenting IHC with UNet 
+    training masks (n=62) generated by segmenting IHC with UNet
     training masks corrected (n=25) with manual annotations
-    wsis (n=62) wsis at 0.48 \mu/pix
+    wsis (n=62) at 0.48 \mu/pix
 
     testing data:
     testset regions collection of xml files with outlines of test regions
@@ -178,14 +177,14 @@ def __init__(self,
         assert data_dir.isdir(), f"Error: data not found at {data_dir}"
 
         if not any(fname.endswith('.h5') for fname in os.listdir(self.data_dir / Path('h5'))):
-            raise Exception('must download dataset from pathml.datasets') 
-         
+            raise Exception('must download dataset from pathml.datasets')
+
         getitemdict = {}
         items = 0
         for h5 in os.listdir(self.data_dir / Path('h5')):
             wsi = read(self.data_dir / Path('h5') / Path(file))
             for i in range(len(wsi.tile)):
-                getitemdict[items] = (wsi.name, i) 
+                getitemdict[items] = (wsi.name, i)
                 items = items + 1
         self.getitemdict = getitemdict
         self.wsi = None
@@ -197,5 +196,5 @@ def __getitem__(self, ix):
         wsiname, index = self.getitemdict[ix]
         if self.wsi is None or self.wsi.name != wsiname:
             self.wsi = read(self.data_dir / Path('h5') / Path(wsiname + '.h5'))
-        tile = self.wsi.tiles[index] 
+        tile = self.wsi.tiles[index]
         return tile