plaknit 0.3.2__tar.gz → 0.3.4__tar.gz

{plaknit-0.3.2/plaknit.egg-info → plaknit-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plaknit
-Version: 0.3.2
+Version: 0.3.4
 Summary: Processing Large-Scale PlanetScope Data
 Author-email: Dryver Finch <dryver2206@gmail.com>
 License: MIT License

{plaknit-0.3.2 → plaknit-0.3.4}/plaknit/__init__.py

@@ -10,7 +10,7 @@ from .models.rf import predict_rf, smooth_probs, train_rf
 
 __author__ = """Dryver Finch"""
 __email__ = "dryver2206@gmail.com"
-__version__ = "0.3.2"
+__version__ = "0.3.4"
 
 __all__ = [
     "BRTEnsemble",

{plaknit-0.3.2 → plaknit-0.3.4}/plaknit/cli/ensemble_cli.py

@@ -263,7 +263,18 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
         "--jobs",
         type=int,
         default=1,
-        help="Parallel workers for block prediction (default: 1). Use -1 for all cores.",
+        help=(
+            "Parallel workers for block prediction (default: 1). Use -1 for all "
+            "cores on CPU, or up to the number of visible GPUs with --gpu."
+        ),
+    )
+    predict_parser.add_argument(
+        "--gpu",
+        action="store_true",
+        help=(
+            "Use CUDA for XGBoost prediction when available. With --jobs > 1, "
+            "models are distributed across persistent GPU workers."
+        ),
     )
     predict_parser.add_argument(
         "--block-overlap",
@@ -338,6 +349,7 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
             feature_importance_out=args.feature_importance_out,
             block_overlap=args.block_overlap,
             jobs=args.jobs,
+            gpu=args.gpu,
         )
         return 0
 

{plaknit-0.3.2 → plaknit-0.3.4}/plaknit/data/sampling.py

@@ -189,33 +189,63 @@ def _collect_pseudo_absence_candidate_pool(
     *,
     buffer_meters: float,
 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Collect valid background features and coordinates eligible for sampling."""
+    """Collect valid background features and coordinates eligible for sampling.
+
+    Raster data is read in bounded windows so large VRTs do not require a full
+    raster-sized float32 allocation during candidate collection.
+    """
 
     if buffer_meters < 0:
         raise ValueError("training buffer must be non-negative.")
 
-    full_window = windows.Window(0, 0, stack.width, stack.height)
-    data = stack.read(window=full_window, out_dtype="float32")
-    samples = data.reshape(stack.count, -1).T
-    valid = ~_nodata_pixel_mask(samples, stack.nodata_values)
-
     geometries = _buffer_training_geometries(stack, gdf, buffer_meters)
-    if geometries:
-        excluded = features.rasterize(
-            [(geom, 1) for geom in geometries],
-            out_shape=(stack.height, stack.width),
-            transform=stack.transform,
-            fill=0,
-            dtype="uint8",
-        ).reshape(-1)
-        valid &= excluded == 0
-
-    candidate_indices = np.flatnonzero(valid)
-    rows, cols = np.unravel_index(candidate_indices, (stack.height, stack.width))
+    feature_chunks: List[np.ndarray] = []
+    row_chunks: List[np.ndarray] = []
+    col_chunks: List[np.ndarray] = []
+    block_size = 512
+    for row_off in range(0, stack.height, block_size):
+        for col_off in range(0, stack.width, block_size):
+            win = windows.Window(
+                col_off=col_off,
+                row_off=row_off,
+                width=min(block_size, stack.width - col_off),
+                height=min(block_size, stack.height - row_off),
+            )
+            data = stack.read(window=win, out_dtype="float32")
+            samples = data.reshape(stack.count, -1).T
+            valid = ~_nodata_pixel_mask(samples, stack.nodata_values)
+
+            if geometries:
+                block_transform = windows.transform(win, stack.transform)
+                excluded = features.rasterize(
+                    [(geom, 1) for geom in geometries],
+                    out_shape=(int(win.height), int(win.width)),
+                    transform=block_transform,
+                    fill=0,
+                    dtype="uint8",
+                ).reshape(-1)
+                valid &= excluded == 0
+
+            candidate_indices = np.flatnonzero(valid)
+            if candidate_indices.size == 0:
+                continue
+            rows, cols = np.unravel_index(
+                candidate_indices, (int(win.height), int(win.width))
+            )
+            feature_chunks.append(samples[candidate_indices])
+            row_chunks.append((rows + int(win.row_off)).astype("int32", copy=False))
+            col_chunks.append((cols + int(win.col_off)).astype("int32", copy=False))
+
+    if not feature_chunks:
+        return (
+            np.empty((0, stack.count), dtype="float32"),
+            np.empty(0, dtype="int32"),
+            np.empty(0, dtype="int32"),
+        )
     return (
-        samples[candidate_indices],
-        rows.astype("int32", copy=False),
-        cols.astype("int32", copy=False),
+        np.vstack(feature_chunks),
+        np.concatenate(row_chunks),
+        np.concatenate(col_chunks),
     )
 
 

{plaknit-0.3.2 → plaknit-0.3.4}/plaknit/models/brt.py

@@ -9,6 +9,7 @@ import os
 import warnings
 from dataclasses import dataclass
 from pathlib import Path
+from functools import lru_cache
 from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Union
 
 import geopandas as gpd
@@ -151,12 +152,37 @@ def _prepare_brt_training_data(
 
 def _detect_gpu_available() -> bool:
     """Detect if GPU is available for XGBoost."""
+    return bool(_available_cuda_devices())
+
+
+@lru_cache(maxsize=1)
+def _available_cuda_devices() -> List[int]:
+    """Return logical CUDA device IDs visible to the current process."""
     try:
-        import cupy as cp  # noqa: F401
+        import cupy as cp
+
+        count = int(cp.cuda.runtime.getDeviceCount())
+    except (ImportError, RuntimeError, AttributeError):
+        return []
+    return list(range(max(0, count)))
+
 
-        return True
-    except ImportError:
-        return False
+def _configure_prediction_device(
+    model: xgb.XGBClassifier, gpu: bool, device_id: Optional[int] = None
+) -> None:
+    """Configure the loaded XGBoost model for CPU or CUDA inference."""
+    if not gpu:
+        return
+    if not _detect_gpu_available():
+        raise RuntimeError(
+            "GPU prediction requested, but CuPy/CUDA is unavailable. "
+            "Install the GPU dependencies or omit --gpu."
+        )
+    device = "cuda" if device_id is None else f"cuda:{device_id}"
+    try:
+        model.set_params(device=device, tree_method="hist", n_jobs=1)
+    except (AttributeError, TypeError, ValueError):
+        model.set_params(predictor="gpu_predictor", tree_method="gpu_hist", n_jobs=1)
 
 
 def train_brt(
@@ -658,6 +684,7 @@ def _init_predict_worker(
     out_dtype: str,
     nodata_value: Union[int, float],
     return_probs: bool,
+    gpu: bool,
 ) -> None:
     global _PREDICT_STACK
     global _PREDICT_MODEL
@@ -669,6 +696,7 @@ def _init_predict_worker(
     stack.__enter__()
     _PREDICT_STACK = stack
     _PREDICT_MODEL = joblib.load(model_path)
+    _configure_prediction_device(_PREDICT_MODEL, gpu)
     _PREDICT_OUT_DTYPE = out_dtype
     _PREDICT_NODATA = nodata_value
     _PREDICT_RETURN_PROBS = return_probs
@@ -706,6 +734,7 @@ def predict_brt(
     binary_out: Optional[PathLike] = None,
     block_overlap: int = 0,
     jobs: int = 1,
+    gpu: bool = False,
 ) -> Path:
     """Apply a trained BRT (XGBoost) model to a raster stack and write a classified GeoTIFF.
 
@@ -724,6 +753,7 @@ def predict_brt(
 
     _log("[bold cyan]Loading model...")
     model: xgb.XGBClassifier = joblib.load(model_path)
+    _configure_prediction_device(model, gpu)
     classes = getattr(model, "classes_", None)
     classes_dtype = getattr(classes, "dtype", np.int32)
     if np.issubdtype(classes_dtype, np.integer):
@@ -828,6 +858,7 @@ def predict_brt(
                         out_dtype,
                         nodata_value,
                         return_probs,
+                        gpu,
                     ),
                 ) as executor:
                     futures: Dict[concurrent.futures.Future, windows.Window] = {}