Fix cihx parsing bug, add tests for save_mraw.

pyMRAW.py

@@ -29,7 +29,7 @@
 import warnings
 import xmltodict
 
-__version__ = '0.23'
+__version__ = '0.3'
 
 SUPPORTED_FILE_FORMATS = ['mraw', 'tiff']
 SUPPORTED_EFFECTIVE_BIT_SIDE = ['lower', 'higher']
@@ -59,8 +59,8 @@ def get_cih(filename):
     elif ext == '.cihx':
         with open(filename, 'r', encoding='utf-8', errors='ignore') as f:
             lines = f.readlines()
-            l0 = lines.pop(0)
-            xml = '\n'.join(lines)
+            first_last_line = [ i for i in range(len(lines)) if '<cih>' in lines[i] or '</cih>' in lines[i] ]
+            xml = ''.join(lines[first_last_line[0]:first_last_line[-1]+1])
 
         raw_cih_dict = xmltodict.parse(xml)
         cih = {
@@ -76,7 +76,7 @@ def get_cih(filename):
             'EffectiveBit Depth': int(raw_cih_dict['cih']['imageDataInfo']['effectiveBit']['depth']),
             'EffectiveBit Side': raw_cih_dict['cih']['imageDataInfo']['effectiveBit']['side'],
             'Color Bit': int(raw_cih_dict['cih']['imageDataInfo']['colorInfo']['bit']),
-            'Comment Text': raw_cih_dict['cih']['basicInfo']['comment'],
+            'Comment Text': raw_cih_dict['cih']['basicInfo'].get('comment', ''),
         }
 
     else:
@@ -89,8 +89,6 @@ def get_cih(filename):
     bits = cih['Color Bit']
     if bits < 12:
         warnings.warn('Not 12bit ({:g} bits)! clipped values?'.format(bits))
-    elif (bits > 12):
-        warnings.warn('Not 12bit ({:g} bits)! Values may/will be divided by /16->12bit (during operation)'.format(bits))
                 # - may cause overflow')
                 # 12-bit values are spaced over the 16bit resolution - in case of photron filming at 12bit
                 # this can be meanded by dividing images with //16
@@ -102,7 +100,7 @@ def get_cih(filename):
     if (cih['File Format'].lower() == 'mraw') & (cih['Color Bit'] not in [8, 16]):
         raise Exception('pyMRAW only works for 8-bit and 16-bit files!')
     if cih['Original Total Frame'] > cih['Total Frame']:
-        warnings.warn('Clipped footage!')
+        warnings.warn('Clipped footage! (Total frame: {}, Original total frame: {})'.format(cih['Total Frame'], cih['Original Total Frame'] ))
 
     return cih
 
@@ -158,64 +156,81 @@ def load_video(cih_file):
     return images, cih
 
 
-def save_mraw(images, save_path, ext='mraw', info_dict={}):
+def save_mraw(images, save_path, bit_depth=16, ext='mraw', info_dict={}):
     """
-    Saves given sequence of images into a 16-bit .mraw file.
+    Saves given sequence of images into .mraw file.
 
     Inputs:
     sequence : array_like of shape (n, h, w), sequence of `n` grayscale images
         of shape (h, w) to save.
-    save_path : str, path to saved file. 
-    comment: str, comment text to be saved into the .cih file.
+    save_path : str, path to saved cih file. 
+    bit_depth: int, bit depth of the image data. Currently supported bit depths are 8 and 16.
     ext : str, generated file extension ('mraw' or 'npy'). If set to 'mraw', it can be viewed in
         PFV. Defaults to '.mraw'.
     info_dict : dict, mraw video information to go into the .cih file. The info keys have to match
         .cih properties descriptions exactly (example common keys: 'Record Rate(fps)', 
         'Shutter Speed(s)', 'Comment Text' etc.).
 
     Outputs:
-    out_file : str, path to output or .mraw (or .npy) file.
-    cih_file : str, path to generated .cih file
+    mraw_path : str, path to output or .mraw (or .npy) file.
+    cih_path : str, path to generated .cih file
     """
 
     filename, extension = path.splitext(save_path)
-    if not extension:
-        save_path = '{:s}.{:s}'.format(filename, ext)
+    mraw_path = '{:s}.{:s}'.format(filename, ext)
+    cih_path = '{:s}.{:s}'.format(filename, '.cih')
 
     directory_path = path.split(save_path)[0]
     if not path.exists(directory_path):
         os.makedirs(directory_path)
 
+    bit_depth_dtype_map = {
+        8: np.uint8,
+        16: np.uint16
+    }
+    if bit_depth not in bit_depth_dtype_map.keys():
+        raise ValueError('Currently supported bit depths are 8 and 16.')
+
+    if bit_depth < 16:
+        effective_bit = bit_depth
+    else:
+        effective_bit = 12
+    if np.max(images) > 2**bit_depth-1:
+        raise ValueError(
+            'The input image data does not match the selected bit depth. ' +
+            'Consider normalizing the image data before saving.')
+
     # Generate .mraw file
-    with open(save_path, 'wb') as file:
+    with open(mraw_path, 'wb') as file:
         for image in images:
