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/01-object-detection/run_detection.py | 14 ++++++++++----
1 files changed, 10 insertions(+), 4 deletions(-)
diff --git a/capabilities/01-object-detection/run_detection.py b/capabilities/01-object-detection/run_detection.py
index e337b07..5852544 100644
--- a/capabilities/01-object-detection/run_detection.py
+++ b/capabilities/01-object-detection/run_detection.py
@@ -36,6 +36,7 @@
parser.add_argument("--image-size", type=int, default=1024, help="Model input size for each tile.")
parser.add_argument("--tile-size", type=int, default=1024, help="Pixel size of the sliding detection window.")
parser.add_argument("--tile-overlap", type=float, default=0.20, help="Overlap ratio between adjacent windows.")
+ parser.add_argument("--device", choices={"auto", "cpu", "cuda"}, default="auto")
return parser.parse_args()
@@ -116,7 +117,7 @@
return kept
-def predict_tiled(model: Any, image_path: Path, args: argparse.Namespace, np: Any) -> tuple[list[dict[str, Any]], int, int]:
+def predict_tiled(model: Any, image_path: Path, args: argparse.Namespace, np: Any, device: str) -> tuple[list[dict[str, Any]], int, int]:
with Image.open(image_path) as source:
image = source.convert("RGB")
width, height = image.size
@@ -127,7 +128,7 @@
tile = image.crop((x0, y0, min(x0 + args.tile_size, width), min(y0 + args.tile_size, height)))
result = model.predict(
source=np.asarray(tile),
- device="cpu",
+ device=device,
imgsz=args.image_size,
conf=args.confidence,
classes=COCO_TARGET_CLASS_IDS,
@@ -195,11 +196,15 @@
import ultralytics
from ultralytics import YOLO
+ if args.device == "cuda" and not torch.cuda.is_available():
+ raise SystemExit("CUDA was requested but is unavailable.")
+ device = "cuda:0" if args.device == "cuda" or (args.device == "auto" and torch.cuda.is_available()) else "cpu"
+
started = time.perf_counter()
model = YOLO(args.model)
detections: list[dict[str, Any]] = []
for image_path in image_paths:
- image_detections, width, height = predict_tiled(model, image_path, args, np)
+ image_detections, width, height = predict_tiled(model, image_path, args, np, device)
annotated_path = annotated_dir / image_path.name
draw_detections(image_path, image_detections, annotated_path)
detections.append(
@@ -221,7 +226,8 @@
"geoai_package": getattr(geoai, "__version__", "unknown"),
"ultralytics": ultralytics.__version__,
"torch": torch.__version__,
- "device": "cpu",
+ "requested_device": args.device,
+ "device": device,
"cuda_available": bool(torch.cuda.is_available()),
"model": args.model,
"confidence": args.confidence,
--
Gitblit v1.9.3