add look_one_level_down (#220)

MouseLand · May 1, 2021 · ff46c5f · ff46c5f
1 parent fa6b841
commit ff46c5f
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 11 deletions.
diff --git a/cellpose/__main__.py b/cellpose/__main__.py
@@ -30,6 +30,8 @@ def main():
     parser.add_argument('--train', action='store_true', help='train network using images in dir')
     parser.add_argument('--dir', required=False, 
                         default=[], type=str, help='folder containing data to run or train on')
+    parser.add_argument('--look_one_level_down', action='store_true', 
+                        help='')
     parser.add_argument('--mxnet', action='store_true', help='use mxnet')
     parser.add_argument('--img_filter', required=False, 
                         default=[], type=str, help='end string for images to run on')
@@ -131,7 +133,10 @@ def main():
                     logger.warning('model path does not exist, using cyto model')
                     args.pretrained_model = 'cyto'
 
-            image_names = io.get_image_files(args.dir, args.mask_filter, imf=imf)
+            image_names = io.get_image_files(args.dir, 
+                                             args.mask_filter, 
+                                             imf=imf,
+                                             look_one_level_down=args.look_one_level_down)
             nimg = len(image_names)
             if args.diameter==0:
                 if args.pretrained_model=='cyto' or args.pretrained_model=='nuclei':

diff --git a/cellpose/io.py b/cellpose/io.py
@@ -80,16 +80,24 @@ def imsave(filename, arr):
     else:
         cv2.imwrite(filename, arr)
 
-def get_image_files(folder, mask_filter, imf=None):
+def get_image_files(folder, mask_filter, imf=None, look_one_level_down=False):
+    """ find all images in a folder and if look_one_level_down all subfolders """
     mask_filters = ['_cp_masks', '_cp_output', '_flows', mask_filter]
     image_names = []
     if imf is None:
         imf = ''
-    image_names.extend(glob.glob(folder + '/*%s.png'%imf))
-    image_names.extend(glob.glob(folder + '/*%s.jpg'%imf))
-    image_names.extend(glob.glob(folder + '/*%s.jpeg'%imf))
-    image_names.extend(glob.glob(folder + '/*%s.tif'%imf))
-    image_names.extend(glob.glob(folder + '/*%s.tiff'%imf))
+
+    folders = []
+    if look_one_level_down:
+        folders = natsorted(glob.glob(os.path.join(folder, "*/")))
+    folders.append(folder)
+
+    for folder in folders:
+        image_names.extend(glob.glob(folder + '/*%s.png'%imf))
+        image_names.extend(glob.glob(folder + '/*%s.jpg'%imf))
+        image_names.extend(glob.glob(folder + '/*%s.jpeg'%imf))
+        image_names.extend(glob.glob(folder + '/*%s.tif'%imf))
+        image_names.extend(glob.glob(folder + '/*%s.tiff'%imf))
     image_names = natsorted(image_names)
     imn = []
     for im in image_names:

diff --git a/cellpose/models.py b/cellpose/models.py
@@ -453,8 +453,11 @@ def eval(self, x, batch_size=8, channels=None, channel_axis=None,
         """
         if isinstance(x, list) or x.squeeze().ndim==5:
             masks, styles, flows = [], [], []
-            for i, xi in enumerate(x):
-                maski, stylei, flowi = self.eval(xi, 
+            tqdm_out = utils.TqdmToLogger(models_logger, level=logging.INFO)
+            nimg = len(x)
+            iterator = trange(nimg, file=tqdm_out) if nimg>1 else range(nimg)
+            for i in iterator:
+                maski, stylei, flowi = self.eval(x[i], 
                                                  batch_size=batch_size, 
                                                  channels=channels[i] if (len(channels)==len(x) and 
                                                                           (isinstance(channels[i], list) and isinstance(channels[i], np.ndarray)) and 
@@ -794,8 +797,10 @@ def eval(self, x, channels=None, channel_axis=None,
 
         if isinstance(x, list):
             diams, diams_style = [], []
-            for i, xi in enumerate(x):
-                diam, diam_style = self.eval(xi, 
+            nimg = len(x)
+            iterator = trange(nimg, file=tqdm_out) if nimg>1 else range(nimg)
+            for i in iterator:
+                diam, diam_style = self.eval(x[i], 
                                             channels=channels[i] if (len(channels)==len(x) and 
                                                                      (isinstance(channels[i], list) and isinstance(channels[i], np.ndarray)) and 
                                                                      len(channels[i])==2) else channels, 
@@ -811,6 +816,10 @@ def eval(self, x, channels=None, channel_axis=None,
 
             return diams, diams_style
 
+        if x.squeeze().ndim > 3:
+            models_logger.warning('image is not 2D cannot compute diameter')
+            return self.diam_mean, self.diam_mean
+
         models_logger.info('computing styles from images')
         styles = self.cp.eval(x, 
                               channels=channels,