finish survival prediction, test pass!

Author: yifanfeng97

HyperG/hyedge/gather_neighbor.py

@@ -77,7 +77,7 @@ def neighbor_distance(x: torch.Tensor, k_nearest, dis_metric=pairwise_euclidean_
     :return:
     """
 
-    assert len(x.shape) == 2, 'should be a tensor with (N x C) or (B x C x M x N)'
+    assert len(x.shape) == 2, 'should be a tensor with dimension (N x C)'
 
     # N x C
     node_num = x.size(0)

HyperG/models/BaseCNNs.py

@@ -4,9 +4,10 @@
 
 class ResNetFeature(nn.Module):
 
-    def __init__(self, depth=34, pretrained=True):
+    def __init__(self, depth=34, pooling=False, pretrained=True):
         super().__init__()
         assert depth in [18, 34, 50, 101, 152]
+        self.pooling = pooling
 
         if depth == 18:
             base_model = torchvision.models.resnet18(pretrained=pretrained)
@@ -34,8 +35,15 @@ def __init__(self, depth=34, pretrained=True):
     def forward(self, x):
         x = self.features(x)
 
-        # Attention! No reshape!
-        return x
+        if self.pooling:
+            # -> batch_size x C x N
+            x = x.view(x.size(0), x.size(1), -1)
+            # -> batch_size x C
+            x = x.mean(dim=-1)
+            return x
+        else:
+            # Attention! No reshape!
+            return x
 
 
 class ResNetClassifier(nn.Module):

HyperG/utils/data/pathology/overview_patches.py

@@ -13,17 +13,17 @@ def draw_patches_on_slide(slide_dir, patch_coors, mini_frac=32):
     mini_size = np.ceil(np.array(slide.level_dimensions[0]) / mini_frac).astype(np.int)
     mini_level = get_just_gt_level(slide, mini_size)
 
-    img = slide.read_region((0, 0), mini_level, slide.level_dimensions[mini_level])
+    img = slide.read_region((0, 0), mini_level, slide.level_dimensions[mini_level]).convert('RGB')
     img = img.resize(mini_size)
 
     sampled_mask = gather_sampled_patches(patch_coors, mini_size, mini_frac)
     sampled_patches_img = fuse_img_mask(np.asarray(img), sampled_mask)
 
     img.close()
-    return Image.fromarray(sampled_patches_img)
+    return sampled_patches_img
 
 
-def gather_sampled_patches(patch_coors, mini_size, mini_frac):
+def gather_sampled_patches(patch_coors, mini_size, mini_frac) -> np.array:
     # generate sampled area mask
     sampled_mask = np.zeros((mini_size[1], mini_size[0]), np.uint8)
     for _coor in patch_coors:
@@ -36,8 +36,8 @@ def gather_sampled_patches(patch_coors, mini_size, mini_frac):
     return sampled_mask
 
 
-def fuse_img_mask(img: np.array, mask: np.array, alpha=0.7):
-    assert img.shape == mask.shape
+def fuse_img_mask(img: np.array, mask: np.array, alpha=0.7) -> Image:
+    assert img.shape[:2] == mask.shape
     img = img.copy()
     if (mask != 0).any():
         img[mask != 0] = alpha * img[mask != 0] + \

HyperG/utils/data/pathology/sample_patches.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 import openslide
-import tqdm
+from tqdm import tqdm
 from scipy import ndimage
 from skimage.filters import threshold_otsu
 from skimage.morphology import dilation, star
@@ -44,7 +44,8 @@ def sample_patch_coors(slide_dir, num_sample=2000, patch_size=256):
     # half of the center
     th_num = int(np.ceil((mini_patch_size * 3 / 4 * mini_patch_size * 3 / 4)))
 
-    for row, col in tqdm(row_col):
+    pbar = tqdm(total=num_sample)
+    for row, col in row_col:
         if cnt >= num_sample:
             break
         mini_patch = bg_mask[row:row + mini_patch_size, col: col + mini_patch_size]
@@ -55,6 +56,8 @@ def sample_patch_coors(slide_dir, num_sample=2000, patch_size=256):
             #     continue
             patch_coors.append(origin)
             cnt += 1
+            pbar.update(1)
+    pbar.close()
 
     return patch_coors
 

HyperG/utils/meter/__init__.py

@@ -1,4 +1,4 @@
-from .class_error import ClassErrorMeter
 from .transductive import trans_class_acc, trans_iou_socre
+from .inductive import CIndexMeter
 
-__all__ = ['trans_class_acc', 'trans_iou_socre', 'ClassErrorMeter']
+__all__ = ['trans_class_acc', 'trans_iou_socre', 'CIndexMeter']

HyperG/utils/meter/class_error.py

@@ -0,0 +1,31 @@
+import numpy as np
+import torch
+
+
+class CIndexMeter:
+    def __init__(self):
+        super(CIndexMeter, self).__init__()
+        self.reset()
+
+    def reset(self):
+        self.output = np.array([])
+        self.target = np.array([])
+
+    def add(self, output: torch.tensor, target: torch.tensor):
+        output = output.cpu().detach().squeeze().numpy()[np.newaxis]
+        target = target.cpu().detach().squeeze().numpy()[np.newaxis]
+
+        assert output.ndim == target.ndim, 'target and output do not match'
+        assert output.ndim == 1
+
+        self.output = np.hstack([self.output, output])
+        self.target = np.hstack([self.target, target])
+
+    def value(self):
+        output = self.output[np.newaxis]
+        target = self.target[np.newaxis]
+
+        num_sample = output.shape[-1]
+        num_hit = (~((output.T > output) ^ (target.T > target))).sum()
+
+        return float(num_hit - num_sample) / float(num_sample * num_sample - num_sample)

HyperG/utils/meter/inductive.py

@@ -8,7 +8,6 @@
 import torch
 from extract_patch_feature import extract_ft
 from torch.utils.data import Dataset, DataLoader
-from torch.utils.data.dataset import T_co
 
 from HyperG.utils.data import split_id
 from HyperG.utils.data.pathology import sample_patch_coors, draw_patches_on_slide
@@ -28,15 +27,19 @@ def split_train_val(data_root, ratio=0.8, save_split_dir=None, resplit=False):
     survival_time_max = 0
     for full_dir in all_list:
         _id = get_id(full_dir)
+        all_dict[_id] = {}
+        st = int(lbls[_id])
         all_dict[_id]['img_dir'] = full_dir
-        all_dict[_id]['survival_time'] = lbls[_id]
+        all_dict[_id]['survival_time'] = st
         survival_time_max = survival_time_max \
-            if survival_time_max > lbls[_id] else lbls[_id]
+            if survival_time_max > st else st
 
     id_list = list(all_dict.keys())
     train_list, val_list = split_id(id_list, ratio)
 
-    result = {'survival_time_max': survival_time_max}
+    result = {'survival_time_max': survival_time_max,
+              'train': {},
+              'val': {}}
     for _id in train_list:
         result['train'][_id] = all_dict[_id]
     for _id in val_list:
@@ -52,40 +55,52 @@ def split_train_val(data_root, ratio=0.8, save_split_dir=None, resplit=False):
     return result
 
 
+# def tmp_get_split(data_root):
+#     def tmp_get_id(_dir):
+#         _num = int(osp.splitext(osp.split(_dir)[1])[0].split('_')[1])
+#         return f'TCGA-GBM-{_num}'
+#
+#     result = {'train': {}, 'val': {}}
+#     for phase in ['train', 'val']:
+#         glob.glob(osp.join(data_root, phase, '*.npy'))
+
+
 def preprocess(data_dict, patch_ft_dir, patch_coors_dir, num_sample=2000,
                patch_size=256, sampled_vis=None, mini_frac=32):
     # check if each slide patch feature exists
     all_dir_list = []
     for phase in ['train', 'val']:
-        for _dir in data_dict[phase]:
-            all_dir_list.append(_dir['img_dir'])
+        for _id in data_dict[phase].keys():
+            all_dir_list.append(data_dict[phase][_id]['img_dir'])
     to_do_list = check_patch_ft(all_dir_list, patch_ft_dir)
 
     if to_do_list is not None:
         for _idx, _dir in enumerate(to_do_list):
-            print(f'processing {_idx + 1}/{len(to_do_list)}...')
+            print(f'{_idx + 1}/{len(to_do_list)}: processing slide {_dir}...')
+
+            print(f'sampling patch...')
             _id = get_id(_dir)
             _patch_coors = sample_patch_coors(_dir, num_sample=2000, patch_size=256)
 
             # save sampled patch coordinates
-            with open(osp.join(patch_coors_dir, f'{_id}_coors.pkl')) as fp:
+            with open(osp.join(patch_coors_dir, f'{_id}_coors.pkl'), 'wb') as fp:
                 pickle.dump(_patch_coors, fp)
 
             # visualize sampled patches on slide
             if sampled_vis is not None:
+                _vis_img_dir = osp.join(sampled_vis, f'{_id}_sampled_patches.jpg')
+                print(f'saving sampled patch_slide visualization {_vis_img_dir}...')
                 _vis_img = draw_patches_on_slide(_dir, _patch_coors, mini_frac=32)
-                with open(osp.join(sampled_vis, f'{_id}_sampled_patches.jpg')) as fp:
+                with open(_vis_img_dir, 'w') as fp:
                     _vis_img.save(fp)
 
-    # extract patch feature for each slide
-    for _dir in all_dir_list:
-        _id = get_id(_dir)
-        _patch_coors = None
-        fts = extract_ft(_dir, _patch_coors)
-        np.save(osp.join(patch_ft_dir, f'{_id}_fts.npy'), fts.cpu().numpy())
+            # extract patch feature for each slide
+            print(f'extracting feature...')
+            fts = extract_ft(_dir, _patch_coors, depth=34, batch_size=512)
+            np.save(osp.join(patch_ft_dir, f'{_id}_fts.npy'), fts.cpu().numpy())
 
 
-def get_dataloader(data_dict, patch_ft_dir):
+def get_dataloaders(data_dict, patch_ft_dir):
     all_ft_list = glob.glob(osp.join(patch_ft_dir, '*_fts.npy'))
 
     ft_dict = {}
@@ -98,7 +113,8 @@ def get_dataloader(data_dict, patch_ft_dir):
                                         shuffle=True, num_workers=4)
                       for phase in ['train', 'val']}
     dataset_size = {phase: len(SP_datasets[phase]) for phase in ['train', 'val']}
-    return SP_dataloaders, dataset_size
+    len_ft = SP_datasets['train'][0][0].size(1)
+    return SP_dataloaders, dataset_size, len_ft
 
 
 class SlidePatch(Dataset):
@@ -113,7 +129,7 @@ def __init__(self, data_dict: dict, ft_dict, survival_time_max):
     def __getitem__(self, idx: int):
         id = self.id_list[idx]
         fts = torch.tensor(np.load(self.ft_dict[id])).float()
-        st = torch.tensor(self.data_dict[id]['survival_time_max']).float()
+        st = torch.tensor(self.data_dict[id]['survival_time']).float()
         return fts, st / self.st_max
 
     def __len__(self) -> int:
@@ -122,7 +138,7 @@ def __len__(self) -> int:
 
 def check_patch_ft(dir_list, patch_ft_dir):
     to_do_list = []
-    done_list = glob.glob(osp.join(patch_ft_dir, '*_ft.npy'))
+    done_list = glob.glob(osp.join(patch_ft_dir, '*_fts.npy'))
     done_list = [get_id(_dir) for _dir in done_list]
     for _dir in dir_list:
         id = get_id(_dir)

examples/regression/survival_prediction/data_helper.py

@@ -0,0 +1,54 @@
+import openslide
+import torch
+from tqdm import tqdm
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+
+from HyperG.models import ResNetFeature
+
+
+def extract_ft(slide_dir: str, patch_coors, depth=34, batch_size=16):
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    slide = openslide.open_slide(slide_dir)
+
+    model_ft = ResNetFeature(depth=depth, pooling=True, pretrained=True)
+    model_ft = model_ft.to(device)
+    model_ft.eval()
+
+    dataset = Patches(slide, patch_coors)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4)
+
+    fts = []
+    with tqdm(total=len(dataset)) as pbar:
+        for _patches in dataloader:
+            _patches = _patches.to(device)
+            with torch.no_grad():
+                _fts = model_ft(_patches)
+            fts.append(_fts)
+            pbar.update(_patches.size(0))
+
+    fts = torch.cat(fts, dim=0)
+    assert fts.size(0) == len(patch_coors)
+    return fts
+
+
+class Patches(Dataset):
+
+    def __init__(self, slide: openslide, patch_coors) -> None:
+        super().__init__()
+        self.slide = slide
+        self.patch_coors = patch_coors
+        self.transform = transforms.Compose([
+            transforms.Resize(224),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                 std=[0.229, 0.224, 0.225])
+        ])
+
+    def __getitem__(self, idx: int):
+        coor = self.patch_coors[idx]
+        img = self.slide.read_region((coor[0], coor[1]), 0, (coor[2], coor[3])).convert('RGB')
+        return self.transform(img)
+
+    def __len__(self) -> int:
+        return len(self.patch_coors)