adding multiple dataset support, updating gitignore

Author: nauyan

.gitignore

@@ -3,4 +3,6 @@ notebooks
 .DS_Store
 checkpoints
 run_logs
-data
+data
+.ipynb_checkpoints
+temp.ipynb

config.py

@@ -1,21 +1,30 @@
 version = 1.0
 
-dataset_name = "consep"
-dataset_raw = "data/archives/CoNSeP/"
-train_dir = "data/consep/train/"
-test_dir = "data/consep/test/"
+# dataset_name = "consep"
+# dataset_raw = "data/archives/CoNSeP/"
+# train_dir = "data/consep/train/"
+# test_dir = "data/consep/test/"
 
+# Combined Dataset CoNSep, PanNuke, Lizard
+dataset_name = "all"
+dataset_raw = None
+train_dir = "data/all/train/"
+test_dir = "data/all/test/"
 
 checkpoints_dir = "checkpoints"
 tensorboard_logs = "run_logs"
 
 encoder = "tu-tf_efficientnet_b5_ns"
 encoder_weights = None
 
-inference_weights = "NC-Net_consep_metric.pth"
+# CoNSep Weights
+# inference_weights = "NC-Net_consep_metric.pth"
+
+# All Weights
+inference_weights = "NC-Net_all_metric.pth"
 
 device = "cuda"
 batch_size = 16
 epochs = 1000
 learning_rate = 0.01
-watershed_threshold = 0.2
+watershed_threshold = 0.2

create_dataset.py

@@ -48,7 +48,6 @@ def read_annotation(annotPath):
     return centroid_map, inst_map, nuclear_map, type_map
 
 
-
 test_folder = glob.glob(os.path.join(config.dataset_raw, "Test/Images"))[0]
 train_folder = glob.glob(os.path.join(config.dataset_raw, "Train/Images"))[0]
 
@@ -70,8 +69,7 @@ def read_annotation(annotPath):
     print(img_path)
     image = read_image(img_path)
     centroid_map, inst_map, nuclear_map, type_map = read_annotation(
-        img_path.replace("Images", "Labels").replace(".png", ".mat")
-    )
+        img_path.replace("Images", "Labels").replace(".png", ".mat"))
 
     output_mask = np.zeros((image.shape[0], image.shape[1], 7))
     output_mask[:, :, :3] = image / 255.0
@@ -80,17 +78,14 @@ def read_annotation(annotPath):
     output_mask[:, :, 5] = centroid_map  # centroids_mask
     output_mask[:, :, 6] = type_map  # type mask
 
-    output = util.view_as_windows(
-        output_mask, (256, 256, 7), step=(245, 245, 7)
-    ).reshape(
-        -1, 256, 256, 7
-    ) 
+    output = util.view_as_windows(output_mask, (256, 256, 7),
+                                  step=(245, 245, 7)).reshape(-1, 256, 256, 7)
 
     for idx in range(output.shape[0]):
         np.save(
             "{}/{}_{}.npy".format(
-                config.test_dir, os.path.basename(img_path).replace(".png", ""), idx
-            ),
+                config.test_dir,
+                os.path.basename(img_path).replace(".png", ""), idx),
             output[idx],
         )
 
@@ -99,8 +94,7 @@ def read_annotation(annotPath):
 
     image = read_image(img_path)
     centroid_map, inst_map, nuclear_map, type_map = read_annotation(
-        img_path.replace("Images", "Labels").replace(".png", ".mat")
-    )
+        img_path.replace("Images", "Labels").replace(".png", ".mat"))
 
     output_mask = np.zeros((image.shape[0], image.shape[1], 7))
     output_mask[:, :, :3] = image / 255.0
@@ -109,16 +103,13 @@ def read_annotation(annotPath):
     output_mask[:, :, 5] = centroid_map  # centroids_mask
     output_mask[:, :, 6] = type_map  # type mask
 
