Merge pull request #25 from oubino/cellpose

Cellpose

src/locpix/img_processing/data_loading/dataset.py

@@ -17,37 +17,33 @@ class ImgDataset(Dataset):
     """
 
     def __init__(
-        self, input_root, label_root, files, input_type, label_type, transform
+        self, input_root, files, input_type, transform
     ):
         """
         Args:
 
-            input_root (string) : Directory containing the histograms
-                of the SMLM data in .npy form
-            label_root (string) : Directory containing the labeled
-                histograms labels for training
+            input_root (string) : Directory containing the input annotated
+                SMLM data
             files (list) : List of the files to include from
                 the directory in this dataset
             input_type (string) : String representing the data format of the
                 input
-            label_type (string) : String representing the data format of the
-                label
             transform (pytorch transform) : Transforms to apply to
                 the dataset
         """
         self.input_data = [
             os.path.join(input_root, file + input_type) for file in files
         ]
-        self.label_data = [
-            os.path.join(label_root, file + label_type) for file in files
-        ]
+        #self.label_data = [
+        #    os.path.join(label_root, file + label_type) for file in files
+        #]
         print("input root", input_root)
-        self.input_data, self.label_data = zip(
-            *sorted(zip(self.input_data, self.label_data))
-        )
-        print("input and label data")
+        #self.input_data, self.label_data = zip(
+        #    *sorted(zip(self.input_data, self.label_data))
+        #)
+        #print("input and label data")
         print(self.input_data)
-        print(self.label_data)
+        #print(self.label_data)
         self.transform = transform
 
     def preprocess(self, folder):
@@ -58,47 +54,71 @@ def preprocess(self, folder):
         """
 
         # join data
-        join_data = zip(self.input_data, self.label_data)
+        #join_data = zip(self.input_data, self.label_data)
+
+        # make folders to save data at if not already present
+        img_folder = os.path.join(folder, 'imgs')
+        label_folder = os.path.join(folder, 'labels')
+        for folder in [img_folder, label_folder]:
+            if os.path.exists(folder):
+                raise ValueError(f"Cannot proceed as {folder} already exists")
+            else:
+                os.makedirs(folder)
+
 
         self.img_data = []
         self.label_data = []
 
         # for file in input
-        for img, label in join_data:
+        for datum in self.input_data:
 
             # load img and label
-            with open(img, "rb") as f:
-                img_bad = pkl.load(f)
+            #with open(img, "rb") as f:
+            #    histo_bad = pkl.load(f)
+
+            # check img and label and check img the same
+            #histos = []
+            #print(type(histo_bad))
+            #for key, value in histo_bad.items():
+            #    histos.append(value)
+            #histo_bad = np.stack(histos)
+            #print("histo")
+            #print(axis_2_chan)
+            #print(histo.shape)
+            #print(histo_bad.shape)
+
+            #np.testing.assert_array_equal(histo, histo_bad)
 
             item = datastruc.item(None, None, None, None)
-            item.load_from_parquet(os.path.join(label))
-            print(item.df)
+            item.load_from_parquet(os.path.join(datum))
+            #print(item.df)
 
             # convert
-            img, axis_2_chan = item.render_histo()
+            histo, axis_2_chan = item.render_histo()
             label = item.render_seg()
 
-            # check img and label and check img the same
-            histos = []
-            print(type(img_bad))
-            for key, value in img_bad.items():
-                histos.append(value)
-            img_bad = np.stack(histos)
-            print("img")
-            print(axis_2_chan)
-            print(img.shape)
-            print(img_bad.shape)
+            # transpose to img space
+            img = np.transpose(histo, (0,2,1))
+            label = label.T
 
-            np.testing.assert_array_equal(img, img_bad)
-            print("label", label.shape)
+            #import matplotlib.pyplot as plt
+            #plt.imshow(np.log2(img[0,:,:]), origin = 'upper', cmap = 'Greys', alpha=1)
+            #plt.imshow(label, origin='upper', cmap='Reds', alpha=.4)
+            #plt.show()
+            #print("label", label.shape)
+            #print('img', img.shape)
 
-            # img path
+            # img & label path
+            img_path = os.path.join(img_folder, item.name + '.npy')
+            label_path = os.path.join(label_folder, item.name + '.npy')
 
-            # label path
+            # add img and label path to lists
+            self.img_data.append(img_path)
+            self.label_data.append(label_path)
 
             # save img and label
