deepliif 1.1.5__py3-none-any.whl → 1.1.7__py3-none-any.whl

cli.py

@@ -11,7 +11,7 @@ from PIL import Image
 
 from deepliif.data import create_dataset, transform
 from deepliif.models import inference, postprocess, compute_overlap, init_nets, DeepLIIFModel, infer_modalities, infer_results_for_wsi
-from deepliif.util import allowed_file, Visualizer, get_information
+from deepliif.util import allowed_file, Visualizer, get_information, test_diff_original_serialized, disable_batchnorm_tracking_stats
 from deepliif.util.util import mkdirs, check_multi_scale
 # from deepliif.util import infer_results_for_wsi
 
@@ -461,21 +461,51 @@ def trainlaunch(**kwargs):
 @cli.command()
 @click.option('--models-dir', default='./model-server/DeepLIIF_Latest_Model', help='reads models from here')
 @click.option('--output-dir', help='saves results here.')
-def serialize(models_dir, output_dir):
+@click.option('--device', default='cpu', type=str, help='device to load model, either cpu or gpu')
+@click.option('--verbose', default=0, type=int,help='saves results here.')
+def serialize(models_dir, output_dir, device, verbose):
     """Serialize DeepLIIF models using Torchscript
     """
     output_dir = output_dir or models_dir
+    ensure_exists(output_dir)
 
     sample = transform(Image.new('RGB', (512, 512)))
-
+    
     with click.progressbar(
             init_nets(models_dir, eager_mode=True).items(),
             label='Tracing nets',
             item_show_func=lambda n: n[0] if n else n
     ) as bar:
         for name, net in bar:
-            traced_net = torch.jit.trace(net, sample)
+            # the model should be in eval model so that there won't be randomness in tracking brought by dropout etc. layers
+            # https://github.com/pytorch/pytorch/issues/23999#issuecomment-747832122
+            net = net.eval()
+            net = disable_batchnorm_tracking_stats(net)
+            net = net.cpu()
+            if name.startswith('GS'):
+                traced_net = torch.jit.trace(net, torch.cat([sample, sample, sample], 1))
+            else:
+                traced_net = torch.jit.trace(net, sample)
+                # traced_net = torch.jit.script(net)
             traced_net.save(f'{output_dir}/{name}.pt')
+    
+    # test: whether the original and the serialized model produces highly similar predictions
+    print('testing similarity between prediction from original vs serialized models...')
+    models_original = init_nets(models_dir,eager_mode=True)
+    models_serialized = init_nets(output_dir,eager_mode=False)
+    if device == 'gpu':
+        sample = sample.cuda()
+    else:
+        sample = sample.cpu()
+    for name in models_serialized.keys():
+        print(name,':')
+        model_original = models_original[name].cuda().eval() if device=='gpu' else models_original[name].cpu().eval()
+        model_serialized = models_serialized[name].cuda() if device=='gpu' else models_serialized[name].cpu().eval()
+        if name.startswith('GS'):
+            test_diff_original_serialized(model_original,model_serialized,torch.cat([sample, sample, sample], 1),verbose)
+        else:
+            test_diff_original_serialized(model_original,model_serialized,sample,verbose)
+        print('PASS')
          
 
 @cli.command()
@@ -486,7 +516,11 @@ def serialize(models_dir, output_dir):
 @click.option('--region-size', default=20000, help='Due to limits in the resources, the whole slide image cannot be processed in whole.'
                                                    'So the WSI image is read region by region. '
                                                    'This parameter specifies the size each region to be read into GPU for inferrence.')
-def test(input_dir, output_dir, tile_size, model_dir, region_size):
+@click.option('--eager-mode', is_flag=True, help='use eager mode (loading original models, otherwise serialized ones)')
+@click.option('--color-dapi', is_flag=True, help='color dapi image to produce the same coloring as in the paper')
+@click.option('--color-marker', is_flag=True, help='color marker image to produce the same coloring as in the paper')
+def test(input_dir, output_dir, tile_size, model_dir, region_size, eager_mode,
+         color_dapi, color_marker):
     
     """Test trained models
     """
@@ -507,7 +541,7 @@ def test(input_dir, output_dir, tile_size, model_dir, region_size):
                 print(time.time() - start_time)
             else:
                 img = Image.open(os.path.join(input_dir, filename)).convert('RGB')
