Added option to specify custom image file type

pykitti/odometry.py

@@ -27,6 +27,9 @@ def __init__(self, base_path, sequence, **kwargs):
         # easy use with OpenCV.
         self.imformat = kwargs.get('imformat', None)
 
+        # Default image file extension is 'png'
+        self.imtype = kwargs.get('imtype', 'png')
+
         # Pre-load data that isn't returned as a generator
         self._load_calib()
         self._load_timestamps()
@@ -46,7 +49,7 @@ def poses(self):
                 lines = f.readlines()
                 if self.frames is not None:
                     lines = [lines[i] for i in self.frames]
-                    
+
                 for line in lines:
                     T_w_cam0 = np.fromstring(line, dtype=float, sep=' ')
                     T_w_cam0 = T_w_cam0.reshape(3, 4)
@@ -60,7 +63,8 @@ def poses(self):
     @property
     def cam0(self):
         """Generator to read image files for cam0 (monochrome left)."""
-        impath = os.path.join(self.sequence_path, 'image_0', '*.png')
+        impath = os.path.join(self.sequence_path, 'image_0',
+                              '*.{}'.format(self.imtype))
         imfiles = sorted(glob.glob(impath))
         # Subselect the chosen range of frames, if any
         if self.frames is not None:
@@ -72,7 +76,8 @@ def cam0(self):
     @property
     def cam1(self):
         """Generator to read image files for cam1 (monochrome right)."""
-        impath = os.path.join(self.sequence_path, 'image_1', '*.png')
+        impath = os.path.join(self.sequence_path, 'image_1',
+                              '*.{}'.format(self.imtype))
         imfiles = sorted(glob.glob(impath))
         # Subselect the chosen range of frames, if any
         if self.frames is not None:
@@ -84,7 +89,8 @@ def cam1(self):
     @property
     def cam2(self):
         """Generator to read image files for cam2 (RGB left)."""
-        impath = os.path.join(self.sequence_path, 'image_2', '*.png')
+        impath = os.path.join(self.sequence_path, 'image_2',
+                              '*.{}'.format(self.imtype))
         imfiles = sorted(glob.glob(impath))
         # Subselect the chosen range of frames, if any
         if self.frames is not None:
@@ -96,7 +102,8 @@ def cam2(self):
     @property
     def cam3(self):
         """Generator to read image files for cam0 (RGB right)."""
-        impath = os.path.join(self.sequence_path, 'image_3', '*.png')
+        impath = os.path.join(self.sequence_path, 'image_3',
+                              '*.{}'.format(self.imtype))
         imfiles = sorted(glob.glob(impath))
         # Subselect the chosen range of frames, if any
         if self.frames is not None:

pykitti/raw.py

@@ -27,6 +27,9 @@ def __init__(self, base_path, date, drive, **kwargs):
         # easy use with OpenCV.
         self.imformat = kwargs.get('imformat', None)
 
+        # Default image file extension is '.png'
+        self.imtype = kwargs.get('imtype', 'png')
+
         # Pre-load data that isn't returned as a generator
         self._load_calib()
         self._load_timestamps()
@@ -52,7 +55,8 @@ def oxts(self):
     @property
     def cam0(self):
         """Generator to read image files for cam0 (monochrome left)."""
-        impath = os.path.join(self.data_path, 'image_00', 'data', '*.png')
+        impath = os.path.join(self.data_path, 'image_00',
+                              'data', '*.{}'.format(self.imtype))
         imfiles = sorted(glob.glob(impath))
         # Subselect the chosen range of frames, if any
         if self.frames is not None:
@@ -64,7 +68,8 @@ def cam0(self):
     @property
     def cam1(self):
         """Generator to read image files for cam1 (monochrome right)."""
-        impath = os.path.join(self.data_path, 'image_01', 'data', '*.png')
+        impath = os.path.join(self.data_path, 'image_01',
+                              'data', '*.{}'.format(self.imtype))
         imfiles = sorted(glob.glob(impath))
         # Subselect the chosen range of frames, if any
         if self.frames is not None:
@@ -76,7 +81,8 @@ def cam1(self):
     @property
     def cam2(self):
         """Generator to read image files for cam2 (RGB left)."""
-        impath = os.path.join(self.data_path, 'image_02', 'data', '*.png')
+        impath = os.path.join(self.data_path, 'image_02',
+                              'data', '*.{}'.format(self.imtype))
         imfiles = sorted(glob.glob(impath))
         # Subselect the chosen range of frames, if any
         if self.frames is not None:
@@ -88,7 +94,8 @@ def cam2(self):
     @property
     def cam3(self):
         """Generator to read image files for cam0 (RGB right)."""
-        impath = os.path.join(self.data_path, 'image_03', 'data', '*.png')
+        impath = os.path.join(self.data_path, 'image_03',
+                              'data', '*.{}'.format(self.imtype))
         imfiles = sorted(glob.glob(impath))
         # Subselect the chosen range of frames, if any
         if self.frames is not None:

setup.py

@@ -2,12 +2,12 @@
 
 setup(
     name='pykitti',
-    version='0.2.3',
+    version='0.2.4',
     description='A minimal set of tools for working with the KITTI dataset in Python',
     author='Lee Clement',
     author_email='[email protected]',
     url='https://github.com/utiasSTARS/pykitti',
-    download_url='https://github.com/utiasSTARS/pykitti/tarball/0.2.3',
+    download_url='https://github.com/utiasSTARS/pykitti/tarball/0.2.4',
     license='MIT',
     packages=['pykitti'],
     install_requires=['numpy', 'matplotlib']