From 2e2d674dc8fa3dbf4a33e99979740d8aebb9a352 Mon Sep 17 00:00:00 2001
From: Piero Toffanin <pt@masseranolabs.com>
Date: Tue, 15 Sep 2020 17:57:36 +0000
Subject: [PATCH] Cleanup, image mask support (initial, POC)

---
 SuperBuild/CMakeLists.txt |  2 +-
 opendm/io.py              |  9 ---------
 opendm/osfm.py            | 12 ++++++++++++
 opendm/photo.py           | 15 ++++++++++++++-
 stages/dataset.py         | 10 +++++-----
 stages/mve.py             |  3 ++-
 6 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/SuperBuild/CMakeLists.txt b/SuperBuild/CMakeLists.txt
index adfc53f2..d3180c9d 100644
--- a/SuperBuild/CMakeLists.txt
+++ b/SuperBuild/CMakeLists.txt
@@ -129,7 +129,7 @@ endforeach()
 
 externalproject_add(mve
     GIT_REPOSITORY  https://github.com/OpenDroneMap/mve.git
-    GIT_TAG         099
+    GIT_TAG         200
     UPDATE_COMMAND  ""
     SOURCE_DIR      ${SB_SOURCE_DIR}/elibs/mve
     CONFIGURE_COMMAND ""
diff --git a/opendm/io.py b/opendm/io.py
index 66ee5f98..bb23a1ed 100644
--- a/opendm/io.py
+++ b/opendm/io.py
@@ -2,19 +2,10 @@ import os
 import shutil, errno
 import json
 
-def get_files_list(path_dir):
-    return os.listdir(path_dir)
-
-
 def absolute_path_file(path_file):
     return os.path.abspath(path_file)
 
 
-def extract_file_from_path_file(path_file):
-    path, file = os.path.split(path_file)
-    return file
-
-
 def extract_path_from_file(file):
     path_file = os.path.abspath(os.path.dirname(file))
     path, file = os.path.split(path_file)
diff --git a/opendm/osfm.py b/opendm/osfm.py
index 21fca2c5..d35d96ce 100644
--- a/opendm/osfm.py
+++ b/opendm/osfm.py
@@ -129,6 +129,18 @@ class OSFMContext:
 
             with open(os.path.join(self.opensfm_project_path, "exif_overrides.json"), 'w') as f:
                 f.write(json.dumps(exif_overrides))
+
+            # Check image masks
+            masks = []
+            for p in photos:
+                if p.mask is not None:
+                    masks.append((p.filename, os.path.join(images_path, p.mask)))
+            
+            if masks:
+                log.ODM_INFO("Found %s image masks" % len(masks))
+                with open(os.path.join(self.opensfm_project_path, "mask_list.txt"), 'w') as f:
+                    for fname, mask in masks:
+                        f.write("{} {}\n".format(fname, mask))
             
             # Compute feature_process_size
             feature_process_size = 2048 # default
diff --git a/opendm/photo.py b/opendm/photo.py
index 55f044c3..28bfb991 100644
--- a/opendm/photo.py
+++ b/opendm/photo.py
@@ -1,5 +1,6 @@
 import logging
 import re
+import os
 
 import exifread
 import numpy as np
@@ -14,12 +15,24 @@ 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)
+        
         # Standard tags (virtually all photos have these)
-        self.filename = io.extract_file_from_path_file(path_file)
         self.width = None
         self.height = None
         self.camera_make = ''
diff --git a/stages/dataset.py b/stages/dataset.py
index 525a7d4b..41caca26 100644
--- a/stages/dataset.py
+++ b/stages/dataset.py
@@ -49,16 +49,16 @@ class ODMLoadDatasetStage(types.ODM_Stage):
             with open(tree.benchmarking, 'a') as b:
                 b.write('ODM Benchmarking file created %s\nNumber of Cores: %s\n\n' % (system.now(), context.num_cores))
     
-        # check if the extension is supported
-        def supported_extension(file_name):
-            (pathfn, ext) = os.path.splitext(file_name)
-            return ext.lower() in context.supported_extensions
+        # check if the image filename is supported
+        def valid_image_filename(filename):
+            (pathfn, ext) = os.path.splitext(filename)
+            return ext.lower() in context.supported_extensions and pathfn[-5:] != "_mask"
 
         # 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 io.get_files_list(in_dir) if supported_extension(f)]
+            return [f for f in os.listdir(path_dir) if valid_image_filename(f)]
 
         # get images directory
         input_dir = tree.input_images
diff --git a/stages/mve.py b/stages/mve.py
index 9828d848..d775cd45 100644
--- a/stages/mve.py
+++ b/stages/mve.py
@@ -77,7 +77,8 @@ class ODMMveStage(types.ODM_Stage):
             self.update_progress(90)
 
             scene2pset_config = [
-                "-F%s" % mve_output_scale
+                "-F%s" % mve_output_scale,
+                '-mmask'
             ]
 
             # run scene2pset