Metadataset improvements

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
 '''
 
 

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:

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")