updated normalize/preprocess to support writing multiple image formats

topaz/commands/normalize.py

@@ -8,6 +8,7 @@
 
 from topaz.transform import ScaledGaussianMixture
 from topaz.utils.data.loader import load_image
+from topaz.utils.image import save_image
 
 name = 'normalize'
 help = 'normalize a set of images using a per image scaled 2-component Gaussian mixture model'
@@ -17,8 +18,13 @@ def add_arguments(parser):
     parser.add_argument('-s', '--sample', default=25, type=int, help='pixel sampling factor for model fit (default: 25)')
     parser.add_argument('--niters', default=200, type=int, help='number of iterations to run for model fit (default: 200)')
     parser.add_argument('--seed', default=1, type=int, help='random seed for model initialization (default: 1)')
+
     parser.add_argument('-o', '--destdir', help='output directory')
+
+    parser.add_argument('--format', dest='format_', default='mrc', help='image format(s) to write. choices are mrc, tiff, and png. images can be written in multiple formats by specifying each in a comma separated list, e.g. mrc,png would write mrc and png format images (default: mrc)')
+
     parser.add_argument('-v', '--verbose', action='store_true', help='verbose output')
+
     return parser
 
 
@@ -80,12 +86,13 @@ def main(args):
     if not os.path.exists(destdir):
         os.makedirs(destdir)
 
-    for name,im in zip(names, images):
-        im = Image.fromarray(im) 
-        path = os.path.join(destdir, name) + '.tiff'
-        if args.verbose:
-            print('# saving:', path)
-        im.save(path, 'tiff')
+    ## what image formats are we writing
+    verbose = args.verbose
+    formats = args.format_.split(',')
+    for name,x in zip(names, images):
+        base = os.path.join(destdir, name)
+        for f in formats:
+            save_image(x, base, f=f, verbose=verbose)
 
     ## save the metadata in json format
     path = os.path.join(destdir, 'metadata.json')

topaz/commands/preprocess.py

@@ -7,7 +7,7 @@
 from PIL import Image # for saving images
 
 from topaz.transform import ScaledGaussianMixture
-from topaz.utils.image import downsample
+from topaz.utils.image import downsample, save_image
 from topaz.utils.data.loader import load_image
 
 
@@ -21,6 +21,8 @@ def add_arguments(parser):
     parser.add_argument('-s', '--scale', default=4, type=int, help='rescaling factor for image downsampling (default: 4)')
     parser.add_argument('-t', '--num-workers', default=0, type=int, help='number of processes to use for parallel image downsampling (default: 0)')
 
+    parser.add_argument('--format', dest='format_', default='mrc', help='image format(s) to write. choices are mrc, tiff, and png. images can be written in multiple formats by specifying each in a comma separated list, e.g. mrc,png would write mrc and png format images (default: mrc)')
+
     parser.add_argument('--pixel-sampling', default=25, type=int, help='pixel sampling factor for model fit (default: 25)')
     parser.add_argument('--niters', default=200, type=int, help='number of iterations to run for model fit (default: 200)')
     parser.add_argument('--seed', default=1, type=int, help='random seed for model initialization (default: 1)')
@@ -130,12 +132,12 @@ def main(args):
     if not os.path.exists(destdir):
         os.makedirs(destdir)
 
-    for name,im in zip(names, images):
-        im = Image.fromarray(im) 
-        path = os.path.join(destdir, name) + '.tiff'
-        if verbose:
-            print('# saving:', path)
-        im.save(path, 'tiff')
+    ## what image formats are we writing
+    formats = args.format_.split(',')
+    for name,x in zip(names, images):
+        base = os.path.join(destdir, name)
+        for f in formats:
+            save_image(x, base, f=f, verbose=verbose)
 
     ## save the metadata in json format
     path = os.path.join(destdir, 'metadata.json')

topaz/utils/image.py

@@ -1,6 +1,10 @@
 from __future__ import division, print_function
 
 import numpy as np
+from PIL import Image
+import os
+
+import topaz.mrc as mrc
 
 def downsample(x, factor):
     """ Downsample 2d array using fourier transform """
@@ -30,5 +34,36 @@ def unquantize(x, mi=-4, ma=4, dtype=np.float32):
     y = x*(ma-mi)/255 + mi
     return y
 
-
+def save_image(x, path, f=None, verbose=False):
+    if f is None:
+        f = os.path.splitext(path)[1]
+        f = f[1:] # remove the period
+    else:
+        path = path + '.' + f
+
+    if verbose:
+        print('# saving:', path)
+
+    if f == 'mrc':
+        save_mrc(x, path)
+    elif f == 'tiff':
+        save_tiff(x, path)
+    elif f == 'png':
+        save_png(x, path)
+
+def save_mrc(x, path):
+    with open(path, 'wb') as f:
+        x = x[np.newaxis] # need to add z-axis for mrc write
+        mrc.write(f, x)
+
+def save_tiff(x, path):
+    im = Image.fromarray(x) 
+    im.save(path, 'tiff')
+
+def save_png(x, path):
+    # byte encode the image
+    im = Image.fromarray(quantize(x))
+    im.save(path, 'png')
+
+