satcube 0.1.16__py3-none-any.whl → 0.1.18__py3-none-any.whl

satcube/__init__.py

@@ -1,9 +1,7 @@
-from satcube.cloud_detection import cloud_masking
 from satcube.download import download
-from satcube.align import align
-import importlib.metadata
 from satcube.objects  import SatCubeMetadata  
+import importlib.metadata
 
-__all__ = ["cloud_masking", "download", "align", "SatCubeMetadata"]
+__all__ = ["download", "SatCubeMetadata"]
 # __version__ = importlib.metadata.version("satcube")
 

satcube/align.py

@@ -1,73 +1,91 @@
 from __future__ import annotations
 
 import pathlib
-from typing import List, Tuple
-import pickle
+from typing import Tuple
 import pandas as pd
 import satalign
-import shutil
-
 import numpy as np
 import rasterio as rio
-import xarray as xr
-from affine import Affine
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from tqdm import tqdm
 
+def _process_image(
+    image: np.ndarray, 
+    reference: np.ndarray, 
+    profile: dict,
+    output_path: pathlib.Path,
+) -> Tuple[float, float]:
+    
+    image_float = image.astype(np.float32) / 10000
+    image_float = image_float[np.newaxis, ...]
+    
+    image, M = satalign.LGM(
+        datacube=image_float,
+        reference=reference
+    ).run_multicore()
+    
+    image = (image * 10000).astype(np.uint16).squeeze()
 
-def process_row(row: pd.Series, reference: np.ndarray, input_dir: pathlib.Path, output_dir: pathlib.Path) -> None:
+    with rio.open(output_path, "w", **profile) as dst:
+        dst.write(image)
+    
+    return M[0][0, 2], M[0][1, 2]
+
+def _process_row(
+    row: pd.Series, 
+    reference: np.ndarray, 
+    input_dir: pathlib.Path, 
+    output_dir: pathlib.Path
+) -> Tuple[str, float, float]:
+    
     row_path = input_dir / (row["id"] + ".tif")
     output_path = output_dir / (row["id"] + ".tif")
-    with rio.open(row_path) as src:
-        row_image = src.read()
-        profile_image = src.profile
     
-    row_image_float = row_image.astype(np.float32) / 10000
-    row_image_float = row_image_float[np.newaxis, ...]
+    with rio.open(row_path) as src:
+        image = src.read()
+        profile = src.profile
     
