Diff of /pathflowai/stain_norm.py [000000] .. [e9500f]

Switch to side-by-side view

--- a
+++ b/pathflowai/stain_norm.py
@@ -0,0 +1,65 @@
+import cv2
+import sys
+import fire
+import histomicstk
+import histomicstk as htk
+import openslide
+import dask
+import tqdm
+import numpy as np
+from dask.diagnostics import ProgressBar
+from pathflowai.utils import generate_tissue_mask
+from histomicstk.preprocessing.color_normalization.\
+    deconvolution_based_normalization import deconvolution_based_normalization
+
+W_target = np.array([
+            [0.6185391,  0.1576997,  -0.01119131],
+            [0.7012888,  0.8638838,  0.45586256],
+            [0.3493163,  0.4657428, -0.85597752]
+        ])
+
+def return_norm_image(img,mask,W_source=None,W_target=None):
+    img=deconvolution_based_normalization(
+        img, W_source=W_source, W_target=W_target, im_target=None,
+        stains=['hematoxylin', 'eosin'], mask_out=~mask,
+        stain_unmixing_routine_params={"I_0":215})
+    return img
+
+def check_ext(image_file):
+    return any([image_file.endswith(ext) for ext in ['.svs','.png','.jpg','.jpeg','.tiff','.tif']])
+
+def stain_norm(image_file,compression=10,patch_size=1024):
+    if check_ext(image_file):
+        img = openslide.open_slide(image_file)
+        image = np.array(img.read_region((0,0), 0, img.level_dimensions[0]))[...,:3]
+    elif image_file.endswith(".npy"):
+        image=np.load(image_file)
+    else: raise NotImplementedError
+    mask=generate_tissue_mask(image,compression=compression,keep_holes=False)
+    img_small=cv2.resize(image,None,fx=1/compression,fy=1/compression)
+    mask_small=cv2.resize(mask.astype(int),None,fx=1/compression,fy=1/compression,interpolation=cv2.INTER_NEAREST).astype(bool)
+    W_source = htk.preprocessing.color_deconvolution.rgb_separate_stains_macenko_pca(img_small, 215)
+    W_source = htk.preprocessing.color_deconvolution._reorder_stains(W_source)
+    res=[]
+    coords=[]
+    for i in np.arange(0,image.shape[0]-patch_size,patch_size):
+        for j in np.arange(0,image.shape[1]-patch_size,patch_size):
+            if mask[i:i+patch_size,j:j+patch_size].mean():
+                coords.append((i,j))
+                res.append(dask.delayed(return_norm_image)(image[i:i+patch_size,j:j+patch_size],mask[i:i+patch_size,j:j+patch_size],W_source,W_target))
+    with ProgressBar():
+        res_returned=dask.compute(*res,scheduler="processes")
+    img_new=np.ones(image.shape).astype(np.uint8)*255
+    for k in tqdm.trange(len(coords)):
+        i,j=coords[k]
+        img_new[i:i+patch_size,j:j+patch_size]=res_returned[k]
+    return img_new
+
+def stain_norm_pipeline(image_file="stain_in.svs",
+                        npy_out='stain_out.npy',
+                        compression=10,
+                        patch_size=1024):
+    np.save(npy_out,stain_norm(image_file,compression,patch_size))
+
+if __name__=="__main__":
+    fire.Fire(stain_norm_pipeline)