satcube 0.1.13__py3-none-any.whl → 0.1.15__py3-none-any.whl

satcube/__init__.py

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

satcube/align.py

@@ -0,0 +1,98 @@
+from __future__ import annotations
+
+import pathlib
+from typing import List, Tuple
+import pickle
+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_row(row: pd.Series, reference: np.ndarray, input_dir: pathlib.Path, output_dir: pathlib.Path) -> None:
+    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, ...]
+    
+    pcc_model = satalign.LGM(
+        datacube     = row_image_float,
+        reference    = reference
+    )
+    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)
+
+def align(
+    input_dir: str | pathlib.Path = "raw",
+    output_dir: str | pathlib.Path = "aligned",
+    nworks: int = 4,
+    cache: bool = False
+) -> 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():
+        raise FileNotFoundError(
+            f"Metadata file not found: {metadata_path}. "
+            "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"],
+        ascending=False,
+    ).iloc[0]["id"]
+    
+    reference_path = input_dir / (id_reference + ".tif")
+    
+    with rio.open(reference_path) as ref_src:
+        reference = ref_src.read()
+        
+    reference_float = reference.astype(np.float32) / 10000
+
+    with ThreadPoolExecutor(max_workers=nworks) as executor:
+        futures = {
+            executor.submit(process_row, row, reference_float, input_dir, output_dir)
+            for _, row in metadata.iterrows()
+        }
+        for future in tqdm(
+            as_completed(futures), 
+            total=len(futures), 
+            desc="Aligning images",
+            unit="image",
+            leave=True
+        ):
+            try:
+                future.result()
+            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)

satcube/cloud_detection.py

@@ -12,36 +12,148 @@ Example
 
 from __future__ import annotations
 
-import time
-from pathlib import Path
-from typing import List
+import pathlib
 
 import mlstac
 import numpy as np
 import rasterio as rio
+from rasterio.windows import Window
 import torch
+import pandas as pd
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from tqdm import tqdm
+import rasterio as rio
+from rasterio.merge import merge
+import shutil
 
-from satcube.utils import DeviceManager, _reset_gpu
-import warnings, re
-
-
+from satcube.utils import define_iteration, DeviceManager
+import warnings
 warnings.filterwarnings(
     "ignore",
-    message=re.escape("The secret `HF_TOKEN` does not exist in your Colab secrets."),
+    message="The secret HF_TOKEN does not exist in your Colab secrets.",
     category=UserWarning,
-    module="huggingface_hub.utils._auth",
+    module=r"huggingface_hub\.utils\._.*",
 )
 
