Update folder structure

deepprofiler/__main__.py

@@ -14,6 +14,7 @@
 import deepprofiler.dataset.image_dataset
 import deepprofiler.dataset.sampling
 import deepprofiler.learning.training
+import deepprofiler.learning.tf2train
 import deepprofiler.learning.profiling
 import deepprofiler.download.normalize_bbbc021_metadata
 
@@ -142,12 +143,16 @@ def prepare(context):
 
 # Second tool: Sample single cells for training
 @cli.command()
+@click.option("--mode", default="sample")
 @click.pass_context
-def sample_sc(context):
+def sample_sc(context, mode):
     if context.parent.obj["config"]["prepare"]["compression"]["implement"]:
         context.parent.obj["config"]["paths"]["images"] = context.obj["config"]["paths"]["compressed_images"]
-    dset = deepprofiler.dataset.image_dataset.read_dataset(context.obj["config"], mode='train')
-    deepprofiler.dataset.sampling.sample_dataset(context.obj["config"], dset)
+    dset = deepprofiler.dataset.image_dataset.read_dataset(context.obj["config"])
+    if mode == "sample":
+        deepprofiler.dataset.sampling.sample_dataset(context.obj["config"], dset)
+    elif mode == "export_all":
+        deepprofiler.dataset.sampling.export_dataset(context.obj["config"], dset)
     print("Single-cell sampling complete.")
 
 
@@ -163,11 +168,19 @@ def train(context, epoch, seed):
     deepprofiler.learning.training.learn_model(context.obj["config"], dset, epoch, seed)
 
 
+# Third tool (b): Train a network with TF dataset
+@cli.command()
+@click.option("--epoch", default=1)
+@click.pass_context
+def traintf2(context, epoch):
+    deepprofiler.learning.tf2train.learn_model(context.obj["config"], epoch)
+
+
 # Fourth tool: Profile cells and extract features
 @cli.command()
 @click.pass_context
 @click.option("--part",
-              help="Part of index to process", 
+               help="Part of index to process",
               default=-1, 
               type=click.INT)
 def profile(context, part):

deepprofiler/dataset/image_dataset.py

@@ -211,7 +211,7 @@ def read_dataset(config, mode = 'train'):
     print(metadata.data.info())
 
     # Split training data
-    if mode == 'train':
+    if mode == 'train' and config["train"]["model"]["crop_generator"] == 'crop_generator':
         split_field = config["train"]["partition"]["split_field"]
         trainingFilter = lambda df: df[split_field].isin(config["train"]["partition"]["training_values"])
         validationFilter = lambda df: df[split_field].isin(config["train"]["partition"]["validation_values"])
@@ -238,7 +238,8 @@ def read_dataset(config, mode = 'train'):
     if config["dataset"]["locations"]["mask_objects"]:
         dset.outlines = outlines
 
-    if mode == 'train':
+    # For training with sampled_crop_generator, no need to read locations again necessary.
+    if mode == 'train' and config["train"]["model"]["crop_generator"] == 'crop_generator':
         dset.prepare_training_locations()
 
     return dset

deepprofiler/dataset/sampling.py

@@ -3,6 +3,7 @@
 import threading
 import tqdm
 import os
+import shutil
 
 import tensorflow as tf
 
@@ -13,6 +14,7 @@
 class SingleCellSampler(deepprofiler.imaging.cropping.CropGenerator):
 
     def start(self, session):
+        self.all_metadata = []
         self.session = session
         # Define input data batches
         with tf.compat.v1.variable_scope("train_inputs"):
@@ -21,23 +23,23 @@ def start(self, session):
 
     def process_batch(self, batch):
         for i in range(len(batch["keys"])):
-            batch["locations"][i]["Key"] = batch["keys"][i]
+            batch["locations"][i]["Key"] = batch["keys"][i].replace('-', '/')
             batch["locations"][i]["Target"] = batch["targets"][i][0]
             batch["locations"][i]["Class_Name"] = self.dset.targets[0].values[batch["targets"][i][0]]
         metadata = pd.concat(batch["locations"])
