Support arbitrary mask extensions

Former-commit-id: efff487706

opendm/photo.py

@@ -15,22 +15,12 @@ import xmltodict as x2d
 from opendm import get_image_size
 from xml.parsers.expat import ExpatError
 
-def find_mask(photo_path):
-    (pathfn, ext) = os.path.splitext(photo_path)
-    mask_path = "{}_mask{}".format(pathfn, ext)
-    if os.path.exists(mask_path):
-        # Spaces are not supported due to OpenSfM's mask_list.txt format reqs
-        if not " " in mask_path:
-            return os.path.basename(mask_path)
-        else:
-            log.ODM_WARNING("Image mask {} has a space. Spaces are currently not supported for image masks.".format(os.path.basename(mask_path)))
-
 class ODM_Photo:
     """ODMPhoto - a class for ODMPhotos"""
 
     def __init__(self, path_file):
         self.filename = os.path.basename(path_file)
-        self.mask = find_mask(path_file)
+        self.mask = None
         
         # Standard tags (virtually all photos have these)
         self.width = None
@@ -88,6 +78,9 @@ class ODM_Photo:
                             self.filename, self.camera_make, self.camera_model, self.width, self.height, 
                             self.latitude, self.longitude, self.altitude, self.band_name, self.band_index)
 
+    def set_mask(self, mask):
+        self.mask = mask
+
     def update_with_geo_entry(self, geo_entry):
         self.latitude = geo_entry.y
         self.longitude = geo_entry.x

stages/dataset.py

@@ -56,19 +56,31 @@ class ODMLoadDatasetStage(types.ODM_Stage):
 
         # Get supported images from dir
         def get_images(in_dir):
-            # filter images for its extension type
             log.ODM_DEBUG(in_dir)
-            return [f for f in os.listdir(in_dir) if valid_image_filename(f)]
+            entries = os.listdir(in_dir)
+            valid, rejects = [], []
+            for f in entries:
+                if valid_image_filename(f):
+                    valid.append(f)
+                else:
+                    rejects.append(f)
+            return valid, rejects
+
+        def find_mask(photo_path, masks):
+            (pathfn, ext) = os.path.splitext(os.path.basename(photo_path))
+            k = "{}_mask".format(pathfn)
+            
+            mask = masks.get(k)
+            if mask:
+                # Spaces are not supported due to OpenSfM's mask_list.txt format reqs
+                if not " " in mask:
+                    return mask
+                else:
+                    log.ODM_WARNING("Image mask {} has a space. Spaces are currently not supported for image masks.".format(mask))
 
         # get images directory
-        input_dir = tree.input_images
         images_dir = tree.dataset_raw
 
-        if not io.dir_exists(images_dir):
-            log.ODM_INFO("Project directory %s doesn't exist. Creating it now. " % images_dir)
-            system.mkdir_p(images_dir)
-            copied = [copyfile(io.join_paths(input_dir, f), io.join_paths(images_dir, f)) for f in get_images(input_dir)]
-
         # define paths and create working directories
         system.mkdir_p(tree.odm_georeferencing)
         if not args.use_3dmesh: system.mkdir_p(tree.odm_25dgeoreferencing)
@@ -78,16 +90,25 @@ class ODMLoadDatasetStage(types.ODM_Stage):
         # check if we rerun cell or not
         images_database_file = io.join_paths(tree.root_path, 'images.json')
         if not io.file_exists(images_database_file) or self.rerun():
-            files = get_images(images_dir)
+            files, rejects = get_images(images_dir)
             if files:
                 # create ODMPhoto list
                 path_files = [io.join_paths(images_dir, f) for f in files]
 
+                # Lookup table for masks
+                masks = {}
+                for r in rejects:
+                    (p, ext) = os.path.splitext(r)
+                    if p[-5:] == "_mask":
+                        masks[p] = r
+
                 photos = []
                 with open(tree.dataset_list, 'w') as dataset_list:
                     log.ODM_INFO("Loading %s images" % len(path_files))
                     for f in path_files:
-                        photos += [types.ODM_Photo(f)]
+                        p = types.ODM_Photo(f)
+                        p.set_mask(find_mask(f, masks))
+                        photos += [p]
                         dataset_list.write(photos[-1].filename + '\n')
 
                 # Check if a geo file is available