writing extracting patch feature functions.

Author: yifanfeng97

.idea/.gitignore

@@ -1,2 +1 @@
-from .mri.io import read_mri_series, save_mri_series
 from .data_helper import split_id

HyperG/utils/data/__init__.py

@@ -1,4 +1,6 @@
 from random import shuffle
+import pickle
+import os.path as osp
 
 
 def split_id(id_list, ratio):
@@ -7,4 +9,7 @@ def split_id(id_list, ratio):
 
     id_train = id_list[:train_len]
     id_val = id_list[train_len:]
-    return id_train, id_val
+    return id_train, id_val
+
+
+

HyperG/utils/data/data_helper.py

@@ -0,0 +1,3 @@
+from .io import read_mri_series, save_mri_series
+
+__all__ = ['read_mri_series', 'save_mri_series']

HyperG/utils/data/mri/__init__.py

@@ -0,0 +1,4 @@
+from .overview_patches import draw_patches_on_slide
+from .sample_patches import sample_patch_coors
+
+__all__ = ['sample_patch_coors', 'draw_patches_on_slide']

HyperG/utils/data/pathology/__init__.py

@@ -0,0 +1,45 @@
+import numpy as np
+import openslide
+from PIL import Image
+
+from .sample_patches import get_just_gt_level
+
+SAMPLED = 2
+SAMPLED_COLOR = [0, 0, 255]
+
+
+def draw_patches_on_slide(slide_dir, patch_coors, mini_frac=32):
+    slide = openslide.open_slide(slide_dir)
+    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 = 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)
+
+
+def gather_sampled_patches(patch_coors, mini_size, mini_frac):
+    # generate sampled area mask
+    sampled_mask = np.zeros((mini_size[1], mini_size[0]), np.uint8)
+    for _coor in patch_coors:
+        _mini_coor = (int(_coor[0] / mini_frac), int(_coor[1] / mini_frac))
+        _mini_patch_size = (int(_coor[2] / mini_frac), int(_coor[3] / mini_frac))
+        sampled_mask[_mini_coor[1]:_mini_coor[1] + _mini_patch_size[1],
+        _mini_coor[0]:_mini_coor[0] + _mini_patch_size[0]] = SAMPLED
+    sampled_mask = np.asarray(Image.fromarray(sampled_mask).resize(mini_size))
+
+    return sampled_mask
+
+
+def fuse_img_mask(img: np.array, mask: np.array, alpha=0.7):
+    assert img.shape == mask.shape
+    img = img.copy()
+    if (mask != 0).any():
+        img[mask != 0] = alpha * img[mask != 0] + \
+                         (1 - alpha) * np.array(SAMPLED_COLOR)
+    return Image.fromarray(img)

HyperG/utils/data/pathology/overview_patches.py

@@ -0,0 +1,107 @@
+import os.path as osp
+from itertools import product
+from random import shuffle
+
+import numpy as np
+import openslide
+import tqdm
+from scipy import ndimage
+from skimage.filters import threshold_otsu
+from skimage.morphology import dilation, star
+
+BACKGROUND = 0
+FOREGROUND = 1
+
+
+def sample_patch_coors(slide_dir, num_sample=2000, patch_size=256):
+    slide = openslide.open_slide(slide_dir)
+    slide_name = osp.basename(slide_dir)
+    slide_name = slide_name[:slide_name.rfind('.')]
+
+    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)
+    mini_patch_size = patch_size // mini_frac
+
+    bg_mask = generate_background_mask(slide, mini_level, mini_size)
+    assert bg_mask.shape == (mini_size[1], mini_size[0])
+
+    # extract patches from available area
+    patch_coors = []
+    num_row, num_col = bg_mask.shape
+    num_row = num_row - mini_patch_size
+    num_col = num_col - mini_patch_size
+
+    row_col = list(product(range(num_row), range(num_col)))
+    shuffle(row_col)
+    cnt = 0
+
+    # attention center
+    H_min = int(np.ceil(mini_patch_size / 8))
+    H_max = int(np.ceil(mini_patch_size / 8 * 7))
+    W_min = int(np.ceil(mini_patch_size / 8))
+    W_max = int(np.ceil(mini_patch_size / 8 * 7))
+    # 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):
+        if cnt >= num_sample:
+            break
+        mini_patch = bg_mask[row:row + mini_patch_size, col: col + mini_patch_size]
+        origin = (int(col * mini_frac), int(row * mini_frac), patch_size, patch_size)
+        if np.count_nonzero(mini_patch[H_min:H_max, W_min:W_max]) >= th_num:
+            # # filter those white background
+            # if is_bg(slide, origin, patch_size):
+            #     continue
+            patch_coors.append(origin)
+            cnt += 1
+
+    return patch_coors
+
+
+# get the just size that equal to mask_size
+def get_just_gt_level(slide: openslide, size):
+    level = slide.level_count - 1
+    while level >= 0 and slide.level_dimensions[level][0] < size[0] and \
+            slide.level_dimensions[level][1] < size[1]:
+        level -= 1
+    return level
+
+
+def generate_background_mask(slide: openslide, mini_level, mini_size):
+    img = slide.read_region((0, 0), mini_level, slide.level_dimensions[mini_level])
+    img = img.resize(mini_size)
+    bg_mask = threshold_segmentation(img)
+    img.close()
+    return bg_mask
+
+
+# background segmentation algorithm
+def threshold_segmentation(img):
+    # calculate the overview level size and retrieve the image
+    img_hsv = img.convert('HSV')
+    img_hsv_np = np.array(img_hsv)
+
+    # dilate image and then threshold the image
+    schannel = img_hsv_np[:, :, 1]
+    mask = np.zeros(schannel.shape)
+
+    schannel = dilation(schannel, star(3))
+    schannel = ndimage.gaussian_filter(schannel, sigma=(5, 5), order=0)
+    threshold_global = threshold_otsu(schannel)
+
+    mask[schannel > threshold_global] = FOREGROUND
+    mask[schannel <= threshold_global] = BACKGROUND
+
+    return mask
+
+
+def is_bg(slide, origin, patch_size):
+    img = slide.read_region(origin, 0, (patch_size, patch_size))
+    # bad case is background
+    if np.array(img)[:, :, 1].mean() > 200:  # is bg
+        img.close()
+        return True
+    else:
+        img.close()
+        return False

