shuishen
6 days ago deac984180e54dcb904f415c8f2e095b8b1661a7
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,