-                images, scoring = infer_modalities(img, tile_size, model_dir)
+                images, scoring = infer_modalities(img, tile_size, model_dir, eager_mode, color_dapi, color_marker)
 
                 for name, i in images.items():
                     i.save(os.path.join(
@@ -589,6 +623,15 @@ def prepare_testing_data(input_dir, dataset_dir):
             cv2.imwrite(os.path.join(test_dir, img), np.concatenate([image, image, image, image, image, image], 1))
 
 
+# to load pickle file saved from gpu in a cpu environment: https://github.com/pytorch/pytorch/issues/16797#issuecomment-633423219
+from io import BytesIO
+class CPU_Unpickler(pickle.Unpickler):
+    def find_class(self, module, name):
+        if module == 'torch.storage' and name == '_load_from_bytes':
+            return lambda b: torch.load(BytesIO(b), map_location='cpu')
+        else: return super().find_class(module, name)
+
+
 @cli.command()
 @click.option('--pickle-dir', required=True, help='directory where the pickled snapshots are stored')
 def visualize(pickle_dir):
@@ -599,8 +642,8 @@ def visualize(pickle_dir):
         time.sleep(1)
 
     params_opt = pickle.load(open(path_init,'rb'))
-    params_opt['remote'] = False
-    visualizer = Visualizer(**params_opt)   # create a visualizer that display/save images and plots
+    params_opt.remote = False
+    visualizer = Visualizer(params_opt)   # create a visualizer that display/save images and plots
 
     paths_plot = {'display_current_results':os.path.join(pickle_dir,'display_current_results.pickle'),
                 'plot_current_losses':os.path.join(pickle_dir,'plot_current_losses.pickle')}
@@ -612,7 +655,7 @@ def visualize(pickle_dir):
             try:
                 last_modified_time_plot = os.path.getmtime(path_plot)
                 if last_modified_time_plot > last_modified_time[method]:
-                    params_plot = pickle.load(open(path_plot,'rb'))
+                    params_plot = CPU_Unpickler(open(path_plot,'rb')).load()
                     last_modified_time[method] = last_modified_time_plot
                     getattr(visualizer,method)(**params_plot)
                     print(f'{method} refreshed, last modified time {time.ctime(last_modified_time[method])}')

deepliif/models/__init__.py

@@ -88,7 +88,10 @@ def create_model(opt):
 
 
 def load_torchscript_model(model_pt_path, device):
-    return torch.jit.load(model_pt_path, map_location=device)
+    net = torch.jit.load(model_pt_path, map_location=device)
+    net = disable_batchnorm_tracking_stats(net)
+    net.eval()
+    return net
 
 
 def read_model_params(file_addr):
@@ -132,7 +135,8 @@ def load_eager_models(model_dir, devices):
             os.path.join(model_dir, f'latest_net_{n}.pth'),
             map_location=devices[n]
         ))
-        nets[n] = net
+        nets[n] = disable_batchnorm_tracking_stats(net)
+        nets[n].eval()
 
     for n in ['G51', 'G52', 'G53', 'G54', 'G55']:
         net = UnetGenerator(input_nc, output_nc, 9, ngf, norm_layer=norm_layer, use_dropout=use_dropout)
@@ -140,7 +144,8 @@ def load_eager_models(model_dir, devices):
             os.path.join(model_dir, f'latest_net_{n}.pth'),
             map_location=devices[n]
         ))
-        nets[n] = net
+        nets[n] = disable_batchnorm_tracking_stats(net)
+        nets[n].eval()
 
     return nets
 
@@ -185,7 +190,12 @@ def compute_overlap(img_size, tile_size):
     return tile_size // 4
 
 
-def run_torchserve(img, model_path=None):
+def run_torchserve(img, model_path=None, eager_mode=False):
+    """
+    eager_mode: not used in this function; put in place to be consistent with run_dask
+           so that run_wrapper() could call either this function or run_dask with
+           same syntax
+    """
     buffer = BytesIO()
     torch.save(transform(img.resize((512, 512))), buffer)
 
@@ -203,9 +213,9 @@ def run_torchserve(img, model_path=None):
     return {k: tensor_to_pil(deserialize_tensor(v)) for k, v in res.json().items()}
 
 
-def run_dask(img, model_path):
+def run_dask(img, model_path, eager_mode=False):
     model_dir = os.getenv('DEEPLIIF_MODEL_DIR', model_path)
-    nets = init_nets(model_dir)
+    nets = init_nets(model_dir, eager_mode)
 
     ts = transform(img.resize((512, 512)))
 
@@ -237,7 +247,7 @@ def is_empty(tile):
     return True if calculate_background_area(tile) > 98 else False
 
 
-def run_wrapper(tile, run_fn, model_path):
+def run_wrapper(tile, run_fn, model_path, eager_mode=False):
     if is_empty(tile):
         return {
             'G1': Image.new(mode='RGB', size=(512, 512), color=(201, 211, 208)),
@@ -247,17 +257,17 @@ def run_wrapper(tile, run_fn, model_path):
             'G5': Image.new(mode='RGB', size=(512, 512), color=(0, 0, 0))
         }
     else:
-        return run_fn(tile, model_path)
-
+        return run_fn(tile, model_path, eager_mode)
 