-        cols = ["Key","Target","Nuclei_Location_Center_X","Nuclei_Location_Center_Y"]
-        seps = ["+","@","x",".png"]
+        cols = ["Key", "Target", "Nuclei_Location_Center_X", "Nuclei_Location_Center_Y"]
+        seps = ["/", "@", "x", ".png"]
         metadata["Image_Name"] = ""
         for c in range(len(cols)):
-            metadata["Image_Name"] += metadata[cols[c]].astype(str).str.replace("/","-") + seps[c]
-        
+            metadata["Image_Name"] += metadata[cols[c]].astype(str) + seps[c]
+
         boxes, box_ind, targets, masks = deepprofiler.imaging.boxes.prepare_boxes(batch, self.config)
 
         feed_dict = {
-            self.input_variables["image_ph"]:batch["images"],
-            self.input_variables["boxes_ph"]:boxes,
-            self.input_variables["box_ind_ph"]:box_ind,
-            self.input_variables["mask_ind_ph"]:masks
+            self.input_variables["image_ph"]: batch["images"],
+            self.input_variables["boxes_ph"]: boxes,
+            self.input_variables["box_ind_ph"]: box_ind,
+            self.input_variables["mask_ind_ph"]: masks
         }
         for i in range(len(targets)):
             tname = "target_" + str(i)
@@ -46,6 +48,20 @@ def process_batch(self, batch):
         output = self.session.run(self.input_variables["labeled_crops"], feed_dict)
         return output[0], metadata.reset_index(drop=True)
 
+    def export_single_cells(self, key, image_array, meta):
+        outdir = self.config["paths"]["single_cell_sample"]
+        key = self.dset.keyGen(meta)
+        batch = {"keys": [key], "images": [image_array], "targets": [], "locations": []}
+        batch["locations"].append(deepprofiler.imaging.boxes.get_locations(key, self.config))
+        batch["targets"].append([t.get_values(meta) for t in self.dset.targets])
+        crops, metadata = self.process_batch(batch)
+        for j in range(crops.shape[0]):
+            image = deepprofiler.imaging.cropping.unfold_channels(crops[j,:,:,:])
+            skimage.io.imsave(os.path.join(outdir, metadata.loc[j, "Image_Name"]), image)
+
+        self.all_metadata.append(metadata)
+        print("{}: {} single cells".format(key, crops.shape[0]))
+
 
 def start_session():
     configuration = tf.compat.v1.ConfigProto()
@@ -69,8 +85,7 @@ def is_directory_empty(outdir):
             return False
         elif erase == "y":
             print("Removing previous sampled files")
-            for f in tqdm.tqdm(files):
-                os.remove(os.path.join(outdir, f))
+            shutil.rmtree(outdir)
     return True
 
 
@@ -99,7 +114,9 @@ def sample_dataset(config, dset):
         if len(batch["keys"]) > 0:
             crops, metadata = cropper.process_batch(batch)
             for j in range(crops.shape[0]):
-                image = deepprofiler.imaging.cropping.unfold_channels(crops[j,:,:,:])
+                plate, well, site, name = metadata.loc[j, "Image_Name"].split('/')
+                os.makedirs(os.path.join(outdir, plate, well, site), exist_ok=True)
+                image = deepprofiler.imaging.cropping.unfold_channels(crops[j, :, :, :])
                 skimage.io.imsave(os.path.join(outdir, metadata.loc[j, "Image_Name"]), image)
             all_metadata.append(metadata)
 
@@ -112,3 +129,17 @@ def sample_dataset(config, dset):
     all_metadata = pd.concat(all_metadata).reset_index(drop=True)
     all_metadata.to_csv(os.path.join(outdir, "sc-metadata.csv"), index=False)
 