-    pcc_model = satalign.LGM(
-        datacube     = row_image_float,
-        reference    = reference
+    dx_px, dy_px = _process_image(
+        image=image,
+        reference=reference,
+        profile=profile,
+        output_path=output_path
     )
-    image, _ = pcc_model.run_multicore()
-    image = (image * 10000).astype(np.uint16).squeeze()
-
-    with rio.open(output_path, "w", **profile_image) as dst:
-        dst.write(image)
+    
+    return row["id"], dx_px, dy_px 
 
-def align(
+def align_fn(
+    metadata: pd.DataFrame | None = None,
     input_dir: str | pathlib.Path = "raw",
     output_dir: str | pathlib.Path = "aligned",
     nworks: int = 4,
     cache: bool = False
-) -> None:
+) -> pd.DataFrame | None:
     
     input_dir = pathlib.Path(input_dir).expanduser().resolve()
     output_dir = pathlib.Path(output_dir).expanduser().resolve()
     output_dir.mkdir(parents=True, exist_ok=True)
-    
-    metadata_path = input_dir / "metadata.csv"
-    
-    if not metadata_path.exists():
+        
+    if metadata is None:
         raise FileNotFoundError(
-            f"Metadata file not found: {metadata_path}. "
+            f"Add metadata file to do alignment."
             "Please run the download step first."
         )
-    else:
-        metadata = pd.read_csv(metadata_path)
-        
-    if cache:
-        exist_files = [file.stem for file in output_dir.glob("*.tif")]
-        metadata = metadata[~metadata["id"].isin(exist_files)]
         
-        if metadata.empty:
-            return
 
     id_reference = metadata.sort_values(
-        by=["cs_cdf", "date"],
+        by=["cs_cdf"],
         ascending=False,
     ).iloc[0]["id"]
     
+    df = metadata.copy()
+    
+    if cache:
+        exist_files = [file.stem for file in output_dir.glob("*.tif")]
+        df = df[~df["id"].isin(exist_files)]
+        if df.empty:
+            return metadata
+        
     reference_path = input_dir / (id_reference + ".tif")
     
     with rio.open(reference_path) as ref_src:
@@ -75,11 +93,20 @@ def align(
         
     reference_float = reference.astype(np.float32) / 10000
 
+    results = []
+    
     with ThreadPoolExecutor(max_workers=nworks) as executor:
         futures = {
-            executor.submit(process_row, row, reference_float, input_dir, output_dir)
-            for _, row in metadata.iterrows()
+            executor.submit(
+                _process_row, 
+                row=row, 
+                reference=reference_float, 
+                input_dir=input_dir, 
+                output_dir=output_dir
+            ): row["id"]
+            for _, row in df.iterrows()
         }
+        
         for future in tqdm(
             as_completed(futures), 
             total=len(futures), 
@@ -88,11 +115,25 @@ def align(
             leave=True
         ):
             try:
-                future.result()
+                img_id, dx_px, dy_px = future.result()
+                results.append({"id": img_id,
+                                "dx_px": dx_px,
+                                "dy_px": dy_px})
             except Exception as e:
-                print(f"Error processing image: {e}")
-                
-    metadata = input_dir / "metadata.csv"
-    if metadata.exists():
-        metadata_dst = output_dir / "metadata.csv"
-        shutil.copy(metadata, metadata_dst)
+                print(f"Error processing image: {e} {futures[future]}")
+    
+    shift_df = pd.DataFrame(results)
+    
+    metadata = metadata.drop(
+        columns=["dx_px","dy_px"], 
+        errors="ignore"
+    )
+    
+    metadata = metadata.merge(
+        shift_df,
+        on="id",
+        how="left",
+        suffixes=('', '') 
+    )
+    
+    return metadata

satcube/archive_cloud_detection.py

@@ -0,0 +1,23 @@
+import torch
+
+class LandsatCloudDetector(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Define bit flags for clouds based on the
+        # Landsat QA band documentation
+        cloud_flags = (1 << 3) | (1 << 4) | (1 << 1)
+
+        ## Get the QA band
+        qa_band = x[6]
+        mask_band = x[:6].mean(axis=0)
+        mask_band[~torch.isnan(mask_band)] = 1
+
+        ## Create a cloud mask
+        cloud_mask = torch.bitwise_and(qa_band.int(), cloud_flags) == 0
+        cloud_mask = cloud_mask.float()
+        cloud_mask[cloud_mask == 0] = torch.nan
+        cloud_mask[cloud_mask == 0] = 1
+        final_mask = cloud_mask * mask_band
+        return final_mask

satcube/archive_dataclass.py

@@ -0,0 +1,39 @@
+import pathlib
+from datetime import datetime
+from typing import List, Optional
+
+import pydantic
+
+
+class Sensor(pydantic.BaseModel):
+    start_date: str
+    end_date: str
+    edge_size: int
+    bands: List[str]
+
+
+class Sentinel2(Sensor):
+    weight_path: pathlib.Path
+    start_date: Optional[str] = "2015-06-27"
+    end_date: Optional[str] = datetime.now().strftime("%Y-%m-%d")
+    resolution: Optional[int] = 10
+    edge_size: Optional[int] = 384
+    embedding_universal: Optional[str] = "s2_embedding_model_universal.pt"
+    cloud_model_universal: str = "s2_cloud_model_universal.pt"
+    cloud_model_specific: str = "s2_cloud_model_specific.pt"
+    super_model_specific: str = "s2_super_model_specific.pt"
+    bands: List[str] = [
+        "B01",
+        "B02",
+        "B03",
+        "B04",
+        "B05",
+        "B06",
+        "B07",
+        "B08",
+        "B8A",
+        "B09",
+        "B10",
+        "B11",
+        "B12",
+    ]