fix imports

won't error if cv2 is not avalaible when not necessary
imageio is used instead of scipy

Author: Clément Pinard

datasets/KITTI.py

@@ -3,20 +3,31 @@
 import glob
 from .listdataset import ListDataset
 from .util import split2list
-import cv2
 import numpy as np
 import flow_transforms
 
+try:
+    import cv2
+except ImportError as e:
+    import warnings
+    with warnings.catch_warnings():
+        warnings.filterwarnings("default", category=ImportWarning)
+        warnings.warn("failed to load openCV, which is needed"
+                      "for KITTI which uses 16bit PNG images", ImportWarning)
+
 '''
 Dataset routines for KITTI_flow, 2012 and 2015.
 http://www.cvlibs.net/datasets/kitti/eval_flow.php
 The dataset is not very big, you might want to only finetune on it for flownet
 EPE are not representative in this dataset because of the sparsity of the GT.
+OpenCV is needed to load 16bit png images
 '''
 
 
 def load_flow_from_png(png_path):
-    flo_file = cv2.imread(png_path,cv2.IMREAD_UNCHANGED)
+    # The -1 is here to specify not to change the image depth (16bit), and is compatible
+    # with both OpenCV2 and OpenCV3
+    flo_file = cv2.imread(png_path, -1)
     flo_img = flo_file[:,:,2:0:-1].astype(np.float32)
     invalid = (flo_file[:,:,0] == 0)
     flo_img = flo_img - 32768
@@ -27,24 +38,25 @@ def load_flow_from_png(png_path):
 
 
 def make_dataset(dir, split, occ=True):
-    '''Will search in training folder for folders 'flow_noc' or 'flow_occ' and 'colored_0' (KITTI 2012) or 'image_2' (KITTI 2015) '''
+    '''Will search in training folder for folders 'flow_noc' or 'flow_occ'
+       and 'colored_0' (KITTI 2012) or 'image_2' (KITTI 2015) '''
     flow_dir = 'flow_occ' if occ else 'flow_noc'
-    assert(os.path.isdir(os.path.join(dir,flow_dir)))
+    assert(os.path.isdir(os.path.join(dir, flow_dir)))
     img_dir = 'colored_0'
-    if not os.path.isdir(os.path.join(dir,img_dir)):
+    if not os.path.isdir(os.path.join(dir, img_dir)):
         img_dir = 'image_2'
-    assert(os.path.isdir(os.path.join(dir,img_dir)))
+    assert(os.path.isdir(os.path.join(dir, img_dir)))
 
     images = []
-    for flow_map in glob.iglob(os.path.join(dir,flow_dir,'*.png')):
+    for flow_map in glob.iglob(os.path.join(dir, flow_dir, '*.png')):
         flow_map = os.path.basename(flow_map)
         root_filename = flow_map[:-7]
-        flow_map = os.path.join(flow_dir,flow_map)
-        img1 = os.path.join(img_dir,root_filename+'_10.png')
-        img2 = os.path.join(img_dir,root_filename+'_11.png')
-        if not (os.path.isfile(os.path.join(dir,img1)) or os.path.isfile(os.path.join(dir,img2))):
+        flow_map = os.path.join(flow_dir, flow_map)
+        img1 = os.path.join(img_dir, root_filename+'_10.png')
+        img2 = os.path.join(img_dir, root_filename+'_11.png')
+        if not (os.path.isfile(os.path.join(dir, img1)) or os.path.isfile(os.path.join(dir, img2))):
             continue
-        images.append([[img1,img2],flow_map])
+        images.append([[img1, img2], flow_map])
 
     return split2list(images, split, default_split=0.9)
 
@@ -58,9 +70,13 @@ def KITTI_loader(root,path_imgs, path_flo):
 def KITTI_occ(root, transform=None, target_transform=None,
               co_transform=None, split=None):
     train_list, test_list = make_dataset(root, split, True)
-    train_dataset = ListDataset(root, train_list, transform, target_transform, co_transform, loader=KITTI_loader)
+    train_dataset = ListDataset(root, train_list, transform,
+                                target_transform, co_transform,
+                                loader=KITTI_loader)
     # All test sample are cropped to lowest possible size of KITTI images
-    test_dataset = ListDataset(root, test_list, transform, target_transform, flow_transforms.CenterCrop((370,1224)), loader=KITTI_loader)
+    test_dataset = ListDataset(root, test_list, transform,
+                               target_transform, flow_transforms.CenterCrop((370,1224)),
+                               loader=KITTI_loader)
 
     return train_dataset, test_dataset
 

datasets/listdataset.py

@@ -1,7 +1,7 @@
 import torch.utils.data as data
 import os
 import os.path
-from scipy.ndimage import imread
+from imageio import imread
 import numpy as np