Merge pull request #33 from FrontierDevelopmentLab/feature/bqa-masks

Feature/bqa masks

Author: frandorr

conf/config.yaml

@@ -25,6 +25,7 @@ defaults:
   - builder: gcp
   - cloud: gcp
   - preparer: gcp
+  - plugins: gcp
   - _self_
 tasks:
   - build
@@ -33,6 +34,7 @@ tasks:
   - schedule
   - prepare
   - deploy
+  - plugins
 
 hydra:
   run:

conf/plugins/gcp.yaml

@@ -0,0 +1 @@
+_target_: satextractor.plugins.gcp_mtl_plugin.copy_mtl_files

src/satextractor/cli.py

@@ -146,19 +146,19 @@ def deployer(cfg):
         topic,
     )
 
-    # extraction_tasks_path = os.path.join(
-    #     ".", cfg.dataset_name + "_extraction_tasks.pkl",
-    # )
 
-    # logger.info(f"deploying on {cfg.deployer._target_} to {cfg.cloud.storage_root}",)
+def plugins(cfg):
 
-    # extraction_tasks = pickle.load(open(extraction_tasks_path, "rb"))
+    logger.info(f"using {cfg.plugins._target_} as plugin")
 
-    # # check tiles meet spec
-    # for t in extraction_tasks:
-    #     assert isinstance(t, ExtractionTask), "Task does not match ExtractionTask spec"
+    extraction_tasks = pickle.load(open(cfg.extraction_tasks, "rb"))
 
-    # hydra.utils.instantiate(cfg.deployer, extraction_tasks)
+    hydra.utils.call(
+        cfg.plugins,
+        cfg.credentials,
+        extraction_tasks,
+        f"{cfg.cloud.storage_prefix}/{cfg.cloud.storage_root}/{cfg.dataset_name}",
+    )
 
 
 @hydra.main(config_path="./../../conf", config_name="config")
@@ -181,7 +181,8 @@ def main(cfg: DictConfig):
             "schedule",
             "prepare",
             "deploy",
-        ], "valid tasks are [build, stac, tile, schedule, prepare, deploy]"
+            "plugins",
+        ], "valid tasks are [build, stac, tile, schedule, prepare, deploy, plugins]"
 
     logger.info(f"Running tasks {cfg.tasks}")
 
@@ -203,6 +204,9 @@ def main(cfg: DictConfig):
     if "deploy" in cfg.tasks:
         deployer(cfg)
 
+    if "plugins" in cfg.tasks:
+        plugins(cfg)
+
     return 0
 
 

src/satextractor/plugins/__init__.py

@@ -0,0 +1 @@
+from .gcp_mtl_plugin import copy_mtl_files  # noqa

src/satextractor/plugins/gcp_mtl_plugin.py

@@ -0,0 +1,55 @@
+from typing import List
+
+from gcsfs import GCSFileSystem
+from joblib import delayed
+from joblib import Parallel
+from satextractor.models import ExtractionTask
+from satextractor.utils import tqdm_joblib
+from tqdm import tqdm
+
+
+def copy_if_exists(fs, src, dst):
+    if fs.exists(src):
+        fs.copy(src, dst)
+
+
+def copy_mtl_files(
+    credentials: str,
+    tasks: List[ExtractionTask],
+    storage_path: str,
+    n_jobs: int = -1,
+) -> bool:
+    fs = GCSFileSystem(token=credentials)
+    mtl_files_info = set(
+        [
+            (
+                task.constellation,
+                tile.id,
+                str(task.sensing_time)[:10],
+                task.item_collection[0].assets["B1"].href.replace("B1.TIF", "MTL.txt"),
+            )
+            for task in tasks
+            for tile in task.tiles
+            if "landsat" in task.constellation
+        ],
+    )
+
+    to_copy = []
+    for mtl_info in mtl_files_info:
+        src = mtl_info[3]
+        dst = (
+            f"{storage_path}/{mtl_info[1]}/{mtl_info[0]}/metadata/{mtl_info[2]}_MTL.txt"
+        )
+        to_copy.append((src, dst))
+
+    with tqdm_joblib(
+        tqdm(
+            desc=f"parallel copying MTL files on {storage_path}",
+            total=len(to_copy),
+        ),
+    ):
+        Parallel(n_jobs=n_jobs, verbose=0, prefer="threads")(
+            [delayed(copy_if_exists)(fs, src, dst) for src, dst in to_copy],
+        )
+
+    return True