diff --git a/opendm/photo.py b/opendm/photo.py index 28bfb991..efbdb3f1 100644 --- a/opendm/photo.py +++ b/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 diff --git a/stages/dataset.py b/stages/dataset.py index a42dbd9f..16262284 100644 --- a/stages/dataset.py +++ b/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