feat(console): auto-select cuda for detection workflows

shuishen

7 days ago deac984180e54dcb904f415c8f2e095b8b1661a7

 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