smartpool-examples 0.1.6__tar.gz → 0.1.7__tar.gz

{smartpool_examples-0.1.6/smartpool_examples.egg-info → smartpool_examples-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: smartpool-examples
-Version: 0.1.6
+Version: 0.1.7
 Summary: Examples for smartpool.
 Author-email: "王炳辉 (Bing-Hui WANG)" <binghui.wang@foxmail.com>
 License: MIT
@@ -11,7 +11,7 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: pysmartpool>=0.1.6
+Requires-Dist: pysmartpool>=0.1.7
 Requires-Dist: matplotlib
 Requires-Dist: scikit-learn
 Requires-Dist: numpy

{smartpool_examples-0.1.6 → smartpool_examples-0.1.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "smartpool-examples"
-version = "0.1.6"
+version = "0.1.7"
 description = "Examples for smartpool."
 readme = "README.md"
 authors = [
@@ -17,7 +17,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "pysmartpool>=0.1.6",
+    "pysmartpool>=0.1.7",
     "matplotlib",
     "scikit-learn",
     "numpy",

smartpool_examples-0.1.7/smartpool_examples/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.7"

{smartpool_examples-0.1.6 → smartpool_examples-0.1.7}/smartpool_examples/count_prime/__main__.py

@@ -8,12 +8,28 @@ if __name__ == "__main__":
     target_folder = os.path.abspath(self_folder + "/../../../smartpool").replace("\\", "/")
     sys.path.append(target_folder)
 
+import typer
 from count_prime import count_prime
 
-from smartpool import ProcessPool
+from smartpool import InterpreterPool, ProcessPool, ThreadPool
 
-if __name__ == "__main__":
-    print("Use ProcessPool to count prime numbers lower than 10000.")
+app = typer.Typer()
+
+POOL_MAP = {
+    "ThreadPool": ThreadPool,
+    "ProcessPool": ProcessPool,
+    "InterpreterPool": InterpreterPool,
+}
+
+
+@app.command()
+def main(pool: str = typer.Option("ProcessPool", "--pool", help="Pool type to use")):
+    pool_cls = POOL_MAP.get(pool)
+    if pool_cls is None:
+        print(f"Unknown pool: {pool}. Choose from {list(POOL_MAP.keys())}")
+        raise typer.Exit(code=1)
+
+    print(f"Use {pool} to count prime numbers lower than 10000.")
     print(f"See source code at folder {self_folder}")
 
     tasks = []
@@ -23,7 +39,7 @@ if __name__ == "__main__":
         tasks.append((start, stop))
         start = stop
 
-    with ProcessPool() as pool:
+    with pool_cls() as pool:
         futures = []
         for task in tasks:
             future = pool.submit(count_prime, args=task)
@@ -31,3 +47,7 @@ if __name__ == "__main__":
 
         total_primes_count = sum(future.result() for future in futures)
         print(total_primes_count)
+
+
+if __name__ == "__main__":
+    app()

{smartpool_examples-0.1.6 → smartpool_examples-0.1.7}/smartpool_examples/cross_validation/__main__.py

@@ -47,8 +47,8 @@ def main(
         help="max number of workers to use, 0 to use all available cores"
     )
 ):
-    import os
     import importlib
+    import os
     os.environ["RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO"] = "0"
 
     print(f"Use {pool} to do 5-fold cross validatation for 7 deep learning models for handwritten digit recognition task.")

smartpool_examples-0.1.7/smartpool_examples/onnx_infer/__main__.py

@@ -0,0 +1,158 @@
+if __name__ == "__main__":
+    import os
+    import sys
+    self_folder = os.path.dirname(os.path.abspath(__file__)).replace("\\", "/")
+    target_folder = os.path.abspath(self_folder + "/../../../smartpool").replace("\\", "/")
+    sys.path.append(target_folder)
+
+import typer
+
+app = typer.Typer(help="Use smartpool to do YOLOv8n ONNX inference on COCO val2017 with real-time progress.")
+
+@app.command()
+def main(
+    thread_pool_max_workers: int = typer.Option(
+        4,
+        "--thread_pool_max_workers",
+        help="max number of thread pool (for preprocess and postprocess) workers to use, 0 to use all available cores"
+    ),
+    session_pool_max_workers: int = typer.Option(
+        0,
+        "--session_pool_max_workers",
+        help="max number of infer session pool workers to use, 0 to use all available cores"
+    ),
+):
+    import os
+
+    print("Use `python -m smartpool_examples.onnx_infer --help` to see all options.")
+    print(f"See source code at folder {os.path.dirname(os.path.abspath(__file__))}")
+    print()
+
+    try:
+        import onnxruntime
+    except ImportError:
+        print("ONNX Runtime is not installed. Follow https://onnxruntime.ai/docs/install/ instructions to install ONNX Runtime.")
+        exit(1)
+
+    import queue
+    import time
+    from functools import partial
+
+    from rich.console import Group
+    from rich.live import Live
+    from rich.progress import (
+        BarColumn,
+        Progress,
+        TextColumn,
+        TimeRemainingColumn,
+    )
+    from rich.text import Text
+
+    from smartpool import DataSize, GPUInfo, InferSessionPool, Resource, ThreadPool
+
+    self_folder = os.path.dirname(os.path.abspath(__file__)).replace("\\", "/")
+    sys.path.append(self_folder)
+
+    from concurrent.futures import Future
+
+    from config import DATASET_DIR, MODEL_PATH, OUTPUT_DIR
+    from data_utils import download_dataset, download_model
+    from inference import postprocess, preprocess
+
+    download_model()
+    download_dataset()
+    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+    image_paths = sorted(DATASET_DIR.glob("*.jpg"))
+    model_path_str = str(MODEL_PATH.resolve())
+    output_dir_str = str(OUTPUT_DIR.resolve())
+
+    thread_pool = ThreadPool(max_workers=thread_pool_max_workers)
+    infer_session_pool = InferSessionPool(model_path_str, max_workers=session_pool_max_workers)
+
+    cpu_res = Resource(
+        cpu_cores_in_python=1,
+        cpu_cores_out_of_python=1,
+        cpu_mem=81*DataSize.MB
+    )
+    gpu_res = Resource(
+        cpu_cores_in_python=1,
+        cpu_cores_out_of_python=1,
+        cpu_mem=200*DataSize.MB,
+        gpu_cores=100,
+        gpu_mem=200*DataSize.MB
+    )
+
+    result_queue = queue.Queue()
+
+    with Progress(
+        TextColumn("[progress.description]{task.description}"),
+        BarColumn(),
+        TextColumn("{task.completed}/{task.total}"),
+        TimeRemainingColumn(),
+    ) as progress, Live(refresh_per_second=4) as live:
+        n_tasks = len(image_paths)
+        preprocess_submit_progress_task = progress.add_task("preprocess submit", total=n_tasks)
+        preprocess_done_progress_task = progress.add_task("preprocess done", total=n_tasks)
+        infer_submit_progress_task = progress.add_task("infer submit", total=n_tasks)
+        infer_done_progress_task = progress.add_task("infer done", total=n_tasks)
+        postprocess_submit_progress_task = progress.add_task("postprocess submit", total=n_tasks)
+        postprocess_done_progress_task = progress.add_task("postprocess done", total=n_tasks)
+
+        def update_task_count():
+            lines = []
+            for provider in GPUInfo.supported_onnx_providers():
+                lines.append(Text(f"task count infer with {provider}: {infer_session_pool.task_count_with_provider(provider)}"))
+            live.update(Group(*lines))
+
+        def postprocess_done_callback(postprocess_future: Future):
+            update_task_count()
+            postprocess_future.result()
+            progress.update(postprocess_done_progress_task, advance=1)
+            result_queue.put(1)
+
+        def infer_done_callback(image_path, img, scale, pad, infer_future: Future):
+            update_task_count()
+            outputs = infer_future.result()[0]
+            progress.update(infer_done_progress_task, advance=1)
+
+            output_path = output_dir_str + "/" + os.path.basename(image_path)
+            postprocess_future: Future = thread_pool.submit(
+                postprocess,
+                args=(img, outputs, output_path, scale, pad),
+                cpu_mode_res=cpu_res
+            )
+            progress.update(postprocess_submit_progress_task, advance=1)
+            postprocess_future.add_done_callback(postprocess_done_callback)
+
+        def preprocess_done_callback(image_path, preprocess_future: Future):
+            update_task_count()
+            progress.update(preprocess_done_progress_task, advance=1)
+            img, blob, scale, pad = preprocess_future.result()
+            infer_future: Future = infer_session_pool.submit(args=(blob,), cpu_mode_res=cpu_res, gpu_mode_res=gpu_res)
+            progress.update(infer_submit_progress_task, advance=1)
+            infer_future.add_done_callback(partial(infer_done_callback, image_path, img, scale, pad))
+
+        for image_path in image_paths:
+            update_task_count()
+            image_path = str(image_path)
+            preprocess_future: Future = thread_pool.submit(
+                preprocess,
+                args=(image_path,),
+                cpu_mode_res=cpu_res
+            )
+            progress.update(preprocess_submit_progress_task, advance=1)
+            preprocess_future.add_done_callback(partial(preprocess_done_callback, image_path))
+
+        start_time = time.perf_counter()
+
+        for _ in range(n_tasks):
+            result_queue.get()
+
+    elapsed = time.perf_counter() - start_time
+    print(f"\ninference completed in {elapsed:.2f} seconds")
+    print(f"Output: {OUTPUT_DIR.resolve()}")
+
+
+if __name__ == "__main__":
+    app()

{smartpool_examples-0.1.6 → smartpool_examples-0.1.7}/smartpool_examples/onnx_infer/config.py

@@ -8,12 +8,8 @@ OUTPUT_DIR = PACKAGE_DIR / "data" / "output"
 MODEL_URL = "https://hf-mirror.com/Kalray/yolov8/resolve/main/yolov8n.onnx"
 MODEL_PATH = MODEL_DIR / "yolov8n.onnx"
 
-COCO_IMAGE_URL = "http://images.cocodataset.org/val2017/{}.jpg"
-COCO_IMAGE_IDS = [
-    "000000000139", "000000000285", "000000000632", "000000000724", "000000000776",
-    "000000000785", "000000000802", "000000000872", "000000001000", "000000001268",
-    "000000001296", "000000001353", "000000001490", "000000001503", "000000001532",
-    "000000001584", "000000001675", "000000001818", "000000096001", "000000202001",
+COCO_ZIP_URLS = [
+    "http://images.cocodataset.org/zips/val2017.zip",
 ]
 
 COCO_CLASSES = [

{smartpool_examples-0.1.6 → smartpool_examples-0.1.7}/smartpool_examples/onnx_infer/data_utils.py

@@ -1,15 +1,11 @@
+import contextlib
 import time
 import urllib.error
 import urllib.request
+import zipfile
 from pathlib import Path
 
-from config import (
-    COCO_IMAGE_IDS,
-    COCO_IMAGE_URL,
-    DATASET_DIR,
-    MODEL_PATH,
-    MODEL_URL,
-)
+from config import COCO_ZIP_URLS, DATASET_DIR, MODEL_PATH, MODEL_URL
 from rich.progress import (
     BarColumn,
     DownloadColumn,
@@ -27,7 +23,7 @@ def download_file(url: str, dest: Path, desc: str = "", max_retries: int = 3):
     for attempt in range(1, max_retries + 1):
         try:
             req = urllib.request.Request(url, method="GET")
-            resp = urllib.request.urlopen(req, timeout=60)
+            resp = urllib.request.urlopen(req, timeout=120)
             total = int(resp.headers.get("Content-Length", 0))
             chunk_size = 8192
             with Progress(
@@ -48,6 +44,13 @@ def download_file(url: str, dest: Path, desc: str = "", max_retries: int = 3):
                         pbar.update(task, advance=len(chunk))
             tmp.rename(dest)
             return
+        except urllib.error.HTTPError as e:
+            if e.code == 404:
+                raise
+            last_err = e
+            if attempt < max_retries:
+                delay = 2 ** attempt
+                time.sleep(delay)
         except (urllib.error.URLError, TimeoutError, OSError) as e:
             last_err = e
             if attempt < max_retries:
@@ -56,23 +59,58 @@ def download_file(url: str, dest: Path, desc: str = "", max_retries: int = 3):
     raise last_err
 
 
+def _count_images():
+    return len(list(DATASET_DIR.rglob("*.jpg")))
+
+
+def _flatten_images():
+    for p in list(DATASET_DIR.rglob("*.jpg")):
+        if p.parent is not DATASET_DIR:
+            p.replace(DATASET_DIR / p.name)
+    for p in list(DATASET_DIR.iterdir()):
+        if p.is_dir() and p.name != "val2017":
+            for f in list(p.rglob("*")):
+                if f.is_file():
+                    f.unlink()
+
+            with contextlib.suppress(OSError):
+                p.rmdir()
+
+
+def download_dataset():
+    DATASET_DIR.mkdir(parents=True, exist_ok=True)
+
+    _flatten_images()
+    if _count_images() > 4000:
+        print(f"[SKIP DOWNLOAD] Dataset exists ({_count_images()} images)")
+        return
+
+    print("[DOWNLOAD] val2017.zip (~1GB) ...")
+    zip_path = DATASET_DIR / "val2017.zip"
+
+    for url in COCO_ZIP_URLS:
+        try:
+            download_file(url, zip_path, "COCO val2017")
+            break
+        except Exception:
+            print(f"  mirror failed: {url}")
+            continue
+    else:
+        raise RuntimeError("all mirrors failed")
+
+    print(f"[EXTRACT] val2017.zip -> {DATASET_DIR} ...")
+    with zipfile.ZipFile(zip_path) as zf:
+        zf.extractall(DATASET_DIR)
+
+    zip_path.unlink()
+    _flatten_images()
+    n = _count_images()
+    print(f"[DONE] {n} images ready")
+
+
 def download_model():
     if MODEL_PATH.exists():
         print(f"[SKIP DOWNLOAD] Model {MODEL_PATH.name} exists")
         return
     print("[DOWNLOAD] YOLOv8n ONNX...")
     download_file(MODEL_URL, MODEL_PATH, "YOLOv8n")
-
-
-def download_dataset():
-    DATASET_DIR.mkdir(parents=True, exist_ok=True)
-    existing = {p.stem for p in DATASET_DIR.glob("*.jpg")}
-    to_dl = [iid for iid in COCO_IMAGE_IDS if iid not in existing]
-    if not to_dl:
-        print(f"[SKIP DOWNLOAD] Dataset exists ({len(COCO_IMAGE_IDS)} images)")
-        return
-    print(f"[DOWNLOAD] {len(to_dl)} COCO images...")
-    for iid in to_dl:
-        url = COCO_IMAGE_URL.format(iid)
-        dest = DATASET_DIR / f"{iid}.jpg"
-        download_file(url, dest, f"COCO {iid}")

{smartpool_examples-0.1.6 → smartpool_examples-0.1.7}/smartpool_examples/onnx_infer/inference.py

@@ -1,130 +1,133 @@
-import os
-from typing import Dict, List, Tuple
-
-import cv2
-import numpy as np
-from config import COCO_CLASSES
-
-from smartpool import InferSessionPool, Resource
-
-
-def letterbox(img: np.ndarray, new_shape=(640, 640), color=(114, 114, 114)):
-    shape = img.shape[:2]
-    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
-    new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
-    dw = new_shape[1] - new_unpad[0]
-    dh = new_shape[0] - new_unpad[1]
-    dw, dh = dw / 2, dh / 2
-    top = int(round(dh - 0.1))
-    bottom = int(round(dh + 0.1))
-    left = int(round(dw - 0.1))
-    right = int(round(dw + 0.1))
-    img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
-    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
-    return img, r, (left, top)
-
-
-def preprocess(img: np.ndarray):
-    img, scale, pad = letterbox(img)
-    img = img.astype(np.float32) / 255.0
-    img = img.transpose(2, 0, 1)[np.newaxis, ...]
-    return img, scale, pad
-
-
-def nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float = 0.5):
-    if len(boxes) == 0:
-        return []
-    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
-    areas = (x2 - x1) * (y2 - y1)
-    order = scores.argsort()[::-1]
-    keep = []
-    while order.size > 0:
-        i = order[0]
-        keep.append(i)
-        xx1 = np.maximum(x1[i], x1[order[1:]])
-        yy1 = np.maximum(y1[i], y1[order[1:]])
-        xx2 = np.minimum(x2[i], x2[order[1:]])
-        yy2 = np.minimum(y2[i], y2[order[1:]])
-        w = np.maximum(0.0, xx2 - xx1)
-        h = np.maximum(0.0, yy2 - yy1)
-        inter = w * h
-        iou = inter / (areas[i] + areas[order[1:]] - inter)
-        inds = np.where(iou <= iou_threshold)[0]
-        order = order[inds + 1]
-    return keep
-
-
-def postprocess(output: np.ndarray, orig_shape: Tuple[int, int], scale: float,
-                pad: Tuple[int, int], conf_thresh: float = 0.5, iou_thresh: float = 0.5):
-    predictions = np.squeeze(output).T
-    box_data = predictions[:, :4]
-    class_scores = predictions[:, 4:]
-    scores = class_scores.max(axis=1)
-    class_ids = class_scores.argmax(axis=1)
-    mask = scores > conf_thresh
-    if not mask.any():
-        return []
-    boxes = box_data[mask]
-    scores = scores[mask]
-    class_ids = class_ids[mask]
-
-    boxes_xyxy = np.zeros_like(boxes)
-    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
-    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
-    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
-    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
-
-    keep = nms(boxes_xyxy, scores, iou_thresh)
-    if not keep:
-        return []
-
-    boxes_xyxy = boxes_xyxy[keep]
-    scores = scores[keep]
-    class_ids = class_ids[keep]
-
-    left, top = pad
-    boxes_xyxy = (boxes_xyxy - np.array([left, top, left, top])) / scale
-    boxes_xyxy = np.clip(boxes_xyxy, 0, [orig_shape[1], orig_shape[0], orig_shape[1], orig_shape[0]])
-
-    return [
-        {"bbox": boxes_xyxy[i].tolist(), "score": float(scores[i]), "class_id": int(class_ids[i])}
-        for i in range(len(boxes_xyxy))
-    ]
-
-
-def draw_detections(img: np.ndarray, detections: List[Dict]) -> np.ndarray:
-    colors = [
-        (56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255),
-        (49, 210, 207), (10, 249, 72), (23, 204, 146), (134, 219, 61),
-        (52, 147, 26), (187, 135, 1), (132, 145, 137), (168, 229, 11),
-    ]
-    for det in detections:
-        x1, y1, x2, y2 = map(int, det["bbox"])
-        cls_id = det["class_id"]
-        score = det["score"]
-        color = colors[cls_id % len(colors)]
-        label = f"{COCO_CLASSES[cls_id]} {score:.2f}"
-        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
-        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-        cv2.rectangle(img, (x1, y1 - th - 4), (x1 + tw, y1), color, -1)
-        cv2.putText(img, label, (x1, y1 - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
-    return img
-
-
-def infer_task(image_path: str, model_path_str: str, output_dir_str: str, infer_session_pool: InferSessionPool) -> str:
-    img = cv2.imread(image_path)
-    if img is None:
-        raise ValueError(f"Failed to read image: {image_path}")
-
-    h, w = img.shape[:2]
-    blob, scale, pad = preprocess(img)
-    gpu_res = Resource(cpu_cores_in_python=1, cpu_mem=512*1024**2,
-                       gpu_cores=1000, gpu_mem=2*1024**3)
-    cpu_res = Resource(cpu_cores_in_python=2, cpu_mem=2*1024**3)
-    future = infer_session_pool.submit(model_path_str, args=(blob,), cpu_mode_res=cpu_res, gpu_mode_res=gpu_res)
-    outputs = future.result()[0]
-    detections = postprocess(outputs, (h, w), scale, pad)
-    result_img = draw_detections(img, detections)
-
-    out_path = os.path.join(output_dir_str, os.path.basename(image_path))
-    cv2.imwrite(out_path, result_img)
+from __future__ import annotations
+
+import os
+from concurrent.futures import Future
+from typing import TYPE_CHECKING, Dict, List, Tuple
+
+import cv2
+import numpy as np
+from config import COCO_CLASSES
+
+if TYPE_CHECKING:
+    from smartpool import InferSessionPool, Resource
+
+
+def letterbox(img: np.ndarray, new_shape=(640, 640), color=(114, 114, 114)):
+    shape = img.shape[:2]
+    r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
+    new_unpad = (int(round(shape[1] * r)), int(round(shape[0] * r)))
+    dw = new_shape[1] - new_unpad[0]
+    dh = new_shape[0] - new_unpad[1]
+    dw, dh = dw / 2, dh / 2
+    top = int(round(dh - 0.1))
+    bottom = int(round(dh + 0.1))
+    left = int(round(dw - 0.1))
+    right = int(round(dw + 0.1))
+    img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
+    img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)
+    return img, r, (left, top)
+
+
+def preprocess(image_path: str):
+    img = cv2.imread(image_path)
+    if img is None:
+        raise ValueError(f"failed to read image: {image_path}")
+
+    blob, scale, pad = letterbox(img)
+    blob = blob.astype(np.float32) / 255.0
+    blob = np.ascontiguousarray(blob.transpose(2, 0, 1)[np.newaxis, ...])
+    return img, blob, scale, pad
+
+
+def nms(boxes: np.ndarray, scores: np.ndarray, iou_threshold: float = 0.5):
+    if len(boxes) == 0:
+        return []
+    x1, y1, x2, y2 = boxes[:, 0], boxes[:, 1], boxes[:, 2], boxes[:, 3]
+    areas = (x2 - x1) * (y2 - y1)
+    order = scores.argsort()[::-1]
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+        w = np.maximum(0.0, xx2 - xx1)
+        h = np.maximum(0.0, yy2 - yy1)
+        inter = w * h
+        iou = inter / (areas[i] + areas[order[1:]] - inter)
+        inds = np.where(iou <= iou_threshold)[0]
+        order = order[inds + 1]
+    return keep
+
+
+def draw_detections(src_img: np.ndarray, output_path: str, detections: List[Dict]):
+    colors = [
+        (56, 56, 255), (151, 157, 255), (31, 112, 255), (29, 178, 255),
+        (49, 210, 207), (10, 249, 72), (23, 204, 146), (134, 219, 61),
+        (52, 147, 26), (187, 135, 1), (132, 145, 137), (168, 229, 11),
+    ]
+    for det in detections:
+        x1, y1, x2, y2 = map(int, det["bbox"])
+        cls_id = det["class_id"]
+        score = det["score"]
+        color = colors[cls_id % len(colors)]
+        label = f"{COCO_CLASSES[cls_id]} {score:.2f}"
+        cv2.rectangle(src_img, (x1, y1), (x2, y2), color, 2)
+        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
+        cv2.rectangle(src_img, (x1, y1 - th - 4), (x1 + tw, y1), color, -1)
+        cv2.putText(src_img, label, (x1, y1 - 2), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
+
+    cv2.imwrite(output_path, src_img)
+
+
+def postprocess(
+    src_img:np.ndarray, output: np.ndarray, output_path:str,
+    scale: float, pad: Tuple[int, int], conf_thresh: float = 0.5,
+    iou_thresh: float = 0.5
+):
+    orig_shape = (src_img.shape[0], src_img.shape[1])
+    predictions = np.squeeze(output).T
+    box_data = predictions[:, :4]
+    class_scores = predictions[:, 4:]
+    scores = class_scores.max(axis=1)
+    class_ids = class_scores.argmax(axis=1)
+    mask = scores > conf_thresh
+    if not mask.any():
+        return []
+    boxes = box_data[mask]
+    scores = scores[mask]
+    class_ids = class_ids[mask]
+
+    boxes_xyxy = np.zeros_like(boxes)
+    boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2] / 2
+    boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3] / 2
+    boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2] / 2
+    boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3] / 2
+
+    keep = nms(boxes_xyxy, scores, iou_thresh)
+    if not keep:
+        return []
+
+    boxes_xyxy = boxes_xyxy[keep]
+    scores = scores[keep]
+    class_ids = class_ids[keep]
+
+    left, top = pad
+    boxes_xyxy = (boxes_xyxy - np.array([left, top, left, top])) / scale
+    boxes_xyxy = np.clip(boxes_xyxy, 0, [orig_shape[1], orig_shape[0], orig_shape[1], orig_shape[0]])
+
+    detections: List[Dict] = [
+        {"bbox": boxes_xyxy[i].tolist(), "score": float(scores[i]), "class_id": int(class_ids[i])}
+        for i in range(len(boxes_xyxy))
+    ]
+
+    draw_detections(src_img, output_path, detections)
+
+def infer_task(image_path:str, output_dir_path:str, infer_session_pool: InferSessionPool, cpu_res: Resource, gpu_res: Resource)->None:
+    img, blob, scale, pad = preprocess(image_path)
+    infer_future: Future = infer_session_pool.submit(args=(blob,), cpu_mode_res=cpu_res, gpu_mode_res=gpu_res)
+    outputs = infer_future.result()[0]
+    output_path = output_dir_path + "/" + os.path.basename(image_path)
+    postprocess(img, outputs, output_path, scale, pad)

{smartpool_examples-0.1.6 → smartpool_examples-0.1.7/smartpool_examples.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: smartpool-examples
-Version: 0.1.6
+Version: 0.1.7
 Summary: Examples for smartpool.
 Author-email: "王炳辉 (Bing-Hui WANG)" <binghui.wang@foxmail.com>
 License: MIT
@@ -11,7 +11,7 @@ Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: pysmartpool>=0.1.6
+Requires-Dist: pysmartpool>=0.1.7
 Requires-Dist: matplotlib
 Requires-Dist: scikit-learn
 Requires-Dist: numpy

{smartpool_examples-0.1.6 → smartpool_examples-0.1.7}/smartpool_examples.egg-info/requires.txt

@@ -1,4 +1,4 @@
-pysmartpool>=0.1.6
+pysmartpool>=0.1.7
 matplotlib
 scikit-learn
 numpy

smartpool_examples-0.1.6/smartpool_examples/onnx_infer/__main__.py

@@ -1,88 +0,0 @@
-if __name__ == "__main__":
-    import os
-    import sys
-    self_folder = os.path.dirname(os.path.abspath(__file__)).replace("\\", "/")
-    target_folder = os.path.abspath(self_folder + "/../../../smartpool").replace("\\", "/")
-    sys.path.append(target_folder)
-
-import os
-import time
-from concurrent.futures import as_completed
-
-import typer
-from rich.progress import (
-    BarColumn,
-    Progress,
-    TextColumn,
-    TimeRemainingColumn,
-)
-
-from smartpool import InferSessionPool, Resource, ThreadPool
-
-app = typer.Typer(help="Use smartpool to do YOLOv8n ONNX inference on COCO val2017 with real-time progress.")
-
-
-@app.command()
-def main(
-    max_workers: int = typer.Option(
-        0,
-        "--max_workers",
-        help="max number of workers to use, 0 to use all available cores"
-    ),
-):
-    self_folder = os.path.dirname(os.path.abspath(__file__)).replace("\\", "/")
-    sys.path.append(self_folder)
-
-    from config import DATASET_DIR, MODEL_PATH, OUTPUT_DIR
-    from data_utils import download_dataset, download_model
-    from inference import infer_task
-
-    print("Use `python -m smartpool_examples.onnx_infer --help` to see all options.")
-    print(f"See source code at folder {os.path.dirname(os.path.abspath(__file__))}")
-    print()
-
-    download_model()
-    download_dataset()
-    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-
-    image_paths = sorted(DATASET_DIR.glob("*.jpg"))
-    cpu_res = Resource(cpu_cores_in_python=1, cpu_mem=2*1024**3)
-
-    model_path_str = str(MODEL_PATH.resolve())
-    output_dir_str = str(OUTPUT_DIR.resolve())
-    n_workers = max_workers
-
-    print(f"\nSubmit {len(image_paths)} tasks to ThreadPool({n_workers})...")
-    thread_pool = ThreadPool(max_workers=n_workers)
-    infer_session_pool = InferSessionPool(max_workers=n_workers)
-    infer_session_pool.print_info = True
-
-    futures = []
-    for path in image_paths:
-        future = thread_pool.submit(
-            infer_task,
-            args=(str(path), model_path_str, output_dir_str, infer_session_pool),
-            cpu_mode_res=cpu_res
-        )
-        futures.append(future)
-
-    start_time = time.perf_counter()
-    with Progress(
-        TextColumn("[progress.description]{task.description}"),
-        BarColumn(),
-        TextColumn("{task.completed}/{task.total}"),
-        TimeRemainingColumn(),
-    ) as progress:
-        task = progress.add_task("infer", total=len(futures))
-        for f in as_completed(futures):
-            f.result()
-            progress.update(task, advance=1)
-            progress.update(task, description=f"infer [{progress.tasks[0].completed}/{progress.tasks[0].total}]")
-
-    elapsed = time.perf_counter() - start_time
-    print(f"\ninference completed in {elapsed:.2f} seconds")
-    print(f"Output: {OUTPUT_DIR.resolve()}")
-
-
-if __name__ == "__main__":
-    app()