-def cloud_masking(
-    input: str | Path,              # noqa: A002 (shadowing built-in is OK here)
-    output: str | Path,
+
+
+def infer_cloudmask(
+    input_path: str | pathlib.Path,
+    output_path: str | pathlib.Path,
+    cloud_model: torch.nn.Module,
     *,
-    tile: int = 512,
-    pad: int = 64,
+    chunk_size: int = 512,
+    overlap: int = 32,
+    device: str = "cpu",
     save_mask: bool = False,
+    prefix: str = ""          
+) -> pathlib.Path:
+    """
+    Predict 'image_path' in overlapping patches of 'chunk_size' x 'chunk_size',
+    but only write the valid (inner) region to avoid seam artifacts.
+
+    This uses partial overlap logic:
+      - For interior tiles, skip overlap//2 on each side.
+      - For boundary tiles, we skip only the interior side to avoid losing data at the edges.
+
+    Parameters
+    ----------
+    image_path : Path to input image.
+    output_path : Path to output single-band mask.
+    cloud_model : PyTorch model (already loaded with weights).
+    chunk_size : Size of each tile to read from the source image (default 512).
+    overlap : Overlap in pixels between adjacent tiles (default 32).
+    device : "cpu" or "cuda:0".
+
+    Returns
+    -------
+    pathlib.Path : The path to the created output image.
+    """
+
+    input_path = pathlib.Path(input_path)
+    output_path = pathlib.Path(output_path)
+
+    with rio.open(input_path) as src:
+        meta = src.profile
+        if not meta.get("tiled", False):
+            raise ValueError("The input image is not marked as tiled in its metadata.")
+        # Ensure the internal blocksize matches chunk_size
+        if chunk_size % meta["blockxsize"] != 0 and meta["blockxsize"] <= chunk_size:
+            raise ValueError(f"Image blocks must be {chunk_size}x{chunk_size}, "
+                            f"got {meta['blockxsize']}x{meta['blockysize']}")
+        height, width = meta["height"], meta["width"]
+
+    full_mask = np.zeros((height, width), dtype=np.float32)
+
+    coords = define_iteration((height, width), chunk_size, overlap)
+
+    with rio.open(input_path) as src:
+
+        for (row_off, col_off) in coords: 
+
+            window = Window(col_off, row_off, chunk_size, chunk_size)
+            patch = src.read(window=window) / 1e4
+            patch_tensor = torch.from_numpy(patch).float().unsqueeze(0).to(device)
+            result = cloud_model(patch_tensor).cpu().numpy().astype(np.uint8)
+      
+            if col_off == 0:
+                offset_x = 0
+            else:
+                offset_x = col_off + overlap // 2
+            if row_off == 0:
+                offset_y = 0
+            else:
+                offset_y = row_off + overlap // 2
+            if (offset_x + chunk_size) == width:
+                length_x = chunk_size
+                sub_x_start = 0
+            else:
+                length_x = chunk_size - (overlap // 2)
+                sub_x_start = overlap // 2 if col_off != 0 else 0
+
+            if (offset_y + chunk_size) == height:
+                length_y = chunk_size
+                sub_y_start = 0
+            else:
+                length_y = chunk_size - (overlap // 2)
+                sub_y_start = overlap // 2 if row_off != 0 else 0
+
+            full_mask[
+                offset_y : offset_y + length_y,
+                offset_x : offset_x + length_x
+            ] = result[
+                sub_y_start : sub_y_start + length_y,
+                sub_x_start : sub_x_start + length_x
+            ]
+
+        if save_mask:
+            out_meta = meta.copy()
+            out_meta.update(count=1, dtype="uint8", nodata=255)
+            output_mask = output_path.parent / (output_path.stem + "_mask.tif")
+            with rio.open(output_mask, "w", **out_meta) as dst:
+                dst.write(full_mask, 1)
+            
+
+        data = src.read()
+        img_prof = src.profile.copy()
+
+        masked = data.copy()
+        masked[:, full_mask != 0] = 65535
+        img_prof.update(dtype="uint16", nodata=65535)
+
+        with rio.open(output_path, "w", **img_prof) as dst:
+            dst.write(masked)
+
+    return output_path
+
+def cloud_masking( 
+    input: str | pathlib.Path = "raw",
+    output: str | pathlib.Path = "masked",
+    model_path: str | pathlib.Path = "SEN2CloudEnsemble",
     device: str = "cpu",
-    max_pix_cpu: float = 7.0e7
-) -> List[Path]:
+    save_mask: bool = False,
+    nworks: int = 4,
+) -> list[pathlib.Path]:
     """Write cloud-masked Sentinel-2 images.
 
     Parameters
@@ -49,7 +161,8 @@ def cloud_masking(
     input
         Path to a single ``.tif`` file **or** a directory containing them.
     output
-        Destination directory (created if missing).
+        Destination directory (created i
+        f missing).
     tile, pad
         Tile size and padding (pixels) when tiling is required.
     save_mask
@@ -64,14 +177,12 @@ def cloud_masking(
     list[pathlib.Path]
         Paths to the generated masked images.
     """
-    t_start = time.perf_counter()
-
-    src = Path(input).expanduser().resolve()
-    dst_dir = Path(output).expanduser().resolve()
+    src = pathlib.Path(input).expanduser().resolve()
+    dst_dir = pathlib.Path(output).expanduser().resolve()
     dst_dir.mkdir(parents=True, exist_ok=True)
 
     # Collect files to process -------------------------------------------------
-    tif_paths: list[Path]
+    tif_paths = []
     if src.is_dir():
         tif_paths = [p for p in src.rglob("*.tif")]
     elif src.is_file() and src.suffix.lower() == ".tif":
@@ -84,95 +195,44 @@ def cloud_masking(
         print(f"[cloud_masking] No .tif files found in {src}")
         return []
     
-    dir = Path("SEN2CloudEnsemble")
-
-    if not dir.exists():
-
+    if not pathlib.Path(model_path).exists():
         mlstac.download(
             file = "https://huggingface.co/tacofoundation/CloudSEN12-models/resolve/main/SEN2CloudEnsemble/mlm.json",
-            output_dir = "SEN2CloudEnsemble",
+            output_dir = model_path
         )
 
-    experiment = mlstac.load(dir.as_posix())
-
-    dm = DeviceManager(experiment, init_device=device)
-
-    masked_paths: list[Path] = []
-
-    # -------------------------------------------------------------------------
-    for idx, tif_path in enumerate(tif_paths, 1):
-        rel = tif_path.relative_to(src)
-        out_dir = dst_dir / rel.parent
-        out_dir.mkdir(parents=True, exist_ok=True)
-
-        mask_path = out_dir / f"{tif_path.stem}_cloudmask.tif"
-        masked_path = out_dir / f"{tif_path.stem}_masked.tif"
-
-        with rio.open(tif_path) as src_img:
-            profile = src_img.profile
-            h, w = src_img.height, src_img.width
-
-        mask_prof = profile.copy()
-        mask_prof.update(driver="GTiff", count=1, dtype="uint8", nodata=255)
-
-        do_tiling = (dm.device == "cuda") or (h * w > max_pix_cpu)
-        full_mask = np.full((h, w), 255, np.uint8)
-
-        t0 = time.perf_counter()
-
-        # ----------------------- inference -----------------------------------
-        if not do_tiling:  # full frame
-            with rio.open(tif_path) as src_img, torch.inference_mode():
-                img = src_img.read().astype(np.float32) / 1e4
-                h32, w32 = (h + 31) // 32 * 32, (w + 31) // 32 * 32
-                pad_b, pad_r = h32 - h, w32 - w
-                tensor = torch.from_numpy(img).unsqueeze(0)
-                if pad_b or pad_r:
-                    tensor = torch.nn.functional.pad(tensor, (0, pad_r, 0, pad_b))
-                mask = dm.model(tensor.to(dm.device)).squeeze(0)
-                full_mask[:] = mask[..., :h, :w].cpu().numpy().astype(np.uint8)
-        else:  # tiled
-            with rio.open(tif_path) as src_img, torch.inference_mode():
-                for y0 in range(0, h, tile):
-                    for x0 in range(0, w, tile):
-                        y0r, x0r = max(0, y0 - pad), max(0, x0 - pad)
-                        y1r, x1r = min(h, y0 + tile + pad), min(w, x0 + tile + pad)
-                        win = rio.windows.Window(x0r, y0r, x1r - x0r, y1r - y0r)
-
-                        patch = src_img.read(window=win).astype(np.float32) / 1e4
-                        tensor = torch.from_numpy(patch).unsqueeze(0).to(dm.device)
-                        mask = dm.model(tensor).squeeze(0).cpu().numpy().astype(np.uint8)
-
-                        y_in0 = pad if y0r else 0
-                        x_in0 = pad if x0r else 0
-                        y_in1 = mask.shape[0] - (pad if y1r < h else 0)
-                        x_in1 = mask.shape[1] - (pad if x1r < w else 0)
-                        core = mask[y_in0:y_in1, x_in0:x_in1]
-                        full_mask[y0 : y0 + core.shape[0], x0 : x0 + core.shape[1]] = core
-
-        # ----------------------- output --------------------------------------
-        if save_mask:
-            with rio.open(mask_path, "w", **mask_prof) as dst:
-                dst.write(full_mask, 1)
-
-        with rio.open(tif_path) as src_img:
-            data = src_img.read()
-            img_prof = src_img.profile.copy()
-
-        masked = data.copy()
-        masked[:, full_mask != 0] = 65535
-        img_prof.update(dtype="uint16", nodata=65535)
-
-        with rio.open(masked_path, "w", **img_prof) as dst:
-            dst.write(masked)
-
-        masked_paths.append(masked_path)
-        dt = time.perf_counter() - t0
-        print(f"[{idx}/{len(tif_paths)}] {rel} → done in {dt:.1f}s")
-
-    if dm.device == "cuda":
-        _reset_gpu()
-
-    total_time = time.perf_counter() - t_start
-    print(f"Processed {len(masked_paths)} image(s) in {total_time:.1f}s.")
-    return masked_paths
+    model = mlstac.load(model_path)
+    cloud_model = DeviceManager(model, init_device=device).model
+    cloud_model.eval()
+    
+    with ThreadPoolExecutor(max_workers=nworks) as executor:
+        futures = { 
+            executor.submit(
+                infer_cloudmask,
+                input_path=p,
+                output_path=dst_dir / p.name,
+                cloud_model=cloud_model,
+                device=device,
+                save_mask=save_mask,
+                prefix=f"[{i+1}/{len(tif_paths)}] "
+            ): p for i, p in enumerate(tif_paths)
+        }
+        
+        for future in tqdm(
+            as_completed(futures),
+            total=len(futures),
+            desc="Cloud Masking",
+            position=0,
+            leave=True
+        ):
+            p = futures[future]
+            try:
+                result = future.result()
+                print(f"{result} processed successfully.")
+            except Exception as e:
+                print(f"Error processing {p}: {e}")
+                
+    metadata = src / "metadata.csv"
+    if metadata.exists():
+        metadata_dst = dst_dir / "metadata.csv"
+        shutil.copy(metadata, metadata_dst)

satcube/download.py

@@ -1,57 +1,68 @@
-import pathlib
-import ee
-import cubexpress
+import sys, time, threading, itertools
+import cubexpress as ce
 import pandas as pd
+from satcube.objects import SatCubeMetadata
+import pathlib
 
-
-def download_data(
-    *,                     # keyword-only
+def download(
     lon: float,
     lat: float,
-    cloud_max: int = 40,
-    edge_size: int = 2_048,
+    edge_size: int,
     start: str,
     end: str,
-    output: str = "raw",
-    scale: int = 10,
-    nworks: int = 4,
-    mosaic: bool = True
-) -> pd.DataFrame:
-    """
-    Download a Sentinel cube for (lon, lat) and return its metadata.
+    *,
+    max_cscore: float = 1,
+    min_cscore: float = 0,
+    outfolder: str = "raw",
+    nworks: int = 4
+) -> "SatCubeMetadata":
+        
+        
+    outfolder = pathlib.Path(outfolder).resolve()
 
-    Parameters
-    ----------
-    lon, lat        Center point in degrees.
-    cloud_max       Max cloud cover (%).
-    edge_size       Square side length (m).
-    start, end      YYYY-MM-DD date range.
-    output          Folder for GeoTIFFs.
-    scale           Pixel size (m).
-    nworks          Parallel workers.
-    mosaic          Merge scenes per date.
-    auto_init_gee   Call ee.Initialize() if needed.
-
-    Returns
-    -------
-    pandas.DataFrame
-        Scene catalogue used for the request.
-    """
-    # Filter scenes
-    df = cubexpress.cloud_table(
+    table = ce.s2_table(
         lon=lon,
         lat=lat,
         edge_size=edge_size,
-        scale=scale,
-        cloud_max=cloud_max,
         start=start,
         end=end,
+        max_cscore=max_cscore,
+        min_cscore=min_cscore
+    )   
+        
+    requests = ce.table_to_requestset(
+        table=table,
+        mosaic=True
+    )
+    
+    ce.get_cube(
+        requests=requests,
+        outfolder=outfolder,
+        nworks=nworks
+    )
+
+    table_req = (
+        requests._dataframe.copy()
+        .drop(columns=['geotransform', 'manifest', 'outname', 'width', 'height', 'scale_x', 'scale_y'])
+    )
+    
+    table_req['date'] = table_req['id'].str.split('_').str[0]
+
+    result_table = (
+        table.groupby('date')
+        .agg(
+            id=('id', lambda x: '-'.join(x)),
+            cs_cdf=('cs_cdf', 'first')
+        )
+        .reset_index()  
     )
+    
+    table_final = table_req.merge(
+        result_table,
+        on='date',
+        how='left'
+    ).rename(columns={'id_x': 'id', 'id_y': 'gee_ids'})
 
-    # Build requests + ensure dir
-    requests = cubexpress.table_to_requestset(df, mosaic=mosaic)
-    pathlib.Path(output).mkdir(parents=True, exist_ok=True)
+    table_final.to_csv(outfolder / "metadata.csv", index=False)
 
-    # Download cube
-    cubexpress.get_cube(requests, output, nworks)
-    return df
+    return SatCubeMetadata(df=table_final, raw_dir=outfolder)

satcube/objects.py

@@ -0,0 +1,71 @@
+# satcube/objects.py
+from __future__ import annotations
+from dataclasses import dataclass, field
+import pathlib
+import pandas as pd
+
+from satcube.align  import align as _align_fn
+from satcube.cloud_detection  import cloud_masking as _cloud_fn
+
+@dataclass
+class SatCubeMetadata:
+    df: pd.DataFrame = field(repr=False)
+    raw_dir: pathlib.Path = field(repr=False)
+
+    def __repr__(self) -> str: 
+        return self.df.__repr__()
+    
+    __str__ = __repr__
+    
+    def _repr_html_(self) -> str:
+        html = getattr(self.df, "_repr_html_", None)
+        return html() if callable(html) else self.df.__repr__()
+
+    def align(
+        self,
+        input_dir: str | pathlib.Path | None = None,
+        output_dir: str | pathlib.Path = "aligned",
+        nworks: int = 4,
+        cache: bool = False
+    ) -> "SatCubeMetadata":
+        
+        if input_dir is None:
+            input_dir = self.raw_dir
+        else:
+            input_dir = pathlib.Path(input_dir).expanduser().resolve()
+
+        _align_fn(
+            input_dir=input_dir, 
+            output_dir=output_dir, 
+            nworks=nworks, 
+            cache=cache
+        )
+        self.aligned_dir = pathlib.Path(output_dir).resolve()
+        return self
+
+    def cloud_masking(
+        self,
+        output_dir: str | pathlib.Path = "masked",
+        model_path: str | pathlib.Path = "SEN2CloudEnsemble",
+        device: str = "cpu"
+    ) -> "SatCubeMetadata":
+        if not hasattr(self, "aligned_dir"):
+            raise RuntimeError("You must run .align() first")
+        _cloud_fn(
+            input=self.aligned_dir,
+            output=output_dir,
+            model_path=model_path,
+            device=device
+        )
+        self.masked_dir = pathlib.Path(output_dir).resolve()
+        return self
+    
+    def __getattr__(self, item):
+        return getattr(self.df, item)
+
+    def __getitem__(self, key):
+        return self.df.__getitem__(key)
+    
+    def __len__(self):
+        return len(self.df)
+