diff --git a/opensfm/config.py b/opensfm/config.py index 1ec6b4332..ed7df0c02 100644 --- a/opensfm/config.py +++ b/opensfm/config.py @@ -119,6 +119,7 @@ submodels_relpath: "submodels" # Relative path to the submodels directory submodel_relpath_template: "submodels/submodel_%04d" # Template to generate the relative path to a submodel directory submodel_images_relpath_template: "submodels/submodel_%04d/images" # Template to generate the relative path to a submodel images directory +submodel_masks_relpath_template: "submodels/submodel_%04d/masks" # Template to generate the relative path to a submodel masks directory ''' diff --git a/opensfm/dataset.py b/opensfm/dataset.py index 6ed712555..8d79fe537 100644 --- a/opensfm/dataset.py +++ b/opensfm/dataset.py @@ -28,7 +28,7 @@ def __init__(self, data_path): :param data_path: Path to directory containing dataset """ - self.data_path = data_path + self.data_path = os.path.abspath(data_path) self._load_config() @@ -91,10 +91,18 @@ def masks(self): """Return list of file names of all masks in this dataset""" return self.mask_list - def mask_as_array(self, image): - """Given an image, returns the associated mask as an array if it exists, otherwise returns None""" + def mask(self, image): + """Given an image, returns the associated mask's name if it exists, otherwise returns None""" mask_name = image + '.png' if mask_name in self.masks(): + return mask_name + else: + return None + + def mask_as_array(self, image): + """Given an image, returns the associated mask as an array if it exists, otherwise returns None""" + mask_name = self.mask(image) + if mask_name is not None: mask_path = self.mask_files[mask_name] mask = cv2.imread(mask_path) if len(mask.shape) == 3: diff --git a/opensfm/large/metadataset.py b/opensfm/large/metadataset.py index 936dc33a0..d48dc742b 100644 --- a/opensfm/large/metadataset.py +++ b/opensfm/large/metadataset.py @@ -16,7 +16,7 @@ def __init__(self, data_path): :param data_path: Path to directory containing meta dataset ''' - self.data_path = data_path + self.data_path = os.path.abspath(data_path) config_file = os.path.join(self.data_path, 'config.yaml') self.config = config.load_config(config_file) @@ -40,6 +40,11 @@ def _submodel_images_path(self, i): template = self.config['submodel_images_relpath_template'] return os.path.join(self.data_path, template % i) + def _submodel_masks_path(self, i): + """Path to submodel i masks folder.""" + template = self.config['submodel_masks_relpath_template'] + return os.path.join(self.data_path, template % i) + def _image_list_path(self): return os.path.join(self._submodels_path(), self._image_list_file_name) @@ -118,8 +123,10 @@ def create_submodels(self, clusters, no_symlinks=False): # create sub model dirs submodel_path = self._submodel_path(i) submodel_images_path = self._submodel_images_path(i) + submodel_masks_path = self._submodel_masks_path(i) io.mkdir_p(submodel_path) io.mkdir_p(submodel_images_path) + io.mkdir_p(submodel_masks_path) # link images and create image list file image_list_path = os.path.join(submodel_path, 'image_list.txt') @@ -129,6 +136,12 @@ def create_submodels(self, clusters, no_symlinks=False): dst = os.path.join(submodel_images_path, image) if not os.path.isfile(dst): os.symlink(src, dst) + mask = data.mask(image) + if mask is not None: + mask_src = data.mask_files[mask] + mask_dst = os.path.join(submodel_masks_path, mask) + if not os.path.isfile(mask_dst): + os.symlink(mask_src, mask_dst) dst_relpath = os.path.relpath(dst, submodel_path) txtfile.write(dst_relpath + "\n")