-def inference(img, tile_size, overlap_size, model_path, use_torchserve=False):
 
+def inference_old(img, tile_size, overlap_size, model_path, use_torchserve=False, eager_mode=False,
+                  color_dapi=False, color_marker=False):
     
     tiles = list(generate_tiles(img, tile_size, overlap_size))
 
     run_fn = run_torchserve if use_torchserve else run_dask
     # res = [Tile(t.i, t.j, run_fn(t.img, model_path)) for t in tiles]
-    res = [Tile(t.i, t.j, run_wrapper(t.img, run_fn, model_path)) for t in tiles]
+    res = [Tile(t.i, t.j, run_wrapper(t.img, run_fn, model_path, eager_mode)) for t in tiles]
 
     def get_net_tiles(n):
         return [Tile(t.i, t.j, t.img[n]) for t in res]
@@ -276,12 +286,14 @@ def inference(img, tile_size, overlap_size, model_path, use_torchserve=False):
 
     images['DAPI'] = stitch(get_net_tiles('G2'), tile_size, overlap_size).resize(img.size)
     dapi_pix = np.array(images['DAPI'].convert('L').convert('RGB'))
-    dapi_pix[:, :, 0] = 0
+    if color_dapi:
+      dapi_pix[:, :, 0] = 0
     images['DAPI'] = Image.fromarray(dapi_pix)
     images['Lap2'] = stitch(get_net_tiles('G3'), tile_size, overlap_size).resize(img.size)
     images['Marker'] = stitch(get_net_tiles('G4'), tile_size, overlap_size).resize(img.size)
     marker_pix = np.array(images['Marker'].convert('L').convert('RGB'))
-    marker_pix[:, :, 2] = 0
+    if color_marker:
+      marker_pix[:, :, 2] = 0
     images['Marker'] = Image.fromarray(marker_pix)
 
     # images['Marker'] = stitch(
@@ -294,6 +306,52 @@ def inference(img, tile_size, overlap_size, model_path, use_torchserve=False):
     return images
 
 
+def inference(img, tile_size, overlap_size, model_path, use_torchserve=False, eager_mode=False,
+              color_dapi=False, color_marker=False):
+
+    rescaled, rows, cols = format_image_for_tiling(img, tile_size, overlap_size)
+
+    run_fn = run_torchserve if use_torchserve else run_dask
+
+    images = {}
+    images['Hema'] = create_image_for_stitching(tile_size, rows, cols)
+    images['DAPI'] = create_image_for_stitching(tile_size, rows, cols)
+    images['Lap2'] = create_image_for_stitching(tile_size, rows, cols)
+    images['Marker'] = create_image_for_stitching(tile_size, rows, cols)
+    images['Seg'] = create_image_for_stitching(tile_size, rows, cols)
+
+    for i in range(cols):
+        for j in range(rows):
+            tile = extract_tile(rescaled, tile_size, overlap_size, i, j)
+            res = run_wrapper(tile, run_fn, model_path, eager_mode)
+
+            stitch_tile(images['Hema'], res['G1'], tile_size, overlap_size, i, j)
+            stitch_tile(images['DAPI'], res['G2'], tile_size, overlap_size, i, j)
+            stitch_tile(images['Lap2'], res['G3'], tile_size, overlap_size, i, j)
+            stitch_tile(images['Marker'], res['G4'], tile_size, overlap_size, i, j)
+            stitch_tile(images['Seg'], res['G5'], tile_size, overlap_size, i, j)
+
+    images['Hema'] = images['Hema'].resize(img.size)
+    images['DAPI'] = images['DAPI'].resize(img.size)
+    images['Lap2'] = images['Lap2'].resize(img.size)
+    images['Marker'] = images['Marker'].resize(img.size)
+    images['Seg'] = images['Seg'].resize(img.size)
+
+    if color_dapi:
+        matrix = (       0,        0,        0, 0,
+                  299/1000, 587/1000, 114/1000, 0,
+                  299/1000, 587/1000, 114/1000, 0)
+        images['DAPI'] = images['DAPI'].convert('RGB', matrix)
+
+    if color_marker:
+        matrix = (299/1000, 587/1000, 114/1000, 0,
+                  299/1000, 587/1000, 114/1000, 0,
+                         0,        0,        0, 0)
+        images['Marker'] = images['Marker'].convert('RGB', matrix)
+
+    return images
+
+
 def postprocess(img, seg_img, thresh=80, noise_objects_size=20, small_object_size=50):
     mask_image = create_basic_segmentation_mask(np.array(img), np.array(seg_img),
                                                 thresh, noise_objects_size, small_object_size)