+def export_dataset(config, dset):
+    outdir = config["paths"]["single_cell_sample"]
+    if not is_directory_empty(outdir):
+        return
+
+    session = start_session()
+    cropper = SingleCellSampler(config, dset)
+    cropper.start(session)
+    dset.scan(cropper.export_single_cells, frame="all")
+    df = pd.concat(cropper.all_metadata).reset_index(drop=True)
+    df.to_csv(os.path.join(outdir, "sc-metadata.csv"), index=False)
+    print("Exporting: done")    
+
+

deepprofiler/imaging/augmentations.py

@@ -1,6 +1,7 @@
 import numpy as np
 import tensorflow as tf
-import tensorflow_addons
+import tensorflow_addons as tfa
+import sys
 
 tf.compat.v1.disable_v2_behavior()
 
@@ -9,51 +10,87 @@
 # CROPPING AND TRANSFORMATION OPERATIONS
 #################################################
 
+def random_crop(image):
+    w,h,c = image.shape
+
+    size = tf.random.uniform([1], minval=int(w.value*0.6), maxval=w, dtype=tf.int32)
+    crop = tf.image.random_crop(image, [size[0],size[0],c])
+
+    result = tf.image.resize(
+        tf.expand_dims(crop, 0), [w,h], method="bicubic"
+    )
+
+    return result[0,...]
+
+
+def random_illumination(image):
+    # Make channels independent images
+    numchn = image.shape[-1].value
+    source = tf.transpose(image, [2, 1, 0])
+    source = tf.expand_dims(source, -1)
+    source = tf.image.grayscale_to_rgb(source)
+
+    # Apply illumination augmentations
+    bright = tf.random.uniform([numchn], minval=-0.1, maxval=0.1, dtype=tf.float32)
+    channels = [tf.image.adjust_brightness(source[s,...], bright[s]) for s in range(numchn)]
+    contrast = tf.random.uniform([numchn], minval=0.8, maxval=1.2, dtype=tf.float32)
+    channels = [tf.image.adjust_contrast(channels[s], contrast[s]) for s in range(numchn)]
+    result = tf.concat([tf.expand_dims(t, 0) for t in channels], axis=0)
+
+    # Recover multi-channel image
+    result = tf.image.rgb_to_grayscale(result)
+    result = tf.transpose(result[:,:,:,0], [2, 1, 0])
+    #result = result / tf.math.reduce_max(result)
+    return result
+
+
+def random_flips(image):
+    # Horizontal flips
+    augmented = tf.image.random_flip_left_right(image)
+
+    # 90 degree rotations
+    angle = tf.random.uniform([1], minval=0, maxval=4, dtype=tf.int32)
+    augmented = tf.image.rot90(augmented, angle[0])
+
+    return augmented
 
-def augment(crop):
-    with tf.compat.v1.variable_scope("augmentation"):
-        # Horizontal flips
-        augmented = tf.image.random_flip_left_right(crop)
-
-        # 90 degree rotations
-        angle = tf.compat.v1.random_uniform([1], minval=0, maxval=4, dtype=tf.int32)
-        augmented = tf.image.rot90(augmented, angle[0])
+def random_rotate(image):
+    w, h, c = image.shape
+    image = tfa.image.rotate(image, np.pi / tf.random.uniform(shape=[], minval=1, maxval=10, dtype=tf.float32))
+    image = tf.image.central_crop(image, 0.7)
+    return tf.image.resize(image, (w, h))
 
-        # 5 degree inclinations
-        angle = tf.compat.v1.random_normal([1], mean=0.0, stddev=0.03 * np.pi, dtype=tf.float32)
-        augmented = tensorflow_addons.image.rotate(augmented, angle[0], interpolation="BILINEAR")
 
-        # Translations (3% movement in x and y)
-        offsets = tf.compat.v1.random_normal([2],
-                                             mean=0,
-                                             stddev=int(crop.shape[0].value * 0.03)
-                                             )
-        augmented = tensorflow_addons.image.translate(augmented, translations=offsets)
+def augment(image):
+    #if tf.less(tf.random.uniform([], minval=0, maxval=1, dtype=tf.float32), tf.cast(0.5, tf.float32)):
+    #    augm = random_crop(image)
+    #else:
+    #    augm = random_rotate(image)
 