-            image = image.astype(np.uint16)
+            image = image.astype(bit_depth_dtype_map[bit_depth])
             image.tofile(file)
     file_shape = (int(len(images)), image.shape[0], image.shape[1])
     file_format = 'MRaw'
 
     image_info = {'Record Rate(fps)': '{:d}'.format(1),
                 'Shutter Speed(s)': '{:.6f}'.format(1),
                 'Total Frame': '{:d}'.format(file_shape[0]),
+                'Original Total Frame': '{:d}'.format(file_shape[0]),
                 'Start Frame': '{:d}'.format(0),
                 'Image Width': '{:d}'.format(file_shape[2]),
                 'Image Height': '{:d}'.format(file_shape[1]),
                 'Color Type': 'Mono', 
-                'Color Bit': '16',
+                'Color Bit': bit_depth,
                 'File Format' : file_format,
-                'EffectiveBit Depth': '12',
+                'EffectiveBit Depth': effective_bit,
                 'Comment Text': 'Generated sequence. Modify measurement info in created .cih file if necessary.',
                 'EffectiveBit Side': 'Lower'}
 
     image_info.update(info_dict)
 
-    cih_file = save_path = '{:s}.{:s}'.format(filename, 'cih')
-    with open(cih_file, 'w') as file:
+    cih_path = '{:s}.{:s}'.format(filename, 'cih')
+    with open(cih_path, 'w') as file:
         file.write('#Camera Information Header\n')
         for key in image_info.keys():
             file.write('{:s} : {:s}\n'.format(key, str(image_info[key])))
 
-    return save_path, cih_file
+    return mraw_path, cih_path
 
 
 def show_UI():

setup.py

@@ -33,9 +33,9 @@
 from pyMRAW import __version__
 setup(name='pyMRAW',
       version=__version__,
-      author='Jaka Javh, Janko Slavič',
-      author_email='jaka.javh@fs.uni-lj.si,janko.slavic@fs.uni-lj.si',
-      description='Module for reading Photron MRAW image sequences.',
+      author='Jaka Javh, Janko Slavič, Domen Gorjup',
+      author_email='jaka.javh@fs.uni-lj.si,janko.slavic@fs.uni-lj.si, domen.gorjup@fs.uni-lj.si',
+      description='Module for reading and writing Photron MRAW image sequences.',
       url='https://github.com/ladisk/pyMRAW',
       py_modules=['pyMRAW'],
       #ext_modules=[Extension('lvm_read', ['data/short.lvm'])],

tests/test_all.py

@@ -1,13 +1,17 @@
 """
 Unit test for pyMRAW.py
 """
+import pytest
 import numpy as np
 import sys, os
+import tempfile
+
 myPath = os.path.dirname(os.path.abspath(__file__))
 sys.path.insert(0, myPath + '/../')
 
 import pyMRAW
 
+@pytest.mark.filterwarnings('ignore')
 def test():
     filename = './data/sample_60k_16bit.cih'
     cih = pyMRAW.get_cih(filename)
@@ -28,6 +32,7 @@ def test():
     np.testing.assert_allclose(image_data[0,0,0],1889, atol=1e-8)
 
 
+@pytest.mark.filterwarnings('ignore')
 def test_cihx():
     filename = './data/beam.cihx'
     cih = pyMRAW.get_cih(filename)
@@ -38,5 +43,36 @@ def test_cihx():
     np.testing.assert_equal(mraw.shape, (4, 80, 1024))
 
 
-if __name__ == '__mains__':
+@pytest.mark.filterwarnings('ignore')
+def test_save_mraw():
+    root_dir = './tests'
+    with tempfile.TemporaryDirectory(dir=root_dir) as tmpdir:
+
+            images8 = np.ones((3, 4, 5), dtype=np.uint8)
+            images16 = np.ones((2, 3, 4), dtype=np.uint16) * (2**12-1)
+
+            mraw8, cih8 = pyMRAW.save_mraw(images8, os.path.join(tmpdir, 'test8.cih'), bit_depth=8, ext='mraw', info_dict={'Record Rate(fps)':10})
+            mraw8_16, cih8_16 = pyMRAW.save_mraw(images8, os.path.join(tmpdir, 'test8_16.cih'), bit_depth=16, ext='mraw', info_dict={'Shutter Speed(s)':0.001})
+            mraw16, cih16 = pyMRAW.save_mraw(images16, os.path.join(tmpdir, 'test16.cih'), bit_depth=16, ext='mraw', info_dict={'Comment Text':'Test saving 16 bit images.'})
+
+            loaded_images8, info8 =  pyMRAW.load_video(cih8)
+            loaded_images8_16, info8_16 =  pyMRAW.load_video(cih8_16)
+            loaded_images16, info16 =  pyMRAW.load_video(cih16)
+
+            assert loaded_images8.shape == images8.shape
+            assert loaded_images8_16.shape == images8.shape
+            assert loaded_images16.shape == images16.shape
+            assert loaded_images8.dtype == np.uint8
+            assert loaded_images8_16.dtype == np.uint16
+            assert loaded_images16.dtype == np.uint16
+            assert info8['Record Rate(fps)'] == 10
+            assert info8_16['Shutter Speed(s)'] == 0.001
+            assert info16['Comment Text'] == 'Test saving 16 bit images.'
+
+            loaded_images8._mmap.close()
+            loaded_images8_16._mmap.close()
+            loaded_images16._mmap.close()
+
+
+if __name__ == '__main__':
     np.testing.run_module_suite()