-    output = util.view_as_windows(
-        output_mask, (256, 256, 7), step=(128, 128, 7)
-    ).reshape(
-        -1, 256, 256, 7
-    )  
+    output = util.view_as_windows(output_mask, (256, 256, 7),
+                                  step=(128, 128, 7)).reshape(-1, 256, 256, 7)
 
     for idx in range(output.shape[0]):
         np.save(
             "{}/{}_{}.npy".format(
-                config.train_dir, os.path.basename(img_path).replace(".png", ""), idx
-            ),
+                config.train_dir,
+                os.path.basename(img_path).replace(".png", ""), idx),
             output[idx],
         )

create_dataset_all.py

@@ -0,0 +1,21 @@
+import shutil
+import os
+
+from tqdm import tqdm
+from glob import glob
+from sklearn.model_selection import train_test_split
+
+samples = glob("data/consep/*/*.npy")
+samples.extend(glob("data/liz1/*/*.npy"))
+samples.extend(glob("data/pan1/*/*.npy"))
+
+X_train, X_test, y_train, y_test = train_test_split(samples,
+                                                    samples,
+                                                    test_size=0.1,
+                                                    random_state=42)
+
+for sample in tqdm(X_train, total=len(X_train)):
+    shutil.copyfile(sample, f"data/all/train/{os.path.basename(sample)}")
+
+for sample in tqdm(X_test, total=len(X_test)):
+    shutil.copyfile(sample, f"data/all/test/{os.path.basename(sample)}")

create_dataset_consep.py

@@ -0,0 +1,94 @@
+from skimage import data, io, util
+from matplotlib import pyplot as plt
+
+import scipy.io as sio
+import numpy as np
+
+import glob
+import cv2
+import os
+import shutil
+import config
+
+
+def create_centroid_mask(centroids, height, width):
+    cent_img = np.zeros((height, width))
+    cent_img[tuple(centroids.T)] = 1
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    centroids_mask = cv2.dilate(cent_img, kernel, iterations=2)
+
+    return centroids_mask
+
+
+def read_image(imagePath):
+    image = io.imread(imagePath)[:, :, :3]
+
+    return image
+
+
+def fix_type(type_map):
+    type_map[(type_map == 3) | (type_map == 4)] = 3
+    type_map[(type_map == 5) | (type_map == 6) | (type_map == 7)] = 4
+    return type_map
+
+
+def read_annotation(annotPath):
+    mask = sio.loadmat(annotPath)
+
+    type_map = fix_type(mask["type_map"])
+    centroid_map = create_centroid_mask(
+        mask["inst_centroid"].astype(int)[:, ::-1],
+        mask["inst_map"].shape[0],
+        mask["inst_map"].shape[1],
+    )
+    inst_map = mask["inst_map"]
+    nuclear_map = inst_map + 0
+    nuclear_map[nuclear_map != 0] = 1
+
+    return centroid_map, inst_map, nuclear_map, type_map
+
+
+def init_dir(input_dir):
+    if not os.path.exists(input_dir):
+        os.makedirs(input_dir)
+    else:
+        shutil.rmtree(input_dir)
+        os.makedirs(input_dir)
+
+
+def process_subset(input_dir, output_dir, image_size=256, step_size=245):
+
+    for img_path in glob.glob("{}/*".format(input_dir)):
+        print(img_path)
+        image = read_image(img_path)
+        centroid_map, inst_map, nuclear_map, type_map = read_annotation(
+            img_path.replace("Images", "Labels").replace(".png", ".mat"))
+
+        output_mask = np.zeros((image.shape[0], image.shape[1], 7))
+        output_mask[:, :, :3] = image / 255.0
+        output_mask[:, :, 3] = inst_map  # inst_mask
+        output_mask[:, :, 4] = nuclear_map  # binary mask
+        output_mask[:, :, 5] = centroid_map  # centroids_mask
+        output_mask[:, :, 6] = type_map  # type mask
+
+        output = util.view_as_windows(output_mask, (image_size, image_size, 7),
+                                      step=(step_size, step_size, 7)).reshape(
+                                          -1, image_size, image_size, 7)
+
+        for idx in range(output.shape[0]):
+            np.save(
+                "{}/{}_{}.npy".format(
+                    output_dir,
+                    os.path.basename(img_path).replace(".png", ""), idx),
+                output[idx],
+            )
+
+
+test_folder = glob.glob(os.path.join(config.dataset_raw, "Test/Images"))[0]
+train_folder = glob.glob(os.path.join(config.dataset_raw, "Train/Images"))[0]
+
+init_dir(config.train_dir)
+init_dir(config.test_dir)
+
+process_subset(train_folder, config.train_dir)
+process_subset(test_folder, config.test_dir)