HyperG/utils/data/pathology/patch_funs.py

@@ -22,7 +22,7 @@
 batch_size = 1
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-save_dir = osp.join(result_root, 'model_save.pth')
+model_save_dir = osp.join(result_root, 'model_best.pth')
 
 # check directions
 assert check_dir(data_root, False)
@@ -115,4 +115,4 @@ def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
     exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)
 
     model_ft = train_model(model_ft, criterion, optimizer_ft, exp_lr_scheduler, num_epochs=20)
-    torch.save(model_ft.cpu().state_dict(), save_dir)
+    torch.save(model_ft.cpu().state_dict(), model_save_dir)

HyperG/utils/data/pathology/sample_patches.py

@@ -0,0 +1,107 @@
+import glob
+import json
+import os
+import os.path as osp
+import pickle
+
+import numpy as np
+from extract_patch_feature import extract_ft
+from torch.utils.data import Dataset
+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
+
+
+def split_train_val(data_root, ratio=0.8, save_split_dir=None, resplit=False):
+    if not resplit and save_split_dir is not None and osp.exists(save_split_dir):
+        with open(save_split_dir, 'rb') as f:
+            result = pickle.load(f)
+        return result
+
+    all_list = glob.glob(osp.join(data_root, '*.svs'))
+    with open(osp.join(data_root, 'opti_survival.json'), 'r') as fp:
+        lbls = json.load(fp)
+
+    all_dict = {}
+    for full_dir in all_list:
+        id = get_id(full_dir)
+        all_dict[id]['img_dir'] = full_dir
+        all_dict[id]['survival_time'] = lbls[id]
+
+    id_list = list(all_dict.keys())
+    train_list, val_list = split_id(id_list, ratio)
+
+    train_list = [all_dict[_id] for _id in train_list]
+    val_list = [all_dict[_id] for _id in val_list]
+
+    result = {'train': train_list, 'val': val_list}
+    if save_split_dir is not None:
+        save_folder = osp.split(save_split_dir)[0]
+        if not osp.exists(save_folder):
+            os.makedirs(save_folder)
+        with open(save_split_dir, 'wb') as f:
+            pickle.dump(result, f)
+
+    return result
+
+
+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'])
+    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)}...')
+            _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:
+                pickle.dump(_patch_coors, fp)
+
+            # visualize sampled patches on slide
+            if sampled_vis is not None:
+                _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:
+                    _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())
+
+
+def get_dataloader(data_dict, patch_ft_dir):
+    pass
+
+
+class slide_patch(Dataset):
+
+    def __getitem__(self, index: int) -> T_co:
+        return super().__getitem__(index)
+
+    def __len__(self) -> int:
+        return super().__len__()
+
+
+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 = [get_id(_dir).split('_ft.')[0] for _dir in done_list]
+    for _dir in dir_list:
+        id = get_id(_dir)
+        if id not in done_list:
+            to_do_list.append(_dir)
+    return to_do_list
+
+
+def get_id(_dir):
+    return osp.splitext(osp.split(_dir)[1])[0]