From deac984180e54dcb904f415c8f2e095b8b1661a7 Mon Sep 17 00:00:00 2001
From: shuishen <1109946754@qq.com>
Date: Tue, 25 Aug 2026 14:56:08 +0800
Subject: [PATCH] feat(console): auto-select cuda for detection workflows
---
capabilities/00-change-detection/run_change_detection.py | 18 +++++++++++++-----
1 files changed, 13 insertions(+), 5 deletions(-)
diff --git a/capabilities/00-change-detection/run_change_detection.py b/capabilities/00-change-detection/run_change_detection.py
index dc16af6..9e9d73c 100644
--- a/capabilities/00-change-detection/run_change_detection.py
+++ b/capabilities/00-change-detection/run_change_detection.py
@@ -304,12 +304,14 @@
vector_path.write_text(json.dumps({"type": "FeatureCollection", "features": []}, indent=2), encoding="utf-8")
-def run_change_detection(before_path: Path, after_path: Path, output_dir: Path, *, processed_dir: Path | None = None, model_name: str = MODEL_NAME, threshold: float = DEFAULT_THRESHOLD, tile_size: int = DEFAULT_TILE_SIZE, overlap: int = DEFAULT_OVERLAP, max_dimension: int = AUTO_MAX_DIMENSION, processing_mode: str = DEFAULT_PROCESSING_MODE) -> dict[str, Any]:
+def run_change_detection(before_path: Path, after_path: Path, output_dir: Path, *, processed_dir: Path | None = None, model_name: str = MODEL_NAME, threshold: float = DEFAULT_THRESHOLD, tile_size: int = DEFAULT_TILE_SIZE, overlap: int = DEFAULT_OVERLAP, max_dimension: int = AUTO_MAX_DIMENSION, processing_mode: str = DEFAULT_PROCESSING_MODE, device: str = "auto") -> dict[str, Any]:
started = time.perf_counter()
if not math.isfinite(threshold) or not MIN_THRESHOLD <= threshold <= MAX_THRESHOLD:
raise ValueError(f"threshold must be between {MIN_THRESHOLD} and {MAX_THRESHOLD}")
if processing_mode not in PROCESSING_MODES:
raise ValueError(f"processing_mode must be one of {sorted(PROCESSING_MODES)}")
+ if device not in {"auto", "cpu", "cuda"}:
+ raise ValueError("device must be auto, cpu, or cuda")
if not isinstance(max_dimension, int) or max_dimension != AUTO_MAX_DIMENSION and not MIN_MAX_DIMENSION <= max_dimension <= MAX_MAX_DIMENSION:
raise ValueError(f"max_dimension must be 0 or between {MIN_MAX_DIMENSION} and {MAX_MAX_DIMENSION}")
before, before_info = _read_rgb(before_path, processing_mode)
@@ -362,8 +364,12 @@
Image.fromarray(after_small).save(output_dir / "after_registered_preview.jpg", quality=92)
from geoai import ChangeStarDetection
+ import torch
- detector = ChangeStarDetection(model_name=model_name, device="cpu")
+ if device == "cuda" and not torch.cuda.is_available():
+ raise RuntimeError("CUDA was requested but is unavailable.")
+ device_name = "cuda" if device == "cuda" or (device == "auto" and torch.cuda.is_available()) else "cpu"
+ detector = ChangeStarDetection(model_name=model_name, device=device_name)
result = detector.predict(str(work_dir / "before.tif"), str(work_dir / "after_registered.tif"), tile_size=tile_size, overlap=overlap, threshold=threshold)
probability = np.asarray(result["change_prob"], dtype=np.float32)
raw_mask = ((probability >= threshold) & valid_small).astype(np.uint8) * 255
@@ -402,7 +408,8 @@
"geoai_version": "0.42.0",
"method": "geoai.ChangeStarDetection + rasterio.features.shapes",
"model": model_name,
- "device": "cpu",
+ "requested_device": device,
+ "device": device_name,
"processing_mode": actual_mode,
"requested_processing_mode": processing_mode,
"thresholds": {"change_probability": threshold, "minimum_component_pixels": min_area},
@@ -460,7 +467,7 @@
def build_parser() -> argparse.ArgumentParser:
- parser = argparse.ArgumentParser(description="Run CPU ChangeStar change detection on a pair of images.")
+ parser = argparse.ArgumentParser(description="Run ChangeStar change detection on a pair of images.")
parser.add_argument("--before", type=Path, required=True)
parser.add_argument("--after", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
@@ -471,12 +478,13 @@
parser.add_argument("--overlap", type=int, default=DEFAULT_OVERLAP)
parser.add_argument("--max-dimension", type=int, default=AUTO_MAX_DIMENSION, help="Long-edge cap in pixels; 0 keeps a valid GeoTIFF at native resolution and uses 1024 for ordinary images.")
parser.add_argument("--processing-mode", choices=sorted(PROCESSING_MODES), default=DEFAULT_PROCESSING_MODE)
+ parser.add_argument("--device", choices={"auto", "cpu", "cuda"}, default="auto")
return parser
if __name__ == "__main__":
args = build_parser().parse_args()
try:
- print(json.dumps(run_change_detection(args.before, args.after, args.output, processed_dir=args.processed_output, model_name=args.model, threshold=args.threshold, tile_size=args.tile_size, overlap=args.overlap, max_dimension=args.max_dimension, processing_mode=args.processing_mode), ensure_ascii=False, indent=2))
+ print(json.dumps(run_change_detection(args.before, args.after, args.output, processed_dir=args.processed_output, model_name=args.model, threshold=args.threshold, tile_size=args.tile_size, overlap=args.overlap, max_dimension=args.max_dimension, processing_mode=args.processing_mode, device=args.device), ensure_ascii=False, indent=2))
except (FileNotFoundError, ValueError, RuntimeError) as exc:
raise SystemExit(f"变化检测失败: {exc}") from exc
--
Gitblit v1.9.3