-
-            # add img and label path to lists
+            np.save(img_path, img)
+            np.save(label_path, label)
 
     def __getitem__(self, idx):
         """Returns an item from the dataset, according to index idx

src/locpix/preprocessing/datastruc.py

@@ -633,6 +633,7 @@ def render_histo(self):
             axis_2_chan (list) : List where the first value is the
                 channel in the first axis of the histogram, second value
                 is the channel in the second axis of the histogram etc.
+                e.g. [1,3] : 1st channel is in 1st axis, 3rd channel in 2nd axis
         """
 
         histos = []

src/locpix/preprocessing/functions.py

@@ -39,11 +39,9 @@ def csv_to_datastruc(
     # Check dimensions correctly specified
     if dim != 2 and dim != 3:
         raise ValueError("Dimensions must be 2 or 3")
-    if z_col is not None:
-        print("here", z_col)
-    if dim == 2 and z_col is not None and z_col:
+    if dim == 2 and z_col:
         raise ValueError("If dimensions are two no z should be specified")
-    if dim == 3 and z_col is None:
+    if dim == 3 and not z_col:
         raise ValueError("If dimensions are 3 then z_col must be specified")
 
     # Load in data

src/locpix/scripts/img_seg/cellpose_eval.py

@@ -44,7 +44,7 @@ def main():
     else:
         root = tk.Tk()
         root.withdraw()
-        project_folder = filedialog.askdirectory()
+        project_folder = filedialog.askdirectory(title="Project directory")
 
     # configuration
     if args.config is not None:
@@ -70,8 +70,9 @@ def main():
     # if output directory not present create it
     output_directories = [output_membrane_prob, output_cell_df, output_cell_img]
     for directory in output_directories:
-        if not os.path.exists(directory):
-            print("Making folder")
+        if os.path.exists(directory):
+            raise ValueError(f"Cannot proceed as {directory} already exists")
+        else:
             os.makedirs(directory)
 
     for file in files:
@@ -167,6 +168,11 @@ def main():
             four_colour=True,
         )
 
+        # save yaml file
+        yaml_save_loc = os.path.join(project_folder, "cellpose_eval.yaml")
+        with open(yaml_save_loc, "w") as outfile:
+            yaml.dump(config, outfile)
+
 
 if __name__ == "__main__":
     main()

src/locpix/scripts/img_seg/cellpose_eval_config.py

@@ -154,10 +154,10 @@ def set_config(self, config):
         Args:
             config (dictionary) : Configuration dict"""
 
-        config["vis_threshold"] = int(self.vis_threshold.text())
+        config["vis_threshold"] = float(self.vis_threshold.text())
         config["vis_interpolate"] = self.selectedItems()[0].text()
         config["model"] = self.cellpose_model.text()
-        config["diameter"] = int(self.cellpose_diameter.text())
+        config["diameter"] = float(self.cellpose_diameter.text())
         config["channels"] = [
             int(self.first_channel.text()),
             int(self.second_channel.text()),

src/locpix/scripts/img_seg/cellpose_train.py

@@ -49,7 +49,7 @@ def main():
     else:
         root = tk.Tk()
         root.withdraw()
-        project_folder = filedialog.askdirectory()
+        project_folder = filedialog.askdirectory(title="Project directory")
 
     if args.config is not None:
         # load yaml
@@ -63,8 +63,8 @@ def main():
         # config = cellpose_train_config.config_gui(gt_file_path)
 
     # load in config
-    input_root = os.path.join(project_folder, "annotate/histos")
-    label_root = os.path.join(project_folder, "annotate/annotated")
+    input_root = os.path.join(project_folder, "annotate/annotated")
+    #label_root = os.path.join(project_folder, "annotate/annotated")
     batch_size = config["batch_size"]
     epochs = config["epochs"]
     gpu = config["gpu"]
@@ -83,9 +83,10 @@ def main():
         raise ValueError("There should be some files to open")
 
     # make necessary folders if not present
-    preprocessed_folder = os.path.join(project_folder, "cellose_train")
-    if not os.path.exists(preprocessed_folder):
-        print("Making folder")
+    preprocessed_folder = os.path.join(project_folder, "cellpose_train")
+    if os.path.exists(preprocessed_folder):
+        raise ValueError(f"Cannot proceed as {preprocessed_folder} already exists")
+    else:
         os.makedirs(preprocessed_folder)
 
     print("files", files)
@@ -121,12 +122,14 @@ def main():
 
     # Initialise train and val dataset
     train_set = dataset.ImgDataset(
-        input_root, label_root, train_files, ".pkl", ".parquet", train_transform
+        input_root, train_files, ".parquet", train_transform
     )
     val_set = dataset.ImgDataset(
-        input_root, label_root, val_files, ".pkl", ".parquet", val_transform
+        input_root, val_files, ".parquet", val_transform
     )
 
+    print('Preprocessing datasets')
+
     # Pre-process train and val dataset
     train_set.preprocess(os.path.join(preprocessed_folder, "train"))
     val_set.preprocess(os.path.join(preprocessed_folder, "val"))
@@ -292,6 +295,11 @@ def main():
             save_loc=save_loc,
             four_colour=True,
         )
+
+        # save yaml file
+        yaml_save_loc = os.path.join(project_folder, "cellpose_train.yaml")
+        with open(yaml_save_loc, "w") as outfile:
+            yaml.dump(config, outfile)
 """
 
 if __name__ == "__main__":