create_dataset_lizard.py

@@ -0,0 +1,104 @@
+import glob
+from skimage import io
+import scipy.io as sio
+from skimage import data, io, util
+from matplotlib import pyplot as plt
+import numpy as np
+
+import cv2
+import os
+import random
+from skimage.transform import rescale, resize, downscale_local_mean
+from tqdm import tqdm
+import shutil
+import pandas as pd
+
+
+def create_centroid_mask(centroids, height, width):
+    cent_img = np.zeros((height, width))
+    cent_img[tuple(centroids.T)] = 1
+    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (5, 5))
+    centroids_mask = cv2.dilate(cent_img, kernel, iterations=2)
+
+    return centroids_mask
+
+
+def rescale_sample(sample):
+    sample = rescale(sample,
+                     2,
+                     anti_aliasing=False,
+                     multichannel=True,
+                     preserve_range=True,
+                     order=0)
+    return sample
+
+
+def read_image(imagePath):
+    image = io.imread(imagePath)[:, :, :3]
+    return image
+
+
+def read_annotation(annotPath):
+    mask = sio.loadmat(annotPath)
+
+    centroid_map = create_centroid_mask(mask['centroid'].astype(int)[:, ::-1],
+                                        mask['inst_map'].shape[0],
+                                        mask['inst_map'].shape[1])
+    inst_map = mask['inst_map']
+    nuclear_map = inst_map + 0
+    nuclear_map[nuclear_map != 0] = 1
+
+    return centroid_map, inst_map, nuclear_map
+
+
+def init_dir(input_dir):
+    if not os.path.exists(input_dir):
+        os.makedirs(input_dir)
+    else:
+        shutil.rmtree(input_dir)
+        os.makedirs(input_dir)
+
+
+meta = pd.read_csv("../datasets/lizard/Lizard_Labels/info.csv")
+# print(meta.columns)
+for fold_idx in range(1, 4):
+    print("Generating Fold", fold_idx)
+
+    train_dir = "data/liz{}/train".format(fold_idx)
+    test_dir = "data/liz{}/test".format(fold_idx)
+
+    init_dir(train_dir)
+    init_dir(test_dir)
+
+    for idx in tqdm(range(meta.shape[0]), total=meta.shape[0]):
+
+        image = read_image(
+            glob.glob("../datasets/lizard/Lizard_Images*/{}.png".format(
+                meta['Filename'][idx]))[0])
+
+        centroid_map, inst_map, nuclear_map = read_annotation(
+            "../datasets/lizard/Lizard_Labels/Labels/{}.mat".format(
+                meta['Filename'][idx]))
+        type_map = np.zeros((image.shape[0], image.shape[1]))
+
+        output_mask = np.zeros((image.shape[0], image.shape[1], 7))
+        output_mask[:, :, :3] = image / 255.0
+        output_mask[:, :, 3] = inst_map  # inst_mask
+        output_mask[:, :, 4] = nuclear_map  # binary mask
+        output_mask[:, :, 5] = centroid_map  # centroids_mask
+        output_mask[:, :, 6] = type_map  # type mask
+
+        output_mask = rescale_sample(output_mask)
+        output = util.view_as_windows(output_mask, (256, 256, 6),
+                                      step=(245, 245,
+                                            6)).reshape(-1, 256, 256, 6)
+
+        for p_idx in range(output.shape[0]):
+            if fold_idx == meta['Split'][idx]:
+                np.save("{}/{}_{}.npy".format(test_dir, meta['Filename'][idx],
+                                              p_idx),
+                        output[p_idx])  #.astype('int32')
+            else:
+                np.save("{}/{}_{}.npy".format(train_dir, meta['Filename'][idx],
+                                              p_idx),
+                        output[p_idx])  #.astype('int32')