@@ -312,7 +370,8 @@ def postprocess(img, seg_img, thresh=80, noise_objects_size=20, small_object_siz
     return images, scoring
 
 
-def infer_modalities(img, tile_size, model_dir):
+def infer_modalities(img, tile_size, model_dir, eager_mode=False,
+                     color_dapi=False, color_marker=False):
     """
     This function is used to infer modalities for the given image using a trained model.
     :param img: The input image.
@@ -329,7 +388,10 @@ def infer_modalities(img, tile_size, model_dir):
         img,
         tile_size=tile_size,
         overlap_size=compute_overlap(img.size, tile_size),
-        model_path=model_dir
+        model_path=model_dir,
+        eager_mode=eager_mode,
+        color_dapi=color_dapi,
+        color_marker=color_marker
     )
 
     post_images, scoring = postprocess(img, images['Seg'], small_object_size=20)

deepliif/models/base_model.py

@@ -3,6 +3,7 @@ import torch
 from collections import OrderedDict
 from abc import ABC, abstractmethod
 from . import networks
+from ..util import disable_batchnorm_tracking_stats
 
 
 class BaseModel(ABC):
@@ -90,6 +91,7 @@ class BaseModel(ABC):
             if isinstance(name, str):
                 net = getattr(self, 'net' + name)
                 net.eval()
+                net = disable_batchnorm_tracking_stats(net)
 
     def test(self):
         """Forward function used in test time.

deepliif/util/__init__.py

@@ -88,6 +88,36 @@ def stitch(tiles, tile_size, overlap_size):
     return new_im
 
 
+def format_image_for_tiling(img, tile_size, overlap_size):
+    mean_background_val = calculate_background_mean_value(img)
+    img = img.resize(output_size(img, tile_size))
+    # Adding borders with size of given overlap around the whole slide image
+    img = ImageOps.expand(img, border=overlap_size, fill=tuple(mean_background_val))
+    rows = int(img.height / tile_size)
+    cols = int(img.width / tile_size)
+    return img, rows, cols
+
+
+def extract_tile(img, tile_size, overlap_size, i, j):
+    return img.crop((
+        i * tile_size, j * tile_size,
+        i * tile_size + tile_size + 2 * overlap_size,
+        j * tile_size + tile_size + 2 * overlap_size
+    ))
+
+
+def create_image_for_stitching(tile_size, rows, cols):
+    width = tile_size * cols
+    height = tile_size * rows
+    return Image.new('RGB', (width, height))
+
+
+def stitch_tile(img, tile, tile_size, overlap_size, i, j):
+    tile = tile.resize((tile_size + 2 * overlap_size, tile_size + 2 * overlap_size))
+    tile = tile.crop((overlap_size, overlap_size, overlap_size + tile_size, overlap_size + tile_size))
+    img.paste(tile, (i * tile_size, j * tile_size))
+
+
 def calculate_background_mean_value(img):
     img = cv2.fastNlMeansDenoisingColored(np.array(img), None, 10, 10, 7, 21)
     img = np.array(img, dtype=float)
@@ -349,3 +379,39 @@ def read_results_from_pickle_file(input_addr):
     pickle_obj.close()
     return results
 
+def test_diff_original_serialized(model_original,model_serialized,example,verbose=0):
+    threshold = 10
+
+    orig_res = model_original(example)
+    if verbose > 0:
+        print('Original:')
+        print(orig_res.shape)
+        print(orig_res[0, 0:10])
+        print('min abs value:{}'.format(torch.min(torch.abs(orig_res))))
+
+    ts_res = model_serialized(example)
+    if verbose > 0:
+        print('Torchscript:')
+        print(ts_res.shape)
+        print(ts_res[0, 0:10])
+        print('min abs value:{}'.format(torch.min(torch.abs(ts_res))))
+
+    abs_diff = torch.abs(orig_res-ts_res)
+    if verbose > 0:
+        print('Dif sum:')
+        print(torch.sum(abs_diff))
+        print('max dif:{}'.format(torch.max(abs_diff)))
+
+    assert torch.sum(abs_diff) <= threshold, f"Sum of difference in predicted values {torch.sum(abs_diff)} is larger than threshold {threshold}"
+
+def disable_batchnorm_tracking_stats(model):
+    # https://discuss.pytorch.org/t/performance-highly-degraded-when-eval-is-activated-in-the-test-phase/3323/16
+    # https://discuss.pytorch.org/t/performance-highly-degraded-when-eval-is-activated-in-the-test-phase/3323/67
+    # https://github.com/pytorch/pytorch/blob/ca39c5b04e30a67512589cafbd9d063cc17168a5/torch/nn/modules/batchnorm.py#L158
+    for m in model.modules():
+        for child in m.children():
+            if type(child) == torch.nn.BatchNorm2d:
+                child.track_running_stats = False
+                child.running_mean = None
+                child.running_var = None
+    return model