src/locpix/scripts/img_seg/classic.py

@@ -44,7 +44,7 @@ def main():
     else:
         root = tk.Tk()
         root.withdraw()
-        project_folder = filedialog.askdirectory()
+        project_folder = filedialog.askdirectory(title="Project directory")
 
     # configuration folder
     if args.config is not None:
@@ -64,20 +64,23 @@ def main():
 
     # if output directory not present create it
     output_membrane_prob = os.path.join(project_folder, "classic/membrane/prob_map")
-    if not os.path.exists(output_membrane_prob):
-        print("Making folder")
+    if os.path.exists(output_membrane_prob):
+            raise ValueError(f"Cannot proceed as {output_membrane_prob} already exists")
+    else:
         os.makedirs(output_membrane_prob)
 
     # if output directory not present create it
     output_cell_df = os.path.join(project_folder, "classic/cell/seg_dataframes")
-    if not os.path.exists(output_cell_df):
-        print("Making folder")
+    if os.path.exists(output_cell_df):
+            raise ValueError(f"Cannot proceed as {output_cell_df} already exists")
+    else:
         os.makedirs(output_cell_df)
 
     # if output directory not present create it
     output_cell_img = os.path.join(project_folder, "classic/cell/seg_img")
-    if not os.path.exists(output_cell_img):
-        print("Making folder")
+    if os.path.exists(output_cell_img):
+            raise ValueError(f"Cannot proceed as {output_cell_img} already exists")
+    else:
         os.makedirs(output_cell_img)
 
     for file in files:
@@ -163,6 +166,11 @@ def main():
             four_colour=True,
         )
 
+        # save yaml file
+        yaml_save_loc = os.path.join(project_folder, "classic.yaml")
+        with open(yaml_save_loc, "w") as outfile:
+            yaml.dump(config, outfile)
+
 
 if __name__ == "__main__":
     main()

src/locpix/scripts/img_seg/classic_config.py

@@ -114,7 +114,7 @@ def set_config(self, config):
         Args:
             config (dictionary) : Configuration dict"""
 
-        config["vis_threshold"] = int(self.vis_threshold.text())
+        config["vis_threshold"] = float(self.vis_threshold.text())
         config["vis_interpolate"] = self.vis_interpolation.selectedItems()[0].text()
         config["sum_chan"] = self.sum_chan.isChecked()
 

src/locpix/scripts/img_seg/get_markers.py

@@ -43,7 +43,7 @@ def main():
     else:
         root = tk.Tk()
         root.withdraw()
-        project_folder = filedialog.askdirectory()
+        project_folder = filedialog.askdirectory(title="Project directory")
 
     if args.config is not None:
         # load yaml
@@ -95,6 +95,11 @@ def main():
         # save
         np.save(markers_loc, markers)
 
+    # save yaml file
+    yaml_save_loc = os.path.join(project_folder, "get_markers.yaml")
+    with open(yaml_save_loc, "w") as outfile:
+        yaml.dump(config, outfile)
+
 
 if __name__ == "__main__":
     main()

src/locpix/scripts/img_seg/get_markers_config.py

@@ -104,7 +104,7 @@ def set_config(self, config):
         Args:
             config (dictionary) : Configuration dict"""
 
-        config["vis_threshold"] = self.vis_threshold
+        config["vis_threshold"] = float(self.vis_threshold)
         config["vis_interpolate"] = self.vis_interpolation.selectedItems()[0].text()
 
         # check config is correct