-        # Illumination changes (10% changes in intensity)
-        illum_s = tf.compat.v1.random_normal([1], mean=1.0, stddev=0.1, dtype=tf.float32)
-        illum_t = tf.compat.v1.random_normal([1], mean=0.0, stddev=0.1, dtype=tf.float32)
-        augmented = augmented * illum_s + illum_t
+    augm = random_flips(image)
+    augm = random_illumination(augm)
 
-    return augmented
+    return augm
 
 
 def augment_multiple(crops, parallel=None):
-    with tf.compat.v1.variable_scope("augmentation"):
-        return tf.map_fn(augment, crops, parallel_iterations=parallel, dtype=tf.float32)
+    return tf.map_fn(augment, crops, parallel_iterations=parallel, dtype=tf.float32)
 
 
 ## A layer for GPU accelerated augmentations
 
+#AugmentationLayer = tf.keras.layers.Lambda(augment_multiple)
+
 class AugmentationLayer(tf.compat.v1.keras.layers.Layer):
     def __init__(self, **kwargs):
+        self.is_training = True
         super(AugmentationLayer, self).__init__(**kwargs)
 
     def build(self, input_shape):
         return
 
-    def call(self, input_tensor, training=False):
-        if training:
+    def call(self, input_tensor):
+        if self.is_training:
             return augment_multiple(input_tensor)
         else:
             return input_tensor

deepprofiler/imaging/cropping.py

@@ -11,6 +11,7 @@
 import deepprofiler.imaging.boxes
 
 tf.compat.v1.disable_v2_behavior()
+tf.config.run_functions_eagerly(False)
 
 
 def crop_graph(image_ph, boxes_ph, box_ind_ph, mask_ind_ph, box_size, mask_boxes=False):
@@ -22,9 +23,12 @@ def crop_graph(image_ph, boxes_ph, box_ind_ph, mask_ind_ph, box_size, mask_boxes
             mask_values = tf.ones_like(crops[:, :, :, -1], dtype=tf.float32) * tf.cast(mask_ind, dtype=tf.float32)
             masks = tf.compat.v1.to_float(tf.equal(crops[:, :, :, -1], mask_values))
             crops = crops[:, :, :, 0:-1] * tf.expand_dims(masks, -1)
-        mean = tf.math.reduce_mean(crops, axis=[1, 2], keepdims=True)
-        std = tf.math.reduce_std(crops, axis=[1, 2], keepdims=True)
-        crops = (crops - mean)/std
+        #mean = tf.math.reduce_mean(crops, axis=[1, 2], keepdims=True)
+        #std = tf.math.reduce_std(crops, axis=[1, 2], keepdims=True)
+        #crops = (crops - mean)/std
+        mini = tf.math.reduce_min(crops, axis=[1, 2], keepdims=True)
+        maxi = tf.math.reduce_max(crops, axis=[1, 2], keepdims=True)
+        crops = (crops - mini) / maxi
     return crops
 
 
@@ -43,11 +47,11 @@ def fold_channels(crop):
     # Expected input image shape: (h, w * c), with h = w
     # Output image shape: (h, w, c), with h = w
     output = np.reshape(crop, (crop.shape[0], crop.shape[0], -1), order="F").astype(np.float)
-    for i in range(output.shape[-1]):
-        mean = np.mean(output[:, :, i])
-        std = np.std(output[:, :, i])
-        output[:, :, i] = (output[:, :, i] - mean) / std
-    return output
+    #for i in range(output.shape[-1]):
+    #    mean = np.mean(output[:, :, i])
+    #    std = np.std(output[:, :, i])
+    #    output[:, :, i] = (output[:, :, i] - mean) / std
+    return output / 255.
 
 
 # TODO: implement abstract crop generator