diff --git a/.gitignore b/.gitignore old mode 100644 new mode 100755 diff --git a/.vscode/launch.json b/.vscode/launch.json old mode 100644 new mode 100755 index 0ee4859d..2b9f3e87 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -12,6 +12,21 @@ "program": "${file}", "console": "integratedTerminal", "cwd" : "${fileDirname}", // run at the location of the runfile + // hook terminal cmd to debug file + "args": [ + // "--gpu=0", + "--nr_types=6", + "--type_info_path=type_info.json", + "--batch_size=64", + "--model_mode=fast", + "--model_path=../../pretrained/hovernet_fast_pannuke_pytorch.tar", + "--nr_inference_workers=4", + "--nr_post_proc_workers=8", + "wsi", + "--ambiguous_size=328", + "--input_dir=exp_output/wsi/samples/", + "--output_dir=exp_output/wsi/pred/", + ] }, { "name": "Python: Remote Attach", diff --git a/CHANGELOG.md b/CHANGELOG.md old mode 100644 new mode 100755 diff --git a/LICENSE b/LICENSE old mode 100644 new mode 100755 diff --git a/README.md b/README.md old mode 100644 new mode 100755 diff --git a/__init__.py b/__init__.py old mode 100644 new mode 100755 diff --git a/compute_stats.py b/compute_stats.py old mode 100644 new mode 100755 diff --git a/config.py b/config.py old mode 100644 new mode 100755 diff --git a/convert_chkpt_tf2pytorch.py b/convert_chkpt_tf2pytorch.py old mode 100644 new mode 100755 diff --git a/convert_format.py b/convert_format.py old mode 100644 new mode 100755 index 88839608..7e32676c --- a/convert_format.py +++ b/convert_format.py @@ -55,7 +55,7 @@ def rgb2int(rgb): target_format = "qupath" # to rescale the coordinate set to match with lv0 mag of the wsi scale_factor = 1.0 - root_dir = "dataset/dummy/out/" + root_dir = "exp_output/wsi/pred/" # to define the name, and color conversion code for each target format type_info_dict = { diff --git a/dataloader/__init__.py b/dataloader/__init__.py old mode 100644 new mode 100755 diff --git a/dataloader/augs.py b/dataloader/augs.py old mode 100644 new mode 100755 diff --git a/dataloader/infer_loader.py b/dataloader/infer_loader.py old mode 100644 new mode 100755 diff --git a/dataloader/train_loader.py b/dataloader/train_loader.py old mode 100644 new mode 100755 diff --git a/dataset.py b/dataset.py old mode 100644 new mode 100755 diff --git a/diagram.png b/diagram.png old mode 100644 new mode 100755 diff --git a/docs/diagram.png b/docs/diagram.png old mode 100644 new mode 100755 diff --git a/docs/seg.gif b/docs/seg.gif old mode 100644 new mode 100755 diff --git a/environment.yml b/environment.yml old mode 100644 new mode 100755 diff --git a/examples/.ipynb_checkpoints/usage-checkpoint.ipynb b/examples/.ipynb_checkpoints/usage-checkpoint.ipynb old mode 100644 new mode 100755 diff --git a/examples/usage.ipynb b/examples/usage.ipynb old mode 100644 new mode 100755 diff --git a/extract_patches.py b/extract_patches.py old mode 100644 new mode 100755 diff --git a/infer/__init__.py b/infer/__init__.py old mode 100644 new mode 100755 diff --git a/infer/base.py b/infer/base.py old mode 100644 new mode 100755 index 81165ec2..64e9993b --- a/infer/base.py +++ b/infer/base.py @@ -67,11 +67,11 @@ def __load_model(self): net.load_state_dict(saved_state_dict, strict=True) net = torch.nn.DataParallel(net) - net = net.to("cuda") + self.model = net.to("cuda") module_lib = import_module("models.hovernet.run_desc") - run_step = getattr(module_lib, "infer_step") - self.run_step = lambda input_batch: run_step(input_batch, net) + self.run_step = getattr(module_lib, "infer_step") + # self.run_step = lambda input_batch: run_step(input_batch, net) module_lib = import_module("models.hovernet.post_proc") self.post_proc_func = getattr(module_lib, "process") diff --git a/infer/super_wsi.py b/infer/super_wsi.py new file mode 100755 index 00000000..913abf0f --- /dev/null +++ b/infer/super_wsi.py @@ -0,0 +1,955 @@ +import multiprocessing as mp +from concurrent.futures import (FIRST_EXCEPTION, ProcessPoolExecutor, + as_completed, wait) + +mp.set_start_method("spawn", True) # ! must be at top for VScode debugging + +import collections +import glob +import json +import logging +import math +import os +import pathlib +import copy +import re +import shutil +import sys +import time +from functools import reduce +from importlib import import_module + +import cv2 +import numpy as np +import psutil +import scipy.io as sio +import torch +import torch.multiprocessing as torch_mp +import torch.utils.data as data +import tqdm + +from misc.utils import (cropping_center, get_bounding_box, log_debug, log_info, + rm_n_mkdir) +from misc.wsi_handler import get_file_handler + +from . import base + + +#### +class SerializeArray(data.Dataset): + """ + `mp_shared_space` must be from torch.multiprocessing, for example + + mp_manager = torch_mp.Manager() + mp_shared_space = mp_manager.Namespace() + mp_shared_space.image = torch.from_numpy(image) + """ + def __init__(self, mp_shared_space, preproc=None): + super().__init__() + self.mp_shared_space = mp_shared_space + self.preproc = preproc + return + + def __len__(self): + return len(self.mp_shared_space.patch_info_list) + + def __getitem__(self, idx): + patch_info = self.mp_shared_space.patch_info_list[idx] + tl, br = patch_info[0] # retrieve input placement, [1] is output + patch_data = self.mp_shared_space.tile_img[tl[0] : br[0], tl[1] : br[1]] + if self.preproc is not None: + patch_dat = patch_data.copy() + patch_data = self.preproc(patch_data) + return patch_data, patch_info + +#### +def _remove_inst(inst_map, remove_id_list): + """Remove instances with id in remove_id_list. + + Args: + inst_map: map of instances + remove_id_list: list of ids to remove from inst_map + """ + for inst_id in remove_id_list: + inst_map[inst_map == inst_id] = 0 + return inst_map + +#### +def _get_patch_info(img_shape, input_size, output_size): + """Get top left coordinate information of patches from original image. + + Args: + img_shape: input image shape + input_size: patch input shape + output_size: patch output shape + + """ + def flat_mesh_grid_coord(y, x): + y, x = np.meshgrid(y, x) + return np.stack([y.flatten(), x.flatten()], axis=-1) + + in_out_diff = input_size - output_size + nr_step = np.floor((img_shape - in_out_diff) / output_size) + 1 + last_output_coord = (in_out_diff // 2) + (nr_step) * output_size + # generating subpatches index from orginal + output_tl_y_list = np.arange( + in_out_diff[0] // 2, last_output_coord[0], output_size[0], dtype=np.int32 + ) + output_tl_x_list = np.arange( + in_out_diff[1] // 2, last_output_coord[1], output_size[1], dtype=np.int32 + ) + output_tl = flat_mesh_grid_coord(output_tl_y_list, output_tl_x_list) + output_br = output_tl + output_size + + input_tl = output_tl - in_out_diff // 2 + input_br = input_tl + input_size + # exclude any patch where the input exceed the range of image, + # can comment this out if do padding in reading + sel = np.any(input_br > img_shape, axis=-1) + + info_list = np.stack( + [ + np.stack([ input_tl[~sel], input_br[~sel]], axis=1), + np.stack([output_tl[~sel], output_br[~sel]], axis=1), + ], axis=1) + # print(info_list.shape) + return info_list + +#### +def _get_tile_info(img_shape, input_size, output_size, margin_size, unit_size): + """Get top left coordinate information of patches from original image. + + Args: + img_shape: input image shape + input_size: patch input shape + output_size: patch output shape + + """ + # ! ouput tile size must be multiple of unit + assert np.sum(output_size % unit_size) == 0 + assert np.sum((margin_size*2) % unit_size) == 0 + + def flat_mesh_grid_coord(y, x): + y, x = np.meshgrid(y, x) + return np.stack([y.flatten(), x.flatten()], axis=-1) + + in_out_diff = input_size - output_size + nr_step = np.ceil((img_shape - in_out_diff) / output_size) + last_output_coord = (in_out_diff // 2) + (nr_step) * output_size + + assert np.sum(output_size % unit_size) == 0 + nr_unit_step = np.floor((img_shape - in_out_diff) / unit_size) + last_unit_output_coord = (in_out_diff // 2) + (nr_unit_step) * unit_size + + # generating subpatches index from orginal + def get_top_left_1d(axis): + o_tl_list = np.arange( + in_out_diff[axis] // 2, + last_output_coord[axis], + output_size[axis], dtype=np.int32 + ) + o_br_list = o_tl_list + output_size[axis] + o_br_list[-1] = last_unit_output_coord[axis] + # in default behavior, last pos >= last multiple of unit + # hence may cause duplication, do a check and remove if necessary + if o_br_list[-1] == o_br_list[-2]: + o_br_list = o_br_list[:-1] + o_tl_list = o_tl_list[:-1] + return o_tl_list, o_br_list + + output_tl_y_list, output_br_y_list = get_top_left_1d(axis=0) + output_tl_x_list, output_br_x_list = get_top_left_1d(axis=1) + + output_tl = flat_mesh_grid_coord(output_tl_y_list, output_tl_x_list) + output_br = flat_mesh_grid_coord(output_br_y_list, output_br_x_list) + + def get_info_stack(output_tl, output_br): + input_tl = output_tl - (in_out_diff // 2) + input_br = output_br + (in_out_diff // 2) + + info_list = np.stack( + [ + np.stack([ input_tl, input_br], axis=1), + np.stack([output_tl, output_br], axis=1), + ], axis=1) + return info_list + + # * Full Tile Grid + info_list = get_info_stack(output_tl, output_br).astype(np.int64) + + # flag surrounding ambiguous (left margin, right margin) + # |----|------------|----| + # |\\\\\\\\\\\\\\\\\\\\\\| + # |\\\\ \\\\| + # |\\\\ \\\\| + # |\\\\ \\\\| + # |\\\\\\\\\\\\\\\\\\\\\\| + # |----|------------|----| + removal_flag = np.full((info_list.shape[0], 4,), 1) # left, right, top, bot + # exclude those contain left most boundary + removal_flag[(info_list[:,1,0,1] == np.min(output_tl[:,1])),0] = 0 + # exclude those contain right most boundary + removal_flag[(info_list[:,1,1,1] == np.max(output_br[:,1])),1] = 0 + # exclude those contain top most boundary + removal_flag[(info_list[:,1,0,0] == np.min(output_tl[:,0])),2] = 0 + # exclude those contain bot most boundary + removal_flag[(info_list[:,1,1,0] == np.max(output_br[:,0])),3] = 0 + mode_list = np.full(info_list.shape[0], 0) + all_info = [[info_list, removal_flag, mode_list]] + + br_most = np.max(output_br, axis=0) + tl_most = np.min(output_tl, axis=0) + # * Tile Boundary Redo with Margin + # get the fix grid tile info + y_fix_output_tl = output_tl - np.array([margin_size[0], 0])[None,:] + y_fix_output_br = np.stack([output_tl[:,0], output_br[:,1]], axis=-1) + y_fix_output_br = y_fix_output_br + np.array([margin_size[0], 0])[None,:] + # bound reassignment + # ? do we need to do bound check for tl ? (extreme case of 1 tile of size < margin size ?) + y_fix_output_br[y_fix_output_br[:,0] > br_most[0], 0] = br_most[0] + y_fix_output_br[y_fix_output_br[:,1] > br_most[1], 1] = br_most[1] + # sel position not on the image boundary + sel = (output_tl[:,0] == np.min(output_tl[:,0])) + y_info_list = get_info_stack(y_fix_output_tl[~sel], y_fix_output_br[~sel]).astype(np.int64) + + # flag horizontal ambiguous region for y (left margin, right margin) + # |----|------------|----| + # |\\\\| |\\\\| + # |----|------------|----| + # <----> ambiguous (margin size) + removal_flag = np.zeros((y_info_list.shape[0], 4,)) # left, right, top, bot + removal_flag[:,[0,1]] = 1 + # exclude the left most boundary + removal_flag[(y_info_list[:,1,0,1] == np.min(output_tl[:,1])),0] = 0 + # exclude the right most boundary + removal_flag[(y_info_list[:,1,1,1] == np.max(output_br[:,1])),1] = 0 + mode_list = np.full(y_info_list.shape[0], 1) + all_info.append([y_info_list, removal_flag, mode_list]) + + x_fix_output_br = output_br + np.array([0, margin_size[1]])[None,:] + x_fix_output_tl = np.stack([output_tl[:,0], output_br[:,1]], axis=-1) + x_fix_output_tl = x_fix_output_tl - np.array([0, margin_size[1]])[None,:] + # bound reassignment + x_fix_output_br[x_fix_output_br[:,0] > br_most[0], 0] = br_most[0] + x_fix_output_br[x_fix_output_br[:,1] > br_most[1], 1] = br_most[1] + # sel position not on the image boundary + sel = (output_br[:,1] == np.max(output_br[:,1])) + x_info_list = get_info_stack(x_fix_output_tl[~sel], x_fix_output_br[~sel]).astype(np.int64) + # flag vertical ambiguous region for x (top margin, bottom margin) + # |----| ^ + # |\\\\| | ambiguous + # |----| V + # | | + # | | + # |----| + # |\\\\| + # |----| + removal_flag = np.zeros((x_info_list.shape[0], 4,)) # left, right, top, bot + removal_flag[:,[2,3]] = 1 + # exclude the left most boundary + removal_flag[(x_info_list[:,1,0,0] == np.min(output_tl[:,0])),2] = 0 + # exclude the right most boundary + removal_flag[(x_info_list[:,1,1,0] == np.max(output_br[:,0])),3] = 0 + mode_list = np.full(x_info_list.shape[0], 2) + all_info.append([x_info_list, removal_flag, mode_list]) + + sel = np.any(output_br == br_most, axis=-1) + xsect = output_br[~sel] + xsect_tl = xsect - margin_size * 2 + xsect_br = xsect + margin_size * 2 + # do the bound check to ensure range stay within + xsect_br[xsect_br[:,0] > br_most[0], 0] = br_most[0] + xsect_br[xsect_br[:,1] > br_most[1], 1] = br_most[1] + xsect_tl[xsect_tl[:,0] < tl_most[0], 0] = tl_most[0] + xsect_tl[xsect_tl[:,1] < tl_most[1], 1] = tl_most[1] + xsect_info_list = get_info_stack(xsect_tl, xsect_br).astype(np.int64) + mode_list = np.full(xsect_info_list.shape[0], 3) + removal_flag = np.full((xsect_info_list.shape[0], 4,), 0) # left, right, top, bot + all_info.append([xsect_info_list, removal_flag, mode_list]) + + return all_info +#### +# ! seem to be 1 pix off at cross section or sthg +def get_inst_in_margin(arr, margin_size, tile_pp_info): + """ + include the margin line itself + """ + assert margin_size > 0 + tile_pp_info = np.array(tile_pp_info) + + inst_in_margin = [] + # extract those lie within margin region + if tile_pp_info[0] == 1: # left edge + inst_in_margin.append(arr[:,:(margin_size+1)]) + if tile_pp_info[1] == 1: # right edge + inst_in_margin.append(arr[:,-(margin_size+1):]) + if tile_pp_info[2] == 1: # top edge + inst_in_margin.append(arr[:(margin_size+1),:]) + if tile_pp_info[3] == 1: # bottom edge + inst_in_margin.append(arr[-(margin_size+1):,:]) + inst_in_margin = [v.flatten() for v in inst_in_margin] + if len(inst_in_margin) > 0: + inst_in_margin = np.concatenate(inst_in_margin, axis=0) + inst_in_margin = np.unique(inst_in_margin) + else: + inst_in_margin = np.array([]) # empty array + return inst_in_margin + +#### +def get_inst_on_margin(arr, margin_size, tile_pp_info): + """ + """ + assert margin_size > 0 + # extract those lie on the margin line + tile_pp_info = np.array(tile_pp_info) + + def line_intersection(line1, line2): + ydiff = (line1[0][0] - line1[1][0], line2[0][0] - line2[1][0]) + xdiff = (line1[0][1] - line1[1][1], line2[0][1] - line2[1][1]) + + def det(a, b): + return a[0] * b[1] - a[1] * b[0] + + div = det(xdiff, ydiff) + if div == 0: + return False # not intersect + + d = (det(*line1), det(*line2)) + x = det(d, xdiff) / div + y = det(d, ydiff) / div + # ! positive region only (due to indexing line) + return int(abs(y)), int(abs(x)) + + last_h, last_w = arr.shape[0]-1, arr.shape[1]-1 + line_list = [ + [[0, 0] , [last_h, 0] ], # left line + [[0, last_w], [last_h, last_w]], # right line + [[0, 0] , [0, last_w] ], # top line + [[last_h, 0], [last_h, last_w]], # bottom line + ] + + if tile_pp_info[0] == 1: + line_list[0] = [[0 , margin_size], + [last_h, margin_size]] + if tile_pp_info[1] == 1: + line_list[1] = [[0 , last_w-margin_size], + [last_h, last_w-margin_size]] + if tile_pp_info[2] == 1: + line_list[2] = [[margin_size, 0], + [margin_size, last_w]] + if tile_pp_info[3] == 1: + line_list[3] = [[last_h-margin_size, 0], + [last_h-margin_size, last_w]] + # x1 x2 + # x3 x4 + # all pts need to be valid idx ! + pts_list = [ + line_intersection(line_list[2], line_list[0]), # x1 + line_intersection(line_list[2], line_list[1]), # x2 + line_intersection(line_list[3], line_list[0]), # x3 + line_intersection(line_list[3], line_list[1]), # x4 + ] + + pt_index = lambda p1, p2: arr[p1[0]:p2[0]+1, + p1[1]:p2[1]+1] + line_pix_list = [] + if tile_pp_info[0] == 1: + line_pix_list.append(pt_index(pts_list[0], pts_list[2])), + if tile_pp_info[1] == 1: + line_pix_list.append(pt_index(pts_list[1], pts_list[3])), + if tile_pp_info[2] == 1: + line_pix_list.append(pt_index(pts_list[0], pts_list[1])), + if tile_pp_info[3] == 1: + line_pix_list.append(pt_index(pts_list[2], pts_list[3])), + + inst_on_margin = [v.flatten() for v in line_pix_list] + + if len(inst_on_margin) > 0: + inst_on_margin = np.concatenate(inst_on_margin, axis=0) + inst_on_margin = np.unique(inst_on_margin) + else: + inst_on_margin = np.array([]) # empty array + + return inst_on_margin +#### +def get_inst_on_edge(arr, tile_pp_info): + inst_on_edge = [] + if tile_pp_info[0] == 1: + inst_on_edge.append(arr[:,0]) + if tile_pp_info[1] == 1: + inst_on_edge.append(arr[:,-1]) + if tile_pp_info[2] == 1: + inst_on_edge.append(arr[0,:]) + if tile_pp_info[3] == 1: + inst_on_edge.append(arr[-1,:]) + + inst_on_edge = [v.flatten() for v in inst_on_edge] + + if len(inst_on_edge) > 0: + inst_on_edge = np.concatenate(inst_on_edge, axis=0) + inst_on_edge = np.unique(inst_on_edge) + else: + inst_on_edge = np.array([]) # empty array + return inst_on_edge +#### + +#### +def run_model_forward( + forward_output_queue, + tile_info_list, + wsi_path, wsi_ext, wsi_proc_mag, wsi_cache_path, + run_step, model, loader_kwargs): + + wsi_handler = get_file_handler(wsi_path, backend=wsi_ext) + wsi_handler.prepare_reading(read_mag=wsi_proc_mag) + + # using shared memory namespace so all the loader workers use same + # underlying image data, also allow persistent worker and fast data switching + mp_manager = torch_mp.Manager() + mp_shared_space = mp_manager.Namespace() + + ds = SerializeArray(mp_shared_space) + loader = data.DataLoader(ds, **loader_kwargs, + drop_last=False, + persistent_workers=True, + ) + + all_time = 0 + for tile_idx, tile_info in enumerate(tile_info_list): + + tile_info, patch_info_list = tile_info + + tile_input_info = tile_info[0] + tile_input_tl, tile_input_br = tile_input_info + # shift from wsi system to tile input system + # ! (tile output is within tile input system) + # ! this will shift both patch input and output placement to tile input system + # ! hence, output placement need to be shifted (corrected) later for post proc + patch_info_list -= np.reshape(tile_input_tl, [1, 1, 1, 2]) + + tile_img = wsi_handler.read_region(tile_input_tl[::-1], + (tile_input_br - tile_input_tl)[::-1]) + + # change the data in namespace to sync across persistent loader worker + # also no need to do locking as these are assumed to be read only from worker + start = time.perf_counter() + mp_shared_space.tile_img = torch.from_numpy(tile_img).share_memory_() + mp_shared_space.patch_info_list = torch.from_numpy(patch_info_list).share_memory_() + end = time.perf_counter() + all_time += (end - start) + + accumulated_patch_output = [] + for batch_idx, batch_data in enumerate(loader): + sample_data_list, sample_info_list = batch_data + sample_output_list = run_step(sample_data_list, model) + sample_info_list = sample_info_list.numpy() + accumulated_patch_output.append([sample_info_list, sample_output_list]) + forward_output_queue.put([tile_idx, accumulated_patch_output]) + log_info('Load Time: {0}'.format(all_time)) + return + +#### +def postproc_tile(tile_io_info, tile_pp_info, tile_mode, + margin_size, patch_info_list, func_opt, + prev_wsi_inst_dict=None): + # output pos of the tile within the source wsi + tile_input_tl, tile_output_br = tile_io_info[0] + tile_output_tl, tile_output_br = tile_io_info[1] # Y, X + offset = tile_output_tl - tile_input_tl + + # ! shape may be uneven hence just detach all into a big list + patch_pos_list = [] + patch_feat_list = [] + split_inst = lambda x : np.split(x, x.shape[0], axis=0) + for batch_pos, batch_feat in patch_info_list: + patch_pos_list.extend(split_inst(batch_pos)) + patch_feat_list.extend(split_inst(batch_feat)) + + # * assemble patch to tile + nr_ch = patch_feat_list[-1].shape[-1] + tile_shape = (tile_output_br - tile_output_tl).tolist() + pred_map = np.zeros(tile_shape + [nr_ch], dtype=np.float32) + for idx in range(len(patch_pos_list)): + # zero idx to remove singleton, squeeze may kill h/w/c + patch_pos = patch_pos_list[idx][0].copy() + + # # ! assume patch pos alrd aligned to be within tile input system + patch_pos = patch_pos - offset # shift from wsi to tile output system + + pos_tl, pos_br = patch_pos[1] # retrieve ouput placement + pred_map[ + pos_tl[0] : pos_br[0], + pos_tl[1] : pos_br[1] + ] = patch_feat_list[idx][0] + del patch_pos_list, patch_feat_list + + # * retrieve actual output + postproc_func, postproc_kwargs = func_opt + pred_inst, inst_info_dict = postproc_func(pred_map, **postproc_kwargs) + del pred_map + + # * perform removal for ambiguous region + + # Consider each symbol as 1 pixel + # This is margin inner area //// + # ----------------------------- ^ ^ + # |///////////////////////////| | | margin area + # |///////////////////////////| | | (yes including the inner and outer edge) + # |///|-------------------|///| | V + # |///| ^margin_size |///| | + # |///| |///| | + # |///| <-- margin line |///| | Image area + # |///| | |///| | + # |///| v |///| | + # |///|-------------------|///| | + # |///////////////////////////| | + # |///////////////////////////| | + # ----------------------------| V + + # ! extreme slow down may happen because the aggregated results are huge + def draw_prev_pred_inst(): + tile_output = tile_io_info[1] + tile_canvas = np.zeros(tile_output[1] - tile_output[0], dtype=np.int32) + if len(prev_wsi_inst_dict) == 0: return tile_canvas + + wsi_inst_uid_list = np.array(list(prev_wsi_inst_dict.keys())) + wsi_inst_com_list = np.array([v['centroid'] for v in prev_wsi_inst_dict.values()]) + wsi_inst_com_list = wsi_inst_com_list[:,::-1] # XY to YX + sel = (wsi_inst_com_list[:,0] > tile_output[0,0]) + sel &= (wsi_inst_com_list[:,0] < tile_output[1,0]) + sel &= (wsi_inst_com_list[:,1] > tile_output[0,1]) + sel &= (wsi_inst_com_list[:,1] < tile_output[1,1]) + sel_idx = np.nonzero(sel.flatten())[0] + + for inst_idx in sel_idx: + inst_uid = wsi_inst_uid_list[inst_idx] + # shift from wsi system to tile output system + inst_info = prev_wsi_inst_dict[inst_uid] + inst_cnt = np.array(inst_info['contour']) + inst_cnt = inst_cnt - tile_output[0,[1,0]][None] + tile_canvas = cv2.drawContours(tile_canvas, [inst_cnt], -1, int(inst_uid), -1) + return tile_canvas + + output_remove_inst_set = None + # tile_mode = -1 # ! no fix, debud mode + if tile_mode == 0: + # for `full grid tile` + # -- extend from the boundary by the margin size, remove + # nuclei lie within the margin area but exclude those + # lie on the margin line + # also contain those lying on the edges + inst_in_margin = get_inst_in_margin(pred_inst, margin_size, tile_pp_info) + # those lying on the margin line, check 2pix toward margin area for sanity + inst_on_margin1 = get_inst_on_margin(pred_inst, margin_size-1, tile_pp_info) + inst_on_margin2 = get_inst_on_margin(pred_inst, margin_size , tile_pp_info) + inst_on_margin = np.union1d(inst_on_margin1, inst_on_margin2) + inst_within_margin = np.setdiff1d(inst_in_margin, inst_on_margin, assume_unique=True) + remove_inst_set = inst_within_margin.tolist() + elif tile_mode == 1 or tile_mode == 2: + # for `horizontal/vertical strip tiles` for fixing artifacts + # -- extend from the marked edges (top/bot or left/right) by the margin size, + # remove all nuclei lie within the margin area (including on the margin line) + # -- nuclei on all edges are removed (as these are alrd within `full grid tile`) + inst_in_margin = get_inst_in_margin(pred_inst, margin_size, tile_pp_info) # also contain those lying on the edges + if np.sum(tile_pp_info) == 1: + holder_flag = tile_pp_info.copy() + if tile_mode == 1: + holder_flag[[2, 3]] = 1 + else: + holder_flag[[0, 1]] = 1 + else: + holder_flag = [1, 1, 1, 1] + inst_on_edge = get_inst_on_edge(pred_inst, holder_flag) + remove_inst_set = np.union1d(inst_in_margin, inst_on_edge) + remove_inst_set = remove_inst_set.tolist() + elif tile_mode == 3: + # inst within the tile after excluding margin area out + # only for a tile at cross-section, which is designed such that + # their shape >= 3* margin size + + # ! for removal of current pred, just like tile_mode = 0 but with all side + inst_in_margin = get_inst_in_margin(pred_inst, margin_size, [1, 1, 1, 1]) + # those lying on the margin line, check 2pix toward margin area for sanity + inst_on_margin1 = get_inst_on_margin(pred_inst, margin_size-1, [1, 1, 1, 1]) + inst_on_margin2 = get_inst_on_margin(pred_inst, margin_size , [1, 1, 1, 1]) + inst_on_margin = np.union1d(inst_on_margin1, inst_on_margin2) + inst_within_margin = np.setdiff1d(inst_in_margin, inst_on_margin, assume_unique=True) + remove_inst_set = inst_within_margin.tolist() + # ! but we also need to remove prev inst exist in the global space of entire wsi + # ! on the margin + prev_pred_inst = draw_prev_pred_inst() + inst_on_margin1 = get_inst_on_margin(prev_pred_inst, margin_size-1, [1, 1, 1, 1]) + inst_on_margin2 = get_inst_on_margin(prev_pred_inst, margin_size , [1, 1, 1, 1]) + inst_on_margin = np.union1d(inst_on_margin1, inst_on_margin2) + output_remove_inst_set = inst_on_margin.tolist() + else: + remove_inst_set = [] + + remove_inst_set = set(remove_inst_set) + # * move pos back to wsi position + renew_id = 0 + new_inst_info_dict = {} + for k, v in inst_info_dict.items(): + if k not in remove_inst_set: + v['bbox'] += tile_output_tl[::-1] + v['centroid'] += tile_output_tl[::-1] + v['contour'] += tile_output_tl[::-1] + new_inst_info_dict[renew_id] = v + renew_id += 1 + + return new_inst_info_dict, output_remove_inst_set +#### +class InferManager(base.InferManager): + + def __get_valid_patch_idx(self, patch_info_list): + """Select valid patches from the list of input patch information. + + Args: + patch_info_list: patch input coordinate information + has_output_info: whether output information is given + + """ + def check_valid(info, wsi_mask): + output_bbox = np.rint(info[1]).astype(np.int64) + output_roi = wsi_mask[ + output_bbox[0][0] : output_bbox[1][0], + output_bbox[0][1] : output_bbox[1][1], + ] + return (torch.sum(output_roi) > 0).item() + + down_sample_ratio = self.wsi_mask.shape[0] / self.wsi_proc_shape[0] + torch_mask = torch.from_numpy(self.wsi_mask).share_memory_() + valid_indices = [check_valid(info * down_sample_ratio, torch_mask) + for info in patch_info_list] + # somehow multiproc is slower than single thread + valid_indices = np.array(valid_indices) + return valid_indices + + def __get_valid_patch_idx_in_tile(self, tile_info, patch_info_list): + # checking basing on the output alignment + tile_tl, tile_br = tile_info[1] + patch_tl_list = patch_info_list[:,1,0] + patch_br_list = patch_info_list[:,1,1] + sel = (patch_tl_list[:,0] >= tile_tl[0]) & (patch_tl_list[:,1] >= tile_tl[1]) + sel &= (patch_br_list[:,0] <= tile_br[0]) & (patch_br_list[:,1] <= tile_br[1]) + return sel + + def _parse_args(self, run_args): + """Parse command line arguments and set as instance variables.""" + for variable, value in run_args.items(): + self.__setattr__(variable, value) + # to tuple + make_shape_array = lambda x : np.array([x, x]).astype(np.int64) + self.tile_shape = make_shape_array(self.tile_shape) + self.ambiguous_size = make_shape_array(self.ambiguous_size) + self.patch_input_shape = make_shape_array(self.patch_input_shape) + self.patch_output_shape = make_shape_array(self.patch_output_shape) + return + + def _get_wsi_mask(self, wsi_handler, mask_path): + if mask_path is not None and os.path.isfile(mask_path): + wsi_mask = cv2.imread(mask_path) + wsi_mask = cv2.cvtColor(self.wsi_mask, cv2.COLOR_BGR2GRAY) + wsi_mask[wsi_mask > 0] = 1 + else: + log_info( + "WARNING: No mask found, generating mask via thresholding at 1.25x!" + ) + from skimage import morphology + + # simple method to extract tissue regions using intensity thresholding and morphological operations + def simple_get_mask(): + scaled_wsi_mag = 1.25 # ! hard coded + wsi_thumb_rgb = wsi_handler.get_full_img(read_mag=scaled_wsi_mag) + gray = cv2.cvtColor(wsi_thumb_rgb, cv2.COLOR_RGB2GRAY) + _, mask = cv2.threshold(gray, 0, 255, cv2.THRESH_OTSU) + mask = morphology.remove_small_objects( + mask == 0, min_size=16 * 16, connectivity=2 + ) + mask = morphology.remove_small_holes(mask, area_threshold=128 * 128) + mask = morphology.binary_dilation(mask, morphology.disk(16)) + return mask + + # ! using entire wsi, debugging purpose ! + # wsi_mask = np.array(simple_get_mask() > 0, dtype=np.uint8) + wsi_mask = simple_get_mask() + return wsi_mask + + def process_single_file(self, wsi_path, mask_path, output_dir): + """Process a single whole-slide image and save the results. + + Args: + wsi_path: path to input whole-slide image + msk_path: path to input mask. If not supplied, mask will be automatically generated. + output_dir: path where output will be saved + + """ + path_obj = pathlib.Path(wsi_path) + wsi_ext = path_obj.suffix + wsi_name = path_obj.stem + + # TODO: expose read mpp mode + start = time.perf_counter() + self.wsi_handler = get_file_handler(wsi_path, backend=wsi_ext) + self.wsi_proc_shape = self.wsi_handler.get_dimensions(self.proc_mag) + # ! cache here is for legacy and to deal with esoteric internal wsi format + self.wsi_handler.prepare_reading(read_mag=self.proc_mag) + self.wsi_proc_shape = np.array(self.wsi_proc_shape[::-1]) # to Y, X + + self.wsi_mask = self._get_wsi_mask(self.wsi_handler, mask_path) + if np.sum(self.wsi_mask) == 0: + log_info("Skip due to empty mask!") + return + if self.save_mask: + cv2.imwrite("%s/mask/%s.png" % (output_dir, wsi_name), + self.wsi_mask * 255) + if self.save_thumb: + wsi_thumb_rgb = self.wsi_handler.get_full_img(read_mag=1.25) + cv2.imwrite( + "%s/thumb/%s.png" % (output_dir, wsi_name), + cv2.cvtColor(wsi_thumb_rgb, cv2.COLOR_RGB2BGR), + ) + + # * retrieve patch and tile placement + patch_info_list = _get_patch_info( + self.wsi_proc_shape, self.patch_input_shape, self.patch_output_shape, + ) + patch_diff_shape = self.patch_input_shape - self.patch_output_shape + # derive tile output placement as consecutive tiling with step size of 0 + # and tile output will have shape being of multiple of patch_output_shape (round down) + tile_output_shape = np.floor(self.tile_shape / self.patch_output_shape) * self.patch_output_shape + tile_input_shape = tile_output_shape + patch_diff_shape + + # [full_grid, vert/horiz, xsect] + all_tile_info = _get_tile_info( + self.wsi_proc_shape, tile_input_shape, tile_output_shape, + self.ambiguous_size, self.patch_output_shape + ) + end = time.perf_counter() + log_info("Preparing Input Output Placement: {0}".format(end - start)) + + # * Async Inference + # * launch a seperate process to do forward and store the result in a queue. + # * main thread will poll for forward result and launch separate process for + # * doing the postproc for every newly predicted tiles + # + # / forward \ (loop) + # main ------------- main----------------main + # \ postproc (loop)/ + + sel_index = self.__get_valid_patch_idx(patch_info_list) + patch_info_list = patch_info_list[sel_index] + + pbar_creator = lambda x, y, pos=0, leave=False: tqdm.tqdm( + desc=x, leave=leave, total=y, ncols=80, ascii=True, position=pos + ) + + def run_once(tile_io_info_list, tile_pp_info_list, tile_mode_list, + prev_wsi_inst_dict=None): + + sel_index = self.__get_valid_patch_idx(tile_io_info_list) + tile_io_info_list = tile_io_info_list[sel_index] + tile_pp_info_list = tile_pp_info_list[sel_index] + tile_mode_list = tile_mode_list[sel_index] + + mp_manager = torch_mp.Manager() + # contain at most 5 tile ouput before polling forward func + mp_forward_output_queue = mp_manager.Queue(maxsize=16) + + forward_info_list = collections.deque() + nr_tile = tile_io_info_list.shape[0] + for tile_idx in range(nr_tile): + tile_info = tile_io_info_list[tile_idx] + sel_index = self.__get_valid_patch_idx_in_tile(tile_info, patch_info_list) + patch_in_tile_info_list = patch_info_list[sel_index] + forward_info_list.append([tile_info, patch_in_tile_info_list]) + + loader_kwargs = dict( + num_workers=self.nr_inference_workers, + batch_size=self.batch_size, + ) + + wsi_cache_path = "%s/src_wsi.npy" % self.cache_path + forward_process = mp.Process(target=run_model_forward, + args=(mp_forward_output_queue, forward_info_list, + wsi_path, wsi_ext, self.proc_mag, wsi_cache_path, + self.run_step, self.model, loader_kwargs)) + forward_process.start() + + post_proc_kwargs = { + "nr_types": self.method["model_args"]["nr_types"], + "return_centroids": True, + } + + proc_pool = None + if self.nr_post_proc_workers > 0: + proc_pool = ProcessPoolExecutor(self.nr_post_proc_workers) + + offset_id = 0 # offset id to increment overall saving dict + wsi_inst_info = {} + if prev_wsi_inst_dict is not None: + wsi_inst_info = copy.deepcopy(prev_wsi_inst_dict) + if len(wsi_inst_info) > 0: + offset_id = max(wsi_inst_info.keys()) + 1 + + foward_pbar = pbar_creator('Forward', nr_tile, pos=0) + postpr_pbar = pbar_creator('PostPro', nr_tile, pos=1) + + future_list = collections.deque() + # will this lead to infinite loop ? + while True: + if forward_process.exitcode is not None \ + and mp_forward_output_queue.empty(): + break + + if not mp_forward_output_queue.empty(): + tile_idx, forward_output = mp_forward_output_queue.get() + foward_pbar.update() + + tile_io_info = tile_io_info_list[tile_idx] + tile_pp_info = tile_pp_info_list[tile_idx] + tile_mode = tile_mode_list[tile_idx] + args = [tile_io_info, tile_pp_info, tile_mode, + self.ambiguous_size[0], forward_output, + (self.post_proc_func, post_proc_kwargs), + prev_wsi_inst_dict] + + if proc_pool is not None: + future = proc_pool.submit(postproc_tile, *args) + future_list.append(future) + else: + new_inst_dict, remove_uid_list = postproc_tile(*args) + + # * aggregate + # ! the return id should be contiguous to maximuize + # ! counting range in int32 + inst_wsi_id = offset_id # barrier in case no output in tile! + for inst_id, inst_info in new_inst_dict.items(): + inst_wsi_id = offset_id + inst_id + 1 + wsi_inst_info[inst_wsi_id] = inst_info + offset_id = inst_wsi_id + 1 + + if remove_uid_list is not None: + for inst_uid in remove_uid_list: + if inst_uid in wsi_inst_info: + # faster than pop + del wsi_inst_info[inst_uid] + postpr_pbar.update() + + if forward_process.exitcode > 0: + raise ValueError(f'Forward process exited with code {forward_process.exitcode}') + forward_process.join() + + while len(future_list) > 0: + if not future_list[0].done(): + future_list.rotate() + continue + proc_future = future_list.popleft() + if proc_future.exception() is not None: + print(proc_future.exception()) + + # * aggregate + # ! the return id should be contiguous to maximize + # ! counting range in int32 + new_inst_dict, remove_uid_list = proc_future.result() + inst_wsi_id = offset_id # barrier in case no output in tile! + for inst_id, inst_info in new_inst_dict.items(): + inst_wsi_id = offset_id + inst_id + 1 + wsi_inst_info[inst_wsi_id] = inst_info + offset_id = inst_wsi_id + 1 + + if remove_uid_list is not None: + for inst_uid in remove_uid_list: + if inst_uid in wsi_inst_info: + # faster than pop + del wsi_inst_info[inst_uid] + postpr_pbar.update() + foward_pbar.close() + postpr_pbar.close() + proc_pool.shutdown() + return wsi_inst_info + + # + ## ** process full grid and vert/horiz fixing at the same time + start = time.perf_counter() + info_list = list(zip(*all_tile_info[:3])) + tile_io_info_list = np.concatenate(info_list[0], axis=0).astype(np.int32) + tile_pp_info_list = np.concatenate(info_list[1], axis=0) + tile_mode_list = np.concatenate(info_list[2], axis=0) + wsi_inst_info = run_once(tile_io_info_list, + tile_pp_info_list, + tile_mode_list) + end = time.perf_counter() + log_info("Proc Grid: {0}".format(end - start)) + + # wsi_inst_info = {} + ## ** re-infer and redo postproc for xsect alone + start = time.perf_counter() + tile_io_info_list = all_tile_info[-1][0] + tile_pp_info_list = all_tile_info[-1][1] + tile_mode_list = all_tile_info[-1][2] + wsi_inst_info = run_once(tile_io_info_list, + tile_pp_info_list, + tile_mode_list, + wsi_inst_info) + end = time.perf_counter() + log_info("Proc XSect: {0}".format(end - start)) + + if self.save_mask or self.save_thumb: + json_path = "%s/json/%s.json" % (output_dir, wsi_name) + else: + json_path = "%s/%s.json" % (output_dir, wsi_name) + self.__save_json(json_path, wsi_inst_info) + return + + def process_wsi_list(self, run_args): + """Process a list of whole-slide images. + + Args: + run_args: arguments as defined in run_infer.py + + """ + self._parse_args(run_args) + + if not os.path.exists(self.cache_path): + rm_n_mkdir(self.cache_path) + + if not os.path.exists(self.output_dir + "/json/"): + rm_n_mkdir(self.output_dir + "/json/") + if self.save_thumb: + if not os.path.exists(self.output_dir + "/thumb/"): + rm_n_mkdir(self.output_dir + "/thumb/") + if self.save_mask: + if not os.path.exists(self.output_dir + "/mask/"): + rm_n_mkdir(self.output_dir + "/mask/") + + wsi_path_list = glob.glob(self.input_dir + "/*") + wsi_path_list.sort() # ensure ordering + for wsi_path in wsi_path_list: + wsi_base_name = pathlib.Path(wsi_path).stem + + msk_path = "%s/%s.png" % (self.input_mask_dir, wsi_base_name) + if self.save_thumb or self.save_mask: + output_file = "%s/json/%s.json" % (self.output_dir, wsi_base_name) + else: + output_file = "%s/%s.json" % (self.output_dir, wsi_base_name) + + # ! cache in case of loading across network + # start = time.perf_counter() + # wsi_holder_path = self.cache_src_wsi_path + # shutil.copyfile(wsi_path, wsi_holder_path) + # end = time.perf_counter() + # log_info('Move WSI: {0}'.format(end - start)) + + if os.path.exists(output_file): + log_info("Skip: %s" % wsi_base_name) + continue + try: + log_info("Process: %s" % wsi_base_name) + self.process_single_file(wsi_path, msk_path, self.output_dir) + log_info("Finish") + except: + logging.exception("Crash") + + rm_n_mkdir(self.cache_path) # clean up all cache + return diff --git a/infer/tile.py b/infer/tile.py old mode 100644 new mode 100755 index f27c7d3d..242e84f7 --- a/infer/tile.py +++ b/infer/tile.py @@ -299,7 +299,7 @@ def detach_items_of_uid(items_list, uid, nr_expected_items): accumulated_patch_output = [] for batch_idx, batch_data in enumerate(dataloader): sample_data_list, sample_info_list = batch_data - sample_output_list = self.run_step(sample_data_list) + sample_output_list = self.run_step(sample_data_list, self.model) sample_info_list = sample_info_list.numpy() curr_batch_size = sample_output_list.shape[0] sample_output_list = np.split( diff --git a/infer/wsi.py b/infer/wsi.py old mode 100644 new mode 100755 index a8409cdc..4e8f9db1 --- a/infer/wsi.py +++ b/infer/wsi.py @@ -2,6 +2,7 @@ from concurrent.futures import FIRST_EXCEPTION, ProcessPoolExecutor, as_completed, wait from multiprocessing import Lock, Pool +# also to ensure similar behavior between linux and windows mp.set_start_method("spawn", True) # ! must be at top for VScode debugging import argparse @@ -24,9 +25,9 @@ import scipy.io as sio import torch import torch.utils.data as data +import torch.multiprocessing as torch_mp import tqdm -from dataloader.infer_loader import SerializeArray, SerializeFileList -from docopt import docopt + from misc.utils import ( cropping_center, get_bounding_box, @@ -41,11 +42,6 @@ thread_lock = Lock() -#### -def _init_worker_child(lock_): - global lock - lock = lock_ - #### def _remove_inst(inst_map, remove_id_list): @@ -255,48 +251,38 @@ def _assemble_and_flush(wsi_pred_map_mmap_path, chunk_info, patch_output_list): # print(chunk_info.flatten(), 'pass') return - #### -class InferManager(base.InferManager): - def __run_model(self, patch_top_left_list, pbar_desc): - # TODO: the cost of creating dataloader may not be cheap ? - dataset = SerializeArray( - "%s/cache_chunk.npy" % self.cache_path, - patch_top_left_list, - self.patch_input_shape, - ) +class SerializeArray(data.Dataset): + """ + `mp_shared_space` must be from torch.multiprocessing, for example - dataloader = data.DataLoader( - dataset, - num_workers=self.nr_inference_workers, - batch_size=self.batch_size, - drop_last=False, - ) + mp_manager = torch_mp.Manager() + mp_shared_space = mp_manager.Namespace() + mp_shared_space.image = torch.from_numpy(image) + """ + def __init__(self, mp_shared_space, patch_size, preproc=None): + super().__init__() + self.patch_size = patch_size + self.mp_shared_space = mp_shared_space + self.preproc = preproc + return - pbar = tqdm.tqdm( - desc=pbar_desc, - leave=True, - total=int(len(dataloader)), - ncols=80, - ascii=True, - position=0, - ) + def __len__(self): + return len(self.mp_shared_space.patch_info_list) - # run inference on input patches - accumulated_patch_output = [] - for batch_idx, batch_data in enumerate(dataloader): - sample_data_list, sample_info_list = batch_data - sample_output_list = self.run_step(sample_data_list) - sample_info_list = sample_info_list.numpy() - curr_batch_size = sample_output_list.shape[0] - sample_output_list = np.split(sample_output_list, curr_batch_size, axis=0) - sample_info_list = np.split(sample_info_list, curr_batch_size, axis=0) - sample_output_list = list(zip(sample_info_list, sample_output_list)) - accumulated_patch_output.extend(sample_output_list) - pbar.update() - pbar.close() - return accumulated_patch_output + def __getitem__(self, idx): + patch_info = self.mp_shared_space.patch_info_list[idx] + tl, br = patch_info[0] # retrieve input placement, [1] is output + patch_data = self.mp_shared_space.tile_img[ + tl[0] : tl[0] + self.patch_size[0], + tl[1] : tl[1] + self.patch_size[1]] + if self.preproc is not None: + patch_dat = patch_data.copy() + patch_data = self.preproc(patch_data) + return patch_data, patch_info[0][0] +#### +class InferManager(base.InferManager): def __select_valid_patches(self, patch_info_list, has_output_info=True): """Select valid patches from the list of input patch information. @@ -334,12 +320,53 @@ def __get_raw_prediction(self, chunk_info_list, patch_info_list): patch_info_list: list of patch coordinate information """ + pbar_creator = lambda x, y, pos=0, leave=False: tqdm.tqdm( + desc=x, leave=leave, total=y, ncols=80, ascii=True, position=pos + ) + + mp_manager = torch_mp.Manager() + mp_shared_space = mp_manager.Namespace() + + dataset = SerializeArray(mp_shared_space, self.patch_input_shape) + loader = data.DataLoader( + dataset, + num_workers=self.nr_inference_workers, + batch_size=self.batch_size, + drop_last=False, + persistent_workers=self.nr_inference_workers > 0 + ) + def run_model_forward(mp_manager, chunk_image, patch_info_list): + + mp_shared_space.tile_img = torch.from_numpy(chunk_image).share_memory_() + mp_shared_space.patch_info_list = torch.from_numpy(patch_info_list).share_memory_() + + nr_batch = int(patch_info_list.shape[0] / self.batch_size) + batch_pbar = pbar_creator('Proc Batch', nr_batch, pos=1) + + # run inference on input patches + accumulated_patch_output = [] + for batch_idx, batch_data in enumerate(loader): + sample_data_list, sample_info_list = batch_data + sample_output_list = self.run_step(sample_data_list, self.model) + sample_info_list = sample_info_list.numpy() + curr_batch_size = sample_output_list.shape[0] + sample_output_list = np.split(sample_output_list, curr_batch_size, axis=0) + sample_info_list = np.split(sample_info_list, curr_batch_size, axis=0) + sample_output_list = list(zip(sample_info_list, sample_output_list)) + accumulated_patch_output.extend(sample_output_list) + batch_pbar.update() + batch_pbar.close() + return accumulated_patch_output + # 1 dedicated thread just to write results back to disk proc_pool = Pool(processes=1) wsi_pred_map_mmap_path = "%s/pred_map.npy" % self.cache_path - + + nr_chunk = chunk_info_list.shape[0] + chunk_pbar = pbar_creator('Proc Chunk', nr_chunk, pos=0) + masking = lambda x, a, b: (a <= x) & (x <= b) - for idx in range(0, chunk_info_list.shape[0]): + for idx in range(0, nr_chunk): chunk_info = chunk_info_list[idx] # select patch basing on top left coordinate of input start_coord = chunk_info[0, 0] @@ -369,15 +396,16 @@ def __get_raw_prediction(self, chunk_info_list, patch_info_list): chunk_data = np.array(chunk_data)[..., :3] np.save("%s/cache_chunk.npy" % self.cache_path, chunk_data) - pbar_desc = "Process Chunk %d/%d" % (idx, chunk_info_list.shape[0]) - patch_output_list = self.__run_model( - chunk_patch_info_list[:, 0, 0], pbar_desc + patch_output_list = run_model_forward( + mp_manager, chunk_data, chunk_patch_info_list ) proc_pool.apply_async( _assemble_and_flush, args=(wsi_pred_map_mmap_path, chunk_info, patch_output_list), ) + chunk_pbar.update() + chunk_pbar.close() proc_pool.close() proc_pool.join() return @@ -469,9 +497,7 @@ def process_single_file(self, wsi_path, msk_path, output_dir): start = time.perf_counter() self.wsi_handler = get_file_handler(wsi_path, backend=wsi_ext) self.wsi_proc_shape = self.wsi_handler.get_dimensions(self.proc_mag) - self.wsi_handler.prepare_reading( - read_mag=self.proc_mag, cache_path="%s/src_wsi.npy" % self.cache_path - ) + self.wsi_handler.prepare_reading(read_mag=self.proc_mag) self.wsi_proc_shape = np.array(self.wsi_proc_shape[::-1]) # to Y, X if msk_path is not None and os.path.isfile(msk_path): diff --git a/metrics/README.md b/metrics/README.md old mode 100644 new mode 100755 diff --git a/metrics/__init__.py b/metrics/__init__.py old mode 100644 new mode 100755 diff --git a/metrics/stats_utils.py b/metrics/stats_utils.py old mode 100644 new mode 100755 index bef1a827..729cd168 --- a/metrics/stats_utils.py +++ b/metrics/stats_utils.py @@ -1,9 +1,10 @@ import warnings -import numpy as np -from scipy.optimize import linear_sum_assignment import cv2 import matplotlib.pyplot as plt +import numpy as np +import scipy +from scipy.optimize import linear_sum_assignment # --------------------------Optimised for Speed @@ -407,32 +408,23 @@ def pair_coordinates(setA, setB, radius): """ # * Euclidean distance as the cost matrix - setA_tile = np.expand_dims(setA, axis=1) - setB_tile = np.expand_dims(setB, axis=0) - setA_tile = np.repeat(setA_tile, setB.shape[0], axis=1) - setB_tile = np.repeat(setB_tile, setA.shape[0], axis=0) - pair_distance = (setA_tile - setB_tile) ** 2 - # set A is row, and set B is paired against set A - pair_distance = np.sqrt(np.sum(pair_distance, axis=-1)) + pair_distance = scipy.spatial.distance.cdist(setA, setB, metric='euclidean') + print(setA.shape, setB.shape) # * Munkres pairing with scipy library # the algorithm return (row indices, matched column indices) - # if there is multiple same cost in a row, index of first occurence + # if there is multiple same cost in a row, index of first occurence # is return, thus the unique pairing is ensured indicesA, paired_indicesB = linear_sum_assignment(pair_distance) - # extract the paired cost and remove instances + # extract the paired cost and remove instances # outside of designated radius pair_cost = pair_distance[indicesA, paired_indicesB] pairedA = indicesA[pair_cost <= radius] pairedB = paired_indicesB[pair_cost <= radius] - unpairedA = [idx for idx in range(setA.shape[0]) if idx not in list(pairedA)] - unpairedB = [idx for idx in range(setB.shape[0]) if idx not in list(pairedB)] - - pairing = np.array(list(zip(pairedA, pairedB))) - unpairedA = np.array(unpairedA, dtype=np.int64) - unpairedB = np.array(unpairedB, dtype=np.int64) - + pairing = np.concatenate([pairedA[:,None], pairedB[:,None]], axis=-1) + unpairedA = np.delete(np.arange(setA.shape[0]), pairedA) + unpairedB = np.delete(np.arange(setB.shape[0]), pairedB) return pairing, unpairedA, unpairedB diff --git a/misc/__init__.py b/misc/__init__.py old mode 100644 new mode 100755 diff --git a/misc/patch_extractor.py b/misc/patch_extractor.py old mode 100644 new mode 100755 diff --git a/misc/utils.py b/misc/utils.py old mode 100644 new mode 100755 diff --git a/misc/viz_utils.py b/misc/viz_utils.py old mode 100644 new mode 100755 diff --git a/misc/wsi_handler.py b/misc/wsi_handler.py old mode 100644 new mode 100755 index 49aa4ea0..8376491d --- a/misc/wsi_handler.py +++ b/misc/wsi_handler.py @@ -1,194 +1,176 @@ -from collections import OrderedDict -import cv2 -import numpy as np -from skimage import img_as_ubyte -from skimage import color import re import subprocess +import warnings +from collections import OrderedDict +import cv2 +import glymur +import numpy as np import openslide +import tifffile +from skimage import color, img_as_ubyte class FileHandler(object): def __init__(self): - """The handler is responsible for storing the processed data, parsing + """ + The handler is responsible for storing the processed data, parsing the metadata from original file, and reading it from storage. """ self.metadata = { - ("available_mag", None), - ("base_mag", None), - ("vendor", None), - ("mpp ", None), - ("base_shape", None), + ('available_mag', None), + ('base_mag' , None), + ('vendor' , None), + ('base_mpp' , None), + ('base_shape' , None), } + pass def __load_metadata(self): - raise NotImplementedError - - def get_full_img(self, read_mag=None, read_mpp=None): - """Only use `read_mag` or `read_mpp`, not both, prioritize `read_mpp`. + pass - `read_mpp` is in X, Y format + def read_region(self): + pass + + def get_dimensions(self, read_mag=None, read_mpp=None): """ - raise NotImplementedError - - def read_region(self, coords, size): - """Must call `prepare_reading` before hand. - - Args: - coords (tuple): (dims_x, dims_y), - top left coordinates of image region at selected - `read_mag` or `read_mpp` from `prepare_reading` - size (tuple): (dims_x, dims_y) - width and height of image region at selected - `read_mag` or `read_mpp` from `prepare_reading` - + Will be in X, Y """ - raise NotImplementedError - - def get_dimensions(self, read_mag=None, read_mpp=None): - """Will be in X, Y.""" if read_mpp is not None: - read_scale = (self.metadata["base_mpp"] / read_mpp)[0] - read_mag = read_scale * self.metadata["base_mag"] - scale = read_mag / self.metadata["base_mag"] + read_scale = (self.metadata['base_mpp'] / read_mpp)[0] + read_mag = read_scale * self.metadata['base_mag'] + scale = read_mag / self.metadata['base_mag'] # may off some pixels wrt existing mag - return (self.metadata["base_shape"] * scale).astype(np.int32) + return (self.metadata['base_shape'] * scale).astype(np.int32) - def prepare_reading(self, read_mag=None, read_mpp=None, cache_path=None): - """Only use `read_mag` or `read_mpp`, not both, prioritize `read_mpp`. + def prepare_reading(self, read_mag=None, read_mpp=None): + """ + Only use either of these parameter, prioritize `read_mpp` - `read_mpp` is in X, Y format. + `read_mpp` is in X, Y format """ - read_lv, scale_factor = self._get_read_info( - read_mag=read_mag, read_mpp=read_mpp - ) - - if scale_factor is None: - self.image_ptr = None - self.read_lv = read_lv - else: - np.save(cache_path, self.get_full_img(read_mag=read_mag)) - self.image_ptr = np.load(cache_path, mmap_mode="r") + hires_lv, scale_from_hires_lv, scale_from_lv0 = self._get_read_info(read_mag=read_mag, read_mpp=read_mpp) + + self.read_lv = hires_lv + self.scale_from_hires_lv = scale_from_hires_lv + self.scale_from_lv0 = scale_from_lv0 return def _get_read_info(self, read_mag=None, read_mpp=None): if read_mpp is not None: - assert read_mpp[0] == read_mpp[1], "Not supported uneven `read_mpp`" - read_scale = (self.metadata["base_mpp"] / read_mpp)[0] - read_mag = read_scale * self.metadata["base_mag"] + assert read_mpp[0] == read_mpp[1], 'Not supported uneven `read_mpp`' + read_scale = (self.metadata['base_mpp'] / read_mpp)[0] + read_mag = read_scale * self.metadata['base_mag'] hires_mag = read_mag - scale_factor = None - if read_mag not in self.metadata["available_mag"]: - if read_mag > self.metadata["base_mag"]: - scale_factor = read_mag / self.metadata["base_mag"] - hires_mag = self.metadata["base_mag"] - else: - mag_list = np.array(self.metadata["available_mag"]) + scale_from_lv0 = read_mag / self.metadata['base_mag'] + scale_from_hires_lv = scale_from_lv0 + if read_mag not in self.metadata['available_mag']: + if read_mag > self.metadata['base_mag']: + scale_from_hires_lv = scale_from_lv0 + hires_mag = self.metadata['base_mag'] + else: + mag_list = np.array(self.metadata['available_mag']) mag_list = np.sort(mag_list)[::-1] hires_mag = mag_list - read_mag # only use higher mag as base for loading hires_mag = hires_mag[hires_mag > 0] # use the immediate higher to save compuration hires_mag = mag_list[np.argmin(hires_mag)] - scale_factor = read_mag / hires_mag - - hires_lv = self.metadata["available_mag"].index(hires_mag) - return hires_lv, scale_factor + scale_from_hires_lv = read_mag / hires_mag + hires_lv = self.metadata['available_mag'].index(hires_mag) + return hires_lv, scale_from_hires_lv, scale_from_lv0 class OpenSlideHandler(FileHandler): - """Class for handling OpenSlide supported whole-slide images.""" - + """ + Class for handling OpenSlide supported whole-slide images + """ def __init__(self, file_path): - """file_path (string): path to single whole-slide image.""" + """ + file_path (string): path to single whole-slide image + """ super().__init__() - self.file_ptr = openslide.OpenSlide(file_path) # load OpenSlide object + self.file_ptr = openslide.OpenSlide(file_path) # load OpenSlide object self.metadata = self.__load_metadata() # only used for cases where the read magnification is different from - self.image_ptr = None # the existing modes of the read file + self.image_ptr = None # the existing modes of the read file self.read_level = None def __load_metadata(self): metadata = {} wsi_properties = self.file_ptr.properties - level_0_magnification = wsi_properties[openslide.PROPERTY_NAME_OBJECTIVE_POWER] - level_0_magnification = float(level_0_magnification) + mpp = [float(wsi_properties[openslide.PROPERTY_NAME_MPP_X]), + float(wsi_properties[openslide.PROPERTY_NAME_MPP_Y])] + mpp = np.array(mpp) - downsample_level = self.file_ptr.level_downsamples + try : + level_0_magnification = wsi_properties[openslide.PROPERTY_NAME_OBJECTIVE_POWER] + level_0_magnification = float(level_0_magnification) + downsample_level = self.file_ptr.level_downsamples + magnification_level = [level_0_magnification / lv for lv in downsample_level] + except: + if mpp[0] > 0.1 and mpp[1] < 0.4: + level_0_magnification = 40.0 + if mpp[0] >= 0.4 and mpp[1] < 0.6: + level_0_magnification = 20.0 + downsample_level = self.file_ptr.level_downsamples + warnings.warn('Could not detect magnification, guess from `mpp`.') magnification_level = [level_0_magnification / lv for lv in downsample_level] - mpp = [ - wsi_properties[openslide.PROPERTY_NAME_MPP_X], - wsi_properties[openslide.PROPERTY_NAME_MPP_Y], - ] - mpp = np.array(mpp) - metadata = [ - ("available_mag", magnification_level), # highest to lowest mag - ("base_mag", magnification_level[0]), - ("vendor", wsi_properties[openslide.PROPERTY_NAME_VENDOR]), - ("mpp ", mpp), - ("base_shape", np.array(self.file_ptr.dimensions)), + ('available_mag', magnification_level), # highest to lowest mag + ('base_mag' , magnification_level[0]), + ('vendor' , wsi_properties[openslide.PROPERTY_NAME_VENDOR]), + ('base_mpp' , mpp), + ('base_shape' , np.array(self.file_ptr.dimensions)), ] return OrderedDict(metadata) - + def read_region(self, coords, size): - """Must call `prepare_reading` before hand. + """ + Must call `prepare_image` before hand Args: - coords (tuple): (dims_x, dims_y), - top left coordinates of image region at selected - `read_mag` or `read_mpp` from `prepare_reading` - size (tuple): (dims_x, dims_y) - width and height of image region at selected - `read_mag` or `read_mpp` from `prepare_reading` - + coords (tuple): top left coordinates of image region at requested mag/mpp in `prepare_reading` + read_level_size (tuple): dimensions of image region at requested mag/mpp in `prepare_reading` (dims_x, dims_y) """ - if self.image_ptr is None: - # convert coord from read lv to lv zero - lv_0_shape = np.array(self.file_ptr.level_dimensions[0]) - lv_r_shape = np.array(self.file_ptr.level_dimensions[self.read_lv]) - up_sample = (lv_0_shape / lv_r_shape)[0] - new_coord = [0, 0] - new_coord[0] = int(coords[0] * up_sample) - new_coord[1] = int(coords[1] * up_sample) - region = self.file_ptr.read_region(new_coord, self.read_lv, size) - else: - region = self.image_ptr[ - coords[1] : coords[1] + size[1], coords[0] : coords[0] + size[0] - ] - return np.array(region)[..., :3] - def get_full_img(self, read_mag=None, read_mpp=None): - """Only use `read_mag` or `read_mpp`, not both, prioritize `read_mpp`. + # convert coord from read lv to lv zero + coord_lv0 = [0, 0] + coord_lv0[0] = int(coords[0] / self.scale_from_lv0) + coord_lv0[1] = int(coords[1] / self.scale_from_lv0) - `read_mpp` is in X, Y format. - """ + size_at_read_lv = (size / self.scale_from_hires_lv).astype(np.int32) + region = self.file_ptr.read_region(coord_lv0, self.read_lv, size_at_read_lv) + + region = np.array(region)[...,:3] + if self.scale_from_hires_lv is not None: + interp = cv2.INTER_LINEAR + region = cv2.resize(region, tuple(size), interpolation=interp) + return region - read_lv, scale_factor = self._get_read_info( - read_mag=read_mag, read_mpp=read_mpp - ) + def get_full_img(self, read_mag=None, read_mpp=None): + + read_lv, scale_from_hires_lv, scale_from_lv0 = self._get_read_info( + read_mag=read_mag, + read_mpp=read_mpp) read_size = self.file_ptr.level_dimensions[read_lv] wsi_img = self.file_ptr.read_region((0, 0), read_lv, read_size) - wsi_img = np.array(wsi_img)[..., :3] # remove alpha channel - if scale_factor is not None: + wsi_img = np.array(wsi_img)[...,:3] # remove alpha channel + if scale_from_hires_lv is not None: # now rescale then return - if scale_factor > 1.0: - interp = cv2.INTER_CUBIC - else: - interp = cv2.INTER_LINEAR - wsi_img = cv2.resize( - wsi_img, (0, 0), fx=scale_factor, fy=scale_factor, interpolation=interp - ) - return wsi_img - + interp = cv2.INTER_LINEAR + wsi_img = cv2.resize(wsi_img, (0, 0), + fx=scale_from_hires_lv, + fy=scale_from_hires_lv, + interpolation=interp) + return wsi_img def get_file_handler(path, backend): if backend in [ diff --git a/models/__init__.py b/models/__init__.py old mode 100644 new mode 100755 diff --git a/models/hovernet/__init__.py b/models/hovernet/__init__.py old mode 100644 new mode 100755 diff --git a/models/hovernet/net_desc.py b/models/hovernet/net_desc.py old mode 100644 new mode 100755 diff --git a/models/hovernet/net_utils.py b/models/hovernet/net_utils.py old mode 100644 new mode 100755 diff --git a/models/hovernet/opt.py b/models/hovernet/opt.py old mode 100644 new mode 100755 diff --git a/models/hovernet/post_proc.py b/models/hovernet/post_proc.py old mode 100644 new mode 100755 diff --git a/models/hovernet/run_desc.py b/models/hovernet/run_desc.py old mode 100644 new mode 100755 diff --git a/models/hovernet/targets.py b/models/hovernet/targets.py old mode 100644 new mode 100755 diff --git a/models/hovernet/utils.py b/models/hovernet/utils.py old mode 100644 new mode 100755 diff --git a/requirements.txt b/requirements.txt old mode 100644 new mode 100755 diff --git a/run_infer.py b/run_infer.py old mode 100644 new mode 100755 index 15477e23..61be79f8 --- a/run_infer.py +++ b/run_infer.py @@ -181,6 +181,7 @@ infer = InferManager(**method_args) infer.process_file_list(run_args) else: - from infer.wsi import InferManager + # from infer.wsi import InferManager + from infer.super_wsi import InferManager infer = InferManager(**method_args) infer.process_wsi_list(run_args) diff --git a/run_tile.sh b/run_tile.sh old mode 100644 new mode 100755 diff --git a/run_train.py b/run_train.py old mode 100644 new mode 100755 diff --git a/run_wsi.sh b/run_wsi.sh old mode 100644 new mode 100755 index 766421c6..863d0f28 --- a/run_wsi.sh +++ b/run_wsi.sh @@ -4,12 +4,14 @@ python run_infer.py \ --type_info_path=type_info.json \ --batch_size=64 \ --model_mode=fast \ ---model_path=../pretrained/hovernet_fast_pannuke_type_tf2pytorch.tar \ ---nr_inference_workers=8 \ ---nr_post_proc_workers=16 \ +--model_path=../../pretrained/hovernet_fast_pannuke_pytorch.tar \ +--nr_inference_workers=4 \ +--nr_post_proc_workers=32 \ wsi \ ---input_dir=dataset/sample_wsis/wsi/ \ ---output_dir=dataset/sample_wsis/out/ \ ---input_mask_dir=dataset/sample_wsis/msk/ \ +--input_dir=exp_output/wsi/samples/full/ \ +--output_dir=exp_output/wsi/pred/ \ +--input_mask_dir=exp_output/wsi/samples/mask/ \ --save_thumb \ ---save_mask +--save_mask \ +--ambiguous_size=328 \ +--tile_shape=2048 \ No newline at end of file diff --git a/seg.gif b/seg.gif old mode 100644 new mode 100755 diff --git a/type_info.json b/type_info.json old mode 100644 new mode 100755 diff --git a/variables_tf2pytorch.csv b/variables_tf2pytorch.csv old mode 100644 new mode 100755