diff --git a/.github/workflows/docker.jetson.6.2.0.yml b/.github/workflows/docker.jetson.6.2.0.yml index 218d784b9e..57ae23c2a4 100644 --- a/.github/workflows/docker.jetson.6.2.0.yml +++ b/.github/workflows/docker.jetson.6.2.0.yml @@ -52,11 +52,14 @@ jobs: base_image: ${{ env.BASE_IMAGE }} force_push: ${{ github.event.inputs.force_push }} token: ${{ secrets.GITHUB_TOKEN }} + - name: Fetch TensorRT local repository + run: docker/scripts/fetch_jetson_6_2_tensorrt.sh - name: Set up Depot CLI uses: depot/setup-action@v1 - name: Build and Push uses: depot/build-push-action@v1 with: + context: . push: ${{ github.event_name == 'release' || (github.event.inputs.force_push == 'true')}} project: 2rp7mfjw7q tags: ${{ steps.tags.outputs.image_tags }} diff --git a/.github/workflows/docker.jetson.7.1.0.yml b/.github/workflows/docker.jetson.7.1.0.yml deleted file mode 100644 index db4bc169ef..0000000000 --- a/.github/workflows/docker.jetson.7.1.0.yml +++ /dev/null @@ -1,61 +0,0 @@ -name: Build and Push Jetson 7.1.0 Container -permissions: - contents: read -on: - release: - types: [created] - push: - branches: [main] - workflow_dispatch: - inputs: - force_push: - type: boolean - description: "Do you want to push image after build?" - default: false - custom_tag: - type: string - description: "Custom tag to use for the image (overrides VERSION)" - default: "" - -env: - VERSION: "0.0.0" # Default version, will be overwritten - BASE_IMAGE: "roboflow/roboflow-inference-server-jetson-7.1.0" - -jobs: - docker: - runs-on: - labels: depot-ubuntu-24.04-4 - group: public-depot - timeout-minutes: 360 - permissions: - id-token: write - contents: read - steps: - - name: Login to Docker Hub - uses: docker/login-action@v4 - with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_TOKEN }} - - name: ๐Ÿ›Ž๏ธ Checkout - uses: actions/checkout@v6 - - name: Read version from file - run: echo "VERSION=$(DISABLE_VERSION_CHECK=true python ./inference/core/version.py)" >> $GITHUB_ENV - - name: Determine Image Tags - id: tags - uses: ./.github/actions/determine-tags - with: - custom_tag: ${{ github.event.inputs.custom_tag }} - version: ${{ env.VERSION }} - base_image: ${{ env.BASE_IMAGE }} - force_push: ${{ github.event.inputs.force_push }} - token: ${{ secrets.GITHUB_TOKEN }} - - name: Set up Depot CLI - uses: depot/setup-action@v1 - - name: Build and Push - uses: depot/build-push-action@v1 - with: - push: ${{ github.event_name == 'release' || (github.event.inputs.force_push == 'true')}} - project: v1xzfwkc4b - tags: ${{ steps.tags.outputs.image_tags }} - platforms: linux/arm64 - file: ./docker/dockerfiles/Dockerfile.onnx.jetson.7.1.0 diff --git a/.github/workflows/docker.jetson.7.2.0.yml b/.github/workflows/docker.jetson.7.2.0.yml index d219facba0..a96cdea8eb 100644 --- a/.github/workflows/docker.jetson.7.2.0.yml +++ b/.github/workflows/docker.jetson.7.2.0.yml @@ -26,7 +26,7 @@ jobs: runs-on: labels: depot-ubuntu-24.04-4 group: public-depot - timeout-minutes: 360 + timeout-minutes: 1440 permissions: id-token: write contents: read diff --git a/.github/workflows/docker.media.jetson.7.2.0.yml b/.github/workflows/docker.media.jetson.7.2.0.yml new file mode 100644 index 0000000000..c32ecbb257 --- /dev/null +++ b/.github/workflows/docker.media.jetson.7.2.0.yml @@ -0,0 +1,74 @@ +name: Build and Push Jetson 7.2.0 Media Container + +permissions: + contents: read + +on: + release: + types: [created] + push: + branches: [main] + paths: + - ".github/workflows/docker.media.jetson.7.2.0.yml" + - "docker/dockerfiles/Dockerfile.media.jetson.7.2.0" + - "docker/native/jetson_tensor_bridge/**" + - "docker/scripts/**" + - "LICENSE" + - "LICENSE.core" + workflow_dispatch: + inputs: + force_push: + type: boolean + description: "Do you want to push the image after building?" + default: false + custom_tag: + type: string + description: "Custom tag to use for the image" + default: "" + +env: + VERSION: "0.0.0" + BASE_IMAGE: "roboflow/roboflow-inference-media-jetson-7.2.0" + +jobs: + docker: + runs-on: + labels: depot-ubuntu-24.04-4 + group: public-depot + timeout-minutes: 180 + permissions: + id-token: write + contents: read + steps: + - name: Login to Docker Hub + if: ${{ github.event_name == 'release' || inputs.force_push }} + uses: docker/login-action@v4 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + - name: Checkout + uses: actions/checkout@v6 + - name: Read version from file + run: echo "VERSION=$(DISABLE_VERSION_CHECK=true python ./inference/core/version.py)" >> "$GITHUB_ENV" + - name: Determine image tags + id: tags + uses: ./.github/actions/determine-tags + with: + custom_tag: ${{ inputs.custom_tag }} + version: ${{ env.VERSION }} + base_image: ${{ env.BASE_IMAGE }} + force_push: ${{ inputs.force_push }} + token: ${{ secrets.GITHUB_TOKEN }} + - name: Set up Depot CLI + uses: depot/setup-action@v1 + - name: Build image + uses: depot/build-push-action@v1 + with: + project: v1xzfwkc4b + context: . + file: ./docker/dockerfiles/Dockerfile.media.jetson.7.2.0 + platforms: linux/arm64 + push: ${{ github.event_name == 'release' || inputs.force_push }} + save: ${{ github.event_name != 'release' && !inputs.force_push }} + save-tag: ${{ github.event_name != 'release' && !inputs.force_push && format('media-jp72-{0}', github.sha) || '' }} + tags: ${{ steps.tags.outputs.image_tags }} diff --git a/.github/workflows/docker.wheels.jetson.7.2.0.yml b/.github/workflows/docker.wheels.jetson.7.2.0.yml new file mode 100644 index 0000000000..7509960ac1 --- /dev/null +++ b/.github/workflows/docker.wheels.jetson.7.2.0.yml @@ -0,0 +1,108 @@ +name: Build and Push Jetson 7.2.0 Wheels Container + +permissions: + contents: read + +on: + release: + types: [created] + push: + branches: + - main + # Bootstrap trigger while the MVP branch is in flight; remove after merge. + - mvp/new-inference-pipeline + paths: + - ".github/workflows/docker.wheels.jetson.7.2.0.yml" + - "docker/dockerfiles/Dockerfile.wheels.jetson.7.2.0" + - "docker/docker-bake.jetson.7.2.0.hcl" + - "docker/patches/**" + workflow_dispatch: + inputs: + force_push: + type: boolean + description: "Do you want to push the image after building?" + default: false + custom_tag: + type: string + description: "Custom tag to use for the image" + default: "" + +env: + VERSION: "0.0.0" + BASE_IMAGE: "roboflow/roboflow-inference-wheels-jetson-7.2.0" + +jobs: + # The expensive compiles fan out to separate Depot builders; each populates the + # shared project cache, and the assemble job below cache-hits all of them. + components: + strategy: + fail-fast: false + matrix: + include: + - name: torch-stack + targets: jetson-wheels-torch-stack-jp72 + - name: ort-stack + targets: jetson-wheels-ort-stack-jp72 + name: Build ${{ matrix.name }} + runs-on: + labels: depot-ubuntu-24.04-4 + group: public-depot + timeout-minutes: 1440 + permissions: + id-token: write + contents: read + steps: + - name: Checkout + uses: actions/checkout@v6 + - name: Set up Depot CLI + uses: depot/setup-action@v1 + - name: Build component + uses: depot/bake-action@v1 + with: + files: ./docker/docker-bake.jetson.7.2.0.hcl + targets: ${{ matrix.targets }} + project: v1xzfwkc4b + + assemble: + name: Assemble wheels image + needs: components + runs-on: + labels: depot-ubuntu-24.04-4 + group: public-depot + timeout-minutes: 180 + permissions: + id-token: write + contents: read + steps: + - name: Login to Docker Hub + if: ${{ github.event_name == 'release' || inputs.force_push }} + uses: docker/login-action@v4 + with: + username: ${{ secrets.DOCKERHUB_USERNAME }} + password: ${{ secrets.DOCKERHUB_TOKEN }} + - name: Checkout + uses: actions/checkout@v6 + - name: Read version from file + run: echo "VERSION=$(DISABLE_VERSION_CHECK=true python ./inference/core/version.py)" >> "$GITHUB_ENV" + - name: Determine image tags + id: tags + uses: ./.github/actions/determine-tags + with: + custom_tag: ${{ inputs.custom_tag }} + version: ${{ env.VERSION }} + base_image: ${{ env.BASE_IMAGE }} + force_push: ${{ inputs.force_push }} + token: ${{ secrets.GITHUB_TOKEN }} + - name: Set up Depot CLI + uses: depot/setup-action@v1 + - name: Build image + uses: depot/bake-action@v1 + env: + JETSON_WHEELS_TAGS: ${{ steps.tags.outputs.image_tags }} + with: + files: ./docker/docker-bake.jetson.7.2.0.hcl + targets: jetson-wheels-jp72 + project: v1xzfwkc4b + push: ${{ github.event_name == 'release' || inputs.force_push }} + save: ${{ github.event_name != 'release' && !inputs.force_push }} + save-tag: ${{ github.event_name != 'release' && !inputs.force_push && format('wheels-jp72-{0}', github.sha) || '' }} diff --git a/.github/workflows/integration_tests_workflows_tensor_native_x86.yml b/.github/workflows/integration_tests_workflows_tensor_native_x86.yml new file mode 100644 index 0000000000..256e5fdb64 --- /dev/null +++ b/.github/workflows/integration_tests_workflows_tensor_native_x86.yml @@ -0,0 +1,64 @@ +name: INTEGRATION TESTS - workflows (tensor-native) +permissions: + contents: read +on: + pull_request: + branches: [main] + push: + branches: [main] + workflow_dispatch: + inputs: + python-version: + description: "Python version to run (or all)" + type: choice + default: all + options: + - all + - "3.10" + - "3.11" + - "3.12" + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + call_is_mergeable: + uses: ./.github/workflows/check_if_branch_is_mergeable.yml + secrets: inherit + build-dev-test: + needs: call_is_mergeable + if: ${{ (github.event.pull_request.head.repo.fork == false) && (github.event_name != 'pull_request' || needs.call_is_mergeable.outputs.mergeable_state != 'not_clean') }} + runs-on: + labels: depot-ubuntu-22.04-16 + group: public-depot + strategy: + matrix: + python-version: ${{ fromJSON((github.event.inputs.python-version || 'all') == 'all' && '["3.10","3.11","3.12"]' || format('["{0}"]', github.event.inputs.python-version)) }} + timeout-minutes: 15 + steps: + - name: ๐Ÿ›Ž๏ธ Checkout + uses: actions/checkout@v6 + - name: ๐Ÿ Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + check-latest: true + - name: ๐Ÿ“ฆ Cache Python packages + uses: actions/cache@v5 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('requirements/**') }} + restore-keys: | + ${{ runner.os }}-pip-${{ matrix.python-version }}- + - name: ๐Ÿšง Install GDAL + run: sudo apt-get update && sudo apt-get install -y libgdal-dev + - name: ๐Ÿ“ฆ Install dependencies + run: | + python -m pip install --upgrade pip + pip install --upgrade setuptools + pip install --extra-index-url https://download.pytorch.org/whl/cpu -r requirements/_requirements.txt -r requirements/requirements.sam.txt -r requirements/requirements.cpu.txt -r requirements/requirements.http.txt -r requirements/requirements.test.unit.txt -r requirements/requirements.doctr.txt -r requirements/requirements.yolo_world.txt -r requirements/requirements.transformers.txt -r requirements/requirements.sdk.http.txt -r requirements/requirements.easyocr.txt + - name: ๐Ÿงช Integration Tests of Workflows (tensor-native) + # Tensor-native path only: ENABLE_TENSOR_DATA_REPRESENTATION is AND-gated with + # USE_INFERENCE_MODELS, which defaults to True on Linux runners, so it is not set here. + run: ENABLE_TENSOR_DATA_REPRESENTATION=True ROBOFLOW_API_KEY=${{ secrets.API_KEY }} DEEP_LAB_V3_API_KEY=${{ secrets.DEEP_LAB_V3_API_KEY }} MODAL_TOKEN_ID=${{ secrets.MODAL_TOKEN_ID }} MODAL_TOKEN_SECRET=${{ secrets.MODAL_TOKEN_SECRET }} MODAL_ALLOW_ANONYMOUS_EXECUTION=True MODAL_ANONYMOUS_WORKSPACE_NAME=github-test MODAL_SKIP_FLORENCE2_TEST=FALSE LOAD_ENTERPRISE_BLOCKS=TRUE python -m pytest tests/workflows/integration_tests diff --git a/.github/workflows/test.nvidia_t4.yml b/.github/workflows/test.nvidia_t4.yml index 6c1e08e4c2..9d85250de3 100644 --- a/.github/workflows/test.nvidia_t4.yml +++ b/.github/workflows/test.nvidia_t4.yml @@ -56,7 +56,35 @@ jobs: - name: ๏ฟฝ๐Ÿ”จ Build and Push Test Docker - GPU run: | docker build -t roboflow/roboflow-inference-server-gpu:test -f docker/dockerfiles/Dockerfile.onnx.gpu . - + + - name: Verify CUDA OpenCV - GPU + if: ${{ github.event.inputs.test_name == '' }} + run: | + docker run --rm --gpus all --entrypoint sh \ + -e OPENCV_REQUIRE_CUDA_BUILD=true \ + -e OPENCV_REQUIRE_CUDA_RUNTIME=true \ + -v "$PWD/docker/scripts/verify_opencv.sh:/tmp/verify_opencv:ro" \ + roboflow/roboflow-inference-server-gpu:test \ + /tmp/verify_opencv + + - name: Verify NVIDIA GStreamer codecs - GPU + if: ${{ github.event.inputs.test_name == '' }} + run: | + docker run --rm --gpus all --entrypoint sh \ + -e GSTREAMER_REQUIRE_NVCODEC=true \ + -e GSTREAMER_REQUIRE_NVCODEC_RUNTIME=true \ + -v "$PWD/docker/scripts/verify_gstreamer.sh:/tmp/verify_gstreamer:ro" \ + roboflow/roboflow-inference-server-gpu:test \ + /tmp/verify_gstreamer + + - name: Verify CUDA GStreamer tensor bridge - GPU + if: ${{ github.event.inputs.test_name == '' }} + run: | + docker run --rm --gpus all --entrypoint python3 \ + -v "$PWD/docker/scripts/verify_gstreamer_cuda_tensor_runtime.py:/tmp/verify_gstreamer_cuda_tensor_runtime.py:ro" \ + roboflow/roboflow-inference-server-gpu:test \ + /tmp/verify_gstreamer_cuda_tensor_runtime.py + - name: ๐Ÿ”‹ Start Test Docker without Torch Preprocessing - GPU if: ${{ github.event.inputs.test_name == '' || github.event.inputs.test_name == 'regression_without_torch' }} run: | diff --git a/.github/workflows/unit_tests_workflows_tensor_native_x86.yml b/.github/workflows/unit_tests_workflows_tensor_native_x86.yml new file mode 100644 index 0000000000..7d7b7e0ad9 --- /dev/null +++ b/.github/workflows/unit_tests_workflows_tensor_native_x86.yml @@ -0,0 +1,55 @@ +name: UNIT TESTS - workflows (tensor-native) +permissions: + contents: read + +on: + pull_request: + branches: [main] + push: + branches: [main] + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + call_is_mergeable: + uses: ./.github/workflows/check_if_branch_is_mergeable.yml + secrets: inherit + build-dev-test: + needs: call_is_mergeable + if: ${{ github.event_name != 'pull_request' || needs.call_is_mergeable.outputs.mergeable_state != 'not_clean' }} + runs-on: + labels: depot-ubuntu-22.04-4 + group: public-depot + timeout-minutes: 10 + strategy: + matrix: + python-version: ["3.10", "3.11", "3.12"] + steps: + - name: ๐Ÿ›Ž๏ธ Checkout + uses: actions/checkout@v6 + - name: ๐Ÿ“ฆ Cache Python packages + uses: actions/cache@v5 + with: + path: ~/.cache/pip + key: ${{ runner.os }}-pip-${{ matrix.python-version }}-${{ hashFiles('requirements/**') }} + restore-keys: | + ${{ runner.os }}-pip-${{ matrix.python-version }}- + - name: ๐Ÿ Set up Python ${{ matrix.python-version }} + uses: actions/setup-python@v6 + with: + python-version: ${{ matrix.python-version }} + check-latest: true + + - name: ๐Ÿ“ฆ Install dependencies + run: | + python -m pip install --upgrade pip + pip install --upgrade setuptools + pip install -r requirements/_requirements.txt -r requirements/requirements.cpu.txt -r requirements/requirements.easyocr.txt -r requirements/requirements.sdk.http.txt -r requirements/requirements.test.unit.txt -r requirements/requirements.http.txt -r requirements/requirements.transformers.txt + - name: ๐Ÿงช Unit Tests of Workflows (tensor-native) + # ENABLE_TENSOR_DATA_REPRESENTATION is AND-gated with USE_INFERENCE_MODELS, which + # defaults to True on Linux runners, so the tensor-native code path is exercised. + timeout-minutes: 30 + run: ENABLE_TENSOR_DATA_REPRESENTATION=True python -m pytest tests/workflows/unit_tests diff --git a/.gitignore b/.gitignore index 45a2429198..90078ac6b9 100644 --- a/.gitignore +++ b/.gitignore @@ -20,6 +20,7 @@ build/ develop-eggs/ dist/ downloads/ +docker/vendor/tensorrt/ eggs/ .eggs/ lib64/ diff --git a/development/stream_interface/benchmark_engine_throughput.py b/development/stream_interface/benchmark_engine_throughput.py new file mode 100644 index 0000000000..f7b1dbdeb3 --- /dev/null +++ b/development/stream_interface/benchmark_engine_throughput.py @@ -0,0 +1,252 @@ +#!/usr/bin/env python3 +"""Micro-benchmark that isolates Execution Engine throughput from the per-frame image +decode + numpy->GPU conversion. + +It decodes `--decoded-frames` frames from a video ONCE, up front, into the format the +active data representation wants: + * ENABLE_TENSOR_DATA_REPRESENTATION=True -> CHW RGB uint8 tensors on the GPU + (WORKFLOWS_IMAGE_TENSOR_DEVICE), i.e. exactly what WorkflowImageData.tensor_image + builds โ€” so the conversion is NOT paid inside the loop; + * ENABLE_TENSOR_DATA_REPRESENTATION=False -> numpy BGR HWC frames. + +Then it runs a single ExecutionEngine instance `--engine-runs` times in a loop, cycling +through the pre-decoded frames by modulo. The measured FPS therefore reflects EE + model +throughput only, with the per-frame image materialization removed from the hot path. + +Run it twice on the same clip to A/B the data path: + ENABLE_TENSOR_DATA_REPRESENTATION=True python ...benchmark_engine_throughput.py ... + ENABLE_TENSOR_DATA_REPRESENTATION=False python ...benchmark_engine_throughput.py ... +If the tensor run now matches (or beats) the numpy run here, while the live pipeline +showed tensor at half the FPS, the gap is the per-frame numpy->GPU conversion. +""" +import argparse +import os +import sys +import time +from typing import List + +import cv2 +import numpy as np +import torch + +# The top-level `inference` package is not pip-installed in dev; resolve it from the repo +# root so this script runs from any working dir. +_REPO_ROOT = os.path.dirname( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +) +if _REPO_ROOT not in sys.path: + sys.path.insert(0, _REPO_ROOT) + +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + MAX_ACTIVE_MODELS, + WORKFLOWS_IMAGE_TENSOR_DEVICE, +) +from inference.core.managers.base import ModelManager +from inference.core.managers.decorators.fixed_size_cache import WithFixedSizeCache +from inference.core.registries.roboflow import RoboflowModelRegistry +from inference.core.roboflow_api import get_workflow_specification +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference.models.utils import ROBOFLOW_MODEL_TYPES + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark Execution Engine throughput over pre-decoded frames.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument("--video", required=True, help="Path to a local video file.") + parser.add_argument( + "--decoded-frames", + type=int, + default=64, + help="Number of frames to decode up front (tensor mode: onto the GPU).", + ) + parser.add_argument( + "--engine-runs", + type=int, + default=1000, + help="Number of timed ExecutionEngine.run() calls (frames cycled by modulo).", + ) + parser.add_argument( + "--batch-size", + type=int, + default=1, + help="Frames fed per ExecutionEngine.run() call (a batch of images).", + ) + parser.add_argument( + "--warmup", + type=int, + default=5, + help="Untimed warm-up runs (model load / TRT build / cache warm-up).", + ) + parser.add_argument( + "--workspace", required=True, help="Roboflow workspace name (the URL slug)." + ) + parser.add_argument( + "--workflow-id", + required=True, + help="Registered workflow id (fetched from the Roboflow platform).", + ) + parser.add_argument( + "--workflow-version-id", + default=None, + help="Optional specific workflow version id.", + ) + parser.add_argument( + "--model-id", + default=None, + help="Optional value for the workflow's `model_id` parameter.", + ) + parser.add_argument( + "--api-key", + default=os.environ.get("ROBOFLOW_API_KEY") or os.environ.get("API_KEY"), + help="Roboflow API key (defaults to the ROBOFLOW_API_KEY / API_KEY env var).", + ) + parser.add_argument( + "--image-input-name", default="image", help="Workflow image input name." + ) + return parser.parse_args() + + +def build_model_manager() -> ModelManager: + # Same construction as tests/workflows/integration_tests/conftest.py::model_manager. + registry = RoboflowModelRegistry(ROBOFLOW_MODEL_TYPES) + return WithFixedSizeCache(ModelManager(model_registry=registry), max_size=MAX_ACTIVE_MODELS) + + +def decode_frames(video_path: str, count: int) -> List[np.ndarray]: + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise SystemExit(f"Could not open video: {video_path}") + frames: List[np.ndarray] = [] + while len(frames) < count: + ok, frame = cap.read() + if not ok: + break + frames.append(frame) # numpy HWC BGR (uint8) + cap.release() + if not frames: + raise SystemExit("Decoded 0 frames โ€” empty or unreadable video.") + return frames + + +def to_gpu_tensor(bgr_hwc: np.ndarray, device: torch.device) -> torch.Tensor: + # Exactly what WorkflowImageData.tensor_image builds: CHW RGB uint8 on the device. + rgb = bgr_hwc[:, :, ::-1].copy() + return torch.from_numpy(rgb).permute(2, 0, 1).contiguous().to(device) + + +def _cuda_sync_if_needed() -> None: + if ENABLE_TENSOR_DATA_REPRESENTATION and WORKFLOWS_IMAGE_TENSOR_DEVICE.type == "cuda": + torch.cuda.synchronize() + + +def main() -> None: + args = parse_args() + if not os.path.isfile(args.video): + raise SystemExit(f"Video file not found: {args.video}") + if not args.api_key: + raise SystemExit( + "A Roboflow API key is required to fetch the workflow from the platform. " + "Pass --api-key or set ROBOFLOW_API_KEY." + ) + + # Fetch the workflow definition from the Roboflow platform (by workspace + id). + workflow = get_workflow_specification( + api_key=args.api_key, + workspace_id=args.workspace, + workflow_id=args.workflow_id, + workflow_version_id=args.workflow_version_id, + use_cache=True, + ) + + print( + f"torch.cuda.is_available() = {torch.cuda.is_available()} | " + f"WORKFLOWS_IMAGE_TENSOR_DEVICE = {WORKFLOWS_IMAGE_TENSOR_DEVICE} | " + f"ENABLE_TENSOR_DATA_REPRESENTATION = {ENABLE_TENSOR_DATA_REPRESENTATION}" + ) + if ENABLE_TENSOR_DATA_REPRESENTATION and WORKFLOWS_IMAGE_TENSOR_DEVICE.type != "cuda": + print( + "WARNING: tensor mode is on but the tensor device is not CUDA โ€” frames live on " + f"'{WORKFLOWS_IMAGE_TENSOR_DEVICE}', so the model will still copy to GPU itself." + ) + + raw = decode_frames(args.video, args.decoded_frames) + + # Pre-build the runtime image inputs ONCE, in the active representation's format, so + # the per-frame decode/conversion is out of the timed loop. + if ENABLE_TENSOR_DATA_REPRESENTATION: + device = WORKFLOWS_IMAGE_TENSOR_DEVICE + images = [to_gpu_tensor(f, device) for f in raw] + _cuda_sync_if_needed() # make sure the H2D copies are done before timing + print( + f"pre-decoded {len(images)} frames -> GPU tensors on {device}, " + f"shape {tuple(images[0].shape)} dtype {images[0].dtype}" + ) + else: + images = raw # numpy HWC BGR โ€” the legacy path + print( + f"pre-decoded {len(images)} frames -> numpy, " + f"shape {images[0].shape} dtype {images[0].dtype}" + ) + + model_manager = build_model_manager() + engine = ExecutionEngine.init( + workflow_definition=workflow, + init_parameters={ + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": args.api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + }, + workflow_id=args.workflow_id, + ) + + extra_params = {"model_id": args.model_id} if args.model_id else {} + n = len(images) + batch_size = max(1, args.batch_size) + if batch_size > n: + print( + f"NOTE: batch size {batch_size} > pre-decoded frames {n}; frames repeat within a batch." + ) + + def run_once(i: int): + # A batch is `batch_size` frames pulled from the pre-decoded pool (modulo). + batch = [images[(i * batch_size + j) % n] for j in range(batch_size)] + return engine.run( + runtime_parameters={args.image_input_name: batch, **extra_params} + ) + + print(f"warming up ({args.warmup} runs โ€” model load / TRT build)...", flush=True) + for i in range(args.warmup): + run_once(i) + _cuda_sync_if_needed() + + print( + f"running {args.engine_runs} EE iterations (batch size {batch_size})...", + flush=True, + ) + start = time.monotonic() + for i in range(args.engine_runs): + run_once(i) + _cuda_sync_if_needed() # flush async GPU work before stopping the clock + elapsed = time.monotonic() - start + + total_frames = args.engine_runs * batch_size + runs_per_s = args.engine_runs / elapsed if elapsed > 0 else 0.0 + frames_per_s = total_frames / elapsed if elapsed > 0 else 0.0 + print("\n=== Summary ===") + print(f"data representation : {'TENSOR' if ENABLE_TENSOR_DATA_REPRESENTATION else 'NUMPY'}") + print(f"pre-decoded frames : {n}") + print(f"batch size : {batch_size}") + print(f"engine runs (timed) : {args.engine_runs}") + print(f"frames processed : {total_frames}") + print(f"elapsed : {elapsed:.3f} s") + print(f"throughput : {frames_per_s:.1f} frames/s ({runs_per_s:.1f} runs/s)") + print(f"per-run latency : {elapsed / args.engine_runs * 1000:.3f} ms") + print(f"per-frame latency : {elapsed / total_frames * 1000:.3f} ms") + + +if __name__ == "__main__": + main() diff --git a/development/stream_interface/run_workflow_on_video.py b/development/stream_interface/run_workflow_on_video.py new file mode 100644 index 0000000000..ba6a7507cb --- /dev/null +++ b/development/stream_interface/run_workflow_on_video.py @@ -0,0 +1,580 @@ +#!/usr/bin/env python3 +"""Run an InferencePipeline against a registered Roboflow workflow over one or more +video sources (local files or live streams) and report the processing FPS (throughput). + +Every frame of every source is streamed through the workflow as fast as the workflow +allows (no max-fps cap by default), and FPS is counted synchronously in the prediction +sink โ€” both a rolling FPS (supervision's FPSMonitor) and a cumulative average. + +The workflow is fetched from the Roboflow platform by workspace + id, so an API key is +required (via --api-key or the ROBOFLOW_API_KEY / API_KEY environment variable). + +Sources: `--video` accepts a local file path or a stream reference (rtsp://..., +/dev/video0, ...) and may be repeated for multiple concurrent sources. `--n` instead +replicates a SINGLE --video into N identical streams โ€” the two forms are mutually +exclusive. Live streams run until Ctrl+C. + +Examples: + # minimal + python development/stream_interface/run_workflow_on_video.py \\ + --workspace my-workspace --workflow-id my-workflow --video /path/clip.mp4 + + # the same file decoded by 4 concurrent streams (multi-camera throughput test) + python development/stream_interface/run_workflow_on_video.py \\ + --workspace my-workspace --workflow-id my-workflow --video /path/clip.mp4 --n 4 + + # two different live sources + python development/stream_interface/run_workflow_on_video.py \\ + --workspace my-workspace --workflow-id my-workflow \\ + --video rtsp://camera-1/stream --video rtsp://camera-2/stream + + # steer an optional `model_id` workflow parameter from the CLI + python development/stream_interface/run_workflow_on_video.py \\ + --workspace my-workspace --workflow-id my-workflow --video /path/clip.mp4 \\ + --model-id yolov8n-640 +""" + +import argparse +import json +import os +import time +from collections import deque +from typing import Any, Dict, List, Optional + +import cv2 +import numpy as np +import supervision as sv + +# Modules whose module-level globals we patch to wire in the workflow profiler. +# Both functions read these names as module globals at call time, so patching the +# attribute after import is enough to steer the behaviour. +import inference.core.interfaces.stream.inference_pipeline as _ip_module +import inference.core.interfaces.stream.utils as _ip_utils +from inference import InferencePipeline +from inference.core.interfaces.stream.watchdog import BasePipelineWatchDog +from inference.core.utils.drawing import create_tiles +from inference.core.workflows.execution_engine.profiling.core import ( + BaseWorkflowsProfiler, +) + + +def _write_trace(path: str, trace: List[dict]) -> None: + os.makedirs(os.path.dirname(os.path.abspath(path)) or ".", exist_ok=True) + with open(path, "w") as f: + json.dump(trace, f) + + +class _FirstNRunsProfiler(BaseWorkflowsProfiler): + """Captures only the FIRST `max_runs_in_buffer` workflow runs (~ first N frames), + then freezes: every later event is dropped, so the trace stays small and stable + no matter how long the pipeline keeps running. + + Unlike the stock profiler (a rolling `deque` that keeps the *last* N runs), this + keeps the *first* N. It is flushed to disk the moment it freezes โ€” so a later hard + kill still keeps the trace โ€” and again on pipeline end via `on_pipeline_end`. + """ + + # Set by `_enable_workflow_profiler` before the pipeline starts. + _trace_path: Optional[str] = None + + @classmethod + def init(cls, max_runs_in_buffer: int = 50, **kwargs) -> "_FirstNRunsProfiler": + inst = cls(runs_buffer=deque()) # unbounded; we cap ourselves by freezing + inst._max_runs = max_runs_in_buffer + inst._frozen = False + return inst + + def start_workflow_run(self) -> None: + if self._frozen: + return + super().start_workflow_run() + + def end_workflow_run(self) -> None: + if self._frozen: + return + super().end_workflow_run() + if len(self._runs_buffer) >= self._max_runs: + self._frozen = True + self._current_run_events = [] + if self._trace_path is not None: + _write_trace(self._trace_path, self.export_trace()) + print( + f"[profiler] captured first {self._max_runs} frame(s) -> " + f"{self._trace_path}", + flush=True, + ) + + def _add_event(self, *args, **kwargs) -> None: + if self._frozen: + return + super()._add_event(*args, **kwargs) + + +def _enable_workflow_profiler(trace_path: str, max_frames: int) -> None: + """Turn on the workflow profiler so `init_with_workflow` builds our first-N-runs + profiler and dumps the Chrome trace to the exact `trace_path` we ask for.""" + _FirstNRunsProfiler._trace_path = trace_path + # `init_with_workflow` reads these two as module globals at call time. + _ip_module.ENABLE_WORKFLOWS_PROFILING = True + _ip_module.WORKFLOWS_PROFILER_BUFFER_SIZE = max_frames + _ip_module.BaseWorkflowsProfiler = _FirstNRunsProfiler + # `on_pipeline_end` (in utils) gates on this flag and calls this save function on + # shutdown โ€” point it at our exact path instead of the timestamped default name. + _ip_utils.ENABLE_WORKFLOWS_PROFILING = True + _ip_utils.save_workflows_profiler_trace = ( + lambda directory, profiler_trace: _write_trace(trace_path, profiler_trace) + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run a registered Roboflow workflow over video files / streams and count FPS.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--workspace", required=True, help="Roboflow workspace name (the URL slug)." + ) + parser.add_argument("--workflow-id", required=True, help="Registered workflow id.") + parser.add_argument( + "--video", + required=True, + action="append", + help="Video source: a local file path or a stream reference (rtsp://..., " + "/dev/video0, ...). Repeat the flag for multiple concurrent sources " + "(mutually exclusive with --n).", + ) + parser.add_argument( + "-n", + "--n", + type=int, + default=None, + dest="n", + help="Number of concurrent streams to run, all decoding the SAME single " + "--video (multi-source / multi-camera throughput test). The source is passed " + "N times as the pipeline's video_reference; aggregate (and per-stream) FPS is " + "reported. Mutually exclusive with passing multiple --video arguments.", + ) + parser.add_argument( + "--model-id", + default=None, + help="Optional value for the workflow's `model_id` parameter " + "(only sent when provided; otherwise the workflow's own default is used).", + ) + parser.add_argument( + "--api-key", + default=os.environ.get("ROBOFLOW_API_KEY") or os.environ.get("API_KEY"), + help="Roboflow API key (defaults to the ROBOFLOW_API_KEY / API_KEY env var).", + ) + parser.add_argument( + "--image-input-name", + default="image", + help="Name of the workflow's image input that frames are injected into.", + ) + parser.add_argument( + "--max-fps", + type=float, + default=None, + help="Cap processing FPS. Leave unset to measure full throughput.", + ) + parser.add_argument( + "--report-every", + type=int, + default=30, + help="Print a rolling FPS line every N processed frames (0 to disable).", + ) + parser.add_argument( + "--profile-trace", + default=None, + metavar="PATH", + help="Enable the workflow profiler and dump a Chrome trace (JSON) to this exact " + "file path (e.g. ./traces/my_run.json). View it in chrome://tracing or " + "https://ui.perfetto.dev. Leave unset to disable profiling.", + ) + parser.add_argument( + "--profile-frames", + type=int, + default=50, + help="Number of leading frames (workflow runs) to capture before the profiler " + "freezes. Keeps the trace small; later frames are not recorded. Only used when " + "--profile-trace is given.", + ) + parser.add_argument( + "--output-video", + default=None, + metavar="PATH", + help="Optionally write the visualization video to this exact path (e.g. " + "./out.mp4). Requires --output-key. With a single stream, source [0] is " + "written verbatim; with multiple streams, every source's output is tiled " + "into one grid frame (inference's create_tiles) per processing round.", + ) + parser.add_argument( + "--output-key", + default=None, + help="Name of the workflow output field that holds the visualization image " + "(e.g. 'label_visualization_output'). Required when --output-video is given.", + ) + parser.add_argument( + "--output-fps", + type=float, + default=None, + help="FPS for the written output video. Defaults to the source video's FPS " + "(falling back to 30 if it cannot be read). Only used with --output-video.", + ) + return parser.parse_args() + + +class FPSCounter: + """Counts processed frames and reports a rolling FPS (supervision's FPSMonitor) + plus a cumulative average over the whole run.""" + + def __init__(self) -> None: + self._monitor = sv.FPSMonitor() + self.frames = 0 + self._start: Optional[float] = None + + def tick(self) -> None: + if self._start is None: + self._start = time.monotonic() + self._monitor.tick() + self.frames += 1 + + @property + def rolling_fps(self) -> float: + # supervision >= 0.18 exposes `.fps`; older versions are callable. + return float( + self._monitor.fps if hasattr(self._monitor, "fps") else self._monitor() + ) + + @property + def overall_fps(self) -> float: + if self._start is None: + return 0.0 + elapsed = time.monotonic() - self._start + return self.frames / elapsed if elapsed > 0 else 0.0 + + +class VisualizationVideoWriter: + """Lazily opens a ``cv2.VideoWriter`` on the first frame (so it can size itself to + the visualization output) and writes subsequent BGR frames to a video file. + + Frame size is fixed at the first frame; any later frame of a different size is + resized to match (cv2 requires a constant frame size). + """ + + def __init__(self, path: str, fps: float) -> None: + self._path = path + self._fps = fps if fps and fps > 0 else 30.0 + self._writer: Optional[cv2.VideoWriter] = None + self._size: Optional[tuple] = None + self.frames_written = 0 + + def write(self, image_bgr: np.ndarray) -> None: + height, width = image_bgr.shape[:2] + if self._writer is None: + os.makedirs( + os.path.dirname(os.path.abspath(self._path)) or ".", exist_ok=True + ) + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + self._writer = cv2.VideoWriter( + self._path, fourcc, self._fps, (width, height) + ) + if not self._writer.isOpened(): + raise SystemExit( + f"Could not open a VideoWriter for '{self._path}' " + f"({width}x{height} @ {self._fps:.2f} fps). Check the path/extension " + "(try .mp4) and codec availability." + ) + self._size = (width, height) + elif (width, height) != self._size: + image_bgr = cv2.resize(image_bgr, self._size) + self._writer.write(np.ascontiguousarray(image_bgr, dtype=np.uint8)) + self.frames_written += 1 + + def release(self) -> None: + if self._writer is not None: + self._writer.release() + self._writer = None + + +def _to_bgr_image(value: Any) -> Optional[np.ndarray]: + """Extract an HWC BGR uint8 image from a workflow output value. + + Visualization outputs arrive (unserialized) as ``WorkflowImageData`` whose + ``.numpy_image`` is HWC BGR uint8 (cv2-native, exactly what VideoWriter wants); + a bare ndarray is taken as-is. Anything else returns ``None``. + """ + numpy_image = getattr(value, "numpy_image", None) + if numpy_image is not None: + return numpy_image + if isinstance(value, np.ndarray): + return value + return None + + +class OutputCapture: + """Pulls the named visualization output out of workflow results and feeds it to + the video writer. + + With a single stream, source [0]'s frames are written directly (legacy + behaviour). With multiple streams, the latest image of EVERY source is kept and + one `create_tiles` mosaic is written per sink round (see `end_round`) โ€” the + first frame waits until every source has delivered once so the grid never + changes size mid-video, and a source that ends early keeps showing its last + frame. Warns once (then stays quiet) on a missing key or a non-image value so a + misconfigured run is obvious without spamming per frame.""" + + def __init__( + self, + writer: VisualizationVideoWriter, + output_key: str, + num_streams: int = 1, + ) -> None: + self.writer = writer + self.output_key = output_key + self._num_streams = num_streams + self._latest: Dict[int, np.ndarray] = {} + self._round_dirty = False + self._warned_missing = False + self._warned_type = False + + def handle(self, prediction: Any, source_id: int) -> None: + if not isinstance(prediction, dict): + return + if self.output_key not in prediction: + if not self._warned_missing: + self._warned_missing = True + print( + f"[output-capture] WARNING: output '{self.output_key}' not found in " + f"the workflow result. Available outputs: {sorted(prediction.keys())}. " + "No video will be written.", + flush=True, + ) + return + image = _to_bgr_image(prediction[self.output_key]) + if image is None: + if not self._warned_type: + self._warned_type = True + print( + f"[output-capture] WARNING: output '{self.output_key}' is " + f"{type(prediction[self.output_key]).__name__}, not an image; " + "cannot write video.", + flush=True, + ) + return + if self._num_streams == 1: + self.writer.write(image) + return + self._latest[source_id] = image + self._round_dirty = True + + def end_round(self) -> None: + """Multi-stream only: write ONE tiled frame per sink round in which any + source updated, once all sources have been seen (stable tile grid).""" + if self._num_streams == 1 or not self._round_dirty: + return + self._round_dirty = False + if len(self._latest) < self._num_streams: + return + tile = create_tiles( + images=[self._latest[source_id] for source_id in sorted(self._latest)] + ) + self.writer.write(tile) + + +def build_sink( + counter: FPSCounter, + report_every: int, + capture: Optional[OutputCapture] = None, +): + def sink(predictions: Any, video_frames: Any) -> None: + # A single video source delivers one (prediction, frame) per call; the list + # form is used for multi-source pipelines. Normalise to a list either way. + if not isinstance(predictions, list): + predictions, video_frames = [predictions], [video_frames] + for prediction, frame in zip(predictions, video_frames): + if frame is None: + continue + counter.tick() + if capture is not None: + # source_id is an int for multi-source pipelines and None for a + # single source - normalise so the capture can key its tiles. + source_id = getattr(frame, "source_id", None) + capture.handle(prediction, source_id if source_id is not None else 0) + if report_every and counter.frames % report_every == 0: + print( + f"[{counter.frames:>7} frames] rolling FPS: {counter.rolling_fps:6.1f} " + f"| avg FPS: {counter.overall_fps:6.1f}", + flush=True, + ) + if capture is not None: + capture.end_round() + + return sink + + +def _looks_like_stream(reference: str) -> bool: + """Anything with a URI scheme (rtsp://, http://, ...) or a V4L2 device path is a + live stream/camera; only bare paths are validated (and FPS-probed) as local + files.""" + return "://" in reference or reference.startswith("/dev/video") + + +def _source_video_fps(video_path: str) -> Optional[float]: + cap = cv2.VideoCapture(video_path) + fps = cap.get(cv2.CAP_PROP_FPS) + cap.release() + return fps if fps and fps > 0 else None + + +def main() -> None: + args = parse_args() + if not args.api_key: + raise SystemExit( + "A Roboflow API key is required (registered workflows are fetched from the " + "platform). Pass --api-key or set ROBOFLOW_API_KEY." + ) + references: List[str] = args.video + if args.n is not None: + if len(references) > 1: + raise SystemExit( + "--n and multiple --video arguments are mutually exclusive: --n " + "replicates a SINGLE --video into N identical streams, while " + "repeating --video enumerates distinct sources explicitly." + ) + if args.n < 1: + raise SystemExit("--n must be >= 1") + for reference in references: + if not _looks_like_stream(reference) and not os.path.isfile(reference): + raise SystemExit(f"Video file not found: {reference}") + if args.profile_trace: + if args.profile_frames < 1: + raise SystemExit("--profile-frames must be >= 1") + _enable_workflow_profiler( + trace_path=args.profile_trace, max_frames=args.profile_frames + ) + print( + f"Profiler enabled: capturing first {args.profile_frames} frame(s) -> " + f"'{args.profile_trace}'" + ) + + # --n multiplies the single source into N concurrent copies; repeated --video + # enumerates distinct sources. A single reference keeps the exact legacy + # behaviour; a list makes InferencePipeline run one VideoSource per entry. + video_references = ( + references * args.n if args.n is not None and args.n > 1 else list(references) + ) + num_streams = len(video_references) + video_reference = video_references[0] if num_streams == 1 else video_references + + # Optional capture of the visualization output: source [0] verbatim for a single + # stream, a create_tiles mosaic of every source for multi-stream runs. + capture: Optional[OutputCapture] = None + if args.output_video: + if not args.output_key: + raise SystemExit( + "--output-key (the workflow output field holding the visualization " + "image) is required when --output-video is given." + ) + # cv2.VideoCapture can block on live sources, so only local files are + # FPS-probed for the writer default; streams fall back to --output-fps / 30. + probed_fps = ( + _source_video_fps(references[0]) + if not _looks_like_stream(references[0]) + else None + ) + out_fps = args.output_fps or probed_fps or 30.0 + capture = OutputCapture( + writer=VisualizationVideoWriter(args.output_video, out_fps), + output_key=args.output_key, + num_streams=num_streams, + ) + target_desc = ( + "source [0]" + if num_streams == 1 + else f"all {num_streams} sources (create_tiles mosaic)" + ) + print( + f"Capturing workflow output '{args.output_key}' from {target_desc} -> " + f"'{args.output_video}' @ {out_fps:.2f} fps" + ) + elif args.output_key: + raise SystemExit("--output-key requires --output-video to also be given.") + + # Only steer `model_id` when explicitly provided, so the workflow keeps its own + # default otherwise. + workflows_parameters = ( + {"model_id": args.model_id} if args.model_id is not None else None + ) + + counter = FPSCounter() + watchdog = BasePipelineWatchDog() + pipeline = InferencePipeline.init_with_workflow( + video_reference=video_reference, + workspace_name=args.workspace, + workflow_id=args.workflow_id, + api_key=args.api_key, + image_input_name=args.image_input_name, + workflows_parameters=workflows_parameters, + on_prediction=build_sink(counter, args.report_every, capture), + max_fps=args.max_fps, # None => run as fast as the workflow allows + watchdog=watchdog, + ) + + cap = f"capped at {args.max_fps} FPS" if args.max_fps else "uncapped" + model = f" model_id='{args.model_id}'" if args.model_id else "" + sources_desc = ( + f"'{video_references[0]}' x{num_streams} stream(s)" + if len(set(video_references)) == 1 + else f"{num_streams} sources: " + ", ".join(f"'{r}'" for r in video_references) + ) + print( + f"Running '{args.workspace}/{args.workflow_id}' over {sources_desc} " + f"({cap}){model}\n" + ) + wall_start = time.monotonic() + try: + try: + # Blocks (use_main_thread) until file sources are exhausted; live + # streams run until Ctrl+C. + pipeline.start() + pipeline.join() + except KeyboardInterrupt: + print("\nInterrupted โ€” terminating pipeline...") + pipeline.terminate() + pipeline.join() + finally: + # Always finalise the output video โ€” even on interrupt โ€” so it stays playable. + if capture is not None: + capture.writer.release() + elapsed = time.monotonic() - wall_start + + per_stream = ( + f" ({counter.overall_fps / num_streams:.2f}/stream)" if num_streams > 1 else "" + ) + same_video = ( + " (same source)" if num_streams > 1 and len(set(video_references)) == 1 else "" + ) + print("\n=== Summary ===") + print(f"streams : {num_streams}{same_video}") + print(f"frames processed : {counter.frames} (aggregate across streams)") + print(f"total wall time : {elapsed:.2f} s (includes model load / warm-up)") + print(f"average processing FPS: {counter.overall_fps:.2f} aggregate{per_stream}") + if args.profile_trace: + print( + f"profiler trace : {os.path.abspath(args.profile_trace)} " + f"(first {args.profile_frames} frame(s))" + ) + if capture is not None: + captured_desc = ( + "from source [0]" + if num_streams == 1 + else f"tiled across {num_streams} sources" + ) + print( + f"output video : {os.path.abspath(args.output_video)} " + f"({capture.writer.frames_written} frame(s) {captured_desc})" + ) + + +if __name__ == "__main__": + main() diff --git a/docker/docker-bake.jetson.7.2.0.hcl b/docker/docker-bake.jetson.7.2.0.hcl new file mode 100644 index 0000000000..0b4edb2fd3 --- /dev/null +++ b/docker/docker-bake.jetson.7.2.0.hcl @@ -0,0 +1,72 @@ +variable "JETSON_SERVER_TAGS" { + default = "roboflow/roboflow-inference-server-jetson-7.2.0:local" +} + +variable "JETSON_WHEELS_TAGS" { + default = "roboflow/roboflow-inference-wheels-jetson-7.2.0:local" +} + +target "jetson-media-jp72" { + context = "." + dockerfile = "docker/dockerfiles/Dockerfile.media.jetson.7.2.0" + platforms = ["linux/arm64"] +} + +target "jetson-wheels-jp72" { + context = "." + dockerfile = "docker/dockerfiles/Dockerfile.wheels.jetson.7.2.0" + platforms = ["linux/arm64"] + tags = split(",", JETSON_WHEELS_TAGS) +} + +# Component stages of the wheels image, grouped by dependency chain so CI can +# fan the compiles out to separate builders; they meet in the shared project +# cache and jetson-wheels-jp72 assembles the result. +target "jetson-wheels-torchvision-jp72" { + inherits = ["jetson-wheels-jp72"] + target = "torchvision-builder" + tags = [] +} + +target "jetson-wheels-flash-attn-jp72" { + inherits = ["jetson-wheels-jp72"] + target = "flash-attn-builder" + tags = [] +} + +target "jetson-wheels-ort-jp72" { + inherits = ["jetson-wheels-jp72"] + target = "ort-builder" + tags = [] +} + +target "jetson-wheels-triton-jp72" { + inherits = ["jetson-wheels-jp72"] + target = "triton-fetcher" + tags = [] +} + +target "jetson-wheels-pycuda-jp72" { + inherits = ["jetson-wheels-jp72"] + target = "pycuda-builder" + tags = [] +} + +group "jetson-wheels-torch-stack-jp72" { + targets = ["jetson-wheels-torchvision-jp72", "jetson-wheels-flash-attn-jp72"] +} + +group "jetson-wheels-ort-stack-jp72" { + targets = ["jetson-wheels-ort-jp72", "jetson-wheels-triton-jp72", "jetson-wheels-pycuda-jp72"] +} + +target "jetson-server-jp72" { + context = "." + dockerfile = "docker/dockerfiles/Dockerfile.onnx.jetson.7.2.0" + contexts = { + jetson-media = "target:jetson-media-jp72" + jetson-wheels = "target:jetson-wheels-jp72" + } + platforms = ["linux/arm64"] + tags = split(",", JETSON_SERVER_TAGS) +} diff --git a/docker/dockerfiles/Dockerfile.media.jetson.7.2.0 b/docker/dockerfiles/Dockerfile.media.jetson.7.2.0 new file mode 100644 index 0000000000..cd30516218 --- /dev/null +++ b/docker/dockerfiles/Dockerfile.media.jetson.7.2.0 @@ -0,0 +1,341 @@ +# syntax=docker/dockerfile:1.7 +ARG CUDA_IMAGE_VERSION=13.2.0 +ARG FFMPEG_VERSION=7.1.3 +ARG FFMPEG_SHA256=f0bf043299db9e3caacb435a712fc541fbb07df613c4b893e8b77e67baf3adbe +ARG GSTREAMER_VERSION=1.24.12 +ARG GSTREAMER_COMMIT=88e31216479a3f0c25096ca406369f3983efec14 +ARG GLIB_NETWORKING_SHA256=cd2a084c7bb91d78e849fb55d40e472f6d8f6862cddc9f12c39149359ba18268 +ARG OPENCV_VERSION=4.13.0 +ARG OPENCV_SHA256=1d40ca017ea51c533cf9fd5cbde5b5fe7ae248291ddf2af99d4c17cf8e13017d +ARG OPENCV_CONTRIB_SHA256=1e0077a4fd2960a7d2f4c9e49d6ba7bb891cac2d1be36d7e8e47aa97a9d1039b +ARG OPENCV_CUDA_ARCH_BIN=8.7;11.0 +ARG OPENCV_CUDA_ARCH_PTX= +ARG JETSON_TENSOR_BRIDGE_CUDA_ARCHITECTURES=87;110 +ARG L4T_REPOSITORY=r39.2 + +FROM ubuntu:24.04 AS ffmpeg-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG FFMPEG_SHA256 +ARG FFMPEG_VERSION +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-ffmpeg-jp72-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-ffmpeg-jp72-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + libbz2-dev \ + liblzma-dev \ + libssl-dev \ + nasm \ + pkg-config \ + xz-utils \ + zlib1g-dev + +COPY docker/scripts/build_ffmpeg.sh /usr/local/bin/build_ffmpeg +COPY docker/scripts/verify_ffmpeg.sh /usr/local/bin/verify_ffmpeg + +ENV PATH=/opt/ffmpeg/bin:$PATH \ + LD_LIBRARY_PATH=/opt/ffmpeg/lib \ + PKG_CONFIG_PATH=/opt/ffmpeg/lib/pkgconfig + +RUN chmod +x /usr/local/bin/build_ffmpeg /usr/local/bin/verify_ffmpeg && \ + FFMPEG_VERSION="${FFMPEG_VERSION}" \ + FFMPEG_SHA256="${FFMPEG_SHA256}" \ + build_ffmpeg && \ + verify_ffmpeg && \ + du -sh /opt/ffmpeg /opt/ffmpeg-runtime + +FROM nvcr.io/nvidia/cuda:${CUDA_IMAGE_VERSION}-devel-ubuntu24.04 AS gstreamer-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG GLIB_NETWORKING_SHA256 +ARG GSTREAMER_COMMIT +ARG GSTREAMER_VERSION +ARG TARGETARCH + +COPY --from=ffmpeg-builder /opt/ffmpeg /opt/ffmpeg + +ENV PATH=/opt/ffmpeg/bin:$PATH \ + LD_LIBRARY_PATH=/opt/ffmpeg/lib \ + PKG_CONFIG_PATH=/opt/ffmpeg/lib/pkgconfig + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-gstreamer-jp72-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-gstreamer-jp72-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + bison \ + build-essential \ + ca-certificates \ + cmake \ + curl \ + flex \ + git \ + gobject-introspection \ + libcurl4-openssl-dev \ + libgl-dev \ + libglib2.0-dev \ + libgnutls28-dev \ + libgirepository1.0-dev \ + libjpeg-dev \ + libopus-dev \ + libpng-dev \ + libsrtp2-dev \ + libssl-dev \ + libv4l-dev \ + libvpx-dev \ + ninja-build \ + pkg-config \ + python3 \ + python3-pip \ + python3-venv + +RUN python3 -m venv /opt/gstreamer-build-tools && \ + /opt/gstreamer-build-tools/bin/pip install --no-cache-dir \ + "meson>=1.4,<2" \ + "ninja>=1.11,<2" + +ENV PATH=/opt/gstreamer/bin:/opt/gstreamer-build-tools/bin:/opt/ffmpeg/bin:$PATH \ + GIO_EXTRA_MODULES=/opt/gstreamer/lib/gio/modules \ + GI_TYPELIB_PATH=/opt/gstreamer/lib/girepository-1.0 \ + GST_PLUGIN_SYSTEM_PATH_1_0=/opt/gstreamer/lib/gstreamer-1.0 \ + LD_LIBRARY_PATH=/opt/gstreamer/lib:/opt/ffmpeg/lib \ + PKG_CONFIG_PATH=/opt/gstreamer/lib/pkgconfig:/opt/ffmpeg/lib/pkgconfig + +COPY docker/scripts/build_gstreamer.sh /usr/local/bin/build_gstreamer +COPY docker/scripts/verify_gstreamer.sh /usr/local/bin/verify_gstreamer + +RUN chmod +x /usr/local/bin/build_gstreamer /usr/local/bin/verify_gstreamer && \ + GSTREAMER_VERSION="${GSTREAMER_VERSION}" \ + GSTREAMER_COMMIT="${GSTREAMER_COMMIT}" \ + GLIB_NETWORKING_SHA256="${GLIB_NETWORKING_SHA256}" \ + GSTREAMER_NVCODEC=disabled \ + build_gstreamer && \ + verify_gstreamer && \ + test ! -e /opt/gstreamer/lib/gstreamer-1.0/libgstnvcodec.so && \ + du -sh /opt/gstreamer /opt/gstreamer-runtime + +FROM gstreamer-builder AS opencv-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG OPENCV_VERSION +ARG OPENCV_SHA256 +ARG OPENCV_CONTRIB_SHA256 +ARG OPENCV_CUDA_ARCH_BIN +ARG OPENCV_CUDA_ARCH_PTX +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-opencv-jp72-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-opencv-jp72-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + libtiff-dev \ + python3-dev + +RUN python3 -m venv /opt/opencv-build-tools && \ + /opt/opencv-build-tools/bin/pip install --no-cache-dir \ + "numpy>=2.0.0,<2.4.0" + +COPY docker/scripts/build_opencv.sh /usr/local/bin/build_opencv +COPY docker/scripts/verify_opencv.sh /usr/local/bin/verify_opencv + +RUN chmod +x /usr/local/bin/build_opencv /usr/local/bin/verify_opencv && \ + OPENCV_VERSION="${OPENCV_VERSION}" \ + OPENCV_SHA256="${OPENCV_SHA256}" \ + OPENCV_CONTRIB_SHA256="${OPENCV_CONTRIB_SHA256}" \ + OPENCV_PYTHON_EXECUTABLE=/opt/opencv-build-tools/bin/python3 \ + OPENCV_WITH_CUDA=ON \ + OPENCV_CUDA_ARCH_BIN="${OPENCV_CUDA_ARCH_BIN}" \ + OPENCV_CUDA_ARCH_PTX="${OPENCV_CUDA_ARCH_PTX}" \ + build_opencv && \ + PATH=/opt/opencv-build-tools/bin:$PATH \ + PYTHONPATH=/opt/opencv/python \ + OPENCV_EXPECTED_CUDA_ARCH_BIN="${OPENCV_CUDA_ARCH_BIN}" \ + OPENCV_REQUIRE_CUDA_BUILD=true \ + verify_opencv && \ + du -sh /opt/opencv /opt/opencv-runtime + +FROM ubuntu:24.04 AS nvidia-repository-key + +ARG DEBIAN_FRONTEND=noninteractive + +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends ca-certificates curl && \ + curl -fsSL https://repo.download.nvidia.com/jetson/jetson-ota-public.asc \ + -o /nvidia-jetson.asc + +FROM gstreamer-builder AS jetson-tensor-bridge-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG JETSON_TENSOR_BRIDGE_CUDA_ARCHITECTURES +ARG L4T_REPOSITORY +ARG TARGETARCH + +COPY --from=nvidia-repository-key /nvidia-jetson.asc /usr/share/keyrings/nvidia-jetson.asc +COPY --from=nvidia-repository-key /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jetson-tensor-bridge-jp72-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jetson-tensor-bridge-jp72-${TARGETARCH} \ + echo "deb [signed-by=/usr/share/keyrings/nvidia-jetson.asc] https://repo.download.nvidia.com/jetson/common ${L4T_REPOSITORY} main" > /etc/apt/sources.list.d/nvidia-jetson.list && \ + echo "deb [signed-by=/usr/share/keyrings/nvidia-jetson.asc] https://repo.download.nvidia.com/jetson/som ${L4T_REPOSITORY} main" >> /etc/apt/sources.list.d/nvidia-jetson.list && \ + apt-get update -y && \ + mkdir -p /tmp/l4t-api && \ + cd /tmp/l4t-api && \ + apt-get download nvidia-l4t-jetson-multimedia-api && \ + dpkg-deb -x ./*.deb / && \ + cd / && \ + rm -rf /tmp/l4t-api + +COPY docker/native/jetson_tensor_bridge /src/jetson_tensor_bridge +COPY docker/scripts/build_jetson_tensor_bridge.sh /usr/local/bin/build_jetson_tensor_bridge +COPY docker/scripts/verify_jetson_tensor_bridge.sh /usr/local/bin/verify_jetson_tensor_bridge + +RUN chmod +x \ + /usr/local/bin/build_jetson_tensor_bridge \ + /usr/local/bin/verify_jetson_tensor_bridge && \ + JETSON_TENSOR_BRIDGE_CUDA_ARCHITECTURES="${JETSON_TENSOR_BRIDGE_CUDA_ARCHITECTURES}" \ + build_jetson_tensor_bridge && \ + verify_jetson_tensor_bridge + +FROM gstreamer-builder AS cuda-runtime-export + +ARG CUDA_IMAGE_VERSION + +RUN cuda_major="${CUDA_IMAGE_VERSION%%.*}" && \ + cuda_minor="$(echo "${CUDA_IMAGE_VERSION}" | cut -d. -f2)" && \ + cuda_package_suffix="${cuda_major}-${cuda_minor}" && \ + mkdir -p /opt/cuda-runtime/lib /opt/cuda-runtime/share/licenses && \ + for library in \ + "libcudart.so.${cuda_major}" \ + "libnvjpeg.so.${cuda_major}" \ + "libnppc.so.${cuda_major}" \ + "libnppial.so.${cuda_major}" \ + "libnppicc.so.${cuda_major}" \ + "libnppidei.so.${cuda_major}" \ + "libnppif.so.${cuda_major}" \ + "libnppig.so.${cuda_major}" \ + "libnppim.so.${cuda_major}" \ + "libnppist.so.${cuda_major}" \ + "libnppitc.so.${cuda_major}"; do \ + cp -L "/usr/local/cuda/lib64/${library}" \ + "/opt/cuda-runtime/lib/${library}"; \ + done && \ + cp "/usr/share/doc/cuda-cudart-${cuda_package_suffix}/copyright" \ + /opt/cuda-runtime/share/licenses/cuda-cudart && \ + cp "/usr/share/doc/libnvjpeg-${cuda_package_suffix}/copyright" \ + /opt/cuda-runtime/share/licenses/libnvjpeg && \ + cp "/usr/share/doc/libnpp-${cuda_package_suffix}/copyright" \ + /opt/cuda-runtime/share/licenses/libnpp && \ + test "$(du -sm /opt/cuda-runtime | cut -f1)" -lt 165 + +FROM ubuntu:24.04 AS runtime + +ARG DEBIAN_FRONTEND=noninteractive +ARG L4T_REPOSITORY +ARG OPENCV_CUDA_ARCH_BIN +ARG TARGETARCH + +COPY --from=nvidia-repository-key /nvidia-jetson.asc /usr/share/keyrings/nvidia-jetson.asc +COPY --from=nvidia-repository-key /etc/ssl/certs/ca-certificates.crt /etc/ssl/certs/ca-certificates.crt + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-media-runtime-jp72-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-media-runtime-jp72-${TARGETARCH} \ + echo "deb [signed-by=/usr/share/keyrings/nvidia-jetson.asc] https://repo.download.nvidia.com/jetson/common ${L4T_REPOSITORY} main" > /etc/apt/sources.list.d/nvidia-jetson.list && \ + echo "deb [signed-by=/usr/share/keyrings/nvidia-jetson.asc] https://repo.download.nvidia.com/jetson/som ${L4T_REPOSITORY} main" >> /etc/apt/sources.list.d/nvidia-jetson.list && \ + echo "deb [signed-by=/usr/share/keyrings/nvidia-jetson.asc] https://repo.download.nvidia.com/jetson/ffmpeg ${L4T_REPOSITORY} main" >> /etc/apt/sources.list.d/nvidia-jetson.list && \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + ca-certificates \ + libasound2t64 \ + libbz2-1.0 \ + libcurl4t64 \ + libdrm2 \ + libglib2.0-0t64 \ + libgles2 \ + libgnutls30t64 \ + libjpeg-turbo8 \ + kmod \ + liblzma5 \ + libopus0 \ + libpng16-16t64 \ + libsrtp2-1 \ + libssl3t64 \ + libtiff6 \ + libv4l-0t64 \ + libwebp7 \ + libwebpdemux2 \ + libwebpmux3 \ + libvpx9 \ + libwayland-client0 \ + libwayland-egl1 \ + libx11-6 \ + libxext6 \ + python3 \ + zlib1g && \ + mkdir -p /tmp/l4t-runtime && \ + cd /tmp/l4t-runtime && \ + apt-get download \ + libegl1 \ + nvidia-l4t-adaruntime \ + nvidia-l4t-camera \ + nvidia-l4t-core \ + nvidia-l4t-gstreamer \ + nvidia-l4t-multimedia \ + nvidia-l4t-multimedia-utils \ + nvidia-l4t-nvsci \ + nvidia-l4t-openwfd && \ + for package in ./*.deb; do dpkg-deb -x "${package}" /; done && \ + cd / && \ + rm -rf \ + /etc/apt/sources.list.d/nvidia-jetson.list \ + /tmp/l4t-runtime \ + /usr/share/keyrings/nvidia-jetson.asc + +COPY --from=ffmpeg-builder /opt/ffmpeg-runtime /opt/ffmpeg +COPY --from=gstreamer-builder /opt/gstreamer-runtime /opt/gstreamer +COPY --from=opencv-builder /opt/opencv-runtime /opt/opencv +COPY --from=cuda-runtime-export /opt/cuda-runtime /opt/cuda-runtime +COPY --from=jetson-tensor-bridge-builder /opt/roboflow/lib /opt/roboflow/lib +COPY LICENSE LICENSE.core /opt/roboflow/share/licenses/roboflow-media-bridges/ + +ENV PATH=/opt/gstreamer/bin:/opt/ffmpeg/bin:$PATH \ + GIO_EXTRA_MODULES=/opt/gstreamer/lib/gio/modules \ + GI_TYPELIB_PATH=/opt/gstreamer/lib/girepository-1.0 \ + GST_PLUGIN_PATH_1_0=/usr/lib/aarch64-linux-gnu/gstreamer-1.0 \ + GST_PLUGIN_SYSTEM_PATH_1_0=/opt/gstreamer/lib/gstreamer-1.0 \ + LD_LIBRARY_PATH=/opt/opencv/lib:/opt/gstreamer/lib:/opt/ffmpeg/lib:/opt/cuda-runtime/lib:/usr/lib/aarch64-linux-gnu/tegra:/usr/lib/aarch64-linux-gnu/nvidia \ + NVIDIA_DRIVER_CAPABILITIES=compute,utility,video \ + NVIDIA_VISIBLE_DEVICES=all \ + PYTHONPATH=/opt/opencv/python + +RUN printf '/opt/opencv/lib\n/opt/gstreamer/lib\n/opt/ffmpeg/lib\n/opt/cuda-runtime/lib\n/usr/lib/aarch64-linux-gnu/tegra\n/usr/lib/aarch64-linux-gnu/nvidia\n' \ + > /etc/ld.so.conf.d/roboflow-media.conf && \ + ldconfig && \ + rm -rf /root/.cache/gstreamer-1.0 + +RUN --mount=type=bind,source=docker/scripts/verify_ffmpeg.sh,target=/tmp/verify_ffmpeg,ro \ + --mount=type=bind,source=docker/scripts/verify_gstreamer.sh,target=/tmp/verify_gstreamer,ro \ + --mount=type=bind,source=docker/scripts/verify_opencv.sh,target=/tmp/verify_opencv,ro \ + --mount=type=bind,source=docker/scripts/verify_media_runtime.sh,target=/tmp/verify_media_runtime,ro \ + --mount=from=opencv-builder,source=/opt/opencv-build-tools/lib/python3.12/site-packages,target=/tmp/opencv-python-deps,ro \ + export GST_REGISTRY=/tmp/roboflow-gstreamer-registry.bin && \ + sh /tmp/verify_ffmpeg && \ + sh /tmp/verify_gstreamer && \ + test ! -e /opt/gstreamer/lib/gstreamer-1.0/libgstnvcodec.so && \ + PYTHONPATH=/tmp/opencv-python-deps:/opt/opencv/python \ + OPENCV_EXPECTED_CUDA_ARCH_BIN="${OPENCV_CUDA_ARCH_BIN}" \ + OPENCV_REQUIRE_CUDA_BUILD=true \ + sh /tmp/verify_opencv && \ + MEDIA_RUNTIME_REQUIRE_JETSON_PLUGINS=true sh /tmp/verify_media_runtime && \ + test -s /usr/lib/aarch64-linux-gnu/libEGL.so.1 && \ + test -s /usr/lib/aarch64-linux-gnu/libGLESv2.so.2 && \ + test -z "$(find /usr/lib/aarch64-linux-gnu -maxdepth 1 -type f \ + \( -name 'libEGL_mesa.so*' -o -name 'libgbm.so*' \ + -o -name 'libgallium-*.so' -o -name 'libLLVM.so*' \) -print -quit)" && \ + test -s /opt/roboflow/lib/libroboflow_jetson_tensor.so.1 && \ + rm -f "${GST_REGISTRY}" + +LABEL org.opencontainers.image.description="Roboflow Jetson 7.2 media runtime" diff --git a/docker/dockerfiles/Dockerfile.onnx.gpu b/docker/dockerfiles/Dockerfile.onnx.gpu index 02734c0f72..3b8ee243db 100644 --- a/docker/dockerfiles/Dockerfile.onnx.gpu +++ b/docker/dockerfiles/Dockerfile.onnx.gpu @@ -1,4 +1,225 @@ -FROM nvcr.io/nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 as builder +# syntax=docker/dockerfile:1.7 +ARG FFMPEG_VERSION=7.1.3 +ARG FFMPEG_SHA256=f0bf043299db9e3caacb435a712fc541fbb07df613c4b893e8b77e67baf3adbe +ARG GSTREAMER_VERSION=1.24.12 +ARG GSTREAMER_COMMIT=88e31216479a3f0c25096ca406369f3983efec14 +ARG GLIB_NETWORKING_SHA256=cd2a084c7bb91d78e849fb55d40e472f6d8f6862cddc9f12c39149359ba18268 +ARG OPENCV_VERSION=4.13.0 +ARG OPENCV_SHA256=1d40ca017ea51c533cf9fd5cbde5b5fe7ae248291ddf2af99d4c17cf8e13017d +ARG OPENCV_CONTRIB_SHA256=1e0077a4fd2960a7d2f4c9e49d6ba7bb891cac2d1be36d7e8e47aa97a9d1039b +ARG OPENCV_CUDA_ARCH_BIN=7.5;8.9 +ARG OPENCV_CUDA_ARCH_PTX=7.5 + +FROM ubuntu:22.04 AS ffmpeg-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG FFMPEG_SHA256 +ARG FFMPEG_VERSION +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-ffmpeg-gpu-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-ffmpeg-gpu-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + libbz2-dev \ + liblzma-dev \ + libssl-dev \ + nasm \ + pkg-config \ + xz-utils \ + zlib1g-dev + +COPY docker/scripts/build_ffmpeg.sh /usr/local/bin/build_ffmpeg + +ENV PATH=/opt/ffmpeg/bin:$PATH \ + LD_LIBRARY_PATH=/opt/ffmpeg/lib \ + PKG_CONFIG_PATH=/opt/ffmpeg/lib/pkgconfig + +RUN chmod +x /usr/local/bin/build_ffmpeg && \ + FFMPEG_VERSION="${FFMPEG_VERSION}" \ + FFMPEG_SHA256="${FFMPEG_SHA256}" \ + build_ffmpeg && \ + du -sh /opt/ffmpeg /opt/ffmpeg-runtime + +COPY docker/scripts/verify_ffmpeg.sh /usr/local/bin/verify_ffmpeg + +RUN chmod +x /usr/local/bin/verify_ffmpeg && verify_ffmpeg + +FROM nvcr.io/nvidia/cuda:12.4.1-devel-ubuntu22.04 AS gstreamer-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG GLIB_NETWORKING_SHA256 +ARG GSTREAMER_COMMIT +ARG GSTREAMER_VERSION +ARG TARGETARCH + +COPY --from=ffmpeg-builder /opt/ffmpeg /opt/ffmpeg + +ENV PATH=/opt/ffmpeg/bin:$PATH \ + LD_LIBRARY_PATH=/opt/ffmpeg/lib \ + PKG_CONFIG_PATH=/opt/ffmpeg/lib/pkgconfig + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-gstreamer-gpu-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-gstreamer-gpu-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + bison \ + build-essential \ + ca-certificates \ + cmake \ + curl \ + flex \ + gobject-introspection \ + git \ + libbz2-dev \ + libcurl4-openssl-dev \ + libglib2.0-dev \ + libgnutls28-dev \ + libgirepository1.0-dev \ + libgl-dev \ + libjpeg-dev \ + libopus-dev \ + libpng-dev \ + libsrtp2-dev \ + libssl-dev \ + libv4l-dev \ + libvpx-dev \ + liblzma-dev \ + ninja-build \ + pkg-config \ + python3 \ + python3-pip \ + python3-venv \ + zlib1g-dev + +RUN python3 -m venv /opt/gstreamer-build-tools && \ + /opt/gstreamer-build-tools/bin/pip install --no-cache-dir \ + "meson>=1.4,<2" \ + "ninja>=1.11,<2" + +ENV PATH=/opt/gstreamer-build-tools/bin:$PATH + +COPY docker/scripts/build_gstreamer.sh /usr/local/bin/build_gstreamer + +RUN chmod +x /usr/local/bin/build_gstreamer && \ + GSTREAMER_VERSION="${GSTREAMER_VERSION}" \ + GSTREAMER_COMMIT="${GSTREAMER_COMMIT}" \ + GLIB_NETWORKING_SHA256="${GLIB_NETWORKING_SHA256}" \ + GSTREAMER_NVCODEC=enabled \ + build_gstreamer + +ENV PATH=/opt/gstreamer/bin:/opt/ffmpeg/bin:$PATH \ + GIO_EXTRA_MODULES=/opt/gstreamer/lib/gio/modules \ + GI_TYPELIB_PATH=/opt/gstreamer/lib/girepository-1.0 \ + LD_LIBRARY_PATH=/opt/gstreamer/lib:/opt/ffmpeg/lib \ + PKG_CONFIG_PATH=/opt/gstreamer/lib/pkgconfig:/opt/ffmpeg/lib/pkgconfig \ + GST_PLUGIN_SYSTEM_PATH_1_0=/opt/gstreamer/lib/gstreamer-1.0 + +COPY docker/scripts/verify_gstreamer.sh /usr/local/bin/verify_gstreamer + +RUN chmod +x /usr/local/bin/verify_gstreamer && \ + GSTREAMER_REQUIRE_NVCODEC=true verify_gstreamer && \ + du -sh /opt/gstreamer /opt/gstreamer-runtime + +FROM gstreamer-builder AS gstreamer-cuda-tensor-bridge-builder + +COPY docker/native/gstreamer_cuda_tensor_bridge /src/gstreamer_cuda_tensor_bridge +COPY docker/scripts/build_gstreamer_cuda_tensor_bridge.sh /usr/local/bin/build_gstreamer_cuda_tensor_bridge +COPY docker/scripts/verify_gstreamer_cuda_tensor_bridge.sh /usr/local/bin/verify_gstreamer_cuda_tensor_bridge + +RUN chmod +x \ + /usr/local/bin/build_gstreamer_cuda_tensor_bridge \ + /usr/local/bin/verify_gstreamer_cuda_tensor_bridge && \ + build_gstreamer_cuda_tensor_bridge && \ + verify_gstreamer_cuda_tensor_bridge && \ + du -sh /opt/roboflow/lib + +FROM gstreamer-builder AS opencv-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG OPENCV_VERSION +ARG OPENCV_SHA256 +ARG OPENCV_CONTRIB_SHA256 +ARG OPENCV_CUDA_ARCH_BIN +ARG OPENCV_CUDA_ARCH_PTX +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-opencv-gpu-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-opencv-gpu-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + libopenblas-dev \ + libtiff-dev \ + python3-dev + +RUN python3 -m venv /opt/opencv-build-tools && \ + /opt/opencv-build-tools/bin/pip install --no-cache-dir \ + "numpy>=2.0.0,<2.4.0" + +COPY docker/scripts/build_opencv.sh /usr/local/bin/build_opencv +COPY docker/scripts/verify_opencv.sh /usr/local/bin/verify_opencv + +RUN chmod +x /usr/local/bin/build_opencv /usr/local/bin/verify_opencv && \ + OPENCV_VERSION="${OPENCV_VERSION}" \ + OPENCV_SHA256="${OPENCV_SHA256}" \ + OPENCV_CONTRIB_SHA256="${OPENCV_CONTRIB_SHA256}" \ + OPENCV_PYTHON_EXECUTABLE=/opt/opencv-build-tools/bin/python3 \ + OPENCV_WITH_CUDA=ON \ + OPENCV_CUDA_ARCH_BIN="${OPENCV_CUDA_ARCH_BIN}" \ + OPENCV_CUDA_ARCH_PTX="${OPENCV_CUDA_ARCH_PTX}" \ + build_opencv && \ + du -sh /opt/opencv /opt/opencv-runtime + +RUN PATH=/opt/opencv-build-tools/bin:$PATH \ + PYTHONPATH=/opt/opencv/python \ + OPENCV_EXPECTED_CUDA_ARCH_BIN="${OPENCV_CUDA_ARCH_BIN}" \ + OPENCV_REQUIRE_CUDA_BUILD=true \ + verify_opencv + +FROM ubuntu:22.04 AS gstreamer-runtime + +ARG DEBIAN_FRONTEND=noninteractive +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-gstreamer-runtime-gpu-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-gstreamer-runtime-gpu-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + ca-certificates \ + libbz2-1.0 \ + libcurl4 \ + libglib2.0-0 \ + libgnutls30 \ + libjpeg-turbo8 \ + libopus0 \ + libsrtp2-1 \ + libssl3 \ + libv4l-0 \ + libvpx7 \ + liblzma5 \ + zlib1g + +COPY --from=ffmpeg-builder /opt/ffmpeg-runtime /opt/ffmpeg +COPY --from=gstreamer-builder /opt/gstreamer-runtime /opt/gstreamer + +ENV PATH=/opt/gstreamer/bin:/opt/ffmpeg/bin:$PATH \ + GIO_EXTRA_MODULES=/opt/gstreamer/lib/gio/modules \ + GI_TYPELIB_PATH=/opt/gstreamer/lib/girepository-1.0 \ + LD_LIBRARY_PATH=/opt/gstreamer/lib:/opt/ffmpeg/lib \ + GST_PLUGIN_SYSTEM_PATH_1_0=/opt/gstreamer/lib/gstreamer-1.0 + +RUN --mount=type=bind,source=docker/scripts/verify_ffmpeg.sh,target=/tmp/verify_ffmpeg,ro \ + --mount=type=bind,source=docker/scripts/verify_gstreamer.sh,target=/tmp/verify_gstreamer,ro \ + --mount=type=bind,source=docker/scripts/verify_media_runtime.sh,target=/tmp/verify_media_runtime,ro \ + sh /tmp/verify_ffmpeg && \ + GSTREAMER_REQUIRE_NVCODEC=true sh /tmp/verify_gstreamer && \ + sh /tmp/verify_media_runtime && \ + rm -rf /root/.cache/gstreamer-1.0 + +FROM nvcr.io/nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04 AS builder WORKDIR /app @@ -33,12 +254,6 @@ COPY requirements/requirements.sam.txt \ RUN python3 -m pip install -U pip uv -# Well - this is good and bad, dependent on how we see this - on one end, we resolve exactly what -# is in inference-models lock, but also install that prior other requirements and w/o the `inference-models` -# package content, which lands in the `runtime` part - here we just install all dependencies and it works -# so far. Given that there is more to come regading new inference, I would assume that the moment we fix -# the main inference dependencies management is when we should structure build differently - but, given how -# world works, sooner or later I will regret that. COPY inference_models/uv.lock inference_models/pyproject.toml build/inference_models/ WORKDIR build/inference_models/ RUN UV_PROJECT_ENVIRONMENT=/usr uv sync --locked --extra torch-cu124 --extra onnx-cu12 --extra trt10 @@ -62,51 +277,96 @@ RUN uv pip install --system \ "setuptools~=83.0.0" \ && rm -rf ~/.cache/pip -# Install setup.py requirements for flash_attn RUN python3 -m pip install packaging==24.1 && rm -rf ~/.cache/pip -# Install flash_attn required for Paligemma and Florence2 RUN python3 -m pip install -r requirements.pali.flash_attn.txt --no-dependencies --no-build-isolation && rm -rf ~/.cache/pip -# Start runtime stage -FROM nvcr.io/nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04 as runtime +COPY . . + +# The `create_wheels_for_gpu_notebook` Makefile target invokes bare `python`; +# the CUDA devel base only ships `python3`, so alias it here in the builder +# stage (the runtime stage has its own symlink further down). +RUN ln -sf /usr/bin/python3 /usr/bin/python + +RUN /bin/make create_wheels_for_gpu_notebook && \ + python3 -m pip install --no-cache-dir \ + dist/inference_cli*.whl \ + dist/inference_core*.whl \ + dist/inference_gpu*.whl \ + dist/inference_sdk*.whl \ + "setuptools~=83.0.0" + +RUN find /usr/local/lib/python3.10 -type d \ + \( -name cv2 -o -name 'opencv_python.libs' -o -name 'opencv_contrib_python.libs' -o -name 'opencv_python_headless.libs' -o -name 'opencv_contrib_python_headless.libs' \) \ + -prune -exec rm -rf {} + && \ + python3 -m pip uninstall -y uv && \ + rm -f /usr/local/bin/uv /usr/local/bin/uvx \ + /usr/local/bin/ninja /usr/local/bin/cmake /usr/local/bin/meson && \ + python3 -m pip uninstall -y pip + +FROM nvcr.io/nvidia/cuda:12.4.1-cudnn-runtime-ubuntu22.04 AS runtime + +ARG DEBIAN_FRONTEND=noninteractive +ARG TARGETARCH WORKDIR /app -# Copy Python and installed packages from builder +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-runtime-gpu-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-runtime-gpu-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + ca-certificates \ + libbz2-1.0 \ + libcurl4 \ + libegl1 \ + libexpat1 \ + libgdal30 \ + libgl1 \ + libglib2.0-0 \ + libgnutls30 \ + libgomp1 \ + libjpeg-turbo8 \ + liblzma5 \ + libopenblas0 \ + libopus0 \ + libpng16-16 \ + libsm6 \ + libsrtp2-1 \ + libssl3 \ + libtiff5 \ + libv4l-0 \ + libvips42 \ + libvpx7 \ + libxext6 \ + python3 \ + zlib1g + +COPY --from=gstreamer-runtime /opt/ffmpeg /opt/ffmpeg +COPY --from=gstreamer-runtime /opt/gstreamer /opt/gstreamer +COPY --from=opencv-builder /opt/opencv-runtime /opt/opencv +COPY --from=gstreamer-cuda-tensor-bridge-builder /opt/roboflow/lib /opt/roboflow/lib COPY --from=builder /usr/local/lib/python3.10 /usr/local/lib/python3.10 COPY --from=builder /usr/local/bin /usr/local/bin -# Install runtime dependencies -ADD https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb /tmp/cuda-keyring.deb -RUN set -eux; \ - rm -rf /var/lib/apt/lists/*; apt-get clean; \ - dpkg -i /tmp/cuda-keyring.deb || true; \ - rm -f /tmp/cuda-keyring.deb; \ - apt-get update -y; \ - DEBIAN_FRONTEND=noninteractive apt-get install -y \ - libxext6 \ - libopencv-dev \ - uvicorn \ - python3-pip \ - git \ - libgdal-dev \ - libvips-dev \ - wget \ - rustc \ - cargo; \ - rm -rf /var/lib/apt/lists/*; \ -# apt's uvicorn drags in python3-h11 0.13.0 (CVE-2025-43859); the served process uses -# pip's h11 from /usr/local, but scanners flag the apt copy - upgrade it in place in the -# same layer so the vulnerable version never lands in any image layer - pip3 install --no-cache-dir --upgrade --target=/usr/lib/python3/dist-packages "h11>=0.16.0"; \ - rm -rf /usr/lib/python3/dist-packages/h11-*.egg-info - -WORKDIR /build -COPY . . -RUN ln -s /usr/bin/python3 /usr/bin/python -RUN /bin/make create_wheels_for_gpu_notebook -RUN pip3 install --no-cache-dir dist/inference_cli*.whl dist/inference_core*.whl dist/inference_gpu*.whl dist/inference_sdk*.whl "setuptools~=83.0.0" +ENV PATH=/opt/gstreamer/bin:/opt/ffmpeg/bin:/usr/local/cuda/bin:$PATH \ + GIO_EXTRA_MODULES=/opt/gstreamer/lib/gio/modules \ + GI_TYPELIB_PATH=/opt/gstreamer/lib/girepository-1.0 \ + LD_LIBRARY_PATH=/opt/opencv/lib:/opt/gstreamer/lib:/opt/ffmpeg/lib:/usr/local/cuda/lib64:/usr/local/nvidia/lib:/usr/local/nvidia/lib64 \ + PYTHONPATH=/opt/opencv/python \ + GST_PLUGIN_SYSTEM_PATH_1_0=/opt/gstreamer/lib/gstreamer-1.0 \ + NVIDIA_DRIVER_CAPABILITIES=compute,utility,video + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +RUN --mount=type=bind,source=docker/scripts/verify_ffmpeg.sh,target=/tmp/verify_ffmpeg,ro \ + --mount=type=bind,source=docker/scripts/verify_gstreamer.sh,target=/tmp/verify_gstreamer,ro \ + --mount=type=bind,source=docker/scripts/verify_opencv.sh,target=/tmp/verify_opencv,ro \ + --mount=type=bind,source=docker/scripts/verify_media_runtime.sh,target=/tmp/verify_media_runtime,ro \ + sh /tmp/verify_ffmpeg && \ + GSTREAMER_REQUIRE_NVCODEC=true sh /tmp/verify_gstreamer && \ + sh /tmp/verify_media_runtime && \ + OPENCV_REQUIRE_CUDA_BUILD=true sh /tmp/verify_opencv && \ + rm -rf /root/.cache/gstreamer-1.0 # Instal Cosmos 3 preview (requires override of transformers, since 21.07.2026 it was not released yet). COPY requirements/requirements.cosmos.txt . @@ -115,7 +375,7 @@ RUN python3 -m pip install -r requirements.cosmos.txt WORKDIR /notebooks COPY examples/notebooks . -WORKDIR /app/ +WORKDIR /app COPY inference inference COPY build_scripts build_scripts RUN python3 build_scripts/download_fonts.py @@ -124,20 +384,20 @@ COPY docker/config/gpu_http.py gpu_http.py COPY docker/entrypoint/run_uvicorn.sh /usr/local/bin/run_uvicorn.sh RUN chmod +x /usr/local/bin/run_uvicorn.sh -ENV VERSION_CHECK_MODE=continuous -ENV PROJECT=roboflow-platform -ENV NUM_WORKERS=1 -ENV HOST=0.0.0.0 -ENV PORT=9001 -ENV WORKFLOWS_STEP_EXECUTION_MODE=local -ENV WORKFLOWS_MAX_CONCURRENT_STEPS=4 -ENV API_LOGGING_ENABLED=True -ENV LMM_ENABLED=True -ENV CORE_MODEL_SAM2_ENABLED=True -ENV CORE_MODEL_SAM3_ENABLED=True -ENV CORE_MODEL_OWLV2_ENABLED=True -ENV ENABLE_STREAM_API=True -ENV ENABLE_PROMETHEUS=True -ENV STREAM_API_PRELOADED_PROCESSES=2 - -ENTRYPOINT ["/usr/local/bin/run_uvicorn.sh", "gpu_http:app"] \ No newline at end of file +ENV VERSION_CHECK_MODE=continuous \ + PROJECT=roboflow-platform \ + NUM_WORKERS=1 \ + HOST=0.0.0.0 \ + PORT=9001 \ + WORKFLOWS_STEP_EXECUTION_MODE=local \ + WORKFLOWS_MAX_CONCURRENT_STEPS=4 \ + API_LOGGING_ENABLED=True \ + LMM_ENABLED=True \ + CORE_MODEL_SAM2_ENABLED=True \ + CORE_MODEL_SAM3_ENABLED=True \ + CORE_MODEL_OWLV2_ENABLED=True \ + ENABLE_STREAM_API=True \ + ENABLE_PROMETHEUS=True \ + STREAM_API_PRELOADED_PROCESSES=2 + +ENTRYPOINT ["/usr/local/bin/run_uvicorn.sh", "gpu_http:app"] diff --git a/docker/dockerfiles/Dockerfile.onnx.jetson.5.1.1 b/docker/dockerfiles/Dockerfile.onnx.jetson.5.1.1 index 5c02e70e16..0b311ecee4 100644 --- a/docker/dockerfiles/Dockerfile.onnx.jetson.5.1.1 +++ b/docker/dockerfiles/Dockerfile.onnx.jetson.5.1.1 @@ -1,8 +1,38 @@ # syntax=docker/dockerfile:1.7 -FROM roboflow/l4t-ml:r35.2.1-py3.12-cu118-trt-10-v0.0.3 +FROM roboflow/l4t-ml:r35.2.1-py3.12-cu118-trt-10-v0.0.4-experimental AS jetson-tensor-bridge-builder + +ARG DEBIAN_FRONTEND=noninteractive + +# JP5 (r35 / CUDA 11.8) build of the native tensor bridge - same source and +# scripts as the JP6 image. JetPack 5 runs on Xavier (sm_72) and Orin (sm_87), +# so fatbin both. The base ships nvcc, the jetson_multimedia_api headers and +# the GStreamer dev stack; nvv4l2decoder and libnvbufsurface come from the +# host L4T mount at runtime (--runtime nvidia), and the bridge's dlopen +# candidates already include the r35 /tegra/ library path. +RUN apt-get update -y && \ + apt-get install -y --no-install-recommends \ + build-essential \ + libglib2.0-dev \ + pkg-config && \ + rm -rf /var/lib/apt/lists/* + +COPY docker/native/jetson_tensor_bridge /src/jetson_tensor_bridge +COPY docker/scripts/build_jetson_tensor_bridge.sh /usr/local/bin/build_jetson_tensor_bridge +COPY docker/scripts/verify_jetson_tensor_bridge.sh /usr/local/bin/verify_jetson_tensor_bridge + +RUN chmod +x \ + /usr/local/bin/build_jetson_tensor_bridge \ + /usr/local/bin/verify_jetson_tensor_bridge && \ + JETSON_TENSOR_BRIDGE_CUDA_ARCHITECTURES="72;87" build_jetson_tensor_bridge && \ + verify_jetson_tensor_bridge && \ + du -sh /opt/roboflow/lib + +FROM roboflow/l4t-ml:r35.2.1-py3.12-cu118-trt-10-v0.0.4-experimental ARG TRITON_VERSION=3.0.0 +COPY --from=jetson-tensor-bridge-builder /opt/roboflow/lib /opt/roboflow/lib + RUN apt-get update -y && apt-get install -y ffmpeg COPY requirements/requirements.clip.txt \ @@ -102,7 +132,10 @@ ENV VERSION_CHECK_MODE=continuous \ HOST=0.0.0.0 \ PORT=9001 \ OPENBLAS_CORETYPE=ARMV8 \ - LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1 \ + # libGLdispatch must claim its static TLS at process start: loaded late + # (torch/ort/cv2 exhaust the r35 glibc surplus first) the tensor bridge's + # deferred EGL load dies with 'cannot allocate memory in static TLS block'. + LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1:/usr/lib/aarch64-linux-gnu/libGLdispatch.so.0 \ WORKFLOWS_STEP_EXECUTION_MODE=local \ WORKFLOWS_MAX_CONCURRENT_STEPS=2 \ API_LOGGING_ENABLED=True \ diff --git a/docker/dockerfiles/Dockerfile.onnx.jetson.6.2.0 b/docker/dockerfiles/Dockerfile.onnx.jetson.6.2.0 index 87fe6c9610..fc762e768c 100644 --- a/docker/dockerfiles/Dockerfile.onnx.jetson.6.2.0 +++ b/docker/dockerfiles/Dockerfile.onnx.jetson.6.2.0 @@ -1,31 +1,253 @@ # syntax=docker/dockerfile:1.7 -# Jetson 6.2.0 with PyTorch/OpenCV compiled from source for numpy 2.x support -# Stage 1: Builder - Compile everything from source -FROM nvcr.io/nvidia/l4t-jetpack:r36.4.0 AS builder +ARG FFMPEG_VERSION=7.1.3 +ARG FFMPEG_SHA256=f0bf043299db9e3caacb435a712fc541fbb07df613c4b893e8b77e67baf3adbe +ARG GSTREAMER_VERSION=1.24.12 +ARG GSTREAMER_COMMIT=88e31216479a3f0c25096ca406369f3983efec14 +ARG GLIB_NETWORKING_SHA256=cd2a084c7bb91d78e849fb55d40e472f6d8f6862cddc9f12c39149359ba18268 +ARG OPENCV_VERSION=4.13.0 +ARG OPENCV_SHA256=1d40ca017ea51c533cf9fd5cbde5b5fe7ae248291ddf2af99d4c17cf8e13017d +ARG OPENCV_CONTRIB_SHA256=1e0077a4fd2960a7d2f4c9e49d6ba7bb891cac2d1be36d7e8e47aa97a9d1039b +ARG OPENCV_CUDA_ARCH_BIN=8.7 +ARG OPENCV_CUDA_ARCH_PTX= + +FROM ubuntu:22.04 AS ffmpeg-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG FFMPEG_SHA256 +ARG FFMPEG_VERSION +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-ffmpeg-jp62-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-ffmpeg-jp62-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + libbz2-dev \ + liblzma-dev \ + libssl-dev \ + nasm \ + pkg-config \ + xz-utils \ + zlib1g-dev + +COPY docker/scripts/build_ffmpeg.sh /usr/local/bin/build_ffmpeg + +ENV PATH=/opt/ffmpeg/bin:$PATH \ + LD_LIBRARY_PATH=/opt/ffmpeg/lib \ + PKG_CONFIG_PATH=/opt/ffmpeg/lib/pkgconfig + +RUN chmod +x /usr/local/bin/build_ffmpeg && \ + FFMPEG_VERSION="${FFMPEG_VERSION}" \ + FFMPEG_SHA256="${FFMPEG_SHA256}" \ + build_ffmpeg && \ + du -sh /opt/ffmpeg /opt/ffmpeg-runtime + +COPY docker/scripts/verify_ffmpeg.sh /usr/local/bin/verify_ffmpeg + +RUN chmod +x /usr/local/bin/verify_ffmpeg && verify_ffmpeg + +FROM ffmpeg-builder AS gstreamer-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG GLIB_NETWORKING_SHA256 +ARG GSTREAMER_COMMIT +ARG GSTREAMER_VERSION +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-gstreamer-jp62-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-gstreamer-jp62-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + bison \ + build-essential \ + ca-certificates \ + cmake \ + flex \ + gobject-introspection \ + git \ + libcurl4-openssl-dev \ + libglib2.0-dev \ + libgnutls28-dev \ + libgirepository1.0-dev \ + libjpeg-dev \ + libopus-dev \ + libpng-dev \ + libsrtp2-dev \ + libssl-dev \ + libv4l-dev \ + libvpx-dev \ + ninja-build \ + pkg-config \ + python3 \ + python3-pip \ + python3-venv + +RUN python3 -m venv /opt/gstreamer-build-tools && \ + /opt/gstreamer-build-tools/bin/pip install --no-cache-dir \ + "meson>=1.4,<2" \ + "ninja>=1.11,<2" + +ENV PATH=/opt/gstreamer-build-tools/bin:$PATH + +COPY docker/scripts/build_gstreamer.sh /usr/local/bin/build_gstreamer + +RUN chmod +x /usr/local/bin/build_gstreamer && \ + GSTREAMER_VERSION="${GSTREAMER_VERSION}" \ + GSTREAMER_COMMIT="${GSTREAMER_COMMIT}" \ + GLIB_NETWORKING_SHA256="${GLIB_NETWORKING_SHA256}" \ + GSTREAMER_NVCODEC=disabled \ + build_gstreamer + +ENV PATH=/opt/gstreamer/bin:/opt/ffmpeg/bin:$PATH \ + GIO_EXTRA_MODULES=/opt/gstreamer/lib/gio/modules \ + GI_TYPELIB_PATH=/opt/gstreamer/lib/girepository-1.0 \ + LD_LIBRARY_PATH=/opt/gstreamer/lib:/opt/ffmpeg/lib \ + PKG_CONFIG_PATH=/opt/gstreamer/lib/pkgconfig:/opt/ffmpeg/lib/pkgconfig \ + GST_PLUGIN_SYSTEM_PATH_1_0=/opt/gstreamer/lib/gstreamer-1.0 + +COPY docker/scripts/verify_gstreamer.sh /usr/local/bin/verify_gstreamer + +RUN chmod +x /usr/local/bin/verify_gstreamer && \ + verify_gstreamer && \ + du -sh /opt/gstreamer /opt/gstreamer-runtime + +FROM ubuntu:22.04 AS gstreamer-runtime + +ARG DEBIAN_FRONTEND=noninteractive +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-gstreamer-runtime-jp62-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-gstreamer-runtime-jp62-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + ca-certificates \ + libbz2-1.0 \ + libcurl4 \ + libglib2.0-0 \ + libgnutls30 \ + libjpeg-turbo8 \ + liblzma5 \ + libopus0 \ + libsrtp2-1 \ + libssl3 \ + libv4l-0 \ + libvpx7 \ + zlib1g + +COPY --from=ffmpeg-builder /opt/ffmpeg-runtime /opt/ffmpeg +COPY --from=gstreamer-builder /opt/gstreamer-runtime /opt/gstreamer + +ENV PATH=/opt/gstreamer/bin:/opt/ffmpeg/bin:$PATH \ + GIO_EXTRA_MODULES=/opt/gstreamer/lib/gio/modules \ + GI_TYPELIB_PATH=/opt/gstreamer/lib/girepository-1.0 \ + LD_LIBRARY_PATH=/opt/gstreamer/lib:/opt/ffmpeg/lib \ + GST_PLUGIN_SYSTEM_PATH_1_0=/opt/gstreamer/lib/gstreamer-1.0 + +RUN --mount=type=bind,source=docker/scripts/verify_ffmpeg.sh,target=/tmp/verify_ffmpeg,ro \ + --mount=type=bind,source=docker/scripts/verify_gstreamer.sh,target=/tmp/verify_gstreamer,ro \ + --mount=type=bind,source=docker/scripts/verify_media_runtime.sh,target=/tmp/verify_media_runtime,ro \ + sh /tmp/verify_ffmpeg && \ + sh /tmp/verify_gstreamer && \ + sh /tmp/verify_media_runtime && \ + rm -rf /root/.cache/gstreamer-1.0 + +FROM nvcr.io/nvidia/l4t-jetpack:r36.4.0 AS jetson-build-base + +COPY --from=ffmpeg-builder /opt/ffmpeg /opt/ffmpeg +COPY --from=gstreamer-builder /opt/gstreamer /opt/gstreamer + +ENV PATH=/opt/gstreamer/bin:/opt/ffmpeg/bin:/usr/local/cuda/bin:$PATH \ + GIO_EXTRA_MODULES=/opt/gstreamer/lib/gio/modules \ + GI_TYPELIB_PATH=/opt/gstreamer/lib/girepository-1.0 \ + LD_LIBRARY_PATH=/usr/local/cuda/lib64:/opt/gstreamer/lib:/opt/ffmpeg/lib \ + PKG_CONFIG_PATH=/opt/gstreamer/lib/pkgconfig:/opt/ffmpeg/lib/pkgconfig \ + GST_PLUGIN_SYSTEM_PATH_1_0=/opt/gstreamer/lib/gstreamer-1.0 + +FROM jetson-build-base AS jetson-tensor-bridge-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jetson-tensor-bridge-jp62-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jetson-tensor-bridge-jp62-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + build-essential \ + libglib2.0-dev \ + pkg-config + +COPY docker/native/jetson_tensor_bridge /src/jetson_tensor_bridge +COPY docker/scripts/build_jetson_tensor_bridge.sh /usr/local/bin/build_jetson_tensor_bridge +COPY docker/scripts/verify_jetson_tensor_bridge.sh /usr/local/bin/verify_jetson_tensor_bridge + +RUN chmod +x \ + /usr/local/bin/build_jetson_tensor_bridge \ + /usr/local/bin/verify_jetson_tensor_bridge && \ + build_jetson_tensor_bridge && \ + verify_jetson_tensor_bridge && \ + du -sh /opt/roboflow/lib + +FROM jetson-build-base AS opencv-builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG OPENCV_VERSION +ARG OPENCV_SHA256 +ARG OPENCV_CONTRIB_SHA256 +ARG OPENCV_CUDA_ARCH_BIN +ARG OPENCV_CUDA_ARCH_PTX +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-opencv-jp62-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-opencv-jp62-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + cmake \ + curl \ + libjpeg-dev \ + libopenblas-dev \ + libpng-dev \ + libtiff-dev \ + libv4l-dev \ + ninja-build \ + pkg-config \ + python3-dev \ + python3-pip \ + python3-venv + +RUN python3 -m venv /opt/opencv-build-tools && \ + /opt/opencv-build-tools/bin/pip install --no-cache-dir \ + "numpy>=2.0.0,<2.4.0" + +COPY docker/scripts/build_opencv.sh /usr/local/bin/build_opencv +COPY docker/scripts/verify_opencv.sh /usr/local/bin/verify_opencv + +RUN chmod +x /usr/local/bin/build_opencv /usr/local/bin/verify_opencv && \ + OPENCV_VERSION="${OPENCV_VERSION}" \ + OPENCV_SHA256="${OPENCV_SHA256}" \ + OPENCV_CONTRIB_SHA256="${OPENCV_CONTRIB_SHA256}" \ + OPENCV_PYTHON_EXECUTABLE=/opt/opencv-build-tools/bin/python3 \ + OPENCV_PYTHON3_LIMITED_API=ON \ + OPENCV_BUILD_PYTHON3=OFF \ + OPENCV_NUMPY_TARGET_VERSION=NPY_2_0_API_VERSION \ + OPENCV_WITH_CUDA=ON \ + OPENCV_CUDA_ARCH_BIN="${OPENCV_CUDA_ARCH_BIN}" \ + OPENCV_CUDA_ARCH_PTX="${OPENCV_CUDA_ARCH_PTX}" \ + build_opencv && \ + du -sh /opt/opencv /opt/opencv-runtime + +FROM jetson-build-base AS builder ARG DEBIAN_FRONTEND=noninteractive ARG CMAKE_VERSION=4.2.0 ARG PYTORCH_VERSION=2.6.0 ARG TORCHVISION_VERSION=0.21.0 ARG TRITON_VERSION=3.2.0 -ARG OPENCV_VERSION=4.13.0 -# Per-Jetson-platform OpenCV build parameters. To replicate this from-source -# OpenCV section in another Jetson dockerfile (jp51, jp71), copy the OpenCV -# block below verbatim and update these four ARGs to match the target base. -# Cross-reference the values already used by the same dockerfile's PyTorch / -# torchvision / onnxruntime build steps so all components target the same SoC: -# OPENCV_CUDA_ARCH: jp51 (Orin on L4T r35): 8.7 -# jp62 (Orin on L4T r36): 8.7 -# jp71 (Thor on L4T r38): 11.0 -# OPENCV_PYTHON_INSTALL_PATH: dist-packages dir of the base's Python -# (jp51: python3.12, jp62: python3.10, jp71: python3.12) -# OPENCV_PYTHON_INCLUDE_DIR: matches the base's Python C headers -# OPENCV_PYTHON_VERSION: cmake's PYTHON_VERSION flag (major+minor, no dot) -ARG OPENCV_CUDA_ARCH=8.7 -ARG OPENCV_PYTHON_INSTALL_PATH=/usr/local/lib/python3.10/dist-packages -ARG OPENCV_PYTHON_INCLUDE_DIR=/usr/include/python3.10 -ARG OPENCV_PYTHON_VERSION=310 ARG ONNXRUNTIME_VERSION=1.20.0 +ARG TARGETARCH ENV LANG=en_US.UTF-8 WORKDIR /build @@ -58,18 +280,10 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp62-bu pkg-config \ libjpeg-dev \ libpng-dev \ - libavcodec-dev \ - libavformat-dev \ - libswscale-dev \ libv4l-dev \ - libavutil-dev \ - libgstreamer1.0-dev \ - libgstreamer-plugins-base1.0-dev \ - unzip \ libatlas-base-dev \ liblapack-dev \ gfortran \ - ffmpeg \ libhdf5-dev \ libgflags-dev \ libgoogle-glog-dev \ @@ -113,6 +327,11 @@ RUN --mount=type=cache,target=/build/pytorch/pytorch/build,sharing=locked,id=pyt --mount=type=cache,target=/root/.cache/pip,sharing=locked,id=pip-cache-jp62-${TARGETARCH} \ cd pytorch && \ python3 -m pip install -r requirements.txt && \ + # The v2.6 source checkout omits generated AOTInductor C shims. Generate + # them before setup.py configures CMake, then clear this cache mount: the + # generator's generic ATen outputs must not be mixed with setup.py's own. + PYTHONPATH=. python3 torchgen/gen.py --source-path aten/src/ATen && \ + rm -rf build/* && \ export USE_CUDA=1 USE_CUDNN=1 CUDA_HOME=/usr/local/cuda && \ export CUDNN_LIB_DIR=/usr/lib/aarch64-linux-gnu CUDNN_INCLUDE_DIR=/usr/include && \ export TORCH_CUDA_ARCH_LIST="8.7" && \ @@ -150,19 +369,17 @@ ENV CUDA_HOME=/usr/local/cuda \ LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH # Install Tensorrt +# For clean-checkout builds, use docker/scripts/build_jetson_6_2_image.sh; +# it fetches and verifies this package before invoking BuildKit. RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp62-builder-${TARGETARCH} \ --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jp62-builder-${TARGETARCH} \ apt remove -y 'libnvinfer*' 'libnvonnxparsers*' 'libnvparsers*' 'libnvinfer-plugin*' 'python3-libnvinfer*' 'tensorrt*' WORKDIR /build/tensorrt-10.x +COPY docker/vendor/tensorrt/nv-tensorrt-local-tegra-repo-ubuntu2204-10.7.0-cuda-12.6_1.0-1_arm64.deb ./_dl/ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp62-builder-${TARGETARCH} \ --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jp62-builder-${TARGETARCH} \ - --mount=type=cache,target=/build/tensorrt-10.x/_dl,sharing=locked,id=trt-deb-10.7.0-${TARGETARCH} \ - if [ ! -f _dl/nv-tensorrt-local-tegra-repo-ubuntu2204-10.7.0-cuda-12.6_1.0-1_arm64.deb ]; then \ - mkdir -p _dl && \ - wget -O _dl/nv-tensorrt-local-tegra-repo-ubuntu2204-10.7.0-cuda-12.6_1.0-1_arm64.deb \ - https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.7.0/local_repo/nv-tensorrt-local-tegra-repo-ubuntu2204-10.7.0-cuda-12.6_1.0-1_arm64.deb; \ - fi && \ - dpkg -i _dl/nv-tensorrt-local-tegra-repo-ubuntu2204-10.7.0-cuda-12.6_1.0-1_arm64.deb && \ + deb="_dl/nv-tensorrt-local-tegra-repo-ubuntu2204-10.7.0-cuda-12.6_1.0-1_arm64.deb" && \ + dpkg -i "${deb}" && \ cp /var/nv-tensorrt-local-tegra-repo-ubuntu2204-10.7.0-cuda-12.6/nv-tensorrt-local-tegra-C50F04B9-keyring.gpg /usr/share/keyrings/ && \ apt-get update && \ apt-get install -y tensorrt @@ -268,7 +485,11 @@ COPY requirements/_requirements.txt \ ./ # Install Python dependencies (torch/torchvision/numpy/onnxruntime already installed from source) -RUN uv pip install --system --break-system-packages --index-strategy unsafe-best-match \ +RUN echo "opencv-python==4.12.0.88" > /tmp/jetson-opencv-overrides.txt && \ + echo "opencv-contrib-python==4.12.0.88" >> /tmp/jetson-opencv-overrides.txt && \ + echo "opencv-python-headless==4.12.0.88" >> /tmp/jetson-opencv-overrides.txt && \ + uv pip install --system --break-system-packages --index-strategy unsafe-best-match \ + --override /tmp/jetson-opencv-overrides.txt \ -r _requirements.txt \ -r requirements.http.txt \ -r requirements.clip.txt \ @@ -284,7 +505,7 @@ RUN uv pip install --system --break-system-packages --index-strategy unsafe-best "setuptools~=83.0.0" \ "transformers>=4.57.3,<5.9.0" \ packaging \ - && rm -rf ~/.cache/uv + && rm -rf ~/.cache/uv /tmp/jetson-opencv-overrides.txt # Rebuild the bitsandbytes CUDA backend for Orin (sm87). The PyPI aarch64 # wheel ships sm90-only kernels, so every GPU quantization op (the VLM @@ -316,73 +537,6 @@ RUN ln -sf /usr/bin/python3 /usr/bin/python && \ dist/inference_sdk*.whl \ "setuptools~=83.0.0" -# Compile OpenCV from source with GStreamer / FFmpeg / CUDA support. -# The pip opencv-python wheel ships without GStreamer, which blocks -# cv2.VideoCapture from dispatching to NVIDIA's nv* GStreamer elements -# (nvv4l2decoder, nvvidconv, nvjpegenc/dec). Built LAST in the builder so -# nothing reinstalls the wheel-version cv2 over our from-source one. -WORKDIR /build/opencv -RUN wget -q https://github.com/opencv/opencv/archive/refs/tags/${OPENCV_VERSION}.zip -O opencv.zip && \ - wget -q https://github.com/opencv/opencv_contrib/archive/refs/tags/${OPENCV_VERSION}.zip -O opencv_contrib.zip && \ - unzip -q opencv.zip && unzip -q opencv_contrib.zip && \ - rm opencv.zip opencv_contrib.zip && \ - python3 -m pip uninstall -y opencv-python opencv-contrib-python opencv-python-headless || true && \ - mkdir -p /build/opencv/opencv-${OPENCV_VERSION}/release -WORKDIR /build/opencv/opencv-${OPENCV_VERSION}/release -# release/ is cache-mounted, so the rebuild forced by `COPY . .` invalidation -# above completes as a fast ninja no-op (CMakeCache + objects survive) instead -# of a 30-60 min recompile. Cache id includes OPENCV_VERSION and CUDA_ARCH so -# version/SoC bumps start clean. -RUN --mount=type=cache,target=/build/opencv/opencv-${OPENCV_VERSION}/release,sharing=locked,id=opencv-build-${OPENCV_VERSION}-${TARGETARCH}-sm${OPENCV_CUDA_ARCH} \ - --mount=type=cache,target=/root/.cache/pip,sharing=locked,id=pip-cache-jp62-${TARGETARCH} \ - cmake -GNinja \ - -D CMAKE_BUILD_TYPE=RELEASE \ - -D CMAKE_INSTALL_PREFIX=/usr/local \ - -D WITH_CUDA=ON \ - -D WITH_CUDNN=ON \ - -D CUDA_ARCH_BIN="${OPENCV_CUDA_ARCH}" \ - -D CUDA_ARCH_PTX="" \ - -D OPENCV_GENERATE_PKGCONFIG=ON \ - -D OPENCV_EXTRA_MODULES_PATH=/build/opencv/opencv_contrib-${OPENCV_VERSION}/modules \ - -D WITH_FFMPEG=ON \ - -D WITH_GSTREAMER=ON \ - -D WITH_LIBV4L=ON \ - -D BUILD_opencv_python3=ON \ - -D BUILD_TESTS=OFF \ - -D BUILD_PERF_TESTS=OFF \ - -D BUILD_EXAMPLES=OFF \ - -D BUILD_DOCS=OFF \ - -D PYTHON3_EXECUTABLE=/usr/bin/python3 \ - -D PYTHON3_INCLUDE_DIR=${OPENCV_PYTHON_INCLUDE_DIR} \ - -D PYTHON_VERSION=${OPENCV_PYTHON_VERSION} \ - -D OPENCV_PYTHON3_INSTALL_PATH=${OPENCV_PYTHON_INSTALL_PATH} \ - -D BUILD_SHARED_LIBS=OFF \ - -D WITH_OPENCLAMDFFT=OFF \ - -D WITH_OPENCLAMDBLAS=OFF \ - -D WITH_VA_INTEL=OFF \ - .. && \ - ninja && \ - ninja install && \ - ldconfig && \ - # Rescue install-tree config files before pip install clobbers them. - # `ninja install` writes ${OPENCV_PYTHON_INSTALL_PATH}/cv2/config*.py with - # install paths (correct). The `pip install` below installs a wheel built - # from the BUILD-TREE python_loader directory, whose config files reference - # /build/... paths that don't survive into the runtime stage. Save the - # install-tree versions now, restore them after pip install registers - # opencv with pip's metadata. Without this, cv2 import in the runtime image - # fails with "recursion is detected during loading of cv2 binary extensions". - mkdir -p /tmp/cv2-install-tree && \ - cp ${OPENCV_PYTHON_INSTALL_PATH}/cv2/config*.py /tmp/cv2-install-tree/ && \ - python3 -m pip wheel /build/opencv/opencv-${OPENCV_VERSION}/release/python_loader --wheel-dir /build/out/wheels && \ - python3 -m pip install --break-system-packages /build/out/wheels/opencv-*.whl && \ - cp /tmp/cv2-install-tree/config*.py ${OPENCV_PYTHON_INSTALL_PATH}/cv2/ && \ - rm -rf /tmp/cv2-install-tree && \ - # Sanity check: simulate the runtime stage by removing /build/ from sys.path. - # If config files still leak build-tree paths, the loader bootstrap fails here - # with a recursion error rather than at runtime on the device. - env -u PYTHONPATH python3 -c "import sys; sys.path = [p for p in sys.path if '/build' not in p]; import cv2; info = cv2.getBuildInformation(); assert 'YES' in info.split('GStreamer:')[1].split(chr(10))[0], 'GStreamer not enabled'; assert 'YES' in info.split('FFMPEG:')[1].split(chr(10))[0], 'FFMPEG not enabled'; print('cv2 ' + cv2.__version__ + ' OK from install path; GStreamer + FFmpeg + CUDA enabled')" - # Pin boto3/botocore last so the final versions win regardless of what the # requirement sets resolved to (previously done as an extra layer in the # runtime stage, which shipped both copies). @@ -452,38 +606,41 @@ FROM nvcr.io/nvidia/l4t-cuda:12.6.11-runtime AS cuda_rt # in the image). Net ~300MB smaller. GPU/BSP access is unchanged: on JetPack 6 # the nvidia-container-runtime injects libcuda/tegra libs from the host at # container start, driven by the NVIDIA_* envs replicated below. -FROM ubuntu:22.04 +FROM gstreamer-runtime AS runtime ARG DEBIAN_FRONTEND=noninteractive +ARG TARGETARCH ENV LANG=en_US.UTF-8 # Replicated verbatim from the l4t-cuda:12.6.11-runtime base / published image. ENV NVIDIA_VISIBLE_DEVICES=all \ - NVIDIA_DRIVER_CAPABILITIES=compute,utility \ + NVIDIA_DRIVER_CAPABILITIES=compute,utility,video \ NVIDIA_REQUIRE_CUDA="cuda>=12.6" \ NVIDIA_REQUIRE_JETPACK_HOST_MOUNTS=base-only \ NVIDIA_PRODUCT_NAME=CUDA \ - PATH=/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin + PATH=/opt/gstreamer/bin:/opt/ffmpeg/bin:/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin # Library search paths the l4t-cuda base provided via ld.so.conf.d (tegra/nvidia # dirs are populated by the host mounts injected at container start). -RUN printf '/usr/local/cuda/lib64\n/usr/lib/aarch64-linux-gnu/tegra\n/usr/lib/aarch64-linux-gnu/tegra-egl\n/usr/lib/aarch64-linux-gnu/nvidia\n/usr/lib/aarch64-linux-gnu/gstreamer-1.0\n' > /etc/ld.so.conf.d/nvidia-tegra.conf +RUN mkdir -p /usr/lib/aarch64-linux-gnu/gstreamer-1.0 && \ + printf '/usr/local/cuda/lib64\n/usr/lib/aarch64-linux-gnu/tegra\n/usr/lib/aarch64-linux-gnu/tegra-egl\n/usr/lib/aarch64-linux-gnu/nvidia\n/usr/lib/aarch64-linux-gnu/gstreamer-1.0\n' > /etc/ld.so.conf.d/nvidia-tegra.conf WORKDIR /app -RUN ln -sf /usr/bin/python3 /usr/bin/python - # Install runtime dependencies (cache id is runtime-specific, distinct from builder). # Trimmed relative to the original list (all verified unreferenced on-device): # file (libmagic; nothing executes it), libproj22 (+proj-data; only served the -# removed system GDAL - rasterio bundles its own proj), python3-pip (pip already -# ships in dist-packages and the runtime stage no longer runs pip), -# gstreamer1.0-tools (debug CLIs), python3-gi + gir1.2-* (no `import gi` -# anywhere; cv2 links the GStreamer C libs directly), libatlas3-base (zero .so -# NEEDs ATLAS; BLAS is libopenblas0). +# removed system GDAL - rasterio bundles its own proj), # libegl1/libgles2 (glvnd loaders, ~1.8MB) are required by the host-injected NV # GStreamer plugins (nvarguscamerasrc/nvivafilter/nveglglessink/...); the actual # EGL/GLES drivers come from the host's tegra-egl mount at container start. +# libgl1 (glvnd desktop-GL loader) provides libGL.so.1, linked by the Qt5 libs +# bundled inside the opencv-python wheel (manylinux forbids vendoring libGL). +# kmod (lsmod/modprobe) is shelled out to by the host-injected NV BSP +# v4l2/GStreamer libs to check the codec kernel modules; without it every +# nvv4l2decoder plugin load spams "sh: 1: lsmod: not found" and the module +# check degrades. The media.jetson.7.2.0 image installs it for the same reason +# (verify_media_runtime.sh requires lsmod whenever BSP plugins are baked). # The dpkg path-exclude keeps docs/man/locale out of every package this layer installs. RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp62-runtime-${TARGETARCH} \ --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jp62-runtime-${TARGETARCH} \ @@ -496,42 +653,41 @@ RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp62-ru libopenblas0 \ libsqlite3-0 \ libtiff5 \ - libcurl4 \ - libssl3 \ - zlib1g \ libgomp1 \ python3 \ libxext6 \ libvips42 \ - libglib2.0-0 \ libsm6 \ - libjpeg-turbo8 \ libpng16-16 \ libexpat1 \ - ca-certificates \ curl \ - libv4l-0 \ - libavcodec58 \ - libavformat58 \ - libswscale5 \ - ffmpeg \ - libgstreamer1.0-0 \ - libgstreamer-plugins-base1.0-0 \ - gstreamer1.0-plugins-base \ - gstreamer1.0-plugins-good \ + kmod \ + libasound2 \ + libdrm2 \ libegl1 \ + libgl1 \ libgles2 \ + libpcre3 \ + libwayland-client0 \ + libwayland-egl1 \ + libwebp7 \ + libwebpdemux2 \ + libwebpmux3 \ + libx11-6 \ build-essential \ python3-dev \ zlib1g-dev \ - libxml2-dev + libxml2-dev && \ + ln -sf /usr/bin/python3 /usr/bin/python && \ + python --version && \ + python3 --version # No system GDAL: rasterio (the only GDAL consumer, via requirements.sam.txt) # installs as a manylinux wheel with its own bundled libgdal/libproj/gdal-data; # zero .so in the image links the from-source libgdal (verified via ldd scan + # live deletion test). GDAL_DATA must stay unset or it would override # rasterio's bundled data directory and break CRS handling. -ENV LD_LIBRARY_PATH=/usr/local/lib:$LD_LIBRARY_PATH +ENV LD_LIBRARY_PATH=/usr/local/lib:/usr/local/cuda/lib64:/opt/opencv/lib:/opt/gstreamer/lib:/opt/ffmpeg/lib # CUDA runtime libraries, exact files (not globs: the source dir mixes real # files and symlinks, and COPY dereferences symlinks, which would triple-copy @@ -594,14 +750,6 @@ COPY --from=builder /usr/lib/aarch64-linux-gnu/libnvinfer.so.*.*.* \ COPY --from=builder /usr/lib/aarch64-linux-gnu/libnvonnxparser*.so.*.*.* /usr/local/cuda/lib64/ COPY --from=builder /usr/lib/aarch64-linux-gnu/nvidia/libnvdla_compiler.so* /usr/local/cuda/lib64/ -# NOTE: NVIDIA's gstreamer plugins (nvv4l2decoder, nvvidconv, nvjpegenc/dec) and -# their tegra runtime deps (libnvbufsurface, libnvjpeg, libnvargus, ...) are NOT -# baked into the image -- they are tied to the host's JetPack BSP and are injected -# at container start by `nvidia-container-runtime` via the CSVs under -# /etc/nvidia-container-runtime/host-files-for-container.d/. The host system must -# be configured to mount them; without that, cv2 here can dispatch a GStreamer -# pipeline but `nvv4l2decoder` / `nvvidconv` / `nvjpegenc` won't resolve. - # Create symlinks for CUDA/cuDNN/TensorRT (saves ~2GB by not duplicating files). # libnvrtc-builtins needs its two-level soname on top of the .so/.so.MAJ links # the loop creates: libnvrtc dlopens it as libnvrtc-builtins.so.12.6. @@ -614,12 +762,11 @@ RUN cd /usr/local/cuda/lib64 && \ done && \ ln -sf libnvrtc-builtins.so.12.6.68 libnvrtc-builtins.so.12.6 -ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH RUN ldconfig # GStreamer plugin discovery: point at the standard arch dir (where libgstnv*.so live). # Clear stale plugin registry so it's rebuilt on first use against the copied plugins. -ENV GST_PLUGIN_PATH=/usr/lib/aarch64-linux-gnu/gstreamer-1.0 +ENV GST_PLUGIN_PATH_1_0=/usr/lib/aarch64-linux-gnu/gstreamer-1.0 RUN rm -rf /root/.cache/gstreamer-1.0 /home/*/.cache/gstreamer-1.0 2>/dev/null || true # build-essential/python3-dev/zlib1g-dev/libxml2-dev: required for Triton JIT kernel @@ -628,6 +775,8 @@ RUN rm -rf /root/.cache/gstreamer-1.0 /home/*/.cache/gstreamer-1.0 2>/dev/null | # JIT toolchains are copied separately (llvm cache excluded at build time). # Copy Python packages +COPY --from=opencv-builder /opt/opencv-runtime /opt/opencv +COPY --from=jetson-tensor-bridge-builder /opt/roboflow/lib /opt/roboflow/lib COPY --from=builder /usr/local/lib/python3.10/dist-packages /usr/local/lib/python3.10/dist-packages COPY --from=builder /root/.triton/nvidia /root/.triton/nvidia COPY --from=builder /root/.triton/json /root/.triton/json @@ -638,8 +787,11 @@ COPY --from=builder /usr/lib/python3.10/dist-packages/tensorrt-10.7.0.dist-info COPY --from=builder /usr/lib/python3.10/dist-packages/tensorrt_lean /usr/local/lib/python3.10/dist-packages/tensorrt_lean COPY --from=builder /usr/lib/python3.10/dist-packages/tensorrt_lean-10.7.0.dist-info /usr/local/lib/python3.10/dist-packages/tensorrt_lean-10.7.0.dist-info COPY --from=builder /usr/local/bin/inference /usr/local/bin/inference +COPY --from=builder /usr/local/bin/pip /usr/local/bin/pip3 /usr/local/bin/pip3.10 /usr/local/bin/ -ENV PYTHONPATH=/usr/local/lib/python3.10/dist-packages:$PYTHONPATH +RUN pip --version && pip3 --version && python3 -m pip --version + +ENV PYTHONPATH=/usr/local/lib/python3.10/dist-packages # Copy application code COPY inference inference @@ -652,6 +804,26 @@ COPY docker/config/gpu_http.py gpu_http.py COPY docker/entrypoint/run_uvicorn.sh /usr/local/bin/run_uvicorn.sh RUN chmod +x /usr/local/bin/run_uvicorn.sh +RUN python3 -c "import cv2, numpy; print(cv2.__version__); print(numpy.__version__)" + +# The Jetson tensor bridge is built and verified in its builder stage, but this +# final stage only COPYs it in โ€” assert it actually landed (a bridge-less image +# silently degrades HW video decode to the cv2 CPU path at runtime) and that +# every linked dependency resolves in THIS stage. libcuda.so.1 is excluded from +# the ldd check: the NVIDIA container runtime injects it at container start, so +# it is legitimately absent at build time. +RUN --mount=type=bind,source=docker/scripts/verify_ffmpeg.sh,target=/tmp/verify_ffmpeg,ro \ + --mount=type=bind,source=docker/scripts/verify_gstreamer.sh,target=/tmp/verify_gstreamer,ro \ + sh /tmp/verify_ffmpeg && \ + sh /tmp/verify_gstreamer && \ + python3 -c "import cv2, numpy; assert cv2.__version__.startswith('4.12.'); assert numpy.__version__.startswith('2.')" && \ + test -s /opt/roboflow/lib/libroboflow_jetson_tensor.so.1 && \ + missing="$(ldd /opt/roboflow/lib/libroboflow_jetson_tensor.so.1 | awk '/not found/ { print $1 }' | grep -v -x libcuda.so.1 || true)" && \ + test -z "${missing}" && \ + command -v lsmod >/dev/null && \ + command -v modprobe >/dev/null && \ + rm -rf /root/.cache/gstreamer-1.0 + # Environment variables ENV VERSION_CHECK_MODE=once \ CORE_MODEL_SAM2_ENABLED=True \ diff --git a/docker/dockerfiles/Dockerfile.onnx.jetson.7.2.0 b/docker/dockerfiles/Dockerfile.onnx.jetson.7.2.0 new file mode 100644 index 0000000000..9c9169ed16 --- /dev/null +++ b/docker/dockerfiles/Dockerfile.onnx.jetson.7.2.0 @@ -0,0 +1,298 @@ +# syntax=docker/dockerfile:1.7 +ARG JETSON_MEDIA_IMAGE=jetson-media +ARG JETSON_WHEELS_IMAGE=jetson-wheels + +FROM ${JETSON_WHEELS_IMAGE} AS jetson-wheels-base + +FROM ubuntu:24.04 AS nvidia-repository-key + +ARG DEBIAN_FRONTEND=noninteractive +ARG TARGETARCH + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp72-server-key-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jp72-server-key-${TARGETARCH} \ + apt-get update -y && \ + apt-get install -y --no-install-recommends ca-certificates curl && \ + curl -fsSL https://repo.download.nvidia.com/jetson/jetson-ota-public.asc \ + -o /nvidia-jetson.asc + +FROM ${JETSON_MEDIA_IMAGE} AS jetson-runtime-base + +ARG DEBIAN_FRONTEND=noninteractive +ARG L4T_REPOSITORY=r39.2 +ARG TARGETARCH + +WORKDIR /app + +COPY --from=nvidia-repository-key /nvidia-jetson.asc /usr/share/keyrings/nvidia-jetson.asc + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp72-server-runtime-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jp72-server-runtime-${TARGETARCH} \ + echo "deb [signed-by=/usr/share/keyrings/nvidia-jetson.asc] https://repo.download.nvidia.com/jetson/common ${L4T_REPOSITORY} main" > /etc/apt/sources.list.d/nvidia-jetson.list && \ + echo "deb [signed-by=/usr/share/keyrings/nvidia-jetson.asc] https://repo.download.nvidia.com/jetson/som ${L4T_REPOSITORY} main" >> /etc/apt/sources.list.d/nvidia-jetson.list && \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + cuda-nvrtc-13-2 \ + gcc \ + libatlas3-base \ + libatomic1 \ + libcublas-13-2 \ + libcudnn9-cuda-13 \ + libcufft-13-2 \ + libcufile-13-2 \ + libcudla-13-2 \ + libcuobjclient-13-2 \ + libcurand-13-2 \ + libcusolver-13-2 \ + libcusparse-13-2 \ + libexpat1 \ + libgomp1 \ + libnvfatbin-13-2 \ + libnvjitlink-13-2 \ + libopenblas0 \ + libpcre3 \ + libproj25 \ + libsm6 \ + libsqlite3-0 \ + libvips42 \ + python3.12-dev \ + python3-libnvinfer \ + python3-libnvinfer-dispatch \ + python3-libnvinfer-lean \ + tensorrt-libs && \ + rm -f \ + /etc/apt/sources.list.d/nvidia-jetson.list \ + /usr/share/keyrings/nvidia-jetson.asc + +FROM ubuntu:24.04 AS builder + +ARG DEBIAN_FRONTEND=noninteractive +ARG PYTORCH_VERSION=2.10.0 +ARG TORCHVISION_VERSION=0.25.0 +ARG TARGETARCH +ENV LANG=en_US.UTF-8 + +WORKDIR /build + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp72-onnx-server-builder-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jp72-onnx-server-builder-${TARGETARCH} \ + rm -f /etc/apt/apt.conf.d/docker-clean && \ + echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' > /etc/apt/apt.conf.d/keep-cache && \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + build-essential \ + ca-certificates \ + curl \ + git \ + python3-dev \ + python3-pip \ + python3-venv + +RUN python3 -m pip install --break-system-packages --ignore-installed pip setuptools wheel + +RUN curl -LsSf https://astral.sh/uv/install.sh | env INSTALLER_NO_MODIFY_PATH=1 sh && \ + ln -s /root/.local/bin/uv /usr/local/bin/uv + +COPY --from=jetson-wheels-base /opt/wheels /opt/wheels + +# The framework wheels are compiled in Dockerfile.wheels.jetson.7.2.0; the version +# ARGs above must match its output or resolution would silently fall back to PyPI. +RUN set -- /opt/wheels/torch-${PYTORCH_VERSION}*.whl && test -e "$1" && \ + set -- /opt/wheels/torchvision-${TORCHVISION_VERSION}*.whl && test -e "$1" && \ + set -- /opt/wheels/onnxruntime_gpu-*.whl && test -e "$1" && \ + set -- /opt/wheels/flash_attn-*.whl && test -e "$1" && \ + set -- /opt/wheels/triton-*.whl && test -e "$1" && \ + set -- /opt/wheels/pycuda-*.whl && test -e "$1" + +WORKDIR /build/reqs +COPY requirements/_requirements.txt \ + requirements/requirements.cli.txt \ + requirements/requirements.http.txt \ + requirements/requirements.clip.txt \ + requirements/requirements.transformers.txt \ + requirements/requirements.sam.txt \ + requirements/requirements.groundingdino.txt \ + requirements/requirements.yolo_world.txt \ + requirements/requirements.doctr.txt \ + requirements/requirements.sdk.http.txt \ + requirements/requirements.easyocr.txt \ + requirements/requirements.jetson.txt \ + ./ + +RUN echo "torch==${PYTORCH_VERSION}" > /tmp/overrides.txt && \ + echo "torchvision==${TORCHVISION_VERSION}" >> /tmp/overrides.txt && \ + uv pip install --system --break-system-packages --index-strategy unsafe-best-match \ + --find-links /opt/wheels \ + --override /tmp/overrides.txt \ + -r _requirements.txt \ + -r requirements.cli.txt \ + -r requirements.http.txt \ + -r requirements.clip.txt \ + -r requirements.transformers.txt \ + -r requirements.sam.txt \ + -r requirements.groundingdino.txt \ + -r requirements.yolo_world.txt \ + -r requirements.doctr.txt \ + -r requirements.sdk.http.txt \ + -r requirements.easyocr.txt \ + -r requirements.jetson.txt \ + /opt/wheels/flash_attn-*.whl \ + /opt/wheels/onnxruntime_gpu-*.whl \ + /opt/wheels/pycuda-*.whl \ + /opt/wheels/triton-*.whl \ + "setuptools<=75.5.0" \ + "transformers>=4.57.3,<5.9.0" \ + packaging \ + && rm -rf ~/.cache/uv /tmp/overrides.txt + +WORKDIR /build/inference +COPY . . +RUN ln -sf /usr/bin/python3 /usr/bin/python && \ + python -m pip install --break-system-packages wheel twine requests && \ + rm -f dist/* && \ + python .release/pypi/inference.core.setup.py bdist_wheel && \ + python .release/pypi/inference.gpu.setup.py bdist_wheel && \ + python .release/pypi/inference.cli.setup.py bdist_wheel && \ + python .release/pypi/inference.sdk.setup.py bdist_wheel && \ + python -m pip install --break-system-packages --no-deps dist/inference_gpu*.whl && \ + python -m pip install --break-system-packages \ + dist/inference_core*.whl \ + dist/inference_cli*.whl \ + dist/inference_sdk*.whl \ + "setuptools<=75.5.0" && \ + cp dist/*.whl /opt/wheels/ && \ + find /usr/local/lib/python3.12/dist-packages -type d \ + \( -name cv2 -o -name 'opencv_python.libs' -o -name 'opencv_contrib_python.libs' -o -name 'opencv_python_headless.libs' -o -name 'opencv_contrib_python_headless.libs' \) \ + -prune -exec rm -rf {} + + +RUN cd /usr/local/lib/python3.12/dist-packages && \ + find . -type d -name __pycache__ -exec rm -rf {} + 2>/dev/null || true && \ + rm -rf debugpy* jupyterlab* jupyter_* notebook* ipython* ipykernel* || true && \ + rm -rf onnx/backend/test onnx/test || true && \ + rm -rf scipy/*/tests pandas/tests || true && \ + rm -rf */examples */benchmarks */docs || true && \ + rm -rf skimage/data || true + +RUN echo "torch==${PYTORCH_VERSION}" > /tmp/runtime-overrides.txt && \ + echo "torchvision==${TORCHVISION_VERSION}" >> /tmp/runtime-overrides.txt && \ + uv pip install --target /opt/python-runtime \ + --index-strategy unsafe-best-match \ + --find-links /opt/wheels \ + --override /tmp/runtime-overrides.txt \ + -r /build/reqs/_requirements.txt \ + -r /build/reqs/requirements.cli.txt \ + -r /build/reqs/requirements.http.txt \ + -r /build/reqs/requirements.clip.txt \ + -r /build/reqs/requirements.transformers.txt \ + -r /build/reqs/requirements.sam.txt \ + -r /build/reqs/requirements.groundingdino.txt \ + -r /build/reqs/requirements.yolo_world.txt \ + -r /build/reqs/requirements.doctr.txt \ + -r /build/reqs/requirements.sdk.http.txt \ + -r /build/reqs/requirements.easyocr.txt \ + -r /build/reqs/requirements.jetson.txt \ + /opt/wheels/flash_attn-*.whl \ + /opt/wheels/onnxruntime_gpu-*.whl \ + /opt/wheels/pycuda-*.whl \ + /opt/wheels/torch-*.whl \ + /opt/wheels/torchvision-*.whl \ + /opt/wheels/triton-*.whl \ + "setuptools<=75.5.0" \ + "transformers>=4.57.3,<5.9.0" \ + packaging && \ + python3 -m pip install --no-deps --target /opt/python-runtime --upgrade \ + /opt/wheels/inference_gpu-*.whl && \ + find /opt/python-runtime -type d \ + \( -name cv2 -o -name 'opencv_python.libs' -o -name 'opencv_contrib_python.libs' -o -name 'opencv_python_headless.libs' -o -name 'opencv_contrib_python_headless.libs' \) \ + -prune -exec rm -rf {} + && \ + find /opt/python-runtime -type d -name __pycache__ -prune -exec rm -rf {} + && \ + rm -rf \ + /opt/python-runtime/debugpy* \ + /opt/python-runtime/ipykernel* \ + /opt/python-runtime/ipython* \ + /opt/python-runtime/jupyter_* \ + /opt/python-runtime/jupyterlab* \ + /opt/python-runtime/notebook* \ + /opt/python-runtime/ninja \ + /opt/python-runtime/ninja-*.dist-info \ + /opt/python-runtime/onnx/backend/test \ + /opt/python-runtime/onnx/test \ + /opt/python-runtime/pandas/tests \ + /opt/python-runtime/scipy/*/tests \ + /opt/python-runtime/skimage/data \ + /opt/python-runtime/bin/ninja \ + /opt/python-runtime/triton/backends/nvidia/bin/ptxas \ + /opt/python-runtime/triton/backends/nvidia/bin/ptxas-blackwell \ + /tmp/runtime-overrides.txt \ + /root/.cache/uv && \ + test -x /opt/python-runtime/bin/inference && \ + du -sh /opt/python-runtime /opt/wheels + +FROM jetson-runtime-base AS runtime + +COPY --from=jetson-wheels-base /opt/cuda/bin/ptxas /usr/local/cuda/bin/ptxas +COPY --from=jetson-wheels-base /opt/cuda/nvvm/libdevice /usr/local/cuda/nvvm/libdevice +COPY --from=jetson-wheels-base /opt/cuda/licenses/cuda-ptxas /usr/local/share/licenses/cuda-ptxas + +ENV LD_LIBRARY_PATH=/usr/local/lib:/usr/local/cuda/lib64:/opt/opencv/lib:/opt/gstreamer/lib:/opt/ffmpeg/lib:/opt/cuda-runtime/lib:/usr/lib/aarch64-linux-gnu/tegra:/usr/lib/aarch64-linux-gnu/nvidia + +RUN ldconfig + +COPY --from=builder /opt/python-runtime /usr/local/lib/python3.12/dist-packages + +ENV PYTHONPATH=/opt/opencv/python:/usr/local/lib/python3.12/dist-packages:/usr/lib/python3.12/dist-packages + +COPY --from=builder /opt/python-runtime/bin/inference /usr/local/bin/inference + +COPY inference inference +COPY inference_cli inference_cli +COPY inference_sdk inference_sdk +COPY docker/config/gpu_http.py gpu_http.py +COPY docker/entrypoint/run_uvicorn.sh /usr/local/bin/run_uvicorn.sh +RUN chmod +x /usr/local/bin/run_uvicorn.sh + +ENV CUDA_HOME=/usr/local/cuda \ + PATH=/usr/local/cuda/bin:$PATH \ + TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas \ + TRITON_PTXAS_BLACKWELL_PATH=/usr/local/cuda/bin/ptxas \ + NVIDIA_VISIBLE_DEVICES=all \ + NVIDIA_DRIVER_CAPABILITIES=compute,utility,video \ + VERSION_CHECK_MODE=once \ + CORE_MODEL_SAM2_ENABLED=True \ + NUM_WORKERS=1 \ + HOST=0.0.0.0 \ + PORT=9001 \ + ORT_TENSORRT_FP16_ENABLE=1 \ + ORT_TENSORRT_ENGINE_CACHE_ENABLE=1 \ + ORT_TENSORRT_ENGINE_CACHE_PATH=/tmp/ort_cache \ + ORT_TENSORRT_MAX_WORKSPACE_SIZE=4294967296 \ + ORT_TENSORRT_BUILDER_OPTIMIZATION_LEVEL=5 \ + OPENBLAS_CORETYPE=ARMV8 \ + GST_PLUGIN_PATH_1_0=/usr/lib/aarch64-linux-gnu/gstreamer-1.0 \ + LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libgomp.so.1 \ + WORKFLOWS_STEP_EXECUTION_MODE=local \ + WORKFLOWS_MAX_CONCURRENT_STEPS=4 \ + API_LOGGING_ENABLED=True \ + DISABLE_WORKFLOW_ENDPOINTS=false \ + ALLOW_INFERENCE_EXP_UNTRUSTED_MODELS=True \ + USE_INFERENCE_EXP_MODELS=False \ + RUNNING_ON_JETSON=True \ + L4T_VERSION=39.2.0 + +RUN --mount=type=bind,source=docker/scripts/verify_ffmpeg.sh,target=/tmp/verify_ffmpeg,ro \ + --mount=type=bind,source=docker/scripts/verify_gstreamer.sh,target=/tmp/verify_gstreamer,ro \ + --mount=type=bind,source=docker/scripts/verify_opencv.sh,target=/tmp/verify_opencv,ro \ + --mount=type=bind,source=docker/scripts/verify_media_runtime.sh,target=/tmp/verify_media_runtime,ro \ + --mount=type=bind,source=docker/scripts/verify_jetson_server_runtime.sh,target=/tmp/verify_jetson_server_runtime,ro \ + sh /tmp/verify_ffmpeg && \ + sh /tmp/verify_gstreamer && \ + OPENCV_REQUIRE_CUDA_BUILD=true sh /tmp/verify_opencv && \ + VERIFY_MEDIA_RUNTIME_DEVELOPMENT_TOOLS=false sh /tmp/verify_media_runtime && \ + sh /tmp/verify_jetson_server_runtime && \ + rm -rf /root/.cache/gstreamer-1.0 + +LABEL org.opencontainers.image.description="Inference Server - Jetson 7.2.0 (PyTorch from source, numpy 2.x)" + +EXPOSE 9001 + +ENTRYPOINT ["/usr/local/bin/run_uvicorn.sh", "gpu_http:app"] diff --git a/docker/dockerfiles/Dockerfile.wheels.jetson.7.2.0 b/docker/dockerfiles/Dockerfile.wheels.jetson.7.2.0 new file mode 100644 index 0000000000..9e3461436f --- /dev/null +++ b/docker/dockerfiles/Dockerfile.wheels.jetson.7.2.0 @@ -0,0 +1,297 @@ +# syntax=docker/dockerfile:1.7 +# +# Prebuilt framework wheels for the JetPack 7.2 server image. Component builds run +# as independent stages so BuildKit parallelizes them; the dependency graph is +# +# toolchain โ”€โ”ฌโ”€ torch โ”€โ”ฌโ”€ torchvision +# โ”‚ โ””โ”€ flash-attn (critical path) +# โ”œโ”€ onnxruntime +# โ”œโ”€ triton (fetch only) +# โ””โ”€ pycuda +# +# Each stage has a private pip-cache mount id: a shared sharing=locked mount would +# serialize the parallel stages for the full duration of each RUN. + +ARG CMAKE_VERSION=4.2.0 +ARG NINJA_VERSION=1.11.1.4 +ARG PYTORCH_VERSION=2.10.0 +ARG TORCHVISION_VERSION=0.25.0 +ARG TRITON_VERSION=3.6.0 +ARG TRITON_WHEEL_SHA256=374f52c11a711fd062b4bfbb201fd9ac0a5febd28a96fb41b4a0f51dde3157f4 +ARG ONNXRUNTIME_VERSION=1.24.2 +ARG PYTORCH_MAX_JOBS=12 +ARG ORT_BUILD_PARALLEL=8 +ARG FLASH_ATTN_MAX_JOBS=2 +ARG TORCH_CUDA_ARCH_LIST=8.7;11.0 +ARG ORT_CUDA_ARCHITECTURES=87-real;110-real;110-virtual +ARG FLASH_ATTN_CUDA_ARCHS=87;110 +ARG FLASH_ATTN_REF=d16e381f6f8041422bd36f27adaf61fbf8ca872a + +FROM ubuntu:24.04 AS toolchain + +ARG DEBIAN_FRONTEND=noninteractive +ARG CMAKE_VERSION +ARG NINJA_VERSION +ARG TARGETARCH +ENV LANG=en_US.UTF-8 + +WORKDIR /build + +RUN mkdir -p /opt/wheels + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp72-onnx-builder-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jp72-onnx-builder-${TARGETARCH} \ + rm -f /etc/apt/apt.conf.d/docker-clean && \ + echo 'Binary::apt::APT::Keep-Downloaded-Packages "true";' > /etc/apt/apt.conf.d/keep-cache && \ + apt-get update -y && \ + apt-get install -y --no-install-recommends \ + ca-certificates \ + gnupg \ + wget \ + && wget -qO - https://repo.download.nvidia.com/jetson/jetson-ota-public.asc | gpg --dearmor -o /usr/share/keyrings/nvidia-jetson.gpg \ + && echo "deb [signed-by=/usr/share/keyrings/nvidia-jetson.gpg] https://repo.download.nvidia.com/jetson/common r39.2 main" > /etc/apt/sources.list.d/nvidia-jetson.list \ + && echo "deb [signed-by=/usr/share/keyrings/nvidia-jetson.gpg] https://repo.download.nvidia.com/jetson/som r39.2 main" >> /etc/apt/sources.list.d/nvidia-jetson.list \ + && apt-get update -y + +RUN --mount=type=cache,target=/var/cache/apt,sharing=locked,id=apt-cache-jp72-onnx-builder-${TARGETARCH} \ + --mount=type=cache,target=/var/lib/apt,sharing=locked,id=apt-lib-jp72-onnx-builder-${TARGETARCH} \ + apt-get install -y --no-install-recommends \ + cuda-toolkit-13-2 \ + libcudnn9-cuda-13 \ + libcudnn9-dev-cuda-13 \ + tensorrt \ + python3-libnvinfer \ + python3-libnvinfer-dev \ + python3-libnvinfer-dispatch \ + python3-libnvinfer-lean \ + build-essential \ + cmake \ + ninja-build \ + git \ + curl \ + python3-dev \ + python3-pip \ + python3-venv \ + libopenblas-dev \ + libproj-dev \ + libsqlite3-dev \ + libtiff-dev \ + libcurl4-openssl-dev \ + libssl-dev \ + zlib1g-dev \ + libxext6 \ + libvips-dev \ + pkg-config \ + libjpeg-dev \ + libpng-dev \ + libavcodec-dev \ + libavformat-dev \ + libswscale-dev \ + libv4l-dev \ + libatlas-base-dev \ + liblapack-dev \ + gfortran \ + libhdf5-dev \ + libgflags-dev \ + libgoogle-glog-dev \ + liblmdb-dev \ + libleveldb-dev \ + libsnappy-dev \ + libprotobuf-dev \ + protobuf-compiler \ + libxml2-dev + +RUN wget -q https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/cmake-${CMAKE_VERSION}-linux-aarch64.sh && \ + chmod +x cmake-${CMAKE_VERSION}-linux-aarch64.sh && \ + ./cmake-${CMAKE_VERSION}-linux-aarch64.sh --prefix=/usr/local --skip-license && \ + rm cmake-${CMAKE_VERSION}-linux-aarch64.sh && \ + printf 'cmake==%s\nninja==%s\n' "${CMAKE_VERSION}" "${NINJA_VERSION}" > /opt/builder-constraints.txt + +ENV PIP_CONSTRAINT=/opt/builder-constraints.txt + +RUN python3 -m pip install --break-system-packages --ignore-installed pip setuptools wheel && \ + python3 -m pip install --break-system-packages \ + "cmake==${CMAKE_VERSION}" \ + "ninja==${NINJA_VERSION}" \ + "numpy>=2.0.0,<2.4.0" + +ENV CUDA_HOME=/usr/local/cuda \ + PATH=/usr/local/cuda/bin:$PATH \ + LD_LIBRARY_PATH=/usr/local/cuda/lib64 + +FROM toolchain AS torch-builder + +ARG PYTORCH_VERSION +ARG PYTORCH_MAX_JOBS +ARG TORCH_CUDA_ARCH_LIST +ARG CMAKE_VERSION +ARG NINJA_VERSION +ARG TARGETARCH + +WORKDIR /build/pytorch +RUN git clone --recursive --branch v${PYTORCH_VERSION} https://github.com/pytorch/pytorch +RUN --mount=type=cache,target=/build/pytorch/pytorch/build,sharing=locked,id=pytorch-build-${PYTORCH_VERSION}-${TARGETARCH}-sm87-sm110-cu132-cmake420-ninja111 \ + --mount=type=cache,target=/root/.cache/pip,sharing=locked,id=pip-cache-jp72-torch-${TARGETARCH} \ + cd pytorch && \ + python3 -m pip install --break-system-packages -r requirements.txt && \ + test "$(python3 -c 'import importlib.metadata as m; print(m.version("cmake"))')" = "${CMAKE_VERSION}" && \ + test "$(python3 -c 'import importlib.metadata as m; print(m.version("ninja"))')" = "${NINJA_VERSION}" && \ + export USE_CUDA=1 USE_CUDNN=1 CUDA_HOME=/usr/local/cuda && \ + export CUDNN_LIB_DIR=/usr/lib/aarch64-linux-gnu CUDNN_INCLUDE_DIR=/usr/include && \ + export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" && \ + export USE_MKLDNN=0 USE_OPENMP=0 && \ + export USE_DISTRIBUTED=0 USE_GLOO=0 USE_MPI=0 USE_TENSORPIPE=0 USE_NCCL=0 && \ + export BUILD_TEST=0 && \ + export USE_QNNPACK=0 USE_PYTORCH_QNNPACK=0 USE_XNNPACK=0 USE_NNPACK=0 && \ + export USE_FBGEMM=0 USE_KINETO=0 USE_CUPTI_SO=0 && \ + export USE_FLASH_ATTENTION=1 USE_MEM_EFF_ATTENTION=1 && \ + export USE_MPS=0 USE_ROCM=0 && \ + export PYTORCH_BUILD_VERSION=${PYTORCH_VERSION} PYTORCH_BUILD_NUMBER=1 && \ + export CMAKE_BUILD_TYPE=Release BUILD_SHARED_LIBS=ON USE_PRIORITIZED_TEXT_FOR_LD=1 && \ + export MAX_JOBS=${PYTORCH_MAX_JOBS} && \ + export CMAKE_POLICY_VERSION_MINIMUM=3.5 && \ + python3 setup.py bdist_wheel && \ + cp dist/torch-*.whl /opt/wheels/ && \ + python3 -m pip install --break-system-packages dist/torch-*.whl + +FROM torch-builder AS torchvision-builder + +ARG TORCHVISION_VERSION +ARG TORCH_CUDA_ARCH_LIST +ARG TARGETARCH + +WORKDIR /build/torchvision +RUN for i in 1 2 3; do \ + git clone --branch v${TORCHVISION_VERSION} https://github.com/pytorch/vision && break || sleep 60; \ + done +RUN --mount=type=cache,target=/build/torchvision/vision/build,sharing=locked,id=torchvision-build-${TORCHVISION_VERSION}-${TARGETARCH}-sm87-sm110-cu132 \ + --mount=type=cache,target=/root/.cache/pip,sharing=locked,id=pip-cache-jp72-torchvision-${TARGETARCH} \ + cd vision && \ + export BUILD_VERSION=${TORCHVISION_VERSION} TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST}" && \ + export FORCE_CUDA=1 CMAKE_BUILD_TYPE=Release USE_PRIORITIZED_TEXT_FOR_LD=1 && \ + python3 setup.py bdist_wheel && \ + cp dist/torchvision-*.whl /opt/wheels/ && \ + python3 -m pip install --break-system-packages dist/torchvision-*.whl && \ + torchvision_image="$(find /usr/local/lib/python3.12/dist-packages/torchvision \ + -type f -name 'image.so' -print -quit)" && \ + test -n "${torchvision_image}" && \ + readelf -d "${torchvision_image}" | grep -q 'Shared library: \[libnvjpeg.so.13\]' + +FROM toolchain AS ort-builder + +ARG ONNXRUNTIME_VERSION +ARG ORT_BUILD_PARALLEL +ARG ORT_CUDA_ARCHITECTURES +ARG CMAKE_VERSION +ARG NINJA_VERSION +ARG TARGETARCH + +WORKDIR /build/onnxruntime +RUN git clone --recursive --branch v${ONNXRUNTIME_VERSION} https://github.com/microsoft/onnxruntime.git +RUN --mount=type=cache,target=/build/onnxruntime/onnxruntime/build/Linux,sharing=locked,id=ort-build-${ONNXRUNTIME_VERSION}-${TARGETARCH}-sm87-sm110-cu132 \ + --mount=type=cache,target=/root/.cache/pip,sharing=locked,id=pip-cache-jp72-ort-${TARGETARCH} \ + cd onnxruntime && \ + python3 -m pip install --break-system-packages -r requirements-dev.txt && \ + test "$(python3 -c 'import importlib.metadata as m; print(m.version("cmake"))')" = "${CMAKE_VERSION}" && \ + test "$(python3 -c 'import importlib.metadata as m; print(m.version("ninja"))')" = "${NINJA_VERSION}" && \ + ./build.sh --config Release --update --build --parallel ${ORT_BUILD_PARALLEL} \ + --build_wheel \ + --allow_running_as_root \ + --compile_no_warning_as_error \ + --use_cuda \ + --cuda_home /usr/local/cuda \ + --cudnn_home /usr \ + --use_tensorrt \ + --tensorrt_home /usr \ + --skip_tests \ + --cmake_extra_defines \ + CMAKE_POLICY_VERSION_MINIMUM=3.5 \ + CMAKE_CUDA_ARCHITECTURES="${ORT_CUDA_ARCHITECTURES}" \ + onnxruntime_BUILD_UNIT_TESTS=OFF && \ + cp build/Linux/Release/dist/onnxruntime_gpu*.whl /opt/wheels/ && \ + python3 -m pip install --break-system-packages build/Linux/Release/dist/onnxruntime_gpu*.whl && \ + ort_cuda_provider="$(find /usr/local/lib/python3.12/dist-packages \ + -type f -name 'libonnxruntime_providers_cuda.so' -print -quit)" && \ + test -n "${ort_cuda_provider}" && \ + cuobjdump --list-elf "${ort_cuda_provider}" | grep -q 'sm_87' && \ + cuobjdump --list-elf "${ort_cuda_provider}" | grep -q 'sm_110' + +# FlashAttention is cloned in its own layer and compiled with the build tree on a +# cache mount so an interrupted build resumes instead of restarting from zero. +FROM torch-builder AS flash-attn-builder + +ARG FLASH_ATTN_REF +ARG FLASH_ATTN_MAX_JOBS +ARG FLASH_ATTN_CUDA_ARCHS +ARG TARGETARCH + +WORKDIR /build/flash-attention +RUN --mount=type=bind,source=docker/patches/flash-attention-sm87.patch,target=/tmp/flash-attention-sm87.patch,ro \ + git clone https://github.com/Dao-AILab/flash-attention.git flash-attention && \ + cd flash-attention && \ + git checkout ${FLASH_ATTN_REF} && \ + git apply /tmp/flash-attention-sm87.patch && \ + git submodule update --init --recursive +RUN --mount=type=cache,target=/build/flash-attention/flash-attention/build,sharing=locked,id=flash-attn-build-${FLASH_ATTN_REF}-${TARGETARCH}-sm87-sm110-cu132 \ + --mount=type=cache,target=/root/.cache/pip,sharing=locked,id=pip-cache-jp72-flash-attn-${TARGETARCH} \ + cd flash-attention && \ + export FLASH_ATTENTION_FORCE_BUILD=TRUE && \ + export FLASH_ATTN_CUDA_ARCHS="${FLASH_ATTN_CUDA_ARCHS}" && \ + export MAX_JOBS=${FLASH_ATTN_MAX_JOBS} NVCC_THREADS=1 && \ + python3 -m pip wheel --no-build-isolation --no-deps \ + --wheel-dir /opt/wheels . && \ + python3 -m pip install --break-system-packages --no-deps \ + /opt/wheels/flash_attn-*.whl && \ + flash_extension="$(find /usr/local/lib/python3.12/dist-packages -type f \ + -name 'flash_attn_2_cuda*.so' -print -quit)" && \ + test -n "${flash_extension}" && \ + cuobjdump --list-elf "${flash_extension}" | grep -q 'sm_87' && \ + cuobjdump --list-elf "${flash_extension}" | grep -q 'sm_110' + +FROM toolchain AS triton-fetcher + +ARG TRITON_VERSION +ARG TRITON_WHEEL_SHA256 +ARG TARGETARCH + +WORKDIR /build/triton +RUN --mount=type=cache,target=/root/.cache/pip,sharing=locked,id=pip-cache-jp72-triton-${TARGETARCH} \ + set -eux; \ + mkdir -p /tmp/triton-wheels; \ + python3 -m pip download --no-cache-dir \ + --no-deps \ + --only-binary=:all: \ + --dest /tmp/triton-wheels \ + "triton==${TRITON_VERSION}"; \ + set -- /tmp/triton-wheels/triton-${TRITON_VERSION}*.whl; \ + if [ ! -e "$1" ]; then \ + echo "Expected Triton wheel for ${TRITON_VERSION} was not found" >&2; \ + exit 1; \ + fi; \ + printf '%s %s\n' "${TRITON_WHEEL_SHA256}" "$1" | sha256sum -c -; \ + cp "$1" /opt/wheels/; \ + python3 -m pip install --break-system-packages --no-cache-dir "$1"; \ + rm -rf /tmp/triton-wheels; \ + python3 -c "import triton; print('triton', triton.__version__)" + +FROM toolchain AS pycuda-builder + +ARG TARGETARCH + +WORKDIR /build/pycuda +RUN --mount=type=cache,target=/root/.cache/pip,sharing=locked,id=pip-cache-jp72-pycuda-${TARGETARCH} \ + python3 -m pip wheel --no-deps --wheel-dir /opt/wheels \ + "pycuda>=2025.0.0,<2026.0.0" + +FROM scratch AS wheels + +COPY --from=torchvision-builder /opt/wheels /opt/wheels +COPY --from=ort-builder /opt/wheels /opt/wheels +COPY --from=flash-attn-builder /opt/wheels /opt/wheels +COPY --from=triton-fetcher /opt/wheels /opt/wheels +COPY --from=pycuda-builder /opt/wheels /opt/wheels +COPY --from=toolchain /usr/local/cuda/bin/ptxas /opt/cuda/bin/ptxas +COPY --from=toolchain /usr/local/cuda/nvvm/libdevice /opt/cuda/nvvm/libdevice +COPY --from=toolchain /usr/share/doc/cuda-nvcc-13-2/copyright /opt/cuda/licenses/cuda-ptxas + +LABEL org.opencontainers.image.description="Roboflow Jetson 7.2 prebuilt framework wheels (PyTorch, torchvision, ONNX Runtime, FlashAttention, Triton, pycuda) + ptxas export" diff --git a/docker/native/gstreamer_cuda_tensor_bridge/gstreamer_cuda_tensor_bridge.cpp b/docker/native/gstreamer_cuda_tensor_bridge/gstreamer_cuda_tensor_bridge.cpp new file mode 100644 index 0000000000..296e5cba92 --- /dev/null +++ b/docker/native/gstreamer_cuda_tensor_bridge/gstreamer_cuda_tensor_bridge.cpp @@ -0,0 +1,706 @@ +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr const char* kSinkName = "rf_tensor_sink"; + +enum DLDeviceType : int32_t { + kDLCUDA = 2, +}; + +enum DLDataTypeCode : uint8_t { + kDLUInt = 1, +}; + +struct DLDevice { + DLDeviceType device_type; + int32_t device_id; +}; + +struct DLDataType { + uint8_t code; + uint8_t bits; + uint16_t lanes; +}; + +struct DLTensor { + void* data; + DLDevice device; + int32_t ndim; + DLDataType dtype; + int64_t* shape; + int64_t* strides; + uint64_t byte_offset; +}; + +struct DLManagedTensor { + DLTensor dl_tensor; + void* manager_ctx; + void (*deleter)(DLManagedTensor* self); +}; + +struct RfFrameInfo { + uint32_t width; + uint32_t height; + int32_t fps_numerator; + int32_t fps_denominator; + int64_t duration_ns; +}; + +struct RfCudaBridgeStats { + uint64_t frames; + uint64_t cuda_maps; + // The next three counters must stay zero on the zero-copy path; the + // verify scripts assert on them to prove no host fallback executed. + uint64_t host_pixel_maps; + uint64_t host_to_device_copies; + uint64_t device_to_host_copies; + uint64_t stream_synchronizations; + uint64_t active_leases; + int64_t last_channel_stride; + int64_t last_row_stride; +}; + +struct RfLeaseCounter { + std::atomic references{1}; + std::atomic active{0}; +}; + +struct RfCudaTensorContext { + GstSample* sample = nullptr; + GstMemory* memory = nullptr; + GstMapInfo map = GST_MAP_INFO_INIT; + bool mapped = false; + CUdevice device = 0; + bool primary_context_retained = false; + RfLeaseCounter* leases = nullptr; + int64_t shape[3] = {0, 0, 0}; + int64_t strides[3] = {0, 0, 0}; + DLManagedTensor managed{}; +}; + +struct RfCudaPipeline { + GstElement* pipeline = nullptr; + GstAppSink* sink = nullptr; + GstSample* sample = nullptr; + GstCudaContext* cuda_context = nullptr; + CUcontext primary_context = nullptr; + CUdevice device = 0; + int device_id = 0; + RfCudaBridgeStats stats{}; + RfLeaseCounter* leases = nullptr; + std::atomic interrupted{false}; + std::mutex mutex; + std::mutex stats_mutex; +}; + +std::once_flag g_gstreamer_once; + +void write_error(char* destination, size_t capacity, const char* format, ...) { + if (destination == nullptr || capacity == 0) { + return; + } + va_list args; + va_start(args, format); + std::vsnprintf(destination, capacity, format, args); + va_end(args); + destination[capacity - 1] = '\0'; +} + +void initialize_gstreamer() { + gst_init(nullptr, nullptr); + gst_cuda_memory_init_once(); +} + +void delete_managed_tensor(DLManagedTensor* managed) { + if (managed == nullptr) { + return; + } + auto* context = static_cast(managed->manager_ctx); + if (context == nullptr) { + return; + } + if (context->primary_context_retained && context->mapped) { + // Consumers (torch) may still have kernels queued against this + // zero-copy mapping when the last tensor reference drops; the buffer + // returns to the decoder pool on unmap/unref, so finish device work + // first. + CUcontext primary_context = nullptr; + if (cuDevicePrimaryCtxRetain(&primary_context, context->device) == + CUDA_SUCCESS) { + if (cuCtxPushCurrent(primary_context) == CUDA_SUCCESS) { + cuCtxSynchronize(); + cuCtxPopCurrent(nullptr); + } + cuDevicePrimaryCtxRelease(context->device); + } + } + if (context->mapped && context->memory != nullptr) { + gst_memory_unmap(context->memory, &context->map); + } + if (context->sample != nullptr) { + gst_sample_unref(context->sample); + } + if (context->primary_context_retained) { + cuDevicePrimaryCtxRelease(context->device); + } + if (context->leases != nullptr) { + context->leases->active.fetch_sub(1, std::memory_order_relaxed); + if (context->leases->references.fetch_sub(1, std::memory_order_acq_rel) == 1) { + delete context->leases; + } + } + delete context; +} + +bool sample_has_cuda_caps(GstSample* sample) { + GstCaps* caps = gst_sample_get_caps(sample); + if (caps == nullptr || gst_caps_is_empty(caps)) { + return false; + } + for (guint index = 0; index < gst_caps_get_size(caps); ++index) { + const GstCapsFeatures* features = gst_caps_get_features(caps, index); + if (features != nullptr && gst_caps_features_contains( + features, GST_CAPS_FEATURE_MEMORY_CUDA_MEMORY)) { + return true; + } + } + return false; +} + +bool read_sample_info(GstSample* sample, RfFrameInfo* info) { + if (sample == nullptr || info == nullptr) { + return false; + } + GstCaps* caps = gst_sample_get_caps(sample); + if (caps == nullptr || gst_caps_is_empty(caps)) { + return false; + } + GstVideoInfo video_info{}; + if (!gst_video_info_from_caps(&video_info, caps)) { + return false; + } + info->width = GST_VIDEO_INFO_WIDTH(&video_info); + info->height = GST_VIDEO_INFO_HEIGHT(&video_info); + info->fps_numerator = GST_VIDEO_INFO_FPS_N(&video_info); + info->fps_denominator = GST_VIDEO_INFO_FPS_D(&video_info); + if (info->fps_denominator <= 0) { + info->fps_denominator = 1; + } + return true; +} + +bool read_bus_error( + RfCudaPipeline* handle, + char* error, + size_t error_capacity) { + GstBus* bus = gst_element_get_bus(handle->pipeline); + if (bus == nullptr) { + return false; + } + GstMessage* message = gst_bus_pop_filtered(bus, GST_MESSAGE_ERROR); + bool has_error = false; + if (message != nullptr) { + has_error = true; + GError* gst_error = nullptr; + gchar* debug = nullptr; + gst_message_parse_error(message, &gst_error, &debug); + write_error( + error, + error_capacity, + "%s", + gst_error != nullptr ? gst_error->message : "GStreamer pipeline error"); + if (gst_error != nullptr) { + g_error_free(gst_error); + } + g_free(debug); + gst_message_unref(message); + } + gst_object_unref(bus); + return has_error; +} + +bool pipeline_has_factory(RfCudaPipeline* handle, const char* factory_name) { + if (handle == nullptr || handle->pipeline == nullptr || factory_name == nullptr || + !GST_IS_BIN(handle->pipeline)) { + return false; + } + GstIterator* iterator = gst_bin_iterate_recurse(GST_BIN(handle->pipeline)); + GValue item = G_VALUE_INIT; + bool found = false; + bool done = false; + while (!done && !found) { + switch (gst_iterator_next(iterator, &item)) { + case GST_ITERATOR_OK: { + auto* element = GST_ELEMENT(g_value_get_object(&item)); + GstElementFactory* factory = gst_element_get_factory(element); + const gchar* name = factory == nullptr + ? nullptr + : gst_plugin_feature_get_name(GST_PLUGIN_FEATURE(factory)); + found = name != nullptr && std::strcmp(name, factory_name) == 0; + g_value_reset(&item); + break; + } + case GST_ITERATOR_RESYNC: + gst_iterator_resync(iterator); + break; + default: + done = true; + break; + } + } + if (G_VALUE_TYPE(&item) != 0) { + g_value_unset(&item); + } + gst_iterator_free(iterator); + return found; +} + +} // namespace + +extern "C" { + +__attribute__((visibility("default"))) +const char* rf_gstreamer_cuda_tensor_bridge_version() { + return "3"; +} + +__attribute__((visibility("default"))) +RfCudaPipeline* rf_gstreamer_cuda_pipeline_create( + const char* pipeline_description, + int device_id, + char* error, + size_t error_capacity) { + if (pipeline_description == nullptr || pipeline_description[0] == '\0') { + write_error(error, error_capacity, "GStreamer pipeline is empty"); + return nullptr; + } + std::call_once(g_gstreamer_once, initialize_gstreamer); + + CUresult cuda_status = cuInit(0); + CUdevice device = 0; + CUcontext primary_context = nullptr; + if (cuda_status == CUDA_SUCCESS) { + cuda_status = cuDeviceGet(&device, device_id); + } + if (cuda_status == CUDA_SUCCESS) { + cuda_status = cuDevicePrimaryCtxRetain(&primary_context, device); + } + if (cuda_status != CUDA_SUCCESS) { + const char* cuda_error = nullptr; + cuGetErrorString(cuda_status, &cuda_error); + write_error( + error, + error_capacity, + "CUDA primary context is unavailable: %s", + cuda_error == nullptr ? "unknown error" : cuda_error); + return nullptr; + } + + GError* parse_error = nullptr; + GstElement* pipeline = gst_parse_launch(pipeline_description, &parse_error); + if (pipeline == nullptr || parse_error != nullptr) { + write_error( + error, + error_capacity, + "GStreamer pipeline parse failed: %s", + parse_error == nullptr ? "unknown error" : parse_error->message); + if (parse_error != nullptr) { + g_error_free(parse_error); + } + if (pipeline != nullptr) { + gst_object_unref(pipeline); + } + cuDevicePrimaryCtxRelease(device); + return nullptr; + } + if (!GST_IS_BIN(pipeline)) { + write_error(error, error_capacity, "GStreamer pipeline is not a bin"); + gst_object_unref(pipeline); + cuDevicePrimaryCtxRelease(device); + return nullptr; + } + GstElement* sink_element = gst_bin_get_by_name(GST_BIN(pipeline), kSinkName); + if (sink_element == nullptr || !GST_IS_APP_SINK(sink_element)) { + write_error( + error, + error_capacity, + "GStreamer pipeline requires appsink name=%s", + kSinkName); + if (sink_element != nullptr) { + gst_object_unref(sink_element); + } + gst_object_unref(pipeline); + cuDevicePrimaryCtxRelease(device); + return nullptr; + } + + GstCudaContext* cuda_context = + gst_cuda_context_new_wrapped(primary_context, device); + if (cuda_context == nullptr) { + write_error(error, error_capacity, "Could not wrap the CUDA primary context"); + gst_object_unref(sink_element); + gst_object_unref(pipeline); + cuDevicePrimaryCtxRelease(device); + return nullptr; + } + GstContext* gst_context = gst_context_new_cuda_context(cuda_context); + gst_element_set_context(pipeline, gst_context); + gst_context_unref(gst_context); + + auto* handle = new (std::nothrow) RfCudaPipeline(); + if (handle == nullptr) { + write_error(error, error_capacity, "Could not allocate pipeline state"); + gst_object_unref(cuda_context); + gst_object_unref(sink_element); + gst_object_unref(pipeline); + cuDevicePrimaryCtxRelease(device); + return nullptr; + } + handle->pipeline = pipeline; + handle->sink = GST_APP_SINK(sink_element); + handle->cuda_context = cuda_context; + handle->primary_context = primary_context; + handle->device = device; + handle->device_id = device_id; + handle->leases = new (std::nothrow) RfLeaseCounter(); + if (handle->leases == nullptr) { + write_error(error, error_capacity, "Could not allocate tensor lease state"); + gst_object_unref(cuda_context); + gst_object_unref(sink_element); + gst_object_unref(pipeline); + cuDevicePrimaryCtxRelease(device); + delete handle; + return nullptr; + } + + const GstStateChangeReturn state_status = + gst_element_set_state(pipeline, GST_STATE_PLAYING); + if (state_status == GST_STATE_CHANGE_FAILURE) { + write_error(error, error_capacity, "GStreamer could not enter PLAYING state"); + // A failed PLAYING transition can leave elements in READY/PAUSED with + // running task threads; GStreamer refuses to dispose a non-NULL + // pipeline, so reset it before dropping the reference. + gst_element_set_state(pipeline, GST_STATE_NULL); + gst_object_unref(cuda_context); + gst_object_unref(sink_element); + gst_object_unref(pipeline); + cuDevicePrimaryCtxRelease(device); + delete handle->leases; + delete handle; + return nullptr; + } + return handle; +} + +__attribute__((visibility("default"))) +int rf_gstreamer_cuda_pipeline_grab( + RfCudaPipeline* handle, + uint64_t timeout_ns, + char* error, + size_t error_capacity) { + if (handle == nullptr) { + write_error(error, error_capacity, "Pipeline handle is null"); + return -1; + } + if (handle->interrupted.load(std::memory_order_acquire)) { + return 0; + } + std::lock_guard lock(handle->mutex); + if (handle->interrupted.load(std::memory_order_acquire)) { + return 0; + } + if (handle->sample != nullptr) { + gst_sample_unref(handle->sample); + handle->sample = nullptr; + } + handle->sample = gst_app_sink_try_pull_sample(handle->sink, timeout_ns); + if (handle->sample == nullptr) { + if (handle->interrupted.load(std::memory_order_acquire) || + gst_app_sink_is_eos(handle->sink)) { + return 0; + } + if (read_bus_error(handle, error, error_capacity)) { + return -1; + } + // No frame, no error, no EOS: the finite timeout expired while the + // stream is still live. + return 2; + } + if (!sample_has_cuda_caps(handle->sample)) { + write_error( + error, + error_capacity, + "GStreamer appsink frame is not memory:CUDAMemory"); + gst_sample_unref(handle->sample); + handle->sample = nullptr; + return -1; + } + return 1; +} + +__attribute__((visibility("default"))) +int rf_gstreamer_cuda_pipeline_get_frame_info( + RfCudaPipeline* handle, + RfFrameInfo* info, + char* error, + size_t error_capacity) { + if (handle == nullptr || info == nullptr) { + write_error(error, error_capacity, "Frame-info arguments are invalid"); + return -1; + } + std::lock_guard lock(handle->mutex); + if (handle->sample == nullptr || !read_sample_info(handle->sample, info)) { + write_error(error, error_capacity, "Frame caps do not contain video metadata"); + return -1; + } + gint64 duration = GST_CLOCK_TIME_NONE; + info->duration_ns = gst_element_query_duration( + handle->pipeline, GST_FORMAT_TIME, &duration) ? duration : 0; + return 1; +} + +__attribute__((visibility("default"))) +int rf_gstreamer_cuda_pipeline_has_factory( + RfCudaPipeline* handle, + const char* factory_name) { + if (handle == nullptr) { + return 0; + } + std::lock_guard lock(handle->mutex); + return pipeline_has_factory(handle, factory_name) ? 1 : 0; +} + +__attribute__((visibility("default"))) +DLManagedTensor* rf_gstreamer_cuda_pipeline_retrieve( + RfCudaPipeline* handle, + char* error, + size_t error_capacity) { + if (handle == nullptr) { + write_error(error, error_capacity, "Pipeline handle is null"); + return nullptr; + } + std::lock_guard lock(handle->mutex); + if (handle->sample == nullptr) { + write_error(error, error_capacity, "No grabbed frame is available"); + return nullptr; + } + + GstSample* sample = handle->sample; + handle->sample = nullptr; + GstBuffer* buffer = gst_sample_get_buffer(sample); + if (buffer == nullptr || gst_buffer_n_memory(buffer) != 1) { + write_error(error, error_capacity, "CUDA frame must contain one GstMemory"); + gst_sample_unref(sample); + return nullptr; + } + GstMemory* memory = gst_buffer_peek_memory(buffer, 0); + if (memory == nullptr || !gst_is_cuda_memory(memory)) { + write_error(error, error_capacity, "Frame memory is not GstCudaMemory"); + gst_sample_unref(sample); + return nullptr; + } + if (GST_MEMORY_FLAG_IS_SET( + memory, + static_cast(GST_CUDA_MEMORY_TRANSFER_NEED_UPLOAD))) { + write_error(error, error_capacity, "CUDA frame requires a host upload"); + gst_sample_unref(sample); + return nullptr; + } + + auto* cuda_memory = GST_CUDA_MEMORY_CAST(memory); + CUcontext memory_context = reinterpret_cast( + gst_cuda_context_get_handle(cuda_memory->context)); + if (memory_context != handle->primary_context) { + write_error(error, error_capacity, "CUDA frame uses a different context"); + gst_sample_unref(sample); + return nullptr; + } + const GstVideoInfo* video_info = &cuda_memory->info; + if (GST_VIDEO_INFO_FORMAT(video_info) != GST_VIDEO_FORMAT_RGBP || + GST_VIDEO_INFO_N_PLANES(video_info) != 3 || + GST_VIDEO_INFO_COMP_PSTRIDE(video_info, 0) != 1 || + GST_VIDEO_INFO_COMP_PSTRIDE(video_info, 1) != 1 || + GST_VIDEO_INFO_COMP_PSTRIDE(video_info, 2) != 1) { + write_error(error, error_capacity, "CUDA frame is not planar uint8 RGB"); + gst_sample_unref(sample); + return nullptr; + } + const int64_t row_stride = GST_VIDEO_INFO_PLANE_STRIDE(video_info, 0); + const int64_t second_row_stride = GST_VIDEO_INFO_PLANE_STRIDE(video_info, 1); + const int64_t third_row_stride = GST_VIDEO_INFO_PLANE_STRIDE(video_info, 2); + const int64_t channel_stride = + GST_VIDEO_INFO_PLANE_OFFSET(video_info, 1) - + GST_VIDEO_INFO_PLANE_OFFSET(video_info, 0); + const int64_t second_channel_stride = + GST_VIDEO_INFO_PLANE_OFFSET(video_info, 2) - + GST_VIDEO_INFO_PLANE_OFFSET(video_info, 1); + if (row_stride <= 0 || row_stride != second_row_stride || + row_stride != third_row_stride || channel_stride <= 0 || + channel_stride != second_channel_stride) { + write_error(error, error_capacity, "CUDA RGB plane strides are incompatible"); + gst_sample_unref(sample); + return nullptr; + } + + auto* tensor = new (std::nothrow) RfCudaTensorContext(); + if (tensor == nullptr) { + write_error(error, error_capacity, "Could not allocate tensor lease"); + gst_sample_unref(sample); + return nullptr; + } + tensor->sample = sample; + tensor->memory = memory; + tensor->device = handle->device; + tensor->managed.manager_ctx = tensor; + tensor->managed.deleter = delete_managed_tensor; + CUresult cuda_status = cuDevicePrimaryCtxRetain( + &memory_context, handle->device); + if (cuda_status != CUDA_SUCCESS) { + write_error(error, error_capacity, "Could not retain the CUDA primary context"); + delete_managed_tensor(&tensor->managed); + return nullptr; + } + tensor->primary_context_retained = true; + + gst_cuda_memory_sync(cuda_memory); + if (!gst_memory_map( + memory, + &tensor->map, + static_cast(GST_MAP_READ | GST_MAP_CUDA))) { + write_error(error, error_capacity, "Could not map CUDA frame memory"); + delete_managed_tensor(&tensor->managed); + return nullptr; + } + tensor->mapped = true; + + tensor->shape[0] = 3; + tensor->shape[1] = GST_VIDEO_INFO_HEIGHT(video_info); + tensor->shape[2] = GST_VIDEO_INFO_WIDTH(video_info); + tensor->strides[0] = channel_stride; + tensor->strides[1] = row_stride; + tensor->strides[2] = 1; + // Fold the plane offset into the pointer: torch's fromDLPack has + // historically ignored nonzero byte_offset. + tensor->managed.dl_tensor.data = + tensor->map.data + GST_VIDEO_INFO_PLANE_OFFSET(video_info, 0); + tensor->managed.dl_tensor.device = {kDLCUDA, handle->device_id}; + tensor->managed.dl_tensor.ndim = 3; + tensor->managed.dl_tensor.dtype = {kDLUInt, 8, 1}; + tensor->managed.dl_tensor.shape = tensor->shape; + tensor->managed.dl_tensor.strides = tensor->strides; + tensor->managed.dl_tensor.byte_offset = 0; + + { + std::lock_guard stats_lock(handle->stats_mutex); + handle->stats.frames += 1; + handle->stats.cuda_maps += 1; + handle->stats.stream_synchronizations += 1; + handle->stats.last_channel_stride = channel_stride; + handle->stats.last_row_stride = row_stride; + } + handle->leases->references.fetch_add(1, std::memory_order_relaxed); + handle->leases->active.fetch_add(1, std::memory_order_relaxed); + tensor->leases = handle->leases; + return &tensor->managed; +} + +__attribute__((visibility("default"))) +int rf_gstreamer_cuda_pipeline_get_stats( + RfCudaPipeline* handle, + RfCudaBridgeStats* stats) { + if (handle == nullptr || stats == nullptr) { + return -1; + } + std::lock_guard lock(handle->stats_mutex); + *stats = handle->stats; + stats->active_leases = handle->leases->active.load(std::memory_order_relaxed); + return 1; +} + +__attribute__((visibility("default"))) +void rf_gstreamer_cuda_dlpack_delete(DLManagedTensor* tensor) { + if (tensor != nullptr && tensor->deleter != nullptr) { + tensor->deleter(tensor); + } +} + +__attribute__((visibility("default"))) +int rf_gstreamer_cuda_pipeline_interrupt(RfCudaPipeline* handle) { + if (handle == nullptr) { + return -1; + } + handle->interrupted.store(true, std::memory_order_release); + if (handle->sink != nullptr) { + gst_app_sink_set_drop(handle->sink, TRUE); + GstSample* queued_sample = nullptr; + while ((queued_sample = gst_app_sink_try_pull_sample(handle->sink, 0)) != + nullptr) { + gst_sample_unref(queued_sample); + } + } + if (handle->pipeline != nullptr) { + gst_element_set_state(handle->pipeline, GST_STATE_NULL); + } + return 1; +} + +__attribute__((visibility("default"))) +void rf_gstreamer_cuda_pipeline_release(RfCudaPipeline* handle) { + if (handle == nullptr) { + return; + } + rf_gstreamer_cuda_pipeline_interrupt(handle); + { + std::lock_guard lock(handle->mutex); + if (handle->sample != nullptr) { + gst_sample_unref(handle->sample); + handle->sample = nullptr; + } + if (handle->sink != nullptr) { + gst_app_sink_set_drop(handle->sink, TRUE); + GstSample* queued_sample = nullptr; + while ((queued_sample = + gst_app_sink_try_pull_sample(handle->sink, 0)) != + nullptr) { + gst_sample_unref(queued_sample); + } + } + if (handle->pipeline != nullptr) { + gst_element_set_state(handle->pipeline, GST_STATE_NULL); + } + if (handle->sink != nullptr) { + gst_object_unref(handle->sink); + } + if (handle->pipeline != nullptr) { + gst_object_unref(handle->pipeline); + } + if (handle->cuda_context != nullptr) { + gst_object_unref(handle->cuda_context); + } + // Null the element pointers before the handle is freed so a + // contract-violating late interrupt() dereferences null instead of + // freed objects. The Python wrapper serializes interrupt()/close(). + handle->sink = nullptr; + handle->pipeline = nullptr; + handle->cuda_context = nullptr; + cuDevicePrimaryCtxRelease(handle->device); + if (handle->leases->references.fetch_sub(1, std::memory_order_acq_rel) == 1) { + delete handle->leases; + } + } + delete handle; +} + +} // extern "C" diff --git a/docker/native/jetson_tensor_bridge/jetson_tensor_bridge.cu b/docker/native/jetson_tensor_bridge/jetson_tensor_bridge.cu new file mode 100644 index 0000000000..1070ad1356 --- /dev/null +++ b/docker/native/jetson_tensor_bridge/jetson_tensor_bridge.cu @@ -0,0 +1,1822 @@ +#include +#include +#include + +#pragma push_macro("__noinline__") +#undef __noinline__ +#include +#include +#include +#include +#pragma pop_macro("__noinline__") + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr const char* kSinkName = "rf_tensor_sink"; +constexpr const char* kSrtpCapsElementName = "rf_srtp_caps"; +constexpr const char* kNvmmCapsFeature = "memory:NVMM"; +// Upper bound on device buffers the per-pipeline tensor pool recycles. The +// steady state holds only the frames a consumer keeps in flight (typically +// one or two); this cap only bounds a slow or bursty consumer. Buffers are +// allocated lazily on demand, so a small cap costs nothing until it is hit. +constexpr size_t kJetsonTensorPoolBuffers = 8; +// Handoff queue depth in lossless (file) mode: enough to absorb consumer +// jitter without noticeably delaying backpressure. Must stay below +// kJetsonTensorPoolBuffers or the pool would starve the converter. +constexpr size_t kLosslessHandoffCapacity = 4; + +enum DLDeviceType : int32_t { + kDLCUDA = 2, +}; + +enum DLDataTypeCode : uint8_t { + kDLUInt = 1, +}; + +struct DLDevice { + DLDeviceType device_type; + int32_t device_id; +}; + +struct DLDataType { + uint8_t code; + uint8_t bits; + uint16_t lanes; +}; + +struct DLTensor { + void* data; + DLDevice device; + int32_t ndim; + DLDataType dtype; + int64_t* shape; + int64_t* strides; + uint64_t byte_offset; +}; + +struct DLManagedTensor { + DLTensor dl_tensor; + void* manager_ctx; + void (*deleter)(DLManagedTensor* self); +}; + +struct RfFrameInfo { + uint32_t width; + uint32_t height; + int32_t fps_numerator; + int32_t fps_denominator; + int64_t duration_ns; +}; + +struct RfBridgeStats { + uint64_t frames; + uint64_t descriptor_maps; + // The next four counters must stay zero on the zero-copy path; the + // verify scripts assert on them to prove no host fallback executed. + uint64_t host_pixel_maps; + uint64_t host_to_device_copies; + uint64_t device_to_host_copies; + uint64_t array_flatten_copies; + uint64_t conversion_kernels; + uint64_t nvmm_frames; + // Pending live samples replaced before consumer-side conversion + // (latest-wins handoff slot). + uint64_t frames_dropped_by_consumer; + int32_t last_nvbuf_memory_type; + int32_t last_egl_frame_type; + int32_t last_egl_color_format; + // --- ABI v5: per-phase timing of CUDA conversion (ns). --- + // Under concurrent TRT/torch GPU load individual native calls in + // convert_sample_to_tensor() intermittently stall for hundreds of ms; + // these accumulators (total + observed max per phase) name the stalling + // call. Totals include failed conversion attempts (phases that ran). + uint64_t egl_map_ns; + uint64_t egl_map_max_ns; + uint64_t cuda_register_ns; + uint64_t cuda_register_max_ns; + uint64_t texture_create_ns; + uint64_t texture_create_max_ns; + uint64_t kernel_launch_ns; + uint64_t kernel_launch_max_ns; + uint64_t sync_ns; + uint64_t sync_max_ns; + uint64_t cleanup_ns; + uint64_t cleanup_max_ns; + // Distinct dmabuf fds (NvBufSurfaceParams.bufferDesc) seen on this + // pipeline โ€” the decoder capture-pool size. A small, stable set is the + // prerequisite for caching EGL/CUDA registrations per pool surface. + uint64_t unique_buffer_fds; + // --- ABI v7: EGL registration cache effectiveness. Hits skip the + // per-frame NvBufSurfaceMapEglImage/cuGraphicsEGLRegisterImage/texture + // creation/unregister sequence whose process-global locks measured + // 13-19 ms mean (max 97 ms) under GPU saturation. Expected hit rate + // after warmup: >99% (decoder pools recycle ~11 stable fds). + uint64_t egl_cache_hits; + uint64_t egl_cache_misses; +}; + +// Recycles fixed-size device allocations so the hot retrieve()/free path never +// calls cudaMalloc/cudaFree. On the Jetson unified-memory allocator a per-frame +// cudaFree synchronizes the whole device, stalling the consumer that drops a +// tensor; returning the buffer to a free list instead makes release a pure CPU +// push. The pool is reference-counted (shared_ptr) so it outlives the pipeline +// whenever a consumer is still holding a tensor. +class RfBufferPool { + public: + RfBufferPool(int device_id, size_t max_buffers) + : device_id_(device_id), max_buffers_(max_buffers) {} + + ~RfBufferPool() { + // The pool is only destroyed once the pipeline and every outstanding + // tensor have released, so all pooled buffers are back on the free + // list here. Bind the device because the destructor may run on a + // consumer thread that never touched CUDA. + int previous_device = -1; + cudaGetDevice(&previous_device); + cudaSetDevice(device_id_); + for (void* buffer : free_list_) { + cudaFree(buffer); + } + if (previous_device >= 0 && previous_device != device_id_) { + cudaSetDevice(previous_device); + } + } + + // Hands out a device buffer of `size` bytes. The caller (retrieve()) has + // already bound the device. `*pooled` reports whether release() should + // recycle the buffer (true) or free it immediately (false, for one-off + // over-budget or off-size allocations). Returns nullptr on cudaMalloc + // failure. + void* acquire(size_t size, bool* pooled) { + std::lock_guard lock(mutex_); + if (buffer_size_ == 0) { + buffer_size_ = size; // Adopt the first frame's size as the pool size. + } + if (size == buffer_size_ && !free_list_.empty()) { + void* buffer = free_list_.back(); + free_list_.pop_back(); + *pooled = true; + return buffer; + } + // Cold start or pool miss: allocate. cudaMalloc under the lock is fine + // here โ€” this is the slow path, not the steady-state recycle. + void* buffer = nullptr; + if (cudaMalloc(&buffer, size) != cudaSuccess) { + return nullptr; + } + // Only track pool-sized buffers within budget; anything else is a + // one-off (resolution change mid-stream, or a consumer holding more + // frames than the cap) that release() frees directly. + if (size == buffer_size_ && total_pooled_ < max_buffers_) { + total_pooled_ += 1; + *pooled = true; + } else { + *pooled = false; + } + return buffer; + } + + // Returns a buffer handed out by acquire(). Pooled buffers go back on the + // free list (no CUDA call, hence no device-wide sync); one-offs are freed. + // For the one-off path the caller must have bound the device. + void release(void* buffer, bool pooled) { + if (buffer == nullptr) { + return; + } + if (pooled) { + std::lock_guard lock(mutex_); + free_list_.push_back(buffer); + return; + } + cudaFree(buffer); + } + + private: + int device_id_; + size_t max_buffers_; + size_t buffer_size_ = 0; + size_t total_pooled_ = 0; + std::vector free_list_; + std::mutex mutex_; +}; + +struct RfTensorContext { + void* allocation = nullptr; + int device_id = 0; + int64_t shape[3] = {0, 0, 0}; + // Non-null when `allocation` came from the pool; release routes back to it + // instead of cudaFree. Held by shared_ptr so the pool survives a pipeline + // that closes while this tensor is still alive. + std::shared_ptr pool; + bool pooled = false; + DLManagedTensor managed{}; +}; + +// One cached EGL/CUDA registration for a decoder-pool surface, keyed by its +// dmabuf fd. `egl_image` doubles as the validity token: if the surface comes +// back with a different (or cleared) mappedAddr.eglImage, the fd was reused +// for a new surface and the entry is rebuilt. +struct RfEglCacheEntry { + uint64_t buffer_desc = 0; + void* egl_image = nullptr; + CUgraphicsResource resource = nullptr; + CUeglFrame frame{}; + cudaTextureObject_t textures[2] = {0, 0}; + uint32_t texture_count = 0; +}; + +struct RfJetsonPipeline { + GstElement* pipeline = nullptr; + GstAppSink* sink = nullptr; + cudaStream_t stream = nullptr; + int device_id = 0; + std::shared_ptr tensor_pool; + RfBridgeStats stats{}; + std::atomic interrupted{false}; + // Serializes consumer-side EGL/CUDA conversion with interrupt()/release() + // transitioning the decoder pipeline to NULL. + std::mutex conversion_mutex; + std::mutex mutex; + // The appsink callback only removes samples from GStreamer's streaming + // thread and publishes them here. grab() performs the EGL/CUDA conversion + // on the consumer thread; doing that work inside nvv4l2decoder's callback + // can deadlock before the first frame on JetPack 6.2. Two modes: + // * live (lossless_handoff=false, capacity 1): a newer pending sample + // replaces an uncollected pending sample (latest-wins); + // * lossless (lossless_handoff=true, capacity kLosslessHandoffCapacity): + // the callback blocks while pending samples plus converted tensors fill + // the queue, backpressuring file decode without dropping frames. + std::condition_variable frame_ready; + std::condition_variable handoff_space; + std::deque ready_samples; + std::deque ready_tensors; + bool lossless_handoff = false; + size_t handoff_capacity = 1; + std::atomic eos{false}; + bool conversion_failed = false; + char conversion_error[1024] = {0}; + RfFrameInfo last_frame_info{}; + bool frame_info_valid = false; + // dmabuf fds observed on this pipeline (guarded by `mutex`); its size is + // exported as stats.unique_buffer_fds. Decoder capture pools are small + // (single digits), so an ordered set is fine. + std::set seen_buffer_fds; + // EGL/CUDA registration cache keyed by dmabuf fd (the DeepStream + // pattern): decoder pools recycle a small stable set of surfaces, so + // the EGL image mapping, CUDA graphics registration and texture objects + // survive across frames instead of paying the process-global-lock + // sequence per frame. Touched only by the consumer thread inside + // rf_jetson_pipeline_grab(); release() tears it down after the GStreamer + // streaming thread is joined (GST_STATE_NULL). + std::vector egl_cache; + GstPad* srtp_probe_pad = nullptr; + gulong srtp_probe_id = 0; +}; + +RfEglCacheEntry* find_egl_cache_entry( + RfJetsonPipeline* handle, uint64_t buffer_desc) { + for (auto& entry : handle->egl_cache) { + if (entry.buffer_desc == buffer_desc) { + return &entry; + } + } + return nullptr; +} + +// Frees the CUDA-side resources of an entry (textures + graphics +// registration). The surface's EGL mapping itself is deliberately left in +// place: it belongs to the decoder-pool surface and dies with it. +void destroy_egl_cache_entry_resources(RfEglCacheEntry* entry) { + for (uint32_t plane = 0; plane < entry->texture_count; ++plane) { + if (entry->textures[plane] != 0) { + cudaDestroyTextureObject(entry->textures[plane]); + } + } + if (entry->resource != nullptr) { + cuGraphicsUnregisterResource(entry->resource); + } + entry->resource = nullptr; + entry->egl_image = nullptr; + entry->textures[0] = 0; + entry->textures[1] = 0; + entry->texture_count = 0; +} + +void signal_pipeline_error(RfJetsonPipeline* handle, const char* message) { + std::lock_guard lock(handle->mutex); + handle->conversion_failed = true; + std::snprintf( + handle->conversion_error, + sizeof(handle->conversion_error), + "%s", + message); + handle->frame_ready.notify_all(); +} + +GstPadProbeReturn configure_sdes_srtp_caps( + GstPad* pad, + GstPadProbeInfo* info, + gpointer user_data) { + auto* handle = static_cast(user_data); + if ((GST_PAD_PROBE_INFO_TYPE(info) & GST_PAD_PROBE_TYPE_EVENT_DOWNSTREAM) == 0) { + return GST_PAD_PROBE_OK; + } + GstEvent* event = GST_PAD_PROBE_INFO_EVENT(info); + if (event == nullptr || GST_EVENT_TYPE(event) != GST_EVENT_CAPS) { + return GST_PAD_PROBE_OK; + } + GstCaps* input_caps = nullptr; + gst_event_parse_caps(event, &input_caps); + if (input_caps == nullptr || gst_caps_is_empty(input_caps) || + gst_caps_is_any(input_caps)) { + signal_pipeline_error(handle, "SRTP stream caps are missing"); + return GST_PAD_PROBE_DROP; + } + const GstStructure* input = gst_caps_get_structure(input_caps, 0); + const char* media = gst_structure_get_string(input, "media"); + const char* encoding = gst_structure_get_string(input, "encoding-name"); + const bool supported_media = + media != nullptr && std::strcmp(media, "video") == 0 && + encoding != nullptr && + (std::strcmp(encoding, "H264") == 0 || + std::strcmp(encoding, "H265") == 0); + if (!supported_media) { + signal_pipeline_error(handle, "SRTP caps do not describe H.264/H.265 video"); + return GST_PAD_PROBE_DROP; + } + const char* crypto = gst_structure_get_string(input, "a-crypto"); + constexpr const char* kCryptoSuite = "AES_CM_128_HMAC_SHA1_80"; + constexpr const char* kInlinePrefix = "inline:"; + const char* tag_end = crypto == nullptr ? nullptr : std::strchr(crypto, ' '); + bool valid_tag = tag_end != nullptr && tag_end != crypto; + for (const char* cursor = crypto; valid_tag && cursor < tag_end; ++cursor) { + valid_tag = *cursor >= '0' && *cursor <= '9'; + } + const char* suite_start = tag_end; + while (suite_start != nullptr && *suite_start == ' ') { + ++suite_start; + } + const char* suite_end = + suite_start == nullptr ? nullptr : std::strchr(suite_start, ' '); + const bool supported_suite = + suite_end != nullptr && + static_cast(suite_end - suite_start) == std::strlen(kCryptoSuite) && + std::strncmp(suite_start, kCryptoSuite, std::strlen(kCryptoSuite)) == 0; + const char* key_start = suite_end; + while (key_start != nullptr && *key_start == ' ') { + ++key_start; + } + if (!valid_tag || !supported_suite || key_start == nullptr || + std::strncmp(key_start, kInlinePrefix, std::strlen(kInlinePrefix)) != 0) { + signal_pipeline_error( + handle, + "SRTP stream does not advertise supported SDES AES-128/HMAC-SHA1-80"); + return GST_PAD_PROBE_DROP; + } + key_start += std::strlen(kInlinePrefix); + const size_t encoded_length = std::strcspn(key_start, "| \t\r\n"); + bool valid_base64 = encoded_length == 40; + for (size_t index = 0; valid_base64 && index < encoded_length; ++index) { + const char value = key_start[index]; + valid_base64 = (value >= 'A' && value <= 'Z') || + (value >= 'a' && value <= 'z') || + (value >= '0' && value <= '9') || value == '+' || value == '/'; + } + if (!valid_base64) { + signal_pipeline_error(handle, "SRTP SDES key is not valid base64"); + return GST_PAD_PROBE_DROP; + } + const char* suffix = key_start + encoded_length; + if (*suffix == '|') { + signal_pipeline_error( + handle, + "SRTP SDES lifetime and MKI parameters are not supported"); + return GST_PAD_PROBE_DROP; + } + while (*suffix == ' ' || *suffix == '\t' || *suffix == '\r' || + *suffix == '\n') { + ++suffix; + } + if (*suffix != '\0') { + signal_pipeline_error( + handle, + "SRTP SDES session parameters are not supported"); + return GST_PAD_PROBE_DROP; + } + std::string encoded_key(key_start, encoded_length); + gsize key_length = 0; + guchar* key_bytes = g_base64_decode(encoded_key.c_str(), &key_length); + std::fill(encoded_key.begin(), encoded_key.end(), '\0'); + constexpr gsize kAes128MasterKeyAndSaltBytes = 30; + if (key_bytes == nullptr || key_length != kAes128MasterKeyAndSaltBytes) { + if (key_bytes != nullptr) { + std::fill_n(key_bytes, key_length, 0); + } + g_free(key_bytes); + signal_pipeline_error(handle, "SRTP SDES key has an invalid length"); + return GST_PAD_PROBE_DROP; + } + GstBuffer* key_buffer = gst_buffer_new_allocate(nullptr, key_length, nullptr); + if (key_buffer == nullptr || + gst_buffer_fill(key_buffer, 0, key_bytes, key_length) != key_length) { + std::fill_n(key_bytes, key_length, 0); + g_free(key_bytes); + if (key_buffer != nullptr) { + gst_buffer_unref(key_buffer); + } + signal_pipeline_error(handle, "Could not allocate SRTP key buffer"); + return GST_PAD_PROBE_DROP; + } + std::fill_n(key_bytes, key_length, 0); + g_free(key_bytes); + + GstCaps* secure_caps = gst_caps_copy(input_caps); + for (guint index = 0; index < gst_caps_get_size(secure_caps); ++index) { + GstStructure* structure = gst_caps_get_structure(secure_caps, index); + gst_structure_set_name(structure, "application/x-srtp"); + gst_structure_remove_field(structure, "a-crypto"); + gst_structure_set( + structure, + "srtp-key", + GST_TYPE_BUFFER, + key_buffer, + "srtp-cipher", + G_TYPE_STRING, + "aes-128-icm", + "srtp-auth", + G_TYPE_STRING, + "hmac-sha1-80", + "srtcp-cipher", + G_TYPE_STRING, + "aes-128-icm", + "srtcp-auth", + G_TYPE_STRING, + "hmac-sha1-80", + nullptr); + } + GstEvent* secure_event = gst_event_new_caps(secure_caps); + gst_caps_unref(secure_caps); + gst_buffer_unref(key_buffer); + if (secure_event == nullptr) { + signal_pipeline_error(handle, "Could not create secure RTP caps event"); + return GST_PAD_PROBE_DROP; + } + gst_event_unref(event); + GST_PAD_PROBE_INFO_DATA(info) = secure_event; + (void)pad; + return GST_PAD_PROBE_OK; +} + +using NvBufSurfaceMapEglImageFn = int (*)(NvBufSurface*, int); +using NvBufSurfaceUnMapEglImageFn = int (*)(NvBufSurface*, int); + +struct NvBufSurfaceApi { + void* library = nullptr; + NvBufSurfaceMapEglImageFn map_egl_image = nullptr; + NvBufSurfaceUnMapEglImageFn unmap_egl_image = nullptr; + const char* error = nullptr; +}; + +struct ChannelMap { + int red; + int green; + int blue; +}; + +std::once_flag g_gstreamer_once; +std::once_flag g_nvbufsurface_once; +NvBufSurfaceApi g_nvbufsurface; +char g_dlerror_message[256]; + +void write_error(char* destination, size_t capacity, const char* format, ...) { + if (destination == nullptr || capacity == 0) { + return; + } + va_list args; + va_start(args, format); + std::vsnprintf(destination, capacity, format, args); + va_end(args); + destination[capacity - 1] = '\0'; +} + +void initialize_gstreamer() { + gst_init(nullptr, nullptr); +} + +const char* capture_dlerror() { + // dlerror() returns a pointer into a thread-local buffer that later dl* + // calls (Python's ctypes dlopens constantly) overwrite; snapshot it. + const char* message = dlerror(); + if (message == nullptr) { + return nullptr; + } + std::snprintf(g_dlerror_message, sizeof(g_dlerror_message), "%s", message); + return g_dlerror_message; +} + +void initialize_nvbufsurface() { + const char* candidates[] = { + "libnvbufsurface.so.1.0.0", + "libnvbufsurface.so", + "/usr/lib/aarch64-linux-gnu/nvidia/libnvbufsurface.so.1.0.0", + "/usr/lib/aarch64-linux-gnu/tegra/libnvbufsurface.so.1.0.0", + }; + for (const char* candidate : candidates) { + g_nvbufsurface.library = dlopen(candidate, RTLD_NOW | RTLD_LOCAL); + if (g_nvbufsurface.library != nullptr) { + break; + } + } + if (g_nvbufsurface.library == nullptr) { + g_nvbufsurface.error = capture_dlerror(); + return; + } + g_nvbufsurface.map_egl_image = + reinterpret_cast( + dlsym(g_nvbufsurface.library, "NvBufSurfaceMapEglImage")); + g_nvbufsurface.unmap_egl_image = + reinterpret_cast( + dlsym(g_nvbufsurface.library, "NvBufSurfaceUnMapEglImage")); + if (g_nvbufsurface.map_egl_image == nullptr || + g_nvbufsurface.unmap_egl_image == nullptr) { + g_nvbufsurface.error = capture_dlerror(); + } +} + +bool get_channel_map(CUeglColorFormat format, ChannelMap* result) { + if (result == nullptr) { + return false; + } + switch (format) { + case CU_EGL_COLOR_FORMAT_ABGR: + *result = {0, 1, 2}; + return true; + case CU_EGL_COLOR_FORMAT_RGBA: + *result = {3, 2, 1}; + return true; + case CU_EGL_COLOR_FORMAT_ARGB: + *result = {2, 1, 0}; + return true; + case CU_EGL_COLOR_FORMAT_BGRA: + *result = {1, 2, 3}; + return true; + default: + return false; + } +} + +__device__ inline void store_rgb_chw( + uchar4 pixel, + ChannelMap channels, + uint8_t* destination, + uint32_t index, + uint32_t plane_size) { + const uint8_t values[4] = {pixel.x, pixel.y, pixel.z, pixel.w}; + destination[index] = values[channels.red]; + destination[plane_size + index] = values[channels.green]; + destination[2 * plane_size + index] = values[channels.blue]; +} + +__global__ void rgba_pitch_to_rgb_chw( + const uint8_t* source, + size_t source_pitch, + uint8_t* destination, + uint32_t width, + uint32_t height, + ChannelMap channels) { + const uint32_t x = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= width || y >= height) { + return; + } + const uchar4* row = reinterpret_cast(source + y * source_pitch); + const uint32_t index = y * width + x; + store_rgb_chw(row[x], channels, destination, index, width * height); +} + +__global__ void rgba_array_to_rgb_chw( + cudaTextureObject_t texture, + uint8_t* destination, + uint32_t width, + uint32_t height, + ChannelMap channels) { + const uint32_t x = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= width || y >= height) { + return; + } + const uchar4 pixel = tex2D(texture, x + 0.5f, y + 0.5f); + const uint32_t index = y * width + x; + store_rgb_chw(pixel, channels, destination, index, width * height); +} + +// YUV -> RGB conversion coefficients for the direct-NV12 path (no nvvidconv +// RGBA hop): R = y_scale*(Y - y_offset) + r_v*(V-128), etc. +struct YuvCoeffs { + float y_scale; + float y_offset; + float r_v; + float g_u; + float g_v; + float b_u; +}; + +bool select_nv12_coeffs(NvBufSurfaceColorFormat format, YuvCoeffs* result) { + switch (format) { + case NVBUF_COLOR_FORMAT_NV12: // BT.601 limited range (decoder default) + *result = {1.1644f, 16.0f, 1.5960f, -0.3918f, -0.8130f, 2.0172f}; + return true; + case NVBUF_COLOR_FORMAT_NV12_ER: // BT.601 full range + *result = {1.0f, 0.0f, 1.4020f, -0.3441f, -0.7141f, 1.7720f}; + return true; + case NVBUF_COLOR_FORMAT_NV12_709: // BT.709 limited range + *result = {1.1644f, 16.0f, 1.7927f, -0.2132f, -0.5329f, 2.1124f}; + return true; + case NVBUF_COLOR_FORMAT_NV12_709_ER: // BT.709 full range + *result = {1.0f, 0.0f, 1.5748f, -0.1873f, -0.4681f, 1.8556f}; + return true; + default: + return false; + } +} + +__device__ inline void store_yuv_as_rgb_chw( + float y_value, + float u_value, + float v_value, + YuvCoeffs coeffs, + uint8_t* destination, + uint32_t index, + uint32_t plane_size) { + const float luma = coeffs.y_scale * (y_value - coeffs.y_offset); + const float d = u_value - 128.0f; + const float e = v_value - 128.0f; + const float red = luma + coeffs.r_v * e; + const float green = luma + coeffs.g_u * d + coeffs.g_v * e; + const float blue = luma + coeffs.b_u * d; + destination[index] = + static_cast(fminf(fmaxf(red, 0.0f), 255.0f)); + destination[plane_size + index] = + static_cast(fminf(fmaxf(green, 0.0f), 255.0f)); + destination[2 * plane_size + index] = + static_cast(fminf(fmaxf(blue, 0.0f), 255.0f)); +} + +__global__ void nv12_pitch_to_rgb_chw( + const uint8_t* y_plane, + const uint8_t* uv_plane, + size_t pitch, + uint8_t* destination, + uint32_t width, + uint32_t height, + YuvCoeffs coeffs) { + const uint32_t x = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= width || y >= height) { + return; + } + const float y_value = static_cast(y_plane[y * pitch + x]); + const uint8_t* uv_row = uv_plane + (y / 2) * pitch; + const uint32_t uv_x = (x / 2) * 2; + const float u_value = static_cast(uv_row[uv_x]); + const float v_value = static_cast(uv_row[uv_x + 1]); + const uint32_t index = y * width + x; + store_yuv_as_rgb_chw( + y_value, u_value, v_value, coeffs, destination, index, width * height); +} + +__global__ void nv12_array_to_rgb_chw( + cudaTextureObject_t y_texture, + cudaTextureObject_t uv_texture, + uint8_t* destination, + uint32_t width, + uint32_t height, + YuvCoeffs coeffs) { + const uint32_t x = blockIdx.x * blockDim.x + threadIdx.x; + const uint32_t y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= width || y >= height) { + return; + } + const float y_value = + static_cast(tex2D(y_texture, x + 0.5f, y + 0.5f)); + const uchar2 uv = + tex2D(uv_texture, (x / 2) + 0.5f, (y / 2) + 0.5f); + const uint32_t index = y * width + x; + store_yuv_as_rgb_chw( + y_value, + static_cast(uv.x), + static_cast(uv.y), + coeffs, + destination, + index, + width * height); +} + +void delete_managed_tensor(DLManagedTensor* managed) { + if (managed == nullptr) { + return; + } + auto* context = static_cast(managed->manager_ctx); + if (context == nullptr) { + return; + } + int previous_device = -1; + cudaGetDevice(&previous_device); + cudaSetDevice(context->device_id); + if (context->allocation != nullptr) { + if (context->pool != nullptr) { + // Pooled buffers return without a CUDA call; only one-offs hit + // cudaFree here, and both need the device bound (above) in case + // this runs on a consumer thread that never touched CUDA. + context->pool->release(context->allocation, context->pooled); + } else { + cudaFree(context->allocation); + } + } + if (previous_device >= 0 && previous_device != context->device_id) { + cudaSetDevice(previous_device); + } + delete context; +} + +bool sample_has_nvmm_caps(GstSample* sample) { + GstCaps* caps = gst_sample_get_caps(sample); + if (caps == nullptr || gst_caps_is_empty(caps)) { + return false; + } + for (guint index = 0; index < gst_caps_get_size(caps); ++index) { + const GstCapsFeatures* features = gst_caps_get_features(caps, index); + if (features != nullptr && + gst_caps_features_contains(features, kNvmmCapsFeature)) { + return true; + } + } + return false; +} + +bool read_sample_info(GstSample* sample, RfFrameInfo* info) { + if (sample == nullptr || info == nullptr) { + return false; + } + GstCaps* caps = gst_sample_get_caps(sample); + if (caps == nullptr || gst_caps_is_empty(caps)) { + return false; + } + const GstStructure* structure = gst_caps_get_structure(caps, 0); + int width = 0; + int height = 0; + int numerator = 0; + int denominator = 1; + if (!gst_structure_get_int(structure, "width", &width) || + !gst_structure_get_int(structure, "height", &height)) { + return false; + } + gst_structure_get_fraction( + structure, "framerate", &numerator, &denominator); + info->width = static_cast(width); + info->height = static_cast(height); + info->fps_numerator = numerator; + info->fps_denominator = denominator > 0 ? denominator : 1; + return true; +} + +bool read_bus_error( + RfJetsonPipeline* handle, + char* error, + size_t error_capacity) { + GstBus* bus = gst_element_get_bus(handle->pipeline); + if (bus == nullptr) { + return false; + } + GstMessage* message = gst_bus_pop_filtered(bus, GST_MESSAGE_ERROR); + bool has_error = false; + if (message != nullptr) { + has_error = true; + GError* gst_error = nullptr; + gchar* debug = nullptr; + gst_message_parse_error(message, &gst_error, &debug); + write_error( + error, + error_capacity, + "%s", + gst_error != nullptr ? gst_error->message : "GStreamer pipeline error"); + if (gst_error != nullptr) { + g_error_free(gst_error); + } + g_free(debug); + gst_message_unref(message); + } + gst_object_unref(bus); + return has_error; +} + +bool pipeline_has_factory(RfJetsonPipeline* handle, const char* factory_name) { + if (handle == nullptr || handle->pipeline == nullptr || factory_name == nullptr || + !GST_IS_BIN(handle->pipeline)) { + return false; + } + GstIterator* iterator = gst_bin_iterate_recurse(GST_BIN(handle->pipeline)); + GValue item = G_VALUE_INIT; + bool found = false; + bool done = false; + while (!done && !found) { + switch (gst_iterator_next(iterator, &item)) { + case GST_ITERATOR_OK: { + auto* element = GST_ELEMENT(g_value_get_object(&item)); + GstElementFactory* factory = gst_element_get_factory(element); + const gchar* name = factory == nullptr + ? nullptr + : gst_plugin_feature_get_name(GST_PLUGIN_FEATURE(factory)); + found = name != nullptr && std::strcmp(name, factory_name) == 0; + g_value_reset(&item); + break; + } + case GST_ITERATOR_RESYNC: + gst_iterator_resync(iterator); + break; + default: + done = true; + break; + } + } + if (G_VALUE_TYPE(&item) != 0) { + g_value_unset(&item); + } + gst_iterator_free(iterator); + return found; +} + +struct RfEglDiagnostics { + int32_t memory_type = -1; + int32_t frame_type = -1; + int32_t color_format = -1; +}; + +// Per-attempt phase durations measured inside convert_sample_to_tensor() and +// accumulated into RfBridgeStats under the handle mutex by the caller. +struct RfPhaseTimings { + uint64_t egl_map_ns = 0; + uint64_t cuda_register_ns = 0; + uint64_t texture_create_ns = 0; + uint64_t kernel_launch_ns = 0; + uint64_t sync_ns = 0; + uint64_t cleanup_ns = 0; + uint64_t buffer_fd = 0; + bool buffer_fd_valid = false; + uint32_t egl_cache_hit = 0; + uint32_t egl_cache_miss = 0; +}; + +uint64_t monotonic_ns() { + return static_cast( + std::chrono::duration_cast( + std::chrono::steady_clock::now().time_since_epoch()) + .count()); +} + +// Convert one appsink sample into a pooled CHW RGB CUDA tensor on the consumer +// thread. The appsink callback has already removed the sample from GStreamer's +// streaming thread, avoiding the JP6.2 deadlock caused by performing EGL/CUDA +// work synchronously inside nvv4l2decoder's callback. Supports both direct +// decoder output (NV12 semi-planar) and nvvidconv output (RGBA). +RfTensorContext* convert_sample_to_tensor( + RfJetsonPipeline* handle, + GstSample* sample, + RfFrameInfo* frame_info_out, + RfEglDiagnostics* diagnostics, + RfPhaseTimings* timings, + char* error, + size_t error_capacity) { + GstBuffer* buffer = gst_sample_get_buffer(sample); + GstMapInfo map = GST_MAP_INFO_INIT; + bool buffer_mapped = false; + bool egl_mapped = false; + CUgraphicsResource graphics_resource = nullptr; + cudaTextureObject_t textures[2] = {0, 0}; + RfTensorContext* tensor = nullptr; + RfTensorContext* result = nullptr; + uint64_t phase_start = 0; + bool egl_resources_cached = false; + RfEglCacheEntry* egl_cache_entry = nullptr; + + // Bind the device on this consumer thread; the driver-API EGL registration + // below needs a current CUDA context. + cudaError_t bind_status = cudaSetDevice(handle->device_id); + if (bind_status == cudaSuccess) { + bind_status = cudaFree(nullptr); + } + if (bind_status != cudaSuccess) { + write_error( + error, + error_capacity, + "CUDA device %d binding failed: %s", + handle->device_id, + cudaGetErrorString(bind_status)); + return nullptr; + } + + RfFrameInfo frame_info{}; + if (!read_sample_info(sample, &frame_info) || !sample_has_nvmm_caps(sample)) { + write_error(error, error_capacity, "Frame caps are invalid"); + goto cleanup; + } + if (buffer == nullptr || !gst_buffer_map(buffer, &map, GST_MAP_READ) || + map.data == nullptr) { + write_error(error, error_capacity, "Could not map the NvBufSurface descriptor"); + goto cleanup; + } + buffer_mapped = true; + + { + auto* surface = reinterpret_cast(map.data); + diagnostics->memory_type = static_cast(surface->memType); + if (surface->memType != NVBUF_MEM_SURFACE_ARRAY || + surface->surfaceList == nullptr || surface->batchSize == 0 || + surface->numFilled == 0) { + write_error( + error, + error_capacity, + "Frame is not a populated NvBufSurface SURFACE_ARRAY"); + goto cleanup; + } + timings->buffer_fd = surface->surfaceList[0].bufferDesc; + timings->buffer_fd_valid = true; + const NvBufSurfaceColorFormat surface_format = + surface->surfaceList[0].colorFormat; + const bool is_rgba = surface_format == NVBUF_COLOR_FORMAT_RGBA; + YuvCoeffs yuv_coeffs{}; + if (!is_rgba && !select_nv12_coeffs(surface_format, &yuv_coeffs)) { + write_error( + error, + error_capacity, + "NvBufSurface format is unsupported (format=%d, expected RGBA " + "or an NV12 variant)", + static_cast(surface_format)); + goto cleanup; + } + egl_cache_entry = find_egl_cache_entry(handle, timings->buffer_fd); + if (egl_cache_entry != nullptr) { + void* current_egl_image = surface->surfaceList[0].mappedAddr.eglImage; + if (current_egl_image != nullptr && + egl_cache_entry->egl_image == current_egl_image) { + egl_resources_cached = true; + timings->egl_cache_hit = 1; + } else { + // The fd was reused for a different surface (pool + // reallocation) - free the stale registration and rebuild + // into the same slot below. + destroy_egl_cache_entry_resources(egl_cache_entry); + } + } + CUeglFrame egl_frame{}; + if (egl_resources_cached) { + egl_frame = egl_cache_entry->frame; + } else { + timings->egl_cache_miss = 1; + phase_start = monotonic_ns(); + if (g_nvbufsurface.map_egl_image(surface, 0) != 0) { + write_error( + error, error_capacity, "NvBufSurface EGL mapping failed"); + goto cleanup; + } + timings->egl_map_ns = monotonic_ns() - phase_start; + egl_mapped = true; + void* egl_image = surface->surfaceList[0].mappedAddr.eglImage; + if (egl_image == nullptr) { + write_error( + error, error_capacity, "NvBufSurface EGL image is null"); + goto cleanup; + } + phase_start = monotonic_ns(); + CUresult driver_status = cuGraphicsEGLRegisterImage( + &graphics_resource, + reinterpret_cast(egl_image), + CU_GRAPHICS_MAP_RESOURCE_FLAGS_READ_ONLY); + if (driver_status != CUDA_SUCCESS) { + const char* driver_error = nullptr; + cuGetErrorString(driver_status, &driver_error); + write_error( + error, + error_capacity, + "CUDA EGL registration failed: %s", + driver_error == nullptr ? "unknown error" : driver_error); + goto cleanup; + } + driver_status = cuGraphicsResourceGetMappedEglFrame( + &egl_frame, graphics_resource, 0, 0); + if (driver_status != CUDA_SUCCESS) { + const char* driver_error = nullptr; + cuGetErrorString(driver_status, &driver_error); + write_error( + error, + error_capacity, + "CUDA EGL frame mapping failed: %s", + driver_error == nullptr ? "unknown error" : driver_error); + goto cleanup; + } + timings->cuda_register_ns = monotonic_ns() - phase_start; + } + diagnostics->frame_type = static_cast(egl_frame.frameType); + diagnostics->color_format = static_cast(egl_frame.eglColorFormat); + const uint32_t expected_planes = is_rgba ? 1 : 2; + if (egl_frame.width < frame_info.width || + egl_frame.height < frame_info.height || + egl_frame.planeCount != expected_planes || + egl_frame.cuFormat != CU_AD_FORMAT_UNSIGNED_INT8 || + (is_rgba && egl_frame.numChannels != 4)) { + write_error(error, error_capacity, "CUDA EGL frame layout is invalid"); + goto cleanup; + } + + ChannelMap channels{}; + if (is_rgba && !get_channel_map(egl_frame.eglColorFormat, &channels)) { + write_error( + error, + error_capacity, + "CUDA EGL color format is unsupported: %d", + static_cast(egl_frame.eglColorFormat)); + goto cleanup; + } + tensor = new (std::nothrow) RfTensorContext(); + if (tensor == nullptr) { + write_error(error, error_capacity, "Could not allocate tensor state"); + goto cleanup; + } + tensor->device_id = handle->device_id; + tensor->pool = handle->tensor_pool; + tensor->managed.manager_ctx = tensor; + tensor->managed.deleter = delete_managed_tensor; + const size_t output_size = + static_cast(frame_info.width) * frame_info.height * 3; + tensor->allocation = handle->tensor_pool->acquire(output_size, &tensor->pooled); + if (tensor->allocation == nullptr) { + write_error( + error, + error_capacity, + "CUDA tensor allocation failed for %zu bytes", + output_size); + goto cleanup; + } + + cudaError_t cuda_status = cudaSuccess; + const dim3 threads(32, 8); + const dim3 blocks( + (frame_info.width + threads.x - 1) / threads.x, + (frame_info.height + threads.y - 1) / threads.y); + if (egl_frame.frameType == CU_EGL_FRAME_TYPE_PITCH) { + phase_start = monotonic_ns(); + if (is_rgba) { + rgba_pitch_to_rgb_chw<<stream>>>( + static_cast(egl_frame.frame.pPitch[0]), + egl_frame.pitch, + static_cast(tensor->allocation), + frame_info.width, + frame_info.height, + channels); + } else { + nv12_pitch_to_rgb_chw<<stream>>>( + static_cast(egl_frame.frame.pPitch[0]), + static_cast(egl_frame.frame.pPitch[1]), + egl_frame.pitch, + static_cast(tensor->allocation), + frame_info.width, + frame_info.height, + yuv_coeffs); + } + timings->kernel_launch_ns = monotonic_ns() - phase_start; + } else if (egl_frame.frameType == CU_EGL_FRAME_TYPE_ARRAY) { + const uint32_t texture_count = is_rgba ? 1 : 2; + if (egl_resources_cached) { + textures[0] = egl_cache_entry->textures[0]; + textures[1] = egl_cache_entry->textures[1]; + } + phase_start = monotonic_ns(); + for (uint32_t plane = 0; + !egl_resources_cached && plane < texture_count; + ++plane) { + cudaResourceDesc resource_description{}; + resource_description.resType = cudaResourceTypeArray; + resource_description.res.array.array = + reinterpret_cast(egl_frame.frame.pArray[plane]); + cudaTextureDesc texture_description{}; + texture_description.addressMode[0] = cudaAddressModeClamp; + texture_description.addressMode[1] = cudaAddressModeClamp; + texture_description.filterMode = cudaFilterModePoint; + texture_description.readMode = cudaReadModeElementType; + texture_description.normalizedCoords = 0; + cuda_status = cudaCreateTextureObject( + &textures[plane], + &resource_description, + &texture_description, + nullptr); + if (cuda_status != cudaSuccess) { + write_error( + error, + error_capacity, + "CUDA texture creation failed: %s", + cudaGetErrorString(cuda_status)); + goto cleanup; + } + } + timings->texture_create_ns = monotonic_ns() - phase_start; + phase_start = monotonic_ns(); + if (is_rgba) { + rgba_array_to_rgb_chw<<stream>>>( + textures[0], + static_cast(tensor->allocation), + frame_info.width, + frame_info.height, + channels); + } else { + nv12_array_to_rgb_chw<<stream>>>( + textures[0], + textures[1], + static_cast(tensor->allocation), + frame_info.width, + frame_info.height, + yuv_coeffs); + } + timings->kernel_launch_ns = monotonic_ns() - phase_start; + } else { + write_error(error, error_capacity, "CUDA EGL frame storage is unsupported"); + goto cleanup; + } + // Timed as one phase: the sync is where the consumer thread parks + // behind any concurrent TRT/torch work. + phase_start = monotonic_ns(); + cuda_status = cudaGetLastError(); + if (cuda_status == cudaSuccess) { + cuda_status = cudaStreamSynchronize(handle->stream); + } + timings->sync_ns = monotonic_ns() - phase_start; + if (cuda_status != cudaSuccess) { + write_error( + error, + error_capacity, + "CUDA frame conversion failed: %s", + cudaGetErrorString(cuda_status)); + goto cleanup; + } + + tensor->shape[0] = 3; + tensor->shape[1] = frame_info.height; + tensor->shape[2] = frame_info.width; + tensor->managed.dl_tensor.data = tensor->allocation; + tensor->managed.dl_tensor.device = {kDLCUDA, handle->device_id}; + tensor->managed.dl_tensor.ndim = 3; + tensor->managed.dl_tensor.dtype = {kDLUInt, 8, 1}; + tensor->managed.dl_tensor.shape = tensor->shape; + tensor->managed.dl_tensor.strides = nullptr; + tensor->managed.dl_tensor.byte_offset = 0; + if (!egl_resources_cached) { + RfEglCacheEntry fresh_entry{}; + fresh_entry.buffer_desc = timings->buffer_fd; + fresh_entry.egl_image = surface->surfaceList[0].mappedAddr.eglImage; + fresh_entry.resource = graphics_resource; + fresh_entry.frame = egl_frame; + fresh_entry.textures[0] = textures[0]; + fresh_entry.textures[1] = textures[1]; + fresh_entry.texture_count = + egl_frame.frameType == CU_EGL_FRAME_TYPE_ARRAY + ? (is_rgba ? 1u : 2u) + : 0u; + if (egl_cache_entry != nullptr) { + *egl_cache_entry = fresh_entry; + } else { + handle->egl_cache.push_back(fresh_entry); + } + // Ownership moved to the cache: neutralise the locals so the + // cleanup section leaves the registration/textures alive, and + // keep the surface EGL-mapped for reuse on the next frame. + graphics_resource = nullptr; + textures[0] = 0; + textures[1] = 0; + egl_mapped = false; + egl_resources_cached = true; + } + *frame_info_out = frame_info; + result = tensor; + tensor = nullptr; + } + +cleanup: + phase_start = monotonic_ns(); + if (!egl_resources_cached) { + // On a cache hit `textures` alias the cached objects - they must + // survive this frame (even a failed one) for reuse. + for (cudaTextureObject_t texture : textures) { + if (texture != 0) { + cudaDestroyTextureObject(texture); + } + } + } + if (graphics_resource != nullptr) { + cuGraphicsUnregisterResource(graphics_resource); + } + if (egl_mapped) { + auto* surface = reinterpret_cast(map.data); + g_nvbufsurface.unmap_egl_image(surface, 0); + } + if (buffer_mapped) { + gst_buffer_unmap(buffer, &map); + } + if (tensor != nullptr) { + delete_managed_tensor(&tensor->managed); + } + timings->cleanup_ns = monotonic_ns() - phase_start; + return result; +} + +GstFlowReturn handle_new_sample(GstAppSink* sink, gpointer user_data) { + auto* handle = static_cast(user_data); + GstSample* sample = gst_app_sink_pull_sample(sink); + if (sample == nullptr) { + return GST_FLOW_OK; + } + if (handle->interrupted.load(std::memory_order_acquire)) { + gst_sample_unref(sample); + return GST_FLOW_OK; + } + { + std::unique_lock lock(handle->mutex); + // Lossless file mode backpressures decode at the bounded sample/tensor + // queue. Live mode keeps only the newest pending sample. + handle->handoff_space.wait(lock, [handle]() { + return !handle->lossless_handoff || + handle->ready_samples.size() + + handle->ready_tensors.size() < + handle->handoff_capacity || + handle->interrupted.load(std::memory_order_acquire); + }); + if (handle->interrupted.load(std::memory_order_acquire)) { + lock.unlock(); + gst_sample_unref(sample); + return GST_FLOW_OK; + } + if (!handle->lossless_handoff) { + while (!handle->ready_samples.empty()) { + gst_sample_unref(handle->ready_samples.front()); + handle->ready_samples.pop_front(); + handle->stats.frames_dropped_by_consumer += 1; + } + } + handle->ready_samples.push_back(sample); + } + handle->frame_ready.notify_all(); + return GST_FLOW_OK; +} + +GstFlowReturn handle_new_preroll(GstAppSink* sink, gpointer user_data) { + // Live pipelines re-deliver the preroll buffer as the first sample once + // PLAYING; consuming it here just keeps the sink from holding it. + GstSample* sample = gst_app_sink_pull_preroll(sink); + if (sample != nullptr) { + gst_sample_unref(sample); + } + (void)user_data; + return GST_FLOW_OK; +} + +void handle_appsink_eos(GstAppSink* sink, gpointer user_data) { + auto* handle = static_cast(user_data); + (void)sink; + handle->eos.store(true, std::memory_order_release); + handle->frame_ready.notify_all(); +} + +} // namespace + +extern "C" { + +__attribute__((visibility("default"))) +const char* rf_jetson_tensor_bridge_version() { + // v8: appsink callbacks hand off NVMM samples without doing EGL/CUDA work; + // grab() converts on the consumer thread, avoiding a JP6.2 first-frame + // deadlock inside nvv4l2decoder's streaming callback. + // v7: per-fd EGL/CUDA registration cache (map/register/texture objects + // survive across frames; the per-frame global-lock sequence is paid only + // on pool warmup) + egl_cache_hits/egl_cache_misses appended to + // RfBridgeStats (ABI change โ€” python mirror must match). + // v6: lossless (file) handoff mode โ€” rf_jetson_pipeline_create() gained + // the lossless_handoff parameter; bounded blocking FIFO replaces the + // latest-wins slot for non-live sources so every file frame is served, + // live behavior unchanged. + // v5: per-phase conversion timing (egl_map/cuda_register/texture_create/ + // kernel_launch/sync/cleanup, total+max ns each) and unique_buffer_fds + // appended to RfBridgeStats. + // v4: streaming-thread conversion + tensor handoff (jetson-utils consume + // model), direct NV12 path, frames_dropped_by_consumer added to + // RfBridgeStats. + return "8"; +} + +__attribute__((visibility("default"))) +RfJetsonPipeline* rf_jetson_pipeline_create( + const char* pipeline_description, + int device_id, + int lossless_handoff, + char* error, + size_t error_capacity) { + if (pipeline_description == nullptr || pipeline_description[0] == '\0') { + write_error(error, error_capacity, "GStreamer pipeline is empty"); + return nullptr; + } + std::call_once(g_gstreamer_once, initialize_gstreamer); + std::call_once(g_nvbufsurface_once, initialize_nvbufsurface); + if (g_nvbufsurface.map_egl_image == nullptr || + g_nvbufsurface.unmap_egl_image == nullptr) { + write_error( + error, + error_capacity, + "NvBufSurface EGL API is unavailable: %s", + g_nvbufsurface.error == nullptr ? "unknown error" : g_nvbufsurface.error); + return nullptr; + } + cudaError_t cuda_status = cudaSetDevice(device_id); + if (cuda_status == cudaSuccess) { + cuda_status = cudaFree(nullptr); + } + if (cuda_status != cudaSuccess) { + write_error( + error, + error_capacity, + "CUDA device %d is unavailable: %s", + device_id, + cudaGetErrorString(cuda_status)); + return nullptr; + } + CUresult driver_status = cuInit(0); + if (driver_status != CUDA_SUCCESS) { + const char* driver_error = nullptr; + cuGetErrorString(driver_status, &driver_error); + write_error( + error, + error_capacity, + "CUDA driver initialization failed: %s", + driver_error == nullptr ? "unknown error" : driver_error); + return nullptr; + } + + GError* parse_error = nullptr; + GstElement* pipeline = gst_parse_launch(pipeline_description, &parse_error); + if (pipeline == nullptr || parse_error != nullptr) { + write_error( + error, + error_capacity, + "GStreamer pipeline parse failed: %s", + parse_error == nullptr ? "unknown error" : parse_error->message); + if (parse_error != nullptr) { + g_error_free(parse_error); + } + if (pipeline != nullptr) { + gst_object_unref(pipeline); + } + return nullptr; + } + if (!GST_IS_BIN(pipeline)) { + write_error(error, error_capacity, "GStreamer pipeline is not a bin"); + gst_object_unref(pipeline); + return nullptr; + } + GstElement* sink_element = gst_bin_get_by_name(GST_BIN(pipeline), kSinkName); + if (sink_element == nullptr || !GST_IS_APP_SINK(sink_element)) { + write_error( + error, + error_capacity, + "GStreamer pipeline requires appsink name=%s", + kSinkName); + if (sink_element != nullptr) { + gst_object_unref(sink_element); + } + gst_object_unref(pipeline); + return nullptr; + } + + auto* handle = new (std::nothrow) RfJetsonPipeline(); + if (handle == nullptr) { + write_error(error, error_capacity, "Could not allocate pipeline state"); + gst_object_unref(sink_element); + gst_object_unref(pipeline); + return nullptr; + } + handle->pipeline = pipeline; + handle->sink = GST_APP_SINK(sink_element); + handle->device_id = device_id; + handle->lossless_handoff = lossless_handoff != 0; + handle->handoff_capacity = + handle->lossless_handoff ? kLosslessHandoffCapacity : 1; + handle->tensor_pool = + std::make_shared(device_id, kJetsonTensorPoolBuffers); + if (handle->tensor_pool == nullptr) { + write_error(error, error_capacity, "Could not allocate tensor buffer pool"); + gst_object_unref(sink_element); + gst_object_unref(pipeline); + delete handle; + return nullptr; + } + cuda_status = cudaStreamCreateWithFlags(&handle->stream, cudaStreamNonBlocking); + if (cuda_status != cudaSuccess) { + write_error( + error, + error_capacity, + "CUDA stream creation failed: %s", + cudaGetErrorString(cuda_status)); + gst_object_unref(sink_element); + gst_object_unref(pipeline); + delete handle; + return nullptr; + } + GstElement* srtp_caps_element = + gst_bin_get_by_name(GST_BIN(pipeline), kSrtpCapsElementName); + if (srtp_caps_element != nullptr) { + GstPad* source_pad = + gst_element_get_static_pad(srtp_caps_element, "src"); + const gulong probe_id = source_pad == nullptr + ? 0 + : gst_pad_add_probe( + source_pad, + GST_PAD_PROBE_TYPE_EVENT_DOWNSTREAM, + configure_sdes_srtp_caps, + handle, + nullptr); + gst_object_unref(srtp_caps_element); + if (probe_id == 0) { + write_error(error, error_capacity, "Could not configure SRTP caps probe"); + if (source_pad != nullptr) { + gst_object_unref(source_pad); + } + cudaStreamDestroy(handle->stream); + gst_object_unref(sink_element); + gst_object_unref(pipeline); + delete handle; + return nullptr; + } + handle->srtp_probe_pad = source_pad; + handle->srtp_probe_id = probe_id; + } + // Remove samples from appsink on GStreamer's streaming thread and hand + // them to grab() for consumer-thread conversion. Live sources stay + // latest-wins, while files backpressure at the bounded FIFO. Must be + // installed before the PLAYING transition. + GstAppSinkCallbacks sink_callbacks{}; + sink_callbacks.eos = handle_appsink_eos; + sink_callbacks.new_preroll = handle_new_preroll; + sink_callbacks.new_sample = handle_new_sample; + gst_app_sink_set_callbacks(handle->sink, &sink_callbacks, handle, nullptr); + const GstStateChangeReturn state_status = + gst_element_set_state(pipeline, GST_STATE_PLAYING); + if (state_status == GST_STATE_CHANGE_FAILURE) { + write_error(error, error_capacity, "GStreamer could not enter PLAYING state"); + // A failed PLAYING transition can leave elements in READY/PAUSED with + // running task threads; GStreamer refuses to dispose a non-NULL + // pipeline, so reset it before dropping the reference. + gst_element_set_state(pipeline, GST_STATE_NULL); + if (handle->srtp_probe_pad != nullptr) { + gst_pad_remove_probe(handle->srtp_probe_pad, handle->srtp_probe_id); + gst_object_unref(handle->srtp_probe_pad); + } + cudaStreamDestroy(handle->stream); + gst_object_unref(sink_element); + gst_object_unref(pipeline); + delete handle; + return nullptr; + } + return handle; +} + +__attribute__((visibility("default"))) +int rf_jetson_pipeline_grab( + RfJetsonPipeline* handle, + uint64_t timeout_ns, + char* error, + size_t error_capacity) { + if (handle == nullptr) { + write_error(error, error_capacity, "Pipeline handle is null"); + return -1; + } + if (handle->interrupted.load(std::memory_order_acquire)) { + return 0; + } + GstSample* sample = nullptr; + { + std::unique_lock lock(handle->mutex); + const auto frame_or_terminal = [handle]() { + return !handle->ready_samples.empty() || + !handle->ready_tensors.empty() || handle->conversion_failed || + handle->interrupted.load(std::memory_order_acquire) || + handle->eos.load(std::memory_order_acquire); + }; + if (!frame_or_terminal()) { + handle->frame_ready.wait_for( + lock, std::chrono::nanoseconds(timeout_ns), frame_or_terminal); + } + if (handle->interrupted.load(std::memory_order_acquire)) { + return 0; + } + if (handle->conversion_failed) { + write_error(error, error_capacity, "%s", handle->conversion_error); + handle->conversion_failed = false; + return -1; + } + if (!handle->ready_tensors.empty()) { + // In lossless mode queued frames are served BEFORE EOS is + // reported, so the tail of a video file is never lost. + return 1; + } + if (!handle->ready_samples.empty()) { + sample = handle->ready_samples.front(); + handle->ready_samples.pop_front(); + } + } + if (sample != nullptr) { + std::lock_guard conversion_lock(handle->conversion_mutex); + if (handle->interrupted.load(std::memory_order_acquire)) { + gst_sample_unref(sample); + handle->handoff_space.notify_all(); + return 0; + } + char conversion_error[1024] = {0}; + RfFrameInfo frame_info{}; + RfEglDiagnostics diagnostics{}; + RfPhaseTimings timings{}; + RfTensorContext* tensor = convert_sample_to_tensor( + handle, + sample, + &frame_info, + &diagnostics, + &timings, + conversion_error, + sizeof(conversion_error)); + gst_sample_unref(sample); + { + std::lock_guard lock(handle->mutex); + handle->stats.descriptor_maps += 1; + handle->stats.last_nvbuf_memory_type = diagnostics.memory_type; + handle->stats.last_egl_frame_type = diagnostics.frame_type; + handle->stats.last_egl_color_format = diagnostics.color_format; + const auto accumulate_phase = + [](uint64_t* total, uint64_t* max_value, uint64_t sample_ns) { + *total += sample_ns; + if (sample_ns > *max_value) { + *max_value = sample_ns; + } + }; + accumulate_phase( + &handle->stats.egl_map_ns, + &handle->stats.egl_map_max_ns, + timings.egl_map_ns); + accumulate_phase( + &handle->stats.cuda_register_ns, + &handle->stats.cuda_register_max_ns, + timings.cuda_register_ns); + accumulate_phase( + &handle->stats.texture_create_ns, + &handle->stats.texture_create_max_ns, + timings.texture_create_ns); + accumulate_phase( + &handle->stats.kernel_launch_ns, + &handle->stats.kernel_launch_max_ns, + timings.kernel_launch_ns); + accumulate_phase( + &handle->stats.sync_ns, + &handle->stats.sync_max_ns, + timings.sync_ns); + accumulate_phase( + &handle->stats.cleanup_ns, + &handle->stats.cleanup_max_ns, + timings.cleanup_ns); + if (timings.buffer_fd_valid && + handle->seen_buffer_fds.insert(timings.buffer_fd).second) { + handle->stats.unique_buffer_fds = + static_cast(handle->seen_buffer_fds.size()); + } + handle->stats.egl_cache_hits += timings.egl_cache_hit; + handle->stats.egl_cache_misses += timings.egl_cache_miss; + if (tensor != nullptr) { + handle->ready_tensors.push_back(tensor); + handle->last_frame_info = frame_info; + handle->frame_info_valid = true; + handle->stats.frames += 1; + handle->stats.conversion_kernels += 1; + handle->stats.nvmm_frames += 1; + } + } + if (tensor == nullptr) { + write_error(error, error_capacity, "%s", conversion_error); + handle->handoff_space.notify_all(); + return -1; + } + return 1; + } + // Nothing ready. Check the bus BEFORE the EOS flag: a pipeline that died + // during startup (RTSP connect/auth failure, autoplug failure) never + // delivers a frame โ€” the real error would otherwise be misclassified as a + // silent end-of-stream. A genuine EOS posts no ERROR, so it still + // returns 0. + if (read_bus_error(handle, error, error_capacity)) { + return -1; + } + if (handle->eos.load(std::memory_order_acquire) || + gst_app_sink_is_eos(handle->sink)) { + return 0; + } + // No frame, no error, no EOS: the finite timeout expired while the + // stream is still live. + return 2; +} + +__attribute__((visibility("default"))) +int rf_jetson_pipeline_get_frame_info( + RfJetsonPipeline* handle, + RfFrameInfo* info, + char* error, + size_t error_capacity) { + if (handle == nullptr || info == nullptr) { + write_error(error, error_capacity, "Frame-info arguments are invalid"); + return -1; + } + std::lock_guard lock(handle->mutex); + if (!handle->frame_info_valid) { + write_error(error, error_capacity, "Frame caps do not contain dimensions"); + return -1; + } + *info = handle->last_frame_info; + gint64 duration = GST_CLOCK_TIME_NONE; + if (gst_element_query_duration(handle->pipeline, GST_FORMAT_TIME, &duration)) { + info->duration_ns = duration; + } else { + info->duration_ns = 0; + } + return 1; +} + +__attribute__((visibility("default"))) +int rf_jetson_pipeline_has_factory( + RfJetsonPipeline* handle, + const char* factory_name) { + if (handle == nullptr) { + return 0; + } + std::lock_guard lock(handle->mutex); + return pipeline_has_factory(handle, factory_name) ? 1 : 0; +} + +__attribute__((visibility("default"))) +DLManagedTensor* rf_jetson_pipeline_retrieve( + RfJetsonPipeline* handle, + char* error, + size_t error_capacity) { + if (handle == nullptr) { + write_error(error, error_capacity, "Pipeline handle is null"); + return nullptr; + } + RfTensorContext* tensor = nullptr; + { + std::lock_guard lock(handle->mutex); + if (handle->ready_tensors.empty()) { + write_error(error, error_capacity, "No grabbed frame is available"); + return nullptr; + } + // grab() already converted this sample on the consumer thread; hand + // the finished tensor to Python. The DLPack deleter binds the device + // when the consumer eventually drops the tensor. + tensor = handle->ready_tensors.front(); + handle->ready_tensors.pop_front(); + } + // Wake a lossless-mode callback parked on a full handoff queue. + handle->handoff_space.notify_all(); + return &tensor->managed; +} + +__attribute__((visibility("default"))) +int rf_jetson_pipeline_get_stats( + RfJetsonPipeline* handle, + RfBridgeStats* stats) { + if (handle == nullptr || stats == nullptr) { + return -1; + } + std::lock_guard lock(handle->mutex); + *stats = handle->stats; + return 1; +} + +__attribute__((visibility("default"))) +void rf_jetson_dlpack_delete(DLManagedTensor* tensor) { + if (tensor != nullptr && tensor->deleter != nullptr) { + tensor->deleter(tensor); + } +} + +__attribute__((visibility("default"))) +int rf_jetson_pipeline_interrupt(RfJetsonPipeline* handle) { + if (handle == nullptr) { + return -1; + } + handle->interrupted.store(true, std::memory_order_release); + // Wake a grab() parked on the handoff queue and a lossless-mode callback + // parked on a full queue so interrupt is prompt. + handle->frame_ready.notify_all(); + handle->handoff_space.notify_all(); + if (handle->sink != nullptr) { + gst_app_sink_set_drop(handle->sink, TRUE); + GstSample* queued_sample = nullptr; + while ((queued_sample = gst_app_sink_try_pull_sample(handle->sink, 0)) != + nullptr) { + gst_sample_unref(queued_sample); + } + } + if (handle->pipeline != nullptr) { + std::lock_guard conversion_lock(handle->conversion_mutex); + gst_element_set_state(handle->pipeline, GST_STATE_NULL); + } + return 1; +} + +__attribute__((visibility("default"))) +void rf_jetson_pipeline_release(RfJetsonPipeline* handle) { + if (handle == nullptr) { + return; + } + rf_jetson_pipeline_interrupt(handle); + if (handle->srtp_probe_pad != nullptr) { + gst_pad_remove_probe(handle->srtp_probe_pad, handle->srtp_probe_id); + gst_object_unref(handle->srtp_probe_pad); + handle->srtp_probe_pad = nullptr; + handle->srtp_probe_id = 0; + } + { + std::lock_guard lock(handle->mutex); + while (!handle->ready_samples.empty()) { + gst_sample_unref(handle->ready_samples.front()); + handle->ready_samples.pop_front(); + } + while (!handle->ready_tensors.empty()) { + // interrupt() already reached GST_STATE_NULL, so the streaming + // thread is joined and no further callback can repopulate the + // queue; return the uncollected frames' buffers to the pool. + delete_managed_tensor(&handle->ready_tensors.front()->managed); + handle->ready_tensors.pop_front(); + } + if (handle->sink != nullptr) { + gst_app_sink_set_drop(handle->sink, TRUE); + GstSample* queued_sample = nullptr; + while ((queued_sample = + gst_app_sink_try_pull_sample(handle->sink, 0)) != + nullptr) { + gst_sample_unref(queued_sample); + } + } + if (handle->pipeline != nullptr) { + gst_element_set_state(handle->pipeline, GST_STATE_NULL); + } + if (!handle->egl_cache.empty()) { + // The streaming thread is joined (GST_STATE_NULL above), so the + // cache is exclusively ours. Free the CUDA-side registrations; + // the surfaces' EGL mappings die with the decoder pool. + int previous_device = -1; + cudaGetDevice(&previous_device); + cudaSetDevice(handle->device_id); + for (auto& entry : handle->egl_cache) { + destroy_egl_cache_entry_resources(&entry); + } + handle->egl_cache.clear(); + if (previous_device >= 0 && previous_device != handle->device_id) { + cudaSetDevice(previous_device); + } + } + if (handle->stream != nullptr) { + cudaStreamDestroy(handle->stream); + } + if (handle->sink != nullptr) { + gst_object_unref(handle->sink); + } + if (handle->pipeline != nullptr) { + gst_object_unref(handle->pipeline); + } + // Null the element pointers before the handle is freed so a + // contract-violating late interrupt() dereferences null instead of + // freed objects. The Python wrapper serializes interrupt()/close(). + handle->stream = nullptr; + handle->sink = nullptr; + handle->pipeline = nullptr; + } + delete handle; +} + +} // extern "C" diff --git a/docker/patches/flash-attention-sm87.patch b/docker/patches/flash-attention-sm87.patch new file mode 100644 index 0000000000..caef91e7c6 --- /dev/null +++ b/docker/patches/flash-attention-sm87.patch @@ -0,0 +1,15 @@ +diff --git a/setup.py b/setup.py +index f9f95fa..898c7fb 100644 +--- a/setup.py ++++ b/setup.py +@@ -111,6 +111,9 @@ def add_cuda_gencodes(cc_flag, archs, bare_metal_version): + if "80" in archs: + cc_flag += ["-gencode", "arch=compute_80,code=sm_80"] +- ++ ++ if "87" in archs: ++ cc_flag += ["-gencode", "arch=compute_87,code=sm_87"] ++ + # Hopper 9.0 needs >= 11.8 + if bare_metal_version >= Version("11.8") and "90" in archs: + cc_flag += ["-gencode", "arch=compute_90,code=sm_90"] diff --git a/docker/scripts/build_ffmpeg.sh b/docker/scripts/build_ffmpeg.sh new file mode 100644 index 0000000000..bf03bcd445 --- /dev/null +++ b/docker/scripts/build_ffmpeg.sh @@ -0,0 +1,59 @@ +#!/bin/sh + +set -eux + +FFMPEG_VERSION="${FFMPEG_VERSION:-7.1.3}" +FFMPEG_SHA256="${FFMPEG_SHA256:?FFMPEG_SHA256 is required}" +FFMPEG_PREFIX="${FFMPEG_PREFIX:-/opt/ffmpeg}" +FFMPEG_SOURCE_DIR="${FFMPEG_SOURCE_DIR:-/tmp/ffmpeg-src}" + +mkdir -p "${FFMPEG_SOURCE_DIR}" +ffmpeg_archive="/tmp/ffmpeg-${FFMPEG_VERSION}.tar.xz" +curl \ + --fail \ + --location \ + --retry 5 \ + --retry-all-errors \ + --retry-delay 2 \ + "https://ffmpeg.org/releases/ffmpeg-${FFMPEG_VERSION}.tar.xz" \ + -o "${ffmpeg_archive}" +echo "${FFMPEG_SHA256} ${ffmpeg_archive}" | sha256sum -c - +tar -xJf "${ffmpeg_archive}" \ + --strip-components=1 \ + -C "${FFMPEG_SOURCE_DIR}" + +cd "${FFMPEG_SOURCE_DIR}" +./configure \ + --prefix="${FFMPEG_PREFIX}" \ + --enable-shared \ + --disable-static \ + --enable-pic \ + --disable-autodetect \ + --enable-bzlib \ + --enable-lzma \ + --enable-openssl \ + --enable-pthreads \ + --enable-zlib \ + --disable-debug \ + --disable-doc \ + --disable-ffplay \ + --extra-cflags=-O3 \ + --extra-ldflags="-Wl,-rpath,${FFMPEG_PREFIX}/lib" + +make -j"$(nproc)" +make install + +mkdir -p "${FFMPEG_PREFIX}/share/licenses/ffmpeg" +cp COPYING.LGPLv2.1 "${FFMPEG_PREFIX}/share/licenses/ffmpeg/" +find "${FFMPEG_PREFIX}/bin" -type f -exec strip --strip-unneeded {} + +find "${FFMPEG_PREFIX}/lib" -type f -name '*.so*' -exec strip --strip-unneeded {} + + +runtime_prefix="${FFMPEG_PREFIX}-runtime" +cp -a "${FFMPEG_PREFIX}" "${runtime_prefix}" +rm -rf \ + "${runtime_prefix}/include" \ + "${runtime_prefix}/lib/pkgconfig" \ + "${runtime_prefix}/share/ffmpeg/examples" +find "${runtime_prefix}" -type f \( -name '*.a' -o -name '*.la' \) -delete + +rm -rf "${ffmpeg_archive}" "${FFMPEG_SOURCE_DIR}" diff --git a/docker/scripts/build_gstreamer.sh b/docker/scripts/build_gstreamer.sh new file mode 100644 index 0000000000..2799c5b3f5 --- /dev/null +++ b/docker/scripts/build_gstreamer.sh @@ -0,0 +1,202 @@ +#!/bin/sh + +set -eux + +GSTREAMER_VERSION="${GSTREAMER_VERSION:-1.24.12}" +GSTREAMER_COMMIT="${GSTREAMER_COMMIT:?GSTREAMER_COMMIT is required}" +GSTREAMER_PREFIX="${GSTREAMER_PREFIX:-/opt/gstreamer}" +GSTREAMER_NVCODEC="${GSTREAMER_NVCODEC:-disabled}" +GSTREAMER_NVRTC_COMPUTE_CAPABILITY="${GSTREAMER_NVRTC_COMPUTE_CAPABILITY:-75}" +GSTREAMER_SOURCE_DIR="${GSTREAMER_SOURCE_DIR:-/tmp/gstreamer-src}" +GSTREAMER_BUILD_DIR="${GSTREAMER_BUILD_DIR:-/tmp/gstreamer-build}" +GLIB_NETWORKING_VERSION="${GLIB_NETWORKING_VERSION:-2.72.2}" +GLIB_NETWORKING_SHA256="${GLIB_NETWORKING_SHA256:?GLIB_NETWORKING_SHA256 is required}" +GLIB_NETWORKING_SERIES="${GLIB_NETWORKING_VERSION%.*}" +GLIB_NETWORKING_SOURCE_DIR="${GLIB_NETWORKING_SOURCE_DIR:-/tmp/glib-networking-src}" +GLIB_NETWORKING_BUILD_DIR="${GLIB_NETWORKING_BUILD_DIR:-/tmp/glib-networking-build}" + +case "${GSTREAMER_NVCODEC}" in + enabled|disabled) ;; + *) + echo "GSTREAMER_NVCODEC must be enabled or disabled" >&2 + exit 2 + ;; +esac + +git clone \ + --branch "${GSTREAMER_VERSION}" \ + --depth 1 \ + https://gitlab.freedesktop.org/gstreamer/gstreamer.git \ + "${GSTREAMER_SOURCE_DIR}" +test "$(git -C "${GSTREAMER_SOURCE_DIR}" rev-parse HEAD)" = \ + "${GSTREAMER_COMMIT}" + +if [ "${GSTREAMER_NVCODEC}" = "enabled" ]; then + nvrtc_source="${GSTREAMER_SOURCE_DIR}/subprojects/gst-plugins-bad/gst-libs/gst/cuda/gstcudanvrtc.cpp" + test "$(grep -c -- '--gpu-architecture=compute_30' "${nvrtc_source}")" -eq 1 + test "$(grep -c -- '--gpu-architecture=compute_52' "${nvrtc_source}")" -eq 1 + sed -i \ + -e "s/--gpu-architecture=compute_30/--gpu-architecture=compute_${GSTREAMER_NVRTC_COMPUTE_CAPABILITY}/" \ + -e "s/--gpu-architecture=compute_52/--gpu-architecture=compute_${GSTREAMER_NVRTC_COMPUTE_CAPABILITY}/" \ + "${nvrtc_source}" +fi + +meson setup "${GSTREAMER_BUILD_DIR}" "${GSTREAMER_SOURCE_DIR}" \ + --buildtype=release \ + --prefix="${GSTREAMER_PREFIX}" \ + --libdir=lib \ + --strip \ + --wrap-mode=default \ + --force-fallback-for=libnice \ + -Dauto_features=disabled \ + -Db_ndebug=true \ + -Dbase=enabled \ + -Dgood=enabled \ + -Dbad=enabled \ + -Dugly=disabled \ + -Dlibav=enabled \ + -Drtsp_server=enabled \ + -Ddevtools=disabled \ + -Dges=disabled \ + -Drs=disabled \ + -Dvaapi=disabled \ + -Dpython=disabled \ + -Dsharp=disabled \ + -Dlibnice=enabled \ + -Dgst-full=disabled \ + -Dgst-examples=disabled \ + -Dgpl=disabled \ + -Dtests=disabled \ + -Dtools=enabled \ + -Dexamples=disabled \ + -Dintrospection=enabled \ + -Dnls=disabled \ + -Dorc=enabled \ + -Dorc-source=subproject \ + -Dqt5=disabled \ + -Dqt6=disabled \ + -Dwebrtc=enabled \ + -Dgst-plugins-base:app=enabled \ + -Dgst-plugins-base:audioconvert=enabled \ + -Dgst-plugins-base:audioresample=enabled \ + -Dgst-plugins-base:encoding=enabled \ + -Dgst-plugins-base:gio=enabled \ + -Dgst-plugins-base:gio-typefinder=enabled \ + -Dgst-plugins-base:opus=enabled \ + -Dgst-plugins-base:playback=enabled \ + -Dgst-plugins-base:rawparse=enabled \ + -Dgst-plugins-base:tcp=enabled \ + -Dgst-plugins-base:typefind=enabled \ + -Dgst-plugins-base:videoconvertscale=enabled \ + -Dgst-plugins-base:videorate=enabled \ + -Dgst-plugins-base:videotestsrc=enabled \ + -Dgst-plugins-good:autodetect=enabled \ + -Dgst-plugins-good:avi=enabled \ + -Dgst-plugins-good:debugutils=enabled \ + -Dgst-plugins-good:flv=enabled \ + -Dgst-plugins-good:isomp4=enabled \ + -Dgst-plugins-good:jpeg=enabled \ + -Dgst-plugins-good:matroska=enabled \ + -Dgst-plugins-good:multifile=enabled \ + -Dgst-plugins-good:multipart=enabled \ + -Dgst-plugins-good:rtp=enabled \ + -Dgst-plugins-good:rtpmanager=enabled \ + -Dgst-plugins-good:rtsp=enabled \ + -Dgst-plugins-good:udp=enabled \ + -Dgst-plugins-good:v4l2=enabled \ + -Dgst-plugins-good:vpx=enabled \ + -Dgst-plugins-good:wavparse=enabled \ + -Dgst-plugins-bad:curl=enabled \ + -Dgst-plugins-bad:bayer=enabled \ + -Dgst-plugins-bad:codectimestamper=enabled \ + -Dgst-plugins-bad:dtls=enabled \ + -Dgst-plugins-bad:jpegformat=enabled \ + -Dgst-plugins-bad:mpegtsdemux=enabled \ + -Dgst-plugins-bad:nvcodec="${GSTREAMER_NVCODEC}" \ + -Dgst-plugins-bad:rtmp2=enabled \ + -Dgst-plugins-bad:rtp=enabled \ + -Dgst-plugins-bad:sctp=enabled \ + -Dgst-plugins-bad:sdp=enabled \ + -Dgst-plugins-bad:srtp=enabled \ + -Dgst-plugins-bad:videoparsers=enabled \ + -Dgst-plugins-bad:webrtc=enabled \ + -Dgst-rtsp-server:rtspclientsink=enabled \ + -Dlibnice:crypto-library=openssl \ + -Dlibnice:gupnp=disabled + +meson compile -C "${GSTREAMER_BUILD_DIR}" +meson install -C "${GSTREAMER_BUILD_DIR}" + +mkdir -p "${GLIB_NETWORKING_SOURCE_DIR}" +glib_networking_archive="/tmp/glib-networking-${GLIB_NETWORKING_VERSION}.tar.xz" +curl \ + --fail \ + --location \ + --retry 5 \ + --retry-all-errors \ + --retry-delay 2 \ + "https://download.gnome.org/sources/glib-networking/${GLIB_NETWORKING_SERIES}/glib-networking-${GLIB_NETWORKING_VERSION}.tar.xz" \ + -o "${glib_networking_archive}" +echo "${GLIB_NETWORKING_SHA256} ${glib_networking_archive}" | sha256sum -c - +tar -xJf "${glib_networking_archive}" \ + --strip-components=1 \ + -C "${GLIB_NETWORKING_SOURCE_DIR}" + +meson setup "${GLIB_NETWORKING_BUILD_DIR}" "${GLIB_NETWORKING_SOURCE_DIR}" \ + --buildtype=release \ + --prefix="${GSTREAMER_PREFIX}" \ + --libdir=lib \ + --strip \ + -Db_ndebug=true \ + -Dgnome_proxy=disabled \ + -Dgnutls=enabled \ + -Dinstalled_tests=false \ + -Dlibproxy=disabled \ + -Dopenssl=disabled \ + -Dstatic_modules=false +meson compile -C "${GLIB_NETWORKING_BUILD_DIR}" +meson install -C "${GLIB_NETWORKING_BUILD_DIR}" +gio-querymodules "${GSTREAMER_PREFIX}/lib/gio/modules" + +license_dir="${GSTREAMER_PREFIX}/share/licenses/gstreamer" +mkdir -p "${license_dir}" +cp "${GSTREAMER_SOURCE_DIR}/LICENSE" "${license_dir}/LICENSE" +for project in \ + gst-libav \ + gst-plugins-bad \ + gst-plugins-base \ + gst-plugins-good \ + gst-rtsp-server \ + gstreamer +do + cp "${GSTREAMER_SOURCE_DIR}/subprojects/${project}/COPYING" \ + "${license_dir}/COPYING.${project}" +done +cp "${GSTREAMER_SOURCE_DIR}/subprojects/libnice/COPYING" \ + "${license_dir}/COPYING.libnice" +cp "${GSTREAMER_SOURCE_DIR}/subprojects/orc/COPYING" \ + "${license_dir}/COPYING.orc" +cp "${GSTREAMER_SOURCE_DIR}/subprojects/gst-plugins-bad/ext/sctp/usrsctp/LICENSE.md" \ + "${license_dir}/LICENSE.usrsctp" +cp "${GLIB_NETWORKING_SOURCE_DIR}/COPYING" \ + "${license_dir}/COPYING.glib-networking" +cp "${GLIB_NETWORKING_SOURCE_DIR}/LICENSE_EXCEPTION" \ + "${license_dir}/LICENSE_EXCEPTION.glib-networking" + +runtime_prefix="${GSTREAMER_PREFIX}-runtime" +cp -a "${GSTREAMER_PREFIX}" "${runtime_prefix}" +rm -rf \ + "${runtime_prefix}/include" \ + "${runtime_prefix}/lib/pkgconfig" \ + "${runtime_prefix}/share/aclocal" \ + "${runtime_prefix}/share/gir-1.0" \ + "${runtime_prefix}/share/gtk-doc" \ + "${runtime_prefix}/share/man" +find "${runtime_prefix}" -type f \( -name '*.a' -o -name '*.la' \) -delete + +rm -rf \ + "${GSTREAMER_SOURCE_DIR}" \ + "${GSTREAMER_BUILD_DIR}" \ + "${glib_networking_archive}" \ + "${GLIB_NETWORKING_SOURCE_DIR}" \ + "${GLIB_NETWORKING_BUILD_DIR}" diff --git a/docker/scripts/build_gstreamer_cuda_tensor_bridge.sh b/docker/scripts/build_gstreamer_cuda_tensor_bridge.sh new file mode 100644 index 0000000000..6fedfad36c --- /dev/null +++ b/docker/scripts/build_gstreamer_cuda_tensor_bridge.sh @@ -0,0 +1,40 @@ +#!/bin/sh + +set -eu + +SOURCE_DIR="${GSTREAMER_CUDA_TENSOR_BRIDGE_SOURCE_DIR:-/src/gstreamer_cuda_tensor_bridge}" +OUTPUT_DIR="${GSTREAMER_CUDA_TENSOR_BRIDGE_OUTPUT_DIR:-/opt/roboflow/lib}" +CUDA_INCLUDE_DIR="${CUDA_INCLUDE_DIR:-$( + find /usr/local/cuda/targets -path '*/include/cuda.h' -print -quit | + sed 's|/cuda.h$||' +)}" +CUDA_STUB_LIBRARY_DIR="${CUDA_STUB_LIBRARY_DIR:-$( + find /usr/local/cuda/targets -path '*/lib/stubs/libcuda.so' -print -quit | + sed 's|/libcuda.so$||' +)}" + +test -f "${SOURCE_DIR}/gstreamer_cuda_tensor_bridge.cpp" +test -f "${CUDA_INCLUDE_DIR}/cuda.h" +test -f "${CUDA_STUB_LIBRARY_DIR}/libcuda.so" +pkg-config --exists gstreamer-app-1.0 gstreamer-cuda-1.0 gstreamer-video-1.0 +mkdir -p "${OUTPUT_DIR}" + +c++ \ + -std=c++17 \ + -O3 \ + -shared \ + -fPIC \ + -fvisibility=hidden \ + -I"${CUDA_INCLUDE_DIR}" \ + $(pkg-config --cflags gstreamer-app-1.0 gstreamer-cuda-1.0 gstreamer-video-1.0) \ + "${SOURCE_DIR}/gstreamer_cuda_tensor_bridge.cpp" \ + -o "${OUTPUT_DIR}/libroboflow_gstreamer_cuda_tensor.so.1" \ + $(pkg-config --libs gstreamer-app-1.0 gstreamer-cuda-1.0 gstreamer-video-1.0) \ + -L"${CUDA_STUB_LIBRARY_DIR}" \ + -Wl,-rpath,/opt/gstreamer/lib:/usr/local/cuda/lib64 \ + -lcuda + +strip --strip-unneeded \ + "${OUTPUT_DIR}/libroboflow_gstreamer_cuda_tensor.so.1" +ln -sf libroboflow_gstreamer_cuda_tensor.so.1 \ + "${OUTPUT_DIR}/libroboflow_gstreamer_cuda_tensor.so" diff --git a/docker/scripts/build_jetson_6_2_image.sh b/docker/scripts/build_jetson_6_2_image.sh new file mode 100755 index 0000000000..f462700356 --- /dev/null +++ b/docker/scripts/build_jetson_6_2_image.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash + +set -euo pipefail + +script_directory="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +repository_root="$(cd -- "${script_directory}/../.." && pwd)" + +"${script_directory}/fetch_jetson_6_2_tensorrt.sh" + +cd -- "${repository_root}" +exec docker buildx build \ + --platform linux/arm64 \ + --file docker/dockerfiles/Dockerfile.onnx.jetson.6.2.0 \ + "$@" \ + . diff --git a/docker/scripts/build_jetson_tensor_bridge.sh b/docker/scripts/build_jetson_tensor_bridge.sh new file mode 100644 index 0000000000..9cf8dfa7b2 --- /dev/null +++ b/docker/scripts/build_jetson_tensor_bridge.sh @@ -0,0 +1,60 @@ +#!/bin/sh + +set -eu + +SOURCE_DIR="${JETSON_TENSOR_BRIDGE_SOURCE_DIR:-/src/jetson_tensor_bridge}" +OUTPUT_DIR="${JETSON_TENSOR_BRIDGE_OUTPUT_DIR:-/opt/roboflow/lib}" +CUDA_ARCHITECTURES="${JETSON_TENSOR_BRIDGE_CUDA_ARCHITECTURES:-87}" +NVBUF_SURFACE_INCLUDE_DIR="${NVBUF_SURFACE_INCLUDE_DIR:-/usr/src/jetson_multimedia_api/include}" +CUDA_STUB_LIBRARY_DIR="${CUDA_STUB_LIBRARY_DIR:-$( + find /usr/local/cuda/targets -path '*/lib/stubs/libcuda.so' -print -quit | + sed 's|/libcuda.so$||' +)}" + +test -f "${SOURCE_DIR}/jetson_tensor_bridge.cu" +test -f "${NVBUF_SURFACE_INCLUDE_DIR}/nvbufsurface.h" +test -f "${CUDA_STUB_LIBRARY_DIR}/libcuda.so" +pkg-config --exists gstreamer-app-1.0 + +mkdir -p "${OUTPUT_DIR}" + +gencode_flags="" +old_ifs="${IFS}" +IFS=';, ' +for architecture in ${CUDA_ARCHITECTURES}; do + test -n "${architecture}" + gencode_flags="${gencode_flags} -gencode arch=compute_${architecture},code=sm_${architecture}" +done +IFS="${old_ifs}" + +# shellcheck disable=SC2086 +nvcc \ + -std=c++17 \ + -O3 \ + --shared \ + --compiler-options=-fPIC,-fvisibility=hidden,-pthread \ + ${gencode_flags} \ + -I"${NVBUF_SURFACE_INCLUDE_DIR}" \ + $(pkg-config --cflags-only-I gstreamer-app-1.0) \ + "${SOURCE_DIR}/jetson_tensor_bridge.cu" \ + -o "${OUTPUT_DIR}/libroboflow_jetson_tensor.so.1" \ + $(pkg-config --libs-only-L --libs-only-l gstreamer-app-1.0) \ + -L"${CUDA_STUB_LIBRARY_DIR}" \ + -Xlinker=-rpath \ + -Xlinker=/opt/gstreamer/lib:/usr/local/cuda/lib64 \ + -lcuda \ + -lcudart \ + -ldl + +old_ifs="${IFS}" +IFS=';, ' +for architecture in ${CUDA_ARCHITECTURES}; do + cuobjdump --list-elf \ + "${OUTPUT_DIR}/libroboflow_jetson_tensor.so.1" | + grep -q "\.sm_${architecture}\.cubin$" +done +IFS="${old_ifs}" + +strip --strip-unneeded "${OUTPUT_DIR}/libroboflow_jetson_tensor.so.1" +ln -sf libroboflow_jetson_tensor.so.1 \ + "${OUTPUT_DIR}/libroboflow_jetson_tensor.so" diff --git a/docker/scripts/build_opencv.sh b/docker/scripts/build_opencv.sh new file mode 100644 index 0000000000..17b2aa22a6 --- /dev/null +++ b/docker/scripts/build_opencv.sh @@ -0,0 +1,187 @@ +#!/bin/sh + +set -eux + +OPENCV_VERSION="${OPENCV_VERSION:-4.13.0}" +OPENCV_SHA256="${OPENCV_SHA256:?OPENCV_SHA256 is required}" +OPENCV_CONTRIB_SHA256="${OPENCV_CONTRIB_SHA256:?OPENCV_CONTRIB_SHA256 is required}" +OPENCV_PREFIX="${OPENCV_PREFIX:-/opt/opencv}" +OPENCV_PYTHON_EXECUTABLE="${OPENCV_PYTHON_EXECUTABLE:-python3}" +OPENCV_PYTHON3_LIMITED_API="${OPENCV_PYTHON3_LIMITED_API:-OFF}" +OPENCV_BUILD_PYTHON3="${OPENCV_BUILD_PYTHON3:-ON}" +OPENCV_WITH_CUDA="${OPENCV_WITH_CUDA:-OFF}" +OPENCV_CUDA_ARCH_BIN="${OPENCV_CUDA_ARCH_BIN:-}" +OPENCV_CUDA_ARCH_PTX="${OPENCV_CUDA_ARCH_PTX:-}" +OPENCV_NUMPY_TARGET_VERSION="${OPENCV_NUMPY_TARGET_VERSION:-}" +OPENCV_WITH_CUDNN="${OPENCV_WITH_CUDNN:-OFF}" +OPENCV_WITH_CUBLAS="${OPENCV_WITH_CUBLAS:-OFF}" +OPENCV_WITH_CUFFT="${OPENCV_WITH_CUFFT:-OFF}" +OPENCV_DNN_CUDA="${OPENCV_DNN_CUDA:-OFF}" +OPENCV_SOURCE_DIR="${OPENCV_SOURCE_DIR:-/tmp/opencv-src}" +OPENCV_CONTRIB_SOURCE_DIR="${OPENCV_CONTRIB_SOURCE_DIR:-/tmp/opencv-contrib-src}" +OPENCV_BUILD_DIR="${OPENCV_BUILD_DIR:-/tmp/opencv-build}" +OPENCV_PYTHON_INCLUDE_DIR="$( + "${OPENCV_PYTHON_EXECUTABLE}" -c \ + 'import sysconfig; print(sysconfig.get_path("include"))' +)" +OPENCV_NUMPY_INCLUDE_DIR="$( + "${OPENCV_PYTHON_EXECUTABLE}" -c \ + 'import numpy; print(numpy.get_include())' +)" +OPENCV_C_FLAGS="${CFLAGS:-}" +OPENCV_CXX_FLAGS="${CXXFLAGS:-}" + +# NumPy 2 defaults extension modules to the NumPy 1.x C API for backwards +# compatibility. That produces a cv2 module which NumPy 2 refuses to import. +# Jetson images ship NumPy 2, so explicitly opt the Python bindings into its +# 2.0 API when requested by the image Dockerfile. +if [ -n "${OPENCV_NUMPY_TARGET_VERSION}" ]; then + OPENCV_C_FLAGS="${OPENCV_C_FLAGS} -DNPY_TARGET_VERSION=${OPENCV_NUMPY_TARGET_VERSION}" + OPENCV_CXX_FLAGS="${OPENCV_CXX_FLAGS} -DNPY_TARGET_VERSION=${OPENCV_NUMPY_TARGET_VERSION}" +fi + +case "${OPENCV_WITH_CUDA}" in + ON|OFF) ;; + *) + echo "OPENCV_WITH_CUDA must be ON or OFF" >&2 + exit 2 + ;; +esac + +if [ "${OPENCV_WITH_CUDA}" = "ON" ] && [ -z "${OPENCV_CUDA_ARCH_BIN}" ]; then + echo "OPENCV_CUDA_ARCH_BIN is required when CUDA support is enabled" >&2 + exit 2 +fi + +mkdir -p "${OPENCV_SOURCE_DIR}" "${OPENCV_CONTRIB_SOURCE_DIR}" +opencv_archive="/tmp/opencv-${OPENCV_VERSION}.tar.gz" +opencv_contrib_archive="/tmp/opencv_contrib-${OPENCV_VERSION}.tar.gz" +curl \ + --fail \ + --location \ + --retry 5 \ + --retry-all-errors \ + --retry-delay 2 \ + "https://github.com/opencv/opencv/archive/refs/tags/${OPENCV_VERSION}.tar.gz" \ + -o "${opencv_archive}" +echo "${OPENCV_SHA256} ${opencv_archive}" | sha256sum -c - +tar -xzf "${opencv_archive}" --strip-components=1 -C "${OPENCV_SOURCE_DIR}" +curl \ + --fail \ + --location \ + --retry 5 \ + --retry-all-errors \ + --retry-delay 2 \ + "https://github.com/opencv/opencv_contrib/archive/refs/tags/${OPENCV_VERSION}.tar.gz" \ + -o "${opencv_contrib_archive}" +echo "${OPENCV_CONTRIB_SHA256} ${opencv_contrib_archive}" | sha256sum -c - +tar -xzf "${opencv_contrib_archive}" \ + --strip-components=1 \ + -C "${OPENCV_CONTRIB_SOURCE_DIR}" + +if [ "${OPENCV_WITH_CUDA}" = "ON" ]; then + cuda_release="$( + nvcc --version | + sed -n 's/.*release \([0-9][0-9]*\)\.\([0-9][0-9]*\).*/\1 \2/p' | + head -n 1 + )" + set -- ${cuda_release} + if [ "$#" -ne 2 ]; then + echo "Could not parse the CUDA compiler version" >&2 + exit 2 + fi + if [ "$1" -gt 13 ] || { [ "$1" -eq 13 ] && [ "$2" -ge 2 ]; }; then + cuda_13_2_patch=/tmp/opencv-contrib-cuda-13.2.patch + curl \ + --fail \ + --location \ + --retry 5 \ + --retry-all-errors \ + --retry-delay 2 \ + https://github.com/opencv/opencv_contrib/commit/f2854f4f5e7b67d4e073ea002ae0174d437e2962.patch \ + -o "${cuda_13_2_patch}" + echo "d436681d50c61837f82f4f70b835b298e781491beadee43d11914bbe793c983e ${cuda_13_2_patch}" | + sha256sum -c - + patch -d "${OPENCV_CONTRIB_SOURCE_DIR}" -p1 < "${cuda_13_2_patch}" + rm "${cuda_13_2_patch}" + fi +fi + +cmake \ + -S "${OPENCV_SOURCE_DIR}" \ + -B "${OPENCV_BUILD_DIR}" \ + -GNinja \ + -D CMAKE_BUILD_TYPE=Release \ + -D CMAKE_INSTALL_PREFIX="${OPENCV_PREFIX}" \ + -D 'CMAKE_INSTALL_RPATH=/opt/gstreamer/lib;/opt/ffmpeg/lib' \ + -D OPENCV_EXTRA_MODULES_PATH="${OPENCV_CONTRIB_SOURCE_DIR}/modules" \ + -D OPENCV_PYTHON3_INSTALL_PATH="${OPENCV_PREFIX}/python" \ + -D PYTHON3_LIMITED_API="${OPENCV_PYTHON3_LIMITED_API}" \ + -D PYTHON3_EXECUTABLE="${OPENCV_PYTHON_EXECUTABLE}" \ + -D Python3_EXECUTABLE="${OPENCV_PYTHON_EXECUTABLE}" \ + -D PYTHON3_INCLUDE_DIR="${OPENCV_PYTHON_INCLUDE_DIR}" \ + -D PYTHON3_NUMPY_INCLUDE_DIRS="${OPENCV_NUMPY_INCLUDE_DIR}" \ + -D CMAKE_C_FLAGS="${OPENCV_C_FLAGS}" \ + -D CMAKE_CXX_FLAGS="${OPENCV_CXX_FLAGS}" \ + -D WITH_GSTREAMER=ON \ + -D WITH_FFMPEG=ON \ + -D WITH_LIBV4L=ON \ + -D WITH_CUDA="${OPENCV_WITH_CUDA}" \ + -D WITH_CUDNN="${OPENCV_WITH_CUDNN}" \ + -D WITH_CUBLAS="${OPENCV_WITH_CUBLAS}" \ + -D WITH_CUFFT="${OPENCV_WITH_CUFFT}" \ + -D OPENCV_DNN_CUDA="${OPENCV_DNN_CUDA}" \ + -D CUDA_ARCH_BIN="${OPENCV_CUDA_ARCH_BIN}" \ + -D CUDA_ARCH_PTX="${OPENCV_CUDA_ARCH_PTX}" \ + -D BUILD_opencv_cudacodec=OFF \ + -D BUILD_opencv_cudafeatures2d=OFF \ + -D BUILD_opencv_cudalegacy=OFF \ + -D BUILD_opencv_cudaobjdetect=OFF \ + -D BUILD_opencv_cudaoptflow=OFF \ + -D BUILD_opencv_cudastereo=OFF \ + -D WITH_NVCUVID=OFF \ + -D WITH_NVCUVENC=OFF \ + -D WITH_GTK=OFF \ + -D WITH_QT=OFF \ + -D VIDEOIO_ENABLE_PLUGINS=OFF \ + -D BUILD_SHARED_LIBS=OFF \ + -D BUILD_JAVA=OFF \ + -D BUILD_opencv_apps=OFF \ + -D BUILD_opencv_python3="${OPENCV_BUILD_PYTHON3}" \ + -D BUILD_opencv_videoio=ON \ + -D BUILD_TESTS=OFF \ + -D BUILD_PERF_TESTS=OFF \ + -D BUILD_EXAMPLES=OFF \ + -D BUILD_DOCS=OFF + +cmake --build "${OPENCV_BUILD_DIR}" --parallel "$(nproc)" +cmake --install "${OPENCV_BUILD_DIR}" + +license_dir="${OPENCV_PREFIX}/share/licenses/opencv" +mkdir -p "${license_dir}" +cp "${OPENCV_SOURCE_DIR}/LICENSE" "${license_dir}/LICENSE.opencv" +cp "${OPENCV_CONTRIB_SOURCE_DIR}/LICENSE" \ + "${license_dir}/LICENSE.opencv_contrib" + +if [ "${OPENCV_BUILD_PYTHON3}" = "ON" ]; then + PYTHONPATH="${OPENCV_PREFIX}/python" \ + "${OPENCV_PYTHON_EXECUTABLE}" -c \ + "import cv2; info = cv2.getBuildInformation(); assert cv2.videoio_registry.hasBackend(cv2.CAP_GSTREAMER); assert 'YES' in info.split('GStreamer:')[1].split(chr(10))[0]; assert 'YES' in info.split('FFMPEG:')[1].split(chr(10))[0]; print(info)" +fi + +runtime_prefix="${OPENCV_PREFIX}-runtime" +cp -a "${OPENCV_PREFIX}" "${runtime_prefix}" +rm -rf \ + "${runtime_prefix}/include" \ + "${runtime_prefix}/lib/cmake" \ + "${runtime_prefix}/lib/pkgconfig" +find "${runtime_prefix}" -type f \( -name '*.a' -o -name '*.la' \) -delete +find "${runtime_prefix}" -type f -name '*.so*' \ + -exec strip --strip-unneeded {} + + +rm -rf \ + "${opencv_archive}" \ + "${opencv_contrib_archive}" \ + "${OPENCV_SOURCE_DIR}" \ + "${OPENCV_CONTRIB_SOURCE_DIR}" \ + "${OPENCV_BUILD_DIR}" diff --git a/docker/scripts/fetch_jetson_6_2_tensorrt.sh b/docker/scripts/fetch_jetson_6_2_tensorrt.sh new file mode 100755 index 0000000000..f72861c216 --- /dev/null +++ b/docker/scripts/fetch_jetson_6_2_tensorrt.sh @@ -0,0 +1,56 @@ +#!/usr/bin/env bash + +set -euo pipefail + +script_directory="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)" +repository_root="$(cd -- "${script_directory}/../.." && pwd)" + +tensorrt_deb_name="nv-tensorrt-local-tegra-repo-ubuntu2204-10.7.0-cuda-12.6_1.0-1_arm64.deb" +tensorrt_deb_url="${TENSORRT_DEB_URL:-https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.7.0/local_repo/${tensorrt_deb_name}}" +tensorrt_deb_sha256="${TENSORRT_DEB_SHA256:-cf8bd26b3b9c0f65ee8f3358bbc7abfcd4bbed4940b2220140fb108f0852e5c0}" +tensorrt_deb_destination="${TENSORRT_DEB_DESTINATION:-${repository_root}/docker/vendor/tensorrt/${tensorrt_deb_name}}" + +verify_sha256() { + local file_path="$1" + + if command -v sha256sum >/dev/null 2>&1; then + echo "${tensorrt_deb_sha256} ${file_path}" | sha256sum --check - + return + fi + if command -v shasum >/dev/null 2>&1; then + echo "${tensorrt_deb_sha256} ${file_path}" | shasum --algorithm 256 --check + return + fi + echo "Neither sha256sum nor shasum is available; cannot verify TensorRT." >&2 + return 1 +} + +mkdir -p -- "$(dirname -- "${tensorrt_deb_destination}")" + +if [[ -f "${tensorrt_deb_destination}" ]] && + verify_sha256 "${tensorrt_deb_destination}"; then + echo "Using verified TensorRT package at ${tensorrt_deb_destination}" + exit 0 +fi + +temporary_file="$(mktemp "${tensorrt_deb_destination}.part.XXXXXX")" +cleanup() { + rm -f -- "${temporary_file}" +} +trap cleanup EXIT + +curl \ + --fail \ + --location \ + --retry 10 \ + --retry-all-errors \ + --continue-at - \ + --connect-timeout 30 \ + --max-time 600 \ + --output "${temporary_file}" \ + "${tensorrt_deb_url}" +verify_sha256 "${temporary_file}" +mv -f -- "${temporary_file}" "${tensorrt_deb_destination}" +trap - EXIT + +echo "Fetched and verified TensorRT package at ${tensorrt_deb_destination}" diff --git a/docker/scripts/verify_ffmpeg.sh b/docker/scripts/verify_ffmpeg.sh new file mode 100644 index 0000000000..100f6df7db --- /dev/null +++ b/docker/scripts/verify_ffmpeg.sh @@ -0,0 +1,48 @@ +#!/bin/sh + +set -eu + +decoders="aac av1 h264 hevc mjpeg mpeg4 opus vp8 vp9" +decoder_list="$(ffmpeg -hide_banner -decoders 2>/dev/null)" +for decoder in ${decoders} +do + printf '%s\n' "${decoder_list}" | + awk -v name="${decoder}" '$2 == name { found = 1 } END { exit !found }' +done + +demuxers="avi flv matroska mov mpegts rtsp" +demuxer_list="$(ffmpeg -hide_banner -demuxers 2>/dev/null)" +for demuxer in ${demuxers} +do + printf '%s\n' "${demuxer_list}" | + awk -v name="${demuxer}" ' + { + count = split($2, names, ",") + for (i = 1; i <= count; i++) { + if (names[i] == name) { + found = 1 + } + } + } + END { exit !found } + ' +done + +protocol_list="$(ffmpeg -hide_banner -protocols 2>/dev/null)" +for protocol in file http https rtp tcp tls udp +do + printf '%s\n' "${protocol_list}" | grep -Eq "^[[:space:]]+${protocol}$" +done + +ffmpeg -hide_banner -loglevel error \ + -f lavfi \ + -i testsrc=size=16x16:rate=1 \ + -frames:v 1 \ + -c:v ffv1 \ + -y /tmp/ffmpeg-smoke.mkv +ffprobe -v error \ + -select_streams v:0 \ + -show_entries stream=codec_name,width,height \ + -of default=noprint_wrappers=1 \ + /tmp/ffmpeg-smoke.mkv | grep -q '^codec_name=ffv1$' +rm /tmp/ffmpeg-smoke.mkv diff --git a/docker/scripts/verify_gstreamer.sh b/docker/scripts/verify_gstreamer.sh new file mode 100644 index 0000000000..e31eb29518 --- /dev/null +++ b/docker/scripts/verify_gstreamer.sh @@ -0,0 +1,151 @@ +#!/bin/sh + +set -eu + +GSTREAMER_REQUIRE_NVCODEC="${GSTREAMER_REQUIRE_NVCODEC:-false}" +GSTREAMER_REQUIRE_NVCODEC_RUNTIME="${GSTREAMER_REQUIRE_NVCODEC_RUNTIME:-false}" + +for element in \ + appsink \ + appsrc \ + avdec_h264 \ + avdec_h265 \ + avidemux \ + bayer2rgb \ + capssetter \ + capsfilter \ + curlhttpsrc \ + decodebin \ + decodebin3 \ + dtlssrtpdec \ + dtlssrtpenc \ + fakesink \ + filesrc \ + h264parse \ + h264timestamper \ + h265parse \ + h265timestamper \ + jpegdec \ + jpegenc \ + jpegparse \ + matroskademux \ + nicesink \ + nicesrc \ + opusdec \ + opusenc \ + parsebin \ + qtdemux \ + queue \ + rtpbin \ + rtph264depay \ + rtph264pay \ + rtph265depay \ + rtph265pay \ + rtpjpegdepay \ + rtpjpegpay \ + rtpjitterbuffer \ + rtpopusdepay \ + rtpopuspay \ + rtpvp8depay \ + rtpvp8pay \ + rtpvp9depay \ + rtpvp9pay \ + rtmp2sink \ + rtmp2src \ + rtspclientsink \ + rtspsrc \ + sctpdec \ + sctpenc \ + srtpdec \ + srtpenc \ + tcpclientsrc \ + tcpserversink \ + tee \ + udpsink \ + udpsrc \ + uridecodebin \ + videoconvert \ + videorate \ + videoscale \ + videotestsrc \ + vp8dec \ + vp8enc \ + vp9dec \ + vp9enc \ + webrtcbin +do + gst-inspect-1.0 "${element}" >/dev/null +done + +test -e /opt/gstreamer/lib/libgstrtspserver-1.0.so.0 +for typelib in Gst GstAllocators GstApp GstRtp GstRtsp GstSdp GstVideo GstWebRTC; do + test -s "/opt/gstreamer/lib/girepository-1.0/${typelib}-1.0.typelib" +done +ldd /opt/gstreamer/lib/gstreamer-1.0/libgstlibav.so | + grep -q '/opt/ffmpeg/lib/libavcodec' +test -s /etc/ssl/certs/ca-certificates.crt +test -e /opt/gstreamer/lib/gio/modules/libgiognutls.so +grep -q 'libgiognutls' /opt/gstreamer/lib/gio/modules/giomodule.cache +ldd /opt/gstreamer/lib/gio/modules/libgiognutls.so | + grep -q 'libgnutls.so' +gst-inspect-1.0 rtspsrc | grep -Eq '^[[:space:]]+rtsps$' + +if [ "${GSTREAMER_REQUIRE_NVCODEC}" = "true" ]; then + test -e /opt/gstreamer/lib/gstreamer-1.0/libgstnvcodec.so + if ldd /opt/gstreamer/lib/gstreamer-1.0/libgstnvcodec.so | grep -q 'not found'; then + exit 1 + fi +fi + +if [ "${GSTREAMER_REQUIRE_NVCODEC_RUNTIME}" = "true" ]; then + for element in \ + cudaconvertscale \ + cudadownload \ + cudaupload \ + nvh264dec \ + nvh264enc \ + nvh265dec \ + nvh265enc \ + nvjpegdec \ + nvjpegenc + do + gst-inspect-1.0 "${element}" >/dev/null + done + gst-launch-1.0 -q \ + videotestsrc num-buffers=1 ! \ + video/x-raw,format=I420,width=64,height=64 ! \ + nvjpegenc ! \ + filesink location=/tmp/nvjpeg-smoke.jpg + gst-launch-1.0 -q \ + filesrc location=/tmp/nvjpeg-smoke.jpg ! \ + jpegparse ! \ + nvjpegdec ! \ + fakesink + gst-launch-1.0 -q \ + videotestsrc num-buffers=1 ! \ + video/x-raw,format=I420,width=64,height=64 ! \ + cudaupload ! \ + cudaconvertscale ! \ + 'video/x-raw(memory:CUDAMemory),format=RGBP' ! \ + cudadownload ! \ + fakesink + rm /tmp/nvjpeg-smoke.jpg +fi + +gst-launch-1.0 -q \ + videotestsrc num-buffers=1 ! \ + video/x-raw,format=BGR,width=16,height=16 ! \ + appsink max-buffers=1 drop=true sync=false wait-on-eos=false + +for pattern in bggr gbrg grbg rggb; do + gst-launch-1.0 -q \ + filesrc location=/dev/zero blocksize=64 num-buffers=1 ! \ + "video/x-bayer,format=${pattern},width=8,height=8,framerate=1/1" ! \ + bayer2rgb ! \ + fakesink + gst-launch-1.0 -q \ + filesrc location=/dev/zero blocksize=128 num-buffers=1 ! \ + "video/x-bayer,format=${pattern}16le,width=8,height=8,framerate=1/1" ! \ + bayer2rgb ! \ + fakesink +done diff --git a/docker/scripts/verify_gstreamer_cuda_tensor_bridge.sh b/docker/scripts/verify_gstreamer_cuda_tensor_bridge.sh new file mode 100644 index 0000000000..5cd82efa99 --- /dev/null +++ b/docker/scripts/verify_gstreamer_cuda_tensor_bridge.sh @@ -0,0 +1,21 @@ +#!/bin/sh + +set -eu + +BRIDGE_LIBRARY="${GSTREAMER_CUDA_TENSOR_BRIDGE_LIBRARY:-/opt/roboflow/lib/libroboflow_gstreamer_cuda_tensor.so.1}" + +test -s "${BRIDGE_LIBRARY}" +readelf -Ws "${BRIDGE_LIBRARY}" | grep -q 'rf_gstreamer_cuda_pipeline_create' +readelf -Ws "${BRIDGE_LIBRARY}" | grep -q 'rf_gstreamer_cuda_pipeline_retrieve' +readelf -Ws "${BRIDGE_LIBRARY}" | grep -q 'rf_gstreamer_cuda_pipeline_get_stats' +readelf -Ws "${BRIDGE_LIBRARY}" | grep -q 'rf_gstreamer_cuda_pipeline_interrupt' + +if readelf -d "${BRIDGE_LIBRARY}" | grep -Eq 'libopencv|libtorch|libpython'; then + exit 1 +fi + +missing="$(ldd "${BRIDGE_LIBRARY}" 2>/dev/null | awk '/not found/ { print $1 }' | sort -u)" +case "${missing}" in + ''|libcuda.so.1) ;; + *) printf '%s\n' "${missing}" >&2; exit 1 ;; +esac diff --git a/docker/scripts/verify_gstreamer_cuda_tensor_runtime.py b/docker/scripts/verify_gstreamer_cuda_tensor_runtime.py new file mode 100644 index 0000000000..fd1536420c --- /dev/null +++ b/docker/scripts/verify_gstreamer_cuda_tensor_runtime.py @@ -0,0 +1,229 @@ +import gc +import subprocess +import tempfile +import threading +import time +from pathlib import Path +from typing import List + +import numpy as np +import torch + +from inference.core.interfaces.camera.gstreamer_cuda_producer import ( + GstreamerCudaVideoFrameProducer, +) +from inference.core.interfaces.camera.gstreamer_cuda_tensor_bridge import ( + NativeGstreamerCudaTensorPipeline, +) + + +def _run_gstreamer(*arguments: str) -> None: + subprocess.run( + ["gst-launch-1.0", "-q", *arguments], + check=True, + timeout=60, + ) + + +def _create_h26x(path: Path, encoder: str, parser: str, pattern: str = "red") -> None: + _run_gstreamer( + "videotestsrc", + "num-buffers=8", + f"pattern={pattern}", + "!", + "video/x-raw,format=I420,width=320,height=180,framerate=30/1", + "!", + encoder, + "!", + parser, + "!", + "filesink", + f"location={path}", + ) + + +def _create_jpeg(path: Path) -> None: + _run_gstreamer( + "videotestsrc", + "num-buffers=1", + "pattern=red", + "!", + "video/x-raw,format=I420,width=320,height=180,framerate=1/1", + "!", + "nvjpegenc", + "!", + "filesink", + f"location={path}", + ) + + +def _validate_source(path: Path, minimum_frames: int) -> None: + producer = GstreamerCudaVideoFrameProducer(str(path), gpu_id=0) + properties = producer.discover_source_properties() + assert properties.width == 320 + assert properties.height == 180 + + success, first = producer.retrieve() + assert success and first is not None + assert first.is_cuda + assert first.dtype == torch.uint8 + assert tuple(first.shape) == (3, 180, 320) + assert first.device.index == 0 + assert first.data_ptr() != 0 + + channel_means = first.float().mean(dim=(1, 2)) + assert channel_means[0] > channel_means[1] + 80 + assert channel_means[0] > channel_means[2] + 80 + first_snapshot = first.clone() + tensors: List[torch.Tensor] = [first] + tensor = None + + while len(tensors) < 5 and producer.grab(): + success, tensor = producer.retrieve() + assert success and tensor is not None + assert tensor.is_cuda + assert tuple(tensor.shape) == (3, 180, 320) + tensors.append(tensor) + + assert len(tensors) >= minimum_frames + stats = producer.tensor_bridge_stats + assert stats["frames"] == len(tensors) + assert stats["cuda_maps"] == len(tensors) + assert stats["host_pixel_maps"] == 0 + assert stats["host_to_device_copies"] == 0 + assert stats["device_to_host_copies"] == 0 + assert stats["stream_synchronizations"] == len(tensors) + assert stats["active_leases"] == len(tensors) + assert tuple(first.stride()) == ( + stats["last_channel_stride"], + stats["last_row_stride"], + 1, + ) + + producer.release() + assert torch.equal(first, first_snapshot) + del first_snapshot + del first + tensor = None + tensors.clear() + gc.collect() + torch.cuda.synchronize() + + +def _validate_numpy_source(path: Path) -> None: + producer = GstreamerCudaVideoFrameProducer(str(path), gpu_id=0, output_tensor=False) + properties = producer.discover_source_properties() + assert properties.width == 320 + assert properties.height == 180 + + success, image = producer.retrieve() + assert success and image is not None + assert isinstance(image, np.ndarray) + assert image.dtype == np.uint8 + assert image.shape == (180, 320, 3) + assert image.flags.c_contiguous + channel_means = image.mean(axis=(0, 1)) + assert channel_means[2] > channel_means[1] + 80 + assert channel_means[2] > channel_means[0] + 80 + producer.release() + + +def _validate_repeated_grab_advances(path: Path) -> None: + baseline = GstreamerCudaVideoFrameProducer(str(path), gpu_id=0) + baseline.discover_source_properties() + success, first_frame = baseline.retrieve() + assert success and first_frame is not None + assert baseline.grab() + success, second_frame = baseline.retrieve() + assert success and second_frame is not None + assert not torch.equal(first_frame, second_frame) + baseline.release() + + producer = GstreamerCudaVideoFrameProducer(str(path), gpu_id=0) + producer.discover_source_properties() + assert producer.grab() + assert producer.grab() + success, frame_after_second_grab = producer.retrieve() + assert success and frame_after_second_grab is not None + assert torch.equal(frame_after_second_grab, second_frame) + producer.release() + + +def _validate_threaded_retrieve(path: Path) -> None: + # Mirrors VideoSource's threading: the producer is constructed (and + # discovery runs) on the caller thread while retrieve happens on a + # dedicated consumer thread that has made no prior CUDA call. This is the + # configuration production uses and single-threaded checks miss. + producer = GstreamerCudaVideoFrameProducer(str(path), gpu_id=0) + producer.discover_source_properties() + result = {} + + def consume() -> None: + try: + success, tensor = producer.retrieve() + result["success"] = success + result["shape"] = tuple(tensor.shape) + except Exception as error: # noqa: BLE001 - surfaced via assert below + result["error"] = error + + thread = threading.Thread(target=consume) + thread.start() + thread.join(timeout=30.0) + + assert not thread.is_alive() + assert "error" not in result, f"retrieve failed off-thread: {result.get('error')!r}" + assert result["success"] + assert result["shape"] == (3, 180, 320) + producer.release() + + +def _validate_interrupt_unblocks_pull() -> None: + pipeline = NativeGstreamerCudaTensorPipeline( + "appsrc is-live=true ! appsink name=rf_tensor_sink wait-on-eos=false", + device_id=0, + ) + entered_grab = threading.Event() + result = {} + + def pull_sample() -> None: + entered_grab.set() + result["grabbed"] = pipeline.grab() + + thread = threading.Thread(target=pull_sample) + thread.start() + assert entered_grab.wait(timeout=1.0) + time.sleep(0.1) + started = time.monotonic() + pipeline.interrupt() + thread.join(timeout=2.0) + elapsed = time.monotonic() - started + + assert not thread.is_alive() + assert result["grabbed"] is False + assert elapsed < 2.0 + pipeline.close() + + +def main() -> None: + assert torch.cuda.is_available() + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + h264_path = root / "test.h264" + h265_path = root / "test.h265" + jpeg_path = root / "test.jpg" + changing_h264_path = root / "changing.h264" + _create_h26x(h264_path, "nvh264enc", "h264parse") + _create_h26x(h265_path, "nvh265enc", "h265parse") + _create_h26x(changing_h264_path, "nvh264enc", "h264parse", pattern="ball") + _create_jpeg(jpeg_path) + _validate_source(h264_path, minimum_frames=5) + _validate_source(h265_path, minimum_frames=5) + _validate_source(jpeg_path, minimum_frames=1) + _validate_numpy_source(h264_path) + _validate_repeated_grab_advances(changing_h264_path) + _validate_threaded_retrieve(h264_path) + _validate_interrupt_unblocks_pull() + + +if __name__ == "__main__": + main() diff --git a/docker/scripts/verify_jetson_nvjpeg_runtime.sh b/docker/scripts/verify_jetson_nvjpeg_runtime.sh new file mode 100755 index 0000000000..3fe472d881 --- /dev/null +++ b/docker/scripts/verify_jetson_nvjpeg_runtime.sh @@ -0,0 +1,26 @@ +#!/bin/sh + +set -eu + +working_directory="$(mktemp -d)" +trap 'rm -rf "${working_directory}"' 0 HUP INT TERM + +jpeg_path="${working_directory}/nvjpeg-smoke.jpg" + +timeout --signal=TERM --kill-after=5 30 \ + gst-launch-1.0 -q \ + videotestsrc num-buffers=1 pattern=red ! \ + video/x-raw,format=I420,width=320,height=180,framerate=1/1 ! \ + nvvidconv ! \ + 'video/x-raw(memory:NVMM),format=I420' ! \ + nvjpegenc ! \ + filesink location="${jpeg_path}" + +test -s "${jpeg_path}" + +timeout --signal=TERM --kill-after=5 30 \ + gst-launch-1.0 -q \ + filesrc location="${jpeg_path}" ! \ + jpegparse ! \ + nvjpegdec ! \ + fakesink diff --git a/docker/scripts/verify_jetson_server_runtime.sh b/docker/scripts/verify_jetson_server_runtime.sh new file mode 100755 index 0000000000..294487ad80 --- /dev/null +++ b/docker/scripts/verify_jetson_server_runtime.sh @@ -0,0 +1,152 @@ +#!/bin/sh +set -eu + +/usr/local/cuda/bin/ptxas --version >/dev/null +test -f /usr/local/cuda/nvvm/libdevice/libdevice.10.bc +test -f /usr/local/share/licenses/cuda-ptxas +test "${TRITON_PTXAS_PATH}" = /usr/local/cuda/bin/ptxas +test "${TRITON_PTXAS_BLACKWELL_PATH}" = /usr/local/cuda/bin/ptxas +test ! -e /usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/bin/ptxas +test ! -e /usr/local/lib/python3.12/dist-packages/triton/backends/nvidia/bin/ptxas-blackwell +test -x /usr/local/bin/inference +test -x /usr/local/bin/run_uvicorn.sh +test -s /opt/roboflow/lib/libroboflow_jetson_tensor.so.1 +test -f /usr/include/python3.12/Python.h +command -v gcc >/dev/null +command -v python3-config >/dev/null +DISABLE_VERSION_CHECK=true inference --help >/dev/null +torchvision_image="$( + python3 -c 'from torchvision._internally_replaced_utils import _get_extension_path; print(_get_extension_path("image"))' +)" +ldd "${torchvision_image}" | grep -q 'libnvjpeg.so.13' + +for command_name in \ + c++ \ + cmake \ + curl \ + g++ \ + git \ + make \ + meson \ + ninja \ + nvcc \ + pip \ + pip3; do + if command -v "${command_name}" >/dev/null 2>&1; then + echo "Build command present in runtime: ${command_name}" >&2 + exit 1 + fi +done + +for package in \ + build-essential \ + cmake \ + cuda-nvcc-13-2 \ + cuda-toolkit-13-2 \ + g++ \ + make \ + ninja-build \ + pkg-config; do + if dpkg-query -W -f='${db:Status-Status}' "${package}" 2>/dev/null | + grep -qx installed; then + echo "Build dependency present in runtime: ${package}" >&2 + exit 1 + fi +done + +development_packages="$( + dpkg-query -W -f='${binary:Package}\n' | + grep -E -- '-dev(:[^[:space:]]+)?$' || true +)" +for package in ${development_packages}; do + package="${package%%:*}" + case "${package}" in + libc6-dev | libcrypt-dev | libexpat1-dev | libgcc-13-dev | \ + libpython3.12-dev | linux-libc-dev | python3.12-dev | zlib1g-dev) + ;; + *) + echo "Development package present in runtime: ${package}" >&2 + exit 1 + ;; + esac +done + +python3 - <<'PY' +import cv2 +import flash_attn +import importlib.metadata +import importlib.util +import onnxruntime +import tensorrt +import torch +import torchvision +import triton +from torchvision import io as torchvision_io + +arch_flags = set(torch._C._cuda_getArchFlags().split()) +assert {"sm_87", "sm_110"}.issubset(arch_flags), arch_flags + +providers = set(onnxruntime.get_available_providers()) +assert {"CUDAExecutionProvider", "TensorrtExecutionProvider"}.issubset(providers) + +assert cv2.cuda.getCudaEnabledDeviceCount() >= 0 +assert flash_attn.__version__ +assert tensorrt.__version__ +assert torch.__version__ +assert torchvision.__version__ +assert triton.__version__ +assert torchvision_io +assert hasattr(torch.ops.image, "decode_jpegs_cuda") + +torch_requirements = importlib.metadata.requires("torch") or [] +assert not any( + requirement.partition(";")[0].strip().lower().startswith("nvidia-") + for requirement in torch_requirements +), torch_requirements + +pip_cuda_distributions = sorted( + name + for distribution in importlib.metadata.distributions() + if (name := distribution.metadata.get("Name")) + and name.lower().startswith("nvidia-") +) +assert not pip_cuda_distributions, pip_cuda_distributions + +for module in ( + "build", + "cmake", + "expecttest", + "hypothesis", + "lintrunner", + "lit", + "ninja", + "parameterized", + "pydocstyle", + "twine", +): + assert importlib.util.find_spec(module) is None, module +PY + +extension_directory="$(mktemp -d)" +trap 'rm -rf "${extension_directory}"' EXIT +cat > "${extension_directory}/runtime_probe.c" <<'C' +#include + +static struct PyModuleDef module = { + PyModuleDef_HEAD_INIT, + "runtime_probe", + NULL, + -1, + NULL, +}; + +PyMODINIT_FUNC PyInit_runtime_probe(void) { + return PyModule_Create(&module); +} +C +python_include="$(python3 -c 'import sysconfig; print(sysconfig.get_path("include"))')" +gcc -O2 -shared -fPIC \ + -I"${python_include}" \ + "${extension_directory}/runtime_probe.c" \ + -o "${extension_directory}/runtime_probe$(python3-config --extension-suffix)" +PYTHONPATH="${extension_directory}" python3 -c "import runtime_probe" diff --git a/docker/scripts/verify_jetson_tensor_bridge.sh b/docker/scripts/verify_jetson_tensor_bridge.sh new file mode 100644 index 0000000000..bcd1b119d3 --- /dev/null +++ b/docker/scripts/verify_jetson_tensor_bridge.sh @@ -0,0 +1,21 @@ +#!/bin/sh + +set -eu + +BRIDGE_LIBRARY="${JETSON_TENSOR_BRIDGE_LIBRARY:-/opt/roboflow/lib/libroboflow_jetson_tensor.so.1}" + +test -s "${BRIDGE_LIBRARY}" +readelf -Ws "${BRIDGE_LIBRARY}" | grep -q 'rf_jetson_pipeline_create' +readelf -Ws "${BRIDGE_LIBRARY}" | grep -q 'rf_jetson_pipeline_retrieve' +readelf -Ws "${BRIDGE_LIBRARY}" | grep -q 'rf_jetson_pipeline_get_stats' +readelf -Ws "${BRIDGE_LIBRARY}" | grep -q 'rf_jetson_pipeline_interrupt' + +if readelf -d "${BRIDGE_LIBRARY}" | grep -Eq 'libopencv|libtorch|libpython'; then + exit 1 +fi + +missing="$(ldd "${BRIDGE_LIBRARY}" 2>/dev/null | awk '/not found/ { print $1 }' | sort -u)" +case "${missing}" in + ''|libcuda.so.1) ;; + *) printf '%s\n' "${missing}" >&2; exit 1 ;; +esac diff --git a/docker/scripts/verify_jetson_tensor_runtime.py b/docker/scripts/verify_jetson_tensor_runtime.py new file mode 100644 index 0000000000..5127f2c1e7 --- /dev/null +++ b/docker/scripts/verify_jetson_tensor_runtime.py @@ -0,0 +1,246 @@ +"""On-device Jetson tensor-runtime check. No CI job runs this: it needs a +Jetson with the BSP media stack, torch, and the inference package installed. +Run it manually inside a Jetson image, e.g.: + + docker run --rm --runtime nvidia --entrypoint python3 \ + -v "$PWD/docker/scripts/verify_jetson_tensor_runtime.py:/tmp/verify.py:ro" \ + /tmp/verify.py +""" + +import gc +import subprocess +import tempfile +import threading +import time +from pathlib import Path +from typing import List + +import numpy as np +import torch + +from inference.core.interfaces.camera.jetson_producer import JetsonVideoFrameProducer +from inference.core.interfaces.camera.jetson_tensor_bridge import ( + NativeJetsonTensorPipeline, +) + + +def _run_gstreamer(*arguments: str) -> None: + subprocess.run( + ["gst-launch-1.0", "-q", *arguments], + check=True, + timeout=60, + ) + + +def _create_h26x(path: Path, encoder: str, parser: str, pattern: str = "red") -> None: + _run_gstreamer( + "videotestsrc", + "num-buffers=8", + f"pattern={pattern}", + "!", + "video/x-raw,format=I420,width=320,height=180,framerate=30/1", + "!", + "nvvidconv", + "!", + "video/x-raw(memory:NVMM),format=NV12", + "!", + encoder, + "!", + parser, + "!", + "filesink", + f"location={path}", + ) + + +def _create_jpeg(path: Path) -> None: + _run_gstreamer( + "videotestsrc", + "num-buffers=1", + "pattern=red", + "!", + "video/x-raw,format=I420,width=320,height=180,framerate=1/1", + "!", + "nvvidconv", + "!", + "video/x-raw(memory:NVMM),format=I420", + "!", + "nvjpegenc", + "!", + "filesink", + f"location={path}", + ) + + +def _validate_source(path: Path, minimum_frames: int) -> None: + producer = JetsonVideoFrameProducer(str(path), output_tensor=True) + properties = producer.discover_source_properties() + assert properties.width == 320 + assert properties.height == 180 + + success, first = producer.retrieve() + assert success and first is not None + assert first.is_cuda + assert first.dtype == torch.uint8 + assert tuple(first.shape) == (3, 180, 320) + assert first.device.index == 0 + assert first.data_ptr() != 0 + assert first.is_contiguous() + + channel_means = first.float().mean(dim=(1, 2)) + assert channel_means[0] > channel_means[1] + 80 + assert channel_means[0] > channel_means[2] + 80 + first_snapshot = first.clone() + tensors: List[torch.Tensor] = [first] + tensor = None + + while len(tensors) < 5 and producer.grab(): + success, tensor = producer.retrieve() + assert success and tensor is not None + assert tensor.is_cuda + assert tuple(tensor.shape) == (3, 180, 320) + tensors.append(tensor) + + assert len(tensors) >= minimum_frames + stats = producer.tensor_bridge_stats + assert stats["frames"] == len(tensors) + assert stats["descriptor_maps"] == len(tensors) + assert stats["host_pixel_maps"] == 0 + assert stats["host_to_device_copies"] == 0 + assert stats["device_to_host_copies"] == 0 + assert stats["array_flatten_copies"] == 0 + assert stats["conversion_kernels"] == len(tensors) + assert stats["nvmm_frames"] == len(tensors) + + producer.release() + assert torch.equal(first, first_snapshot) + del first_snapshot + del first + tensor = None + tensors.clear() + gc.collect() + torch.cuda.synchronize() + + +def _validate_numpy_source(path: Path) -> None: + producer = JetsonVideoFrameProducer(str(path), output_tensor=False) + properties = producer.discover_source_properties() + assert properties.width == 320 + assert properties.height == 180 + + success, image = producer.retrieve() + assert success and image is not None + assert isinstance(image, np.ndarray) + assert image.dtype == np.uint8 + assert image.shape == (180, 320, 3) + assert image.flags.c_contiguous + channel_means = image.mean(axis=(0, 1)) + assert channel_means[2] > channel_means[1] + 80 + assert channel_means[2] > channel_means[0] + 80 + producer.release() + + +def _validate_repeated_grab_advances(path: Path) -> None: + baseline = JetsonVideoFrameProducer(str(path), output_tensor=True) + baseline.discover_source_properties() + success, first_frame = baseline.retrieve() + assert success and first_frame is not None + assert baseline.grab() + success, second_frame = baseline.retrieve() + assert success and second_frame is not None + assert not torch.equal(first_frame, second_frame) + baseline.release() + + producer = JetsonVideoFrameProducer(str(path), output_tensor=True) + producer.discover_source_properties() + assert producer.grab() + assert producer.grab() + success, frame_after_second_grab = producer.retrieve() + assert success and frame_after_second_grab is not None + assert torch.equal(frame_after_second_grab, second_frame) + producer.release() + + +def _validate_threaded_retrieve(path: Path) -> None: + # Mirrors VideoSource's threading: the producer is constructed (and + # discovery runs) on the caller thread while retrieve happens on a + # dedicated consumer thread that has made no prior CUDA call. This is the + # configuration production uses and single-threaded checks miss. + producer = JetsonVideoFrameProducer(str(path), output_tensor=True) + producer.discover_source_properties() + result = {} + + def consume() -> None: + try: + success, tensor = producer.retrieve() + result["success"] = success + result["shape"] = tuple(tensor.shape) + except Exception as error: # noqa: BLE001 - surfaced via assert below + result["error"] = error + + thread = threading.Thread(target=consume) + thread.start() + thread.join(timeout=30.0) + + assert not thread.is_alive() + assert "error" not in result, f"retrieve failed off-thread: {result.get('error')!r}" + assert result["success"] + assert result["shape"] == (3, 180, 320) + producer.release() + + +def _validate_interrupt_unblocks_pull() -> None: + pipeline = NativeJetsonTensorPipeline( + "appsrc is-live=true ! appsink name=rf_tensor_sink wait-on-eos=false", + device_id=0, + ) + entered_grab = threading.Event() + result = {} + + def pull_sample() -> None: + entered_grab.set() + result["grabbed"] = pipeline.grab() + + thread = threading.Thread(target=pull_sample) + thread.start() + assert entered_grab.wait(timeout=1.0) + time.sleep(0.1) + started = time.monotonic() + pipeline.interrupt() + thread.join(timeout=2.0) + elapsed = time.monotonic() - started + + assert not thread.is_alive() + assert result["grabbed"] is False + assert elapsed < 2.0 + pipeline.close() + + +def main() -> None: + assert torch.cuda.is_available() + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + h264_path = root / "test.h264" + h265_path = root / "test.h265" + jpeg_path = root / "test.jpg" + changing_h264_path = root / "changing.h264" + _create_h26x(h264_path, "nvv4l2h264enc", "h264parse") + _create_h26x(h265_path, "nvv4l2h265enc", "h265parse") + _create_h26x( + changing_h264_path, + "nvv4l2h264enc", + "h264parse", + pattern="ball", + ) + _create_jpeg(jpeg_path) + _validate_source(h264_path, minimum_frames=5) + _validate_source(h265_path, minimum_frames=5) + _validate_source(jpeg_path, minimum_frames=1) + _validate_numpy_source(h264_path) + _validate_repeated_grab_advances(changing_h264_path) + _validate_threaded_retrieve(h264_path) + _validate_interrupt_unblocks_pull() + + +if __name__ == "__main__": + main() diff --git a/docker/scripts/verify_media_runtime.sh b/docker/scripts/verify_media_runtime.sh new file mode 100644 index 0000000000..4335a1f678 --- /dev/null +++ b/docker/scripts/verify_media_runtime.sh @@ -0,0 +1,129 @@ +#!/bin/sh + +set -eu + +# Shared runtime hygiene check for every media-enabled image. Jetson-only +# artifacts are validated when present; set +# MEDIA_RUNTIME_REQUIRE_JETSON_PLUGINS=true to make the baked L4T plugin set +# a hard requirement (only the Jetson media image bakes them: JetPack 6 +# injects the BSP plugins at container start, and dGPU images have none). + +if [ "${VERIFY_MEDIA_RUNTIME_DEVELOPMENT_TOOLS:-true}" = true ]; then + for command_name in \ + cc \ + c++ \ + gcc \ + g++ \ + make \ + cmake \ + meson \ + ninja \ + git \ + curl \ + pip \ + pip3 + do + if command -v "${command_name}" >/dev/null 2>&1; then + echo "Unexpected build command in media runtime: ${command_name}" >&2 + exit 1 + fi + done + + development_packages="$( + dpkg-query -W -f='${binary:Package}\n' | + grep -E -- '-dev(:[^[:space:]]+)?$' || true + )" + if [ -n "${development_packages}" ]; then + printf 'Unexpected development packages in media runtime:\n%s\n' \ + "${development_packages}" >&2 + exit 1 + fi +fi + +set -- /opt/ffmpeg /opt/gstreamer +if [ -d /opt/roboflow ]; then + set -- "$@" /opt/roboflow +fi +if [ -d /opt/opencv ]; then + set -- "$@" /opt/opencv +fi +if [ -d /opt/cuda-runtime ]; then + set -- "$@" /opt/cuda-runtime +fi +if [ -d /usr/lib/aarch64-linux-gnu/gstreamer-1.0 ]; then + set -- "$@" /usr/lib/aarch64-linux-gnu/gstreamer-1.0 +fi +if [ -d /usr/lib/aarch64-linux-gnu/nvidia ]; then + set -- "$@" /usr/lib/aarch64-linux-gnu/nvidia +fi +if [ -d /usr/lib/aarch64-linux-gnu/tegra ]; then + set -- "$@" /usr/lib/aarch64-linux-gnu/tegra +fi + +development_artifact="$( + find "$@" -type f \ + \( -name '*.a' -o -name '*.la' -o -name '*.h' -o -name '*.pc' \) \ + -print -quit +)" +if [ -n "${development_artifact}" ]; then + echo "Unexpected development artifact in media runtime: ${development_artifact}" >&2 + exit 1 +fi + +set -- /opt/ffmpeg /opt/gstreamer +if [ -d /opt/roboflow ]; then + set -- "$@" /opt/roboflow +fi +if [ -d /opt/opencv ]; then + set -- "$@" /opt/opencv +fi +if [ -d /opt/cuda-runtime ]; then + set -- "$@" /opt/cuda-runtime +fi + +required_jetson_plugins=" +/usr/lib/aarch64-linux-gnu/gstreamer-1.0/libgstnvarguscamerasrc.so +/usr/lib/aarch64-linux-gnu/gstreamer-1.0/libgstnvjpeg.so +/usr/lib/aarch64-linux-gnu/gstreamer-1.0/libgstnvvidconv.so +/usr/lib/aarch64-linux-gnu/gstreamer-1.0/libgstnvvideo4linux2.so +" +if [ "${MEDIA_RUNTIME_REQUIRE_JETSON_PLUGINS:-false}" = "true" ]; then + for plugin in ${required_jetson_plugins}; do + test -s "${plugin}" + done +fi + +if [ -s /usr/lib/aarch64-linux-gnu/gstreamer-1.0/libgstnvjpeg.so ]; then + test -s /opt/cuda-runtime/lib/libnvjpeg.so.13 + command -v lsmod >/dev/null +fi + +find "$@" -type f -name '*.so*' -print0 | + xargs -0 sh -c ' + for library do + missing="$( + ldd "${library}" 2>/dev/null | + awk '\''/not found/ { print $1 }'\'' | + grep -v -x '\''libcuda.so.1'\'' || true + )" + if [ -n "${missing}" ]; then + printf "%s\n%s\n" "${library}" "${missing}" >&2 + exit 1 + fi + done + ' sh + +for library in ${required_jetson_plugins}; do + if [ ! -s "${library}" ]; then + continue + fi + missing="$( + ldd "${library}" 2>/dev/null | + awk '/not found/ { print $1 }' | + grep -v -x 'libcuda.so.1' || true + )" + if [ -n "${missing}" ]; then + printf '%s\n%s\n' "${library}" "${missing}" >&2 + exit 1 + fi +done diff --git a/docker/scripts/verify_opencv.sh b/docker/scripts/verify_opencv.sh new file mode 100644 index 0000000000..e9ba280692 --- /dev/null +++ b/docker/scripts/verify_opencv.sh @@ -0,0 +1,138 @@ +#!/bin/sh + +set -eu + +OPENCV_REQUIRE_CUDA_BUILD="${OPENCV_REQUIRE_CUDA_BUILD:-false}" +OPENCV_REQUIRE_CUDA_RUNTIME="${OPENCV_REQUIRE_CUDA_RUNTIME:-false}" +OPENCV_EXPECTED_CUDA_ARCH_BIN="${OPENCV_EXPECTED_CUDA_ARCH_BIN:-}" +export \ + OPENCV_EXPECTED_CUDA_ARCH_BIN \ + OPENCV_REQUIRE_CUDA_BUILD \ + OPENCV_REQUIRE_CUDA_RUNTIME + +python3 - <<'PY' +import os + +import cv2 +import numpy as np + +assert cv2.videoio_registry.hasBackend(cv2.CAP_GSTREAMER) +info = cv2.getBuildInformation() +assert "YES" in info.split("GStreamer:", 1)[1].splitlines()[0] +assert "YES" in info.split("FFMPEG:", 1)[1].splitlines()[0] + +require_cuda_build = os.environ["OPENCV_REQUIRE_CUDA_BUILD"].lower() == "true" +require_cuda_runtime = os.environ["OPENCV_REQUIRE_CUDA_RUNTIME"].lower() == "true" +expected_cuda_arches = os.environ["OPENCV_EXPECTED_CUDA_ARCH_BIN"] +cuda_section = info.split("NVIDIA CUDA:", 1) +cuda_built = len(cuda_section) == 2 and "YES" in cuda_section[1].splitlines()[0] +if require_cuda_build: + assert cuda_built + assert hasattr(cv2, "cuda_GpuMat") + assert hasattr(cv2.cuda, "cvtColor") + assert hasattr(cv2.cuda, "demosaicing") + assert hasattr(cv2.cuda, "resize") + assert hasattr(cv2.cuda, "createGaussianFilter") +if expected_cuda_arches: + built_arches = ( + info.split("NVIDIA GPU arch:", 1)[1].splitlines()[0].split() + ) + for expected_arch in expected_cuda_arches.replace(",", ";").split(";"): + expected_arch = expected_arch.strip().replace(".", "") + if expected_arch: + assert expected_arch in built_arches + +gray = np.arange(16, dtype=np.uint8).reshape(4, 4) +assert cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR).shape == (4, 4, 3) +assert cv2.cvtColor( + np.zeros((4, 4, 4), dtype=np.uint8), cv2.COLOR_BGRA2BGR +).shape == (4, 4, 3) +for dtype in (np.uint8, np.uint16): + bayer = np.arange(64, dtype=dtype).reshape(8, 8) + for conversion in ( + cv2.COLOR_BAYER_BG2BGR, + cv2.COLOR_BAYER_GB2BGR, + cv2.COLOR_BAYER_RG2BGR, + cv2.COLOR_BAYER_GR2BGR, + ): + converted = cv2.cvtColor(bayer, conversion) + assert converted.shape == (8, 8, 3) + assert converted.dtype == dtype +assert cv2.cvtColor( + np.zeros((2, 4, 2), dtype=np.uint8), cv2.COLOR_YUV2BGR_YUY2 +).shape == (2, 4, 3) +assert cv2.resize(gray, (2, 2), interpolation=cv2.INTER_AREA).shape == (2, 2) + +frame = np.zeros((4, 4, 3), dtype=np.uint8) +lut = np.arange(256, dtype=np.uint8) +assert cv2.LUT(frame, lut).shape == frame.shape +assert cv2.imencode(".jpg", frame)[0] +assert cv2.SIFT_create() is not None +assert cv2.QRCodeDetector() is not None +assert cv2.createBackgroundSubtractorMOG2() is not None + +cuda_device_count = cv2.cuda.getCudaEnabledDeviceCount() if cuda_built else 0 +if require_cuda_runtime: + assert cuda_device_count > 0 +if cuda_device_count > 0: + gray = np.arange(64, dtype=np.uint8).reshape(8, 8) + gray_cuda = cv2.cuda_GpuMat() + gray_cuda.upload(gray) + assert np.array_equal(gray_cuda.download(), gray) + + color_cuda = cv2.cuda.cvtColor(gray_cuda, cv2.COLOR_GRAY2BGR) + color = color_cuda.download() + assert color.shape == (8, 8, 3) + assert np.array_equal(color[:, :, 0], gray) + assert np.array_equal(color[:, :, 1], gray) + assert np.array_equal(color[:, :, 2], gray) + + for dtype in (np.uint8, np.uint16): + bayer_cuda = cv2.cuda_GpuMat() + bayer_cuda.upload(np.arange(64, dtype=dtype).reshape(8, 8)) + for conversion in ( + cv2.COLOR_BAYER_BG2BGR, + cv2.COLOR_BAYER_GB2BGR, + cv2.COLOR_BAYER_RG2BGR, + cv2.COLOR_BAYER_GR2BGR, + ): + demosaiced_cuda = cv2.cuda.demosaicing( + bayer_cuda, conversion + ) + converted = demosaiced_cuda.download() + assert converted.shape == (8, 8, 3) + assert converted.dtype == dtype + expected = cv2.cvtColor(bayer_cuda.download(), conversion) + assert np.allclose(converted, expected, atol=1) + + resized_cuda = cv2.cuda.resize( + gray_cuda, (4, 4), interpolation=cv2.INTER_NEAREST + ) + resized = resized_cuda.download() + assert resized.shape == (4, 4) + assert np.array_equal( + resized, + cv2.resize(gray, (4, 4), interpolation=cv2.INTER_NEAREST), + ) + + gaussian = cv2.cuda.createGaussianFilter( + cv2.CV_8UC1, cv2.CV_8UC1, (3, 3), 0 + ) + filtered = gaussian.apply(gray_cuda).download() + assert filtered.shape == (8, 8) + assert np.allclose( + filtered, + cv2.GaussianBlur(gray, (3, 3), 0, borderType=cv2.BORDER_REFLECT_101), + atol=1, + ) + +capture = cv2.VideoCapture( + "videotestsrc num-buffers=1 ! " + "video/x-raw,format=BGR,width=16,height=16 ! appsink", + cv2.CAP_GSTREAMER, +) +assert capture.isOpened() +ok, captured_frame = capture.read() +capture.release() +assert ok and captured_frame.shape == (16, 16, 3) +PY diff --git a/docs/workflows/custom_python_code_blocks.md b/docs/workflows/custom_python_code_blocks.md index 2f9432d552..5f21e11234 100644 --- a/docs/workflows/custom_python_code_blocks.md +++ b/docs/workflows/custom_python_code_blocks.md @@ -405,6 +405,117 @@ as if that was normal block exposed through static plugin: } ``` +## Tensor data representation and the `tensor_compatibility` knob + +Servers running the tensor data representation (`ENABLE_TENSOR_DATA_REPRESENTATION=True`) +pass predictions between blocks as native `inference_models` objects (torch tensors on +device) instead of `sv.Detections` / numpy. Custom Python blocks declare how they want to +meet that world through an optional manifest field: + +```json +{ + "type": "ManifestDescription", + "block_type": "MyBlock", + "tensor_compatibility": "legacy_compatibility", + ... +} +``` + +### `legacy_compatibility` (the default) + +Your existing blocks keep working **unmodified**. The execution engine converts native +tensor predictions into the documented `sv.Detections` / numpy representations right +before your `run(...)` receives them, and converts returned legacy objects back into +native ones afterwards. Conversion is driven by the kinds you declare on inputs/outputs; +undeclared (wildcard) values are recognised by type where possible, and values that +cannot cross the boundary raise a clear error naming the block, the input/output and a +remediation. + +Modal remote execution is covered by the same boundary: inputs are converted before +they are serialized to the sandbox (predictions travel as the standard detections wire +payloads), and results are converted back to native objects when they return โ€” your +code sees exactly what it would see locally. + +Be aware of the conversion cost on GPU servers: every prediction crossing the boundary +pays a deviceโ†’host copy, and instance segmentation additionally materialises dense +masks (native masks may travel in a compact RLE form). For pipelines where the custom +block sits between heavy segmentation steps, consider `tensor_native`. + +### `tensor_native` + +Your `run(...)` receives and must return native `inference_models` objects โ€” no +conversion happens on either side. The standard imports are extended (only on +tensor-enabled servers) so your code can operate on and mint native predictions: + +- `torch` +- `Detections`, `InstanceDetections`, `KeyPoints`, `ClassificationPrediction`, + `MultiLabelClassificationPrediction` (the `inference_models` base types) +- `WORKFLOWS_IMAGE_TENSOR_DEVICE` โ€” the device native tensors pin to +- `build_native_image_metadata`, `attach_native_detection_metadata` โ€” helpers that + attach the metadata the engine requires + +The numpy/`sv` imports remain available โ€” mixing representations inside your own code is +fine, as long as what you *return* is native. + +**The native output contract.** Detections you return MUST carry +`image_metadata["class_names"]` (a `class_id -> name` map) and a per-box +`detection_id` โ€” the serializer raises otherwise. The helpers guarantee both: + +```python +def run(self, image) -> BlockResult: + detections = Detections( + xyxy=torch.tensor([[10.0, 10.0, 60.0, 60.0]], device=WORKFLOWS_IMAGE_TENSOR_DEVICE), + class_id=torch.tensor([0], device=WORKFLOWS_IMAGE_TENSOR_DEVICE), + confidence=torch.tensor([0.9], device=WORKFLOWS_IMAGE_TENSOR_DEVICE), + ) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names={0: "widget"}, + prediction_type="object-detection", + ) + return {"predictions": detections} +``` + +with the block declared as: + +```json +{ + "type": "DynamicBlockDefinition", + "manifest": { + "type": "ManifestDescription", + "block_type": "NativeMinter", + "tensor_compatibility": "tensor_native", + "inputs": { + "image": { + "type": "DynamicInputDefinition", + "selector_types": ["input_image"] + } + }, + "outputs": { + "predictions": { + "type": "DynamicOutputDefinition", + "kind": ["object_detection_prediction"] + } + } + }, + "code": { + "type": "PythonCode", + "run_function_code": "..." + } +} +``` + +### Compile-time errors + +Declaring `tensor_native` in the wrong environment fails fast at workflow compilation +(the block's tensor-expecting code would break mid-run anyway): + +- on a server running the numpy representation (flag off) โ€” use `legacy_compatibility` + or enable `ENABLE_TENSOR_DATA_REPRESENTATION`; +- with `WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE=modal` โ€” `tensor_native` is not yet + supported for remote execution (`legacy_compatibility` over Modal is fully supported). + ## Debugging dynamic blocks Set `debug=True` on the `/workflows/run` request to capture diagnostics from custom Python diff --git a/inference/core/env.py b/inference/core/env.py index e455cd431d..58e1714a39 100644 --- a/inference/core/env.py +++ b/inference/core/env.py @@ -828,7 +828,9 @@ def _reset_offline_mode_lock_after_fork() -> None: os.getenv("INFERENCE_PIPELINE_PREDICTIONS_QUEUE_SIZE", 512) ) RESTART_ATTEMPT_DELAY = int(os.getenv("INFERENCE_PIPELINE_RESTART_ATTEMPT_DELAY", 1)) -DEFAULT_BUFFER_SIZE = int(os.getenv("VIDEO_SOURCE_BUFFER_SIZE", "64")) +# DEFAULT_BUFFER_SIZE (VIDEO_SOURCE_BUFFER_SIZE) is defined further down - its +# default depends on ENABLE_TENSOR_DATA_REPRESENTATION, which is not parsed yet +# at this point of the module. DEFAULT_ADAPTIVE_MODE_STREAM_PACE_TOLERANCE = float( os.getenv("VIDEO_SOURCE_ADAPTIVE_MODE_STREAM_PACE_TOLERANCE", "0.1") ) @@ -1037,6 +1039,22 @@ def _reset_offline_mode_lock_after_fork() -> None: RUNS_ON_JETSON = str2bool(os.getenv("RUNS_ON_JETSON", "False")) +# Opt-out from the GStreamer-based legacy video sources (e.g. the Jetson RTSP +# producer) โ€” when True, plain (non-tensor) source references always decode +# through the cv2 producer. Default is False. +DISABLE_GSTREAMER_VIDEO_SOURCES = str2bool( + os.getenv("DISABLE_GSTREAMER_VIDEO_SOURCES", "False") +) + +# Instance-segmentation tensor blocks request dense (on-device) masks from the +# inference_models adapter instead of the default RLE carrier. Dense masks let +# GPU consumers (e.g. the mask-visualization compositor) skip the host-side +# RLE decode entirely; RLE remains preferable when predictions are mostly +# serialized to the wire. Default is False (RLE). +WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS = str2bool( + os.getenv("WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS", "False") +) + DOCKER_SOCKET_PATH: Optional[str] = os.getenv("DOCKER_SOCKET_PATH") ENABLE_WORKFLOWS_PROFILING = str2bool(os.getenv("ENABLE_WORKFLOWS_PROFILING", "False")) @@ -1408,3 +1426,63 @@ def _reset_offline_mode_lock_after_fork() -> None: ) else: DISABLED_INFERENCE_MODELS_BACKENDS = set() + +ENABLE_TENSOR_DATA_REPRESENTATION = ( + str2bool(os.getenv("ENABLE_TENSOR_DATA_REPRESENTATION", "False")) + and USE_INFERENCE_MODELS +) + + +WORKFLOWS_IMAGE_TENSOR_DEVICE_STR: Optional[str] = os.getenv( + "WORKFLOWS_IMAGE_TENSOR_DEVICE" +) +# `torch` is an OPTIONAL dependency: the slim `inference-core` artifact ships +# without it, and `import inference.core.env` must succeed when torch is ABSENT. +# Only the tensor-native path (ENABLE_TENSOR_DATA_REPRESENTATION on) needs a torch +# device, so both the import and the device materialisation are deferred behind +# the flag AND guarded on torch's presence โ€” this code never touches torch when +# the flag is off or torch is missing. Every consumer of this value is a +# tensor-only (flag-on) code path, so leaving it ``None`` off-flag is safe. +WORKFLOWS_IMAGE_TENSOR_DEVICE = None +if ENABLE_TENSOR_DATA_REPRESENTATION: + try: + import torch + + if WORKFLOWS_IMAGE_TENSOR_DEVICE_STR is None: + WORKFLOWS_IMAGE_TENSOR_DEVICE_STR = ( + "cuda" if torch.cuda.is_available() else "cpu" + ) + WORKFLOWS_IMAGE_TENSOR_DEVICE = torch.device(WORKFLOWS_IMAGE_TENSOR_DEVICE_STR) + except ImportError: + # Flag on but torch not installed: keep env import working; the tensor + # path surfaces a targeted error when it actually tries to use a device. + pass + +# VideoSource decode-buffer depth. 64 buffered frames is a sane host-RAM +# default, but under ENABLE_TENSOR_DATA_REPRESENTATION the hardware decoders +# emit CLONED GPU tensors (dgpu/jetson producers), so every buffered 1080p RGB +# frame holds ~6 MB of VRAM (~25 MB at 4K): a 64-deep queue costs ~0.4 GB per +# 1080p stream (~1.6 GB at 4K) before any model allocates, multiplied across +# sources. Default to a shallow queue when the flag is on; an explicit +# VIDEO_SOURCE_BUFFER_SIZE always wins. +DEFAULT_BUFFER_SIZE = int( + os.getenv( + "VIDEO_SOURCE_BUFFER_SIZE", + "8" if ENABLE_TENSOR_DATA_REPRESENTATION else "64", + ) +) + +# Instance-mask carrier for the tensor-native SAM video-tracker blocks +# ("rle" = compact COCO RLE, "dense" = boolean torch tensors). This is an +# execution-level flag, NOT a block manifest field (manifests stay identical +# across the flag swap); GCP_SERVERLESS forces "rle" at run time regardless. +WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION = ( + os.getenv("WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION", "rle").strip().lower() +) +if WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION not in {"rle", "dense"}: + warnings.warn( + "Invalid value of `WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION` variable: " + f"{WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION!r} - allowed values are 'rle' " + "and 'dense'. Falling back to 'rle'." + ) + WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION = "rle" diff --git a/inference/core/interfaces/camera/collection_policy.py b/inference/core/interfaces/camera/collection_policy.py new file mode 100644 index 0000000000..f37d88c66e --- /dev/null +++ b/inference/core/interfaces/camera/collection_policy.py @@ -0,0 +1,373 @@ +"""Collection policies for multi-source video consumption in InferencePipeline. + +Motivation (measured on Jetson AGX Orin, 8x 2K@15 RTSP): consumer cameras emit +frames in bursts separated by encoder pauses of 400-500 ms around every +I-frame. The legacy blocking batch collection waits for a fresh frame from +EVERY source per cycle, so with several staggered cameras almost every batch +stalls on whichever source is inside its pause - throughput collapses to a +fraction of the aggregate frame rate while decoded frames are silently +discarded. The policies in this module remove that coupling: + +* the batch-collection window self-tunes from the pipeline's own rhythms + instead of a hand-picked ``batch_collection_timeout``: once per-source + arrival rates are measured, the round period is matched to the fastest + live source's frame period (full batches whenever the model has headroom, + floor-window under saturation); before estimates exist it falls back to a + fraction of the measured execution time, +* consumption is FIFO with a bounded staleness budget - every frame is served + while the consumer keeps up, and under overload served frames are never + older than ``max_staleness`` (drops are counted and reported, not silent). + +File sources are exempt from the staleness budget by design: file decoding is +demand-paced (a "stale" frame only means the consumer was busy) and dropping +frames from a file would silently corrupt every-frame processing guarantees. +""" + +import logging +from collections import deque +from datetime import datetime +from enum import Enum +from time import monotonic +from typing import TYPE_CHECKING, Callable, Deque, Dict, Optional, Union + +from inference.core import env as core_env +from inference.core.interfaces.camera.entities import VideoFrame + +if TYPE_CHECKING: # pragma: no cover - typing only + from inference.core.interfaces.camera.video_source import VideoSource + +logger = logging.getLogger(__name__) + +DEFAULT_MAX_STALENESS_SECONDS = 0.5 +MIN_COLLECTION_WINDOW_SECONDS = 0.002 +MAX_COLLECTION_WINDOW_SECONDS = 0.030 +INITIAL_COLLECTION_WINDOW_SECONDS = 0.005 +COLLECTION_WINDOW_EXECUTION_FRACTION = 0.2 +EXECUTION_GAP_EMA_ALPHA = 0.2 +FRESHEST_MODE_BATCH_COLLECTION_TIMEOUT = 0.02 +STALENESS_DROP_CAUSE = "STALENESS_BUDGET_EXCEEDED" +LEGACY_MODE_ALIASES = frozenset({"legacy", "none"}) +# Rate-matched window: cap and the arrival-period estimator's shape. The +# estimator is count-over-span (never an EMA of gaps - bursty encoders like +# consumer RTSP cameras emit 2-frame clusters around GOP pauses, and gap +# EMAs oscillate at the burst frequency while a span over several burst +# cycles converges on the true rate). +RATE_MATCHED_WINDOW_CAP_SECONDS = 0.1 +ARRIVAL_PERIOD_SAMPLE_WINDOW = 64 +MIN_ARRIVAL_SAMPLES_TO_TRUST = 16 +SOURCE_ACTIVITY_HORIZON_SECONDS = 2.0 + + +class VideoProcessingMode(str, Enum): + """High-level intent for live multi-source consumption. + + * AUTO - FIFO consumption with a staleness budget and a self-tuning + collection window. Serves every frame while the consumer keeps up; + under overload degrades to freshest-at-capacity with bounded, + reported drops. + * EVERY_FRAME - AUTO's machinery with the staleness budget disabled: + strict FIFO for live sources (under sustained overload latency pins + at the decoding-buffer depth). + * FRESHEST - legacy latest-wins semantics (EAGER consumption) with a + small fixed collection timeout; minimal latency, silent skipping. + """ + + AUTO = "auto" + EVERY_FRAME = "every_frame" + FRESHEST = "freshest" + + +def resolve_video_processing_mode( + explicit_mode: Optional[Union[str, VideoProcessingMode]], +) -> Optional[VideoProcessingMode]: + """Resolve the effective processing mode. + + Order: explicit argument > tensor-representation cohort default (AUTO + when ``ENABLE_TENSOR_DATA_REPRESENTATION`` is set - the same opt-in + boundary that already gates decoding-buffer depth) > ``None`` meaning + the legacy collection behavior, byte-for-byte. The explicit strings + ``"legacy"`` / ``"none"`` force the legacy behavior even inside the + tensor cohort - the escape hatch from the flag-driven AUTO default. + """ + if explicit_mode is not None: + if ( + isinstance(explicit_mode, str) + and explicit_mode.lower() in LEGACY_MODE_ALIASES + ): + return None + return VideoProcessingMode(explicit_mode) + if core_env.ENABLE_TENSOR_DATA_REPRESENTATION: + return VideoProcessingMode.AUTO + return None + + +class _SourceArrivalEstimator: + """Count-over-span estimate of one live source's true frame period. + + Fed with ingress capture timestamps (``VideoFrame.frame_timestamp`` is + stamped on the decode thread), so downstream queueing cannot distort the + span. A gap longer than the activity horizon (reconnect, stall) clears + the window - a rejoin gap is not a frame period - and the estimate stays + ``None`` until enough fresh samples accumulate again. + """ + + def __init__(self, sample_window: int = ARRIVAL_PERIOD_SAMPLE_WINDOW): + self._timestamps: Deque[datetime] = deque(maxlen=sample_window) + self._last_seen_at: Optional[float] = None + + def observe(self, frame_timestamp: datetime, now: float) -> None: + if self._timestamps: + gap = (frame_timestamp - self._timestamps[-1]).total_seconds() + if gap > SOURCE_ACTIVITY_HORIZON_SECONDS or gap < 0: + self._timestamps.clear() + self._timestamps.append(frame_timestamp) + self._last_seen_at = now + + def period(self, now: float) -> Optional[float]: + if ( + self._last_seen_at is None + or now - self._last_seen_at > SOURCE_ACTIVITY_HORIZON_SECONDS + ): + return None + if len(self._timestamps) < MIN_ARRIVAL_SAMPLES_TO_TRUST: + return None + span = (self._timestamps[-1] - self._timestamps[0]).total_seconds() + if span <= 0: + return None + return span / (len(self._timestamps) - 1) + + +class AdaptiveWindowController: + """Self-tunes the batch-collection window from the collection rhythm. + + Two regimes, picked per round by whether a trustworthy arrival-period + estimate exists: + + * RATE-MATCHED (estimate available): ``window = clamp(min live source + period - exec EMA, floor, cap)``. The round period then equals the + fastest source's frame period, so in the light-load regime every + round finds each source with exactly one fresh frame - full batches + by construction. Under saturation (exec >= period) the subtraction + goes negative and the window floors, which is the correct move: every + source already has frames queued when the round starts. The added + collection wait is offset by removed queue wait (a frame that misses + a round today sits in the LAZY queue for a full round period), so + end-to-end latency stays roughly flat while occupancy rises. + * FALLBACK (startup, no live estimates): a fraction of the exec-time + EMA, clamped - light models get a near-zero window, heavy models a + larger alignment window. + + The gap between a non-empty collection finishing and the next collection + starting is, by construction, the batch execution time. + """ + + def __init__( + self, + alpha: float = EXECUTION_GAP_EMA_ALPHA, + execution_fraction: float = COLLECTION_WINDOW_EXECUTION_FRACTION, + min_window: float = MIN_COLLECTION_WINDOW_SECONDS, + max_window: float = MAX_COLLECTION_WINDOW_SECONDS, + initial_window: float = INITIAL_COLLECTION_WINDOW_SECONDS, + rate_matched_cap: float = RATE_MATCHED_WINDOW_CAP_SECONDS, + clock: Callable[[], float] = monotonic, + ): + self._alpha = alpha + self._execution_fraction = execution_fraction + self._min_window = min_window + self._max_window = max_window + self._rate_matched_cap = rate_matched_cap + self._window = initial_window + self._clock = clock + self._execution_gap_ema: Optional[float] = None + self._last_non_empty_collection_end: Optional[float] = None + + def on_collection_start( + self, minimum_arrival_period: Optional[float] = None + ) -> float: + now = self._clock() + if self._last_non_empty_collection_end is not None: + execution_gap = now - self._last_non_empty_collection_end + if self._execution_gap_ema is None: + self._execution_gap_ema = execution_gap + else: + self._execution_gap_ema = ( + 1 - self._alpha + ) * self._execution_gap_ema + self._alpha * execution_gap + if minimum_arrival_period is not None: + rate_matched = minimum_arrival_period - self._execution_gap_ema + self._window = min( + max(rate_matched, self._min_window), + self._rate_matched_cap, + ) + else: + self._window = min( + max( + self._execution_fraction * self._execution_gap_ema, + self._min_window, + ), + self._max_window, + ) + return self._window + + def on_collection_end(self, collected_any_frame: bool) -> None: + self._last_non_empty_collection_end = ( + self._clock() if collected_any_frame else None + ) + + @property + def window(self) -> float: + return self._window + + @property + def execution_gap_ema(self) -> Optional[float]: + return self._execution_gap_ema + + +class CollectionPolicy: + """Per-round collection behavior for AUTO / EVERY_FRAME modes. + + Live sources are read FIFO with frames older than ``max_staleness`` + dropped (counted, optionally reported via ``on_frame_dropped``). File + sources - and sources whose properties are not yet known - are read + as-is and NEVER dropped; liveness is resolved lazily from + ``VideoSource.describe_source()`` and cached once known. + """ + + def __init__( + self, + mode: VideoProcessingMode, + max_staleness: Optional[float] = None, + on_frame_dropped: Optional[Callable[[VideoFrame], None]] = None, + window_controller: Optional[AdaptiveWindowController] = None, + ): + if mode is VideoProcessingMode.FRESHEST: + raise ValueError( + "FRESHEST mode is realised through EAGER buffer consumption " + "and does not use CollectionPolicy" + ) + self._mode = mode + if mode is VideoProcessingMode.EVERY_FRAME: + self._max_staleness = None + else: + self._max_staleness = ( + DEFAULT_MAX_STALENESS_SECONDS + if max_staleness is None + else max_staleness + ) + self._on_frame_dropped = on_frame_dropped + self._window_controller = window_controller or AdaptiveWindowController() + self._source_is_file: Dict[int, bool] = {} + self._frames_dropped_on_staleness: Dict[int, int] = {} + self._arrival_estimators: Dict[int, _SourceArrivalEstimator] = {} + + @property + def mode(self) -> VideoProcessingMode: + return self._mode + + @property + def max_staleness(self) -> Optional[float]: + return self._max_staleness + + @property + def frames_dropped_on_staleness(self) -> Dict[int, int]: + return dict(self._frames_dropped_on_staleness) + + def collection_window(self) -> float: + return self._window_controller.on_collection_start( + minimum_arrival_period=self.minimum_live_arrival_period() + ) + + def minimum_live_arrival_period(self) -> Optional[float]: + """Smallest trustworthy frame period across ACTIVE live sources. + + The fastest source is the binding constraint for the rate-matched + window: a round period above ANY source's frame period makes that + source queue unboundedly. Files never contribute (demand-paced, + always ready) and dormant sources drop out via the activity horizon + so a dying camera cannot pin the window while it flaps. + """ + now = monotonic() + periods = [ + period + for period in ( + estimator.period(now) for estimator in self._arrival_estimators.values() + ) + if period is not None + ] + if not periods: + return None + return min(periods) + + def note_collection_result(self, batch_frames: list) -> None: + self._window_controller.on_collection_end( + collected_any_frame=bool(batch_frames) + ) + + def read_frame( + self, + source_ord: int, + source: "VideoSource", + timeout: Optional[float], + ) -> Optional[VideoFrame]: + treated_as_file = self._source_treated_as_file( + source_ord=source_ord, source=source + ) + if self._max_staleness is None or treated_as_file: + frame = source.read_frame(timeout=timeout) + if frame is not None and not treated_as_file: + self._observe_arrival(source_ord=source_ord, frame=frame) + return frame + deadline = None if timeout is None else monotonic() + timeout + while True: + remaining = None if deadline is None else max(deadline - monotonic(), 0.0) + frame = source.read_frame(timeout=remaining) + if frame is None: + return None + # Staleness-drained frames feed the estimator too: they are real + # arrivals, and skipping them would bias the period estimate + # upward exactly when the pipeline is busiest. + self._observe_arrival(source_ord=source_ord, frame=frame) + frame_age = (datetime.now() - frame.frame_timestamp).total_seconds() + if frame_age <= self._max_staleness: + return frame + self._register_staleness_drop(source_ord=source_ord, frame=frame) + + def _observe_arrival(self, source_ord: int, frame: VideoFrame) -> None: + estimator = self._arrival_estimators.get(source_ord) + if estimator is None: + estimator = _SourceArrivalEstimator() + self._arrival_estimators[source_ord] = estimator + estimator.observe(frame_timestamp=frame.frame_timestamp, now=monotonic()) + + def _source_treated_as_file(self, source_ord: int, source: "VideoSource") -> bool: + cached = self._source_is_file.get(source_ord) + if cached is not None: + return cached + is_file: Optional[bool] = None + try: + source_metadata = source.describe_source() + source_properties = source_metadata.source_properties + if source_properties is not None: + is_file = source_properties.is_file + except Exception: # noqa: BLE001 - metadata probe must never break reads + is_file = None + if is_file is None: + # Liveness unknown (source still initialising) - never drop. + return True + self._source_is_file[source_ord] = is_file + return is_file + + def _register_staleness_drop(self, source_ord: int, frame: VideoFrame) -> None: + self._frames_dropped_on_staleness[source_ord] = ( + self._frames_dropped_on_staleness.get(source_ord, 0) + 1 + ) + if self._on_frame_dropped is None: + return + try: + self._on_frame_dropped(frame) + except Exception: # noqa: BLE001 - reporting must never break collection + logger.warning( + "on_frame_dropped callback raised while reporting a staleness " + "drop for source %s", + frame.source_id, + ) diff --git a/inference/core/interfaces/camera/dgpu_producer.py b/inference/core/interfaces/camera/dgpu_producer.py new file mode 100644 index 0000000000..758a87b232 --- /dev/null +++ b/inference/core/interfaces/camera/dgpu_producer.py @@ -0,0 +1,114 @@ +"""GPU-native ``VideoFrameProducer`` for discrete NVIDIA GPUs, backed by ``PyNvVideoCodec`` (NVDEC). + +Decodes a video file through the NVIDIA Video Codec SDK and yields each frame as a +``torch.Tensor`` on CUDA. DLPack wraps the decoded surface, then ``clone()`` gives the +consumer independent storage. ``SimpleDecoder`` accepts file sources. Its runtime library +closure includes ``libnvidia-encode`` and requires the NVIDIA ``video`` driver capability. +""" + +from typing import TYPE_CHECKING, Dict, Optional, Tuple + +from inference.core.interfaces.camera.entities import ( + FrameImage, + SourceProperties, + VideoFrameProducer, +) + +if TYPE_CHECKING: + import torch + +_DEFAULT_OUTPUT_COLOR_TYPE = "RGBP" + + +class PyNvVideoCodecFrameProducer(VideoFrameProducer): + """``VideoFrameProducer`` backed by dGPU NVDEC via ``PyNvVideoCodec.SimpleDecoder``.""" + + def __init__( + self, + video: str, + *, + gpu_id: int = 0, + output_color_type: str = _DEFAULT_OUTPUT_COLOR_TYPE, + ): + try: + import PyNvVideoCodec as nvc + import torch # noqa: F401 - needed at retrieve() time; imported here to fail fast + except Exception as error: # noqa: BLE001 + raise ImportError( + "PyNvVideoCodecFrameProducer requires `PyNvVideoCodec` + `torch` with a " + "working CUDA/NVDEC stack (note: PyNvVideoCodec hard-links libnvidia-encode " + "even for decode). Probe via " + "inference.core.interfaces.camera.discoverability.check_pynvvideocodec()." + ) from error + + self._source_ref = video + self._gpu_id = gpu_id + self._decoder = nvc.SimpleDecoder( + video, + gpu_id=gpu_id, + use_device_memory=True, + output_color_type=getattr(nvc.OutputColorType, output_color_type), + ) + self._iterator = iter(self._decoder) + self._last_frame = None + self._opened = True + + def isOpened(self) -> bool: + return self._opened + + def grab(self) -> bool: + try: + self._last_frame = next(self._iterator) + return True + except StopIteration: + self._last_frame = None + self._opened = False + return False + + def retrieve(self) -> Tuple[bool, Optional[FrameImage]]: + import torch + + frame = self._last_frame + self._last_frame = None + if frame is None: + return False, None + tensor = torch.from_dlpack(frame) + return True, tensor.clone() + + def initialize_source_properties(self, properties: Dict[str, float]) -> None: + return None + + def discover_source_properties(self) -> SourceProperties: + width, height, fps, total_frames = self._read_stream_metadata() + return SourceProperties( + width=width, + height=height, + total_frames=total_frames, + is_file=True, + fps=fps, + is_reconnectable=False, + timestamp_created=None, + ) + + def _read_stream_metadata(self) -> Tuple[int, int, float, int]: + try: + meta = self._decoder.get_stream_metadata() + width = int(getattr(meta, "width", 0) or 0) + height = int(getattr(meta, "height", 0) or 0) + fps = float(getattr(meta, "average_fps", 0.0) or 0.0) + total_frames = int(getattr(meta, "num_frames", 0) or 0) + except Exception: # noqa: BLE001 - metadata API differs across versions + width = height = total_frames = 0 + fps = 0.0 + if not total_frames: + try: + total_frames = len(self._decoder) + except Exception: # noqa: BLE001 + total_frames = 0 + return width, height, fps, total_frames + + def release(self) -> None: + self._last_frame = None + self._iterator = None + self._decoder = None + self._opened = False diff --git a/inference/core/interfaces/camera/discoverability.py b/inference/core/interfaces/camera/discoverability.py new file mode 100644 index 0000000000..48b04eb380 --- /dev/null +++ b/inference/core/interfaces/camera/discoverability.py @@ -0,0 +1,336 @@ +"""Runtime discovery of hardware video frame producers.""" + +import logging +from dataclasses import dataclass +from typing import Dict, List, Optional, Union + +from inference.core.interfaces.camera.entities import VideoFrameProducer + +logger = logging.getLogger(__name__) + +JETSON = "jetson" +DGPU = "dgpu" +GSTREAMER_CUDA = "gstreamer_cuda" + + +@dataclass(frozen=True) +class ProducerAvailability: + """Result of probing one hardware-decode backend in this environment.""" + + name: str + available: bool + reason: str + + +def check_jetson_gstreamer( + video: Optional[Union[str, int]] = None, + *, + require_cuda_tensor: bool = True, +) -> ProducerAvailability: + """Probe the in-repo Jetson GStreamer producer and its required elements.""" + + try: + from inference.core.interfaces.camera.jetson_producer import ( + probe_gstreamer_elements, + required_gstreamer_elements, + ) + except Exception as error: # noqa: BLE001 + return ProducerAvailability( + JETSON, False, f"Jetson GStreamer producer import failed: {error!r}" + ) + gst_ok, gst_reason = probe_gstreamer_elements( + required_gstreamer_elements(video, output_tensor=True) + ) + if not gst_ok: + return ProducerAvailability(JETSON, False, gst_reason) + try: + from inference.core.interfaces.camera.jetson_tensor_bridge import ( + jetson_tensor_bridge_available, + ) + + bridge_ok, bridge_reason = jetson_tensor_bridge_available() + except Exception as error: # noqa: BLE001 + return ProducerAvailability( + JETSON, False, f"Jetson tensor bridge probe failed: {error!r}" + ) + if not bridge_ok: + return ProducerAvailability(JETSON, False, bridge_reason) + try: + import torch + except Exception as error: # noqa: BLE001 + return ProducerAvailability(JETSON, False, f"torch import failed: {error!r}") + try: + cuda_ok = torch.cuda.is_available() + except Exception as error: # noqa: BLE001 + return ProducerAvailability( + JETSON, False, f"torch.cuda probe failed: {error!r}" + ) + if not cuda_ok: + return ProducerAvailability(JETSON, False, "torch.cuda.is_available() is False") + return ProducerAvailability(JETSON, True, "ok") + + +def check_pynvvideocodec() -> ProducerAvailability: + """Probe whether the dGPU (``PyNvVideoCodec``) producer is usable here.""" + try: + import PyNvVideoCodec # noqa: F401 + except Exception as error: # noqa: BLE001 + # Common failure: `libnvidia-encode.so.1: cannot open shared object file` โ€” the + # NVENC lib that PyNvVideoCodec hard-links even for decode. Treated as unavailable. + return ProducerAvailability( + DGPU, False, f"PyNvVideoCodec import failed: {error!r}" + ) + try: + import torch + except Exception as error: # noqa: BLE001 + return ProducerAvailability(DGPU, False, f"torch import failed: {error!r}") + try: + cuda_ok = torch.cuda.is_available() + except Exception as error: # noqa: BLE001 + return ProducerAvailability(DGPU, False, f"torch.cuda probe failed: {error!r}") + if not cuda_ok: + return ProducerAvailability(DGPU, False, "torch.cuda.is_available() is False") + return ProducerAvailability(DGPU, True, "ok") + + +def check_gstreamer_cuda( + video: Optional[Union[str, int]] = None, +) -> ProducerAvailability: + if video is not None and ( + not isinstance(video, str) + or video.startswith("/dev/video") + or video.lower().startswith("csi://") + ): + return ProducerAvailability( + GSTREAMER_CUDA, + False, + "GStreamer CUDA producer requires a URI or file path", + ) + try: + from inference.core.interfaces.camera.gstreamer_cuda_producer import ( + probe_gstreamer_cuda_elements, + required_gstreamer_cuda_elements, + ) + except Exception as error: # noqa: BLE001 + return ProducerAvailability( + GSTREAMER_CUDA, + False, + f"GStreamer CUDA producer import failed: {error!r}", + ) + gst_ok, gst_reason = probe_gstreamer_cuda_elements( + required_gstreamer_cuda_elements(video) + ) + if not gst_ok: + return ProducerAvailability(GSTREAMER_CUDA, False, gst_reason) + try: + from inference.core.interfaces.camera.gstreamer_cuda_tensor_bridge import ( + gstreamer_cuda_tensor_bridge_available, + ) + + bridge_ok, bridge_reason = gstreamer_cuda_tensor_bridge_available() + except Exception as error: # noqa: BLE001 + return ProducerAvailability( + GSTREAMER_CUDA, + False, + f"GStreamer CUDA tensor bridge probe failed: {error!r}", + ) + if not bridge_ok: + return ProducerAvailability(GSTREAMER_CUDA, False, bridge_reason) + try: + import torch + except Exception as error: # noqa: BLE001 + return ProducerAvailability( + GSTREAMER_CUDA, False, f"torch import failed: {error!r}" + ) + try: + cuda_ok = torch.cuda.is_available() + except Exception as error: # noqa: BLE001 + return ProducerAvailability( + GSTREAMER_CUDA, False, f"torch.cuda probe failed: {error!r}" + ) + if not cuda_ok: + return ProducerAvailability( + GSTREAMER_CUDA, False, "torch.cuda.is_available() is False" + ) + return ProducerAvailability(GSTREAMER_CUDA, True, "ok") + + +def available_producers( + video: Optional[Union[str, int]] = None, + *, + require_cuda_tensor: bool = True, +) -> Dict[str, ProducerAvailability]: + """Probe every hardware-decode backend; map name -> :class:`ProducerAvailability`.""" + if require_cuda_tensor and _dgpu_supports_source(video): + dgpu_availability = check_pynvvideocodec() + elif not require_cuda_tensor: + dgpu_availability = ProducerAvailability( + DGPU, False, "PyNvVideoCodec produces tensor frames" + ) + else: + dgpu_availability = ProducerAvailability( + DGPU, False, "PyNvVideoCodec SimpleDecoder requires a file source" + ) + gstreamer_cuda_availability = check_gstreamer_cuda(video) + return { + GSTREAMER_CUDA: gstreamer_cuda_availability, + JETSON: check_jetson_gstreamer( + video=video, require_cuda_tensor=require_cuda_tensor + ), + DGPU: dgpu_availability, + } + + +def build_hw_producer( + video: Union[str, int], + *, + prefer: Optional[str] = None, + output_tensor: bool = True, + **producer_kwargs, +) -> Optional[VideoFrameProducer]: + """Best-effort factory: return an instantiated GPU producer for ``video``, or ``None``. + + Resolution order: ``prefer`` (``"jetson"``/``"dgpu"``) first if given; otherwise + Jetson is preferred on ``aarch64`` Linux, dGPU elsewhere. Only backends that pass their + ``check_*`` probe are attempted, and instantiation failures fall through to the next + candidate. The producer modules are imported locally so this stays import-safe. + + Source capability is included in discovery. The PyNvVideoCodec backend accepts file + references, while the Jetson GStreamer backend accepts files, cameras, and network + streams. + """ + checks = available_producers( + video=video, + require_cuda_tensor=output_tensor, + ) + for name in _resolution_order(prefer, video): + if not checks[name].available: + continue + if name in (GSTREAMER_CUDA, DGPU): + # The onnx.gpu image's CUDA media stack is built for CC>=7.5; fail fast + # with a clear message on older GPUs instead of a cryptic runtime crash. + # (Raised here so it propagates to the caller, which logs it and falls + # back to the cv2 CPU decoder.) + _require_media_compute_capability() + try: + if name == GSTREAMER_CUDA: + from inference.core.interfaces.camera.gstreamer_cuda_producer import ( + GstreamerCudaVideoFrameProducer, + ) + + return GstreamerCudaVideoFrameProducer( + video, + output_tensor=output_tensor, + **producer_kwargs, + ) + if name == JETSON: + from inference.core.interfaces.camera.jetson_producer import ( + JetsonVideoFrameProducer, + ) + + return JetsonVideoFrameProducer( + video, + output_tensor=output_tensor, + **producer_kwargs, + ) + if name == DGPU and output_tensor: + from inference.core.interfaces.camera.dgpu_producer import ( + PyNvVideoCodecFrameProducer, + ) + + return PyNvVideoCodecFrameProducer(video, **producer_kwargs) + except Exception as error: # noqa: BLE001 - try the next candidate + # Without this line the caller only ever sees the (passing) probe + # results and the fallback to cv2 looks inexplicable. + logger.warning( + f"Constructing the '{name}' hardware decoder for source " + f"reference {video} failed: {error!r}. Trying the next " + "candidate decoder." + ) + continue + return None + + +def _resolution_order( + prefer: Optional[str], video: Optional[Union[str, int]] = None +) -> List[str]: + """Source-type-aware producer routing (per the media-decode plan). + + An explicit ``prefer`` always wins, with the remaining backends kept as + fallbacks. Otherwise the backend is chosen from (platform, source type): + + - **Jetson** (aarch64 Linux): GStreamer for every source type. Live sources + use the bridge's latest-wins slot; local files use its lossless handoff + (bridge ABI v6+), so NVDEC decode is backpressured to consumption speed + and every file frame is served - cv2 remains only the construction-failure + fallback. + - **dGPU / x86**: GStreamer for live/stream sources; PyNvVideoCodec (``dgpu``) for + local FILES (its ``SimpleDecoder`` is seekable-file only). + + ``http(s)`` and other URI schemes are treated as streams (GStreamer handles them; + PyNvVideoCodec cannot seek them). + """ + producer_names = (GSTREAMER_CUDA, JETSON, DGPU) + if prefer in producer_names: + return [prefer] + [name for name in producer_names if name != prefer] + import platform + + is_file = _is_file_source(video) + if platform.machine() == "aarch64" and platform.system() == "Linux": + return [JETSON] + # dGPU / x86: streams -> GStreamer; local files -> PyNvVideoCodec. + return [DGPU] if is_file else [GSTREAMER_CUDA] + + +def _is_file_source(video: Optional[Union[str, int]]) -> bool: + """A seekable local FILE path (routed to PyNvVideoCodec on dGPU; on Jetson + files share the GStreamer producer via its lossless handoff), as opposed to a + live/stream/camera source (routed to GStreamer). URI schemes (rtsp/http/...), + ``/dev/video*``, ``csi://`` and integer camera indices are NOT files.""" + return ( + isinstance(video, str) + and "://" not in video + and not video.startswith("/dev/video") + and not video.lower().startswith("csi://") + ) + + +def _require_media_compute_capability(minimum: tuple = (7, 5)) -> None: + """Guard the dGPU CUDA media path (GStreamer-CUDA / PyNvVideoCodec) against GPUs + below the compiled compute-capability floor. + + ``Dockerfile.onnx.gpu`` builds OpenCV-CUDA + nvcodec for CC>=7.5 only, so on + older GPUs (V100 7.0, Pascal 6.x) the nvcodec conversion kernels and cv2.cuda ops + fail at runtime with an opaque 'no kernel image' error. Detect this at producer + selection and raise a clear, actionable error; the caller (``VideoSource``) logs + it and falls back to the cv2 CPU decode path. + + Robust in non-CUDA / test environments: if CUDA is absent or the capability can't + be read, this does nothing (the ``check_*`` probe already gates on CUDA, and the + runtime would surface any real issue).""" + try: + import torch + + if not torch.cuda.is_available(): + return + capability = tuple(torch.cuda.get_device_capability()) + except Exception: # noqa: BLE001 - can't determine CC -> don't block here + return + if capability < minimum: + raise RuntimeError( + f"Hardware GPU video decode requires CUDA compute capability " + f">= {minimum[0]}.{minimum[1]}, but this GPU is " + f"{capability[0]}.{capability[1]}. The onnx.gpu image's OpenCV-CUDA + " + f"nvcodec media stack is compiled for >= {minimum[0]}.{minimum[1]} only; " + f"use the CPU (cv2) decode path on this hardware." + ) + + +def _dgpu_supports_source(video: Optional[Union[str, int]]) -> bool: + if video is None: + return True + return ( + isinstance(video, str) + and "://" not in video + and not video.startswith("/dev/video") + ) diff --git a/inference/core/interfaces/camera/entities.py b/inference/core/interfaces/camera/entities.py index 9b9973cad6..15aee2e6f3 100644 --- a/inference/core/interfaces/camera/entities.py +++ b/inference/core/interfaces/camera/entities.py @@ -2,13 +2,22 @@ from dataclasses import dataclass from datetime import datetime from enum import Enum -from typing import Callable, Dict, Optional, Tuple, Union +from typing import TYPE_CHECKING, Callable, Dict, Optional, Tuple, Union import numpy as np +if TYPE_CHECKING: + import torch + FrameTimestamp = datetime FrameID = int +# A decoded frame image is either a host numpy array (CPU โ€” the default cv2 path) or a +# torch.Tensor (typically GPU-resident, produced by a hardware decoder such as the +# Jetson / PyNvVideoCodec producers). Consumers must branch on the concrete type. The +# torch reference is a forward ref (string) so this module imports without torch present. +FrameImage = Union[np.ndarray, "torch.Tensor"] + class UpdateSeverity(Enum): """Enumeration for defining different levels of update severity. @@ -50,7 +59,8 @@ class VideoFrame: """Represents a single frame of video data. Attributes: - image (np.ndarray): The image data of the frame as a NumPy array. + image (FrameImage): The image data of the frame โ€” a host ``np.ndarray`` (CPU, + cv2 path) or a ``torch.Tensor`` (GPU, hardware-decoder path). frame_id (FrameID): A unique identifier for the frame. frame_timestamp (FrameTimestamp): The timestamp when the frame was captured. source_id (int): The index of the video_reference element which was passed to InferencePipeline for this frame @@ -60,7 +70,7 @@ class VideoFrame: comes_from_video_file (Optional[bool]): flag to determine if frame comes from video file """ - image: np.ndarray + image: FrameImage frame_id: FrameID frame_timestamp: FrameTimestamp # TODO: in next major version of inference replace `fps` with `declared_fps` @@ -85,7 +95,7 @@ class VideoFrameProducer: def grab(self) -> bool: raise NotImplementedError - def retrieve(self) -> Tuple[bool, np.ndarray]: + def retrieve(self) -> Tuple[bool, FrameImage]: raise NotImplementedError def release(self): diff --git a/inference/core/interfaces/camera/gstreamer_cuda_producer.py b/inference/core/interfaces/camera/gstreamer_cuda_producer.py new file mode 100644 index 0000000000..b82f03cd3d --- /dev/null +++ b/inference/core/interfaces/camera/gstreamer_cuda_producer.py @@ -0,0 +1,393 @@ +import ctypes +import ctypes.util +import os +from datetime import datetime, timedelta +from pathlib import Path +from typing import Dict, Iterable, Optional, Sequence, Tuple, Union +from urllib.parse import unquote, urlparse + +from inference.core.interfaces.camera.entities import ( + FrameImage, + SourceProperties, + VideoFrameProducer, +) + +_GST_RANK_PRIMARY = 256 +_NVIDIA_DECODER_RANK = _GST_RANK_PRIMARY + 100 +# A live/RTSP source that yields no frame within this window is treated as +# stalled: the native grab() raises TimeoutError instead of blocking forever, so +# VideoSource surfaces an error (reconnect / cv2 fallback) rather than deadlock +# on its state-change lock. Applies to first-frame discovery and steady-state +# consumption alike. Tune per deployment via the env var below; the value must +# be validated on-device against real RTSP connect + keyframe latency. +_DEFAULT_GRAB_TIMEOUT_NS = 15_000_000_000 +_GRAB_TIMEOUT_ENV_VAR = "ROBOFLOW_GSTREAMER_CUDA_GRAB_TIMEOUT_SECONDS" +_BASE_ELEMENTS = ("appsink", "cudaconvertscale", "queue", "uridecodebin") +_RTSP_ELEMENTS = ( + "h264parse", + "h265parse", + "rtph264depay", + "rtph265depay", + "rtspsrc", +) +_FILE_DEMUXERS = { + ".avi": "avidemux", + ".m4v": "qtdemux", + ".mkv": "matroskademux", + ".mov": "qtdemux", + ".mp4": "qtdemux", + ".webm": "matroskademux", +} +_DECODER_FACTORY_NAMES = tuple( + dict.fromkeys( + [ + "nvh264dec", + "nvh264sldec", + "nvh265dec", + "nvh265sldec", + "nvjpegdec", + "nvvp8dec", + "nvvp9dec", + "nvav1dec", + ] + + [ + f"{decoder}device{device_id}dec" + for decoder in ( + "nvh264", + "nvh265", + "nvjpeg", + "nvvp8", + "nvvp9", + "nvav1", + ) + for device_id in range(16) + ] + ) +) +_FORBIDDEN_FACTORY_NAMES = ( + "avdec_h264", + "avdec_h265", + "avdec_mjpeg", + "avdec_av1", + "avdec_vp8", + "avdec_vp9", + "cudadownload", + "cudaupload", + "jpegdec", + "videoconvert", +) + + +def _resolve_grab_timeout_ns() -> int: + raw = os.getenv(_GRAB_TIMEOUT_ENV_VAR) + if raw is None: + return _DEFAULT_GRAB_TIMEOUT_NS + try: + seconds = float(raw) + except ValueError: + return _DEFAULT_GRAB_TIMEOUT_NS + if seconds <= 0: + return _DEFAULT_GRAB_TIMEOUT_NS + return int(seconds * 1_000_000_000) + + +def probe_gstreamer_cuda_elements( + elements: Iterable[str], *, boost_ranks: bool = False +) -> Tuple[bool, str]: + """Verify the required GStreamer-CUDA elements exist (including at least one NVIDIA + decoder). When ``boost_ranks`` is True (only at producer BUILD time, not during + availability probing) the NVIDIA decoders are ranked above software decoders so + ``uridecodebin`` auto-selects them; kept OFF during probing so discovery does not + perturb decoder selection process-wide for unrelated paths (e.g. a later + cv2-GStreamer decode).""" + library_name = ctypes.util.find_library("gstreamer-1.0") + if not library_name: + library_name = "libgstreamer-1.0.so.0" + try: + gst = ctypes.CDLL(library_name) + except OSError as error: + return False, f"could not load {library_name}: {error}" + + gst.gst_init_check.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_void_p), + ] + gst.gst_init_check.restype = ctypes.c_int + gst.gst_element_factory_find.argtypes = [ctypes.c_char_p] + gst.gst_element_factory_find.restype = ctypes.c_void_p + gst.gst_plugin_feature_set_rank.argtypes = [ctypes.c_void_p, ctypes.c_uint] + gst.gst_plugin_feature_set_rank.restype = None + gst.gst_object_unref.argtypes = [ctypes.c_void_p] + gst.gst_object_unref.restype = None + + error = ctypes.c_void_p() + if not gst.gst_init_check(None, None, ctypes.byref(error)): + return False, "GStreamer initialisation failed" + + factories = {} + missing = [] + for name in sorted(set(elements)): + factory = gst.gst_element_factory_find(name.encode("utf-8")) + if not factory: + missing.append(name) + else: + factories[name] = factory + try: + if missing: + return False, f"missing GStreamer elements: {', '.join(missing)}" + + decoder_found = False + for decoder_name in _DECODER_FACTORY_NAMES: + decoder = gst.gst_element_factory_find(decoder_name.encode("utf-8")) + if decoder: + decoder_found = True + if boost_ranks: + gst.gst_plugin_feature_set_rank(decoder, _NVIDIA_DECODER_RANK) + gst.gst_object_unref(decoder) + if not decoder_found: + return False, "no NVIDIA GStreamer decoder factory is available" + return True, "ok" + finally: + for factory in factories.values(): + gst.gst_object_unref(factory) + + +def required_gstreamer_cuda_elements( + video: Optional[Union[str, int]] = None, +) -> Sequence[str]: + elements = list(_BASE_ELEMENTS) + if video is None: + return tuple(elements) + if not isinstance(video, str): + return tuple(elements) + if _is_rtsp_source(video): + elements.extend(_RTSP_ELEMENTS) + local_file_path = _local_file_path(video) + if local_file_path is not None: + demuxer = _FILE_DEMUXERS.get(Path(local_file_path).suffix.lower()) + if demuxer is not None: + elements.append(demuxer) + return tuple(elements) + + +def build_gstreamer_cuda_pipeline(video: str, *, device_id: int = 0) -> str: + is_live = _local_file_path(video) is None + queue_options = ( + "max-size-buffers=2 max-size-bytes=0 max-size-time=0 leaky=downstream" + if is_live + else "max-size-buffers=4 max-size-bytes=0 max-size-time=0" + ) + appsink_options = ( + "max-buffers=1 drop=true sync=false" + if is_live + else "max-buffers=4 drop=false sync=false" + ) + uri = _source_uri(video) + return ( + f'uridecodebin uri="{_quote_gstreamer_value(uri)}" ' + 'caps="video/x-raw(memory:CUDAMemory)" ! ' + f"queue {queue_options} ! " + f"cudaconvertscale cuda-device-id={device_id} ! " + "video/x-raw(memory:CUDAMemory),format=RGBP ! " + f"appsink name=rf_tensor_sink {appsink_options} wait-on-eos=false" + ) + + +class GstreamerCudaVideoFrameProducer(VideoFrameProducer): + def __init__( + self, + video: str, + *, + gpu_id: int = 0, + output_tensor: bool = True, + ) -> None: + if not _supports_uri_source(video): + raise TypeError("GStreamer CUDA producer requires a URI or file path") + + gst_ok, gst_reason = probe_gstreamer_cuda_elements( + required_gstreamer_cuda_elements(video), + boost_ranks=True, + ) + if not gst_ok: + raise RuntimeError(gst_reason) + + try: + import torch + except Exception as error: # noqa: BLE001 - optional runtime capability + raise ImportError("GStreamer CUDA decoding requires torch") from error + if not torch.cuda.is_available(): + raise RuntimeError("GStreamer CUDA decoding requires CUDA") + if gpu_id < 0 or gpu_id >= torch.cuda.device_count(): + raise ValueError(f"CUDA device {gpu_id} is unavailable") + + from inference.core.interfaces.camera.gstreamer_cuda_tensor_bridge import ( + NativeGstreamerCudaTensorPipeline, + gstreamer_cuda_tensor_bridge_available, + ) + + bridge_ok, bridge_reason = gstreamer_cuda_tensor_bridge_available() + if not bridge_ok: + raise RuntimeError(bridge_reason) + + self._source_ref = video + self._output_tensor = output_tensor + self._pipeline_description = build_gstreamer_cuda_pipeline( + video, device_id=gpu_id + ) + self._native_pipeline = NativeGstreamerCudaTensorPipeline( + self._pipeline_description, device_id=gpu_id + ) + self._decoder_validated = False + self._prerolled_frame_pending = False + self._cached_source_properties: Optional[SourceProperties] = None + self._grab_timeout_ns = _resolve_grab_timeout_ns() + self._closed = False + self._eos = False + + @property + def pipeline(self) -> str: + return self._pipeline_description + + def isOpened(self) -> bool: + return not self._closed and not self._eos + + def grab(self) -> bool: + if self._closed or self._eos: + return False + if self._prerolled_frame_pending: + self._prerolled_frame_pending = False + return True + grabbed = self._native_pipeline.grab(timeout_ns=self._grab_timeout_ns) + if grabbed and not self._decoder_validated: + self._validate_pipeline() + if not grabbed: + self._eos = True + return grabbed + + def retrieve(self) -> Tuple[bool, Optional[FrameImage]]: + if self._closed or self._eos: + return False, None + self._prerolled_frame_pending = False + rgb_tensor = self._native_pipeline.retrieve() + if self._output_tensor: + return True, rgb_tensor + return True, _rgb_tensor_to_bgr_numpy(rgb_tensor) + + def initialize_source_properties(self, properties: Dict[str, float]) -> None: + return None + + def discover_source_properties(self) -> SourceProperties: + if self._cached_source_properties is not None: + return self._cached_source_properties + if not self.grab(): + raise RuntimeError( + "GStreamer CUDA pipeline did not produce source metadata" + ) + self._prerolled_frame_pending = True + frame_info = self._native_pipeline.frame_info() + fps = ( + frame_info.fps_numerator / frame_info.fps_denominator + if frame_info.fps_denominator > 0 + else 0.0 + ) + total_frames = ( + int(frame_info.duration_ns * fps / 1_000_000_000) + if frame_info.duration_ns > 0 and fps > 0 + else 0 + ) + local_path = _local_file_path(self._source_ref) + is_file = local_path is not None + timestamp_created = None + if is_file and fps > 0 and total_frames > 0 and os.path.isfile(local_path): + file_length_seconds = total_frames / fps + last_modified = datetime.fromtimestamp(os.path.getmtime(local_path)) + timestamp_created = last_modified - timedelta(seconds=file_length_seconds) + properties = SourceProperties( + width=frame_info.width, + height=frame_info.height, + total_frames=total_frames, + is_file=is_file, + fps=fps, + is_reconnectable=not is_file, + timestamp_created=timestamp_created, + ) + self._cached_source_properties = properties + return properties + + def interrupt(self) -> None: + if self._closed or self._eos: + return + self._eos = True + self._native_pipeline.interrupt() + + def release(self) -> None: + if self._closed: + return + self._closed = True + self._prerolled_frame_pending = False + self._native_pipeline.close() + + @property + def tensor_bridge_stats(self) -> Dict[str, int]: + return self._native_pipeline.stats() + + def _validate_pipeline(self) -> None: + if not self._native_pipeline.has_factory("cudaconvertscale"): + raise RuntimeError( + "GStreamer pipeline did not instantiate cudaconvertscale" + ) + forbidden = [ + factory + for factory in _FORBIDDEN_FACTORY_NAMES + if self._native_pipeline.has_factory(factory) + ] + if forbidden: + raise RuntimeError( + "GStreamer CUDA pipeline instantiated forbidden elements: " + + ", ".join(forbidden) + ) + if not any( + self._native_pipeline.has_factory(factory) + for factory in _DECODER_FACTORY_NAMES + ): + raise RuntimeError( + "GStreamer CUDA pipeline did not instantiate an NVIDIA decoder" + ) + self._decoder_validated = True + + +def _source_uri(video: str) -> str: + local_path = _local_file_path(video) + if local_path is not None: + return Path(local_path).resolve().as_uri() + return video + + +def _local_file_path(video: str) -> Optional[str]: + if video.startswith("file://"): + parsed = urlparse(video) + return unquote(parsed.path) + if "://" in video: + return None + return video + + +def _is_rtsp_source(video: str) -> bool: + return video.lower().startswith(("rtsp://", "rtsps://")) + + +def _supports_uri_source(video: object) -> bool: + return ( + isinstance(video, str) + and not video.startswith("/dev/video") + and not video.lower().startswith("csi://") + ) + + +def _rgb_tensor_to_bgr_numpy(rgb_tensor): + return rgb_tensor.permute(1, 2, 0).flip(-1).contiguous().cpu().numpy() + + +def _quote_gstreamer_value(value: str) -> str: + return value.replace("\\", "\\\\").replace('"', '\\"') diff --git a/inference/core/interfaces/camera/gstreamer_cuda_tensor_bridge.py b/inference/core/interfaces/camera/gstreamer_cuda_tensor_bridge.py new file mode 100644 index 0000000000..4d65c2647a --- /dev/null +++ b/inference/core/interfaces/camera/gstreamer_cuda_tensor_bridge.py @@ -0,0 +1,297 @@ +import ctypes +import ctypes.util +import os +import threading +import time +from dataclasses import dataclass +from typing import Dict, Optional, Tuple + +_ERROR_CAPACITY = 1024 +# Upper bound on a single native pull. grab() loops on this window so that an +# interrupt() from another thread, a GStreamer bus error, and the caller's +# overall deadline are re-checked between pulls instead of the pull blocking +# forever inside gst_app_sink_try_pull_sample (e.g. an unreachable RTSP source +# or a mid-stream drop that never posts EOS). +_GRAB_POLL_TIMEOUT_NS = 250_000_000 +_DEFAULT_LIBRARY_PATH = "/opt/roboflow/lib/libroboflow_gstreamer_cuda_tensor.so.1" + + +class _FrameInfo(ctypes.Structure): + _fields_ = [ + ("width", ctypes.c_uint32), + ("height", ctypes.c_uint32), + ("fps_numerator", ctypes.c_int32), + ("fps_denominator", ctypes.c_int32), + ("duration_ns", ctypes.c_int64), + ] + + +class _BridgeStats(ctypes.Structure): + _fields_ = [ + ("frames", ctypes.c_uint64), + ("cuda_maps", ctypes.c_uint64), + ("host_pixel_maps", ctypes.c_uint64), + ("host_to_device_copies", ctypes.c_uint64), + ("device_to_host_copies", ctypes.c_uint64), + ("stream_synchronizations", ctypes.c_uint64), + ("active_leases", ctypes.c_uint64), + ("last_channel_stride", ctypes.c_int64), + ("last_row_stride", ctypes.c_int64), + ] + + +@dataclass(frozen=True) +class GstreamerCudaFrameInfo: + width: int + height: int + fps_numerator: int + fps_denominator: int + duration_ns: int + + +def gstreamer_cuda_tensor_bridge_available() -> Tuple[bool, str]: + try: + library = _load_bridge_library() + version = library.rf_gstreamer_cuda_tensor_bridge_version() + except Exception as error: # noqa: BLE001 - runtime capability probe + return False, f"GStreamer CUDA tensor bridge is unavailable: {error!r}" + if version != b"3": + return False, f"Unsupported GStreamer CUDA tensor bridge version: {version!r}" + return True, "ok" + + +class NativeGstreamerCudaTensorPipeline: + def __init__(self, pipeline: str, *, device_id: int = 0) -> None: + # Created before anything that can raise so __del__ -> close() can always + # acquire it. This lock serializes interrupt()/close() so that a native + # release() (which unrefs sink/pipeline and frees the handle) can never + # run concurrently with a native interrupt() that dereferences the same + # handle. The native release() explicitly relies on the Python wrapper + # providing this serialization. + self._lifecycle_lock = threading.Lock() + self._handle = None + self._library = _load_bridge_library() + error = ctypes.create_string_buffer(_ERROR_CAPACITY) + self._handle = self._library.rf_gstreamer_cuda_pipeline_create( + pipeline.encode("utf-8"), + device_id, + error, + len(error), + ) + if not self._handle: + raise RuntimeError(_decode_error(error)) + + def grab(self, timeout_ns: Optional[int] = None) -> bool: + """Poll the native pipeline for the next frame. + + ``timeout_ns`` is the overall deadline: when it elapses without a frame + the source is treated as stalled and ``TimeoutError`` is raised so the + caller can surface an error / reconnect instead of blocking forever. + ``None`` preserves the historic unbounded wait, but still returns + promptly on interrupt()/EOS/bus errors because the pull is chunked. + """ + error = ctypes.create_string_buffer(_ERROR_CAPACITY) + deadline_ns: Optional[int] = None + if timeout_ns is not None: + deadline_ns = time.monotonic_ns() + int(timeout_ns) + while True: + self._ensure_open() + poll_ns = _GRAB_POLL_TIMEOUT_NS + if deadline_ns is not None: + remaining_ns = deadline_ns - time.monotonic_ns() + if remaining_ns <= 0: + raise TimeoutError( + "GStreamer CUDA pipeline produced no frame within " + f"{timeout_ns} ns; the source appears stalled or " + "unreachable" + ) + poll_ns = min(poll_ns, remaining_ns) + status = self._library.rf_gstreamer_cuda_pipeline_grab( + self._handle, + poll_ns, + error, + len(error), + ) + if status < 0: + raise RuntimeError(_decode_error(error)) + if status == 1: + return True + if status == 0: + # End of stream, or interrupt() flipped the native flag from + # another thread. + return False + # status == 2: the poll window expired while the stream is still + # live. Loop so interrupt(), bus errors and the overall deadline are + # re-checked instead of blocking on a single unbounded pull. + + def retrieve(self): + self._ensure_open() + error = ctypes.create_string_buffer(_ERROR_CAPACITY) + managed_tensor = self._library.rf_gstreamer_cuda_pipeline_retrieve( + self._handle, + error, + len(error), + ) + if not managed_tensor: + raise RuntimeError(_decode_error(error)) + + capsule = _create_dlpack_capsule(managed_tensor) + try: + import torch + + tensor = torch.utils.dlpack.from_dlpack(capsule) + except Exception: + self._library.rf_gstreamer_cuda_dlpack_delete(managed_tensor) + raise + if not tensor.is_cuda or tensor.dtype != torch.uint8 or tensor.ndim != 3: + raise RuntimeError( + "GStreamer CUDA bridge returned an invalid CUDA uint8 CHW tensor" + ) + return tensor + + def frame_info(self) -> GstreamerCudaFrameInfo: + self._ensure_open() + info = _FrameInfo() + error = ctypes.create_string_buffer(_ERROR_CAPACITY) + status = self._library.rf_gstreamer_cuda_pipeline_get_frame_info( + self._handle, + ctypes.byref(info), + error, + len(error), + ) + if status < 0: + raise RuntimeError(_decode_error(error)) + return GstreamerCudaFrameInfo( + width=info.width, + height=info.height, + fps_numerator=info.fps_numerator, + fps_denominator=info.fps_denominator, + duration_ns=info.duration_ns, + ) + + def has_factory(self, factory_name: str) -> bool: + self._ensure_open() + return bool( + self._library.rf_gstreamer_cuda_pipeline_has_factory( + self._handle, factory_name.encode("utf-8") + ) + ) + + def stats(self) -> Dict[str, int]: + self._ensure_open() + stats = _BridgeStats() + status = self._library.rf_gstreamer_cuda_pipeline_get_stats( + self._handle, ctypes.byref(stats) + ) + if status < 0: + raise RuntimeError("Could not read GStreamer CUDA bridge statistics") + return {name: int(getattr(stats, name)) for name, _ in stats._fields_} + + def interrupt(self) -> None: + with self._lifecycle_lock: + handle = getattr(self, "_handle", None) + if handle: + status = self._library.rf_gstreamer_cuda_pipeline_interrupt(handle) + if status < 0: + raise RuntimeError("Could not interrupt GStreamer CUDA pipeline") + + def close(self) -> None: + with self._lifecycle_lock: + handle = getattr(self, "_handle", None) + if handle: + self._library.rf_gstreamer_cuda_pipeline_release(handle) + self._handle = None + + def _ensure_open(self) -> None: + if not getattr(self, "_handle", None): + raise RuntimeError("GStreamer CUDA tensor pipeline is closed") + + def __del__(self) -> None: + try: + self.close() + except Exception: # noqa: BLE001 - interpreter shutdown + pass + + +def _load_bridge_library(): + configured_path = os.getenv("ROBOFLOW_GSTREAMER_CUDA_TENSOR_BRIDGE_LIBRARY") + candidates = [configured_path] if configured_path else [] + candidates.extend( + [ + _DEFAULT_LIBRARY_PATH, + ctypes.util.find_library("roboflow_gstreamer_cuda_tensor"), + ] + ) + last_error = None + for candidate in candidates: + if not candidate: + continue + try: + library = ctypes.CDLL(candidate) + _configure_library(library) + return library + except OSError as error: + last_error = error + if last_error is not None: + raise last_error + raise OSError("GStreamer CUDA tensor bridge library was not found") + + +def _configure_library(library) -> None: + library.rf_gstreamer_cuda_tensor_bridge_version.argtypes = [] + library.rf_gstreamer_cuda_tensor_bridge_version.restype = ctypes.c_char_p + library.rf_gstreamer_cuda_pipeline_create.argtypes = [ + ctypes.c_char_p, + ctypes.c_int, + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.rf_gstreamer_cuda_pipeline_create.restype = ctypes.c_void_p + library.rf_gstreamer_cuda_pipeline_grab.argtypes = [ + ctypes.c_void_p, + ctypes.c_uint64, + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.rf_gstreamer_cuda_pipeline_grab.restype = ctypes.c_int + library.rf_gstreamer_cuda_pipeline_get_frame_info.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(_FrameInfo), + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.rf_gstreamer_cuda_pipeline_get_frame_info.restype = ctypes.c_int + library.rf_gstreamer_cuda_pipeline_has_factory.argtypes = [ + ctypes.c_void_p, + ctypes.c_char_p, + ] + library.rf_gstreamer_cuda_pipeline_has_factory.restype = ctypes.c_int + library.rf_gstreamer_cuda_pipeline_retrieve.argtypes = [ + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.rf_gstreamer_cuda_pipeline_retrieve.restype = ctypes.c_void_p + library.rf_gstreamer_cuda_pipeline_get_stats.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(_BridgeStats), + ] + library.rf_gstreamer_cuda_pipeline_get_stats.restype = ctypes.c_int + library.rf_gstreamer_cuda_pipeline_interrupt.argtypes = [ctypes.c_void_p] + library.rf_gstreamer_cuda_pipeline_interrupt.restype = ctypes.c_int + library.rf_gstreamer_cuda_dlpack_delete.argtypes = [ctypes.c_void_p] + library.rf_gstreamer_cuda_dlpack_delete.restype = None + library.rf_gstreamer_cuda_pipeline_release.argtypes = [ctypes.c_void_p] + library.rf_gstreamer_cuda_pipeline_release.restype = None + + +def _create_dlpack_capsule(managed_tensor): + py_capsule_new = ctypes.pythonapi.PyCapsule_New + py_capsule_new.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + py_capsule_new.restype = ctypes.py_object + return py_capsule_new(managed_tensor, b"dltensor", None) + + +def _decode_error(error_buffer) -> str: + message = error_buffer.value.decode("utf-8", errors="replace") + return message or "GStreamer CUDA tensor bridge failed" diff --git a/inference/core/interfaces/camera/jetson_producer.py b/inference/core/interfaces/camera/jetson_producer.py new file mode 100644 index 0000000000..e7e36ba4d0 --- /dev/null +++ b/inference/core/interfaces/camera/jetson_producer.py @@ -0,0 +1,610 @@ +"""Jetson video capture with NVIDIA GStreamer and CUDA tensor output.""" + +import ctypes +import ctypes.util +import os +from datetime import datetime, timedelta +from pathlib import Path +from typing import Dict, Iterable, Optional, Sequence, Tuple, Union +from urllib.parse import parse_qsl, unquote, urlparse + +from inference.core.interfaces.camera.entities import ( + FrameImage, + SourceProperties, + VideoFrameProducer, +) + +_GST_RANK_PRIMARY = 256 +_NVIDIA_DECODER_RANK = _GST_RANK_PRIMARY + 100 +# A live/RTSP source that yields no frame within this window is treated as +# stalled: the native grab() raises TimeoutError instead of blocking forever, so +# VideoSource surfaces an error (reconnect / cv2 fallback) rather than deadlock +# on its state-change lock. Applies to first-frame discovery and steady-state +# consumption alike. Tune per deployment via the env var below; the value must +# be validated on-device against real RTSP connect + keyframe latency. +_DEFAULT_GRAB_TIMEOUT_NS = 15_000_000_000 +_GRAB_TIMEOUT_ENV_VAR = "ROBOFLOW_JETSON_GRAB_TIMEOUT_SECONDS" + + +def _resolve_grab_timeout_ns() -> int: + raw = os.getenv(_GRAB_TIMEOUT_ENV_VAR) + if raw is None: + return _DEFAULT_GRAB_TIMEOUT_NS + try: + seconds = float(raw) + except ValueError: + return _DEFAULT_GRAB_TIMEOUT_NS + if seconds <= 0: + return _DEFAULT_GRAB_TIMEOUT_NS + return int(seconds * 1_000_000_000) + + +_RTSP_CODEC_ENV_VAR = "ROBOFLOW_RTSP_VIDEO_CODEC" +_RTSP_PROTOCOLS_ENV_VAR = "ROBOFLOW_RTSP_PROTOCOLS" +_RTSP_LATENCY_ENV_VAR = "ROBOFLOW_RTSP_LATENCY_MS" +_RTSP_TLS_VALIDATION_FLAGS_ENV_VAR = "ROBOFLOW_RTSP_TLS_VALIDATION_FLAGS" +_DEFAULT_RTSP_PROTOCOLS = "tcp" +_DEFAULT_RTSP_LATENCY_MS = 200 +_RTSP_VIDEO_CODECS = ("h264", "h265") + +_COMMON_ELEMENTS = ( + "appsink", + "nvvidconv", + "queue", +) +_URI_DECODE_ELEMENTS = ( + "decodebin", + "h264parse", + "h265parse", + "jpegparse", + "nvjpegdec", + "nvv4l2decoder", + "uridecodebin", +) +_RTSP_ELEMENTS = ( + "rtspsrc", +) +_SRTP_ELEMENTS = ("capssetter", "srtpdec") +_SOFTWARE_DECODER_ELEMENTS = ( + "avdec_h264", + "avdec_h265", + "avdec_mjpeg", + "jpegdec", + "libde265dec", + "openh264dec", +) +_FILE_DEMUXERS = { + ".avi": "avidemux", + ".m4v": "qtdemux", + ".mkv": "matroskademux", + ".mov": "qtdemux", + ".mp4": "qtdemux", + ".webm": "matroskademux", +} + + +def probe_gstreamer_elements( + elements: Iterable[str], *, boost_ranks: bool = False +) -> Tuple[bool, str]: + """Verify the required GStreamer element factories exist. + + When ``boost_ranks`` is True (only at producer BUILD time, not during availability + probing) the NVIDIA HW decoders (``nvv4l2decoder``/``nvjpegdec``) are ranked above + software decoders so ``decodebin``/``uridecodebin`` auto-selects them. Keeping the + boost OFF during probing means merely discovering this backend no longer perturbs + decoder selection process-wide for unrelated paths (e.g. a later cv2-GStreamer + decode) - the boost is applied only when a Jetson pipeline is actually built.""" + + library_name = ctypes.util.find_library("gstreamer-1.0") + if not library_name: + library_name = "libgstreamer-1.0.so.0" + try: + gst = ctypes.CDLL(library_name) + except OSError as error: + return False, f"could not load {library_name}: {error}" + + gst.gst_init_check.argtypes = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_void_p), + ] + gst.gst_init_check.restype = ctypes.c_int + gst.gst_element_factory_find.argtypes = [ctypes.c_char_p] + gst.gst_element_factory_find.restype = ctypes.c_void_p + gst.gst_plugin_feature_set_rank.argtypes = [ctypes.c_void_p, ctypes.c_uint] + gst.gst_plugin_feature_set_rank.restype = None + gst.gst_object_unref.argtypes = [ctypes.c_void_p] + gst.gst_object_unref.restype = None + + error = ctypes.c_void_p() + if not gst.gst_init_check(None, None, ctypes.byref(error)): + return False, "GStreamer initialisation failed" + + factories = {} + missing = [] + for name in sorted(set(elements)): + factory = gst.gst_element_factory_find(name.encode("utf-8")) + if not factory: + missing.append(name) + else: + factories[name] = factory + try: + if missing: + return False, f"missing GStreamer elements: {', '.join(missing)}" + + if boost_ranks: + for decoder_name in ("nvv4l2decoder", "nvjpegdec"): + decoder = factories.get(decoder_name) + if decoder: + gst.gst_plugin_feature_set_rank(decoder, _NVIDIA_DECODER_RANK) + return True, "ok" + finally: + for factory in factories.values(): + gst.gst_object_unref(factory) + + +def required_gstreamer_elements( + video: Optional[Union[str, int]] = None, + *, + output_tensor: bool = False, +) -> Sequence[str]: + """Return the element set needed by a source, or the common baseline.""" + + elements = list(_COMMON_ELEMENTS) + if video is None: + return tuple(elements + list(_URI_DECODE_ELEMENTS)) + if _is_csi_source(video): + return tuple(elements + ["nvarguscamerasrc"]) + if _is_v4l2_source(video): + return tuple( + elements + + [ + "decodebin", + "h264parse", + "h265parse", + "jpegparse", + "nvjpegdec", + "nvv4l2decoder", + "v4l2src", + ] + ) + if _is_rtsp_source(video): + # RTSP uses an explicit rtspsrc ! depay ! parse ! nvv4l2decoder chain + # (no uridecodebin autoplugging), so only those elements are required. + codec = _rtsp_video_codec() + srtp_elements = list(_SRTP_ELEMENTS) if _rtsp_uses_srtp(video) else [] + return tuple( + elements + + srtp_elements + + [ + f"rtp{codec}depay", + f"{codec}parse", + f"{codec}timestamper", + "nvv4l2decoder", + ] + + list(_RTSP_ELEMENTS) + ) + elements.extend(_URI_DECODE_ELEMENTS) + local_file_path = _local_file_path(video) + if local_file_path is not None: + demuxer = _FILE_DEMUXERS.get(Path(local_file_path).suffix.lower()) + if demuxer is not None: + elements.append(demuxer) + return tuple(elements) + + +def build_gstreamer_pipeline( + video: Union[str, int], + *, + output_tensor: bool = False, + rtsp_tls_validation_flags: Optional[int] = None, +) -> str: + """Build a Jetson GStreamer pipeline ending in an NVMM appsink.""" + + is_live = _is_live_source(video) + sink = _build_sink(is_live=is_live) + if _is_csi_source(video): + sensor_id = _csi_sensor_id(video) + return ( + f"nvarguscamerasrc sensor-id={sensor_id} ! " + "video/x-raw(memory:NVMM),format=NV12 ! " + f"{sink}" + ) + if _is_v4l2_source(video): + device = _v4l2_device(video) + return ( + f'v4l2src device="{_quote_gstreamer_value(device)}" ! ' + f"decodebin ! {sink}" + ) + + if _is_rtsp_source(video): + # Explicit video-only chain (the jetson-utils shape) instead of + # uridecodebin: autoplugging an RTSP source decodes EVERY track, and a + # camera that muxes audio (e.g. A-Law) poisons the bus with a + # missing-decoder error โ€” fatal at startup when it races preroll, and a + # mid-run grab() failure otherwise. A codec-specific depayloader only + # ever links the video stream, so the audio track is never plugged. + # protocols defaults to tcp: RTP-over-UDP needs raised kernel buffers + # and a NAT-free path that containers typically lack, and a failed UDP + # SETUP can make cameras drop the whole control connection. + # + # Further jetson-utils parity (v4 bridge): + # - the queue buffers COMPRESSED data before the depayloader (cheap, + # non-leaky) instead of leaking decoded NVMM frames after the decoder; + # - no nvvidconv: the decoder's NV12 NVMM output goes straight to the + # appsink and the bridge converts NV12->RGB CHW in CUDA, removing the + # per-frame VIC pass and its extra buffer pool; + # - enable-max-performance keeps the decoder clocks pinned; + # - the appsink never accumulates (the bridge's new-sample callback + # drains it on the streaming thread), so a small non-dropping queue + # is enough. + codec = _rtsp_video_codec() + tls_validation_flags = _rtsp_tls_validation_flags( + explicit_flags=rtsp_tls_validation_flags + ) + source = ( + f'rtspsrc location="{_quote_gstreamer_value(str(video))}" ' + f"protocols={_rtsp_protocols()} latency={_rtsp_latency_ms()}" + " drop-on-latency=true teardown-timeout=0" + f"{tls_validation_flags} ! " + "application/x-rtp,media=video ! " + "queue ! " + ) + if _rtsp_uses_srtp(video): + # UniFi and other SDES endpoints advertise the master key through + # the RTP caps' a-crypto field. The native bridge rewrites the + # named capssetter's CAPS event before srtpdec sees the first + # packet; key material never crosses Python or appears in the + # launch string. + source += ( + "capssetter name=rf_srtp_caps caps=application/x-srtp " + "join=false replace=false ! srtpdec ! " + ) + depayloader = f"rtp{codec}depay" + if codec == "h264": + # rtph264depay can request a fresh keyframe over RTCP when startup + # begins mid-GOP. GStreamer 1.24's rtph265depay does not expose + # these properties, so keep its launch string portable. + depayloader += " request-keyframe=true wait-for-keyframe=true" + return ( + source + # RTSP commonly supplies PTS without DTS. Reconstruct DTS from the + # codec's SPS reordering metadata before NVDEC. + + f"{depayloader} ! " + f"{codec}parse ! {codec}timestamper ! " + "nvv4l2decoder enable-max-performance=1 ! " + "video/x-raw(memory:NVMM),format=NV12 ! " + "appsink name=rf_tensor_sink max-buffers=4 drop=false sync=false " + "wait-on-eos=false" + ) + + uri = _source_uri(video) + return ( + f'uridecodebin uri="{_quote_gstreamer_value(uri)}" ' + 'caps="video/x-raw(memory:NVMM)" ! ' + f"{sink}" + ) + + +def _rtsp_video_codec() -> str: + codec = os.getenv(_RTSP_CODEC_ENV_VAR, _RTSP_VIDEO_CODECS[0]).strip().lower() + if codec not in _RTSP_VIDEO_CODECS: + raise ValueError( + f"Unsupported RTSP video codec {codec!r} in {_RTSP_CODEC_ENV_VAR} " + f"(supported: {', '.join(_RTSP_VIDEO_CODECS)})" + ) + return codec + + +def _rtsp_uses_srtp(video: Union[str, int]) -> bool: + """Return whether an RTSP URL explicitly requests encrypted RTP media.""" + + if not isinstance(video, str): + return False + for key, value in parse_qsl(urlparse(video).query, keep_blank_values=True): + if key.lower() != "enablesrtp": + continue + return value.strip().lower() not in {"0", "false", "no", "off"} + return False + + +def _rtsp_protocols() -> str: + return os.getenv(_RTSP_PROTOCOLS_ENV_VAR, _DEFAULT_RTSP_PROTOCOLS).strip() or ( + _DEFAULT_RTSP_PROTOCOLS + ) + + +def _rtsp_latency_ms() -> int: + raw = os.getenv(_RTSP_LATENCY_ENV_VAR) + if raw is None: + return _DEFAULT_RTSP_LATENCY_MS + try: + latency = int(raw) + except ValueError: + return _DEFAULT_RTSP_LATENCY_MS + return latency if latency >= 0 else _DEFAULT_RTSP_LATENCY_MS + + +def _rtsp_tls_validation_flags( + explicit_flags: Optional[int] = None, +) -> str: + """Return an explicit rtspsrc TLS-validation setting when requested. + + ``0`` disables certificate validation for cameras with private/self-signed + certificates. Keeping this unset by default avoids weakening RTSPS + validation for normal deployments. + """ + + raw = ( + explicit_flags + if explicit_flags is not None + else os.getenv(_RTSP_TLS_VALIDATION_FLAGS_ENV_VAR) + ) + if raw is None or (isinstance(raw, str) and not raw.strip()): + return "" + try: + flags = int(raw) + except (TypeError, ValueError) as error: + raise ValueError( + f"{_RTSP_TLS_VALIDATION_FLAGS_ENV_VAR} must be a non-negative integer" + ) from error + if flags < 0: + raise ValueError( + f"{_RTSP_TLS_VALIDATION_FLAGS_ENV_VAR} must be a non-negative integer" + ) + return f" tls-validation-flags={flags}" + + +def _build_sink(is_live: bool) -> str: + queue_options = ( + "max-size-buffers=2 max-size-bytes=0 max-size-time=0 leaky=downstream" + if is_live + else "max-size-buffers=4 max-size-bytes=0 max-size-time=0" + ) + appsink_options = ( + "max-buffers=1 drop=true sync=false" + if is_live + else "max-buffers=4 drop=false sync=false" + ) + return ( + f"queue {queue_options} ! " + "nvvidconv ! video/x-raw(memory:NVMM),format=RGBA ! " + f"appsink name=rf_tensor_sink {appsink_options} wait-on-eos=false" + ) + + +class JetsonVideoFrameProducer(VideoFrameProducer): + """Decode Jetson file/camera/RTSP sources with NVIDIA GStreamer elements.""" + + def __init__( + self, + video: Union[str, int], + *, + output_tensor: bool = True, + tensor_device: str = "cuda", + pin_host_memory: bool = True, + rtsp_tls_validation_flags: Optional[int] = None, + ): + gst_ok, gst_reason = probe_gstreamer_elements( + required_gstreamer_elements(video, output_tensor=True), + boost_ranks=True, + ) + if not gst_ok: + raise RuntimeError(gst_reason) + + self._source_ref = video + self._output_tensor = output_tensor + self._pipeline = build_gstreamer_pipeline( + video, + output_tensor=True, + rtsp_tls_validation_flags=rtsp_tls_validation_flags, + ) + self._decoder_validated = not _source_requires_decoder(video) + self._prerolled_frame_pending = False + self._cached_source_properties: Optional[SourceProperties] = None + self._grab_timeout_ns = _resolve_grab_timeout_ns() + del pin_host_memory + + import torch + + from inference.core.interfaces.camera.jetson_tensor_bridge import ( + NativeJetsonTensorPipeline, + jetson_tensor_bridge_available, + ) + + device = torch.device(tensor_device) + if device.type != "cuda" or not torch.cuda.is_available(): + raise RuntimeError("Jetson decoding requires an available CUDA device") + device_id = ( + torch.cuda.current_device() if device.index is None else device.index + ) + bridge_ok, bridge_reason = jetson_tensor_bridge_available() + if not bridge_ok: + raise RuntimeError(bridge_reason) + # Files (and other non-live sources) use the bridge's lossless + # handoff: decode is backpressured to consumption speed and every + # frame is served. Live sources keep the latest-wins slot. + self._native_pipeline = NativeJetsonTensorPipeline( + self._pipeline, + device_id=device_id, + lossless_handoff=not _is_live_source(video), + ) + self._closed = False + self._eos = False + + @property + def pipeline(self) -> str: + """Pipeline string exposed for diagnostics and on-device tests.""" + + return self._pipeline + + def isOpened(self) -> bool: + return not self._closed and not self._eos + + def grab(self) -> bool: + if self._closed or self._eos: + return False + if self._prerolled_frame_pending: + self._prerolled_frame_pending = False + return True + grabbed = self._native_pipeline.grab(timeout_ns=self._grab_timeout_ns) + if grabbed and not self._decoder_validated: + self._validate_hardware_decoder() + if not grabbed: + self._eos = True + return grabbed + + def retrieve(self) -> Tuple[bool, Optional[FrameImage]]: + if self._closed or self._eos: + return False, None + self._prerolled_frame_pending = False + rgb_tensor = self._native_pipeline.retrieve() + if self._output_tensor: + return True, rgb_tensor + return True, _rgb_tensor_to_bgr_numpy(rgb_tensor) + + def initialize_source_properties(self, properties: Dict[str, float]) -> None: + return None + + def discover_source_properties(self) -> SourceProperties: + if self._cached_source_properties is not None: + return self._cached_source_properties + if not self.grab(): + raise RuntimeError("Jetson pipeline did not produce source metadata") + self._prerolled_frame_pending = True + frame_info = self._native_pipeline.frame_info() + width = frame_info.width + height = frame_info.height + fps = ( + frame_info.fps_numerator / frame_info.fps_denominator + if frame_info.fps_denominator > 0 + else 0.0 + ) + total_frames = ( + int(frame_info.duration_ns * fps / 1_000_000_000) + if frame_info.duration_ns > 0 and fps > 0 + else 0 + ) + local_path = _local_file_path(self._source_ref) + is_file = local_path is not None + timestamp_created = None + if is_file and fps > 0 and total_frames > 0 and os.path.isfile(local_path): + file_length_seconds = total_frames / fps + last_modified = datetime.fromtimestamp(os.path.getmtime(local_path)) + timestamp_created = last_modified - timedelta(seconds=file_length_seconds) + properties = SourceProperties( + width=width, + height=height, + total_frames=total_frames, + is_file=is_file, + fps=fps, + is_reconnectable=not is_file, + timestamp_created=timestamp_created, + ) + self._cached_source_properties = properties + return properties + + def interrupt(self) -> None: + if self._closed or self._eos: + return + self._eos = True + self._native_pipeline.interrupt() + + def release(self) -> None: + if self._closed: + return + self._closed = True + self._prerolled_frame_pending = False + self._native_pipeline.close() + + @property + def tensor_bridge_stats(self) -> Dict[str, int]: + return self._native_pipeline.stats() + + def _validate_hardware_decoder(self) -> None: + if any( + self._native_pipeline.has_factory(factory) + for factory in ("nvv4l2decoder", "nvjpegdec") + ): + self._decoder_validated = True + return + software_decoders = [ + factory + for factory in _SOFTWARE_DECODER_ELEMENTS + if self._native_pipeline.has_factory(factory) + ] + if software_decoders: + raise RuntimeError( + "Jetson pipeline instantiated software decoders: " + + ", ".join(software_decoders) + ) + if _is_v4l2_source(self._source_ref): + self._decoder_validated = True + return + raise RuntimeError("Jetson pipeline did not instantiate an NVIDIA decoder") + + +def _source_uri(video: Union[str, int]) -> str: + local_path = _local_file_path(video) + if local_path is not None: + return Path(local_path).resolve().as_uri() + return str(video) + + +def _local_file_path(video: Union[str, int]) -> Optional[str]: + if not isinstance(video, str): + return None + if video.startswith("file://"): + parsed = urlparse(video) + return unquote(parsed.path) + if "://" in video: + return None + return video + + +def _is_rtsp_source(video: Union[str, int]) -> bool: + # rtspt:// / rtspst:// are rtspsrc's force-TCP variants of rtsp:// / rtsps://. + return isinstance(video, str) and video.lower().startswith( + ("rtsp://", "rtsps://", "rtspt://", "rtspst://") + ) + + +def _is_csi_source(video: Union[str, int]) -> bool: + return isinstance(video, str) and video.lower().startswith("csi://") + + +def _csi_sensor_id(video: Union[str, int]) -> int: + try: + return int(str(video).split("://", 1)[1] or 0) + except ValueError as error: + raise ValueError(f"Invalid CSI source reference: {video!r}") from error + + +def _is_v4l2_source(video: Union[str, int]) -> bool: + return isinstance(video, int) or ( + isinstance(video, str) and video.startswith("/dev/video") + ) + + +def _v4l2_device(video: Union[str, int]) -> str: + return f"/dev/video{video}" if isinstance(video, int) else video + + +def _is_live_source(video: Union[str, int]) -> bool: + return ( + _is_csi_source(video) + or _is_v4l2_source(video) + or _local_file_path(video) is None + ) + + +def _source_requires_decoder(video: Union[str, int]) -> bool: + return not _is_csi_source(video) + + +def _rgb_tensor_to_bgr_numpy(rgb_tensor): + return rgb_tensor.permute(1, 2, 0).flip(-1).contiguous().cpu().numpy() + + +def _quote_gstreamer_value(value: str) -> str: + return value.replace("\\", "\\\\").replace('"', '\\"') diff --git a/inference/core/interfaces/camera/jetson_tensor_bridge.py b/inference/core/interfaces/camera/jetson_tensor_bridge.py new file mode 100644 index 0000000000..cd36834ff8 --- /dev/null +++ b/inference/core/interfaces/camera/jetson_tensor_bridge.py @@ -0,0 +1,339 @@ +import ctypes +import ctypes.util +import os +import threading +import time +from dataclasses import dataclass +from typing import Dict, Optional, Tuple + +_ERROR_CAPACITY = 1024 +# Upper bound on a single native pull. grab() loops on this window so that an +# interrupt() from another thread, a GStreamer bus error, and the caller's +# overall deadline are re-checked between pulls instead of the pull blocking +# forever inside gst_app_sink_try_pull_sample (e.g. an unreachable RTSP source +# or a mid-stream drop that never posts EOS). +_GRAB_POLL_TIMEOUT_NS = 250_000_000 +_DEFAULT_LIBRARY_PATH = "/opt/roboflow/lib/libroboflow_jetson_tensor.so.1" + + +class _FrameInfo(ctypes.Structure): + _fields_ = [ + ("width", ctypes.c_uint32), + ("height", ctypes.c_uint32), + ("fps_numerator", ctypes.c_int32), + ("fps_denominator", ctypes.c_int32), + ("duration_ns", ctypes.c_int64), + ] + + +class _BridgeStats(ctypes.Structure): + _fields_ = [ + ("frames", ctypes.c_uint64), + ("descriptor_maps", ctypes.c_uint64), + ("host_pixel_maps", ctypes.c_uint64), + ("host_to_device_copies", ctypes.c_uint64), + ("device_to_host_copies", ctypes.c_uint64), + ("array_flatten_copies", ctypes.c_uint64), + ("conversion_kernels", ctypes.c_uint64), + ("nvmm_frames", ctypes.c_uint64), + ("frames_dropped_by_consumer", ctypes.c_uint64), + ("last_nvbuf_memory_type", ctypes.c_int32), + ("last_egl_frame_type", ctypes.c_int32), + ("last_egl_color_format", ctypes.c_int32), + # ABI v5: per-phase conversion timing (total + max + # ns per phase) and the count of distinct decoder dmabuf fds. Field + # order and widths mirror RfBridgeStats in jetson_tensor_bridge.cu. + ("egl_map_ns", ctypes.c_uint64), + ("egl_map_max_ns", ctypes.c_uint64), + ("cuda_register_ns", ctypes.c_uint64), + ("cuda_register_max_ns", ctypes.c_uint64), + ("texture_create_ns", ctypes.c_uint64), + ("texture_create_max_ns", ctypes.c_uint64), + ("kernel_launch_ns", ctypes.c_uint64), + ("kernel_launch_max_ns", ctypes.c_uint64), + ("sync_ns", ctypes.c_uint64), + ("sync_max_ns", ctypes.c_uint64), + ("cleanup_ns", ctypes.c_uint64), + ("cleanup_max_ns", ctypes.c_uint64), + ("unique_buffer_fds", ctypes.c_uint64), + # ABI v7: per-fd EGL registration cache effectiveness (hits skip the + # per-frame map/register/texture-create/unregister sequence). + ("egl_cache_hits", ctypes.c_uint64), + ("egl_cache_misses", ctypes.c_uint64), + ] + + +@dataclass(frozen=True) +class JetsonFrameInfo: + width: int + height: int + fps_numerator: int + fps_denominator: int + duration_ns: int + + +def jetson_tensor_bridge_available() -> Tuple[bool, str]: + try: + library = _load_bridge_library() + version = library.rf_jetson_tensor_bridge_version() + except Exception as error: # noqa: BLE001 - runtime capability probe + return False, f"Jetson tensor bridge is unavailable: {error!r}" + # v8 = consumer-thread conversion (the ABI is unchanged from v7). + # v7 = per-fd EGL registration cache; RfBridgeStats gained + # egl_cache_hits/egl_cache_misses (struct layout change) on top of v6's + # lossless_handoff create() parameter, so older .so versions must be + # refused. Requiring v8 also prevents pairing this wrapper with the + # JP6.2-deadlocking streaming-callback implementation. + if version != b"8": + return False, f"Unsupported Jetson tensor bridge version: {version!r}" + return True, "ok" + + +class NativeJetsonTensorPipeline: + def __init__( + self, + pipeline: str, + *, + device_id: int = 0, + lossless_handoff: bool = False, + ) -> None: + # ``lossless_handoff`` selects the file-mode handoff: a bounded + # blocking FIFO in the native bridge that backpressures decode so no + # frame is ever dropped (required for every-frame video-file + # processing). Live sources keep the latest-wins slot (False). + # + # Created before anything that can raise so __del__ -> close() can always + # acquire it. This lock serializes interrupt()/close() so that a native + # release() (which unrefs sink/pipeline and frees the handle) can never + # run concurrently with a native interrupt() that dereferences the same + # handle. The native release() explicitly relies on the Python wrapper + # providing this serialization. + self._lifecycle_lock = threading.Lock() + self._handle = None + self._library = _load_bridge_library() + error = ctypes.create_string_buffer(_ERROR_CAPACITY) + self._handle = self._library.rf_jetson_pipeline_create( + pipeline.encode("utf-8"), + device_id, + 1 if lossless_handoff else 0, + error, + len(error), + ) + if not self._handle: + raise RuntimeError(_decode_error(error)) + + def grab(self, timeout_ns: Optional[int] = None) -> bool: + """Poll the native pipeline for the next frame. + + ``timeout_ns`` is the overall deadline: when it elapses without a frame + the source is treated as stalled and ``TimeoutError`` is raised so the + caller can surface an error / reconnect instead of blocking forever. + ``None`` preserves the historic unbounded wait, but still returns + promptly on interrupt()/EOS/bus errors because the pull is chunked. + """ + error = ctypes.create_string_buffer(_ERROR_CAPACITY) + deadline_ns: Optional[int] = None + if timeout_ns is not None: + deadline_ns = time.monotonic_ns() + int(timeout_ns) + while True: + self._ensure_open() + poll_ns = _GRAB_POLL_TIMEOUT_NS + if deadline_ns is not None: + remaining_ns = deadline_ns - time.monotonic_ns() + if remaining_ns <= 0: + raise TimeoutError( + "Jetson pipeline produced no frame within " + f"{timeout_ns} ns; the source appears stalled or " + "unreachable" + ) + poll_ns = min(poll_ns, remaining_ns) + status = self._library.rf_jetson_pipeline_grab( + self._handle, + poll_ns, + error, + len(error), + ) + if status < 0: + raise RuntimeError(_decode_error(error)) + if status == 1: + return True + if status == 0: + # End of stream, or interrupt() flipped the native flag from + # another thread. + return False + # status == 2: the poll window expired while the stream is still + # live. Loop so interrupt(), bus errors and the overall deadline are + # re-checked instead of blocking on a single unbounded pull. + + def retrieve(self): + self._ensure_open() + error = ctypes.create_string_buffer(_ERROR_CAPACITY) + managed_tensor = self._library.rf_jetson_pipeline_retrieve( + self._handle, + error, + len(error), + ) + if not managed_tensor: + raise RuntimeError(_decode_error(error)) + + capsule = _create_dlpack_capsule(managed_tensor) + try: + import torch + + tensor = torch.utils.dlpack.from_dlpack(capsule) + except Exception: + self._library.rf_jetson_dlpack_delete(managed_tensor) + raise + if not tensor.is_cuda or tensor.dtype != torch.uint8 or tensor.ndim != 3: + raise RuntimeError( + "Jetson tensor bridge returned an invalid CUDA uint8 CHW tensor" + ) + return tensor + + def frame_info(self) -> JetsonFrameInfo: + self._ensure_open() + info = _FrameInfo() + error = ctypes.create_string_buffer(_ERROR_CAPACITY) + status = self._library.rf_jetson_pipeline_get_frame_info( + self._handle, + ctypes.byref(info), + error, + len(error), + ) + if status < 0: + raise RuntimeError(_decode_error(error)) + return JetsonFrameInfo( + width=info.width, + height=info.height, + fps_numerator=info.fps_numerator, + fps_denominator=info.fps_denominator, + duration_ns=info.duration_ns, + ) + + def has_factory(self, factory_name: str) -> bool: + self._ensure_open() + return bool( + self._library.rf_jetson_pipeline_has_factory( + self._handle, factory_name.encode("utf-8") + ) + ) + + def stats(self) -> Dict[str, int]: + self._ensure_open() + stats = _BridgeStats() + status = self._library.rf_jetson_pipeline_get_stats( + self._handle, ctypes.byref(stats) + ) + if status < 0: + raise RuntimeError("Could not read Jetson tensor bridge statistics") + return {name: int(getattr(stats, name)) for name, _ in stats._fields_} + + def interrupt(self) -> None: + with self._lifecycle_lock: + handle = getattr(self, "_handle", None) + if handle: + status = self._library.rf_jetson_pipeline_interrupt(handle) + if status < 0: + raise RuntimeError("Could not interrupt Jetson tensor pipeline") + + def close(self) -> None: + with self._lifecycle_lock: + handle = getattr(self, "_handle", None) + if handle: + self._library.rf_jetson_pipeline_release(handle) + self._handle = None + + def _ensure_open(self) -> None: + if not getattr(self, "_handle", None): + raise RuntimeError("Jetson tensor pipeline is closed") + + def __del__(self) -> None: + try: + self.close() + except Exception: # noqa: BLE001 - interpreter shutdown + pass + + +def _load_bridge_library(): + configured_path = os.getenv("ROBOFLOW_JETSON_TENSOR_BRIDGE_LIBRARY") + candidates = [configured_path] if configured_path else [] + candidates.extend( + [ + _DEFAULT_LIBRARY_PATH, + ctypes.util.find_library("roboflow_jetson_tensor"), + ] + ) + last_error = None + for candidate in candidates: + if not candidate: + continue + try: + library = ctypes.CDLL(candidate) + _configure_library(library) + return library + except OSError as error: + last_error = error + if last_error is not None: + raise last_error + raise OSError("Jetson tensor bridge library was not found") + + +def _configure_library(library) -> None: + library.rf_jetson_tensor_bridge_version.argtypes = [] + library.rf_jetson_tensor_bridge_version.restype = ctypes.c_char_p + library.rf_jetson_pipeline_create.argtypes = [ + ctypes.c_char_p, + ctypes.c_int, + ctypes.c_int, + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.rf_jetson_pipeline_create.restype = ctypes.c_void_p + library.rf_jetson_pipeline_grab.argtypes = [ + ctypes.c_void_p, + ctypes.c_uint64, + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.rf_jetson_pipeline_grab.restype = ctypes.c_int + library.rf_jetson_pipeline_get_frame_info.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(_FrameInfo), + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.rf_jetson_pipeline_get_frame_info.restype = ctypes.c_int + library.rf_jetson_pipeline_has_factory.argtypes = [ + ctypes.c_void_p, + ctypes.c_char_p, + ] + library.rf_jetson_pipeline_has_factory.restype = ctypes.c_int + library.rf_jetson_pipeline_retrieve.argtypes = [ + ctypes.c_void_p, + ctypes.c_char_p, + ctypes.c_size_t, + ] + library.rf_jetson_pipeline_retrieve.restype = ctypes.c_void_p + library.rf_jetson_pipeline_get_stats.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(_BridgeStats), + ] + library.rf_jetson_pipeline_get_stats.restype = ctypes.c_int + library.rf_jetson_pipeline_interrupt.argtypes = [ctypes.c_void_p] + library.rf_jetson_pipeline_interrupt.restype = ctypes.c_int + library.rf_jetson_dlpack_delete.argtypes = [ctypes.c_void_p] + library.rf_jetson_dlpack_delete.restype = None + library.rf_jetson_pipeline_release.argtypes = [ctypes.c_void_p] + library.rf_jetson_pipeline_release.restype = None + + +def _create_dlpack_capsule(managed_tensor): + py_capsule_new = ctypes.pythonapi.PyCapsule_New + py_capsule_new.argtypes = [ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p] + py_capsule_new.restype = ctypes.py_object + return py_capsule_new(managed_tensor, b"dltensor", None) + + +def _decode_error(error_buffer) -> str: + message = error_buffer.value.decode("utf-8", errors="replace") + return message or "Jetson tensor bridge failed" diff --git a/inference/core/interfaces/camera/utils.py b/inference/core/interfaces/camera/utils.py index ccbf2cf0ab..ac18879ac0 100644 --- a/inference/core/interfaces/camera/utils.py +++ b/inference/core/interfaces/camera/utils.py @@ -19,6 +19,7 @@ from inference.core import logger from inference.core.env import RESTART_ATTEMPT_DELAY +from inference.core.interfaces.camera.collection_policy import CollectionPolicy from inference.core.interfaces.camera.entities import VideoFrame from inference.core.interfaces.camera.exceptions import ( EndOfStreamError, @@ -174,6 +175,49 @@ def retrieve_frames_from_sources( self._last_batch_yielded_time = datetime.now() return batch_frames + def retrieve_frames_from_sources_with_policy( + self, + collection_policy: CollectionPolicy, + ) -> Optional[List[VideoFrame]]: + """Policy-driven sibling of `retrieve_frames_from_sources`. + + Mirrors the legacy loop (stop signal, inactive sources, EOS + registration, reconnection-thread joins, the + `_last_batch_yielded_time`-anchored budget) but the per-round budget + comes from the policy's self-tuning window and per-source reads go + through the policy (bounded-staleness FIFO for live sources, plain + reads for files). + """ + batch_frames = [] + window = collection_policy.collection_window() + batch_timeout_moment = self._last_batch_yielded_time + timedelta(seconds=window) + for source_ord, (source, source_should_reconnect) in enumerate( + zip(self._video_sources.all_sources, self._video_sources.allow_reconnection) + ): + if self._external_should_stop(): + self.join_all_reconnection_threads(include_not_finished=True) + collection_policy.note_collection_result(batch_frames=[]) + return None + if self._is_source_inactive(source_ord=source_ord): + continue + batch_time_left = max( + (batch_timeout_moment - datetime.now()).total_seconds(), 0.0 + ) + try: + frame = collection_policy.read_frame( + source_ord=source_ord, + source=source, + timeout=batch_time_left, + ) + if frame is not None: + batch_frames.append(frame) + except EndOfStreamError: + self._register_end_of_stream(source_ord=source_ord) + self.join_all_reconnection_threads() + self._last_batch_yielded_time = datetime.now() + collection_policy.note_collection_result(batch_frames=batch_frames) + return batch_frames + def all_sources_ended(self) -> bool: return len(self._ended_sources) >= len(self._video_sources.all_sources) @@ -246,6 +290,7 @@ def multiplex_videos( on_reconnection_error: Callable[ [Optional[int], SourceConnectionError], None ] = log_error, + collection_policy: Optional[CollectionPolicy] = None, ) -> Generator[List[VideoFrame], None, None]: """ Function that is supposed to provide a generator over frames from multiple video sources. It is capable to @@ -284,6 +329,10 @@ def multiplex_videos( on_reconnection_error (Callable[[Optional[int], SourceConnectionError], None]): Function that will be called whenever source cannot re-connect after disconnection. First parameter is source_id, second is connection error instance. + collection_policy (Optional[CollectionPolicy]): When given, batch collection is driven by the policy + (self-tuning collection window + bounded-staleness FIFO reads - see + `inference.core.interfaces.camera.collection_policy`) and `batch_collection_timeout` is ignored. + When `None` (default) - the legacy collection behavior is used, unchanged. Returns Generator[List[VideoFrame], None, None]: allowing to iterate through frames from multiple video sources. @@ -310,6 +359,7 @@ def multiplex_videos( batch_collection_timeout=batch_collection_timeout, should_stop=should_stop, on_reconnection_error=on_reconnection_error, + collection_policy=collection_policy, ) if max_fps is None: yield from generator @@ -329,6 +379,7 @@ def _multiplex_videos( batch_collection_timeout: Optional[float], should_stop: Callable[[], bool], on_reconnection_error: Callable[[Optional[int], SourceConnectionError], None], + collection_policy: Optional[CollectionPolicy] = None, ) -> Generator[List[VideoFrame], None, None]: sources_manager = VideoSourcesManager.init( video_sources=video_sources, @@ -336,9 +387,14 @@ def _multiplex_videos( on_reconnection_error=on_reconnection_error, ) while not sources_manager.all_sources_ended(): - batch_frames = sources_manager.retrieve_frames_from_sources( - batch_collection_timeout=batch_collection_timeout, - ) + if collection_policy is not None: + batch_frames = sources_manager.retrieve_frames_from_sources_with_policy( + collection_policy=collection_policy, + ) + else: + batch_frames = sources_manager.retrieve_frames_from_sources( + batch_collection_timeout=batch_collection_timeout, + ) if batch_frames is None: break if len(batch_frames) > 0: diff --git a/inference/core/interfaces/camera/video_source.py b/inference/core/interfaces/camera/video_source.py index 9241a5f7f9..021579e4ab 100644 --- a/inference/core/interfaces/camera/video_source.py +++ b/inference/core/interfaces/camera/video_source.py @@ -19,6 +19,8 @@ DEFAULT_BUFFER_SIZE, DEFAULT_MAXIMUM_ADAPTIVE_FRAMES_DROPPED_IN_ROW, DEFAULT_MINIMUM_ADAPTIVE_MODE_SAMPLES, + DISABLE_GSTREAMER_VIDEO_SOURCES, + ENABLE_TENSOR_DATA_REPRESENTATION, RUNS_ON_JETSON, ) from inference.core.interfaces.camera.entities import ( @@ -31,7 +33,6 @@ ) from inference.core.interfaces.camera.exceptions import ( EndOfStreamError, - SourceConnectionError, StreamOperationNotAllowedError, ) from inference.core.interfaces.camera.stream_error_classifier import ( @@ -208,8 +209,67 @@ def _is_test_pattern_reference(video: Union[str, int]) -> bool: ) +def _build_default_producer( + stream_reference: Union[str, int], + *, + output_tensor: bool = False, + producer_options: Optional[Dict[str, object]] = None, +) -> VideoFrameProducer: + """Pick the decoder for a plain (non-callable, non-test-pattern) source reference. + + When ``ENABLE_TENSOR_DATA_REPRESENTATION`` is on, hardware decoders are discovered + and verified at runtime. ``output_tensor`` selects the representation requested by + the consumer. Construction failures fall back to the general cv2 decoder. + """ + if not ENABLE_TENSOR_DATA_REPRESENTATION: + logger.debug( + "Using legacy decoder for source " f"reference: {stream_reference}" + ) + return _create_video_frame_producer(stream_reference) + # Local import: the discoverability layer pulls optional GPU-decode deps lazily. + from inference.core.interfaces.camera.discoverability import ( + available_producers, + build_hw_producer, + ) + + producer = None + try: + producer = build_hw_producer( + stream_reference, + output_tensor=output_tensor, + **(producer_options or {}), + ) + except ( + Exception + ) as error: # noqa: BLE001 - decoder selection must never break startup + logger.warning( + "Initialising a hardware decoder for source reference " + f"{stream_reference} raised: " + f"{error!r}. Falling back to the cv2 CPU decoder." + ) + if producer is not None: + logger.info( + "Selected hardware decoder " + f"'{type(producer).__name__}' for source reference: {stream_reference}" + ) + return producer + probe_reasons = { + name: availability.reason + for name, availability in available_producers( + video=stream_reference, + require_cuda_tensor=output_tensor, + ).items() + } + logger.warning( + "No hardware decoder is " + f"usable for source reference {stream_reference} (probes: {probe_reasons}). " + "Falling back to the cv2 CPU decoder." + ) + return CV2VideoFrameProducer(stream_reference) + + def _create_video_frame_producer(video: Union[str, int]) -> VideoFrameProducer: - if isinstance(video, str): + if isinstance(video, str) and not DISABLE_GSTREAMER_VIDEO_SOURCES: from inference.core.interfaces.camera.gstreamer_rtsp_producer import ( GStreamerRtspVideoFrameProducer, gstreamer_rtsp_capture_available, @@ -241,6 +301,8 @@ def init( video_source_properties: Optional[Dict[str, float]] = None, source_id: Optional[int] = None, desired_fps: Optional[Union[float, int]] = None, + allow_tensor_frames: bool = False, + video_source_options: Optional[Dict[str, object]] = None, ): """ This class is meant to represent abstraction over video sources - both video files and @@ -372,7 +434,9 @@ def init( buffer_consumption_strategy=buffer_consumption_strategy, video_consumer=video_consumer, video_source_properties=video_source_properties, + video_source_options=video_source_options, source_id=source_id, + allow_tensor_frames=allow_tensor_frames, ) def __init__( @@ -384,6 +448,8 @@ def __init__( video_consumer: "VideoConsumer", video_source_properties: Optional[Dict[str, float]], source_id: Optional[int], + allow_tensor_frames: bool = False, + video_source_options: Optional[Dict[str, object]] = None, ): self._stream_reference = stream_reference self._video: Optional[VideoFrameProducer] = None @@ -398,7 +464,9 @@ def __init__( self._stream_consumption_thread: Optional[Thread] = None self._state_change_lock = Lock() self._video_source_properties = video_source_properties or {} + self._video_source_options = video_source_options or {} self._source_id = source_id + self._allow_tensor_frames = allow_tensor_frames self._last_frame_timestamp: int = time.time_ns() self._fps: Optional[float] = None self._is_file: Optional[bool] = None @@ -640,6 +708,7 @@ def _restart( def _start(self) -> None: self._change_state(target_state=StreamState.INITIALISING) + uses_default_producer = False try: if callable(self._stream_reference): self._video = self._stream_reference() @@ -650,20 +719,33 @@ def _start(self) -> None: self._video = TestPatternStreamProducer() else: - self._video = _create_video_frame_producer(self._stream_reference) - if not self._video.isOpened(): - source_reference = self._stream_reference - if callable(source_reference): - source_reference = str(self._stream_reference) - raise wrap_source_connection_error( - build_source_connection_error_message( - source_reference=str(source_reference), - underlying_error=self._video.connection_error_message(), - ), - source_reference=str(source_reference), + uses_default_producer = True + producer_kwargs = { + "output_tensor": self._allow_tensor_frames, + } + if self._video_source_options: + producer_kwargs["producer_options"] = self._video_source_options + self._video = _build_default_producer( + self._stream_reference, + **producer_kwargs, + ) + try: + self._initialise_selected_video() + except Exception as hardware_error: + can_fall_back_to_cv2 = ( + uses_default_producer + and type(self._video) is not CV2VideoFrameProducer + ) + if not can_fall_back_to_cv2: + raise + self._release_video() + logger.warning( + "Hardware video source initialisation failed for " + f"{self._stream_reference}: {hardware_error!r}. " + "Falling back to the cv2 decoder." ) - self._video.initialize_source_properties(self._video_source_properties) - self._source_properties = self._video.discover_source_properties() + self._video = CV2VideoFrameProducer(self._stream_reference) + self._initialise_selected_video() self._video_consumer.reset(source_properties=self._source_properties) if self._source_properties.is_file: self._set_file_mode_consumption_strategies() @@ -677,8 +759,41 @@ def _start(self) -> None: # discovery fails) leaves the source in ERROR, not INITIALISING, so # recovery (which keys off ERROR) fires. Re-raise for the caller. self._change_state(target_state=StreamState.ERROR) + self._release_video() raise + def _initialise_selected_video(self) -> None: + if not self._video.isOpened(): + source_reference = self._stream_reference + if callable(source_reference): + source_reference = str(self._stream_reference) + raise wrap_source_connection_error( + build_source_connection_error_message( + source_reference=str(source_reference), + underlying_error=self._video.connection_error_message(), + ), + source_reference=str(source_reference), + ) + self._video.initialize_source_properties(self._video_source_properties) + self._source_properties = self._video.discover_source_properties() + + def _interrupt_video(self) -> None: + interrupt = getattr(self._video, "interrupt", None) + if not callable(interrupt): + return + try: + interrupt() + except Exception as error: # noqa: BLE001 + logger.warning(f"Could not interrupt video source: {error}") + + def _release_video(self) -> None: + if self._video is None: + return + try: + self._video.release() + except Exception as error: # noqa: BLE001 + logger.warning(f"Could not release video source: {error}") + def _terminate( self, wait_on_frames_consumption: bool, purge_frames_buffer: bool ) -> None: @@ -689,6 +804,7 @@ def _terminate( if purge_frames_buffer: _ = get_from_queue(queue=self._frames_buffer, timeout=0.0, purge=True) if self._stream_consumption_thread is not None: + self._interrupt_video() self._stream_consumption_thread.join() if wait_on_frames_consumption: self._frames_buffer.join() @@ -750,7 +866,7 @@ def _consume_video(self) -> None: if not success: break self._frames_buffer.put(POISON_PILL) - self._video.release() + self._release_video() self._change_state(target_state=StreamState.ENDED) send_video_source_status_update( severity=UpdateSeverity.INFO, @@ -761,6 +877,7 @@ def _consume_video(self) -> None: logger.info(f"Video consumption finished") except Exception as error: self._change_state(target_state=StreamState.ERROR) + self._release_video() # Emit a poison pill so a consumer blocked on read_frame gets # end-of-stream instead of timing out forever; ERROR is # restart-eligible, so reconnection fires. diff --git a/inference/core/interfaces/http/orjson_utils.py b/inference/core/interfaces/http/orjson_utils.py index bf870e6723..c2550223e4 100644 --- a/inference/core/interfaces/http/orjson_utils.py +++ b/inference/core/interfaces/http/orjson_utils.py @@ -8,9 +8,34 @@ from inference.core.entities.responses.inference import InferenceResponse from inference.core.utils.function import deprecated from inference.core.utils.image_utils import ImageType -from inference.core.workflows.core_steps.common.serializers import ( - serialize_wildcard_kind, -) + + +def _resolve_wildcard_serializer(): + """Pick the wildcard serialiser at call time (not import time). + + The tensor-native path yields torch (often CUDA) tensors that must be moved to + CPU during serialisation before results cross the stream-manager process + boundary via a multiprocessing queue: pickling a live CUDA tensor relies on + CUDA IPC, which is unsupported on Jetson/Tegra and fails with "CUDA error: + invalid argument". The tensor serialiser calls .detach().cpu(); the numpy one + passes tensor-native objects through untouched. + + Resolution is deferred to call time because a load-time ``if FLAG: import`` in + this module binds unreliably depending on whether ``inference.core.env`` was + imported before this module; by the time a workflow result is serialised the + flag is settled, and sys.modules caches the import after the first call. + """ + from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION + + if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialize_wildcard_kind, + ) + else: + from inference.core.workflows.core_steps.common.serializers import ( + serialize_wildcard_kind, + ) + return serialize_wildcard_kind class ORJSONResponseBytes(ORJSONResponse): @@ -88,6 +113,7 @@ def serialise_single_workflow_result_element( if excluded_fields is None: excluded_fields = [] excluded_fields = set(excluded_fields) + serialize_wildcard_kind = _resolve_wildcard_serializer() serialised_result = {} for key, value in result_element.items(): if key in excluded_fields: diff --git a/inference/core/interfaces/stream/inference_pipeline.py b/inference/core/interfaces/stream/inference_pipeline.py index 9006b2e574..b146ce4ff5 100644 --- a/inference/core/interfaces/stream/inference_pipeline.py +++ b/inference/core/interfaces/stream/inference_pipeline.py @@ -19,12 +19,20 @@ DEFAULT_BUFFER_SIZE, DISABLE_PREPROC_AUTO_ORIENT, ENABLE_FRAME_DROP_ON_VIDEO_FILE_RATE_LIMITING, + ENABLE_TENSOR_DATA_REPRESENTATION, ENABLE_WORKFLOWS_PROFILING, MAX_ACTIVE_MODELS, PREDICTIONS_QUEUE_SIZE, WORKFLOWS_PROFILER_BUFFER_SIZE, ) from inference.core.exceptions import CannotInitialiseModelError, MissingApiKeyError +from inference.core.interfaces.camera.collection_policy import ( + FRESHEST_MODE_BATCH_COLLECTION_TIMEOUT, + STALENESS_DROP_CAUSE, + CollectionPolicy, + VideoProcessingMode, + resolve_video_processing_mode, +) from inference.core.interfaces.camera.entities import ( StatusUpdate, UpdateSeverity, @@ -33,6 +41,7 @@ ) from inference.core.interfaces.camera.utils import multiplex_videos from inference.core.interfaces.camera.video_source import ( + FRAME_DROPPED_EVENT, BufferConsumptionStrategy, BufferFillingStrategy, VideoSource, @@ -49,6 +58,7 @@ ) from inference.core.interfaces.stream.sinks import active_learning_sink, multi_sink from inference.core.interfaces.stream.utils import ( + VideoSourceOptions, on_pipeline_end, prepare_video_sources, ) @@ -116,10 +126,13 @@ def init( ] = None, active_learning_target_dataset: Optional[str] = None, batch_collection_timeout: Optional[float] = None, + video_processing_mode: Optional[Union[str, VideoProcessingMode]] = None, + max_staleness: Optional[float] = None, sink_mode: SinkMode = SinkMode.ADAPTIVE, predictions_queue_size: int = PREDICTIONS_QUEUE_SIZE, decoding_buffer_size: int = DEFAULT_BUFFER_SIZE, exec_session_id: Optional[str] = None, + video_source_options: Optional[VideoSourceOptions] = None, ) -> "InferencePipeline": """ This class creates the abstraction for making inferences from Roboflow models against video stream. @@ -215,12 +228,27 @@ def init( as list of configs. Then the list must be of length of `video_reference` and may also contain None values to denote that specific source should remain not configured. Example valid properties are: {"frame_width": 1920, "frame_height": 1080, "fps": 30.0} + video_source_options (Optional[VideoSourceOptions]): Optional + producer-specific settings. A single dictionary applies to all video + sources; a list must align with `video_reference` and may contain None + for sources that need no special configuration. active_learning_target_dataset (Optional[str]): Parameter to be used when Active Learning data registration should happen against different dataset than the one pointed by model_id batch_collection_timeout (Optional[float]): Parameter of multiplex_videos(...) dictating how long process to grab frames from multiple sources can wait for batch to be filled before yielding already collected frames. Please set this value in PRODUCTION to avoid performance drops when specific sources shows unstable latency. Visit `multiplex_videos(...)` for more information about multiplexing process. + video_processing_mode (Optional[Union[str, VideoProcessingMode]]): High-level intent for live + multi-source consumption: "auto" (FIFO with a staleness budget and self-tuning collection + window), "every_frame" (strict FIFO) or "freshest" (legacy latest-wins with a small fixed + collection timeout). Defaults to "auto" when ENABLE_TENSOR_DATA_REPRESENTATION is set, + otherwise the legacy collection behavior is preserved unchanged; pass "legacy" to force + the legacy behavior explicitly (the escape hatch from the flag-driven default). File + sources always keep every-frame semantics regardless of mode. See + `inference.core.interfaces.camera.collection_policy` for details. + max_staleness (Optional[float]): Staleness budget (seconds) of "auto" mode - live frames older + than this are dropped (reported as FRAME_DROPPED status updates with cause + STALENESS_BUDGET_EXCEEDED) instead of being served late. Default: 0.5. sink_mode (SinkMode): Parameter that controls how video frames and predictions will be passed to sink handler. With SinkMode.SEQUENTIAL - each frame and prediction triggers separate call for sink, in case of SinkMode.BATCH - list of frames and predictions will be provided to sink, always aligned @@ -314,7 +342,10 @@ def init( source_buffer_filling_strategy=source_buffer_filling_strategy, source_buffer_consumption_strategy=source_buffer_consumption_strategy, video_source_properties=video_source_properties, + video_source_options=video_source_options, batch_collection_timeout=batch_collection_timeout, + video_processing_mode=video_processing_mode, + max_staleness=max_staleness, sink_mode=sink_mode, predictions_queue_size=predictions_queue_size, decoding_buffer_size=decoding_buffer_size, @@ -340,10 +371,13 @@ def init_with_yolo_world( max_detections: Optional[int] = None, video_source_properties: Optional[Dict[str, float]] = None, batch_collection_timeout: Optional[float] = None, + video_processing_mode: Optional[Union[str, VideoProcessingMode]] = None, + max_staleness: Optional[float] = None, sink_mode: SinkMode = SinkMode.ADAPTIVE, predictions_queue_size: int = PREDICTIONS_QUEUE_SIZE, decoding_buffer_size: int = DEFAULT_BUFFER_SIZE, exec_session_id: Optional[str] = None, + video_source_options: Optional[VideoSourceOptions] = None, ) -> "InferencePipeline": """ This class creates the abstraction for making inferences from YoloWorld against video stream. @@ -400,10 +434,25 @@ def init_with_yolo_world( as list of configs. Then the list must be of length of `video_reference` and may also contain None values to denote that specific source should remain not configured. Example valid properties are: {"frame_width": 1920, "frame_height": 1080, "fps": 30.0} + video_source_options (Optional[VideoSourceOptions]): Optional + producer-specific settings. A single dictionary applies to all video + sources; a list must align with `video_reference` and may contain None + for sources that need no special configuration. batch_collection_timeout (Optional[float]): Parameter of multiplex_videos(...) dictating how long process to grab frames from multiple sources can wait for batch to be filled before yielding already collected frames. Please set this value in PRODUCTION to avoid performance drops when specific sources shows unstable latency. Visit `multiplex_videos(...)` for more information about multiplexing process. + video_processing_mode (Optional[Union[str, VideoProcessingMode]]): High-level intent for live + multi-source consumption: "auto" (FIFO with a staleness budget and self-tuning collection + window), "every_frame" (strict FIFO) or "freshest" (legacy latest-wins with a small fixed + collection timeout). Defaults to "auto" when ENABLE_TENSOR_DATA_REPRESENTATION is set, + otherwise the legacy collection behavior is preserved unchanged; pass "legacy" to force + the legacy behavior explicitly (the escape hatch from the flag-driven default). File + sources always keep every-frame semantics regardless of mode. See + `inference.core.interfaces.camera.collection_policy` for details. + max_staleness (Optional[float]): Staleness budget (seconds) of "auto" mode - live frames older + than this are dropped (reported as FRAME_DROPPED status updates with cause + STALENESS_BUDGET_EXCEEDED) instead of being served late. Default: 0.5. sink_mode (SinkMode): Parameter that controls how video frames and predictions will be passed to sink handler. With SinkMode.SEQUENTIAL - each frame and prediction triggers separate call for sink, in case of SinkMode.BATCH - list of frames and predictions will be provided to sink, always aligned @@ -466,7 +515,10 @@ def init_with_yolo_world( source_buffer_filling_strategy=source_buffer_filling_strategy, source_buffer_consumption_strategy=source_buffer_consumption_strategy, video_source_properties=video_source_properties, + video_source_options=video_source_options, batch_collection_timeout=batch_collection_timeout, + video_processing_mode=video_processing_mode, + max_staleness=max_staleness, sink_mode=sink_mode, predictions_queue_size=predictions_queue_size, decoding_buffer_size=decoding_buffer_size, @@ -500,6 +552,8 @@ def init_with_workflow( cancel_thread_pool_tasks_on_exit: bool = True, video_metadata_input_name: str = "video_metadata", batch_collection_timeout: Optional[float] = None, + video_processing_mode: Optional[Union[str, VideoProcessingMode]] = None, + max_staleness: Optional[float] = None, profiling_directory: str = "./inference_profiling", use_workflow_definition_cache: bool = True, serialize_results: bool = False, @@ -510,6 +564,7 @@ def init_with_workflow( workflow_version_id: Optional[str] = None, exec_session_id: Optional[str] = None, workflows_dependencies_pre_init: Optional[List[str]] = None, + video_source_options: Optional[VideoSourceOptions] = None, ) -> "InferencePipeline": """ This class creates the abstraction for making inferences from given workflow against video stream. @@ -562,6 +617,10 @@ def init_with_workflow( corresponding to cv2 VideoCapture properties cv2.CAP_PROP_*. If not given, defaults for the video source will be used. Example valid properties are: {"frame_width": 1920, "frame_height": 1080, "fps": 30.0} + video_source_options (Optional[VideoSourceOptions]): Optional + producer-specific settings. A single dictionary applies to all video + sources; a list must align with `video_reference` and may contain None + for sources that need no special configuration. workflow_init_parameters (Optional[Dict[str, Any]]): Additional init parameters to be used by workflows Execution Engine to init steps of your workflow - may be required when running workflows with custom plugins. @@ -578,6 +637,17 @@ def init_with_workflow( to grab frames from multiple sources can wait for batch to be filled before yielding already collected frames. Please set this value in PRODUCTION to avoid performance drops when specific sources shows unstable latency. Visit `multiplex_videos(...)` for more information about multiplexing process. + video_processing_mode (Optional[Union[str, VideoProcessingMode]]): High-level intent for live + multi-source consumption: "auto" (FIFO with a staleness budget and self-tuning collection + window), "every_frame" (strict FIFO) or "freshest" (legacy latest-wins with a small fixed + collection timeout). Defaults to "auto" when ENABLE_TENSOR_DATA_REPRESENTATION is set, + otherwise the legacy collection behavior is preserved unchanged; pass "legacy" to force + the legacy behavior explicitly (the escape hatch from the flag-driven default). File + sources always keep every-frame semantics regardless of mode. See + `inference.core.interfaces.camera.collection_policy` for details. + max_staleness (Optional[float]): Staleness budget (seconds) of "auto" mode - live frames older + than this are dropped (reported as FRAME_DROPPED status updates with cause + STALENESS_BUDGET_EXCEEDED) instead of being served late. Default: 0.5. profiling_directory (str): Directory where workflows profiler traces will be dumped. To enable profiling export `ENABLE_WORKFLOWS_PROFILING=True` environmental variable. You may specify number of workflow runs in a buffer with environmental variable `WORKFLOWS_PROFILER_BUFFER_SIZE=n` - making last `n` @@ -724,9 +794,13 @@ def init_with_workflow( source_buffer_filling_strategy=source_buffer_filling_strategy, source_buffer_consumption_strategy=source_buffer_consumption_strategy, video_source_properties=video_source_properties, + video_source_options=video_source_options, batch_collection_timeout=batch_collection_timeout, + video_processing_mode=video_processing_mode, + max_staleness=max_staleness, predictions_queue_size=predictions_queue_size, decoding_buffer_size=decoding_buffer_size, + allow_tensor_frames=ENABLE_TENSOR_DATA_REPRESENTATION, exec_session_id=exec_session_id, ) @@ -745,10 +819,14 @@ def init_with_custom_logic( source_buffer_consumption_strategy: Optional[BufferConsumptionStrategy] = None, video_source_properties: Optional[Dict[str, float]] = None, batch_collection_timeout: Optional[float] = None, + video_processing_mode: Optional[Union[str, VideoProcessingMode]] = None, + max_staleness: Optional[float] = None, sink_mode: SinkMode = SinkMode.ADAPTIVE, predictions_queue_size: int = PREDICTIONS_QUEUE_SIZE, decoding_buffer_size: int = DEFAULT_BUFFER_SIZE, exec_session_id: Optional[str] = None, + allow_tensor_frames: bool = False, + video_source_options: Optional[VideoSourceOptions] = None, ) -> "InferencePipeline": """ This class creates the abstraction for making inferences from given workflow against video stream. @@ -800,10 +878,25 @@ def init_with_custom_logic( as list of configs. Then the list must be of length of `video_reference` and may also contain None values to denote that specific source should remain not configured. Example valid properties are: {"frame_width": 1920, "frame_height": 1080, "fps": 30.0} + video_source_options (Optional[VideoSourceOptions]): Optional + producer-specific settings. A single dictionary applies to all video + sources; a list must align with `video_reference` and may contain None + for sources that need no special configuration. batch_collection_timeout (Optional[float]): Parameter of multiplex_videos(...) dictating how long process to grab frames from multiple sources can wait for batch to be filled before yielding already collected frames. Please set this value in PRODUCTION to avoid performance drops when specific sources shows unstable latency. Visit `multiplex_videos(...)` for more information about multiplexing process. + video_processing_mode (Optional[Union[str, VideoProcessingMode]]): High-level intent for live + multi-source consumption: "auto" (FIFO with a staleness budget and self-tuning collection + window), "every_frame" (strict FIFO) or "freshest" (legacy latest-wins with a small fixed + collection timeout). Defaults to "auto" when ENABLE_TENSOR_DATA_REPRESENTATION is set, + otherwise the legacy collection behavior is preserved unchanged; pass "legacy" to force + the legacy behavior explicitly (the escape hatch from the flag-driven default). File + sources always keep every-frame semantics regardless of mode. See + `inference.core.interfaces.camera.collection_policy` for details. + max_staleness (Optional[float]): Staleness budget (seconds) of "auto" mode - live frames older + than this are dropped (reported as FRAME_DROPPED status updates with cause + STALENESS_BUDGET_EXCEEDED) instead of being served late. Default: 0.5. sink_mode (SinkMode): Parameter that controls how video frames and predictions will be passed to sink handler. With SinkMode.SEQUENTIAL - each frame and prediction triggers separate call for sink, in case of SinkMode.BATCH - list of frames and predictions will be provided to sink, always aligned @@ -836,17 +929,49 @@ def init_with_custom_logic( watchdog = NullPipelineWatchdog() status_update_handlers = list(status_update_handlers or []) status_update_handlers.append(watchdog.on_status_update) + resolved_processing_mode = resolve_video_processing_mode( + explicit_mode=video_processing_mode + ) + collection_policy = None + if resolved_processing_mode is VideoProcessingMode.FRESHEST: + if source_buffer_consumption_strategy is None: + source_buffer_consumption_strategy = BufferConsumptionStrategy.EAGER + if batch_collection_timeout is None: + batch_collection_timeout = FRESHEST_MODE_BATCH_COLLECTION_TIMEOUT + elif resolved_processing_mode is not None: + if source_buffer_consumption_strategy is None: + source_buffer_consumption_strategy = BufferConsumptionStrategy.LAZY + + def _report_stale_frame_dropped(frame: VideoFrame) -> None: + send_inference_pipeline_status_update( + severity=UpdateSeverity.DEBUG, + event_type=FRAME_DROPPED_EVENT, + status_update_handlers=status_update_handlers, + payload={ + "source_id": frame.source_id, + "frame_id": frame.frame_id, + "cause": STALENESS_DROP_CAUSE, + }, + ) + + collection_policy = CollectionPolicy( + mode=resolved_processing_mode, + max_staleness=max_staleness, + on_frame_dropped=_report_stale_frame_dropped, + ) desired_source_fps = None if ENABLE_FRAME_DROP_ON_VIDEO_FILE_RATE_LIMITING: desired_source_fps = max_fps video_sources = prepare_video_sources( video_reference=video_reference, video_source_properties=video_source_properties, + video_source_options=video_source_options, status_update_handlers=status_update_handlers, source_buffer_filling_strategy=source_buffer_filling_strategy, source_buffer_consumption_strategy=source_buffer_consumption_strategy, desired_source_fps=desired_source_fps, decoding_buffer_size=decoding_buffer_size, + allow_tensor_frames=allow_tensor_frames, ) watchdog.register_video_sources(video_sources=video_sources) try: @@ -874,6 +999,7 @@ def init_with_custom_logic( on_pipeline_end=on_pipeline_end, batch_collection_timeout=batch_collection_timeout, sink_mode=sink_mode, + collection_policy=collection_policy, exec_session_id=exec_session_id, ) @@ -890,6 +1016,7 @@ def __init__( max_fps: Optional[float] = None, batch_collection_timeout: Optional[float] = None, sink_mode: SinkMode = SinkMode.ADAPTIVE, + collection_policy: Optional[CollectionPolicy] = None, exec_session_id: Optional[str] = None, ): self._on_video_frame = on_video_frame @@ -909,6 +1036,7 @@ def __init__( self._batch_collection_timeout = batch_collection_timeout self._sink_mode = sink_mode self._stream_session_id = exec_session_id or mint_stream_session_id() + self._collection_policy = collection_policy def start(self, use_main_thread: bool = True) -> None: self._stop = False @@ -1154,7 +1282,15 @@ def _use_sink( video_frames: Union[VideoFrame, List[Optional[VideoFrame]]], ) -> None: try: - self._on_prediction(predictions, video_frames) + # Frames are handed to the sink AS-IS: under + # ENABLE_TENSOR_DATA_REPRESENTATION that is the original on-device + # tensor frame (no per-frame device-to-host materialisation here). + # Pixel-consuming sinks materialise at their own boundary via + # stream.utils.materialise_video_frame_for_sink. + self._on_prediction( + predictions, + video_frames, + ) except Exception as error: payload = { "error_type": error.__class__.__name__, @@ -1182,6 +1318,7 @@ def _generate_frames( max_fps=max_fps, batch_collection_timeout=self._batch_collection_timeout, should_stop=lambda: self._stop, + collection_policy=self._collection_policy, ) diff --git a/inference/core/interfaces/stream/model_handlers/workflows.py b/inference/core/interfaces/stream/model_handlers/workflows.py index 410785ab94..2625f62690 100644 --- a/inference/core/interfaces/stream/model_handlers/workflows.py +++ b/inference/core/interfaces/stream/model_handlers/workflows.py @@ -2,6 +2,9 @@ from dataclasses import dataclass from typing import Any, Dict, List, Optional, Set +import torch + +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.interfaces.camera.entities import VideoFrame from inference.core.interfaces.stream.entities import InferenceHandlerResult from inference.core.workflows.execution_engine.core import ExecutionEngine @@ -117,7 +120,16 @@ def _build_workflows_parameters( ] workflows_parameters[self._image_input_name] = [ { - "type": "numpy_object", + # GPU-tensor decoding is best-effort under the tensor flag + # (the cv2 CPU fallback emits numpy frames, and sources may + # mix within one batch), so each frame declares its actual + # payload type instead of a fixed flag-derived one. + "type": ( + "tensor" + if ENABLE_TENSOR_DATA_REPRESENTATION + and isinstance(video_frame.image, torch.Tensor) + else "numpy_object" + ), "value": video_frame.image, "video_metadata": video_metadata, } diff --git a/inference/core/interfaces/stream/sinks.py b/inference/core/interfaces/stream/sinks.py index 7134f4e0ed..45b4f0272a 100644 --- a/inference/core/interfaces/stream/sinks.py +++ b/inference/core/interfaces/stream/sinks.py @@ -14,7 +14,10 @@ from inference.core.active_learning.middlewares import ActiveLearningMiddleware from inference.core.interfaces.camera.entities import VideoFrame from inference.core.interfaces.stream.entities import SinkHandler -from inference.core.interfaces.stream.utils import wrap_in_list +from inference.core.interfaces.stream.utils import ( + materialise_video_frame_for_sink, + wrap_in_list, +) from inference.core.utils.drawing import create_tiles from inference.core.utils.preprocess import letterbox_image @@ -163,6 +166,9 @@ def _handle_frame_rendering( if frame is None: image = np.zeros((256, 256, 3), dtype=np.uint8) else: + # This sink draws on CPU pixels โ€” materialise a tensor frame here, at + # the consumer boundary (dispatch hands tensor frames through as-is). + frame = materialise_video_frame_for_sink(frame) try: labels = [p["class"] for p in prediction["predictions"]] if hasattr(sv.Detections, "from_inference"): @@ -393,7 +399,11 @@ def active_learning_sink( """ video_frame = wrap_in_list(element=video_frame) predictions = wrap_in_list(element=predictions) - images = [f.image for f in video_frame if f is not None] + # Active learning ships CPU pixels to the backend โ€” materialise tensor + # frames at this consumer boundary. + images = [ + materialise_video_frame_for_sink(f).image for f in video_frame if f is not None + ] predictions = [p for p in predictions if p is not None] active_learning_middleware.register_batch( inference_inputs=images, diff --git a/inference/core/interfaces/stream/utils.py b/inference/core/interfaces/stream/utils.py index 5dadabbaba..f9b4fc9a4a 100644 --- a/inference/core/interfaces/stream/utils.py +++ b/inference/core/interfaces/stream/utils.py @@ -1,13 +1,17 @@ import json import os from concurrent.futures import ThreadPoolExecutor +from dataclasses import replace from datetime import datetime from typing import Callable, Dict, List, Optional, TypeVar, Union +import numpy as np + from inference.core import logger from inference.core.env import DEFAULT_BUFFER_SIZE, ENABLE_WORKFLOWS_PROFILING from inference.core.interfaces.camera.entities import ( StatusUpdate, + VideoFrame, VideoSourceIdentifier, ) from inference.core.interfaces.camera.video_source import ( @@ -18,6 +22,10 @@ from inference.core.workflows.execution_engine.profiling.core import WorkflowsProfiler T = TypeVar("T") +VideoSourceOptions = Union[ + Dict[str, object], + List[Optional[Dict[str, object]]], +] def prepare_video_sources( @@ -30,6 +38,8 @@ def prepare_video_sources( source_buffer_consumption_strategy: Optional[BufferConsumptionStrategy], desired_source_fps: Optional[Union[float, int]] = None, decoding_buffer_size: int = DEFAULT_BUFFER_SIZE, + allow_tensor_frames: bool = False, + video_source_options: Optional[VideoSourceOptions] = None, ) -> List[VideoSource]: video_reference = wrap_in_list(element=video_reference) if len(video_reference) < 1: @@ -43,14 +53,23 @@ def prepare_video_sources( error_description="Cannot apply `video_source_properties` to video sources due to missmatch in " "number of entries in properties configuration.", ) + video_source_options = wrap_in_list(element=video_source_options) + video_source_options = broadcast_elements( + elements=video_source_options, + desired_length=len(video_reference), + error_description="Cannot apply `video_source_options` to video sources due to mismatch in " + "number of entries in options configuration.", + ) return initialise_video_sources( video_reference=video_reference, video_source_properties=video_source_properties, + video_source_options=video_source_options, status_update_handlers=status_update_handlers, source_buffer_filling_strategy=source_buffer_filling_strategy, source_buffer_consumption_strategy=source_buffer_consumption_strategy, desired_source_fps=desired_source_fps, decoding_buffer_size=decoding_buffer_size, + allow_tensor_frames=allow_tensor_frames, ) @@ -80,6 +99,8 @@ def initialise_video_sources( source_buffer_consumption_strategy: Optional[BufferConsumptionStrategy], desired_source_fps: Optional[Union[float, int]] = None, decoding_buffer_size: int = DEFAULT_BUFFER_SIZE, + allow_tensor_frames: bool = False, + video_source_options: Optional[List[Optional[Dict[str, object]]]] = None, ) -> List[VideoSource]: if isinstance(source_buffer_filling_strategy, str): source_buffer_filling_strategy = BufferFillingStrategy( @@ -89,6 +110,8 @@ def initialise_video_sources( source_buffer_consumption_strategy = BufferConsumptionStrategy( source_buffer_consumption_strategy ) + if video_source_options is None: + video_source_options = [None] * len(video_reference) return [ VideoSource.init( video_reference=reference, @@ -96,16 +119,59 @@ def initialise_video_sources( buffer_filling_strategy=source_buffer_filling_strategy, buffer_consumption_strategy=source_buffer_consumption_strategy, video_source_properties=source_properties, + video_source_options=source_options, source_id=i, desired_fps=desired_source_fps, buffer_size=decoding_buffer_size, + allow_tensor_frames=allow_tensor_frames, ) - for i, (reference, source_properties) in enumerate( - zip(video_reference, video_source_properties) + for i, (reference, source_properties, source_options) in enumerate( + zip(video_reference, video_source_properties, video_source_options) ) ] +def materialise_video_frames_for_sink( + video_frames: Union[VideoFrame, List[Optional[VideoFrame]]], +) -> Union[VideoFrame, List[Optional[VideoFrame]]]: + if isinstance(video_frames, list): + return [ + materialise_video_frame_for_sink(video_frame) + for video_frame in video_frames + ] + return materialise_video_frame_for_sink(video_frames) + + +def materialise_video_frame_for_sink( + video_frame: Optional[VideoFrame], +) -> Optional[VideoFrame]: + """Convert a tensor video frame to a host BGR ``np.ndarray`` frame. + + Under ENABLE_TENSOR_DATA_REPRESENTATION the pipeline hands sinks the + original on-device tensor frame โ€” nothing materialises in dispatch. Sinks + that actually consume pixels on CPU (visualisation, active learning, ...) + call this at their own boundary and pay the device-to-host copy only when + the pixels are genuinely needed. A numpy frame passes through untouched. + """ + if video_frame is None or isinstance(video_frame.image, np.ndarray): + return video_frame + + tensor_image = video_frame.image.detach().cpu() + if tensor_image.ndim == 2: + numpy_image = tensor_image.contiguous().numpy() + elif tensor_image.ndim == 3 and tensor_image.shape[0] == 1: + numpy_image = tensor_image[0].contiguous().numpy() + elif tensor_image.ndim == 3 and tensor_image.shape[0] in {3, 4}: + numpy_image = tensor_image.permute(1, 2, 0).contiguous().numpy() + channel_order = [2, 1, 0] + if tensor_image.shape[0] == 4: + channel_order.append(3) + numpy_image = np.ascontiguousarray(numpy_image[..., channel_order]) + else: + raise ValueError("Tensor video frames must use HW, 1CHW, 3CHW, or 4CHW layout") + return replace(video_frame, image=numpy_image) + + def on_pipeline_end( thread_pool_executor: ThreadPoolExecutor, cancel_thread_pool_tasks_on_exit: bool, diff --git a/inference/core/interfaces/stream_manager/manager_app/inference_pipeline_manager.py b/inference/core/interfaces/stream_manager/manager_app/inference_pipeline_manager.py index 54ff94f7dd..0fe6c4ef99 100644 --- a/inference/core/interfaces/stream_manager/manager_app/inference_pipeline_manager.py +++ b/inference/core/interfaces/stream_manager/manager_app/inference_pipeline_manager.py @@ -32,6 +32,7 @@ ) from inference.core.interfaces.stream.inference_pipeline import InferencePipeline from inference.core.interfaces.stream.sinks import InMemoryBufferSink, multi_sink +from inference.core.interfaces.stream.utils import materialise_video_frame_for_sink from inference.core.interfaces.stream.watchdog import ( BasePipelineWatchDog, PipelineWatchDog, @@ -413,7 +414,12 @@ def webrtc_sink( "Please try to adjust the scene so models detect objects" ) errors.append("or stop preview, update workflow and try again.") - frame = video_frame.image.copy() + # The WebRTC preview needs CPU pixels; the pipeline + # hands tensor frames through unmaterialised, so + # convert at this consumer boundary. + frame = materialise_video_frame_for_sink( + video_frame + ).image.copy() for row, error in enumerate(errors): frame = cv.putText( diff --git a/inference/core/managers/base.py b/inference/core/managers/base.py index c643bd56b1..3aa74f01f3 100644 --- a/inference/core/managers/base.py +++ b/inference/core/managers/base.py @@ -1,7 +1,7 @@ import time from contextlib import contextmanager from threading import Lock -from typing import Dict, Generator, List, Optional, Tuple, Union +from typing import Any, Dict, Generator, List, Optional, Tuple, Union import numpy as np from fastapi.encoders import jsonable_encoder @@ -455,6 +455,24 @@ def model_infer_sync( model = self._get_model_reference(model_id=model_id) return model.infer_from_request(request) + def run_tensor_native_inference(self, model_id: str, **kwargs) -> Any: + with start_span( + "model.infer", + { + "model.id": model_id, + "model.infer.caller": "run_tensor_native_inference", + }, + ): + try: + t_infer_start = time.perf_counter() + model = self._get_model_reference(model_id=model_id) + result = model.run_tensor_native_inference(**kwargs) + record_inference(model_id, time.perf_counter() - t_infer_start) + return result + except Exception as error: + record_error(error) + raise + def make_response( self, model_id: str, predictions: List[List[float]], *args, **kwargs ) -> InferenceResponse: diff --git a/inference/core/managers/decorators/base.py b/inference/core/managers/decorators/base.py index 0c6b39886a..d20854b7bf 100644 --- a/inference/core/managers/decorators/base.py +++ b/inference/core/managers/decorators/base.py @@ -1,4 +1,4 @@ -from typing import List, Optional, Tuple +from typing import Any, List, Optional, Tuple import numpy as np @@ -143,6 +143,9 @@ def infer_only(self, model_id: str, request, img_in, img_dims, batch_size=None): model_id, request, img_in, img_dims, batch_size ) + def run_tensor_native_inference(self, model_id: str, **kwargs) -> Any: + return self.model_manager.run_tensor_native_inference(model_id, **kwargs) + def preprocess(self, model_id: str, request: InferenceRequest): """Processes the preprocessing part of a request. diff --git a/inference/core/managers/decorators/fixed_size_cache.py b/inference/core/managers/decorators/fixed_size_cache.py index b04e731ff3..d4e8c66347 100644 --- a/inference/core/managers/decorators/fixed_size_cache.py +++ b/inference/core/managers/decorators/fixed_size_cache.py @@ -1,7 +1,7 @@ import gc from collections import deque from threading import Lock -from typing import List, Optional +from typing import Any, List, Optional from inference.core import logger from inference.core.entities.requests.inference import InferenceRequest @@ -226,6 +226,10 @@ def infer_from_request_sync( self._refresh_model_position_in_a_queue(model_id=model_id) return super().infer_from_request_sync(model_id, request, **kwargs) + def run_tensor_native_inference(self, model_id: str, **kwargs) -> Any: + self._refresh_model_position_in_a_queue(model_id=model_id) + return super().run_tensor_native_inference(model_id, **kwargs) + def infer_only(self, model_id: str, request, img_in, img_dims, batch_size=None): """Performs only the inference part of a request and updates the cache. diff --git a/inference/core/managers/decorators/logger.py b/inference/core/managers/decorators/logger.py index 5bcbd13bd5..cc00b775f8 100644 --- a/inference/core/managers/decorators/logger.py +++ b/inference/core/managers/decorators/logger.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import Any, Optional from inference.core.entities.requests.inference import InferenceRequest from inference.core.entities.responses.inference import InferenceResponse @@ -73,6 +73,12 @@ def infer_from_request_sync( logger.info(f"๐Ÿ“ฅ [{model_id}] res={res}.") return res + def run_tensor_native_inference(self, model_id: str, **kwargs) -> Any: + logger.info(f"๐Ÿ“ฅ [{model_id}] request={kwargs}.") + res = super().run_tensor_native_inference(model_id, **kwargs) + logger.info(f"๐Ÿ“ฅ [{model_id}] res={res}.") + return res + def remove(self, model_id: str, delete_from_disk: bool = True) -> Model: """Removes a model from the manager and logs the action. diff --git a/inference/core/models/base.py b/inference/core/models/base.py index b478782395..e988b400f9 100644 --- a/inference/core/models/base.py +++ b/inference/core/models/base.py @@ -40,6 +40,9 @@ def infer(self, image: Any, **kwargs) -> Any: return postprocessed + def run_tensor_native_inference(self, **kwargs) -> Any: + raise NotImplementedError + def preprocess( self, image: Any, **kwargs ) -> Tuple[np.ndarray, PreprocessReturnMetadata]: diff --git a/inference/core/models/inference_models_adapters.py b/inference/core/models/inference_models_adapters.py index 0ce0abb848..17bc510122 100644 --- a/inference/core/models/inference_models_adapters.py +++ b/inference/core/models/inference_models_adapters.py @@ -235,6 +235,16 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): ) self.class_names = list(self._model.class_names) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[Detections]: + caller_color_format = kwargs.pop("input_color_format", None) + kwargs = self.map_inference_kwargs(kwargs) + kwargs["input_color_format"] = caller_color_format + return self._model(images, **kwargs) + def map_inference_kwargs(self, kwargs: dict) -> dict: kwargs["input_color_format"] = "bgr" pre_processing_overrides = PreProcessingOverrides( @@ -426,6 +436,29 @@ def _model_supports_stream_pipeline(self) -> bool: return bool(supports_stream_pipeline()) return bool(supports_stream_pipeline) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[InstanceDetections]: + enforce_dense_masks = ( + False + if GCP_SERVERLESS + else kwargs.get("enforce_dense_masks_in_inference_models", False) + ) + if not enforce_dense_masks and "rle" not in self._model.supported_mask_formats: + raise PostProcessingError( + "RLE masks are required on the tensor-native instance-segmentation " + "path (enforce_dense_masks_in_inference_models is False) but the loaded " + f"model only supports mask formats {self._model.supported_mask_formats}. " + "Either use a model that supports 'rle' or set " + "enforce_dense_masks_in_inference_models=True to receive dense masks." + ) + caller_color_format = kwargs.pop("input_color_format", None) + kwargs = self.map_inference_kwargs(kwargs) + kwargs["input_color_format"] = caller_color_format + return self._model(images, **kwargs) + def map_inference_kwargs(self, kwargs: dict) -> dict: kwargs["input_color_format"] = "bgr" pre_processing_overrides = PreProcessingOverrides( @@ -440,6 +473,9 @@ def map_inference_kwargs(self, kwargs: dict) -> dict: "enforce_dense_masks_in_inference_models", False, ) + # Consumed here โ€” must not leak into the model call, whose deeper + # pre/post stages do not accept arbitrary kwargs. + kwargs.pop("enforce_dense_masks_in_inference_models", None) kwargs["pre_processing_overrides"] = pre_processing_overrides if ( "rle" in self._model.supported_mask_formats @@ -1067,6 +1103,17 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) self.class_names = list(self._model.class_names) + self.key_points_classes = list(self._model.key_points_classes) + + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> Tuple[List[KeyPoints], Optional[List[Detections]]]: + caller_color_format = kwargs.pop("input_color_format", None) + kwargs = self.map_inference_kwargs(kwargs) + kwargs["input_color_format"] = caller_color_format + return self._model(images, **kwargs) def map_inference_kwargs(self, kwargs: dict) -> dict: kwargs["input_color_format"] = "bgr" @@ -1276,6 +1323,16 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): ) self.class_names = list(self._model.class_names) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> Union[ClassificationPrediction, List[MultiLabelClassificationPrediction]]: + caller_color_format = kwargs.pop("input_color_format", None) + kwargs = self.map_inference_kwargs(kwargs) + kwargs["input_color_format"] = caller_color_format + return self._model(images, **kwargs) + def map_inference_kwargs(self, kwargs: dict) -> dict: kwargs["input_color_format"] = "bgr" pre_processing_overrides = PreProcessingOverrides( @@ -1308,7 +1365,7 @@ def predict(self, img_in, **kwargs): def postprocess( self, - predictions: Tuple[List[KeyPoints], Optional[List[Detections]]], + predictions: torch.Tensor, returned_metadata: List[Tuple[int, int]], **kwargs, ) -> Union[ @@ -1607,6 +1664,16 @@ def class_map(self): # match segment.roboflow.com return {str(k): v for k, v in enumerate(self.class_names)} + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[SemanticSegmentationResult]: + caller_color_format = kwargs.pop("input_color_format", None) + kwargs = self.map_inference_kwargs(kwargs) + kwargs["input_color_format"] = caller_color_format + return self._model(images, **kwargs) + def map_inference_kwargs(self, kwargs: dict) -> dict: kwargs["input_color_format"] = "bgr" pre_processing_overrides = PreProcessingOverrides( @@ -1743,6 +1810,33 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[torch.Tensor]: + caller_color_format = kwargs.pop("input_color_format", None) + kwargs = self.map_inference_kwargs(kwargs) + kwargs["input_color_format"] = caller_color_format + depth_maps = self._model(images, **kwargs) + # Models behind this adapter (YOLO26-depth) emit metric depth, where + # larger means FARTHER. The tensor-native depth contract mirrors the + # DepthAnything adapters: raw per-image maps in which larger means + # CLOSER, with the caller (depth-estimation block) applying + # `(map - min) / (max - min)` itself. Negate so that formula reproduces + # this adapter's numpy-path `(max - map) / (max - min)` exactly. + return [-depth_map for depth_map in depth_maps] + + def map_inference_kwargs(self, kwargs: dict) -> dict: + kwargs["input_color_format"] = "bgr" + pre_processing_overrides = PreProcessingOverrides( + disable_contrast_enhancement=kwargs.get("disable_preproc_contrast", False), + disable_grayscale=kwargs.get("disable_preproc_grayscale", False), + disable_static_crop=kwargs.get("disable_preproc_static_crop", False), + ) + kwargs["pre_processing_overrides"] = pre_processing_overrides + return kwargs + def preprocess(self, image: Any, **kwargs): if isinstance(image, list): raise ValueError("Depth estimation does not support batched inference.") diff --git a/inference/core/version.py b/inference/core/version.py index 985364d0d4..58bca5c41b 100644 --- a/inference/core/version.py +++ b/inference/core/version.py @@ -1,4 +1,4 @@ -__version__ = "1.3.8" +__version__ = "1.4.0" if __name__ == "__main__": diff --git a/inference/core/workflows/core_steps/analytics/_zone_geometry.py b/inference/core/workflows/core_steps/analytics/_zone_geometry.py new file mode 100644 index 0000000000..e4bf376abb --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/_zone_geometry.py @@ -0,0 +1,234 @@ +from collections import Counter, defaultdict, deque +from math import sqrt +from typing import DefaultDict, Deque, Optional, Sequence, Tuple, Union + +import numpy as np +import supervision as sv + +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks + + +def empty_detections_like( + detections: Union[Detections, InstanceDetections], +) -> Union[Detections, InstanceDetections]: + bboxes_metadata = [] if isinstance(detections.bboxes_metadata, list) else None + if isinstance(detections, InstanceDetections): + mask = detections.mask + if isinstance(mask, InstancesRLEMasks): + mask = InstancesRLEMasks(image_size=mask.image_size, masks=[]) + else: + mask = mask[0:0] + return InstanceDetections( + xyxy=detections.xyxy[0:0], + class_id=detections.class_id[0:0], + confidence=detections.confidence[0:0], + mask=mask, + image_metadata=detections.image_metadata, + bboxes_metadata=bboxes_metadata, + ) + return Detections( + xyxy=detections.xyxy[0:0], + class_id=detections.class_id[0:0], + confidence=detections.confidence[0:0], + image_metadata=detections.image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +_DEFAULT_LINE_ANCHORS = ( + sv.Position.TOP_LEFT, + sv.Position.TOP_RIGHT, + sv.Position.BOTTOM_LEFT, + sv.Position.BOTTOM_RIGHT, +) + + +def _stack_anchor_coordinates( + xyxy: np.ndarray, anchors: Sequence[sv.Position] +) -> np.ndarray: + center_x = (xyxy[:, 0] + xyxy[:, 2]) / 2 + center_y = (xyxy[:, 1] + xyxy[:, 3]) / 2 + result = np.empty((len(anchors), len(xyxy), 2), dtype=xyxy.dtype) + for index, anchor in enumerate(anchors): + if anchor == sv.Position.CENTER: + result[index, :, 0] = center_x + result[index, :, 1] = center_y + elif anchor == sv.Position.CENTER_OF_MASS: + raise ValueError( + "Cannot use `Position.CENTER_OF_MASS` without a detection mask." + ) + elif anchor == sv.Position.CENTER_LEFT: + result[index, :, 0] = xyxy[:, 0] + result[index, :, 1] = center_y + elif anchor == sv.Position.CENTER_RIGHT: + result[index, :, 0] = xyxy[:, 2] + result[index, :, 1] = center_y + elif anchor == sv.Position.BOTTOM_CENTER: + result[index, :, 0] = center_x + result[index, :, 1] = xyxy[:, 3] + elif anchor == sv.Position.BOTTOM_LEFT: + result[index, :, 0] = xyxy[:, 0] + result[index, :, 1] = xyxy[:, 3] + elif anchor == sv.Position.BOTTOM_RIGHT: + result[index, :, 0] = xyxy[:, 2] + result[index, :, 1] = xyxy[:, 3] + elif anchor == sv.Position.TOP_CENTER: + result[index, :, 0] = center_x + result[index, :, 1] = xyxy[:, 1] + elif anchor == sv.Position.TOP_LEFT: + result[index, :, 0] = xyxy[:, 0] + result[index, :, 1] = xyxy[:, 1] + elif anchor == sv.Position.TOP_RIGHT: + result[index, :, 0] = xyxy[:, 2] + result[index, :, 1] = xyxy[:, 1] + else: + raise ValueError(f"{anchor} is not supported.") + return result + + +def anchor_coordinates(xyxy: np.ndarray, anchor: sv.Position) -> np.ndarray: + return _stack_anchor_coordinates(xyxy, (anchor,))[0] + + +class LeanLineZone: + def __init__( + self, + start: Tuple[float, float], + end: Tuple[float, float], + triggering_anchors: Optional[Sequence[sv.Position]] = None, + ) -> None: + self.start = start + self.end = end + self.triggering_anchors = tuple( + _DEFAULT_LINE_ANCHORS if triggering_anchors is None else triggering_anchors + ) + if not self.triggering_anchors: + raise ValueError("Triggering anchors cannot be empty.") + + minimum_crossing_threshold = 1 + self.crossing_history_length = max(2, minimum_crossing_threshold + 1) + self.crossing_state_history: DefaultDict[ + Tuple[int, Optional[int]], Deque[bool] + ] = defaultdict(lambda: deque(maxlen=self.crossing_history_length)) + self._in_count_per_class: Counter = Counter() + self._out_count_per_class: Counter = Counter() + + start_x, start_y = start + end_x, end_y = end + delta_x = end_x - start_x + delta_y = end_y - start_y + magnitude = sqrt(delta_x**2 + delta_y**2) + if magnitude == 0: + raise ValueError("The magnitude of the vector cannot be zero.") + + unit_vector_x = delta_x / magnitude + unit_vector_y = delta_y / magnitude + perpendicular_vector_x = -unit_vector_y + perpendicular_vector_y = unit_vector_x + self._limit_starts = np.array([start, end], dtype=float) + self._limit_vectors = np.array( + [ + [perpendicular_vector_x, perpendicular_vector_y], + [-perpendicular_vector_x, -perpendicular_vector_y], + ], + dtype=float, + ) + self._line_vector = np.array([delta_x, delta_y], dtype=float) + + @property + def in_count(self) -> int: + return sum(self._in_count_per_class.values()) + + @property + def out_count(self) -> int: + return sum(self._out_count_per_class.values()) + + def trigger( + self, xyxy: np.ndarray, tracker_ids: np.ndarray + ) -> Tuple[np.ndarray, np.ndarray]: + crossed_in = np.full(len(xyxy), False) + crossed_out = np.full(len(xyxy), False) + if len(xyxy) == 0: + return crossed_in, crossed_out + + all_anchors = _stack_anchor_coordinates(xyxy, self.triggering_anchors) + limit_deltas = all_anchors[None, :, :, :] - self._limit_starts[:, None, None, :] + limit_cross_products = ( + self._limit_vectors[:, 0, None, None] * limit_deltas[:, :, :, 1] + - self._limit_vectors[:, 1, None, None] * limit_deltas[:, :, :, 0] + ) + in_limits = np.all( + (limit_cross_products[0] > 0) == (limit_cross_products[1] > 0), + axis=0, + ) + + line_deltas = all_anchors - np.asarray(self.start) + triggers = ( + self._line_vector[0] * line_deltas[:, :, 1] + - self._line_vector[1] * line_deltas[:, :, 0] + ) < 0 + has_any_left_trigger = np.any(triggers, axis=0) + has_any_right_trigger = np.any(~triggers, axis=0) + + class_id = None + for index, tracker_id in enumerate(tracker_ids): + if not in_limits[index]: + continue + if has_any_left_trigger[index] and has_any_right_trigger[index]: + continue + + tracker_state = has_any_left_trigger[index] + crossing_history = self.crossing_state_history[(tracker_id, class_id)] + crossing_history.append(tracker_state) + if len(crossing_history) < self.crossing_history_length: + continue + + oldest_state = crossing_history[0] + if crossing_history.count(oldest_state) > 1: + continue + + if tracker_state: + self._in_count_per_class[class_id] += 1 + crossed_in[index] = True + else: + self._out_count_per_class[class_id] += 1 + crossed_out[index] = True + + return crossed_in, crossed_out + + +class LeanPolygonZone: + def __init__( + self, + polygon: np.ndarray, + triggering_anchors: Sequence[sv.Position], + ) -> None: + self.polygon = polygon.astype(int) + self.triggering_anchors = tuple(triggering_anchors) + if not self.triggering_anchors: + raise ValueError("Triggering anchors cannot be empty.") + + self.current_count = 0 + x_max, y_max = np.max(polygon, axis=0) + self.mask = sv.polygon_to_mask( + polygon=polygon, resolution_wh=(x_max + 2, y_max + 2) + ) + + def trigger(self, xyxy: np.ndarray) -> np.ndarray: + if len(xyxy) == 0: + self.current_count = 0 + return np.array([], dtype=bool) + + all_anchors = np.ceil( + _stack_anchor_coordinates(xyxy, self.triggering_anchors) + ).astype(int) + mask_h, mask_w = self.mask.shape + x, y = all_anchors[:, :, 0], all_anchors[:, :, 1] + in_bounds = (x >= 0) & (y >= 0) & (x < mask_w) & (y < mask_h) + x_safe = np.clip(x, 0, mask_w - 1) + y_safe = np.clip(y, 0, mask_h - 1) + is_in_zone = np.all(in_bounds & self.mask[y_safe, x_safe], axis=0) + self.current_count = int(np.sum(is_in_zone)) + return is_in_zone.astype(bool) diff --git a/inference/core/workflows/core_steps/analytics/detection_event_log/v1_tensor.py b/inference/core/workflows/core_steps/analytics/detection_event_log/v1_tensor.py new file mode 100644 index 0000000000..34b81ca15c --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/detection_event_log/v1_tensor.py @@ -0,0 +1,479 @@ +import heapq +import time +from dataclasses import dataclass +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +from pydantic import ConfigDict, Field + +from inference.core import logger +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + VideoMetadata, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + DICTIONARY_KIND, + FLOAT_KIND, + INTEGER_KIND, + Selector, + WorkflowImageSelector, + WorkflowParameterSelector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY = "event_log" +DETECTIONS_OUTPUT_KEY = "detections" +MAX_VIDEOS = 100 # Maximum number of video streams to track before evicting oldest + + +@dataclass +class DetectionEvent: + """Stores event data for a tracked detection.""" + + tracker_id: int + class_name: str + first_seen_frame: int + first_seen_timestamp: float # Unix wall-clock time (frame_timestamp or time.time()) + last_seen_frame: int + last_seen_timestamp: float # Unix wall-clock time (frame_timestamp or time.time()) + first_seen_relative: float = 0.0 # seconds since video start + last_seen_relative: float = 0.0 # seconds since video start + frame_count: int = 1 + logged: bool = False + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Detection Event Log", + "version": "v1", + "short_description": "Tracks detection events over time, logging when objects first appear and persist.", + "long_description": ( + "This block maintains a log of detection events from tracked objects. " + "For each tracked object it records: class name, first and last seen frame numbers, " + "absolute wall-clock timestamps (Unix epoch floats derived from frame_timestamp metadata, " + "or time.time() as fallback), and relative timestamps in seconds since the video started. " + "Objects must be seen for a minimum number of frames (frame_threshold) before being moved " + "from 'pending' to 'logged' status. " + "Stale events (not seen for stale_frames frames) are removed during periodic cleanup " + "(every flush_interval frames). When a logged event goes stale it is emitted in the " + "complete_events output, which contains the full event data for objects that were tracked " + "long enough to be logged and have since left the scene. " + "The reference_timestamp parameter is deprecated and no longer used." + ), + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "analytics", + "icon": "fal fa-list-timeline", + "blockPriority": 3, + }, + } + ) + type: Literal["roboflow_core/detection_event_log@v1"] + + image: WorkflowImageSelector = Field( + description="Reference to the image for video metadata (frame number, timestamp).", + examples=["$inputs.image"], + ) + + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( + description="Tracked detections from byte tracker (must have tracker_id).", + examples=["$steps.byte_tracker.tracked_detections"], + ) + + frame_threshold: Union[int, WorkflowParameterSelector(kind=[INTEGER_KIND])] = Field( + default=30, + description="Number of frames an object must be seen before being logged.", + examples=[5, 10], + ge=1, + ) + + flush_interval: Union[int, WorkflowParameterSelector(kind=[INTEGER_KIND])] = Field( + default=30, + description="How often (in frames) to run the cleanup operation for stale events.", + examples=[30, 60], + ge=1, + ) + + stale_frames: Union[int, WorkflowParameterSelector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Remove events that haven't been seen for this many frames.", + examples=[150, 300], + ge=1, + ) + + reference_timestamp: Optional[Union[float, Selector(kind=[FLOAT_KIND])]] = Field( + default=None, + description="Deprecated, no longer used. Absolute timestamps are now taken directly from frame_timestamp metadata (or time.time() as fallback).", + examples=[1726570875.0], + deprecated=True, + ) + + fallback_fps: Union[float, WorkflowParameterSelector(kind=[FLOAT_KIND])] = Field( + default=1.0, + description="Fallback FPS to use when video metadata does not provide FPS information. Used to calculate relative timestamps.", + examples=[1.0, 30.0], + gt=0, + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[DICTIONARY_KIND], + ), + OutputDefinition( + name=DETECTIONS_OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + OutputDefinition( + name="total_logged", + kind=[INTEGER_KIND], + ), + OutputDefinition( + name="total_pending", + kind=[INTEGER_KIND], + ), + OutputDefinition( + name="complete_events", + kind=[DICTIONARY_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class DetectionEventLogBlockV1(WorkflowBlock): + """ + Block that tracks detection events over time. + + Maintains a dictionary of tracked objects with: + - First seen timestamp and frame + - Last seen timestamp and frame + - Class name + - Frame count (number of frames the object has been seen) + + Only logs objects that have been seen for at least frame_threshold frames. + Runs cleanup every flush_interval frames, removing events not seen for stale_frames. + """ + + def __init__(self): + # Dict[video_id, Dict[tracker_id, DetectionEvent]] + self._event_logs: Dict[str, Dict[int, DetectionEvent]] = {} + # Dict[video_id, last_flush_frame] + self._last_flush_frame: Dict[str, int] = {} + # Dict[video_id, frame_count] - internal frame counter (increments each run) + self._frame_count: Dict[str, int] = {} + # Dict[video_id, last_access_frame] - tracks when each video was last accessed (global frame count) + self._last_access: Dict[str, int] = {} + # Dict[video_id, first_frame_timestamp] - stores the first frame's wall-clock timestamp + # Used as the anchor for frame_timestamp-based relative time calculation + self._first_frame_timestamps: Dict[str, float] = {} + # Global frame counter for tracking video access order + self._global_frame: int = 0 + # Min-heap of (last_access_frame, video_id) for efficient oldest video lookup + self._access_heap: List[Tuple[int, str]] = [] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def _get_relative_time( + self, + video_id: str, + metadata: VideoMetadata, + fallback_fps: float, + ) -> float: + """Calculate relative time in seconds since video started. + + Uses frame_timestamp from metadata when available for accurate timing, + even when inference doesn't run at the camera's reported FPS (e.g. due + to dropped frames or processing lag). Falls back to metadata.frame_number + / FPS when frame_timestamp is not available. + """ + if metadata.frame_timestamp is not None: + frame_ts = metadata.frame_timestamp.timestamp() + if video_id not in self._first_frame_timestamps: + self._first_frame_timestamps[video_id] = frame_ts + return frame_ts - self._first_frame_timestamps[video_id] + + # Fallback: use actual video frame number (not internal counter) to + # correctly account for dropped/skipped frames during inference. + # frame_number=0 is a sentinel for static/non-video images, treat as first frame. + fps = metadata.fps if metadata.fps and metadata.fps != 0 else fallback_fps + return max(metadata.frame_number - 1, 0) / fps + + def _evict_oldest_video(self) -> None: + """Remove the oldest video stream data when MAX_VIDEOS is exceeded.""" + if len(self._event_logs) <= MAX_VIDEOS: + return + + # Rebuild heap if out of sync with current state + if len(self._access_heap) < len(self._last_access): + self._access_heap[:] = [ + (frame, vid) for vid, frame in self._last_access.items() + ] + heapq.heapify(self._access_heap) + + # Pop stale entries until we find a valid current entry + while self._access_heap: + frame, vid = heapq.heappop(self._access_heap) + if self._last_access.get(vid) == frame and vid in self._event_logs: + oldest_video_id = vid + break + else: + # If heap is empty but we have event_logs, use fallback + oldest_video_id = min(self._last_access, key=self._last_access.get) + + # Remove all data for this video + self._event_logs.pop(oldest_video_id, None) + self._last_flush_frame.pop(oldest_video_id, None) + self._frame_count.pop(oldest_video_id, None) + self._last_access.pop(oldest_video_id, None) + self._first_frame_timestamps.pop(oldest_video_id, None) + + def _remove_stale_events( + self, + event_log: Dict[int, DetectionEvent], + current_frame: int, + stale_frames: int, + frame_threshold: int, + ) -> List[DetectionEvent]: + """Remove events that haven't been seen for stale_frames. + + Returns list of removed LOGGED events (events that met frame_threshold). + These are "complete" events - objects that were tracked long enough + to be logged and have now left the scene. + """ + stale_tracker_ids = [] + complete_events = [] + + for tracker_id, event in event_log.items(): + frames_since_seen = current_frame - event.last_seen_frame + if frames_since_seen > stale_frames: + stale_tracker_ids.append(tracker_id) + # Only return logged events as "complete" - pending events are just discarded + if event.frame_count >= frame_threshold: + complete_events.append(event) + + for tracker_id in stale_tracker_ids: + del event_log[tracker_id] + + return complete_events + + def run( + self, + image: WorkflowImageData, + detections: Union[Detections, InstanceDetections], + frame_threshold: int, + flush_interval: int, + stale_frames: int, + fallback_fps: float = 1.0, + reference_timestamp: Optional[float] = None, + ) -> BlockResult: + """Process detections and update the event log. + + Args: + image: Workflow image data containing video metadata. + detections: Tracked tensor-native detections with tracker_id (in + `bboxes_metadata`) from ByteTracker. + frame_threshold: Minimum frames an object must be seen before logging. + flush_interval: How often to run stale event cleanup. + stale_frames: Remove events not seen for this many frames. + fallback_fps: FPS to use when video metadata doesn't provide FPS. + reference_timestamp: Unused, kept for backward compatibility. + + Returns: + Dictionary containing event_log, detections, total_logged, and total_pending. + """ + metadata = image.video_metadata + video_id = metadata.video_identifier + + # Track global frame count and video access for eviction + self._global_frame += 1 + self._last_access[video_id] = self._global_frame + + # Increment internal frame counter + current_frame = self._frame_count.get(video_id, 0) + 1 + self._frame_count[video_id] = current_frame + + current_time = self._get_relative_time(video_id, metadata, fallback_fps) + + # Use frame_timestamp for absolute time when available (reflects actual capture + # time, not inference processing time). Falls back to time.time(). + current_absolute_time = ( + metadata.frame_timestamp.timestamp() + if metadata.frame_timestamp is not None + else time.time() + ) + + # Initialize event log for this video if needed + event_log = self._event_logs.setdefault(video_id, {}) + + # Evict oldest video if we've exceeded MAX_VIDEOS (after adding current video) + self._evict_oldest_video() + + # Initialize last flush frame if not set + if video_id not in self._last_flush_frame: + self._last_flush_frame[video_id] = current_frame + + # Check if it's time to run cleanup + complete_events_list = [] + last_flush = self._last_flush_frame.get(video_id, 0) + if (current_frame - last_flush) >= flush_interval: + complete_events_list = self._remove_stale_events( + event_log, current_frame, stale_frames, frame_threshold + ) + self._last_flush_frame[video_id] = current_frame + + # Format complete events + complete_events = self._format_complete_events(complete_events_list) + + # Process detections. Tensor-native predictions carry tracker_id in + # per-detection `bboxes_metadata` and class names in `image_metadata`; + # detections without a tracker_id are skipped. + class_names_map = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + tracked_objects: List[Tuple[int, str]] = [] + for _xyxy, _mask, class_id, _conf, tracker_id, _data, _meta in detections: + if tracker_id is None: + continue + class_id_int = int(class_id) + class_name = class_names_map.get(class_id_int, f"class_{class_id_int}") + tracked_objects.append((int(tracker_id), str(class_name))) + + if not tracked_objects: + # No tracked detections, return current log + event_log_dict, total_logged, total_pending = self._format_event_log( + event_log, frame_threshold + ) + return { + OUTPUT_KEY: event_log_dict, + DETECTIONS_OUTPUT_KEY: detections, + "total_logged": total_logged, + "total_pending": total_pending, + "complete_events": complete_events, + } + + # Update event log for each tracked detection + for tracker_id, class_name in tracked_objects: + if tracker_id in event_log: + # Update existing event + event = event_log[tracker_id] + event.last_seen_frame = current_frame + event.last_seen_timestamp = current_absolute_time + event.last_seen_relative = current_time + event.frame_count += 1 + + # Mark as logged once threshold is reached + if event.frame_count >= frame_threshold and not event.logged: + event.logged = True + logger.debug( + f"Object {tracker_id} ({event.class_name}) logged after {event.frame_count} frames" + ) + else: + # Create new event + event_log[tracker_id] = DetectionEvent( + tracker_id=tracker_id, + class_name=class_name, + first_seen_frame=current_frame, + first_seen_timestamp=current_absolute_time, + last_seen_frame=current_frame, + last_seen_timestamp=current_absolute_time, + first_seen_relative=current_time, + last_seen_relative=current_time, + frame_count=1, + logged=False, + ) + + event_log_dict, total_logged, total_pending = self._format_event_log( + event_log, frame_threshold + ) + return { + OUTPUT_KEY: event_log_dict, + DETECTIONS_OUTPUT_KEY: detections, + "total_logged": total_logged, + "total_pending": total_pending, + "complete_events": complete_events, + } + + def _format_complete_events( + self, + complete_events: List[DetectionEvent], + ) -> Dict[str, Any]: + """Format complete events for output. + + Args: + complete_events: List of DetectionEvent objects that have completed (gone stale). + + Returns: + Dictionary with tracker_id as key and event data as value. + """ + formatted = {} + for event in complete_events: + event_data = event.__dict__.copy() + del event_data["logged"] + formatted[str(event.tracker_id)] = event_data + + return formatted + + def _format_event_log( + self, + event_log: Dict[int, DetectionEvent], + frame_threshold: int, + ) -> tuple: + """Format the event log for output. + + Returns: + Tuple of (event_log_dict, total_logged, total_pending) + """ + logged_events = {} + pending_events = {} + + for tracker_id, event in event_log.items(): + event_data = event.__dict__.copy() + del event_data["logged"] + + if event.frame_count >= frame_threshold: + logged_events[str(tracker_id)] = event_data + else: + pending_events[str(tracker_id)] = event_data + + event_log_dict = { + "logged": logged_events, + "pending": pending_events, + } + + return event_log_dict, len(logged_events), len(pending_events) diff --git a/inference/core/workflows/core_steps/analytics/line_counter/v1_tensor.py b/inference/core/workflows/core_steps/analytics/line_counter/v1_tensor.py new file mode 100644 index 0000000000..97fbd00106 --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/line_counter/v1_tensor.py @@ -0,0 +1,221 @@ +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field +from typing_extensions import Literal, Type + +from inference.core.workflows.core_steps.analytics._zone_geometry import LeanLineZone +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + VideoMetadata, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + VIDEO_METADATA_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY_COUNT_IN: str = "count_in" +OUTPUT_KEY_COUNT_OUT: str = "count_out" +IN: str = "in" +OUT: str = "out" +DETECTIONS_IN_OUT_PARAM: str = "in_out" +SHORT_DESCRIPTION = "Count detections passing a line." +LONG_DESCRIPTION = """ +Count objects crossing a defined line segment in video using tracked detections, maintaining separate counts for objects crossing in opposite directions (in and out) for traffic analysis, people counting, entry/exit monitoring, and directional flow measurement workflows. + +## How This Block Works + +This block counts objects that cross a line segment by tracking their movement across video frames. The block: + +1. Receives tracked detection predictions with unique tracker IDs and video metadata +2. Validates that detections have tracker IDs (required for tracking object movement across frames) +3. Initializes or retrieves a line zone for the video: + - Creates a LineZone from two coordinate points defining the line segment + - Stores line zone configuration per video using video_identifier + - Maintains separate counting state for each video +4. Monitors object positions across frames: + - Tracks each object's position using its unique tracker_id + - Detects when an object's triggering anchor point (default: CENTER of bounding box) crosses the line + - Determines crossing direction based on which side of the line the object approaches from +5. Counts line crossings: + - **In Direction**: Objects crossing the line in one direction increment the count_in counter + - **Out Direction**: Objects crossing the line in the opposite direction increment the count_out counter + - Each unique tracker_id is counted only once per crossing (prevents duplicate counting if object oscillates near line) +6. Maintains persistent counting state: + - Counts accumulate across frames for the entire video + - State persists for each video until workflow execution completes + - Separate counters for each unique video_identifier +7. Returns two count values: + - **count_in**: Total number of objects that crossed the line in the "in" direction + - **count_out**: Total number of objects that crossed the line in the "out" direction + +The line segment defines a virtual boundary in the video frame. The direction (in/out) is determined by which side of the line objects approach from - for a horizontal line, objects coming from above might count as "in" while objects from below count as "out" (or vice versa, depending on line orientation). The triggering anchor determines which point on the bounding box must cross the line for the crossing to be counted - using CENTER ensures the object is substantially across the line before counting. + +## Common Use Cases + +- **People Counting**: Count people entering and exiting buildings, stores, or events (e.g., count visitors entering store, track people entering/exiting building, monitor event attendance), enabling entry/exit counting workflows +- **Traffic Analysis**: Count vehicles passing through intersections or road segments (e.g., count vehicles crossing intersection, track traffic flow in specific directions, monitor vehicle passage at checkpoints), enabling traffic flow analysis workflows +- **Retail Analytics**: Track customer movement and foot traffic in retail spaces (e.g., count customers entering store sections, track movement between departments, monitor shopping flow patterns), enabling retail foot traffic analytics workflows +- **Security Monitoring**: Monitor entry and exit at secure areas or checkpoints (e.g., track entries to restricted areas, count people at access points, monitor checkpoint crossings), enabling security access monitoring workflows +- **Occupancy Management**: Track occupancy changes by counting objects entering and leaving spaces (e.g., count entries/exits to manage room capacity, track vehicle arrivals/departures in parking, monitor space occupancy changes), enabling occupancy tracking workflows +- **Wildlife Monitoring**: Count animals crossing defined paths or boundaries (e.g., track animal migration patterns, count wildlife crossing roads, monitor animal movement in habitats), enabling wildlife behavior analysis workflows + +## Connecting to Other Blocks + +This block receives tracked detections and video metadata, and produces count_in and count_out values: + +- **After Byte Tracker blocks** to count tracked objects crossing lines (e.g., count tracked people crossing line, track vehicle crossings with consistent IDs, monitor tracked object movements), enabling tracking-to-counting workflows +- **After object detection or instance segmentation blocks** with tracking enabled to count detected objects (e.g., count detected vehicles, track people crossings, monitor object movements), enabling detection-to-counting workflows +- **Before visualization blocks** to display line counter information (e.g., visualize line and counts, display crossing statistics, show counting results), enabling counting visualization workflows +- **Before data storage blocks** to record counting data (e.g., log entry/exit counts, store traffic statistics, record occupancy metrics), enabling counting data logging workflows +- **Before notification blocks** to alert on count thresholds or events (e.g., alert when count exceeds limit, notify on occupancy changes, trigger actions based on counts), enabling count-based notification workflows +- **Before analysis blocks** to process counting metrics (e.g., analyze traffic patterns, process occupancy data, work with counting statistics), enabling counting analysis workflows + +## Requirements + +This block requires tracked detections with tracker_id information (detections must come from a tracking block like Byte Tracker). The line must be defined as a list of exactly 2 points, where each point is a list or tuple of exactly 2 coordinates (x, y). The block requires video metadata with video_identifier to maintain separate counting state for different videos. The block maintains persistent counting state across frames for each video, so it should be used in video workflows where frames are processed sequentially. For accurate counting, detections should be provided consistently across frames with valid tracker IDs. +""" + + +class LineCounterManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Line Counter", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "video", + "icon": "far fa-arrow-down-up-across-line", + "blockPriority": 2, + }, + } + ) + type: Literal["roboflow_core/line_counter@v1"] + metadata: Selector(kind=[VIDEO_METADATA_KIND]) = Field( + description="Video metadata containing video_identifier to maintain separate counting state for different videos. Required for persistent counting across frames.", + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Tracked object detection or instance segmentation predictions. Must include tracker_id information from a tracking block. Objects are counted when their triggering anchor point crosses the line segment.", + examples=["$steps.object_detection_model.predictions"], + ) + + line_segment: Union[list, Selector(kind=[LIST_OF_VALUES_KIND]), Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + description="Line segment defined by exactly two points, each with [x, y] coordinates. Objects crossing from one side count as 'in', objects crossing from the other side count as 'out'. Example: [[0, 100], [500, 100]] creates a horizontal line at y=100. Crossing direction depends on which side objects approach from.", + examples=[[[0, 50], [500, 50]], "$inputs.zones"], + ) + triggering_anchor: Union[str, Selector(kind=[STRING_KIND]), Literal[tuple(sv.Position.list())]] = Field( # type: ignore + description="Point on the bounding box that must cross the line for counting. Options: CENTER (default), BOTTOM_CENTER, TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, etc. CENTER ensures the object is substantially across the line before counting, reducing false positives from objects near but not fully crossing the line.", + default="CENTER", + examples=["CENTER"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY_COUNT_IN, + kind=[INTEGER_KIND], + ), + OutputDefinition( + name=OUTPUT_KEY_COUNT_OUT, + kind=[INTEGER_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class LineCounterBlockV1(WorkflowBlock): + def __init__(self): + self._batch_of_line_zones: Dict[str, LeanLineZone] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return LineCounterManifest + + def run( + self, + detections: Union[Detections, InstanceDetections], + metadata: VideoMetadata, + line_segment: List[Tuple[int, int]], + triggering_anchor: str = "CENTER", + ) -> BlockResult: + n = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + tracker_ids = [ + box_metadata.get("tracker_id") for box_metadata in bboxes_metadata + ] + if n > 0 and any(tracker_id is None for tracker_id in tracker_ids): + raise ValueError( + f"tracker_id not initialized, {self.__class__.__name__} requires detections to be tracked" + ) + if metadata.video_identifier not in self._batch_of_line_zones: + if not isinstance(line_segment, list) or len(line_segment) != 2: + raise ValueError( + f"{self.__class__.__name__} requires line zone to be a list containing exactly 2 points" + ) + if any(not isinstance(e, list) or len(e) != 2 for e in line_segment): + raise ValueError( + f"{self.__class__.__name__} requires each point of line zone to be a list containing exactly 2 coordinates" + ) + if any( + not isinstance(e[0], (int, float)) or not isinstance(e[1], (int, float)) + for e in line_segment + ): + raise ValueError( + f"{self.__class__.__name__} requires each coordinate of line zone to be a number" + ) + self._batch_of_line_zones[metadata.video_identifier] = LeanLineZone( + start=tuple(line_segment[0]), + end=tuple(line_segment[1]), + triggering_anchors=[sv.Position(triggering_anchor)], + ) + line_zone = self._batch_of_line_zones[metadata.video_identifier] + + # The crossing state is host-based, so transfer box geometry once and + # run the lean vectorized anchor computation directly on numpy arrays. + xyxy_host = detections.xyxy.detach().to("cpu").numpy().astype(float) + tracker_np = np.array( + [int(tracker_id) for tracker_id in tracker_ids], dtype=int + ) + line_zone.trigger(xyxy_host, tracker_np) + + return { + OUTPUT_KEY_COUNT_IN: line_zone.in_count, + OUTPUT_KEY_COUNT_OUT: line_zone.out_count, + } diff --git a/inference/core/workflows/core_steps/analytics/line_counter/v2_tensor.py b/inference/core/workflows/core_steps/analytics/line_counter/v2_tensor.py new file mode 100644 index 0000000000..4c43925af9 --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/line_counter/v2_tensor.py @@ -0,0 +1,282 @@ +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field +from typing_extensions import Literal, Type + +from inference.core.workflows.core_steps.analytics._zone_geometry import ( + LeanLineZone, + empty_detections_like, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, + WorkflowImageSelector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY_COUNT_IN: str = "count_in" +OUTPUT_KEY_COUNT_OUT: str = "count_out" +OUTPUT_KEY_DETECTIONS_IN: str = "detections_in" +OUTPUT_KEY_DETECTIONS_OUT: str = "detections_out" +IN: str = "in" +OUT: str = "out" +DETECTIONS_IN_OUT_PARAM: str = "in_out" +SHORT_DESCRIPTION = "Count detections passing a line." +LONG_DESCRIPTION = """ +Count objects crossing a defined line segment in video using tracked detections, maintaining separate counts for objects crossing in opposite directions (in and out), and outputting both count values and the actual detection objects that crossed the line for traffic analysis, people counting, entry/exit monitoring, and directional flow measurement workflows. + +## How This Block Works + +This block counts objects that cross a line segment by tracking their movement across video frames. The block: + +1. Receives tracked detection predictions with unique tracker IDs and an image with embedded video metadata +2. Extracts video metadata from the image: + - Accesses video_metadata from the WorkflowImageData object + - Extracts video_identifier to maintain separate counting state for different videos + - Uses video metadata to initialize and manage line zone state per video +3. Validates that detections have tracker IDs (required for tracking object movement across frames) +4. Initializes or retrieves a line zone for the video: + - Creates a LineZone from two coordinate points defining the line segment + - Configures triggering anchor point if specified (optional - if not specified, uses default anchor behavior) + - Stores line zone configuration per video using video_identifier + - Maintains separate counting state for each video +5. Monitors object positions across frames: + - Tracks each object's position using its unique tracker_id + - Detects when an object's triggering anchor point (if specified) or default anchor crosses the line + - Determines crossing direction based on which side of the line the object approaches from +6. Counts line crossings: + - **In Direction**: Objects crossing the line in one direction increment the count_in counter + - **Out Direction**: Objects crossing the line in the opposite direction increment the count_out counter + - Each unique tracker_id is counted only once per crossing (prevents duplicate counting if object oscillates near line) +7. Identifies crossing detections: + - Creates masks identifying which detections crossed in each direction in the current frame + - Filters detections to separate those that crossed "in" from those that crossed "out" + - Returns the actual detection objects (not just counts) for further processing +8. Maintains persistent counting state: + - Counts accumulate across frames for the entire video + - State persists for each video until workflow execution completes + - Separate counters for each unique video_identifier +9. Returns four outputs: + - **count_in**: Total number of objects that crossed the line in the "in" direction (cumulative across video) + - **count_out**: Total number of objects that crossed the line in the "out" direction (cumulative across video) + - **detections_in**: Detection objects that crossed the line in the "in" direction (current frame crossings) + - **detections_out**: Detection objects that crossed the line in the "out" direction (current frame crossings) + +The line segment defines a virtual boundary in the video frame. The direction (in/out) is determined by which side of the line objects approach from - for a horizontal line, objects coming from above might count as "in" while objects from below count as "out" (or vice versa, depending on line orientation). The triggering anchor (if specified) determines which point on the bounding box must cross the line for the crossing to be counted - if not specified, the line zone uses its default anchor behavior. The count outputs provide cumulative totals across the video, while the detection outputs provide the actual objects that crossed in the current frame, enabling further analysis or visualization of crossing events. + +## Common Use Cases + +- **People Counting**: Count people entering and exiting buildings, stores, or events (e.g., count visitors entering store, track people entering/exiting building, monitor event attendance), enabling entry/exit counting workflows +- **Traffic Analysis**: Count vehicles passing through intersections or road segments (e.g., count vehicles crossing intersection, track traffic flow in specific directions, monitor vehicle passage at checkpoints), enabling traffic flow analysis workflows +- **Retail Analytics**: Track customer movement and foot traffic in retail spaces (e.g., count customers entering store sections, track movement between departments, monitor shopping flow patterns), enabling retail foot traffic analytics workflows +- **Security Monitoring**: Monitor entry and exit at secure areas or checkpoints (e.g., track entries to restricted areas, count people at access points, monitor checkpoint crossings), enabling security access monitoring workflows +- **Occupancy Management**: Track occupancy changes by counting objects entering and leaving spaces (e.g., count entries/exits to manage room capacity, track vehicle arrivals/departures in parking, monitor space occupancy changes), enabling occupancy tracking workflows +- **Wildlife Monitoring**: Count animals crossing defined paths or boundaries (e.g., track animal migration patterns, count wildlife crossing roads, monitor animal movement in habitats), enabling wildlife behavior analysis workflows + +## Connecting to Other Blocks + +This block receives tracked detections and an image with embedded video metadata, and produces count_in, count_out, detections_in, and detections_out: + +- **After Byte Tracker blocks** to count tracked objects crossing lines (e.g., count tracked people crossing line, track vehicle crossings with consistent IDs, monitor tracked object movements), enabling tracking-to-counting workflows +- **After object detection or instance segmentation blocks** with tracking enabled to count detected objects (e.g., count detected vehicles, track people crossings, monitor object movements), enabling detection-to-counting workflows +- **Using detections_in or detections_out outputs** to process or visualize objects that crossed the line (e.g., visualize objects that crossed, analyze crossing objects, filter for crossing events), enabling crossing object analysis workflows +- **Before visualization blocks** to display line counter information and crossing objects (e.g., visualize line and counts, display crossing statistics, show crossing objects with annotations), enabling counting visualization workflows +- **Before data storage blocks** to record counting data and crossing events (e.g., log entry/exit counts, store traffic statistics, record crossing objects with metadata), enabling counting data logging workflows +- **Before notification blocks** to alert on count thresholds or crossing events (e.g., alert when count exceeds limit, notify on specific object crossings, trigger actions based on counts), enabling count-based notification workflows + +## Version Differences + +**Enhanced from v1:** + +- **Detection Outputs**: Adds two new outputs (`detections_in` and `detections_out`) that provide the actual detection objects that crossed the line in each direction, not just count totals, enabling downstream processing and visualization of crossing objects +- **Simplified Input**: Uses `image` input that contains embedded video metadata instead of requiring a separate `metadata` field, simplifying workflow connections and reducing input complexity +- **Optional Triggering Anchor**: Makes `triggering_anchor` optional (default None) instead of required, allowing the line zone to use its default anchor behavior when no specific anchor is needed +- **Improved Integration**: Better integration with image-based workflows since video metadata is accessed directly from the image object rather than requiring separate metadata input + +## Requirements + +This block requires tracked detections with tracker_id information (detections must come from a tracking block like Byte Tracker). The line must be defined as a list of exactly 2 points, where each point is a list or tuple of exactly 2 coordinates (x, y). The image's video_metadata should include video_identifier to maintain separate counting state for different videos. The block maintains persistent counting state across frames for each video, so it should be used in video workflows where frames are processed sequentially. For accurate counting, detections should be provided consistently across frames with valid tracker IDs. +""" + + +class LineCounterManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Line Counter", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "video", + "icon": "far fa-arrow-down-up-across-line", + "blockPriority": 2, + }, + } + ) + type: Literal["roboflow_core/line_counter@v2"] + image: WorkflowImageSelector = Field( + description="Image with embedded video metadata. The video_metadata contains video_identifier to maintain separate counting state for different videos. Required for persistent counting across frames.", + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Tracked object detection or instance segmentation predictions. Must include tracker_id information from a tracking block. Objects are counted when their triggering anchor point (if specified) crosses the line segment. The detections_in and detections_out outputs provide the actual detection objects that crossed in each direction.", + examples=["$steps.object_detection_model.predictions"], + ) + + line_segment: Union[list, Selector(kind=[LIST_OF_VALUES_KIND]), Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + description="Line segment defined by exactly two points, each with [x, y] coordinates. Objects crossing from one side count as 'in', objects crossing from the other side count as 'out'. Example: [[0, 100], [500, 100]] creates a horizontal line at y=100. Crossing direction depends on which side objects approach from.", + examples=[[[0, 50], [500, 50]], "$inputs.zones"], + ) + triggering_anchor: Optional[Union[str, Selector(kind=[STRING_KIND]), Literal[tuple(sv.Position.list())]]] = Field( # type: ignore + description="Optional point on the bounding box that must cross the line for counting. If not specified (None), the line zone uses its default anchor behavior. Options when specified: CENTER, BOTTOM_CENTER, TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, etc. Specifying CENTER ensures the object is substantially across the line before counting, reducing false positives from objects near but not fully crossing the line.", + default=None, + examples=["CENTER", None], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY_COUNT_IN, + kind=[INTEGER_KIND], + ), + OutputDefinition( + name=OUTPUT_KEY_COUNT_OUT, + kind=[INTEGER_KIND], + ), + OutputDefinition( + name=OUTPUT_KEY_DETECTIONS_IN, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + OutputDefinition( + name=OUTPUT_KEY_DETECTIONS_OUT, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class LineCounterBlockV2(WorkflowBlock): + def __init__(self): + self._batch_of_line_zones: Dict[str, LeanLineZone] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return LineCounterManifest + + def run( + self, + detections: Union[Detections, InstanceDetections], + image: WorkflowImageData, + line_segment: List[Tuple[int, int]], + triggering_anchor: Optional[str] = None, + ) -> BlockResult: + n = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + tracker_ids = [ + box_metadata.get("tracker_id") for box_metadata in bboxes_metadata + ] + if n > 0 and any(tracker_id is None for tracker_id in tracker_ids): + raise ValueError( + f"tracker_id not initialized, {self.__class__.__name__} requires detections to be tracked" + ) + metadata = image.video_metadata + if metadata.video_identifier not in self._batch_of_line_zones: + if not isinstance(line_segment, list) or len(line_segment) != 2: + raise ValueError( + f"{self.__class__.__name__} requires line zone to be a list containing exactly 2 points" + ) + if any(not isinstance(e, list) or len(e) != 2 for e in line_segment): + raise ValueError( + f"{self.__class__.__name__} requires each point of line zone to be a list containing exactly 2 coordinates" + ) + if any( + not isinstance(e[0], (int, float)) or not isinstance(e[1], (int, float)) + for e in line_segment + ): + raise ValueError( + f"{self.__class__.__name__} requires each coordinate of line zone to be a number" + ) + if triggering_anchor is not None: + self._batch_of_line_zones[metadata.video_identifier] = LeanLineZone( + start=tuple(line_segment[0]), + end=tuple(line_segment[1]), + triggering_anchors=[sv.Position(triggering_anchor)], + ) + else: + self._batch_of_line_zones[metadata.video_identifier] = LeanLineZone( + start=tuple(line_segment[0]), + end=tuple(line_segment[1]), + ) + line_zone = self._batch_of_line_zones[metadata.video_identifier] + + # Transfer box geometry once; the returned host masks remain aligned to + # the tensor-native detections and can slice them directly. + xyxy_host = detections.xyxy.detach().to("cpu").numpy().astype(float) + tracker_np = np.array( + [int(tracker_id) for tracker_id in tracker_ids], dtype=int + ) + mask_in, mask_out = line_zone.trigger(xyxy_host, tracker_np) + detections_in = ( + empty_detections_like(detections) + if n and not mask_in.any() + else take_prediction_by_mask(detections, mask_in) + ) + detections_out = ( + empty_detections_like(detections) + if n and not mask_out.any() + else take_prediction_by_mask(detections, mask_out) + ) + + return { + OUTPUT_KEY_COUNT_IN: line_zone.in_count, + OUTPUT_KEY_COUNT_OUT: line_zone.out_count, + OUTPUT_KEY_DETECTIONS_IN: detections_in, + OUTPUT_KEY_DETECTIONS_OUT: detections_out, + } diff --git a/inference/core/workflows/core_steps/analytics/overlap/v1_tensor.py b/inference/core/workflows/core_steps/analytics/overlap/v1_tensor.py new file mode 100644 index 0000000000..9a00e3dc3f --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/overlap/v1_tensor.py @@ -0,0 +1,191 @@ +from functools import lru_cache +from typing import List, Optional, Union + +from pydantic import ConfigDict, Field +from typing_extensions import Literal, Type + +from inference.core.workflows.core_steps.common.tensor_native import ( + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import Selector +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "overlaps" +SHORT_DESCRIPTION = "Filter objects overlapping some other class" +LONG_DESCRIPTION = """ +Filter detection predictions to keep only objects that overlap with instances of a specified class, enabling spatial relationship filtering to identify objects that are positioned relative to other objects (e.g., people on bicycles, items on pallets, objects in containers). + +## How This Block Works + +This block filters detections based on spatial overlap relationships with a specified overlap class. The block: + +1. Takes detection predictions (object detection or instance segmentation) and an overlap class name as input +2. Separates detections into two groups: + - **Overlap class detections**: Objects matching the specified `overlap_class_name` (e.g., "bicycle", "pallet", "car") + - **Other detections**: All remaining objects that may overlap with the overlap class +3. For each overlap class detection, identifies other detections that spatially overlap with it using one of two overlap modes: + - **Center Overlap**: Checks if the center point of other detections falls within the overlap class bounding box (more precise, requires the center to be inside) + - **Any Overlap**: Checks if there's any spatial intersection between bounding boxes (more lenient, any overlap counts) +4. Collects all detections that overlap with any overlap class instance +5. Filters out the overlap class detections themselves from the output +6. Returns only the overlapping detections (objects that are positioned relative to the overlap class) + +The block effectively answers: "Which objects are overlapping with instances of class X?" For example, if you specify "bicycle" as the overlap class, the block finds people or other objects that overlap with bicycles, but removes the bicycles themselves from the output. This enables workflows to identify objects that have spatial relationships with specific reference classes, such as identifying items on surfaces, objects in containers, or people on vehicles. + +## Common Use Cases + +- **Person-on-Vehicle Detection**: Identify people on bicycles, motorcycles, or other vehicles by using the vehicle class as the overlap class (e.g., filter for people overlapping with "bicycle" detections), enabling detection of riders, passengers, or people using vehicles +- **Items on Surfaces**: Find objects positioned on pallets, tables, or shelves by using the surface class as the overlap class (e.g., filter for items overlapping with "pallet" detections), enabling inventory tracking, object counting on surfaces, or surface occupancy analysis +- **Objects in Containers**: Identify items inside containers, boxes, or vehicles by using the container class as the overlap class (e.g., filter for objects overlapping with "container" detections), enabling content detection, loading verification, or container monitoring +- **Spatial Relationship Filtering**: Filter detections based on proximity or containment relationships (e.g., find all objects that are inside or overlapping with a specific class), enabling conditional processing based on spatial arrangements +- **Nested Object Detection**: Identify objects that are part of or attached to other objects (e.g., find equipment attached to vehicles, accessories on people), enabling detection of composite objects or object relationships +- **Zone-Based Filtering**: Use overlap class as a reference zone to find objects that intersect with specific regions (e.g., filter objects overlapping with "parking_space" class), enabling zone-based analysis and conditional detection filtering + +## Connecting to Other Blocks + +The filtered overlapping detections from this block can be connected to: + +- **Detection model blocks** (e.g., Object Detection Model, Instance Segmentation Model) to receive predictions that are filtered to show only objects overlapping with a specified reference class, enabling spatial relationship analysis +- **Visualization blocks** (e.g., Bounding Box Visualization, Polygon Visualization, Label Visualization) to display only the overlapping objects, highlighting objects that have spatial relationships with the reference class +- **Counting and analytics blocks** (e.g., Line Counter, Time in Zone, Velocity) to count or analyze only overlapping objects (e.g., count people on bicycles, track items on pallets), providing metrics for spatially related objects +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload, Webhook Sink) to save or transmit filtered overlapping detection results, storing data about objects with specific spatial relationships +- **Filtering blocks** (e.g., Detections Filter) to apply additional filtering criteria to the overlapping detections, enabling multi-stage filtering workflows +- **Flow control blocks** (e.g., Continue If) to conditionally trigger downstream processing based on whether overlapping objects are detected, enabling workflows that respond to spatial relationships +""" + + +class OverlapManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Overlap Filter", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "flow_control", + "icon": "far fa-square-o", + "blockPriority": 1.5, + }, + } + ) + type: Literal["roboflow_core/overlap@v1"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Detection predictions (object detection or instance segmentation) containing objects that may overlap with the specified overlap class. The block identifies detections matching the overlap_class_name and finds other detections that spatially overlap with them. Only the overlapping detections (not the overlap class itself) are returned in the output.", + examples=["$steps.object_detection_model.predictions"], + ) + overlap_type: Literal["Center Overlap", "Any Overlap"] = Field( + default="Center Overlap", + description="Method for determining spatial overlap between detections. 'Center Overlap' checks if the center point of other detections falls within the overlap class bounding box (more precise, requires center to be inside). 'Any Overlap' checks if there's any spatial intersection between bounding boxes (more lenient, any overlap counts). Center Overlap is stricter and better for containment relationships, while Any Overlap is more inclusive and better for detecting any proximity or partial overlap.", + examples=["Center Overlap", "Any Overlap"], + ) + overlap_class_name: Union[str] = Field( + description="Class name of the reference objects used for overlap detection. Detections matching this class name are used as reference points, and other detections that overlap with these reference objects are kept in the output. The overlap class detections themselves are removed from the results. Example: Use 'bicycle' to find people or objects overlapping with bicycles; use 'pallet' to find items on pallets.", + json_schema_extra={ + "hide_description": True, + }, + ) + + @classmethod + @lru_cache(maxsize=None) + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class OverlapBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return OverlapManifest + + @classmethod + def coords_overlap( + cls, + overlap: list[int], + other: list[int], + overlap_type: Literal["Center Overlap", "Any Overlap"], + ): + + # coords are [x1, y1, x2, y2] + if overlap_type == "Center Overlap": + size = [other[2] - other[0], other[3] - other[1]] + x, y = [other[0] + size[0] / 2, other[1] + size[1] / 2] + return ( + x > overlap[0] and x < overlap[2] and y > overlap[1] and y < overlap[3] + ) + else: + return not ( + other[2] < overlap[0] + or other[0] > overlap[2] + or other[3] < overlap[1] + or other[1] > overlap[3] + ) + + def run( + self, + predictions: Union[Detections, InstanceDetections], + overlap_type: Literal["Center Overlap", "Any Overlap"], + overlap_class_name: str, + ) -> BlockResult: + + # Class names resolve through the image_metadata {class_id: name} map + # (fallback `class_` when absent). + class_names_map = (predictions.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + xyxy = predictions.xyxy.detach().to("cpu").numpy().astype(float) + class_id = predictions.class_id.detach().to("cpu").numpy() + overlaps = [] + others = {} + for i in range(len(predictions.xyxy)): + class_id_int = int(class_id[i]) + class_name = class_names_map.get(class_id_int, f"class_{class_id_int}") + if class_name == overlap_class_name: + overlaps.append(xyxy[i]) + else: + others[i] = xyxy[i] + + # set of indices representing the overlapped objects + idx = set() + for overlap in overlaps: + if not others: + break + overlapped = { + k + for k in others + if OverlapBlockV1.coords_overlap(overlap, others[k], overlap_type) + } + # once it's overlapped we don't need to check again + for k in overlapped: + del others[k] + + idx = idx.union(overlapped) + + return {OUTPUT_KEY: take_prediction_by_indices(predictions, sorted(idx))} diff --git a/inference/core/workflows/core_steps/analytics/path_deviation/v1_tensor.py b/inference/core/workflows/core_steps/analytics/path_deviation/v1_tensor.py new file mode 100644 index 0000000000..c4126f0a9a --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/path_deviation/v1_tensor.py @@ -0,0 +1,271 @@ +import copy +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field +from typing_extensions import Literal, Type + +from inference.core.workflows.execution_engine.constants import ( + PATH_DEVIATION_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + VideoMetadata, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + LIST_OF_VALUES_KIND, + STRING_KIND, + VIDEO_METADATA_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "path_deviation_detections" +SHORT_DESCRIPTION = "Calculate Frรฉchet distance of object from the reference path." +LONG_DESCRIPTION = """ +Measure how closely tracked objects follow a reference path by calculating the Frรฉchet distance between the object's actual trajectory and the expected reference path, enabling path compliance monitoring, route deviation detection, quality control in automated systems, and behavioral analysis workflows. + +## How This Block Works + +This block compares the actual movement path of tracked objects against a predefined reference path to measure deviation. The block: + +1. Receives tracked detection predictions with unique tracker IDs, video metadata, and a reference path definition +2. Validates that detections have tracker IDs (required for tracking object movement across frames) +3. Initializes or retrieves path tracking state for the video: + - Maintains a history of positions for each tracked object per video + - Stores object paths using video_identifier to separate state for different videos + - Creates new path tracking entries for objects appearing for the first time +4. Extracts anchor point coordinates for each detection: + - Uses the triggering_anchor to determine which point on the bounding box to track (default: CENTER) + - Gets the (x, y) coordinates of the anchor point for each detection in the current frame + - The anchor point represents the position of the object used for path comparison +5. Accumulates object paths over time: + - Appends each object's anchor point to its path history as frames are processed + - Maintains separate path histories for each unique tracker_id + - Builds complete trajectory paths by accumulating positions across all processed frames +6. Calculates Frรฉchet distance for each tracked object: + - **Frรฉchet Distance**: Measures the similarity between two curves (paths) considering both location and ordering of points + - Compares the object's accumulated path (actual trajectory) against the reference path (expected trajectory) + - Uses dynamic programming to compute the minimum "leash length" required to traverse both paths simultaneously + - Accounts for the order of points along each path, not just point-to-point distances + - Lower values indicate the object follows the reference path closely, higher values indicate greater deviation +7. Stores path deviation in detection metadata: + - Adds the Frรฉchet distance value to each detection's metadata + - Each detection includes path_deviation representing how much it deviates from the reference path + - Distance is measured in pixels (same units as image coordinates) +8. Maintains persistent path tracking: + - Path histories accumulate across frames for the entire video + - Each object's deviation is calculated based on its complete path from the start of tracking + - Separate tracking state maintained for each video_identifier +9. Returns detections enhanced with path deviation information: + - Outputs detection objects with added path_deviation metadata + - Each detection now includes the Frรฉchet distance measuring its deviation from the reference path + +The Frรฉchet distance is a metric that measures the similarity between two curves by finding the minimum length of a "leash" that connects a point moving along one curve to a point moving along the other curve, where both points move forward along their respective curves. Unlike simple Euclidean distance, Frรฉchet distance considers the ordering and continuity of points along paths, making it ideal for comparing trajectories where the sequence of movement matters. An object that follows the reference path exactly will have a Frรฉchet distance of 0, while objects that deviate significantly will have larger distances. + +## Common Use Cases + +- **Path Compliance Monitoring**: Monitor whether vehicles, robots, or objects follow predefined routes (e.g., verify vehicles stay in lanes, check robots follow programmed paths, ensure objects follow expected routes), enabling compliance monitoring workflows +- **Quality Control**: Detect deviations in manufacturing or assembly processes where objects should follow specific paths (e.g., detect conveyor belt deviations, monitor assembly line paths, check product movement patterns), enabling quality control workflows +- **Traffic Analysis**: Analyze vehicle movement patterns and detect lane departures or route deviations (e.g., detect vehicles leaving lanes, monitor route adherence, analyze traffic pattern compliance), enabling traffic analysis workflows +- **Security Monitoring**: Detect suspicious movement patterns or deviations from expected paths in security scenarios (e.g., detect unauthorized route deviations, monitor perimeter breach attempts, track movement compliance), enabling security monitoring workflows +- **Automated Systems**: Monitor and validate that automated systems (robots, AGVs, drones) follow expected paths correctly (e.g., verify robot navigation accuracy, check automated vehicle paths, validate drone flight paths), enabling automated system validation workflows +- **Behavioral Analysis**: Study movement patterns and path adherence in behavioral research (e.g., analyze animal movement patterns, study path following behavior, measure route preference deviations), enabling behavioral research workflows + +## Connecting to Other Blocks + +This block receives tracked detections, video metadata, and a reference path, and produces detections enhanced with path_deviation metadata: + +- **After Byte Tracker blocks** to measure path deviation for tracked objects (e.g., measure tracked vehicle path compliance, analyze tracked person route adherence, monitor tracked object path deviations), enabling tracking-to-path-analysis workflows +- **After object detection or instance segmentation blocks** with tracking enabled to analyze movement paths (e.g., analyze vehicle paths, track object route compliance, measure path deviations), enabling detection-to-path-analysis workflows +- **Before visualization blocks** to display path deviation information (e.g., visualize paths and deviations, display reference and actual paths, show deviation metrics), enabling path deviation visualization workflows +- **Before logic blocks** like Continue If to make decisions based on path deviation thresholds (e.g., continue if deviation exceeds limit, filter based on path compliance, trigger actions on route violations), enabling path-based decision workflows +- **Before notification blocks** to alert on path deviations or compliance violations (e.g., alert on route deviations, notify on path compliance issues, trigger deviation-based alerts), enabling path-based notification workflows +- **Before data storage blocks** to record path deviation measurements (e.g., log path compliance data, store deviation statistics, record route adherence metrics), enabling path deviation data logging workflows + +## Requirements + +This block requires tracked detections with tracker_id information (detections must come from a tracking block like Byte Tracker). The reference path must be defined as a list of at least 2 points, where each point is a tuple or list of exactly 2 coordinates (x, y). The block requires video metadata with video_identifier to maintain separate path tracking state for different videos. The block maintains persistent path tracking across frames for each video, accumulating complete trajectories, so it should be used in video workflows where frames are processed sequentially. For accurate path deviation measurement, detections should be provided consistently across frames with valid tracker IDs. The Frรฉchet distance is calculated in pixels (same units as image coordinates). +""" + + +class PathDeviationManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Path Deviation", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "video", + "icon": "far fa-tower-observation", + }, + } + ) + type: Literal["roboflow_core/path_deviation_analytics@v1"] + metadata: Selector(kind=[VIDEO_METADATA_KIND]) = Field( + description="Video metadata containing video_identifier to maintain separate path tracking state for different videos. Required for persistent path accumulation across frames.", + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Tracked object detection or instance segmentation predictions. Must include tracker_id information from a tracking block. The block tracks anchor point positions across frames to build object trajectories and compares them against the reference path. Output detections include path_deviation metadata containing the Frรฉchet distance from the reference path.", + examples=["$steps.object_detection_model.predictions"], + ) + triggering_anchor: Union[str, Selector(kind=[STRING_KIND]), Literal[tuple(sv.Position.list())]] = Field( # type: ignore + description="Point on the bounding box used to track object position for path calculation. Options: CENTER (default), BOTTOM_CENTER, TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, etc. This anchor point's coordinates are accumulated over frames to build the object's trajectory path, which is compared against the reference path using Frรฉchet distance.", + default="CENTER", + examples=["CENTER"], + ) + reference_path: Union[list, Selector(kind=[LIST_OF_VALUES_KIND]), Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + description="Expected reference path as a list of at least 2 points, where each point is a tuple or list of [x, y] coordinates. Example: [(100, 200), (200, 300), (300, 400)] defines a path with 3 points. The Frรฉchet distance measures how closely tracked objects follow this reference path. Points should be ordered along the expected trajectory.", + examples=[[(100, 200), (200, 300), (300, 400)], "$inputs.expected_path"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class PathDeviationAnalyticsBlockV1(WorkflowBlock): + def __init__(self): + self._object_paths: Dict[ + str, Dict[Union[int, str], List[Tuple[float, float]]] + ] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return PathDeviationManifest + + def run( + self, + detections: Union[Detections, InstanceDetections], + metadata: VideoMetadata, + triggering_anchor: str, + reference_path: List[Tuple[int, int]], + ) -> BlockResult: + n = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + tracker_ids = [ + box_metadata.get("tracker_id") for box_metadata in bboxes_metadata + ] + if n > 0 and any(tracker_id is None for tracker_id in tracker_ids): + raise ValueError( + f"tracker_id not initialized, {self.__class__.__name__} requires detections to be tracked" + ) + + video_id = metadata.video_identifier + if video_id not in self._object_paths: + self._object_paths[video_id] = {} + + # sv anchor resolution is numpy-based: materialise a minimal sv.Detections + # (only box coordinates are needed) and reuse get_anchors_coordinates. + # The returned points are aligned to the input order. + sv_input = sv.Detections( + xyxy=detections.xyxy.detach().to("cpu").numpy().astype(float), + ) + anchor_points = sv_input.get_anchors_coordinates(anchor=triggering_anchor) + new_bboxes_metadata = [] + for i, tracker_id in enumerate(tracker_ids): + box_metadata = dict(bboxes_metadata[i]) + anchor_point = anchor_points[i] + if tracker_id not in self._object_paths[video_id]: + self._object_paths[video_id][tracker_id] = [] + self._object_paths[video_id][tracker_id].append(anchor_point) + + object_path = np.array(self._object_paths[video_id][tracker_id]) + ref_path = np.array(reference_path) + + frechet_distance = self._calculate_frechet_distance(object_path, ref_path) + box_metadata[PATH_DEVIATION_KEY_IN_SV_DETECTIONS] = float(frechet_distance) + new_bboxes_metadata.append(box_metadata) + + # Do not mutate the caller's native object: a shallow copy keeps the + # tensors/masks shared by reference and only swaps the bboxes_metadata list. + result_detections = copy.copy(detections) + result_detections.bboxes_metadata = new_bboxes_metadata + return {OUTPUT_KEY: result_detections} + + def _calculate_frechet_distance( + self, path1: np.ndarray, path2: np.ndarray + ) -> float: + dist_matrix = np.ones((len(path1), len(path2))) * -1 + return self._compute_distance( + dist_matrix, len(path1) - 1, len(path2) - 1, path1, path2 + ) + + def _compute_distance( + self, + dist_matrix: np.ndarray, + i: int, + j: int, + path1: np.ndarray, + path2: np.ndarray, + ) -> float: + if dist_matrix[i, j] > -1: + return dist_matrix[i, j] + elif i == 0 and j == 0: + dist_matrix[i, j] = self._euclidean_distance(path1[0], path2[0]) + elif i > 0 and j == 0: + dist_matrix[i, j] = max( + self._compute_distance(dist_matrix, i - 1, 0, path1, path2), + self._euclidean_distance(path1[i], path2[0]), + ) + elif i == 0 and j > 0: + dist_matrix[i, j] = max( + self._compute_distance(dist_matrix, 0, j - 1, path1, path2), + self._euclidean_distance(path1[0], path2[j]), + ) + elif i > 0 and j > 0: + dist_matrix[i, j] = max( + min( + self._compute_distance(dist_matrix, i - 1, j, path1, path2), + self._compute_distance(dist_matrix, i - 1, j - 1, path1, path2), + self._compute_distance(dist_matrix, i, j - 1, path1, path2), + ), + self._euclidean_distance(path1[i], path2[j]), + ) + else: + dist_matrix[i, j] = float("inf") + return dist_matrix[i, j] + + def _euclidean_distance(self, point1: np.ndarray, point2: np.ndarray) -> float: + return np.sqrt(np.sum((point1 - point2) ** 2)) diff --git a/inference/core/workflows/core_steps/analytics/path_deviation/v2_tensor.py b/inference/core/workflows/core_steps/analytics/path_deviation/v2_tensor.py new file mode 100644 index 0000000000..6a50b96fab --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/path_deviation/v2_tensor.py @@ -0,0 +1,284 @@ +from dataclasses import replace as dataclass_replace +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field +from typing_extensions import Literal, Type + +from inference.core.workflows.execution_engine.constants import ( + PATH_DEVIATION_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, + WorkflowImageSelector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "path_deviation_detections" +SHORT_DESCRIPTION = "Calculate Frรฉchet distance of object from the reference path." +LONG_DESCRIPTION = """ +Measure how closely tracked objects follow a reference path by calculating the Frรฉchet distance between the object's actual trajectory and the expected reference path, enabling path compliance monitoring, route deviation detection, quality control in automated systems, and behavioral analysis workflows. + +## How This Block Works + +This block compares the actual movement path of tracked objects against a predefined reference path to measure deviation. The block: + +1. Receives tracked detection predictions with unique tracker IDs, an image with embedded video metadata, and a reference path definition +2. Extracts video metadata from the image: + - Accesses video_metadata from the WorkflowImageData object + - Extracts video_identifier to maintain separate path tracking state for different videos + - Uses video metadata to initialize and manage path tracking state per video +3. Validates that detections have tracker IDs (required for tracking object movement across frames) +4. Initializes or retrieves path tracking state for the video: + - Maintains a history of positions for each tracked object per video + - Stores object paths using video_identifier to separate state for different videos + - Creates new path tracking entries for objects appearing for the first time +5. Extracts anchor point coordinates for each detection: + - Uses the triggering_anchor to determine which point on the bounding box to track (default: CENTER) + - Gets the (x, y) coordinates of the anchor point for each detection in the current frame + - The anchor point represents the position of the object used for path comparison +6. Accumulates object paths over time: + - Appends each object's anchor point to its path history as frames are processed + - Maintains separate path histories for each unique tracker_id + - Builds complete trajectory paths by accumulating positions across all processed frames +7. Calculates Frรฉchet distance for each tracked object: + - **Frรฉchet Distance**: Measures the similarity between two curves (paths) considering both location and ordering of points + - Compares the object's accumulated path (actual trajectory) against the reference path (expected trajectory) + - Uses dynamic programming to compute the minimum "leash length" required to traverse both paths simultaneously + - Accounts for the order of points along each path, not just point-to-point distances + - Lower values indicate the object follows the reference path closely, higher values indicate greater deviation +8. Stores path deviation in detection metadata: + - Adds the Frรฉchet distance value to each detection's metadata + - Each detection includes path_deviation representing how much it deviates from the reference path + - Distance is measured in pixels (same units as image coordinates) +9. Maintains persistent path tracking: + - Path histories accumulate across frames for the entire video + - Each object's deviation is calculated based on its complete path from the start of tracking + - Separate tracking state maintained for each video_identifier +10. Returns detections enhanced with path deviation information: + - Outputs detection objects with added path_deviation metadata + - Each detection now includes the Frรฉchet distance measuring its deviation from the reference path + +The Frรฉchet distance is a metric that measures the similarity between two curves by finding the minimum length of a "leash" that connects a point moving along one curve to a point moving along the other curve, where both points move forward along their respective curves. Unlike simple Euclidean distance, Frรฉchet distance considers the ordering and continuity of points along paths, making it ideal for comparing trajectories where the sequence of movement matters. An object that follows the reference path exactly will have a Frรฉchet distance of 0, while objects that deviate significantly will have larger distances. + +## Common Use Cases + +- **Path Compliance Monitoring**: Monitor whether vehicles, robots, or objects follow predefined routes (e.g., verify vehicles stay in lanes, check robots follow programmed paths, ensure objects follow expected routes), enabling compliance monitoring workflows +- **Quality Control**: Detect deviations in manufacturing or assembly processes where objects should follow specific paths (e.g., detect conveyor belt deviations, monitor assembly line paths, check product movement patterns), enabling quality control workflows +- **Traffic Analysis**: Analyze vehicle movement patterns and detect lane departures or route deviations (e.g., detect vehicles leaving lanes, monitor route adherence, analyze traffic pattern compliance), enabling traffic analysis workflows +- **Security Monitoring**: Detect suspicious movement patterns or deviations from expected paths in security scenarios (e.g., detect unauthorized route deviations, monitor perimeter breach attempts, track movement compliance), enabling security monitoring workflows +- **Automated Systems**: Monitor and validate that automated systems (robots, AGVs, drones) follow expected paths correctly (e.g., verify robot navigation accuracy, check automated vehicle paths, validate drone flight paths), enabling automated system validation workflows +- **Behavioral Analysis**: Study movement patterns and path adherence in behavioral research (e.g., analyze animal movement patterns, study path following behavior, measure route preference deviations), enabling behavioral research workflows + +## Connecting to Other Blocks + +This block receives tracked detections, an image with embedded video metadata, and a reference path, and produces detections enhanced with path_deviation metadata: + +- **After Byte Tracker blocks** to measure path deviation for tracked objects (e.g., measure tracked vehicle path compliance, analyze tracked person route adherence, monitor tracked object path deviations), enabling tracking-to-path-analysis workflows +- **After object detection or instance segmentation blocks** with tracking enabled to analyze movement paths (e.g., analyze vehicle paths, track object route compliance, measure path deviations), enabling detection-to-path-analysis workflows +- **Before visualization blocks** to display path deviation information (e.g., visualize paths and deviations, display reference and actual paths, show deviation metrics), enabling path deviation visualization workflows +- **Before logic blocks** like Continue If to make decisions based on path deviation thresholds (e.g., continue if deviation exceeds limit, filter based on path compliance, trigger actions on route violations), enabling path-based decision workflows +- **Before notification blocks** to alert on path deviations or compliance violations (e.g., alert on route deviations, notify on path compliance issues, trigger deviation-based alerts), enabling path-based notification workflows +- **Before data storage blocks** to record path deviation measurements (e.g., log path compliance data, store deviation statistics, record route adherence metrics), enabling path deviation data logging workflows + +## Version Differences + +**Enhanced from v1:** + +- **Simplified Input**: Uses `image` input that contains embedded video metadata instead of requiring a separate `metadata` field, simplifying workflow connections and reducing input complexity +- **Improved Integration**: Better integration with image-based workflows since video metadata is accessed directly from the image object rather than requiring separate metadata input + +## Requirements + +This block requires tracked detections with tracker_id information (detections must come from a tracking block like Byte Tracker). The reference path must be defined as a list of at least 2 points, where each point is a tuple or list of exactly 2 coordinates (x, y). The image's video_metadata should include video_identifier to maintain separate path tracking state for different videos. The block maintains persistent path tracking across frames for each video, accumulating complete trajectories, so it should be used in video workflows where frames are processed sequentially. For accurate path deviation measurement, detections should be provided consistently across frames with valid tracker IDs. The Frรฉchet distance is calculated in pixels (same units as image coordinates). +""" + + +class PathDeviationManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Path Deviation", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "video", + "icon": "far fa-road", + "blockPriority": 3, + }, + } + ) + type: Literal["roboflow_core/path_deviation_analytics@v2"] + image: WorkflowImageSelector = Field( + description="Image with embedded video metadata. The video_metadata contains video_identifier to maintain separate path tracking state for different videos. Required for persistent path accumulation across frames.", + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Tracked object detection or instance segmentation predictions. Must include tracker_id information from a tracking block. The block tracks anchor point positions across frames to build object trajectories and compares them against the reference path. Output detections include path_deviation metadata containing the Frรฉchet distance from the reference path.", + examples=["$steps.object_detection_model.predictions"], + ) + triggering_anchor: Union[str, Selector(kind=[STRING_KIND]), Literal[tuple(sv.Position.list())]] = Field( # type: ignore + description="Point on the bounding box used to track object position for path calculation. Options include CENTER (default), BOTTOM_CENTER, TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, etc. This anchor point's coordinates are accumulated over frames to build the object's trajectory path, which is compared against the reference path using Frรฉchet distance.", + default="CENTER", + examples=["CENTER"], + ) + reference_path: Union[list, Selector(kind=[LIST_OF_VALUES_KIND]), Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + description="Expected reference path as a list of at least 2 points, where each point is a tuple or list of [x, y] coordinates. Example: [(100, 200), (200, 300), (300, 400)] defines a path with 3 points. The Frรฉchet distance measures how closely tracked objects follow this reference path. Points should be ordered along the expected trajectory.", + examples=[[(100, 200), (200, 300), (300, 400)], "$inputs.expected_path"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.2.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class PathDeviationAnalyticsBlockV2(WorkflowBlock): + def __init__(self): + self._object_paths: Dict[ + str, Dict[Union[int, str], List[Tuple[float, float]]] + ] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return PathDeviationManifest + + def run( + self, + detections: Union[Detections, InstanceDetections], + image: WorkflowImageData, + triggering_anchor: str, + reference_path: List[Tuple[int, int]], + ) -> BlockResult: + n = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + tracker_ids = [ + box_metadata.get("tracker_id") for box_metadata in bboxes_metadata + ] + if n > 0 and any(tracker_id is None for tracker_id in tracker_ids): + raise ValueError( + f"tracker_id not initialized, {self.__class__.__name__} requires detections to be tracked" + ) + metadata = image.video_metadata + video_id = metadata.video_identifier + if video_id not in self._object_paths: + self._object_paths[video_id] = {} + + # sv anchor resolution is numpy-based: materialise a minimal sv.Detections + # (only box coordinates are needed) and reuse get_anchors_coordinates. + # The returned points are aligned to the input order. + sv_input = sv.Detections( + xyxy=detections.xyxy.detach().to("cpu").numpy().astype(float), + ) + anchor_points = sv_input.get_anchors_coordinates(anchor=triggering_anchor) + new_bboxes_metadata = [] + for i, tracker_id in enumerate(tracker_ids): + box_metadata = dict(bboxes_metadata[i]) + anchor_point = anchor_points[i] + if tracker_id not in self._object_paths[video_id]: + self._object_paths[video_id][tracker_id] = [] + self._object_paths[video_id][tracker_id].append(anchor_point) + + object_path = np.array(self._object_paths[video_id][tracker_id]) + ref_path = np.array(reference_path) + + frechet_distance = self._calculate_frechet_distance(object_path, ref_path) + box_metadata[PATH_DEVIATION_KEY_IN_SV_DETECTIONS] = float(frechet_distance) + new_bboxes_metadata.append(box_metadata) + + # Do not mutate the caller's native object: dataclasses.replace keeps the + # tensors/masks shared by reference and only swaps the bboxes_metadata list. + result_detections = dataclass_replace( + detections, bboxes_metadata=new_bboxes_metadata + ) + return {OUTPUT_KEY: result_detections} + + def _calculate_frechet_distance( + self, path1: np.ndarray, path2: np.ndarray + ) -> float: + dist_matrix = np.ones((len(path1), len(path2))) * -1 + return self._compute_distance( + dist_matrix, len(path1) - 1, len(path2) - 1, path1, path2 + ) + + def _compute_distance( + self, + dist_matrix: np.ndarray, + i: int, + j: int, + path1: np.ndarray, + path2: np.ndarray, + ) -> float: + if dist_matrix[i, j] > -1: + return dist_matrix[i, j] + elif i == 0 and j == 0: + dist_matrix[i, j] = self._euclidean_distance(path1[0], path2[0]) + elif i > 0 and j == 0: + dist_matrix[i, j] = max( + self._compute_distance(dist_matrix, i - 1, 0, path1, path2), + self._euclidean_distance(path1[i], path2[0]), + ) + elif i == 0 and j > 0: + dist_matrix[i, j] = max( + self._compute_distance(dist_matrix, 0, j - 1, path1, path2), + self._euclidean_distance(path1[0], path2[j]), + ) + elif i > 0 and j > 0: + dist_matrix[i, j] = max( + min( + self._compute_distance(dist_matrix, i - 1, j, path1, path2), + self._compute_distance(dist_matrix, i - 1, j - 1, path1, path2), + self._compute_distance(dist_matrix, i, j - 1, path1, path2), + ), + self._euclidean_distance(path1[i], path2[j]), + ) + else: + dist_matrix[i, j] = float("inf") + return dist_matrix[i, j] + + def _euclidean_distance(self, point1: np.ndarray, point2: np.ndarray) -> float: + return np.sqrt(np.sum((point1 - point2) ** 2)) diff --git a/inference/core/workflows/core_steps/analytics/time_in_zone/v1_tensor.py b/inference/core/workflows/core_steps/analytics/time_in_zone/v1_tensor.py new file mode 100644 index 0000000000..e0f0d174b7 --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/time_in_zone/v1_tensor.py @@ -0,0 +1,302 @@ +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field +from typing_extensions import Literal, Type + +from inference.core.workflows.core_steps.analytics._zone_geometry import ( + LeanPolygonZone, + empty_detections_like, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + TIME_IN_ZONE_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + VideoMetadata, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + VIDEO_METADATA_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "timed_detections" +SHORT_DESCRIPTION = "Track object time in zone." +LONG_DESCRIPTION = """ +Calculate and track the time spent by tracked objects within a defined polygon zone, measure duration of object presence in specific areas, filter detections based on zone membership, reset time tracking when objects leave zones, and enable zone-based analytics, dwell time analysis, and presence monitoring workflows. + +## How This Block Works + +This block measures how long each tracked object has been inside a defined polygon zone by tracking entry and exit times for each unique track ID. The block: + +1. Receives tracked detection predictions with track IDs, an image, video metadata, and a polygon zone definition +2. Validates that detections have track IDs (tracker_id must be present): + - Requires detections to come from a tracking block (e.g., Byte Tracker) + - Each object must have a unique tracker_id that persists across frames + - Raises an error if tracker_id is missing +3. Initializes or retrieves a polygon zone for the video: + - Creates a PolygonZone object from zone coordinates for each unique video + - Validates zone coordinates (must be a list of at least 3 points, each with 2 coordinates) + - Stores zone configuration per video using video_identifier + - Configures triggering anchor point (e.g., CENTER, BOTTOM_CENTER) for zone detection +4. Initializes or retrieves time tracking state for the video: + - Maintains a dictionary tracking when each track_id entered the zone + - Stores entry timestamps per video using video_identifier + - Maintains separate tracking state for each video +5. Calculates current timestamp for time measurement: + - For video files: Calculates timestamp as frame_number / fps + - For streamed video: Uses frame_timestamp from metadata + - Provides accurate time measurement for duration calculation +6. Checks which detections are in the zone: + - Uses polygon zone trigger to test if each detection's anchor point is inside the zone + - The triggering_anchor determines which point on the bounding box is checked (CENTER, BOTTOM_CENTER, etc.) + - Returns boolean for each detection indicating zone membership +7. Updates time tracking for each tracked object: + - **For objects entering the zone**: Records entry timestamp if not already tracked + - **For objects in the zone**: Calculates time spent as current_timestamp - entry_timestamp + - **For objects leaving the zone**: + - If reset_out_of_zone_detections is True: Removes entry timestamp (resets to 0) + - If reset_out_of_zone_detections is False: Keeps entry timestamp (continues tracking) +8. Handles out-of-zone detections: + - **If remove_out_of_zone_detections is True**: Filters out detections outside the zone from output + - **If remove_out_of_zone_detections is False**: Includes out-of-zone detections with time = 0 +9. Adds time_in_zone information to each detection: + - Attaches time_in_zone value (in seconds) to each detection as metadata + - Objects in zone: Time represents duration spent in zone + - Objects outside zone: Time is 0 (if not reset) or undefined (if removed) +10. Returns detections with time_in_zone information: + - Outputs tracked detections enhanced with time_in_zone metadata + - Filtered or unfiltered based on remove_out_of_zone_detections setting + - Maintains all original detection properties plus time tracking information + +The block maintains persistent tracking state across frames, allowing accurate cumulative time measurement for objects that remain in the zone over multiple frames. Time is measured from when an object first enters the zone (based on its track_id) until the current frame, providing real-time duration tracking. The zone is defined as a polygon with multiple points, allowing flexible area definitions. The triggering anchor determines which part of the bounding box is used for zone detection, enabling different zone entry/exit behaviors based on object position. + +## Common Use Cases + +- **Dwell Time Analysis**: Measure how long objects remain in specific areas for behavior analysis (e.g., measure customer dwell time in store sections, track time spent in parking spaces, analyze time in waiting areas), enabling dwell time analytics workflows +- **Zone-Based Monitoring**: Monitor object presence in defined areas for security and safety (e.g., detect loitering in restricted areas, monitor time in danger zones, track presence in secure zones), enabling zone monitoring workflows +- **Retail Analytics**: Track customer time in different store sections for retail insights (e.g., measure time in product aisles, analyze shopping patterns, track department engagement), enabling retail analytics workflows +- **Occupancy Management**: Measure time objects spend in spaces for space utilization (e.g., track vehicle parking duration, measure table occupancy time, analyze space usage patterns), enabling occupancy management workflows +- **Safety Compliance**: Monitor time violations in restricted or time-limited zones (e.g., detect extended stays in hazardous areas, monitor time limit violations, track safety compliance), enabling safety monitoring workflows +- **Traffic Analysis**: Measure time vehicles spend in traffic zones or intersections (e.g., track time at intersections, measure queue waiting time, analyze traffic flow patterns), enabling traffic analytics workflows + +## Connecting to Other Blocks + +This block receives tracked detections, image, video metadata, and zone coordinates, and produces timed_detections with time_in_zone metadata: + +- **After Byte Tracker blocks** to measure time for tracked objects (e.g., track time in zones for tracked objects, measure dwell time with consistent IDs, analyze tracked object presence), enabling tracking-to-time workflows +- **After zone definition blocks** to apply time tracking to defined areas (e.g., measure time in polygon zones, track duration in custom zones, analyze zone-based presence), enabling zone-to-time workflows +- **Before logic blocks** like Continue If to make decisions based on time in zone (e.g., continue if time exceeds threshold, filter based on dwell time, trigger actions on time violations), enabling time-based decision workflows +- **Before analysis blocks** to analyze time-based metrics (e.g., analyze dwell time patterns, process time-in-zone data, work with duration metrics), enabling time analysis workflows +- **Before notification blocks** to alert on time violations or thresholds (e.g., alert on extended stays, notify on time limit violations, trigger time-based alerts), enabling time-based notification workflows +- **Before data storage blocks** to record time metrics (e.g., store dwell time data, log time-in-zone metrics, record duration measurements), enabling time metrics logging workflows + +## Requirements + +This block requires tracked detections with tracker_id information (detections must come from a tracking block like Byte Tracker). The zone must be defined as a list of at least 3 points, where each point is a list or tuple of exactly 2 coordinates (x, y). The block requires video metadata with frame rate (fps) for video files or frame timestamps for streamed video to calculate accurate time measurements. The block maintains persistent tracking state across frames for each video, so it should be used in video workflows where frames are processed sequentially. For accurate time measurement, detections should be provided consistently across frames with valid track IDs. +""" + + +class TimeInZoneManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Time in Zone", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "video", + "icon": "far fa-timer", + "blockPriority": 1, + }, + } + ) + type: Literal["roboflow_core/time_in_zone@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + title="Image", + description="Input image for the current video frame. Used for zone visualization and reference. The block uses the image dimensions to validate zone coordinates. The image metadata may be used for time calculation if frame timestamps are needed.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + metadata: Selector(kind=[VIDEO_METADATA_KIND]) = Field( + description="Video metadata containing frame rate (fps), frame number, frame timestamp, video identifier, and video source information required for time calculation and state management. The fps and frame_number are used for video files to calculate timestamps (timestamp = frame_number / fps). For streamed video, frame_timestamp is used directly. The video_identifier is used to maintain separate tracking state and zone configurations for different videos. The metadata must include valid fps for video files or frame_timestamp for streams to enable accurate time measurement.", + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Tracked detection predictions (object detection or instance segmentation) with tracker_id information. Detections must come from a tracking block (e.g., Byte Tracker) that has assigned unique tracker_id values that persist across frames. Each detection must have a tracker_id to enable time tracking. The block calculates time_in_zone for each tracked object based on when its track_id first entered the zone. The output will include the same detections enhanced with time_in_zone metadata (duration in seconds). If remove_out_of_zone_detections is True, only detections inside the zone are included in the output.", + examples=["$steps.object_detection_model.predictions"], + ) + zone: Union[list, Selector(kind=[LIST_OF_VALUES_KIND]), Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + description="Polygon zone coordinates defining the area for time measurement. Must be a list of at least 3 points, where each point is a list or tuple of exactly 2 coordinates [x, y] or (x, y). Coordinates should be in pixel space matching the image dimensions. Example: [[x1, y1], [x2, y2], [x3, y3], [x4, y4]] for a quadrilateral zone. The zone defines the polygon area where time tracking occurs. Objects are considered 'in zone' when their triggering_anchor point is inside this polygon. Zone coordinates are validated and a PolygonZone object is created for each video.", + examples=["$inputs.zones"], + ) + triggering_anchor: Union[str, Selector(kind=[STRING_KIND]), Literal[tuple(sv.Position.list())]] = Field( # type: ignore + description="Point on the detection bounding box that must be inside the zone to consider the object 'in zone'. Options include: 'CENTER' (default, center of bounding box), 'BOTTOM_CENTER' (bottom center point), 'TOP_CENTER' (top center point), 'CENTER_LEFT' (center left point), 'CENTER_RIGHT' (center right point), and other Position enum values. The triggering anchor determines which part of the object's bounding box is checked against the zone polygon. Use CENTER for standard zone detection, BOTTOM_CENTER for ground-level zones (e.g., tracking feet/vehicle base), or other anchors based on detection needs. Default is 'CENTER'.", + default="CENTER", + examples=["CENTER"], + ) + remove_out_of_zone_detections: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + description="If True (default), detections found outside the zone are filtered out and not included in the output. Only detections inside the zone are returned. If False, all detections are included in the output, with time_in_zone = 0 for objects outside the zone. Use True to focus analysis only on objects in the zone, or False to maintain all detections with zone status. Default is True for cleaner output focused on zone activity.", + default=True, + examples=[True, False], + ) + reset_out_of_zone_detections: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + description="If True (default), when a tracked object leaves the zone, its time tracking is reset (entry timestamp is cleared). When the object re-enters the zone, time tracking starts from 0 again. If False, time tracking continues even after leaving the zone, and re-entry maintains cumulative time. Use True to measure current continuous time in zone (resets on exit), or False to measure cumulative time across multiple entries. Default is True for measuring continuous presence duration.", + default=True, + examples=[True, False], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class TimeInZoneBlockV1(WorkflowBlock): + def __init__(self): + self._batch_of_tracked_ids_in_zone: Dict[str, Dict[Union[int, str], float]] = {} + self._batch_of_polygon_zones: Dict[str, LeanPolygonZone] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return TimeInZoneManifest + + def run( + self, + image: WorkflowImageData, + detections: Union[Detections, InstanceDetections], + metadata: VideoMetadata, + zone: List[Tuple[int, int]], + triggering_anchor: str, + remove_out_of_zone_detections: bool, + reset_out_of_zone_detections: bool, + ) -> BlockResult: + n = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + tracker_ids = [ + box_metadata.get("tracker_id") for box_metadata in bboxes_metadata + ] + if n > 0 and any(tracker_id is None for tracker_id in tracker_ids): + raise ValueError( + f"tracker_id not initialized, {self.__class__.__name__} requires detections to be tracked" + ) + if metadata.video_identifier not in self._batch_of_polygon_zones: + if not isinstance(zone, list) or len(zone) < 3: + raise ValueError( + f"{self.__class__.__name__} requires zone to be a list containing more than 2 points" + ) + if any( + (not isinstance(e, list) and not isinstance(e, tuple)) or len(e) != 2 + for e in zone + ): + raise ValueError( + f"{self.__class__.__name__} requires each point of zone to be a list containing exactly 2 coordinates" + ) + if any( + not isinstance(e[0], (int, float)) or not isinstance(e[1], (int, float)) + for e in zone + ): + raise ValueError( + f"{self.__class__.__name__} requires each coordinate of zone to be a number" + ) + self._batch_of_polygon_zones[metadata.video_identifier] = LeanPolygonZone( + polygon=np.array(zone), + triggering_anchors=(sv.Position(triggering_anchor),), + ) + polygon_zone = self._batch_of_polygon_zones[metadata.video_identifier] + tracked_ids_in_zone = self._batch_of_tracked_ids_in_zone.setdefault( + metadata.video_identifier, {} + ) + if metadata.comes_from_video_file and metadata.fps != 0: + ts_end = metadata.frame_number / metadata.fps + else: + ts_end = metadata.frame_timestamp.timestamp() + + # Transfer box geometry once and evaluate zone membership with the lean + # vectorized host implementation. + xyxy_host = detections.xyxy.detach().to("cpu").numpy().astype(float) + is_in_zone_mask = polygon_zone.trigger(xyxy_host) + surviving_mask = np.zeros(n, dtype=bool) + surviving_times: Dict[int, float] = {} + for i, is_in_zone, tracker_id in zip( + range(n), + is_in_zone_mask, + tracker_ids, + ): + if ( + not is_in_zone + and tracker_id in tracked_ids_in_zone + and reset_out_of_zone_detections + ): + del tracked_ids_in_zone[tracker_id] + if not is_in_zone and remove_out_of_zone_detections: + continue + + time_in_zone = 0.0 + if is_in_zone: + ts_start = tracked_ids_in_zone.setdefault(tracker_id, ts_end) + time_in_zone = ts_end - ts_start + elif tracker_id in tracked_ids_in_zone: + del tracked_ids_in_zone[tracker_id] + surviving_mask[i] = True + surviving_times[i] = time_in_zone + result_detections = ( + empty_detections_like(detections) + if n and not surviving_mask.any() + else take_prediction_by_mask(detections, surviving_mask) + ) + result_bboxes_metadata = result_detections.bboxes_metadata or [ + {} for _ in range(int(result_detections.xyxy.shape[0])) + ] + for position, source_index in enumerate(np.nonzero(surviving_mask)[0].tolist()): + result_bboxes_metadata[position][TIME_IN_ZONE_KEY_IN_SV_DETECTIONS] = ( + surviving_times[source_index] + ) + result_detections.bboxes_metadata = result_bboxes_metadata + return {OUTPUT_KEY: result_detections} diff --git a/inference/core/workflows/core_steps/analytics/time_in_zone/v2_tensor.py b/inference/core/workflows/core_steps/analytics/time_in_zone/v2_tensor.py new file mode 100644 index 0000000000..7578a324e5 --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/time_in_zone/v2_tensor.py @@ -0,0 +1,309 @@ +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field +from typing_extensions import Literal, Type + +from inference.core.workflows.core_steps.analytics._zone_geometry import ( + LeanPolygonZone, + empty_detections_like, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + TIME_IN_ZONE_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, + WorkflowImageSelector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "timed_detections" +SHORT_DESCRIPTION = "Track object time in zone." +LONG_DESCRIPTION = """ +Calculate and track the time spent by tracked objects within a defined polygon zone, measure duration of object presence in specific areas, filter detections based on zone membership, reset time tracking when objects leave zones, and enable zone-based analytics, dwell time analysis, and presence monitoring workflows. + +## How This Block Works + +This block measures how long each tracked object has been inside a defined polygon zone by tracking entry and exit times for each unique track ID. The block: + +1. Receives tracked detection predictions with track IDs, an image with embedded video metadata, and a polygon zone definition +2. Extracts video metadata from the image: + - Accesses video_metadata from the WorkflowImageData object + - Extracts fps, frame_number, frame_timestamp, video_identifier, and video source information + - Uses video_identifier to maintain separate tracking state for different videos +3. Validates that detections have track IDs (tracker_id must be present): + - Requires detections to come from a tracking block (e.g., Byte Tracker) + - Each object must have a unique tracker_id that persists across frames + - Raises an error if tracker_id is missing +4. Initializes or retrieves a polygon zone for the video: + - Creates a PolygonZone object from zone coordinates for each unique video + - Validates zone coordinates (must be a list of at least 3 points, each with 2 coordinates) + - Stores zone configuration per video using video_identifier + - Configures triggering anchor point (e.g., CENTER, BOTTOM_CENTER) for zone detection +5. Initializes or retrieves time tracking state for the video: + - Maintains a dictionary tracking when each track_id entered the zone + - Stores entry timestamps per video using video_identifier + - Maintains separate tracking state for each video +6. Calculates current timestamp for time measurement: + - For video files: Calculates timestamp as frame_number / fps + - For streamed video: Uses frame_timestamp from metadata + - Provides accurate time measurement for duration calculation +7. Checks which detections are in the zone: + - Uses polygon zone trigger to test if each detection's anchor point is inside the zone + - The triggering_anchor determines which point on the bounding box is checked (CENTER, BOTTOM_CENTER, etc.) + - Returns boolean for each detection indicating zone membership +8. Updates time tracking for each tracked object: + - **For objects entering the zone**: Records entry timestamp if not already tracked + - **For objects in the zone**: Calculates time spent as current_timestamp - entry_timestamp + - **For objects leaving the zone**: + - If reset_out_of_zone_detections is True: Removes entry timestamp (resets to 0) + - If reset_out_of_zone_detections is False: Keeps entry timestamp (continues tracking) +9. Handles out-of-zone detections: + - **If remove_out_of_zone_detections is True**: Filters out detections outside the zone from output + - **If remove_out_of_zone_detections is False**: Includes out-of-zone detections with time = 0 +10. Adds time_in_zone information to each detection: + - Attaches time_in_zone value (in seconds) to each detection as metadata + - Objects in zone: Time represents duration spent in zone + - Objects outside zone: Time is 0 (if not reset) or undefined (if removed) +11. Returns detections with time_in_zone information: + - Outputs tracked detections enhanced with time_in_zone metadata + - Filtered or unfiltered based on remove_out_of_zone_detections setting + - Maintains all original detection properties plus time tracking information + +The block maintains persistent tracking state across frames, allowing accurate cumulative time measurement for objects that remain in the zone over multiple frames. Time is measured from when an object first enters the zone (based on its track_id) until the current frame, providing real-time duration tracking. The zone is defined as a polygon with multiple points, allowing flexible area definitions. The triggering anchor determines which part of the bounding box is used for zone detection, enabling different zone entry/exit behaviors based on object position. + +## Common Use Cases + +- **Dwell Time Analysis**: Measure how long objects remain in specific areas for behavior analysis (e.g., measure customer dwell time in store sections, track time spent in parking spaces, analyze time in waiting areas), enabling dwell time analytics workflows +- **Zone-Based Monitoring**: Monitor object presence in defined areas for security and safety (e.g., detect loitering in restricted areas, monitor time in danger zones, track presence in secure zones), enabling zone monitoring workflows +- **Retail Analytics**: Track customer time in different store sections for retail insights (e.g., measure time in product aisles, analyze shopping patterns, track department engagement), enabling retail analytics workflows +- **Occupancy Management**: Measure time objects spend in spaces for space utilization (e.g., track vehicle parking duration, measure table occupancy time, analyze space usage patterns), enabling occupancy management workflows +- **Safety Compliance**: Monitor time violations in restricted or time-limited zones (e.g., detect extended stays in hazardous areas, monitor time limit violations, track safety compliance), enabling safety monitoring workflows +- **Traffic Analysis**: Measure time vehicles spend in traffic zones or intersections (e.g., track time at intersections, measure queue waiting time, analyze traffic flow patterns), enabling traffic analytics workflows + +## Connecting to Other Blocks + +This block receives an image with embedded video metadata, tracked detections, and zone coordinates, and produces timed_detections with time_in_zone metadata: + +- **After Byte Tracker blocks** to measure time for tracked objects (e.g., track time in zones for tracked objects, measure dwell time with consistent IDs, analyze tracked object presence), enabling tracking-to-time workflows +- **After zone definition blocks** to apply time tracking to defined areas (e.g., measure time in polygon zones, track duration in custom zones, analyze zone-based presence), enabling zone-to-time workflows +- **Before logic blocks** like Continue If to make decisions based on time in zone (e.g., continue if time exceeds threshold, filter based on dwell time, trigger actions on time violations), enabling time-based decision workflows +- **Before analysis blocks** to analyze time-based metrics (e.g., analyze dwell time patterns, process time-in-zone data, work with duration metrics), enabling time analysis workflows +- **Before notification blocks** to alert on time violations or thresholds (e.g., alert on extended stays, notify on time limit violations, trigger time-based alerts), enabling time-based notification workflows +- **Before data storage blocks** to record time metrics (e.g., store dwell time data, log time-in-zone metrics, record duration measurements), enabling time metrics logging workflows + +## Version Differences + +**Enhanced from v1:** + +- **Simplified Input**: Uses `image` input that contains embedded video metadata instead of requiring a separate `metadata` field, simplifying workflow connections and reducing input complexity +- **Improved Integration**: Better integration with image-based workflows since video metadata is accessed directly from the image object rather than requiring separate metadata input +- **Streamlined Workflow**: Reduces the number of inputs needed, making it easier to connect in workflows where image and metadata come from the same source + +## Requirements + +This block requires tracked detections with tracker_id information (detections must come from a tracking block like Byte Tracker). The zone must be defined as a list of at least 3 points, where each point is a list or tuple of exactly 2 coordinates (x, y). The image's video_metadata should include frame rate (fps) for video files or frame timestamps for streamed video to calculate accurate time measurements. The block maintains persistent tracking state across frames for each video, so it should be used in video workflows where frames are processed sequentially. For accurate time measurement, detections should be provided consistently across frames with valid track IDs. +""" + + +class TimeInZoneManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Time in Zone", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "video", + "icon": "far fa-timer", + "blockPriority": 1, + }, + } + ) + type: Literal["roboflow_core/time_in_zone@v2"] + image: Union[WorkflowImageSelector] = Field( + title="Image", + description="Input image for the current video frame containing embedded video metadata (fps, frame_number, frame_timestamp, video_identifier, video source) required for time calculation and state management. The block extracts video_metadata from the WorkflowImageData object. The fps and frame_number are used for video files to calculate timestamps (timestamp = frame_number / fps). For streamed video, frame_timestamp is used directly. The video_identifier is used to maintain separate tracking state and zone configurations for different videos. Used for zone visualization and reference. The image dimensions are used to validate zone coordinates. This version simplifies input by embedding metadata in the image object rather than requiring a separate metadata field.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Tracked detection predictions (object detection or instance segmentation) with tracker_id information. Detections must come from a tracking block (e.g., Byte Tracker) that has assigned unique tracker_id values that persist across frames. Each detection must have a tracker_id to enable time tracking. The block calculates time_in_zone for each tracked object based on when its track_id first entered the zone. The output will include the same detections enhanced with time_in_zone metadata (duration in seconds). If remove_out_of_zone_detections is True, only detections inside the zone are included in the output.", + examples=["$steps.object_detection_model.predictions"], + ) + zone: Union[list, Selector(kind=[LIST_OF_VALUES_KIND]), Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + description="Polygon zone coordinates defining the area for time measurement. Must be a list of at least 3 points, where each point is a list or tuple of exactly 2 coordinates [x, y] or (x, y). Coordinates should be in pixel space matching the image dimensions. Example: [(100, 100), (100, 200), (300, 200), (300, 100)] for a quadrilateral zone. The zone defines the polygon area where time tracking occurs. Objects are considered 'in zone' when their triggering_anchor point is inside this polygon. Zone coordinates are validated and a PolygonZone object is created for each video.", + examples=[[(100, 100), (100, 200), (300, 200), (300, 100)], "$inputs.zones"], + ) + triggering_anchor: Union[str, Selector(kind=[STRING_KIND]), Literal[tuple(sv.Position.list())]] = Field( # type: ignore + description="Point on the detection bounding box that must be inside the zone to consider the object 'in zone'. Options include: 'CENTER' (default, center of bounding box), 'BOTTOM_CENTER' (bottom center point), 'TOP_CENTER' (top center point), 'CENTER_LEFT' (center left point), 'CENTER_RIGHT' (center right point), and other Position enum values. The triggering anchor determines which part of the object's bounding box is checked against the zone polygon. Use CENTER for standard zone detection, BOTTOM_CENTER for ground-level zones (e.g., tracking feet/vehicle base), or other anchors based on detection needs. Default is 'CENTER'.", + default="CENTER", + examples=["CENTER"], + ) + remove_out_of_zone_detections: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + description="If True (default), detections found outside the zone are filtered out and not included in the output. Only detections inside the zone are returned. If False, all detections are included in the output, with time_in_zone = 0 for objects outside the zone. Use True to focus analysis only on objects in the zone, or False to maintain all detections with zone status. Default is True for cleaner output focused on zone activity.", + default=True, + examples=[True, False], + ) + reset_out_of_zone_detections: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + description="If True (default), when a tracked object leaves the zone, its time tracking is reset (entry timestamp is cleared). When the object re-enters the zone, time tracking starts from 0 again. If False, time tracking continues even after leaving the zone, and re-entry maintains cumulative time. Use True to measure current continuous time in zone (resets on exit), or False to measure cumulative time across multiple entries. Default is True for measuring continuous presence duration.", + default=True, + examples=[True, False], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class TimeInZoneBlockV2(WorkflowBlock): + def __init__(self): + self._batch_of_tracked_ids_in_zone: Dict[str, Dict[Union[int, str], float]] = {} + self._batch_of_polygon_zones: Dict[str, LeanPolygonZone] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return TimeInZoneManifest + + def run( + self, + image: WorkflowImageData, + detections: Union[Detections, InstanceDetections], + zone: List[Tuple[int, int]], + triggering_anchor: str, + remove_out_of_zone_detections: bool, + reset_out_of_zone_detections: bool, + ) -> BlockResult: + n = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + tracker_ids = [ + box_metadata.get("tracker_id") for box_metadata in bboxes_metadata + ] + if n > 0 and any(tracker_id is None for tracker_id in tracker_ids): + raise ValueError( + f"tracker_id not initialized, {self.__class__.__name__} requires detections to be tracked" + ) + metadata = image.video_metadata + if metadata.video_identifier not in self._batch_of_polygon_zones: + if not isinstance(zone, list) or len(zone) < 3: + raise ValueError( + f"{self.__class__.__name__} requires zone to be a list containing more than 2 points" + ) + if any( + (not isinstance(e, list) and not isinstance(e, tuple)) or len(e) != 2 + for e in zone + ): + raise ValueError( + f"{self.__class__.__name__} requires each point of zone to be a list containing exactly 2 coordinates" + ) + if any( + not isinstance(e[0], (int, float)) or not isinstance(e[1], (int, float)) + for e in zone + ): + raise ValueError( + f"{self.__class__.__name__} requires each coordinate of zone to be a number" + ) + self._batch_of_polygon_zones[metadata.video_identifier] = LeanPolygonZone( + polygon=np.array(zone), + triggering_anchors=(sv.Position(triggering_anchor),), + ) + polygon_zone = self._batch_of_polygon_zones[metadata.video_identifier] + tracked_ids_in_zone = self._batch_of_tracked_ids_in_zone.setdefault( + metadata.video_identifier, {} + ) + if metadata.comes_from_video_file and metadata.fps != 0: + ts_end = metadata.frame_number / metadata.fps + else: + ts_end = metadata.frame_timestamp.timestamp() + + # Transfer box geometry once and evaluate zone membership with the lean + # vectorized host implementation. + xyxy_host = detections.xyxy.detach().to("cpu").numpy().astype(float) + is_in_zone_mask = polygon_zone.trigger(xyxy_host) + surviving_mask = np.zeros(n, dtype=bool) + surviving_times: Dict[int, float] = {} + for i, is_in_zone, tracker_id in zip( + range(n), + is_in_zone_mask, + tracker_ids, + ): + if ( + not is_in_zone + and tracker_id in tracked_ids_in_zone + and reset_out_of_zone_detections + ): + del tracked_ids_in_zone[tracker_id] + if not is_in_zone and remove_out_of_zone_detections: + continue + + time_in_zone = 0.0 + if is_in_zone: + ts_start = tracked_ids_in_zone.setdefault(tracker_id, ts_end) + time_in_zone = ts_end - ts_start + elif tracker_id in tracked_ids_in_zone: + del tracked_ids_in_zone[tracker_id] + surviving_mask[i] = True + surviving_times[i] = time_in_zone + result_detections = ( + empty_detections_like(detections) + if n and not surviving_mask.any() + else take_prediction_by_mask(detections, surviving_mask) + ) + result_bboxes_metadata = result_detections.bboxes_metadata or [ + {} for _ in range(int(result_detections.xyxy.shape[0])) + ] + for position, source_index in enumerate(np.nonzero(surviving_mask)[0].tolist()): + result_bboxes_metadata[position][TIME_IN_ZONE_KEY_IN_SV_DETECTIONS] = ( + surviving_times[source_index] + ) + result_detections.bboxes_metadata = result_bboxes_metadata + return {OUTPUT_KEY: result_detections} diff --git a/inference/core/workflows/core_steps/analytics/time_in_zone/v3_tensor.py b/inference/core/workflows/core_steps/analytics/time_in_zone/v3_tensor.py new file mode 100644 index 0000000000..909c314b4b --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/time_in_zone/v3_tensor.py @@ -0,0 +1,412 @@ +import itertools +from collections import OrderedDict +from typing import Dict, List, Optional, Tuple, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field +from typing_extensions import Literal, Type + +from inference.core.workflows.core_steps.analytics._zone_geometry import ( + LeanPolygonZone, + empty_detections_like, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + TIME_IN_ZONE_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, + WorkflowImageSelector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +PolygonAsNestedList = List[List[int]] +PolygonAsArray = np.ndarray +PolygonAsListOfArrays = List[np.ndarray] +PolygonAsListOfTuples = List[Tuple[int, int]] +Polygon = Union[ + PolygonAsNestedList, PolygonAsArray, PolygonAsListOfArrays, PolygonAsListOfTuples +] + + +OUTPUT_KEY: str = "timed_detections" +SHORT_DESCRIPTION = "Track object time in zone." +LONG_DESCRIPTION = """ +Calculate and track the time spent by tracked objects within one or more defined polygon zones, measure duration of object presence in specific areas (supporting multiple zones where objects are considered 'in zone' if present in any zone), filter detections based on zone membership, reset time tracking when objects leave zones, and enable zone-based analytics, dwell time analysis, and presence monitoring workflows. + +## How This Block Works + +This block measures how long each tracked object has been inside one or more defined polygon zones by tracking entry and exit times for each unique track ID. The block supports multiple zones, treating objects as 'in zone' if they are present in any of the defined zones. The block: + +1. Receives tracked detection predictions with track IDs, an image with embedded video metadata, and polygon zone definition(s) (single zone or list of zones) +2. Extracts video metadata from the image: + - Accesses video_metadata from the WorkflowImageData object + - Extracts fps, frame_number, frame_timestamp, video_identifier, and video source information + - Uses video_identifier to maintain separate tracking state for different videos +3. Validates that detections have track IDs (tracker_id must be present): + - Requires detections to come from a tracking block (e.g., Byte Tracker) + - Each object must have a unique tracker_id that persists across frames + - Raises an error if tracker_id is missing +4. Normalizes zone input to a list of polygons: + - Accepts a single polygon zone or a list of polygon zones + - Automatically wraps single polygons in a list for consistent processing + - Validates nesting depth and coordinate format for all zones + - Enables flexible zone input formats (single zone or multiple zones) +5. Initializes or retrieves polygon zones for the video: + - Creates a list of PolygonZone objects from zone coordinates for each unique zone combination + - Validates zone coordinates (each zone must be a list of at least 3 points, each with 2 coordinates) + - Stores zone configurations in an OrderedDict with zone cache (max 100 zone combinations) + - Uses zone key combining video_identifier and zone coordinates for cache lookup + - Implements FIFO eviction when cache exceeds 100 zone combinations + - Configures triggering anchor point (e.g., CENTER, BOTTOM_CENTER) for zone detection +6. Initializes or retrieves time tracking state for the video: + - Maintains a dictionary tracking when each track_id entered any zone + - Stores entry timestamps per video using video_identifier + - Maintains separate tracking state for each video +7. Calculates current timestamp for time measurement: + - For video files: Calculates timestamp as frame_number / fps + - For streamed video: Uses frame_timestamp from metadata + - Provides accurate time measurement for duration calculation +8. Checks which detections are in any zone: + - Tests each detection against all polygon zones using polygon zone triggers + - Creates a matrix of zone membership (zones x detections) + - Uses logical OR operation: objects are considered 'in zone' if they're in ANY of the zones + - The triggering_anchor determines which point on the bounding box is checked (CENTER, BOTTOM_CENTER, etc.) + - Returns boolean for each detection indicating zone membership in any zone +9. Updates time tracking for each tracked object: + - **For objects entering any zone**: Records entry timestamp if not already tracked + - **For objects in any zone**: Calculates time spent as current_timestamp - entry_timestamp + - **For objects leaving all zones**: + - If reset_out_of_zone_detections is True: Removes entry timestamp (resets to 0) + - If reset_out_of_zone_detections is False: Keeps entry timestamp (continues tracking) +10. Handles out-of-zone detections: + - **If remove_out_of_zone_detections is True**: Filters out detections outside all zones from output + - **If remove_out_of_zone_detections is False**: Includes out-of-zone detections with time = 0 +11. Adds time_in_zone information to each detection: + - Attaches time_in_zone value (in seconds) to each detection as metadata + - Objects in any zone: Time represents duration spent in any zone + - Objects outside all zones: Time is 0 (if not reset) or undefined (if removed) +12. Returns detections with time_in_zone information: + - Outputs tracked detections enhanced with time_in_zone metadata + - Filtered or unfiltered based on remove_out_of_zone_detections setting + - Maintains all original detection properties plus time tracking information + +The block maintains persistent tracking state across frames, allowing accurate cumulative time measurement for objects that remain in any zone over multiple frames. Time is measured from when an object first enters any zone (based on its track_id) until the current frame, providing real-time duration tracking. When multiple zones are provided, objects are considered 'in zone' if their anchor point is inside any of the zones, allowing tracking across multiple areas as a single combined zone. The zone cache efficiently manages multiple zone configurations per video using FIFO eviction to limit memory usage. The triggering anchor determines which part of the bounding box is used for zone detection, enabling different zone entry/exit behaviors based on object position. + +## Common Use Cases + +- **Multi-Zone Dwell Time Analysis**: Measure how long objects remain in any of multiple areas for behavior analysis (e.g., measure customer time in any store section, track time spent in multiple parking areas, analyze time in overlapping zones), enabling multi-zone dwell time analytics workflows +- **Zone-Based Monitoring**: Monitor object presence across multiple defined areas for security and safety (e.g., detect loitering in any restricted area, monitor time in multiple danger zones, track presence across secure zones), enabling multi-zone monitoring workflows +- **Retail Analytics**: Track customer time across multiple store sections for retail insights (e.g., measure time in any product aisle, analyze shopping patterns across departments, track engagement in multiple zones), enabling multi-zone retail analytics workflows +- **Occupancy Management**: Measure time objects spend in any of multiple spaces for space utilization (e.g., track vehicle parking duration in multiple lots, measure table occupancy across zones, analyze space usage in multiple areas), enabling multi-zone occupancy management workflows +- **Safety Compliance**: Monitor time violations across multiple restricted or time-limited zones (e.g., detect extended stays in any hazardous area, monitor time limit violations across zones, track safety compliance in multiple areas), enabling multi-zone safety monitoring workflows +- **Traffic Analysis**: Measure time vehicles spend in any of multiple traffic zones or intersections (e.g., track time at multiple intersections, measure queue waiting time across zones, analyze traffic flow in multiple areas), enabling multi-zone traffic analytics workflows + +## Connecting to Other Blocks + +This block receives an image with embedded video metadata, tracked detections, and zone coordinates (single or multiple zones), and produces timed_detections with time_in_zone metadata: + +- **After Byte Tracker blocks** to measure time for tracked objects across multiple zones (e.g., track time in multiple zones for tracked objects, measure dwell time with consistent IDs across areas, analyze tracked object presence in multiple zones), enabling tracking-to-time workflows +- **After zone definition blocks** to apply time tracking to multiple defined areas (e.g., measure time across multiple polygon zones, track duration in custom multi-zone configurations, analyze zone-based presence across areas), enabling zone-to-time workflows +- **Before logic blocks** like Continue If to make decisions based on time in any zone (e.g., continue if time exceeds threshold in any zone, filter based on dwell time across zones, trigger actions on time violations in multiple areas), enabling time-based decision workflows +- **Before analysis blocks** to analyze time-based metrics across multiple zones (e.g., analyze dwell time patterns across zones, process time-in-zone data for multiple areas, work with duration metrics across zones), enabling time analysis workflows +- **Before notification blocks** to alert on time violations or thresholds in any zone (e.g., alert on extended stays in any zone, notify on time limit violations across areas, trigger time-based alerts for multiple zones), enabling time-based notification workflows +- **Before data storage blocks** to record time metrics across multiple zones (e.g., store dwell time data for multiple areas, log time-in-zone metrics across zones, record duration measurements for multiple zones), enabling time metrics logging workflows + +## Version Differences + +**Enhanced from v2:** + +- **Multiple Zone Support**: Supports tracking time across multiple polygon zones simultaneously, where objects are considered 'in zone' if they're present in any of the defined zones, enabling multi-zone time tracking and analysis +- **Flexible Zone Input**: Accepts either a single polygon zone or a list of polygon zones, automatically normalizing the input to handle both formats seamlessly +- **Zone Cache Management**: Implements a zone cache with FIFO eviction (max 100 zone combinations) to efficiently manage multiple zone configurations per video while limiting memory usage +- **Combined Zone Logic**: Uses logical OR operation across all zones, allowing tracking across multiple areas as a unified zone system for comprehensive presence monitoring +- **Enhanced Zone Key System**: Uses combined zone keys (video_identifier + zone coordinates) for cache lookup, enabling efficient storage and retrieval of zone configurations + +## Requirements + +This block requires tracked detections with tracker_id information (detections must come from a tracking block like Byte Tracker). The zone can be a single polygon or a list of polygons, where each polygon must be defined as a list of at least 3 points, with each point being a list or tuple of exactly 2 coordinates (x, y). The image's video_metadata should include frame rate (fps) for video files or frame timestamps for streamed video to calculate accurate time measurements. The block maintains persistent tracking state across frames for each video, so it should be used in video workflows where frames are processed sequentially. For accurate time measurement, detections should be provided consistently across frames with valid track IDs. When multiple zones are provided, objects are considered 'in zone' if they're present in any of the zones. +""" +ZONE_CACHE_SIZE = 100 + + +class TimeInZoneManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Time in Zone", + "version": "v3", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "video", + "icon": "far fa-timer", + "blockPriority": 1, + }, + } + ) + type: Literal["roboflow_core/time_in_zone@v3"] + image: Union[WorkflowImageSelector] = Field( + title="Image", + description="Input image for the current video frame containing embedded video metadata (fps, frame_number, frame_timestamp, video_identifier, video source) required for time calculation and state management. The block extracts video_metadata from the WorkflowImageData object. The fps and frame_number are used for video files to calculate timestamps (timestamp = frame_number / fps). For streamed video, frame_timestamp is used directly. The video_identifier is used to maintain separate tracking state and zone configurations for different videos. Used for zone visualization and reference. The image dimensions are used to validate zone coordinates. This version supports multiple zones per video with efficient zone cache management.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Tracked detection predictions (object detection or instance segmentation) with tracker_id information. Detections must come from a tracking block (e.g., Byte Tracker) that has assigned unique tracker_id values that persist across frames. Each detection must have a tracker_id to enable time tracking. The block calculates time_in_zone for each tracked object based on when its track_id first entered any of the zones. Objects are considered 'in zone' if their anchor point is inside any of the provided zones. The output will include the same detections enhanced with time_in_zone metadata (duration in seconds). If remove_out_of_zone_detections is True, only detections inside any zone are included in the output.", + examples=["$steps.object_detection_model.predictions"], + ) + zone: Union[list, Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + description="Polygon zone coordinates defining one or more areas for time measurement. Can be a single polygon zone or a list of polygon zones. Each zone must be a list of at least 3 points, where each point is a list or tuple of exactly 2 coordinates [x, y] or (x, y). Coordinates should be in pixel space matching the image dimensions. Example for single zone: [(100, 100), (100, 200), (300, 200), (300, 100)]. Example for multiple zones: [[(100, 100), (100, 200), (300, 200), (300, 100)], [(400, 400), (400, 500), (600, 500), (600, 400)]]. Objects are considered 'in zone' if their triggering_anchor point is inside ANY of the provided zones. Zone coordinates are validated and PolygonZone objects are created for each zone. Zone configurations are cached (max 100 combinations) with FIFO eviction.", + examples=[[(100, 100), (100, 200), (300, 200), (300, 100)], "$inputs.zones"], + ) + triggering_anchor: Union[ + str, Selector(kind=[STRING_KIND]), Literal[tuple(sv.Position.list())] + ] = Field( # type: ignore + description="Point on the detection bounding box that must be inside the zone to consider the object 'in zone'. Options include: 'CENTER' (default, center of bounding box), 'BOTTOM_CENTER' (bottom center point), 'TOP_CENTER' (top center point), 'CENTER_LEFT' (center left point), 'CENTER_RIGHT' (center right point), and other Position enum values. The triggering anchor determines which part of the object's bounding box is checked against the zone polygon(s). When multiple zones are provided, the object is considered 'in zone' if its anchor point is inside ANY of the zones. Use CENTER for standard zone detection, BOTTOM_CENTER for ground-level zones (e.g., tracking feet/vehicle base), or other anchors based on detection needs. Default is 'CENTER'.", + default="CENTER", + examples=["CENTER"], + ) + remove_out_of_zone_detections: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + description="If True (default), detections found outside all zones are filtered out and not included in the output. Only detections inside at least one zone are returned. If False, all detections are included in the output, with time_in_zone = 0 for objects outside all zones. Use True to focus analysis only on objects in any zone, or False to maintain all detections with zone status. When multiple zones are provided, objects are considered 'in zone' if present in any zone. Default is True for cleaner output focused on zone activity.", + default=True, + examples=[True, False], + ) + reset_out_of_zone_detections: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + description="If True (default), when a tracked object leaves all zones, its time tracking is reset (entry timestamp is cleared). When the object re-enters any zone, time tracking starts from 0 again. If False, time tracking continues even after leaving all zones, and re-entry maintains cumulative time. Use True to measure current continuous time in any zone (resets on exit from all zones), or False to measure cumulative time across multiple entries. When multiple zones are provided, time is reset only when the object leaves all zones. Default is True for measuring continuous presence duration.", + default=True, + examples=[True, False], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class TimeInZoneBlockV3(WorkflowBlock): + def __init__(self): + self._batch_of_tracked_ids_in_zone: Dict[str, Dict[Union[int, str], float]] = {} + self._batch_of_polygon_zones: OrderedDict[str, List[LeanPolygonZone]] = ( + OrderedDict() + ) + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return TimeInZoneManifest + + def run( + self, + image: WorkflowImageData, + detections: Union[Detections, InstanceDetections], + zone: List[List[Tuple[int, int]]], + triggering_anchor: str, + remove_out_of_zone_detections: bool, + reset_out_of_zone_detections: bool, + ) -> BlockResult: + n = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + tracker_ids = [ + box_metadata.get("tracker_id") for box_metadata in bboxes_metadata + ] + if n > 0 and any(tracker_id is None for tracker_id in tracker_ids): + raise ValueError( + f"tracker_id not initialized, {self.__class__.__name__} requires detections to be tracked" + ) + metadata = image.video_metadata + zones = ensure_zone_is_list_of_polygons(zone) + zone_key = f"{metadata.video_identifier}_{str(zones)}" + if zone_key not in self._batch_of_polygon_zones: + if len(zones) > 0 and (not isinstance(zones[0], list) or len(zones[0]) < 3): + raise ValueError( + f"{self.__class__.__name__} requires zone to be a list containing more than 2 points" + ) + if any( + (not isinstance(e, list) and not isinstance(e, tuple)) or len(e) != 2 + for e in itertools.chain.from_iterable(zones) + ): + raise ValueError( + f"{self.__class__.__name__} requires each point of zone to be a list containing exactly 2 coordinates" + ) + if any( + not isinstance(e[0], (int, float)) or not isinstance(e[1], (int, float)) + for e in itertools.chain.from_iterable(zones) + ): + raise ValueError( + f"{self.__class__.__name__} requires each coordinate of zone to be a number" + ) + self._batch_of_polygon_zones[zone_key] = [ + LeanPolygonZone( + polygon=np.array(zone), + triggering_anchors=(sv.Position(triggering_anchor),), + ) + for zone in zones + ] + # keeps the cache size at ZONE_CACHE_SIZE + if len(self._batch_of_polygon_zones) > ZONE_CACHE_SIZE: + self._batch_of_polygon_zones.popitem(last=False) + polygon_zones = self._batch_of_polygon_zones[zone_key] + tracked_ids_in_zone = self._batch_of_tracked_ids_in_zone.setdefault( + metadata.video_identifier, {} + ) + if metadata.comes_from_video_file and metadata.fps != 0: + ts_end = metadata.frame_number / metadata.fps + else: + ts_end = metadata.frame_timestamp.timestamp() + + # Transfer box geometry once and reuse it across every cached zone. + xyxy_host = detections.xyxy.detach().to("cpu").numpy().astype(float) + + # get trigger for all zones. It is a matrix of shape (len(zones), len(detections)) + polygon_triggers = [ + polygon_zone.trigger(xyxy_host) for polygon_zone in polygon_zones + ] + is_in_any_zone = ( + np.any(polygon_triggers, axis=0) + if len(polygon_triggers) > 0 + else np.array([False] * n) + ) + + surviving_mask = np.zeros(n, dtype=bool) + surviving_times: Dict[int, float] = {} + for i, is_in_zone, tracker_id in zip( + range(n), + is_in_any_zone, + tracker_ids, + ): + if ( + not is_in_zone + and tracker_id in tracked_ids_in_zone + and reset_out_of_zone_detections + ): + del tracked_ids_in_zone[tracker_id] + if not is_in_zone and remove_out_of_zone_detections: + continue + + time_in_zone = 0.0 + if is_in_zone: + ts_start = tracked_ids_in_zone.setdefault(tracker_id, ts_end) + time_in_zone = ts_end - ts_start + elif tracker_id in tracked_ids_in_zone: + del tracked_ids_in_zone[tracker_id] + surviving_mask[i] = True + surviving_times[i] = time_in_zone + result_detections = ( + empty_detections_like(detections) + if n and not surviving_mask.any() + else take_prediction_by_mask(detections, surviving_mask) + ) + result_bboxes_metadata = result_detections.bboxes_metadata or [ + {} for _ in range(int(result_detections.xyxy.shape[0])) + ] + for position, source_index in enumerate(np.nonzero(surviving_mask)[0].tolist()): + result_bboxes_metadata[position][TIME_IN_ZONE_KEY_IN_SV_DETECTIONS] = ( + surviving_times[source_index] + ) + result_detections.bboxes_metadata = result_bboxes_metadata + return {OUTPUT_KEY: result_detections} + + +def ensure_zone_is_list_of_polygons( + zone: Union[Polygon, List[Polygon]], +) -> List[Polygon]: + nesting_depth = calculate_nesting_depth(zone=zone, max_depth=3) + if nesting_depth > 3: + raise ValueError( + "roboflow_core/time_in_zone@v3 block requires `zone` input to be list of points, but " + "input with excessive nesting depth found. If you created the `zone` input manually, verify it's " + "correctness. If the input is constructed by another Workflow block - raise an issue: " + "https://github.com/roboflow/inference/issues" + ) + if nesting_depth == 2: + return [zone] + return zone + + +def calculate_nesting_depth( + zone: Union[Polygon, List[Polygon]], max_depth: int, current_depth: int = 0 +) -> int: + remaining_depth = max_depth - current_depth + if isinstance(zone, np.ndarray): + array_depth = len(zone.shape) + if array_depth > remaining_depth: + raise ValueError( + "While processing polygon zone detected an instance of the zone which is invalid, as " + "the input is nested beyond limits - the block supports single and multiple " + "lists of zone points. If you created the `zone` input manually, verify it's correctness. If " + "the input is constructed by another Workflow block - raise an issue: " + "https://github.com/roboflow/inference/issues" + ) + return current_depth + array_depth + if isinstance(zone, (list, tuple)): + if remaining_depth < 1: + raise ValueError( + "While processing polygon zone detected an instance of the zone which is invalid, as " + "the input is nested beyond limits - the block supports single and multiple " + "lists of zone points. If you created the `zone` input manually, verify it's correctness. If " + "the input is constructed by another Workflow block - raise an issue: " + "https://github.com/roboflow/inference/issues" + ) + depths = { + calculate_nesting_depth( + zone=e, max_depth=max_depth, current_depth=current_depth + 1 + ) + for e in zone + } + if not depths: + return current_depth + 1 + if len(depths) != 1: + raise ValueError( + "While processing polygon zone detected an instance of the zone which is invalid, as " + "the input is nested in irregular way. If you created the `zone` input manually, verify it's correctness. " + "If the input is constructed by another Workflow block - raise an issue: " + "https://github.com/roboflow/inference/issues" + ) + return min(depths) + return current_depth diff --git a/inference/core/workflows/core_steps/analytics/velocity/v1_tensor.py b/inference/core/workflows/core_steps/analytics/velocity/v1_tensor.py new file mode 100644 index 0000000000..7b0514044c --- /dev/null +++ b/inference/core/workflows/core_steps/analytics/velocity/v1_tensor.py @@ -0,0 +1,392 @@ +from typing import Dict, List, Optional, Tuple, Union + +import torch +from pydantic import ConfigDict, Field +from typing_extensions import Literal, Type + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + split_key_point_prediction, +) +from inference.core.workflows.execution_engine.constants import ( + SMOOTHED_SPEED_KEY_IN_SV_DETECTIONS, + SMOOTHED_VELOCITY_KEY_IN_SV_DETECTIONS, + SPEED_KEY_IN_SV_DETECTIONS, + VELOCITY_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + Selector, + StepOutputSelector, + WorkflowImageSelector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "velocity_detections" +SHORT_DESCRIPTION = "Calculate the velocity and speed of tracked objects with smoothing and unit conversion." +LONG_DESCRIPTION = """ +Calculate the velocity and speed of tracked objects across video frames by measuring displacement of object centers over time, applying exponential moving average smoothing to reduce noise, and converting measurements from pixels per second to meters per second for traffic speed monitoring, movement analysis, behavior tracking, and performance measurement workflows. + +## How This Block Works + +This block measures how fast objects are moving by tracking their positions across video frames. The block: + +1. Receives tracked detection predictions with unique tracker IDs and an image with embedded video metadata +2. Extracts video metadata from the image: + - Accesses video_metadata to get frame timestamps or frame numbers and frame rate (fps) + - Extracts video_identifier to maintain separate tracking state for different videos + - Determines current timestamp using frame_number/fps for video files or frame_timestamp for streams +3. Validates that detections have tracker IDs (required for tracking object movement across frames) +4. Calculates object center positions: + - Computes the center point (x, y) of each bounding box in the current frame + - Uses bounding box coordinates to find geometric centers +5. Retrieves or initializes tracking state: + - Maintains previous positions and timestamps for each tracker_id per video + - Stores smoothed velocity history for each tracker_id per video + - Creates new tracking entries for objects appearing for the first time +6. Calculates velocity and speed for each tracked object: + - **For objects with previous positions**: Computes displacement (change in position) and time delta (change in time) between current and previous frames + - **Velocity**: Calculates velocity vector as displacement divided by time delta (pixels per second) + - **Speed**: Computes speed as the magnitude (length) of the velocity vector (total pixels per second regardless of direction) + - **For new objects**: Sets velocity and speed to zero (no movement data available yet) +7. Applies exponential moving average smoothing: + - Smooths velocity measurements using exponential moving average with configurable smoothing factor (alpha) + - Reduces noise and jitter in velocity calculations from detection variations + - Lower alpha values provide more smoothing (slower response to changes), higher alpha values provide less smoothing (faster response to changes) + - Calculates smoothed velocity and smoothed speed for each object +8. Converts units to meters per second: + - Divides pixel-based velocities and speeds by pixels_per_meter conversion factor + - Converts all measurements (velocity, speed, smoothed_velocity, smoothed_speed) to real-world units + - Enables comparison with real-world speed measurements (e.g., km/h, mph) +9. Stores velocity data in detection metadata: + - Adds four velocity metrics to each detection: velocity (m/s), speed (m/s), smoothed_velocity (m/s), smoothed_speed (m/s) + - Velocity is a 2D vector [vx, vy] representing direction and magnitude of movement + - Speed is a scalar value representing total speed regardless of direction + - All measurements are stored in detections.data for downstream use +10. Updates tracking state for next frame: + - Saves current positions and timestamps for all tracked objects + - Stores smoothed velocities for next frame's smoothing calculations +11. Returns detections enhanced with velocity information: + - Outputs the same detection objects with added velocity metadata + - Each detection now includes velocity and speed data in its metadata + +Velocity is calculated based on the displacement of object centers (bounding box centers) over time. The block maintains separate tracking state for each video, allowing velocity calculation across multiple video streams. Due to perspective distortion and camera positioning, calculated velocity may vary depending on where objects appear in the frame - objects closer to the camera or at different depths will have different pixel-per-second values for the same real-world speed. The smoothing helps reduce noise from detection inaccuracies and frame-to-frame variations. + +## Common Use Cases + +- **Traffic Speed Monitoring**: Measure vehicle speeds on roads and highways (e.g., monitor traffic speeds, detect speeding violations, analyze traffic flow rates), enabling traffic enforcement and analysis workflows +- **Sports Performance Analysis**: Track athlete movement and speed during sports activities (e.g., measure player speeds, analyze sprint performance, track movement patterns), enabling sports analytics workflows +- **Security and Surveillance**: Monitor movement speed of people or objects in security scenarios (e.g., detect running or suspicious rapid movement, monitor crowd flow speeds, track object movement rates), enabling security monitoring workflows +- **Retail Analytics**: Analyze customer movement patterns and walking speeds in retail spaces (e.g., measure customer flow rates, analyze shopping behavior patterns, track movement efficiency), enabling retail behavior analysis workflows +- **Wildlife Behavior Studies**: Track animal movement speeds and patterns in natural habitats (e.g., measure animal speeds, analyze migration patterns, study movement behavior), enabling wildlife research workflows +- **Industrial Monitoring**: Monitor speeds of vehicles, equipment, or products in industrial settings (e.g., track conveyor speeds, measure vehicle speeds in facilities, monitor production line movement rates), enabling industrial automation workflows + +## Connecting to Other Blocks + +This block receives tracked detections and an image with embedded video metadata, and produces detections enhanced with velocity metadata: + +- **After Byte Tracker blocks** to calculate velocity for tracked objects (e.g., measure speeds of tracked vehicles, analyze tracked person movement, monitor tracked object velocities), enabling tracking-to-velocity workflows +- **After object detection or instance segmentation blocks** with tracking enabled to measure movement speeds (e.g., calculate vehicle speeds, track person movement rates, monitor object velocities), enabling detection-to-velocity workflows +- **Before visualization blocks** to display velocity information (e.g., visualize speed overlays, display velocity vectors, show movement speed annotations), enabling velocity visualization workflows +- **Before logic blocks** like Continue If to make decisions based on speed thresholds (e.g., continue if speed exceeds limit, filter based on velocity ranges, trigger actions on speed violations), enabling speed-based decision workflows +- **Before notification blocks** to alert on speed violations or threshold events (e.g., alert on speeding violations, notify on rapid movement, trigger speed-based alerts), enabling velocity-based notification workflows +- **Before data storage blocks** to record velocity measurements (e.g., log speed data, store velocity statistics, record movement metrics), enabling velocity data logging workflows + +## Requirements + +This block requires tracked detections with tracker_id information (detections must come from a tracking block like Byte Tracker). The image's video_metadata should include frame rate (fps) for video files or frame timestamps for streamed video to calculate accurate time deltas. The block maintains persistent tracking state across frames for each video using video_identifier, so it should be used in video workflows where frames are processed sequentially. For accurate velocity measurement, detections should be provided consistently across frames with valid tracker IDs. The pixels_per_meter conversion factor should be calibrated based on camera setup and scene geometry for accurate real-world speed measurements. Note that velocity accuracy may vary due to perspective distortion depending on object position in the frame. +""" + + +class VelocityManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Velocity", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "analytics", + "ui_manifest": { + "section": "video", + "icon": "far fa-gauge", + "blockPriority": 2.5, + }, + } + ) + type: Literal["roboflow_core/velocity@v1"] + image: WorkflowImageSelector = Field( + description="Image with embedded video metadata. The video_metadata contains fps, frame_number, frame_timestamp, and video_identifier. Required for calculating time deltas and maintaining separate velocity tracking state for different videos.", + ) + detections: StepOutputSelector( + kind=[ + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Tracked object detection or instance segmentation predictions. Must include tracker_id information from a tracking block. Velocity is calculated based on displacement of bounding box centers over time. Output detections include velocity (m/s vector), speed (m/s scalar), smoothed_velocity (m/s vector), and smoothed_speed (m/s scalar) in detection metadata.", + examples=["$steps.object_detection_model.predictions"], + ) + smoothing_alpha: Union[float, Selector(kind=[FLOAT_KIND])] = Field( # type: ignore + default=0.5, + description="Smoothing factor (alpha) for exponential moving average, range 0 < alpha <= 1. Controls how much smoothing is applied to velocity measurements. Lower values (closer to 0) provide more smoothing - slower response to changes, less noise. Higher values (closer to 1) provide less smoothing - faster response to changes, more noise. Default 0.5 balances smoothness and responsiveness.", + examples=[0.5, 0.3, 0.7], + ) + pixels_per_meter: Union[float, Selector(kind=[FLOAT_KIND])] = Field( # type: ignore + default=1.0, + description="Conversion factor from pixels to meters for real-world speed calculation. Velocity measurements in pixels per second are divided by this value to convert to meters per second. Must be greater than 0. For accurate real-world speeds, calibrate based on camera height, angle, and scene geometry. Example: if 1 pixel = 0.01 meters (1cm), use 0.01. Default 1.0 means no conversion (results in pixels per second).", + examples=[0.01, 0.1, 1.0], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.0.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class _VideoVelocityState: + """Device-resident velocity state for one video. + + Positions and smoothed velocities stay torch tensors on the prediction's + device across frames. Row bookkeeping (tracker_id -> row) and absolute + timestamps stay on host: unix epochs do not fit float32 (~100 s resolution + at 1.7e9) and MPS has no float64 tensors. + """ + + def __init__(self) -> None: + self.rows_by_tracker: Dict[int, int] = {} + self.timestamps_by_tracker: Dict[int, float] = {} + self.positions: Optional[torch.Tensor] = None # (K, 2) + self.smoothed_velocities: Optional[torch.Tensor] = None # (K, 2) + + def to_device(self, device: torch.device) -> None: + if self.positions is not None and self.positions.device != device: + self.positions = self.positions.to(device) + self.smoothed_velocities = self.smoothed_velocities.to(device) + + +class VelocityBlockV1(WorkflowBlock): + def __init__(self): + self._states: Dict[str, _VideoVelocityState] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return VelocityManifest + + def run( + self, + image: WorkflowImageData, + detections: Union[ + Detections, InstanceDetections, Tuple[KeyPoints, Optional[Detections]] + ], + smoothing_alpha: float, + pixels_per_meter: float, + ) -> BlockResult: + # The keypoint kind is a (KeyPoints, Detections) tuple; velocity uses the + # bbox component and preserves row order, so the components stay aligned. + key_points, detections = split_key_point_prediction(detections) + num_detections = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(num_detections)] + else: + bboxes_metadata = [ + dict(box_metadata) if box_metadata is not None else {} + for box_metadata in bboxes_metadata + ] + tracker_ids = [ + box_metadata.get("tracker_id") for box_metadata in bboxes_metadata + ] + if num_detections > 0 and any(tracker_id is None for tracker_id in tracker_ids): + raise ValueError( + "tracker_id not initialized, VelocityBlock requires detections to be tracked" + ) + if not (0 < smoothing_alpha <= 1): + raise ValueError( + "smoothing_alpha must be between 0 (exclusive) and 1 (inclusive)" + ) + if not (pixels_per_meter > 0): + raise ValueError("pixels_per_meter must be greater than 0") + + if image.video_metadata.comes_from_video_file and image.video_metadata.fps != 0: + ts_current = image.video_metadata.frame_number / image.video_metadata.fps + else: + ts_current = image.video_metadata.frame_timestamp.timestamp() + + video_id = image.video_metadata.video_identifier + state = self._states.setdefault(video_id, _VideoVelocityState()) + + if num_detections > 0: + self._compute_and_attach_velocities( + detections=detections, + bboxes_metadata=bboxes_metadata, + tracker_ids=[int(tracker_id) for tracker_id in tracker_ids], + state=state, + ts_current=float(ts_current), + smoothing_alpha=float(smoothing_alpha), + pixels_per_meter=float(pixels_per_meter), + ) + detections.bboxes_metadata = bboxes_metadata + + if key_points is not None: + return {OUTPUT_KEY: (key_points, detections)} + return {OUTPUT_KEY: detections} + + @staticmethod + def _compute_and_attach_velocities( + detections: TensorNativeDetections, + bboxes_metadata: List[dict], + tracker_ids: List[int], + state: _VideoVelocityState, + ts_current: float, + smoothing_alpha: float, + pixels_per_meter: float, + ) -> None: + """Vectorised on-device velocity update. + + Geometry (centers, displacement, EMA smoothing, unit conversion) runs + as torch ops on the prediction's device. The only device->host transfer + is the single batched ``.cpu()`` of the (N, 6) result block: the output + contract stores velocities as python floats in ``bboxes_metadata``. + """ + xyxy = detections.xyxy.detach() + device = xyxy.device + state.to_device(device) + centers = ( + xyxy[:, :2].to(torch.float32) + xyxy[:, 2:].to(torch.float32) + ) * 0.5 # (N, 2) + + # Host-side bookkeeping: tracker ids and timestamps are host data. + known_flags = [ + tracker_id in state.rows_by_tracker for tracker_id in tracker_ids + ] + previous_rows = [ + state.rows_by_tracker.get(tracker_id, 0) for tracker_id in tracker_ids + ] + time_deltas = [ + (ts_current - state.timestamps_by_tracker[tracker_id] if known else 0.0) + for tracker_id, known in zip(tracker_ids, known_flags) + ] + + if state.positions is not None and any(known_flags): + row_index = torch.as_tensor(previous_rows, dtype=torch.long, device=device) + previous_positions = state.positions[row_index] # (N, 2) + previous_smoothed = state.smoothed_velocities[row_index] # (N, 2) + else: + previous_positions = centers + previous_smoothed = torch.zeros_like(centers) + + known = torch.as_tensor(known_flags, dtype=torch.bool, device=device) + deltas = torch.as_tensor(time_deltas, dtype=torch.float32, device=device) + has_movement_window = known & (deltas > 0) + safe_deltas = torch.where(deltas > 0, deltas, torch.ones_like(deltas)) + velocity = torch.where( + has_movement_window.unsqueeze(1), + (centers - previous_positions) / safe_deltas.unsqueeze(1), + torch.zeros_like(centers), + ) # pixels per second, (N, 2) + speed = torch.linalg.vector_norm(velocity, dim=1) # (N,) + # EMA smoothing: known trackers blend with their previous smoothed + # velocity; new trackers initialise with the current velocity. State + # keeps unconverted pixels/s values; unit conversion applies to outputs only. + smoothed_velocity = torch.where( + known.unsqueeze(1), + smoothing_alpha * velocity + (1 - smoothing_alpha) * previous_smoothed, + velocity, + ) + smoothed_speed = torch.linalg.vector_norm(smoothed_velocity, dim=1) + + # The single batched device->host hop. + results = ( + torch.cat( + [ + velocity, + speed.unsqueeze(1), + smoothed_velocity, + smoothed_speed.unsqueeze(1), + ], + dim=1, + ) + / pixels_per_meter + ) + results_host = results.cpu().tolist() # (N, 6) + for box_metadata, row in zip(bboxes_metadata, results_host): + box_metadata[VELOCITY_KEY_IN_SV_DETECTIONS] = [row[0], row[1]] + box_metadata[SPEED_KEY_IN_SV_DETECTIONS] = row[2] + box_metadata[SMOOTHED_VELOCITY_KEY_IN_SV_DETECTIONS] = [row[3], row[4]] + box_metadata[SMOOTHED_SPEED_KEY_IN_SV_DETECTIONS] = row[5] + + # State update stays on device. Trackers absent from this frame keep + # their previous position/timestamp/smoothing (they may reappear). + current_set = set(tracker_ids) + survivors = [ + (tracker_id, row) + for tracker_id, row in state.rows_by_tracker.items() + if tracker_id not in current_set + ] + if survivors and state.positions is not None: + survivor_index = torch.as_tensor( + [row for _, row in survivors], dtype=torch.long, device=device + ) + new_positions = torch.cat([state.positions[survivor_index], centers], dim=0) + new_smoothed = torch.cat( + [state.smoothed_velocities[survivor_index], smoothed_velocity], dim=0 + ) + else: + survivors = [] + new_positions = centers + new_smoothed = smoothed_velocity + state.positions = new_positions + state.smoothed_velocities = new_smoothed + rows_by_tracker = { + tracker_id: row for row, (tracker_id, _) in enumerate(survivors) + } + timestamps_by_tracker = { + tracker_id: state.timestamps_by_tracker[tracker_id] + for tracker_id, _ in survivors + } + base = len(survivors) + for offset, tracker_id in enumerate(tracker_ids): + rows_by_tracker[tracker_id] = base + offset + timestamps_by_tracker[tracker_id] = ts_current + state.rows_by_tracker = rows_by_tracker + state.timestamps_by_tracker = timestamps_by_tracker diff --git a/inference/core/workflows/core_steps/classical_cv/contrast_enhancement/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/contrast_enhancement/v1_tensor.py new file mode 100644 index 0000000000..5624ef08c5 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/contrast_enhancement/v1_tensor.py @@ -0,0 +1,116 @@ +"""Tensor-native sibling of ``contrast_enhancement/v1``. + +Tensor-materialised images run the per-channel chain (percentile/min-max +stretch, linear contrast around 128, optional gamma, clip, uint8) as torch ops +on the image's device; numpy/base64-born images delegate to the v1 numpy +implementation. Channel independence makes the math identical for RGB (tensor +layout) and BGR (numpy layout). + +Numpy parity: +- ``astype(np.uint8)`` truncates, so the tensor path truncates too; +- flat-histogram channels (``v_max <= v_min``) pass through untouched; a flat + single-channel image returns the input object, like the numpy grayscale + branch; +- percentiles use numpy's linear interpolation, implemented sort-based + (``torch.quantile`` has an input-size limit and patchy MPS support). +""" + +import math +from typing import Type + +import torch + +from inference.core.workflows.core_steps.classical_cv.contrast_enhancement.v1 import ( + ContrastEnhancementManifest, + enhance_contrast, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlock + + +class ContrastEnhancementBlock(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[ContrastEnhancementManifest]: + return ContrastEnhancementManifest + + def run( + self, + image: WorkflowImageData, + clip_limit: int, + contrast_multiplier: float, + normalize_brightness: bool, + *args, + **kwargs, + ) -> BlockResult: + if not image.is_tensor_materialised(): + return { + "image": enhance_contrast( + image=image, + clip_limit=clip_limit, + contrast_multiplier=contrast_multiplier, + normalize_brightness=normalize_brightness, + ) + } + return { + "image": _enhance_contrast_tensor( + image=image, + clip_limit=clip_limit, + contrast_multiplier=contrast_multiplier, + normalize_brightness=normalize_brightness, + ) + } + + +def _enhance_contrast_tensor( + image: WorkflowImageData, + clip_limit: int, + contrast_multiplier: float, + normalize_brightness: bool, +) -> WorkflowImageData: + """Device-resident mirror of ``v1.enhance_contrast`` for CHW tensors.""" + chw = image.tensor_image # (C, H, W) uint8, C in {1, 3} + clip_pct = float(clip_limit) / 100.0 + contrast_mult = float(contrast_multiplier) + gamma_val = 1.0 / 1.3 if normalize_brightness else 1.0 + + channels = int(chw.shape[0]) + flat = chw.detach().reshape(channels, -1).to(torch.float32) + if clip_pct > 0: + sorted_values, _ = torch.sort(flat, dim=1) + v_min = _linear_percentile(sorted_values, clip_pct) + v_max = _linear_percentile(sorted_values, 1.0 - clip_pct) + else: + v_min = flat.amin(dim=1) + v_max = flat.amax(dim=1) + + valid = v_max > v_min # (C,) - flat-histogram channels pass through raw + if channels == 1 and not bool(valid.any()): + # Numpy grayscale parity: nothing to stretch -> the input is returned. + return image + + spread = torch.where(valid, v_max - v_min, torch.ones_like(v_max)) + normalized = (flat - v_min.unsqueeze(1)) / spread.unsqueeze(1) * 255.0 + if contrast_mult != 1.0: + normalized = 128.0 + (normalized - 128.0) * contrast_mult + if gamma_val != 1.0: + normalized = normalized.clamp(0.0, 255.0) + normalized = torch.pow(normalized / 255.0, gamma_val) * 255.0 + normalized = normalized.clamp(0.0, 255.0) + enhanced = torch.where(valid.unsqueeze(1), normalized, flat) + # Numpy parity: astype(np.uint8) truncates, and so does .to(torch.uint8). + enhanced_chw = enhanced.to(torch.uint8).reshape(chw.shape) + return WorkflowImageData.copy_and_replace( + origin_image_data=image, + tensor_image=enhanced_chw, + ) + + +def _linear_percentile(sorted_values: torch.Tensor, quantile: float) -> torch.Tensor: + """Per-row percentile of pre-sorted ``(C, N)`` values with numpy's default + linear interpolation - equivalent to ``np.percentile(..., method="linear")``.""" + n = sorted_values.shape[1] + position = quantile * (n - 1) + lower = int(math.floor(position)) + upper = min(int(math.ceil(position)), n - 1) + weight = position - lower + return sorted_values[:, lower] * (1.0 - weight) + sorted_values[:, upper] * weight diff --git a/inference/core/workflows/core_steps/classical_cv/contrast_equalization/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/contrast_equalization/v1_tensor.py new file mode 100644 index 0000000000..e29ffd8337 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/contrast_equalization/v1_tensor.py @@ -0,0 +1,127 @@ +"""Tensor-native sibling of ``contrast_equalization/v1``. + +'Contrast Stretching' and 'Histogram Equalization' are global value maps driven +by the image-wide 256-bin histogram (layout-independent: the same multiset of +values whether CHW/RGB or HWC/BGR). The tensor path runs ``torch.bincount`` on +the device, syncs the 256 counts to host to build a 256-entry LUT with the +identical numpy/skimage arithmetic of v1 (bit-exact result versus the numpy +block), and applies one device gather (``lut[image]``). + +'Adaptive Equalization' (CLAHE) is position-dependent (tile-local mappings), so +it cannot be expressed as a value LUT and delegates to v1; numpy/base64-born +images delegate as well. +""" + +import math +from typing import Type + +import numpy as np +import torch + +from inference.core.workflows.core_steps.classical_cv.contrast_equalization.v1 import ( + ContrastEqualizationManifest, + update_image, +) +from inference.core.workflows.core_steps.visualizations.common.base import ( + OUTPUT_IMAGE_KEY, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlock + + +class ContrastEqualizationBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[ContrastEqualizationManifest]: + return ContrastEqualizationManifest + + def run( + self, + image: WorkflowImageData, + equalization_type: str, + ) -> BlockResult: + if ( + not image.is_tensor_materialised() + or equalization_type == "Adaptive Equalization" + ): + updated_image = update_image(image.numpy_image, equalization_type) + output = WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=updated_image, + ) + return {OUTPUT_IMAGE_KEY: output} + return { + OUTPUT_IMAGE_KEY: _equalize_tensor( + image=image, equalization_type=equalization_type + ) + } + + +def _equalize_tensor( + image: WorkflowImageData, equalization_type: str +) -> WorkflowImageData: + if equalization_type == "Contrast Stretching": + build_lut = _contrast_stretching_lut + elif equalization_type == "Histogram Equalization": + build_lut = _histogram_equalization_lut + else: + raise ValueError( + f"contrast equalization type `{equalization_type}` not implemented!" + ) + chw = image.tensor_image + flat_indices = chw.detach().reshape(-1).long() + counts = torch.bincount(flat_indices, minlength=256).cpu().numpy() + lut = torch.from_numpy(build_lut(counts)).to(chw.device) + equalized = lut[flat_indices].reshape(chw.shape) + return WorkflowImageData.copy_and_replace( + origin_image_data=image, + tensor_image=equalized, + ) + + +def _contrast_stretching_lut(counts: np.ndarray) -> np.ndarray: + """256-entry LUT of ``exposure.rescale_intensity(img, in_range=(p2, p98))`` + for uint8 images, replicating skimage's float64 arithmetic exactly.""" + p2 = _percentile_of_uint8_counts(counts, 2.0) + p98 = _percentile_of_uint8_counts(counts, 98.0) + values = np.arange(256, dtype=np.uint8) + clipped = np.clip(values, p2, p98) # float64, matching skimage + if p2 != p98: + scaled = (clipped - p2) / (p98 - p2) + return (scaled * 255.0).astype(np.uint8) + return np.clip(clipped, 0.0, 255.0).astype(np.uint8) + + +def _histogram_equalization_lut(counts: np.ndarray) -> np.ndarray: + """256-entry LUT of v1's ``equalize_hist(img.astype(float32) / 255) * 255`` + -> uint8 chain, replicating skimage's dtype handling exactly.""" + values = np.arange(256, dtype=np.float32) / 255 # v1's float32 division + present = counts > 0 + # np.histogram derives bin edges from data min/max; min/max of the present + # distinct values equal those of the full image, so the float32 bin edges + # (and each value's bin) match skimage's histogram of all pixels. + hist, bin_edges = np.histogram(values[present], bins=256, weights=counts[present]) + bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2.0 + img_cdf = hist.cumsum() + img_cdf = img_cdf / float(img_cdf[-1]) + img_cdf = img_cdf.astype(np.float32) # cumulative_distribution's cast + equalized = np.interp(values, bin_centers, img_cdf) # float64 output + equalized = equalized.astype(np.float32) # equalize_hist's cast + return (equalized * 255).astype(np.uint8) + + +def _percentile_of_uint8_counts(counts: np.ndarray, q: float) -> float: + """``np.percentile(values, q)`` (default linear method) computed from exact + 256-bin counts of the uint8 values, including numpy's ``_lerp`` fixup.""" + cumulative = np.cumsum(counts) + n = int(cumulative[-1]) + position = (q / 100.0) * (n - 1) + lower_rank = int(math.floor(position)) + fraction = position - lower_rank + upper_rank = min(lower_rank + 1, n - 1) + # k-th order statistic = smallest value whose cumulative count reaches k+1 + lower_value = float(np.searchsorted(cumulative, lower_rank + 1, side="left")) + upper_value = float(np.searchsorted(cumulative, upper_rank + 1, side="left")) + difference = upper_value - lower_value + if fraction >= 0.5: + return upper_value - difference * (1.0 - fraction) + return lower_value + difference * fraction diff --git a/inference/core/workflows/core_steps/classical_cv/convert_grayscale/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/convert_grayscale/v1_tensor.py new file mode 100644 index 0000000000..54590f416f --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/convert_grayscale/v1_tensor.py @@ -0,0 +1,73 @@ +"""Tensor-native sibling of ``convert_grayscale/v1``. + +OpenCV 4.x's uint8 BGR2GRAY is the fixed-point map (``color_rgb2gray`` +``RY15``/``GY15``/``BY15`` - not the widely quoted legacy +``4899/9617/1868 >> 14`` constants, which disagree with modern cv2): + + gray = (9798*R + 19235*G + 3735*B + (1 << 14)) >> 15 + +The tensor path computes it in int32 on the device (max accumulator +``255*32768 + 2^14 = 8_372_224``, inside int32) and emits the ``(1, H, W)`` +contract shape, bit-exact versus the numpy block. + +Delegation to v1: +- numpy/base64-born images (no materialised tensor); also covers 4-channel + BGRA input, which can only arrive numpy-born; +- tensor-born single-channel images: the materialised 2-D numpy view makes + ``cv2.cvtColor`` raise the same ``cv2.error`` as v1. +""" + +from typing import Type + +import cv2 +import torch + +from inference.core.workflows.core_steps.classical_cv.convert_grayscale.v1 import ( + ConvertGrayscaleManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base import ( + OUTPUT_IMAGE_KEY, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlock + +# OpenCV 4.x u8 BGR2GRAY fixed-point constants (color_rgb2gray RY15/GY15/BY15). +_RY15 = 9798 +_GY15 = 19235 +_BY15 = 3735 +_GRAY_SHIFT = 15 +_GRAY_ROUND = 1 << (_GRAY_SHIFT - 1) + + +class ConvertGrayscaleBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[ConvertGrayscaleManifest]: + return ConvertGrayscaleManifest + + def run( + self, + image: WorkflowImageData, + *args, + **kwargs, + ) -> BlockResult: + if not image.is_tensor_materialised() or image.tensor_image.shape[0] != 3: + gray = cv2.cvtColor(image.numpy_image, cv2.COLOR_BGR2GRAY) + output = WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=gray + ) + return {OUTPUT_IMAGE_KEY: output} + return {OUTPUT_IMAGE_KEY: _convert_grayscale_tensor(image=image)} + + +def _convert_grayscale_tensor(image: WorkflowImageData) -> WorkflowImageData: + """Device-resident mirror of ``cv2.cvtColor(..., cv2.COLOR_BGR2GRAY)`` for + ``(3, H, W)`` RGB uint8 tensors, emitting the ``(1, H, W)`` contract shape.""" + chw = image.tensor_image.detach().to(torch.int32) + weighted = ( + chw[0] * _RY15 + chw[1] * _GY15 + chw[2] * _BY15 + _GRAY_ROUND + ) >> _GRAY_SHIFT + gray_chw = weighted.to(torch.uint8).unsqueeze(0) + return WorkflowImageData.copy_and_replace( + origin_image_data=image, + tensor_image=gray_chw, + ) diff --git a/inference/core/workflows/core_steps/classical_cv/detections_nearest_neighbor/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/detections_nearest_neighbor/v1_tensor.py new file mode 100644 index 0000000000..ad8290c6a5 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/detections_nearest_neighbor/v1_tensor.py @@ -0,0 +1,316 @@ +"""Tensor-native sibling of detections_nearest_neighbor/v1.py. + +The numpy block performs a nearest-neighbor spatial join on ``sv.Detections``: +anchor points are resolved per detection (bbox anchors or a named keypoint), +a brute-force pairwise Euclidean pixel-distance matrix is built, self-matches +(shared ``detection_id``) and targets beyond ``max_distance`` are excluded, and +each query's minimum distance (with a 1 px tie epsilon) selects the matched +target rows. The nearest distance lands in +``query_predictions.data["nearest_target_distance"]``. + +This sibling keeps exactly those semantics on the native inference_models +dataclasses: anchor points and the distance matrix are computed as torch ops on +the prediction's device (invalid candidates are masked instead of NaN-assigned, +with the same nan-propagation behaviour for missing keypoints), and the single +device->host hops are the batched ``.cpu()`` of the per-query minima and the +tie-mask indices. The computed per-box scalar is written to +``bboxes_metadata[i][NEAREST_TARGET_DISTANCE_KEY]`` (``None`` when a query has +no eligible target) - the native equivalent of the ``sv.Detections.data`` +column, mutating the input prediction in place like the numpy block. Keypoint +input arrives as a ``(KeyPoints, Detections)`` tuple; the bbox component +carries the flattened per-box ``keypoints_xy`` / ``keypoints_class_name`` +payloads in ``bboxes_metadata`` (the native mirror of the sv data columns) and +the matched outputs are sliced with ``take_prediction_by_indices`` so keypoint +tensors and instance masks ride along, duplicated rows included on ties. +""" + +import math +from typing import List, Optional, Tuple, Type, Union + +import torch + +from inference.core.workflows.core_steps.classical_cv.detections_nearest_neighbor.v1 import ( + KEYPOINT_POINT_OPTION, + OUTPUT_KEY_MATCHED_QUERY_DETECTIONS, + OUTPUT_KEY_MATCHED_TARGET_DETECTIONS, + OUTPUT_KEY_QUERY_PREDICTIONS, + TIE_EPSILON_PX, + BlockManifest, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + split_key_point_prediction, + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.constants import ( + DETECTION_ID_KEY, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + NEAREST_TARGET_DISTANCE_KEY, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +TensorNativeDetections = Union[Detections, InstanceDetections] +KeyPointPrediction = Tuple[KeyPoints, Optional[Detections]] +NearestNeighborInput = Union[Detections, InstanceDetections, KeyPointPrediction] + + +class DetectionsNearestNeighborBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + query_predictions: NearestNeighborInput, + target_predictions: NearestNeighborInput, + query_point: str, + target_point: str, + query_keypoint_name: Optional[str], + target_keypoint_name: Optional[str], + max_distance: Optional[int], + ) -> BlockResult: + if query_point == KEYPOINT_POINT_OPTION and not query_keypoint_name: + raise ValueError( + "`query_keypoint_name` must be provided when `query_point` is set to 'KEYPOINT'." + ) + if target_point == KEYPOINT_POINT_OPTION and not target_keypoint_name: + raise ValueError( + "`target_keypoint_name` must be provided when `target_point` is set to 'KEYPOINT'." + ) + + # The keypoint-detection kind is a (KeyPoints, Detections) tuple; the + # bbox component carries xyxy, detection ids and the flattened per-box + # keypoint payloads used below. + _, query_detections = split_key_point_prediction(query_predictions) + _, target_detections = split_key_point_prediction(target_predictions) + + query_points = resolve_anchor_points( + detections=query_detections, + point=query_point, + keypoint_name=query_keypoint_name, + ) + target_points = resolve_anchor_points( + detections=target_detections, + point=target_point, + keypoint_name=target_keypoint_name, + ) + + distances, matched_query_indices, matched_target_indices = ( + match_query_to_targets( + query_detections=query_detections, + target_detections=target_detections, + query_points=query_points, + target_points=target_points, + max_distance=max_distance, + ) + ) + # Mutates the input prediction in place (same convention as the numpy + # block and Velocity/Time in Zone): the per-box scalar lands in + # `bboxes_metadata`, so `matched_query_detections` below picks it up + # for free via the index-slice. Entry dicts are copied so the write + # cannot leak into other references of the same metadata dicts. + number_of_queries = int(query_detections.xyxy.shape[0]) + bboxes_metadata = query_detections.bboxes_metadata + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(number_of_queries)] + else: + bboxes_metadata = [ + dict(box_metadata) if box_metadata is not None else {} + for box_metadata in bboxes_metadata + ] + for box_metadata, distance in zip(bboxes_metadata, distances): + box_metadata[NEAREST_TARGET_DISTANCE_KEY] = distance + query_detections.bboxes_metadata = bboxes_metadata + + return { + OUTPUT_KEY_QUERY_PREDICTIONS: query_predictions, + OUTPUT_KEY_MATCHED_QUERY_DETECTIONS: take_prediction_by_indices( + query_predictions, matched_query_indices + ), + OUTPUT_KEY_MATCHED_TARGET_DETECTIONS: take_prediction_by_indices( + target_predictions, matched_target_indices + ), + } + + +# Bbox anchor resolution: (x, y) per option as (column selectors) over +# [x_min, y_min, x_max, y_max, center_x, center_y] - the torch mirror of +# ``sv.Detections.get_anchors_coordinates(anchor=sv.Position[point])``. +_X_MIN, _Y_MIN, _X_MAX, _Y_MAX, _CENTER_X, _CENTER_Y = range(6) +_BBOX_ANCHOR_COLUMNS = { + "CENTER": (_CENTER_X, _CENTER_Y), + "CENTER_LEFT": (_X_MIN, _CENTER_Y), + "CENTER_RIGHT": (_X_MAX, _CENTER_Y), + "TOP_CENTER": (_CENTER_X, _Y_MIN), + "TOP_LEFT": (_X_MIN, _Y_MIN), + "TOP_RIGHT": (_X_MAX, _Y_MIN), + "BOTTOM_LEFT": (_X_MIN, _Y_MAX), + "BOTTOM_CENTER": (_CENTER_X, _Y_MAX), + "BOTTOM_RIGHT": (_X_MAX, _Y_MAX), +} + + +def resolve_anchor_points( + detections: TensorNativeDetections, + point: str, + keypoint_name: Optional[str], +) -> torch.Tensor: + """Resolve one (x, y) anchor per detection as an ``(N, 2)`` float32 tensor + on the prediction's device.""" + if point == KEYPOINT_POINT_OPTION: + return resolve_keypoint_anchor_points( + detections=detections, keypoint_name=keypoint_name + ) + if point not in _BBOX_ANCHOR_COLUMNS: + raise ValueError( + f"Invalid anchor point option '{point}'. Supported options: " + f"{sorted(_BBOX_ANCHOR_COLUMNS)} or '{KEYPOINT_POINT_OPTION}'." + ) + xyxy = detections.xyxy.detach().to(dtype=torch.float32) + centers = (xyxy[:, :2] + xyxy[:, 2:]) * 0.5 + columns = torch.cat([xyxy, centers], dim=1) # (N, 6) + x_column, y_column = _BBOX_ANCHOR_COLUMNS[point] + return torch.stack([columns[:, x_column], columns[:, y_column]], dim=1) + + +def resolve_keypoint_anchor_points( + detections: TensorNativeDetections, + keypoint_name: str, +) -> torch.Tensor: + """Read the named keypoint per detection from the flattened + ``bboxes_metadata`` payloads (``keypoints_xy`` / ``keypoints_class_name`` - + the native mirror of the sv data columns the numpy block reads). + + Detections whose keypoint set does not include ``keypoint_name`` (e.g. an + occluded joint) get a NaN anchor point, which propagates through the + distance matrix and naturally excludes them from matching - the same + graceful degradation as the numpy block. + """ + number_of_detections = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [] + has_keypoint_payload = any( + box_metadata is not None + and KEYPOINTS_XY_KEY_IN_SV_DETECTIONS in box_metadata + and KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS in box_metadata + for box_metadata in bboxes_metadata + ) + if number_of_detections > 0 and not has_keypoint_payload: + raise ValueError( + "`query_point`/`target_point` set to 'KEYPOINT' but the corresponding " + "predictions do not contain keypoint data. Provide keypoint detection " + "predictions to use this option." + ) + points = torch.full((number_of_detections, 2), float("nan"), dtype=torch.float32) + for index in range(number_of_detections): + box_metadata = ( + bboxes_metadata[index] if index < len(bboxes_metadata) else None + ) or {} + keypoint_names = box_metadata.get(KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS) + keypoints_xy = box_metadata.get(KEYPOINTS_XY_KEY_IN_SV_DETECTIONS) + if keypoint_names is None or keypoints_xy is None: + continue + for name, coordinates in zip(keypoint_names, keypoints_xy): + if str(name) == keypoint_name: + points[index, 0] = float(coordinates[0]) + points[index, 1] = float(coordinates[1]) + break + return points.to(device=detections.xyxy.device) + + +def _detection_ids(detections: TensorNativeDetections) -> List[Optional[str]]: + """Per-box ``detection_id`` (``None`` when absent) - the native mirror of + the ``sv.Detections.data["detection_id"]`` column.""" + number_of_detections = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [] + return [ + ( + (bboxes_metadata[index] or {}).get(DETECTION_ID_KEY) + if index < len(bboxes_metadata) + else None + ) + for index in range(number_of_detections) + ] + + +def match_query_to_targets( + query_detections: TensorNativeDetections, + target_detections: TensorNativeDetections, + query_points: torch.Tensor, + target_points: torch.Tensor, + max_distance: Optional[int], +) -> Tuple[List[Optional[float]], List[int], List[int]]: + num_query = int(query_points.shape[0]) + if int(target_points.shape[0]) == 0: + return [None] * num_query, [], [] + + device = query_points.device + target_points = target_points.to(device) + + # Plain brute-force pairwise distance matrix (not a KD-tree), same as the + # numpy block: typical detection counts here are tens per set. + diff = query_points[:, None, :] - target_points[None, :, :] + distance_matrix = torch.sqrt((diff * diff).sum(dim=-1)) # (Q, T) + + # Invalid candidates are tracked in a boolean mask instead of NaN-assigning + # the matrix (torch mirror of the numpy block's NaN bookkeeping). NaN + # distances from missing-keypoint anchors start out invalid. + valid = ~torch.isnan(distance_matrix) + + # A target counts as a self-match (and is excluded) whenever it shares the + # query detection's `detection_id`. Boxes without a `detection_id` are + # never treated as self-matches - when neither side carries ids the + # exclusion is skipped entirely, mirroring the numpy block's behaviour for + # a missing `detection_id` column. + query_ids = _detection_ids(query_detections) + target_ids = _detection_ids(target_detections) + if any(query_id is not None for query_id in query_ids) and any( + target_id is not None for target_id in target_ids + ): + self_match = torch.tensor( + [ + [ + query_id is not None and query_id == target_id + for target_id in target_ids + ] + for query_id in query_ids + ], + dtype=torch.bool, + device=device, + ) + valid &= ~self_match + + if max_distance is not None: + # Candidates beyond the limit are dropped before ranking (boundary + # inclusive), so a tie can only form among in-range targets. + valid &= distance_matrix <= max_distance + + infinity = float("inf") + filled = torch.where( + valid, distance_matrix, torch.full_like(distance_matrix, infinity) + ) + # An all-invalid query row keeps +inf as its minimum and is omitted from + # the paired outputs (no placeholder row); `query_predictions` still + # carries it, with `nearest_target_distance` set to `None`. + min_per_row = filled.min(dim=1).values # (Q,) + + # A tie duplicates the query row once per tied target; `torch.nonzero` is + # row-major like `np.where`, keeping the two paired outputs index-aligned. + tie_mask = valid & (filled <= (min_per_row[:, None] + TIE_EPSILON_PX)) + + # The only device->host hops: the batched minima and the tie indices. + matched_pairs = torch.nonzero(tie_mask, as_tuple=False).cpu().tolist() + matched_query_indices = [pair[0] for pair in matched_pairs] + matched_target_indices = [pair[1] for pair in matched_pairs] + distances = [ + None if math.isinf(distance) else float(distance) + for distance in min_per_row.cpu().tolist() + ] + return distances, matched_query_indices, matched_target_indices diff --git a/inference/core/workflows/core_steps/classical_cv/distance_measurement/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/distance_measurement/v1_tensor.py new file mode 100644 index 0000000000..e307299db5 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/distance_measurement/v1_tensor.py @@ -0,0 +1,516 @@ +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +from pydantic import ConfigDict, Field + +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + INTEGER_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +SHORT_DESCRIPTION = "Calculate the distance between two bounding boxes on a 2D plane." + +LONG_DESCRIPTION = """ +Calculate the distance between two detected objects on a 2D plane using bounding box coordinates, supporting horizontal or vertical distance measurement along a specified axis, with two calibration methods (reference object with known dimensions or pixel-to-centimeter ratio) to convert pixel distances to real-world measurements for spatial analysis, object spacing assessment, safety monitoring, and measurement workflows. + +## How This Block Works + +This block measures the distance between two detected objects by analyzing their bounding box positions and converting pixel distances to real-world units (centimeters). The block: + +1. Receives detection predictions containing bounding boxes and class names for objects in the image +2. Identifies the two target objects using their class names (object_1_class_name and object_2_class_name): + - Searches through all detections to find bounding boxes matching the specified class names + - Extracts bounding box coordinates (x_min, y_min, x_max, y_max) for both objects + - Validates that both objects are found in the detections +3. Validates object positioning for distance measurement: + - Checks if bounding boxes overlap (if they overlap, distance is set to 0) + - Verifies objects have a gap along the specified reference axis (horizontal or vertical) + - Returns 0 distance if objects overlap or are positioned incorrectly for the selected axis +4. Determines the calibration method and performs calibration: + + **For Reference Object Calibration:** + - Searches detections for a reference object with known real-world dimensions (reference_object_class_name) + - Extracts the reference object's bounding box coordinates + - Measures reference object dimensions in pixels (width and height) + - Calculates pixel-to-centimeter ratios: + - Width ratio: reference_width_pixels / reference_width (cm) + - Height ratio: reference_height_pixels / reference_height (cm) + - Computes average pixel ratio from width and height ratios for more accurate scaling + - Uses the average ratio to convert all pixel measurements to centimeters + + **For Pixel-to-Centimeter Ratio Calibration:** + - Uses the provided pixel_ratio directly (e.g., 100 pixels = 1 centimeter) + - Applies the ratio to convert pixel distances to centimeter distances + - Suitable when the pixel-to-real-world scale is already known or calibrated + +5. Measures pixel distance between the two objects along the specified axis: + - **For Vertical Distance**: Calculates distance along the Y-axis (vertical separation) + - Finds the gap between bounding boxes vertically + - Measures distance from bottom of upper object to top of lower object (or vice versa) + - Accounts for bounding box positions to find the actual gap distance + - **For Horizontal Distance**: Calculates distance along the X-axis (horizontal separation) + - Finds the gap between bounding boxes horizontally + - Measures distance from right edge of left object to left edge of right object (or vice versa) + - Accounts for bounding box positions to find the actual gap distance +6. Converts pixel distance to centimeter distance: + - Divides pixel distance by the pixel-to-centimeter ratio (from calibration) + - Produces real-world distance measurement in centimeters +7. Returns both pixel distance and centimeter distance values + +The block assumes a perpendicular camera view (top-down or frontal view) where perspective distortion is minimal, ensuring accurate 2D distance measurements. Distance is measured as the gap between bounding boxes along the specified axis (horizontal or vertical), not the diagonal distance between object centers. The calibration process converts pixel measurements to real-world units using either a reference object with known dimensions (more flexible, works with different scales) or a direct pixel ratio (simpler, requires pre-calibration). This enables accurate spatial measurements for monitoring, analysis, and control applications. + +## Common Use Cases + +- **Safety Monitoring**: Measure distances between objects to ensure safe spacing (e.g., measure distance between people for social distancing, monitor spacing between vehicles, ensure safe gaps in industrial settings), enabling safety monitoring workflows +- **Warehouse Management**: Measure spacing between items or objects in storage and logistics (e.g., measure gaps between packages, assess shelf spacing, monitor object placement), enabling warehouse management workflows +- **Quality Control**: Verify spacing and positioning of objects in manufacturing and assembly (e.g., measure gaps between components, verify spacing in assembly lines, check positioning accuracy), enabling quality control workflows +- **Traffic Analysis**: Measure distances between vehicles or objects in traffic monitoring (e.g., measure vehicle spacing, assess safe following distances, monitor traffic gaps), enabling traffic analysis workflows +- **Retail Analytics**: Measure spacing between products or customers in retail environments (e.g., measure product spacing on shelves, assess customer spacing, monitor display arrangements), enabling retail analytics workflows +- **Agricultural Monitoring**: Measure spacing between crops, plants, or agricultural objects (e.g., measure crop spacing, assess plant gaps, monitor field arrangements), enabling agricultural monitoring workflows + +## Connecting to Other Blocks + +This block receives detection predictions and produces distance_cm and distance_pixel values: + +- **After object detection or instance segmentation blocks** to measure distances between detected objects (e.g., measure distance between detected objects, calculate spacing from detections, analyze object relationships), enabling detection-to-measurement workflows +- **Before logic blocks** like Continue If to make decisions based on distance measurements (e.g., continue if distance is safe, filter based on spacing requirements, make decisions using distance thresholds), enabling distance-based decision workflows +- **Before analysis blocks** to analyze spatial relationships between objects (e.g., analyze object spacing, process distance measurements, work with spatial data), enabling spatial analysis workflows +- **Before notification blocks** to alert when distances violate thresholds (e.g., send alerts when spacing is too close, notify on distance violations, trigger actions based on measurements), enabling distance-based notification workflows +- **Before data storage blocks** to record distance measurements (e.g., store distance measurements, log spacing data, record spatial metrics), enabling distance measurement logging workflows +- **In measurement pipelines** where distance calculation is part of a larger spatial analysis workflow (e.g., measure distances in analysis pipelines, calculate spacing in monitoring systems, process spatial measurements in chains), enabling spatial measurement pipeline workflows + +## Requirements + +This block requires detection predictions with bounding boxes and class names. The image should be captured from a perpendicular camera view (top-down or frontal) to minimize perspective distortion and ensure accurate 2D distance measurements. For reference object calibration, a reference object with known dimensions must be present in the detections. For pixel-to-centimeter ratio calibration, the pixel ratio must be pre-calibrated or known for the camera setup. Objects must not overlap and must have a gap along the specified measurement axis (horizontal or vertical). The block assumes objects are on the same plane for accurate 2D measurement. +""" + +OUTPUT_KEY_CENTIMETER = "distance_cm" +OUTPUT_KEY_PIXEL = "distance_pixel" + +TensorNativeDetections = Union[Detections, InstanceDetections] + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Distance Measurement", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "classical_computer_vision", + "ui_manifest": { + "section": "classical_cv", + "icon": "far fa-ruler-triangle", + "blockPriority": 11, + "opencv": True, + }, + } + ) + + type: Literal["roboflow_core/distance_measurement@v1"] + + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( + title="Object Detections", + description="Detection predictions containing bounding boxes and class names for objects in the image. Must include detections for the two objects to measure (object_1_class_name and object_2_class_name) and optionally a reference object (if using reference object calibration method). The bounding boxes will be used to calculate distances between objects. Both object detection and instance segmentation predictions are supported. The detections must contain class_name information to identify objects.", + examples=["$steps.model.predictions"], + ) + + object_1_class_name: str = Field( + title="First Object Class Name", + description="Class name of the first object to measure distance from. Must match exactly the class name in the detection predictions. The block searches for this class name in the detections and uses its bounding box for distance calculation. Example: if detections contain objects labeled 'person', 'car', 'bicycle', use 'person' to measure distance from a person to another object. The class name is case-sensitive and must match exactly.", + examples=["car"], + ) + + object_2_class_name: str = Field( + title="Second Object Class Name", + description="Class name of the second object to measure distance to. Must match exactly the class name in the detection predictions. The block searches for this class name in the detections and uses its bounding box for distance calculation. Example: if detections contain objects labeled 'person', 'car', 'bicycle', use 'person' to measure distance to a person from another object. The class name is case-sensitive and must match exactly. The block measures the gap between object_1 and object_2 along the specified reference_axis.", + examples=["person"], + ) + + reference_axis: Literal["horizontal", "vertical"] = Field( + title="Reference Axis", + description="Axis along which to measure the distance between the two objects. Options: 'horizontal' measures distance along the X-axis (left-right gap between objects, useful when objects are side-by-side), or 'vertical' measures distance along the Y-axis (top-bottom gap between objects, useful when objects are stacked vertically). The distance is measured as the gap between bounding boxes along the selected axis. Objects must have a gap along this axis (not overlap) for accurate measurement. Choose based on object orientation: horizontal for side-by-side objects, vertical for stacked objects.", + examples=["vertical", "horizontal", "$inputs.reference_axis"], + ) + + calibration_method: Literal["reference object", "pixel to centimeter"] = Field( + title="Calibration Method", + description="Method to calibrate pixel measurements to real-world units (centimeters). Options: 'reference object' (uses a reference object with known dimensions in the image to calculate pixel-to-centimeter ratio automatically, more flexible for different scales), or 'pixel to centimeter' (uses a pre-calibrated pixel ratio directly, simpler but requires known scale). For reference object method, a reference object must be present in detections with known width and height. For pixel ratio method, the pixel_ratio must be pre-calibrated for your camera setup.", + ) + + reference_object_class_name: Union[str, Selector(kind=[STRING_KIND])] = Field( + title="Reference Object Class Name", + description="Class name of the reference object used for calibration (only used when calibration_method is 'reference object'). Must match exactly the class name in the detection predictions. The reference object must have known real-world dimensions (reference_width and reference_height). The block measures the reference object's pixel dimensions and calculates a pixel-to-centimeter ratio to convert all distance measurements. Default is 'reference-object'. The reference object must be present in the detections and should be clearly visible and correctly detected.", + default="reference-object", + examples=["reference-object", "$inputs.reference_object_class_name"], + json_schema_extra={ + "relevant_for": { + "calibration_method": { + "values": ["reference object"], + "required": True, + }, + }, + }, + ) + + reference_width: Union[float, Selector(kind=[FLOAT_KIND])] = Field( + title="Width", + default=2.5, + description="Real-world width of the reference object in centimeters (only used when calibration_method is 'reference object'). Must be greater than 0. This is the actual physical width of the reference object. The block measures the reference object's width in pixels and divides by this value to calculate the pixel-to-centimeter ratio. Use accurate measurements for best results. Example: if your reference object is a 2.5cm wide card, use 2.5. The reference_width and reference_height are used to calculate separate width and height ratios, then averaged for more accurate scaling.", + examples=[2.5, "$inputs.reference_width"], + gt=0, + json_schema_extra={ + "relevant_for": { + "calibration_method": { + "values": ["reference object"], + "required": True, + }, + }, + }, + ) + + reference_height: Union[float, Selector(kind=[FLOAT_KIND])] = Field( # type: ignore + title="Height", + default=2.5, + description="Real-world height of the reference object in centimeters (only used when calibration_method is 'reference object'). Must be greater than 0. This is the actual physical height of the reference object. The block measures the reference object's height in pixels and divides by this value to calculate the pixel-to-centimeter ratio. Use accurate measurements for best results. Example: if your reference object is a 2.5cm tall card, use 2.5. The reference_width and reference_height are used to calculate separate width and height ratios, then averaged for more accurate scaling.", + examples=[2.5, "$inputs.reference_height"], + gt=0, + json_schema_extra={ + "relevant_for": { + "calibration_method": { + "values": ["reference object"], + "required": True, + }, + }, + }, + ) + + pixel_ratio: Union[float, Selector(kind=[FLOAT_KIND])] = Field( + title="Reference Pixel-to-Centimeter Ratio", + description="Pixel-to-centimeter conversion ratio for the image (only used when calibration_method is 'pixel to centimeter'). Must be greater than 0. This value represents how many pixels equal 1 centimeter. Example: if 100 pixels = 1 centimeter, use 100. The block divides pixel distances by this ratio to convert to centimeters. This ratio must be pre-calibrated for your specific camera setup, viewing distance, and image resolution. Typical values range from 10-500 depending on camera distance and resolution. A higher ratio means more pixels per centimeter (objects appear larger, camera is closer), a lower ratio means fewer pixels per centimeter (objects appear smaller, camera is farther).", + default=100, + examples=[100, "$inputs.pixel_ratio"], + gt=0, + json_schema_extra={ + "relevant_for": { + "calibration_method": { + "values": ["pixel to centimeter"], + "required": True, + }, + }, + }, + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY_CENTIMETER, + kind=[INTEGER_KIND], + ), + OutputDefinition( + name=OUTPUT_KEY_PIXEL, + kind=[INTEGER_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.0.0,<2.0.0" + + +class DistanceMeasurementBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + predictions: TensorNativeDetections, + object_1_class_name: str, + object_2_class_name: str, + reference_axis: Literal["horizontal", "vertical"], + calibration_method: Literal["reference object", "pixel to centimeter"], + reference_object_class_name: str, + reference_width: float, + reference_height: float, + pixel_ratio: float, + ) -> BlockResult: + if calibration_method == "reference object": + reference_predictions = predictions + distances = measure_distance_with_reference_object( + detections=predictions, + object_1_class_name=object_1_class_name, + object_2_class_name=object_2_class_name, + reference_predictions=reference_predictions, + reference_object_class_name=reference_object_class_name, + reference_width=reference_width, + reference_height=reference_height, + reference_axis=reference_axis, + ) + elif calibration_method == "pixel to centimeter": + distances = measure_distance_with_pixel_ratio( + detections=predictions, + pixel_ratio=pixel_ratio, + object_1_class_name=object_1_class_name, + object_2_class_name=object_2_class_name, + reference_axis=reference_axis, + ) + else: + raise ValueError(f"Invalid calibration type: {calibration_method}") + + return distances + + +def _iter_boxes_with_class_names( + detections: TensorNativeDetections, +) -> List[Tuple[Tuple[int, int, int, int], str]]: + """Yield ``((x_min, y_min, x_max, y_max), class_name)`` pairs for each native + detection โ€” the tensor-native equivalent of + ``zip(detections.xyxy.round().astype(int), detections.data["class_name"])``. + + Bounding boxes are materialised host-side and rounded to ``int`` (matching the + numpy block's ``xyxy.round().astype(int)``); per-box class names mirror the + numpy block's ``detections.data["class_name"]`` by preferring an explicit + per-box label (``bboxes_metadata[i][CLASS_NAME_KEY]``, set by producers such as + vlm_as_detector / google_vision_ocr) and falling back to ``class_id`` indexed + into ``image_metadata[CLASS_NAMES_KEY]``. + """ + xyxy = detections.xyxy.detach().to("cpu").numpy().round().astype(dtype=int) + class_ids = detections.class_id.detach().to("cpu").numpy() + class_names_map = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + bboxes_metadata = detections.bboxes_metadata or [] + pairs: List[Tuple[Tuple[int, int, int, int], str]] = [] + for index, (box, class_id) in enumerate(zip(xyxy, class_ids)): + x_min, y_min, x_max, y_max = ( + int(box[0]), + int(box[1]), + int(box[2]), + int(box[3]), + ) + per_box_name = None + if index < len(bboxes_metadata): + per_box_name = bboxes_metadata[index].get(CLASS_NAME_KEY) + if per_box_name is not None: + class_name = str(per_box_name) + else: + class_name = class_names_map.get(int(class_id), f"class_{int(class_id)}") + pairs.append(((x_min, y_min, x_max, y_max), class_name)) + return pairs + + +def measure_distance_with_reference_object( + detections: TensorNativeDetections, + object_1_class_name: str, + object_2_class_name: str, + reference_predictions: TensorNativeDetections, + reference_object_class_name: str, + reference_width: float, + reference_height: float, + reference_axis: Literal["horizontal", "vertical"], +): + reference_bbox_1 = None + reference_bbox_2 = None + + reference_bbox_1, reference_bbox_2 = find_reference_bboxes( + detections, object_1_class_name, object_2_class_name + ) + + if not reference_bbox_1 or not reference_bbox_2: + raise ValueError( + f"Reference class '{object_1_class_name}' or '{object_2_class_name}' not found in predictions." + ) + + if has_overlap(reference_bbox_1, reference_bbox_2) or not has_axis_gap( + reference_bbox_1, reference_bbox_2, reference_axis + ): + return {OUTPUT_KEY_CENTIMETER: 0, OUTPUT_KEY_PIXEL: 0} + + # get the reference object bounding box + reference_bbox = None + for (x_min, y_min, x_max, y_max), class_name in _iter_boxes_with_class_names( + reference_predictions + ): + if class_name == reference_object_class_name: + reference_bbox = (x_min, y_min, x_max, y_max) + break + + if not reference_bbox: + raise ValueError( + f"Reference class '{reference_object_class_name}' not found in predictions." + ) + + # calculate the pixel-to-centimeter ratio + reference_width_pixels = abs(reference_bbox[2] - reference_bbox[0]) + reference_height_pixels = abs(reference_bbox[3] - reference_bbox[1]) + + # Ensure the reference dimensions are positive and non-zero + if reference_width <= 0 or reference_height <= 0: + raise ValueError("Reference object dimensions must be greater than zero.") + + pixel_ratio_width = reference_width_pixels / reference_width + pixel_ratio_height = reference_height_pixels / reference_height + + # get the average pixel ratio + pixel_ratio = (pixel_ratio_width + pixel_ratio_height) / 2 + + distance_pixels = measure_distance_pixels( + reference_axis, reference_bbox_1, reference_bbox_2 + ) + + distance_cm = distance_pixels / pixel_ratio + + return {OUTPUT_KEY_CENTIMETER: distance_cm, OUTPUT_KEY_PIXEL: distance_pixels} + + +def measure_distance_with_pixel_ratio( + detections: TensorNativeDetections, + pixel_ratio: float, + object_1_class_name: str, + object_2_class_name: str, + reference_axis: Literal["horizontal", "vertical"], +) -> List[Dict[str, Union[str, float]]]: + reference_bbox_1 = None + reference_bbox_2 = None + + reference_bbox_1, reference_bbox_2 = find_reference_bboxes( + detections, object_1_class_name, object_2_class_name + ) + + if not reference_bbox_1 or not reference_bbox_2: + raise ValueError( + f"Reference class '{object_1_class_name}' or '{object_2_class_name}' not found in predictions." + ) + + if has_overlap(reference_bbox_1, reference_bbox_2) or not has_axis_gap( + reference_bbox_1, reference_bbox_2, reference_axis + ): + return {OUTPUT_KEY_CENTIMETER: 0, OUTPUT_KEY_PIXEL: 0} + + if pixel_ratio is None: + raise ValueError("Pixel-to-centimeter ratio must be provided.") + + if not isinstance(pixel_ratio, (int, float)): + raise ValueError("Pixel-to-centimeter ratio must be a number.") + + if pixel_ratio <= 0: + raise ValueError("Pixel-to-centimeter ratio must be greater than zero.") + + distance_pixels = measure_distance_pixels( + reference_axis, reference_bbox_1, reference_bbox_2 + ) + + distance_cm = distance_pixels / pixel_ratio + + return {OUTPUT_KEY_CENTIMETER: distance_cm, OUTPUT_KEY_PIXEL: distance_pixels} + + +def has_overlap( + bbox1: Tuple[int, int, int, int], bbox2: Tuple[int, int, int, int] +) -> bool: + """ + Check if two bounding boxes overlap. + + Args: + bbox1: A tuple of (x_min, y_min, x_max, y_max) for the first bounding box. + bbox2: A tuple of (x_min, y_min, x_max, y_max) for the second bounding box. + + Returns: + True if the bounding boxes overlap, False otherwise. + """ + x1_min, y1_min, x1_max, y1_max = bbox1 + x2_min, y2_min, x2_max, y2_max = bbox2 + + if x1_max < x2_min or x2_max < x1_min: + return False + if y1_max < y2_min or y2_max < y1_min: + return False + + return True + + +def has_axis_gap( + reference_bbox_1: Tuple[int, int, int, int], + reference_bbox_2: Tuple[int, int, int, int], + reference_axis: str, +) -> bool: + if reference_axis == "horizontal": + if ( + reference_bbox_1[0] < reference_bbox_2[2] + and reference_bbox_1[2] > reference_bbox_2[0] + ): + return False + else: + if ( + reference_bbox_1[1] < reference_bbox_2[3] + and reference_bbox_1[3] > reference_bbox_2[1] + ): + return False + + return True + + +def find_reference_bboxes( + detections: TensorNativeDetections, + object_1_class_name: str, + object_2_class_name: str, +): + reference_bbox_1 = None + reference_bbox_2 = None + + for (x_min, y_min, x_max, y_max), class_name in _iter_boxes_with_class_names( + detections + ): + if class_name == object_1_class_name: + reference_bbox_1 = (x_min, y_min, x_max, y_max) + elif class_name == object_2_class_name: + reference_bbox_2 = (x_min, y_min, x_max, y_max) + + if reference_bbox_1 and reference_bbox_2: + break + + return reference_bbox_1, reference_bbox_2 + + +def measure_distance_pixels( + reference_axis: str, + reference_bbox_1: Tuple[int, int, int, int], + reference_bbox_2: Tuple[int, int, int, int], +): + if reference_axis == "vertical": + distance_pixels = ( + abs(reference_bbox_2[1] - reference_bbox_1[3]) + if reference_bbox_2[1] > reference_bbox_1[3] + else abs(reference_bbox_1[1] - reference_bbox_2[3]) + ) + else: + distance_pixels = ( + abs(reference_bbox_2[0] - reference_bbox_1[2]) + if reference_bbox_2[0] > reference_bbox_1[2] + else abs(reference_bbox_1[0] - reference_bbox_2[2]) + ) + return distance_pixels diff --git a/inference/core/workflows/core_steps/classical_cv/dominant_color/v1.py b/inference/core/workflows/core_steps/classical_cv/dominant_color/v1.py index 9c901bfea7..deca094a17 100644 --- a/inference/core/workflows/core_steps/classical_cv/dominant_color/v1.py +++ b/inference/core/workflows/core_steps/classical_cv/dominant_color/v1.py @@ -1,4 +1,4 @@ -from typing import List, Literal, Optional, Type, Union +from typing import List, Literal, Optional, Tuple, Type, Union import numpy as np from pydantic import AliasChoices, ConfigDict, Field @@ -155,41 +155,53 @@ def run( scale_factor = max(1, min(width, height) // target_size) np_image = np_image[::scale_factor, ::scale_factor] - pixels = np_image.reshape(-1, 3).astype(np.float32) - - centroids = pixels[ - np.random.choice(pixels.shape[0], color_clusters, replace=False) - ] - - for _ in range(max_iterations): - # Assign pixels to nearest centroid - distances = np.sqrt(((pixels[:, np.newaxis] - centroids) ** 2).sum(axis=2)) - labels = np.argmin(distances, axis=1) - - # Update centroids - new_centroids = np.zeros_like(centroids) - for i in range(color_clusters): - cluster_points = pixels[labels == i] - if len(cluster_points) > 0: - new_centroids[i] = cluster_points.mean(axis=0) - else: - # If cluster is empty, reinitialize to a random point - new_centroids[i] = pixels[np.random.choice(pixels.shape[0])] - - # Check for convergence - if np.allclose(centroids, new_centroids): - break - - centroids = new_centroids - - # Get the colors and their counts - colors = centroids - uniq, counts = np.unique(labels, return_counts=True) - - # Find the most dominant color - dominant_color = colors[uniq[np.argmax(counts)]] - rgb_color = tuple( - int(np.clip(round(x), 0, 255)) for x in reversed(dominant_color) + rgb_color = find_dominant_color( + pixels_image=np_image, + color_clusters=color_clusters, + max_iterations=max_iterations, ) return {"rgb_color": rgb_color} + + +def find_dominant_color( + pixels_image: np.ndarray, color_clusters: int, max_iterations: int +) -> Tuple[int, int, int]: + """K-means dominant color of an already-downsampled HWC BGR uint8 image, + returned as an ``(r, g, b)`` tuple. + + Shared by v1 and the tensor sibling: the clustering trajectory depends on + the exact input bytes (float32 arithmetic) and on the unseeded global + ``np.random`` draw order.""" + pixels = pixels_image.reshape(-1, 3).astype(np.float32) + + centroids = pixels[np.random.choice(pixels.shape[0], color_clusters, replace=False)] + + for _ in range(max_iterations): + # Assign pixels to nearest centroid + distances = np.sqrt(((pixels[:, np.newaxis] - centroids) ** 2).sum(axis=2)) + labels = np.argmin(distances, axis=1) + + # Update centroids + new_centroids = np.zeros_like(centroids) + for i in range(color_clusters): + cluster_points = pixels[labels == i] + if len(cluster_points) > 0: + new_centroids[i] = cluster_points.mean(axis=0) + else: + # If cluster is empty, reinitialize to a random point + new_centroids[i] = pixels[np.random.choice(pixels.shape[0])] + + # Check for convergence + if np.allclose(centroids, new_centroids): + break + + centroids = new_centroids + + # Get the colors and their counts + colors = centroids + uniq, counts = np.unique(labels, return_counts=True) + + # Find the most dominant color + dominant_color = colors[uniq[np.argmax(counts)]] + return tuple(int(np.clip(round(x), 0, 255)) for x in reversed(dominant_color)) diff --git a/inference/core/workflows/core_steps/classical_cv/dominant_color/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/dominant_color/v1_tensor.py new file mode 100644 index 0000000000..b521864c40 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/dominant_color/v1_tensor.py @@ -0,0 +1,71 @@ +"""Tensor-native sibling of ``dominant_color/v1``. + +The k-means stays on the CPU in v1's ``find_dominant_color``: the clustered +array is tiny after downsampling, the convergence check would sync every +iteration, and the trajectory is coupled to unseeded global ``np.random`` +draws. The tensor path only moves the strided downsample to the device, so +the small downsampled block crosses device->host instead of the full frame. + +The host-side array is made byte-identical to v1's +``numpy_image[::scale_factor, ::scale_factor]`` (HWC, BGR) before the shared +k-means: float32 summation order in the distance computation depends on the +channel order, so RGB-ordered input could diverge the trajectory. Identical +bytes + identical RNG state => identical output; the RNG is unseeded, so +run-to-run output is nondeterministic on both paths. + +Numpy/base64-born images delegate to the v1 numpy path (no materialised +tensor). +""" + +from typing import Optional, Type + +import numpy as np +import torch + +from inference.core.workflows.core_steps.classical_cv.dominant_color.v1 import ( + DominantColorManifest, + find_dominant_color, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlock + + +class DominantColorBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[DominantColorManifest]: + return DominantColorManifest + + def run( + self, + image: WorkflowImageData, + color_clusters: Optional[int], + max_iterations: Optional[int], + target_size: Optional[int], + *args, + **kwargs + ) -> BlockResult: + if not image.is_tensor_materialised(): + np_image = image.numpy_image + height, width = np_image.shape[:2] + scale_factor = max(1, min(width, height) // target_size) + downsampled = np_image[::scale_factor, ::scale_factor] + else: + chw = image.tensor_image + height, width = int(chw.shape[-2]), int(chw.shape[-1]) + scale_factor = max(1, min(width, height) // target_size) + downsampled = _downsample_to_bgr_numpy(chw=chw, scale_factor=scale_factor) + rgb_color = find_dominant_color( + pixels_image=downsampled, + color_clusters=color_clusters, + max_iterations=max_iterations, + ) + return {"rgb_color": rgb_color} + + +def _downsample_to_bgr_numpy(chw: torch.Tensor, scale_factor: int) -> np.ndarray: + """Strided downsample on the device, then one small D2H copy yielding an + HWC BGR uint8 array byte-identical to v1's + ``numpy_image[::scale_factor, ::scale_factor]``.""" + downsampled = chw.detach()[:, ::scale_factor, ::scale_factor] + # CHW RGB -> HWC BGR; contiguous so only the downsampled block transfers. + return downsampled.flip(0).permute(1, 2, 0).contiguous().cpu().numpy() diff --git a/inference/core/workflows/core_steps/classical_cv/image_blur/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/image_blur/v1_tensor.py new file mode 100644 index 0000000000..3f8f61bd76 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/image_blur/v1_tensor.py @@ -0,0 +1,183 @@ +"""Tensor-native sibling of ``image_blur/v1``. + +All tensor paths are channel-independent, so RGB (tensor layout) vs BGR +(numpy layout) is irrelevant and grayscale ``(1, H, W)`` tensors run as a +one-channel batch. Bit parity with cv2 per blur type: + +- ``average`` (odd kernel <= 15): ``cv2.blur`` window sums over REFLECT_101 + borders are exact in a float32 convolution (255 * 15^2 < 2^24); k odd -> + k^2 odd -> ``sum / k^2`` has no .5 ties, so nearest-integer rounding is + unambiguous. Even kernels delegate: their exact .5 ties are rounded + differently across cv2 SIMD backends/builds. + +- ``gaussian`` (coerced-odd kernel in {1, 3, 5, 7}): cv2's uint8 GaussianBlur + with sigma=0 and kernel <= 7 uses the hardcoded ``small_gaussian_tab`` + Q8.8 kernels - exact dyadic rationals - so the separable fixed-point + pipeline (horizontal pass exact in uint16, vertical sums < 2^24, final + rounding ``(acc + 2^15) >> 16``) is replicated with float32 convolutions. + Kernels >= 9 delegate: cv2 derives their Q8.8 coefficients with softdouble + arithmetic that a float64 re-derivation diverges from. + +- ``median`` (coerced-odd kernel <= 15): ``cv2.medianBlur`` uses replicate + borders, and the median of an odd-count uint8 window is a pure integer + order statistic - replicate-pad + unfold + sort + middle element matches + bit-for-bit. Windows are materialised in row chunks to bound the k^2-fold + unfold blow-up to ~2^24 float32 elements (~64 MB) per chunk. Kernels > 15 + delegate (cv2 switches algorithms; unfold cost grows quadratically). + +- ``bilateral``: always delegates - data-dependent float weights built from + exp lookup tables cannot be replicated bit-exactly. + +Numpy/base64-born images delegate to v1 (no materialised tensor). Unknown +blur types delegate so the exact v1 ``ValueError`` is raised. Tensor paths +also require the border pad (k // 2) to be smaller than both spatial dims: +beyond that cv2 multi-reflects while ``F.pad(mode="reflect")`` refuses, so +that regime delegates. +""" + +from typing import Type + +import torch +import torch.nn.functional as F + +from inference.core.workflows.core_steps.classical_cv.image_blur.v1 import ( + ImageBlurManifest, + _to_positive_odd, + apply_blur, +) +from inference.core.workflows.core_steps.visualizations.common.base import ( + OUTPUT_IMAGE_KEY, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlock + +# OpenCV small_gaussian_tab (sigma=0, ksize<=7) in Q8.8 fixed point - all +# coefficients are exact dyadic rationals, each row sums to 256. +_SMALL_GAUSSIAN_KERNELS_Q8_8 = { + 1: (256,), + 3: (64, 128, 64), + 5: (16, 64, 96, 64, 16), + 7: (8, 28, 56, 72, 56, 28, 8), +} +_MAX_TENSOR_AVERAGE_KSIZE = 15 +_MAX_TENSOR_MEDIAN_KSIZE = 15 +# Max unfolded float32 window elements per median chunk (floor: one output row). +_MEDIAN_UNFOLD_ELEMENT_BUDGET = 2**24 + + +class ImageBlurBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[ImageBlurManifest]: + return ImageBlurManifest + + def run( + self, + image: WorkflowImageData, + blur_type: str, + kernel_size: int, + *args, + **kwargs, + ) -> BlockResult: + if not image.is_tensor_materialised() or _requires_numpy_delegation( + blur_type=blur_type, + kernel_size=kernel_size, + tensor_image=image.tensor_image, + ): + blurred_image = apply_blur(image.numpy_image, blur_type, kernel_size) + output = WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=blurred_image, + ) + return {OUTPUT_IMAGE_KEY: output} + blurred_tensor = _blur_tensor( + chw=image.tensor_image, blur_type=blur_type, kernel_size=kernel_size + ) + output = WorkflowImageData.copy_and_replace( + origin_image_data=image, + tensor_image=blurred_tensor, + ) + return {OUTPUT_IMAGE_KEY: output} + + +def _requires_numpy_delegation( + blur_type: str, kernel_size: int, tensor_image: torch.Tensor +) -> bool: + """True for every (type, ksize, shape) regime whose bit parity with cv2 + cannot be guaranteed on the tensor path - see the module docstring.""" + if blur_type == "average": + ksize = int(kernel_size) + if ksize < 1 or ksize % 2 == 0 or ksize > _MAX_TENSOR_AVERAGE_KSIZE: + return True + elif blur_type == "gaussian": + ksize = _to_positive_odd(kernel_size) + if ksize not in _SMALL_GAUSSIAN_KERNELS_Q8_8: + return True + elif blur_type == "median": + ksize = _to_positive_odd(kernel_size) + if ksize > _MAX_TENSOR_MEDIAN_KSIZE: + return True + else: + # bilateral and unknown blur types; the latter raise v1's exact + # ValueError inside apply_blur(). + return True + pad = ksize // 2 + return pad >= min(tensor_image.shape[-2], tensor_image.shape[-1]) + + +def _blur_tensor(chw: torch.Tensor, blur_type: str, kernel_size: int) -> torch.Tensor: + if blur_type == "average": + return _average_blur_tensor(chw=chw, ksize=int(kernel_size)) + if blur_type == "gaussian": + return _gaussian_blur_tensor(chw=chw, ksize=_to_positive_odd(kernel_size)) + return _median_blur_tensor(chw=chw, ksize=_to_positive_odd(kernel_size)) + + +def _average_blur_tensor(chw: torch.Tensor, ksize: int) -> torch.Tensor: + """``cv2.blur`` for ODD ksize: exact integer window sums over REFLECT_101 + borders, then unambiguous nearest-integer rounding of ``sum / k^2``.""" + pad = ksize // 2 + x = chw.detach().to(torch.float32).unsqueeze(1) # (C, 1, H, W) + x = F.pad(x, (pad, pad, pad, pad), mode="reflect") + weight = torch.ones((1, 1, ksize, ksize), dtype=torch.float32, device=x.device) + sums = F.conv2d(x, weight) # integer-valued float32, exact (< 2^24) + return torch.round(sums * (1.0 / (ksize * ksize))).to(torch.uint8).squeeze(1) + + +def _gaussian_blur_tensor(chw: torch.Tensor, ksize: int) -> torch.Tensor: + """OpenCV's bit-exact separable Q8.8 fixed-point uint8 GaussianBlur for + the dyadic small_gaussian_tab kernels: horizontal pass exact in uint16, + vertical pass < 2^24, final half-up rounding ``(acc + 2^15) >> 16``.""" + kernel_q8_8 = torch.tensor( + _SMALL_GAUSSIAN_KERNELS_Q8_8[ksize], dtype=torch.float32, device=chw.device + ) + pad = ksize // 2 + x = chw.detach().to(torch.float32).unsqueeze(1) # (C, 1, H, W) + x = F.pad(x, (pad, pad, pad, pad), mode="reflect") + x = F.conv2d(x, kernel_q8_8.view(1, 1, 1, ksize)) # Q8.8, <= 255 * 256 + x = F.conv2d(x, kernel_q8_8.view(1, 1, ksize, 1)) # Q16.16, < 2^24 + return ((x + 32768.0) / 65536.0).floor().to(torch.uint8).squeeze(1) + + +def _median_blur_tensor(chw: torch.Tensor, ksize: int) -> torch.Tensor: + """``cv2.medianBlur``: exact integer median (values 0..255 are exact in + float32 and the median of an odd count is an element of the window) over + replicate borders, unfolded in row chunks to bound the k^2 blow-up.""" + pad = ksize // 2 + channels, height, width = chw.shape + x = chw.detach().to(torch.float32).unsqueeze(0) + x = F.pad(x, (pad, pad, pad, pad), mode="replicate").squeeze(0) + middle = (ksize * ksize) // 2 + rows_per_chunk = max( + 1, _MEDIAN_UNFOLD_ELEMENT_BUDGET // (channels * width * ksize * ksize) + ) + chunks = [] + for row_start in range(0, height, rows_per_chunk): + row_end = min(height, row_start + rows_per_chunk) + windows = ( + x[:, row_start : row_end + 2 * pad, :] + .unfold(1, ksize, 1) + .unfold(2, ksize, 1) + .reshape(channels, row_end - row_start, width, ksize * ksize) + ) + chunks.append(windows.sort(dim=-1).values[..., middle]) + return torch.cat(chunks, dim=1).to(torch.uint8) diff --git a/inference/core/workflows/core_steps/classical_cv/image_preprocessing/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/image_preprocessing/v1_tensor.py new file mode 100644 index 0000000000..f762524e15 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/image_preprocessing/v1_tensor.py @@ -0,0 +1,218 @@ +"""Tensor-native sibling of ``image_preprocessing/v1``. + +Per-task contract: + +* ``flip``: tensor-native, bit-exact - a pure pixel permutation + (``cv2.flip`` <-> ``torch.flip``), independent of channel order. + +* ``rotate`` by right angles (+-90/+-180/+-270/+-360): tensor-native and + bit-exact, but not a naive ``rot90``: v1 rotates about ``(w // 2, h // 2)`` + (off the true pixel centre for odd dims) onto an ``int()``-truncated + canvas. In float64 the residual trig terms (cos 90deg ~ 6.1e-17) vanish in + ``warpAffine``'s 1/1024 fixed-point coordinates, so the inverse map + degenerates to axis-separable sampling at integer or half-integer + coordinates: + + - even source axes sample at integers - a permutation, except one border + index falls outside the source (BORDER_CONSTANT zeros); + - odd source axes sample halfway between pixels - a 2-tap average with + warpAffine's fixed-point rounding: ``(a + b + 1) >> 1`` for one half + axis, a single fused ``(a + b + c + d + 2) >> 2`` when both are half + (composing two rounded 1-D averages would round twice and diverge on + ties). + + +-360deg is a genuine warp in v1 (only ``0``/``None`` early-returns a + copy): identity for even dims, a half-pixel self-average for odd ones. + +* ``rotate`` by arbitrary angles: delegates (full fixed-point bilinear warp). + +* ``resize``: delegates for every real target - ``cv2.INTER_AREA`` rounding + is kernel-size- and SIMD-path-dependent, so no torch pooling stays + bit-exact across platforms. The ``width=height=None`` early-return copy is + tensor-native (a device-side ``clone``). + +Numpy/base64-born images always delegate (no materialised tensor). +""" + +from typing import Callable, Optional, Tuple, Type + +import torch + +from inference.core.workflows.core_steps.classical_cv.image_preprocessing.v1 import ( + ImagePreprocessingManifest, + apply_flip_image, + apply_resize_image, + apply_rotate_image, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlock + + +class ImagePreprocessingBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[ImagePreprocessingManifest]: + return ImagePreprocessingManifest + + def run( + self, + image: WorkflowImageData, + task_type: str, + width: Optional[int], + height: Optional[int], + rotation_degrees: Optional[int], + flip_type: Optional[str], + *args, + **kwargs, + ) -> BlockResult: + # Validation mirrors v1 and precedes any tensor-vs-delegate decision, + # so both siblings raise identically. + if task_type == "resize": + if width is not None and width <= 0: + raise ValueError("Width must be greater than 0") + if height is not None and height <= 0: + raise ValueError("Height must be greater than 0") + if image.is_tensor_materialised() and width is None and height is None: + # v1's early-return copy - replicated as a device-side clone. + return _emit_tensor_output(image=image, chw=image.tensor_image.clone()) + return _delegate_to_numpy(image, apply_resize_image, width, height) + elif task_type == "rotate": + if rotation_degrees is not None and not (-360 <= rotation_degrees <= 360): + raise ValueError("Rotation degrees must be between -360 and 360") + if image.is_tensor_materialised(): + if rotation_degrees is None or rotation_degrees == 0: + # v1's early-return copy - replicated as a device-side clone. + return _emit_tensor_output( + image=image, chw=image.tensor_image.clone() + ) + if rotation_degrees % 90 == 0: + rotated = _rotate_right_angle_tensor( + chw=image.tensor_image, rotation_degrees=rotation_degrees + ) + return _emit_tensor_output(image=image, chw=rotated) + return _delegate_to_numpy(image, apply_rotate_image, rotation_degrees) + elif task_type == "flip": + if flip_type is not None and flip_type not in [ + "vertical", + "horizontal", + "both", + ]: + raise ValueError( + "Flip type must be 'vertical', 'horizontal', or 'both'" + ) + if image.is_tensor_materialised(): + return _emit_tensor_output( + image=image, + chw=_flip_tensor(chw=image.tensor_image, flip_type=flip_type), + ) + return _delegate_to_numpy(image, apply_flip_image, flip_type) + else: + raise ValueError(f"Invalid task type: {task_type}") + + +def _delegate_to_numpy( + image: WorkflowImageData, + transformation: Callable, + *transformation_args, +) -> BlockResult: + response_image = transformation(image.numpy_image, *transformation_args) + output_image = WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=response_image + ) + return {"image": output_image} + + +def _emit_tensor_output(image: WorkflowImageData, chw: torch.Tensor) -> BlockResult: + output_image = WorkflowImageData.copy_and_replace( + origin_image_data=image, tensor_image=chw + ) + return {"image": output_image} + + +def _flip_tensor(chw: torch.Tensor, flip_type: Optional[str]) -> torch.Tensor: + if flip_type == "vertical": + return torch.flip(chw, dims=(-2,)) + if flip_type == "horizontal": + return torch.flip(chw, dims=(-1,)) + if flip_type == "both": + return torch.flip(chw, dims=(-2, -1)) + # v1's apply_flip_image else-branch: pass-through (only flip_type=None + # reaches it via run()). + return chw + + +# Per-axis sampling modes of the right-angle warp (see the module docstring): +# "exact" copies the axis, "shift" moves it by one index with a zero fill +# (the out-of-range border sample of even axes), "half" sums each element +# with its predecessor (the /2 with fixed-point rounding is applied once, +# fused, at the end). +_AXIS_EXACT = "exact" +_AXIS_SHIFT = "shift" +_AXIS_HALF = "half" + + +def _rotate_right_angle_tensor( + chw: torch.Tensor, rotation_degrees: int +) -> torch.Tensor: + """Bit-exact replica of v1's ``apply_rotate_image`` for multiples of 90deg. + + Derived dest->source maps (X = dest column, Y = dest row; ``cx = w // 2``, + ``cy = h // 2``; half offsets appear exactly when the axis length is odd): + + * q=1 (90deg ccw): ``y_src = X - 0.5*(h odd)``, ``x_src = w - Y - 0.5*(w odd)`` + * q=2 (180deg): ``x_src = w - X - 0.5*(w odd)``, ``y_src = h - Y - 0.5*(h odd)`` + * q=3 (270deg): ``y_src = h - X - 0.5*(h odd)``, ``x_src = Y - 0.5*(w odd)`` + * q=0 (+-360deg): ``x_src = X - 0.5*(w odd)``, ``y_src = Y - 0.5*(h odd)`` + + Each axis is therefore (optionally) reversed, then sampled exactly, + shifted by one (out-of-range -> 0), or 2-tap averaged; 90/270 transpose + H<->W at the end. + """ + height, width = int(chw.shape[-2]), int(chw.shape[-1]) + quarter_turns = int(rotation_degrees // 90) % 4 + y_half = _AXIS_HALF if height % 2 else None + x_half = _AXIS_HALF if width % 2 else None + if quarter_turns == 0: + y_spec = (False, y_half or _AXIS_EXACT) + x_spec = (False, x_half or _AXIS_EXACT) + elif quarter_turns == 1: + y_spec = (False, y_half or _AXIS_EXACT) + x_spec = (True, x_half or _AXIS_SHIFT) + elif quarter_turns == 2: + y_spec = (True, y_half or _AXIS_SHIFT) + x_spec = (True, x_half or _AXIS_SHIFT) + else: + y_spec = (True, y_half or _AXIS_SHIFT) + x_spec = (False, x_half or _AXIS_EXACT) + accumulator = chw.to(torch.int32) + accumulator, y_halved = _sample_axis(accumulator, dim=-2, spec=y_spec) + accumulator, x_halved = _sample_axis(accumulator, dim=-1, spec=x_spec) + halvings = y_halved + x_halved + if halvings: + # warpAffine's fused fixed-point rounding: (a+b+1)>>1 / (a+b+c+d+2)>>2. + accumulator = (accumulator + (1 << (halvings - 1))) // (1 << halvings) + if quarter_turns in (1, 3): + accumulator = accumulator.transpose(-2, -1) + return accumulator.to(torch.uint8).contiguous() + + +def _sample_axis( + accumulator: torch.Tensor, dim: int, spec: Tuple[bool, str] +) -> Tuple[torch.Tensor, int]: + reverse, mode = spec + if reverse: + accumulator = torch.flip(accumulator, dims=(dim,)) + if mode == _AXIS_EXACT: + return accumulator, 0 + shifted = _shifted_by_one(accumulator, dim=dim) + if mode == _AXIS_SHIFT: + return shifted, 0 + return accumulator + shifted, 1 + + +def _shifted_by_one(accumulator: torch.Tensor, dim: int) -> torch.Tensor: + shifted = torch.zeros_like(accumulator) + if dim == -1: + shifted[..., 1:] = accumulator[..., :-1] + else: + shifted[..., 1:, :] = accumulator[..., :-1, :] + return shifted diff --git a/inference/core/workflows/core_steps/classical_cv/mask_area_measurement/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/mask_area_measurement/v1_tensor.py new file mode 100644 index 0000000000..512ddb80f4 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/mask_area_measurement/v1_tensor.py @@ -0,0 +1,206 @@ +from typing import List, Literal, Optional, Type, Union + +import cv2 as cv +import numpy as np +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + instance_mask_to_numpy, +) +from inference.core.workflows.execution_engine.constants import ( + AREA_CONVERTED_KEY_IN_SV_DETECTIONS, + AREA_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY = "predictions" +SHORT_DESCRIPTION = ( + "Measure the area of detected objects and optionally convert to real-world units." +) +LONG_DESCRIPTION = """ +Measure the area of detected objects. For instance segmentation masks, the area is computed by counting non-zero mask pixels (correctly handling holes). For bounding-box-only detections, the area is width multiplied by height. Optionally converts pixel areas to real-world units using a `pixels_per_unit` calibration value. + +## How This Block Works + +This block calculates the area of each detected object and stores two values per detection: + +- **`area_px`** โ€” area in square pixels (always computed) +- **`area_converted`** โ€” area in real-world units: `area_px / (pixels_per_unit ** 2)` (equals `area_px` when `pixels_per_unit` is 1.0) + +Both values are attached to each detection and included in the serialized JSON output. The block returns the input detections with these fields added, so downstream blocks (e.g., label visualization) can display the area values. + +### Area Computation + +The block operates in two modes depending on the type of predictions it receives: + +1. **Mask Pixel Area (Instance Segmentation)**: When the input detections include segmentation masks, the block counts the non-zero pixels in each mask using `cv2.countNonZero`. This correctly handles masks with holes โ€” hole pixels are zero and are excluded from the count. + +2. **Bounding Box Area (Object Detection)**: When no segmentation mask is available, the block falls back to computing the area as the bounding box width multiplied by height (`w * h`). + +### Unit Conversion + +Set the `pixels_per_unit` input to convert pixel areas to real-world units (e.g., cmยฒ, inยฒ, mmยฒ). Because area is two-dimensional, the conversion squares the ratio: + +``` +area_converted = area_px / (pixels_per_unit ** 2) +``` + +For example, if your calibration is 130 pixels/cm, a detection with `area_px = 16900` would have `area_converted = 16900 / 16900 = 1.0 cmยฒ`. + +**How to determine pixels_per_unit:** Place an object of known size in the camera's field of view (e.g., a ruler or calibration target). Measure its length in pixels in the image and divide by its real-world length. For instance, if a 10 cm reference object spans 1300 pixels, then `pixels_per_cm = 1300 / 10 = 130`. If you are using perspective correction, the calibration object must be placed on the same plane from which the perspective correction was calculated. + +## Common Use Cases + +- **Size-Based Filtering**: Filter out small noise detections by chaining with a filtering block to keep only detections above a minimum area threshold. +- **Quality Control**: Verify that manufactured components meet size specifications by comparing measured areas against expected ranges. +- **Agricultural Analysis**: Measure leaf area, crop coverage, or canopy extent from aerial or close-up imagery. +- **Medical Imaging**: Quantify the area of wounds, lesions, or anatomical structures. Use `pixels_per_unit` to get real-world measurements for clinical documentation. + +## Connecting to Other Blocks + +- **Upstream -- Detection and Segmentation Models**: Connect the output of an object detection or instance segmentation model to the `predictions` input. Instance segmentation models (which produce masks) yield more accurate area measurements than bounding-box-only detections. +- **Upstream -- Camera Calibration Block**: Use `roboflow_core/camera_calibration@v1` upstream to correct lens distortion before detection. +- **Upstream -- Perspective Correction Block**: Use `roboflow_core/perspective_correction@v1` upstream to transform angled images to a top-down view so that area measurements reflect true object footprints. +- **Downstream -- Visualization**: Pass the output `predictions` to label or polygon visualization blocks. The `area_px` and `area_converted` fields are available for display as labels. +- **Downstream -- Filtering Blocks**: Use the enriched detections with a filtering block to keep only detections whose area meets a threshold. + +## Requirements + +This block requires detection predictions from an object detection or instance segmentation model. No additional environment variables, API keys, or external dependencies are needed beyond OpenCV and NumPy (included with inference). For the most accurate area measurements, use instance segmentation models that produce per-object masks. +""" + + +class MaskAreaMeasurementManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Mask Area Measurement", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "classical_cv", + "icon": "far fa-ruler-combined", + "blockPriority": 12, + "opencv": True, + }, + } + ) + type: Literal["roboflow_core/mask_area_measurement@v1"] + + predictions: Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + description="Detection predictions to measure areas for.", + examples=["$steps.model.predictions"], + ) + + pixels_per_unit: Union[float, Selector(kind=[FLOAT_KIND])] = Field( # type: ignore + default=1.0, + description="Number of pixels per real-world unit of length (e.g., pixels per cm). " + "The converted area is computed as area_px / (pixels_per_unit ** 2). " + "Default 1.0 means no conversion (area_converted equals area_px).", + examples=[1.0, 130.0], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +def compute_detection_areas( + detections: Union[Detections, InstanceDetections], +) -> List[float]: + """Compute the area of all detections in square pixels. + + For bounding-box-only detections, areas are computed in a single vectorized + operation. For detections with segmentation masks, the area is the count of + non-zero mask pixels (via ``cv2.countNonZero``). This correctly handles masks + with holes โ€” hole pixels are zero and are not counted. Falls back to the + bounding box area when the mask pixel count is zero. + + Args: + detections: A supervision Detections object. + + Returns: + List of areas in square pixels, one per detection. + """ + n = len(detections) + if n == 0: + return [] + + areas = [] + for i in range(n): + if isinstance(detections, InstanceDetections): + count = cv.countNonZero( + instance_mask_to_numpy(detections, i).astype(np.uint8) + ) + if count > 0: + areas.append(float(count)) + continue + x1, y1, x2, y2 = detections.xyxy[i] + areas.append(float((x2 - x1) * (y2 - y1))) + + return areas + + +class MaskAreaMeasurementBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return MaskAreaMeasurementManifest + + def run( + self, + predictions: Union[Detections, InstanceDetections], + pixels_per_unit: float = 1.0, + ) -> BlockResult: + areas_px = np.array(compute_detection_areas(predictions)) + scale = pixels_per_unit**2 if pixels_per_unit > 0 else 1.0 + areas_converted = areas_px / scale + num_detections = len(predictions) + bboxes_metadata = predictions.bboxes_metadata + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(num_detections)] + else: + bboxes_metadata = [ + dict(box_metadata) if box_metadata is not None else {} + for box_metadata in bboxes_metadata + ] + for i in range(num_detections): + bboxes_metadata[i][AREA_KEY_IN_SV_DETECTIONS] = float(areas_px[i]) + bboxes_metadata[i][AREA_CONVERTED_KEY_IN_SV_DETECTIONS] = float( + areas_converted[i] + ) + predictions.bboxes_metadata = bboxes_metadata + return {OUTPUT_KEY: predictions} diff --git a/inference/core/workflows/core_steps/classical_cv/mask_edge_snap/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/mask_edge_snap/v1_tensor.py new file mode 100644 index 0000000000..4814a1d26e --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/mask_edge_snap/v1_tensor.py @@ -0,0 +1,625 @@ +from typing import List, Literal, Optional, Type, Union +from uuid import uuid4 + +import cv2 +import numpy as np +import torch +from pydantic import AliasChoices, ConfigDict, Field, field_validator + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, + instance_mask_to_numpy, +) +from inference.core.workflows.execution_engine.constants import DETECTION_ID_KEY +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + ANY_DATA_AS_SELECTED_ELEMENT, + FLOAT_KIND, + IMAGE_KIND, + INTEGER_KIND, + KIND_KEY, + REFERENCE_KEY, + SELECTED_ELEMENT_KEY, + SELECTOR_POINTS_TO_BATCH_KEY, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections + +SHORT_DESCRIPTION: str = ( + "Refine instance segmentation masks by snapping edges to detected boundaries." +) +LONG_DESCRIPTION = """ +Refine instance segmentation masks by snapping contour points to Sobel edges within a band around the predicted boundary. This block improves segmentation accuracy by adjusting mask edges to align with detected image features. + +## How This Block Works + +This block refines segmentation masks through a sophisticated multi-step pipeline: + +1. **Edge Detection**: Computes Sobel gradient magnitudes from the input image to detect edges +2. **Adaptive Thresholding**: Uses per-pixel adaptive thresholding (local mean + sigma * local std) to identify significant edges +3. **Morphological Processing**: Applies closing (dilation + erosion) to bridge small gaps in edge segments +4. **Thinning**: Applies Zhang-Suen single-iteration thinning to reduce edge width to 1-2 pixels while preserving connectivity +5. **Boundary Band Creation**: Builds a search band around each predicted mask's contour +6. **Area Filtering**: Removes small edge components below a minimum area threshold +7. **Contour Snapping**: For each original mask contour point, finds the strongest nearby edge within tolerance and snaps to it + +## Common Use Cases + +- **Medical Image Analysis**: Refine organ/tumor segmentation masks to align with anatomical boundaries +- **Industrial Quality Control**: Improve part boundary detection for precise dimension measurement +- **Autonomous Vehicles**: Refine road/lane segmentation boundaries for improved path planning +- **Agricultural Monitoring**: Enhance crop boundary detection for yield estimation +- **Microscopy Analysis**: Refine cell/nuclei segmentation for morphological analysis +- **Document Processing**: Improve text region boundary detection for OCR + +## Input Parameters + +**image** : Input image (color or grayscale) +- Can be single-channel, 3-channel (BGR), or 4-channel (BGRA) +- Preprocessing (blur, contrast enhancement) should be applied upstream if needed + +**segmentation** : Initial instance segmentation predictions +- Source: from object detection or instance segmentation model +- Must contain populated `mask` field; if empty, passed through unchanged + +**pixel_tolerance** : Maximum perpendicular distance (pixels) for edge snapping +- Range: 5-50 typically +- 5-15: tight predictions with minimal offset +- 20-50: rough predictions needing more forgiveness + +**sigma** : Strictness multiplier for adaptive Sobel threshold +- Range: 0.1-2.0 typically +- 0.1-0.5: permissive, keeps weaker edges, good for low-contrast boundaries +- 1.0-2.0: strict, only strongest edges survive, good for high-contrast images + +**min_contour_area** : Minimum enclosed-polygon area for edge components +- Range: 10-1000 typically +- Small (10-50): keeps fragmented edges +- Large (200-1000): aggressive noise rejection + +**dilation_iterations** : Number of morphological closing iterations +- Range: 0-10 typically +- 0: no closing, only thresholded edges +- 1-2: bridges hairline gaps +- 3-5: bridges visible dashes +- 10+: aggressive, can merge unrelated edges + +**boundary_band_width** : Half-width of search band around mask contour (default: 15) +- Sets maximum distance between predicted and true boundary that can be corrected + +**adaptive_window_size** : Side length of local-statistics window (default: 41) +- Should be roughly 5-10% of smaller image dimension +- Smaller (15-25): fine local contrast sensitivity, can pick up noise +- Larger (81-121): smooth threshold field, closer to global thresholding + +## Outputs + +**refined_segmentation** : Same detections with snapped mask contours +**edges** : Single detection containing union of all surviving edge pixels (debug/visualization) + +## Preprocessing + +**Preprocessing is usually critical for success.** This block does no preprocessing โ€” what you feed in is what Sobel sees. For challenging imagery, chain Roboflow image-processing blocks upstream: + +**Gaussian Blur** + For grainy or noisy surfaces (welds, machined metal, biological tissue), blur before edge detection to suppress per-pixel noise. A 5x5 kernel with sigma 1.0 is a sensible default; increase to 7x7 or 9x9 for very noisy imagery. Don't over-blur โ€” strong blur rounds off corners and softens real boundaries, leading to boundary positions that are biased inward. + +**Bilateral Blur** + Better than Gaussian when the image has both noise AND important sharp edges (e.g. textured fabric on a clean background). Slower, but preserves edges while denoising flat regions. + +**Contrast Enhancement** + Use when boundary contrast is genuinely too low to threshold reliably. The Contrast Enhancement block normalizes the histogram to use the full range, improving edge detection sensitivity without the noise amplification of aggressive methods. Follow with blur to suppress any remaining noise. Avoid on already-high-contrast images. + +**Morphological Opening then Closing** + Opening (erode then dilate) removes small bright specks and thin protrusions from the input before edge detection โ€” useful when the surface has fine debris or hot pixels that would otherwise generate spurious edges. Closing (dilate then erode) fills small dark holes/gaps in bright regions; less commonly needed as preprocessing, since gap filling on the edge map itself is what the `dilation_iterations` parameter already does. Use the Morphological Transformation v2 block with the "Opening then Closing" operation for this preprocessing. + +**Order matters**: Blur first, then contrast adjustment if needed. Reverse causes contrast adjustment to amplify the noise before blur can suppress it. +""" + +# Custom type for segmentation that accepts both selector strings and Detections objects +# This includes the proper json_schema_extra to tell the workflow engine to resolve selectors +_segmentation_json_schema_extra = { + REFERENCE_KEY: True, + SELECTED_ELEMENT_KEY: ANY_DATA_AS_SELECTED_ELEMENT, + KIND_KEY: [TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND.dict()], + SELECTOR_POINTS_TO_BATCH_KEY: "dynamic", +} + + +class MaskEdgeSnapManifest(WorkflowBlockManifest): + type: Literal["roboflow_core/mask_edge_snap@v1"] + model_config = ConfigDict( + arbitrary_types_allowed=True, + json_schema_extra={ + "name": "Mask Edge Snap", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "classical_computer_vision", + "ui_manifest": { + "section": "classical_cv", + "icon": "far fa-scissors", + "blockPriority": 12, + "opencv": True, + }, + }, + ) + + image: Selector(kind=[IMAGE_KIND]) = Field( + title="Input Image", + description="Input image (color or grayscale) for edge detection and snapping. Can be grayscale, single-channel, BGR, or BGRA. No preprocessing is applied internally; use upstream blocks for blur or contrast enhancement if needed.", + examples=["$inputs.image", "$steps.preprocessing.image"], + validation_alias=AliasChoices("image", "images"), + ) + + segmentation: Union[str, InstanceDetections] = Field( + title="Segmentation", + description="Instance segmentation predictions with mask field populated. Each mask contour will be snapped to detected edges. If empty, segmentation is passed through unchanged. Can be a reference string like '$steps.segmentation_model.predictions' or a supervision.Detections object.", + examples=["$steps.segmentation_model.predictions", "$inputs.segmentation"], + json_schema_extra=_segmentation_json_schema_extra, + ) + + @field_validator("segmentation") + @classmethod + def validate_segmentation( + cls, value: Union[str, InstanceDetections] + ) -> Union[str, InstanceDetections]: + if isinstance(value, str): + if not value.startswith("$"): + raise ValueError( + f"segmentation must be a workflow reference starting with '$' or a supervision.Detections object, got: {value}" + ) + return value + + pixel_tolerance: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=15, + description="Maximum perpendicular distance (pixels) from each contour point to candidate edges during snapping. Typical: 5-15 for tight predictions, 20-50 for rough ones. Too small: real edges outside range get missed. Too large: snap can wander to unrelated edges.", + ) + + sigma: Union[float, Selector(kind=[FLOAT_KIND])] = Field( + default=1.0, + description="Strictness multiplier for adaptive Sobel threshold (local_mean + sigma * local_std). Lower (0.1-0.5): permissive, good for low-contrast. Higher (1.0-2.0): strict, only strongest edges. Tune this AFTER other parameters.", + ) + + min_contour_area: Union[float, Selector(kind=[FLOAT_KIND])] = Field( + default=50.0, + description="Minimum enclosed-polygon area for edge components to keep. Small (10-50): keeps fragmented edges. Large (200-1000): aggressive noise rejection. Scales roughly with dilation_iterations.", + ) + + dilation_iterations: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=2, + description="Morphological closing iterations to bridge gaps in thresholded edge map. Each iteration bridges ~2px gaps. 0: no closing. 1-2: hairline gaps. 3-5: visible dashes. 10+: aggressive merging.", + ) + + boundary_band_width: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=15, + description="Half-width (pixels) of search band around segmentation contour. Sets maximum distance between predicted boundary and true boundary that can be corrected. Should generally be >= pixel_tolerance.", + ) + + adaptive_window_size: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=41, + description="Side length of local-statistics window for adaptive threshold. Small (15-25): fine local sensitivity, can pick noise. Default 41: balanced. Large (81-121): smooth field, closer to global thresholding. Should be ~5-10% of smaller image dimension.", + ) + + @field_validator("pixel_tolerance") + @classmethod + def validate_pixel_tolerance(cls, value: Union[int, str]) -> Union[int, str]: + if isinstance(value, int) and (value < 1 or value > 100): + raise ValueError( + "pixel_tolerance must be between 1 and 100, got: {}".format(value) + ) + return value + + @field_validator("sigma") + @classmethod + def validate_sigma(cls, value: Union[float, str]) -> Union[float, str]: + if isinstance(value, float) and (value < 0.01 or value > 10.0): + raise ValueError( + "sigma must be between 0.01 and 10.0, got: {}".format(value) + ) + return value + + @field_validator("min_contour_area") + @classmethod + def validate_min_contour_area(cls, value: Union[float, str]) -> Union[float, str]: + if isinstance(value, float) and (value < 0.0 or value > 10000.0): + raise ValueError( + "min_contour_area must be between 0.0 and 10000.0, got: {}".format( + value + ) + ) + return value + + @field_validator("dilation_iterations") + @classmethod + def validate_dilation_iterations(cls, value: Union[int, str]) -> Union[int, str]: + if isinstance(value, int) and (value < 0 or value > 20): + raise ValueError( + "dilation_iterations must be between 0 and 20, got: {}".format(value) + ) + return value + + @field_validator("boundary_band_width") + @classmethod + def validate_boundary_band_width(cls, value: Union[int, str]) -> Union[int, str]: + if isinstance(value, int) and (value < 1 or value > 100): + raise ValueError( + "boundary_band_width must be between 1 and 100, got: {}".format(value) + ) + return value + + @field_validator("adaptive_window_size") + @classmethod + def validate_adaptive_window_size(cls, value: Union[int, str]) -> Union[int, str]: + if isinstance(value, int) and (value < 3 or value > 201): + raise ValueError( + "adaptive_window_size must be between 3 and 201, got: {}".format(value) + ) + return value + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="refined_segmentation", + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + OutputDefinition( + name="edges", + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class MaskEdgeSnapBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[MaskEdgeSnapManifest]: + return MaskEdgeSnapManifest + + def run( + self, + image: WorkflowImageData, + segmentation: InstanceDetections, + pixel_tolerance: int, + sigma: float, + min_contour_area: float, + dilation_iterations: int, + boundary_band_width: int, + adaptive_window_size: int, + *args, + **kwargs, + ) -> BlockResult: + refined_segmentation, edges = refine_masks( + image=image, + segmentation=segmentation, + pixel_tolerance=pixel_tolerance, + sigma=sigma, + min_contour_area=min_contour_area, + dilation_iterations=dilation_iterations, + boundary_band_width=boundary_band_width, + adaptive_window_size=adaptive_window_size, + ) + + return { + "refined_segmentation": refined_segmentation, + "edges": edges, + } + + +def _zhang_suen_one_iteration(image: np.ndarray) -> np.ndarray: + """Run a single iteration of Zhang-Suen thinning. + + Removes only "simple points" โ€” pixels whose deletion preserves local + connectivity. Repeatedly applying this would converge to a 1-pixel + skeleton; calling it once peels exactly one layer off thick regions + while guaranteeing no narrow section, junction, or endpoint is broken. + + Neighborhood layout: + P9 P2 P3 + P8 P1 P4 + P7 P6 P5 + """ + img = (image > 0).astype(np.uint8) + h, w = img.shape + if h < 3 or w < 3: + return (img * 255).astype(np.uint8) + + def neighbors(arr): + P2 = arr[:-2, 1:-1] + P3 = arr[:-2, 2:] + P4 = arr[1:-1, 2:] + P5 = arr[2:, 2:] + P6 = arr[2:, 1:-1] + P7 = arr[2:, :-2] + P8 = arr[1:-1, :-2] + P9 = arr[:-2, :-2] + return P2, P3, P4, P5, P6, P7, P8, P9 + + def conditions(P1, P2, P3, P4, P5, P6, P7, P8, P9): + # B(p): count of foreground neighbors (must be 2..6) + B = (P2 + P3 + P4 + P5 + P6 + P7 + P8 + P9).astype(np.int16) + # A(p): number of 0->1 transitions in cyclic sequence + A = ( + ((P2 == 0) & (P3 == 1)).astype(np.int16) + + ((P3 == 0) & (P4 == 1)).astype(np.int16) + + ((P4 == 0) & (P5 == 1)).astype(np.int16) + + ((P5 == 0) & (P6 == 1)).astype(np.int16) + + ((P6 == 0) & (P7 == 1)).astype(np.int16) + + ((P7 == 0) & (P8 == 1)).astype(np.int16) + + ((P8 == 0) & (P9 == 1)).astype(np.int16) + + ((P9 == 0) & (P2 == 1)).astype(np.int16) + ) + return (P1 == 1) & (B >= 2) & (B <= 6) & (A == 1) + + # Sub-iteration 1: target south-east border points + P1 = img[1:-1, 1:-1] + P2, P3, P4, P5, P6, P7, P8, P9 = neighbors(img) + base = conditions(P1, P2, P3, P4, P5, P6, P7, P8, P9) + cond1 = base & ((P2 * P4 * P6) == 0) & ((P4 * P6 * P8) == 0) + img[1:-1, 1:-1] = np.where(cond1, 0, P1) + + # Sub-iteration 2: target north-west border points + P1 = img[1:-1, 1:-1] + P2, P3, P4, P5, P6, P7, P8, P9 = neighbors(img) + base = conditions(P1, P2, P3, P4, P5, P6, P7, P8, P9) + cond2 = base & ((P2 * P4 * P8) == 0) & ((P2 * P6 * P8) == 0) + img[1:-1, 1:-1] = np.where(cond2, 0, P1) + + return (img * 255).astype(np.uint8) + + +def refine_masks( + image: WorkflowImageData, + segmentation: InstanceDetections, + pixel_tolerance: int, + sigma: float, + min_contour_area: float, + dilation_iterations: int, + boundary_band_width: int, + adaptive_window_size: int, +) -> tuple: + """Refine instance segmentation masks by snapping edges to detected boundaries.""" + + np_img = image.numpy_image.copy() + H, W = np_img.shape[:2] + + # Convert to grayscale + if len(np_img.shape) == 2: + gray = np_img.copy() + elif np_img.shape[2] == 1: + gray = np_img[:, :, 0] + elif np_img.shape[2] == 4: + gray = cv2.cvtColor(np_img, cv2.COLOR_BGRA2GRAY) + elif np_img.shape[2] == 3: + gray = cv2.cvtColor(np_img, cv2.COLOR_BGR2GRAY) + else: + # For any other case, try to convert assuming BGR + try: + gray = cv2.cvtColor(np_img, cv2.COLOR_BGR2GRAY) + except cv2.error: + # If conversion fails, take first channel + gray = np_img[:, :, 0] if len(np_img.shape) >= 3 else np_img.copy() + + tol = int(pixel_tolerance) + band_radius = max(1, int(boundary_band_width)) + band_kernel = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, (band_radius * 2 + 1, band_radius * 2 + 1) + ) + + # Read the input instance masks one-at-a-time as numpy so cv2/numpy edge math + # operates on plain arrays (never materialising the native input to sv.Detections). + num_detections = len(segmentation) + has_masks = segmentation.mask is not None and num_detections > 0 + segmentation_masks = ( + [instance_mask_to_numpy(segmentation, i) for i in range(num_detections)] + if has_masks + else [] + ) + + # Create boundary band from segmentation masks + if has_masks: + boundary_band_pre = np.zeros((H, W), dtype=np.uint8) + for m in segmentation_masks: + m_uint8 = m.astype(np.uint8) + if m_uint8.shape[:2] != (H, W): + m_uint8 = cv2.resize(m_uint8, (W, H), interpolation=cv2.INTER_NEAREST) + inner_i = cv2.erode(m_uint8, band_kernel) + outer_i = cv2.dilate(m_uint8, band_kernel) + boundary_band_pre = np.maximum( + boundary_band_pre, ((outer_i > 0) & (inner_i == 0)).astype(np.uint8) + ) + else: + boundary_band_pre = np.ones((H, W), dtype=np.uint8) + + # Compute Sobel edges + gx = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3) + gy = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3) + magnitude = cv2.magnitude(gx, gy) + + # Adaptive threshold on Sobel magnitude + win = max(3, int(adaptive_window_size)) + if win % 2 == 0: + win += 1 + mag_mean = cv2.boxFilter(magnitude, cv2.CV_32F, (win, win)) + mag_sq_mean = cv2.boxFilter(magnitude * magnitude, cv2.CV_32F, (win, win)) + mag_var = np.maximum(mag_sq_mean - mag_mean * mag_mean, 0.0) + mag_std = np.sqrt(mag_var) + threshold_field = mag_mean + float(sigma) * mag_std + edges_adaptive = (magnitude > threshold_field).astype(np.uint8) * 255 + + # Morphological closing and thinning + iterations = max(0, int(dilation_iterations)) + if iterations > 0: + close_kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)) + edges_adaptive = cv2.morphologyEx( + edges_adaptive, cv2.MORPH_CLOSE, close_kernel, iterations=iterations + ) + edges_adaptive = _zhang_suen_one_iteration(edges_adaptive) + + edges = edges_adaptive + + # Apply boundary band filter + if has_masks: + edges_to_filter = (edges * boundary_band_pre).astype(np.uint8) + else: + edges_to_filter = edges.copy() + + # Filter by contour area + min_area = max(0.0, float(min_contour_area)) + num_labels, labels, _, _ = cv2.connectedComponentsWithStats( + edges_to_filter, connectivity=8 + ) + edge_filtered = np.zeros((H, W), dtype=np.uint8) + for lbl in range(1, num_labels): + comp_mask = (labels == lbl).astype(np.uint8) + comp_contours, _ = cv2.findContours( + comp_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE + ) + if not comp_contours: + continue + main_contour = max(comp_contours, key=cv2.contourArea) + if cv2.contourArea(main_contour) >= min_area: + edge_filtered[labels == lbl] = 255 + + snap_region = edge_filtered > 0 + + # Build edges detection output (native โ€” single synthetic "edges" instance) + edges_image_metadata = build_native_image_metadata( + image=image, + class_names={0: "edges"}, + prediction_type="object-detection", + ) + if snap_region.any(): + ys, xs = np.where(snap_region) + edges_detections = InstanceDetections( + xyxy=torch.tensor( + [[xs.min(), ys.min(), xs.max() + 1, ys.max() + 1]], + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + class_id=torch.zeros( + (1,), dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.ones( + (1,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + mask=torch.from_numpy(snap_region[None, :, :]).to( + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, dtype=torch.bool + ), + image_metadata=edges_image_metadata, + bboxes_metadata=[{DETECTION_ID_KEY: str(uuid4())}], + ) + else: + edges_detections = InstanceDetections( + xyxy=torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + class_id=torch.zeros( + (0,), dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.zeros( + (0,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + mask=torch.zeros( + (0, H, W), dtype=torch.bool, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + image_metadata=edges_image_metadata, + bboxes_metadata=None, + ) + + # If no segmentation, return early + if not has_masks: + return segmentation, edges_detections + + # Snap contours to edges + refined_masks = [] + TANGENT_WINDOW = 5 + + for mask in segmentation_masks: + mask_uint8 = mask.astype(np.uint8) + if mask_uint8.shape[:2] != (H, W): + mask_uint8 = cv2.resize(mask_uint8, (W, H), interpolation=cv2.INTER_NEAREST) + + dist = cv2.distanceTransform(mask_uint8, cv2.DIST_L2, cv2.DIST_MASK_PRECISE) + mask_dilated = cv2.dilate(mask_uint8, np.ones((3, 3), np.uint8)) + valid_snap = snap_region & ((dist <= tol) | (mask_dilated == 0)) + + mask_contours, _ = cv2.findContours( + mask_uint8, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE + ) + if not mask_contours: + refined_masks.append(mask_uint8.astype(bool)) + continue + + new_mask = np.zeros((H, W), dtype=np.uint8) + + for mc in mask_contours: + mc_pts = mc.reshape(-1, 2).astype(np.float32) + n_pts = mc_pts.shape[0] + refined_pts = mc_pts.copy() + + for i in range(n_pts): + px, py = mc_pts[i] + + prev_pt = mc_pts[(i - TANGENT_WINDOW) % n_pts] + next_pt = mc_pts[(i + TANGENT_WINDOW) % n_pts] + tangent = next_pt - prev_pt + t_len = np.linalg.norm(tangent) + if t_len < 1e-6: + continue + tangent /= t_len + nx, ny = -tangent[1], tangent[0] + + best_score = 0.0 + best_pt = None + + for sign in (1.0, -1.0): + for d in range(1, tol + 1): + sx = int(round(px + sign * nx * d)) + sy = int(round(py + sign * ny * d)) + if 0 <= sx < W and 0 <= sy < H and valid_snap[sy, sx]: + mag = float(magnitude[sy, sx]) + proximity = 1.0 - d / (tol + 1) + score = mag * proximity + if score > best_score: + best_score = score + best_pt = (sx, sy) + + if best_pt is not None: + refined_pts[i] = best_pt + + refined_contour = refined_pts.reshape(-1, 1, 2).astype(np.int32) + cv2.fillPoly(new_mask, [refined_contour], color=1) + + refined_masks.append(new_mask.astype(bool)) + + refined_masks_np = np.stack(refined_masks, axis=0) + refined_detections = InstanceDetections( + xyxy=segmentation.xyxy, + class_id=segmentation.class_id, + confidence=segmentation.confidence, + mask=torch.from_numpy(refined_masks_np).to( + device=segmentation.xyxy.device, dtype=torch.bool + ), + image_metadata=segmentation.image_metadata, + bboxes_metadata=segmentation.bboxes_metadata, + ) + + return refined_detections, edges_detections diff --git a/inference/core/workflows/core_steps/classical_cv/motion_detection/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/motion_detection/v1_tensor.py new file mode 100644 index 0000000000..777ae05b5f --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/motion_detection/v1_tensor.py @@ -0,0 +1,422 @@ +import json +from typing import List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import cv2 +import numpy as np +import torch +from pydantic import AliasChoices, ConfigDict, Field, PositiveInt +from shapely.geometry import Polygon + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.execution_engine.constants import DETECTION_ID_KEY +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ZONE_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections + + +class MotionDetectionManifest(WorkflowBlockManifest): + type: Literal["roboflow_core/motion_detection@v1"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Motion Detection", + "version": "v1", + "short_description": "Detect motion in a video using OpenCV.", + "long_description": ( + """ +Detect motion in video streams using OpenCV's background subtraction algorithm. + +## How This Block Works + +This block uses background subtraction (specifically the MOG2 algorithm) to detect motion in video frames. The block maintains state across frames to build a background model and track motion patterns: + +1. **Initializes background model** - on the first frame, creates a background subtractor using the specified history and threshold parameters +2. **Processes each frame** - applies background subtraction to identify pixels that differ from the learned background model +3. **Filters noise** - applies morphological operations to remove noise and combine nearby motion regions into coherent contours +4. **Extracts motion regions** - finds contours representing motion areas, filters them by minimum size, and optionally clips them to a detection zone +5. **Simplifies contours** - reduces contour complexity to keep detection data manageable +6. **Generates outputs** - creates object detection predictions with bounding boxes, determines motion status, triggers alarms when motion starts, and provides motion zone polygons + +The block tracks motion state across frames - the **alarm** output becomes true only when motion transitions from not detected to detected, making it useful for triggering actions when motion first appears. + +## Common Use Cases + +- **Security Monitoring**: Detect motion in surveillance cameras to trigger alerts, recordings, or notifications when activity is detected +- **Resource Optimization**: Conditionally run expensive inference operations (e.g., object detection, classification) only when motion is detected to save computational resources +- **Activity Detection**: Monitor areas for movement to track occupancy, identify entry/exit events, or detect unauthorized access +- **Video Analytics**: Analyze video streams to identify motion patterns, track activity levels, or detect anomalies in monitored areas +- **Smart Recording**: Trigger video recording or snapshot capture when motion is detected, reducing storage requirements compared to continuous recording +- **Zone Monitoring**: Monitor specific areas within a frame using detection zones to focus motion detection on relevant regions while ignoring busy but irrelevant areas + +## Connecting to Other Blocks + +The motion detection outputs from this block can be connected to: + +- **Conditional logic blocks** (e.g., Continue If) to execute workflow steps only when motion is detected or when alarms trigger +- **Object detection blocks** to run detection models only on frames with motion, saving computational resources +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send alerts when motion is detected or alarms trigger +- **Data storage blocks** (e.g., Roboflow Dataset Upload, CSV Formatter) to log motion events, timestamps, and detection data for analytics +- **Visualization blocks** to draw motion zones, bounding boxes, or annotations on frames showing detected motion +- **Filter blocks** to filter images or data based on motion status before passing to downstream processing +""" + ), + "license": "Apache-2.0", + "block_type": "classical_computer_vision", + "ui_manifest": { + "section": "video", + "icon": "far fa-bell-exclamation", + "blockPriority": 8, + "opencv": True, + }, + } + ) + + image: Selector(kind=[IMAGE_KIND]) = Field( + title="Input Image", + description="The input image or video frame to analyze for motion. The block processes frames sequentially to build a background model - each frame updates the background model and detects motion relative to learned background patterns. Can be connected from workflow inputs or previous steps.", + examples=["$inputs.image", "$steps.cropping.crops"], + validation_alias=AliasChoices("image", "images"), + ) + + minimum_contour_area: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + title="Minimum Contour Area", + description="Minimum area in square pixels for a motion region to be detected. Contours smaller than this threshold are filtered out to ignore noise, small shadows, or minor pixel variations. Lower values increase sensitivity but may detect more false positives (e.g., 100 for very sensitive detection, 500 for only large objects). Default is 200 square pixels.", + examples=[200, 100, 500], + default=200, + ) + + morphological_kernel_size: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = ( + Field( + title="Morphological Kernel Size", + description="Size of the morphological kernel in pixels used to combine nearby motion regions and filter noise. Larger values merge more distant motion regions into single contours but may also merge separate objects. Smaller values preserve more detail but may leave fragmented detections. The kernel uses an elliptical shape. Default is 3 pixels.", + examples=[3, 5, 7], + default=3, + ) + ) + + threshold: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + title="Threshold", + description="Threshold value for the squared Mahalanobis distance used by the MOG2 background subtraction algorithm. Controls sensitivity to motion - smaller values increase sensitivity (detect smaller changes) but may produce more false positives, larger values decrease sensitivity (only detect significant changes) but may miss subtle motion. Recommended range is 8-32. Default is 16.", + examples=[16, 8, 24, 32], + default=16, + ) + + history: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + title="History", + description="Number of previous frames used to build the background model. Controls how quickly the background adapts to changes - larger values (e.g., 50-100) create a more stable background model that's less sensitive to temporary changes but adapts slowly to permanent background changes. Smaller values (e.g., 10-20) allow faster adaptation but may treat moving objects as background if they stop moving. Default is 30 frames.", + examples=[30, 50, 100], + default=30, + ) + + detection_zone: Union[list, str, Selector(kind=[ZONE_KIND]), Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + title="Detection Zone", + description="Optional polygon zone to limit motion detection to a specific area of the frame. Motion is only detected within this zone, ignoring activity outside. Format: [[x1, y1], [x2, y2], [x3, y3], ...] where coordinates are in pixels. The polygon must have more than 3 points. Can be provided as a list, JSON string, or selector referencing zone outputs from other blocks. Useful for focusing on specific regions (e.g., doorways, windows, restricted areas) while ignoring busy but irrelevant areas. If not provided, motion is detected across the entire frame.", + default=None, + ) + + suppress_first_detections: Union[Selector(kind=[BOOLEAN_KIND]), bool] = Field( # type: ignore + title="Don't Detect Until History is Full", + description="If true, suppresses motion detections until the background model has been initialized with enough frames (specified by the history parameter). This prevents false positives from early frames where the background model hasn't learned the scene yet. When false, the block attempts to detect motion immediately, which may produce unreliable results during initialization. Default is true (recommended for most use cases).", + examples=[True, False], + default=True, + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="motion", + kind=[ + BOOLEAN_KIND, + ], + description="Boolean flag indicating whether motion was detected in the current frame. True if any motion regions were found, false otherwise. This flag is true for every frame with detected motion.", + ), + OutputDefinition( + name="alarm", + kind=[ + BOOLEAN_KIND, + ], + description="Boolean flag that becomes true only when motion transitions from not detected (previous frame) to detected (current frame). Useful for triggering actions when motion first appears. Returns false if motion was already detected in the previous frame, even if motion continues in the current frame.", + ), + OutputDefinition( + name="detections", + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + ], + description="Object detection predictions containing bounding boxes for all detected motion regions. Each detection has class name 'motion', confidence 1.0, and bounding box coordinates. Empty detections if no motion is detected. Compatible with other blocks that accept object detection predictions.", + ), + OutputDefinition( + name="motion_zones", + kind=[ + LIST_OF_VALUES_KIND, + ], + description="List of polygon coordinates representing the exact shapes of detected motion regions. Each polygon is a list of [x, y] coordinate pairs defining the contour of a motion region. Useful for visualization or precise motion area analysis. Empty list if no motion is detected.", + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class MotionDetectionBlockV1(WorkflowBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.last_motion = False + self.back_sub = None + self.frame_count = 0 + self._kernel_cache = {} # Cache morphological kernels by size + + @classmethod + def get_manifest(cls) -> Type[MotionDetectionManifest]: + return MotionDetectionManifest + + def run( + self, + image: WorkflowImageData, + minimum_contour_area: int, + morphological_kernel_size: int, + threshold: int, + history: int, + suppress_first_detections: bool, + detection_zone: Optional[Union[str, List[Tuple[int, int]]]], + *args, + **kwargs, + ) -> BlockResult: + + if isinstance(detection_zone, str): + try: + detection_zone = json.loads(detection_zone) + except Exception as e: + raise ValueError(f"Could not parse detection zone as a valid json") + + if not self.back_sub: + self.frame_count = 0 + self.back_sub = cv2.createBackgroundSubtractorMOG2( + history=history, varThreshold=threshold, detectShadows=True + ) + + frame = image.numpy_image + + # apply background subtraction + mask = self.back_sub.apply(frame) + + # if frames aren't initialized yet, return no motion + if self.frame_count < history and suppress_first_detections: + self.frame_count += 1 + return { + "motion": False, + "detections": _native_detections_from_boxes( + xyxy_boxes=[], + image=image, + ), + "alarm": False, + "motion_zones": [], + } + + # apply morphological filtering to ignore changes due to noise + # Use cached kernel to avoid recreating the same kernel repeatedly + if morphological_kernel_size not in self._kernel_cache: + self._kernel_cache[morphological_kernel_size] = cv2.getStructuringElement( + cv2.MORPH_ELLIPSE, + (morphological_kernel_size, morphological_kernel_size), + ) + kernel = self._kernel_cache[morphological_kernel_size] + mask_morph = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel) + + # create contours around filtered areas + contours, hierarchy = cv2.findContours( + mask_morph, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE + ) + + # apply minimum contour size and filter out 0 length contours + # Check length first (cheaper) before computing area + filtered_contours = [ + contour + for contour in contours + if len(contour) > 2 and cv2.contourArea(contour) > minimum_contour_area + ] + + # clip contours if a detection zone is provided + if detection_zone and len(detection_zone) > 0: + filtered_contours = clip_contours_to_contour( + filtered_contours, detection_zone + ) + + # simplify contours by 1% of their perimeter + # this is ideal for keeping the detections to a reasonable size + simplified_contours = [] + for contour in filtered_contours: + perimeter = cv2.arcLength(contour, True) + epsilon = 0.01 * perimeter + approx = cv2.approxPolyDP(contour, epsilon, True) + # Only keep contours with at least 3 vertices + if len(approx) >= 3: + simplified_contours.append(approx) + + # get bounding boxes and polygons + xyxy_boxes = [] + polygons = [] + for cnt in simplified_contours: + x, y, w, h = cv2.boundingRect(cnt) + xyxy_boxes.append([int(x), int(y), int(x + w), int(y + h)]) + # Extract polygon coordinates, handling both squeezed and unsqueezed formats + polygon = np.squeeze(cnt) + if polygon.ndim == 1: # Single point case, skip + continue + polygons.append(polygon.tolist()) + + # convert to native detections + detections = _native_detections_from_boxes( + xyxy_boxes=xyxy_boxes, + image=image, + ) + + # if contours exist, there's motion + motion = len(filtered_contours) > 0 + + # alarm flips to true only if there was no motion before and motion now + alarm = not self.last_motion and motion + self.last_motion = motion + + return { + "motion": motion, + "detections": detections, + "alarm": alarm, + "motion_zones": polygons, + } + + +def _native_detections_from_boxes( + xyxy_boxes: List[List[int]], + image: WorkflowImageData, +) -> Detections: + number_of_detections = len(xyxy_boxes) + image_metadata = build_native_image_metadata( + image=image, + class_names={0: "motion"}, + prediction_type="object-detection", + ) + bboxes_metadata = ( + [{DETECTION_ID_KEY: str(uuid4())} for _ in range(number_of_detections)] + if number_of_detections > 0 + else None + ) + return Detections( + xyxy=torch.as_tensor( + np.asarray(xyxy_boxes, dtype=np.float32), + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ).reshape(-1, 4), + class_id=torch.zeros( + (number_of_detections,), + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.ones( + (number_of_detections,), + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def clip_contours_to_contour( + contours: List[np.ndarray], clip_contour: np.ndarray +) -> List[np.ndarray]: + """ + Clip OpenCV contours to another contour and return clipped OpenCV contours. + + Args: + contours: List of OpenCV contours, each as numpy array of shape (N, 1, 2) + clip_contour: Clip contour as numpy array of shape (M, 2) with xy points + + Returns: + List of clipped OpenCV contours as numpy arrays of shape (N, 1, 2). + Only includes contours that overlap with the clip contour. + """ + + clip_poly = Polygon(clip_contour) + result = [] + + for contour in contours: + # Convert OpenCV contour (N, 1, 2) to xy points (N, 2) + points = contour.reshape(-1, 2) + + if len(points) < 3: + continue + + try: + poly = Polygon(points) + clipped = poly.intersection(clip_poly) + + if clipped.is_empty: + continue + + # Extract coordinates based on geometry type + if clipped.geom_type == "Polygon": + coords = list(clipped.exterior.coords[:-1]) + if len(coords) >= 3: + result.append(list_to_contour(coords)) + + elif clipped.geom_type == "MultiPolygon": + for geom in clipped.geoms: + coords = list(geom.exterior.coords[:-1]) + if len(coords) >= 3: + result.append(list_to_contour(coords)) + + except Exception: + # Silently skip contours that fail shapely operations + # (e.g., self-intersecting polygons) + continue + + return result + + +def list_to_contour(list_of_tuples: List[Tuple]) -> np.ndarray: + """ + Convert a list of (x, y) tuples to an OpenCV contour format. + + Args: + list_of_tuples: List of coordinate tuples [(x1, y1), (x2, y2), ...] + + Returns: + NumPy array of shape (N, 1, 2) suitable for OpenCV operations + """ + points = np.array( + [[int(xy[0]), int(xy[1])] for xy in list_of_tuples], dtype=np.int32 + ) + return points.reshape(-1, 1, 2) diff --git a/inference/core/workflows/core_steps/classical_cv/pixel_color_count/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/pixel_color_count/v1_tensor.py new file mode 100644 index 0000000000..a2046ccf52 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/pixel_color_count/v1_tensor.py @@ -0,0 +1,97 @@ +"""Tensor-native sibling of ``pixel_color_count/v1``. + +On a tensor-materialised ``(3, H, W)`` image the count is one broadcasted +``(px >= lower) & (px <= upper)`` on the device; only the final count crosses +device->host. + +Parity with v1, which feeds ``cv2.inRange`` unclipped int64 +``target +- tolerance`` bound arrays: + +* ``cv2.inRange`` treats 3-element bounds as a per-channel Scalar: each bound + is ``cvRound``-ed (half-to-even) and tested as the inclusive range + ``lower <= px <= upper``, ANDed across channels, without uint8 saturation - + a bound of -10 means "no lower limit", 310 "no upper limit", 256 matches + nothing, inverted ranges (negative tolerance) match nothing. Reproduced + with int16 device comparisons; bounds are pre-clamped to ``[-1, 256]``, + which cannot change any comparison outcome against uint8 pixels but keeps + them int16-safe at any magnitude. (Bounds beyond int32 wrap inside cv2 but + not here; unreachable with the manifest's 0-255 tolerance.) +* Colour parsing uses v1's ``convert_color_to_bgr_tuple``, so malformed + colours raise the identical ``ValueError`` before any image work. The + per-channel AND is layout-independent, so the BGR bounds tuple is reversed + to pair with the CHW RGB channel axis. +* Single-channel ``(1, H, W)`` tensors delegate: v1 raises ``cv2.error`` for + 3-element bounds against a 1-channel image, and delegation fails + identically. +* Numpy/base64-born images delegate to v1's ``count_specific_color_pixels`` + (no materialised tensor). +""" + +from typing import Tuple, Type, Union + +import torch + +from inference.core.workflows.core_steps.classical_cv.pixel_color_count.v1 import ( + ColorPixelCountManifest, + convert_color_to_bgr_tuple, + count_specific_color_pixels, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlock + + +class PixelationCountBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[ColorPixelCountManifest]: + return ColorPixelCountManifest + + def run( + self, + image: WorkflowImageData, + target_color: Union[str, tuple], + tolerance: int, + ) -> BlockResult: + if image.is_tensor_materialised() and image.tensor_image.shape[0] == 3: + color_pixel_count = _count_specific_color_pixels_tensor( + image.tensor_image, target_color, tolerance + ) + else: + color_pixel_count = count_specific_color_pixels( + image.numpy_image, target_color, tolerance + ) + return {"matching_pixels_count": color_pixel_count} + + +def _count_specific_color_pixels_tensor( + image: torch.Tensor, + target_color: Union[str, Tuple[int, int, int]], + tolerance: int, +) -> int: + """Counts pixels of a CHW RGB uint8 tensor that match the target colour + within tolerance, replicating ``cv2.inRange`` + ``cv2.countNonZero`` + exactly; only the final count leaves the device.""" + # Colour conversion first, before any image work - mirrors v1's error + # ordering. + target_color_bgr = convert_color_to_bgr_tuple(color=target_color) + target_color_rgb = target_color_bgr[::-1] + lower_bound = [ + _mirror_cv2_scalar_bound(channel - tolerance) for channel in target_color_rgb + ] + upper_bound = [ + _mirror_cv2_scalar_bound(channel + tolerance) for channel in target_color_rgb + ] + device = image.device + lower = torch.tensor(lower_bound, dtype=torch.int16, device=device).view(3, 1, 1) + upper = torch.tensor(upper_bound, dtype=torch.int16, device=device).view(3, 1, 1) + pixels = image.detach().to(torch.int16) + matching = ((pixels >= lower) & (pixels <= upper)).all(dim=0) + return int(matching.sum().item()) + + +def _mirror_cv2_scalar_bound(value: Union[int, float]) -> int: + """cv2 converts each Scalar bound with ``cvRound`` (half-to-even, matching + Python's ``round``) and compares in integer space without clamping. + Clamping to ``[-1, 256]`` cannot change any comparison outcome against + uint8 pixels while keeping the bound int16-safe at any magnitude.""" + return min(256, max(-1, round(value))) diff --git a/inference/core/workflows/core_steps/classical_cv/size_measurement/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/size_measurement/v1_tensor.py new file mode 100644 index 0000000000..02e6f00535 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/size_measurement/v1_tensor.py @@ -0,0 +1,316 @@ +from typing import List, Literal, Optional, Tuple, Type, Union + +import cv2 as cv +import numpy as np +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + instance_mask_to_numpy, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY = "dimensions" +SHORT_DESCRIPTION = ( + "Measure the dimensions of objects in relation to a reference object." +) +LONG_DESCRIPTION = """ +The [**Size Measurement Block**](https://www. + +## How This Block Works + +youtube.com/watch?v=FQY7TSHfZeI) calculates the dimensions of objects relative to a reference object. It uses one model to detect the reference object and another to detect the objects to measure. The block outputs the dimensions of the objects in terms of the reference object. + +- **Reference Object**: This is the known object used as a baseline for measurements. Its dimensions are known and used to scale the measurements of other objects. +- **Object to Measure**: This is the object whose dimensions are being calculated. The block measures these dimensions relative to the reference object. + +### Block Usage + +To use the Size Measurement Block, follow these steps: + +1. **Select Models**: Choose a model to detect the reference object and another model to detect the objects you want to measure. +2. **Configure Inputs**: Provide the predictions from both models as inputs to the block. +3. **Set Reference Dimensions**: Specify the known dimensions of the reference object in the format 'width,height' or as a tuple (width, height). +4. **Run the Block**: Execute the block to calculate the dimensions of the detected objects relative to the reference object. + +### Example + +Imagine you have a scene with a calibration card and several packages. The calibration card has known dimensions of 5.0 inches by 3.0 inches. You want to measure the dimensions of packages in the scene. + +- **Reference Object**: Calibration card with dimensions 5.0 inches (width) by 3.0 inches (height). +- **Objects to Measure**: Packages detected in the scene. + +The block will use the known dimensions of the calibration card to calculate the dimensions of each package. For example, if a package is detected with a width of 100 pixels and a height of 60 pixels, and the calibration card is detected with a width of 50 pixels and a height of 30 pixels, the block will calculate the package's dimensions as: + +- **Width**: (100 pixels / 50 pixels) * 5.0 inches = 10.0 inches +- **Height**: (60 pixels / 30 pixels) * 3.0 inches = 6.0 inches + +This allows you to obtain the real-world dimensions of the packages based on the reference object's known size. + +[Watch the video tutorial](https://www.youtube.com/watch?v=FQY7TSHfZeI) +""" + + +class SizeMeasurementManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Size Measurement", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "classical_cv", + "icon": "far fa-ruler", + "blockPriority": 10, + "opencv": True, + }, + } + ) + type: Literal["roboflow_core/size_measurement@v1"] + + object_predictions: Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + description="Model predictions to measure the dimensions of.", + examples=["$segmentation.object_predictions"], + ) + + reference_predictions: Union[ + list, + Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + LIST_OF_VALUES_KIND, + ] + ), + ] = Field( + description="Reference object used to calculate the dimensions of the specified objects. If multiple objects are provided, the highest confidence prediction will be used.", + examples=["$segmentation.reference_predictions"], + ) + reference_dimensions: Union[ + str, + Tuple[float, float], + List[float], + Selector( + kind=[STRING_KIND, LIST_OF_VALUES_KIND], + ), + ] = Field( + description="Dimensions of the reference object in desired units, (e.g. inches). Will be used to convert the pixel dimensions of the other objects to real-world units.", + examples=[(4.5, 3.0), "5.0,5.0", "$inputs.reference_dimensions"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=OUTPUT_KEY, kind=[LIST_OF_VALUES_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +def horizontal_score(angle: float) -> float: + """ + Determine how close an angle is to horizontal (0 or 180 degrees). + Lower score means more horizontal. + """ + mod_angle = abs(angle % 180) + return min(mod_angle, 180 - mod_angle) + + +def compute_aligned_dimensions(contour: np.ndarray) -> Tuple[float, float]: + """ + Compute the width and height of an object based on its contour, ensuring proper orientation. + + This function: + 1. Finds the minimum area rectangle that encloses the contour + 2. Determines which edges correspond to width and height by analyzing their angles + 3. Returns dimensions where width is the more horizontal edge and height is the more vertical edge + + Args: + contour (np.ndarray): Array of points representing the object's contour + + Returns: + Tuple[float, float]: A tuple of (width_pixels, height_pixels) where: + - width_pixels: Length of the more horizontal edge + - height_pixels: Length of the more vertical edge + + Note: + The function uses angle analysis to ensure consistent width/height assignment + regardless of the object's rotation. The edge closer to horizontal (0ยฐ or 180ยฐ) + is always considered the width. + """ + rect = cv.minAreaRect(contour) + box = cv.boxPoints(rect) + box = np.array(box, dtype=np.float32) + + edge1 = box[1] - box[0] + edge2 = box[2] - box[1] + + len_edge1 = np.linalg.norm(edge1) + len_edge2 = np.linalg.norm(edge2) + + angle1 = np.degrees(np.arctan2(edge1[1], edge1[0])) + angle2 = np.degrees(np.arctan2(edge2[1], edge2[0])) + + h_score1 = horizontal_score(angle1) + h_score2 = horizontal_score(angle2) + + if h_score1 < h_score2: + width_pixels = len_edge1 + height_pixels = len_edge2 + else: + width_pixels = len_edge2 + height_pixels = len_edge1 + + return float(width_pixels), float(height_pixels) + + +def get_detection_dimensions( + detection: Union[Detections, InstanceDetections], index: int +) -> Tuple[Optional[float], Optional[float]]: + """ + Retrieve the width and height dimensions of a detected object in pixels. + + Args: + detection (sv.Detections): Detection object containing masks and/or bounding boxes + index (int): Index of the specific detection to analyze + + Returns: + Tuple[float, float]: A tuple of (width_pixels, height_pixels) where: + - width_pixels: Width of the object in pixels + - height_pixels: Height of the object in pixels + + Notes: + The function uses two methods to compute dimensions: + 1. If a segmentation mask is available: + - Extracts the largest contour from the mask + - Uses compute_aligned_dimensions() to get orientation-aware measurements + 2. If no mask is available: + - Falls back to using the bounding box dimensions + - Simply computes width and height as box edges + """ + if isinstance(detection, InstanceDetections): + mask = instance_mask_to_numpy(detection, index).astype(np.uint8) + contours, _ = cv.findContours(mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE) + if contours: + largest_contour = max(contours, key=cv.contourArea) + if cv.contourArea(largest_contour) > 0: + return compute_aligned_dimensions(largest_contour) + + else: + bbox = detection.xyxy[index].detach().to("cpu").numpy() + w = bbox[2] - bbox[0] + h = bbox[3] - bbox[1] + return float(w), float(h) + + return None, None + + +def get_polygon_dimensions(polygon: list) -> Tuple[Optional[float], Optional[float]]: + polygon = np.array(polygon) + if len(polygon) >= 3 and cv.contourArea(polygon) > 0: + return compute_aligned_dimensions(polygon) + return None, None + + +def parse_reference_dimensions( + reference_dimensions: Union[str, Tuple[float, float], List[float]], +) -> Tuple[float, float]: + """Parse reference dimensions from various input formats.""" + if isinstance(reference_dimensions, str): + parts = reference_dimensions.split(",") + if len(parts) != 2: + raise ValueError( + "reference_dimensions must be a string in the format 'width,height'" + ) + try: + reference_dimensions = [float(p.strip()) for p in parts] + except ValueError: + raise ValueError("Invalid format for reference_dimensions") + + if len(reference_dimensions) != 2: + raise ValueError("reference_dimensions must have two values (width, height)") + + return tuple(reference_dimensions) + + +class SizeMeasurementBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return SizeMeasurementManifest + + def run( + self, + reference_predictions: Union[list, Detections, InstanceDetections], + object_predictions: Union[Detections, InstanceDetections], + reference_dimensions: Union[str, Tuple[float, float], List[float]], + ) -> BlockResult: + ref_width_actual, ref_height_actual = parse_reference_dimensions( + reference_dimensions + ) + + if hasattr(reference_predictions, "confidence"): + if len(reference_predictions.confidence) == 0: + return {OUTPUT_KEY: None} + + ref_index = int( + np.argmax(reference_predictions.confidence.detach().to("cpu").numpy()) + ) + ref_width_pixels, ref_height_pixels = get_detection_dimensions( + reference_predictions, ref_index + ) + elif isinstance(reference_predictions, list): + if len(reference_predictions) < 2: + return {OUTPUT_KEY: None} + + ref_width_pixels, ref_height_pixels = get_polygon_dimensions( + reference_predictions + ) + + if not ref_width_pixels or not ref_height_pixels: + return {OUTPUT_KEY: None} + + width_scale = ref_width_actual / ref_width_pixels + height_scale = ref_height_actual / ref_height_pixels + + dimensions = [] + for i in range(len(object_predictions)): + obj_w_pixels, obj_h_pixels = get_detection_dimensions(object_predictions, i) + if obj_w_pixels and obj_h_pixels and obj_w_pixels > 0 and obj_h_pixels > 0: + obj_w_actual = obj_w_pixels * width_scale + obj_h_actual = obj_h_pixels * height_scale + dimensions.append( + { + "width": obj_w_actual, + "height": obj_h_actual, + "longer": max(obj_w_actual, obj_h_actual), + "shorter": min(obj_w_actual, obj_h_actual), + } + ) + else: + dimensions.append(None) + + return {OUTPUT_KEY: dimensions} diff --git a/inference/core/workflows/core_steps/classical_cv/template_matching/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/template_matching/v1_tensor.py new file mode 100644 index 0000000000..a3d9f1a3d7 --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/template_matching/v1_tensor.py @@ -0,0 +1,269 @@ +from typing import List, Literal, Optional, Type, Union +from uuid import uuid4 + +import cv2 +import numpy as np +import supervision as sv +import torch +from pydantic import AliasChoices, ConfigDict, Field + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.execution_engine.constants import DETECTION_ID_KEY +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections + +SHORT_DESCRIPTION: str = ( + "Locate instances of a given template within a specified image." +) +LONG_DESCRIPTION: str = """ +Locate instances of a template image within a larger image using template matching with normalized cross-correlation, finding exact or near-exact matches of the template pattern at any location in the image, outputting bounding box detections with optional NMS filtering for object detection, logo detection, pattern recognition, and template-based object localization workflows. + +## How This Block Works + +This block searches for occurrences of a template image within a larger input image using normalized cross-correlation template matching. The block: + +1. Receives an input image and a template image (smaller pattern to search for) +2. Converts both images to grayscale for template matching (template matching typically works on grayscale images for efficiency and robustness) +3. Performs template matching using OpenCV's matchTemplate with TM_CCOEFF_NORMED method: + - Slides the template across the input image at every possible position + - Computes normalized cross-correlation coefficient at each position (measures similarity between template and image region) + - Generates a similarity map showing how well the template matches at each location +4. Identifies match locations where similarity exceeds the matching_threshold: + - Finds all positions where the correlation coefficient is greater than or equal to the threshold + - Threshold values range from 0.0 to 1.0, with higher values requiring closer matches + - Lower thresholds find more potential matches (including partial matches), higher thresholds find only very similar matches +5. Creates bounding boxes for each match: + - Each match location becomes a detection with a bounding box matching the template's dimensions + - All detections have confidence of 1.0 (they met the threshold requirement) + - All detections are assigned class "template_match" and class_id 0 + - Each detection gets a unique detection ID for tracking +6. Optionally applies Non-Maximum Suppression (NMS) to filter overlapping detections: + - Template matching often produces many overlapping detections at the same location (duplicate matches) + - NMS removes overlapping detections, keeping only the best match in each area + - NMS threshold controls how much overlap is allowed before removing detections + - Can be disabled (apply_nms=False) if NMS becomes computationally intractable with very large numbers of matches +7. Attaches metadata to detections: + - Sets parent_id to reference the input image + - Sets prediction_type to "object-detection" + - Stores image dimensions for coordinate reference + - Attaches parent coordinate information for workflow tracking +8. Returns detection predictions in sv.Detections format along with the total number of matches found + +The block uses normalized cross-correlation which is effective for finding exact or near-exact template matches. It works best when the template appears in the image at the same scale, rotation, and lighting conditions. The method tends to produce many overlapping detections for the same match location, which is why NMS filtering is important. However, in cases with extremely large numbers of matches (e.g., repeating patterns), NMS may become computationally expensive and can be disabled if needed. + +## Common Use Cases + +- **Logo and Brand Detection**: Find specific logos or brand elements within images (e.g., detect company logos in photos, find brand markers in images, locate specific logo patterns in scenes), enabling logo detection workflows +- **Exact Pattern Matching**: Locate specific patterns or objects that appear identically in images (e.g., find specific UI elements in screenshots, detect exact patterns in images, locate specific visual elements), enabling exact pattern detection workflows +- **Quality Control and Inspection**: Find reference patterns or features for quality inspection (e.g., detect specific features in manufacturing images, find reference markers for alignment, locate inspection targets), enabling quality control workflows +- **Object Localization**: Locate specific objects or regions when exact appearance is known (e.g., find specific objects with known appearance, locate reference objects in images, detect specific visual elements), enabling template-based object localization +- **Document Processing**: Find specific elements or regions in documents (e.g., locate form fields in documents, detect specific document elements, find reference markers in scanned documents), enabling document processing workflows +- **UI Element Detection**: Detect specific UI components or elements in interface images (e.g., find buttons in UI screenshots, locate specific UI elements, detect interface components), enabling UI analysis workflows + +## Connecting to Other Blocks + +This block receives an image and template, and produces detection predictions: + +- **After image input blocks** to find template patterns in input images (e.g., search for templates in input images, locate patterns in camera feeds, find templates in image streams), enabling template matching workflows +- **After preprocessing blocks** to find templates in preprocessed images (e.g., match templates after image enhancement, find patterns in filtered images, locate templates in normalized images), enabling preprocessed template matching +- **Before visualization blocks** to visualize template match locations (e.g., visualize detected template matches, display bounding boxes for matches, show template match results), enabling template match visualization workflows +- **Before filtering blocks** to filter template matches by criteria (e.g., filter matches by location, select specific match regions, refine template match results), enabling filtered template matching workflows +- **Before crop blocks** to extract regions around template matches (e.g., crop areas around matches, extract match regions for analysis, crop template match locations), enabling template-based region extraction +- **In quality control workflows** where template matching is used for inspection or alignment (e.g., find reference markers for alignment, detect inspection targets, locate quality control features), enabling quality control template matching workflows +""" + + +class TemplateMatchingManifest(WorkflowBlockManifest): + type: Literal["roboflow_core/template_matching@v1"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Template Matching", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "classical_computer_vision", + "ui_manifest": { + "section": "classical_cv", + "icon": "far fa-crosshairs", + "blockPriority": 0.5, + "opencv": True, + }, + } + ) + image: Selector(kind=[IMAGE_KIND]) = Field( + title="Input Image", + description="Large image in which to search for the template pattern. The template will be searched across this entire image at all possible positions. The image is converted to grayscale internally for template matching. Template matching works best when the image and template have similar lighting conditions and the template appears at similar scale and orientation in the image.", + examples=["$inputs.image", "$steps.cropping.crops"], + validation_alias=AliasChoices("image", "images"), + ) + template: Selector(kind=[IMAGE_KIND]) = Field( + title="Template Image", + description="Small template image pattern to search for within the input image. The template should be smaller than the input image. The template is converted to grayscale internally for matching. Template matching finds exact or near-exact matches of this template at any location in the input image. Works best when the template appears in the image at the same scale, rotation, and lighting conditions. The template's dimensions determine the size of the detection bounding boxes.", + examples=["$inputs.template", "$steps.cropping.template"], + validation_alias=AliasChoices("template", "templates"), + ) + matching_threshold: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + title="Matching Threshold", + description="Minimum similarity threshold (0.0 to 1.0) required for a template match. Higher values (closer to 1.0) require very close matches and find fewer but more precise matches. Lower values (closer to 0.0) allow more lenient matches and find more potential matches including partial matches. Default is 0.8, which requires fairly close matches. Use lower thresholds (0.6-0.7) to find more matches or handle slight variations. Use higher thresholds (0.85-0.95) for exact matches only. The threshold compares normalized cross-correlation coefficients from template matching.", + default=0.8, + examples=[0.8, "$inputs.threshold"], + ) + apply_nms: Union[Selector(kind=[BOOLEAN_KIND]), bool] = Field( + title="Apply NMS", + description="Whether to apply Non-Maximum Suppression (NMS) to filter overlapping detections. Template matching often produces many overlapping detections at the same location. NMS removes overlapping detections, keeping only the best match in each area. Default is True (recommended for most cases). Set to False if: (1) the number of matches is extremely large (NMS may become computationally expensive), (2) you want to see all raw matches without filtering, or (3) matches are intentionally close together and should all be kept. When disabled, you may see many duplicate detections for the same match location.", + default=True, + examples=["$inputs.apply_nms", False], + ) + nms_threshold: Union[Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), FloatZeroToOne] = ( + Field( + title="NMS threshold", + description="Intersection over Union (IoU) threshold for Non-Maximum Suppression. Only relevant when apply_nms is True. Detections with IoU overlap greater than this threshold are considered duplicates, and only the detection with highest confidence is kept. Lower values (0.3-0.4) are more aggressive at removing overlaps, removing detections that are only slightly overlapping. Higher values (0.6-0.7) are more lenient, only removing heavily overlapping detections. Default is 0.5, which provides balanced overlap filtering. Adjust based on how much overlap you expect between template matches and how close together valid matches can be.", + default=0.5, + examples=["$inputs.nms_threshold", 0.3], + ) + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition( + name="number_of_matches", + kind=[INTEGER_KIND], + ), + ] + + +class TemplateMatchingBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[TemplateMatchingManifest]: + return TemplateMatchingManifest + + def run( + self, + image: WorkflowImageData, + template: WorkflowImageData, + matching_threshold: float, + apply_nms: bool, + nms_threshold: float, + ) -> BlockResult: + detections = apply_template_matching( + image=image, + template=template.numpy_image, + matching_threshold=matching_threshold, + apply_nms=apply_nms, + nms_threshold=nms_threshold, + ) + return {"predictions": detections, "number_of_matches": len(detections)} + + +def apply_template_matching( + image: WorkflowImageData, + template: np.ndarray, + matching_threshold: float, + apply_nms: bool, + nms_threshold: float, +) -> Detections: + img_gray = cv2.cvtColor(image.numpy_image, cv2.COLOR_BGR2GRAY) + template_gray = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY) + w, h = template_gray.shape[::-1] + res = cv2.matchTemplate(img_gray, template_gray, cv2.TM_CCOEFF_NORMED) + loc = np.where(res >= matching_threshold) + xyxy = [] + for pt in zip(*loc[::-1]): + top_left = pt + bottom_right = (pt[0] + w, pt[1] + h) + xyxy.append(top_left + bottom_right) + if len(xyxy) == 0: + return _native_detections_from_boxes( + xyxy=np.zeros((0, 4), dtype=np.int32), + confidence=np.zeros((0,), dtype=float), + image=image, + ) + xyxy = np.array(xyxy).astype(np.int32) + confidence = np.ones(len(xyxy), dtype=float) + if apply_nms: + # sv.Detections is used here purely as the NMS algorithm (mirroring the + # numpy block's with_nms); the output is built natively below โ€” no + # sv.Detections is returned. + nms_input = sv.Detections( + xyxy=xyxy, + confidence=confidence, + class_id=np.zeros(len(xyxy), dtype=np.uint32), + ).with_nms(threshold=nms_threshold) + xyxy = nms_input.xyxy + confidence = nms_input.confidence + return _native_detections_from_boxes( + xyxy=xyxy, + confidence=confidence, + image=image, + ) + + +def _native_detections_from_boxes( + xyxy: np.ndarray, + confidence: np.ndarray, + image: WorkflowImageData, +) -> Detections: + number_of_detections = len(xyxy) + image_metadata = build_native_image_metadata( + image=image, + class_names={0: "template_match"}, + prediction_type="object-detection", + ) + bboxes_metadata = ( + [{DETECTION_ID_KEY: str(uuid4())} for _ in range(number_of_detections)] + if number_of_detections > 0 + else None + ) + return Detections( + xyxy=torch.as_tensor( + np.asarray(xyxy), + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ).reshape(-1, 4), + class_id=torch.zeros( + (number_of_detections,), + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.as_tensor( + np.asarray(confidence), + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) diff --git a/inference/core/workflows/core_steps/classical_cv/threshold/v1_tensor.py b/inference/core/workflows/core_steps/classical_cv/threshold/v1_tensor.py new file mode 100644 index 0000000000..7c23285a8b --- /dev/null +++ b/inference/core/workflows/core_steps/classical_cv/threshold/v1_tensor.py @@ -0,0 +1,264 @@ +"""Tensor-native sibling of ``threshold/v1``. + +Every path replicates the exact uint8 semantics of the cv2 calls in v1's +``apply_thresholding`` (OpenCV modules/imgproc/src/thresh.cpp) or delegates +to them - parity versus the numpy block is bit-exact: + +* Fixed types (``binary`` / ``binary_inv`` / ``trunc`` / ``tozero`` / + ``tozero_inv``): per-pixel integer comparisons with cv2's parameter + handling - ``ithresh = cvFloor(thresh)`` with a strict ``>`` compare, + ``imaxval = saturate_cast(cvRound(maxval))`` (round half to even, + then clamp to [0, 255]), ``trunc`` ignoring ``maxval`` and filling with the + saturated ``ithresh``, and the degenerate ``ithresh < 0 / >= 255`` branches + that fill or copy without touching pixels. cv2 applies these per-channel, + so any channel count runs tensor-native. + +* ``otsu`` (cv2 requires single-channel 8-bit input): ``torch.bincount`` runs + on the device, one D2H sync moves the 256 counts to the host, and + ``getThreshVal_Otsu_8u``'s double-precision scan is replicated in python + floats operation-for-operation (FLT_EPSILON guards, strict ``>`` + tie-breaking, the ``mu1 *= q1`` pre-multiplication ordering); the found + threshold is applied on the device with the binary semantics above. + Three-channel tensor input delegates so cv2 raises exactly v1's error. + +* ``adaptive_mean``: cv2 computes ``boxFilter(src, CV_8U, 11x11, + normalize=True, BORDER_REPLICATE|BORDER_ISOLATED)`` - the window mean + rounded to uint8 - then ``dst = src > mean - 2 ? imaxval : 0`` (delta 2 is + ``cvCeil``-ed to the exact integer 2 for THRESH_BINARY). Sums of 121 uint8 + pixels stay below 2**24, so a float32 convolution over a replicate-padded + image yields exact integer sums on any device in any summation order, and + ``sum/121`` can never land on a .5 tie (2*sum is even, 121*(2k+1) is odd), + so rounding matches ``cvRound`` bit-exactly. + +* ``adaptive_gaussian``: delegates. cv2 computes the local mean as + ``convertTo(GaussianBlur(float32(src), 11x11, sigma 2.0), CV_8U)``, and the + float32 filter's bit pattern differs across SIMD backends, so no torch + replication is bit-stable across deployment hardware. + +Numpy/base64-born images delegate (no materialised tensor). +""" + +import math +from typing import Type, Union + +import torch +import torch.nn.functional as F + +from inference.core.workflows.core_steps.classical_cv.threshold.v1 import ( + ImageThresholdManifest, + apply_thresholding, +) +from inference.core.workflows.core_steps.visualizations.common.base import ( + OUTPUT_IMAGE_KEY, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlock + +_FIXED_THRESHOLD_TYPES = {"binary", "binary_inv", "trunc", "tozero", "tozero_inv"} +_SINGLE_CHANNEL_TENSOR_TYPES = {"otsu", "adaptive_mean"} +# C's FLT_EPSILON (2**-23) - getThreshVal_Otsu_8u guards with the FLOAT +# epsilon even though the scan itself runs in double precision. +_FLT_EPSILON = 1.1920928955078125e-07 +# v1 hardcodes cv2.adaptiveThreshold(..., blockSize=11, C=2) with THRESH_BINARY, +# for which cv2 turns delta into `idelta = cvCeil(2.0)` - the exact integer 2. +_ADAPTIVE_BLOCK_SIZE = 11 +_ADAPTIVE_DELTA = 2 + + +class ImageThresholdBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[ImageThresholdManifest]: + return ImageThresholdManifest + + def run( + self, + image: WorkflowImageData, + threshold_type: str, + thresh_value: int, + max_value: int, + *args, + **kwargs, + ) -> BlockResult: + if _requires_numpy_delegation(image=image, threshold_type=threshold_type): + thresholded_image = apply_thresholding( + image.numpy_image, threshold_type, thresh_value, max_value + ) + output = WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=thresholded_image, + ) + return {OUTPUT_IMAGE_KEY: output} + thresholded_tensor = _apply_thresholding_tensor( + chw=image.tensor_image, + threshold_type=threshold_type, + thresh_value=thresh_value, + max_value=max_value, + ) + output = WorkflowImageData.copy_and_replace( + origin_image_data=image, + tensor_image=thresholded_tensor, + ) + return {OUTPUT_IMAGE_KEY: output} + + +def _requires_numpy_delegation(image: WorkflowImageData, threshold_type: str) -> bool: + if not image.is_tensor_materialised(): + return True + if threshold_type == "adaptive_gaussian": + # SIMD-dispatch-dependent float32 blur - see the module docstring. + return True + if threshold_type in _SINGLE_CHANNEL_TENSOR_TYPES: + # cv2 asserts CV_8UC1 for otsu/adaptive - delegating multi-channel + # input reproduces exactly the error v1 raises. + return int(image.tensor_image.shape[0]) != 1 + return False + + +def _apply_thresholding_tensor( + chw: torch.Tensor, + threshold_type: str, + thresh_value: Union[int, float], + max_value: Union[int, float], +) -> torch.Tensor: + if threshold_type in _FIXED_THRESHOLD_TYPES: + return _fixed_threshold_tensor( + chw=chw, + threshold_type=threshold_type, + thresh_value=thresh_value, + max_value=max_value, + ) + if threshold_type == "otsu": + return _otsu_threshold_tensor(chw=chw, max_value=max_value) + if threshold_type == "adaptive_mean": + return _adaptive_mean_threshold_tensor(chw=chw, max_value=max_value) + raise ValueError(f"Unknown threshold type: {threshold_type}") + + +def _fixed_threshold_tensor( + chw: torch.Tensor, + threshold_type: str, + thresh_value: Union[int, float], + max_value: Union[int, float], +) -> torch.Tensor: + ithresh = _cv_floor(thresh_value) + imaxval = _cv_round(max_value) + if threshold_type == "trunc": + imaxval = ithresh # trunc ignores maxval and fills with the threshold + imaxval = _saturate_cast_uint8(imaxval) + if ithresh < 0 or ithresh >= 255: + # cv2's degenerate branches fill or copy without reading pixels. + if threshold_type == "binary": + return torch.full_like(chw, 0 if ithresh >= 255 else imaxval) + if threshold_type == "binary_inv": + return torch.full_like(chw, imaxval if ithresh >= 255 else 0) + if threshold_type in ("trunc", "tozero_inv") and ithresh < 0: + return torch.full_like(chw, 0) + if threshold_type == "tozero" and ithresh >= 255: + return torch.full_like(chw, 0) + return chw.clone() + above = chw > ithresh # cv2 compares strictly against the floored threshold + zero = chw.new_zeros(()) + if threshold_type == "binary": + return torch.where(above, chw.new_full((), imaxval), zero) + if threshold_type == "binary_inv": + return torch.where(above, zero, chw.new_full((), imaxval)) + if threshold_type == "trunc": + return torch.clamp(chw, max=imaxval) # imaxval == ithresh in [0, 254] + if threshold_type == "tozero": + return torch.where(above, chw, zero) + return torch.where(above, zero, chw) # tozero_inv + + +def _otsu_threshold_tensor( + chw: torch.Tensor, max_value: Union[int, float] +) -> torch.Tensor: + counts = torch.bincount(chw.detach().reshape(-1).long(), minlength=256) + ithresh = _otsu_threshold_from_counts(counts=counts.cpu().tolist()) + # cv2 feeds the found threshold through the standard binary path; the scan + # only returns values in [0, 254], so the strict-`>` in-range branch + # always applies. + return _fixed_threshold_tensor( + chw=chw, + threshold_type="binary", + thresh_value=ithresh, + max_value=max_value, + ) + + +def _otsu_threshold_from_counts(counts: list) -> int: + """OpenCV's ``getThreshVal_Otsu_8u`` (modules/imgproc/src/thresh.cpp) + replicated operation-for-operation in python floats (IEEE doubles): the + same accumulation order, the FLT_EPSILON validity guards, and the strict + ``sigma > max_sigma`` comparison that keeps the FIRST maximizer.""" + total = sum(counts) + scale = 1.0 / total + mu = 0.0 + for i in range(256): + mu += i * float(counts[i]) + mu *= scale + mu1 = 0.0 + q1 = 0.0 + max_sigma = 0.0 + max_val = 0.0 + for i in range(256): + p_i = counts[i] * scale + mu1 *= q1 # cv2 un-normalizes BEFORE the validity guard + q1 += p_i + q2 = 1.0 - q1 + if min(q1, q2) < _FLT_EPSILON or max(q1, q2) > 1.0 - _FLT_EPSILON: + continue + mu1 = (mu1 + i * p_i) / q1 + mu2 = (mu - q1 * mu1) / q2 + sigma = q1 * q2 * (mu1 - mu2) * (mu1 - mu2) + if sigma > max_sigma: + max_sigma = sigma + max_val = i + return int(max_val) + + +def _adaptive_mean_threshold_tensor( + chw: torch.Tensor, max_value: Union[int, float] +) -> torch.Tensor: + imaxval = _saturate_cast_uint8(_cv_round(max_value)) + pad = _ADAPTIVE_BLOCK_SIZE // 2 + x = chw.detach().to(torch.float32).unsqueeze(0) # (1, 1, H, W) + padded = _replicate_pad(x=x, pad=pad) + kernel = torch.ones( + (1, 1, _ADAPTIVE_BLOCK_SIZE, _ADAPTIVE_BLOCK_SIZE), + dtype=torch.float32, + device=chw.device, + ) + # 121 uint8 addends keep every partial sum an exact float32 integer + # (< 2**24), so the sums equal cv2's integer boxFilter sums bit-for-bit + # on any device and in any accumulation order. + sums = F.conv2d(padded, kernel) + scale = 1.0 / (_ADAPTIVE_BLOCK_SIZE * _ADAPTIVE_BLOCK_SIZE) + # torch.round is half-to-even like cvRound; sum/121 can never tie on .5, + # so the rounded means match cv2's uint8 boxFilter output exactly. + mean = torch.round(sums * scale).clamp_(0, 255).to(torch.int16) + above = chw.to(torch.int16).unsqueeze(0) > mean - _ADAPTIVE_DELTA + thresholded = torch.where(above, chw.new_full((), imaxval), chw.new_zeros(())) + return thresholded.squeeze(0) + + +def _replicate_pad(x: torch.Tensor, pad: int) -> torch.Tensor: + """BORDER_REPLICATE as a clamped-index gather: unlike ``F.pad`` with + ``mode='replicate'`` it also handles images smaller than the padding, + which cv2 supports.""" + height, width = x.shape[-2], x.shape[-1] + rows = torch.arange(-pad, height + pad, device=x.device).clamp_(0, height - 1) + cols = torch.arange(-pad, width + pad, device=x.device).clamp_(0, width - 1) + return x.index_select(-2, rows).index_select(-1, cols) + + +def _cv_floor(value: Union[int, float]) -> int: + return math.floor(value) + + +def _cv_round(value: Union[int, float]) -> int: + # cvRound rounds half to even (IEEE round-to-nearest), like python round(). + return int(round(float(value))) + + +def _saturate_cast_uint8(value: int) -> int: + return min(max(int(value), 0), 255) diff --git a/inference/core/workflows/core_steps/common/deserializers_tensor.py b/inference/core/workflows/core_steps/common/deserializers_tensor.py new file mode 100644 index 0000000000..29d4387d31 --- /dev/null +++ b/inference/core/workflows/core_steps/common/deserializers_tensor.py @@ -0,0 +1,685 @@ +""" +Tensor-native sibling of `common/deserializers.py`. The loader selects this module +when `ENABLE_TENSOR_DATA_REPRESENTATION` is enabled. Functions here add tensor-aware +code paths and delegate other inputs to the NumPy implementations. +""" + +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.deserializers import ( + _parse_optional_parent_metadata, +) +from inference.core.workflows.core_steps.common.deserializers import ( + deserialize_image_kind as _deserialize_image_kind_numpy, +) +from inference.core.workflows.core_steps.common.deserializers import ( + deserialize_video_metadata_kind, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + native_detections_from_inference_predictions, +) +from inference.core.workflows.errors import RuntimeInputError +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_KEY, + CLASSIFICATION_STYLE_MODEL, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + NEAREST_TARGET_DISTANCE_KEY, + PARENT_ID_KEY, + PARENT_ORIGIN_KEY, + PREDICTION_TYPE_KEY, + RLE_MASK_KEY_IN_INFERENCE_RESPONSE, + ROOT_PARENT_ID_KEY, + ROOT_PARENT_ORIGIN_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + OriginCoordinatesSystem, + WorkflowImageData, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks + +DEFAULT_OBJECT_DETECTION_PREDICTION_TYPE = "object-detection" +DEFAULT_INSTANCE_SEGMENTATION_PREDICTION_TYPE = "instance-segmentation" + + +#: Channel counts treated as a "channels" axis when sniffing tensor layout. +_CHANNEL_AXIS_SIZES = (1, 3, 4) + + +def _ensure_chw_layout(image: torch.Tensor) -> torch.Tensor: + """Normalise a single-image tensor to the `WorkflowImageData` CHW contract. + + `WorkflowImageData.tensor_image` uses CHW RGB. Producers may provide a + channels-last HWC tensor, which is normalised before constructing the workflow + image container. + + Detect channels-last with the same heuristic the model preprocessing uses + (`pre_processing.py`: channels not at the front but present at the back) and + permute HWC -> CHW. CHW input (front axis is the channels axis) is returned + untouched. Non-3D tensors are left alone. + """ + if image.ndim == 2: + # Single-channel (H, W) tensors are normalised to the (1, H, W) CHW + # contract WorkflowImageData.tensor_image documents - grayscale carries + # no channel semantics, so no reversal is involved. + return image.unsqueeze(0).contiguous() + if image.ndim != 3: + return image + channels_first = image.shape[0] in _CHANNEL_AXIS_SIZES + channels_last = image.shape[2] in _CHANNEL_AXIS_SIZES + if channels_last and not channels_first: + return image.permute(2, 0, 1).contiguous() + return image + + +def deserialize_image_kind( + parameter: str, + image: Any, + prevent_local_images_loading: bool = False, +) -> WorkflowImageData: + if isinstance(image, WorkflowImageData): + return image + if isinstance(image, torch.Tensor): + parent_metadata, workflow_root_ancestor_metadata, video_metadata = ( + _parse_image_metadata_fields(parameter=parameter, image=None) + ) + return WorkflowImageData( + parent_metadata=parent_metadata, + workflow_root_ancestor_metadata=workflow_root_ancestor_metadata, + tensor_image=_ensure_chw_layout(image), + video_metadata=video_metadata, + ) + if isinstance(image, dict) and image.get("type") == "tensor": + value = image.get("value") + if not isinstance(value, torch.Tensor): + raise RuntimeInputError( + public_message=( + f"Detected runtime parameter `{parameter}` declared with " + f"type='tensor' but its value is of type {type(value)}; " + "expected torch.Tensor." + ), + context="workflow_execution | runtime_input_validation", + ) + parent_metadata, workflow_root_ancestor_metadata, video_metadata = ( + _parse_image_metadata_fields(parameter=parameter, image=image) + ) + return WorkflowImageData( + parent_metadata=parent_metadata, + workflow_root_ancestor_metadata=workflow_root_ancestor_metadata, + tensor_image=_ensure_chw_layout(value), + video_metadata=video_metadata, + ) + return _deserialize_image_kind_numpy( + parameter=parameter, + image=image, + prevent_local_images_loading=prevent_local_images_loading, + ) + + +def _parse_image_metadata_fields(parameter: str, image: Any): + """Parse the parent/root-parent/video metadata block shared by the + image-kind deserializer paths. Mirrors the prefix of + `deserialize_image_kind` in the numpy file.""" + is_image_dict = isinstance(image, dict) + parent_id = image.get(PARENT_ID_KEY, parameter) if is_image_dict else parameter + parent_origin = image.get(PARENT_ORIGIN_KEY) if is_image_dict else None + parent_metadata = _parse_optional_parent_metadata( + parameter=parameter, + parent_id=parent_id, + parent_origin=parent_origin, + ) + root_parent_id = image.get(ROOT_PARENT_ID_KEY) if is_image_dict else None + root_parent_origin = image.get(ROOT_PARENT_ORIGIN_KEY) if is_image_dict else None + workflow_root_ancestor_metadata = _parse_optional_parent_metadata( + parameter=parameter, + parent_id=root_parent_id, + parent_origin=root_parent_origin, + ) + video_metadata = None + if is_image_dict and "video_metadata" in image: + video_metadata = deserialize_video_metadata_kind( + parameter=parameter, video_metadata=image["video_metadata"] + ) + return parent_metadata, workflow_root_ancestor_metadata, video_metadata + + +def deserialize_detections_kind( + parameter: str, + detections: Any, +) -> Detections: + """Tensor-native sibling of the numpy ``deserialize_detections_kind``. + + The numpy path returns ``sv.Detections``; on the tensor branch every consumer + block expects a native ``inference_models.Detections`` (xyxy/class_id/confidence + tensors on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``, ``image_metadata[CLASS_NAMES_KEY]`` + + per-box ``detection_id``). Builds it by re-using + ``native_detections_from_inference_predictions`` from the serialised inference + prediction dicts (center ``x``/``y``/``width``/``height``). Lineage carried by the + serialised predictions is reconstructed onto the placeholder image so crop-aware + coordinate recovery downstream keeps working. + """ + if isinstance(detections, (Detections, InstanceDetections)): + return detections + raw_predictions = _validate_serialized_detections( + parameter=parameter, detections=detections + ) + return _native_detections_from_serialized( + parameter=parameter, + detections=detections, + raw_predictions=raw_predictions, + prediction_type=DEFAULT_OBJECT_DETECTION_PREDICTION_TYPE, + ) + + +def deserialize_rle_detections_kind( + parameter: str, + detections: Any, +) -> InstanceDetections: + """Tensor-native sibling of the numpy ``deserialize_rle_detections_kind``. + + Builds the base ``Detections`` exactly as ``deserialize_detections_kind`` does, + then rebuilds ``InstancesRLEMasks`` from each serialised prediction's COCO RLE + (``RLE_MASK_KEY_IN_INFERENCE_RESPONSE``) and returns a native + ``InstanceDetections``. Used for the rle-instance-seg and semantic-seg kinds. + """ + if isinstance(detections, (Detections, InstanceDetections)): + return detections + raw_predictions = _validate_serialized_detections( + parameter=parameter, detections=detections + ) + base_detections = _native_detections_from_serialized( + parameter=parameter, + detections=detections, + raw_predictions=raw_predictions, + prediction_type=DEFAULT_INSTANCE_SEGMENTATION_PREDICTION_TYPE, + ) + mask = _rebuild_instances_rle_masks( + detections=detections, raw_predictions=raw_predictions + ) + return InstanceDetections( + xyxy=base_detections.xyxy, + class_id=base_detections.class_id, + confidence=base_detections.confidence, + mask=mask, + image_metadata=base_detections.image_metadata, + bboxes_metadata=base_detections.bboxes_metadata, + ) + + +def _validate_serialized_detections(parameter: str, detections: Any) -> List[dict]: + """Validate the serialised-detections dict shape (mirrors the numpy + deserialiser's checks) and return the list of per-box prediction dicts.""" + if not isinstance(detections, dict): + raise RuntimeInputError( + public_message=f"Detected runtime parameter `{parameter}` declared to hold " + f"detections, but invalid type of data found.", + context="workflow_execution | runtime_input_validation", + ) + if "predictions" not in detections or "image" not in detections: + raise RuntimeInputError( + public_message=f"Detected runtime parameter `{parameter}` declared to hold " + f"detections, but dictionary misses required keys.", + context="workflow_execution | runtime_input_validation", + ) + raw_predictions = detections["predictions"] + if not isinstance(raw_predictions, list): + raise RuntimeInputError( + public_message=f"Detected runtime parameter `{parameter}` declared to hold " + f"detections, but `predictions` is not a list.", + context="workflow_execution | runtime_input_validation", + ) + return raw_predictions + + +def _native_detections_from_serialized( + parameter: str, + detections: dict, + raw_predictions: List[dict], + prediction_type: str, +) -> Detections: + placeholder_image = _build_placeholder_image( + parameter=parameter, + detections=detections, + raw_predictions=raw_predictions, + ) + native_detections = native_detections_from_inference_predictions( + image=placeholder_image, + predictions=raw_predictions, + prediction_type=prediction_type, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + _attach_optional_nearest_target_distance( + raw_predictions=raw_predictions, + bboxes_metadata=native_detections.bboxes_metadata, + ) + return native_detections + + +def _attach_optional_nearest_target_distance( + raw_predictions: List[dict], + bboxes_metadata: Optional[List[dict]], +) -> None: + if not bboxes_metadata or NEAREST_TARGET_DISTANCE_KEY not in raw_predictions[0]: + return + for prediction, entry in zip(raw_predictions, bboxes_metadata): + entry[NEAREST_TARGET_DISTANCE_KEY] = prediction[NEAREST_TARGET_DISTANCE_KEY] + + +def _rebuild_instances_rle_masks( + detections: dict, + raw_predictions: List[dict], +) -> InstancesRLEMasks: + """Rebuild an ``InstancesRLEMasks`` from the serialised per-box COCO RLE entries. + + Each prediction carries ``rle_mask`` as ``{"size": [h, w], "counts": ...}``. + The image size is taken from the RLE entries (falling back to the serialised + ``image`` dimensions) so the rebuilt masks decode against the right canvas. + """ + image_height, image_width = _read_serialized_image_shape(detections) + coco_rle_masks: List[dict] = [] + for prediction in raw_predictions: + rle = prediction.get(RLE_MASK_KEY_IN_INFERENCE_RESPONSE) + if rle is None: + raise RuntimeInputError( + public_message=( + "Detected runtime parameter declared to hold RLE instance " + "segmentation, but a prediction is missing the " + f"`{RLE_MASK_KEY_IN_INFERENCE_RESPONSE}` entry." + ), + context="workflow_execution | runtime_input_validation", + ) + coco_rle_masks.append(rle) + image_size: Tuple[int, int] + if coco_rle_masks and coco_rle_masks[0].get("size") is not None: + size = coco_rle_masks[0]["size"] + image_size = (int(size[0]), int(size[1])) + else: + image_size = (int(image_height), int(image_width)) + return InstancesRLEMasks.from_coco_rle_masks( + image_size=image_size, + masks=coco_rle_masks, + ) + + +def _build_placeholder_image( + parameter: str, + detections: dict, + raw_predictions: List[dict], +) -> WorkflowImageData: + """Build a placeholder ``WorkflowImageData`` carrying the serialised image + dimensions and the parent/root lineage embedded in the predictions. + + ``native_detections_from_inference_predictions`` only reads the image to derive + ``image_metadata`` (dimensions + parent/root lineage), so a zero CHW tensor on + ``WORKFLOWS_IMAGE_TENSOR_DEVICE`` is sufficient and avoids materialising any + real pixels. Lineage is shared across detections of a single image, so it is + read from the first prediction that carries it. + """ + image_height, image_width = _read_serialized_image_shape(detections) + tensor_image = torch.zeros( + (3, int(image_height), int(image_width)), + dtype=torch.uint8, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + parent_metadata, root_parent_metadata = _reconstruct_lineage( + parameter=parameter, + raw_predictions=raw_predictions, + image_height=int(image_height), + image_width=int(image_width), + ) + return WorkflowImageData( + parent_metadata=parent_metadata, + workflow_root_ancestor_metadata=root_parent_metadata, + tensor_image=tensor_image, + ) + + +def _read_serialized_image_shape(detections: dict) -> Tuple[int, int]: + image = detections.get("image") or {} + height = image.get("height") + width = image.get("width") + if height is None or width is None: + # numpy parity: a serialised detections payload always carries image dims; + # fall back to a 1x1 canvas so metadata can still be built for empty inputs. + return 1, 1 + return int(height), int(width) + + +def _reconstruct_lineage( + parameter: str, + raw_predictions: List[dict], + image_height: int, + image_width: int, +) -> Tuple[ImageParentMetadata, ImageParentMetadata]: + parent_id = parameter + root_parent_id = parameter + parent_origin = None + root_parent_origin = None + for prediction in raw_predictions: + if not isinstance(prediction, dict): + continue + parent_id = prediction.get(PARENT_ID_KEY, parent_id) + root_parent_id = prediction.get(ROOT_PARENT_ID_KEY, root_parent_id) + parent_origin = prediction.get(PARENT_ORIGIN_KEY, parent_origin) + root_parent_origin = prediction.get(ROOT_PARENT_ORIGIN_KEY, root_parent_origin) + if parent_origin is not None or root_parent_origin is not None: + break + parent_metadata = ImageParentMetadata( + parent_id=parent_id, + origin_coordinates=_origin_coordinates_from_serialized( + origin=parent_origin, + image_height=image_height, + image_width=image_width, + ), + ) + root_parent_metadata = ImageParentMetadata( + parent_id=root_parent_id, + origin_coordinates=_origin_coordinates_from_serialized( + origin=root_parent_origin, + image_height=image_height, + image_width=image_width, + ), + ) + return parent_metadata, root_parent_metadata + + +def _origin_coordinates_from_serialized( + origin: Optional[dict], + image_height: int, + image_width: int, +) -> OriginCoordinatesSystem: + """Map a serialised ``ParentOrigin`` dict back to an ``OriginCoordinatesSystem``. + + When the prediction carried no origin (top-level image, not a crop), the origin + defaults to the image's own dimensions at offset (0, 0) - matching how the + producer-side ``build_native_image_metadata`` reads a non-crop image's lineage. + """ + if isinstance(origin, dict): + return OriginCoordinatesSystem( + left_top_x=int(origin.get("offset_x", 0)), + left_top_y=int(origin.get("offset_y", 0)), + origin_width=int(origin.get("width", image_width)), + origin_height=int(origin.get("height", image_height)), + ) + return OriginCoordinatesSystem( + left_top_x=0, + left_top_y=0, + origin_width=int(image_width), + origin_height=int(image_height), + ) + + +def deserialize_native_embedding_kind(parameter: str, value: Any) -> torch.Tensor: + """Tensor-native deserialiser for the embedding kind. + + The numpy branch registers no deserialiser for embeddings, so a serialised + embedding (a JSON ``List[float]`` โ€” the inverse of ``serialise_native_embedding``, + which emits ``value.detach().cpu().tolist()``) would reach a tensor consumer as a + plain ``list`` and break on ``.shape`` / ``torch.dot`` (e.g. ``cosine_similarity``). + This rebuilds a 1-D ``torch.Tensor`` on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``. + """ + return _tensor_from_serialized(parameter=parameter, value=value) + + +def deserialize_native_tensor_kind(parameter: str, value: Any) -> torch.Tensor: + """Tensor-native deserialiser for the tensor kind. + + Inverse of ``serialise_native_tensor`` (``value.detach().cpu().tolist()``); rebuilds + a (possibly N-D) ``torch.Tensor`` from the serialised nested-list JSON so tensor + consumers receive a real tensor rather than a nested ``list``. Same tensor device as + the producers (``WORKFLOWS_IMAGE_TENSOR_DEVICE``). + """ + return _tensor_from_serialized(parameter=parameter, value=value) + + +def _tensor_from_serialized(parameter: str, value: Any) -> torch.Tensor: + if isinstance(value, torch.Tensor): + return value.to(WORKFLOWS_IMAGE_TENSOR_DEVICE) + if not isinstance(value, (list, tuple)): + raise RuntimeInputError( + public_message=( + f"Detected runtime parameter `{parameter}` declared to hold an " + f"embedding / tensor value, but found {type(value)}; expected a JSON " + f"list of numbers (optionally nested) or a torch.Tensor." + ), + context="workflow_execution | runtime_input_validation", + ) + try: + return torch.as_tensor( + value, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + except (ValueError, TypeError) as error: + raise RuntimeInputError( + public_message=( + f"Detected runtime parameter `{parameter}` declared to hold an " + f"embedding / tensor value, but its list could not be converted to a " + f"torch.Tensor: {error}" + ), + context="workflow_execution | runtime_input_validation", + ) from error + + +def deserialize_native_classification_prediction_kind( + parameter: str, + value: Any, +) -> Union[ClassificationPrediction, MultiLabelClassificationPrediction]: + """Tensor-native sibling of the numpy ``deserialize_classification_prediction_kind``. + + The numpy path returns the classification dict unchanged; on the tensor branch every + consumer (e.g. the UQL ``*_tensor_native`` classification extractors) expects a native + ``inference_models.ClassificationPrediction`` (single-label) or + ``MultiLabelClassificationPrediction`` (multi-label). This rebuilds the native object + from the serialised classification dict โ€” the same shape + ``serialise_native_classification`` emits โ€” so a round-trip is byte-faithful for a + canonical model-emitted prediction (see the module-level caveat below). + + Single-label (``top``/``confidence`` present): a dense ``confidence`` vector + ``(1, num_classes)`` indexed by ``class_id`` is rebuilt from ``predictions`` and the + top-1 ``class_id`` is recovered from ``top``. Multi-label (``predicted_classes`` + present): a dense ``confidence`` vector ``(num_classes,)`` plus the predicted + ``class_ids`` are rebuilt from the ``predictions`` map. Both carry the + ``class_id -> name`` map and image lineage on the metadata (``CLASS_NAMES_KEY`` etc.) + that the serializer and consumers read. + + BYTE-PARITY CAVEAT: perfect ``serialise(deserialize(x)) == x`` identity only holds for + a canonical model-emitted classification prediction โ€” the full class distribution with + contiguous 0-based ``class_id`` values, already sorted desc and rounded. Owner-2's + ``serialise_native_classification`` re-enumerates the whole confidence vector, re-sorts + desc and re-rounds; a sparse / top-K / out-of-list (``class_id == -1``) / non-contiguous + input cannot be represented losslessly in the native (dense, index==class_id) structure + and will normalise rather than preserve. Flagged in the owner report. + """ + if isinstance( + value, (ClassificationPrediction, MultiLabelClassificationPrediction) + ): + return value + value = _validate_classification_dict(parameter=parameter, value=value) + if "predicted_classes" in value: + return _build_multi_label_prediction(value=value) + return _build_single_label_prediction(value=value) + + +def _validate_classification_dict(parameter: str, value: Any) -> dict: + if not isinstance(value, dict): + raise RuntimeInputError( + public_message=( + f"Detected runtime parameter `{parameter}` declared to hold a " + f"classification prediction, but found {type(value)}; expected a dict." + ), + context="workflow_execution | runtime_input_validation", + ) + if "image" not in value or "predictions" not in value: + raise RuntimeInputError( + public_message=( + f"Detected runtime parameter `{parameter}` declared to hold a " + f"classification prediction, but the dict misses required keys " + f"('image', 'predictions')." + ), + context="workflow_execution | runtime_input_validation", + ) + if "predicted_classes" not in value and ( + "top" not in value or "confidence" not in value + ): + raise RuntimeInputError( + public_message=( + f"Detected runtime parameter `{parameter}` declared to hold a " + f"classification prediction, but the value misses prediction details " + f"(neither 'predicted_classes' nor 'top'/'confidence' present)." + ), + context="workflow_execution | runtime_input_validation", + ) + return value + + +def _dense_confidence_vector( + class_id_to_confidence: Dict[int, float], +) -> List[float]: + """Build a dense confidence list indexed by ``class_id``. + + Length is ``max(class_id) + 1`` for the canonical contiguous 0-based case; class ids + outside ``[0, length)`` (e.g. an out-of-list ``-1``) cannot be positioned and are + dropped from the dense vector (their name still lives in ``CLASS_NAMES_KEY``). + """ + if not class_id_to_confidence: + return [] + highest_class_id = max(class_id_to_confidence) + length = highest_class_id + 1 if highest_class_id >= 0 else 0 + vector = [0.0] * length + for class_id, confidence in class_id_to_confidence.items(): + if 0 <= class_id < length: + vector[class_id] = confidence + return vector + + +def _build_classification_image_metadata( + value: dict, + class_names_mapping: Dict[int, str], +) -> dict: + """Assemble the per-image ``image_metadata`` that the serializer and the UQL + ``*_tensor_native`` classification extractors read back.""" + metadata: dict = {CLASS_NAMES_KEY: class_names_mapping} + # Lane 1b: deserialised predictions follow the canonical model contract (decision + # 1a), so tag them "model" rather than relying on which incidental keys survived. + metadata[CLASSIFICATION_STYLE_KEY] = CLASSIFICATION_STYLE_MODEL + image = value.get("image") or {} + height, width = image.get("height"), image.get("width") + if height is not None and width is not None: + metadata[IMAGE_DIMENSIONS_KEY] = [int(height), int(width)] + for key in ( + INFERENCE_ID_KEY, + "time", + PREDICTION_TYPE_KEY, + PARENT_ID_KEY, + ROOT_PARENT_ID_KEY, + ): + if value.get(key) is not None: + metadata[key] = value[key] + return metadata + + +def _build_single_label_prediction(value: dict) -> ClassificationPrediction: + predictions = value["predictions"] + if not isinstance(predictions, list): + raise RuntimeInputError( + public_message=( + "Detected a single-label classification prediction (with 'top'), but " + "'predictions' is not a list of per-class entries." + ), + context="workflow_execution | runtime_input_validation", + ) + class_names_mapping: Dict[int, str] = {} + class_id_to_confidence: Dict[int, float] = {} + for entry in predictions: + class_id = int(entry["class_id"]) + class_names_mapping[class_id] = str(entry[CLASS_NAME_KEY]) + class_id_to_confidence[class_id] = float(entry["confidence"]) + top_name = value.get("top") + top_class_id = next( + ( + class_id + for class_id, name in class_names_mapping.items() + if name == top_name + ), + None, + ) + if top_class_id is None: + top_class_id = ( + max(class_id_to_confidence, key=class_id_to_confidence.get) + if class_id_to_confidence + else 0 + ) + confidence_vector = _dense_confidence_vector(class_id_to_confidence) + image_metadata = _build_classification_image_metadata( + value=value, class_names_mapping=class_names_mapping + ) + return ClassificationPrediction( + class_id=torch.tensor( + [top_class_id], dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.tensor( + [confidence_vector], + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + images_metadata=[image_metadata], + ) + + +def _build_multi_label_prediction(value: dict) -> MultiLabelClassificationPrediction: + predictions = value["predictions"] + if not isinstance(predictions, dict): + raise RuntimeInputError( + public_message=( + "Detected a multi-label classification prediction (with " + "'predicted_classes'), but 'predictions' is not a name->entry map." + ), + context="workflow_execution | runtime_input_validation", + ) + class_names_mapping: Dict[int, str] = {} + class_id_to_confidence: Dict[int, float] = {} + name_to_class_id: Dict[str, int] = {} + for name, entry in predictions.items(): + class_id = int(entry["class_id"]) + class_names_mapping[class_id] = str(name) + class_id_to_confidence[class_id] = float(entry["confidence"]) + name_to_class_id[str(name)] = class_id + predicted_class_ids = [ + name_to_class_id[str(name)] + for name in value.get("predicted_classes", []) + if str(name) in name_to_class_id + ] + confidence_vector = _dense_confidence_vector(class_id_to_confidence) + image_metadata = _build_classification_image_metadata( + value=value, class_names_mapping=class_names_mapping + ) + return MultiLabelClassificationPrediction( + class_ids=torch.tensor( + predicted_class_ids, dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.tensor( + confidence_vector, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + image_metadata=image_metadata, + ) + + +__all__ = [ + "deserialize_image_kind", + "deserialize_detections_kind", + "deserialize_rle_detections_kind", + "deserialize_native_classification_prediction_kind", + "deserialize_native_embedding_kind", + "deserialize_native_tensor_kind", +] diff --git a/inference/core/workflows/core_steps/common/query_language/operations/classification_results/base.py b/inference/core/workflows/core_steps/common/query_language/operations/classification_results/base.py index bee2f8e3fe..65efe56fb8 100644 --- a/inference/core/workflows/core_steps/common/query_language/operations/classification_results/base.py +++ b/inference/core/workflows/core_steps/common/query_language/operations/classification_results/base.py @@ -1,14 +1,21 @@ from typing import Any, List, Union +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.common.query_language.entities.enums import ( ClassificationProperty, ) from inference.core.workflows.core_steps.common.query_language.errors import ( InvalidInputTypeError, + OperationError, ) from inference.core.workflows.core_steps.common.query_language.operations.utils import ( safe_stringify, ) +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY +from inference_models import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) def extract_top_class(prediction: dict) -> Union[str, List[str]]: @@ -17,6 +24,104 @@ def extract_top_class(prediction: dict) -> Union[str, List[str]]: return prediction.get("predicted_classes", []) +def extract_top_class_tensor_native( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> Union[str, List[str]]: + if isinstance(prediction, ClassificationPrediction): + if prediction.images_metadata is None: + raise OperationError( + public_message=( + "Executing extract_top_class_tensor_native(...) on " + "`inference_models.ClassificationPrediction`, but `images_metadata` " + "is missing โ€” class-id-to-name lookup requires the producer block " + "to attach per-image metadata." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + if prediction.class_id.shape[0] != 1: + raise OperationError( + public_message=( + "Executing extract_top_class_tensor_native(...) on " + "`inference_models.ClassificationPrediction` with batch size " + f"{prediction.class_id.shape[0]} โ€” expected a single-image slice. " + "Batch must be unpacked before invoking this extractor." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + class_names = prediction.images_metadata[0].get(CLASS_NAMES_KEY) + if class_names is None: + raise OperationError( + public_message=( + "Executing extract_top_class_tensor_native(...) on " + "`inference_models.ClassificationPrediction`, but " + f"`images_metadata[0]['{CLASS_NAMES_KEY}']` is missing โ€” the " + "producer block must attach the class_id โ†’ name mapping." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + class_id = int(prediction.class_id[0]) + class_name = class_names.get(class_id) + if class_name is None: + raise OperationError( + public_message=( + "Executing extract_top_class_tensor_native(...) on " + "`inference_models.ClassificationPrediction`, predicted " + f"class_id={class_id} is missing from the class_names mapping " + f"(keys present: {sorted(class_names.keys())})." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + return class_name + if isinstance(prediction, MultiLabelClassificationPrediction): + if prediction.image_metadata is None: + raise OperationError( + public_message=( + "Executing extract_top_class_tensor_native(...) on " + "`inference_models.MultiLabelClassificationPrediction`, but " + "`image_metadata` is missing โ€” class-id-to-name lookup requires " + "the producer block to attach metadata." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + class_names = prediction.image_metadata.get(CLASS_NAMES_KEY) + if class_names is None: + raise OperationError( + public_message=( + "Executing extract_top_class_tensor_native(...) on " + "`inference_models.MultiLabelClassificationPrediction`, but " + f"`image_metadata['{CLASS_NAMES_KEY}']` is missing โ€” the producer " + "block must attach the class_id โ†’ name mapping." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + result: List[str] = [] + for class_id_scalar in prediction.class_ids.tolist(): + class_id = int(class_id_scalar) + class_name = class_names.get(class_id) + if class_name is None: + raise OperationError( + public_message=( + "Executing extract_top_class_tensor_native(...) on " + "`inference_models.MultiLabelClassificationPrediction`, " + f"predicted class_id={class_id} is missing from the " + f"class_names mapping (keys present: " + f"{sorted(class_names.keys())})." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + result.append(class_name) + return result + raise InvalidInputTypeError( + public_message=( + "While executing extract_top_class_tensor_native(...) operation it was " + "expected to get `inference_models.ClassificationPrediction` or " + "`inference_models.MultiLabelClassificationPrediction`, but got instance " + f"of type {type(prediction)}" + ), + context="step_execution | roboflow_query_language_evaluation", + ) + + def extract_top_class_confidence(prediction: dict) -> Union[float, List[float]]: if "confidence" in prediction: return prediction["confidence"] @@ -27,6 +132,35 @@ def extract_top_class_confidence(prediction: dict) -> Union[float, List[float]]: ] +def extract_top_class_confidence_tensor_native( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> Union[float, List[float]]: + if isinstance(prediction, ClassificationPrediction): + if prediction.class_id.shape[0] != 1: + raise OperationError( + public_message=( + "Executing extract_top_class_confidence_tensor_native(...) on " + "`inference_models.ClassificationPrediction` with batch size " + f"{prediction.class_id.shape[0]} โ€” expected a single-image slice. " + "Batch must be unpacked before invoking this extractor." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + class_id = int(prediction.class_id[0]) + return float(prediction.confidence[0, class_id]) + if isinstance(prediction, MultiLabelClassificationPrediction): + return [float(c) for c in prediction.confidence[prediction.class_ids].tolist()] + raise InvalidInputTypeError( + public_message=( + "While executing extract_top_class_confidence_tensor_native(...) operation " + "it was expected to get `inference_models.ClassificationPrediction` or " + "`inference_models.MultiLabelClassificationPrediction`, but got instance " + f"of type {type(prediction)}" + ), + context="step_execution | roboflow_query_language_evaluation", + ) + + def extract_top_class_confidence_single(prediction: dict) -> Union[float, List[float]]: if "confidence" in prediction: return prediction["confidence"] @@ -40,6 +174,37 @@ def extract_top_class_confidence_single(prediction: dict) -> Union[float, List[f return max(predicted_confidences) +def extract_top_class_confidence_single_tensor_native( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> float: + if isinstance(prediction, ClassificationPrediction): + if prediction.class_id.shape[0] != 1: + raise OperationError( + public_message=( + "Executing extract_top_class_confidence_single_tensor_native(...) on " + "`inference_models.ClassificationPrediction` with batch size " + f"{prediction.class_id.shape[0]} โ€” expected a single-image slice. " + "Batch must be unpacked before invoking this extractor." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + class_id = int(prediction.class_id[0]) + return float(prediction.confidence[0, class_id]) + if isinstance(prediction, MultiLabelClassificationPrediction): + if prediction.class_ids.shape[0] == 0: + return 0.0 + return float(prediction.confidence[prediction.class_ids].max()) + raise InvalidInputTypeError( + public_message=( + "While executing extract_top_class_confidence_single_tensor_native(...) " + "operation it was expected to get `inference_models.ClassificationPrediction` " + "or `inference_models.MultiLabelClassificationPrediction`, but got instance " + f"of type {type(prediction)}" + ), + context="step_execution | roboflow_query_language_evaluation", + ) + + def extract_all_class_names(prediction: dict) -> List[str]: predictions = prediction["predictions"] if isinstance(predictions, list): @@ -51,6 +216,67 @@ def extract_all_class_names(prediction: dict) -> List[str]: return [class_id2_class_name[class_id] for class_id in sorted_ids] +def extract_all_class_names_tensor_native( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> List[str]: + if isinstance(prediction, ClassificationPrediction): + if prediction.images_metadata is None: + raise OperationError( + public_message=( + "Executing extract_all_class_names_tensor_native(...) on " + "`inference_models.ClassificationPrediction`, but `images_metadata` " + "is missing โ€” class-id-to-name lookup requires the producer block " + "to attach per-image metadata." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + if prediction.class_id.shape[0] != 1: + raise OperationError( + public_message=( + "Executing extract_all_class_names_tensor_native(...) on " + "`inference_models.ClassificationPrediction` with batch size " + f"{prediction.class_id.shape[0]} โ€” expected a single-image slice. " + "Batch must be unpacked before invoking this extractor." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + class_names = prediction.images_metadata[0].get(CLASS_NAMES_KEY) + carrier_path = f"images_metadata[0]['{CLASS_NAMES_KEY}']" + elif isinstance(prediction, MultiLabelClassificationPrediction): + if prediction.image_metadata is None: + raise OperationError( + public_message=( + "Executing extract_all_class_names_tensor_native(...) on " + "`inference_models.MultiLabelClassificationPrediction`, but " + "`image_metadata` is missing โ€” class-id-to-name lookup requires " + "the producer block to attach metadata." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + class_names = prediction.image_metadata.get(CLASS_NAMES_KEY) + carrier_path = f"image_metadata['{CLASS_NAMES_KEY}']" + else: + raise InvalidInputTypeError( + public_message=( + "While executing extract_all_class_names_tensor_native(...) operation " + "it was expected to get `inference_models.ClassificationPrediction` or " + "`inference_models.MultiLabelClassificationPrediction`, but got " + f"instance of type {type(prediction)}" + ), + context="step_execution | roboflow_query_language_evaluation", + ) + if class_names is None: + raise OperationError( + public_message=( + f"Executing extract_all_class_names_tensor_native(...), but " + f"`{carrier_path}` is missing โ€” the producer block must attach the " + "class_id โ†’ name mapping." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + return [class_names[class_id] for class_id in sorted(class_names)] + + def extract_all_classes_confidence(prediction: dict) -> List[float]: predictions = prediction["predictions"] if isinstance(predictions, list): @@ -62,19 +288,84 @@ def extract_all_classes_confidence(prediction: dict) -> List[float]: return [class_id2_class_confidence[class_id] for class_id in sorted_ids] +def extract_all_classes_confidence_tensor_native( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> List[float]: + if isinstance(prediction, ClassificationPrediction): + if prediction.confidence.shape[0] != 1: + raise OperationError( + public_message=( + "Executing extract_all_classes_confidence_tensor_native(...) on " + "`inference_models.ClassificationPrediction` with batch size " + f"{prediction.confidence.shape[0]} โ€” expected a single-image slice. " + "Batch must be unpacked before invoking this extractor." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + return [float(c) for c in prediction.confidence[0].tolist()] + if isinstance(prediction, MultiLabelClassificationPrediction): + return [float(c) for c in prediction.confidence.tolist()] + raise InvalidInputTypeError( + public_message=( + "While executing extract_all_classes_confidence_tensor_native(...) operation " + "it was expected to get `inference_models.ClassificationPrediction` or " + "`inference_models.MultiLabelClassificationPrediction`, but got instance " + f"of type {type(prediction)}" + ), + context="step_execution | roboflow_query_language_evaluation", + ) + + CLASSIFICATION_PROPERTY_EXTRACTORS = { - ClassificationProperty.TOP_CLASS: extract_top_class, - ClassificationProperty.TOP_CLASS_CONFIDENCE: extract_top_class_confidence, - ClassificationProperty.TOP_CLASS_CONFIDENCE_SINGLE: extract_top_class_confidence_single, - ClassificationProperty.ALL_CLASSES: extract_all_class_names, - ClassificationProperty.ALL_CONFIDENCES: extract_all_classes_confidence, + ClassificationProperty.TOP_CLASS: ( + extract_top_class + if not ENABLE_TENSOR_DATA_REPRESENTATION + else extract_top_class_tensor_native + ), + ClassificationProperty.TOP_CLASS_CONFIDENCE: ( + extract_top_class_confidence + if not ENABLE_TENSOR_DATA_REPRESENTATION + else extract_top_class_confidence_tensor_native + ), + ClassificationProperty.TOP_CLASS_CONFIDENCE_SINGLE: ( + extract_top_class_confidence_single + if not ENABLE_TENSOR_DATA_REPRESENTATION + else extract_top_class_confidence_single_tensor_native + ), + ClassificationProperty.ALL_CLASSES: ( + extract_all_class_names + if not ENABLE_TENSOR_DATA_REPRESENTATION + else extract_all_class_names_tensor_native + ), + ClassificationProperty.ALL_CONFIDENCES: ( + extract_all_classes_confidence + if not ENABLE_TENSOR_DATA_REPRESENTATION + else extract_all_classes_confidence_tensor_native + ), } def extract_classification_property( value: Any, property_name: ClassificationProperty, **kwargs ) -> Union[str, float, list]: - if not isinstance(value, dict): + if ENABLE_TENSOR_DATA_REPRESENTATION: + # Native-only under the flag: the *_tensor_native extractors operate on + # `inference_models` prediction dataclasses, not on serialised dicts. + if not isinstance( + value, (ClassificationPrediction, MultiLabelClassificationPrediction) + ): + value_as_str = safe_stringify(value=value) + raise InvalidInputTypeError( + public_message=( + "Executing extract_classification_property(...) under " + "ENABLE_TENSOR_DATA_REPRESENTATION, expected " + "`inference_models.ClassificationPrediction` or " + "`inference_models.MultiLabelClassificationPrediction`, got " + f"{value_as_str} of type {type(value)}" + ), + context="step_execution | roboflow_query_language_evaluation", + ) + elif not isinstance(value, dict): value_as_str = safe_stringify(value=value) raise InvalidInputTypeError( public_message=f"Executing extract_classification_property(...), expected classification results object, " diff --git a/inference/core/workflows/core_steps/common/query_language/operations/detection/base.py b/inference/core/workflows/core_steps/common/query_language/operations/detection/base.py index 3fc75a4cd3..763d9e3c94 100644 --- a/inference/core/workflows/core_steps/common/query_language/operations/detection/base.py +++ b/inference/core/workflows/core_steps/common/query_language/operations/detection/base.py @@ -1,5 +1,6 @@ -from typing import Any +from typing import Any, Optional +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.analytics.line_counter.v2 import ( DETECTIONS_IN_OUT_PARAM, ) @@ -14,6 +15,7 @@ ) from inference.core.workflows.core_steps.common.query_language.errors import ( InvalidInputTypeError, + OperationError, ) from inference.core.workflows.core_steps.common.query_language.operations.utils import ( safe_stringify, @@ -25,6 +27,7 @@ BOUNDING_RECT_HEIGHT_KEY_IN_SV_DETECTIONS, BOUNDING_RECT_RECT_KEY_IN_SV_DETECTIONS, BOUNDING_RECT_WIDTH_KEY_IN_SV_DETECTIONS, + CLASS_NAMES_KEY, IMAGE_DIMENSIONS_KEY, KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, PATH_DEVIATION_KEY_IN_SV_DETECTIONS, @@ -93,7 +96,7 @@ } -def extract_detection_property( +def _extract_detection_property( value: Any, property_name: DetectionsProperty, execution_context: str, @@ -108,3 +111,143 @@ def extract_detection_property( context=f"step_execution | roboflow_query_language_evaluation | {execution_context}", ) return DETECTION_PROPERTY_EXTRACTION[property_name](value) + + +def _extract_class_name_tensor_native(detection: tuple) -> str: + class_names = detection[6].get(CLASS_NAMES_KEY) + if class_names is None: + raise OperationError( + public_message=( + "Executing extract_detection_property_tensor_native(...) for property " + f"`class_name`, but `metadata['{CLASS_NAMES_KEY}']` is missing โ€” the " + "producer block must attach the class_id โ†’ name mapping." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + if detection[2] is None: + raise OperationError( + public_message=( + "Executing extract_detection_property_tensor_native(...) for property " + "`class_name`, but the detection carries no class_id." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + class_id = int(detection[2]) + class_name = class_names.get(class_id) + if class_name is None: + raise OperationError( + public_message=( + "Executing extract_detection_property_tensor_native(...) for property " + f"`class_name`, class_id={class_id} is missing from the class_names " + f"mapping (keys present: {sorted(class_names.keys())})." + ), + context="step_execution | roboflow_query_language_evaluation", + ) + return class_name + + +def _extract_tracker_id_tensor_native(detection: tuple) -> Optional[int]: + if detection[4] is not None: + return int(detection[4]) + tracker_id = detection[5].get("tracker_id") + if tracker_id is None: + return None + return int(tracker_id) + + +DETECTION_PROPERTY_EXTRACTION_TENSOR_NATIVE = { + DetectionsProperty.X_MIN: lambda x: float(x[0][0]), + DetectionsProperty.Y_MIN: lambda x: float(x[0][1]), + DetectionsProperty.X_MAX: lambda x: float(x[0][2]), + DetectionsProperty.Y_MAX: lambda x: float(x[0][3]), + DetectionsProperty.CONFIDENCE: lambda x: (None if x[3] is None else float(x[3])), + DetectionsProperty.CLASS_ID: lambda x: None if x[2] is None else int(x[2]), + DetectionsProperty.CLASS_NAME: _extract_class_name_tensor_native, + DetectionsProperty.SIZE: lambda x: float((x[0][3] - x[0][1]) * (x[0][2] - x[0][0])), + DetectionsProperty.CENTER: lambda x: ( + float(x[0][0] + (x[0][2] - x[0][0]) / 2), + float(x[0][1] + (x[0][3] - x[0][1]) / 2), + ), + DetectionsProperty.TOP_LEFT: lambda x: (float(x[0][0]), float(x[0][1])), + DetectionsProperty.TOP_RIGHT: lambda x: (float(x[0][2]), float(x[0][1])), + DetectionsProperty.BOTTOM_LEFT: lambda x: (float(x[0][0]), float(x[0][3])), + DetectionsProperty.BOTTOM_RIGHT: lambda x: (float(x[0][2]), float(x[0][3])), + DetectionsProperty.IN_OUT: lambda x: x[5].get(DETECTIONS_IN_OUT_PARAM), + DetectionsProperty.PATH_DEVIATION: lambda x: x[5].get( + PATH_DEVIATION_KEY_IN_SV_DETECTIONS + ), + DetectionsProperty.POLYGON: lambda x: x[5].get(POLYGON_KEY_IN_SV_DETECTIONS), + DetectionsProperty.TIME_IN_ZONE: lambda x: x[5].get( + TIME_IN_ZONE_KEY_IN_SV_DETECTIONS + ), + DetectionsProperty.TRACKER_ID: _extract_tracker_id_tensor_native, + DetectionsProperty.VELOCITY: lambda x: x[5].get(VELOCITY_KEY_IN_SV_DETECTIONS), + DetectionsProperty.SPEED: lambda x: x[5].get(SPEED_KEY_IN_SV_DETECTIONS), + DetectionsProperty.SMOOTHED_VELOCITY: lambda x: x[5].get( + SMOOTHED_VELOCITY_KEY_IN_SV_DETECTIONS + ), + DetectionsProperty.SMOOTHED_SPEED: lambda x: x[5].get( + SMOOTHED_SPEED_KEY_IN_SV_DETECTIONS + ), + DetectionsProperty.DIMENSIONS: lambda x: x[6].get(IMAGE_DIMENSIONS_KEY), + DetectionsProperty.PREDICTION_TYPE: lambda x: x[6].get(PREDICTION_TYPE_KEY), + DetectionsProperty.KEYPOINTS_XY: lambda x: x[5].get( + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS + ), + DetectionsProperty.BOUNDING_RECT: lambda x: x[5].get( + BOUNDING_RECT_RECT_KEY_IN_SV_DETECTIONS + ), + DetectionsProperty.BOUNDING_RECT_WIDTH: lambda x: x[5].get( + BOUNDING_RECT_WIDTH_KEY_IN_SV_DETECTIONS + ), + DetectionsProperty.BOUNDING_RECT_HEIGHT: lambda x: x[5].get( + BOUNDING_RECT_HEIGHT_KEY_IN_SV_DETECTIONS + ), + DetectionsProperty.BOUNDING_RECT_ANGLE: lambda x: x[5].get( + BOUNDING_RECT_ANGLE_KEY_IN_SV_DETECTIONS + ), + DetectionsProperty.AREA: lambda x: x[5].get(AREA_KEY_IN_SV_DETECTIONS), + DetectionsProperty.AREA_CONVERTED: lambda x: x[5].get( + AREA_CONVERTED_KEY_IN_SV_DETECTIONS + ), +} + + +def _extract_detection_property_tensor_native( + value: Any, + property_name: DetectionsProperty, + execution_context: str, + **kwargs, +) -> Any: + if not isinstance(value, tuple) or len(value) != 7: + value_as_str = safe_stringify(value=value) + raise InvalidInputTypeError( + public_message=( + "While executing extract_detection_property_tensor_native(...) " + f"operation in context {execution_context} it was expected to get " + "a 7-element tuple (xyxy, mask, class_id, confidence, tracker_id, " + "data, metadata) representing a single tensor-native detection. " + f"Got value: {value_as_str} of type {type(value)}" + ), + context=f"step_execution | roboflow_query_language_evaluation | {execution_context}", + ) + return DETECTION_PROPERTY_EXTRACTION_TENSOR_NATIVE[property_name](value) + + +def extract_detection_property( + value: Any, + property_name: DetectionsProperty, + execution_context: str, + **kwargs, +) -> Any: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _extract_detection_property_tensor_native( + value=value, + property_name=property_name, + execution_context=execution_context, + ) + return _extract_detection_property( + value=value, + property_name=property_name, + execution_context=execution_context, + ) diff --git a/inference/core/workflows/core_steps/common/query_language/operations/detections/base.py b/inference/core/workflows/core_steps/common/query_language/operations/detections/base.py index b305a636b1..d97f9ac2a8 100644 --- a/inference/core/workflows/core_steps/common/query_language/operations/detections/base.py +++ b/inference/core/workflows/core_steps/common/query_language/operations/detections/base.py @@ -3,8 +3,10 @@ import numpy as np import supervision as sv +import torch from supervision import Position +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.common.query_language.entities.enums import ( DetectionsProperty, DetectionsSelectionMode, @@ -24,6 +26,20 @@ from inference.core.workflows.core_steps.common.serializers import ( serialise_sv_detections, ) +from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_sv_detections as serialise_tensor_native_detections, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + strip_host_mirror_metadata, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks def detections_anchor_coordinates( @@ -63,7 +79,7 @@ def detections_anchor_coordinates( } -def extract_detections_property( +def _extract_detections_property( detections: Any, property_name: DetectionsProperty, execution_context: str, @@ -84,7 +100,7 @@ def extract_detections_property( return PROPERTIES_EXTRACTORS[property_name](detections) -def filter_detections( +def _filter_detections( detections: Any, filtering_fun: Callable[[Dict[str, Any]], bool], global_parameters: Dict[str, Any], @@ -105,7 +121,7 @@ def filter_detections( return detections[result] -def offset_detections( +def _offset_detections( value: Any, offset_x: int, offset_y: int, **kwargs ) -> sv.Detections: if not isinstance(value, sv.Detections): @@ -120,7 +136,9 @@ def offset_detections( return detections_copy -def shift_detections(value: Any, shift_x: int, shift_y: int, **kwargs) -> sv.Detections: +def _shift_detections( + value: Any, shift_x: int, shift_y: int, **kwargs +) -> sv.Detections: if not isinstance(value, sv.Detections): value_as_str = safe_stringify(value=value) raise InvalidInputTypeError( @@ -181,7 +199,7 @@ def select_last_detection(detections: sv.Detections) -> sv.Detections: } -def select_detections( +def _select_detections( value: Any, mode: DetectionsSelectionMode, **kwargs ) -> sv.Detections: if not isinstance(value, sv.Detections): @@ -220,7 +238,7 @@ def extract_y_coordinate_of_detections_center(detections: sv.Detections) -> np.n } -def sort_detections( +def _sort_detections( value: Any, mode: DetectionsSortProperties, ascending: bool, **kwargs ) -> sv.Detections: if not isinstance(value, sv.Detections): @@ -248,7 +266,7 @@ def sort_detections( return value[sorted_indices] -def rename_detections( +def _rename_detections( detections: Any, class_map: Union[Dict[str, str], str], strict: Union[bool, str], @@ -372,7 +390,7 @@ def _build_non_strict_class_to_id_mapping( return original_mapping -def detections_to_dictionary( +def _detections_to_dictionary( detections: Any, execution_context: str, **kwargs, @@ -395,7 +413,7 @@ def detections_to_dictionary( ) -def pick_detections_by_parent_class( +def _pick_detections_by_parent_class( detections: Any, parent_class: str, execution_context: str, @@ -409,7 +427,7 @@ def pick_detections_by_parent_class( context=f"step_execution | roboflow_query_language_evaluation | {execution_context}", ) try: - return _pick_detections_by_parent_class( + return _pick_detections_by_parent_class_impl( detections=detections, parent_class=parent_class ) except Exception as error: @@ -421,7 +439,7 @@ def pick_detections_by_parent_class( ) -def _pick_detections_by_parent_class( +def _pick_detections_by_parent_class_impl( detections: sv.Detections, parent_class: str, ) -> sv.Detections: @@ -453,3 +471,1176 @@ def _is_point_within_box(point: np.ndarray, box: np.ndarray) -> bool: px, py = point x1, y1, x2, y2 = box return x1 <= px <= x2 and y1 <= py <= y2 + + +TensorNativeDetections = Union[Detections, InstanceDetections] +TENSOR_NATIVE_DETECTIONS_TYPES = (Detections, InstanceDetections) +# A keypoint-detection prediction kind is carried as a 2-tuple +# (KeyPoints, Optional[Detections]); the bounding-box component is operated on by +# the UQL detections ops and the KeyPoints component is sliced / shifted to match. +KeyPointPrediction = tuple +TensorNativePrediction = Union[Detections, InstanceDetections, "KeyPointPrediction"] + + +def _is_key_point_prediction(value: Any) -> bool: + """True for the keypoint-detection tuple ``(KeyPoints, Optional[Detections])``.""" + return ( + isinstance(value, tuple) + and len(value) == 2 + and isinstance(value[0], KeyPoints) + and (value[1] is None or isinstance(value[1], TENSOR_NATIVE_DETECTIONS_TYPES)) + ) + + +def _split_key_point_prediction( + value: tuple, + operation_name: str, + execution_context: Optional[str] = None, +) -> tuple: + """Return ``(key_points, detections)`` from a keypoint-detection tuple, raising + if the bbox component is missing (the UQL ops operate on bounding boxes).""" + key_points, detections = value + if detections is None: + context = "step_execution | roboflow_query_language_evaluation" + in_context = "" + if execution_context is not None: + context = f"{context} | {execution_context}" + in_context = f" in context {execution_context}" + raise InvalidInputTypeError( + public_message=f"Executing {operation_name}(...){in_context}, the keypoint " + f"prediction is missing the bounding-box `inference_models.Detections` " + f"component required by this operation.", + context=context, + ) + return key_points, detections + + +def _apply_index_op_on_prediction( + value: TensorNativePrediction, + operation_name: str, + index_fn: Callable[[TensorNativeDetections], List[int]], + execution_context: Optional[str] = None, +) -> TensorNativePrediction: + """Apply an index-producing op over the bounding-box component of a prediction, + handling the keypoint-tuple vs plain-``Detections`` split/re-wrap in one place. + + ``index_fn`` receives the bounding-box ``Detections``/``InstanceDetections`` and + returns the absolute index list to keep. For a keypoint-detection tuple + ``(KeyPoints, Detections)`` the same index list slices both components and the + result is re-wrapped as a tuple; for a plain detections object only the boxes are + sliced. + """ + if _is_key_point_prediction(value): + key_points, bboxes = _split_key_point_prediction( + value, + operation_name=operation_name, + execution_context=execution_context, + ) + indices = index_fn(bboxes) + return _take_key_points(key_points, indices), _take_detections(bboxes, indices) + indices = index_fn(value) + return _take_detections(value, indices) + + +def _ensure_tensor_native_detections( + value: Any, + operation_name: str, + execution_context: Optional[str] = None, +) -> None: + if isinstance(value, TENSOR_NATIVE_DETECTIONS_TYPES) or _is_key_point_prediction( + value + ): + return + value_as_str = safe_stringify(value=value) + context = "step_execution | roboflow_query_language_evaluation" + in_context = "" + if execution_context is not None: + context = f"{context} | {execution_context}" + in_context = f" in context {execution_context}" + raise InvalidInputTypeError( + public_message=f"Executing {operation_name}(...){in_context}, expected " + f"`inference_models.Detections`, `inference_models.InstanceDetections` or a " + f"keypoint-detection `(KeyPoints, Detections)` tuple as value, got " + f"{value_as_str} of type {type(value)}", + context=context, + ) + + +def _take_key_points(key_points: KeyPoints, indices: List[int]) -> KeyPoints: + """Slice a ``KeyPoints`` along the instance dimension by index list, carrying + per-instance ``key_points_metadata`` and sharing ``image_metadata`` as-is.""" + index_tensor = torch.as_tensor( + indices, dtype=torch.long, device=key_points.xy.device + ) + key_points_metadata = None + if key_points.key_points_metadata is not None: + key_points_metadata = [key_points.key_points_metadata[i] for i in indices] + return KeyPoints( + xy=key_points.xy[index_tensor], + class_id=key_points.class_id[index_tensor], + confidence=key_points.confidence[index_tensor], + image_metadata=key_points.image_metadata, + key_points_metadata=key_points_metadata, + ) + + +def _detections_count(detections: TensorNativeDetections) -> int: + return int(detections.xyxy.shape[0]) + + +def _bboxes_metadata_list(detections: TensorNativeDetections) -> List[dict]: + if detections.bboxes_metadata is not None: + return detections.bboxes_metadata + return [{} for _ in range(_detections_count(detections))] + + +def _class_names_lookup( + detections: TensorNativeDetections, operation_name: str +) -> Dict[int, str]: + image_metadata = detections.image_metadata or {} + class_names = image_metadata.get(CLASS_NAMES_KEY) + if class_names is None: + raise OperationError( + public_message=f"Executing {operation_name}(...), but " + f"`image_metadata['{CLASS_NAMES_KEY}']` is missing โ€” the producer block " + f"must attach the class_id โ†’ name mapping.", + context="step_execution | roboflow_query_language_evaluation", + ) + return class_names + + +def _extract_class_names_property_tensor_native( + detections: TensorNativeDetections, +) -> List[str]: + """``class_name`` extractor for ``extract_detections_property``. + + Numpy parity: the numpy sibling reads ``detections.data.get("class_name", [])``, + so an absent ``image_metadata[CLASS_NAMES_KEY]`` mapping yields ``[]`` rather + than raising. A present-but-incomplete mapping still raises in + ``_resolve_class_names`` (a class_id without a name is a producer contract + violation, not an absent property). + """ + if _detections_count(detections) == 0: + return [] + image_metadata = detections.image_metadata or {} + if image_metadata.get(CLASS_NAMES_KEY) is None: + return [] + return _resolve_class_names( + detections, operation_name="extract_detections_property" + ) + + +def _resolve_class_names( + detections: TensorNativeDetections, operation_name: str +) -> List[str]: + if _detections_count(detections) == 0: + return [] + class_names = _class_names_lookup(detections, operation_name=operation_name) + result = [] + for class_id_scalar in detections.class_id.tolist(): + class_id = int(class_id_scalar) + class_name = class_names.get(class_id) + if class_name is None: + raise OperationError( + public_message=f"Executing {operation_name}(...), class_id={class_id} " + f"is missing from the class_names mapping " + f"(keys present: {sorted(class_names.keys())}).", + context="step_execution | roboflow_query_language_evaluation", + ) + result.append(class_name) + return result + + +def _resolve_effective_class_names( + detections: TensorNativeDetections, operation_name: str +) -> List[str]: + """Per-row class names โ€” the native analog of numpy's per-row + ``data["class_name"]``: the per-box ``CLASS_NAME_KEY`` override wins (the + serializer's precedence), with fallback to the image-level ``CLASS_NAMES_KEY`` + map. A row with an override never touches the map, so a missing/incomplete + map raises only for rows that actually need it.""" + if _detections_count(detections) == 0: + return [] + bboxes_metadata = detections.bboxes_metadata + class_names_map: Optional[Dict[int, str]] = None + result = [] + for index, class_id_scalar in enumerate(detections.class_id.tolist()): + entry = {} + if bboxes_metadata is not None: + entry = bboxes_metadata[index] or {} + if CLASS_NAME_KEY in entry: + result.append(str(entry[CLASS_NAME_KEY])) + continue + if class_names_map is None: + class_names_map = _class_names_lookup( + detections, operation_name=operation_name + ) + class_id = int(class_id_scalar) + class_name = class_names_map.get(class_id) + if class_name is None: + raise OperationError( + public_message=f"Executing {operation_name}(...), class_id={class_id} " + f"is missing from the class_names mapping " + f"(keys present: {sorted(class_names_map.keys())}).", + context="step_execution | roboflow_query_language_evaluation", + ) + result.append(class_name) + return result + + +def _take_mask( + mask: Union[torch.Tensor, InstancesRLEMasks], indices: List[int] +) -> Union[torch.Tensor, InstancesRLEMasks]: + if isinstance(mask, InstancesRLEMasks): + return InstancesRLEMasks( + image_size=mask.image_size, + masks=[mask.masks[index] for index in indices], + ) + return mask[torch.as_tensor(indices, dtype=torch.long, device=mask.device)] + + +def _take_tensor_field( + field: Optional[torch.Tensor], index_tensor: torch.Tensor +) -> Optional[torch.Tensor]: + """Index an optional per-box tensor field, tolerating ``None``. + + ``inference_models.Detections`` types declare ``class_id``/``confidence`` as + non-optional, but the numpy sibling path slices through + ``sv.Detections.__getitem__`` which silently passes ``None`` optional fields + through unindexed. Guard the same way so a producer that ever emits a native + object without confidence/class_id does not crash here where numpy would not. + """ + if field is None: + return None + return field[index_tensor] + + +def _take_detections( + detections: TensorNativeDetections, indices: List[int] +) -> TensorNativeDetections: + index_tensor = torch.as_tensor( + indices, dtype=torch.long, device=detections.xyxy.device + ) + bboxes_metadata = None + if detections.bboxes_metadata is not None: + bboxes_metadata = [detections.bboxes_metadata[index] for index in indices] + if isinstance(detections, InstanceDetections): + return InstanceDetections( + xyxy=detections.xyxy[index_tensor], + class_id=_take_tensor_field(detections.class_id, index_tensor), + confidence=_take_tensor_field(detections.confidence, index_tensor), + mask=_take_mask(detections.mask, indices), + image_metadata=detections.image_metadata, + bboxes_metadata=bboxes_metadata, + ) + return Detections( + xyxy=detections.xyxy[index_tensor], + class_id=_take_tensor_field(detections.class_id, index_tensor), + confidence=_take_tensor_field(detections.confidence, index_tensor), + image_metadata=detections.image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def _copy_detections(detections: TensorNativeDetections) -> TensorNativeDetections: + if isinstance(detections, InstanceDetections): + mask = detections.mask + if isinstance(mask, InstancesRLEMasks): + mask = InstancesRLEMasks(image_size=mask.image_size, masks=list(mask.masks)) + else: + mask = mask.clone() + return InstanceDetections( + xyxy=detections.xyxy.clone(), + class_id=detections.class_id.clone(), + confidence=detections.confidence.clone(), + mask=mask, + image_metadata=deepcopy(detections.image_metadata), + bboxes_metadata=deepcopy(detections.bboxes_metadata), + ) + return Detections( + xyxy=detections.xyxy.clone(), + class_id=detections.class_id.clone(), + confidence=detections.confidence.clone(), + image_metadata=deepcopy(detections.image_metadata), + bboxes_metadata=deepcopy(detections.bboxes_metadata), + ) + + +def _concatenate_detections( + first: TensorNativeDetections, second: TensorNativeDetections +) -> TensorNativeDetections: + bboxes_metadata = None + if first.bboxes_metadata is not None or second.bboxes_metadata is not None: + bboxes_metadata = _bboxes_metadata_list(first) + _bboxes_metadata_list(second) + if isinstance(first, InstanceDetections): + if isinstance(first.mask, InstancesRLEMasks) != isinstance( + second.mask, InstancesRLEMasks + ): + raise OperationError( + public_message="Cannot concatenate InstanceDetections with mixed mask " + "representations (dense tensor vs RLE).", + context="step_execution | roboflow_query_language_evaluation", + ) + if isinstance(first.mask, InstancesRLEMasks): + mask = InstancesRLEMasks( + image_size=first.mask.image_size, + masks=first.mask.masks + second.mask.masks, + ) + else: + mask = torch.cat([first.mask, second.mask], dim=0) + return InstanceDetections( + xyxy=torch.cat([first.xyxy, second.xyxy], dim=0), + class_id=torch.cat([first.class_id, second.class_id], dim=0), + confidence=torch.cat([first.confidence, second.confidence], dim=0), + mask=mask, + image_metadata=first.image_metadata, + bboxes_metadata=bboxes_metadata, + ) + return Detections( + xyxy=torch.cat([first.xyxy, second.xyxy], dim=0), + class_id=torch.cat([first.class_id, second.class_id], dim=0), + confidence=torch.cat([first.confidence, second.confidence], dim=0), + image_metadata=first.image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def _detections_anchor_coordinates_tensor_native( + detections: TensorNativeDetections, anchor: Position +) -> List[List[int]]: + xyxy = detections.xyxy + if anchor is Position.CENTER: + xs = (xyxy[:, 0] + xyxy[:, 2]) * 0.5 + ys = (xyxy[:, 1] + xyxy[:, 3]) * 0.5 + elif anchor is Position.TOP_LEFT: + xs, ys = xyxy[:, 0], xyxy[:, 1] + elif anchor is Position.TOP_RIGHT: + xs, ys = xyxy[:, 2], xyxy[:, 1] + elif anchor is Position.BOTTOM_LEFT: + xs, ys = xyxy[:, 0], xyxy[:, 3] + else: + xs, ys = xyxy[:, 2], xyxy[:, 3] + return torch.stack([xs, ys], dim=1).round().long().tolist() + + +PROPERTIES_EXTRACTORS_TENSOR_NATIVE = { + DetectionsProperty.CONFIDENCE: lambda detections: detections.confidence.tolist(), + DetectionsProperty.CLASS_NAME: lambda detections: _extract_class_names_property_tensor_native( + detections + ), + DetectionsProperty.X_MIN: lambda detections: detections.xyxy[:, 0].tolist(), + DetectionsProperty.Y_MIN: lambda detections: detections.xyxy[:, 1].tolist(), + DetectionsProperty.X_MAX: lambda detections: detections.xyxy[:, 2].tolist(), + DetectionsProperty.Y_MAX: lambda detections: detections.xyxy[:, 3].tolist(), + DetectionsProperty.CLASS_ID: lambda detections: [ + int(value) for value in detections.class_id.tolist() + ], + DetectionsProperty.SIZE: lambda detections: ( + (detections.xyxy[:, 2] - detections.xyxy[:, 0]) + * (detections.xyxy[:, 3] - detections.xyxy[:, 1]) + ).tolist(), + DetectionsProperty.CENTER: lambda detections: ( + _detections_anchor_coordinates_tensor_native( + detections=detections, anchor=Position.CENTER + ) + ), + DetectionsProperty.TOP_LEFT: lambda detections: ( + _detections_anchor_coordinates_tensor_native( + detections=detections, anchor=Position.TOP_LEFT + ) + ), + DetectionsProperty.TOP_RIGHT: lambda detections: ( + _detections_anchor_coordinates_tensor_native( + detections=detections, anchor=Position.TOP_RIGHT + ) + ), + DetectionsProperty.BOTTOM_LEFT: lambda detections: ( + _detections_anchor_coordinates_tensor_native( + detections=detections, anchor=Position.BOTTOM_LEFT + ) + ), + DetectionsProperty.BOTTOM_RIGHT: lambda detections: ( + _detections_anchor_coordinates_tensor_native( + detections=detections, anchor=Position.BOTTOM_RIGHT + ) + ), +} + + +def _extract_detections_property_tensor_native( + detections: Any, + property_name: DetectionsProperty, + execution_context: str, + **kwargs, +) -> List[Any]: + _ensure_tensor_native_detections( + detections, + operation_name="extract_detections_property", + execution_context=execution_context, + ) + if _is_key_point_prediction(detections): + # Property extraction reads the bounding-box component (mirrors the numpy + # path where keypoints ride alongside boxes in the same sv.Detections). + _, detections = _split_key_point_prediction( + detections, + operation_name="extract_detections_property", + execution_context=execution_context, + ) + if property_name not in PROPERTIES_EXTRACTORS_TENSOR_NATIVE: + bboxes_metadata = _bboxes_metadata_list(detections) + if any(property_name.value in data for data in bboxes_metadata): + return [data.get(property_name.value) for data in bboxes_metadata] + raise OperationError( + public_message=f"Executing extract_detections_property(...) in context " + f"{execution_context}, property `{property_name.value}` is neither " + f"natively supported nor present in `bboxes_metadata` of the detections.", + context=f"step_execution | roboflow_query_language_evaluation | {execution_context}", + ) + return PROPERTIES_EXTRACTORS_TENSOR_NATIVE[property_name](detections) + + +def _filter_detections_indices_tensor_native( + detections: TensorNativeDetections, + filtering_fun: Callable[[Dict[str, Any]], bool], + global_parameters: Dict[str, Any], +) -> List[int]: + local_parameters = copy(global_parameters) + indices_to_keep = [] + for index, detection in enumerate(detections): + local_parameters[DEFAULT_OPERAND_NAME] = detection + if filtering_fun(local_parameters): + indices_to_keep.append(index) + return indices_to_keep + + +def _filter_detections_tensor_native( + detections: Any, + filtering_fun: Callable[[Dict[str, Any]], bool], + global_parameters: Dict[str, Any], +) -> TensorNativePrediction: + _ensure_tensor_native_detections(detections, operation_name="filter_detections") + return _apply_index_op_on_prediction( + detections, + operation_name="filter_detections", + index_fn=lambda bboxes: _filter_detections_indices_tensor_native( + bboxes, filtering_fun=filtering_fun, global_parameters=global_parameters + ), + ) + + +def _offset_detections_impl_tensor_native( + detections: TensorNativeDetections, offset_x: int, offset_y: int +) -> TensorNativeDetections: + detections_copy = _copy_detections(detections) + detections_copy.xyxy = detections_copy.xyxy + torch.tensor( + [-offset_x / 2, -offset_y / 2, offset_x / 2, offset_y / 2], + dtype=detections_copy.xyxy.dtype, + device=detections_copy.xyxy.device, + ) + # Boxes changed -> the per-box host mirror would be stale; drop it so + # consumers fall back to tensor reads. + detections_copy.bboxes_metadata = strip_host_mirror_metadata( + detections_copy.bboxes_metadata + ) + return detections_copy + + +def _offset_detections_tensor_native( + value: Any, offset_x: int, offset_y: int, **kwargs +) -> TensorNativePrediction: + _ensure_tensor_native_detections(value, operation_name="offset_detections") + # Mirrors the numpy path: only bbox `xyxy` is offset; the keypoint `xy` + # coordinates (carried in sv `.data` in numpy mode) are left untouched. + if _is_key_point_prediction(value): + key_points, bboxes = _split_key_point_prediction( + value, operation_name="offset_detections" + ) + return key_points, _offset_detections_impl_tensor_native( + bboxes, offset_x=offset_x, offset_y=offset_y + ) + return _offset_detections_impl_tensor_native( + value, offset_x=offset_x, offset_y=offset_y + ) + + +def _shift_detections_impl_tensor_native( + detections: TensorNativeDetections, shift_x: int, shift_y: int +) -> TensorNativeDetections: + detections_copy = _copy_detections(detections) + detections_copy.xyxy = detections_copy.xyxy + torch.tensor( + [shift_x, shift_y, shift_x, shift_y], + dtype=detections_copy.xyxy.dtype, + device=detections_copy.xyxy.device, + ) + # Boxes changed -> the per-box host mirror would be stale; drop it so + # consumers fall back to tensor reads. + detections_copy.bboxes_metadata = strip_host_mirror_metadata( + detections_copy.bboxes_metadata + ) + return detections_copy + + +def _shift_detections_tensor_native( + value: Any, shift_x: int, shift_y: int, **kwargs +) -> TensorNativePrediction: + _ensure_tensor_native_detections(value, operation_name="shift_detections") + # Mirrors the numpy path: only bbox `xyxy` is shifted; the keypoint `xy` + # coordinates (carried in sv `.data` in numpy mode) are left untouched. + if _is_key_point_prediction(value): + key_points, bboxes = _split_key_point_prediction( + value, operation_name="shift_detections" + ) + return key_points, _shift_detections_impl_tensor_native( + bboxes, shift_x=shift_x, shift_y=shift_y + ) + return _shift_detections_impl_tensor_native(value, shift_x=shift_x, shift_y=shift_y) + + +def _select_top_confidence_index_tensor_native( + detections: TensorNativeDetections, +) -> Optional[int]: + if _detections_count(detections) == 0: + return None + return int(torch.argmax(detections.confidence)) + + +def _select_leftmost_index_tensor_native( + detections: TensorNativeDetections, +) -> Optional[int]: + if _detections_count(detections) == 0: + return None + centers_x = (detections.xyxy[:, 0] + detections.xyxy[:, 2]) * 0.5 + return int(torch.argmin(centers_x)) + + +def _select_rightmost_index_tensor_native( + detections: TensorNativeDetections, +) -> Optional[int]: + if _detections_count(detections) == 0: + return None + centers_x = (detections.xyxy[:, 0] + detections.xyxy[:, 2]) * 0.5 + return int(torch.argmax(centers_x)) + + +def _select_first_index_tensor_native( + detections: TensorNativeDetections, +) -> Optional[int]: + if _detections_count(detections) == 0: + return None + return 0 + + +def _select_last_index_tensor_native( + detections: TensorNativeDetections, +) -> Optional[int]: + count = _detections_count(detections) + if count == 0: + return None + return count - 1 + + +def _select_top_confidence_detection_tensor_native( + detections: TensorNativeDetections, +) -> TensorNativeDetections: + index = _select_top_confidence_index_tensor_native(detections) + if index is None: + return _copy_detections(detections) + return _take_detections(detections, [index]) + + +def _select_leftmost_detection_tensor_native( + detections: TensorNativeDetections, +) -> TensorNativeDetections: + index = _select_leftmost_index_tensor_native(detections) + if index is None: + # Copy on empty so no tensor-native selector aliases the caller's object + # (numpy's leftmost returns the original; value-parity still holds). + return _copy_detections(detections) + return _take_detections(detections, [index]) + + +def _select_rightmost_detection_tensor_native( + detections: TensorNativeDetections, +) -> TensorNativeDetections: + index = _select_rightmost_index_tensor_native(detections) + if index is None: + # See _select_leftmost_detection_tensor_native: copy on empty for selector + # consistency (no caller-object aliasing). + return _copy_detections(detections) + return _take_detections(detections, [index]) + + +def _select_first_detection_tensor_native( + detections: TensorNativeDetections, +) -> TensorNativeDetections: + index = _select_first_index_tensor_native(detections) + if index is None: + return _copy_detections(detections) + return _take_detections(detections, [index]) + + +def _select_last_detection_tensor_native( + detections: TensorNativeDetections, +) -> TensorNativeDetections: + index = _select_last_index_tensor_native(detections) + if index is None: + return _copy_detections(detections) + return _take_detections(detections, [index]) + + +DETECTIONS_SELECTORS_TENSOR_NATIVE = { + DetectionsSelectionMode.FIRST: _select_first_detection_tensor_native, + DetectionsSelectionMode.LAST: _select_last_detection_tensor_native, + DetectionsSelectionMode.LEFT_MOST: _select_leftmost_detection_tensor_native, + DetectionsSelectionMode.RIGHT_MOST: _select_rightmost_detection_tensor_native, + DetectionsSelectionMode.TOP_CONFIDENCE: _select_top_confidence_detection_tensor_native, +} + +DETECTIONS_SELECTOR_INDICES_TENSOR_NATIVE = { + DetectionsSelectionMode.FIRST: _select_first_index_tensor_native, + DetectionsSelectionMode.LAST: _select_last_index_tensor_native, + DetectionsSelectionMode.LEFT_MOST: _select_leftmost_index_tensor_native, + DetectionsSelectionMode.RIGHT_MOST: _select_rightmost_index_tensor_native, + DetectionsSelectionMode.TOP_CONFIDENCE: _select_top_confidence_index_tensor_native, +} + + +def _select_detections_tensor_native( + value: Any, mode: DetectionsSelectionMode, **kwargs +) -> TensorNativePrediction: + _ensure_tensor_native_detections(value, operation_name="select_detections") + if mode not in DETECTIONS_SELECTORS_TENSOR_NATIVE: + raise InvalidInputTypeError( + public_message=f"Executing select_detections(...), expected mode to be one " + f"of {list(DETECTIONS_SELECTORS_TENSOR_NATIVE.keys())}, got {mode}.", + context="step_execution | roboflow_query_language_evaluation", + ) + if _is_key_point_prediction(value): + key_points, bboxes = _split_key_point_prediction( + value, operation_name="select_detections" + ) + index = DETECTIONS_SELECTOR_INDICES_TENSOR_NATIVE[mode](bboxes) + indices = [] if index is None else [index] + return _take_key_points(key_points, indices), _take_detections(bboxes, indices) + return DETECTIONS_SELECTORS_TENSOR_NATIVE[mode](value) + + +SORT_PROPERTIES_EXTRACT_TENSOR_NATIVE = { + DetectionsSortProperties.CONFIDENCE: lambda detections: detections.confidence, + DetectionsSortProperties.X_MIN: lambda detections: detections.xyxy[:, 0], + DetectionsSortProperties.X_MAX: lambda detections: detections.xyxy[:, 2], + DetectionsSortProperties.Y_MIN: lambda detections: detections.xyxy[:, 1], + DetectionsSortProperties.Y_MAX: lambda detections: detections.xyxy[:, 3], + DetectionsSortProperties.SIZE: lambda detections: ( + (detections.xyxy[:, 2] - detections.xyxy[:, 0]) + * (detections.xyxy[:, 3] - detections.xyxy[:, 1]) + ), + DetectionsSortProperties.CENTER_X: lambda detections: ( + (detections.xyxy[:, 0] + detections.xyxy[:, 2]) * 0.5 + ), + DetectionsSortProperties.CENTER_Y: lambda detections: ( + (detections.xyxy[:, 1] + detections.xyxy[:, 3]) * 0.5 + ), +} + + +def _sort_detections_indices_tensor_native( + detections: TensorNativeDetections, + mode: DetectionsSortProperties, + ascending: bool, +) -> List[int]: + extracted_property = SORT_PROPERTIES_EXTRACT_TENSOR_NATIVE[mode](detections) + sorted_indices = torch.argsort(extracted_property, descending=not ascending) + return [int(index) for index in sorted_indices.tolist()] + + +def _sort_detections_tensor_native( + value: Any, mode: DetectionsSortProperties, ascending: bool, **kwargs +) -> TensorNativePrediction: + _ensure_tensor_native_detections(value, operation_name="sort_detections") + if mode not in SORT_PROPERTIES_EXTRACT_TENSOR_NATIVE: + raise InvalidInputTypeError( + public_message=f"Executing sort_detections(...), expected mode to be one of " + f"{list(SORT_PROPERTIES_EXTRACT_TENSOR_NATIVE.keys())}, got {mode}.", + context="step_execution | roboflow_query_language_evaluation", + ) + bboxes_for_count = value + if _is_key_point_prediction(value): + _, bboxes_for_count = _split_key_point_prediction( + value, operation_name="sort_detections" + ) + # Parity with the numpy sibling: return the input object unchanged on empty. + if _detections_count(bboxes_for_count) == 0: + return value + return _apply_index_op_on_prediction( + value, + operation_name="sort_detections", + index_fn=lambda bboxes: _sort_detections_indices_tensor_native( + bboxes, mode=mode, ascending=ascending + ), + ) + + +def _rename_detections_tensor_native( + detections: Any, + class_map: Union[Dict[str, str], str], + strict: Union[bool, str], + new_classes_id_offset: int, + global_parameters: Dict[str, Any], + **kwargs, +) -> TensorNativePrediction: + _ensure_tensor_native_detections(detections, operation_name="rename_detections") + # Rename only affects the bounding-box `Detections` class ids / names (mirrors + # the numpy path, which touches `data["class_name"]` / `class_id`). The + # keypoint component is preserved unchanged in the re-wrapped tuple. + if _is_key_point_prediction(detections): + key_points, bboxes = _split_key_point_prediction( + detections, operation_name="rename_detections" + ) + renamed_bboxes = _rename_detections_tensor_native( + detections=bboxes, + class_map=class_map, + strict=strict, + new_classes_id_offset=new_classes_id_offset, + global_parameters=global_parameters, + **kwargs, + ) + return key_points, renamed_bboxes + if isinstance(class_map, str): + if class_map not in global_parameters: + raise UndeclaredSymbolError( + public_message=f"Attempted to retrieve variable `{class_map}` that was expected to hold " + f"class mapping of rename_detections(...), but that turned out not to be registered.", + context="step_execution | roboflow_query_language_evaluation", + ) + class_map = global_parameters[class_map] + if not isinstance(class_map, dict): + value_as_str = safe_stringify(value=class_map) + raise InvalidInputTypeError( + public_message=f"Executing rename_detections(...), expected dictionary to be given as class map, " + f"got {value_as_str} of type {type(class_map)}", + context="step_execution | roboflow_query_language_evaluation", + ) + if isinstance(strict, str): + if strict not in global_parameters: + raise UndeclaredSymbolError( + public_message=f"Attempted to retrieve variable `{strict}` that was expected to hold " + f"parameter for `strict` flag of rename_detections(...), but that turned out not " + f"to be registered.", + context="step_execution | roboflow_query_language_evaluation", + ) + strict = global_parameters[strict] + if not isinstance(strict, bool): + value_as_str = safe_stringify(value=strict) + raise InvalidInputTypeError( + public_message=f"Executing rename_detections(...), expected dictionary to be given as `strict` flag, " + f"got {value_as_str} of type {type(strict)}", + context="step_execution | roboflow_query_language_evaluation", + ) + # Numpy-parity for detections that carry no class names (main c03c2514b): the + # numpy arm returns the copy unchanged unless strict-and-non-empty, in which + # case the shared helper raises. The native analog of a missing + # ``data["class_name"]`` is a missing ``image_metadata[CLASS_NAMES_KEY]`` map; + # a present-but-incomplete map still raises in ``_resolve_class_names`` below + # (real producer-contract violation, not an "absent property"). + if (detections.image_metadata or {}).get(CLASS_NAMES_KEY) is None: + if strict and _detections_count(detections) > 0: + _ensure_all_classes_covered_in_new_mapping( + original_class_names=None, + class_map=class_map, + ) + return _copy_detections(detections) + # Row-level rename subject, like numpy's per-row ``data["class_name"]``: the + # per-box ``CLASS_NAME_KEY`` override wins over the map (the serializer's C1 + # precedence), so a classes_replacement -> rename chain renames the labels a + # consumer would actually see instead of leaving stale overrides behind. + bboxes_metadata = detections.bboxes_metadata + row_has_override = [ + bboxes_metadata is not None and CLASS_NAME_KEY in (bboxes_metadata[index] or {}) + for index in range(_detections_count(detections)) + ] + original_class_names = _resolve_effective_class_names( + detections, operation_name="rename_detections" + ) + original_class_ids = [int(value) for value in detections.class_id.tolist()] + if strict: + _ensure_all_classes_covered_in_new_mapping( + original_class_names=original_class_names, + class_map=class_map, + ) + new_class_mapping = { + class_name: class_id + for class_id, class_name in enumerate(sorted(set(class_map.values()))) + } + else: + new_class_mapping = _build_non_strict_class_to_id_mapping( + original_class_names=original_class_names, + original_class_ids=original_class_ids, + class_map=class_map, + new_classes_id_offset=new_classes_id_offset, + ) + new_class_names = [ + class_map.get(class_name, class_name) for class_name in original_class_names + ] + new_class_ids = [new_class_mapping[class_name] for class_name in new_class_names] + detections_copy = _copy_detections(detections) + detections_copy.class_id = torch.tensor( + new_class_ids, + dtype=detections.class_id.dtype, + device=detections.class_id.device, + ) + # class_id changed -> the per-box host mirror would be stale; drop it so + # consumers fall back to tensor reads. + detections_copy.bboxes_metadata = strip_host_mirror_metadata( + detections_copy.bboxes_metadata + ) + if detections_copy.bboxes_metadata is not None: + # Present-only rewrite: rows that carried an override keep the channel, + # now pointing at the renamed label; rows without one gain nothing. + for index, entry in enumerate(detections_copy.bboxes_metadata): + if row_has_override[index]: + entry[CLASS_NAME_KEY] = new_class_names[index] + allowed_names = set(new_class_names) + allowed_names.update(class_map.values()) + image_metadata = detections_copy.image_metadata or {} + renamed_class_names_map = { + class_id: class_name + for class_name, class_id in new_class_mapping.items() + if class_name in allowed_names + } + # Row-level serializer guarantee: an override-less row resolves its name via + # the map, so its (id -> name) pair must win over a colliding inverted entry + # (non-strict can hold two names for one id when an override diverged from + # the map name at the same class_id). Identity when no overrides exist, + # because override-less effective names are a function of class_id. + for index, has_override in enumerate(row_has_override): + if has_override: + renamed_class_names_map.setdefault( + new_class_ids[index], new_class_names[index] + ) + else: + renamed_class_names_map[new_class_ids[index]] = new_class_names[index] + image_metadata[CLASS_NAMES_KEY] = renamed_class_names_map + detections_copy.image_metadata = image_metadata + return detections_copy + + +def _detections_to_dictionary_tensor_native( + detections: Any, + execution_context: str, + **kwargs, +) -> dict: + _ensure_tensor_native_detections( + detections, + operation_name="detections_to_dictionary", + execution_context=execution_context, + ) + if _is_key_point_prediction(detections): + # The op-level numpy path serialises the (box-carrying) sv.Detections; here + # we serialise the bounding-box component (keypoint details are not part of + # this dict-serialisation, matching the established kind serialisation). + _, detections = _split_key_point_prediction( + detections, + operation_name="detections_to_dictionary", + execution_context=execution_context, + ) + try: + return serialise_tensor_native_detections(detections=detections) + except Exception as error: + raise OperationError( + public_message=f"While Using operation detections_to_dictionary(...) in context {execution_context} " + f"encountered error: {error}", + context=f"step_execution | roboflow_query_language_evaluation | {execution_context}", + inner_error=error, + ) + + +def _pick_detections_by_parent_class_tensor_native( + detections: Any, + parent_class: str, + execution_context: str, + **kwargs, +) -> TensorNativePrediction: + _ensure_tensor_native_detections( + detections, + operation_name="pick_detections_by_parent_class", + execution_context=execution_context, + ) + try: + if _is_key_point_prediction(detections): + return _apply_index_op_on_prediction( + detections, + operation_name="pick_detections_by_parent_class", + index_fn=lambda bboxes: _pick_detections_by_parent_class_indices_tensor_native( + detections=bboxes, parent_class=parent_class + ), + execution_context=execution_context, + ) + return _pick_detections_by_parent_class_tensor_native_impl( + detections=detections, parent_class=parent_class + ) + except Exception as error: + raise OperationError( + public_message=f"While Using operation pick_detections_by_parent_class(...) in context {execution_context} " + f"encountered error: {error}", + context=f"step_execution | roboflow_query_language_evaluation | {execution_context}", + inner_error=error, + ) + + +def _parent_class_row_flags( + detections: TensorNativeDetections, + parent_class: str, + operation_name: str, +) -> List[bool]: + """Per-row "is this a parent-class detection" flags for + ``pick_detections_by_parent_class`` โ€” resolving names only as far as the query + actually needs. + + Parity with the numpy sibling (``_pick_detections_by_parent_class_impl``): a row + is a parent iff its ``class_id`` maps to ``parent_class`` in the image-level + ``CLASS_NAMES_KEY`` map; every other row is a dependent. Because dependents are + identified by *not* matching ``parent_class`` (never by resolving their own + name), an incomplete map that omits an unrelated dependent's ``class_id`` does + NOT raise here โ€” the numpy sibling likewise never raises for such a row (its + ``data["class_name"]`` simply holds a non-parent string). This is the fix for + the divergence where the old ``_resolve_class_names`` call raised on any + unmapped row, even ones the query doesn't touch. + + A *truly absent* map (no ``CLASS_NAMES_KEY`` at all) remains a producer contract + violation and still raises via ``_class_names_lookup`` โ€” the single accepted + divergence from numpy, which returns empty there. + """ + class_names_map = _class_names_lookup(detections, operation_name=operation_name) + return [ + class_names_map.get(int(class_id_scalar)) == parent_class + for class_id_scalar in detections.class_id.tolist() + ] + + +def _pick_detections_by_parent_class_indices_tensor_native( + detections: TensorNativeDetections, + parent_class: str, +) -> List[int]: + """Absolute index order produced by ``pick_detections_by_parent_class``: every + parent detection first (in original order), then the dependent detections whose + center lies inside any parent (mirrors the ``_concatenate_detections`` ordering + in the impl). Returned so the keypoint component can be sliced consistently.""" + if _detections_count(detections) == 0: + return [] + is_parent = _parent_class_row_flags( + detections, + parent_class=parent_class, + operation_name="pick_detections_by_parent_class", + ) + parent_indices = [index for index, flag in enumerate(is_parent) if flag] + if not parent_indices: + return [] + dependent_indices = [index for index, flag in enumerate(is_parent) if not flag] + dependent_detections = _take_detections(detections, dependent_indices) + parent_detections = _take_detections(detections, parent_indices) + centers_x = ( + (dependent_detections.xyxy[:, 0] + dependent_detections.xyxy[:, 2]) * 0.5 + ).unsqueeze(1) + centers_y = ( + (dependent_detections.xyxy[:, 1] + dependent_detections.xyxy[:, 3]) * 0.5 + ).unsqueeze(1) + parents_x1 = parent_detections.xyxy[:, 0].unsqueeze(0) + parents_y1 = parent_detections.xyxy[:, 1].unsqueeze(0) + parents_x2 = parent_detections.xyxy[:, 2].unsqueeze(0) + parents_y2 = parent_detections.xyxy[:, 3].unsqueeze(0) + inside_any_parent = ( + (centers_x >= parents_x1) + & (centers_x <= parents_x2) + & (centers_y >= parents_y1) + & (centers_y <= parents_y2) + ).any(dim=1) + kept_dependent_indices = [ + dependent_indices[position] + for position, keep in enumerate(inside_any_parent.tolist()) + if keep + ] + return parent_indices + kept_dependent_indices + + +def _pick_detections_by_parent_class_tensor_native_impl( + detections: TensorNativeDetections, + parent_class: str, +) -> TensorNativeDetections: + if _detections_count(detections) == 0: + return _take_detections(detections, []) + is_parent = _parent_class_row_flags( + detections, + parent_class=parent_class, + operation_name="pick_detections_by_parent_class", + ) + parent_indices = [index for index, flag in enumerate(is_parent) if flag] + if not parent_indices: + return _take_detections(detections, []) + dependent_indices = [index for index, flag in enumerate(is_parent) if not flag] + parent_detections = _take_detections(detections, parent_indices) + dependent_detections = _take_detections(detections, dependent_indices) + centers_x = ( + (dependent_detections.xyxy[:, 0] + dependent_detections.xyxy[:, 2]) * 0.5 + ).unsqueeze(1) + centers_y = ( + (dependent_detections.xyxy[:, 1] + dependent_detections.xyxy[:, 3]) * 0.5 + ).unsqueeze(1) + parents_x1 = parent_detections.xyxy[:, 0].unsqueeze(0) + parents_y1 = parent_detections.xyxy[:, 1].unsqueeze(0) + parents_x2 = parent_detections.xyxy[:, 2].unsqueeze(0) + parents_y2 = parent_detections.xyxy[:, 3].unsqueeze(0) + inside_any_parent = ( + (centers_x >= parents_x1) + & (centers_x <= parents_x2) + & (centers_y >= parents_y1) + & (centers_y <= parents_y2) + ).any(dim=1) + dependent_detections_to_keep = [ + index for index, keep in enumerate(inside_any_parent.tolist()) if keep + ] + filtered_dependent_detections = _take_detections( + dependent_detections, dependent_detections_to_keep + ) + return _concatenate_detections(parent_detections, filtered_dependent_detections) + + +def extract_detections_property( + detections: Any, + property_name: DetectionsProperty, + execution_context: str, + **kwargs, +) -> List[Any]: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _extract_detections_property_tensor_native( + detections=detections, + property_name=property_name, + execution_context=execution_context, + **kwargs, + ) + return _extract_detections_property( + detections=detections, + property_name=property_name, + execution_context=execution_context, + **kwargs, + ) + + +def filter_detections( + detections: Any, + filtering_fun: Callable[[Dict[str, Any]], bool], + global_parameters: Dict[str, Any], +) -> Union[sv.Detections, TensorNativeDetections]: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _filter_detections_tensor_native( + detections=detections, + filtering_fun=filtering_fun, + global_parameters=global_parameters, + ) + return _filter_detections( + detections=detections, + filtering_fun=filtering_fun, + global_parameters=global_parameters, + ) + + +def offset_detections( + value: Any, offset_x: int, offset_y: int, **kwargs +) -> Union[sv.Detections, TensorNativeDetections]: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _offset_detections_tensor_native( + value=value, offset_x=offset_x, offset_y=offset_y, **kwargs + ) + return _offset_detections( + value=value, offset_x=offset_x, offset_y=offset_y, **kwargs + ) + + +def shift_detections( + value: Any, shift_x: int, shift_y: int, **kwargs +) -> Union[sv.Detections, TensorNativeDetections]: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _shift_detections_tensor_native( + value=value, shift_x=shift_x, shift_y=shift_y, **kwargs + ) + return _shift_detections(value=value, shift_x=shift_x, shift_y=shift_y, **kwargs) + + +def select_detections( + value: Any, mode: DetectionsSelectionMode, **kwargs +) -> Union[sv.Detections, TensorNativeDetections]: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _select_detections_tensor_native(value=value, mode=mode, **kwargs) + return _select_detections(value=value, mode=mode, **kwargs) + + +def sort_detections( + value: Any, mode: DetectionsSortProperties, ascending: bool, **kwargs +) -> Union[sv.Detections, TensorNativeDetections]: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _sort_detections_tensor_native( + value=value, mode=mode, ascending=ascending, **kwargs + ) + return _sort_detections(value=value, mode=mode, ascending=ascending, **kwargs) + + +def rename_detections( + detections: Any, + class_map: Union[Dict[str, str], str], + strict: Union[bool, str], + new_classes_id_offset: int, + global_parameters: Dict[str, Any], + **kwargs, +) -> Union[sv.Detections, TensorNativeDetections]: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _rename_detections_tensor_native( + detections=detections, + class_map=class_map, + strict=strict, + new_classes_id_offset=new_classes_id_offset, + global_parameters=global_parameters, + **kwargs, + ) + return _rename_detections( + detections=detections, + class_map=class_map, + strict=strict, + new_classes_id_offset=new_classes_id_offset, + global_parameters=global_parameters, + **kwargs, + ) + + +def detections_to_dictionary( + detections: Any, + execution_context: str, + **kwargs, +) -> dict: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _detections_to_dictionary_tensor_native( + detections=detections, execution_context=execution_context, **kwargs + ) + return _detections_to_dictionary( + detections=detections, execution_context=execution_context, **kwargs + ) + + +def pick_detections_by_parent_class( + detections: Any, + parent_class: str, + execution_context: str, + **kwargs, +) -> Union[sv.Detections, TensorNativeDetections]: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return _pick_detections_by_parent_class_tensor_native( + detections=detections, + parent_class=parent_class, + execution_context=execution_context, + **kwargs, + ) + return _pick_detections_by_parent_class( + detections=detections, + parent_class=parent_class, + execution_context=execution_context, + **kwargs, + ) diff --git a/inference/core/workflows/core_steps/common/rle_compact.py b/inference/core/workflows/core_steps/common/rle_compact.py new file mode 100644 index 0000000000..fae01407ee --- /dev/null +++ b/inference/core/workflows/core_steps/common/rle_compact.py @@ -0,0 +1,230 @@ +"""Zero-densification bridge: inference-models full-frame COCO RLE masks +(``InstancesRLEMasks``) -> supervision ``CompactMask`` (per-crop RLE). + +Why this exists +--------------- +``inference_models.InstanceDetections`` can carry masks as ``InstancesRLEMasks``: +a list of **full-frame**, **column-major (Fortran-order)** COCO RLE byte strings +(one per instance), as produced by ``pycocotools``. The visualisation seam used +to turn those into an ``sv.Detections`` by decoding the whole stack to a dense +``(N, H, W)`` boolean array (``coco_rle_masks_to_numpy_mask`` -> +``pycocotools.mask.decode``). That is O(NยทHยทW) memory and time and is the +dominant cost when visualising many instances on high-resolution frames. + +``supervision.CompactMask`` stores each mask as an RLE of its **bounding-box +crop** instead of a full ``(H, W)`` array, and its annotators paint directly +into the crop region (``_paint_masks_by_area``) โ€” no full-frame allocation. + +The key observation that makes a *direct* transcode possible: supervision's +per-crop RLE (``_mask_to_rle_counts``) and COCO's uncompressed counts use the +**same encoding** โ€” column-major run lengths starting with a leading +``False`` (background) run, alternating False/True. They differ only in *scope* +(full image vs. bbox crop). So converting one to the other is a matter of: + +1. decompressing the COCO counts to plain run lengths (run arithmetic, NOT a + pixel decode โ€” see :func:`_decode_coco_counts`); +2. splitting the full-frame run list into per-column run lists + (``_rle_split_cols``); +3. selecting the bbox's columns and trimming each to the bbox's rows + (:func:`_trim_col_runs`); +4. re-joining the selected/trimmed columns into a flat crop RLE + (``_rle_join_cols``). + +No ``(H, W)`` or ``(N, H, W)`` array is ever allocated. + +Parity contract +--------------- +:func:`compact_mask_from_coco_rle` is defined to produce **exactly** what +``CompactMask.from_dense(coco_rle_masks_to_numpy_mask(rle), xyxy, image_shape)`` +would produce โ€” same box clipping (clip to ``[0, dim-1]``, inclusive max +coords), same invalid-box handling (``x2 < x1`` or ``y2 < y1`` -> ``1x1`` +all-False crop) โ€” just without the dense intermediate. The accompanying unit +test asserts decoded-mask equality against that reference. + +Note on layering / upstreaming +------------------------------- +This reuses supervision's private RLE primitives (``_rle_split_cols`` / +``_rle_join_cols``) to guarantee identical junction-merge semantics. The natural +long-term home for :func:`compact_mask_from_coco_rle` is a +``CompactMask.from_coco_rle`` classmethod in supervision itself, at which point +those imports become internal. It lives here for now so the tensor pipeline can +adopt it without waiting on a supervision release. +""" + +from typing import List, Sequence, Tuple, Union + +import numpy as np + +# Private supervision primitives. Isolated here so a single import site breaks +# if supervision relocates them (rather than scattering the coupling). +from supervision.detection.compact_mask import ( + CompactMask, + _rle_join_cols, + _rle_split_cols, +) + +from inference_models.models.base.types import InstancesRLEMasks + + +def _decode_coco_counts(counts: Union[bytes, str]) -> List[int]: + """Decompress a COCO compressed-RLE ``counts`` string to plain run lengths. + + Inverse of ``pycocotools``' ``rleToString`` (``maskApi.c``): a LEB128-style + codec using 5 payload bits per character (ascii offset 48), a continuation + bit (0x20), a sign bit (0x10) on the final char, and a delta against the + value two positions back for every run from index 3 onward. + + Returns the uncompressed, column-major (F-order) run lengths: ``[False_run, + True_run, False_run, ...]`` summing to ``H*W``. This is pure integer + arithmetic โ€” it never materialises pixels. + """ + data = counts.encode("ascii") if isinstance(counts, str) else bytes(counts) + cnts: List[int] = [] + p = 0 + n = len(data) + m = 0 + while p < n: + x = 0 + k = 0 + more = True + while more: + c = data[p] - 48 + x |= (c & 0x1F) << (5 * k) + more = bool(c & 0x20) + p += 1 + k += 1 + if not more and (c & 0x10): + x |= (-1) << (5 * k) + if m > 2: + x += cnts[m - 2] + cnts.append(x) + m += 1 + return cnts + + +def _trim_col_runs(col_runs: Sequence[int], y1: int, y2: int) -> List[int]: + """Restrict one full-height column run list to rows ``[y1, y2]`` inclusive. + + ``col_runs`` (as produced by ``_rle_split_cols``) starts with a ``False`` + count and alternates, summing to the full column height. The returned list + covers ``y2 - y1 + 1`` rows and also starts with a ``False`` count (a + leading ``0`` is inserted when the window begins on a ``True`` pixel), + matching the convention ``_rle_join_cols`` expects. + """ + want = y2 - y1 + 1 + collected: List[Tuple[bool, int]] = [] + row = 0 + for idx, run_len in enumerate(col_runs): + is_true = idx % 2 == 1 + start = row + end = row + int(run_len) + row = end + lo = max(start, y1) + hi = min(end, y2 + 1) + if hi > lo: + collected.append((is_true, hi - lo)) + if row > y2: + break + + if not collected: + return [want] + + out: List[int] = [] + if collected[0][0]: # window starts on True -> leading False count of 0 + out.append(0) + for is_true, length in collected: + last_is_true = bool(out) and ((len(out) - 1) % 2 == 1) + if out and last_is_true == is_true: + out[-1] += length + else: + out.append(length) + return out + + +def compact_mask_from_coco_rle( + image_shape: Tuple[int, int], + masks_counts: Sequence[Union[bytes, str]], + xyxy: np.ndarray, +) -> CompactMask: + """Build a :class:`CompactMask` from full-frame COCO RLE counts, no densify. + + Args: + image_shape: ``(H, W)`` of the full image. + masks_counts: one COCO compressed-RLE ``counts`` per instance, + column-major, full-frame (``InstancesRLEMasks.masks``). + xyxy: ``(N, 4)`` boxes ``[x1, y1, x2, y2]`` (supervision inclusive-max + convention), used as the crop bounds โ€” identical to + ``CompactMask.from_dense``. + + Returns: + A :class:`CompactMask` decode-equal to + ``CompactMask.from_dense(decode(masks_counts), xyxy, image_shape)``. + """ + img_h, img_w = int(image_shape[0]), int(image_shape[1]) + num_masks = len(masks_counts) + + if num_masks == 0: + return CompactMask( + [], + np.empty((0, 2), dtype=np.int32), + np.empty((0, 2), dtype=np.int32), + (img_h, img_w), + ) + + rles: List[np.ndarray] = [] + crop_shapes: List[Tuple[int, int]] = [] + offsets: List[Tuple[int, int]] = [] + + for i in range(num_masks): + x1, y1, x2, y2 = xyxy[i] + x1c = int(max(0, min(int(x1), img_w - 1))) + y1c = int(max(0, min(int(y1), img_h - 1))) + x2c = int(max(0, min(int(x2), img_w - 1))) + y2c = int(max(0, min(int(y2), img_h - 1))) + + # Mirror CompactMask.from_dense's degenerate-box handling exactly. + if x2c < x1c or y2c < y1c: + rles.append(np.array([1], dtype=np.int32)) + crop_shapes.append((1, 1)) + offsets.append((x1c, y1c)) + continue + + crop_h = y2c - y1c + 1 + crop_w = x2c - x1c + 1 + + full_counts = _decode_coco_counts(masks_counts[i]) + # Split the full-frame F-order RLE into one run list per image column + # (each column = img_h pixels). Pure run arithmetic โ€” no pixel buffer. + columns = _rle_split_cols(np.asarray(full_counts, dtype=np.int64), img_h, img_w) + selected = [ + _trim_col_runs(columns[col], y1c, y2c) for col in range(x1c, x2c + 1) + ] + crop_rle = _rle_join_cols(selected, crop_h * crop_w) + + rles.append(crop_rle) + crop_shapes.append((crop_h, crop_w)) + offsets.append((x1c, y1c)) + + return CompactMask( + rles, + np.array(crop_shapes, dtype=np.int32), + np.array(offsets, dtype=np.int32), + (img_h, img_w), + ) + + +def instances_rle_to_compact_mask( + masks: InstancesRLEMasks, + xyxy: np.ndarray, +) -> CompactMask: + """Adapter: ``InstancesRLEMasks`` -> ``CompactMask`` (zero densification). + + ``xyxy`` must be the full-frame boxes for the same instances, in the + supervision inclusive-max convention (the visualisation seam already has + them as a host numpy array). + """ + return compact_mask_from_coco_rle( + image_shape=masks.image_size, + masks_counts=masks.masks, + xyxy=xyxy, + ) diff --git a/inference/core/workflows/core_steps/common/serializers_tensor.py b/inference/core/workflows/core_steps/common/serializers_tensor.py new file mode 100644 index 0000000000..95f67e9698 --- /dev/null +++ b/inference/core/workflows/core_steps/common/serializers_tensor.py @@ -0,0 +1,805 @@ +from datetime import datetime +from typing import Any, List, Optional, Tuple, Union + +import numpy as np +import supervision as sv +import torch + +from inference.core.workflows.core_steps.common.keypoints import real_keypoints_count +from inference.core.workflows.core_steps.common.serializers import ( + _attach_parent_metadata_to_detection_dict, + mask_to_polygon, + serialise_image, +) +from inference.core.workflows.core_steps.common.serializers import ( + serialise_sv_detections as _serialise_legacy_sv_detections, +) +from inference.core.workflows.core_steps.common.serializers import serialize_timestamp +from inference.core.workflows.execution_engine.constants import ( + AREA_CONVERTED_KEY_IN_INFERENCE_RESPONSE, + AREA_CONVERTED_KEY_IN_SV_DETECTIONS, + AREA_KEY_IN_INFERENCE_RESPONSE, + AREA_KEY_IN_SV_DETECTIONS, + BOUNDING_RECT_ANGLE_KEY_IN_INFERENCE_RESPONSE, + BOUNDING_RECT_ANGLE_KEY_IN_SV_DETECTIONS, + BOUNDING_RECT_HEIGHT_KEY_IN_INFERENCE_RESPONSE, + BOUNDING_RECT_HEIGHT_KEY_IN_SV_DETECTIONS, + BOUNDING_RECT_RECT_KEY_IN_INFERENCE_RESPONSE, + BOUNDING_RECT_RECT_KEY_IN_SV_DETECTIONS, + BOUNDING_RECT_WIDTH_KEY_IN_INFERENCE_RESPONSE, + BOUNDING_RECT_WIDTH_KEY_IN_SV_DETECTIONS, + CLASS_ID_KEY, + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_FORMATTER, + CLASSIFICATION_STYLE_KEY, + CLASSIFICATION_STYLE_MODEL, + CONFIDENCE_KEY, + DETECTED_CODE_KEY, + DETECTION_ID_KEY, + HEIGHT_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS, + KEYPOINTS_KEY_IN_INFERENCE_RESPONSE, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + NEAREST_TARGET_DISTANCE_KEY, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + PARENT_ORIGIN_KEY, + PATH_DEVIATION_KEY_IN_INFERENCE_RESPONSE, + PATH_DEVIATION_KEY_IN_SV_DETECTIONS, + POLYGON_KEY, + POLYGON_KEY_IN_INFERENCE_RESPONSE, + POLYGON_KEY_IN_SV_DETECTIONS, + PREDICTION_TYPE_KEY, + RLE_MASK_KEY_IN_INFERENCE_RESPONSE, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, + ROOT_PARENT_ORIGIN_KEY, + SERIALISE_POLYGONS_KEY, + SMOOTHED_SPEED_KEY_IN_INFERENCE_RESPONSE, + SMOOTHED_SPEED_KEY_IN_SV_DETECTIONS, + SMOOTHED_VELOCITY_KEY_IN_INFERENCE_RESPONSE, + SMOOTHED_VELOCITY_KEY_IN_SV_DETECTIONS, + SPEED_KEY_IN_INFERENCE_RESPONSE, + SPEED_KEY_IN_SV_DETECTIONS, + TIME_IN_ZONE_KEY_IN_INFERENCE_RESPONSE, + TIME_IN_ZONE_KEY_IN_SV_DETECTIONS, + TRACKER_ID_KEY, + VELOCITY_KEY_IN_INFERENCE_RESPONSE, + VELOCITY_KEY_IN_SV_DETECTIONS, + WIDTH_KEY, + X_KEY, + Y_KEY, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import coco_rle_masks_to_numpy_mask + +TensorNativeDetections = Union[Detections, InstanceDetections] + +# D4 discriminator signals. A single native ``ClassificationPrediction`` / +# ``MultiLabelClassificationPrediction`` type is emitted by producers whose flag-OFF +# JSON has two INCOMPATIBLE shapes (see ``serialise_native_classification``): the +# numpy ``*InferenceResponse`` "model" shape and the hand-built ``vlm_as_classifier`` +# "formatter" shape. Since the tensors cannot self-describe which shape they want, +# we infer it from metadata the producers already attach - the model producers carry +# a confidence threshold and/or ``root_parent_id`` and/or ``time``; the vlm formatter +# carries none of these. +_CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY = "classification_confidence_threshold" +_TIME_KEY = "time" + + +def serialise_sv_detections( + detections: TensorNativeDetections, +) -> dict: + """Serialise native ``inference_models.Detections`` / ``InstanceDetections`` into + the response dict shape produced by the numpy ``serialise_sv_detections``. + + The name is retained for loader symbol-swap compatibility (it is load-bearing) - + despite the ``sv`` prefix it now consumes native tensor objects, not + ``sv.Detections``. + """ + serialized_image_metadata, serialized_detections, _ = _serialise_sv_detections( + detections=detections + ) + return {"image": serialized_image_metadata, "predictions": serialized_detections} + + +def _serialise_sv_detections( + detections: TensorNativeDetections, + emit_polygons: Optional[bool] = None, +) -> Tuple[dict, List[dict], List[int]]: + """Shared core for the detection serialisers. + + Returns the serialised image metadata, the per-box prediction dicts (after the + polygon-None skip) and the ORIGINAL row index of each surviving prediction so + downstream serialisers (e.g. RLE) can re-align masks despite skipped instances. + + ``emit_polygons`` controls the per-instance mask -> ``points`` polygon work + (the RLE-decode + ``mask_to_polygon``/``findContours`` cost): + + * ``None`` (default): resolve from ``image_metadata[SERIALISE_POLYGONS_KEY]``, + defaulting to ``True`` when the key is ABSENT. Absent-key behaviour is + byte-identical to the historical serialiser (polygons computed & emitted, + instances with an empty contour dropped). A producer that sets the key to + ``False`` gets the polygon work skipped: instances are kept (no drop) and + no ``points`` field is emitted. + * ``False`` (forced): callers that discard the polygon anyway (the native RLE + serialiser) pass this to avoid the wasted O(N*H*W) computation. Overrides the + metadata flag. + """ + if not isinstance(detections, (Detections, InstanceDetections)): + # C2: this typed path only accepts the native tensor objects. ``sv.Detections`` + # is routed to the numpy serialiser by ``serialize_wildcard_kind`` instead, so + # the message must NOT claim this function accepts it. + raise ValueError( + f"serialise_sv_detections(...) expected `inference_models.Detections` " + f"or `inference_models.InstanceDetections`, got {type(detections)}." + ) + image_metadata = detections.image_metadata or {} + if emit_polygons is None: + # Absent key -> True -> historical behaviour (byte-identical). Producer may + # set it to False to disable polygon emission. + emit_polygons = image_metadata.get(SERIALISE_POLYGONS_KEY, True) + detections_number = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(detections_number)] + class_names_mapping = None + if detections_number > 0: + class_names_mapping = image_metadata.get(CLASS_NAMES_KEY) + if class_names_mapping is None: + raise ValueError( + f"Serialising tensor-native detections, but " + f"`image_metadata['{CLASS_NAMES_KEY}']` is missing - the producer " + f"block must attach the class_id -> name mapping." + ) + boxes = detections.xyxy.detach().cpu().tolist() + confidences = detections.confidence.detach().cpu().tolist() + class_ids = [int(value) for value in detections.class_id.detach().cpu().tolist()] + serialized_detections = [] + kept_indices: List[int] = [] + for index in range(detections_number): + data = bboxes_metadata[index] + detection_dict = {} + x1, y1, x2, y2 = (float(coordinate) for coordinate in boxes[index]) + detection_dict[WIDTH_KEY] = abs(x2 - x1) + detection_dict[HEIGHT_KEY] = abs(y2 - y1) + detection_dict[X_KEY] = x1 + detection_dict[WIDTH_KEY] / 2 + detection_dict[Y_KEY] = y1 + detection_dict[HEIGHT_KEY] / 2 + detection_dict[CONFIDENCE_KEY] = float(confidences[index]) + detection_dict[CLASS_ID_KEY] = class_ids[index] + if isinstance(detections, InstanceDetections) and emit_polygons: + polygon = _resolve_instance_polygon( + mask=detections.mask, index=index, data=data + ) + if polygon is None: + # ignoring the whole instance - mirrors the numpy serialiser + continue + detection_dict[POLYGON_KEY] = [] + for x, y in polygon: + detection_dict[POLYGON_KEY].append( + { + X_KEY: float(x), + Y_KEY: float(y), + } + ) + if data.get("tracker_id") is not None: + detection_dict[TRACKER_ID_KEY] = int(data["tracker_id"]) + # C1: a producer may carry an arbitrary per-box label on the box metadata; + # prefer it, otherwise fall back to the class_id -> name mapping. + if CLASS_NAME_KEY in data: + detection_dict[CLASS_NAME_KEY] = str(data[CLASS_NAME_KEY]) + else: + detection_dict[CLASS_NAME_KEY] = _resolve_class_name( + class_id=class_ids[index], class_names_mapping=class_names_mapping + ) + if DETECTION_ID_KEY not in data: + raise ValueError( + f"Serialising tensor-native detections, but " + f"`bboxes_metadata['{DETECTION_ID_KEY}']` is missing for detection " + f"at index {index} - the producer block must attach it." + ) + detection_dict[DETECTION_ID_KEY] = str(data[DETECTION_ID_KEY]) + if PATH_DEVIATION_KEY_IN_SV_DETECTIONS in data: + detection_dict[PATH_DEVIATION_KEY_IN_INFERENCE_RESPONSE] = data[ + PATH_DEVIATION_KEY_IN_SV_DETECTIONS + ] + if TIME_IN_ZONE_KEY_IN_SV_DETECTIONS in data: + detection_dict[TIME_IN_ZONE_KEY_IN_INFERENCE_RESPONSE] = data[ + TIME_IN_ZONE_KEY_IN_SV_DETECTIONS + ] + if POLYGON_KEY_IN_SV_DETECTIONS in data: + detection_dict[POLYGON_KEY_IN_INFERENCE_RESPONSE] = ( + np.asarray(data[POLYGON_KEY_IN_SV_DETECTIONS]) + .astype(float) + .round() + .astype(int) + .tolist() + ) + if ( + BOUNDING_RECT_ANGLE_KEY_IN_SV_DETECTIONS in data + and BOUNDING_RECT_RECT_KEY_IN_SV_DETECTIONS in data + and BOUNDING_RECT_HEIGHT_KEY_IN_SV_DETECTIONS in data + and BOUNDING_RECT_WIDTH_KEY_IN_SV_DETECTIONS in data + ): + detection_dict[BOUNDING_RECT_ANGLE_KEY_IN_INFERENCE_RESPONSE] = data[ + BOUNDING_RECT_ANGLE_KEY_IN_SV_DETECTIONS + ] + detection_dict[BOUNDING_RECT_RECT_KEY_IN_INFERENCE_RESPONSE] = data[ + BOUNDING_RECT_RECT_KEY_IN_SV_DETECTIONS + ] + detection_dict[BOUNDING_RECT_HEIGHT_KEY_IN_INFERENCE_RESPONSE] = data[ + BOUNDING_RECT_HEIGHT_KEY_IN_SV_DETECTIONS + ] + detection_dict[BOUNDING_RECT_WIDTH_KEY_IN_INFERENCE_RESPONSE] = data[ + BOUNDING_RECT_WIDTH_KEY_IN_SV_DETECTIONS + ] + if PARENT_ID_KEY in image_metadata: + detection_dict[PARENT_ID_KEY] = str(image_metadata[PARENT_ID_KEY]) + # Add parent origin metadata if detection is based on a crop/slice + if ( + PARENT_ID_KEY in image_metadata + and ROOT_PARENT_ID_KEY in image_metadata + and str(image_metadata[PARENT_ID_KEY]) + != str(image_metadata[ROOT_PARENT_ID_KEY]) + ): + _attach_parent_metadata_to_detection_dict( + detection_dict=detection_dict, + data=image_metadata, + coordinates_key=PARENT_COORDINATES_KEY, + dimensions_key=PARENT_DIMENSIONS_KEY, + origin_key=PARENT_ORIGIN_KEY, + ) + detection_dict[ROOT_PARENT_ID_KEY] = str(image_metadata[ROOT_PARENT_ID_KEY]) + _attach_parent_metadata_to_detection_dict( + detection_dict=detection_dict, + data=image_metadata, + coordinates_key=ROOT_PARENT_COORDINATES_KEY, + dimensions_key=ROOT_PARENT_DIMENSIONS_KEY, + origin_key=ROOT_PARENT_ORIGIN_KEY, + ) + if ( + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS in data + and KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS in data + and KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS in data + and KEYPOINTS_XY_KEY_IN_SV_DETECTIONS in data + ): + kp_class_id = data[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS] + kp_class_name = data[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS] + kp_confidence = data[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS] + kp_xy = data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] + # Drop the trailing padding slots so we never emit fabricated + # keypoints for detections that carry fewer than the batch maximum. + kp_count = real_keypoints_count(kp_class_name, total=len(kp_xy)) + detection_dict[KEYPOINTS_KEY_IN_INFERENCE_RESPONSE] = [] + for ( + keypoint_class_id, + keypoint_class_name, + keypoint_confidence, + (x, y), + ) in zip( + kp_class_id[:kp_count], + kp_class_name[:kp_count], + kp_confidence[:kp_count], + kp_xy[:kp_count], + ): + detection_dict[KEYPOINTS_KEY_IN_INFERENCE_RESPONSE].append( + { + "class_id": int(keypoint_class_id), + "class": str(keypoint_class_name), + "confidence": float(keypoint_confidence), + "x": float(x), + "y": float(y), + } + ) + if DETECTED_CODE_KEY in data: + detection_dict[DETECTED_CODE_KEY] = data[DETECTED_CODE_KEY] + if VELOCITY_KEY_IN_SV_DETECTIONS in data: + detection_dict[VELOCITY_KEY_IN_INFERENCE_RESPONSE] = _to_plain_list( + data[VELOCITY_KEY_IN_SV_DETECTIONS] + ) + if SPEED_KEY_IN_SV_DETECTIONS in data: + detection_dict[SPEED_KEY_IN_INFERENCE_RESPONSE] = float( + data[SPEED_KEY_IN_SV_DETECTIONS] + ) + if SMOOTHED_VELOCITY_KEY_IN_SV_DETECTIONS in data: + detection_dict[SMOOTHED_VELOCITY_KEY_IN_INFERENCE_RESPONSE] = ( + _to_plain_list(data[SMOOTHED_VELOCITY_KEY_IN_SV_DETECTIONS]) + ) + if SMOOTHED_SPEED_KEY_IN_SV_DETECTIONS in data: + detection_dict[SMOOTHED_SPEED_KEY_IN_INFERENCE_RESPONSE] = float( + data[SMOOTHED_SPEED_KEY_IN_SV_DETECTIONS] + ) + if AREA_KEY_IN_SV_DETECTIONS in data: + detection_dict[AREA_KEY_IN_INFERENCE_RESPONSE] = float( + data[AREA_KEY_IN_SV_DETECTIONS] + ) + if AREA_CONVERTED_KEY_IN_SV_DETECTIONS in data: + detection_dict[AREA_CONVERTED_KEY_IN_INFERENCE_RESPONSE] = float( + data[AREA_CONVERTED_KEY_IN_SV_DETECTIONS] + ) + if NEAREST_TARGET_DISTANCE_KEY in data: + nearest_target_distance = data[NEAREST_TARGET_DISTANCE_KEY] + detection_dict[NEAREST_TARGET_DISTANCE_KEY] = ( + float(nearest_target_distance) + if nearest_target_distance is not None + else None + ) + serialized_detections.append(detection_dict) + kept_indices.append(index) + serialized_image_metadata = { + "width": None, + "height": None, + } + image_dimensions = image_metadata.get(IMAGE_DIMENSIONS_KEY) + if image_dimensions is not None: + serialized_image_metadata = { + "width": int(image_dimensions[1]), + "height": int(image_dimensions[0]), + } + return serialized_image_metadata, serialized_detections, kept_indices + + +def serialise_native_rle_detections(detections: InstanceDetections) -> dict: + """C3 native RLE serialiser for the rle-instance-seg and semantic-seg kinds. + + Emits the same per-box prediction dicts as ``serialise_sv_detections`` (re-using + its class-name / detection_id / image_metadata logic) but attaches the COCO RLE + pulled straight from the carried ``InstancesRLEMasks`` instead of a ``points`` + polygon. Polygon computation is skipped up front (``emit_polygons=False``) so the + O(N*H*W) mask->polygon work is never spent on a value that would only be popped; + as a consequence no instance is dropped for an empty contour and every box keeps + its RLE, attached by ORIGINAL row index. + """ + if not isinstance(detections, InstanceDetections): + raise ValueError( + f"serialise_native_rle_detections(...) expected " + f"`inference_models.InstanceDetections`, got {type(detections)}." + ) + if not isinstance(detections.mask, InstancesRLEMasks): + raise ValueError( + "serialise_native_rle_detections(...) requires the instance masks to be " + f"carried as `InstancesRLEMasks`, got {type(detections.mask)}." + ) + serialized_image_metadata, serialized_detections, kept_indices = ( + _serialise_sv_detections(detections=detections, emit_polygons=False) + ) + rle_masks = detections.mask.to_coco_rle_masks() + for detection_dict, original_index in zip(serialized_detections, kept_indices): + rle = rle_masks[original_index] + counts = rle.get("counts") + if isinstance(counts, bytes): + rle = {"size": rle["size"], "counts": counts.decode("utf-8")} + detection_dict[RLE_MASK_KEY_IN_INFERENCE_RESPONSE] = rle + return {"image": serialized_image_metadata, "predictions": serialized_detections} + + +# Exported under the legacy numpy name so existing loader/imports keep working; now +# points at the native implementation rather than the numpy alias. +serialise_rle_sv_detections = serialise_native_rle_detections + + +def _resolve_instance_polygon( + mask: Union[torch.Tensor, InstancesRLEMasks], + index: int, + data: dict, +) -> Optional[Any]: + declared_polygon = data.get(POLYGON_KEY_IN_SV_DETECTIONS) + if declared_polygon is not None and len(declared_polygon) > 2: + return declared_polygon + if isinstance(mask, InstancesRLEMasks): + instance_mask = coco_rle_masks_to_numpy_mask( + InstancesRLEMasks(image_size=mask.image_size, masks=[mask.masks[index]]) + )[0] + else: + instance_mask = mask[index].detach().cpu().numpy() + return mask_to_polygon(mask=instance_mask) + + +def _resolve_class_name(class_id: int, class_names_mapping: dict) -> str: + class_name = class_names_mapping.get(class_id) + if class_name is None: + raise ValueError( + f"Serialising tensor-native detections, class_id={class_id} is missing " + f"from the class_names mapping " + f"(keys present: {sorted(class_names_mapping.keys())})." + ) + return str(class_name) + + +def _to_plain_list(value: Any) -> list: + if hasattr(value, "tolist"): + return value.tolist() + return list(value) + + +def serialise_native_classification( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> dict: + """Serialise a native single-label ``ClassificationPrediction`` (single-row, + bs=1) or a native ``MultiLabelClassificationPrediction`` into the response dict + the numpy classification blocks produce flag-OFF. + + The ``class_id -> name`` map is read from the prediction's metadata + (``image_metadata[CLASS_NAMES_KEY]``). Single-label carries PLURAL + ``images_metadata`` (list, [0] used); multi-label carries SINGULAR + ``image_metadata`` (dict). + + IMPORTANT - two INCOMPATIBLE flag-OFF shapes share this one serialiser, because + the tensor-native ``CLASSIFICATION_PREDICTION_KIND`` is produced both by model + blocks and by hand-built formatter blocks: + + * "model" shape - roboflow multi_class / multi_label classification model blocks + (and visual_search_classifier). Flag-OFF these ``model_dump(by_alias=True, + exclude_none=True)`` a numpy ``*InferenceResponse`` then append + ``prediction_type`` / ``parent_id`` / ``root_parent_id``. Dict is + ``inference_id``-first (``time`` right after when present); single-label is + confidence-threshold FILTERED, ``round(..., 4)`` and SORTED desc; a + ``prediction_type`` key is appended. + * "formatter" shape - ``vlm_as_classifier`` (D4). Flag-OFF it hand-builds the + dict and passes it through UNSERIALISED, so it is ``image``-first, in INSERTION + (== class_id) order, UNROUNDED, with NO threshold filter and NO + ``prediction_type``. Key order single-label + ``image, predictions, top, confidence, inference_id, parent_id``; multi-label + ``image, predictions, predicted_classes, inference_id, parent_id``. + + A native prediction cannot self-describe which shape it wants, so the producer + stamps an explicit ``CLASSIFICATION_STYLE_KEY`` (``"model"`` / ``"formatter"``) + into ``image_metadata`` and the serialiser reads it (see + ``_wants_model_style_classification``, which falls back to the original + metadata heuristic when the key is absent). NOTE: the formatter shape still + cannot reproduce the numpy vlm ``class_id = -1`` for an out-of-list class, + because the native producer collapses such classes into a dense + ``len(classes)`` id before building the tensors (in-list classes ARE + byte-identical). + """ + if isinstance(prediction, ClassificationPrediction): + metadata_list = prediction.images_metadata or [{}] + image_metadata = metadata_list[0] or {} + elif isinstance(prediction, MultiLabelClassificationPrediction): + image_metadata = prediction.image_metadata or {} + else: + raise ValueError( + f"serialise_native_classification(...) expected " + f"`inference_models.ClassificationPrediction` or " + f"`inference_models.MultiLabelClassificationPrediction`, " + f"got {type(prediction)}." + ) + + class_names_mapping = image_metadata.get(CLASS_NAMES_KEY) + if class_names_mapping is None: + raise ValueError( + f"Serialising tensor-native classification, but " + f"`image_metadata['{CLASS_NAMES_KEY}']` is missing - the producer " + f"block must attach the class_id -> name mapping." + ) + + serialized_image_metadata = {"width": None, "height": None} + image_dimensions = image_metadata.get(IMAGE_DIMENSIONS_KEY) + if image_dimensions is not None: + serialized_image_metadata = { + "width": int(image_dimensions[1]), + "height": int(image_dimensions[0]), + } + + if _wants_model_style_classification(image_metadata): + return _serialise_classification_model_style( + prediction=prediction, + image_metadata=image_metadata, + class_names_mapping=class_names_mapping, + serialized_image_metadata=serialized_image_metadata, + ) + return _serialise_classification_formatter_style( + prediction=prediction, + image_metadata=image_metadata, + class_names_mapping=class_names_mapping, + serialized_image_metadata=serialized_image_metadata, + ) + + +def _wants_model_style_classification(image_metadata: dict) -> bool: + """D4 / lane 1b: choose which flag-OFF classification shape to reproduce. + + ``True`` -> numpy ``*InferenceResponse``-derived "model" shape. ``False`` -> + hand-built ``vlm_as_classifier`` "formatter" shape. + + Primary signal: the explicit ``CLASSIFICATION_STYLE_KEY`` the producer stamps + into ``image_metadata`` (``"model"`` / ``"formatter"``). This replaces the + original brittle heuristic, which inferred the shape from incidental metadata. + + Fallback (key ABSENT or an unrecognised value): the original heuristic. The vlm + formatter attaches only CLASS_NAMES / PREDICTION_TYPE / IMAGE_DIMENSIONS / + INFERENCE_ID / PARENT_ID, while the model producers additionally attach a + confidence threshold (multi_class model blocks), ``root_parent_id`` (all model + blocks + visual_search_classifier) and/or ``time`` - presence of ANY of the + three signals the model shape. Keeping the heuristic as a fallback makes wiring + the explicit key non-breaking: an un-wired producer, and the pinned key-ordering + tests (which set ``root_parent_id``), keep today's output. + """ + style = image_metadata.get(CLASSIFICATION_STYLE_KEY) + if style == CLASSIFICATION_STYLE_MODEL: + return True + if style == CLASSIFICATION_STYLE_FORMATTER: + return False + return ( + image_metadata.get(_CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY) is not None + or image_metadata.get(ROOT_PARENT_ID_KEY) is not None + or image_metadata.get(_TIME_KEY) is not None + ) + + +def _serialise_classification_model_style( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], + image_metadata: dict, + class_names_mapping: dict, + serialized_image_metadata: dict, +) -> dict: + """Numpy ``*InferenceResponse``-derived "model" shape (unchanged behaviour). + + Key order mirrors the numpy ``model_dump(by_alias=True, exclude_none=True)`` + followed by the block's in-place ``prediction_type`` / ``parent_id`` / + ``root_parent_id`` writes: ``inference_id`` first, ``time`` next when present, + then ``image``, the predictions section, and the appended lineage keys. + """ + result: dict = {} + # numpy dump order puts inference_id BEFORE image (InferenceResponse base + # fields precede CvInferenceResponse.image in the pydantic MRO). + if image_metadata.get(INFERENCE_ID_KEY) is not None: + result[INFERENCE_ID_KEY] = image_metadata[INFERENCE_ID_KEY] + # `time` sits between inference_id and image in the numpy dump + # (InferenceResponse declares inference_id, frame_id, time before + # CvInferenceResponse.image; frame_id is None-excluded in workflows). + if image_metadata.get(_TIME_KEY) is not None: + result[_TIME_KEY] = image_metadata[_TIME_KEY] + result["image"] = serialized_image_metadata + + if isinstance(prediction, ClassificationPrediction): + # confidence: (1, num_classes) full softmax; class_id: (1,) + confidence_vector = prediction.confidence.detach().cpu().reshape(-1).tolist() + # Mirror the numpy `prepare_classification_response` cutoff - drop classes + # whose (raw) score is below the resolved threshold when the producer attached + # it; otherwise keep the full softmax. Rounding/sort match the numpy + # `ClassificationInferenceResponse` (round(score, 4); sort desc by confidence). + confidence_threshold = image_metadata.get( + _CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY + ) + individual_classes_predictions = [] + for class_id, score in enumerate(confidence_vector): + class_score = float(score) + if confidence_threshold is not None and class_score < confidence_threshold: + continue + individual_classes_predictions.append( + { + "class": str(class_names_mapping.get(class_id, class_id)), + "class_id": class_id, + "confidence": round(class_score, 4), + } + ) + individual_classes_predictions = sorted( + individual_classes_predictions, + key=lambda item: item["confidence"], + reverse=True, + ) + result["predictions"] = individual_classes_predictions + result["top"] = ( + individual_classes_predictions[0]["class"] + if individual_classes_predictions + else "" + ) + result["confidence"] = ( + individual_classes_predictions[0]["confidence"] + if individual_classes_predictions + else 0.0 + ) + else: + # MultiLabel: confidence (num_classes,) sigmoid; class_ids = predicted ids + confidence_vector = prediction.confidence.detach().cpu().reshape(-1).tolist() + # Same opt-in cutoff as the multi-class branch above: producers that + # gap-fill a dense vector (e.g. visual_search_classifier) attach the + # threshold so their synthetic zero-confidence entries are dropped; + # model predictions attach no threshold and keep every class. + confidence_threshold = image_metadata.get( + _CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY + ) + predictions_dict = { + str(class_names_mapping.get(class_id, class_id)): { + "confidence": float(score), + "class_id": class_id, + } + for class_id, score in enumerate(confidence_vector) + if confidence_threshold is None or float(score) >= confidence_threshold + } + predicted_classes = [ + str(class_names_mapping.get(int(class_id), int(class_id))) + for class_id in prediction.class_ids.detach().cpu().tolist() + ] + result["predictions"] = predictions_dict + result["predicted_classes"] = predicted_classes + + # Append order mirrors the numpy block's in-place writes after model_dump: + # attach_prediction_type_info, then parent_id, then root_parent_id + # (inference_id was already emitted first, matching the dump order). + for source_key in ( + PREDICTION_TYPE_KEY, + PARENT_ID_KEY, + ROOT_PARENT_ID_KEY, + ): + if image_metadata.get(source_key) is not None: + result[source_key] = image_metadata[source_key] + return result + + +def _serialise_classification_formatter_style( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], + image_metadata: dict, + class_names_mapping: dict, + serialized_image_metadata: dict, +) -> dict: + """D4: reproduce the hand-built ``vlm_as_classifier`` flag-OFF dict byte-for-byte. + + Single-label key order: ``image, predictions, top, confidence, inference_id, + parent_id``. Multi-label: ``image, predictions, predicted_classes, inference_id, + parent_id``. Predictions are in INSERTION (== class_id) order, confidences are + UNROUNDED, there is no threshold filter and no ``prediction_type`` key. Per-entry + order matches the numpy vlm path (single-label ``class, class_id, confidence``; + multi-label ``confidence, class_id``). + + Residual divergences (reported, not fixable in the serialiser): (a) an out-of-list + class gets a dense ``len(classes)`` id from the producer instead of the numpy + ``-1`` (in-list classes are byte-identical); (b) confidences are read back through + a float32 tensor, so their least-significant digits can differ from the numpy + float. ``top`` / ``confidence`` are taken from the native argmax id, matching the + numpy path even when the top class was out-of-list. + """ + result: dict = {"image": serialized_image_metadata} + + if isinstance(prediction, ClassificationPrediction): + confidence_vector = prediction.confidence.detach().cpu().reshape(-1).tolist() + result["predictions"] = [ + { + "class": class_names_mapping.get(class_id, class_id), + "class_id": class_id, + "confidence": float(score), + } + for class_id, score in enumerate(confidence_vector) + ] + top_ids = prediction.class_id.detach().cpu().reshape(-1).tolist() + top_id = int(top_ids[0]) if top_ids else None + if top_id is not None and 0 <= top_id < len(confidence_vector): + result["top"] = class_names_mapping.get(top_id, top_id) + result["confidence"] = float(confidence_vector[top_id]) + else: + # Empty prediction - numpy vlm never emits this (there is always a top), + # but keep the numpy `ClassificationInferenceResponse` empty fallback. + result["top"] = "" + result["confidence"] = 0.0 + else: + confidence_vector = prediction.confidence.detach().cpu().reshape(-1).tolist() + result["predictions"] = { + class_names_mapping.get(class_id, class_id): { + "confidence": float(score), + "class_id": class_id, + } + for class_id, score in enumerate(confidence_vector) + } + result["predicted_classes"] = [ + class_names_mapping.get(int(class_id), int(class_id)) + for class_id in prediction.class_ids.detach().cpu().tolist() + ] + + # Trailing keys, matching the numpy vlm dict (always present there). + result[INFERENCE_ID_KEY] = image_metadata.get(INFERENCE_ID_KEY) + result[PARENT_ID_KEY] = image_metadata.get(PARENT_ID_KEY) + return result + + +def serialise_native_keypoint_detection( + prediction: Tuple[KeyPoints, Optional[Detections]], +) -> dict: + """Tuple-aware serialiser for the keypoint-detection kind. The native prediction + is a ``(KeyPoints, Detections)`` tuple; the per-instance keypoint payload is + carried on the bbox ``Detections`` ``bboxes_metadata``, so unwrap and defer to + ``serialise_sv_detections``.""" + if not isinstance(prediction, tuple) or len(prediction) != 2: + raise ValueError( + f"serialise_native_keypoint_detection(...) expected a " + f"Tuple[KeyPoints, Detections], got {type(prediction)}." + ) + _, detections = prediction + if detections is None: + raise ValueError( + "Keypoint prediction is missing the bounding-box component required for " + "serialisation." + ) + return serialise_sv_detections(detections) + + +def serialise_native_embedding(value: torch.Tensor) -> list: + """C4 serialiser for the embedding kind. The native CLIP/PE embedding is a + (possibly CUDA) ``torch.Tensor``; emit the numpy-faithful ``List[float]``.""" + return value.detach().cpu().tolist() + + +def serialise_native_tensor(value: torch.Tensor) -> list: + """C4 serialiser for the tensor kind (e.g. raw model tensors). The native + value is a (possibly CUDA/MPS) ``torch.Tensor``; emit a JSON-serialisable + nested list.""" + return value.detach().cpu().tolist() + + +def serialise_numpy_array_kind(value: Any) -> Any: + """Flag-on serialiser for the ``numpy_array`` kind (e.g. depth-estimation + ``normalized_depth``). Flag-off the kind has NO serialiser at all - + ``np.ndarray`` payloads pass through raw and the HTTP layer serialises + them - so non-tensor values are returned untouched here. Tensor-native + producers carry a ``torch.Tensor`` under the same kind name; materialise it + to the flag-off-identical ``np.ndarray`` so the wire format matches in both + flag directions.""" + if isinstance(value, torch.Tensor): + return value.detach().cpu().numpy() + return value + + +def serialize_wildcard_kind(value: Any) -> Any: + """Tensor-aware wildcard (`*`) serialiser โ€” same public name as the numpy + sibling in ``serializers.py``; the loader swaps the symbol under + ``ENABLE_TENSOR_DATA_REPRESENTATION``. + + Extends the numpy contract with arms for the native tensor-path values that + can legitimately reach a wildcard output โ€” native detections, the + ``(KeyPoints, Detections)`` keypoint tuple, classification predictions + (e.g. wildcard outputs of ``legacy_compatibility`` dynamic blocks after + boundary conversion) and bare ``torch.Tensor`` values (serialised exactly + like the ``tensor`` kind). Legacy values keep the numpy behavior verbatim, + including the ``sv.Detections`` arm (defensive: non-swapped paths can still + route sv under the flag) โ€” routed to the NUMPY detections serialiser, since + this module's same-name ``serialise_sv_detections`` consumes native objects. + Tuples other than the keypoint prediction pass through untouched, mirroring + the numpy contract. + """ + if isinstance(value, WorkflowImageData): + return serialise_image(image=value) + if isinstance(value, (Detections, InstanceDetections)): + return serialise_sv_detections(value) + if _is_native_key_point_prediction(value): + # A (KeyPoints, None) tuple raises inside the kind serialiser โ€” same + # loud behavior as a declared keypoint output missing its bbox component. + return serialise_native_keypoint_detection(prediction=value) + if isinstance( + value, (ClassificationPrediction, MultiLabelClassificationPrediction) + ): + return serialise_native_classification(prediction=value) + if isinstance(value, torch.Tensor): + return serialise_native_tensor(value=value) + if isinstance(value, dict): + return {key: serialize_wildcard_kind(value=item) for key, item in value.items()} + if isinstance(value, list): + return [serialize_wildcard_kind(value=element) for element in value] + if isinstance(value, sv.Detections): + return _serialise_legacy_sv_detections(detections=value) + if isinstance(value, datetime): + return serialize_timestamp(timestamp=value) + return value + + +def _is_native_key_point_prediction(value: Any) -> bool: + return ( + isinstance(value, tuple) + and len(value) == 2 + and isinstance(value[0], KeyPoints) + and (value[1] is None or isinstance(value[1], (Detections, InstanceDetections))) + ) diff --git a/inference/core/workflows/core_steps/common/tensor_native.py b/inference/core/workflows/core_steps/common/tensor_native.py new file mode 100644 index 0000000000..c11fc9d18d --- /dev/null +++ b/inference/core/workflows/core_steps/common/tensor_native.py @@ -0,0 +1,1058 @@ +"""Shared helpers for tensor-native predictions (the inference_models dataclasses +used under ENABLE_TENSOR_DATA_REPRESENTATION). + +These consolidate logic that otherwise gets copy-pasted into every tensor-native +block sibling: selecting a subset of detections by boolean mask or index list, +and normalising the keypoint-detection input shape. + +Supported prediction shapes: +- ``inference_models.Detections`` (object detection) +- ``inference_models.InstanceDetections`` (instance segmentation; dense or RLE masks) +- ``inference_models.KeyPoints`` (keypoints, standalone) +- ``Tuple[KeyPoints, Optional[Detections]]`` (keypoint-detection workflow kind) +""" + +from typing import Any, Dict, List, Optional, Sequence, Tuple, Union +from uuid import uuid4 + +import numpy as np +import torch +from pycocotools import mask as mask_utils +from supervision.config import ORIENTED_BOX_COORDINATES + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.execution_engine.constants import ( + CLASS_ID_KEY, + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + CONFIDENCE_KEY, + DETECTION_ID_KEY, + HEIGHT_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + POLYGON_KEY_IN_SV_DETECTIONS, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, + WIDTH_KEY, + X_KEY, + Y_KEY, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import coco_rle_masks_to_numpy_mask + +TensorNativeDetections = Union[Detections, InstanceDetections] +KeyPointPrediction = Tuple[KeyPoints, Optional[Detections]] +TensorNativePrediction = Union[ + Detections, InstanceDetections, KeyPoints, KeyPointPrediction +] + +# --- Per-box host mirror of the detection tensors ----------------------------- +# +# ``attach_native_detection_metadata`` stores a host-side copy of each box's +# ``xyxy`` / ``class_id`` / ``confidence`` in its ``bboxes_metadata`` entry. +# Rationale: visualization blocks need these tiny values on the host, but a +# ``.cpu()`` there queues behind unrelated kernels on the default CUDA stream +# (measured up to ~65 ms per batch on Jetson under load) even though the values +# are ready; at attach time (model block, execution-engine thread, queue empty) +# the same read costs microseconds. Consumers PREFER the mirror when every box +# carries it and fall back to tensor reads otherwise (remote/deserialized/ +# transformed predictions). +# +# The keys are underscore-private because two sv-view builders generically copy +# arbitrary ``bboxes_metadata`` keys into ``sv.Detections.data`` +# (``base_tensor._materialise_data`` and +# ``representation_boundary._attach_per_box_columns``) โ€” both explicitly +# exclude these keys. The wire serializers (``serializers_tensor``) copy only +# specific keys, so serialized output is unaffected by construction. +# +# STALENESS CONTRACT: any code that changes ``xyxy`` / ``class_id`` / +# ``confidence`` of a prediction while carrying its ``bboxes_metadata`` forward +# MUST either re-run ``attach_native_detection_metadata`` (which overwrites the +# mirror from the current tensors) or drop the mirror via +# ``strip_host_mirror_metadata`` โ€” a stale mirror would be silently preferred. +HOST_XYXY_KEY = "_host_xyxy" +HOST_CLASS_ID_KEY = "_host_class_id" +HOST_CONFIDENCE_KEY = "_host_confidence" +HOST_MIRROR_KEYS = (HOST_XYXY_KEY, HOST_CLASS_ID_KEY, HOST_CONFIDENCE_KEY) + + +def read_host_mirror( + bboxes_metadata: Optional[List[dict]], + expected_rows: int, +) -> Optional[Tuple[np.ndarray, np.ndarray, np.ndarray]]: + """Assemble ``(xyxy float32 (N, 4), class_id int (N,), confidence float32 + (N,))`` from the per-box host mirror WITHOUT touching any tensor. + + Returns ``None`` (caller falls back to device reads) unless ``bboxes_metadata`` + has exactly ``expected_rows`` entries and EVERY entry carries all three + mirror keys with a well-formed 4-element ``xyxy``. The dtypes match the + tensor-read path (``.cpu().numpy().astype(...)``) exactly, so downstream + output is bit-identical either way.""" + if bboxes_metadata is None or len(bboxes_metadata) != expected_rows: + return None + if expected_rows == 0: + return None + if not all( + isinstance(entry, dict) and all(key in entry for key in HOST_MIRROR_KEYS) + for entry in bboxes_metadata + ): + return None + try: + xyxy = np.asarray( + [entry[HOST_XYXY_KEY] for entry in bboxes_metadata], dtype=np.float32 + ).reshape(expected_rows, 4) + class_id = np.asarray( + [int(entry[HOST_CLASS_ID_KEY]) for entry in bboxes_metadata], dtype=int + ) + confidence = np.asarray( + [float(entry[HOST_CONFIDENCE_KEY]) for entry in bboxes_metadata], + dtype=np.float32, + ) + except (TypeError, ValueError): + return None + return xyxy, class_id, confidence + + +def strip_host_mirror_metadata( + bboxes_metadata: Optional[List[dict]], +) -> Optional[List[dict]]: + """Return ``bboxes_metadata`` without the per-box host-mirror keys. + + Call whenever ``xyxy`` / ``class_id`` / ``confidence`` are recomputed + without re-running ``attach_native_detection_metadata`` โ€” consumers then + fall back to tensor reads instead of trusting a stale mirror. Entries that + carry mirror keys are shallow-copied (caller-shared dicts are never + mutated); mirror-free entries are reused as-is. ``None`` passes through.""" + if bboxes_metadata is None: + return None + stripped: List[dict] = [] + for entry in bboxes_metadata: + if entry and any(key in entry for key in HOST_MIRROR_KEYS): + entry = { + key: value + for key, value in entry.items() + if key not in HOST_MIRROR_KEYS + } + stripped.append(entry) + return stripped + + +def mask_to_indices(mask: Union[np.ndarray, Sequence[bool]]) -> List[int]: + """Convert a boolean mask (numpy array, list, or torch tensor) into the list + of surviving row indices, in ascending order.""" + if isinstance(mask, torch.Tensor): + mask = mask.detach().to("cpu").numpy() + return np.nonzero(np.asarray(mask))[0].tolist() + + +class _RowSelection: + """Normalised row selection shared by the ``take_*`` helpers. + + Tensor fields are selected with a torch selector (a boolean mask stays a + boolean mask - it is moved to the field's device ONCE and reused, so a + device-resident mask never round-trips through a python list). Python-list + fields (``bboxes_metadata`` / RLE masks / ``key_points_metadata``) need + concrete row positions, which are pulled to host lazily and cached - a + selection that only touches tensor fields never pays the transfer. + """ + + def __init__( + self, + selector: Optional[torch.Tensor], + host_indices: Optional[List[int]], + is_identity: bool, + total_rows: int, + ) -> None: + self._selector = selector + self._host_indices = host_indices + self.is_identity = is_identity + self._total_rows = total_rows + + @classmethod + def from_indices(cls, indices: Sequence[int], total_rows: int) -> "_RowSelection": + indices = [int(index) for index in indices] + is_identity = indices == list(range(total_rows)) + selector = None if is_identity else torch.as_tensor(indices, dtype=torch.long) + return cls(selector, indices, is_identity, total_rows) + + @classmethod + def from_mask( + cls, + mask: Union[np.ndarray, Sequence[bool], torch.Tensor], + total_rows: int, + ) -> "_RowSelection": + if not isinstance(mask, torch.Tensor) or mask.numel() != total_rows: + # Host-born masks (and the legacy mismatched-length edge) keep the + # exact index-list semantics of the old implementation. + return cls.from_indices(mask_to_indices(mask), total_rows) + mask = mask.detach().to(dtype=torch.bool).reshape(-1) + if bool(mask.all()): + return cls(None, None, True, total_rows) + return cls(mask, None, False, total_rows) + + def select_tensor(self, field: torch.Tensor) -> torch.Tensor: + if self.is_identity: + return field + if self._selector.device != field.device: + # Move once, reuse for every subsequent field of this prediction. + self._selector = self._selector.to(field.device) + return field[self._selector] + + def host_indices(self) -> List[int]: + if self._host_indices is None: + if self.is_identity: + self._host_indices = list(range(self._total_rows)) + else: + # The single (lazy) device->host transfer of this selection. + self._host_indices = ( + torch.nonzero(self._selector, as_tuple=False) + .reshape(-1) + .cpu() + .tolist() + ) + return self._host_indices + + +def _take_detections( + detections: TensorNativeDetections, + selection: _RowSelection, +) -> TensorNativeDetections: + """Shared gather body for ``Detections`` / ``InstanceDetections``. + + Per-detection state (``bboxes_metadata``) and masks (dense torch or RLE) are + carried over for the surviving rows; ``image_metadata`` is shared as-is. + The surviving ``bboxes_metadata`` dicts are COPIED (not shared by reference) + so a downstream block that mutates a selected box's metadata (e.g. assigns a + ``tracker_id``) cannot leak the mutation back into the source prediction. + An identity selection skips the tensor gathers entirely. + """ + bboxes_metadata = None + if detections.bboxes_metadata is not None: + bboxes_metadata = [ + dict(detections.bboxes_metadata[i]) for i in selection.host_indices() + ] + if isinstance(detections, InstanceDetections): + mask_field = detections.mask + if isinstance(mask_field, InstancesRLEMasks): + new_mask: Union[torch.Tensor, InstancesRLEMasks] = InstancesRLEMasks( + image_size=mask_field.image_size, + masks=[mask_field.masks[i] for i in selection.host_indices()], + ) + else: + new_mask = selection.select_tensor(mask_field) + return InstanceDetections( + xyxy=selection.select_tensor(detections.xyxy), + class_id=selection.select_tensor(detections.class_id), + confidence=selection.select_tensor(detections.confidence), + mask=new_mask, + image_metadata=detections.image_metadata, + bboxes_metadata=bboxes_metadata, + ) + return Detections( + xyxy=selection.select_tensor(detections.xyxy), + class_id=selection.select_tensor(detections.class_id), + confidence=selection.select_tensor(detections.confidence), + image_metadata=detections.image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def _take_key_points( + key_points: KeyPoints, + selection: _RowSelection, +) -> KeyPoints: + """Shared gather body for ``KeyPoints`` (instance dimension). + + The surviving ``key_points_metadata`` dicts are COPIED (not shared by + reference) so downstream mutation cannot leak back into the source; the + auxiliary per-instance tensors (populated by RF-DETR) are sliced in + lockstep so lossless selections stay lossless.""" + key_points_metadata = None + if key_points.key_points_metadata is not None: + key_points_metadata = [ + dict(key_points.key_points_metadata[i]) for i in selection.host_indices() + ] + covariance = key_points.covariance + if covariance is not None: + covariance = selection.select_tensor(covariance) + detection_confidence = key_points.detection_confidence + if detection_confidence is not None: + detection_confidence = selection.select_tensor(detection_confidence) + return KeyPoints( + xy=selection.select_tensor(key_points.xy), + class_id=selection.select_tensor(key_points.class_id), + confidence=selection.select_tensor(key_points.confidence), + image_metadata=key_points.image_metadata, + key_points_metadata=key_points_metadata, + covariance=covariance, + detection_confidence=detection_confidence, + ) + + +def take_detections_by_indices( + detections: TensorNativeDetections, + indices: Sequence[int], +) -> TensorNativeDetections: + """Select rows of a ``Detections`` / ``InstanceDetections`` by index list.""" + selection = _RowSelection.from_indices(indices, int(detections.xyxy.shape[0])) + return _take_detections(detections, selection) + + +def take_key_points_by_indices( + key_points: KeyPoints, + indices: Sequence[int], +) -> KeyPoints: + """Select instances of a ``KeyPoints`` by index list.""" + selection = _RowSelection.from_indices(indices, int(key_points.xy.shape[0])) + return _take_key_points(key_points, selection) + + +def _prediction_row_count(prediction: TensorNativePrediction) -> int: + if isinstance(prediction, tuple): + key_points, _ = prediction + return int(key_points.xy.shape[0]) + if isinstance(prediction, KeyPoints): + return int(prediction.xy.shape[0]) + return int(prediction.xyxy.shape[0]) + + +def _take_prediction( + prediction: TensorNativePrediction, + selection: _RowSelection, +) -> TensorNativePrediction: + if isinstance(prediction, tuple): + key_points, detections = prediction + sliced_key_points = _take_key_points(key_points, selection) + sliced_detections = ( + _take_detections(detections, selection) if detections is not None else None + ) + return sliced_key_points, sliced_detections + if isinstance(prediction, KeyPoints): + return _take_key_points(prediction, selection) + return _take_detections(prediction, selection) + + +def take_prediction_by_indices( + prediction: TensorNativePrediction, + indices: Sequence[int], +) -> TensorNativePrediction: + """Select a subset of any tensor-native prediction by index list, returning + the same shape. For the keypoint-detection tuple, both the ``KeyPoints`` and + the bbox ``Detections`` components are sliced consistently.""" + selection = _RowSelection.from_indices(indices, _prediction_row_count(prediction)) + return _take_prediction(prediction, selection) + + +def take_prediction_by_mask( + prediction: TensorNativePrediction, + mask: Union[np.ndarray, Sequence[bool], torch.Tensor], +) -> TensorNativePrediction: + """Select a subset of any tensor-native prediction by boolean mask. + + A torch mask is used AS the selector: tensor fields are boolean-indexed on + their own device with no host round trip; the row positions only reach the + host (once, lazily) when a python-list field of the prediction needs them. + Non-torch masks (and torch masks whose length does not match the + prediction) keep the legacy index-list semantics.""" + selection = _RowSelection.from_mask(mask, _prediction_row_count(prediction)) + return _take_prediction(prediction, selection) + + +def split_key_point_prediction( + prediction: Union[TensorNativeDetections, KeyPointPrediction], +) -> Tuple[Optional[KeyPoints], TensorNativeDetections]: + """Normalise a block input to its bounding-box component. + + Returns ``(key_points, detections)`` where ``key_points`` is the + ``KeyPoints`` component when the input is a keypoint-detection tuple + (``(KeyPoints, Detections)``), otherwise ``None``. Blocks that only need + bounding boxes operate on the returned ``detections`` and re-wrap with the + returned ``key_points`` to preserve keypoints downstream. Raises if a + keypoint tuple lacks its bbox component. + """ + if isinstance(prediction, tuple): + key_points, detections = prediction + if detections is None: + raise ValueError( + "Keypoint prediction is missing the bounding-box component " + "required by this block." + ) + return key_points, detections + return None, prediction + + +def _read_root_coordinates_shift( + image_metadata: Optional[dict], +) -> Optional[Tuple[float, float]]: + """Return the ``(shift_x, shift_y)`` offset that maps this prediction's local + coordinates back to the root (workflow input) image, or ``None`` when no shift + is needed. + + The offset is the crop origin recorded at + ``image_metadata[ROOT_PARENT_COORDINATES_KEY] = [left_top_x, left_top_y]`` by + ``build_native_image_metadata``. ``None`` is returned (a no-op) only when the + key is absent or the prediction is provably root-anchored already. A zero + offset alone is NOT that proof: a crop taken at ``(0, 0)`` still has + crop-sized dimensions and its own lineage, and its masks still need + re-embedding onto the root canvas - so identity additionally requires the + image dimensions to match the root dimensions and the parent to BE the root. + A ``(0.0, 0.0)`` return means "no coordinate shift, but the conversion + (mask re-embedding + metadata rewrite) must still run". + """ + if not image_metadata: + return None + root_coordinates = image_metadata.get(ROOT_PARENT_COORDINATES_KEY) + if not root_coordinates: + return None + shift_x, shift_y = float(root_coordinates[0]), float(root_coordinates[1]) + if shift_x != 0.0 or shift_y != 0.0: + return shift_x, shift_y + image_dimensions = image_metadata.get(IMAGE_DIMENSIONS_KEY) + root_dimensions = image_metadata.get(ROOT_PARENT_DIMENSIONS_KEY) + dimensions_confirm_root = ( + image_dimensions is None + or root_dimensions is None + or list(image_dimensions) == list(root_dimensions) + ) + parent_is_root = image_metadata.get(PARENT_ID_KEY) == image_metadata.get( + ROOT_PARENT_ID_KEY + ) + if dimensions_confirm_root and parent_is_root: + return None + return shift_x, shift_y + + +def _native_image_metadata_in_root_coordinates(image_metadata: dict) -> dict: + """Copy ``image_metadata`` and rewrite the lineage so it describes a + root-anchored prediction: the parent/root origin offsets collapse to ``[0, 0]`` + and the reported image dimensions become the root-parent dimensions. Mirrors + the metadata rewrite that ``sv_detections_to_root_coordinates`` performs via + ``attach_parent_coordinates_to_detections`` (utils.py).""" + new_metadata = dict(image_metadata) + root_dimensions = new_metadata.get(ROOT_PARENT_DIMENSIONS_KEY) + if root_dimensions is not None: + new_metadata[IMAGE_DIMENSIONS_KEY] = list(root_dimensions) + new_metadata[PARENT_COORDINATES_KEY] = [0, 0] + new_metadata[ROOT_PARENT_COORDINATES_KEY] = [0, 0] + if ROOT_PARENT_ID_KEY in new_metadata: + new_metadata[PARENT_ID_KEY] = new_metadata[ROOT_PARENT_ID_KEY] + if root_dimensions is not None: + new_metadata[PARENT_DIMENSIONS_KEY] = list(root_dimensions) + return new_metadata + + +# Per-box geometry payloads shifted alongside xyxy on root conversion - the same +# set the crop side localizes (dynamic_crop subtracts the crop origin from +# keypoints, polygons AND oriented-box corners; root conversion adds it back). +# The numpy ``sv_detections_to_root_coordinates`` shifts the same three keys. +_GEOMETRY_KEYS_SHIFTED_TO_ROOT = ( + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + POLYGON_KEY_IN_SV_DETECTIONS, + ORIENTED_BOX_COORDINATES, +) + + +def _add_offset( + value: Union[List, np.ndarray], offset_xy: np.ndarray +) -> Union[List, np.ndarray]: + """Add ``offset_xy`` ([x, y]) to a coordinate container - the exact inverse of + ``dynamic_crop/v1_tensor._subtract_offset``, with the same container- and + dtype-preservation rules: python lists stay lists, numpy arrays stay arrays, + integer coordinates stay integers (crop origins are integral, so the shift is + lossless), floats keep float precision. Empty/None payloads pass through.""" + if value is None: + return value + was_list = isinstance(value, list) + array = np.asarray(value) + if array.size == 0: + return value + is_integer = np.issubdtype(array.dtype, np.integer) + if is_integer: + shifted = array.astype(np.int64) + offset_xy.astype(np.int64) + else: + shifted = array.astype(float) + offset_xy + return shifted.tolist() if was_list else shifted + + +def _shift_bboxes_metadata_to_root_coordinates( + bboxes_metadata: Optional[List[dict]], + shift_x: float, + shift_y: float, +) -> Optional[List[dict]]: + """Copy the per-box metadata entries, shifting the geometry payloads the + tensor-native serialiser reads back (``keypoints_xy``, ``polygon``) - the + mirror of the per-row shifts numpy's ``sv_detections_to_root_coordinates`` + applies to ``sv.Detections.data``. The per-box host mirror is dropped: the + ``xyxy`` tensor is shifted by the caller, so a carried mirror would be + stale โ€” consumers fall back to tensor reads.""" + if bboxes_metadata is None: + return None + offset_xy = np.asarray([shift_x, shift_y]) + shifted_entries = [] + for entry in bboxes_metadata: + entry = dict(entry) + for key in _GEOMETRY_KEYS_SHIFTED_TO_ROOT: + if key in entry: + entry[key] = _add_offset(entry[key], offset_xy) + for key in HOST_MIRROR_KEYS: + entry.pop(key, None) + shifted_entries.append(entry) + return shifted_entries + + +def _shift_native_masks_to_root_coordinates( + mask: Union[torch.Tensor, InstancesRLEMasks, None], + shift_x: float, + shift_y: float, + root_dimensions: Optional[Sequence[int]], +) -> Union[torch.Tensor, InstancesRLEMasks, None]: + """Re-anchor crop-local instance masks onto the root-image canvas, mirroring + the numpy path's paste into an ``np.full((H_root, W_root), False)`` base. + Dense torch masks are pasted into a zeros canvas on the same device; RLE + masks are re-embedded without densifying the full canvas via + ``embed_rle_masks_in_larger_canvas``.""" + if mask is None: + return None + if root_dimensions is None: + raise ValueError( + "Cannot shift tensor-native instance masks to root coordinates: the " + f"prediction carries a non-zero root offset but no " + f"`image_metadata['{ROOT_PARENT_DIMENSIONS_KEY}']` to size the root " + "canvas. The producer block must attach the root lineage keys " + "(see build_native_image_metadata)." + ) + root_height, root_width = int(root_dimensions[0]), int(root_dimensions[1]) + offset_x, offset_y = int(shift_x), int(shift_y) + if isinstance(mask, InstancesRLEMasks): + return embed_rle_masks_in_larger_canvas( + masks=mask, + offset_xy=(offset_x, offset_y), + target_size_hw=(root_height, root_width), + ) + _, mask_height, mask_width = mask.shape + if offset_x < 0 or offset_y < 0: + raise ValueError( + f"Cannot shift tensor-native instance masks to root coordinates: got " + f"negative crop offset ({offset_x}, {offset_y}); offsets must be " + "non-negative." + ) + if offset_x + mask_width > root_width or offset_y + mask_height > root_height: + raise ValueError( + f"Crop-local masks of size (h={mask_height}, w={mask_width}) at offset " + f"(x={offset_x}, y={offset_y}) do not fit into the root canvas of size " + f"(H={root_height}, W={root_width})." + ) + anchored = torch.zeros( + (mask.shape[0], root_height, root_width), dtype=mask.dtype, device=mask.device + ) + anchored[:, offset_y : offset_y + mask_height, offset_x : offset_x + mask_width] = ( + mask + ) + return anchored + + +def _shift_native_detections_to_root_coordinates( + detections: TensorNativeDetections, + shift_x: float, + shift_y: float, +) -> TensorNativeDetections: + shift = torch.as_tensor( + [shift_x, shift_y, shift_x, shift_y], + dtype=detections.xyxy.dtype, + device=detections.xyxy.device, + ) + shifted_xyxy = detections.xyxy + shift + # Root canvas dims must be read BEFORE the metadata rewrite below collapses + # the lineage to the root frame ([0, 0] offsets, root dims everywhere). + root_dimensions = ( + detections.image_metadata.get(ROOT_PARENT_DIMENSIONS_KEY) + if detections.image_metadata is not None + else None + ) + image_metadata = ( + _native_image_metadata_in_root_coordinates(detections.image_metadata) + if detections.image_metadata is not None + else None + ) + bboxes_metadata = _shift_bboxes_metadata_to_root_coordinates( + bboxes_metadata=detections.bboxes_metadata, + shift_x=shift_x, + shift_y=shift_y, + ) + if isinstance(detections, InstanceDetections): + return InstanceDetections( + xyxy=shifted_xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + mask=_shift_native_masks_to_root_coordinates( + mask=detections.mask, + shift_x=shift_x, + shift_y=shift_y, + root_dimensions=root_dimensions, + ), + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + return Detections( + xyxy=shifted_xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def _shift_native_key_points_to_root_coordinates( + key_points: KeyPoints, + shift_x: float, + shift_y: float, +) -> KeyPoints: + shift = torch.as_tensor( + [shift_x, shift_y], + dtype=key_points.xy.dtype, + device=key_points.xy.device, + ) + image_metadata = ( + _native_image_metadata_in_root_coordinates(key_points.image_metadata) + if key_points.image_metadata is not None + else None + ) + key_points_metadata = ( + [dict(entry) for entry in key_points.key_points_metadata] + if key_points.key_points_metadata is not None + else None + ) + return KeyPoints( + xy=key_points.xy + shift, + class_id=key_points.class_id, + confidence=key_points.confidence, + image_metadata=image_metadata, + key_points_metadata=key_points_metadata, + # Auxiliary RF-DETR tensors ride along unchanged - a pure translation + # affects neither positional covariance nor detection confidence. + covariance=key_points.covariance, + detection_confidence=key_points.detection_confidence, + ) + + +def native_detections_to_root_coordinates( + prediction: TensorNativePrediction, +) -> TensorNativePrediction: + """Shift a tensor-native prediction from its crop-local coordinates back to the + root (workflow input) image coordinates, returning a copy. + + Reads the crop origin from + ``image_metadata[ROOT_PARENT_COORDINATES_KEY] = [shift_x, shift_y]`` and adds + ``[shift_x, shift_y, shift_x, shift_y]`` to ``xyxy`` (and ``[shift_x, shift_y]`` + to keypoint ``xy``). Instance masks are re-anchored onto the root-image canvas + (dense torch masks pasted into a zeros canvas; RLE masks re-embedded without + densifying), and the per-box ``keypoints_xy`` / ``polygon`` payloads in + ``bboxes_metadata`` are shifted - matching the mask paste and the ``data``-key + shifts of numpy's ``sv_detections_to_root_coordinates``. The keypoint-detection + tuple ``(KeyPoints, Detections)`` has both components shifted consistently. + This is the tensor-native mirror of ``sv_detections_to_root_coordinates`` + (utils.py), used by the execution-engine output constructor when an output is + requested in PARENT coordinates. + + No-op (the input is returned unchanged) when the prediction carries no root + offset โ€” i.e. the key is absent or the shift is ``[0, 0]`` (already + root-anchored, e.g. a model run directly on the workflow input image). + """ + if isinstance(prediction, tuple): + key_points, detections = prediction + shift = _read_root_coordinates_shift( + key_points.image_metadata if key_points is not None else None + ) + if shift is None and detections is not None: + shift = _read_root_coordinates_shift(detections.image_metadata) + if shift is None: + return prediction + shift_x, shift_y = shift + shifted_key_points = ( + _shift_native_key_points_to_root_coordinates(key_points, shift_x, shift_y) + if key_points is not None + else None + ) + shifted_detections = ( + _shift_native_detections_to_root_coordinates(detections, shift_x, shift_y) + if detections is not None + else None + ) + return shifted_key_points, shifted_detections + if isinstance(prediction, KeyPoints): + shift = _read_root_coordinates_shift(prediction.image_metadata) + if shift is None: + return prediction + return _shift_native_key_points_to_root_coordinates(prediction, *shift) + shift = _read_root_coordinates_shift(prediction.image_metadata) + if shift is None: + return prediction + return _shift_native_detections_to_root_coordinates(prediction, *shift) + + +def instance_mask_to_numpy( + detections: InstanceDetections, + index: int, +) -> np.ndarray: + """Materialise a single instance's mask as a 2-D ``np.ndarray`` of bool + ``(H, W)``. RLE masks are decoded one instance at a time so the full stack + is never materialised at once (the same convention as the serialiser). + + The dense seg-adapter mask is already ``torch.bool`` (binarised via + ``.gt_(threshold).to(dtype=torch.bool)`` in the post-processing path), so no + ``.astype(bool)`` is applied โ€” the serialiser's dense branch makes the same + assumption. + """ + mask = detections.mask + if isinstance(mask, InstancesRLEMasks): + return coco_rle_masks_to_numpy_mask( + InstancesRLEMasks(image_size=mask.image_size, masks=[mask.masks[index]]) + )[0] + return mask[index].detach().to("cpu").numpy() + + +def build_native_image_metadata( + image: WorkflowImageData, + class_names: Dict[int, str], + prediction_type: str, + inference_id: Optional[str] = None, +) -> dict: + """Build the per-image ``image_metadata`` dict carried by a tensor-native + ``Detections`` prediction produced by a model block. + + Holds the ``class_id -> name`` map (required by the tensor-native serialiser), + the image dimensions, and the parent/root lineage needed for crop-aware + coordinate recovery downstream. The image shape is read without forcing a + device->host materialization (so tensor-only inputs stay on device). + """ + height, width = image._read_shape_without_materialization() + parent = image.parent_metadata + root = image.workflow_root_ancestor_metadata + parent_coordinates = parent.origin_coordinates + root_coordinates = root.origin_coordinates + metadata = { + CLASS_NAMES_KEY: class_names, + PREDICTION_TYPE_KEY: prediction_type, + IMAGE_DIMENSIONS_KEY: [height, width], + PARENT_ID_KEY: parent.parent_id, + PARENT_COORDINATES_KEY: [ + parent_coordinates.left_top_x, + parent_coordinates.left_top_y, + ], + PARENT_DIMENSIONS_KEY: [ + parent_coordinates.origin_height, + parent_coordinates.origin_width, + ], + ROOT_PARENT_ID_KEY: root.parent_id, + ROOT_PARENT_COORDINATES_KEY: [ + root_coordinates.left_top_x, + root_coordinates.left_top_y, + ], + ROOT_PARENT_DIMENSIONS_KEY: [ + root_coordinates.origin_height, + root_coordinates.origin_width, + ], + } + if inference_id is not None: + metadata[INFERENCE_ID_KEY] = inference_id + return metadata + + +def attach_native_detection_metadata( + detections: Detections, + image: WorkflowImageData, + class_names: Dict[int, str], + prediction_type: str, + inference_id: Optional[str] = None, +) -> Detections: + """LOCAL-path helper for model blocks that get a native ``Detections`` straight + from an inference_models adapter's ``run_tensor_native_inference``. + + The adapter fills ``xyxy`` / ``class_id`` / ``confidence`` (and sometimes a few + per-detection ``bboxes_metadata`` fields) but knows nothing about the workflow + image lineage. This attaches the workflow ``image_metadata`` and guarantees each + detection carries a ``detection_id`` (generated when missing), preserving any + keys the model already set (e.g. EasyOCR's per-box ``text``). Mutates and + returns the same object. + + Each entry additionally receives the host mirror of its box + (``HOST_XYXY_KEY`` / ``HOST_CLASS_ID_KEY`` / ``HOST_CONFIDENCE_KEY``) read + via ONE batched device->host transfer per tensor โ€” this runs on the + execution-engine thread right after the model synchronised its streams, so + the read costs microseconds; downstream host consumers (visualization + blocks) then never touch the device for these values. The mirror is always + OVERWRITTEN from the current tensors (never ``setdefault``), so re-attaching + after a geometry rewrite (e.g. perspective correction) refreshes it. + """ + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=class_names, + prediction_type=prediction_type, + inference_id=inference_id, + ) + number_of_detections = int(detections.xyxy.shape[0]) + if number_of_detections == 0: + detections.bboxes_metadata = None + return detections + # One batched read per tensor (`.numpy().tolist()` adds no further device + # traffic) โ€” never a per-box `.item()` loop, which would issue N blocking + # device round-trips each. + xyxy_host = ( + detections.xyxy.detach().reshape(number_of_detections, 4).cpu().numpy().tolist() + ) + class_id_host = detections.class_id.detach().reshape(-1).cpu().numpy().tolist() + confidence_host = detections.confidence.detach().reshape(-1).cpu().numpy().tolist() + existing = detections.bboxes_metadata + bboxes_metadata = [] + for index in range(number_of_detections): + entry = ( + dict(existing[index]) + if existing is not None and index < len(existing) + else {} + ) + entry.setdefault(DETECTION_ID_KEY, str(uuid4())) + entry[HOST_XYXY_KEY] = [float(value) for value in xyxy_host[index]] + entry[HOST_CLASS_ID_KEY] = int(class_id_host[index]) + entry[HOST_CONFIDENCE_KEY] = float(confidence_host[index]) + bboxes_metadata.append(entry) + detections.bboxes_metadata = bboxes_metadata + return detections + + +def native_detections_from_inference_predictions( + image: WorkflowImageData, + predictions: List[dict], + prediction_type: str, + class_names: Optional[Dict[int, str]] = None, + inference_id: Optional[str] = None, + device: Optional[torch.device] = None, +) -> Detections: + """REMOTE-path helper: build a native ``Detections`` from standard inference + object-detection prediction dicts (center ``x``/``y``/``width``/``height``, + ``confidence``, ``class_id``, optional ``class`` name and ``detection_id``). + + Boxes are converted from center form to corner ``xyxy``. When ``class_names`` + is not supplied it is derived from the predictions' own ``class_id``/``class`` + pairs so the serialiser can resolve every id. A ``detection_id`` is preserved + when present, otherwise generated. When ``device`` is not supplied the output + tensors are pinned to ``WORKFLOWS_IMAGE_TENSOR_DEVICE`` (so a REMOTE/HTTP + detection result lands on the same device as the LOCAL-path siblings rather + than silently staying on CPU). + """ + if device is None: + device = WORKFLOWS_IMAGE_TENSOR_DEVICE + xyxy: List[List[float]] = [] + class_id: List[int] = [] + confidence: List[float] = [] + bboxes_metadata: List[dict] = [] + derived_class_names: Dict[int, str] = {} + for prediction in predictions: + center_x = float(prediction[X_KEY]) + center_y = float(prediction[Y_KEY]) + box_width = float(prediction[WIDTH_KEY]) + box_height = float(prediction[HEIGHT_KEY]) + xyxy.append( + [ + center_x - box_width / 2, + center_y - box_height / 2, + center_x + box_width / 2, + center_y + box_height / 2, + ] + ) + prediction_class_id = int(prediction.get(CLASS_ID_KEY, 0)) + class_id.append(prediction_class_id) + confidence.append(float(prediction.get(CONFIDENCE_KEY, 1.0))) + if CLASS_NAME_KEY in prediction: + derived_class_names[prediction_class_id] = str(prediction[CLASS_NAME_KEY]) + bboxes_metadata.append( + {DETECTION_ID_KEY: str(prediction.get(DETECTION_ID_KEY) or uuid4())} + ) + number_of_detections = len(xyxy) + resolved_class_names = ( + class_names if class_names is not None else derived_class_names + ) + image_metadata = build_native_image_metadata( + image=image, + class_names=resolved_class_names, + prediction_type=prediction_type, + inference_id=inference_id, + ) + return Detections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32, device=device).reshape(-1, 4), + class_id=torch.as_tensor(class_id, dtype=torch.long, device=device).reshape(-1), + confidence=torch.as_tensor( + confidence, dtype=torch.float32, device=device + ).reshape(-1), + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata if number_of_detections > 0 else None, + ) + + +def _column_to_runs(column: np.ndarray) -> List[Tuple[int, int]]: + """Split a single 1-D column into (value, run_length) pairs of consecutive equal values.""" + if column.shape[0] == 0: + return [] + change_points = np.flatnonzero(column[1:] != column[:-1]) + 1 + boundaries = np.concatenate(([0], change_points, [column.shape[0]])) + return [ + (int(column[start]), int(end - start)) + for start, end in zip(boundaries[:-1], boundaries[1:]) + ] + + +def _embed_single_mask_counts( + column_major_slice: np.ndarray, + offset_xy: Tuple[int, int], + target_size_hw: Tuple[int, int], +) -> List[int]: + """Build the column-major uncompressed COCO counts for one slice placed onto the canvas. + + ``column_major_slice`` is the dense (h, w) slice already laid out as a list of its w + columns (each of height h). The big (H, W) canvas is never densified: we emit run + lengths directly and merge adjacent equal-value runs. + """ + h, w = column_major_slice.shape + x0, y0 = offset_xy + target_h, target_w = target_size_hw + + bottom_zeros = target_h - h - y0 + leading_zero_pixels = x0 * target_h + trailing_zero_pixels = (target_w - w - x0) * target_h + + # (value, run_length) pairs in column-major order across the whole canvas. + runs: List[Tuple[int, int]] = [] + if leading_zero_pixels > 0: + runs.append((0, leading_zero_pixels)) + for column_index in range(w): + column = column_major_slice[:, column_index] + if y0 > 0: + runs.append((0, y0)) + runs.extend(_column_to_runs(column)) + if bottom_zeros > 0: + runs.append((0, bottom_zeros)) + if trailing_zero_pixels > 0: + runs.append((0, trailing_zero_pixels)) + + # COCO uncompressed counts are alternating run lengths starting with a zero run. + merged_counts: List[int] = [0] + current_value = 0 + for value, length in runs: + if length == 0: + continue + if value == current_value: + merged_counts[-1] += length + else: + merged_counts.append(length) + current_value = value + return merged_counts + + +def embed_rle_masks_in_larger_canvas( + masks: InstancesRLEMasks, + offset_xy: Tuple[int, int], + target_size_hw: Tuple[int, int], +) -> InstancesRLEMasks: + """Place N slice-resolution RLE masks onto a larger all-zeros canvas, in RLE. + + Each input mask sits at slice resolution ``masks.image_size == (h, w)``. The returned + masks live on a ``(H, W)`` canvas with the slice's top-left at ``offset_xy == (x0, y0)``; + everything outside the slice is zero. COCO RLE here is column-major (fortran), matching + ``torch_mask_to_coco_rle``. The big canvas is never densified -- only the small slice is + decoded to dense and run lengths are emitted directly onto the canvas. + """ + h, w = masks.image_size + x0, y0 = offset_xy + target_h, target_w = target_size_hw + + if x0 < 0 or y0 < 0: + raise ValueError( + f"embed_rle_masks_in_larger_canvas got negative offset_xy={offset_xy}; " + f"offsets must be non-negative." + ) + if x0 + w > target_w or y0 + h > target_h: + raise ValueError( + f"Slice of size (h={h}, w={w}) at offset_xy=(x0={x0}, y0={y0}) does not fit " + f"into target canvas of size (H={target_h}, W={target_w}): requires " + f"x0+w={x0 + w} <= W={target_w} and y0+h={y0 + h} <= H={target_h}." + ) + + if len(masks.masks) == 0: + return InstancesRLEMasks(image_size=(target_h, target_w), masks=[]) + + # Decode only the small (h, w) slices to dense; the big canvas stays in run-list form. + dense_slices = coco_rle_masks_to_numpy_mask(masks).astype(np.uint8) + + embedded: List[bytes] = [] + for dense_slice in dense_slices: + # Lay the (h, w) slice out as columns (column j of height h) -> shape (h, w). + counts = _embed_single_mask_counts( + column_major_slice=dense_slice, + offset_xy=offset_xy, + target_size_hw=target_size_hw, + ) + rle = mask_utils.frPyObjects( + {"counts": counts, "size": [target_h, target_w]}, target_h, target_w + ) + embedded.append(rle["counts"]) + + return InstancesRLEMasks(image_size=(target_h, target_w), masks=embedded) + + +def build_native_key_points( + per_instance_xy: List[Optional[List[List[float]]]], + per_instance_confidence: List[Optional[List[float]]], + object_class_ids: List[Any], + image_metadata: dict, +) -> KeyPoints: + """Rebuild a padded native ``KeyPoints`` from per-instance keypoint lists (the + flattened ``keypoints_xy`` / ``keypoints_confidence`` shape the keypoint + producer writes into ``bboxes_metadata``). Mirrors + ``keypoint_detection/v3_tensor._native_key_points_from_inference_predictions``: + ragged per-instance keypoint counts are zero-padded to a uniform ``K`` with + confidence ``0.0`` in the padding rows. ``class_id`` is the per-instance + *object* class id (one per skeleton), matching the bbox ``Detections.class_id``. + Used by the rollup block and the dynamic-block representation boundary to + rebuild the ``(KeyPoints, Detections)`` tuple the visualizer siblings require. + """ + number_of_instances = len(object_class_ids) + normalised_xy = [list(xy) if xy else [] for xy in per_instance_xy] + normalised_confidence = [ + list(conf) if conf else [] for conf in per_instance_confidence + ] + max_key_points = max((len(xy) for xy in normalised_xy), default=0) + xy_tensor = torch.zeros( + (number_of_instances, max_key_points, 2), dtype=torch.float32 + ) + confidence_tensor = torch.zeros( + (number_of_instances, max_key_points), dtype=torch.float32 + ) + for index in range(number_of_instances): + keypoint_count = len(normalised_xy[index]) + if keypoint_count > 0: + xy_tensor[index, :keypoint_count] = torch.as_tensor( + normalised_xy[index], dtype=torch.float32 + ) + confidence_count = len(normalised_confidence[index]) + if confidence_count > 0: + confidence_tensor[index, :confidence_count] = torch.as_tensor( + normalised_confidence[index], dtype=torch.float32 + ) + class_id_tensor = torch.as_tensor( + [int(value) for value in object_class_ids], dtype=torch.long + ).reshape(-1) + return KeyPoints( + xy=xy_tensor, + class_id=class_id_tensor, + confidence=confidence_tensor, + image_metadata=image_metadata, + ) diff --git a/inference/core/workflows/core_steps/common/utils.py b/inference/core/workflows/core_steps/common/utils.py index da35eca799..9297202705 100644 --- a/inference/core/workflows/core_steps/common/utils.py +++ b/inference/core/workflows/core_steps/common/utils.py @@ -6,7 +6,7 @@ import numpy as np import supervision as sv -from supervision.config import CLASS_NAME_DATA_FIELD +from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES from inference.core.entities.requests.clip import ClipCompareRequest from inference.core.entities.requests.doctr import DoctrOCRInferenceRequest @@ -311,6 +311,13 @@ def sv_detections_to_root_coordinates( detections_copy.data[POLYGON_KEY_IN_SV_DETECTIONS] = ( detections_copy.data[POLYGON_KEY_IN_SV_DETECTIONS] + polygon_shift ) + if ORIENTED_BOX_COORDINATES in detections_copy.data: + # crop localization subtracts the crop origin from the OBB corners + # (dynamic_crop), so root conversion must add it back - same as xyxy, + # keypoints and polygons above + detections_copy.data[ORIENTED_BOX_COORDINATES] = detections_copy.data[ + ORIENTED_BOX_COORDINATES + ] + np.asarray([shift_x, shift_y]) if detections_copy.mask is not None: origin_mask_base = np.full((origin_height, origin_width), False) new_anchored_masks = np.array( diff --git a/inference/core/workflows/core_steps/common/utils_tensor.py b/inference/core/workflows/core_steps/common/utils_tensor.py new file mode 100644 index 0000000000..c0c98ef6dd --- /dev/null +++ b/inference/core/workflows/core_steps/common/utils_tensor.py @@ -0,0 +1,70 @@ +""" +Tensor-native mirrors of the image-shape-reading helpers in +`common/utils.py`. Created per the plan's Step 5b. + +The functions here have the same logical contract as their numpy +counterparts. They are currently thin wrappers that delegate to the +numpy implementations because the underlying `attach_parent_coordinates_to_detections` +reads `ImageParentMetadata.origin_coordinates` (populated by +`WorkflowImageData.parent_metadata` / +`workflow_root_ancestor_metadata`, which already use +`_read_shape_without_materialization` to avoid forcing +device->host on tensor-only inputs). Keeping the mirror in its own +module gives tensor-specific optimisations a landing spot that the +loader can swap to without touching the numpy file. +""" + +from typing import Iterable, List + +import supervision as sv + +from inference.core.workflows.core_steps.common.utils import ( + attach_parent_coordinates_to_detections, + attach_parents_coordinates_to_batch_of_sv_detections, + attach_parents_coordinates_to_sv_detections, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) + + +def attach_parents_coordinates_to_batch_of_sv_detections_tensor( + predictions: List[sv.Detections], + images: Iterable[WorkflowImageData], +) -> List[sv.Detections]: + return attach_parents_coordinates_to_batch_of_sv_detections( + predictions=predictions, images=images + ) + + +def attach_parents_coordinates_to_sv_detections_tensor( + detections: sv.Detections, + image: WorkflowImageData, +) -> sv.Detections: + return attach_parents_coordinates_to_sv_detections( + detections=detections, image=image + ) + + +def attach_parent_coordinates_to_detections_tensor( + detections: sv.Detections, + parent_metadata: ImageParentMetadata, + parent_id_key: str, + coordinates_key: str, + dimensions_key: str, +) -> sv.Detections: + return attach_parent_coordinates_to_detections( + detections=detections, + parent_metadata=parent_metadata, + parent_id_key=parent_id_key, + coordinates_key=coordinates_key, + dimensions_key=dimensions_key, + ) + + +__all__ = [ + "attach_parents_coordinates_to_batch_of_sv_detections_tensor", + "attach_parents_coordinates_to_sv_detections_tensor", + "attach_parent_coordinates_to_detections_tensor", +] diff --git a/inference/core/workflows/core_steps/formatters/vlm_as_classifier/v1_tensor.py b/inference/core/workflows/core_steps/formatters/vlm_as_classifier/v1_tensor.py new file mode 100644 index 0000000000..e93491ecbc --- /dev/null +++ b/inference/core/workflows/core_steps/formatters/vlm_as_classifier/v1_tensor.py @@ -0,0 +1,397 @@ +import json +import logging +import re +from typing import Dict, List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_FORMATTER, + CLASSIFICATION_STYLE_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + LANGUAGE_MODEL_OUTPUT_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) + +JSON_MARKDOWN_BLOCK_PATTERN = re.compile(r"```json([\s\S]*?)```", flags=re.IGNORECASE) + +LONG_DESCRIPTION = """ +Parse JSON strings from Visual Language Models (VLMs) and Large Language Models (LLMs) into standardized classification prediction format by extracting class predictions, mapping class names to class IDs, handling both single-class and multi-label formats, and converting VLM/LLM text outputs into workflow-compatible classification results for VLM-based classification, LLM classification parsing, and text-to-classification conversion workflows. + +## How This Block Works + +This block converts VLM/LLM text outputs containing classification predictions into standardized classification prediction format. The block: + +1. Receives image and VLM output string containing classification results in JSON format +2. Parses JSON content from VLM output: + + **Handles Markdown-wrapped JSON:** + - Searches for JSON wrapped in Markdown code blocks (```json ... ```) + - This format is common in LLM/VLM responses + - If multiple markdown JSON blocks are found, only the first block is parsed + - Extracts JSON content from within markdown tags + + **Handles raw JSON strings:** + - If no markdown blocks are found, attempts to parse the entire string as JSON + - Supports standard JSON format strings +3. Detects classification format and parses accordingly: + + **Single-Class Classification Format:** + - Detects format containing "class_name" and "confidence" fields + - Extracts the predicted class name and confidence score + - Creates classification prediction with single top class + - Maps class name to class ID using provided classes list + + **Multi-Label Classification Format:** + - Detects format containing "predicted_classes" array + - Extracts all predicted classes with their confidence scores + - Handles duplicate classes by taking maximum confidence + - Maps all class names to class IDs using provided classes list +4. Creates class name to class ID mapping: + - Uses the provided classes list to create index mapping (class_name โ†’ class_id) + - Maps classes in order (first class = ID 0, second = ID 1, etc.) + - Classes not in the provided list get class_id = -1 +5. Normalizes confidence scores: + - Scales confidence values to valid range [0.0, 1.0] + - Clamps values outside the range to 0.0 or 1.0 +6. Constructs classification prediction: + - Includes image dimensions (width, height) from input image + - For single-class: includes "top" class, confidence, and predictions array + - For multi-label: includes "predicted_classes" list and predictions dictionary + - Includes inference_id and parent_id for tracking + - Formats prediction in standard classification prediction format +7. Handles errors: + - Sets `error_status` to True if JSON parsing fails + - Sets `error_status` to True if classification format cannot be determined + - Returns None for predictions when errors occur + - Always includes inference_id for tracking +8. Returns classification prediction: + - Outputs `predictions` in standard classification format (compatible with classification blocks) + - Outputs `error_status` indicating parsing success/failure + - Outputs `inference_id` for tracking and lineage + +The block enables using VLMs/LLMs for classification by converting their text-based JSON outputs into standardized classification predictions that can be used in workflows like any other classification model output. + +## Common Use Cases + +- **VLM-Based Classification**: Use Visual Language Models for image classification by parsing VLM outputs into classification predictions (e.g., classify images with VLMs, use GPT-4V for classification, parse Claude Vision classifications), enabling VLM classification workflows +- **LLM Classification Parsing**: Parse LLM text outputs containing classification results into standardized format (e.g., parse GPT classification outputs, convert LLM predictions to classification format, use LLMs for classification), enabling LLM classification workflows +- **Text-to-Classification Conversion**: Convert text-based classification outputs from models into workflow-compatible classification predictions (e.g., convert text predictions to classification format, parse text-based classifications, convert model outputs to classifications), enabling text-to-classification workflows +- **Multi-Format Classification Support**: Handle both single-class and multi-label classification formats from VLM/LLM outputs (e.g., support single-label VLM classifications, support multi-label VLM classifications, handle different classification formats), enabling flexible classification workflows +- **VLM Integration**: Integrate VLM outputs into classification workflows (e.g., use VLMs in classification pipelines, integrate VLM predictions with classification blocks, combine VLM and traditional classification), enabling VLM integration workflows +- **Flexible Classification Sources**: Enable classification from various model types that output text/JSON (e.g., use any text-output model for classification, convert model outputs to classifications, parse various classification formats), enabling flexible classification workflows + +## Connecting to Other Blocks + +This block receives images and VLM outputs and produces classification predictions: + +- **After VLM/LLM blocks** to parse classification outputs into standard format (e.g., VLM output to classification, LLM output to classification, parse model outputs), enabling VLM-to-classification workflows +- **Before classification-based blocks** to use parsed classifications (e.g., use parsed classifications in workflows, provide classifications to downstream blocks, use VLM classifications with classification blocks), enabling classification-to-workflow workflows +- **Before filtering blocks** to filter based on VLM classifications (e.g., filter by VLM classification results, use parsed classifications for filtering, apply filters to VLM predictions), enabling classification-to-filter workflows +- **Before analytics blocks** to analyze VLM classification results (e.g., analyze VLM classifications, perform analytics on parsed classifications, track VLM classification metrics), enabling classification analytics workflows +- **Before visualization blocks** to display VLM classification results (e.g., visualize VLM classifications, display parsed classification predictions, show VLM classification outputs), enabling classification visualization workflows +- **In workflow outputs** to provide VLM classifications as final output (e.g., VLM classification outputs, parsed classification results, VLM-based classification outputs), enabling classification output workflows + +## Requirements + +This block requires an image input (for metadata and dimensions) and a VLM output string containing JSON classification data. The JSON can be raw JSON or wrapped in Markdown code blocks (```json ... ```). The block supports two JSON formats: single-class (with "class_name" and "confidence" fields) and multi-label (with "predicted_classes" array). The `classes` parameter must contain a list of all class names used by the model to generate class_id mappings. Classes are mapped to IDs by index (first class = 0, second = 1, etc.). Classes not in the list get class_id = -1. Confidence scores are normalized to [0.0, 1.0] range. The block outputs classification predictions in standard format (compatible with classification blocks), error_status (boolean), and inference_id (string) for tracking. +""" + +SHORT_DESCRIPTION = "Parse a raw string into a classification prediction." + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "VLM As Classifier", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "formatter", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-tags", + "blockPriority": 5, + }, + } + ) + type: Literal["roboflow_core/vlm_as_classifier@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Input image that was used to generate the VLM prediction. Used to extract image dimensions (width, height) and metadata (parent_id) for the classification prediction. The same image that was provided to the VLM/LLM block should be used here to maintain consistency.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + vlm_output: Selector(kind=[LANGUAGE_MODEL_OUTPUT_KIND]) = Field( + title="VLM Output", + description="String output from a VLM or LLM block containing classification prediction in JSON format. Can be raw JSON string (e.g., '{\"class_name\": \"dog\", \"confidence\": 0.95}') or JSON wrapped in Markdown code blocks (e.g., ```json {...} ```). Supports two formats: single-class (with 'class_name' and 'confidence' fields) or multi-label (with 'predicted_classes' array). If multiple markdown blocks exist, only the first is parsed.", + examples=["$steps.lmm.output", "$steps.vlm.output", "$steps.claude.output"], + ) + classes: Union[ + Selector(kind=[LIST_OF_VALUES_KIND]), + Selector(kind=[LIST_OF_VALUES_KIND]), + List[str], + ] = Field( + description="List of all class names used by the classification model, in order. Required to generate mapping between class names (from VLM output) and class IDs (for classification format). Classes are mapped to IDs by index: first class = ID 0, second = ID 1, etc. Classes from VLM output that are not in this list get class_id = -1. Should match the classes the VLM was asked to classify.", + examples=[ + [ + "$steps.lmm.classes", + "$inputs.classes", + ["dog", "cat", "bird"], + ["class_a", "class_b"], + ] + ], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND], + ), + OutputDefinition(name="inference_id", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class VLMAsClassifierBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + vlm_output: str, + classes: List[str], + ) -> BlockResult: + inference_id = f"{uuid4()}" + error_status, parsed_data = string2json( + raw_json=vlm_output, + ) + if error_status: + return { + "error_status": True, + "predictions": None, + "inference_id": inference_id, + } + if "class_name" in parsed_data and "confidence" in parsed_data: + return parse_multi_class_classification_results( + image=image, + results=parsed_data, + classes=classes, + inference_id=inference_id, + ) + if "predicted_classes" in parsed_data: + return parse_multi_label_classification_results( + image=image, + results=parsed_data, + classes=classes, + inference_id=inference_id, + ) + return { + "error_status": True, + "predictions": None, + "inference_id": inference_id, + } + + +def string2json( + raw_json: str, +) -> Tuple[bool, dict]: + json_blocks_found = JSON_MARKDOWN_BLOCK_PATTERN.findall(raw_json) + if len(json_blocks_found) == 0: + return try_parse_json(raw_json) + first_block = json_blocks_found[0] + return try_parse_json(first_block) + + +def try_parse_json(content: str) -> Tuple[bool, dict]: + try: + return False, json.loads(content) + except Exception as error: + logging.warning( + f"Could not parse JSON to dict in `roboflow_core/vlm_as_classifier@v1` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return True, {} + + +def parse_multi_class_classification_results( + image: WorkflowImageData, + results: dict, + classes: List[str], + inference_id: str, +) -> dict: + try: + class2id_mapping = create_classes_index(classes=classes) + height, width = image._read_shape_without_materialization() + top_class = results["class_name"] + confidences = {top_class: scale_confidence(results["confidence"])} + # The native ClassificationPrediction carries a dense confidence vector + # indexed by class_id plus a class_id -> class_name map in image metadata. + if top_class not in class2id_mapping: + # Out-of-list top class is appended at the end: class_id values must + # be non-negative dense indices into the confidence vector and + # class_names map. + class2id_mapping[top_class] = len(class2id_mapping) + class_names = { + class_id: class_name for class_name, class_id in class2id_mapping.items() + } + num_classes = len(class_names) + confidence_vector = [0.0] * num_classes + for class_name, class_id in class2id_mapping.items(): + confidence_vector[class_id] = confidences.get(class_name, 0.0) + top_class_id = class2id_mapping[top_class] + image_metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag -> serialiser reproduces the vlm formatter shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_FORMATTER, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + } + parsed_prediction = ClassificationPrediction( + class_id=torch.tensor( + [top_class_id], + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.tensor( + [confidence_vector], + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + images_metadata=[image_metadata], + ) + return { + "error_status": False, + "predictions": parsed_prediction, + "inference_id": inference_id, + } + except Exception as error: + logging.warning( + f"Could not parse multi-class classification results in `roboflow_core/vlm_as_classifier@v1` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return {"error_status": True, "predictions": None, "inference_id": inference_id} + + +def parse_multi_label_classification_results( + image: WorkflowImageData, + results: dict, + classes: List[str], + inference_id: str, +) -> dict: + try: + class2id_mapping = create_classes_index(classes=classes) + height, width = image._read_shape_without_materialization() + predicted_classes_confidences = {} + for prediction in results["predicted_classes"]: + if prediction["class"] not in class2id_mapping: + # Out-of-list class is appended at the end: class_id values must + # be non-negative dense indices into the confidence vector and + # class_names map. + class2id_mapping[prediction["class"]] = len(class2id_mapping) + if prediction["class"] in predicted_classes_confidences: + old_confidence = predicted_classes_confidences[prediction["class"]] + new_confidence = scale_confidence(value=prediction["confidence"]) + predicted_classes_confidences[prediction["class"]] = max( + old_confidence, new_confidence + ) + else: + predicted_classes_confidences[prediction["class"]] = scale_confidence( + value=prediction["confidence"] + ) + # The native MultiLabelClassificationPrediction carries a dense sigmoid + # confidence vector indexed by class_id, the ids of the above-threshold + # (predicted) classes, and a class_id -> class_name map in image metadata. + class_names = { + class_id: class_name for class_name, class_id in class2id_mapping.items() + } + num_classes = len(class_names) + confidence_vector = [0.0] * num_classes + for class_name, class_id in class2id_mapping.items(): + confidence_vector[class_id] = predicted_classes_confidences.get( + class_name, 0.0 + ) + predicted_class_ids = [ + class2id_mapping[class_name] + for class_name in predicted_classes_confidences.keys() + ] + image_metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag -> serialiser reproduces the vlm formatter shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_FORMATTER, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + } + parsed_prediction = MultiLabelClassificationPrediction( + class_ids=torch.tensor( + predicted_class_ids, + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.tensor( + confidence_vector, + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + image_metadata=image_metadata, + ) + return { + "error_status": False, + "predictions": parsed_prediction, + "inference_id": inference_id, + } + except Exception as error: + logging.warning( + f"Could not parse multi-label classification results in `roboflow_core/vlm_as_classifier@v1` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return {"error_status": True, "predictions": None, "inference_id": inference_id} + + +def create_classes_index(classes: List[str]) -> Dict[str, int]: + return {class_name: idx for idx, class_name in enumerate(classes)} + + +def scale_confidence(value: float) -> float: + return min(max(float(value), 0.0), 1.0) diff --git a/inference/core/workflows/core_steps/formatters/vlm_as_classifier/v2_tensor.py b/inference/core/workflows/core_steps/formatters/vlm_as_classifier/v2_tensor.py new file mode 100644 index 0000000000..f1831c33df --- /dev/null +++ b/inference/core/workflows/core_steps/formatters/vlm_as_classifier/v2_tensor.py @@ -0,0 +1,403 @@ +import json +import logging +import re +from typing import Dict, List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_FORMATTER, + CLASSIFICATION_STYLE_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + LANGUAGE_MODEL_OUTPUT_KIND, + LIST_OF_VALUES_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) + +JSON_MARKDOWN_BLOCK_PATTERN = re.compile(r"```json([\s\S]*?)```", flags=re.IGNORECASE) + +LONG_DESCRIPTION = """ +Parse JSON strings from Visual Language Models (VLMs) and Large Language Models (LLMs) into standardized classification prediction format by extracting class predictions, mapping class names to class IDs, handling both single-class and multi-label formats, and converting VLM/LLM text outputs into workflow-compatible classification results for VLM-based classification, LLM classification parsing, and text-to-classification conversion workflows. + +## How This Block Works + +This block converts VLM/LLM text outputs containing classification predictions into standardized classification prediction format. The block: + +1. Receives image and VLM output string containing classification results in JSON format +2. Parses JSON content from VLM output: + + **Handles Markdown-wrapped JSON:** + - Searches for JSON wrapped in Markdown code blocks (```json ... ```) + - This format is common in LLM/VLM responses + - If multiple markdown JSON blocks are found, only the first block is parsed + - Extracts JSON content from within markdown tags + + **Handles raw JSON strings:** + - If no markdown blocks are found, attempts to parse the entire string as JSON + - Supports standard JSON format strings +3. Detects classification format and parses accordingly: + + **Single-Class Classification Format:** + - Detects format containing "class_name" and "confidence" fields + - Extracts the predicted class name and confidence score + - Creates classification prediction with single top class + - Maps class name to class ID using provided classes list + + **Multi-Label Classification Format:** + - Detects format containing "predicted_classes" array + - Extracts all predicted classes with their confidence scores + - Handles duplicate classes by taking maximum confidence + - Maps all class names to class IDs using provided classes list +4. Creates class name to class ID mapping: + - Uses the provided classes list to create index mapping (class_name โ†’ class_id) + - Maps classes in order (first class = ID 0, second = ID 1, etc.) + - Classes not in the provided list get class_id = -1 +5. Normalizes confidence scores: + - Scales confidence values to valid range [0.0, 1.0] + - Clamps values outside the range to 0.0 or 1.0 +6. Constructs classification prediction: + - Includes image dimensions (width, height) from input image + - For single-class: includes "top" class, confidence, and predictions array + - For multi-label: includes "predicted_classes" list and predictions dictionary + - Includes inference_id and parent_id for tracking + - Formats prediction in standard classification prediction format +7. Handles errors: + - Sets `error_status` to True if JSON parsing fails + - Sets `error_status` to True if classification format cannot be determined + - Returns None for predictions when errors occur + - Always includes inference_id for tracking +8. Returns classification prediction: + - Outputs `predictions` in standard classification format (compatible with classification blocks) + - Outputs `error_status` indicating parsing success/failure + - Outputs `inference_id` with specific type for tracking and lineage + +The block enables using VLMs/LLMs for classification by converting their text-based JSON outputs into standardized classification predictions that can be used in workflows like any other classification model output. + +## Common Use Cases + +- **VLM-Based Classification**: Use Visual Language Models for image classification by parsing VLM outputs into classification predictions (e.g., classify images with VLMs, use GPT-4V for classification, parse Claude Vision classifications), enabling VLM classification workflows +- **LLM Classification Parsing**: Parse LLM text outputs containing classification results into standardized format (e.g., parse GPT classification outputs, convert LLM predictions to classification format, use LLMs for classification), enabling LLM classification workflows +- **Text-to-Classification Conversion**: Convert text-based classification outputs from models into workflow-compatible classification predictions (e.g., convert text predictions to classification format, parse text-based classifications, convert model outputs to classifications), enabling text-to-classification workflows +- **Multi-Format Classification Support**: Handle both single-class and multi-label classification formats from VLM/LLM outputs (e.g., support single-label VLM classifications, support multi-label VLM classifications, handle different classification formats), enabling flexible classification workflows +- **VLM Integration**: Integrate VLM outputs into classification workflows (e.g., use VLMs in classification pipelines, integrate VLM predictions with classification blocks, combine VLM and traditional classification), enabling VLM integration workflows +- **Flexible Classification Sources**: Enable classification from various model types that output text/JSON (e.g., use any text-output model for classification, convert model outputs to classifications, parse various classification formats), enabling flexible classification workflows + +## Connecting to Other Blocks + +This block receives images and VLM outputs and produces classification predictions: + +- **After VLM/LLM blocks** to parse classification outputs into standard format (e.g., VLM output to classification, LLM output to classification, parse model outputs), enabling VLM-to-classification workflows +- **Before classification-based blocks** to use parsed classifications (e.g., use parsed classifications in workflows, provide classifications to downstream blocks, use VLM classifications with classification blocks), enabling classification-to-workflow workflows +- **Before filtering blocks** to filter based on VLM classifications (e.g., filter by VLM classification results, use parsed classifications for filtering, apply filters to VLM predictions), enabling classification-to-filter workflows +- **Before analytics blocks** to analyze VLM classification results (e.g., analyze VLM classifications, perform analytics on parsed classifications, track VLM classification metrics), enabling classification analytics workflows +- **Before visualization blocks** to display VLM classification results (e.g., visualize VLM classifications, display parsed classification predictions, show VLM classification outputs), enabling classification visualization workflows +- **In workflow outputs** to provide VLM classifications as final output (e.g., VLM classification outputs, parsed classification results, VLM-based classification outputs), enabling classification output workflows + +## Version Differences + +This version (v2) includes the following enhancements over v1: + +- **Improved Type System**: The `inference_id` output now uses `INFERENCE_ID_KIND` instead of generic `STRING_KIND`, providing better type safety and semantic clarity for inference ID values in the workflow type system + +## Requirements + +This block requires an image input (for metadata and dimensions) and a VLM output string containing JSON classification data. The JSON can be raw JSON or wrapped in Markdown code blocks (```json ... ```). The block supports two JSON formats: single-class (with "class_name" and "confidence" fields) and multi-label (with "predicted_classes" array). The `classes` parameter must contain a list of all class names used by the model to generate class_id mappings. Classes are mapped to IDs by index (first class = 0, second = 1, etc.). Classes not in the list get class_id = -1. Confidence scores are normalized to [0.0, 1.0] range. The block outputs classification predictions in standard format (compatible with classification blocks), error_status (boolean), and inference_id (INFERENCE_ID_KIND) for tracking. +""" + +SHORT_DESCRIPTION = "Parse a raw string into a classification prediction." + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "VLM As Classifier", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "formatter", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-tags", + "blockPriority": 5, + }, + } + ) + type: Literal["roboflow_core/vlm_as_classifier@v2"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Input image that was used to generate the VLM prediction. Used to extract image dimensions (width, height) and metadata (parent_id) for the classification prediction. The same image that was provided to the VLM/LLM block should be used here to maintain consistency.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + vlm_output: Selector(kind=[LANGUAGE_MODEL_OUTPUT_KIND]) = Field( + title="VLM Output", + description="String output from a VLM or LLM block containing classification prediction in JSON format. Can be raw JSON string (e.g., '{\"class_name\": \"dog\", \"confidence\": 0.95}') or JSON wrapped in Markdown code blocks (e.g., ```json {...} ```). Supports two formats: single-class (with 'class_name' and 'confidence' fields) or multi-label (with 'predicted_classes' array). If multiple markdown blocks exist, only the first is parsed.", + examples=["$steps.lmm.output", "$steps.vlm.output", "$steps.claude.output"], + ) + classes: Union[ + Selector(kind=[LIST_OF_VALUES_KIND]), + Selector(kind=[LIST_OF_VALUES_KIND]), + List[str], + ] = Field( + description="List of all class names used by the classification model, in order. Required to generate mapping between class names (from VLM output) and class IDs (for classification format). Classes are mapped to IDs by index: first class = ID 0, second = ID 1, etc. Classes from VLM output that are not in this list get class_id = -1. Should match the classes the VLM was asked to classify.", + examples=[ + [ + "$steps.lmm.classes", + "$inputs.classes", + ["dog", "cat", "bird"], + ["class_a", "class_b"], + ] + ], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND], + ), + OutputDefinition(name="inference_id", kind=[INFERENCE_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class VLMAsClassifierBlockV2(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + vlm_output: str, + classes: List[str], + ) -> BlockResult: + inference_id = f"{uuid4()}" + error_status, parsed_data = string2json( + raw_json=vlm_output, + ) + if error_status: + return { + "error_status": True, + "predictions": None, + "inference_id": inference_id, + } + if "class_name" in parsed_data and "confidence" in parsed_data: + return parse_multi_class_classification_results( + image=image, + results=parsed_data, + classes=classes, + inference_id=inference_id, + ) + if "predicted_classes" in parsed_data: + return parse_multi_label_classification_results( + image=image, + results=parsed_data, + classes=classes, + inference_id=inference_id, + ) + return { + "error_status": True, + "predictions": None, + "inference_id": inference_id, + } + + +def string2json( + raw_json: str, +) -> Tuple[bool, dict]: + json_blocks_found = JSON_MARKDOWN_BLOCK_PATTERN.findall(raw_json) + if len(json_blocks_found) == 0: + return try_parse_json(raw_json) + first_block = json_blocks_found[0] + return try_parse_json(first_block) + + +def try_parse_json(content: str) -> Tuple[bool, dict]: + try: + return False, json.loads(content) + except Exception as error: + logging.warning( + f"Could not parse JSON to dict in `roboflow_core/vlm_as_classifier@v2` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return True, {} + + +def parse_multi_class_classification_results( + image: WorkflowImageData, + results: dict, + classes: List[str], + inference_id: str, +) -> dict: + try: + class2id_mapping = create_classes_index(classes=classes) + height, width = image._read_shape_without_materialization() + top_class = results["class_name"] + confidences = {top_class: scale_confidence(results["confidence"])} + # The native ClassificationPrediction carries a dense confidence vector + # indexed by class_id plus a class_id -> class_name map in image metadata. + if top_class not in class2id_mapping: + # Out-of-list top class is appended at the end: class_id values must + # be non-negative dense indices into the confidence vector and + # class_names map. + class2id_mapping[top_class] = len(class2id_mapping) + class_names = { + class_id: class_name for class_name, class_id in class2id_mapping.items() + } + num_classes = len(class_names) + confidence_vector = [0.0] * num_classes + for class_name, class_id in class2id_mapping.items(): + confidence_vector[class_id] = confidences.get(class_name, 0.0) + top_class_id = class2id_mapping[top_class] + image_metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag -> serialiser reproduces the vlm formatter shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_FORMATTER, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + } + parsed_prediction = ClassificationPrediction( + class_id=torch.tensor( + [top_class_id], + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.tensor( + [confidence_vector], + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + images_metadata=[image_metadata], + ) + return { + "error_status": False, + "predictions": parsed_prediction, + "inference_id": inference_id, + } + except Exception as error: + logging.warning( + f"Could not parse multi-class classification results in `roboflow_core/vlm_as_classifier@v2` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return {"error_status": True, "predictions": None, "inference_id": inference_id} + + +def parse_multi_label_classification_results( + image: WorkflowImageData, + results: dict, + classes: List[str], + inference_id: str, +) -> dict: + try: + class2id_mapping = create_classes_index(classes=classes) + height, width = image._read_shape_without_materialization() + predicted_classes_confidences = {} + for prediction in results["predicted_classes"]: + if prediction["class"] not in class2id_mapping: + # Out-of-list class is appended at the end: class_id values must + # be non-negative dense indices into the confidence vector and + # class_names map. + class2id_mapping[prediction["class"]] = len(class2id_mapping) + if prediction["class"] in predicted_classes_confidences: + old_confidence = predicted_classes_confidences[prediction["class"]] + new_confidence = scale_confidence(value=prediction["confidence"]) + predicted_classes_confidences[prediction["class"]] = max( + old_confidence, new_confidence + ) + else: + predicted_classes_confidences[prediction["class"]] = scale_confidence( + value=prediction["confidence"] + ) + # The native MultiLabelClassificationPrediction carries a dense sigmoid + # confidence vector indexed by class_id, the ids of the above-threshold + # (predicted) classes, and a class_id -> class_name map in image metadata. + class_names = { + class_id: class_name for class_name, class_id in class2id_mapping.items() + } + num_classes = len(class_names) + confidence_vector = [0.0] * num_classes + for class_name, class_id in class2id_mapping.items(): + confidence_vector[class_id] = predicted_classes_confidences.get( + class_name, 0.0 + ) + predicted_class_ids = [ + class2id_mapping[class_name] + for class_name in predicted_classes_confidences.keys() + ] + image_metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag -> serialiser reproduces the vlm formatter shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_FORMATTER, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + } + parsed_prediction = MultiLabelClassificationPrediction( + class_ids=torch.tensor( + predicted_class_ids, + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.tensor( + confidence_vector, + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + image_metadata=image_metadata, + ) + return { + "error_status": False, + "predictions": parsed_prediction, + "inference_id": inference_id, + } + except Exception as error: + logging.warning( + f"Could not parse multi-label classification results in `roboflow_core/vlm_as_classifier@v2` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return {"error_status": True, "predictions": None, "inference_id": inference_id} + + +def create_classes_index(classes: List[str]) -> Dict[str, int]: + return {class_name: idx for idx, class_name in enumerate(classes)} + + +def scale_confidence(value: float) -> float: + return min(max(float(value), 0.0), 1.0) diff --git a/inference/core/workflows/core_steps/formatters/vlm_as_detector/v1_tensor.py b/inference/core/workflows/core_steps/formatters/vlm_as_detector/v1_tensor.py new file mode 100644 index 0000000000..13de49dce2 --- /dev/null +++ b/inference/core/workflows/core_steps/formatters/vlm_as_detector/v1_tensor.py @@ -0,0 +1,593 @@ +import hashlib +import json +import logging +import re +from functools import partial +from typing import Dict, List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field, model_validator + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.vlms import VLM_TASKS_METADATA +from inference.core.workflows.core_steps.formatters.vlm_as_detector.gemini_detection_parsing import ( + convert_gemini_detection_to_pixel_xyxy, + create_classes_index, + extract_gemini_detection_entries, + get_gemini_detection_class_name, + scale_confidence, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + LANGUAGE_MODEL_OUTPUT_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections + +JSON_MARKDOWN_BLOCK_PATTERN = re.compile(r"```json([\s\S]*?)```", flags=re.IGNORECASE) + +LONG_DESCRIPTION = """ +Parse JSON strings from Visual Language Models (VLMs) and Large Language Models (LLMs) into standardized object detection prediction format by extracting bounding boxes, class names, and confidences, converting normalized coordinates to pixel coordinates, mapping class names to class IDs, and handling multiple model types and task formats to enable VLM-based object detection, LLM detection parsing, and text-to-detection conversion workflows. + +## How This Block Works + +This block converts VLM/LLM text outputs containing object detection predictions into standardized object detection format compatible with workflow detection blocks. The block: + +1. Receives image and VLM output string containing detection results in JSON format +2. Parses JSON content from VLM output: + + **Handles Markdown-wrapped JSON:** + - Searches for JSON wrapped in Markdown code blocks (```json ... ```) + - This format is common in LLM/VLM responses + - If multiple markdown JSON blocks are found, only the first block is parsed + - Extracts JSON content from within markdown tags + + **Handles raw JSON strings:** + - If no markdown blocks are found, attempts to parse the entire string as JSON + - Supports standard JSON format strings +3. Selects appropriate parser based on model type and task type: + - Uses registered parsers that handle different model outputs (google-gemini, anthropic-claude, florence-2) + - Supports multiple task types: object-detection, open-vocabulary-object-detection, object-detection-and-caption, phrase-grounded-object-detection, region-proposal, ocr-with-text-detection + - Each model/task combination uses a specialized parser for that format +4. Parses detection data based on model type: + + **For Gemini/Claude models:** + - Extracts detections array from parsed JSON + - Converts normalized coordinates (0-1 range) to pixel coordinates using image dimensions + - Extracts class names, confidence scores, and bounding box coordinates + - Maps class names to class IDs using provided classes list + - Creates detection objects with bounding boxes, classes, and confidences + + **For Florence-2 model:** + - Uses supervision's built-in LMM parser for Florence-2 format + - Handles different task types with specialized parsing (object detection, open vocabulary, region proposal, OCR, etc.) + - For region proposal tasks: assigns "roi" as class name + - For open vocabulary detection: uses provided classes list for class ID mapping + - For other tasks: uses MD5-based class ID generation or provided classes + - Sets confidence to 1.0 for Florence-2 detections (model doesn't provide confidence) +5. Converts coordinates and normalizes data: + - Converts normalized coordinates (0-1) to absolute pixel coordinates (x_min, y_min, x_max, y_max) + - Scales coordinates using image width and height + - Normalizes confidence scores to valid range [0.0, 1.0] + - Clamps confidence values outside the range +6. Creates class name to class ID mapping: + - For Gemini/Claude: uses provided classes list to create index mapping (class_name โ†’ class_id) + - Classes are mapped in order (first class = ID 0, second = ID 1, etc.) + - Classes not in the provided list get class_id = -1 + - For Florence-2: uses different mapping strategies based on task type +7. Constructs object detection predictions: + - Creates supervision Detections objects with bounding boxes (xyxy format) + - Includes class IDs, class names, and confidence scores + - Adds metadata: detection IDs, inference IDs, image dimensions, prediction type + - Attaches parent coordinates for crop-aware detections + - Formats predictions in standard object detection format +8. Handles errors: + - Sets `error_status` to True if JSON parsing fails + - Sets `error_status` to True if detection parsing fails + - Returns None for predictions when errors occur + - Always includes inference_id for tracking +9. Returns object detection predictions: + - Outputs `predictions` in standard object detection format (compatible with detection blocks) + - Outputs `error_status` indicating parsing success/failure + - Outputs `inference_id` for tracking and lineage + +The block enables using VLMs/LLMs for object detection by converting their text-based JSON outputs into standardized detection predictions that can be used in workflows like any other object detection model output. + +## Common Use Cases + +- **VLM-Based Object Detection**: Use Visual Language Models for object detection by parsing VLM outputs into detection predictions (e.g., detect objects with GPT-4V, use Claude Vision for detection, parse Gemini detection outputs), enabling VLM detection workflows +- **Open-Vocabulary Detection**: Use VLMs for open-vocabulary object detection with custom classes (e.g., detect custom objects with VLMs, use open-vocabulary detection, detect objects not in training set), enabling open-vocabulary detection workflows +- **Multi-Task Detection**: Use VLMs for various detection tasks (e.g., object detection with captions, phrase-grounded detection, region proposal, OCR with detection), enabling multi-task detection workflows +- **LLM Detection Parsing**: Parse LLM text outputs containing detection results into standardized format (e.g., parse GPT detection outputs, convert LLM predictions to detection format, use LLMs for detection), enabling LLM detection workflows +- **Text-to-Detection Conversion**: Convert text-based detection outputs from models into workflow-compatible detection predictions (e.g., convert text predictions to detection format, parse text-based detections, convert model outputs to detections), enabling text-to-detection workflows +- **VLM Integration**: Integrate VLM outputs into detection workflows (e.g., use VLMs in detection pipelines, integrate VLM predictions with detection blocks, combine VLM and traditional detection), enabling VLM integration workflows + +## Connecting to Other Blocks + +This block receives images and VLM outputs and produces object detection predictions: + +- **After VLM/LLM blocks** to parse detection outputs into standard format (e.g., VLM output to detections, LLM output to detections, parse model outputs), enabling VLM-to-detection workflows +- **Before detection-based blocks** to use parsed detections (e.g., use parsed detections in workflows, provide detections to downstream blocks, use VLM detections with detection blocks), enabling detection-to-workflow workflows +- **Before filtering blocks** to filter VLM detections (e.g., filter by class, filter by confidence, apply filters to VLM predictions), enabling detection-to-filter workflows +- **Before analytics blocks** to analyze VLM detection results (e.g., analyze VLM detections, perform analytics on parsed detections, track VLM detection metrics), enabling detection analytics workflows +- **Before visualization blocks** to display VLM detection results (e.g., visualize VLM detections, display parsed detection predictions, show VLM detection outputs), enabling detection visualization workflows +- **In workflow outputs** to provide VLM detections as final output (e.g., VLM detection outputs, parsed detection results, VLM-based detection outputs), enabling detection output workflows + +## Requirements + +This block requires an image input (for metadata and dimensions) and a VLM output string containing JSON detection data. The JSON can be raw JSON or wrapped in Markdown code blocks (```json ... ```). The block supports three model types: "google-gemini", "anthropic-claude", and "florence-2". It supports multiple task types: "object-detection", "open-vocabulary-object-detection", "object-detection-and-caption", "phrase-grounded-object-detection", "region-proposal", and "ocr-with-text-detection". The `classes` parameter is required for Gemini and Claude models (to map class names to IDs) but optional for Florence-2 (some tasks don't require it). Classes are mapped to IDs by index (first class = 0, second = 1, etc.). Classes not in the list get class_id = -1. The block outputs object detection predictions in standard format (compatible with detection blocks), error_status (boolean), and inference_id (string) for tracking. +""" + +SHORT_DESCRIPTION = "Parses raw string into object-detection prediction." + +SUPPORTED_TASKS = { + "object-detection", + "object-detection-and-caption", + "open-vocabulary-object-detection", + "phrase-grounded-object-detection", + "region-proposal", + "ocr-with-text-detection", +} +RELEVANT_TASKS_METADATA = { + k: v for k, v in VLM_TASKS_METADATA.items() if k in SUPPORTED_TASKS +} + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "VLM As Detector", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "formatter", + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/vlm_as_detector@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Input image that was used to generate the VLM prediction. Used to extract image dimensions (width, height) for converting normalized coordinates to pixel coordinates and metadata (parent_id) for the detection predictions. The same image that was provided to the VLM/LLM block should be used here to maintain consistency.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + vlm_output: Selector(kind=[LANGUAGE_MODEL_OUTPUT_KIND]) = Field( + title="VLM Output", + description="String output from a VLM or LLM block containing object detection prediction in JSON format. Can be raw JSON string or JSON wrapped in Markdown code blocks (e.g., ```json {...} ```). Format depends on model_type and task_type - different models and tasks produce different JSON structures. If multiple markdown blocks exist, only the first is parsed.", + examples=[ + ["$steps.lmm.output"], + ["$steps.vlm.output"], + ["$steps.claude.output"], + ], + ) + classes: Optional[ + Union[ + Selector(kind=[LIST_OF_VALUES_KIND]), + Selector(kind=[LIST_OF_VALUES_KIND]), + List[str], + ] + ] = Field( + description="List of all class names used by the detection model, in order. Required for google-gemini and anthropic-claude models to generate mapping between class names (from VLM output) and class IDs (for detection format). Optional for florence-2 model (required only for open-vocabulary-object-detection task). Classes are mapped to IDs by index: first class = ID 0, second = ID 1, etc. Classes from VLM output that are not in this list get class_id = -1. Should match the classes the VLM was asked to detect.", + examples=[ + [ + "$steps.lmm.classes", + "$inputs.classes", + ["dog", "cat", "bird"], + ["class_a", "class_b"], + ] + ], + json_schema_extra={ + "relevant_for": { + "model_type": { + "values": ["google-gemini", "anthropic-claude"], + "required": True, + }, + } + }, + ) + model_type: Literal["google-gemini", "anthropic-claude", "florence-2"] = Field( + description="Type of VLM/LLM model that generated the detection prediction. Determines which parser to use for parsing the JSON output. 'google-gemini': Google Gemini model outputs. 'anthropic-claude': Anthropic Claude model outputs. 'florence-2': Microsoft Florence-2 model outputs. Each model type has different JSON output formats and requires appropriate parsing.", + examples=["google-gemini", "anthropic-claude", "florence-2"], + ) + task_type: Literal[tuple(SUPPORTED_TASKS)] = Field( + description="Task type that was performed by the VLM model. Determines how the JSON output is parsed and what detection format is expected. Supported tasks: 'object-detection' (unprompted detection), 'open-vocabulary-object-detection' (detection with provided classes), 'object-detection-and-caption' (detection with captions), 'phrase-grounded-object-detection' (prompted detection), 'region-proposal' (regions of interest), 'ocr-with-text-detection' (text detection with OCR). Each task type has specific output format requirements.", + json_schema_extra={ + "values_metadata": RELEVANT_TASKS_METADATA, + }, + ) + + @model_validator(mode="after") + def validate(self) -> "BlockManifest": + if (self.model_type, self.task_type) not in REGISTERED_PARSERS: + raise ValueError( + f"Could not parse result of task {self.task_type} for model {self.model_type}" + ) + if self.model_type != "florence-2" and self.classes is None: + raise ValueError( + "Must pass list of classes to this block when using gemini or claude" + ) + + return self + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition(name="inference_id", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class VLMAsDetectorBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + vlm_output: str, + classes: Optional[List[str]], + model_type: str, + task_type: str, + ) -> BlockResult: + inference_id = f"{uuid4()}" + error_status, parsed_data = string2json( + raw_json=vlm_output, + ) + if error_status: + return { + "error_status": True, + "predictions": None, + "inference_id": inference_id, + } + try: + predictions = REGISTERED_PARSERS[(model_type, task_type)]( + image=image, + parsed_data=parsed_data, + classes=classes, + inference_id=inference_id, + ) + return { + "error_status": False, + "predictions": predictions, + "inference_id": inference_id, + } + except Exception as error: + logging.warning( + f"Could not parse VLM prediction for model {model_type} and task {task_type} " + f"in `roboflow_core/vlm_as_detector@v1` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return { + "error_status": True, + "predictions": None, + "inference_id": inference_id, + } + + +def string2json( + raw_json: str, +) -> Tuple[bool, Union[dict, list]]: + json_blocks_found = JSON_MARKDOWN_BLOCK_PATTERN.findall(raw_json) + if len(json_blocks_found) == 0: + return try_parse_json(raw_json) + first_block = json_blocks_found[0] + return try_parse_json(first_block) + + +def try_parse_json(content: str) -> Tuple[bool, Union[dict, list]]: + try: + parsed = json.loads(content) + if isinstance(parsed, (dict, list)): + return False, parsed + logging.warning( + "Could not parse JSON to dict in `roboflow_core/vlm_as_detector@v1` block. " + f"Unexpected JSON root type: {type(parsed).__name__}." + ) + return True, {} + except Exception as error: + logging.warning( + f"Could not parse JSON to dict in `roboflow_core/vlm_as_detector@v1` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return True, {} + + +def build_image_metadata( + image: WorkflowImageData, + image_height: int, + image_width: int, + inference_id: str, + class_names: Dict[int, str], +) -> dict: + # Per-image state (class_names map, lineage, inference_id) lives on + # `image_metadata`. + parent = image.parent_metadata + root = image.workflow_root_ancestor_metadata + parent_coordinates = parent.origin_coordinates + root_coordinates = root.origin_coordinates + return { + CLASS_NAMES_KEY: class_names, + PREDICTION_TYPE_KEY: "object-detection", + IMAGE_DIMENSIONS_KEY: [image_height, image_width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: parent.parent_id, + PARENT_COORDINATES_KEY: [ + parent_coordinates.left_top_x, + parent_coordinates.left_top_y, + ], + PARENT_DIMENSIONS_KEY: [ + parent_coordinates.origin_height, + parent_coordinates.origin_width, + ], + ROOT_PARENT_ID_KEY: root.parent_id, + ROOT_PARENT_COORDINATES_KEY: [ + root_coordinates.left_top_x, + root_coordinates.left_top_y, + ], + ROOT_PARENT_DIMENSIONS_KEY: [ + root_coordinates.origin_height, + root_coordinates.origin_width, + ], + } + + +def native_detections_from_parsed( + image: WorkflowImageData, + image_height: int, + image_width: int, + inference_id: str, + xyxy: np.ndarray, + class_id: np.ndarray, + class_name: List[str], + confidence: np.ndarray, +) -> Detections: + # `class_names` maps each resolved class_id -> the class name string the VLM + # produced; built from the (class_id, class_name) pairs actually present so + # the serialiser can resolve every id, including unmapped ids (class_id == -1). + number_of_detections = len(xyxy) + class_names = { + int(detection_class_id): str(detection_class_name) + for detection_class_id, detection_class_name in zip(class_id, class_name) + } + image_metadata = build_image_metadata( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + class_names=class_names, + ) + # The per-box VLM label string is carried on bboxes_metadata[i]["class"] so + # distinct unmapped labels (all sharing class_id == -1) survive: the + # serialiser prefers this per-box label over the class_id -> name map. + bboxes_metadata = ( + [ + { + DETECTION_ID_KEY: str(uuid4()), + CLASS_NAME_KEY: str(detection_class_name), + } + for detection_class_name in class_name + ] + if number_of_detections > 0 + else None + ) + return Detections( + xyxy=torch.as_tensor( + np.asarray(xyxy), + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ).reshape(-1, 4), + class_id=torch.as_tensor( + np.asarray(class_id), + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ).reshape(-1), + confidence=torch.as_tensor( + np.asarray(confidence), + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ).reshape(-1), + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def empty_native_detections( + image: WorkflowImageData, + image_height: int, + image_width: int, + inference_id: str, + class_names: Dict[int, str], +) -> Detections: + image_metadata = build_image_metadata( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + class_names=class_names, + ) + return Detections( + xyxy=torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + class_id=torch.zeros( + (0,), dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.zeros( + (0,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + image_metadata=image_metadata, + bboxes_metadata=None, + ) + + +def parse_gemini_object_detection_response( + image: WorkflowImageData, + parsed_data: Union[dict, list], + classes: List[str], + inference_id: str, +) -> Detections: + class_name2id = create_classes_index(classes=classes) + image_height, image_width = image._read_shape_without_materialization() + detections = extract_gemini_detection_entries(parsed_data=parsed_data) + if len(detections) == 0: + return empty_native_detections( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + class_names={idx: class_name for class_name, idx in class_name2id.items()}, + ) + + xyxy, class_id, class_name, confidence = [], [], [], [] + for detection in detections: + xyxy.append( + convert_gemini_detection_to_pixel_xyxy( + detection=detection, + image_height=image_height, + image_width=image_width, + ) + ) + label = get_gemini_detection_class_name(detection=detection) + class_id.append(class_name2id.get(label, -1)) + class_name.append(label) + confidence.append(scale_confidence(detection.get("confidence", 1.0))) + + xyxy = np.array(xyxy).round(0) if len(xyxy) > 0 else np.empty((0, 4)) + confidence = np.array(confidence) if len(confidence) > 0 else np.empty(0) + class_id = np.array(class_id).astype(int) if len(class_id) > 0 else np.empty(0) + return native_detections_from_parsed( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + xyxy=xyxy, + class_id=class_id, + class_name=class_name, + confidence=confidence, + ) + + +def parse_florence2_object_detection_response( + image: WorkflowImageData, + parsed_data: dict, + classes: Optional[List[str]], + inference_id: str, + florence_task_type: str, +): + image_height, image_width = image._read_shape_without_materialization() + # sv.Detections.from_lmm is used purely as the Florence-2 parsing algorithm; + # the output is built natively from the arrays it returns. + detections = sv.Detections.from_lmm( + "florence_2", + result={florence_task_type: parsed_data}, + resolution_wh=(image_width, image_height), + ) + detections.class_id = np.array([0] * len(detections)) + if florence_task_type == "": + detections.data["class_name"] = np.array(["roi"] * len(detections)) + if florence_task_type in {"", ""}: + unique_class_names = set(detections.data.get("class_name", [])) + class_name_to_id = { + name: get_4digit_from_md5(name) for name in unique_class_names + } + class_ids = [ + class_name_to_id.get(name, -1) + for name in detections.data.get("class_name", ["unknown"] * len(detections)) + ] + detections.class_id = np.array(class_ids) + if florence_task_type == "": + class_name_to_id = {name: idx for idx, name in enumerate(classes)} + class_ids = [ + class_name_to_id.get(name, -1) + for name in detections.data.get("class_name", ["unknown"] * len(detections)) + ] + detections.class_id = np.array(class_ids) + detections.confidence = np.array([1.0 for _ in detections]) + class_name = list(detections.data.get("class_name", ["unknown"] * len(detections))) + return native_detections_from_parsed( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + xyxy=detections.xyxy, + class_id=detections.class_id, + class_name=class_name, + confidence=detections.confidence, + ) + + +def get_4digit_from_md5(input_string): + md5_hash = hashlib.md5(input_string.encode("utf-8")) + hex_digest = md5_hash.hexdigest() + integer_value = int(hex_digest[:9], 16) + return integer_value % 10000 + + +REGISTERED_PARSERS = { + ("google-gemini", "object-detection"): parse_gemini_object_detection_response, + ("anthropic-claude", "object-detection"): parse_gemini_object_detection_response, + ("florence-2", "object-detection"): partial( + parse_florence2_object_detection_response, florence_task_type="" + ), + ("florence-2", "open-vocabulary-object-detection"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), + ("florence-2", "object-detection-and-caption"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), + ("florence-2", "phrase-grounded-object-detection"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), + ("florence-2", "region-proposal"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), + ("florence-2", "ocr-with-text-detection"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), +} diff --git a/inference/core/workflows/core_steps/formatters/vlm_as_detector/v2_tensor.py b/inference/core/workflows/core_steps/formatters/vlm_as_detector/v2_tensor.py new file mode 100644 index 0000000000..4fb85c3e9c --- /dev/null +++ b/inference/core/workflows/core_steps/formatters/vlm_as_detector/v2_tensor.py @@ -0,0 +1,669 @@ +import hashlib +import json +import logging +import re +from functools import partial +from typing import Dict, List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field, model_validator + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.vlms import VLM_TASKS_METADATA +from inference.core.workflows.core_steps.formatters.vlm_as_detector.gemini_detection_parsing import ( + convert_gemini_detection_to_pixel_xyxy, + create_classes_index, + extract_gemini_detection_entries, + get_gemini_detection_class_name, + scale_confidence, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + LANGUAGE_MODEL_OUTPUT_KIND, + LIST_OF_VALUES_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections + +JSON_MARKDOWN_BLOCK_PATTERN = re.compile(r"```json([\s\S]*?)```", flags=re.IGNORECASE) + +LONG_DESCRIPTION = """ +Parse JSON strings from Visual Language Models (VLMs) and Large Language Models (LLMs) into standardized object detection prediction format by extracting bounding boxes, class names, and confidences, converting normalized coordinates to pixel coordinates, mapping class names to class IDs, and handling multiple model types and task formats to enable VLM-based object detection, LLM detection parsing, and text-to-detection conversion workflows. + +## How This Block Works + +This block converts VLM/LLM text outputs containing object detection predictions into standardized object detection format compatible with workflow detection blocks. The block: + +1. Receives image and VLM output string containing detection results in JSON format +2. Parses JSON content from VLM output: + + **Handles Markdown-wrapped JSON:** + - Searches for JSON wrapped in Markdown code blocks (```json ... ```) + - This format is common in LLM/VLM responses + - If multiple markdown JSON blocks are found, only the first block is parsed + - Extracts JSON content from within markdown tags + + **Handles raw JSON strings:** + - If no markdown blocks are found, attempts to parse the entire string as JSON + - Supports standard JSON format strings +3. Selects appropriate parser based on model type and task type: + - Uses registered parsers that handle different model outputs (google-gemini, anthropic-claude, florence-2, openai) + - Supports multiple task types: object-detection, open-vocabulary-object-detection, object-detection-and-caption, phrase-grounded-object-detection, region-proposal, ocr-with-text-detection + - Each model/task combination uses a specialized parser for that format +4. Parses detection data based on model type: + + **For OpenAI/Gemini/Claude models:** + - Extracts detections array from parsed JSON + - Converts normalized coordinates (0-1 range) to pixel coordinates using image dimensions + - Extracts class names, confidence scores, and bounding box coordinates + - Maps class names to class IDs using provided classes list + - Creates detection objects with bounding boxes, classes, and confidences + + **For Florence-2 model:** + - Uses supervision's built-in LMM parser for Florence-2 format + - Handles different task types with specialized parsing (object detection, open vocabulary, region proposal, OCR, etc.) + - For region proposal tasks: assigns "roi" as class name + - For open vocabulary detection: uses provided classes list for class ID mapping + - For other tasks: uses MD5-based class ID generation or provided classes + - Sets confidence to 1.0 for Florence-2 detections (model doesn't provide confidence) +5. Converts coordinates and normalizes data: + - Converts normalized coordinates (0-1) to absolute pixel coordinates (x_min, y_min, x_max, y_max) + - Scales coordinates using image width and height + - Normalizes confidence scores to valid range [0.0, 1.0] + - Clamps confidence values outside the range +6. Creates class name to class ID mapping: + - For OpenAI/Gemini/Claude: uses provided classes list to create index mapping (class_name โ†’ class_id) + - Classes are mapped in order (first class = ID 0, second = ID 1, etc.) + - Classes not in the provided list get class_id = -1 + - For Florence-2: uses different mapping strategies based on task type +7. Constructs object detection predictions: + - Creates supervision Detections objects with bounding boxes (xyxy format) + - Includes class IDs, class names, and confidence scores + - Adds metadata: detection IDs, inference IDs, image dimensions, prediction type + - Attaches parent coordinates for crop-aware detections + - Formats predictions in standard object detection format +8. Handles errors: + - Sets `error_status` to True if JSON parsing fails + - Sets `error_status` to True if detection parsing fails + - Returns None for predictions when errors occur + - Always includes inference_id for tracking +9. Returns object detection predictions: + - Outputs `predictions` in standard object detection format (compatible with detection blocks) + - Outputs `error_status` indicating parsing success/failure + - Outputs `inference_id` for tracking and lineage + +The block enables using VLMs/LLMs for object detection by converting their text-based JSON outputs into standardized detection predictions that can be used in workflows like any other object detection model output. + +## Common Use Cases + +- **VLM-Based Object Detection**: Use Visual Language Models for object detection by parsing VLM outputs into detection predictions (e.g., detect objects with GPT-4V, use Claude Vision for detection, parse Gemini detection outputs), enabling VLM detection workflows +- **Open-Vocabulary Detection**: Use VLMs for open-vocabulary object detection with custom classes (e.g., detect custom objects with VLMs, use open-vocabulary detection, detect objects not in training set), enabling open-vocabulary detection workflows +- **Multi-Task Detection**: Use VLMs for various detection tasks (e.g., object detection with captions, phrase-grounded detection, region proposal, OCR with detection), enabling multi-task detection workflows +- **LLM Detection Parsing**: Parse LLM text outputs containing detection results into standardized format (e.g., parse GPT detection outputs, convert LLM predictions to detection format, use LLMs for detection), enabling LLM detection workflows +- **Text-to-Detection Conversion**: Convert text-based detection outputs from models into workflow-compatible detection predictions (e.g., convert text predictions to detection format, parse text-based detections, convert model outputs to detections), enabling text-to-detection workflows +- **VLM Integration**: Integrate VLM outputs into detection workflows (e.g., use VLMs in detection pipelines, integrate VLM predictions with detection blocks, combine VLM and traditional detection), enabling VLM integration workflows + +## Connecting to Other Blocks + +This block receives images and VLM outputs and produces object detection predictions: + +- **After VLM/LLM blocks** to parse detection outputs into standard format (e.g., VLM output to detections, LLM output to detections, parse model outputs), enabling VLM-to-detection workflows +- **Before detection-based blocks** to use parsed detections (e.g., use parsed detections in workflows, provide detections to downstream blocks, use VLM detections with detection blocks), enabling detection-to-workflow workflows +- **Before filtering blocks** to filter VLM detections (e.g., filter by class, filter by confidence, apply filters to VLM predictions), enabling detection-to-filter workflows +- **Before analytics blocks** to analyze VLM detection results (e.g., analyze VLM detections, perform analytics on parsed detections, track VLM detection metrics), enabling detection analytics workflows +- **Before visualization blocks** to display VLM detection results (e.g., visualize VLM detections, display parsed detection predictions, show VLM detection outputs), enabling detection visualization workflows +- **In workflow outputs** to provide VLM detections as final output (e.g., VLM detection outputs, parsed detection results, VLM-based detection outputs), enabling detection output workflows + +## Version Differences + +This version (v2) includes the following enhancements over v1: + +- **Improved Type System**: The `inference_id` output now uses `INFERENCE_ID_KIND` instead of `STRING_KIND`, providing better type safety and semantic meaning for inference tracking identifiers in the workflow system +- **OpenAI Model Support**: Added support for OpenAI models in addition to Google Gemini, Anthropic Claude, and Florence-2 models, expanding the range of VLM/LLM models that can be used for object detection +- **Enhanced Type Safety**: Improved type system ensures better integration with workflow execution engine and provides clearer semantic meaning for inference tracking + +## Requirements + +This block requires an image input (for metadata and dimensions) and a VLM output string containing JSON detection data. The JSON can be raw JSON or wrapped in Markdown code blocks (```json ... ```). The block supports four model types: "openai", "google-gemini", "anthropic-claude", and "florence-2". It supports multiple task types: "object-detection", "open-vocabulary-object-detection", "object-detection-and-caption", "phrase-grounded-object-detection", "region-proposal", and "ocr-with-text-detection". The `classes` parameter is required for OpenAI, Gemini, and Claude models (to map class names to IDs) but optional for Florence-2 (some tasks don't require it). Classes are mapped to IDs by index (first class = 0, second = 1, etc.). Classes not in the list get class_id = -1. The block outputs object detection predictions in standard format (compatible with detection blocks), error_status (boolean), and inference_id (INFERENCE_ID_KIND) for tracking. +""" + +SHORT_DESCRIPTION = "Parses raw string into object-detection prediction." + +SUPPORTED_TASKS = { + "object-detection", + "object-detection-and-caption", + "open-vocabulary-object-detection", + "phrase-grounded-object-detection", + "region-proposal", + "ocr-with-text-detection", +} +RELEVANT_TASKS_METADATA = { + k: v for k, v in VLM_TASKS_METADATA.items() if k in SUPPORTED_TASKS +} + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "VLM As Detector", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "formatter", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-object-ungroup", + "blockPriority": 5, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/vlm_as_detector@v2"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Input image that was used to generate the VLM prediction. Used to extract image dimensions (width, height) for converting normalized coordinates to pixel coordinates and metadata (parent_id) for the detection predictions. The same image that was provided to the VLM/LLM block should be used here to maintain consistency.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + vlm_output: Selector(kind=[LANGUAGE_MODEL_OUTPUT_KIND]) = Field( + title="VLM Output", + description="String output from a VLM or LLM block containing object detection prediction in JSON format. Can be raw JSON string or JSON wrapped in Markdown code blocks (e.g., ```json {...} ```). Format depends on model_type and task_type - different models and tasks produce different JSON structures. If multiple markdown blocks exist, only the first is parsed.", + examples=[ + ["$steps.lmm.output"], + ["$steps.vlm.output"], + ["$steps.claude.output"], + ], + ) + classes: Optional[ + Union[ + Selector(kind=[LIST_OF_VALUES_KIND]), + Selector(kind=[LIST_OF_VALUES_KIND]), + List[str], + ] + ] = Field( + description="List of all class names used by the classification model, in order. Required to generate mapping between class names (from VLM output) and class IDs (for detection format). Classes are mapped to IDs by index: first class = ID 0, second = ID 1, etc. Classes from VLM output that are not in this list get class_id = -1. Required for OpenAI, Gemini, and Claude models. Optional for Florence-2 (some tasks don't require it). Should match the classes the VLM was asked to detect.", + examples=[ + [ + "$steps.lmm.classes", + "$inputs.classes", + ["dog", "cat", "bird"], + ["class_a", "class_b"], + ] + ], + json_schema_extra={ + "relevant_for": { + "model_type": { + "values": ["openai", "google-gemini", "anthropic-claude"], + "required": True, + }, + } + }, + ) + model_type: Literal["openai", "google-gemini", "anthropic-claude", "florence-2"] = ( + Field( + description="Type of the VLM/LLM model that generated the prediction. Determines which parser is used to extract detection data from the JSON output. Supported models: 'openai' (GPT-4V), 'google-gemini' (Gemini Vision), 'anthropic-claude' (Claude Vision), 'florence-2' (Microsoft Florence-2). Each model type has different JSON output formats, so the correct model type must be specified for proper parsing.", + examples=[ + ["openai"], + ["google-gemini"], + ["anthropic-claude"], + ["florence-2"], + ], + ) + ) + task_type: Literal[tuple(SUPPORTED_TASKS)] = Field( + description="Task type performed by the VLM/LLM model. Determines which parser and format handler is used. Supported task types: 'object-detection' (standard object detection), 'open-vocabulary-object-detection' (detect objects with custom classes), 'object-detection-and-caption' (detection with captions), 'phrase-grounded-object-detection' (ground phrases to detections), 'region-proposal' (propose regions of interest), 'ocr-with-text-detection' (OCR with text region detection). The task type must match what the VLM/LLM was asked to perform.", + json_schema_extra={ + "values_metadata": RELEVANT_TASKS_METADATA, + }, + ) + + @model_validator(mode="after") + def validate(self) -> "BlockManifest": + if (self.model_type, self.task_type) not in REGISTERED_PARSERS: + raise ValueError( + f"Could not parse result of task {self.task_type} for model {self.model_type}" + ) + if self.model_type != "florence-2" and self.classes is None: + raise ValueError( + "Must pass list of classes to this block when using gemini or claude" + ) + + return self + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition(name="inference_id", kind=[INFERENCE_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class VLMAsDetectorBlockV2(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + vlm_output: str, + classes: Optional[List[str]], + model_type: str, + task_type: str, + ) -> BlockResult: + inference_id = f"{uuid4()}" + error_status, parsed_data = string2json( + raw_json=vlm_output, + ) + if error_status: + return { + "error_status": True, + "predictions": None, + "inference_id": inference_id, + } + try: + predictions = REGISTERED_PARSERS[(model_type, task_type)]( + image=image, + parsed_data=parsed_data, + classes=classes, + inference_id=inference_id, + ) + return { + "error_status": False, + "predictions": predictions, + "inference_id": inference_id, + } + except Exception as error: + logging.warning( + f"Could not parse VLM prediction for model {model_type} and task {task_type} " + f"in `roboflow_core/vlm_as_detector@v2` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return { + "error_status": True, + "predictions": None, + "inference_id": inference_id, + } + + +def string2json( + raw_json: str, +) -> Tuple[bool, Union[dict, list]]: + json_blocks_found = JSON_MARKDOWN_BLOCK_PATTERN.findall(raw_json) + if len(json_blocks_found) == 0: + return try_parse_json(raw_json) + first_block = json_blocks_found[0] + return try_parse_json(first_block) + + +def try_parse_json(content: str) -> Tuple[bool, Union[dict, list]]: + try: + parsed = json.loads(content) + if isinstance(parsed, (dict, list)): + return False, parsed + logging.warning( + "Could not parse JSON to dict in `roboflow_core/vlm_as_detector@v2` block. " + f"Unexpected JSON root type: {type(parsed).__name__}." + ) + return True, {} + except Exception as error: + logging.warning( + f"Could not parse JSON to dict in `roboflow_core/vlm_as_detector@v2` block. " + f"Error type: {error.__class__.__name__}. Details: {error}" + ) + return True, {} + + +def build_image_metadata( + image: WorkflowImageData, + image_height: int, + image_width: int, + inference_id: str, + class_names: Dict[int, str], +) -> dict: + # Per-image state (class_names map, lineage, inference_id) lives on + # `image_metadata`. + parent = image.parent_metadata + root = image.workflow_root_ancestor_metadata + parent_coordinates = parent.origin_coordinates + root_coordinates = root.origin_coordinates + return { + CLASS_NAMES_KEY: class_names, + PREDICTION_TYPE_KEY: "object-detection", + IMAGE_DIMENSIONS_KEY: [image_height, image_width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: parent.parent_id, + PARENT_COORDINATES_KEY: [ + parent_coordinates.left_top_x, + parent_coordinates.left_top_y, + ], + PARENT_DIMENSIONS_KEY: [ + parent_coordinates.origin_height, + parent_coordinates.origin_width, + ], + ROOT_PARENT_ID_KEY: root.parent_id, + ROOT_PARENT_COORDINATES_KEY: [ + root_coordinates.left_top_x, + root_coordinates.left_top_y, + ], + ROOT_PARENT_DIMENSIONS_KEY: [ + root_coordinates.origin_height, + root_coordinates.origin_width, + ], + } + + +def native_detections_from_parsed( + image: WorkflowImageData, + image_height: int, + image_width: int, + inference_id: str, + xyxy: np.ndarray, + class_id: np.ndarray, + class_name: List[str], + confidence: np.ndarray, +) -> Detections: + # `class_names` maps each resolved class_id -> the class name string the VLM + # produced; built from the (class_id, class_name) pairs actually present so + # the serialiser can resolve every id, including unmapped ids (class_id == -1). + number_of_detections = len(xyxy) + class_names = { + int(detection_class_id): str(detection_class_name) + for detection_class_id, detection_class_name in zip(class_id, class_name) + } + image_metadata = build_image_metadata( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + class_names=class_names, + ) + # The per-box VLM label string is carried on bboxes_metadata[i]["class"] so + # distinct unmapped labels (all sharing class_id == -1) survive: the + # serialiser prefers this per-box label over the class_id -> name map. + bboxes_metadata = ( + [ + { + DETECTION_ID_KEY: str(uuid4()), + CLASS_NAME_KEY: str(detection_class_name), + } + for detection_class_name in class_name + ] + if number_of_detections > 0 + else None + ) + return Detections( + xyxy=torch.as_tensor( + np.asarray(xyxy), + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ).reshape(-1, 4), + class_id=torch.as_tensor( + np.asarray(class_id), + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ).reshape(-1), + confidence=torch.as_tensor( + np.asarray(confidence), + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ).reshape(-1), + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def empty_native_detections( + image: WorkflowImageData, + image_height: int, + image_width: int, + inference_id: str, + class_names: Dict[int, str], +) -> Detections: + image_metadata = build_image_metadata( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + class_names=class_names, + ) + return Detections( + xyxy=torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + class_id=torch.zeros( + (0,), dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.zeros( + (0,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + image_metadata=image_metadata, + bboxes_metadata=None, + ) + + +def parse_gemini_object_detection_response( + image: WorkflowImageData, + parsed_data: Union[dict, list], + classes: List[str], + inference_id: str, +) -> Detections: + class_name2id = create_classes_index(classes=classes) + image_height, image_width = image._read_shape_without_materialization() + detections = extract_gemini_detection_entries(parsed_data=parsed_data) + if len(detections) == 0: + return empty_native_detections( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + class_names={idx: class_name for class_name, idx in class_name2id.items()}, + ) + + xyxy, class_id, class_name, confidence = [], [], [], [] + for detection in detections: + xyxy.append( + convert_gemini_detection_to_pixel_xyxy( + detection=detection, + image_height=image_height, + image_width=image_width, + ) + ) + label = get_gemini_detection_class_name(detection=detection) + class_id.append(class_name2id.get(label, -1)) + class_name.append(label) + confidence.append(scale_confidence(detection.get("confidence", 1.0))) + + xyxy = np.array(xyxy).round(0) if len(xyxy) > 0 else np.empty((0, 4)) + confidence = np.array(confidence) if len(confidence) > 0 else np.empty(0) + class_id = np.array(class_id).astype(int) if len(class_id) > 0 else np.empty(0) + return native_detections_from_parsed( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + xyxy=xyxy, + class_id=class_id, + class_name=class_name, + confidence=confidence, + ) + + +def parse_llm_object_detection_response( + image: WorkflowImageData, + parsed_data: dict, + classes: List[str], + inference_id: str, +) -> Detections: + class_name2id = create_classes_index(classes=classes) + image_height, image_width = image._read_shape_without_materialization() + detections = extract_gemini_detection_entries(parsed_data=parsed_data) + if len(detections) == 0: + return empty_native_detections( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + class_names={idx: class_name for class_name, idx in class_name2id.items()}, + ) + xyxy, class_id, class_name, confidence = [], [], [], [] + for detection in detections: + xyxy.append( + [ + detection["x_min"] * image_width, + detection["y_min"] * image_height, + detection["x_max"] * image_width, + detection["y_max"] * image_height, + ] + ) + class_id.append(class_name2id.get(detection["class_name"], -1)) + class_name.append(detection["class_name"]) + confidence.append(scale_confidence(detection.get("confidence", 1.0))) + xyxy = np.array(xyxy).round(0) if len(xyxy) > 0 else np.empty((0, 4)) + confidence = np.array(confidence) if len(confidence) > 0 else np.empty(0) + class_id = np.array(class_id).astype(int) if len(class_id) > 0 else np.empty(0) + return native_detections_from_parsed( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + xyxy=xyxy, + class_id=class_id, + class_name=class_name, + confidence=confidence, + ) + + +def create_classes_index(classes: List[str]) -> Dict[str, int]: + return {class_name: idx for idx, class_name in enumerate(classes)} + + +def scale_confidence(value: float) -> float: + return min(max(float(value), 0.0), 1.0) + + +def parse_florence2_object_detection_response( + image: WorkflowImageData, + parsed_data: dict, + classes: Optional[List[str]], + inference_id: str, + florence_task_type: str, +): + image_height, image_width = image._read_shape_without_materialization() + # sv.Detections.from_lmm is used purely as the Florence-2 parsing algorithm; + # the output is built natively from the arrays it returns. + detections = sv.Detections.from_lmm( + "florence_2", + result={florence_task_type: parsed_data}, + resolution_wh=(image_width, image_height), + ) + detections.class_id = np.array([0] * len(detections)) + if florence_task_type == "": + detections.data["class_name"] = np.array(["roi"] * len(detections)) + if florence_task_type in {"", ""}: + unique_class_names = set(detections.data.get("class_name", [])) + class_name_to_id = { + name: get_4digit_from_md5(name) for name in unique_class_names + } + class_ids = [ + class_name_to_id.get(name, -1) + for name in detections.data.get("class_name", ["unknown"] * len(detections)) + ] + detections.class_id = np.array(class_ids) + if florence_task_type == "": + class_name_to_id = {name: idx for idx, name in enumerate(classes)} + class_ids = [ + class_name_to_id.get(name, -1) + for name in detections.data.get("class_name", ["unknown"] * len(detections)) + ] + detections.class_id = np.array(class_ids) + detections.confidence = np.array([1.0 for _ in detections]) + class_name = list(detections.data.get("class_name", ["unknown"] * len(detections))) + return native_detections_from_parsed( + image=image, + image_height=image_height, + image_width=image_width, + inference_id=inference_id, + xyxy=detections.xyxy, + class_id=detections.class_id, + class_name=class_name, + confidence=detections.confidence, + ) + + +def get_4digit_from_md5(input_string): + md5_hash = hashlib.md5(input_string.encode("utf-8")) + hex_digest = md5_hash.hexdigest() + integer_value = int(hex_digest[:9], 16) + return integer_value % 10000 + + +REGISTERED_PARSERS = { + # LLMs + ("openai", "object-detection"): parse_llm_object_detection_response, + ("google-gemini", "object-detection"): parse_gemini_object_detection_response, + ("anthropic-claude", "object-detection"): parse_llm_object_detection_response, + # Florence 2 + ("florence-2", "object-detection"): partial( + parse_florence2_object_detection_response, florence_task_type="" + ), + ("florence-2", "open-vocabulary-object-detection"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), + ("florence-2", "object-detection-and-caption"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), + ("florence-2", "phrase-grounded-object-detection"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), + ("florence-2", "region-proposal"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), + ("florence-2", "ocr-with-text-detection"): partial( + parse_florence2_object_detection_response, + florence_task_type="", + ), +} diff --git a/inference/core/workflows/core_steps/fusion/detections_classes_replacement/v1_tensor.py b/inference/core/workflows/core_steps/fusion/detections_classes_replacement/v1_tensor.py new file mode 100644 index 0000000000..5865424585 --- /dev/null +++ b/inference/core/workflows/core_steps/fusion/detections_classes_replacement/v1_tensor.py @@ -0,0 +1,592 @@ +import sys +from typing import Dict, List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import torch +from pydantic import ConfigDict, Field + +from inference.core import logger +from inference.core.workflows.core_steps.common.tensor_native import ( + HOST_MIRROR_KEYS, + TensorNativeDetections, + TensorNativePrediction, + split_key_point_prediction, + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections + +LONG_DESCRIPTION = """ +Replace class labels of detection bounding boxes with classes predicted by a classification model applied to cropped regions, combining generic detection results with specialized classification predictions to enable two-stage detection workflows, fine-grained classification, and class refinement workflows where generic detections are refined with specific class labels from specialized classifiers. + +## How This Block Works + +This block combines results from a detection model (with bounding boxes and generic classes) with classification predictions (from a specialized classifier applied to cropped regions) to replace generic class labels with specific ones. The block: + +1. Receives two inputs with different dimensionality levels: + - `object_detection_predictions`: Detection results (dimensionality level 1) containing bounding boxes with generic classes (e.g., "dog", "person", "vehicle") + - `classification_predictions`: Classification results (dimensionality level 2) from a classifier applied to cropped regions of each detection (e.g., "Golden Retriever", "Labrador" for dog detections). Can also be a list of strings (e.g. from OCR). +2. Matches classifications to detections: + - Uses `PARENT_ID_KEY` (detection_id) in classification predictions to link each classification result to its source detection, OR + - Uses positional mapping (order-based) if predictions are raw strings/lists without parent IDs. +3. Extracts leading class from each classification prediction: + + **For single-label classifications:** + - Uses the "top" class (predicted class) from the classification result + - Extracts class name, class ID, and confidence from the classification prediction + + **For multi-label classifications:** + - Finds the class with the highest confidence score + - Uses the most confident label as the replacement class + - Extracts class name, class ID, and confidence from the highest-confidence prediction + + **For string predictions:** + - Uses the string as the class name + - Assigns a default confidence of 1.0 and class ID of 0 + +4. Handles missing classifications: + - Detections without corresponding classification predictions are discarded by default + - If `fallback_class_name` is provided, detections without classifications use the fallback class instead of being discarded + - Fallback class ID is set to the provided value, or `sys.maxsize` if not specified or negative +5. Filters detections: + - Keeps only detections that have classification results (or fallback if specified) + - Removes detections that cannot be matched to classification predictions +6. Replaces class information: + - Replaces class names in detections with classification class names + - Replaces class IDs in detections with classification class IDs + - Replaces confidence scores in detections with classification confidence scores + - Updates all detection metadata to reflect the new class information +7. Generates new detection IDs: + - Creates new unique detection IDs for updated detections (prevents ID conflicts) + - Ensures detection IDs are unique after class replacement +8. Returns updated detections: + - Outputs detections with replaced classes, maintaining bounding box coordinates and other properties + - Output dimensionality matches input detection predictions (dimensionality level 1) + +The block enables two-stage detection workflows where a generic detection model locates objects and a specialized classification model provides fine-grained labels. This is useful when you need generic localization (e.g., "dog") combined with specific classification (e.g., "Golden Retriever", "German Shepherd") without losing spatial information. + +## Common Use Cases + +- **Two-Stage Detection and Classification**: Combine generic detection with specialized classification for fine-grained labeling (e.g., detect "dog" then classify breed, detect "vehicle" then classify type, detect "person" then classify age group), enabling two-stage detection workflows +- **Class Refinement**: Refine generic class labels with specific classifications from specialized models (e.g., refine "animal" to specific species, refine "vehicle" to specific models, refine "food" to specific dishes), enabling class refinement workflows +- **Multi-Model Workflows**: Combine detection and classification models to leverage the strengths of both (e.g., use generic detector for localization and specialist classifier for identification, combine coarse and fine-grained models, leverage specialized classifiers with general detectors), enabling multi-model workflows +- **Hierarchical Classification**: Apply hierarchical classification where detection provides high-level classes and classification provides detailed sub-classes (e.g., detect "mammal" then classify species, detect "plant" then classify variety, detect "structure" then classify type), enabling hierarchical classification workflows +- **Crop-Based Classification**: Use classification results from cropped regions to enhance detection results (e.g., classify crops to improve detection labels, apply specialized classifiers to detected regions, refine detections with crop classifications), enabling crop-based classification workflows +- **Fine-Grained Object Recognition**: Enable fine-grained recognition by combining localization and detailed classification (e.g., recognize specific product models, identify specific animal breeds, classify specific vehicle types), enabling fine-grained recognition workflows + +## Connecting to Other Blocks + +This block receives detection and classification predictions and produces detections with replaced classes: + +- **After detection and classification model blocks** to combine generic detection with specialized classification (e.g., object detection + classification to refined detections, detection model + classifier to labeled detections), enabling detection-classification fusion workflows +- **After crop blocks** that create crops from detections for classification (e.g., crop detections then classify crops, create crops for classification then replace classes), enabling crop-classification workflows +- **Before visualization blocks** to display detections with refined classes (e.g., visualize refined detections, display detections with specific labels, show classification-enhanced detections), enabling refined detection visualization workflows +- **Before filtering blocks** to filter detections with refined classes (e.g., filter by specific classes, filter refined detections, apply filters to classified detections), enabling refined detection filtering workflows +- **Before analytics blocks** to perform analytics on refined detections (e.g., analyze specific classes, perform analytics on classified detections, track refined detection metrics), enabling refined detection analytics workflows +- **In workflow outputs** to provide refined detections as final output (e.g., two-stage detection outputs, classification-enhanced detection outputs, refined detection results), enabling refined detection output workflows + +## Requirements + +This block requires object detection predictions (with bounding boxes) and classification predictions from crops of those bounding boxes. The classification predictions must have `PARENT_ID_KEY` (detection_id) to link classifications to their source detections. The block accepts different dimensionality levels: detection predictions at level 1 and classification predictions at level 2 (from crops). For single-label classifications, the "top" class is used. For multi-label classifications, the most confident class is selected. Detections without classification results are discarded unless `fallback_class_name` is provided. The block outputs detections with replaced classes, class IDs, and confidences, with new detection IDs generated. Output dimensionality matches input detection predictions (level 1). +""" + +SHORT_DESCRIPTION = "Replace classes of detections with classes predicted by a chained classification model." + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Detections Classes Replacement", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "fusion", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-arrow-right-arrow-left", + "blockPriority": 5, + }, + } + ) + type: Literal[ + "roboflow_core/detections_classes_replacement@v1", + "DetectionsClassesReplacement", + ] + object_detection_predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + title="Regions of Interest", + description="Detection predictions (object detection, instance segmentation, or keypoint detection) containing bounding boxes with generic class labels that will be replaced with classification results. These detections should correspond to the regions that were cropped and classified. Detections must have detection IDs that match the PARENT_ID_KEY in classification predictions. Detections at dimensionality level 1.", + examples=[ + "$steps.object_detection_model.predictions", + "$steps.instance_segmentation_model.predictions", + ], + ) + classification_predictions: Selector( + kind=[ + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + ] + ) = Field( + title="Replacement Class Labels", + description="Labels to replace detection class names with. Accepts classification predictions (linked via parent_id), plain strings, or lists of strings (e.g. OCR/LMM output like Gemini). String inputs are matched to detections positionally (1:1 by index). Classification inputs support single-label ('top' class) and multi-label (most confident class).", + examples=[ + "$steps.classification_model.predictions", + "$steps.breed_classifier.predictions", + "$steps.ocr_model.predictions", + ], + ) + fallback_class_name: Union[Optional[str], Selector(kind=[STRING_KIND])] = Field( + default=None, + title="Fallback class name", + description="Optional class name to use for detections that don't have corresponding classification predictions. If not provided (default None), detections without classifications are discarded. If provided, detections without classifications use this fallback class name instead of being removed. Useful for preserving detections when classification fails or is unavailable.", + examples=[None, "unknown", "unclassified"], + ) + fallback_class_id: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( + default=None, + title="Fallback class id", + description="Optional class ID to use with fallback_class_name for detections without classification predictions. If not specified or negative, the class ID is set to sys.maxsize. Only used when fallback_class_name is provided. Should match the class ID mapping used in your model.", + examples=[None, 77, 999], + ) + + @classmethod + def accepts_empty_values(cls) -> bool: + return True + + @classmethod + def get_input_dimensionality_offsets(cls) -> Dict[str, int]: + return { + "object_detection_predictions": 0, + "classification_predictions": 1, + } + + @classmethod + def get_dimensionality_reference_property(cls) -> Optional[str]: + return "object_detection_predictions" + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ], + ) + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class DetectionsClassesReplacementBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + object_detection_predictions: Optional[TensorNativePrediction], + classification_predictions: Optional[ + Batch[ + Optional[ + Union[ + ClassificationPrediction, + MultiLabelClassificationPrediction, + str, + List[str], + ] + ] + ] + ], + fallback_class_name: Optional[str], + fallback_class_id: Optional[int], + ) -> BlockResult: + if object_detection_predictions is None: + return {"predictions": None} + # Keypoint predictions arrive as a (KeyPoints, Detections) tuple; class + # replacement only touches the bbox component (used here for length / + # detection_id reads). The selection below re-slices the original input, + # so the keypoints stay aligned and are preserved on the re-wrapped output. + _, detections = split_key_point_prediction(object_detection_predictions) + if not classification_predictions: + return { + "predictions": _select_native_prediction( + object_detection_predictions=object_detection_predictions, + indices=[], + new_class_ids=[], + new_class_names=[], + new_confidences=[], + ) + } + + # Check if predictions are string-based (e.g. from OCR/LMM models) + # rather than classification dicts with parent_id + first_valid_pred = next( + (p for p in classification_predictions if p is not None), None + ) + is_string_prediction = isinstance(first_valid_pred, (str, list)) + + if is_string_prediction: + if len(detections) != len(classification_predictions): + logger.warning( + "Detections count (%d) does not match classification predictions " + "count (%d). Unmatched detections will use the fallback class " + "(if configured) or be discarded.", + len(detections), + len(classification_predictions), + ) + # Pad classification_predictions with None so every detection is + # processed through the fallback path instead of being silently + # truncated by zip. + padded_predictions = list(classification_predictions) + [None] * ( + len(detections) - len(classification_predictions) + ) + classification_predictions = padded_predictions + + new_class_names = [] + new_class_ids = [] + new_confidences = [] + valid_indices = [] + + for i, (det_idx, prediction) in enumerate( + zip(range(len(detections)), classification_predictions) + ): + if prediction is None: + if fallback_class_name: + resolved_fallback_id = ( + sys.maxsize + if fallback_class_id is None or fallback_class_id < 0 + else fallback_class_id + ) + class_name, class_id, confidence = ( + fallback_class_name, + resolved_fallback_id, + 0.0, + ) + else: + continue + else: + extracted = extract_leading_class_from_prediction( + prediction, + fallback_class_name=fallback_class_name, + fallback_class_id=fallback_class_id, + ) + if extracted is None: + continue + class_name, class_id, confidence = extracted + + new_class_names.append(class_name) + new_class_ids.append(class_id) + new_confidences.append(confidence) + valid_indices.append(i) + + if not valid_indices: + return { + "predictions": _select_native_prediction( + object_detection_predictions=object_detection_predictions, + indices=[], + new_class_ids=[], + new_class_names=[], + new_confidences=[], + ) + } + + # Filter detections to keep only those with valid predictions, then + # replace class_id / confidence tensors, refresh the class_names map + # and mint new detection IDs - all natively, no sv round-trip. + return { + "predictions": _select_native_prediction( + object_detection_predictions=object_detection_predictions, + indices=valid_indices, + new_class_ids=new_class_ids, + new_class_names=new_class_names, + new_confidences=new_confidences, + ) + } + + if all( + _native_classification_prediction_is_empty(p) + for p in classification_predictions + ): + return { + "predictions": _select_native_prediction( + object_detection_predictions=object_detection_predictions, + indices=[], + new_class_ids=[], + new_class_names=[], + new_confidences=[], + ) + } + detection_id_by_class: Dict[str, Optional[Tuple[str, int, float]]] = { + _native_classification_parent_id( + prediction + ): extract_leading_class_from_prediction( + prediction=prediction, + fallback_class_name=fallback_class_name, + fallback_class_id=fallback_class_id, + ) + for prediction in classification_predictions + if prediction is not None + } + bboxes_metadata = detections.bboxes_metadata or [ + {} for _ in range(len(detections)) + ] + detection_ids = [box.get(DETECTION_ID_KEY) for box in bboxes_metadata] + detections_to_remain_indices = [ + index + for index, detection_id in enumerate(detection_ids) + if detection_id_by_class.get(detection_id) is not None + ] + replaced_class_names = [ + detection_id_by_class[detection_ids[index]][0] + for index in detections_to_remain_indices + ] + replaced_class_ids = [ + detection_id_by_class[detection_ids[index]][1] + for index in detections_to_remain_indices + ] + replaced_confidences = [ + detection_id_by_class[detection_ids[index]][2] + for index in detections_to_remain_indices + ] + return { + "predictions": _select_native_prediction( + object_detection_predictions=object_detection_predictions, + indices=detections_to_remain_indices, + new_class_ids=replaced_class_ids, + new_class_names=replaced_class_names, + new_confidences=replaced_confidences, + ) + } + + +def _select_native_prediction( + object_detection_predictions: TensorNativePrediction, + indices: List[int], + new_class_ids: List[int], + new_class_names: List[str], + new_confidences: List[float], +) -> TensorNativePrediction: + """Select surviving detections by index and overwrite their class_id / + confidence with the classification result, refreshing the class_names map and + minting new detection IDs. Returns the same native shape as the input + (re-wrapping the keypoint tuple when keypoints are present). + + NOTE (shared-helper candidate): overwriting class_id/confidence + rebuilding + the class_names map + writing fresh detection_ids on a native prediction is a + pattern that several class-mutating siblings will need - flagged for the + maintainer to consolidate. + """ + selected = take_prediction_by_indices(object_detection_predictions, indices) + selected_detections: TensorNativeDetections = ( + selected[1] if isinstance(selected, tuple) else selected + ) + # `take_detections_by_indices` aliases `xyxy` (and a dense mask) by reference + # on an identity selection (all rows kept in order). The numpy block always + # produces a fresh copy of the box array via sv boolean indexing, so its + # output never shares storage with the input. Clone the spatial tensors here + # to honour that copy-of-data contract: the returned prediction must not alias + # the input's `xyxy` (or dense mask), matching the numpy block byte-for-byte. + selected_detections.xyxy = selected_detections.xyxy.clone() + if isinstance(selected_detections, InstanceDetections) and isinstance( + selected_detections.mask, torch.Tensor + ): + selected_detections.mask = selected_detections.mask.clone() + selected_detections.class_id = torch.as_tensor( + new_class_ids, + dtype=selected_detections.class_id.dtype, + device=selected_detections.class_id.device, + ) + selected_detections.confidence = torch.as_tensor( + new_confidences, + dtype=selected_detections.confidence.dtype, + device=selected_detections.confidence.device, + ) + # Refresh the class_id -> name map so the per-box class_id resolves to the + # replacement class name on serialisation (numpy wrote a per-box class_name + # string column; natively the name lives in image_metadata["class_names"]). + image_metadata = dict(selected_detections.image_metadata or {}) + image_metadata[CLASS_NAMES_KEY] = { + int(class_id): class_name + for class_id, class_name in zip(new_class_ids, new_class_names) + } + selected_detections.image_metadata = image_metadata + # Mint fresh detection IDs (the numpy block regenerates them to avoid + # collisions after class replacement). The per-box host mirror is dropped: + # class_id / confidence were replaced above, so a carried mirror would be + # stale โ€” consumers fall back to tensor reads. + selected_detections.bboxes_metadata = [ + { + **{ + key: value + for key, value in (box or {}).items() + if key not in HOST_MIRROR_KEYS + }, + DETECTION_ID_KEY: f"{uuid4()}", + } + for box in ( + selected_detections.bboxes_metadata + if selected_detections.bboxes_metadata is not None + else [{} for _ in range(len(selected_detections))] + ) + ] + if isinstance(selected, tuple): + return selected[0], selected_detections + return selected_detections + + +def _native_classification_parent_id( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> Optional[str]: + if isinstance(prediction, ClassificationPrediction): + images_metadata = prediction.images_metadata or [{}] + return images_metadata[0].get(PARENT_ID_KEY) + image_metadata = prediction.image_metadata or {} + return image_metadata.get(PARENT_ID_KEY) + + +def _native_classification_prediction_is_empty( + prediction: Optional[ + Union[ClassificationPrediction, MultiLabelClassificationPrediction] + ], +) -> bool: + """Native equivalent of the numpy early-out check + `p is None or "top" in p and not p["top"] or "predictions" not in p`: + a single-label prediction is "empty" when it carries no class distribution + (no "predictions"); a multi-label prediction is only empty when None. + """ + if prediction is None: + return True + if isinstance(prediction, ClassificationPrediction): + return int(prediction.confidence.shape[-1]) == 0 + return False + + +def extract_leading_class_from_prediction( + prediction: Union[ + ClassificationPrediction, MultiLabelClassificationPrediction, str, List[str] + ], + fallback_class_name: Optional[str] = None, + fallback_class_id: Optional[int] = None, +) -> Optional[Tuple[str, int, float]]: + if isinstance(prediction, str): + return prediction, 0, 1.0 + + if isinstance(prediction, list): + if not prediction: + if fallback_class_name: + try: + fallback_class_id = int(fallback_class_id) + except (ValueError, TypeError): + fallback_class_id = None + if fallback_class_id is None or fallback_class_id < 0: + fallback_class_id = sys.maxsize + return fallback_class_name, fallback_class_id, 0.0 + return None + + # Take the first string in the list if it contains strings + # Recursive call would be cleaner but let's be explicit + first_item = prediction[0] + if isinstance(first_item, str): + return first_item, 0, 1.0 + # If it's a list of something else (not expected for now based on user request "nested arrays"), + # we might need recursion or just fail. + # User said: "gemini_out": [ ["K619879"], ... ] which is List[List[str]] + # So prediction passed here is ["K619879"] (List[str]) + + return None + + if isinstance(prediction, ClassificationPrediction): + # Single-label: "top" class is the predicted class_id, "predictions" is + # the full distribution. An empty distribution mirrors the numpy + # `not prediction.get("predictions")` branch. + has_predictions = int(prediction.confidence.shape[-1]) > 0 + if not has_predictions and not fallback_class_name: + return None + elif not has_predictions and fallback_class_name: + try: + fallback_class_id = int(fallback_class_id) + except (ValueError, TypeError): + fallback_class_id = None + if fallback_class_id is None or fallback_class_id < 0: + fallback_class_id = sys.maxsize + return fallback_class_name, fallback_class_id, 0 + images_metadata = prediction.images_metadata or [{}] + class_names = images_metadata[0].get(CLASS_NAMES_KEY) or {} + class_id = int(prediction.class_id[0]) + class_name = class_names.get(class_id, f"class_{class_id}") + confidence = float(prediction.confidence[0, class_id]) + return class_name, class_id, confidence + predicted_classes = ( + prediction.class_ids.tolist() if prediction.class_ids is not None else [] + ) + if not predicted_classes: + return None + image_metadata = prediction.image_metadata or {} + class_names = image_metadata.get(CLASS_NAMES_KEY) or {} + # Mirrors the numpy block: once at least one class is predicted, the + # replacement is the globally most-confident class over the FULL sigmoid + # distribution (numpy iterated `prediction["predictions"].items()`, i.e. all + # classes - not only the above-threshold `predicted_classes`). + max_confidence, max_confidence_class_name, max_confidence_class_id = ( + None, + None, + None, + ) + for current_class_id in range(int(prediction.confidence.shape[-1])): + current_class_confidence = float(prediction.confidence[current_class_id]) + current_class_name = class_names.get( + current_class_id, f"class_{current_class_id}" + ) + if max_confidence is None: + max_confidence = current_class_confidence + max_confidence_class_name = current_class_name + max_confidence_class_id = current_class_id + continue + if max_confidence < current_class_confidence: + max_confidence = current_class_confidence + max_confidence_class_name = current_class_name + max_confidence_class_id = current_class_id + if not max_confidence: + return None + return max_confidence_class_name, max_confidence_class_id, max_confidence diff --git a/inference/core/workflows/core_steps/fusion/detections_consensus/v1_tensor.py b/inference/core/workflows/core_steps/fusion/detections_consensus/v1_tensor.py new file mode 100644 index 0000000000..1d496aae30 --- /dev/null +++ b/inference/core/workflows/core_steps/fusion/detections_consensus/v1_tensor.py @@ -0,0 +1,1178 @@ +import math +import statistics +from collections import Counter +from enum import Enum +from functools import lru_cache +from typing import ( + Any, + Dict, + Generator, + List, + Literal, + Optional, + Set, + Tuple, + Type, + Union, +) +from uuid import uuid4 + +import numpy as np +import supervision as sv +import torch +from pydantic import AliasChoices, ConfigDict, Field, PositiveInt + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.logger import logger +from inference.core.workflows.core_steps.common.tensor_native import ( + instance_mask_to_numpy, + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, + SCALING_RELATIVE_TO_PARENT_KEY, + SCALING_RELATIVE_TO_ROOT_PARENT_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + DICTIONARY_KIND, + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +# Tensor-native detections handled by this block. The consensus pipeline only +# ever needs the bounding-box component, so keypoint predictions (which arrive +# as a `(KeyPoints, Detections)` tuple) are reduced to their `Detections` part +# up-front and the consensus output is always a plain `Detections` / +# `InstanceDetections` (the output kind never includes keypoints, matching the +# numpy block). +TensorNativeDetections = Union[Detections, InstanceDetections] + + +class AggregationMode(Enum): + AVERAGE = "average" + MAX = "max" + MIN = "min" + + +class MaskAggregationMode(Enum): + INTERSECTION = "intersection" + UNION = "union" + MAX = "max" + MIN = "min" + + +LONG_DESCRIPTION = """ +Combine detection predictions from multiple models using a majority vote consensus strategy, merging overlapping detections that receive sufficient votes from different models and aggregating their properties (confidence scores, bounding boxes, masks) into unified consensus detections with improved accuracy and reliability. + +## How This Block Works + +This block fuses predictions from multiple detection models by requiring agreement (consensus) among models before accepting detections. The block: + +1. Takes detection predictions from multiple model sources (object detection, instance segmentation, or keypoint detection) as input +2. Matches detections from different models that overlap spatially by calculating Intersection over Union (IoU) between bounding boxes +3. Compares overlapping detections against an IoU threshold to determine if they represent the same object +4. Counts "votes" for each detection by finding matching detections from other models (subject to class-awareness if enabled) +5. Requires a minimum number of votes (`required_votes`) before accepting a detection as part of the consensus output +6. Aggregates properties of matching detections using configurable modes: + - **Confidence aggregation**: Combines confidence scores using average, max, or min + - **Coordinates aggregation**: Merges bounding boxes using average (mean coordinates), max (largest box), or min (smallest box) + - **Mask aggregation** (for instance segmentation): Combines masks using union, intersection, max, or min + - **Class selection**: Chooses class name based on majority vote (average), highest confidence (max), or lowest confidence (min) +7. Filters detections based on optional criteria (specific classes to consider, minimum confidence threshold) +8. Determines object presence by checking if the required number of objects (per class or total) are present in consensus results +9. Returns merged consensus detections, object presence indicators, and presence confidence scores + +The block enables class-aware or class-agnostic matching: when `class_aware` is true, only detections with matching class names are considered for voting; when false, any overlapping detections (regardless of class) contribute votes. The consensus mechanism helps reduce false positives (detections seen by only one model) and improves reliability by requiring multiple models to agree on object presence. Aggregation modes allow flexibility in how overlapping detections are combined, balancing between conservative (intersection, min) and inclusive (union, max) strategies. + +## Common Use Cases + +- **Multi-Model Ensemble**: Combine predictions from multiple specialized models (e.g., one optimized for people, another for vehicles) to improve overall detection accuracy, leveraging strengths of different models while filtering out detections that only one model sees +- **Reducing False Positives**: Require consensus from multiple models before accepting detections (e.g., require 2 out of 3 models to detect an object), reducing false positives by filtering out detections seen by only one model +- **Improving Detection Reliability**: Use majority voting to increase confidence in detections (e.g., merge overlapping detections from 3 models, keeping only those with 2+ votes), ensuring only high-confidence, multi-model-agreed detections are retained +- **Object Presence Detection**: Determine if specific objects are present based on consensus (e.g., check if at least 2 "person" detections exist across models, use aggregated confidence to determine presence), enabling robust object presence checking with configurable thresholds +- **Class-Specific Consensus**: Apply different consensus requirements per class (e.g., require 3 votes for "car" but only 2 for "person"), allowing stricter criteria for critical objects while being more lenient for common detections +- **Specialized Model Fusion**: Combine general-purpose and specialized models (e.g., general object detector + specialized license plate detector), creating a unified detection system that benefits from both broad coverage and specific expertise + +## Connecting to Other Blocks + +The consensus predictions from this block can be connected to: + +- **Multiple detection model blocks** (e.g., Object Detection Model, Instance Segmentation Model) to receive predictions from different models that are fused into consensus detections based on majority voting and spatial overlap matching +- **Visualization blocks** (e.g., Bounding Box Visualization, Polygon Visualization, Label Visualization) to display the merged consensus detections, showing unified results from multiple models with improved accuracy +- **Counting and analytics blocks** (e.g., Line Counter, Time in Zone, Velocity) to count or analyze consensus detections, providing more reliable metrics based on multi-model agreement +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload, Webhook Sink) to save or transmit consensus detection results, storing fused predictions that represent multi-model agreement +- **Flow control blocks** (e.g., Continue If) to conditionally trigger downstream processing based on `object_present` indicators or `presence_confidence` scores, enabling workflows that respond to multi-model consensus on object presence +- **Filtering blocks** (e.g., Detections Filter) to further refine consensus detections based on additional criteria, enabling multi-stage filtering after consensus fusion +""" + +SHORT_DESCRIPTION = ( + "Combine predictions from multiple detections models to make a " + "decision about object presence." +) + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Detections Consensus", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "fusion", + "ui_manifest": { + "section": "flow_control", + "icon": "fak fa-circles-overlap", + "blockPriority": 4, + }, + } + ) + type: Literal["roboflow_core/detections_consensus@v1", "DetectionsConsensus"] + predictions_batches: List[ + Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ), + ] = Field( + min_items=1, + description="List of references to detection predictions from multiple models. Each model's predictions must be made against the same input image. Predictions can be from object detection, instance segmentation, or keypoint detection models. The block matches overlapping detections across models and requires a minimum number of votes (required_votes) before accepting detections in the consensus output. Requires at least one prediction source. Supports batch processing.", + examples=[["$steps.a.predictions", "$steps.b.predictions"]], + validation_alias=AliasChoices("predictions_batches", "predictions"), + ) + required_votes: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + description="Minimum number of votes (matching detections from different models) required to accept a detection in the consensus output. Detections that receive fewer votes than this threshold are filtered out. For example, if set to 2, at least 2 models must detect an overlapping object (above IoU threshold) for it to appear in the consensus results. Higher values create stricter consensus requirements, reducing false positives but potentially missing detections seen by fewer models.", + examples=[2, "$inputs.required_votes"], + ) + class_aware: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="If true, only detections with matching class names from different models are considered as votes for the same object. If false, any overlapping detections (regardless of class) contribute votes. Class-aware mode is more conservative and ensures class consistency in consensus, while class-agnostic mode allows voting across different classes but may merge detections of different object types.", + examples=[True, "$inputs.class_aware"], + ) + iou_threshold: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = ( + Field( + default=0.3, + description="Intersection over Union (IoU) threshold for considering detections from different models as matching the same object. Detections with IoU above this threshold are considered overlapping and contribute votes to each other. Lower values (e.g., 0.2) are more lenient and match detections with less overlap, while higher values (e.g., 0.5) require stronger spatial overlap for matching. Typical values range from 0.2 to 0.5.", + examples=[0.3, "$inputs.iou_threshold"], + ) + ) + confidence: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( + default=0.0, + description="Confidence threshold applied to merged consensus detections. Only detections with aggregated confidence scores above this threshold are included in the output. Set to 0.0 to disable confidence filtering. Higher values filter out low-confidence consensus detections, improving output quality at the cost of potentially removing valid but lower-confidence detections.", + examples=[0.1, "$inputs.confidence"], + ) + classes_to_consider: Optional[ + Union[List[str], Selector(kind=[LIST_OF_VALUES_KIND])] + ] = Field( + default=None, + description="Optional list of class names to include in the consensus procedure. If provided, only detections of these classes are considered for voting and merging; all other classes are filtered out before consensus matching. Use this to focus consensus on specific object types while ignoring irrelevant detections. If None, all classes participate in consensus.", + examples=[["a", "b"], "$inputs.classes_to_consider"], + ) + required_objects: Optional[ + Union[ + PositiveInt, + Dict[str, PositiveInt], + Selector(kind=[INTEGER_KIND, DICTIONARY_KIND]), + ] + ] = Field( + default=None, + description="Optional minimum number of objects required to determine object presence. Can be an integer (total objects across all classes) or a dictionary mapping class names to per-class minimum counts. Used in conjunction with object_present output to determine if sufficient objects of each class are detected. For example, 3 means at least 3 total objects must be present, while {'person': 2, 'car': 1} requires at least 2 persons and 1 car. If None, object presence is determined solely by whether any consensus detections exist.", + examples=[3, {"a": 7, "b": 2}, "$inputs.required_objects"], + ) + presence_confidence_aggregation: AggregationMode = Field( + default=AggregationMode.MAX, + description="Aggregation mode for calculating presence confidence scores. Determines how confidence values are combined when computing object presence confidence: 'average' (mean confidence), 'max' (highest confidence), or 'min' (lowest confidence). This mode applies to the presence_confidence output which indicates confidence that required objects are present.", + examples=["max", "min"], + ) + detections_merge_confidence_aggregation: AggregationMode = Field( + default=AggregationMode.AVERAGE, + description="Aggregation mode for merging confidence scores of overlapping detections. 'average' computes mean confidence (majority vote approach), 'max' uses the highest confidence among matching detections, 'min' uses the lowest confidence. For class selection, 'average' represents majority vote (most common class), 'max' selects class from detection with highest confidence, 'min' selects class from detection with lowest confidence.", + examples=["min", "max"], + ) + detections_merge_coordinates_aggregation: AggregationMode = Field( + default=AggregationMode.AVERAGE, + description="Aggregation mode for merging bounding box coordinates of overlapping detections. 'average' computes mean coordinates from all matching boxes (balanced approach), 'max' takes the largest box (most inclusive), 'min' takes the smallest box (most conservative). This mode only applies to bounding boxes; mask aggregation uses detections_merge_mask_aggregation instead.", + examples=["min", "max"], + ) + detections_merge_mask_aggregation: MaskAggregationMode = Field( + default=MaskAggregationMode.UNION, + description="Aggregation mode for merging segmentation masks of overlapping detections. 'union' combines all masks into the largest possible area (most inclusive), 'intersection' takes only the overlapping region (most conservative), 'max' selects the largest mask, 'min' selects the smallest mask. This mode applies only to instance segmentation detections with masks; bounding box detections use detections_merge_coordinates_aggregation instead.", + examples=["union", "intersection"], + ) + raise_on_class_name_conflict: bool = Field( + default=False, + description="Controls how a class_id that maps to DIFFERENT class names across the merged prediction sources is handled. Tensor-native predictions carry a single class_id->name map per image, so two sources that reuse the same class_id for different classes (e.g. model A uses id 0 for 'car', model B uses id 0 for 'person') collide when their maps are unioned. When False (default), the later source's name wins and a warning is logged, matching the historical permissive behaviour. When True, a real collision (same class_id, different names) raises an error instead of silently mislabeling detections.", + examples=[False, True], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["predictions_batches"] + + @classmethod + @lru_cache(maxsize=None) + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + OutputDefinition( + name="object_present", kind=[BOOLEAN_KIND, DICTIONARY_KIND] + ), + OutputDefinition( + name="presence_confidence", + kind=[FLOAT_ZERO_TO_ONE_KIND, DICTIONARY_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class DetectionsConsensusBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + predictions_batches: List[Batch[TensorNativeDetections]], + required_votes: int, + class_aware: bool, + iou_threshold: float, + confidence: float, + classes_to_consider: Optional[List[str]], + required_objects: Optional[Union[int, Dict[str, int]]], + presence_confidence_aggregation: AggregationMode, + detections_merge_confidence_aggregation: AggregationMode, + detections_merge_coordinates_aggregation: AggregationMode, + detections_merge_mask_aggregation: MaskAggregationMode, + raise_on_class_name_conflict: bool = False, + ) -> BlockResult: + if len(predictions_batches) < 1: + raise ValueError( + f"Consensus step requires at least one source of predictions." + ) + results = [] + for detections_from_sources in zip(*predictions_batches): + ( + parent_id, + object_present, + presence_confidence, + consensus_detections, + ) = agree_on_consensus_for_all_detections_sources( + detections_from_sources=[ + _to_bbox_component(detections) + for detections in detections_from_sources + ], + required_votes=required_votes, + class_aware=class_aware, + iou_threshold=iou_threshold, + confidence=confidence, + classes_to_consider=classes_to_consider, + required_objects=required_objects, + presence_confidence_aggregation=presence_confidence_aggregation, + detections_merge_confidence_aggregation=detections_merge_confidence_aggregation, + detections_merge_coordinates_aggregation=detections_merge_coordinates_aggregation, + detections_merge_mask_aggregation=detections_merge_mask_aggregation, + raise_on_class_name_conflict=raise_on_class_name_conflict, + ) + results.append( + { + "predictions": consensus_detections, + "object_present": object_present, + "presence_confidence": presence_confidence, + } + ) + return results + + +def _to_bbox_component( + detections: Union[TensorNativeDetections, Tuple], +) -> TensorNativeDetections: + # Keypoint predictions arrive as a (KeyPoints, Detections) tuple; consensus + # only operates on bounding boxes so the bbox component is used directly. + if isinstance(detections, tuple): + _key_points, bbox_detections = detections + if bbox_detections is None: + raise ValueError( + "Keypoint prediction is missing the bounding-box component " + "required by the consensus step." + ) + return bbox_detections + return detections + + +def _resolve_class_name(detection: TensorNativeDetections, index: int = 0) -> str: + # class_name of detection i = image_metadata["class_names"][int(class_id[i])] + # (fallback f"class_{id}") - mirrors the per-detection class name that the + # numpy block read from sv.Detections.data["class_name"]. + class_id = int(detection.class_id[index]) + class_names = (detection.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + return class_names.get(class_id, f"class_{class_id}") + + +def _resolve_detection_id(detection: TensorNativeDetections, index: int = 0) -> str: + bboxes_metadata = detection.bboxes_metadata or [{}] + return bboxes_metadata[index][DETECTION_ID_KEY] + + +def _xyxy_as_numpy(detection: TensorNativeDetections) -> np.ndarray: + return detection.xyxy.detach().to("cpu").numpy().astype(float) + + +def _native_areas(detections: TensorNativeDetections) -> np.ndarray: + xyxy = _xyxy_as_numpy(detections) + return (xyxy[:, 2] - xyxy[:, 0]) * (xyxy[:, 3] - xyxy[:, 1]) + + +def _empty_native_detections(device: Optional[torch.device] = None) -> Detections: + # Empties must live on the same device as the predictions they may later be + # concatenated with (torch.cat rejects mixed devices even for 0-row inputs); + # callers pass the surrounding predictions' device when one is in scope and + # the documented native default is used otherwise. + if device is None: + device = WORKFLOWS_IMAGE_TENSOR_DEVICE + return Detections( + xyxy=torch.zeros((0, 4), dtype=torch.float32, device=device), + class_id=torch.zeros((0,), dtype=torch.long, device=device), + confidence=torch.zeros((0,), dtype=torch.float32, device=device), + image_metadata=None, + bboxes_metadata=None, + ) + + +def _handle_class_name_conflict( + class_id: int, + existing_name: str, + new_name: str, + raise_on_class_name_conflict: bool, +) -> None: + # A class_id that resolves to DIFFERENT names across sources cannot be + # represented by the single class_id->name map that tensor-native predictions + # carry. The numpy block stored a per-row class_name string and stayed correct; + # here the union has to pick one. Permissive (default): keep the later value + # and warn โ€” byte-identical output to the historical last-wins behaviour. + # Strict: raise instead of silently mislabeling every row with this class_id. + message = ( + f"Conflicting class names for class_id={class_id} while merging " + f"tensor-native detections: existing '{existing_name}' vs incoming " + f"'{new_name}'." + ) + if raise_on_class_name_conflict: + raise ValueError( + f"{message} Disable 'raise_on_class_name_conflict' to fall back to the " + f"permissive behaviour (keep the later value and continue)." + ) + logger.warning( + f"{message} Keeping the later value ('{new_name}'); enable " + f"'raise_on_class_name_conflict' to raise on conflict instead." + ) + + +def _concat_metadata( + detections_list: List[TensorNativeDetections], + raise_on_class_name_conflict: bool = False, +) -> Tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, Optional[dict], Optional[List[dict]] +]: + # Concatenate xyxy / class_id / confidence / bboxes_metadata of several + # predictions (masks handled separately). image_metadata is shared within one + # image so the first non-empty one is carried over as the base, BUT the + # class_names map is unioned across all sources: each consensus row supplies + # its own {class_id: class_name} pair, so taking only the first row's map + # would lose names for every other class in a multi-class consensus output. + image_metadata = None + merged_class_names: dict = {} + for d in detections_list: + if d.image_metadata is None: + continue + if image_metadata is None: + image_metadata = d.image_metadata + source_class_names = d.image_metadata.get(CLASS_NAMES_KEY) + if source_class_names: + for class_id_key, class_name in source_class_names.items(): + class_id_int = int(class_id_key) + existing_name = merged_class_names.get(class_id_int) + if existing_name is not None and existing_name != class_name: + _handle_class_name_conflict( + class_id=class_id_int, + existing_name=existing_name, + new_name=class_name, + raise_on_class_name_conflict=raise_on_class_name_conflict, + ) + merged_class_names[class_id_int] = class_name + if image_metadata is not None and merged_class_names: + image_metadata = dict(image_metadata) + image_metadata[CLASS_NAMES_KEY] = merged_class_names + bboxes_metadata: Optional[List[dict]] = None + if any(d.bboxes_metadata is not None for d in detections_list): + bboxes_metadata = [] + for d in detections_list: + if d.bboxes_metadata is not None: + bboxes_metadata.extend(d.bboxes_metadata) + else: + bboxes_metadata.extend({} for _ in range(len(d))) + xyxy = torch.cat([d.xyxy for d in detections_list], dim=0) + class_id = torch.cat([d.class_id for d in detections_list], dim=0) + confidence = torch.cat([d.confidence for d in detections_list], dim=0) + return xyxy, class_id, confidence, image_metadata, bboxes_metadata + + +def _merge_native_detections( + detections_list: List[TensorNativeDetections], + raise_on_class_name_conflict: bool = False, +) -> TensorNativeDetections: + # Native counterpart of sv.Detections.merge(list): concatenate several + # predictions into one native object. When any source carries a mask the + # result is an InstanceDetections whose dense masks are padded to a common + # (H, W) (mask-less / object-detection rows become zero masks, mirroring the + # numpy block's zero padding before sv.Detections.merge). + # SHARED-HELPER CANDIDATE: this inline native merge mirrors the + # _concatenate_detections helper in query_language and should be consolidated. + if not detections_list: + return _empty_native_detections() + xyxy, class_id, confidence, image_metadata, bboxes_metadata = _concat_metadata( + detections_list, + raise_on_class_name_conflict=raise_on_class_name_conflict, + ) + has_masks = any( + isinstance(d, InstanceDetections) and d.mask is not None + for d in detections_list + ) + if not has_masks: + return Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + mask_height, mask_width = 0, 0 + for d in detections_list: + if isinstance(d, InstanceDetections) and d.mask is not None and len(d) > 0: + sample = instance_mask_to_numpy(d, 0) + mask_height, mask_width = sample.shape + break + mask_rows = [] + for d in detections_list: + if isinstance(d, InstanceDetections) and d.mask is not None: + for i in range(len(d)): + mask_rows.append(instance_mask_to_numpy(d, i)) + else: + for _ in range(len(d)): + mask_rows.append(np.zeros((mask_height, mask_width), dtype=bool)) + # Keep the padded mask on the same device as the merged xyxy/class_id/ + # confidence (they come from the input predictions' device); a CPU mask here + # would build a mixed-device InstanceDetections that later fails to + # concatenate with CUDA/MPS model predictions. + mask = torch.from_numpy(np.stack(mask_rows, axis=0)).to(torch.bool).to(xyxy.device) + return InstanceDetections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def does_not_detect_objects_in_any_source( + detections_from_sources: List[TensorNativeDetections], +) -> bool: + return all(len(p) == 0 for p in detections_from_sources) + + +def get_parent_id_of_detections_from_sources( + detections_from_sources: List[TensorNativeDetections], +) -> str: + # Per-image parent_id lives in image_metadata (the numpy block read it from + # the per-detection sv.Detections.data[PARENT_ID_KEY]); a source contributes + # one parent_id per non-empty prediction. + encountered_parent_ids = set() + for detections in detections_from_sources: + image_metadata = detections.image_metadata or {} + if PARENT_ID_KEY in image_metadata and len(detections) > 0: + encountered_parent_ids.add(image_metadata[PARENT_ID_KEY]) + if len(encountered_parent_ids) != 1: + raise ValueError( + "Missmatch in predictions - while executing consensus step, " + "in equivalent batches, detections are assigned different parent " + "identifiers, whereas consensus can only be applied for predictions " + "made against the same input." + ) + return next(iter(encountered_parent_ids)) + + +def filter_predictions( + predictions: List[TensorNativeDetections], + classes_to_consider: Optional[List[str]], +) -> List[TensorNativeDetections]: + if not classes_to_consider: + return predictions + filtered = [] + for detections in predictions: + class_names = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) + if class_names is None: + continue + mask = [ + _resolve_class_name(detections, index) in classes_to_consider + for index in range(len(detections)) + ] + indices = [index for index, keep in enumerate(mask) if keep] + filtered.append(take_prediction_by_indices(detections, indices)) + return filtered + + +def get_detections_from_different_sources_with_max_overlap( + detection: TensorNativeDetections, + source: int, + detections_from_sources: List[TensorNativeDetections], + iou_threshold: float, + class_aware: bool, + detections_already_considered: Set[str], + detection_global_index: int, + iou_matrix: np.ndarray, + class_names: List[str], + detection_ids: List[Any], +) -> Dict[int, Tuple[TensorNativeDetections, float]]: + # IoU and class name are read from the once-computed host structures (see + # `_precompute_pair_data`) keyed by the detections' global index, rather than + # calling `calculate_iou` / `_resolve_class_name` per pair (each of which used + # to trigger device syncs). The greedy per-source max-overlap bookkeeping, + # iteration order, and tie behaviour are preserved exactly; single-row + # predictions are materialised only for the surviving winners, not per pair. + current_max_overlap: Dict[int, Tuple[int, float]] = {} + detection_class_name = class_names[detection_global_index] + for ( + other_source, + other_global_index, + other_local_index, + ) in enumerate_detection_indices( + detections_from_sources=detections_from_sources, + excluded_source_id=source, + ): + if detection_ids[other_global_index] in detections_already_considered: + continue + if class_aware and detection_class_name != class_names[other_global_index]: + continue + iou_value = float(iou_matrix[detection_global_index, other_global_index]) + if iou_value <= iou_threshold: + continue + if current_max_overlap.get(other_source) is None: + current_max_overlap[other_source] = (other_local_index, iou_value) + if current_max_overlap[other_source][1] < iou_value: + current_max_overlap[other_source] = (other_local_index, iou_value) + return { + other_source: ( + take_prediction_by_indices( + detections_from_sources[other_source], [local_index] + ), + iou_value, + ) + for other_source, (local_index, iou_value) in current_max_overlap.items() + } + + +def enumerate_detection_indices( + detections_from_sources: List[TensorNativeDetections], + excluded_source_id: Optional[int] = None, +) -> Generator[Tuple[int, int, int], None, None]: + # Yields (source_id, global_index, local_index) WITHOUT materialising a + # single-row prediction per detection - the O(N^2) matching loop only needs + # indices into the precomputed host structures. `global_index` is the + # absolute position in the source-ordered concatenation (excluded sources + # still advance the offset) so it indexes `iou_matrix` / `class_names` / + # `detection_ids` consistently regardless of `excluded_source_id`. + global_offset = 0 + for source_id, detections in enumerate(detections_from_sources): + n = len(detections) + if excluded_source_id == source_id: + global_offset += n + continue + for i in range(n): + yield source_id, global_offset + i, i + global_offset += n + + +def enumerate_detections( + detections_from_sources: List[TensorNativeDetections], + excluded_source_id: Optional[int] = None, +) -> Generator[Tuple[int, int, TensorNativeDetections], None, None]: + for source_id, global_index, local_index in enumerate_detection_indices( + detections_from_sources=detections_from_sources, + excluded_source_id=excluded_source_id, + ): + yield source_id, global_index, take_prediction_by_indices( + detections_from_sources[source_id], [local_index] + ) + + +def _precompute_pair_data( + detections_from_sources: List[TensorNativeDetections], +) -> Tuple[np.ndarray, List[str], List[Any]]: + """One host download per tensor field per source, then a single vectorised + IoU matrix over every detection. `global index` = position in the + source-ordered concatenation of all detections. + + - `iou_matrix[a, b]` is bit-identical to `calculate_iou(row_a, row_b)`: the + same `sv.box_iou_batch` arithmetic (float32 internally), computed per cell + independently, and the `nan_to_num` mirrors calculate_iou's + `if math.isnan(iou): iou = 0` guard. + - `class_names[g]` mirrors `_resolve_class_name` (each source's own + `image_metadata` class-names map, `f"class_{id}"` fallback). + - `detection_ids[g]` mirrors `_resolve_detection_id` (reads + `bboxes_metadata`; KeyError-compatible when a reached row lacks metadata). + """ + xyxy_arrays: List[np.ndarray] = [] + class_names: List[str] = [] + detection_ids: List[Any] = [] + for detections in detections_from_sources: + n = len(detections) + if n == 0: + continue + xyxy_arrays.append(_xyxy_as_numpy(detections)) + class_id = detections.class_id.detach().to("cpu").numpy() + class_names_map = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + class_names.extend( + class_names_map.get(int(c), f"class_{int(c)}") for c in class_id + ) + bboxes_metadata = detections.bboxes_metadata + for i in range(n): + row_metadata = bboxes_metadata[i] if bboxes_metadata is not None else {} + detection_ids.append(row_metadata[DETECTION_ID_KEY]) + if xyxy_arrays: + all_xyxy = np.concatenate(xyxy_arrays, axis=0) + iou_matrix = np.nan_to_num(sv.box_iou_batch(all_xyxy, all_xyxy)) + else: + iou_matrix = np.zeros((0, 0), dtype=np.float32) + return iou_matrix, class_names, detection_ids + + +def calculate_iou( + detection_a: TensorNativeDetections, detection_b: TensorNativeDetections +) -> float: + # sv.box_iou_batch is used here purely as the IoU algorithm over numpy boxes + # (no sv.Detections involved); the xyxy tensors are read out as float arrays. + iou = float( + sv.box_iou_batch(_xyxy_as_numpy(detection_a), _xyxy_as_numpy(detection_b))[0][0] + ) + if math.isnan(iou): + iou = 0 + return iou + + +def agree_on_consensus_for_all_detections_sources( + detections_from_sources: List[TensorNativeDetections], + required_votes: int, + class_aware: bool, + iou_threshold: float, + confidence: float, + classes_to_consider: Optional[List[str]], + required_objects: Optional[Union[int, Dict[str, int]]], + presence_confidence_aggregation: AggregationMode, + detections_merge_confidence_aggregation: AggregationMode, + detections_merge_coordinates_aggregation: AggregationMode, + detections_merge_mask_aggregation: MaskAggregationMode, + raise_on_class_name_conflict: bool = False, +) -> Tuple[str, bool, Dict[str, float], TensorNativeDetections]: + if does_not_detect_objects_in_any_source( + detections_from_sources=detections_from_sources + ): + return ( + "undefined", + False, + {}, + _empty_native_detections( + device=( + detections_from_sources[0].xyxy.device + if detections_from_sources + else None + ) + ), + ) + parent_id = get_parent_id_of_detections_from_sources( + detections_from_sources=detections_from_sources, + ) + detections_from_sources = filter_predictions( + predictions=detections_from_sources, + classes_to_consider=classes_to_consider, + ) + # One host download per tensor field per source + one vectorised IoU matrix, + # replacing the per-pair `calculate_iou` / `_resolve_class_name` device syncs. + iou_matrix, class_names, detection_ids = _precompute_pair_data( + detections_from_sources + ) + detections_already_considered = set() + consensus_detections = [] + for source_id, detection_global_index, detection in enumerate_detections( + detections_from_sources=detections_from_sources + ): + ( + consensus_detections_update, + detections_already_considered, + ) = get_consensus_for_single_detection( + detection=detection, + source_id=source_id, + detections_from_sources=detections_from_sources, + iou_threshold=iou_threshold, + class_aware=class_aware, + required_votes=required_votes, + confidence=confidence, + detections_merge_confidence_aggregation=detections_merge_confidence_aggregation, + detections_merge_coordinates_aggregation=detections_merge_coordinates_aggregation, + detections_merge_mask_aggregation=detections_merge_mask_aggregation, + detections_already_considered=detections_already_considered, + detection_global_index=detection_global_index, + iou_matrix=iou_matrix, + class_names=class_names, + detection_ids=detection_ids, + raise_on_class_name_conflict=raise_on_class_name_conflict, + ) + consensus_detections += consensus_detections_update + consensus_detections = _merge_native_detections( + consensus_detections, + raise_on_class_name_conflict=raise_on_class_name_conflict, + ) + ( + object_present, + presence_confidence, + ) = check_objects_presence_in_consensus_detections( + consensus_detections=consensus_detections, + aggregation_mode=presence_confidence_aggregation, + class_aware=class_aware, + required_objects=required_objects, + ) + return ( + parent_id, + object_present, + presence_confidence, + consensus_detections, + ) + + +def get_consensus_for_single_detection( + detection: TensorNativeDetections, + source_id: int, + detections_from_sources: List[TensorNativeDetections], + iou_threshold: float, + class_aware: bool, + required_votes: int, + confidence: float, + detections_merge_confidence_aggregation: AggregationMode, + detections_merge_coordinates_aggregation: AggregationMode, + detections_merge_mask_aggregation: MaskAggregationMode, + detections_already_considered: Set[str], + detection_global_index: int, + iou_matrix: np.ndarray, + class_names: List[str], + detection_ids: List[Any], + raise_on_class_name_conflict: bool = False, +) -> Tuple[List[TensorNativeDetections], Set[str]]: + if ( + len(detection) + and _resolve_detection_id(detection) in detections_already_considered + ): + return [], detections_already_considered + consensus_detections = [] + detections_with_max_overlap = ( + get_detections_from_different_sources_with_max_overlap( + detection=detection, + source=source_id, + detections_from_sources=detections_from_sources, + iou_threshold=iou_threshold, + class_aware=class_aware, + detections_already_considered=detections_already_considered, + detection_global_index=detection_global_index, + iou_matrix=iou_matrix, + class_names=class_names, + detection_ids=detection_ids, + ) + ) + + if len(detections_with_max_overlap) < (required_votes - 1): + # Returning empty list of detections + return consensus_detections, detections_already_considered + # Mask handling: object-detection rows have no mask while instance- + # segmentation rows do. The numpy block padded any missing mask with zeros so + # the group could stack; here the per-row masks are materialised to numpy and + # padded to a common shape before aggregation, never round-tripping sv. + detection_mask = _single_mask(detection) + overlap_masks = { + other_source: _single_mask(matched_value[0]) + for other_source, matched_value in detections_with_max_overlap.items() + } + if detection_mask is not None: + for other_source, matched_mask in overlap_masks.items(): + if matched_mask is None: + overlap_masks[other_source] = np.zeros(detection_mask.shape, dtype=bool) + else: + shape = None + for matched_mask in overlap_masks.values(): + if matched_mask is not None: + shape = matched_mask.shape + break + if shape: + for other_source, matched_mask in overlap_masks.items(): + if matched_mask is None: + overlap_masks[other_source] = np.zeros(shape, dtype=bool) + detection_mask = np.zeros(shape, dtype=bool) + group_detections = [detection] + [ + matched_value[0] for matched_value in detections_with_max_overlap.values() + ] + group_masks = ( + [detection_mask] + + [overlap_masks[other_source] for other_source in detections_with_max_overlap] + if detection_mask is not None + else None + ) + merged_detection = merge_detections( + detections=group_detections, + masks=group_masks, + confidence_aggregation_mode=detections_merge_confidence_aggregation, + boxes_aggregation_mode=detections_merge_coordinates_aggregation, + mask_aggregation_mode=detections_merge_mask_aggregation, + raise_on_class_name_conflict=raise_on_class_name_conflict, + ) + if float(merged_detection.confidence[0]) < confidence: + # Returning empty list of detections + return consensus_detections, detections_already_considered + consensus_detections.append(merged_detection) + detections_already_considered.add(_resolve_detection_id(detection)) + for matched_value in detections_with_max_overlap.values(): + detections_already_considered.add(_resolve_detection_id(matched_value[0])) + return consensus_detections, detections_already_considered + + +def _single_mask(detection: TensorNativeDetections) -> Optional[np.ndarray]: + # One detection's mask as a numpy (H, W) bool array, or None for object + # detection. `detection` is always a single-row prediction here. + if not isinstance(detection, InstanceDetections) or detection.mask is None: + return None + if len(detection) == 0: + return None + return instance_mask_to_numpy(detection, 0) + + +def check_objects_presence_in_consensus_detections( + consensus_detections: TensorNativeDetections, + class_aware: bool, + aggregation_mode: AggregationMode, + required_objects: Optional[Union[int, Dict[str, int]]], +) -> Tuple[bool, Dict[str, float]]: + if len(consensus_detections) == 0: + return False, {} + if required_objects is None: + required_objects = 0 + if isinstance(required_objects, dict) and not class_aware: + required_objects = sum(required_objects.values()) + if ( + isinstance(required_objects, int) + and len(consensus_detections) < required_objects + ): + return False, {} + if not class_aware: + aggregated_confidence = aggregate_field_values( + detections=consensus_detections, + field="confidence", + aggregation_mode=aggregation_mode, + ) + return True, {"any_object": aggregated_confidence} + class_names = [ + _resolve_class_name(consensus_detections, index) + for index in range(len(consensus_detections)) + ] + class2detections = {} + for class_name in set(class_names): + indices = [ + index for index, name in enumerate(class_names) if name == class_name + ] + class2detections[class_name] = take_prediction_by_indices( + consensus_detections, indices + ) + if isinstance(required_objects, dict): + for requested_class, required_objects_count in required_objects.items(): + if ( + requested_class not in class2detections + or len(class2detections[requested_class]) < required_objects_count + ): + return False, {} + class2confidence = { + class_name: aggregate_field_values( + detections=class_detections, + field="confidence", + aggregation_mode=aggregation_mode, + ) + for class_name, class_detections in class2detections.items() + } + return True, class2confidence + + +def merge_detections( + detections: List[TensorNativeDetections], + masks: Optional[List[np.ndarray]], + confidence_aggregation_mode: AggregationMode, + boxes_aggregation_mode: AggregationMode, + mask_aggregation_mode: MaskAggregationMode, + raise_on_class_name_conflict: bool = False, +) -> TensorNativeDetections: + # `detections` is the native group of overlapping detections (one single-row + # prediction per voting source). `masks` is the parallel list of numpy masks + # (already padded to a common shape) or None when no source carried a mask. + # The aggregation group only needs xyxy / class_id / confidence / metadata + # (masks come from the `masks` argument), so it is concatenated mask-free. + group_xyxy, group_class_id, group_confidence, group_metadata, _ = _concat_metadata( + detections, + raise_on_class_name_conflict=raise_on_class_name_conflict, + ) + group = Detections( + xyxy=group_xyxy, + class_id=group_class_id, + confidence=group_confidence, + image_metadata=group_metadata, + ) + class_name, class_id = AGGREGATION_MODE2CLASS_SELECTOR[confidence_aggregation_mode]( + group + ) + if masks is not None: + mask_stack = np.stack(masks, axis=0) + aggregated_mask = np.array( + [AGGREGATION_MODE2MASKS_AGGREGATOR[mask_aggregation_mode](mask_stack)] + ) + x1, y1, x2, y2 = sv.mask_to_xyxy(aggregated_mask)[0] + # Keep the merged output on the native inputs' device (the documented + # native-output contract); CPU allocations here would later fail to + # concatenate with CUDA/MPS model predictions. + output_mask = torch.from_numpy(aggregated_mask.astype(bool)).to( + group.xyxy.device + ) + else: + output_mask = None + x1, y1, x2, y2 = AGGREGATION_MODE2BOXES_AGGREGATOR[boxes_aggregation_mode]( + group + ) + # Per-image state is carried in image_metadata; the merged consensus row is a + # fresh detection so it gets a new detection_id in bboxes_metadata while the + # class_names map carries the chosen class. The source image_metadata supplies + # parent/root coordinates which are shared across the same input image. + source_metadata = group.image_metadata or {} + image_metadata = { + CLASS_NAMES_KEY: {int(class_id): class_name}, + PARENT_ID_KEY: source_metadata.get(PARENT_ID_KEY), + PREDICTION_TYPE_KEY: "object-detection", + PARENT_COORDINATES_KEY: source_metadata.get(PARENT_COORDINATES_KEY), + PARENT_DIMENSIONS_KEY: source_metadata.get(PARENT_DIMENSIONS_KEY), + ROOT_PARENT_ID_KEY: source_metadata.get(ROOT_PARENT_ID_KEY), + ROOT_PARENT_COORDINATES_KEY: source_metadata.get(ROOT_PARENT_COORDINATES_KEY), + ROOT_PARENT_DIMENSIONS_KEY: source_metadata.get(ROOT_PARENT_DIMENSIONS_KEY), + IMAGE_DIMENSIONS_KEY: source_metadata.get(IMAGE_DIMENSIONS_KEY), + } + if SCALING_RELATIVE_TO_PARENT_KEY in source_metadata: + image_metadata[SCALING_RELATIVE_TO_PARENT_KEY] = source_metadata[ + SCALING_RELATIVE_TO_PARENT_KEY + ] + else: + image_metadata[SCALING_RELATIVE_TO_PARENT_KEY] = 1.0 + if SCALING_RELATIVE_TO_ROOT_PARENT_KEY in source_metadata: + image_metadata[SCALING_RELATIVE_TO_ROOT_PARENT_KEY] = source_metadata[ + SCALING_RELATIVE_TO_ROOT_PARENT_KEY + ] + else: + image_metadata[SCALING_RELATIVE_TO_ROOT_PARENT_KEY] = 1.0 + device = group.xyxy.device + xyxy = torch.tensor([[x1, y1, x2, y2]], dtype=torch.float32, device=device) + class_id_tensor = torch.tensor([class_id], dtype=torch.long, device=device) + confidence_tensor = torch.tensor( + [ + aggregate_field_values( + detections=group, + field="confidence", + aggregation_mode=confidence_aggregation_mode, + ) + ], + dtype=torch.float32, + device=device, + ) + bboxes_metadata = [{DETECTION_ID_KEY: str(uuid4())}] + if output_mask is not None: + return InstanceDetections( + xyxy=xyxy, + class_id=class_id_tensor, + confidence=confidence_tensor, + mask=output_mask, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + return Detections( + xyxy=xyxy, + class_id=class_id_tensor, + confidence=confidence_tensor, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def get_majority_class(detections: TensorNativeDetections) -> Tuple[str, int]: + class_counts = Counter( + [ + (_resolve_class_name(detections, index), int(detections.class_id[index])) + for index in range(len(detections)) + ] + ) + return class_counts.most_common(1)[0][0] + + +def get_class_of_most_confident_detection( + detections: TensorNativeDetections, +) -> Tuple[str, int]: + confidences: List[float] = ( + detections.confidence.detach().to("cpu").numpy().astype(float).tolist() + ) + max_confidence_index = confidences.index(max(confidences)) + return ( + _resolve_class_name(detections, max_confidence_index), + int(detections.class_id[max_confidence_index]), + ) + + +def get_class_of_least_confident_detection( + detections: TensorNativeDetections, +) -> Tuple[str, int]: + confidences: List[float] = ( + detections.confidence.detach().to("cpu").numpy().astype(float).tolist() + ) + min_confidence_index = confidences.index(min(confidences)) + return ( + _resolve_class_name(detections, min_confidence_index), + int(detections.class_id[min_confidence_index]), + ) + + +AGGREGATION_MODE2CLASS_SELECTOR = { + AggregationMode.MAX: get_class_of_most_confident_detection, + AggregationMode.MIN: get_class_of_least_confident_detection, + AggregationMode.AVERAGE: get_majority_class, +} + + +def get_average_bounding_box( + detections: TensorNativeDetections, +) -> Tuple[int, int, int, int]: + if len(detections) == 0: + return (0.0, 0.0, 0.0, 0.0) + + avg_xyxy = np.mean(_xyxy_as_numpy(detections), axis=0) + return tuple(avg_xyxy) + + +def get_smallest_bounding_box( + detections: TensorNativeDetections, +) -> Tuple[int, int, int, int]: + areas: List[float] = _native_areas(detections).astype(float).tolist() + min_area = min(areas) + min_area_index = areas.index(min_area) + return tuple(_xyxy_as_numpy(detections)[min_area_index]) + + +def get_largest_bounding_box( + detections: TensorNativeDetections, +) -> Tuple[int, int, int, int]: + areas: List[float] = _native_areas(detections).astype(float).tolist() + max_area = max(areas) + max_area_index = areas.index(max_area) + return tuple(_xyxy_as_numpy(detections)[max_area_index]) + + +AGGREGATION_MODE2BOXES_AGGREGATOR = { + AggregationMode.MAX: get_largest_bounding_box, + AggregationMode.MIN: get_smallest_bounding_box, + AggregationMode.AVERAGE: get_average_bounding_box, +} + + +def get_intersection_mask(mask: np.ndarray) -> np.ndarray: + return np.all(mask, axis=0) + + +def get_union_mask(mask: np.ndarray) -> np.ndarray: + return np.any(mask, axis=0) + + +def get_smallest_mask(mask: np.ndarray) -> np.ndarray: + areas: List[float] = [float(m.sum()) for m in mask] + min_area = min(areas) + min_area_index = areas.index(min_area) + return mask[min_area_index] + + +def get_largest_mask(mask: np.ndarray) -> np.ndarray: + areas: List[float] = [float(m.sum()) for m in mask] + max_area = max(areas) + max_area_index = areas.index(max_area) + return mask[max_area_index] + + +AGGREGATION_MODE2MASKS_AGGREGATOR = { + MaskAggregationMode.MAX: get_largest_mask, + MaskAggregationMode.MIN: get_smallest_mask, + MaskAggregationMode.UNION: get_union_mask, + MaskAggregationMode.INTERSECTION: get_intersection_mask, +} + + +AGGREGATION_MODE2FIELD_AGGREGATOR = { + AggregationMode.MAX: max, + AggregationMode.MIN: min, + AggregationMode.AVERAGE: statistics.mean, +} + + +def aggregate_field_values( + detections: TensorNativeDetections, + field: str, + aggregation_mode: AggregationMode = AggregationMode.AVERAGE, +) -> float: + values = [] + if hasattr(detections, field): + values = getattr(detections, field) + if isinstance(values, torch.Tensor): + values = values.detach().to("cpu").numpy().astype(float).tolist() + elif isinstance(values, np.ndarray): + values = values.astype(float).tolist() + return AGGREGATION_MODE2FIELD_AGGREGATOR[aggregation_mode](values) diff --git a/inference/core/workflows/core_steps/fusion/detections_list_rollup/v1_tensor.py b/inference/core/workflows/core_steps/fusion/detections_list_rollup/v1_tensor.py new file mode 100644 index 0000000000..99d37dbbaf --- /dev/null +++ b/inference/core/workflows/core_steps/fusion/detections_list_rollup/v1_tensor.py @@ -0,0 +1,1302 @@ +import uuid +from typing import Any, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import torch +from pydantic import ConfigDict, Field + +from inference.core.logger import logger +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_key_points, + instance_mask_to_numpy, + split_key_point_prediction, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + FLOAT_ZERO_TO_ONE_KIND, + LIST_OF_VALUES_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +LONG_DESCRIPTION = """ +Rolls up dimensionality from children to parent detections + +Useful in scenarios like: +* rolling up results from a secondary model run on crops back to parent images +* rolling up OCR results for dynamically cropped images +""" + +SHORT_DESCRIPTION = ( + "Roll up multiple levels of dimensionality back to a single dimension." +) + +# Per-image metadata fields. In tensor-native predictions these live in +# ``image_metadata`` (one value per prediction) rather than being duplicated +# across every detection in ``bboxes_metadata``. They are read from / written to +# ``image_metadata`` while the rollup machinery keeps treating them as scalar +# "data" values for backward-compatible merge behaviour. +_PER_IMAGE_METADATA_KEYS = ( + PREDICTION_TYPE_KEY, + PARENT_ID_KEY, + INFERENCE_ID_KEY, + ROOT_PARENT_ID_KEY, + IMAGE_DIMENSIONS_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_COORDINATES_KEY, +) + +# Per-detection keypoint fields live in ``bboxes_metadata`` natively (same keys +# the serialiser reads back) โ€” mirrors the legacy ``sv.Detections.data`` keys. +_KEYPOINT_DATA_KEYS = ( + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS, +) + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Detections List Roll-Up", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "fusion", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-bring-front", + "blockPriority": 6, + }, + } + ) + type: Literal["roboflow_core/detections_list_rollup@v1", "DetectionsListRollUp"] + parent_detection: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + description="The parent detection the dimensionality inherits from.", + ) + + child_detections: Selector(kind=[LIST_OF_VALUES_KIND]) = Field( + description="A list of child detections resulting from higher dimensionality," + ' such as predictions made on dynamic crops. Use the "Dimension Collapse" to ' + " reduce the higher dimensionality result to one that can be used with this." + " Example: Prediction -> Dimension Collapse -> Detections List Roll-Up", + ) + + confidence_strategy: Union[ + Selector(kind=[LIST_OF_VALUES_KIND]), Literal["max", "mean", "min"] + ] = Field( + default="max", + title="Confidence Strategy", + description=( + "Strategy to use when merging confidence scores from child detections. " + "Options are 'max', 'mean', or 'min'." + ), + examples=["min", "mean", "max"], + ) + + overlap_threshold: Union[Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), float] = Field( + default=0.0, + title="Overlap Threshold", + description=( + "Minimum overlap ratio (IoU) to consider when merging overlapping " + "detections from child crops. " + "A value of 0.0 merges any overlapping detections, while higher values " + "require greater overlap to merge. Specify between 0.0 and 1.0. A value of 1.0 " + "only merges completely overlapping detections." + ), + examples=[0.0, 0.5], + ge=0.0, + le=1.0, + ) + + keypoint_merge_threshold: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + default=10.0, + title="Keypoint Merge Threshold", + description=( + "Keypoint distance (in pixels) to merge keypoint detections if the child detections contain keypoint data." + ), + examples=[0.0, 20.0], + ge=0.0, + ) + + raise_on_class_name_conflict: bool = Field( + default=False, + title="Raise On Class Name Conflict", + description=( + "Controls how a class_id that maps to DIFFERENT class names across the " + "child predictions is handled. Tensor-native predictions carry a single " + "class_id->name map per image, so two children that reuse the same " + "class_id for different classes (e.g. crop A uses id 0 for 'car', crop B " + "uses id 0 for 'person') collide when their maps are unioned. When False " + "(default), the later child's name wins and a warning is logged, matching " + "the historical permissive behaviour. When True, a real collision (same " + "class_id, different names) raises an error instead of silently " + "mislabeling rolled-up detections." + ), + examples=[False, True], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="rolled_up_detections", + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ], + ), + OutputDefinition( + name="crop_zones", + kind=[LIST_OF_VALUES_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class DetectionsListRollUpBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + parent_detection: Any, + child_detections: Any, + confidence_strategy: str = "max", + overlap_threshold: float = 0.0, + keypoint_merge_threshold: float = 10.0, + raise_on_class_name_conflict: bool = False, + ) -> BlockResult: + + detections, zones = merge_crop_predictions( + parent_detection, + child_detections, + confidence_strategy, + overlap_threshold, + keypoint_merge_threshold, + raise_on_class_name_conflict=raise_on_class_name_conflict, + ) + + return {"rolled_up_detections": detections, "crop_zones": zones} + + +def _native_box_metadata(prediction) -> List[dict]: + """Per-detection metadata list for a tensor-native bbox prediction, always + of length equal to the number of detections (``{}`` where absent).""" + number_of_detections = int(prediction.xyxy.shape[0]) + bboxes_metadata = prediction.bboxes_metadata + if bboxes_metadata is None: + return [{} for _ in range(number_of_detections)] + return [ + dict(box_metadata) if box_metadata is not None else {} + for box_metadata in bboxes_metadata + ] + + +def _handle_class_name_conflict( + class_id: int, + existing_name: str, + new_name: str, + raise_on_class_name_conflict: bool, +) -> None: + # A class_id that resolves to DIFFERENT names across child predictions cannot + # be represented by the single class_id->name map that tensor-native + # predictions carry. The numpy block stored a per-row class_name string and + # stayed correct; here the union has to pick one. Permissive (default): keep + # the later value and warn โ€” byte-identical output to the historical last-wins + # behaviour. Strict: raise instead of silently mislabeling every row with this + # class_id. + message = ( + f"Conflicting class names for class_id={class_id} while rolling up " + f"tensor-native detections: existing '{existing_name}' vs incoming " + f"'{new_name}'." + ) + if raise_on_class_name_conflict: + raise ValueError( + f"{message} Disable 'raise_on_class_name_conflict' to fall back to the " + f"permissive behaviour (keep the later value and continue)." + ) + logger.warning( + f"{message} Keeping the later value ('{new_name}'); enable " + f"'raise_on_class_name_conflict' to raise on conflict instead." + ) + + +def _merge_keypoint_detections( + preds: List[dict], confidence_strategy: str, keypoint_threshold: float +) -> List[dict]: + """ + Merge keypoint detections based on keypoint proximity. + + Args: + preds: List of prediction dicts with 'bbox', 'confidence', 'class_id', 'keypoint_data' + confidence_strategy: How to combine confidences ('max', 'mean', 'min') + keypoint_threshold: Maximum average keypoint distance (in pixels) to merge detections + + Returns: + List of merged prediction dicts + """ + if not preds: + return [] + + # Filter predictions that have keypoint data + preds_with_keypoints = [ + p + for p in preds + if p.get("keypoint_data") and "keypoints_xy" in p["keypoint_data"] + ] + preds_without_keypoints = [ + p + for p in preds + if not (p.get("keypoint_data") and "keypoints_xy" in p["keypoint_data"]) + ] + + if not preds_with_keypoints: + return preds + + merged = [] + used = set() + + for i, pred1 in enumerate(preds_with_keypoints): + if i in used: + continue + + # Start a new merged group with this prediction + group = [pred1] + used.add(i) + + kp1 = np.array(pred1["keypoint_data"]["keypoints_xy"]) + + # Find all predictions that should merge with this one + for j, pred2 in enumerate(preds_with_keypoints[i + 1 :], start=i + 1): + if j in used: + continue + + kp2 = np.array(pred2["keypoint_data"]["keypoints_xy"]) + + # Calculate average distance between corresponding keypoints + if len(kp1) == len(kp2): + distances = np.linalg.norm(kp1 - kp2, axis=1) + avg_distance = np.mean(distances) + + if avg_distance < keypoint_threshold: + group.append(pred2) + used.add(j) + + # Merge the group + if len(group) == 1: + merged.append(group[0]) + else: + # Merge multiple predictions + if confidence_strategy == "max": + best_idx = np.argmax([p["confidence"] for p in group]) + confidence = group[best_idx]["confidence"] + elif confidence_strategy == "mean": + confidence = np.mean([p["confidence"] for p in group]) + else: # 'min' + confidence = np.min([p["confidence"] for p in group]) + + # Average keypoint coordinates + all_kp_xy = [np.array(p["keypoint_data"]["keypoints_xy"]) for p in group] + merged_kp_xy = np.mean(all_kp_xy, axis=0).tolist() + + # Average keypoint confidences if available + merged_kp_data = { + "keypoints_xy": merged_kp_xy, + "keypoints_class_name": group[0]["keypoint_data"].get( + "keypoints_class_name" + ), + "keypoints_class_id": group[0]["keypoint_data"].get( + "keypoints_class_id" + ), + } + + if "keypoints_confidence" in group[0]["keypoint_data"]: + all_kp_conf = [ + np.array(p["keypoint_data"]["keypoints_confidence"]) for p in group + ] + merged_kp_conf = np.mean(all_kp_conf, axis=0).tolist() + merged_kp_data["keypoints_confidence"] = merged_kp_conf + + # Average bbox coordinates + all_bboxes = np.array([p["bbox"] for p in group]) + merged_bbox = np.mean(all_bboxes, axis=0) + + merged.append( + { + "bbox": merged_bbox, + "confidence": confidence, + "class_id": group[0]["class_id"], + "mask": None, + "keypoint_data": merged_kp_data, + "detection_data": group[0].get( + "detection_data", {} + ), # Preserve first detection's metadata + } + ) + + # Add back predictions without keypoints + merged.extend(preds_without_keypoints) + + return merged + + +def merge_crop_predictions( + parent_prediction, + child_predictions: List, + confidence_strategy: str = "max", + overlap_threshold: float = 0.0, + keypoint_merge_threshold: float = 10.0, + raise_on_class_name_conflict: bool = False, +) -> Tuple: + """ + Merge predictions from multiple crops back to parent image coordinates. + + Args: + parent_prediction: Supervision Detections object that defines the crop locations. + Each detection in this prediction represents one crop region. + child_predictions: List of Supervision Detections objects from crops. + Order matches the detection order in parent_prediction. + confidence_strategy: How to handle confidence when merging overlaps. + Options: "max", "mean", "min" + overlap_threshold: Minimum IoU/overlap ratio to merge detections (0.0 to 1.0). + - 0.0: Only merge if detections touch or overlap at all (default) + - >0.0: Only merge if overlap ratio exceeds this threshold + - 1.0: Only merge completely overlapping detections + keypoint_merge_threshold: Maximum distance in pixels to merge keypoints (default: 10). + For keypoint detections, merges detections if their average + keypoint distance is below this threshold. + + Returns: + Tuple of (detections, crop_zones): + - detections: Detections object with merged predictions in parent image coordinates. + Works for both instance segmentation (with masks) and object detection (without masks). + - crop_zones: List of lists of (x, y) tuples. Each inner list defines the rectangular + zone boundary of a crop in parent image coordinates as 4 corner points. + """ + # Keypoint predictions arrive as a (KeyPoints, Detections) tuple; the rollup + # only needs the bbox component (parent crop boxes / child boxes + masks + + # per-detection metadata). Keypoints themselves are carried in the child + # bboxes_metadata (keypoints_* keys), exactly like the legacy data dict. + _parent_key_points, parent_prediction = split_key_point_prediction( + parent_prediction + ) + child_predictions = [ + split_key_point_prediction(child_pred)[1] for child_pred in child_predictions + ] + + if len(parent_prediction) != len(child_predictions): + raise ValueError( + f"Number of detections in parent_prediction ({len(parent_prediction)}) " + f"must match number of child predictions ({len(child_predictions)})" + ) + + # Extract parent image shape from parent prediction's image_metadata. + # root_parent_dimensions is a single per-image value (height, width) for the + # tensor-native prediction (numpy stored it once per detection). + parent_image_metadata = parent_prediction.image_metadata or {} + root_parent_dims = parent_image_metadata.get(ROOT_PARENT_DIMENSIONS_KEY) + + if root_parent_dims is None or len(root_parent_dims) == 0: + raise ValueError( + "parent_prediction must have 'root_parent_dimensions' in its data attribute" + ) + + # The per-image value is already the (height, width) for the parent image. + parent_image_shape = tuple(root_parent_dims) + + # Pre-read per-detection metadata + per-image metadata for every child once. + child_box_metadata = [ + _native_box_metadata(child_pred) for child_pred in child_predictions + ] + child_image_metadata = [ + (child_pred.image_metadata or {}) for child_pred in child_predictions + ] + parent_xyxy = parent_prediction.xyxy.detach().to("cpu").numpy() + + # Merge the class_id -> name maps across all children so the rolled-up + # native prediction can carry a single image_metadata["class_names"] map + # (numpy stored class_name per detection instead). Keys are normalized to + # python int so map lookups never miss on a numpy-scalar class_id. + merged_class_names: dict = {} + for image_metadata in child_image_metadata: + child_class_names = image_metadata.get(CLASS_NAMES_KEY) + if child_class_names: + for class_id_key, class_name in child_class_names.items(): + class_id_int = int(class_id_key) + existing_name = merged_class_names.get(class_id_int) + if existing_name is not None and existing_name != class_name: + _handle_class_name_conflict( + class_id=class_id_int, + existing_name=existing_name, + new_name=class_name, + raise_on_class_name_conflict=raise_on_class_name_conflict, + ) + merged_class_names[class_id_int] = class_name + + # Build crop zones list - one zone per crop/child prediction + crop_zones = [] + for i in range(len(parent_prediction)): + crop_bbox = parent_xyxy[i] # [x_min, y_min, x_max, y_max] + x_min, y_min, x_max, y_max = ( + crop_bbox[0], + crop_bbox[1], + crop_bbox[2], + crop_bbox[3], + ) + + # Create zone as list of 4 corner points: top-left, top-right, bottom-right, bottom-left + zone = [ + (float(x_min), float(y_min)), # top-left + (float(x_max), float(y_min)), # top-right + (float(x_max), float(y_max)), # bottom-right + (float(x_min), float(y_max)), # bottom-left + ] + crop_zones.append(zone) + + # Check if we have instance segmentation (with masks) or object detection (without masks) + has_masks = False + is_keypoint_detection = False + for child_pred in child_predictions: + if ( + isinstance(child_pred, InstanceDetections) + and child_pred.mask is not None + and len(child_pred) > 0 + ): + has_masks = True + break + + for image_metadata in child_image_metadata: + # Check for keypoint detection. Native image_metadata stores + # prediction_type as a scalar string (build_native_image_metadata), + # so no np.ndarray legacy form is possible here. + if PREDICTION_TYPE_KEY in image_metadata: + if image_metadata[PREDICTION_TYPE_KEY] == "keypoint-detection": + is_keypoint_detection = True + break + + # Group predictions by class + class_predictions = {} + + # Iterate through each crop region and its corresponding child predictions + for i, child_pred in enumerate(child_predictions): + box_metadata = child_box_metadata[i] + # Get crop location from parent prediction + crop_bbox = parent_xyxy[i] # [x_min, y_min, x_max, y_max] + x_min, y_min = int(crop_bbox[0]), int(crop_bbox[1]) + + child_class_ids = child_pred.class_id.detach().to("cpu").numpy() + child_confidences = child_pred.confidence.detach().to("cpu").numpy() + child_xyxy = child_pred.xyxy.detach().to("cpu").numpy() + child_has_masks = ( + isinstance(child_pred, InstanceDetections) and child_pred.mask is not None + ) + + # Process each detection in the child prediction + for j in range(len(child_pred)): + detection_metadata = box_metadata[j] + class_id = child_class_ids[j] + confidence = child_confidences[j] + + # Prepare keypoint data if present + keypoint_data = {} + if ( + is_keypoint_detection + and KEYPOINTS_XY_KEY_IN_SV_DETECTIONS in detection_metadata + ): + # Transform keypoint coordinates from crop to parent space + keypoints_xy = detection_metadata[ + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS + ] # Shape: (num_keypoints, 2) + + # Vectorised offset โ€” avoids a per-keypoint Python loop. + # copy=True prevents mutating the source array; reshape handles + # edge cases where keypoints_xy arrives as a flat/empty array. + kp_array = np.array(keypoints_xy, dtype=np.float64, copy=True) + if kp_array.size == 0: + keypoint_data["keypoints_xy"] = [] + else: + kp_array = kp_array.reshape(-1, 2) + kp_array = kp_array + np.array([x_min, y_min], dtype=np.float64) + keypoint_data["keypoints_xy"] = kp_array.tolist() + + # Copy other keypoint data + if KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS in detection_metadata: + keypoint_data["keypoints_class_name"] = detection_metadata[ + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS + ] + if KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS in detection_metadata: + keypoint_data["keypoints_class_id"] = detection_metadata[ + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS + ] + if KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS in detection_metadata: + keypoint_data["keypoints_confidence"] = detection_metadata[ + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS + ] + + # Collect per-detection data fields to preserve individual detection metadata + # This is crucial for preserving class_name and other fields when multiple + # detections have the same class_id but different values + detection_data = {} + for key in detection_metadata.keys(): + if key not in [ + "detection_id", + "parent_id", + "inference_id", + "keypoints_xy", + "keypoints_class_name", + "keypoints_class_id", + "keypoints_confidence", + ]: + detection_data[key] = detection_metadata[key] + + if has_masks and child_has_masks: + # Instance segmentation - transform mask + mask = instance_mask_to_numpy(child_pred, j) + transformed_mask = _transform_mask_to_parent( + mask, x_min, y_min, parent_image_shape + ) + + # Also store the transformed bbox for cheap pre-filtering + raw_bbox = child_xyxy[j] + transformed_bbox = np.array( + [ + raw_bbox[0] + x_min, + raw_bbox[1] + y_min, + raw_bbox[2] + x_min, + raw_bbox[3] + y_min, + ] + ) + + # Store prediction with transformed mask + if class_id not in class_predictions: + class_predictions[class_id] = [] + + class_predictions[class_id].append( + { + "mask": transformed_mask, + "confidence": confidence, + "class_id": class_id, + "bbox": transformed_bbox, + "keypoint_data": keypoint_data, + "detection_data": detection_data, # Store per-detection metadata + } + ) + else: + # Object detection - transform bounding box + bbox = child_xyxy[j] # [x_min, y_min, x_max, y_max] + transformed_bbox = np.array( + [bbox[0] + x_min, bbox[1] + y_min, bbox[2] + x_min, bbox[3] + y_min] + ) + + # Store prediction with transformed bbox + if class_id not in class_predictions: + class_predictions[class_id] = [] + + class_predictions[class_id].append( + { + "bbox": transformed_bbox, + "confidence": confidence, + "class_id": class_id, + "mask": None, + "keypoint_data": keypoint_data, + "detection_data": detection_data, # Store per-detection metadata + } + ) + + # Merge overlapping predictions for each class + merged_masks = [] + merged_bboxes = [] + merged_confidences = [] + merged_class_ids = [] + + # Collect all data field names from child predictions. Tensor-native + # per-detection keys live in bboxes_metadata; the per-image fields below + # (prediction_type, parent/root ids+coords+dims, inference_id) are added + # explicitly so the rollup re-creates them on the output exactly as the + # numpy block did from sv.Detections.data. + all_data_keys = set() + for box_metadata in child_box_metadata: + for detection_metadata in box_metadata: + all_data_keys.update(detection_metadata.keys()) + for image_metadata in child_image_metadata: + for key in _PER_IMAGE_METADATA_KEYS: + if key in image_metadata: + all_data_keys.add(key) + + # Initialize lists for each data field + merged_data = { + key: [] + for key in all_data_keys + if key + not in [ + "keypoints_xy", + "keypoints_class_name", + "keypoints_class_id", + "keypoints_confidence", + ] + } + + # Collect keypoint data separately + all_keypoints_data = { + "keypoints_xy": [], + "keypoints_class_name": [], + "keypoints_class_id": [], + "keypoints_confidence": [], + } + + # Build mapping from class_id to typical data values + class_id_to_data = {} + for i, child_pred in enumerate(child_predictions): + box_metadata = child_box_metadata[i] + image_metadata = child_image_metadata[i] + child_class_ids = child_pred.class_id.detach().to("cpu").numpy() + for index in range(len(child_pred)): + detection_metadata = box_metadata[index] + class_id = child_class_ids[index] + if class_id not in class_id_to_data: + class_id_to_data[class_id] = {} + # Store sample values for this class_id (except ID fields and keypoint fields) + for key in detection_metadata.keys(): + if key not in [ + "detection_id", + "parent_id", + "inference_id", + "keypoints_xy", + "keypoints_class_name", + "keypoints_class_id", + "keypoints_confidence", + ]: + class_id_to_data[class_id][key] = detection_metadata[key] + # Per-image fields are shared across all detections of this child + for key in _PER_IMAGE_METADATA_KEYS: + if key in image_metadata and key not in [ + "parent_id", + "inference_id", + ]: + class_id_to_data[class_id][key] = image_metadata[key] + + # Get a sample inference_id and parent_id from the first child prediction if available + sample_inference_id = None + sample_parent_id = None + if len(child_predictions) > 0 and len(child_predictions[0]) > 0: + first_image_metadata = child_image_metadata[0] + if INFERENCE_ID_KEY in first_image_metadata: + sample_inference_id = first_image_metadata[INFERENCE_ID_KEY] + if PARENT_ID_KEY in first_image_metadata: + sample_parent_id = first_image_metadata[PARENT_ID_KEY] + + for class_id, preds in class_predictions.items(): + if is_keypoint_detection: + # For keypoint detection, merge based on keypoint proximity + merged_preds = _merge_keypoint_detections( + preds, confidence_strategy, keypoint_merge_threshold + ) + elif has_masks: + merged_preds = _merge_overlapping_masks( + preds, confidence_strategy, overlap_threshold + ) + else: + merged_preds = _merge_overlapping_bboxes( + preds, confidence_strategy, overlap_threshold + ) + + for pred in merged_preds: + if has_masks: + merged_masks.append(pred["mask"]) + else: + # For non-mask detections, collect bboxes + if "bbox" in pred and pred["bbox"] is not None: + merged_bboxes.append(pred["bbox"]) + merged_confidences.append(pred["confidence"]) + merged_class_ids.append(pred["class_id"]) + + # Collect keypoint data if present + if "keypoint_data" in pred and pred["keypoint_data"]: + kp_data = pred["keypoint_data"] + all_keypoints_data["keypoints_xy"].append(kp_data.get("keypoints_xy")) + all_keypoints_data["keypoints_class_name"].append( + kp_data.get("keypoints_class_name") + ) + all_keypoints_data["keypoints_class_id"].append( + kp_data.get("keypoints_class_id") + ) + all_keypoints_data["keypoints_confidence"].append( + kp_data.get("keypoints_confidence") + ) + + # Add data fields for this detection + for key in all_data_keys: + # Skip keypoint fields as they're handled separately + if key in [ + "keypoints_xy", + "keypoints_class_name", + "keypoints_class_id", + "keypoints_confidence", + ]: + continue + + if key == "detection_id": + # Generate new UUID for merged detection + merged_data[key].append(str(uuid.uuid4())) + elif key == "parent_id": + # Use sample parent_id or generate new one + merged_data[key].append( + sample_parent_id if sample_parent_id else str(uuid.uuid4()) + ) + elif key == "inference_id": + # Use the same inference_id as inputs (they're from same inference batch) + merged_data[key].append( + sample_inference_id + if sample_inference_id + else str(uuid.uuid4()) + ) + elif key == "root_parent_dimensions": + # Add the parent image shape as a list [height, width] + merged_data[key].append(list(parent_image_shape)) + elif key == "parent_dimensions": + # Parent dimensions should be same as root_parent_dimensions for merged results + merged_data[key].append(list(parent_image_shape)) + elif key == "image_dimensions": + # Image dimensions for this detection + merged_data[key].append(list(parent_image_shape)) + elif key == "root_parent_coordinates": + # The merged geometry lives in the PARENT's frame, so the + # root offset is the parent prediction's root offset - the + # child entries carry their crop origins, which no longer + # apply after rollup (sampling them made the output + # constructor re-shift already-rolled-up predictions and + # crash re-anchoring parent-canvas masks). + merged_data[key].append( + list( + parent_image_metadata.get(ROOT_PARENT_COORDINATES_KEY) + or [0, 0] + ) + ) + elif key == "parent_coordinates": + # Same frame argument as root_parent_coordinates above. + merged_data[key].append( + list( + parent_image_metadata.get(PARENT_COORDINATES_KEY) or [0, 0] + ) + ) + elif key == "root_parent_id": + # Root parent ID + if ( + pred["class_id"] in class_id_to_data + and key in class_id_to_data[pred["class_id"]] + ): + merged_data[key].append(class_id_to_data[pred["class_id"]][key]) + else: + merged_data[key].append("image") + elif key == "prediction_type": + # Prediction type should be 'instance-segmentation' + merged_data[key].append("instance-segmentation") + else: + # For other fields like class_name, check pred dict first (per-detection data) + # then fall back to class_id_to_data (class-level defaults) + if key in pred.get("detection_data", {}): + merged_data[key].append(pred["detection_data"][key]) + elif ( + pred["class_id"] in class_id_to_data + and key in class_id_to_data[pred["class_id"]] + ): + merged_data[key].append(class_id_to_data[pred["class_id"]][key]) + else: + merged_data[key].append(None) + + if not merged_confidences: + # Return empty detections if no detections + return _empty_native_detections(has_masks, merged_class_names), crop_zones + + # Convert to numpy arrays + merged_confidences_array = np.array(merged_confidences, dtype=np.float32) + merged_class_ids_array = np.array(merged_class_ids, dtype=int) + + if has_masks: + # Instance segmentation - stack masks and compute bounding boxes + merged_masks_array = np.stack(merged_masks, axis=0) + + # Compute bounding boxes from masks + xyxy = [] + for mask in merged_masks_array: + rows, cols = np.where(mask) + if len(rows) > 0: + x_min, x_max = cols.min(), cols.max() + y_min, y_max = rows.min(), rows.max() + xyxy.append([x_min, y_min, x_max + 1, y_max + 1]) + else: + xyxy.append([0, 0, 0, 0]) + + xyxy_array = np.array(xyxy, dtype=np.float32) + else: + # Object detection - use bounding boxes directly + if merged_bboxes: + xyxy_array = np.array(merged_bboxes, dtype=np.float32) + else: + # Shouldn't happen, but handle edge case + xyxy_array = np.zeros((len(merged_confidences), 4), dtype=np.float32) + + # Route the merged data fields into native per-image (image_metadata) and + # per-detection (bboxes_metadata) state instead of a flat sv.Detections.data + # dict. Per-image fields are written once (taken from the first merged value, + # which is shared across all rolled-up detections); per-detection fields are + # written per row. class_name is not stored per-detection natively โ€” names + # are resolved from class_id via image_metadata["class_names"]. + number_of_detections = len(merged_confidences) + image_metadata: dict = {} + # Always carry a class_names map covering every merged class_id so the + # serializer's per-row class_id lookup never fails (e.g. the OCR / map-less + # rollup path, where children never declared a class_names map). Names come + # from the unioned child maps; any class_id without an entry falls back to + # f"class_{id}". Keys are normalized to python int to match map lookups. + output_class_names: dict = {} + for class_id_value in merged_class_ids: + class_id_int = int(class_id_value) + output_class_names[class_id_int] = merged_class_names.get( + class_id_int, f"class_{class_id_int}" + ) + image_metadata[CLASS_NAMES_KEY] = output_class_names + bboxes_metadata: List[dict] = [{} for _ in range(number_of_detections)] + + for key, values in merged_data.items(): + if key in _PER_IMAGE_METADATA_KEYS: + # Per-image field: store once in image_metadata (values are identical + # across rows by construction above). + if values: + image_metadata[key] = _coerce_metadata_value(key, values[0]) + else: + # Per-detection field: store per row in bboxes_metadata. + for index in range(number_of_detections): + bboxes_metadata[index][key] = _coerce_metadata_value(key, values[index]) + + # Add keypoint data if it exists (per-detection, into bboxes_metadata). + if is_keypoint_detection: + for key in _KEYPOINT_DATA_KEYS: + values = all_keypoints_data[key] + if values: + for index in range(number_of_detections): + bboxes_metadata[index][key] = values[index] + + if has_masks: + result = InstanceDetections( + xyxy=torch.as_tensor(xyxy_array, dtype=torch.float32).reshape(-1, 4), + mask=torch.from_numpy(merged_masks_array).to(torch.bool), + confidence=torch.as_tensor(merged_confidences_array, dtype=torch.float32), + class_id=torch.as_tensor(merged_class_ids_array, dtype=torch.long), + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + else: + result = Detections( + xyxy=torch.as_tensor(xyxy_array, dtype=torch.float32).reshape(-1, 4), + confidence=torch.as_tensor(merged_confidences_array, dtype=torch.float32), + class_id=torch.as_tensor(merged_class_ids_array, dtype=torch.long), + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + if is_keypoint_detection: + # The keypoint kind's runtime payload is a ``(KeyPoints, Detections)`` tuple + # (the same shape the keypoint model producer emits), so consumers like + # keypoint_visualization can call ``KeyPoints.to_supervision()``. Rebuild the + # native KeyPoints from the merged per-instance keypoint lists; ``result`` is + # the bounding-box component. + key_points = build_native_key_points( + per_instance_xy=all_keypoints_data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS], + per_instance_confidence=all_keypoints_data[ + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS + ], + object_class_ids=merged_class_ids, + image_metadata=image_metadata, + ) + return (key_points, result), crop_zones + + return result, crop_zones + + +def _coerce_metadata_value(key: str, value: Any) -> Any: + """Coerce a merged metadata value into the dtype the numpy block used when + writing into sv.Detections.data, but as a plain python scalar/list (native + per-detection / per-image state is python, not numpy arrays).""" + if value is None: + return value + if key in [ + "prediction_type", + "detection_id", + "parent_id", + "inference_id", + "root_parent_id", + ]: + # String fields. class_name is intentionally omitted: names are never + # stored per-detection natively (resolved from class_id via the + # image_metadata["class_names"] map). + return str(value) + if key in [ + "root_parent_dimensions", + "parent_dimensions", + "image_dimensions", + "root_parent_coordinates", + "parent_coordinates", + ]: + # Array/coordinate fields - integers + return np.array(value, dtype=int).tolist() + return value + + +def _empty_native_detections( + has_masks: bool, class_names: dict +) -> Union[Detections, InstanceDetections]: + """Build an empty tensor-native prediction (mirrors sv.Detections.empty()).""" + image_metadata = {CLASS_NAMES_KEY: class_names} if class_names else {} + if has_masks: + return InstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, 0, 0), dtype=torch.bool), + image_metadata=image_metadata, + bboxes_metadata=None, + ) + return Detections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + image_metadata=image_metadata, + bboxes_metadata=None, + ) + + +def _transform_mask_to_parent( + mask: np.ndarray, x_offset: int, y_offset: int, parent_shape: Tuple[int, int] +) -> np.ndarray: + """ + Transform a mask from crop coordinates to parent image coordinates. + + Args: + mask: Boolean mask array from crop (H, W) + x_offset: X offset of crop in parent image + y_offset: Y offset of crop in parent image + parent_shape: (height, width) of parent image + + Returns: + Boolean mask in parent image coordinates + """ + parent_mask = np.zeros(parent_shape, dtype=bool) + + crop_h, crop_w = mask.shape + parent_h, parent_w = parent_shape + + # Calculate valid region to paste (handle edge cases) + y_start = max(0, y_offset) + y_end = min(parent_h, y_offset + crop_h) + x_start = max(0, x_offset) + x_end = min(parent_w, x_offset + crop_w) + + # Calculate corresponding crop region + crop_y_start = y_start - y_offset + crop_y_end = crop_y_start + (y_end - y_start) + crop_x_start = x_start - x_offset + crop_x_end = crop_x_start + (x_end - x_start) + + # Paste the mask + parent_mask[y_start:y_end, x_start:x_end] = mask[ + crop_y_start:crop_y_end, crop_x_start:crop_x_end + ] + + return parent_mask + + +def _merge_overlapping_masks( + predictions: List[dict], confidence_strategy: str, overlap_threshold: float = 0.0 +) -> List[dict]: + """ + Merge overlapping masks of the same class using union operations. + + Args: + predictions: List of dictionaries with 'mask', 'confidence', and 'class_id' keys + confidence_strategy: How to combine confidence scores + overlap_threshold: Minimum overlap ratio (IoU) to merge (0.0 to 1.0) + + Returns: + List of merged prediction dictionaries + """ + if not predictions: + return [] + + n = len(predictions) + masks = [p["mask"] for p in predictions] + + # Pre-compute bbox arrays for cheap spatial pre-filtering. + # Bboxes are stored alongside masks when predictions are collected. + bboxes_present = all(p.get("bbox") is not None for p in predictions) + if bboxes_present: + bboxes = np.array([p["bbox"] for p in predictions], dtype=np.float64) + bx1, by1, bx2, by2 = bboxes[:, 0], bboxes[:, 1], bboxes[:, 2], bboxes[:, 3] + + parent = list(range(n)) + + def find(x: int) -> int: + root = x + while parent[root] != root: + root = parent[root] + while parent[x] != root: + next_x = parent[x] + parent[x] = root + x = next_x + return root + + def union(x: int, y: int) -> None: + px, py = find(x), find(y) + if px != py: + parent[px] = py + + # Precompute per-mask pixel counts so IoU union can be derived arithmetically + # (area_i + area_j - intersection) instead of materialising a full OR array. + mask_areas = [int(np.count_nonzero(m)) for m in masks] + + for i in range(n): + mask_i = masks[i] + for j in range(i + 1, n): + # Cheap bbox pre-filter: skip pixel-level AND if bboxes don't overlap + if bboxes_present: + if ( + bx2[i] <= bx1[j] + or bx2[j] <= bx1[i] + or by2[i] <= by1[j] + or by2[j] <= by1[i] + ): + continue + + mask_j = masks[j] + intersection_count = int(np.count_nonzero(mask_i & mask_j)) + if overlap_threshold <= 0.0: + if intersection_count > 0: + union(i, j) + else: + # Compute union via arithmetic to avoid allocating a temporary OR array + union_count = mask_areas[i] + mask_areas[j] - intersection_count + iou = intersection_count / union_count if union_count > 0 else 0.0 + if iou >= overlap_threshold: + union(i, j) + + groups_dict: dict = {} + for i in range(n): + root = find(i) + if root not in groups_dict: + groups_dict[root] = [] + groups_dict[root].append(i) + + merged_results = [] + for group_indices in groups_dict.values(): + group = [predictions[idx] for idx in group_indices] + confidences = [p["confidence"] for p in group] + if confidence_strategy == "max": + merged_confidence = max(confidences) + elif confidence_strategy == "mean": + merged_confidence = float(np.mean(confidences)) + else: # min + merged_confidence = min(confidences) + + merged_mask = masks[group_indices[0]].copy() + for idx in group_indices[1:]: + merged_mask |= masks[idx] + + if merged_mask.any(): + merged_results.append( + { + "mask": merged_mask, + "confidence": merged_confidence, + "class_id": group[0]["class_id"], + "detection_data": group[0].get("detection_data", {}), + } + ) + + return merged_results + + +def _merge_overlapping_bboxes( + predictions: List[dict], confidence_strategy: str, overlap_threshold: float = 0.0 +) -> List[dict]: + """ + Merge overlapping bounding boxes of the same class. + + Args: + predictions: List of dictionaries with 'bbox', 'confidence', and 'class_id' keys + confidence_strategy: How to combine confidence scores + overlap_threshold: Minimum overlap ratio (IoU) to merge (0.0 to 1.0) + + Returns: + List of merged prediction dictionaries + """ + if not predictions: + return [] + + # Find connected components (groups of overlapping bboxes) + groups = _find_overlapping_bbox_groups(predictions, overlap_threshold) + + # Merge each group + merged_results = [] + for group in groups: + # Calculate merged confidence + confidences = [item["confidence"] for item in group] + if confidence_strategy == "max": + merged_confidence = max(confidences) + elif confidence_strategy == "mean": + merged_confidence = np.mean(confidences) + elif confidence_strategy == "min": + merged_confidence = min(confidences) + else: + merged_confidence = max(confidences) + + class_id = group[0]["class_id"] + + # Merge bounding boxes - take the union (min/max coordinates) + bboxes_arr = np.array([item["bbox"] for item in group]) + merged_bbox = np.array( + [ + bboxes_arr[:, 0].min(), + bboxes_arr[:, 1].min(), + bboxes_arr[:, 2].max(), + bboxes_arr[:, 3].max(), + ] + ) + + merged_results.append( + { + "bbox": merged_bbox, + "confidence": merged_confidence, + "class_id": class_id, + "detection_data": group[0].get( + "detection_data", {} + ), # Preserve first detection's metadata + } + ) + + return merged_results + + +def _find_overlapping_bbox_groups( + predictions: List[dict], overlap_threshold: float = 0.0 +) -> List[List[dict]]: + """ + Find groups of overlapping bounding boxes using union-find with vectorised numpy ops. + + Performs all intersection/IoU computations as batched numpy operations, avoiding + per-pair Python overhead and eliminating the need for Shapely geometry objects. + + This is a vectorised all-pairs broadphase: for each box it computes intersections + against all subsequent boxes in one numpy call. This removes per-pair Python + overhead and avoids Shapely geometry objects, but worst-case runtime is O(nยฒ) in + the number of predictions. For typical rollup workloads (tens of detections per + class) this is faster than an STRtree due to lower constant overhead, but for + very large detection sets a spatial index would be preferable. + + Args: + predictions: List of dictionaries with 'bbox' key + overlap_threshold: Minimum overlap ratio (IoU) to consider overlap + + Returns: + List of groups, where each group is a list of prediction dicts + """ + n = len(predictions) + if n == 0: + return [] + + bboxes = np.array([p["bbox"] for p in predictions], dtype=np.float64) + x1 = bboxes[:, 0] + y1 = bboxes[:, 1] + x2 = bboxes[:, 2] + y2 = bboxes[:, 3] + areas = (x2 - x1) * (y2 - y1) + + parent = list(range(n)) + + def find(x: int) -> int: + root = x + while parent[root] != root: + root = parent[root] + while parent[x] != root: + next_x = parent[x] + parent[x] = root + x = next_x + return root + + def union(x: int, y: int) -> None: + px, py = find(x), find(y) + if px != py: + parent[px] = py + + for i in range(n - 1): + # Vectorised intersection against all j > i in one shot + inter_x1 = np.maximum(x1[i], x1[i + 1 :]) + inter_y1 = np.maximum(y1[i], y1[i + 1 :]) + inter_x2 = np.minimum(x2[i], x2[i + 1 :]) + inter_y2 = np.minimum(y2[i], y2[i + 1 :]) + inter_w = np.maximum(0.0, inter_x2 - inter_x1) + inter_h = np.maximum(0.0, inter_y2 - inter_y1) + intersection = inter_w * inter_h + + if overlap_threshold <= 0.0: + overlapping = np.where(intersection > 0)[0] + else: + union_areas = areas[i] + areas[i + 1 :] - intersection + iou = np.where(union_areas > 0, intersection / union_areas, 0.0) + overlapping = np.where(iou >= overlap_threshold)[0] + + for j_offset in overlapping: + union(i, i + 1 + int(j_offset)) + + groups_dict: dict = {} + for i in range(n): + root = find(i) + if root not in groups_dict: + groups_dict[root] = [] + groups_dict[root].append(predictions[i]) + + return list(groups_dict.values()) diff --git a/inference/core/workflows/core_steps/fusion/detections_stitch/v1_tensor.py b/inference/core/workflows/core_steps/fusion/detections_stitch/v1_tensor.py new file mode 100644 index 0000000000..9c878db518 --- /dev/null +++ b/inference/core/workflows/core_steps/fusion/detections_stitch/v1_tensor.py @@ -0,0 +1,1170 @@ +from copy import copy +from typing import Dict, List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import numpy as np +import supervision as sv +import torch +import torchvision +from pydantic import ConfigDict, Field +from supervision import OverlapFilter +from supervision.config import ORIENTED_BOX_COORDINATES + +from inference.core import logger +from inference.core.workflows.core_steps.common.tensor_native import ( + embed_rle_masks_in_larger_canvas, + strip_host_mirror_metadata, + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, + SCALING_RELATIVE_TO_PARENT_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + STRING_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import ( + coco_rle_masks_to_numpy_mask, + torch_mask_to_coco_rle, +) + +TensorNativeDetections = Union[Detections, InstanceDetections] + +LONG_DESCRIPTION = """ +Merge detections from multiple image slices or crops back into a single unified detection result by converting coordinates from slice/crop space to original image coordinates, combining all detections, and optionally filtering overlapping detections to enable SAHI workflows, multi-stage detection pipelines, and coordinate-space merging workflows where detections from sub-images need to be reconstructed as if they were detected on the original image. + +## How This Block Works + +This block merges detections that were made on multiple sub-parts (slices or crops) of the same input image, reconstructing them as a single detection result in the original image coordinate space. The block: + +1. Receives reference image and slice/crop predictions: + - Takes the original reference image that was sliced or cropped + - Receives predictions from detection models that processed each slice/crop + - Predictions must contain parent coordinate metadata indicating slice/crop position +2. Retrieves crop offsets for each detection: + - Extracts parent coordinates from each detection's metadata + - Gets the offset (x, y position) indicating where each slice/crop was located in the original image + - Uses this offset to transform coordinates from slice space to original image space +3. Manages crop metadata: + - Updates image dimensions in detection metadata to match reference image dimensions + - Validates that detections were not scaled (scaled detections are not supported) + - Attaches parent coordinate information to detections for proper coordinate transformation +4. Transforms coordinates to original image space: + - Moves bounding box coordinates (xyxy) from slice/crop coordinates to original image coordinates + - Transforms segmentation masks from slice/crop space to original image space (if present) + - Applies offset to align detections with their position in the original image +5. Merges all transformed detections: + - Combines all re-aligned detections from all slices/crops into a single detection result + - Creates unified detection output containing all detections from all sub-images +6. Applies overlap filtering (optional): + - **None strategy**: Returns all merged detections without filtering (may contain duplicates from overlapping slices) + - **NMS (Non-Maximum Suppression)**: Removes lower-confidence detections when IoU exceeds threshold, keeping only the highest confidence detection for each overlapping region + - **NMM (Non-Maximum Merge)**: Combines overlapping detections instead of discarding them, merging detections that exceed IoU threshold +7. Returns merged detections: + - Outputs unified detection result in original image coordinate space + - Reduces dimensionality by 1 (multiple slice detections โ†’ single image detections) + - All detections are now referenced to the original image dimensions and coordinates + +This block is essential for SAHI (Slicing Adaptive Inference) workflows where an image is sliced, each slice is processed separately, and results need to be merged back. Overlapping slices can produce duplicate detections for the same object, so overlap filtering (NMS/NMM) helps clean up these duplicates. The coordinate transformation ensures that detection coordinates are correctly positioned relative to the original image, not the slices. + +## Common Use Cases + +- **SAHI Workflows**: Complete SAHI technique by merging detections from image slices back to original image coordinates (e.g., merge slice detections from SAHI processing, reconstruct full-image detections from slices, combine small object detection results), enabling SAHI detection workflows +- **Multi-Stage Detection**: Merge detections from secondary high-resolution models applied to dynamically cropped regions (e.g., coarse detection โ†’ crop โ†’ precise detection โ†’ merge, two-stage detection pipelines, hierarchical detection workflows), enabling multi-stage detection workflows +- **Small Object Detection**: Combine detection results from sliced images processed separately for small object detection (e.g., merge detections from aerial image slices, combine slice detection results, reconstruct detections from tiled images), enabling small object detection workflows +- **High-Resolution Processing**: Merge detections from high-resolution images processed in smaller chunks (e.g., merge detections from satellite image tiles, combine results from medical image regions, reconstruct detections from large image segments), enabling high-resolution detection workflows +- **Coordinate Space Unification**: Convert detections from multiple coordinate spaces (slice/crop space) to a single unified coordinate space (original image space) for consistent processing (e.g., unify detection coordinates, merge coordinate spaces, standardize detection positions), enabling coordinate unification workflows +- **Overlapping Region Handling**: Handle duplicate detections from overlapping slices or crops by applying overlap filtering (e.g., remove duplicate detections from overlapping slices, merge overlapping detections, clean up overlapping results), enabling overlap resolution workflows + +## Connecting to Other Blocks + +This block receives slice/crop predictions and reference images, and produces merged detections: + +- **After detection models in SAHI workflows** following Image Slicer โ†’ Detection Model โ†’ Detections Stitch pattern to merge slice detections (e.g., merge SAHI slice detections, reconstruct full-image detections, combine slice results), enabling SAHI completion workflows +- **After secondary detection models** in multi-stage pipelines following Dynamic Crop โ†’ Detection Model โ†’ Detections Stitch pattern to merge cropped detections (e.g., merge cropped region detections, combine two-stage detection results, unify multi-stage outputs), enabling multi-stage detection workflows +- **Before visualization blocks** to visualize merged detection results on the original image (e.g., visualize merged detections, display stitched results, show unified detection output), enabling visualization workflows +- **Before filtering or analytics blocks** to process merged detection results (e.g., filter merged detections, analyze stitched results, process unified outputs), enabling analysis workflows +- **Before sink or storage blocks** to store or export merged detection results (e.g., save merged detections, export stitched results, store unified outputs), enabling storage workflows +- **In workflow outputs** to provide merged detections as final workflow output (e.g., return merged detections, output stitched results, provide unified detection output), enabling output workflows + +## Requirements + +This block requires a reference image (the original image that was sliced/cropped) and predictions from detection models that processed slices/crops. The predictions must contain parent coordinate metadata (PARENT_COORDINATES_KEY) indicating the position of each slice/crop in the original image. The block does not support scaled detections (detections that were resized relative to the parent image). Predictions should be from object detection or instance segmentation models. The block supports three overlap filtering strategies: "none" (no filtering, may include duplicates), "nms" (Non-Maximum Suppression, removes lower-confidence overlapping detections, default), and "nmm" (Non-Maximum Merge, combines overlapping detections). The IoU threshold (default 0.3) determines when detections are considered overlapping for filtering purposes. For more information on SAHI technique, see: https://ieeexplore.ieee.org/document/9897990. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Detections Stitch", + "version": "v1", + "short_description": "Merges detections made against multiple pieces of input image into single detection.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "fusion", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-reel", + "blockPriority": 10, + "supervision": True, + }, + } + ) + type: Literal["roboflow_core/detections_stitch@v1"] + reference_image: Selector(kind=[IMAGE_KIND]) = Field( + description="Original reference image that was sliced or cropped to produce the input predictions. This image is used to determine the target coordinate space and image dimensions for the merged detections. All detection coordinates will be transformed to match this reference image's coordinate system. The same image that was provided to Image Slicer or Dynamic Crop blocks should be used here to ensure proper coordinate alignment.", + examples=["$inputs.image", "$steps.input_image.output"], + ) + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( + description="Model predictions (object detection or instance segmentation) from detection models that processed image slices or crops. These predictions must contain parent coordinate metadata indicating the position of each slice/crop in the original image. Predictions are collected from multiple slices/crops and merged into a single unified detection result. The block converts coordinates from slice/crop space to original image space and combines all detections.", + examples=[ + "$steps.object_detection.predictions", + "$steps.instance_segmentation.predictions", + "$steps.slice_model.predictions", + ], + ) + overlap_filtering_strategy: Union[ + Literal["none", "nms", "nmm"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="nms", + description="Strategy for handling overlapping detections when merging results from overlapping slices/crops. 'none': No filtering applied, all detections are kept (may include duplicates from overlapping regions). 'nms' (Non-Maximum Suppression, default): Removes lower-confidence detections when IoU exceeds threshold, keeping only the highest confidence detection for each overlapping region. 'nmm' (Non-Maximum Merge): Combines overlapping detections instead of discarding them, merging detections that exceed IoU threshold. Use 'none' when you want to preserve all detections, 'nms' to remove duplicates (recommended for most cases), or 'nmm' to combine overlapping detections.", + examples=["none", "nms", "nmm", "$inputs.filtering_strategy"], + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Intersection over Union (IoU) threshold for overlap filtering. Range: 0.0 to 1.0. When overlap filtering strategy is 'nms' or 'nmm', detections with IoU above this threshold are considered overlapping. For NMS: overlapping detections with IoU above threshold result in lower-confidence detection being removed. For NMM: overlapping detections with IoU above threshold are merged. Lower values (e.g., 0.2-0.3) are more aggressive, removing/merging more detections. Higher values (e.g., 0.5-0.7) are more permissive, only handling highly overlapping detections. Default 0.3 works well for most use cases with overlapping slices.", + examples=[0.2, 0.3, 0.4, 0.5, "$inputs.iou_threshold"], + ) + + @classmethod + def get_dimensionality_reference_property(cls) -> Optional[str]: + return "reference_image" + + @classmethod + def get_input_dimensionality_offsets(cls) -> Dict[str, int]: + return {"predictions": 1} + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class DetectionsStitchBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + reference_image: WorkflowImageData, + predictions: Batch[TensorNativeDetections], + overlap_filtering_strategy: Optional[Literal["none", "nms", "nmm"]], + iou_threshold: Optional[float], + ) -> BlockResult: + # Use reference image to ensure all masks have the same dimensions + reference_height, reference_width = read_image_shape(reference_image) + resolution_wh = (reference_width, reference_height) + + re_aligned_predictions = [] + # (dx, dy, row_count) per NON-EMPTY crop, in the order merge_detections + # concatenates them - see apply_deferred_box_shifts. + deferred_shifts: List[Tuple[float, float, int]] = [] + for detections in predictions: + detections_copy = shallow_copy_detections(detections) + offset = retrieve_crop_offset(detections=detections_copy) + detections_copy = manage_crops_metadata( + detections=detections_copy, image=reference_image + ) + re_aligned_detections = move_detections( + detections=detections_copy, + offset=offset, + resolution_wh=resolution_wh, + defer_box_shift=True, + ) + re_aligned_predictions.append(re_aligned_detections) + row_count = len(re_aligned_detections) + if row_count > 0: + # move_detections has already rejected a missing offset for a + # non-empty prediction, so offset is guaranteed here. + deferred_shifts.append( + (float(offset[0]), float(offset[1]), row_count) + ) + overlap_filter = choose_overlap_filter_strategy( + overlap_filtering_strategy=overlap_filtering_strategy, + ) + merged = merge_detections(detections_list=re_aligned_predictions) + # Boxes reach reference coordinates HERE, in one op. Masks were already + # embedded onto the reference canvas per crop (they need the offset to + # place RLE runs), so between the loop and this line the two geometries + # are transiently in different spaces - nothing reads them in between. + merged = apply_deferred_box_shifts(detections=merged, shifts=deferred_shifts) + if overlap_filter is OverlapFilter.NONE: + return {"predictions": merged} + if overlap_filter is OverlapFilter.NON_MAX_SUPPRESSION: + return {"predictions": with_nms(detections=merged, threshold=iou_threshold)} + return {"predictions": with_nmm(detections=merged, threshold=iou_threshold)} + + +def shallow_copy_detections( + detections: TensorNativeDetections, +) -> TensorNativeDetections: + """Per-crop working copy that SHARES the caller's tensors and metadata. + + This used to be ``deepcopy``, which cloned every device tensor (an + allocation plus a device copy each) and every per-box dict in + ``bboxes_metadata`` - O(rows) python object work on the hot path. None of + it was needed: the block never writes THROUGH a shared object, it rebinds + fields with freshly built ones (``xyxy`` via an add, + ``image_metadata`` via ``dict(...)`` in manage_crops_metadata, + ``bboxes_metadata`` via strip_host_mirror_metadata / the OBB rebuild, + ``mask`` via _move_native_masks). A shallow copy gives the same isolation + from the upstream step output - which other steps may still be reading - + at a fraction of the cost. + + The invariant this relies on: NOTHING below may mutate a shared field in + place. The OBB branch of move_detections is written as a rebuild for + exactly that reason. + """ + return copy(detections) + + +def apply_deferred_box_shifts( + detections: TensorNativeDetections, + shifts: List[Tuple[float, float, int]], +) -> TensorNativeDetections: + """Move every crop's boxes into reference coordinates in a single op. + + Shifting per crop cost one pageable host->device transfer (a 4-element + offset tensor) plus one add kernel PER CROP. On Jetson that launch and + transfer overhead dwarfs the arithmetic itself - a few hundred boxes moved + by two floats - and it scales with the slice count, which is exactly what + SAHI maximises. Expanding the offsets on host with ``np.repeat`` collapses + the whole thing to one (N, 4) transfer and one add regardless of how many + crops were stitched. + + ``shifts`` must list ``(dx, dy, row_count)`` for the non-empty predictions + in the same order ``merge_detections`` concatenated them. + """ + if not shifts or len(detections) == 0: + return detections + per_crop = np.asarray( + [[dx, dy, dx, dy] for dx, dy, _ in shifts], dtype=np.float64 + ) + row_counts = [row_count for _, _, row_count in shifts] + expanded = np.repeat(per_crop, row_counts, axis=0) + total_rows = int(detections.xyxy.shape[0]) + if expanded.shape[0] != total_rows: + # A silent mismatch would shift boxes by the WRONG crop's offset and + # emit plausible-looking garbage, so fail loudly instead. + raise RuntimeError( + f"Deferred crop offsets cover {expanded.shape[0]} boxes but the merged " + f"prediction has {total_rows}. The offsets collected while re-aligning " + f"crops must stay in lockstep with the rows merge_detections " + f"concatenates." + ) + shift = torch.as_tensor( + expanded, dtype=detections.xyxy.dtype, device=detections.xyxy.device + ) + detections.xyxy = detections.xyxy + shift + return detections + + +def read_image_shape(image: WorkflowImageData) -> Tuple[int, int]: + """(height, width) of the reference image without forcing a representation + conversion. + + Reading ``.numpy_image`` here would materialise the ENTIRE frame - a + device->host transfer plus the host-side RGB->BGR pass - to obtain two + integers. In a tensor-mode SAHI workflow the slicer and the model both stay + on device, so this block is the first consumer to touch numpy and pays that + full cost once per reference image, on every frame. Mirrors the guard used + by ``transformations/dynamic_crop/v1_tensor``. + """ + if image.is_tensor_materialised(): + # tensor_image is CHW -> H=shape[1], W=shape[2] + tensor_image = image.tensor_image + return int(tensor_image.shape[1]), int(tensor_image.shape[2]) + numpy_image = image.numpy_image + return int(numpy_image.shape[0]), int(numpy_image.shape[1]) + + +def retrieve_crop_offset(detections: TensorNativeDetections) -> Optional[np.ndarray]: + if len(detections) == 0: + return None + image_metadata = detections.image_metadata or {} + if PARENT_COORDINATES_KEY not in image_metadata: + raise RuntimeError( + f"Offset for crops is expected to be saved in image_metadata key {PARENT_COORDINATES_KEY} " + f"of the tensor-native detections, but could not be found. Probably block producing the " + f"detections lack this part of implementation or has a bug." + ) + return np.asarray(image_metadata[PARENT_COORDINATES_KEY][:2]).copy() + + +def manage_crops_metadata( + detections: TensorNativeDetections, + image: WorkflowImageData, +) -> TensorNativeDetections: + if len(detections) == 0: + return detections + + image_metadata = detections.image_metadata or {} + if SCALING_RELATIVE_TO_PARENT_KEY in image_metadata: + scale = image_metadata[SCALING_RELATIVE_TO_PARENT_KEY] + if abs(scale - 1.0) > 1e-4: + raise ValueError( + f"Scaled bounding boxes were passed to Detections Stitch block " + f"which is not supported. Block is supposed to merge predictions " + f"from multiple crops of the same image into single prediction, but " + f"scaling cannot be used in the meantime. This error probably indicate " + f"wrong step output plugged as input of this step." + ) + + height, width = read_image_shape(image) + image_metadata = dict(image_metadata) + image_metadata[IMAGE_DIMENSIONS_KEY] = [height, width] + image_metadata = attach_parents_coordinates_to_image_metadata( + image_metadata=image_metadata, + image=image, + ) + detections.image_metadata = image_metadata + + return detections + + +def attach_parents_coordinates_to_image_metadata( + image_metadata: dict, + image: WorkflowImageData, +) -> dict: + root = image.workflow_root_ancestor_metadata + parent = image.parent_metadata + root_coordinates = root.origin_coordinates + parent_coordinates = parent.origin_coordinates + image_metadata[ROOT_PARENT_ID_KEY] = root.parent_id + image_metadata[ROOT_PARENT_COORDINATES_KEY] = [ + root_coordinates.left_top_x, + root_coordinates.left_top_y, + ] + image_metadata[ROOT_PARENT_DIMENSIONS_KEY] = [ + root_coordinates.origin_height, + root_coordinates.origin_width, + ] + image_metadata[PARENT_ID_KEY] = parent.parent_id + image_metadata[PARENT_COORDINATES_KEY] = [ + parent_coordinates.left_top_x, + parent_coordinates.left_top_y, + ] + image_metadata[PARENT_DIMENSIONS_KEY] = [ + parent_coordinates.origin_height, + parent_coordinates.origin_width, + ] + return image_metadata + + +def move_detections( + detections: TensorNativeDetections, + offset: Optional[np.ndarray], + resolution_wh: Optional[Tuple[int, int]], + defer_box_shift: bool = False, +) -> TensorNativeDetections: + """ + Shift detections by ``offset``, keeping every geometry field consistent: + axis-aligned boxes, segmentation masks, and oriented-box corners. + + Mirrors ``supervision.detection.tools.inference_slicer.move_detections``; + kept local since that helper is not part of supervision's public API. + + With ``defer_box_shift=True`` the ``xyxy`` translation is SKIPPED and left + to ``apply_deferred_box_shifts``, which does it for every crop at once + after the merge - one transfer and one kernel instead of two per crop. + Masks and OBB corners are still moved here: both need this crop's offset + at this point (masks are embedded onto the reference canvas, OBB corners + live in host metadata), and neither is worth batching. + """ + if len(detections) == 0: + return detections + if offset is None: + raise ValueError("To move non-empty detections offset is needed, but not given") + if not defer_box_shift: + # Translate xyxy on-device: add [dx, dy, dx, dy] broadcast over the box + # rows, staying on detections.xyxy.device/dtype (no D2H->H2D round-trip). + dx, dy = float(offset[0]), float(offset[1]) + xyxy_shift = torch.as_tensor( + [dx, dy, dx, dy], + dtype=detections.xyxy.dtype, + device=detections.xyxy.device, + ) + detections.xyxy = detections.xyxy + xyxy_shift + # Drop the per-box host mirror: xyxy moves from slice-local to image + # coordinates (here or in apply_deferred_box_shifts), so a carried mirror + # would be stale โ€” consumers fall back to tensor reads. + detections.bboxes_metadata = strip_host_mirror_metadata(detections.bboxes_metadata) + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is not None and any( + ORIENTED_BOX_COORDINATES in (box_metadata or {}) + for box_metadata in bboxes_metadata + ): + # OBB corners live per-detection in `bboxes_metadata[i]["xyxyxyxy"]` with + # shape (4, 2); broadcast `offset` (shape (2,)) over the trailing axis to + # translate each (x, y). Without this, downstream OBB-aware NMS/NMM compares + # corners in tile-local coords against `xyxy` already moved to image coords. + # Rebuilt rather than mutated in place: the caller's copy is shallow + # (shallow_copy_detections), and strip_host_mirror_metadata passes + # mirror-free entries through by reference, so writing into one of these + # dicts would corrupt the upstream step's prediction. + detections.bboxes_metadata = [ + ( + { + **box_metadata, + ORIENTED_BOX_COORDINATES: np.asarray( + box_metadata[ORIENTED_BOX_COORDINATES] + ) + + offset, + } + if box_metadata and ORIENTED_BOX_COORDINATES in box_metadata + else box_metadata + ) + for box_metadata in bboxes_metadata + ] + if isinstance(detections, InstanceDetections) and detections.mask is not None: + if resolution_wh is None: + raise ValueError( + "To move non-empty detections with segmentation mask, resolution_wh is needed, but not given." + ) + detections.mask = _move_native_masks( + mask=detections.mask, offset=offset, resolution_wh=resolution_wh + ) + return detections + + +def _move_native_masks( + mask: Union[torch.Tensor, InstancesRLEMasks], + offset: np.ndarray, + resolution_wh: Tuple[int, int], +) -> InstancesRLEMasks: + # Move masks in RLE form via embed_rle_masks_in_larger_canvas: the slice-resolution + # masks are placed onto the (H, W) reference canvas at the slice's top-left offset, + # entirely in RLE (the big canvas is never densified, no full-frame D2H). The output + # is always RLE โ€” dense input is first converted slice-by-slice to RLE. + target_width, target_height = resolution_wh + target_size_hw = (target_height, target_width) + x0, y0 = int(offset[0]), int(offset[1]) + if isinstance(mask, InstancesRLEMasks): + rle_masks = mask + else: + # DENSE input: convert each slice mask to RLE, wrap as InstancesRLEMasks. + dense_masks = mask.detach().to(dtype=torch.bool) + if dense_masks.shape[0] == 0: + slice_height, slice_width = int(dense_masks.shape[1]), int( + dense_masks.shape[2] + ) + rle_masks = InstancesRLEMasks( + image_size=(slice_height, slice_width), masks=[] + ) + else: + slice_height, slice_width = int(dense_masks.shape[1]), int( + dense_masks.shape[2] + ) + counts = [ + torch_mask_to_coco_rle(single_mask)["counts"] + for single_mask in dense_masks + ] + rle_masks = InstancesRLEMasks( + image_size=(slice_height, slice_width), masks=counts + ) + return embed_rle_masks_in_larger_canvas( + masks=rle_masks, + offset_xy=(x0, y0), + target_size_hw=target_size_hw, + ) + + +def merge_detections( + detections_list: List[TensorNativeDetections], +) -> TensorNativeDetections: + """Concatenate native detections from every crop into a single prediction. + + Mirrors ``sv.Detections.merge`` but operates directly on the + ``inference_models`` dataclasses: ``xyxy`` / ``class_id`` / ``confidence`` + tensors are concatenated, ``bboxes_metadata`` lists are joined, masks (dense + or RLE) are stacked, and per-image ``image_metadata`` (parent coordinates, + dimensions, class-name maps) is taken from the re-aligned crops with their + ``class_names`` maps unioned. NOTE (shared-helper candidate): a native + merge/concat of detections is needed across several blocks and should be + consolidated by the maintainer. + """ + non_empty = [detections for detections in detections_list if len(detections) > 0] + is_instance_segmentation = any( + isinstance(detections, InstanceDetections) for detections in detections_list + ) + image_metadata = _merge_image_metadata(detections_list) + if len(non_empty) == 0: + if is_instance_segmentation: + return InstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, 0, 0), dtype=torch.bool), + image_metadata=image_metadata, + bboxes_metadata=None, + ) + return Detections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + image_metadata=image_metadata, + bboxes_metadata=None, + ) + xyxy = torch.cat([detections.xyxy for detections in non_empty], dim=0) + class_id = torch.cat([detections.class_id for detections in non_empty], dim=0) + confidence = torch.cat([detections.confidence for detections in non_empty], dim=0) + bboxes_metadata: List[dict] = [] + for detections in non_empty: + per_detection = detections.bboxes_metadata + if per_detection is None: + bboxes_metadata.extend({} for _ in range(len(detections))) + continue + bboxes_metadata.extend( + dict(entry) if isinstance(entry, dict) else entry + for entry in per_detection + ) + if is_instance_segmentation: + mask = _merge_masks(non_empty) + return InstanceDetections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + return Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def _merge_image_metadata( + detections_list: List[TensorNativeDetections], +) -> Optional[dict]: + merged: Optional[dict] = None + class_names: Dict[int, str] = {} + for detections in detections_list: + image_metadata = detections.image_metadata + if not image_metadata: + continue + if merged is None: + merged = dict(image_metadata) + per_image_class_names = image_metadata.get(CLASS_NAMES_KEY) + if per_image_class_names: + class_names.update(per_image_class_names) + if merged is None: + return None + if class_names: + merged[CLASS_NAMES_KEY] = class_names + return merged + + +def _ensure_class_names_map(image_metadata: dict, class_id: torch.Tensor) -> dict: + """Guarantee image_metadata carries a class_names map covering every output + class_id so the serializer's per-row class_id lookup never fails when the + merged metadata was None/map-less. Existing names are preserved; any + surviving class_id without an entry falls back to f"class_{id}".""" + image_metadata = dict(image_metadata) + class_names = dict(image_metadata.get(CLASS_NAMES_KEY) or {}) + class_names = {int(key): value for key, value in class_names.items()} + for class_id_value in class_id.detach().to("cpu").tolist(): + class_id_int = int(class_id_value) + if class_id_int not in class_names: + class_names[class_id_int] = f"class_{class_id_int}" + image_metadata[CLASS_NAMES_KEY] = class_names + return image_metadata + + +def _merge_masks( + detections_list: List[InstanceDetections], +) -> Union[torch.Tensor, InstancesRLEMasks]: + any_rle = any( + isinstance(detections.mask, InstancesRLEMasks) for detections in detections_list + ) + if any_rle: + image_size = None + rle_masks: List[bytes] = [] + for detections in detections_list: + mask = detections.mask + if isinstance(mask, InstancesRLEMasks): + image_size = mask.image_size + rle_masks.extend(mask.masks) + else: + numpy_masks = mask.detach().to("cpu").numpy().astype(bool) + for single_mask in numpy_masks: + rle = torch_mask_to_coco_rle( + torch.as_tensor(single_mask, dtype=torch.bool) + ) + image_size = tuple(rle["size"]) + rle_masks.append(rle["counts"]) + return InstancesRLEMasks(image_size=image_size, masks=rle_masks) + return torch.cat([detections.mask for detections in detections_list], dim=0) + + +def with_nms( + detections: TensorNativeDetections, + threshold: float, +) -> TensorNativeDetections: + if len(detections) == 0: + return detections + # NMS runs on-device via torchvision.ops.batched_nms over the merged native + # xyxy/confidence/class_id tensors (box IoU, class-aware โ€” matching the prior + # sv box NMS, which this block fed without masks). The kept indices index the + # native detections directly; no sv.Detections is materialised. batched_nms + # returns kept indices sorted by descending score; sort ascending so the + # surviving rows keep their original relative order. + boxes = detections.xyxy.to(dtype=torch.float32) + scores = detections.confidence.to(dtype=torch.float32) + # torchvision's NMS kernel requires float boxes/scores and an int64 idxs + # (the kernel is not implemented for int32 class ids). + class_ids = detections.class_id.to(dtype=torch.long) + keep = torchvision.ops.batched_nms( + boxes=boxes, + scores=scores, + idxs=class_ids, + iou_threshold=threshold, + ) + surviving_indices = torch.sort(keep).values.detach().to("cpu").tolist() + survivors = take_prediction_by_indices( + prediction=detections, indices=surviving_indices + ) + # Guarantee a class_names map covering the surviving class_ids so a non-empty + # NMS result never trips the serializer when the merged metadata is map-less. + survivors.image_metadata = _ensure_class_names_map( + survivors.image_metadata or {}, survivors.class_id + ) + return survivors + + +def with_nmm( + detections: TensorNativeDetections, + threshold: float, +) -> TensorNativeDetections: + if len(detections) == 0: + return detections + if isinstance(detections, InstanceDetections) and isinstance( + detections.mask, torch.Tensor + ): + # Dense device masks: run the torch NMM port, which keeps the (N, H, W) + # masks on device end-to-end (only the tiny xyxy/confidence/class_id and + # the (N, N) resized-mask intersection matrix cross the PCIe bus). + # Decision- and value-identical to the sv path below; the sv path stays + # as insurance and for RLE / mask-less inputs. + try: + return _with_nmm_dense_masks_torch( + detections=detections, threshold=threshold + ) + except Exception as error: + logger.warning( + f"Torch NMM path of detections_stitch failed ({error}); " + f"falling back to supervision-based NMM." + ) + return _with_nmm_sv(detections=detections, threshold=threshold) + + +def _with_nmm_sv( + detections: TensorNativeDetections, + threshold: float, +) -> TensorNativeDetections: + # sv.Detections is used here purely as the NMM algorithm (mirroring the numpy + # block's `merged.with_nmm`). NMM rewrites box geometry (and union-merges + # masks), so the merged sv output cannot be mapped 1:1 back to native rows; + # the native result is built fresh from the sv output's xyxy/class_id/ + # confidence/mask, with class names preserved from the input image_metadata + # and fresh per-detection ids. No sv.Detections is returned. + image_metadata = detections.image_metadata or {} + is_instance_segmentation = isinstance(detections, InstanceDetections) + masks = None + if is_instance_segmentation and detections.mask is not None: + if isinstance(detections.mask, InstancesRLEMasks): + masks = coco_rle_masks_to_numpy_mask(detections.mask) + else: + masks = detections.mask.detach().to("cpu").numpy().astype(bool) + nmm_output = sv.Detections( + xyxy=detections.xyxy.detach().to("cpu").numpy().astype(float), + confidence=detections.confidence.detach().to("cpu").numpy().astype(float), + class_id=detections.class_id.detach().to("cpu").numpy().astype(int), + mask=masks, + ).with_nmm(threshold=threshold) + number_of_detections = len(nmm_output) + bboxes_metadata = [ + {DETECTION_ID_KEY: str(uuid4())} for _ in range(number_of_detections) + ] + device = detections.xyxy.device + xyxy = torch.as_tensor( + np.asarray(nmm_output.xyxy), dtype=torch.float32, device=device + ).reshape(-1, 4) + class_id = torch.as_tensor( + np.asarray(nmm_output.class_id), dtype=torch.long, device=device + ) + confidence = torch.as_tensor( + np.asarray(nmm_output.confidence), dtype=torch.float32, device=device + ) + # Guarantee a class_names map covering the output class_ids so a non-empty + # NMM result never trips the serializer when the input metadata is None/map-less. + image_metadata = _ensure_class_names_map(image_metadata, class_id) + if is_instance_segmentation: + if nmm_output.mask is not None: + mask = torch.as_tensor( + np.asarray(nmm_output.mask), dtype=torch.bool, device=device + ) + else: + mask = torch.zeros((number_of_detections, 0, 0), dtype=torch.bool) + return InstanceDetections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + return Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +# Constants of the torch NMM port. _NMM_MASK_DIMENSION mirrors the default +# `mask_dimension` of `sv...mask_non_max_merge` (grouping decisions are taken on +# masks resized to this max dimension). The float budget caps the transient +# float32 copy of the flattened resized masks used by the pairwise-intersection +# matmul; above it the matmul is tiled (results are identical either way โ€” the +# counts are exact integers, see `_pairwise_mask_intersection_on_device`). +_NMM_MASK_DIMENSION = 640 +_NMM_PAIRWISE_FLOAT_BUDGET_BYTES = 128 * 1024 * 1024 + + +def _with_nmm_dense_masks_torch( + detections: InstanceDetections, + threshold: float, +) -> InstanceDetections: + """Mask-based NMM equivalent to `sv.Detections.with_nmm` (class-aware, IoU + metric) that never ships the dense (N, H, W) masks through host memory. + + supervision semantics replicated exactly (from supervision 0.29.1 sources): + + * grouping decisions run on masks resized to max-dim 640 with the + nearest-index grid of `sv...resize_masks` (`np.linspace(0, dim - 1, + new_dim).astype(int)` sampling โ€” upsamples when the mask is smaller); + * per class id (ascending `np.unique` order), detections are seeded by + descending confidence (`scores.argsort()` pop-from-the-back) and the + merge candidate is the *union* of already-absorbed resized masks; + absorption is retried against the growing union until a fixed point + (`sv..._group_overlapping_masks`); + * IoU arithmetic mirrors `sv..._mask_iou_batch_split`: float32 + integer-exact intersection counts, `union = area_a + area_b - inter` in + float32, division into a float64 zeros buffer where union != 0; + * each group merges into one output detection (`sv._merge_detection_group`): + union box via float32 min/max, confidence = box-area-weighted mean of + member confidences (`np.dot(f32 areas, f64 confs) / total_area`, cast to + float32; the winner's confidence when total area <= 0), the winning + (highest-confidence) member's class id, mask = logical OR of the ORIGINAL + full-resolution masks; singleton groups pass through unchanged. + + Device traffic: one D2H of xyxy/confidence/class_id, one D2H of the (N, N) + intersection matrix + (N,) areas of the resized masks, one tiny D2H per + union-growth round (rare: only groups that absorb members and re-test), and + one small H2D of the merged xyxy/confidence/class_id. Full-resolution masks + never leave the device. + """ + device = detections.xyxy.device + masks = detections.mask.detach() + if masks.dtype is not torch.bool: + masks = masks.to(dtype=torch.bool) + number_of_input_detections = int(masks.shape[0]) + # Host copies of the small per-detection fields with the same dtypes the + # sv-based path feeds into sv.Detections (float64 boxes / confidences, + # int class ids) so every downstream decision is bit-identical. + xyxy_host = detections.xyxy.detach().to("cpu").numpy().astype(float) + confidence_host = detections.confidence.detach().to("cpu").numpy().astype(float) + class_id_host = detections.class_id.detach().to("cpu").numpy().astype(int) + resized_masks = _resize_masks_like_supervision( + masks=masks, max_dimension=_NMM_MASK_DIMENSION + ) + flat_masks = resized_masks.reshape(number_of_input_detections, -1) + intersection, areas = _pairwise_mask_intersection_on_device(flat_masks=flat_masks) + intersection_host = intersection.to("cpu").numpy() + areas_host = areas.to("cpu").numpy() + merge_groups = _mask_non_max_merge_groups( + confidence=confidence_host, + class_id=class_id_host, + intersection=intersection_host, + areas=areas_host, + flat_masks=flat_masks, + iou_threshold=threshold, + ) + xyxy_out, confidence_out, class_id_out = _merge_detection_groups_bookkeeping( + merge_groups=merge_groups, + xyxy=xyxy_host, + confidence=confidence_host, + class_id=class_id_host, + ) + mask = _union_masks_on_device(masks=masks, merge_groups=merge_groups) + number_of_detections = len(merge_groups) + bboxes_metadata = [ + {DETECTION_ID_KEY: str(uuid4())} for _ in range(number_of_detections) + ] + xyxy = torch.as_tensor(xyxy_out, dtype=torch.float32, device=device).reshape(-1, 4) + class_id = torch.as_tensor(class_id_out, dtype=torch.long, device=device) + confidence = torch.as_tensor(confidence_out, dtype=torch.float32, device=device) + image_metadata = _ensure_class_names_map(detections.image_metadata or {}, class_id) + return InstanceDetections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask.to(device), + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def _resize_masks_like_supervision( + masks: torch.Tensor, max_dimension: int +) -> torch.Tensor: + """Torch replica of `sv...resize_masks`: nearest-index resampling so the + largest mask dimension becomes `max_dimension` (aspect ratio kept). The + integer sampling grids are computed on host with the exact numpy expression + supervision uses, so the resampled masks match sv's pixel-for-pixel.""" + height, width = int(masks.shape[1]), int(masks.shape[2]) + scale = min(max_dimension / height, max_dimension / width) + new_height = int(scale * height) + new_width = int(scale * width) + if new_height == height and new_width == width: + # np.linspace(0, dim - 1, dim).astype(int) is the identity grid. + return masks + y_indices = torch.as_tensor( + np.linspace(0, height - 1, new_height).astype(int), + dtype=torch.long, + device=masks.device, + ) + x_indices = torch.as_tensor( + np.linspace(0, width - 1, new_width).astype(int), + dtype=torch.long, + device=masks.device, + ) + return masks.index_select(1, y_indices).index_select(2, x_indices) + + +def _pairwise_mask_intersection_on_device( + flat_masks: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(N, N) float32 intersection pixel counts and (N,) float32 areas of the + flattened bool masks, computed on the masks' device. + + Counts are exact integers: the matmul operands are 0/1 (exact in float32 + and in TF32's 10-bit mantissa) and every partial sum stays below + 2**24 (P <= 640 * 640 after resize), so any accumulation order โ€” cuBLAS, + TF32-with-fp32-accumulate, CPU BLAS โ€” yields the exact count, matching + supervision's float32 numpy matmul bit-for-bit. Memory: the transient + float32 copy of the flattened masks is tiled once it would exceed + _NMM_PAIRWISE_FLOAT_BUDGET_BYTES (the bool masks themselves stay resident). + """ + number_of_masks, pixels = int(flat_masks.shape[0]), int(flat_masks.shape[1]) + if number_of_masks * pixels * 4 <= _NMM_PAIRWISE_FLOAT_BUDGET_BYTES: + flat_f32 = flat_masks.to(dtype=torch.float32) + return flat_f32 @ flat_f32.T, flat_f32.sum(dim=1) + chunk = max(1, _NMM_PAIRWISE_FLOAT_BUDGET_BYTES // (2 * max(pixels, 1) * 4)) + intersection = torch.empty( + (number_of_masks, number_of_masks), + dtype=torch.float32, + device=flat_masks.device, + ) + areas = torch.empty( + (number_of_masks,), dtype=torch.float32, device=flat_masks.device + ) + for row_start in range(0, number_of_masks, chunk): + row_end = min(row_start + chunk, number_of_masks) + rows_f32 = flat_masks[row_start:row_end].to(dtype=torch.float32) + areas[row_start:row_end] = rows_f32.sum(dim=1) + for col_start in range(row_start, number_of_masks, chunk): + col_end = min(col_start + chunk, number_of_masks) + if col_start == row_start: + cols_f32 = rows_f32 + else: + cols_f32 = flat_masks[col_start:col_end].to(dtype=torch.float32) + block = rows_f32 @ cols_f32.T + intersection[row_start:row_end, col_start:col_end] = block + if col_start != row_start: + intersection[col_start:col_end, row_start:row_end] = block.T + return intersection, areas + + +def _mask_non_max_merge_groups( + confidence: np.ndarray, + class_id: np.ndarray, + intersection: np.ndarray, + areas: np.ndarray, + flat_masks: torch.Tensor, + iou_threshold: float, +) -> List[List[int]]: + """Port of `sv...mask_non_max_merge` (class-aware, IoU metric): per class id + in ascending order, run the greedy growing-union grouping and translate the + class-local groups back to global row indices.""" + merge_groups: List[List[int]] = [] + for category_id in np.unique(class_id): + current_indices = np.where(class_id == category_id)[0] + local_groups = _group_overlapping_masks_greedy( + scores=confidence[current_indices], + global_indices=current_indices, + intersection=intersection, + areas=areas, + flat_masks=flat_masks, + iou_threshold=iou_threshold, + ) + for local_group in local_groups: + merge_groups.append(current_indices[local_group].tolist()) + return merge_groups + + +def _group_overlapping_masks_greedy( + scores: np.ndarray, + global_indices: np.ndarray, + intersection: np.ndarray, + areas: np.ndarray, + flat_masks: torch.Tensor, + iou_threshold: float, +) -> List[List[int]]: + """Port of `sv..._group_overlapping_masks` returning groups of positions + into `global_indices` (sv's class-local indices). The first absorption round + of every group is answered from the precomputed pairwise intersection + matrix (the candidate is still a single mask); only rounds against a grown + union candidate query the device โ€” one small D2H each.""" + merge_groups: List[List[int]] = [] + order = scores.argsort() + while len(order) > 0: + idx = int(order[-1]) + order = order[:-1] + if len(order) == 0: + merge_groups.append([idx]) + break + candidate_group = [idx] + first_round = True + while len(order) > 0: + remaining_global = global_indices[order] + if first_round: + seed_global = int(global_indices[idx]) + intersection_vector = intersection[remaining_global, seed_global] + candidate_area = areas[seed_global] + else: + intersection_vector, candidate_area = _union_candidate_iou_inputs( + flat_masks=flat_masks, + member_global_ids=global_indices[candidate_group], + remaining_global_ids=remaining_global, + ) + # Same arithmetic as sv._mask_iou_batch_split: float32 union, + # division into a float64 zeros buffer where union != 0. + union_area = areas[remaining_global] + candidate_area - intersection_vector + ious = np.divide( + intersection_vector, + union_area, + out=np.zeros_like(intersection_vector, dtype=float), + where=union_area != 0, + ) + ious = np.nan_to_num(ious) + above_threshold = ious >= iou_threshold + if not above_threshold.any(): + break + above_idx = order[above_threshold] + candidate_group.extend(np.flip(above_idx).tolist()) + order = order[~above_threshold] + first_round = False + merge_groups.append(candidate_group) + return merge_groups + + +def _union_candidate_iou_inputs( + flat_masks: torch.Tensor, + member_global_ids: np.ndarray, + remaining_global_ids: np.ndarray, +) -> Tuple[np.ndarray, np.float32]: + """Intersection counts of the remaining resized masks against the union of + the current group members, plus the union's area โ€” computed on device and + shipped as one packed vector (single small D2H sync). Counts are exact + integers, matching sv's float32 numpy arithmetic bit-for-bit.""" + device = flat_masks.device + members = torch.as_tensor(member_global_ids, dtype=torch.long, device=device) + remaining = torch.as_tensor(remaining_global_ids, dtype=torch.long, device=device) + candidate_f32 = ( + flat_masks.index_select(0, members).any(dim=0).to(dtype=torch.float32) + ) + intersection_vector = ( + flat_masks.index_select(0, remaining).to(dtype=torch.float32) @ candidate_f32 + ) + packed = torch.cat([intersection_vector, candidate_f32.sum().reshape(1)]) + packed_host = packed.to("cpu").numpy() + return packed_host[:-1], np.float32(packed_host[-1]) + + +def _merge_detection_groups_bookkeeping( + merge_groups: List[List[int]], + xyxy: np.ndarray, + confidence: np.ndarray, + class_id: np.ndarray, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Port of the AABB branch of `sv._merge_detection_group` (+ the trivial + concatenation of `sv.Detections.merge`), on host numpy with the exact + dtypes sv uses so merged boxes and confidences are bit-identical.""" + number_of_groups = len(merge_groups) + xyxy_out = np.empty((number_of_groups, 4), dtype=np.float64) + confidence_out = np.empty((number_of_groups,), dtype=np.float64) + class_id_out = np.empty((number_of_groups,), dtype=np.int64) + for position, group in enumerate(merge_groups): + if len(group) == 1: + index = group[0] + xyxy_out[position] = xyxy[index] + confidence_out[position] = confidence[index] + class_id_out[position] = class_id[index] + continue + group_confidences = confidence[group] + winner_index = int(np.argmax(group_confidences)) + all_xyxy = xyxy[group].astype(np.float32) + box_areas = (all_xyxy[:, 2] - all_xyxy[:, 0]) * ( + all_xyxy[:, 3] - all_xyxy[:, 1] + ) + total_area = float(box_areas.sum()) + if total_area > 0: + merged_confidence = np.float32( + float(np.dot(box_areas, group_confidences) / total_area) + ) + else: + merged_confidence = group_confidences[winner_index] + xyxy_out[position] = [ + all_xyxy[:, 0].min(), + all_xyxy[:, 1].min(), + all_xyxy[:, 2].max(), + all_xyxy[:, 3].max(), + ] + confidence_out[position] = merged_confidence + class_id_out[position] = class_id[group[winner_index]] + return xyxy_out, confidence_out, class_id_out + + +def _union_masks_on_device( + masks: torch.Tensor, merge_groups: List[List[int]] +) -> torch.Tensor: + """Merged output masks built from the ORIGINAL full-resolution device masks: + singleton groups are gathered with one index_select, multi-member groups are + unioned with one batched index_add_ (member count per pixel > 0 == logical + OR โ€” exact bool result, matching sv's np.logical_or.reduce). Fixed kernel + count regardless of the number of groups.""" + device = masks.device + number_of_groups = len(merge_groups) + height, width = int(masks.shape[1]), int(masks.shape[2]) + mask_out = torch.empty( + (number_of_groups, height, width), dtype=torch.bool, device=device + ) + single_positions = [ + position for position, group in enumerate(merge_groups) if len(group) == 1 + ] + multi_positions = [ + position for position, group in enumerate(merge_groups) if len(group) > 1 + ] + if single_positions: + single_ids = torch.as_tensor( + [merge_groups[position][0] for position in single_positions], + dtype=torch.long, + device=device, + ) + positions = torch.as_tensor(single_positions, dtype=torch.long, device=device) + mask_out[positions] = masks.index_select(0, single_ids) + if multi_positions: + member_ids: List[int] = [] + member_slots: List[int] = [] + for slot, position in enumerate(multi_positions): + for member in merge_groups[position]: + member_ids.append(member) + member_slots.append(slot) + largest_group = max(len(merge_groups[position]) for position in multi_positions) + # uint8 accumulation wraps at 256 members per group; escape to int32 + # for (degenerate) larger groups. + accumulator_dtype = torch.uint8 if largest_group < 256 else torch.int32 + accumulator = torch.zeros( + (len(multi_positions), height, width), + dtype=accumulator_dtype, + device=device, + ) + member_ids_t = torch.as_tensor(member_ids, dtype=torch.long, device=device) + member_slots_t = torch.as_tensor(member_slots, dtype=torch.long, device=device) + accumulator.index_add_( + 0, member_slots_t, masks.index_select(0, member_ids_t).to(accumulator_dtype) + ) + positions = torch.as_tensor(multi_positions, dtype=torch.long, device=device) + mask_out[positions] = accumulator != 0 + return mask_out + + +def choose_overlap_filter_strategy( + overlap_filtering_strategy: Literal["none", "nms", "nmm"], +) -> sv.OverlapFilter: + if overlap_filtering_strategy == "none": + return sv.OverlapFilter.NONE + if overlap_filtering_strategy == "nms": + return sv.OverlapFilter.NON_MAX_SUPPRESSION + elif overlap_filtering_strategy == "nmm": + return sv.OverlapFilter.NON_MAX_MERGE + raise ValueError( + f"Invalid overlap filtering strategy: {overlap_filtering_strategy}" + ) diff --git a/inference/core/workflows/core_steps/fusion/frame_delay/v1_tensor.py b/inference/core/workflows/core_steps/fusion/frame_delay/v1_tensor.py new file mode 100644 index 0000000000..6e16a5b868 --- /dev/null +++ b/inference/core/workflows/core_steps/fusion/frame_delay/v1_tensor.py @@ -0,0 +1,142 @@ +"""Tensor-data-representation sibling of ``fusion/frame_delay/v1``. + +The delay machinery (per-video ring buffer keyed by frame number, LRU stream +tracking, frame-number-restart detection, window eviction) is +representation-agnostic and reused verbatim from the numpy sibling by +subclassing it. The ONLY behavioral change is at buffer-insertion time: image +payloads are spilled to host memory before they are buffered. + +Why: under ``ENABLE_TENSOR_DATA_REPRESENTATION`` a ``WorkflowImageData`` may +carry its pixels as a CHW uint8 CUDA tensor (``tensor_image``). Video frames +delivered by the Jetson tensor bridge borrow buffers from a small (8-slot) +per-pipeline CUDA pool that are recycled only once the tensor reference is +dropped. Buffering such frames verbatim pins ~11 MB of device memory per 2K +frame for the whole delay window, and ``|offset| > 8`` starves the bridge pool +outright, degrading every subsequent frame to fresh ``cudaMalloc`` churn. +Spilling at insertion converts that into ~11 MB of ordinary host memory per +buffered 2K frame and returns the pool buffer immediately. + +What is spilled: every ``WorkflowImageData`` found at the top level of ``data`` +or nested inside list/tuple containers - mirroring what the numpy manifest can +actually be wired with (``IMAGE_KIND`` at the top level; ``LIST_OF_VALUES_KIND`` +e.g. a collapsed list of crops). Each spilled entry is a NEW +``WorkflowImageData`` built via ``WorkflowImageData.copy_and_replace(..., +numpy_image=...)``: reading ``numpy_image`` first materialises host pixels +(the D2H conversion inside ``WorkflowImageData.numpy_image`` copies, so the +array does not alias the tensor storage), and ``copy_and_replace`` called with +an image-representation kwarg rebuilds the instance with every non-passed +representation slot set to ``None`` - the result provably holds no +``_tensor_image``, while ``parent_metadata``, ``workflow_root_ancestor_metadata`` +and ``video_metadata`` are carried over as-is. Images whose tensor was never +materialised (``is_tensor_materialised()`` is ``False``) are buffered +untouched: they hold no device memory, and spilling them would force an eager +decode of base64/reference-born images. On emission the stored host-resident +object is returned directly; downstream tensor-mode blocks re-materialise +``tensor_image`` lazily on first access (re-upload to +``WORKFLOWS_IMAGE_TENSOR_DEVICE``). + +Deliberately NOT spilled: every non-image payload - predictions, numbers, +strings, dicts, arbitrary objects - is buffered untouched, INCLUDING +tensor-native ``Detections`` / ``InstanceDetections`` whose boxes / masks live +on GPU. Delaying mask-carrying predictions therefore still retains their GPU +memory for the delay window; the image spill is the load-bearing fix, because +full frames dominate the footprint and pin the bridge buffer pool. Images +nested deeper than list/tuple containers (e.g. inside dicts) are likewise +buffered as-is. +""" + +from typing import Any, Optional, Type, Union + +from pydantic import ConfigDict + +from inference.core.workflows.core_steps.fusion.frame_delay.v1 import ( + LONG_DESCRIPTION as NUMPY_LONG_DESCRIPTION, +) +from inference.core.workflows.core_steps.fusion.frame_delay.v1 import ( + BlockManifest as NumpyBlockManifest, +) +from inference.core.workflows.core_steps.fusion.frame_delay.v1 import ( + FrameDelayBlockV1 as NumpyFrameDelayBlockV1, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TENSOR_MODE_ADDENDUM = """ +## Tensor Data Representation Behavior + +Under `ENABLE_TENSOR_DATA_REPRESENTATION`, image payloads wired into `data` may +carry their pixels as a CUDA tensor; on Jetson video pipelines that tensor is a +buffer borrowed from a small per-pipeline CUDA pool which is recycled only once +the tensor reference is dropped. This variant therefore spills image payloads +to host memory at buffering time: the buffered (and later emitted) image holds +host-resident `numpy_image` pixels and no tensor reference, so device memory is +released immediately instead of being pinned for the delay window. The cost is +~11 MB of host memory per buffered 2K frame (~6 MB at 1080p, ~25 MB at 4K); +downstream consumers of the delayed image re-upload it to the device lazily on +first `tensor_image` access. Non-image payloads (predictions, numbers, dicts, +...) are buffered untouched - delaying tensor-native predictions that carry GPU +masks still retains their GPU memory for the delay window. +""" + +LONG_DESCRIPTION = NUMPY_LONG_DESCRIPTION + TENSOR_MODE_ADDENDUM + + +class BlockManifest(NumpyBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + **NumpyBlockManifest.model_config["json_schema_extra"], + "long_description": LONG_DESCRIPTION, + } + ) + + +def _spill_images_to_host(data: Any) -> Any: + """Returns `data` with every `WorkflowImageData` found at the top level or + inside list/tuple containers replaced by a host-resident copy that holds no + tensor reference (see the module docstring). Containers are rebuilt only + when an element actually changed; everything else - non-image payloads, + images with no materialised tensor - is returned untouched, preserving + object identity exactly like the numpy sibling does.""" + if isinstance(data, WorkflowImageData): + if not data.is_tensor_materialised(): + # Already host-resident (or a base64/reference-born image that + # never materialised a tensor): buffering it holds no device + # memory, and spilling would force an eager decode. + return data + # Reading `numpy_image` materialises host pixels (the D2H conversion + # copies, so the array does not alias tensor storage); passing it as an + # image-representation kwarg makes `copy_and_replace` null every other + # representation slot, so the copy provably drops the tensor reference + # while keeping parent / root-ancestor / video metadata intact. + return WorkflowImageData.copy_and_replace( + origin_image_data=data, numpy_image=data.numpy_image + ) + if isinstance(data, (list, tuple)): + spilled = [_spill_images_to_host(data=element) for element in data] + if all(new is original for new, original in zip(spilled, data)): + return data + return tuple(spilled) if isinstance(data, tuple) else spilled + return data + + +class FrameDelayBlockV1(NumpyFrameDelayBlockV1): + """Numpy frame-delay block with image payloads spilled to host memory at + buffer-insertion time - see the module docstring for the full policy.""" + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + data: Any, + offset: int, + default_value: Optional[Union[bool, int, float, str]] = None, + ) -> BlockResult: + return super().run( + image=image, + data=_spill_images_to_host(data=data), + offset=offset, + default_value=default_value, + ) diff --git a/inference/core/workflows/core_steps/fusion/overlap_analysis/v1_tensor.py b/inference/core/workflows/core_steps/fusion/overlap_analysis/v1_tensor.py new file mode 100644 index 0000000000..68108e75c0 --- /dev/null +++ b/inference/core/workflows/core_steps/fusion/overlap_analysis/v1_tensor.py @@ -0,0 +1,397 @@ +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field +from shapely.geometry import Polygon, box + +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + DETECTIONS_OVERLAPS_KIND, + FLOAT_ZERO_TO_ONE_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import coco_rle_masks_to_numpy_mask + +LONG_DESCRIPTION = """ +Compute pairwise geometric overlap between two sets of detections. + +For each pair (reference, candidate) drawn from `reference_predictions` x +`candidate_predictions`, the block computes +`intersection_area / reference_polygon.area` and emits one record per +pair whose ratio reaches `min_overlap`. + +When a detection carries a mask, the precise polygon is the longest +contour of that mask (validated via shapely); otherwise the bounding +box polygon is used. A vectorised bbox-IoU prefilter +(`supervision.box_iou_batch`) eliminates non-touching pairs before any +shapely intersection is computed. + +The relation is intentionally **not symmetric** across the two inputs: +the denominator is always the reference detection's area. Swap the two +selectors if you want overlap reported relative to the other set. + +The output is a flat list of dicts (one per accepted pair) attached to +the same dimensionality as the inputs โ€” the block does not increase +dimensionality. See the `detections_overlaps` kind docs for the +per-record schema. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Overlap Analysis", + "version": "v1", + "short_description": ( + "Compute pairwise overlap between two sets of detections." + ), + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "fusion", + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/overlap_analysis@v1", "OverlapAnalysis"] + reference_predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( + description=( + "Detections whose area is the denominator of the overlap ratio. " + "For each reference detection, overlap with every candidate is " + "computed; pairs above `min_overlap` appear in the output." + ), + examples=["$steps.model_a.predictions"], + ) + candidate_predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( + description="Detections checked against each reference.", + examples=["$steps.model_b.predictions"], + ) + min_overlap: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( + default=0.1, + description=( + "Minimum (intersection / reference_area) ratio for a pair to be " + "included in the output." + ), + examples=[0.1, "$inputs.min_overlap"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="overlaps", kind=[DETECTIONS_OVERLAPS_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class OverlapAnalysisBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + reference_predictions: Union[Detections, InstanceDetections], + candidate_predictions: Union[Detections, InstanceDetections], + min_overlap: float, + ) -> BlockResult: + if len(reference_predictions) == 0 or len(candidate_predictions) == 0: + return {"overlaps": []} + + ref_xyxy = reference_predictions.xyxy.detach().to("cpu").numpy() + cand_xyxy = candidate_predictions.xyxy.detach().to("cpu").numpy() + + iou_matrix = sv.box_iou_batch(ref_xyxy, cand_xyxy) + + ref_ids = _detection_ids(reference_predictions) + cand_ids = _detection_ids(candidate_predictions) + ref_class_names = _class_names(reference_predictions) + cand_class_names = _class_names(candidate_predictions) + ref_confidences = _confidence_values(reference_predictions) + cand_confidences = _confidence_values(candidate_predictions) + + if _is_bbox_only(reference_predictions) and _is_bbox_only( + candidate_predictions + ): + ref_boxes = ref_xyxy.astype(np.float64, copy=False) + cand_boxes = cand_xyxy.astype(np.float64, copy=False) + + ref_x1 = np.minimum(ref_boxes[:, 0], ref_boxes[:, 2]) + ref_y1 = np.minimum(ref_boxes[:, 1], ref_boxes[:, 3]) + ref_x2 = np.maximum(ref_boxes[:, 0], ref_boxes[:, 2]) + ref_y2 = np.maximum(ref_boxes[:, 1], ref_boxes[:, 3]) + cand_x1 = np.minimum(cand_boxes[:, 0], cand_boxes[:, 2]) + cand_y1 = np.minimum(cand_boxes[:, 1], cand_boxes[:, 3]) + cand_x2 = np.maximum(cand_boxes[:, 0], cand_boxes[:, 2]) + cand_y2 = np.maximum(cand_boxes[:, 1], cand_boxes[:, 3]) + + intersection_width = np.clip( + np.minimum(ref_x2[:, None], cand_x2[None, :]) + - np.maximum(ref_x1[:, None], cand_x1[None, :]), + 0.0, + None, + ) + intersection_height = np.clip( + np.minimum(ref_y2[:, None], cand_y2[None, :]) + - np.maximum(ref_y1[:, None], cand_y1[None, :]), + 0.0, + None, + ) + intersection = intersection_width * intersection_height + ref_area = (ref_x2 - ref_x1) * (ref_y2 - ref_y1) + overlap_ratios = np.zeros_like(intersection, dtype=np.float64) + np.divide( + intersection, + ref_area[:, None], + out=overlap_ratios, + where=ref_area[:, None] > 0.0, + ) + emitted_pairs = np.argwhere( + (iou_matrix > 0.0) + & (ref_area[:, None] > 0.0) + & (overlap_ratios >= min_overlap) + ) + + results: List[Dict[str, Any]] = [] + for i, j in emitted_pairs: + results.append( + _build_record( + i=int(i), + j=int(j), + overlap_ratio=overlap_ratios[i, j], + ref_class_names=ref_class_names, + cand_class_names=cand_class_names, + ref_confidences=ref_confidences, + cand_confidences=cand_confidences, + ref_ids=ref_ids, + cand_ids=cand_ids, + ) + ) + return {"overlaps": results} + + eligible_pairs = iou_matrix > 0.0 + ref_mask_rows = _materialize_mask_rows( + reference_predictions, + np.flatnonzero(np.any(eligible_pairs, axis=1)), + ) + cand_mask_rows = _materialize_mask_rows( + candidate_predictions, + np.flatnonzero(np.any(eligible_pairs, axis=0)), + ) + + results: List[Dict[str, Any]] = [] + ref_polys: Dict[int, Polygon] = {} + cand_polys: Dict[int, Polygon] = {} + + for i in range(len(reference_predictions)): + for j in range(len(candidate_predictions)): + if not eligible_pairs[i, j]: + continue + if i not in ref_polys: + ref_polys[i] = _detection_to_shapely( + reference_predictions, ref_xyxy, i, ref_mask_rows + ) + if j not in cand_polys: + cand_polys[j] = _detection_to_shapely( + candidate_predictions, cand_xyxy, j, cand_mask_rows + ) + ref_poly = ref_polys[i] + cand_poly = cand_polys[j] + if ref_poly.area <= 0: + continue + intersection_area = ref_poly.intersection(cand_poly).area + overlap_ratio = intersection_area / ref_poly.area + if overlap_ratio < min_overlap: + continue + results.append( + _build_record( + i=i, + j=j, + overlap_ratio=overlap_ratio, + ref_class_names=ref_class_names, + cand_class_names=cand_class_names, + ref_confidences=ref_confidences, + cand_confidences=cand_confidences, + ref_ids=ref_ids, + cand_ids=cand_ids, + ) + ) + return {"overlaps": results} + + +def _is_bbox_only(detections: Union[Detections, InstanceDetections]) -> bool: + return not isinstance(detections, InstanceDetections) or detections.mask is None + + +def _confidence_values( + detections: Union[Detections, InstanceDetections], +) -> Optional[np.ndarray]: + if detections.confidence is None: + return None + return detections.confidence.detach().cpu().numpy() + + +def _materialize_mask_rows( + detections: Union[Detections, InstanceDetections], indices: np.ndarray +) -> Dict[int, np.ndarray]: + if ( + not isinstance(detections, InstanceDetections) + or detections.mask is None + or len(indices) == 0 + ): + return {} + mask = detections.mask + if isinstance(mask, InstancesRLEMasks): + return { + int(index): coco_rle_masks_to_numpy_mask( + InstancesRLEMasks( + image_size=mask.image_size, + masks=[mask.masks[int(index)]], + ) + )[0] + for index in indices + } + indices_tensor = torch.as_tensor(indices, dtype=torch.long, device=mask.device) + materialized = mask[indices_tensor].detach().to("cpu").numpy() + return { + int(index): materialized[position] for position, index in enumerate(indices) + } + + +def _build_record( + i: int, + j: int, + overlap_ratio: float, + ref_class_names: List[Optional[str]], + cand_class_names: List[Optional[str]], + ref_confidences: Optional[np.ndarray], + cand_confidences: Optional[np.ndarray], + ref_ids: Optional[List[Any]], + cand_ids: Optional[List[Any]], +) -> Dict[str, Any]: + record: Dict[str, Any] = { + "reference_class": _safe_get(ref_class_names, i), + "reference_confidence": ( + float(ref_confidences[i]) if ref_confidences is not None else None + ), + "candidate_class": _safe_get(cand_class_names, j), + "candidate_confidence": ( + float(cand_confidences[j]) if cand_confidences is not None else None + ), + "overlap_ratio": float(overlap_ratio), + } + if ref_ids is not None: + record["reference_detection_id"] = _safe_get(ref_ids, i) + if cand_ids is not None: + record["candidate_detection_id"] = _safe_get(cand_ids, j) + return record + + +def _detection_ids( + detections: Union[Detections, InstanceDetections], +) -> Optional[List[Any]]: + """Return the per-detection `detection_id` list (read from `bboxes_metadata`) + or `None` when no detection carries one โ€” mirroring the numpy block's + `detections.data.get(DETECTION_ID_KEY)`.""" + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + return None + ids = [m.get(DETECTION_ID_KEY) for m in bboxes_metadata] + if all(value is None for value in ids): + return None + return ids + + +def _class_names( + detections: Union[Detections, InstanceDetections], +) -> List[Optional[str]]: + """Resolve the per-detection class name from the `image_metadata` class-names + map (the tensor-native equivalent of the numpy block's + `detections.data.get(CLASS_NAME_DATA_FIELD)`).""" + class_names_map = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + class_ids = detections.class_id.detach().to("cpu").numpy() + return [ + class_names_map.get(int(class_id), f"class_{int(class_id)}") + for class_id in class_ids + ] + + +def _detection_to_shapely( + detections: Union[Detections, InstanceDetections], + xyxy: np.ndarray, + idx: int, + materialized_masks: Dict[int, np.ndarray], +) -> Polygon: + """Return the precise polygon for the detection at `idx`. + + When `detections` carries a mask and it is non-empty for this row, the + polygon is the longest contour of the mask (validated via shapely). + Otherwise โ€” and on any invalidity / emptiness โ€” the 4-corner bbox polygon is + returned. Object detection (`inference_models.Detections`) has no mask, so it + always falls back to the bbox polygon. + """ + x1, y1, x2, y2 = xyxy[idx] + bbox_poly = box(float(x1), float(y1), float(x2), float(y2)) + if isinstance(detections, InstanceDetections) and detections.mask is not None: + mask = materialized_masks[idx] + if np.any(mask): + polygons = sv.mask_to_polygons(mask=mask.astype(np.uint8)) + if polygons: + longest = max(polygons, key=len) + if len(longest) >= 3: + candidate = Polygon( + [(float(pt[0]), float(pt[1])) for pt in longest] + ) + if candidate.is_valid and not candidate.is_empty: + return candidate + return bbox_poly + + +def _safe_get(arr: Any, idx: int) -> Optional[Any]: + """Return `arr[idx]` when `arr` is not None and `idx` is within range, + else `None`. Works on both numpy arrays and plain Python sequences.""" + if arr is None: + return None + try: + if len(arr) <= idx: + return None + value = arr[idx] + except (TypeError, IndexError): + return None + # numpy scalar -> Python scalar where applicable; keep strings as-is. + if hasattr(value, "item") and not isinstance(value, (bytes, str)): + try: + return value.item() + except (ValueError, TypeError): + return value + return value diff --git a/inference/core/workflows/core_steps/integrations/roboflow/visual_search_classifier/v1_tensor.py b/inference/core/workflows/core_steps/integrations/roboflow/visual_search_classifier/v1_tensor.py new file mode 100644 index 0000000000..6f545d9e3b --- /dev/null +++ b/inference/core/workflows/core_steps/integrations/roboflow/visual_search_classifier/v1_tensor.py @@ -0,0 +1,617 @@ +"""Tensor-native sibling of ``visual_search_classifier/v1.py``, loaded when +``ENABLE_TENSOR_DATA_REPRESENTATION`` is on. + +The block itself is an external-API integration (Roboflow project image +search) - the search flow, candidate formatting and all non-prediction outputs +are verbatim copies of the numpy source. The single tensor-native change is the +``predictions`` output: under the flag the classification kind carries native +``inference_models`` objects, and ``KINDS_SERIALIZERS`` swaps to +``serialise_native_classification`` (which raises on plain dicts), so this +block builds a native ``ClassificationPrediction`` / +``MultiLabelClassificationPrediction`` instead of the legacy response dict. + +Reconstruction follows the ``_native_classification_from_inference_response`` +precedent (models/roboflow/multi_class_classification/v1_tensor.py): a dense +confidence vector indexed by ``class_id`` with the matched class' confidence, +plus the ``class_id -> name`` map and lineage keys in the metadata. + +KNOWN serialized-output divergences vs the numpy sibling (surface before merge): + +- P1 (rounding): ``serialise_native_classification`` rounds confidences to 4 + decimal places and round-trips through float32; the numpy block emits the + raw float64 confidence unrounded. The native serializer convention matches + the model classification blocks, not this block's hand-built dict. +- P2 (single-label, zero confidence): when the matched candidate has no usable + ``score`` (confidence == 0.0) no serializer threshold can be attached, so + gap class ids ``0..class_id-1`` (named ``str(id)``) appear in the serialized + ``predictions`` list with confidence 0.0; numpy emits only the matched entry. + With a positive confidence the attached + ``classification_confidence_threshold`` filters the gap entries and the + output matches numpy's single-entry shape. +- P3 (multi-label): ``serialise_native_classification`` has no threshold on + the multi-label branch and emits an entry for EVERY dense-vector index, so + non-matched (gap) class ids appear with confidence 0.0; numpy emits only the + matched classes. ``predicted_classes`` matches numpy exactly. Fixing this + cleanly needs a serializer-side rule (e.g. skip ids absent from + ``class_names``), which affects other producers - decide there, not here. +""" + +import base64 +import math +from functools import partial +from typing import Any, Dict, List, Literal, Optional, Type, Union +from uuid import uuid4 + +import torch +from pydantic import ConfigDict, Field +from typing_extensions import Annotated + +from inference.core.env import ( + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.roboflow_api import ( + get_roboflow_workspace, + search_project_images_at_roboflow, +) +from inference.core.utils.image_utils import encode_image_to_jpeg_bytes +from inference.core.utils.preprocess import downscale_image_keeping_aspect_ratio +from inference.core.workflows.core_steps.common.utils import run_in_parallel +from inference.core.workflows.core_steps.integrations.roboflow.visual_search.helpers import ( + build_visual_search_candidate_image, + format_visual_search_candidate, +) +from inference.core.workflows.core_steps.integrations.roboflow.visual_search_classifier.classification_annotations import ( + parse_visual_search_classification, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_KEY, + CLASSIFICATION_STYLE_MODEL, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + DICTIONARY_KIND, + FLOAT_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + AirGappedAvailability, + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_project, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) + +# Same literal as models/roboflow/multi_class_classification/v1_tensor.py - the +# serializer reads it from image metadata to filter sub-threshold classes. +CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY = "classification_confidence_threshold" + +SHORT_DESCRIPTION = ( + "Classify an image by finding the most visually similar annotated image." +) + +LONG_DESCRIPTION = """ +Search a Roboflow classification project for the visually closest image and return +the matched image's classification annotation as a standard classification +prediction. Single-label annotations are returned in single-label classification +shape, and multi-label annotations are returned in multi-label classification +shape. Classification confidence is derived from the best candidate's public +Roboflow project search `score` field when present. + +This block performs an external visual search API call and is not intended for +real-time or high-throughput workloads. To bound latency, query images are +downscaled with aspect ratio preserved to a maximum side length of 224 pixels by +default. Increase `max_image_size` when finer visual matching is more important +than lower latency. + +## How This Block Works + +This block uses Roboflow project image search: + +1. Receives an input image from the workflow +2. Downscales the query image if its larger side exceeds `max_image_size` +3. Sends the query image to Roboflow project search as `image_base64` +4. Requests candidate image fields and classification labels or annotations +5. Uses the best candidate's annotation as the predicted class +6. Maps the API `score` field to classification confidence when present +7. Returns the result through the standard `predictions` classification output + +The target project must expose classification annotation data in visual search +results. This block does not train a model, create annotations, consume raw +search backend relevance scores, or manage the visual search index. +""" + +TopK = Annotated[int, Field(ge=1, le=50)] +MaxImageSize = Annotated[int, Field(ge=1)] + +VISUAL_SEARCH_CLASSIFIER_FIELDS = [ + "id", + "name", + "filename", + "url", + "user_metadata", + "tags", + "width", + "height", + "aspectRatio", + "score", + "labels", + "annotations", +] + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Roboflow Visual Search Classifier", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-tags", + "blockPriority": 4, + "inference": True, + "requires_rf_key": True, + }, + } + ) + type: Literal["roboflow_core/visual_search_classifier@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Image to classify using visual search.", + examples=["$inputs.image", "$steps.crop.crops"], + ) + target_project: Union[Selector(kind=[ROBOFLOW_PROJECT_KIND]), str] = Field( + description="Roboflow classification project URL slug to search.", + examples=["reference-images", "$inputs.target_project"], + ) + workspace: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + description=( + "Optional Roboflow workspace URL slug that owns the target project. " + "If not provided, the workspace is resolved from the request API key." + ), + examples=["my-workspace", "$inputs.workspace"], + ) + top_k: Union[TopK, Selector(kind=[INTEGER_KIND])] = Field( + default=1, + description=( + "Number of visually similar image candidates to request. The nearest " + "candidate is used for classification." + ), + examples=[1, 5, "$inputs.top_k"], + ) + max_image_size: Union[MaxImageSize, Selector(kind=[INTEGER_KIND])] = Field( + default=224, + description=( + "Maximum side length, in pixels, for the visual search query image. " + "Images larger than this are downscaled with aspect ratio preserved " + "before search. Increase this value for finer matching at higher " + "latency." + ), + examples=[224, 640, "$inputs.max_image_size"], + ) + + @classmethod + def get_air_gapped_availability(cls) -> AirGappedAvailability: + return AirGappedAvailability(available=False, reason="requires_internet") + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["image"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND] + ), + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition(name="candidate_found", kind=[BOOLEAN_KIND]), + OutputDefinition(name="class_found", kind=[BOOLEAN_KIND]), + OutputDefinition(name="best_candidate", kind=[DICTIONARY_KIND]), + OutputDefinition(name="candidates", kind=[LIST_OF_VALUES_KIND]), + OutputDefinition(name="best_candidate_image", kind=[IMAGE_KIND]), + OutputDefinition(name="visual_search_score", kind=[FLOAT_KIND]), + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="message", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + return [roboflow_platform_project(project_url=self.target_project)] + + +class RoboflowVisualSearchClassifierBlockV1(WorkflowBlock): + def __init__(self, api_key: Optional[str]): + self._api_key = api_key + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["api_key"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: Union[WorkflowImageData, Batch[WorkflowImageData]], + target_project: str, + workspace: Optional[str] = None, + top_k: int = 1, + max_image_size: int = 224, + ) -> BlockResult: + if self._api_key is None: + raise ValueError( + "Roboflow Visual Search Classifier block cannot run without a " + "Roboflow API key. Visit https://docs.roboflow.com/api-reference/" + "authentication#retrieve-an-api-key to learn how to retrieve one." + ) + + if isinstance(image, Batch): + tasks = [ + partial( + self._classify_single_image, + image=single_image, + workspace=workspace, + target_project=target_project, + top_k=top_k, + max_image_size=max_image_size, + ) + for single_image in image + ] + return run_in_parallel( + tasks=tasks, + max_workers=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + ) + + return self._classify_single_image( + image=image, + workspace=workspace, + target_project=target_project, + top_k=top_k, + max_image_size=max_image_size, + ) + + def _classify_single_image( + self, + image: WorkflowImageData, + target_project: str, + workspace: Optional[str], + top_k: int, + max_image_size: int, + ) -> Dict[str, Any]: + inference_id = f"{uuid4()}" + try: + workspace = self._resolve_workspace(workspace=workspace) + response = search_project_images_at_roboflow( + api_key=self._api_key, + workspace=workspace, + project=target_project, + image_base64=_prepare_query_image_base64( + image=image, max_image_size=max_image_size + ), + limit=top_k, + fields=VISUAL_SEARCH_CLASSIFIER_FIELDS, + ) + candidates = [ + _format_visual_search_classifier_candidate(candidate=candidate) + for candidate in response.get("results", []) + ] + if not candidates: + return _empty_result( + inference_id=inference_id, + error_status=False, + message="No visually similar images found.", + ) + + best_candidate = candidates[0] + classification = parse_visual_search_classification( + candidate=best_candidate + ) + if classification is None: + return _candidate_without_class_result( + inference_id=inference_id, + best_candidate=best_candidate, + candidates=candidates, + ) + + return { + "predictions": _build_classification_prediction( + image=image, + classification=classification, + inference_id=inference_id, + confidence=_normalise_visual_search_confidence( + score=best_candidate.get("score") + ), + ), + INFERENCE_ID_KEY: inference_id, + "candidate_found": True, + "class_found": True, + "best_candidate": best_candidate, + "candidates": candidates, + "best_candidate_image": build_visual_search_candidate_image( + candidate=best_candidate, + fallback_parent_id="visual_search_classifier_candidate", + ), + "visual_search_score": _normalise_visual_search_score( + score=best_candidate.get("score") + ), + "error_status": False, + "message": "Visual search classification completed.", + } + except Exception as error: + return _empty_result( + inference_id=inference_id, + error_status=True, + message=f"Visual search classification failed: {error}", + ) + + def _resolve_workspace(self, workspace: Optional[str]) -> str: + if workspace: + return workspace + return get_roboflow_workspace(api_key=self._api_key) + + +def _prepare_query_image_base64( + image: WorkflowImageData, + max_image_size: int, +) -> str: + # The search API needs JPEG bytes anyway, so a tensor-only input image is + # materialised to numpy here - a deliberate host copy, same as every + # remote-execution path. + source_image = image.numpy_image + query_image = downscale_image_keeping_aspect_ratio( + image=source_image, + desired_size=(max_image_size, max_image_size), + ) + if query_image.shape[:2] == source_image.shape[:2]: + return image.base64_image + return base64.b64encode( + encode_image_to_jpeg_bytes(query_image, jpeg_quality=95) + ).decode("ascii") + + +def _build_classification_prediction( + image: WorkflowImageData, + classification: Dict[str, Any], + inference_id: str, + confidence: float, +) -> Union[ClassificationPrediction, MultiLabelClassificationPrediction]: + height, width = image.numpy_image.shape[:2] + if classification["type"] == "multi_label": + return _build_multi_label_classification_prediction( + image=image, + classification=classification, + inference_id=inference_id, + width=width, + height=height, + confidence=confidence, + ) + + class_entry = classification["classes"][0] + class_name = class_entry["class"] + class_id = int(class_entry["class_id"]) + # Dense vector indexed by class_id; ids below the matched one are gap-filled + # (str(id) names, 0.0 confidence) - the full class list is not available on + # this path (same convention as the classifiers' remote-response converter). + num_classes = class_id + 1 + class_names = { + gap_class_id: str(gap_class_id) for gap_class_id in range(num_classes) + } + class_names[class_id] = str(class_name) + confidence_vector = [0.0] * num_classes + confidence_vector[class_id] = confidence + image_metadata = _build_image_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + width=width, + height=height, + ) + if confidence > 0.0: + # Filters the gap-filled zero-confidence classes at serialization so the + # output keeps numpy's single-entry shape (see P2 in the module docstring). + image_metadata[CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY] = confidence + return ClassificationPrediction( + class_id=torch.tensor( + [class_id], dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + # float64 for the same reason as the multi-label sibling: float32 would + # push the API-provided confidence below its own serialization + # threshold and off the numpy byte-parity value. + confidence=torch.tensor( + [confidence_vector], + dtype=torch.float64, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + images_metadata=[image_metadata], + ) + + +def _build_multi_label_classification_prediction( + image: WorkflowImageData, + classification: Dict[str, Any], + inference_id: str, + width: int, + height: int, + confidence: float, +) -> MultiLabelClassificationPrediction: + class_entries = classification["classes"] + num_classes = ( + max(int(class_entry["class_id"]) for class_entry in class_entries) + 1 + if class_entries + else 0 + ) + class_names = { + gap_class_id: str(gap_class_id) for gap_class_id in range(num_classes) + } + confidence_vector = [0.0] * num_classes + predicted_class_ids = [] + for class_entry in class_entries: + class_id = int(class_entry["class_id"]) + class_names[class_id] = str(class_entry["class"]) + confidence_vector[class_id] = confidence + predicted_class_ids.append(class_id) + image_metadata = _build_image_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + width=width, + height=height, + ) + if confidence > 0.0: + # Filters the gap-filled zero-confidence classes at serialization so the + # output keeps numpy's real-labels-only shape (see P2 in the module + # docstring) - same opt-in as the single-class path below. + image_metadata[CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY] = confidence + return MultiLabelClassificationPrediction( + class_ids=torch.tensor( + predicted_class_ids, dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + # float64: the confidence is an API-provided python float; float32 + # storage would shift it below its own serialization threshold and off + # the numpy byte-parity value (0.82 -> 0.8199999928...). + confidence=torch.tensor( + confidence_vector, + dtype=torch.float64, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + # MultiLabel carries SINGULAR image_metadata (dict) - see + # serialise_native_classification. + image_metadata=image_metadata, + ) + + +def _build_image_metadata( + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + width: int, + height: int, +) -> dict: + return { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag the serialiser reads to pick the flag-OFF shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_MODEL, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + ROOT_PARENT_ID_KEY: image.workflow_root_ancestor_metadata.parent_id, + } + + +def _format_visual_search_classifier_candidate( + candidate: Dict[str, Any], +) -> Dict[str, Any]: + formatted_candidate = format_visual_search_candidate( + candidate=candidate, + extra_fields=_visual_search_classifier_extra_fields(candidate=candidate), + ) + formatted_candidate["score"] = _normalise_visual_search_score( + score=candidate.get("score") + ) + return formatted_candidate + + +def _visual_search_classifier_extra_fields(candidate: Dict[str, Any]) -> List[str]: + extra_fields = ["score", "labels", "annotations"] + if "classification" in candidate: + extra_fields.append("classification") + return extra_fields + + +def _normalise_visual_search_score(score: Any) -> Optional[float]: + try: + raw_score = float(score) + except (TypeError, ValueError): + return None + if not math.isfinite(raw_score): + return None + return raw_score + + +def _normalise_visual_search_confidence(score: Any) -> float: + raw_score = _normalise_visual_search_score(score=score) + if raw_score is None: + return 0.0 + # Convert the project-search match score into a bounded confidence-like value. + return max(0.0, min(1.0, raw_score / 2.0)) + + +def _empty_result( + inference_id: str, + error_status: bool, + message: str, +) -> Dict[str, Any]: + return { + "predictions": None, + INFERENCE_ID_KEY: inference_id, + "candidate_found": False, + "class_found": False, + "best_candidate": {}, + "candidates": [], + "best_candidate_image": None, + "visual_search_score": None, + "error_status": error_status, + "message": message, + } + + +def _candidate_without_class_result( + inference_id: str, + best_candidate: Dict[str, Any], + candidates: List[Dict[str, Any]], +) -> Dict[str, Any]: + return { + "predictions": None, + INFERENCE_ID_KEY: inference_id, + "candidate_found": True, + "class_found": False, + "best_candidate": best_candidate, + "candidates": candidates, + "best_candidate_image": build_visual_search_candidate_image( + candidate=best_candidate, + fallback_parent_id="visual_search_classifier_candidate", + ), + "visual_search_score": _normalise_visual_search_score( + score=best_candidate.get("score") + ), + "error_status": True, + "message": ( + "Best visual search candidate does not include classification labels " + "or annotations." + ), + } diff --git a/inference/core/workflows/core_steps/loader.py b/inference/core/workflows/core_steps/loader.py index 5510ee9204..a16cd0b024 100644 --- a/inference/core/workflows/core_steps/loader.py +++ b/inference/core/workflows/core_steps/loader.py @@ -5,6 +5,7 @@ ALLOW_WORKFLOW_BLOCKS_ACCESSING_ENVIRONMENTAL_VARIABLES, ALLOW_WORKFLOW_BLOCKS_ACCESSING_LOCAL_STORAGE, API_KEY, + ENABLE_TENSOR_DATA_REPRESENTATION, SAM3_3D_OBJECTS_ENABLED, WORKFLOW_BLOCKS_WRITE_DIRECTORY, WORKFLOW_DISABLED_BLOCK_PATTERNS, @@ -14,32 +15,74 @@ from inference.core.workflows.core_steps.analytics.data_aggregator.v1 import ( DataAggregatorBlockV1, ) -from inference.core.workflows.core_steps.analytics.detection_event_log.v1 import ( - DetectionEventLogBlockV1, -) -from inference.core.workflows.core_steps.analytics.line_counter.v1 import ( - LineCounterBlockV1, -) -from inference.core.workflows.core_steps.analytics.line_counter.v2 import ( - LineCounterBlockV2, -) -from inference.core.workflows.core_steps.analytics.overlap.v1 import OverlapBlockV1 -from inference.core.workflows.core_steps.analytics.path_deviation.v1 import ( - PathDeviationAnalyticsBlockV1, -) -from inference.core.workflows.core_steps.analytics.path_deviation.v2 import ( - PathDeviationAnalyticsBlockV2, -) -from inference.core.workflows.core_steps.analytics.time_in_zone.v1 import ( - TimeInZoneBlockV1, -) -from inference.core.workflows.core_steps.analytics.time_in_zone.v2 import ( - TimeInZoneBlockV2, -) -from inference.core.workflows.core_steps.analytics.time_in_zone.v3 import ( - TimeInZoneBlockV3, -) -from inference.core.workflows.core_steps.analytics.velocity.v1 import VelocityBlockV1 + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.analytics.detection_event_log.v1 import ( + DetectionEventLogBlockV1, + ) +else: + from inference.core.workflows.core_steps.analytics.detection_event_log.v1_tensor import ( + DetectionEventLogBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.analytics.line_counter.v1 import ( + LineCounterBlockV1, + ) +else: + from inference.core.workflows.core_steps.analytics.line_counter.v1_tensor import ( + LineCounterBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.analytics.line_counter.v2 import ( + LineCounterBlockV2, + ) +else: + from inference.core.workflows.core_steps.analytics.line_counter.v2_tensor import ( + LineCounterBlockV2, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.analytics.overlap.v1 import OverlapBlockV1 + from inference.core.workflows.core_steps.analytics.path_deviation.v1 import ( + PathDeviationAnalyticsBlockV1, + ) + from inference.core.workflows.core_steps.analytics.path_deviation.v2 import ( + PathDeviationAnalyticsBlockV2, + ) + from inference.core.workflows.core_steps.analytics.time_in_zone.v1 import ( + TimeInZoneBlockV1, + ) + from inference.core.workflows.core_steps.analytics.time_in_zone.v2 import ( + TimeInZoneBlockV2, + ) + from inference.core.workflows.core_steps.analytics.time_in_zone.v3 import ( + TimeInZoneBlockV3, + ) + from inference.core.workflows.core_steps.analytics.velocity.v1 import ( + VelocityBlockV1, + ) +else: + from inference.core.workflows.core_steps.analytics.overlap.v1_tensor import ( + OverlapBlockV1, + ) + from inference.core.workflows.core_steps.analytics.path_deviation.v1_tensor import ( + PathDeviationAnalyticsBlockV1, + ) + from inference.core.workflows.core_steps.analytics.path_deviation.v2_tensor import ( + PathDeviationAnalyticsBlockV2, + ) + from inference.core.workflows.core_steps.analytics.time_in_zone.v1_tensor import ( + TimeInZoneBlockV1, + ) + from inference.core.workflows.core_steps.analytics.time_in_zone.v2_tensor import ( + TimeInZoneBlockV2, + ) + from inference.core.workflows.core_steps.analytics.time_in_zone.v3_tensor import ( + TimeInZoneBlockV3, + ) + from inference.core.workflows.core_steps.analytics.velocity.v1_tensor import ( + VelocityBlockV1, + ) + from inference.core.workflows.core_steps.cache.cache_get.v1 import CacheGetBlockV1 from inference.core.workflows.core_steps.cache.cache_set.v1 import CacheSetBlockV1 from inference.core.workflows.core_steps.classical_cv.auto_rotate_on_edges.v1 import ( @@ -57,48 +100,115 @@ from inference.core.workflows.core_steps.classical_cv.contours.v1 import ( ImageContoursDetectionBlockV1, ) -from inference.core.workflows.core_steps.classical_cv.contrast_enhancement.v1 import ( - ContrastEnhancementBlock, -) -from inference.core.workflows.core_steps.classical_cv.contrast_equalization.v1 import ( - ContrastEqualizationBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.convert_grayscale.v1 import ( - ConvertGrayscaleBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.detections_nearest_neighbor.v1 import ( - DetectionsNearestNeighborBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.distance_measurement.v1 import ( - DistanceMeasurementBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.dominant_color.v1 import ( - DominantColorBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.image_blur.v1 import ( - ImageBlurBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.image_preprocessing.v1 import ( - ImagePreprocessingBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.mask_area_measurement.v1 import ( - MaskAreaMeasurementBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.mask_edge_snap.v1 import ( - MaskEdgeSnapBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.contrast_enhancement.v1 import ( + ContrastEnhancementBlock, + ) +else: + from inference.core.workflows.core_steps.classical_cv.contrast_enhancement.v1_tensor import ( + ContrastEnhancementBlock, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.contrast_equalization.v1 import ( + ContrastEqualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.contrast_equalization.v1_tensor import ( + ContrastEqualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.convert_grayscale.v1 import ( + ConvertGrayscaleBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.convert_grayscale.v1_tensor import ( + ConvertGrayscaleBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.detections_nearest_neighbor.v1 import ( + DetectionsNearestNeighborBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.detections_nearest_neighbor.v1_tensor import ( + DetectionsNearestNeighborBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.distance_measurement.v1 import ( + DistanceMeasurementBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.distance_measurement.v1_tensor import ( + DistanceMeasurementBlockV1, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.dominant_color.v1 import ( + DominantColorBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.dominant_color.v1_tensor import ( + DominantColorBlockV1, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.image_blur.v1 import ( + ImageBlurBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.image_blur.v1_tensor import ( + ImageBlurBlockV1, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.image_preprocessing.v1 import ( + ImagePreprocessingBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.image_preprocessing.v1_tensor import ( + ImagePreprocessingBlockV1, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.mask_area_measurement.v1 import ( + MaskAreaMeasurementBlockV1, + ) + from inference.core.workflows.core_steps.classical_cv.mask_edge_snap.v1 import ( + MaskEdgeSnapBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.mask_area_measurement.v1_tensor import ( + MaskAreaMeasurementBlockV1, + ) + from inference.core.workflows.core_steps.classical_cv.mask_edge_snap.v1_tensor import ( + MaskEdgeSnapBlockV1, + ) + from inference.core.workflows.core_steps.classical_cv.morphological_transformation.v1 import ( MorphologicalTransformationBlockV1, ) from inference.core.workflows.core_steps.classical_cv.morphological_transformation.v2 import ( MorphologicalTransformationBlockV2, ) -from inference.core.workflows.core_steps.classical_cv.motion_detection.v1 import ( - MotionDetectionBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.pixel_color_count.v1 import ( - PixelationCountBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.motion_detection.v1 import ( + MotionDetectionBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.motion_detection.v1_tensor import ( + MotionDetectionBlockV1, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.pixel_color_count.v1 import ( + PixelationCountBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.pixel_color_count.v1_tensor import ( + PixelationCountBlockV1, + ) + from inference.core.workflows.core_steps.classical_cv.sift.v1 import SIFTBlockV1 from inference.core.workflows.core_steps.classical_cv.sift_comparison.v1 import ( SIFTComparisonBlockV1, @@ -106,24 +216,38 @@ from inference.core.workflows.core_steps.classical_cv.sift_comparison.v2 import ( SIFTComparisonBlockV2, ) -from inference.core.workflows.core_steps.classical_cv.size_measurement.v1 import ( - SizeMeasurementBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.template_matching.v1 import ( - TemplateMatchingBlockV1, -) -from inference.core.workflows.core_steps.classical_cv.threshold.v1 import ( - ImageThresholdBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.size_measurement.v1 import ( + SizeMeasurementBlockV1, + ) + from inference.core.workflows.core_steps.classical_cv.template_matching.v1 import ( + TemplateMatchingBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.size_measurement.v1_tensor import ( + SizeMeasurementBlockV1, + ) + from inference.core.workflows.core_steps.classical_cv.template_matching.v1_tensor import ( + TemplateMatchingBlockV1, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.classical_cv.threshold.v1 import ( + ImageThresholdBlockV1, + ) +else: + from inference.core.workflows.core_steps.classical_cv.threshold.v1_tensor import ( + ImageThresholdBlockV1, + ) + from inference.core.workflows.core_steps.common.deserializers import ( deserialize_boolean_kind, deserialize_bytes_kind, deserialize_classification_prediction_kind, - deserialize_detections_kind, deserialize_dictionary_kind, deserialize_float_kind, deserialize_float_zero_to_one_kind, - deserialize_image_kind, deserialize_integer_kind, deserialize_labeled_points_kind, deserialize_list_of_values_kind, @@ -131,7 +255,6 @@ deserialize_optional_string_kind, deserialize_point_kind, deserialize_rgb_color_kind, - deserialize_rle_detections_kind, deserialize_string_kind, deserialize_timestamp, deserialize_video_metadata_kind, @@ -140,13 +263,56 @@ from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.core_steps.common.serializers import ( serialise_image, - serialise_rle_sv_detections, - serialise_sv_detections, serialize_secret, serialize_timestamp, serialize_video_metadata_kind, - serialize_wildcard_kind, ) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_KIND, + TENSOR_NATIVE_BAR_CODE_DETECTION_KIND, + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + TENSOR_NATIVE_DETECTION_KIND, + TENSOR_NATIVE_EMBEDDING_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_QR_CODE_DETECTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_SEMANTIC_SEGMENTATION_PREDICTION_KIND, +) + +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.common.deserializers_tensor import ( + deserialize_detections_kind, + deserialize_image_kind, + deserialize_native_classification_prediction_kind, + deserialize_native_embedding_kind, + deserialize_native_tensor_kind, + deserialize_rle_detections_kind, + ) + from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_native_classification, + serialise_native_embedding, + serialise_native_keypoint_detection, + serialise_native_rle_detections, + serialise_native_tensor, + serialise_numpy_array_kind, + serialise_rle_sv_detections, + serialise_sv_detections, + serialize_wildcard_kind, + ) +else: + from inference.core.workflows.core_steps.common.deserializers import ( + deserialize_detections_kind, + deserialize_image_kind, + deserialize_rle_detections_kind, + ) + from inference.core.workflows.core_steps.common.serializers import ( + serialise_rle_sv_detections, + serialise_sv_detections, + serialize_wildcard_kind, + ) + from inference.core.workflows.core_steps.flow_control.continue_if.v1 import ( ContinueIfBlockV1, ) @@ -178,48 +344,121 @@ from inference.core.workflows.core_steps.formatters.property_definition.v1 import ( PropertyDefinitionBlockV1, ) -from inference.core.workflows.core_steps.formatters.vlm_as_classifier.v1 import ( - VLMAsClassifierBlockV1, -) -from inference.core.workflows.core_steps.formatters.vlm_as_classifier.v2 import ( - VLMAsClassifierBlockV2, -) -from inference.core.workflows.core_steps.formatters.vlm_as_detector.v1 import ( - VLMAsDetectorBlockV1, -) -from inference.core.workflows.core_steps.formatters.vlm_as_detector.v2 import ( - VLMAsDetectorBlockV2, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.formatters.vlm_as_classifier.v1 import ( + VLMAsClassifierBlockV1, + ) +else: + from inference.core.workflows.core_steps.formatters.vlm_as_classifier.v1_tensor import ( + VLMAsClassifierBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.formatters.vlm_as_classifier.v2 import ( + VLMAsClassifierBlockV2, + ) +else: + from inference.core.workflows.core_steps.formatters.vlm_as_classifier.v2_tensor import ( + VLMAsClassifierBlockV2, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.formatters.vlm_as_detector.v1 import ( + VLMAsDetectorBlockV1, + ) +else: + from inference.core.workflows.core_steps.formatters.vlm_as_detector.v1_tensor import ( + VLMAsDetectorBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.formatters.vlm_as_detector.v2 import ( + VLMAsDetectorBlockV2, + ) +else: + from inference.core.workflows.core_steps.formatters.vlm_as_detector.v2_tensor import ( + VLMAsDetectorBlockV2, + ) + from inference.core.workflows.core_steps.fusion.buffer.v1 import BufferBlockV1 -from inference.core.workflows.core_steps.fusion.detections_classes_replacement.v1 import ( - DetectionsClassesReplacementBlockV1, -) -from inference.core.workflows.core_steps.fusion.detections_consensus.v1 import ( - DetectionsConsensusBlockV1, -) -from inference.core.workflows.core_steps.fusion.detections_list_rollup.v1 import ( - DetectionsListRollUpBlockV1, -) -from inference.core.workflows.core_steps.fusion.detections_stitch.v1 import ( - DetectionsStitchBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.fusion.detections_classes_replacement.v1 import ( + DetectionsClassesReplacementBlockV1, + ) +else: + from inference.core.workflows.core_steps.fusion.detections_classes_replacement.v1_tensor import ( + DetectionsClassesReplacementBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.fusion.detections_consensus.v1 import ( + DetectionsConsensusBlockV1, + ) +else: + from inference.core.workflows.core_steps.fusion.detections_consensus.v1_tensor import ( + DetectionsConsensusBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.fusion.detections_list_rollup.v1 import ( + DetectionsListRollUpBlockV1, + ) +else: + from inference.core.workflows.core_steps.fusion.detections_list_rollup.v1_tensor import ( + DetectionsListRollUpBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.fusion.detections_stitch.v1 import ( + DetectionsStitchBlockV1, + ) +else: + from inference.core.workflows.core_steps.fusion.detections_stitch.v1_tensor import ( + DetectionsStitchBlockV1, + ) + from inference.core.workflows.core_steps.fusion.dimension_collapse.v1 import ( DimensionCollapseBlockV1, ) -from inference.core.workflows.core_steps.fusion.frame_delay.v1 import FrameDelayBlockV1 + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.fusion.frame_delay.v1 import ( + FrameDelayBlockV1, + ) +else: + from inference.core.workflows.core_steps.fusion.frame_delay.v1_tensor import ( + FrameDelayBlockV1, + ) + from inference.core.workflows.core_steps.fusion.image_stack.v1 import ImageStackBlockV1 -from inference.core.workflows.core_steps.fusion.overlap_analysis.v1 import ( - OverlapAnalysisBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.fusion.overlap_analysis.v1 import ( + OverlapAnalysisBlockV1, + ) +else: + from inference.core.workflows.core_steps.fusion.overlap_analysis.v1_tensor import ( + OverlapAnalysisBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.math.cosine_similarity.v1 import ( + CosineSimilarityBlockV1, + ) +else: + from inference.core.workflows.core_steps.math.cosine_similarity.v1_tensor import ( + CosineSimilarityBlockV1, + ) + +# visual_search emits only dict/scalar/image outputs, so it needs no _tensor sibling. from inference.core.workflows.core_steps.integrations.roboflow.visual_search.v1 import ( RoboflowVisualSearchBlockV1, ) -from inference.core.workflows.core_steps.integrations.roboflow.visual_search_classifier.v1 import ( - RoboflowVisualSearchClassifierBlockV1, -) -from inference.core.workflows.core_steps.math.cosine_similarity.v1 import ( - CosineSimilarityBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.integrations.roboflow.visual_search_classifier.v1 import ( + RoboflowVisualSearchClassifierBlockV1, + ) +else: + from inference.core.workflows.core_steps.integrations.roboflow.visual_search_classifier.v1_tensor import ( + RoboflowVisualSearchClassifierBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.anthropic_claude.v1 import ( AnthropicClaudeBlockV1, ) @@ -229,33 +468,70 @@ from inference.core.workflows.core_steps.models.foundation.anthropic_claude.v3 import ( AnthropicClaudeBlockV3, ) -from inference.core.workflows.core_steps.models.foundation.clip.v1 import ( - ClipModelBlockV1, -) -from inference.core.workflows.core_steps.models.foundation.clip_comparison.v1 import ( - ClipComparisonBlockV1, -) -from inference.core.workflows.core_steps.models.foundation.clip_comparison.v2 import ( - ClipComparisonBlockV2, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.clip.v1 import ( + ClipModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.clip.v1_tensor import ( + ClipModelBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.clip_comparison.v1 import ( + ClipComparisonBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.clip_comparison.v2 import ( + ClipComparisonBlockV2, + ) +else: + from inference.core.workflows.core_steps.models.foundation.clip_comparison.v1_tensor import ( + ClipComparisonBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.clip_comparison.v2_tensor import ( + ClipComparisonBlockV2, + ) + from inference.core.workflows.core_steps.models.foundation.cog_vlm.v1 import ( CogVLMBlockV1, ) from inference.core.workflows.core_steps.models.foundation.cosmos3.v1 import ( Cosmos3EdgeBlockV1, ) -from inference.core.workflows.core_steps.models.foundation.depth_estimation.v1 import ( - DepthEstimationBlockV1, -) -from inference.core.workflows.core_steps.models.foundation.easy_ocr.v1 import ( - EasyOCRBlockV1, -) -from inference.core.workflows.core_steps.models.foundation.florence2.v1 import ( - Florence2BlockV1, -) -from inference.core.workflows.core_steps.models.foundation.florence2.v2 import ( - Florence2BlockV2, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.depth_estimation.v1 import ( + DepthEstimationBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.depth_estimation.v1_tensor import ( + DepthEstimationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.easy_ocr.v1 import ( + EasyOCRBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.easy_ocr.v1_tensor import ( + EasyOCRBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.florence2.v1 import ( + Florence2BlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.florence2.v1_tensor import ( + Florence2BlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.florence2.v2 import ( + Florence2BlockV2, + ) +else: + from inference.core.workflows.core_steps.models.foundation.florence2.v2_tensor import ( + Florence2BlockV2, + ) + from inference.core.workflows.core_steps.models.foundation.gaze.v1 import GazeBlockV1 from inference.core.workflows.core_steps.models.foundation.glm_ocr.v1 import ( GLMOCRBlockV1, @@ -278,9 +554,16 @@ from inference.core.workflows.core_steps.models.foundation.google_gemma.v2 import ( GoogleGemmaBlockV2, ) -from inference.core.workflows.core_steps.models.foundation.google_vision_ocr.v1 import ( - GoogleVisionOCRBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.google_vision_ocr.v1 import ( + GoogleVisionOCRBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.google_vision_ocr.v1_tensor import ( + GoogleVisionOCRBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.kimi_openrouter.v1 import ( KimiOpenRouterBlockV1, ) @@ -297,10 +580,24 @@ from inference.core.workflows.core_steps.models.foundation.lmm_classifier.v1 import ( LMMForClassificationBlockV1, ) -from inference.core.workflows.core_steps.models.foundation.moondream2.v1 import ( - Moondream2BlockV1, -) -from inference.core.workflows.core_steps.models.foundation.ocr.v1 import OCRModelBlockV1 + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.moondream2.v1 import ( + Moondream2BlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.moondream2.v1_tensor import ( + Moondream2BlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.ocr.v1 import ( + OCRModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.ocr.v1_tensor import ( + OCRModelBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.openai.v1 import ( OpenAIBlockV1, ) @@ -319,142 +616,344 @@ from inference.core.workflows.core_steps.models.foundation.openrouter.v1 import ( OpenRouterBlockV1, ) -from inference.core.workflows.core_steps.models.foundation.perception_encoder.v1 import ( - PerceptionEncoderModelBlockV1, -) -from inference.core.workflows.core_steps.models.foundation.pp_ocr.v1 import PPOCRBlockV1 + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.perception_encoder.v1 import ( + PerceptionEncoderModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.perception_encoder.v1_tensor import ( + PerceptionEncoderModelBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.pp_ocr.v1 import ( + PPOCRBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.pp_ocr.v1_tensor import ( + PPOCRBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.qwen3_5_openrouter.v1 import ( Qwen35OpenRouterBlockV1, ) -from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1 import ( - Qwen35VLBlockV1, -) -from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v2 import ( - Qwen35VLBlockV2, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1 import ( + Qwen35VLBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1_tensor import ( + Qwen35VLBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v2 import ( + Qwen35VLBlockV2, + ) +else: + from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v2_tensor import ( + Qwen35VLBlockV2, + ) + from inference.core.workflows.core_steps.models.foundation.qwen3_6_openrouter.v1 import ( Qwen36OpenRouterBlockV1, ) -from inference.core.workflows.core_steps.models.foundation.qwen3vl.v1 import ( - Qwen3VLBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.qwen3vl.v1 import ( + Qwen3VLBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.qwen3vl.v1_tensor import ( + Qwen3VLBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.qwen.v1 import ( Qwen25VLBlockV1, ) from inference.core.workflows.core_steps.models.foundation.qwen_vlm.v1 import ( QwenVlmBlockV1, ) -from inference.core.workflows.core_steps.models.foundation.seg_preview.v1 import ( - SegPreviewBlockV1, -) -from inference.core.workflows.core_steps.models.foundation.segment_anything2.v1 import ( - SegmentAnything2BlockV1, -) -from inference.core.workflows.core_steps.models.foundation.segment_anything2_video.v1 import ( - SegmentAnything2VideoBlockV1, -) -from inference.core.workflows.core_steps.models.foundation.segment_anything3.v1 import ( - SegmentAnything3BlockV1, -) -from inference.core.workflows.core_steps.models.foundation.segment_anything3.v2 import ( - SegmentAnything3BlockV2, -) -from inference.core.workflows.core_steps.models.foundation.segment_anything3.v3 import ( - SegmentAnything3BlockV3, -) -from inference.core.workflows.core_steps.models.foundation.segment_anything3_interactive.v1 import ( - SegmentAnything3InteractiveBlockV1, -) -from inference.core.workflows.core_steps.models.foundation.segment_anything3_video.v1 import ( - SegmentAnything3VideoBlockV1, -) -from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v4 import ( - RoboflowInstanceSegmentationModelBlockV4, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.seg_preview.v1 import ( + SegPreviewBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.seg_preview.v1_tensor import ( + SegPreviewBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.segment_anything2.v1 import ( + SegmentAnything2BlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.segment_anything2.v1_tensor import ( + SegmentAnything2BlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.segment_anything2_video.v1 import ( + SegmentAnything2VideoBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.segment_anything2_video.v1_tensor import ( + SegmentAnything2VideoBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.segment_anything3.v1 import ( + SegmentAnything3BlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.segment_anything3.v2 import ( + SegmentAnything3BlockV2, + ) + from inference.core.workflows.core_steps.models.foundation.segment_anything3.v3 import ( + SegmentAnything3BlockV3, + ) +else: + from inference.core.workflows.core_steps.models.foundation.segment_anything3.v1_tensor import ( + SegmentAnything3BlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.segment_anything3.v2_tensor import ( + SegmentAnything3BlockV2, + ) + from inference.core.workflows.core_steps.models.foundation.segment_anything3.v3_tensor import ( + SegmentAnything3BlockV3, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v4_tensor import ( + RoboflowInstanceSegmentationModelBlockV4, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v4 import ( + RoboflowInstanceSegmentationModelBlockV4, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.segment_anything3_interactive.v1 import ( + SegmentAnything3InteractiveBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.segment_anything3_video.v1 import ( + SegmentAnything3VideoBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.segment_anything3_interactive.v1_tensor import ( + SegmentAnything3InteractiveBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.segment_anything3_video.v1_tensor import ( + SegmentAnything3VideoBlockV1, + ) if SAM3_3D_OBJECTS_ENABLED: - from inference.core.workflows.core_steps.models.foundation.segment_anything3_3d.v1 import ( - SegmentAnything3_3D_ObjectsBlockV1, + if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.segment_anything3_3d.v1 import ( + SegmentAnything3_3D_ObjectsBlockV1, + ) + else: + from inference.core.workflows.core_steps.models.foundation.segment_anything3_3d.v1_tensor import ( + SegmentAnything3_3D_ObjectsBlockV1, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.smolvlm.v1 import ( + SmolVLM2BlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.smolvlm.v1_tensor import ( + SmolVLM2BlockV1, ) -from inference.core.workflows.core_steps.models.foundation.smolvlm.v1 import ( - SmolVLM2BlockV1, -) from inference.core.workflows.core_steps.models.foundation.stability_ai.image_gen.v1 import ( StabilityAIImageGenBlockV1, ) -from inference.core.workflows.core_steps.models.foundation.stability_ai.inpainting.v1 import ( - StabilityAIInpaintingBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.stability_ai.inpainting.v1 import ( + StabilityAIInpaintingBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.stability_ai.inpainting.v1_tensor import ( + StabilityAIInpaintingBlockV1, + ) + from inference.core.workflows.core_steps.models.foundation.stability_ai.outpainting.v1 import ( StabilityAIOutpaintingBlockV1, ) -from inference.core.workflows.core_steps.models.foundation.yolo_world.v1 import ( - YoloWorldModelBlockV1, -) -from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v1 import ( - RoboflowInstanceSegmentationModelBlockV1, -) -from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v2 import ( - RoboflowInstanceSegmentationModelBlockV2, -) -from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v3 import ( - RoboflowInstanceSegmentationModelBlockV3, -) -from inference.core.workflows.core_steps.models.roboflow.keypoint_detection.v1 import ( - RoboflowKeypointDetectionModelBlockV1, -) -from inference.core.workflows.core_steps.models.roboflow.keypoint_detection.v2 import ( - RoboflowKeypointDetectionModelBlockV2, -) -from inference.core.workflows.core_steps.models.roboflow.keypoint_detection.v3 import ( - RoboflowKeypointDetectionModelBlockV3, -) -from inference.core.workflows.core_steps.models.roboflow.multi_class_classification.v1 import ( - RoboflowClassificationModelBlockV1, -) -from inference.core.workflows.core_steps.models.roboflow.multi_class_classification.v2 import ( - RoboflowClassificationModelBlockV2, -) -from inference.core.workflows.core_steps.models.roboflow.multi_class_classification.v3 import ( - RoboflowClassificationModelBlockV3, -) -from inference.core.workflows.core_steps.models.roboflow.multi_label_classification.v1 import ( - RoboflowMultiLabelClassificationModelBlockV1, -) -from inference.core.workflows.core_steps.models.roboflow.multi_label_classification.v2 import ( - RoboflowMultiLabelClassificationModelBlockV2, -) -from inference.core.workflows.core_steps.models.roboflow.multi_label_classification.v3 import ( - RoboflowMultiLabelClassificationModelBlockV3, -) -from inference.core.workflows.core_steps.models.roboflow.object_detection.v1 import ( - RoboflowObjectDetectionModelBlockV1, -) -from inference.core.workflows.core_steps.models.roboflow.object_detection.v2 import ( - RoboflowObjectDetectionModelBlockV2, -) -from inference.core.workflows.core_steps.models.roboflow.object_detection.v3 import ( - RoboflowObjectDetectionModelBlockV3, -) -from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v1 import ( - RoboflowSemanticSegmentationModelBlockV1, -) -from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v2 import ( - RoboflowSemanticSegmentationModelBlockV2, -) -from inference.core.workflows.core_steps.models.third_party.barcode_detection.v1 import ( - BarcodeDetectorBlockV1, -) -from inference.core.workflows.core_steps.models.third_party.qr_code_detection.v1 import ( - QRCodeDetectorBlockV1, -) -from inference.core.workflows.core_steps.sampling.identify_changes.v1 import ( - IdentifyChangesBlockV1, -) -from inference.core.workflows.core_steps.sampling.identify_outliers.v1 import ( - IdentifyOutliersBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.foundation.yolo_world.v1 import ( + YoloWorldModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.foundation.yolo_world.v1_tensor import ( + YoloWorldModelBlockV1, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v1_tensor import ( + RoboflowInstanceSegmentationModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v1 import ( + RoboflowInstanceSegmentationModelBlockV1, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v2_tensor import ( + RoboflowInstanceSegmentationModelBlockV2, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v2 import ( + RoboflowInstanceSegmentationModelBlockV2, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v3_tensor import ( + RoboflowInstanceSegmentationModelBlockV3, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.instance_segmentation.v3 import ( + RoboflowInstanceSegmentationModelBlockV3, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.keypoint_detection.v1_tensor import ( + RoboflowKeypointDetectionModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.keypoint_detection.v1 import ( + RoboflowKeypointDetectionModelBlockV1, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.keypoint_detection.v2_tensor import ( + RoboflowKeypointDetectionModelBlockV2, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.keypoint_detection.v2 import ( + RoboflowKeypointDetectionModelBlockV2, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.keypoint_detection.v3_tensor import ( + RoboflowKeypointDetectionModelBlockV3, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.keypoint_detection.v3 import ( + RoboflowKeypointDetectionModelBlockV3, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.multi_class_classification.v1_tensor import ( + RoboflowClassificationModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.multi_class_classification.v1 import ( + RoboflowClassificationModelBlockV1, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.multi_class_classification.v2_tensor import ( + RoboflowClassificationModelBlockV2, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.multi_class_classification.v2 import ( + RoboflowClassificationModelBlockV2, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.multi_class_classification.v3_tensor import ( + RoboflowClassificationModelBlockV3, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.multi_class_classification.v3 import ( + RoboflowClassificationModelBlockV3, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.multi_label_classification.v1_tensor import ( + RoboflowMultiLabelClassificationModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.multi_label_classification.v1 import ( + RoboflowMultiLabelClassificationModelBlockV1, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.multi_label_classification.v2_tensor import ( + RoboflowMultiLabelClassificationModelBlockV2, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.multi_label_classification.v2 import ( + RoboflowMultiLabelClassificationModelBlockV2, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.multi_label_classification.v3_tensor import ( + RoboflowMultiLabelClassificationModelBlockV3, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.multi_label_classification.v3 import ( + RoboflowMultiLabelClassificationModelBlockV3, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.object_detection.v1_tensor import ( + RoboflowObjectDetectionModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.object_detection.v1 import ( + RoboflowObjectDetectionModelBlockV1, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.object_detection.v2_tensor import ( + RoboflowObjectDetectionModelBlockV2, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.object_detection.v2 import ( + RoboflowObjectDetectionModelBlockV2, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.object_detection.v3_tensor import ( + RoboflowObjectDetectionModelBlockV3, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.object_detection.v3 import ( + RoboflowObjectDetectionModelBlockV3, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v1_tensor import ( + RoboflowSemanticSegmentationModelBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v1 import ( + RoboflowSemanticSegmentationModelBlockV1, + ) +if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v2_tensor import ( + RoboflowSemanticSegmentationModelBlockV2, + ) +else: + from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v2 import ( + RoboflowSemanticSegmentationModelBlockV2, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.third_party.barcode_detection.v1 import ( + BarcodeDetectorBlockV1, + ) + from inference.core.workflows.core_steps.models.third_party.qr_code_detection.v1 import ( + QRCodeDetectorBlockV1, + ) +else: + from inference.core.workflows.core_steps.models.third_party.barcode_detection.v1_tensor import ( + BarcodeDetectorBlockV1, + ) + from inference.core.workflows.core_steps.models.third_party.qr_code_detection.v1_tensor import ( + QRCodeDetectorBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.sampling.identify_changes.v1 import ( + IdentifyChangesBlockV1, + ) +else: + from inference.core.workflows.core_steps.sampling.identify_changes.v1_tensor import ( + IdentifyChangesBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.sampling.identify_outliers.v1 import ( + IdentifyOutliersBlockV1, + ) +else: + from inference.core.workflows.core_steps.sampling.identify_outliers.v1_tensor import ( + IdentifyOutliersBlockV1, + ) + from inference.core.workflows.core_steps.secrets_providers.environment_secrets_store.v1 import ( EnvironmentSecretsStoreBlockV1, ) @@ -465,25 +964,61 @@ EmailNotificationBlockV2, ) from inference.core.workflows.core_steps.sinks.local_file.v1 import LocalFileSinkBlockV1 -from inference.core.workflows.core_steps.sinks.onvif_movement.v1 import ONVIFSinkBlockV1 + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.sinks.onvif_movement.v1 import ( + ONVIFSinkBlockV1, + ) +else: + from inference.core.workflows.core_steps.sinks.onvif_movement.v1_tensor import ( + ONVIFSinkBlockV1, + ) + from inference.core.workflows.core_steps.sinks.roboflow.asset_library_attributes.v1 import ( RoboflowAssetLibraryAttributesBlockV1, ) -from inference.core.workflows.core_steps.sinks.roboflow.custom_metadata.v1 import ( - RoboflowCustomMetadataBlockV1, -) -from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1 import ( - RoboflowDatasetUploadBlockV1, -) -from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v2 import ( - RoboflowDatasetUploadBlockV2, -) -from inference.core.workflows.core_steps.sinks.roboflow.model_monitoring_inference_aggregator.v1 import ( - ModelMonitoringInferenceAggregatorBlockV1, -) -from inference.core.workflows.core_steps.sinks.roboflow.vision_events.v1 import ( - RoboflowVisionEventsBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.sinks.roboflow.custom_metadata.v1 import ( + RoboflowCustomMetadataBlockV1, + ) +else: + from inference.core.workflows.core_steps.sinks.roboflow.custom_metadata.v1_tensor import ( + RoboflowCustomMetadataBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1 import ( + RoboflowDatasetUploadBlockV1, + ) +else: + from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1_tensor import ( + RoboflowDatasetUploadBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v2 import ( + RoboflowDatasetUploadBlockV2, + ) +else: + from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v2_tensor import ( + RoboflowDatasetUploadBlockV2, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.sinks.roboflow.model_monitoring_inference_aggregator.v1 import ( + ModelMonitoringInferenceAggregatorBlockV1, + ) +else: + from inference.core.workflows.core_steps.sinks.roboflow.model_monitoring_inference_aggregator.v1_tensor import ( + ModelMonitoringInferenceAggregatorBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.sinks.roboflow.vision_events.v1 import ( + RoboflowVisionEventsBlockV1, + ) +else: + from inference.core.workflows.core_steps.sinks.roboflow.vision_events.v1_tensor import ( + RoboflowVisionEventsBlockV1, + ) + from inference.core.workflows.core_steps.sinks.s3.v1 import S3SinkBlockV1 from inference.core.workflows.core_steps.sinks.slack.notification.v1 import ( SlackNotificationBlockV1, @@ -495,185 +1030,439 @@ TwilioSMSNotificationBlockV2, ) from inference.core.workflows.core_steps.sinks.webhook.v1 import WebhookSinkBlockV1 -from inference.core.workflows.core_steps.trackers.botsort.v1 import ( - BoTSORTBlockV1 as TrackerBoTSORTBlockV1, -) -from inference.core.workflows.core_steps.trackers.bytetrack.v1 import ( - ByteTrackBlockV1 as TrackerByteTrackBlockV1, -) -from inference.core.workflows.core_steps.trackers.ocsort.v1 import ( - OCSORTBlockV1 as TrackerOCSORTBlockV1, -) -from inference.core.workflows.core_steps.trackers.sort.v1 import ( - SORTBlockV1 as TrackerSORTBlockV1, -) -from inference.core.workflows.core_steps.transformations.absolute_static_crop.v1 import ( - AbsoluteStaticCropBlockV1, -) -from inference.core.workflows.core_steps.transformations.bounding_rect.v1 import ( - BoundingRectBlockV1, -) -from inference.core.workflows.core_steps.transformations.byte_tracker.v1 import ( - ByteTrackerBlockV1, -) -from inference.core.workflows.core_steps.transformations.byte_tracker.v2 import ( - ByteTrackerBlockV2, -) -from inference.core.workflows.core_steps.transformations.byte_tracker.v3 import ( - ByteTrackerBlockV3, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.trackers.botsort.v1 import ( + BoTSORTBlockV1 as TrackerBoTSORTBlockV1, + ) +else: + from inference.core.workflows.core_steps.trackers.botsort.v1_tensor import ( + BoTSORTBlockV1 as TrackerBoTSORTBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.trackers.bytetrack.v1 import ( + ByteTrackBlockV1 as TrackerByteTrackBlockV1, + ) +else: + from inference.core.workflows.core_steps.trackers.bytetrack.v1_tensor import ( + ByteTrackBlockV1 as TrackerByteTrackBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.trackers.ocsort.v1 import ( + OCSORTBlockV1 as TrackerOCSORTBlockV1, + ) +else: + from inference.core.workflows.core_steps.trackers.ocsort.v1_tensor import ( + OCSORTBlockV1 as TrackerOCSORTBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.trackers.sort.v1 import ( + SORTBlockV1 as TrackerSORTBlockV1, + ) +else: + from inference.core.workflows.core_steps.trackers.sort.v1_tensor import ( + SORTBlockV1 as TrackerSORTBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.absolute_static_crop.v1 import ( + AbsoluteStaticCropBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.absolute_static_crop.v1_tensor import ( + AbsoluteStaticCropBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.bounding_rect.v1 import ( + BoundingRectBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.bounding_rect.v1_tensor import ( + BoundingRectBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.byte_tracker.v1 import ( + ByteTrackerBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.byte_tracker.v1_tensor import ( + ByteTrackerBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.byte_tracker.v2 import ( + ByteTrackerBlockV2, + ) +else: + from inference.core.workflows.core_steps.transformations.byte_tracker.v2_tensor import ( + ByteTrackerBlockV2, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.byte_tracker.v3 import ( + ByteTrackerBlockV3, + ) +else: + from inference.core.workflows.core_steps.transformations.byte_tracker.v3_tensor import ( + ByteTrackerBlockV3, + ) + from inference.core.workflows.core_steps.transformations.camera_calibration.v1 import ( CameraCalibrationBlockV1, ) -from inference.core.workflows.core_steps.transformations.detection_offset.v1 import ( - DetectionOffsetBlockV1, -) -from inference.core.workflows.core_steps.transformations.detections_combine.v1 import ( - DetectionsCombineBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.detection_offset.v1 import ( + DetectionOffsetBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.detection_offset.v1_tensor import ( + DetectionOffsetBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.detections_combine.v1 import ( + DetectionsCombineBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.detections_combine.v1_tensor import ( + DetectionsCombineBlockV1, + ) + from inference.core.workflows.core_steps.transformations.detections_filter.v1 import ( DetectionsFilterBlockV1, ) -from inference.core.workflows.core_steps.transformations.detections_merge.v1 import ( - DetectionsMergeBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.detections_merge.v1 import ( + DetectionsMergeBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.detections_merge.v1_tensor import ( + DetectionsMergeBlockV1, + ) + from inference.core.workflows.core_steps.transformations.detections_transformation.v1 import ( DetectionsTransformationBlockV1, ) -from inference.core.workflows.core_steps.transformations.dynamic_crop.v1 import ( - DynamicCropBlockV1, -) -from inference.core.workflows.core_steps.transformations.dynamic_zones.v1 import ( - DynamicZonesBlockV1, -) -from inference.core.workflows.core_steps.transformations.geotag_detection.v1 import ( - GeoTagDetectionBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.dynamic_crop.v1 import ( + DynamicCropBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.dynamic_crop.v1_tensor import ( + DynamicCropBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.dynamic_zones.v1 import ( + DynamicZonesBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.dynamic_zones.v1_tensor import ( + DynamicZonesBlockV1, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.geotag_detection.v1 import ( + GeoTagDetectionBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.geotag_detection.v1_tensor import ( + GeoTagDetectionBlockV1, + ) + from inference.core.workflows.core_steps.transformations.image_slicer.v1 import ( ImageSlicerBlockV1, ) from inference.core.workflows.core_steps.transformations.image_slicer.v2 import ( ImageSlicerBlockV2, ) -from inference.core.workflows.core_steps.transformations.per_class_confidence_filter.v1 import ( - PerClassConfidenceFilterBlockV1, -) -from inference.core.workflows.core_steps.transformations.perspective_correction.v1 import ( - PerspectiveCorrectionBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.per_class_confidence_filter.v1 import ( + PerClassConfidenceFilterBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.per_class_confidence_filter.v1_tensor import ( + PerClassConfidenceFilterBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.perspective_correction.v1 import ( + PerspectiveCorrectionBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.perspective_correction.v1_tensor import ( + PerspectiveCorrectionBlockV1, + ) + from inference.core.workflows.core_steps.transformations.qr_code_generator.v1 import ( QRCodeGeneratorBlockV1, ) from inference.core.workflows.core_steps.transformations.relative_static_crop.v1 import ( RelativeStaticCropBlockV1, ) -from inference.core.workflows.core_steps.transformations.stabilize_detections.v1 import ( - StabilizeTrackedDetectionsBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.stabilize_detections.v1 import ( + StabilizeTrackedDetectionsBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.stabilize_detections.v1_tensor import ( + StabilizeTrackedDetectionsBlockV1, + ) + from inference.core.workflows.core_steps.transformations.stitch_images.v1 import ( StitchImagesBlockV1, ) -from inference.core.workflows.core_steps.transformations.stitch_ocr_detections.v1 import ( - StitchOCRDetectionsBlockV1, -) -from inference.core.workflows.core_steps.transformations.stitch_ocr_detections.v2 import ( - StitchOCRDetectionsBlockV2, -) -from inference.core.workflows.core_steps.transformations.track_class_lock.v1 import ( - TrackClassLockBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.stitch_ocr_detections.v1 import ( + StitchOCRDetectionsBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.stitch_ocr_detections.v1_tensor import ( + StitchOCRDetectionsBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.stitch_ocr_detections.v2 import ( + StitchOCRDetectionsBlockV2, + ) +else: + from inference.core.workflows.core_steps.transformations.stitch_ocr_detections.v2_tensor import ( + StitchOCRDetectionsBlockV2, + ) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.transformations.track_class_lock.v1 import ( + TrackClassLockBlockV1, + ) +else: + from inference.core.workflows.core_steps.transformations.track_class_lock.v1_tensor import ( + TrackClassLockBlockV1, + ) # Visualizers -from inference.core.workflows.core_steps.visualizations.background_color.v1 import ( - BackgroundColorVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.blur.v1 import ( - BlurVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.bounding_box.v1 import ( - BoundingBoxVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.circle.v1 import ( - CircleVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.classification_label.v1 import ( - ClassificationLabelVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.color.v1 import ( - ColorVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.corner.v1 import ( - CornerVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.crop.v1 import ( - CropVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.dot.v1 import ( - DotVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.ellipse.v1 import ( - EllipseVisualizationBlockV1, -) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.background_color.v1 import ( + BackgroundColorVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.background_color.v1_tensor import ( + BackgroundColorVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.blur.v1 import ( + BlurVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.blur.v1_tensor import ( + BlurVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.bounding_box.v1 import ( + BoundingBoxVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.bounding_box.v1_tensor import ( + BoundingBoxVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.circle.v1 import ( + CircleVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.circle.v1_tensor import ( + CircleVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.classification_label.v1 import ( + ClassificationLabelVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.classification_label.v1_tensor import ( + ClassificationLabelVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.color.v1 import ( + ColorVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.color.v1_tensor import ( + ColorVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.corner.v1 import ( + CornerVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.corner.v1_tensor import ( + CornerVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.crop.v1 import ( + CropVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.crop.v1_tensor import ( + CropVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.dot.v1 import ( + DotVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.dot.v1_tensor import ( + DotVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.ellipse.v1 import ( + EllipseVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.ellipse.v1_tensor import ( + EllipseVisualizationBlockV1, + ) + from inference.core.workflows.core_steps.visualizations.grid.v1 import ( GridVisualizationBlockV1, ) -from inference.core.workflows.core_steps.visualizations.halo.v1 import ( - HaloVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.halo.v2 import ( - HaloVisualizationBlockV2, -) -from inference.core.workflows.core_steps.visualizations.heatmap.v1 import ( - HeatmapVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.icon.v1 import ( - IconVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.keypoint.v1 import ( - KeypointVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.label.v1 import ( - LabelVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.label.v2 import ( - LabelVisualizationBlockV2, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.halo.v1 import ( + HaloVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.halo.v1_tensor import ( + HaloVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.halo.v2 import ( + HaloVisualizationBlockV2, + ) +else: + from inference.core.workflows.core_steps.visualizations.halo.v2_tensor import ( + HaloVisualizationBlockV2, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.heatmap.v1 import ( + HeatmapVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.heatmap.v1_tensor import ( + HeatmapVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.icon.v1 import ( + IconVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.icon.v1_tensor import ( + IconVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.keypoint.v1 import ( + KeypointVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.keypoint.v1_tensor import ( + KeypointVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.label.v1 import ( + LabelVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.label.v1_tensor import ( + LabelVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.label.v2 import ( + LabelVisualizationBlockV2, + ) +else: + from inference.core.workflows.core_steps.visualizations.label.v2_tensor import ( + LabelVisualizationBlockV2, + ) + from inference.core.workflows.core_steps.visualizations.line_zone.v1 import ( LineCounterZoneVisualizationBlockV1, ) -from inference.core.workflows.core_steps.visualizations.mask.v1 import ( - MaskVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.model_comparison.v1 import ( - ModelComparisonVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.pixelate.v1 import ( - PixelateVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.polygon.v1 import ( - PolygonVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.polygon.v2 import ( - PolygonVisualizationBlockV2, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.mask.v1 import ( + MaskVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.mask.v1_tensor import ( + MaskVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.model_comparison.v1 import ( + ModelComparisonVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.model_comparison.v1_tensor import ( + ModelComparisonVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.pixelate.v1 import ( + PixelateVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.pixelate.v1_tensor import ( + PixelateVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.polygon.v1 import ( + PolygonVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.polygon.v1_tensor import ( + PolygonVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.polygon.v2 import ( + PolygonVisualizationBlockV2, + ) +else: + from inference.core.workflows.core_steps.visualizations.polygon.v2_tensor import ( + PolygonVisualizationBlockV2, + ) + from inference.core.workflows.core_steps.visualizations.polygon_zone.v1 import ( PolygonZoneVisualizationBlockV1, ) from inference.core.workflows.core_steps.visualizations.reference_path.v1 import ( ReferencePathVisualizationBlockV1, ) -from inference.core.workflows.core_steps.visualizations.rich_label.v1 import ( - RichLabelVisualizationBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.rich_label.v1 import ( + RichLabelVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.rich_label.v1_tensor import ( + RichLabelVisualizationBlockV1, + ) + from inference.core.workflows.core_steps.visualizations.text_display.v1 import ( TextDisplayVisualizationBlockV1, ) -from inference.core.workflows.core_steps.visualizations.trace.v1 import ( - TraceVisualizationBlockV1, -) -from inference.core.workflows.core_steps.visualizations.triangle.v1 import ( - TriangleVisualizationBlockV1, -) + +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.trace.v1 import ( + TraceVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.trace.v1_tensor import ( + TraceVisualizationBlockV1, + ) +if not ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.visualizations.triangle.v1 import ( + TriangleVisualizationBlockV1, + ) +else: + from inference.core.workflows.core_steps.visualizations.triangle.v1_tensor import ( + TriangleVisualizationBlockV1, + ) + from inference.core.workflows.execution_engine.entities.types import ( BAR_CODE_DETECTION_KIND, BOOLEAN_KIND, @@ -750,6 +1539,38 @@ DETECTIONS_OVERLAPS_KIND.name: serialize_wildcard_kind, TIMESTAMP_KIND.name: serialize_timestamp, } +if ENABLE_TENSOR_DATA_REPRESENTATION: + # Tensor-native producers emit native dataclasses for these kinds, which the + # numpy serialisers cannot handle. Classification has no numpy serialiser at all; + # keypoint produces a (KeyPoints, Detections) tuple; semantic-seg uses the + # per-class RLE InstanceDetections carrier. + KINDS_SERIALIZERS[CLASSIFICATION_PREDICTION_KIND.name] = ( + serialise_native_classification + ) + KINDS_SERIALIZERS[KEYPOINT_DETECTION_PREDICTION_KIND.name] = ( + serialise_native_keypoint_detection + ) + # Emit the native COCO RLE per box (no polygon collapse). + KINDS_SERIALIZERS[SEMANTIC_SEGMENTATION_PREDICTION_KIND.name] = ( + serialise_native_rle_detections + ) + # instance-seg blocks declare the RLE kind too; the numpy `serialise_rle_sv_detections` + # cannot handle a native InstanceDetections. + KINDS_SERIALIZERS[RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND.name] = ( + serialise_native_rle_detections + ) + # Tensor-native embedding/tensor kinds serialise to plain Python lists. + KINDS_SERIALIZERS[EMBEDDING_KIND.name] = serialise_native_embedding + KINDS_SERIALIZERS[TENSOR_KIND.name] = serialise_native_tensor + # `numpy_array` has NO numpy-side serialiser (ndarrays pass through raw to + # the HTTP layer). Flag-on, the depth-estimation block carries a + # `torch.Tensor` under the same kind name, which the HTTP layer cannot + # serialise - materialise it to the flag-off-identical ndarray. Numpy + # producers of this kind (SIFT, contours) keep the raw pass-through. + KINDS_SERIALIZERS[NUMPY_ARRAY_KIND.name] = serialise_numpy_array_kind + # The wildcard (`*`) serialiser needs no override here: the same-name symbol + # swap above already selects serializers_tensor.serialize_wildcard_kind, whose + # native arms serialise native values routed to `*` outputs to the standard dicts. KINDS_DESERIALIZERS = { IMAGE_KIND.name: deserialize_image_kind, VIDEO_METADATA_KIND.name: deserialize_video_metadata_kind, @@ -787,6 +1608,16 @@ INFERENCE_ID_KIND.name: deserialize_string_kind, TIMESTAMP_KIND.name: deserialize_timestamp, } +if ENABLE_TENSOR_DATA_REPRESENTATION: + # Tensor-native consumers expect native dataclasses / torch tensors for these + # kinds. The numpy classification deserialiser returns a plain dict; embedding / + # tensor kinds have NO numpy deserialiser at all (a serialised JSON list would + # reach a tensor consumer as a list and break on `.shape` / `torch.dot`). + KINDS_DESERIALIZERS[CLASSIFICATION_PREDICTION_KIND.name] = ( + deserialize_native_classification_prediction_kind + ) + KINDS_DESERIALIZERS[EMBEDDING_KIND.name] = deserialize_native_embedding_kind + KINDS_DESERIALIZERS[TENSOR_KIND.name] = deserialize_native_tensor_kind def _should_filter_block(block_class: Type[WorkflowBlock]) -> bool: @@ -1087,23 +1918,61 @@ def load_kinds() -> List[Kind]: TOP_CLASS_KIND, FLOAT_KIND, DICTIONARY_KIND, - DETECTION_KIND, - CLASSIFICATION_PREDICTION_KIND, + ( + DETECTION_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_DETECTION_KIND + ), + ( + CLASSIFICATION_PREDICTION_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND + ), DETECTIONS_OVERLAPS_KIND, POINT_KIND, LABELED_POINTS_KIND, ZONE_KIND, - OBJECT_DETECTION_PREDICTION_KIND, - INSTANCE_SEGMENTATION_PREDICTION_KIND, - KEYPOINT_DETECTION_PREDICTION_KIND, - SEMANTIC_SEGMENTATION_PREDICTION_KIND, + ( + OBJECT_DETECTION_PREDICTION_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND + ), + ( + INSTANCE_SEGMENTATION_PREDICTION_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND + ), + ( + RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND + ), + ( + KEYPOINT_DETECTION_PREDICTION_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND + ), + ( + SEMANTIC_SEGMENTATION_PREDICTION_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_SEMANTIC_SEGMENTATION_PREDICTION_KIND + ), RGB_COLOR_KIND, IMAGE_KEYPOINTS_KIND, CONTOURS_KIND, LANGUAGE_MODEL_OUTPUT_KIND, NUMPY_ARRAY_KIND, - QR_CODE_DETECTION_KIND, - BAR_CODE_DETECTION_KIND, + TENSOR_KIND, + ( + QR_CODE_DETECTION_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_QR_CODE_DETECTION_KIND + ), + ( + BAR_CODE_DETECTION_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_BAR_CODE_DETECTION_KIND + ), PREDICTION_TYPE_KIND, ROBOFLOW_MANAGED_KEY, PARENT_ID_KIND, @@ -1111,6 +1980,10 @@ def load_kinds() -> List[Kind]: BYTES_KIND, INFERENCE_ID_KIND, SECRET_KIND, - EMBEDDING_KIND, + ( + EMBEDDING_KIND + if not ENABLE_TENSOR_DATA_REPRESENTATION + else TENSOR_NATIVE_EMBEDDING_KIND + ), TIMESTAMP_KIND, ] diff --git a/inference/core/workflows/core_steps/math/cosine_similarity/v1_tensor.py b/inference/core/workflows/core_steps/math/cosine_similarity/v1_tensor.py new file mode 100644 index 0000000000..1b204b6fa1 --- /dev/null +++ b/inference/core/workflows/core_steps/math/cosine_similarity/v1_tensor.py @@ -0,0 +1,135 @@ +from typing import List, Literal, Optional, Type + +import torch +from pydantic import ConfigDict, Field + +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_EMBEDDING_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) + +LONG_DESCRIPTION = """ +Calculate the cosine similarity between two embedding vectors by computing the cosine of the angle between them, measuring directional similarity regardless of magnitude to enable similarity comparison, semantic matching, embedding-based search, and similarity-based filtering workflows. + +## How This Block Works + +This block computes cosine similarity, a measure of similarity between two vectors based on the cosine of the angle between them. The block: + +1. Receives two embedding vectors from workflow steps (e.g., from CLIP, Perception Encoder, or other embedding models) +2. Validates embedding dimensions: + - Ensures both embeddings have the same dimensionality (same number of elements) + - Raises an error if dimensions don't match +3. Computes cosine similarity: + - Calculates the dot product of the two embedding vectors + - Computes the L2 norm (magnitude) of each embedding vector + - Divides the dot product by the product of the two norms: similarity = (a ยท b) / (||a|| ร— ||b||) + - This measures the cosine of the angle between the vectors, indicating directional similarity +4. Returns similarity score: + - Outputs a similarity value ranging from -1 to 1 + - Value of 1: Vectors point in the same direction (identical or proportional) - maximum similarity + - Value of 0: Vectors are orthogonal (perpendicular) - no similarity + - Value of -1: Vectors point in opposite directions - maximum dissimilarity + - Greater values (closer to 1) indicate greater similarity + +Cosine similarity is magnitude-invariant, meaning it measures similarity in direction rather than size. Two vectors that point in the same direction will have high cosine similarity even if they have different magnitudes. This makes it ideal for comparing embeddings where magnitude may vary but semantic meaning (direction) is what matters. + +## Common Use Cases + +- **Semantic Similarity Comparison**: Compare semantic similarity between images, text, or other data types using embeddings (e.g., compare image embeddings, match text to images, find similar content), enabling similarity comparison workflows +- **Embedding-Based Search**: Use similarity scores for embedding-based search and retrieval (e.g., find similar images, search by embedding similarity, retrieve similar content), enabling embedding search workflows +- **Cross-Modal Matching**: Match embeddings across different modalities (e.g., match images to text, find images matching text descriptions, match text to images), enabling cross-modal matching workflows +- **Similarity-Based Filtering**: Filter data based on similarity thresholds (e.g., filter similar items, find duplicates using similarity, identify near-duplicates), enabling similarity filtering workflows +- **Content Recommendation**: Use similarity scores for content recommendation and matching (e.g., recommend similar content, match related items, suggest similar products), enabling recommendation workflows +- **Quality Control and Validation**: Validate embeddings or compare embeddings for quality control (e.g., validate embedding quality, compare embeddings for consistency, check embedding similarity), enabling quality control workflows + +## Connecting to Other Blocks + +This block receives embeddings from embedding model blocks and produces similarity scores: + +- **After embedding model blocks** (CLIP, Perception Encoder, etc.) to compare embeddings (e.g., compare image and text embeddings, compare multiple embeddings, compute similarity scores), enabling embedding-to-similarity workflows +- **Before logic blocks** like Continue If to use similarity scores in conditions (e.g., continue if similarity exceeds threshold, filter based on similarity, make decisions using similarity), enabling similarity-based decision workflows +- **Before filtering blocks** to filter based on similarity (e.g., filter by similarity threshold, remove low-similarity items, keep high-similarity matches), enabling similarity-to-filter workflows +- **Before data storage blocks** to store similarity scores (e.g., store similarity metrics, log similarity comparisons, save similarity results), enabling similarity storage workflows +- **Before notification blocks** to send similarity-based alerts (e.g., notify on high similarity matches, alert on similarity changes, send similarity reports), enabling similarity notification workflows +- **In workflow outputs** to provide similarity scores as final output (e.g., similarity comparison outputs, matching results, similarity metrics), enabling similarity output workflows + +## Requirements + +This block requires two embedding vectors with the same dimensionality (same number of elements). Embeddings can be from any embedding model (CLIP, Perception Encoder, etc.) and can represent images, text, or other data types. The embeddings are passed as lists of floats. The block computes cosine similarity using the dot product divided by the product of L2 norms, producing a similarity score between -1 and 1. Values closer to 1 indicate greater similarity, values closer to 0 indicate orthogonal vectors, and values closer to -1 indicate opposite directions. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Cosine Similarity", + "version": "v1", + "short_description": "Calculate the cosine similarity between two embeddings.", + "long_description": LONG_DESCRIPTION, + "license": "MIT", + "block_type": "math", + "ui_manifest": { + "section": "advanced", + "icon": "far fa-calculator-simple", + "blockPriority": 3, + }, + } + ) + type: Literal["roboflow_core/cosine_similarity@v1"] + name: str = Field(description="Unique name of step in workflows") + embedding_1: Selector(kind=[TENSOR_NATIVE_EMBEDDING_KIND]) = Field( + description="First embedding vector to compare. Must have the same dimensionality (same number of elements) as embedding_2. Can be from any embedding model (CLIP, Perception Encoder, etc.) and can represent images, text, or other data types. Embedding vectors are lists of floats representing high-dimensional feature representations.", + examples=[ + "$steps.clip_image.embedding", + "$steps.perception_encoder.embedding", + "$steps.clip_text.embedding", + ], + ) + embedding_2: Selector(kind=[TENSOR_NATIVE_EMBEDDING_KIND]) = Field( + description="Second embedding vector to compare. Must have the same dimensionality (same number of elements) as embedding_1. Can be from any embedding model (CLIP, Perception Encoder, etc.) and can represent images, text, or other data types. Embedding vectors are lists of floats representing high-dimensional feature representations. The cosine similarity measures the similarity between embedding_1 and embedding_2.", + examples=[ + "$steps.clip_text.embedding", + "$steps.clip_image.embedding", + "$steps.perception_encoder.embedding", + ], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [OutputDefinition(name="similarity", kind=[FLOAT_KIND])] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class CosineSimilarityBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run(self, embedding_1: torch.Tensor, embedding_2: torch.Tensor) -> BlockResult: + # Tensor-native embeddings stay on-device: compute cosine similarity in + # torch (aยทb / (||a||ยท||b||)) and only extract the final scalar. + # Validate the vector dimension via shape[-1] and flatten to 1-D so both + # 1-D `(D,)` and batched `(1, D)` embeddings work with torch.dot. + if embedding_1.shape[-1] != embedding_2.shape[-1]: + raise RuntimeError( + f"roboflow_core/cosine_similarity@v1 block feed with different shape of embeddings. " + f"`embedding_1`: (N, {embedding_1.shape[-1]}), `embedding_2`: (N, {embedding_2.shape[-1]})" + ) + embedding_1 = embedding_1.reshape(-1) + embedding_2 = embedding_2.reshape(-1) + similarity = torch.dot(embedding_1, embedding_2) / torch.sqrt( + torch.dot(embedding_1, embedding_1) * torch.dot(embedding_2, embedding_2) + ) + return {"similarity": float(similarity)} diff --git a/inference/core/workflows/core_steps/models/foundation/clip/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/clip/v1_tensor.py new file mode 100644 index 0000000000..afb1089bf7 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/clip/v1_tensor.py @@ -0,0 +1,228 @@ +import hashlib +from typing import List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field + +from inference.core.cache.lru_cache import LRUCache +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import ModelEndpointType +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_EMBEDDING_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_sdk import InferenceHTTPClient + +LONG_DESCRIPTION = """ +Use a CLIP model to create semantic embeddings of text and images. + +This block accepts an image or string and returns an embedding. +The embedding can be used to compare the similarity between different +images or between images and text. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "CLIP Embedding Model", + "version": "v1", + "short_description": "Generate an embedding of an image or string.", + "long_description": LONG_DESCRIPTION, + "license": "MIT", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-paperclip", + "blockPriority": 9.9, + }, + } + ) + type: Literal["roboflow_core/clip@v1"] + name: str = Field(description="Unique name of step in workflows") + data: Union[Selector(kind=[IMAGE_KIND, STRING_KIND]), str] = Field( + title="Data", + description="The string or image to generate an embedding for.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + + version: Union[ + Literal[ + "RN101", + "RN50", + "RN50x16", + "RN50x4", + "RN50x64", + "ViT-B-16", + "ViT-B-32", + "ViT-L-14-336px", + "ViT-L-14", + ], + Selector(kind=[STRING_KIND]), + ] = Field( + default="ViT-B-32", + description="Variant of CLIP model", + examples=["ViT-B-16", "$inputs.variant"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [OutputDefinition(name="embedding", kind=[TENSOR_NATIVE_EMBEDDING_KIND])] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + return list(CLIP_CACHE_MODEL_IDS) + + +# All CLIP model_id cache paths. Shared with clip_comparison blocks. +CLIP_CACHE_MODEL_IDS = [ + "clip/RN101", + "clip/RN50", + "clip/RN50x16", + "clip/RN50x4", + "clip/RN50x64", + "clip/ViT-B-16", + "clip/ViT-B-32", + "clip/ViT-L-14-336px", + "clip/ViT-L-14", +] + + +text_cache = LRUCache() + + +class ClipModelBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + data: Union[WorkflowImageData, str], + version: str, + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally(data=data, version=version) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely(data=data, version=version) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + data: Union[WorkflowImageData, str], + version: str, + ) -> BlockResult: + clip_model_id = f"clip/{version}" + self._model_manager.add_model( + clip_model_id, + self._api_key, + endpoint_type=ModelEndpointType.CORE_MODEL, + ) + if isinstance(data, str): + hash_key = hashlib.md5((version + data).encode("utf-8")).hexdigest() + + cached_value = text_cache.get(hash_key) + if cached_value is not None: + # Cache holds CPU tensors so the process-wide LRU never pins GPU memory. + return {"embedding": cached_value.to(WORKFLOWS_IMAGE_TENSOR_DEVICE)} + + embeddings = self._model_manager.run_tensor_native_inference( + clip_model_id, + action="embed-text", + texts=[data], + ) + embedding = embeddings[0] + + text_cache.set(hash_key, embedding.detach().cpu()) + + return {"embedding": embedding.to(WORKFLOWS_IMAGE_TENSOR_DEVICE)} + else: + if data.is_tensor_materialised(): + model_image, image_color_format = data.tensor_image, "rgb" + else: + model_image, image_color_format = data.numpy_image, "bgr" + embeddings = self._model_manager.run_tensor_native_inference( + clip_model_id, + action="embed-image", + images=[model_image], + input_color_format=image_color_format, + ) + return {"embedding": embeddings[0].to(WORKFLOWS_IMAGE_TENSOR_DEVICE)} + + def run_remotely( + self, + data: Union[WorkflowImageData, str], + version: str, + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + if isinstance(data, str): + result = client.get_clip_text_embeddings( + text=data, + clip_version=version, + ) + else: + result = client.get_clip_image_embeddings( + inference_input=data.base64_image, + clip_version=version, + ) + + # Remote returns the embedding as a JSON List[float]. + return { + "embedding": torch.tensor( + result["embeddings"][0], device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + } diff --git a/inference/core/workflows/core_steps/models/foundation/clip_comparison/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/clip_comparison/v1_tensor.py new file mode 100644 index 0000000000..6cdf3078f2 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/clip_comparison/v1_tensor.py @@ -0,0 +1,236 @@ +from functools import partial +from typing import List, Literal, Optional, Type, Union + +import torch +import torch.nn.functional as F +from pydantic import AliasChoices, ConfigDict, Field + +from inference.core.env import ( + CLIP_VERSION_ID, + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import ModelEndpointType +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.utils import ( + remove_unexpected_keys_from_dictionary, + run_in_parallel, +) +from inference.core.workflows.execution_engine.constants import ( + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + LIST_OF_VALUES_KIND, + PARENT_ID_KIND, + PREDICTION_TYPE_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_sdk import InferenceHTTPClient + +LONG_DESCRIPTION = """ +Use the OpenAI CLIP zero-shot classification model to classify images. + +This block accepts an image and a list of text prompts. The block then returns the +similarity of each text label to the provided image. + +This block is useful for classifying images without having to train a fine-tuned +classification model. For example, you could use CLIP to classify the type of vehicle +in an image, or if an image contains NSFW material. +""" + +EXPECTED_OUTPUT_KEYS = {"similarity", "parent_id", "root_parent_id", "prediction_type"} + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Clip Comparison", + "version": "v1", + "short_description": "Compare CLIP image and text embeddings.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "fak fa-message-image", + "blockPriority": 10, + "inference": True, + }, + } + ) + type: Literal["roboflow_core/clip_comparison@v1", "ClipComparison"] + name: str = Field(description="Unique name of step in workflows") + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + texts: Union[Selector(kind=[LIST_OF_VALUES_KIND]), List[str]] = Field( + description="List of texts to calculate similarity against each input image", + examples=[["a", "b", "c"], "$inputs.texts"], + validation_alias=AliasChoices("texts", "text"), + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="similarity", kind=[LIST_OF_VALUES_KIND]), + OutputDefinition(name="parent_id", kind=[PARENT_ID_KIND]), + OutputDefinition(name="root_parent_id", kind=[PARENT_ID_KIND]), + OutputDefinition(name="prediction_type", kind=[PREDICTION_TYPE_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + from inference.core.workflows.core_steps.models.foundation.clip.v1 import ( + CLIP_CACHE_MODEL_IDS, + ) + + return list(CLIP_CACHE_MODEL_IDS) + + +class ClipComparisonBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + texts: List[str], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally(images=images, texts=texts) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely(images=images, texts=texts) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + texts: List[str], + ) -> BlockResult: + clip_model_id = f"clip/{CLIP_VERSION_ID}" + self._model_manager.add_model( + clip_model_id, + self._api_key, + endpoint_type=ModelEndpointType.CORE_MODEL, + ) + text_embeddings = F.normalize( + self._model_manager.run_tensor_native_inference( + clip_model_id, + action="embed-text", + texts=texts, + ), + dim=1, + ) + predictions = [] + for single_image in images: + if single_image.is_tensor_materialised(): + model_image, image_color_format = single_image.tensor_image, "rgb" + else: + model_image, image_color_format = single_image.numpy_image, "bgr" + image_embedding = F.normalize( + self._model_manager.run_tensor_native_inference( + clip_model_id, + action="embed-image", + images=[model_image], + input_color_format=image_color_format, + ), + dim=1, + ) + similarities = (image_embedding @ text_embeddings.T)[0] + predictions.append({"similarity": similarities.detach().to("cpu").tolist()}) + return self._post_process_result( + images=images, + predictions=predictions, + ) + + def run_remotely( + self, + images: Batch[WorkflowImageData], + texts: List[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + tasks = [ + partial( + client.clip_compare, + subject=single_image.numpy_image, + prompt=texts, + ) + for single_image in images + ] + predictions = run_in_parallel( + tasks=tasks, + max_workers=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + ) + return self._post_process_result(images=images, predictions=predictions) + + def _post_process_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + ) -> List[dict]: + for prediction, image in zip(predictions, images): + prediction[PREDICTION_TYPE_KEY] = "classification" + prediction[PARENT_ID_KEY] = image.parent_metadata.parent_id + prediction[ROOT_PARENT_ID_KEY] = ( + image.workflow_root_ancestor_metadata.parent_id + ) + # removing fields from `inference` response model + # that are not registered as outputs + _ = remove_unexpected_keys_from_dictionary( + dictionary=prediction, + expected_keys=EXPECTED_OUTPUT_KEYS, + ) + return predictions diff --git a/inference/core/workflows/core_steps/models/foundation/clip_comparison/v2_tensor.py b/inference/core/workflows/core_steps/models/foundation/clip_comparison/v2_tensor.py new file mode 100644 index 0000000000..e2dc2830a9 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/clip_comparison/v2_tensor.py @@ -0,0 +1,306 @@ +from functools import partial +from typing import List, Literal, Optional, Type, Union + +import numpy as np +import torch +import torch.nn.functional as F +from pydantic import ConfigDict, Field + +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import ModelEndpointType +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.utils import run_in_parallel +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + IMAGE_DIMENSIONS_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + PARENT_ID_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models import ClassificationPrediction +from inference_sdk import InferenceHTTPClient + +LONG_DESCRIPTION = """ +Use the OpenAI CLIP zero-shot classification model to classify images. + +This block accepts an image and a list of text prompts. The block then returns the +similarity of each text label to the provided image. + +This block is useful for classifying images without having to train a fine-tuned +classification model. For example, you could use CLIP to classify the type of vehicle +in an image, or if an image contains NSFW material. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Clip Comparison", + "version": "v2", + "short_description": "Compare CLIP image and text embeddings.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "fak fa-message-image", + "blockPriority": 10, + "inference": True, + }, + } + ) + type: Literal["roboflow_core/clip_comparison@v2"] + name: str = Field(description="Unique name of step in workflows") + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + classes: Union[Selector(kind=[LIST_OF_VALUES_KIND]), List[str]] = Field( + description="List of classes to calculate similarity against each input image", + examples=[["a", "b", "c"], "$inputs.texts"], + min_items=1, + ) + version: Union[ + Literal[ + "RN101", + "RN50", + "RN50x16", + "RN50x4", + "RN50x64", + "ViT-B-16", + "ViT-B-32", + "ViT-L-14-336px", + "ViT-L-14", + ], + Selector(kind=[STRING_KIND]), + ] = Field( + default="ViT-B-16", + description="Variant of CLIP model", + examples=["ViT-B-16", "$inputs.variant"], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="similarities", kind=[LIST_OF_VALUES_KIND]), + OutputDefinition(name="max_similarity", kind=[FLOAT_ZERO_TO_ONE_KIND]), + OutputDefinition(name="most_similar_class", kind=[STRING_KIND]), + OutputDefinition(name="min_similarity", kind=[FLOAT_ZERO_TO_ONE_KIND]), + OutputDefinition(name="least_similar_class", kind=[STRING_KIND]), + OutputDefinition( + name="classification_predictions", + kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND], + ), + OutputDefinition(name="parent_id", kind=[PARENT_ID_KIND]), + OutputDefinition(name="root_parent_id", kind=[PARENT_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + from inference.core.workflows.core_steps.models.foundation.clip.v1 import ( + CLIP_CACHE_MODEL_IDS, + ) + + return list(CLIP_CACHE_MODEL_IDS) + + +class ClipComparisonBlockV2(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + classes: List[str], + version: str, + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally(images=images, classes=classes, version=version) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely(images=images, classes=classes, version=version) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + classes: List[str], + version: str, + ) -> BlockResult: + clip_model_id = f"clip/{version}" + self._model_manager.add_model( + clip_model_id, + self._api_key, + endpoint_type=ModelEndpointType.CORE_MODEL, + ) + class_embeddings = F.normalize( + self._model_manager.run_tensor_native_inference( + clip_model_id, + action="embed-text", + texts=classes, + ), + dim=1, + ) + predictions = [] + for single_image in images: + if single_image.is_tensor_materialised(): + model_image, image_color_format = single_image.tensor_image, "rgb" + else: + model_image, image_color_format = single_image.numpy_image, "bgr" + image_embedding = F.normalize( + self._model_manager.run_tensor_native_inference( + clip_model_id, + action="embed-image", + images=[model_image], + input_color_format=image_color_format, + ), + dim=1, + ) + similarities = (image_embedding @ class_embeddings.T)[0] + predictions.append({"similarity": similarities.detach().to("cpu").tolist()}) + return self._post_process_result( + images=images, + predictions=predictions, + classes=classes, + ) + + def run_remotely( + self, + images: Batch[WorkflowImageData], + classes: List[str], + version: str, + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + tasks = [ + partial( + client.clip_compare, + subject=single_image.base64_image, + prompt=classes, + clip_version=version, + ) + for single_image in images + ] + predictions = run_in_parallel( + tasks=tasks, + max_workers=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + ) + return self._post_process_result( + images=images, + predictions=predictions, + classes=classes, + ) + + def _post_process_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + classes: List[str], + ) -> List[dict]: + results = [] + for prediction, image in zip(predictions, images): + similarities = prediction["similarity"] + max_similarity = float(np.max(similarities)) + max_similarity_id = np.argmax(similarities) + min_similarity = float(np.min(similarities)) + min_similarity_id = np.argmin(similarities) + most_similar_class_name = classes[max_similarity_id] + least_similar_class_name = classes[min_similarity_id] + image_height, image_width = image._read_shape_without_materialization() + # `confidence` carries per-class cosine scores, not softmax probabilities. + classification_predictions = ClassificationPrediction( + class_id=torch.tensor([int(max_similarity_id)]), + confidence=torch.tensor([similarities], dtype=torch.float32), + images_metadata=[ + # Lane 1b NOTE: intentionally NOT stamped with CLASSIFICATION_STYLE_KEY. + # clip_comparison v2's flag-OFF `classification_predictions` is a + # bespoke dict ({predictions, top, confidence, parent_id}) that matches + # NEITHER the serialiser's "model" nor "formatter" shape (both add + # `image`/`inference_id`; model also sorts/rounds/filters + prediction_type). + # This is a PRE-EXISTING byte-parity gap, not a lane-1b regression, so the + # prediction is left on the fallback heuristic (-> formatter, unchanged from + # before lane 1b). Byte-parity for clip_comparison needs its own decision. + { + CLASS_NAMES_KEY: { + class_id: class_name + for class_id, class_name in enumerate(classes) + }, + PREDICTION_TYPE_KEY: "classification", + PARENT_ID_KEY: image.parent_metadata.parent_id, + IMAGE_DIMENSIONS_KEY: [image_height, image_width], + } + ], + ) + result = { + PARENT_ID_KEY: image.parent_metadata.parent_id, + ROOT_PARENT_ID_KEY: image.workflow_root_ancestor_metadata.parent_id, + "similarities": similarities, + "max_similarity": max_similarity, + "most_similar_class": most_similar_class_name, + "min_similarity": min_similarity, + "least_similar_class": least_similar_class_name, + "classification_predictions": classification_predictions, + } + results.append(result) + return results diff --git a/inference/core/workflows/core_steps/models/foundation/depth_estimation/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/depth_estimation/v1_tensor.py new file mode 100644 index 0000000000..8633399490 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/depth_estimation/v1_tensor.py @@ -0,0 +1,316 @@ +from typing import List, Literal, Optional, Type, Union + +import matplotlib.pyplot as plt +import numpy as np +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import ( + DEPTH_ESTIMATION_ENABLED, + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + NUMPY_ARRAY_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_sdk import InferenceHTTPClient + + +class BlockManifest(WorkflowBlockManifest): + # Standard model configuration for UI, schema, etc. + model_config = ConfigDict( + json_schema_extra={ + "name": "Depth Estimation", + "version": "v1", + "short_description": "Estimate relative scene depth in an image.", + "long_description": ( + """ + ๐ŸŽฏ This workflow block performs monocular depth estimation with a + selected Depth Anything or YOLO26 depth model. + + The model outputs: + 1. ๐Ÿ—บ๏ธ A visualization of the estimated scene depth + 2. ๐Ÿ“Š `normalized_depth`, an image-sized ordinal proximity map + + `normalized_depth` is normalized independently for every image: + ๐Ÿ” 1.0 indicates the nearest prediction + ๐Ÿ”ญ 0.0 indicates the farthest prediction + + Intermediate values preserve relative near-to-far ordering. They are + not physical distances and should not be compared numerically across + images or model families. YOLO26's metric output is normalized by this + block; use `inference_models.AutoModel` directly when meters are needed. + + This is particularly useful for: + - ๐Ÿ—๏ธ Understanding 3D structure from 2D images + - ๐ŸŽจ Creating depth-aware visualizations + - ๐Ÿ“ Analyzing relative spatial relationships in scenes + - ๐Ÿ•ถ๏ธ Depth-aware image processing + """ + ), + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": [ + "Depth Estimation", + "Depth Anything", + "Depth Anything V2", + "Depth Anything V3", + "YOLO26", + "YOLO26 Depth", + "Hugging Face", + "HuggingFace", + ], + "is_vlm_block": True, + "ui_manifest": { + "section": "model", + "icon": "fal fa-atom", + "blockPriority": 5.5, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/depth_estimation@v1"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + + model_version: Union[ + Literal[ + "depth-anything-v2/small", + "depth-anything-v3/small", + "depth-anything-v3/base", + "yolo26n-depth-768", + "yolo26s-depth-768", + "yolo26m-depth-768", + "yolo26l-depth-768", + "yolo26x-depth-768", + ], + Selector(kind=[STRING_KIND]), + ] = Field( + default="depth-anything-v3/small", + description="The Depth Estimation model to be used for inference.", + examples=["depth-anything-v2/small", "$inputs.variant"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + # Same-name kind as the numpy sibling; flag-on the payload is a + # torch.Tensor materialised back to an ndarray at the serialization + # boundary (see serializers_tensor.serialise_numpy_array_kind). + return [ + OutputDefinition(name="image", kind=[IMAGE_KIND]), + OutputDefinition(name="normalized_depth", kind=[NUMPY_ARRAY_KIND]), + ] + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + # Only images can be passed in as a list/batch + return ["images"] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restrictions = [ + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; run_locally() loads a model that needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + ] + if not DEPTH_ESTIMATION_ENABLED: + restrictions.append( + RuntimeRestriction( + severity=Severity.HARD, + note=( + "DEPTH_ESTIMATION_ENABLED=False on Roboflow Hosted " + "Serverless: the depth-estimation endpoint is not " + "registered, so run_remotely() returns 404." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + ) + return restrictions + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + return [ + "depth-anything-v2/small", + "depth-anything-v3/small", + "depth-anything-v3/base", + "yolo26n-depth-768", + "yolo26s-depth-768", + "yolo26m-depth-768", + "yolo26l-depth-768", + "yolo26x-depth-768", + ] + + +def _depth_to_visualization( + origin_image: WorkflowImageData, + normalized_depth: torch.Tensor, +) -> WorkflowImageData: + depth_for_viz = (normalized_depth * 255.0).to(torch.uint8).detach().cpu().numpy() + cmap = plt.get_cmap("viridis") + colored_depth = (cmap(depth_for_viz)[:, :, :3] * 255).astype(np.uint8) + return WorkflowImageData.copy_and_replace( + origin_image_data=origin_image, + numpy_image=colored_depth, + ) + + +class DepthEstimationBlockV1(WorkflowBlock): + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_version: str = "depth-anything-v3/small", + ) -> BlockResult: + if self._step_execution_mode == StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_version=model_version, + ) + elif self._step_execution_mode == StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_version=model_version, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_version: str = "depth-anything-v3/small", + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + predictions = [] + for single_image in images: + result = client.depth_estimation( + inference_input=single_image.base64_image, + model_id=model_version, + model_id_in_path=True, + depth_map_format="png16", + ) + # Convert the result back to the expected format + # Remote returns {"normalized_depth": ..., "image": hex_string}; the + # depth map is an ndarray decoded from png16, or a nested float list + # from servers that predate depth_map_format - np.array handles both + image_output = WorkflowImageData.copy_and_replace( + origin_image_data=single_image, + base64_image=result.get("image", ""), + ) + + normalized_depth = torch.as_tensor( + np.array(result.get("normalized_depth", [])), dtype=torch.float32 + ) + + # Return in the same format as local execution expects + predictions.append( + { + "image": image_output, + "normalized_depth": normalized_depth, + } + ) + + return predictions + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_version: str = "depth-anything-v3/small", + ) -> BlockResult: + # Register Depth Estimation with the model manager. + self._model_manager.add_model(model_id=model_version, api_key=self._api_key) + + predictions = [] + for single_image in images: + if single_image.is_tensor_materialised(): + model_image, image_color_format = single_image.tensor_image, "rgb" + else: + model_image, image_color_format = single_image.numpy_image, "bgr" + depth_maps = self._model_manager.run_tensor_native_inference( + model_version, + images=[model_image], + input_color_format=image_color_format, + ) + # Tensor-native depth models return raw per-image maps in which + # larger means CLOSER (DepthAnything emits disparity-like maps + # natively; the YOLO26 adapter negates its metric output to match - + # see InferenceModelsDepthEstimationAdapter). Min-max normalization + # therefore yields the endpoint's ordinal proximity map: 1.0 is the + # nearest prediction, 0.0 the farthest. + depth_map = depth_maps[0] + depth_min = depth_map.min() + depth_max = depth_map.max() + if depth_max == depth_min: + raise ValueError("Depth map has no variation (min equals max)") + normalized_depth = (depth_map - depth_min) / (depth_max - depth_min) + normalized_depth = normalized_depth.to(WORKFLOWS_IMAGE_TENSOR_DEVICE) + image_output = _depth_to_visualization( + origin_image=single_image, + normalized_depth=normalized_depth, + ) + predictions.append( + { + "image": image_output, + "normalized_depth": normalized_depth, + } + ) + + return predictions diff --git a/inference/core/workflows/core_steps/models/foundation/easy_ocr/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/easy_ocr/v1_tensor.py new file mode 100644 index 0000000000..ddeaa2f597 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/easy_ocr/v1_tensor.py @@ -0,0 +1,344 @@ +from typing import Dict, List, Literal, Optional, Tuple, Type + +from pydantic import ConfigDict, Field + +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import ModelEndpointType +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + native_detections_from_inference_predictions, +) +from inference.core.workflows.execution_engine.constants import CLASS_NAME_KEY +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + PARENT_ID_KIND, + PREDICTION_TYPE_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections +from inference_sdk import InferenceHTTPClient +from inference_sdk.http.entities import InferenceConfiguration + +# These are the displayed languages in the UI dropdown +LANGUAGES = Literal[ + "English", + "Japanese", + "Kannada", + "Korean", + "Latin", + "Telugu", + "Simplified Chinese", +] + +# Dictionary of displayed_language: (model, language_code) +# This is not an extensive list of supported languages, more codes can be added +MODELS: Dict[str, Tuple[str, List[str]]] = { + "English": ("english_g2", ["en"]), + "Japanese": ("japanese_g2", ["en", "ja"]), + "Kannada": ("kannada_g2", ["en", "kn"]), + "Korean": ("korean_g2", ["en", "ko"]), + "Latin": ("latin_g2", ["en", "la", "es", "fr", "it", "pt", "de", "pl", "nl"]), + "Telugu": ("telugu_g2", ["en", "te"]), + "Simplified Chinese": ("zh_sim_g2", ["en", "ch_sim"]), +} + +# Fixed, non-empty `class_id -> name` fallback map. The EasyOCR inference_models model +# emits `class_id == 0` for every box; the recognised text is attached per box onto +# `bboxes_metadata[i][CLASS_NAME_KEY]` (see `_attach_local_ocr_text` / +# `_attach_remote_ocr_text`), so the tensor serializer's per-box override emits the +# recognised text as `class` - byte-parity with the numpy sibling, whose +# `sv.Detections.from_inference` carries the recognised text as the class name. This +# map only satisfies the serializer's "CLASS_NAMES_KEY must be present" check. +CLASS_NAMES: Dict[int, str] = {0: "text-region"} +PREDICTION_TYPE = "ocr" + +LONG_DESCRIPTION = """ + Retrieve the characters in an image using EasyOCR Optical Character Recognition (OCR). + +This block returns the text within an image. + +You may want to use this block in combination with a detections-based block (i.e. +ObjectDetectionBlock). An object detection model could isolate specific regions from an +image (i.e. a shipping container ID in a logistics use case) for further processing. +You can then use a DynamicCropBlock to crop the region of interest before running OCR. + +Using a detections model then cropping detections allows you to isolate your analysis +on particular regions of an image. + +Note that EasyOCR has limitations running within containers on Apple Silicon. +""" + + +class BlockManifest(WorkflowBlockManifest): + + model_config = ConfigDict( + json_schema_extra={ + "name": "EasyOCR", + "version": "v1", + "short_description": "Extract text from an image using EasyOCR optical character recognition.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-text", + "blockPriority": 11, + "inDevelopment": False, + "inference": True, + }, + } + ) + type: Literal["roboflow_core/easy_ocr@v1", "EasyOCR"] + name: str = Field(description="Unique name of step in workflows") + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + language: LANGUAGES = Field( + title="Language", + description="Language model to use for OCR", + default="English", + ) + quantize: bool = Field( + title="Use Quantized Model", + description="Quantized models are smaller and faster, but may be less accurate and won't work correctly on all hardware.", + default=False, + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="result", kind=[STRING_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition(name="parent_id", kind=[PARENT_ID_KIND]), + OutputDefinition(name="root_parent_id", kind=[PARENT_ID_KIND]), + OutputDefinition(name="prediction_type", kind=[PREDICTION_TYPE_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + return [ + "easy_ocr/english_g2", + "easy_ocr/japanese_g2", + "easy_ocr/kannada_g2", + "easy_ocr/korean_g2", + "easy_ocr/latin_g2", + "easy_ocr/telugu_g2", + "easy_ocr/zh_sim_g2", + ] + + +class EasyOCRBlockV1(WorkflowBlock): + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + language: LANGUAGES = "English", + quantize: bool = False, + ) -> BlockResult: + + if language not in MODELS: + raise ValueError(f"Unsupported language: {language}") + + version, language_codes = MODELS.get(language, "english_g2") + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + language_codes=language_codes, + version=version, + quantize=quantize, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + language_codes=language_codes, + version=version, + quantize=quantize, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + language_codes: List[str], + version: str = "english_g2", + quantize: bool = False, + ) -> BlockResult: + model_id = f"easy_ocr/{version}" + self._model_manager.add_model( + model_id, + self._api_key, + endpoint_type=ModelEndpointType.CORE_MODEL, + ) + results = [] + for single_image in images: + if single_image.is_tensor_materialised(): + model_image, image_color_format = single_image.tensor_image, "rgb" + else: + model_image, image_color_format = single_image.numpy_image, "bgr" + texts, dets = self._model_manager.run_tensor_native_inference( + model_id, + images=[model_image], + input_color_format=image_color_format, + confidence=0.0, # no confidence filtering + ) + detections = attach_native_detection_metadata( + dets[0], + single_image, + class_names=CLASS_NAMES, + prediction_type=PREDICTION_TYPE, + ) + detections = _attach_local_ocr_text(detections) + results.append( + { + "result": texts[0], + "predictions": detections, + "parent_id": single_image.parent_metadata.parent_id, + "root_parent_id": single_image.workflow_root_ancestor_metadata.parent_id, + "prediction_type": PREDICTION_TYPE, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + language_codes: List[str], + version: str = "english_g2", + quantize: bool = False, + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + configuration = InferenceConfiguration( + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + ) + client.configure(configuration) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.ocr_image( + inference_input=non_empty_inference_images, + model="easy_ocr", + version=version, + quantize=quantize, + language_codes=language_codes, + ) + if len(images) == 1: + predictions = [predictions] + # Remote returns OCRInferenceResponse dicts: result text + ObjectDetectionPrediction list. + results = [] + for single_image, prediction in zip(images, predictions): + raw_predictions = prediction.get("predictions", []) or [] + detections = native_detections_from_inference_predictions( + single_image, + predictions=raw_predictions, + prediction_type=PREDICTION_TYPE, + class_names=CLASS_NAMES, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + detections = _attach_remote_ocr_text(detections, raw_predictions) + results.append( + { + "result": prediction.get("result", ""), + "predictions": detections, + "parent_id": single_image.parent_metadata.parent_id, + "root_parent_id": single_image.workflow_root_ancestor_metadata.parent_id, + "prediction_type": PREDICTION_TYPE, + } + ) + return results + + +def _attach_local_ocr_text(detections: Detections) -> Detections: + """Byte-parity (`class`) for the LOCAL path. + + The EasyOCR inference_models model carries each box's recognised text on + ``bboxes_metadata[i]["text"]`` (and ``class_id == 0`` for every box). The tensor + serializer emits ``bboxes_metadata[i][CLASS_NAME_KEY]`` as the per-box ``class`` + when present, else it falls back to the fixed ``CLASS_NAMES`` label. So promote the + recognised ``text`` onto ``CLASS_NAME_KEY`` to match the numpy sibling, whose + ``sv.Detections.from_inference`` carries the recognised text as the class name. + Mutates and returns the same object. + """ + if detections.bboxes_metadata is not None: + for entry in detections.bboxes_metadata: + entry[CLASS_NAME_KEY] = str(entry.get("text", "")) + return detections + + +def _attach_remote_ocr_text( + detections: Detections, predictions: List[dict] +) -> Detections: + """Byte-parity (`class`) for the REMOTE path. + + The inference OCR response carries each box's recognised text under ``class``, but + ``native_detections_from_inference_predictions`` keeps only ``detection_id`` on + ``bboxes_metadata``. Re-attach the per-box ``class`` (mirrors ``pp_ocr/v1_tensor``) + so the serializer emits the recognised text, matching the numpy sibling. Mutates + and returns the same object. + """ + if detections.bboxes_metadata is not None: + for entry, prediction in zip(detections.bboxes_metadata, predictions): + entry[CLASS_NAME_KEY] = str(prediction.get(CLASS_NAME_KEY, "")) + return detections diff --git a/inference/core/workflows/core_steps/models/foundation/florence2/v1.py b/inference/core/workflows/core_steps/models/foundation/florence2/v1.py index 10ea7e4d00..76c2d27c86 100644 --- a/inference/core/workflows/core_steps/models/foundation/florence2/v1.py +++ b/inference/core/workflows/core_steps/models/foundation/florence2/v1.py @@ -638,7 +638,7 @@ def _prepare_grounding_bounding_box_from_detections( COORDINATES_EXTRACTION = { "first": lambda detections: detections.xyxy[0].tolist(), - "last": lambda detections: detections.xyxy[0].tolist(), + "last": lambda detections: detections.xyxy[-1].tolist(), "biggest": lambda detections: detections.xyxy[np.argmax(detections.area)].tolist(), "smallest": lambda detections: detections.xyxy[np.argmin(detections.area)].tolist(), "most-confident": lambda detections: detections.xyxy[ diff --git a/inference/core/workflows/core_steps/models/foundation/florence2/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/florence2/v1_tensor.py new file mode 100644 index 0000000000..f8db934476 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/florence2/v1_tensor.py @@ -0,0 +1,375 @@ +"""Tensor-native `roboflow_core/florence_2@v1` block. + +Florence-2 outputs (`raw_output` / `parsed_output` / `classes`) are text/dict/list, +not prediction kinds; the only tensor-native surface is the `grounding_detection` +input, which arrives in the `TensorNativeGrounding` shapes. +""" + +import json +from typing import List, Optional, Tuple, Union + +import torch + +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + split_key_point_prediction, +) +from inference.core.workflows.core_steps.models.foundation.florence2.v1 import ( + TASK_TYPE_TO_FLORENCE_TASK, + TASKS_REQUIRING_DETECTION_GROUNDING, + TASKS_TO_EXTRACT_LABELS_AS_CLASSES, + BlockManifest, + GroundingSelectionMode, + TaskType, + _coordinate_to_loc, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlock +from inference_models import Detections, InstanceDetections, KeyPoints +from inference_sdk import InferenceHTTPClient + +# inference_models native prediction shapes accepted on the grounding input. +TensorNativeGrounding = Union[ + Detections, # OD + InstanceDetections, # IS + Tuple[KeyPoints, Optional[Detections]], # KP dual representation +] + + +class Florence2BlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls): + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_version: str, + task_type: TaskType, + prompt: Optional[str], + classes: Optional[List[str]], + grounding_detection: Optional[ + Union[Batch[TensorNativeGrounding], List[int], List[float]] + ], + grounding_selection_mode: GroundingSelectionMode, + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + task_type=task_type, + model_version=model_version, + prompt=prompt, + classes=classes, + grounding_detection=grounding_detection, + grounding_selection_mode=grounding_selection_mode, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + task_type=task_type, + model_version=model_version, + prompt=prompt, + classes=classes, + grounding_detection=grounding_detection, + grounding_selection_mode=grounding_selection_mode, + ) + raise ValueError(f"Unknown step execution mode: {self._step_execution_mode}") + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_version: str, + task_type: TaskType, + prompt: Optional[str], + classes: Optional[List[str]], + grounding_detection: Optional[ + Union[Batch[TensorNativeGrounding], List[int], List[float]] + ], + grounding_selection_mode: GroundingSelectionMode, + ) -> BlockResult: + requires_detection_grounding = task_type in TASKS_REQUIRING_DETECTION_GROUNDING + is_not_florence_task = task_type == "custom" + florence_task = TASK_TYPE_TO_FLORENCE_TASK[task_type] + + prompts = _build_prompts( + images=images, + classes=classes, + base_prompt=prompt, + grounding_detection=grounding_detection, + grounding_selection_mode=grounding_selection_mode, + ) + if classes is None: + classes = [] + + self._model_manager.add_model(model_id=model_version, api_key=self._api_key) + + predictions = [] + for image, single_prompt in zip(images, prompts): + if single_prompt is None and requires_detection_grounding: + predictions.append( + {"raw_output": None, "parsed_output": None, "classes": None} + ) + continue + if is_not_florence_task: + final_prompt = single_prompt or "" + else: + final_prompt = florence_task + (single_prompt or "") + + # The adapter derives the Florence task from the prompt string. + if image.is_tensor_materialised(): + model_image, image_color_format = image.tensor_image, "rgb" + else: + model_image, image_color_format = image.numpy_image, "bgr" + result = self._model_manager.run_tensor_native_inference( + model_version, + images=[model_image], + input_color_format=image_color_format, + prompt=final_prompt, + ) + # `run_tensor_native_inference` -> `Florence2HF.prompt` returns one + # `post_process_generation` dict per image, each shaped + # `{task: answer}`. Unwrap the per-task answer exactly like the numpy + # sibling (v1.py `run_locally`) so `raw_output`/`parsed_output`/`classes` + # match flag-off; leaving the wrapper on makes `raw_output` a wrapper + # dict and `.get("labels")` miss the nested labels (classes == []). + prediction_dict = result[0] + if is_not_florence_task: + prediction_data = prediction_dict[list(prediction_dict.keys())[0]] + else: + prediction_data = prediction_dict[florence_task] + + extracted_classes = classes + if florence_task in TASKS_TO_EXTRACT_LABELS_AS_CLASSES and isinstance( + prediction_data, dict + ): + extracted_classes = prediction_data.get("labels", []) + + predictions.append( + { + "raw_output": json.dumps(prediction_data), + "parsed_output": ( + prediction_data if isinstance(prediction_data, dict) else None + ), + "classes": extracted_classes, + } + ) + return predictions + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_version: str, + task_type: TaskType, + prompt: Optional[str], + classes: Optional[List[str]], + grounding_detection: Optional[ + Union[Batch[TensorNativeGrounding], List[int], List[float]] + ], + grounding_selection_mode: GroundingSelectionMode, + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient(api_url=api_url, api_key=self._api_key) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + requires_detection_grounding = task_type in TASKS_REQUIRING_DETECTION_GROUNDING + is_not_florence_task = task_type == "custom" + florence_task = TASK_TYPE_TO_FLORENCE_TASK[task_type] + + prompts = _build_prompts( + images=images, + classes=classes, + base_prompt=prompt, + grounding_detection=grounding_detection, + grounding_selection_mode=grounding_selection_mode, + ) + if classes is None: + classes = [] + + predictions = [] + for image, single_prompt in zip(images, prompts): + if single_prompt is None and requires_detection_grounding: + predictions.append( + {"raw_output": None, "parsed_output": None, "classes": None} + ) + continue + if is_not_florence_task: + final_prompt = single_prompt or "" + else: + final_prompt = florence_task + (single_prompt or "") + + result = client.infer_lmm( + inference_input=image.base64_image, + model_id=model_version, + prompt=final_prompt, + model_id_in_path=True, + ) + response = result.get("response", {}) + if is_not_florence_task: + if isinstance(response, dict) and len(response) > 0: + prediction_data = response[list(response.keys())[0]] + else: + prediction_data = response + else: + prediction_data = response.get(florence_task, response) + + extracted_classes = classes + if florence_task in TASKS_TO_EXTRACT_LABELS_AS_CLASSES and isinstance( + prediction_data, dict + ): + extracted_classes = prediction_data.get("labels", []) + + predictions.append( + { + "raw_output": json.dumps(prediction_data), + "parsed_output": ( + prediction_data if isinstance(prediction_data, dict) else None + ), + "classes": extracted_classes, + } + ) + return predictions + + +def _build_prompts( + images: Batch[WorkflowImageData], + classes: Optional[List[str]], + base_prompt: Optional[str], + grounding_detection, + grounding_selection_mode: GroundingSelectionMode, +) -> List[Optional[str]]: + if grounding_detection is not None: + return prepare_detection_grounding_prompts( + images=images, + grounding_detection=grounding_detection, + grounding_selection_mode=grounding_selection_mode, + ) + if classes is not None: + return ["".join(classes)] * len(images) + return [base_prompt] * len(images) + + +def prepare_detection_grounding_prompts( + images: Batch[WorkflowImageData], + grounding_detection, + grounding_selection_mode: GroundingSelectionMode, +) -> List[Optional[str]]: + if isinstance(grounding_detection, list): + return [ + _location_prompt_from_static_box( + image=image, bounding_box=grounding_detection + ) + for image in images + ] + return [ + _location_prompt_from_tensor_detections( + image=image, + prediction=prediction, + grounding_selection_mode=grounding_selection_mode, + ) + for image, prediction in zip(images, grounding_detection) + ] + + +def _location_prompt_from_tensor_detections( + image: WorkflowImageData, + prediction, + grounding_selection_mode: GroundingSelectionMode, +) -> Optional[str]: + # Pull the bbox-bearing Detections out of OD / IS / (KeyPoints, Detections). + _, detections = split_key_point_prediction(prediction) + if len(detections) == 0: + return None + height, width = image._read_shape_without_materialization() + left_top_x, left_top_y, right_bottom_x, right_bottom_y = _select_box( + detections=detections, mode=grounding_selection_mode + ) + return ( + f"" + f"" + f"" + f"" + ) + + +def _select_box(detections, mode: GroundingSelectionMode) -> List[float]: + xyxy = detections.xyxy + if mode == "first": + index = 0 + elif mode == "last": + index = len(detections) - 1 + elif mode in ("biggest", "smallest"): + area = (xyxy[:, 2] - xyxy[:, 0]) * ( + xyxy[:, 3] - xyxy[:, 1] + ) # no .area on inference_models + index = int(torch.argmax(area) if mode == "biggest" else torch.argmin(area)) + elif mode in ("most-confident", "least-confident"): + confidence = detections.confidence + index = int( + torch.argmax(confidence) + if mode == "most-confident" + else torch.argmin(confidence) + ) + else: + raise ValueError(f"Unknown grounding selection mode: {mode}") + return xyxy[index].detach().to("cpu").tolist() + + +def _location_prompt_from_static_box( + image: WorkflowImageData, + bounding_box: Union[List[float], List[int]], +) -> str: + height, width = image._read_shape_without_materialization() + coordinates = bounding_box[:4] + if len(coordinates) != 4: + raise ValueError( + "Could not extract 4 coordinates of bounding box to perform detection " + "grounded Florence 2 prediction." + ) + left_top_x, left_top_y, right_bottom_x, right_bottom_y = coordinates + if all(isinstance(c, float) for c in coordinates): + return ( + f"" + f"" + f"" + f"" + ) + if all(isinstance(c, int) for c in coordinates): + return ( + f"" + f"" + f"" + f"" + ) + raise ValueError( + "Provided coordinates in mixed format - coordinates must be all integers " + "or all floats in range [0.0-1.0]" + ) diff --git a/inference/core/workflows/core_steps/models/foundation/florence2/v2_tensor.py b/inference/core/workflows/core_steps/models/foundation/florence2/v2_tensor.py new file mode 100644 index 0000000000..f0009e59e1 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/florence2/v2_tensor.py @@ -0,0 +1,48 @@ +"""Tensor-native `roboflow_core/florence_2@v2` block.""" + +from typing import List, Optional, Type, Union + +from inference.core.workflows.core_steps.models.foundation.florence2.v1 import ( + GroundingSelectionMode, + TaskType, +) +from inference.core.workflows.core_steps.models.foundation.florence2.v1_tensor import ( + Florence2BlockV1, + TensorNativeGrounding, +) +from inference.core.workflows.core_steps.models.foundation.florence2.v2 import ( + V2BlockManifest, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + + +class Florence2BlockV2(Florence2BlockV1): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return V2BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + task_type: TaskType, + prompt: Optional[str], + classes: Optional[List[str]], + grounding_detection: Optional[ + Union[Batch[TensorNativeGrounding], List[int], List[float]] + ], + grounding_selection_mode: GroundingSelectionMode, + ) -> BlockResult: + return super().run( + images=images, + model_version=model_id, + task_type=task_type, + prompt=prompt, + classes=classes, + grounding_detection=grounding_detection, + grounding_selection_mode=grounding_selection_mode, + ) diff --git a/inference/core/workflows/core_steps/models/foundation/google_vision_ocr/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/google_vision_ocr/v1_tensor.py new file mode 100644 index 0000000000..d0d5126662 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/google_vision_ocr/v1_tensor.py @@ -0,0 +1,335 @@ +"""Tensor-native `roboflow_core/google_vision_ocr@v1` block.""" + +from typing import List, Literal, Optional, Type, Union +from uuid import uuid4 + +import requests +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.roboflow_api import post_to_roboflow_api +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + ROBOFLOW_MANAGED_KEY, + SECRET_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + AirGappedAvailability, + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections + +# The recognized text per block is carried per box in +# `bboxes_metadata[i][CLASS_NAME_KEY]`, which the serializer emits as +# `predictions[i]["class"]`; this image-level class map must stay non-empty as the +# serializer's class-id fallback (every box carries `class_id == 0`). +CLASS_NAMES = {0: "text-region"} +PREDICTION_TYPE = "ocr" + +LONG_DESCRIPTION = """ +Detect text in images using Google Vision OCR. + +Supported types of text detection: + +- `text_detection`: optimized for areas of text within a larger image. +- `ocr_text_detection`: optimized for dense text documents. + +Provide your Google Vision API key or set the value to ``rf_key:account`` (or +``rf_key:user:``) to proxy requests through Roboflow's API. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Google Vision OCR", + "version": "v1", + "short_description": "Detect text in images using Google Vision API", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "fa-brands fa-google", + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/google_vision_ocr@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Image to run OCR", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + ocr_type: Literal["text_detection", "ocr_text_detection"] = Field( + description="Type of OCR to use", + json_schema_extra={ + "values_metadata": { + "text_detection": { + "name": "Any Scene Text Detection", + "description": "Detects and extracts text from any image, including photographs that contain blocks of text.", + }, + "ocr_text_detection": { + "name": "Document Text Detection", + "description": "Optimized for dense text documents, such as scanned pages or photographs of printed text.", + }, + }, + }, + ) + api_key: Union[ + Selector(kind=[STRING_KIND, SECRET_KIND, ROBOFLOW_MANAGED_KEY]), str + ] = Field( + default="rf_key:account", + description="Your Google Vision API key", + examples=["xxx-xxx", "$inputs.google_api_key"], + private=True, + ) + language_hints: Optional[List[str]] = Field( + default=None, + description="Optional list of language codes to pass to the OCR API. If not provided, the API will attempt to detect the language automatically." + "If provided, language codes must be supported by the OCR API, visit https://cloud.google.com/vision/docs/languages for list of supported language codes.", + examples=[["en", "fr"], ["de"]], + ) + + @classmethod + def get_air_gapped_availability(cls) -> AirGappedAvailability: + return AirGappedAvailability(available=False, reason="requires_internet") + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="text", kind=[STRING_KIND]), + OutputDefinition(name="language", kind=[STRING_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.4.0,<2.0.0" + + +class GoogleVisionOCRBlockV1(WorkflowBlock): + + def __init__( + self, + api_key: Optional[str], + ): + self._roboflow_api_key = api_key + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["api_key"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + ocr_type: Literal["text_detection", "ocr_text_detection"], + language_hints: Optional[List[str]], + api_key: str = "rf_key:account", + ) -> BlockResult: + # Decide which type of OCR to use + if ocr_type == "text_detection": + detection_type = "TEXT_DETECTION" + elif ocr_type == "ocr_text_detection": + detection_type = "DOCUMENT_TEXT_DETECTION" + else: + raise ValueError(f"Invalid ocr_type: {ocr_type}") + + request_json = _build_request_json( + image=image, + detection_type=detection_type, + language_hints=language_hints, + ) + + # Route to proxy or direct API based on api_key format + if api_key.startswith(("rf_key:account", "rf_key:user:")): + result = _execute_proxied_google_vision_request( + roboflow_api_key=self._roboflow_api_key, + google_vision_api_key=api_key, + request_json=request_json, + ) + else: + result = _execute_google_vision_request( + api_key=api_key, + request_json=request_json, + ) + + return _parse_google_vision_response(result=result, image=image) + + +def _build_request_json( + image: WorkflowImageData, + detection_type: str, + language_hints: Optional[List[str]], +) -> dict: + ocr_request = { + "image": {"content": image.base64_image}, + "features": [{"type": detection_type}], + } + + if language_hints is not None: + ocr_request["imageContext"] = {"languageHints": language_hints} + + return {"requests": [ocr_request]} + + +def _execute_proxied_google_vision_request( + roboflow_api_key: str, + google_vision_api_key: str, + request_json: dict, +) -> dict: + payload = { + "google_vision_api_key": google_vision_api_key, + "request_json": request_json, + } + + try: + response_data = post_to_roboflow_api( + endpoint="apiproxy/google_vision_ocr", + api_key=roboflow_api_key, + payload=payload, + ) + return response_data["responses"][0] + except requests.exceptions.RequestException as e: + raise RuntimeError(f"Failed to connect to Roboflow proxy: {e}") from e + except (KeyError, IndexError) as e: + raise RuntimeError( + f"Invalid response structure from Roboflow proxy: {e}" + ) from e + + +def _execute_google_vision_request( + api_key: str, + request_json: dict, +) -> dict: + response = requests.post( + "https://vision.googleapis.com/v1/images:annotate", + params={"key": api_key}, + json=request_json, + ) + + if response.status_code != 200: + raise RuntimeError( + f"Request to Google Cloud Vision API failed: {str(response.json())}" + ) + + return response.json()["responses"][0] + + +def _parse_google_vision_response( + result: dict, + image: WorkflowImageData, +) -> BlockResult: + # Check for image without text + if "textAnnotations" not in result or not result["textAnnotations"]: + return { + "text": "", + "language": "", + "predictions": _build_native_detections( + image=image, xyxy=[], confidence=[], texts=[] + ), + } + + # Extract predictions from the response + text = result["textAnnotations"][0]["description"] + language = result["textAnnotations"][0]["locale"] + + xyxy: List[List[float]] = [] + confidence: List[float] = [] + texts: List[str] = [] + + for page in result["fullTextAnnotation"]["pages"]: + for block in page["blocks"]: + # Get bounding box coordinates + box = block["boundingBox"]["vertices"] + x_min = min(v.get("x", 0) for v in box) + y_min = min(v.get("y", 0) for v in box) + x_max = max(v.get("x", 0) for v in box) + y_max = max(v.get("y", 0) for v in box) + xyxy.append([float(x_min), float(y_min), float(x_max), float(y_max)]) + + # Only DOCUMENT_TEXT_DETECTION provides confidence score, use 1.0 otherwise + confidence.append(float(block.get("confidence", 1.0))) + + # Get block text + block_text = [] + for paragraph in block["paragraphs"]: + for word in paragraph["words"]: + word_text = "".join(symbol["text"] for symbol in word["symbols"]) + block_text.append(word_text) + texts.append(" ".join(block_text)) + + return { + "text": text, + "language": language, + "predictions": _build_native_detections( + image=image, xyxy=xyxy, confidence=confidence, texts=texts + ), + } + + +def _build_native_detections( + image: WorkflowImageData, + xyxy: List[List[float]], + confidence: List[float], + texts: List[str], +) -> Detections: + """Build an `inference_models.Detections` from the parsed Google Vision blocks.""" + number_of_detections = len(xyxy) + detections = Detections( + xyxy=( + torch.tensor( + xyxy, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + if number_of_detections + else torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + ), + class_id=torch.zeros( + (number_of_detections,), + dtype=torch.int64, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.tensor( + confidence, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=CLASS_NAMES, + prediction_type=PREDICTION_TYPE, + ) + if number_of_detections == 0: + detections.bboxes_metadata = None + return detections + detections.bboxes_metadata = [ + {DETECTION_ID_KEY: str(uuid4()), CLASS_NAME_KEY: texts[index]} + for index in range(number_of_detections) + ] + return detections diff --git a/inference/core/workflows/core_steps/models/foundation/moondream2/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/moondream2/v1_tensor.py new file mode 100644 index 0000000000..a4b5f09c99 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/moondream2/v1_tensor.py @@ -0,0 +1,260 @@ +from typing import Dict, List, Literal, Optional, Type, Union + +from pydantic import ConfigDict, Field + +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + MOONDREAM2_ENABLED, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + native_detections_from_inference_predictions, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + ROBOFLOW_MODEL_ID_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, +) +from inference_sdk import InferenceHTTPClient + + +class BlockManifest(WorkflowBlockManifest): + # SmolVLM needs an image and a text prompt. + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + prompt: Union[ + Selector(kind=[STRING_KIND]), + str, + ] = Field( + description="Optional text prompt to provide additional context to Moondream2.", + examples=["my prompt", "$inputs.prompt"], + default=None, + ) + + # Standard model configuration for UI, schema, etc. + model_config = ConfigDict( + json_schema_extra={ + "name": "Moondream2", + "version": "v1", + "short_description": "Run Moondream2 on an image.", + "long_description": "This workflow block runs Moondream2, a multimodal vision-language model. You can use this block to run zero-shot object detection.", + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": [ + "Moondream2", + "moondream", + "vision language model", + "VLM", + "object detection", + ], + "ui_manifest": { + "section": "model", + "icon": "fal fa-atom", + "blockPriority": 5.5, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/moondream2@v1"] + + model_version: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = Field( + default="moondream2/moondream2_2b_jul24", + description="The Moondream2 model to be used for inference.", + examples=["moondream2/moondream2_2b_jul24", "moondream2/moondream2-2b"], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restrictions = [ + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; run_locally() loads a model that needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + ] + if not MOONDREAM2_ENABLED: + restrictions.append( + RuntimeRestriction( + severity=Severity.HARD, + note=( + "MOONDREAM2_ENABLED=False on Roboflow Hosted Serverless: " + "the Moondream2 endpoint is not registered, so " + "run_remotely() returns 404." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + ) + return restrictions + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + return ["moondream2/moondream2_2b_jul24"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + return [roboflow_platform_model(model_id=self.model_version)] + + +class Moondream2BlockV1(WorkflowBlock): + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + ) -> BlockResult: + if self._step_execution_mode == StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_version=model_version, + prompt=prompt, + ) + elif self._step_execution_mode == StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_version=model_version, + prompt=prompt, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + prompt = prompt or "" + class_names = _build_class_names(prompt) + predictions = [] + for image in images: + result = client.infer_lmm( + inference_input=image.base64_image, + model_id=model_version, + prompt=prompt, + model_id_in_path=True, + ) + # `result["predictions"]` holds inference detection dicts in center + # x/y/width/height format with `class` / `class_id` set. + detections = native_detections_from_inference_predictions( + image=image, + predictions=result.get("predictions", []), + prediction_type="object-detection", + class_names=class_names, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + predictions.append({"predictions": detections}) + return predictions + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + ) -> BlockResult: + # Use the provided prompt (or an empty string if None) for every image. + prompt = prompt or "" + class_names = _build_class_names(prompt) + + # Register Moondream2 with the model manager. + self._model_manager.add_model(model_id=model_version, api_key=self._api_key) + + predictions = [] + for image in images: + if image.is_tensor_materialised(): + model_image, image_color_format = image.tensor_image, "rgb" + else: + model_image, image_color_format = image.numpy_image, "bgr" + dets = self._model_manager.run_tensor_native_inference( + model_id=model_version, + images=[model_image], + input_color_format=image_color_format, + classes=[prompt], + ) + detections = attach_native_detection_metadata( + dets[0], + image=image, + class_names=class_names, + prediction_type="object-detection", + ) + predictions.append({"predictions": detections}) + return predictions + + +def _build_class_names(prompt: str) -> Dict[int, str]: + # Moondream2 detection is single-object prompted: class_id 0 -> prompt text. + return {0: prompt} diff --git a/inference/core/workflows/core_steps/models/foundation/ocr/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/ocr/v1_tensor.py new file mode 100644 index 0000000000..124ad8df0e --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/ocr/v1_tensor.py @@ -0,0 +1,232 @@ +from typing import Dict, List, Literal, Optional, Type + +from pydantic import ConfigDict, Field + +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import ModelEndpointType +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + native_detections_from_inference_predictions, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + PARENT_ID_KIND, + PREDICTION_TYPE_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +LONG_DESCRIPTION = """ + Retrieve the characters in an image using DocTR Optical Character Recognition (OCR). + +This block returns the text within an image. + +You may want to use this block in combination with a detections-based block (i.e. +ObjectDetectionBlock). An object detection model could isolate specific regions from an +image (i.e. a shipping container ID in a logistics use case) for further processing. +You can then use a DynamicCropBlock to crop the region of interest before running OCR. + +Using a detections model then cropping detections allows you to isolate your analysis +on particular regions of an image. +""" + +# Must match the inference_models DocTR model's `class_names` order; the serializer +# resolves each detection's class from this map on image_metadata. +DOCTR_CLASS_NAMES: Dict[int, str] = {0: "block", 1: "line", 2: "word"} + +PREDICTION_TYPE = "ocr" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "OCR Model", + "version": "v1", + "short_description": "Extract text from an image using DocTR optical character recognition.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-text", + "blockPriority": 11, + "inference": True, + }, + } + ) + type: Literal["roboflow_core/ocr_model@v1", "OCRModel"] + name: str = Field(description="Unique name of step in workflows") + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="result", kind=[STRING_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition(name="parent_id", kind=[PARENT_ID_KIND]), + OutputDefinition(name="root_parent_id", kind=[PARENT_ID_KIND]), + OutputDefinition(name="prediction_type", kind=[PREDICTION_TYPE_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + return ["doctr/default"] + + +class OCRModelBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally(images=images) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely(images=images) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + ) -> BlockResult: + doctr_model_id = "doctr/default" + self._model_manager.add_model( + doctr_model_id, + self._api_key, + endpoint_type=ModelEndpointType.CORE_MODEL, + ) + predictions = [] + for single_image in images: + if single_image.is_tensor_materialised(): + model_image, image_color_format = single_image.tensor_image, "rgb" + else: + model_image, image_color_format = single_image.numpy_image, "bgr" + # The returned Detections carry per-box {"text": ...} on bboxes_metadata. + texts, detections_batch = self._model_manager.run_tensor_native_inference( + doctr_model_id, + images=[model_image], + input_color_format=image_color_format, + ) + detections = attach_native_detection_metadata( + detections=detections_batch[0], + image=single_image, + class_names=DOCTR_CLASS_NAMES, + prediction_type=PREDICTION_TYPE, + ) + predictions.append( + { + "result": texts[0], + "predictions": detections, + "parent_id": single_image.parent_metadata.parent_id, + "root_parent_id": single_image.workflow_root_ancestor_metadata.parent_id, + "prediction_type": PREDICTION_TYPE, + } + ) + return predictions + + def run_remotely( + self, + images: Batch[WorkflowImageData], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + configuration = InferenceConfiguration( + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + ) + client.configure(configuration) + non_empty_inference_images = [i.base64_image for i in images] + responses = client.ocr_image( + inference_input=non_empty_inference_images, + generate_bounding_boxes=True, + ) + if len(images) == 1: + responses = [responses] + # The "predictions" key may be absent/empty when generate_bounding_boxes + # yields nothing. + predictions = [] + for single_image, response in zip(images, responses): + raw_predictions = response.get("predictions") or [] + detections = native_detections_from_inference_predictions( + image=single_image, + predictions=raw_predictions, + prediction_type=PREDICTION_TYPE, + class_names=DOCTR_CLASS_NAMES, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + predictions.append( + { + "result": response.get("result", ""), + "predictions": detections, + "parent_id": single_image.parent_metadata.parent_id, + "root_parent_id": single_image.workflow_root_ancestor_metadata.parent_id, + "prediction_type": PREDICTION_TYPE, + } + ) + return predictions diff --git a/inference/core/workflows/core_steps/models/foundation/perception_encoder/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/perception_encoder/v1_tensor.py new file mode 100644 index 0000000000..48e42f0ab9 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/perception_encoder/v1_tensor.py @@ -0,0 +1,238 @@ +import hashlib +from typing import List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field + +from inference.core.cache.lru_cache import LRUCache +from inference.core.env import ( + CORE_MODEL_PE_ENABLED, + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import ModelEndpointType +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_EMBEDDING_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_sdk import InferenceHTTPClient + +LONG_DESCRIPTION = """ +Use the Meta Perception Encoder model to create semantic embeddings of text and images. + +This block accepts an image or string and returns an embedding. The embedding can be used to compare +similarity between different images or between images and text. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Perception Encoder Embedding Model", + "version": "v1", + "short_description": "Generate an embedding of an image or string.", + "long_description": LONG_DESCRIPTION, + "license": "MIT", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-paperclip", + "blockPriority": 9.9, + }, + } + ) + type: Literal["roboflow_core/perception_encoder@v1"] + name: str = Field(description="Unique name of step in workflows") + data: Union[Selector(kind=[IMAGE_KIND, STRING_KIND]), str] = Field( + title="Data", + description="The string or image to generate an embedding for.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + version: Union[ + Literal[ + "PE-Core-B16-224", + "PE-Core-L14-336", + "PE-Core-G14-448", + ], + Selector(kind=[STRING_KIND]), + ] = Field( + default="PE-Core-L14-336", + description="Variant of Perception Encoder model", + examples=["PE-Core-B16-224", "$inputs.variant"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [OutputDefinition(name="embedding", kind=[TENSOR_NATIVE_EMBEDDING_KIND])] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restrictions = [ + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; run_locally() loads a model that needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + ] + if not CORE_MODEL_PE_ENABLED: + restrictions.append( + RuntimeRestriction( + severity=Severity.HARD, + note=( + "CORE_MODEL_PE_ENABLED=False on Roboflow Hosted Serverless: " + "the Perception Encoder endpoint is not registered, so " + "run_remotely() returns 404." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + ) + return restrictions + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + return [ + "perception_encoder/PE-Core-B16-224", + "perception_encoder/PE-Core-L14-336", + "perception_encoder/PE-Core-G14-448", + ] + + +text_cache = LRUCache() + + +class PerceptionEncoderModelBlockV1(WorkflowBlock): + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + data: Union[WorkflowImageData, str], + version: str, + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally(data=data, version=version) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely(data=data, version=version) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + data: Union[WorkflowImageData, str], + version: str, + ) -> BlockResult: + # Model registration happens after the text-cache check so a cache hit + # never loads the model. + pe_model_id = f"perception_encoder/{version}" + if isinstance(data, str): + hash_key = hashlib.md5((version + data).encode("utf-8")).hexdigest() + cached_value = text_cache.get(hash_key) + if cached_value is not None: + # Cache holds CPU tensors so the process-wide LRU never pins GPU memory. + return {"embedding": cached_value.to(WORKFLOWS_IMAGE_TENSOR_DEVICE)} + self._model_manager.add_model( + pe_model_id, + self._api_key, + endpoint_type=ModelEndpointType.CORE_MODEL, + ) + embeddings = self._model_manager.run_tensor_native_inference( + pe_model_id, + action="embed-text", + texts=[data], + ) + embedding = embeddings[0] + text_cache.set(hash_key, embedding.detach().cpu()) + return {"embedding": embedding.to(WORKFLOWS_IMAGE_TENSOR_DEVICE)} + else: + self._model_manager.add_model( + pe_model_id, + self._api_key, + endpoint_type=ModelEndpointType.CORE_MODEL, + ) + if data.is_tensor_materialised(): + model_image, image_color_format = data.tensor_image, "rgb" + else: + model_image, image_color_format = data.numpy_image, "bgr" + embeddings = self._model_manager.run_tensor_native_inference( + pe_model_id, + action="embed-image", + images=[model_image], + input_color_format=image_color_format, + ) + return {"embedding": embeddings[0].to(WORKFLOWS_IMAGE_TENSOR_DEVICE)} + + def run_remotely( + self, + data: Union[WorkflowImageData, str], + version: str, + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient(api_url=api_url, api_key=self._api_key) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + if isinstance(data, str): + result = client.get_perception_encoder_text_embeddings( + text=data, + perception_encoder_version=version, + ) + else: + result = client.get_perception_encoder_image_embeddings( + inference_input=data.base64_image, + perception_encoder_version=version, + ) + # Remote returns embeddings as a JSON List[List[float]]. + return { + "embedding": torch.tensor( + result["embeddings"][0], + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + } diff --git a/inference/core/workflows/core_steps/models/foundation/pp_ocr/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/pp_ocr/v1_tensor.py new file mode 100644 index 0000000000..a87f93ccf1 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/pp_ocr/v1_tensor.py @@ -0,0 +1,329 @@ +"""Tensor-native sibling of `roboflow_core/pp_ocr@v1`. + +The loader import-swaps `pp_ocr.v1.PPOCRBlockV1` for this class when +`ENABLE_TENSOR_DATA_REPRESENTATION` is enabled. The manifest is identical to the +numpy sibling (same `type`/version, params, validators); only the `predictions` +output kind changes to `TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND` and the +run-body emits a native `inference_models.Detections` instead of `sv.Detections`. + +PP-OCR has no `run_tensor_native_inference` adapter, so both execution modes keep +the exact model-calling machinery of `pp_ocr.v1` (a `PPOCRInferenceRequest` served +by the local `ModelManager`, or `InferenceHTTPClient.ocr_image` remotely) and +produce the standard inference-format `OCRInferenceResponse` dicts. Those dicts are +converted to a native `Detections` HERE (no shared util is touched) so the numpy +`post_process_ocr_result` / `sv.Detections.from_inference` path is bypassed. + +Byte-parity note (`class`): a PP-OCR box carries the recognised line text under +`class` with a fixed `class_id == 0`. `native_detections_from_inference_predictions` +folds `class` into a `class_id -> name` map, which for OCR would collapse every box +to a single label. So the recognised text is re-attached per box onto +`bboxes_metadata[i][CLASS_NAME_KEY]`; the tensor serializer's per-box override then +emits it as `class`, matching the numpy path's per-box `sv.Detections` class name. +The image-level `class_names` map is only a non-empty fallback required by the +serializer (never consulted, since every box carries the per-box override). +""" + +from typing import List, Literal, Optional, Type + +from pydantic import ConfigDict, Field, model_validator + +from inference.core.entities.requests.pp_ocr import PPOCRInferenceRequest +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + native_detections_from_inference_predictions, +) +from inference.core.workflows.core_steps.common.utils import load_core_model +from inference.core.workflows.execution_engine.constants import CLASS_NAME_KEY +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + PARENT_ID_KIND, + PREDICTION_TYPE_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections +from inference_sdk import InferenceHTTPClient +from inference_sdk.http.entities import InferenceConfiguration + +# Fixed, non-empty `class_id -> name` fallback map. PP-OCR emits `class_id == 0` for +# every box; the recognised text is attached per box (see module docstring), so this +# map only satisfies the tensor serializer's "CLASS_NAMES_KEY must be present" check. +CLASS_NAMES = {0: "text-region"} +PREDICTION_TYPE = "ocr" + +LONG_DESCRIPTION = """ + Retrieve the characters in an image using PP-OCR (PaddleOCR) Optical Character Recognition (OCR). + +This block returns the text within an image. + +You may want to use this block in combination with a detections-based block (i.e. +ObjectDetectionBlock). An object detection model could isolate specific regions from an +image (i.e. a shipping container ID in a logistics use case) for further processing. +You can then use a DynamicCropBlock to crop the region of interest before running OCR. + +Using a detections model then cropping detections allows you to isolate your analysis +on particular regions of an image. +""" + + +class BlockManifest(WorkflowBlockManifest): + + model_config = ConfigDict( + json_schema_extra={ + "name": "PP-OCR", + "version": "v1", + "short_description": "Extract text from an image using PP-OCR (PaddleOCR) optical character recognition.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["ocr", "text", "paddle", "paddleocr"], + "ui_manifest": { + "section": "model", + "icon": "far fa-text", + "blockPriority": 11, + "inDevelopment": False, + "inference": True, + }, + } + ) + type: Literal["roboflow_core/pp_ocr@v1"] + name: str = Field(description="Unique name of step in workflows") + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + text_detection: Literal["none", "tiny", "small", "medium"] = Field( + title="Text Detection Model Size", + description="Text detection model size. Set to `none` to disable detection and run " + "recognition on each input image as a single, pre-cropped text line.", + default="small", + ) + text_recognition: Literal["none", "tiny", "small", "medium"] = Field( + title="Text Recognition Model Size", + description="Text recognition model size. Set to `none` to disable recognition and " + "return detected text boxes only, without transcribed text.", + default="small", + ) + + @model_validator(mode="after") + def validate_stages(self) -> "BlockManifest": + if self.text_detection == "none" and self.text_recognition == "none": + raise ValueError( + "PP-OCR requires at least one of text detection or text recognition to be enabled" + ) + return self + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="result", kind=[STRING_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition(name="parent_id", kind=[PARENT_ID_KIND]), + OutputDefinition(name="root_parent_id", kind=[PARENT_ID_KIND]), + OutputDefinition(name="prediction_type", kind=[PREDICTION_TYPE_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block. + + These are core-model ids in the standard ``/`` + form; for PP-OCR the version token encodes the two stage sizes as + ``-`` and is parsed back into the + stage sizes by ``PPOCRInferenceRequest``. + """ + return [ + "pp_ocr/small-small", + "pp_ocr/tiny-tiny", + "pp_ocr/medium-medium", + "pp_ocr/small-none", + "pp_ocr/none-small", + ] + + +class PPOCRBlockV1(WorkflowBlock): + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + text_detection: str = "small", + text_recognition: str = "small", + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + text_detection=text_detection, + text_recognition=text_recognition, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + text_detection=text_detection, + text_recognition=text_recognition, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + text_detection: str = "small", + text_recognition: str = "small", + ) -> BlockResult: + predictions = [] + for single_image in images: + inference_request = PPOCRInferenceRequest( + text_detection=text_detection, + text_recognition=text_recognition, + image=single_image.to_inference_format(numpy_preferred=True), + api_key=self._api_key, + ) + model_id = load_core_model( + model_manager=self._model_manager, + inference_request=inference_request, + core_model="pp_ocr", + ) + result = self._model_manager.infer_from_request_sync( + model_id, inference_request + ) + predictions.append( + _build_native_prediction( + image=single_image, + response=result.model_dump(by_alias=True, exclude_none=True), + ) + ) + return predictions + + def run_remotely( + self, + images: Batch[WorkflowImageData], + text_detection: str = "small", + text_recognition: str = "small", + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + configuration = InferenceConfiguration( + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + ) + client.configure(configuration) + non_empty_inference_images = [i.base64_image for i in images] + responses = client.ocr_image( + inference_input=non_empty_inference_images, + model="pp_ocr", + version=f"{text_detection}-{text_recognition}", + ) + if len(images) == 1: + responses = [responses] + return [ + _build_native_prediction(image=single_image, response=response) + for single_image, response in zip(images, responses) + ] + + +def _build_native_prediction( + image: WorkflowImageData, + response: dict, +) -> dict: + """Convert a single-image ``OCRInferenceResponse`` dict into the block output. + + Mirrors the output dict shape of the numpy ``post_process_ocr_result`` + (``result`` / ``predictions`` / ``parent_id`` / ``root_parent_id`` / + ``prediction_type``) but with ``predictions`` as a native ``Detections``. + """ + raw_predictions = response.get("predictions") or [] + detections = _native_detections_from_ocr_predictions( + image=image, + predictions=raw_predictions, + ) + return { + "result": response.get("result", ""), + "predictions": detections, + "parent_id": image.parent_metadata.parent_id, + "root_parent_id": image.workflow_root_ancestor_metadata.parent_id, + "prediction_type": PREDICTION_TYPE, + } + + +def _native_detections_from_ocr_predictions( + image: WorkflowImageData, + predictions: List[dict], +) -> Detections: + """Build a native ``Detections`` from PP-OCR inference-format boxes. + + Reuses the shared ``native_detections_from_inference_predictions`` (center->corner + conversion, ``detection_id`` preservation, image lineage) with the fixed + ``CLASS_NAMES`` fallback map, then re-attaches each box's recognised text onto + ``bboxes_metadata[i][CLASS_NAME_KEY]`` so the tensor serializer emits it as the + per-box ``class`` (byte-parity with the numpy path). + """ + detections = native_detections_from_inference_predictions( + image=image, + predictions=predictions, + prediction_type=PREDICTION_TYPE, + class_names=CLASS_NAMES, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + if detections.bboxes_metadata is not None: + for entry, prediction in zip(detections.bboxes_metadata, predictions): + entry[CLASS_NAME_KEY] = str(prediction.get(CLASS_NAME_KEY, "")) + return detections diff --git a/inference/core/workflows/core_steps/models/foundation/qwen3_5vl/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/qwen3_5vl/v1_tensor.py new file mode 100644 index 0000000000..21ccea6224 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/qwen3_5vl/v1_tensor.py @@ -0,0 +1,174 @@ +"""Tensor-native `roboflow_core/qwen3_5vl@v1` block.""" + +from typing import List, Optional, Type + +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1 import ( + BlockManifest, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_sdk import InferenceHTTPClient + + +class Qwen35VLBlockV1(WorkflowBlock): + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + system_prompt: Optional[str], + enable_thinking: bool = False, + max_new_tokens: Optional[int] = None, + ) -> BlockResult: + if self._step_execution_mode == StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_version=model_version, + prompt=prompt, + system_prompt=system_prompt, + enable_thinking=enable_thinking, + max_new_tokens=max_new_tokens, + ) + elif self._step_execution_mode == StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_version=model_version, + prompt=prompt, + system_prompt=system_prompt, + enable_thinking=enable_thinking, + max_new_tokens=max_new_tokens, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + system_prompt: Optional[str], + enable_thinking: bool = False, + max_new_tokens: Optional[int] = None, + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + prompt = prompt or "Describe what's in this image." + system_prompt = ( + system_prompt + or "You are a Qwen3.5-VL model that can answer questions about any image." + ) + combined_prompt = prompt + "" + system_prompt + + predictions = [] + for image in images: + result = client.infer_lmm( + inference_input=image.base64_image, + model_id=model_version, + prompt=combined_prompt, + model_id_in_path=True, + enable_thinking=enable_thinking, + max_new_tokens=max_new_tokens, + ) + response_text = result.get("response", result) + predictions.append({"parsed_output": response_text, "thinking": ""}) + + return predictions + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + system_prompt: Optional[str], + enable_thinking: bool = False, + max_new_tokens: Optional[int] = None, + ) -> BlockResult: + prompt = prompt or "Describe what's in this image." + system_prompt = system_prompt or "You are a helpful assistant." + # The model splits the combined prompt back apart on "". + combined_prompt = prompt + "" + system_prompt + + self._model_manager.add_model(model_id=model_version, api_key=self._api_key) + + predictions = [] + for image in images: + kwargs = dict( + prompt=combined_prompt, + enable_thinking=enable_thinking, + ) + if max_new_tokens is not None: + kwargs["max_new_tokens"] = max_new_tokens + if image.is_tensor_materialised(): + model_image, image_color_format = image.tensor_image, "rgb" + else: + model_image, image_color_format = image.numpy_image, "bgr" + result = self._model_manager.run_tensor_native_inference( + model_version, + images=[model_image], + input_color_format=image_color_format, + **kwargs, + ) + # One entry per image: str, or a {"thinking", "answer"} dict when + # thinking is enabled. + response_text = result[0] + if enable_thinking and isinstance(response_text, dict): + thinking = response_text.get("thinking", "") + answer = response_text.get("answer", "") + predictions.append( + { + "parsed_output": answer, + "thinking": thinking, + } + ) + else: + predictions.append( + { + "parsed_output": response_text, + "thinking": "", + } + ) + return predictions diff --git a/inference/core/workflows/core_steps/models/foundation/qwen3_5vl/v2_tensor.py b/inference/core/workflows/core_steps/models/foundation/qwen3_5vl/v2_tensor.py new file mode 100644 index 0000000000..888b605d42 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/qwen3_5vl/v2_tensor.py @@ -0,0 +1,151 @@ +"""Tensor-native sibling of `roboflow_core/qwen3_5vl@v2`. + +SCRATCH โ€” first pass for review. v2 is a variant of v1 with a different model +selector / variant list and NO thinking mode: it always runs with +`enable_thinking=False` and emits only `parsed_output` (no `thinking` output). + +Following florence2/v2_tensor: subclass the tensor-native `Qwen35VLBlockV1` +(from `v1_tensor`) and reuse the verbatim v2 `BlockManifest`. v2's run signature +has no `enable_thinking`, and its output dict has no `thinking` key, so the +`run` / `run_locally` / `run_remotely` shapes are overridden here (they cannot be +delegated to v1's, which always carry a `thinking` key). The LOCAL path still +routes through the inference_models adapter exactly as v1_tensor does, with +`enable_thinking=False` hard-wired. +""" + +from typing import List, Optional, Type + +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v1_tensor import ( + Qwen35VLBlockV1, +) +from inference.core.workflows.core_steps.models.foundation.qwen3_5vl.v2 import ( + BlockManifest, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest +from inference_sdk import InferenceHTTPClient + + +class Qwen35VLBlockV2(Qwen35VLBlockV1): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + system_prompt: Optional[str], + max_new_tokens: Optional[int] = None, + ) -> BlockResult: + if self._step_execution_mode == StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_version=model_version, + prompt=prompt, + system_prompt=system_prompt, + max_new_tokens=max_new_tokens, + ) + elif self._step_execution_mode == StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_version=model_version, + prompt=prompt, + system_prompt=system_prompt, + max_new_tokens=max_new_tokens, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + system_prompt: Optional[str], + max_new_tokens: Optional[int] = None, + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + prompt = prompt or "Describe what's in this image." + system_prompt = ( + system_prompt + or "You are a Qwen3.5 model that can answer questions about any image." + ) + combined_prompt = prompt + "" + system_prompt + + predictions = [] + for image in images: + result = client.infer_lmm( + inference_input=image.base64_image, + model_id=model_version, + prompt=combined_prompt, + model_id_in_path=True, + enable_thinking=False, + max_new_tokens=max_new_tokens, + ) + response_text = result.get("response", result) + predictions.append({"parsed_output": response_text}) + + return predictions + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + system_prompt: Optional[str], + max_new_tokens: Optional[int] = None, + ) -> BlockResult: + # Local exec goes through the inference_models adapter with the CHW RGB + # tensor_image. v2 has no thinking mode, so enable_thinking is hard-wired + # to False โ€” the adapter then returns one str per image. + prompt = prompt or "Describe what's in this image." + system_prompt = system_prompt or "You are a helpful assistant." + combined_prompt = prompt + "" + system_prompt + + self._model_manager.add_model(model_id=model_version, api_key=self._api_key) + + predictions = [] + for image in images: + kwargs = dict( + prompt=combined_prompt, + enable_thinking=False, + ) + if max_new_tokens is not None: + kwargs["max_new_tokens"] = max_new_tokens + if image.is_tensor_materialised(): + model_image, image_color_format = image.tensor_image, "rgb" + else: + model_image, image_color_format = image.numpy_image, "bgr" + result = self._model_manager.run_tensor_native_inference( + model_version, + images=[model_image], + input_color_format=image_color_format, + **kwargs, + ) + response_text = result[0] + predictions.append({"parsed_output": response_text}) + return predictions diff --git a/inference/core/workflows/core_steps/models/foundation/qwen3vl/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/qwen3vl/v1_tensor.py new file mode 100644 index 0000000000..0799b0a753 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/qwen3vl/v1_tensor.py @@ -0,0 +1,168 @@ +"""Tensor-native sibling of `roboflow_core/qwen3vl@v1`. + +SCRATCH โ€” first pass for review. Qwen3-VL's *output* is TEXT (parsed_output is a +dict wrapping the model's response string), NOT a prediction kind, so this block +does not PRODUCE tensor-native predictions. The only tensor-native surface is the +local image path: under ENABLE_TENSOR_DATA_REPRESENTATION the block reads +`image.tensor_image` (CHW uint8 RGB on WORKFLOWS_IMAGE_TENSOR_DEVICE) and routes +local inference through the inference_models adapter's run_tensor_native_inference +instead of LMMInferenceRequest. + +Manifest, class name, type literal and describe_outputs are IDENTICAL to v1 and +imported verbatim. run_remotely is unchanged from v1 (the model output is text +either way; there is nothing tensor-native to convert on the remote side). +""" + +from typing import List, Optional, Type + +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode + +# Unchanged from v1 โ€” verbatim manifest, class name, type literal, outputs. +from inference.core.workflows.core_steps.models.foundation.qwen3vl.v1 import ( + BlockManifest, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_sdk import InferenceHTTPClient + + +class Qwen3VLBlockV1(WorkflowBlock): + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + system_prompt: Optional[str], + ) -> BlockResult: + if self._step_execution_mode == StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_version=model_version, + prompt=prompt, + system_prompt=system_prompt, + ) + elif self._step_execution_mode == StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_version=model_version, + prompt=prompt, + system_prompt=system_prompt, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + system_prompt: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + prompt = prompt or "Describe what's in this image." + system_prompt = ( + system_prompt + or "You are a Qwen3-VL model that can answer questions about any image." + ) + combined_prompt = prompt + "" + system_prompt + + predictions = [] + for image in images: + result = client.infer_lmm( + inference_input=image.base64_image, + model_id=model_version, + prompt=combined_prompt, + model_id_in_path=True, + ) + response_text = result.get("response", result) + predictions.append({"parsed_output": response_text}) + + return predictions + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + system_prompt: Optional[str], + ) -> BlockResult: + # Use the provided prompt or default to a generic image description request. + prompt = prompt or "Describe what's in this image." + system_prompt = ( + system_prompt + or "You are a Qwen3-VL model that can answer questions about any image." + ) + # The adapter forwards `prompt` straight to Qwen3VLHF.prompt(), which splits + # on "" internally โ€” same combined-prompt convention as the + # numpy block's LMMInferenceRequest path. + combined_prompt = prompt + "" + system_prompt + + # Register Qwen3-VL with the model manager. + self._model_manager.add_model(model_id=model_version, api_key=self._api_key) + + predictions = [] + for image in images: + # Local exec goes through the inference_models adapter (CHW RGB tensor). + # run_tensor_native_inference -> Qwen3VLHF.prompt returns List[str], + # one entry per image; we pass one image at a time. + if image.is_tensor_materialised(): + model_image, image_color_format = image.tensor_image, "rgb" + else: + model_image, image_color_format = image.numpy_image, "bgr" + result = self._model_manager.run_tensor_native_inference( + model_version, + images=[model_image], + input_color_format=image_color_format, + prompt=combined_prompt, + ) + response_text = result[0] + predictions.append( + { + "parsed_output": response_text, + } + ) + return predictions diff --git a/inference/core/workflows/core_steps/models/foundation/seg_preview/v1.py b/inference/core/workflows/core_steps/models/foundation/seg_preview/v1.py index 2706d65cab..62175ca3b5 100644 --- a/inference/core/workflows/core_steps/models/foundation/seg_preview/v1.py +++ b/inference/core/workflows/core_steps/models/foundation/seg_preview/v1.py @@ -96,7 +96,7 @@ class BlockManifest(WorkflowBlockManifest): @classmethod def get_parameters_accepting_batches(cls) -> List[str]: - return ["images", "boxes"] + return ["images"] @classmethod def describe_outputs(cls) -> List[OutputDefinition]: diff --git a/inference/core/workflows/core_steps/models/foundation/seg_preview/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/seg_preview/v1_tensor.py new file mode 100644 index 0000000000..8da963e855 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/seg_preview/v1_tensor.py @@ -0,0 +1,363 @@ +"""Tensor-native sibling of `roboflow_core/seg-preview@v1`. + +Seg Preview is a text-prompted open-vocabulary instance-segmentation PRODUCER. In +numpy mode (`seg_preview/v1.py`) it POSTs to the Roboflow-internal +`API_BASE_URL/inferenceproxy/seg-preview` endpoint, which returns POLYGON +point-lists per mask, and builds `sv.Detections`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this producer must instead emit a native +`inference_models.InstanceDetections` (torch tensors) under the tensor-native +kind. The remote-request path (base64 image, text prompts, proxy endpoint, +threshold, air-gapped/hosted restriction, the empty-class `None` append, and the +`except Exception -> empty results` fallback) is preserved verbatim from v1. Only +the post-processing tail changes: each polygon point-list is converted straight +to a compact COCO RLE via `pycocotools.frPyObjects` (in C, one instance at a time, +with no dense H x W mask ever allocated) and assembled into an InstanceDetections. +See `_build_instance_detections_from_polygons`. +""" + +import uuid +from typing import Dict, List, Literal, Optional, Type, Union + +import requests +import torch +from pycocotools import mask as mask_utils +from pydantic import ConfigDict, Field + +from inference.core.env import ( + API_BASE_URL, + ROBOFLOW_INTERNAL_SERVICE_NAME, + ROBOFLOW_INTERNAL_SERVICE_SECRET, + WORKFLOWS_IMAGE_TENSOR_DEVICE, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import build_roboflow_api_headers +from inference.core.utils.url_utils import wrap_url +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.offline import ensure_builtin_remote_execution_allowed +from inference.core.workflows.prototypes.block import ( + AirGappedAvailability, + BlockResult, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks + +PREDICTION_TYPE = "instance-segmentation" + + +LONG_DESCRIPTION = "Seg Preview" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Seg Preview", + "version": "v1", + "short_description": "Seg Preview", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["Seg Preview"], + "ui_manifest": { + "section": "model", + "icon": "fa-solid fa-eye", + "blockPriority": 9.49, + "needsGPU": True, + "inference": True, + }, + }, + protected_namespaces=(), + ) + + type: Literal["roboflow_core/seg-preview@v1"] + + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + + class_names: Union[ + List[str], str, Selector(kind=[LIST_OF_VALUES_KIND, STRING_KIND]) + ] = Field( + title="Class Names", + default=None, + description="List of classes to recognise", + examples=[["car", "person"], "$inputs.classes"], + ) + threshold: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + default=0.5, description="Threshold for predicted mask scores", examples=[0.3] + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + RuntimeRestriction( + severity=Severity.HARD, + note=( + "Seg Preview calls the Roboflow-internal " + "API_BASE_URL/inferenceproxy/seg-preview endpoint, which is " + "only reachable from Roboflow-hosted runtimes " + "(HOSTED_SERVERLESS, DEDICATED_DEPLOYMENT). Self-hosted " + "deployments cannot run this block." + ), + applies_to_runtimes=[ + Runtime.SELF_HOSTED_CPU, + Runtime.SELF_HOSTED_GPU, + Runtime.INFERENCE_PIPELINE, + ], + ), + ] + + @classmethod + def get_air_gapped_availability(cls) -> AirGappedAvailability: + """This block requires internet access to the remote inference proxy.""" + return AirGappedAvailability(available=False, reason="requires_internet") + + +class SegPreviewBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + class_names: Optional[Union[List[str], str]], + threshold: float, + ) -> BlockResult: + + if isinstance(class_names, str): + class_names = class_names.split(",") + elif isinstance(class_names, list): + class_names = class_names + else: + raise ValueError(f"Invalid class names type: {type(class_names)}") + + return self.run_via_request( + images=images, + class_names=class_names, + threshold=threshold, + ) + + def run_via_request( + self, + images: Batch[WorkflowImageData], + class_names: Optional[List[str]], + threshold: float, + ) -> BlockResult: + ensure_builtin_remote_execution_allowed("Seg Preview remote execution") + results: List[dict] = [] + if class_names is None: + class_names = [] + if len(class_names) == 0: + class_names.append(None) + + endpoint = f"{API_BASE_URL}/inferenceproxy/seg-preview" + api_key = self._api_key + + for single_image in images: + # Build unified prompt list payloads for HTTP + http_prompts: List[dict] = [] + for class_name in class_names: + http_prompts.append({"type": "text", "text": class_name}) + + # Prepare image for remote API (base64) + http_image = {"type": "base64", "value": single_image.base64_image} + + payload = { + "image": http_image, + "prompts": http_prompts, + "output_prob_thresh": threshold, + } + + try: + headers = {"Content-Type": "application/json"} + if ROBOFLOW_INTERNAL_SERVICE_NAME: + headers["X-Roboflow-Internal-Service-Name"] = ( + ROBOFLOW_INTERNAL_SERVICE_NAME + ) + if ROBOFLOW_INTERNAL_SERVICE_SECRET: + headers["X-Roboflow-Internal-Service-Secret"] = ( + ROBOFLOW_INTERNAL_SERVICE_SECRET + ) + + headers = build_roboflow_api_headers(explicit_headers=headers) + + response = requests.post( + wrap_url(f"{endpoint}?api_key={api_key}"), + json=payload, + headers=headers, + timeout=60, + ) + response.raise_for_status() + resp_json = response.json() + except Exception: + resp_json = {"prompt_results": []} + + results.append( + { + "predictions": _build_instance_detections_from_polygons( + prompt_results=resp_json.get("prompt_results", []), + class_names=class_names, + image=single_image, + threshold=threshold, + ) + } + ) + return results + + +def _build_instance_detections_from_polygons( + prompt_results: List[dict], + class_names: List[Optional[str]], + image: WorkflowImageData, + threshold: float, +) -> InstanceDetections: + """Stream the proxy's polygon point-lists straight into RLE, one instance at a + time, without ever allocating a dense ``H x W`` mask. + + ``pycocotools.frPyObjects`` rasterises each polygon into a compact COCO RLE in + C (no full-image numpy array), and ``InstancesRLEMasks.from_coco_rle_masks`` + only keeps the ``counts`` payload. Peak mask memory is therefore the sum of the + compressed run-length strings (kilobytes) rather than ``N x H x W`` bools + (megabytes), and the slow ``cv2.fillPoly`` -> ``np.where`` -> torch round-trip + is gone. The bbox is taken from the polygon point min/max, matching the numpy + seg_preview block (``seg_preview/v1.py``). + """ + height, width = image._read_shape_without_materialization() + xyxy: List[List[float]] = [] + confidences: List[float] = [] + class_ids: List[int] = [] + class_names_map: Dict[int, str] = {} + bboxes_metadata: List[dict] = [] + rle_dicts: List[dict] = [] + + for prompt_result in prompt_results: + # Key class identity off the response's prompt_index field (not loop order), + # matching the numpy block. + idx = prompt_result.get("prompt_index", 0) + class_name = class_names[idx] if idx < len(class_names) else None + class_name = class_name or "foreground" + for prediction in prompt_result.get("predictions", []): + confidence = float(prediction.get("confidence", 0.0)) + if confidence < threshold: + continue + for polygon in prediction.get("masks", []): + if polygon is None or len(polygon) < 3: + # degenerate polygon cannot enclose area - skip (as numpy block does) + continue + xs = [float(point[0]) for point in polygon] + ys = [float(point[1]) for point in polygon] + x_min, y_min, x_max, y_max = min(xs), min(ys), max(xs), max(ys) + if x_max <= x_min or y_max <= y_min: + # zero-area polygon - skip + continue + flat_polygon = [ + coord + for point in polygon + for coord in (float(point[0]), float(point[1])) + ] + # polygon -> compressed COCO RLE entirely in C; no dense mask is built. + rle = mask_utils.frPyObjects([flat_polygon], height, width)[0] + xyxy.append([x_min, y_min, x_max, y_max]) + confidences.append(confidence) + class_ids.append(idx) + class_names_map[idx] = class_name + bboxes_metadata.append( + {DETECTION_ID_KEY: str(uuid.uuid4()), CLASS_NAME_KEY: class_name} + ) + rle_dicts.append(rle) + + n = len(rle_dicts) + if n == 0: + xyxy_t = torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + class_id_t = torch.zeros( + (0,), dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + confidence_t = torch.zeros( + (0,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + else: + xyxy_t = torch.tensor( + xyxy, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + class_id_t = torch.tensor( + class_ids, dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + confidence_t = torch.tensor( + confidences, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + mask = InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), masks=rle_dicts + ) + + detections = InstanceDetections( + xyxy=xyxy_t, class_id=class_id_t, confidence=confidence_t, mask=mask + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=class_names_map, + prediction_type=PREDICTION_TYPE, + inference_id=str(uuid.uuid4()), + ) + detections.bboxes_metadata = bboxes_metadata if n else None + return detections diff --git a/inference/core/workflows/core_steps/models/foundation/segment_anything2/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/segment_anything2/v1_tensor.py new file mode 100644 index 0000000000..2f6787a42e --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/segment_anything2/v1_tensor.py @@ -0,0 +1,623 @@ +import uuid +from typing import Dict, List, Optional, Tuple, Union + +import torch +from pydantic import ConfigDict, Field +from typing_extensions import Literal + +from inference.core.env import ( + CORE_MODEL_SAM2_ENABLED, + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import ModelEndpointType +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, + split_key_point_prediction, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + IMAGE_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import ( + coco_rle_masks_to_numpy_mask, + torch_mask_to_coco_rle, +) +from inference_sdk import InferenceHTTPClient + +# SAM2 mask-binarisation threshold in logit space, mirroring the numpy adapter's +# MASK_THRESHOLD (inference/models/sam2/segment_anything2_inference_models.py): a pixel +# is foreground when its mask logit is >= 0.0 (equivalently sigmoid(logit) >= 0.5). +MASK_THRESHOLD = 0.0 + +LONG_DESCRIPTION = """ +Run Segment Anything 2, a zero-shot instance segmentation model, on an image. + +** Dedicated inference server required (GPU recomended) ** + +You can use pass in boxes/predictions from other models to Segment Anything 2 to use as prompts for the model. +If you pass in box detections from another model, the class names of the boxes will be forwarded to the predicted masks. If using the model unprompted, the model will assign integers as class names / ids. +""" + + +PREDICTION_TYPE = "instance-segmentation" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Segment Anything 2 Model", + "version": "v1", + "short_description": "Convert bounding boxes to polygons, or run SAM2 on an entire image to generate a mask.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["SAM2", "META"], + "ui_manifest": { + "section": "model", + "icon": "fa-brands fa-meta", + "blockPriority": 9.5, + "needsGPU": True, + "inference": True, + }, + }, + protected_namespaces=(), + ) + + type: Literal["roboflow_core/segment_anything@v1"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + boxes: Optional[ + Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) + ] = Field( # type: ignore + description="Bounding boxes (from another model) to convert to polygons", + examples=["$steps.object_detection_model.predictions"], + default=None, + json_schema_extra={"always_visible": True}, + ) + version: Union[ + Selector(kind=[STRING_KIND]), + Literal["hiera_large", "hiera_small", "hiera_tiny", "hiera_b_plus"], + ] = Field( + default="hiera_tiny", + description="Model to be used. One of hiera_large, hiera_small, hiera_tiny, hiera_b_plus", + examples=["hiera_large", "$inputs.openai_model"], + ) + threshold: Union[ + Selector(kind=[FLOAT_KIND]), + float, + ] = Field( + default=0.0, description="Threshold for predicted masks scores", examples=[0.3] + ) + multimask_output: Union[Optional[bool], Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Flag to determine whether to use sam2 internal multimask or single mask mode. For ambiguous prompts setting to True is recomended.", + examples=[True, "$inputs.multimask_output"], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images", "boxes"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restrictions = [ + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; run_locally() loads a model that needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + ] + if not CORE_MODEL_SAM2_ENABLED: + restrictions.append( + RuntimeRestriction( + severity=Severity.HARD, + note=( + "CORE_MODEL_SAM2_ENABLED=False on Roboflow Hosted " + "Serverless: the SAM2 endpoint is not registered, so " + "run_remotely() returns 404." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + ) + return restrictions + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + return [ + "sam2/hiera_large", + "sam2/hiera_small", + "sam2/hiera_tiny", + "sam2/hiera_b_plus", + ] + + +class SegmentAnything2BlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls): + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + boxes: Optional[Batch], + version: str, + threshold: float, + multimask_output: bool, + ) -> BlockResult: + # RLE mask output is enforced ALWAYS on the tensor path. The numpy SAM2 + # sibling exposes no mask-format field, so neither does this manifest; the + # shared internal helpers (with SAM3 v1/v2/v3) still take a + # `mask_representation` argument, so it is pinned to "rle" here. + mask_representation = "rle" + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + boxes=boxes, + version=version, + threshold=threshold, + multimask_output=multimask_output, + mask_representation=mask_representation, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + boxes=boxes, + version=version, + threshold=threshold, + multimask_output=multimask_output, + mask_representation=mask_representation, + ) + raise ValueError(f"Unknown step execution mode: {self._step_execution_mode}") + + def run_locally( + self, + images: Batch[WorkflowImageData], + boxes: Optional[Batch], + version: str, + threshold: float, + multimask_output: bool, + mask_representation: Literal["rle", "dense"], + ) -> BlockResult: + sam_model_id = f"sam2/{version}" + self._model_manager.add_model( + sam_model_id, self._api_key, endpoint_type=ModelEndpointType.CORE_MODEL + ) + + boxes_iter = boxes if boxes is not None else [None] * len(images) + results: List[dict] = [] + for image, boxes_for_image in zip(images, boxes_iter): + prompt_detections = _prompt_detections( + boxes_for_image + ) # raises if KP-no-bbox + box_tensor = ( + prompt_detections.xyxy if prompt_detections is not None else None + ) + if image.is_tensor_materialised(): + model_image, image_color_format = image.tensor_image, "rgb" + else: + model_image, image_color_format = image.numpy_image, "bgr" + sam2_predictions = self._model_manager.run_tensor_native_inference( + sam_model_id, + action="segment", + images=[model_image], + boxes=[box_tensor] if box_tensor is not None else None, + multi_mask_output=multimask_output, + input_color_format=image_color_format, + # Return raw mask logits and binarise explicitly below (legacy style), + # rather than relying on segment_images' internal default threshold. + return_logits=True, + ) + instance_detections = _sam2_prediction_to_instance_detections( + sam2_prediction=sam2_predictions[0], + image=image, + prompt_detections=prompt_detections, + threshold=threshold, + mask_representation=mask_representation, + ) + results.append({"predictions": instance_detections}) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + boxes: Optional[Batch], + version: str, + threshold: float, + multimask_output: bool, + mask_representation: Literal["rle", "dense"], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient(api_url=api_url, api_key=self._api_key) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + # The SAM2 server always collapses to the most-confident mask per prompt + # (segment_image -> choose_most_confident_sam_prediction), matching the numpy + # sibling; there is no non-collapsed mode. + boxes_iter = boxes if boxes is not None else [None] * len(images) + results: List[dict] = [] + for image, boxes_for_image in zip(images, boxes_iter): + prompt_detections = _prompt_detections(boxes_for_image) + prompts = _box_prompts_payload(prompt_detections) + # C8: request RLE for compact transfer (mask_input_format="rle", NOT + # response_mask_format). threshold (score filter) and mask_representation + # (rle/dense output) are applied in the converter. + response = client.sam2_segment_image( + inference_input=image.base64_image, + sam2_version_id=version, + prompts=prompts, + multimask_output=multimask_output, + mask_input_format="rle", + ) + instance_detections = _rle_response_to_instance_detections( + response=response, + image=image, + prompt_detections=prompt_detections, + threshold=threshold, + mask_representation=mask_representation, + ) + results.append({"predictions": instance_detections}) + return results + + +def _prompt_detections(boxes_for_image): + """Normalize the tensor-native ``boxes`` prompt into a bbox-bearing Detections. + + Per the prediction kinds, ``boxes`` arrives as one of: + * ``Detections`` (object detection) + * ``InstanceDetections`` (instance segmentation) + * ``Tuple[KeyPoints, Optional[Detections]]`` (keypoint detection โ€” see + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND.internal_data_type) + + The keypoint kind is a TUPLE whose instances (``Detections``) component may be + missing (``None``). SAM2 can only be box-prompted, so missing instances is a + hard runtime error โ€” raised by ``split_key_point_prediction`` ("Keypoint + prediction is missing the bounding-box component required by this block."). + + Returns ``None`` only when ``boxes`` is entirely absent (unprompted whole-image + segmentation) or the prediction carries zero instances. + """ + if boxes_for_image is None: + return None + # Raises on a keypoint tuple whose Detections component is None (instances + # missing); returns the prediction as-is for OD / IS. + _key_points, detections = split_key_point_prediction(boxes_for_image) + if len(detections) == 0: + return None + return detections + + +def _choose_most_confident_torch( + masks: torch.Tensor, scores: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Tensor-native equivalent of choose_most_confident_sam_prediction: pick the + highest-scoring proposed mask per prompt, on-device, no numpy bridge. + + masks: (prompt_set, proposed, H, W) or (proposed, H, W) โ€” SAM2 squeezes the + prompt_set dimension when prompt_set == 1, so we re-expand it. + Returns (selected_masks (prompt_set, H, W), selected_scores (prompt_set,)). + """ + if masks.dim() == 3: + masks = masks.unsqueeze(0) + scores = scores.unsqueeze(0) + prompt_set, _proposed, height, width = masks.shape + best = torch.argmax(scores, dim=1) # (prompt_set,) + gather_index = best.view(prompt_set, 1, 1, 1).expand(prompt_set, 1, height, width) + selected_masks = torch.gather(masks, 1, gather_index).squeeze(1) + selected_scores = torch.gather(scores, 1, best.unsqueeze(1)).squeeze(1) + return selected_masks, selected_scores + + +def _sam2_prediction_to_instance_detections( + sam2_prediction, + *, + image: WorkflowImageData, + prompt_detections, + threshold: float, + mask_representation: str, +) -> InstanceDetections: + # SAM2 always collapses to the single most-confident mask per prompt, matching + # the numpy sibling (choose_most_confident_sam_prediction). + selected_masks, selected_scores = _choose_most_confident_torch( + sam2_prediction.masks, sam2_prediction.scores + ) # (m, H, W) mask logits, (m,) + # Binarise the mask logits at the SAM2 threshold, mirroring the numpy adapter's + # `masks >= MASK_THRESHOLD` (MASK_THRESHOLD == 0.0). + binary = (selected_masks >= MASK_THRESHOLD).to(torch.bool) + source_indices, keep_index = _score_filter(selected_scores, threshold) + return _assemble_instance_detections( + binary=binary[keep_index], + confidence=selected_scores[keep_index], + image=image, + prompt_detections=prompt_detections, + source_indices=source_indices, + mask_representation=mask_representation, + ) + + +def _score_filter( + scores: torch.Tensor, threshold: float +) -> Tuple[List[int], torch.Tensor]: + """Confidence filter โ€” mirrors the numpy block's `if confidence < threshold: + continue`. Returns (surviving source indices, index tensor for tensor selection).""" + source_indices = [ + i for i, keep in enumerate((scores >= threshold).tolist()) if keep + ] + return source_indices, torch.tensor(source_indices, dtype=torch.long) + + +def _assemble_instance_detections( + *, + binary: torch.Tensor, # (n, H, W) bool, already score-filtered + confidence: torch.Tensor, # (n,) + image: WorkflowImageData, + prompt_detections, + source_indices: List[int], + mask_representation: str, +) -> InstanceDetections: + n = len(source_indices) + height, width = image._read_shape_without_materialization() + if n == 0: + xyxy = torch.zeros((0, 4), dtype=torch.float32) + class_id = torch.zeros((0,), dtype=torch.int64) + else: + xyxy = torch.stack([_mask_to_xyxy(binary[i]) for i in range(n)]) + class_id = _prompt_class_ids(prompt_detections, source_indices) + + if mask_representation == "rle": + rle_dicts = [torch_mask_to_coco_rle(binary[i]) for i in range(n)] + mask = InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), masks=rle_dicts + ) + else: + mask = binary + + detections = InstanceDetections( + xyxy=xyxy.to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + class_id=class_id.to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + confidence=confidence.to(torch.float32) + .reshape(-1) + .to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + mask=( + mask + if isinstance(mask, InstancesRLEMasks) + else mask.to(WORKFLOWS_IMAGE_TENSOR_DEVICE) + ), + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=_class_names_map(prompt_detections, source_indices), + prediction_type=PREDICTION_TYPE, + inference_id=str(uuid.uuid4()), + ) + detections.bboxes_metadata = [ + _per_instance_metadata(prompt_detections, src) for src in source_indices + ] + return detections + + +def _mask_to_xyxy(mask: torch.Tensor) -> torch.Tensor: + """Tight xyxy bbox of a 2-D boolean mask, on the mask's device. + + Uses inclusive min/max with NO +1, matching the numpy SAM2 sibling's + polygon-derived bbox (``np.min``/``np.max`` on the contour points) and the + other SAM tensor siblings (``segment_anything2_video``, ``segment_anything3``, + ``seg_preview``), so masks/areas/IoU stay numpy-faithful across the family. + """ + nonzero = torch.nonzero(mask, as_tuple=False) + if nonzero.numel() == 0: + return torch.zeros((4,), dtype=torch.float32, device=mask.device) + ys, xs = nonzero[:, 0], nonzero[:, 1] + return torch.stack([xs.min(), ys.min(), xs.max(), ys.max()]).to(torch.float32) + + +def _prompted(prompt_detections) -> bool: + return prompt_detections is not None and len(prompt_detections) > 0 + + +def _prompt_class_ids(prompt_detections, source_indices: List[int]) -> torch.Tensor: + # Forward each surviving instance's prompt class_id; unprompted -> 0 (the numpy + # block uses class_id 0 / "foreground" when there is no prompt). + if _prompted(prompt_detections): + ids = [int(prompt_detections.class_id[src]) for src in source_indices] + return torch.tensor(ids, dtype=torch.int64) + return torch.zeros((len(source_indices),), dtype=torch.int64) + + +def _class_names_map(prompt_detections, source_indices: List[int]) -> Dict[int, str]: + # class_id -> name map carried on image_metadata (used by the serializer). + names: Dict[int, str] = {} + for src in source_indices: + class_id = ( + int(prompt_detections.class_id[src]) if _prompted(prompt_detections) else 0 + ) + names[class_id] = _resolve_prompt_class_name(prompt_detections, src) + return names + + +def _per_instance_metadata(prompt_detections, source_index: int) -> dict: + entry = {DETECTION_ID_KEY: str(uuid.uuid4())} + if ( + prompt_detections is not None + and prompt_detections.bboxes_metadata is not None + and source_index < len(prompt_detections.bboxes_metadata) + ): + src = prompt_detections.bboxes_metadata[source_index] + if DETECTION_ID_KEY in src: + entry[DETECTION_ID_KEY] = src[DETECTION_ID_KEY] + # Forward a per-box class override only if the prompt explicitly carried one + # (vlm/ocr prompts); standard OD prompts resolve the name from the + # class_id -> name map on image_metadata, so no spurious per-box `class`. + if CLASS_NAME_KEY in src: + entry[CLASS_NAME_KEY] = src[CLASS_NAME_KEY] + return entry + + +def _resolve_prompt_class_name(prompt_detections, source_index: int) -> str: + """Forward the prompt box's class name, mirroring the numpy block. Native OD + producers carry names in ``image_metadata['class_names']`` keyed by ``class_id`` + (NOT a per-box field); a per-box ``class`` override (vlm/ocr prompts) wins if + present. Unprompted SAM (no boxes) -> ``foreground``.""" + if not _prompted(prompt_detections): + return "foreground" + if prompt_detections.bboxes_metadata is not None and source_index < len( + prompt_detections.bboxes_metadata + ): + override = prompt_detections.bboxes_metadata[source_index].get(CLASS_NAME_KEY) + if override is not None: + return str(override) + class_names = (prompt_detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + return class_names.get(int(prompt_detections.class_id[source_index]), "foreground") + + +def _box_prompts_payload(prompt_detections) -> Optional[List[dict]]: + if prompt_detections is None: + return None + prompts = [] + for i in range(len(prompt_detections)): + x1, y1, x2, y2 = prompt_detections.xyxy[i].tolist() + w, h = x2 - x1, y2 - y1 + prompts.append( + {"box": {"x": x1 + w / 2, "y": y1 + h / 2, "width": w, "height": h}} + ) + return prompts + + +def _rle_response_to_instance_detections( + response, + image: WorkflowImageData, + prompt_detections, + *, + threshold: float, + mask_representation: str, +) -> InstanceDetections: + """Build InstanceDetections from the SAM2 server's rle response. Schema (from + turn_segmentation_results_into_rle_response -> Sam2SegmentationPrediction): + {"predictions": [{"masks": {"size": [h, w], "counts": ""}, + "confidence": float, "format": "rle"}, ...]} + The response carries no bbox, so xyxy is derived from the decoded mask; classes + are forwarded from the prompt by index (response preserves prompt order). The + response is decoded to dense, score-filtered, then re-assembled per + mask_representation โ€” uniform with the local path. + """ + if isinstance(response, list): + response = response[0] + predictions = response.get("predictions", []) or [] + # Class identity is forwarded from the prompt by RESPONSE ROW POSITION + # (source_indices index into prompt_detections after the score filter). That is + # only sound if the SAM2 server returns exactly one prediction per prompt in + # prompt order. Guard the positional contract before forwarding classes: if the + # server ever drops/reorders a prompt's prediction the indices would silently + # pull the wrong class_id/detection_id for every subsequent instance. + if _prompted(prompt_detections) and len(predictions) != len(prompt_detections): + raise ValueError( + "SAM2 remote response returned " + f"{len(predictions)} prediction(s) for {len(prompt_detections)} box " + "prompt(s); positional class forwarding requires one prediction per " + "prompt in prompt order." + ) + height, width = image._read_shape_without_materialization() + + counts: List[bytes] = [] + scores: List[float] = [] + for prediction in predictions: + raw_counts = prediction["masks"]["counts"] # {"size":[h,w],"counts":""} + # Normalize to bytes (local frPyObjects path stores counts as bytes; the remote + # rle response decodes them to a utf-8 string). + counts.append( + raw_counts.encode("utf-8") if isinstance(raw_counts, str) else raw_counts + ) + scores.append(float(prediction["confidence"])) + + if not counts: + binary = torch.zeros((0, height, width), dtype=torch.bool) + else: + dense = coco_rle_masks_to_numpy_mask( + InstancesRLEMasks(image_size=(height, width), masks=counts) + ) + binary = torch.from_numpy(dense).to(torch.bool) + + selected_scores = torch.tensor(scores, dtype=torch.float32) + source_indices, keep_index = _score_filter(selected_scores, threshold) + return _assemble_instance_detections( + binary=binary[keep_index], + confidence=selected_scores[keep_index], + image=image, + prompt_detections=prompt_detections, + source_indices=source_indices, + mask_representation=mask_representation, + ) diff --git a/inference/core/workflows/core_steps/models/foundation/segment_anything2_video/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/segment_anything2_video/v1_tensor.py new file mode 100644 index 0000000000..967f000867 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/segment_anything2_video/v1_tensor.py @@ -0,0 +1,514 @@ +"""Tensor-native sibling of `roboflow_core/segment_anything_2_video@v1`. + +SAM2 Video Tracker is a STATEFUL, LOCAL-only streaming instance-segmentation +producer. It differs from the SAM2 image block: + +- No remote path (per-video session state cannot be shipped per-frame). +- The model is loaded directly via AutoModel (inference_models). Its HF + Sam2VideoProcessor expects HWC RGB; the model's `_ensure_numpy_image` + (hf_streaming_video.py) now permutes a CHW tensor to HWC before the host + transfer, so the block passes `WorkflowImageData.tensor_image` (CHW RGB) + tensor-natively โ€” which also feeds the processor correct RGB (v1's numpy_image + was HWC BGR). + +The tensor-native surface is the IMAGE (CHW RGB tensor), the INPUT boxes +(tensor-native prediction prompts), and the OUTPUT +(inference_models.InstanceDetections with tracker ids), plus the RLE-by-default +mask carriage (an execution-level choice driven by the +WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION env variable โ€” NOT a manifest field, +so the manifest stays identical to the numpy sibling; GCP_SERVERLESS forces +"rle"). The layout-agnostic session/state helpers +(VideoSessionBookkeeping, decide_prompt_vs_track, build_obj_id_metadata_from_boxes, +BoxPromptMetadata) are reused verbatim; only extract_box_prompts and +masks_to_sv_detections get tensor-native equivalents here. +""" + +import uuid +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import ( + GCP_SERVERLESS, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import get_extra_weights_provider_headers +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.models.foundation._streaming_video_common import ( + BoxPromptMetadata, + VideoSessionBookkeeping, + build_obj_id_metadata_from_boxes, + decide_prompt_vs_track, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + DETECTION_ID_KEY, + TRACKER_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + IMAGE_KIND, + INTEGER_KIND, + ROBOFLOW_MODEL_ID_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import torch_mask_to_coco_rle + +PromptMode = Literal["first_frame", "every_n_frames", "every_frame"] +PREDICTION_TYPE = "instance-segmentation" + + +def _resolve_mask_representation() -> str: + """Execution-level selection of the instance-mask carrier ("rle"/"dense"). + + Deliberately NOT a manifest field: the numpy sibling has no such knob and + manifests must stay identical across the flag swap. Driven by the + ``WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION`` env variable ("rle" default); + ``GCP_SERVERLESS`` forces the compact RLE carrier regardless. + """ + if GCP_SERVERLESS: + return "rle" + return WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION + + +SHORT_DESCRIPTION = ( + "Segment and track objects across video frames with SAM2's streaming " + "camera predictor." +) +LONG_DESCRIPTION = """ +Run Segment Anything 2 on a live video stream frame by frame, keeping +per-video temporal memory so object identities are preserved across +frames. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "SAM2 Video Tracker", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": [ + "SAM2", + "segment anything", + "video", + "tracking", + "META", + ], + "ui_manifest": { + "section": "video", + "icon": "fa-brands fa-meta", + "blockPriority": 9.4, + "needsGPU": True, + "inference": True, + "trackers": True, + }, + }, + protected_namespaces=(), + ) + + type: Literal["roboflow_core/segment_anything_2_video@v1"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + boxes: Optional[ + Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) + ] = Field( + description=( + "Bounding boxes to use as SAM2 prompts. Only read on frames " + "where the block re-prompts (see `prompt_mode`)." + ), + examples=["$steps.object_detection_model.predictions"], + default=None, + json_schema_extra={"always_visible": True}, + ) + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = Field( + default="sam2video/small", + description="Streaming SAM2 model id resolved by `inference_models`.", + examples=[ + "sam2video/tiny", + "sam2video/small", + "sam2video/base-plus", + "sam2video/large", + ], + ) + prompt_mode: PromptMode = Field( + default="first_frame", + description=( + "When to consume `boxes` as SAM2 prompts. `first_frame` prompts " + "once per session and then tracks; `every_n_frames` re-seeds every " + "`prompt_interval` frames; `every_frame` re-seeds every frame." + ), + ) + prompt_interval: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=30, + description="For `prompt_mode=every_n_frames`: re-prompt every N frames.", + examples=[30], + ) + threshold: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + default=0.0, + description="Minimum confidence for emitted masks.", + examples=[0.0], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images", "boxes"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; the streaming SAM2 video model needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + return [ + "sam2video/small", + "sam2video/tiny", + "sam2video/base-plus", + "sam2video/large", + ] + + +class SegmentAnything2VideoBlockV1(WorkflowBlock): + """Stateful SAM2 streaming video tracking block (tensor-native output).""" + + _REMOTE_EXECUTION_NOT_SUPPORTED_MESSAGE = ( + "SAM2 Video Tracker only supports LOCAL workflow step " + "execution. Remote execution would ship each frame to a " + "separate process and break the per-video SAM2 session " + "that holds the temporal memory. Set " + "WORKFLOWS_STEP_EXECUTION_MODE=local (or run on a " + "dedicated deployment) to use this block." + ) + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + self._model = None # lazily loaded + self._current_model_id: Optional[str] = None + self._sessions: Dict[str, VideoSessionBookkeeping] = {} + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def _get_model(self, model_id: str): + if self._model is None or self._current_model_id != model_id: + from inference_models import AutoModel + + extra_weights_provider_headers = get_extra_weights_provider_headers() + self._model = AutoModel.from_pretrained( + model_id_or_path=model_id, + api_key=self._api_key, + weights_provider_extra_headers=extra_weights_provider_headers, + ) + self._current_model_id = model_id + self._sessions.clear() + return self._model + + def run( + self, + images: Batch[WorkflowImageData], + boxes: Optional[Batch], + model_id: str, + prompt_mode: PromptMode, + prompt_interval: int, + threshold: float, + ) -> BlockResult: + if self._step_execution_mode is not StepExecutionMode.LOCAL: + raise NotImplementedError(self._REMOTE_EXECUTION_NOT_SUPPORTED_MESSAGE) + mask_representation = _resolve_mask_representation() + model = self._get_model(model_id=model_id) + boxes_iter = boxes if boxes is not None else [None] * len(images) + + results: List[dict] = [] + for single_image, boxes_for_image in zip(images, boxes_iter): + metadata = single_image.video_metadata + video_id = metadata.video_identifier + frame_number = metadata.frame_number or 0 + + session = self._sessions.setdefault(video_id, VideoSessionBookkeeping()) + has_box_prompts = boxes_for_image is not None and len(boxes_for_image) > 0 + should_reset, should_prompt = decide_prompt_vs_track( + session=session, + frame_number=frame_number, + prompt_mode=prompt_mode, + prompt_interval=prompt_interval, + has_prompts=has_box_prompts, + ) + if should_reset: + session.state_dict = None + session.obj_id_metadata = {} + session.frames_since_prompt = 0 + + # CROSS-REPO CONTRACT (intentional divergence from the numpy sibling): + # - This block forwards the tensor-native CHW *RGB* frame. + # - inference_models hf_streaming_video._ensure_numpy_image is + # contracted to permute CHW->HWC WITHOUT reinterpreting channel + # order, so the HF Sam2VideoProcessor receives HWC RGB. + # - The numpy v1 sibling fed HWC *BGR* (numpy_image); the HF + # processor expects RGB, so the tensor path is the correct one. + # This correctness now hinges on _ensure_numpy_image NOT swapping + # channels: a regression there silently corrupts every mask. The + # `frame.dim() == 3` guard below catches a layout regression (a + # non-CHW frame) early instead of letting a malformed tensor reach + # the model. Channel-order parity itself is pinned by a focused test + # in inference_models (see hf_streaming_video tests); it cannot be + # asserted from a workflow block without materialising the frame. + # Use the tensor when already materialised, otherwise hand the model an RGB + # host frame: _ensure_numpy_image permutes CHW->HWC WITHOUT swapping channels + # and the HF processor expects RGB, so the BGR numpy frame is flipped here + # (avoids a forced CHW transpose + H2D the model would undo anyway). + if single_image.is_tensor_materialised(): + frame = single_image.tensor_image + if frame.dim() != 3: + raise ValueError( + "SAM2 video tracker expects a CHW (3-D) RGB frame tensor; got " + f"a tensor with {frame.dim()} dim(s). The model's " + "_ensure_numpy_image permutes CHW->HWC and assumes this layout." + ) + else: + frame = np.ascontiguousarray(single_image.numpy_image[:, :, ::-1]) + + if should_prompt: + boxes_xyxy, per_box_meta = _extract_box_prompts_tensor(boxes_for_image) + masks, obj_ids, new_state = model.prompt( + image=frame, + bboxes=boxes_xyxy, + state_dict=session.state_dict, + clear_old_prompts=True, + frame_idx=frame_number, + ) + session.obj_id_metadata = build_obj_id_metadata_from_boxes( + obj_ids=obj_ids, box_metas=per_box_meta + ) + session.state_dict = new_state + session.frames_since_prompt = 0 + elif session.state_dict is not None: + masks, obj_ids, new_state = model.track( + image=frame, state_dict=session.state_dict + ) + session.state_dict = new_state + session.frames_since_prompt += 1 + else: + height, width = single_image._read_shape_without_materialization() + masks = np.zeros((0, height, width), dtype=bool) + obj_ids = np.zeros((0,), dtype=np.int64) + + session.last_frame_number = frame_number + + results.append( + { + "predictions": _masks_to_instance_detections( + masks=masks, + obj_ids=obj_ids, + image=single_image, + obj_id_metadata=session.obj_id_metadata, + threshold=threshold, + mask_representation=mask_representation, + ) + } + ) + return results + + +def _extract_box_prompts_tensor( + boxes_for_image, +) -> Tuple[List[Tuple[float, float, float, float]], List[BoxPromptMetadata]]: + """Tensor-native equivalent of extract_box_prompts: flatten a tensor-native + prediction prompt into xyxy tuples + per-box metadata. + + `boxes` follows the prediction kinds (OD Detections, IS InstanceDetections, or + the keypoint tuple Tuple[KeyPoints, Optional[Detections]]). A keypoint tuple with + a missing Detections component is a hard runtime error (split_key_point_prediction + raises). Empty / absent input returns two empty lists. + """ + if boxes_for_image is None: + return [], [] + _key_points, detections = split_key_point_prediction(boxes_for_image) + n = len(detections) + if n == 0: + return [], [] + bboxes_metadata = detections.bboxes_metadata + boxes_xyxy: List[Tuple[float, float, float, float]] = [] + metas: List[BoxPromptMetadata] = [] + for i in range(n): + x1, y1, x2, y2 = detections.xyxy[i].tolist() + boxes_xyxy.append((float(x1), float(y1), float(x2), float(y2))) + meta = ( + bboxes_metadata[i] + if bboxes_metadata is not None and i < len(bboxes_metadata) + else {} + ) + parent_id = meta.get(DETECTION_ID_KEY) + metas.append( + BoxPromptMetadata( + class_id=int(detections.class_id[i]), + class_name=str(meta.get(CLASS_NAME_KEY, "foreground")), + confidence=( + float(detections.confidence[i]) + if detections.confidence is not None + else 1.0 + ), + parent_id=str(parent_id) if parent_id is not None else None, + ) + ) + return boxes_xyxy, metas + + +def _masks_to_instance_detections( + masks: np.ndarray, + obj_ids: np.ndarray, + image: WorkflowImageData, + obj_id_metadata: Dict[int, BoxPromptMetadata], + threshold: float, + mask_representation: str, +) -> InstanceDetections: + """Tensor-native equivalent of masks_to_sv_detections: one InstanceDetections per + SAM-assigned object, tracker id carried in bboxes_metadata. Masks with no positive + pixels, or whose forwarded confidence < threshold, are dropped.""" + height, width = image._read_shape_without_materialization() + xyxy: List[List[float]] = [] + confidences: List[float] = [] + class_ids: List[int] = [] + class_names_map: Dict[int, str] = {} + bboxes_metadata: List[dict] = [] + kept_masks: List[np.ndarray] = [] + + for mask, obj_id in zip(masks, obj_ids.tolist()): + meta = obj_id_metadata.get(int(obj_id)) + confidence = meta.confidence if meta is not None else 1.0 + if confidence < threshold: + continue + ys, xs = np.where(mask) + if xs.size == 0: + continue + class_id = int(meta.class_id) if meta is not None else 0 + class_name = meta.class_name if meta is not None else "foreground" + xyxy.append( + [float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max())] + ) + confidences.append(float(confidence)) + class_ids.append(class_id) + class_names_map[class_id] = class_name + bboxes_metadata.append( + { + DETECTION_ID_KEY: str(uuid.uuid4()), + CLASS_NAME_KEY: class_name, + TRACKER_ID_KEY: int(obj_id), + } + ) + kept_masks.append(mask.astype(bool)) + + n = len(kept_masks) + if n == 0: + xyxy_t = torch.zeros((0, 4), dtype=torch.float32) + class_id_t = torch.zeros((0,), dtype=torch.int64) + confidence_t = torch.zeros((0,), dtype=torch.float32) + mask = ( + InstancesRLEMasks(image_size=(height, width), masks=[]) + if mask_representation == "rle" + else torch.zeros((0, height, width), dtype=torch.bool) + ) + else: + xyxy_t = torch.tensor(xyxy, dtype=torch.float32) + class_id_t = torch.tensor(class_ids, dtype=torch.int64) + confidence_t = torch.tensor(confidences, dtype=torch.float32) + if mask_representation == "rle": + rle_dicts = [ + torch_mask_to_coco_rle(torch.from_numpy(m)) for m in kept_masks + ] + mask = InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), masks=rle_dicts + ) + else: + mask = torch.from_numpy(np.stack(kept_masks, axis=0)) + + detections = InstanceDetections( + xyxy=xyxy_t.to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + class_id=class_id_t.to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + confidence=confidence_t.to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + mask=( + mask + if isinstance(mask, InstancesRLEMasks) + else mask.to(WORKFLOWS_IMAGE_TENSOR_DEVICE) + ), + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=class_names_map, + prediction_type=PREDICTION_TYPE, + inference_id=str(uuid.uuid4()), + ) + detections.bboxes_metadata = bboxes_metadata if n else None + return detections diff --git a/inference/core/workflows/core_steps/models/foundation/segment_anything3/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/segment_anything3/v1_tensor.py new file mode 100644 index 0000000000..72a3712828 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/segment_anything3/v1_tensor.py @@ -0,0 +1,648 @@ +"""Tensor-native sibling of `roboflow_core/sam3@v1`. + +SAM3 v1 is a TEXT-prompted open-vocabulary +instance-segmentation producer (one text prompt per `class_names` entry; SAM3 +returns N masks per prompt with per-mask scores). Output is tensor-native +`inference_models.InstanceDetections` (RLE masks by default), combining all +prompts/classes into one prediction per image, with class_id = prompt index and +the class name forwarded onto each mask (mirrors the numpy block's +convert_sam3_segmentation_response_to_inference_instances_seg_response). + +Paths: +- LOCAL: `run_tensor_native_inference(model_id, images=[tensor_image], + prompts=[{"text": class_name}], output_prob_thresh=threshold)` -> raw per-prompt + `{"masks": (N,1,H,W) bool ndarray, "scores": [...]}` (PostProcessImage leaves a + channel dim from mask.unsqueeze(1); squeezed to (N,H,W) on ingestion). Masks are + already binary (sigmoid>0.5) and threshold-filtered; built into InstanceDetections here. +- REMOTE / proxy: the server returns polygon predictions; streamed straight to + compact COCO RLE via pycocotools.frPyObjects (no dense mask) then built into + InstanceDetections โ€” parity with v1. +""" + +import uuid +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import requests +import torch +from pycocotools import mask as mask_utils +from pydantic import ConfigDict, Field + +from inference.core.env import ( + API_BASE_URL, + CORE_MODEL_SAM3_ENABLED, + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + ROBOFLOW_INTERNAL_SERVICE_NAME, + ROBOFLOW_INTERNAL_SERVICE_SECRET, + SAM3_EXEC_MODE, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import build_roboflow_api_headers +from inference.core.utils.url_utils import wrap_url +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.offline import ensure_builtin_remote_execution_allowed +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import torch_mask_to_coco_rle +from inference_sdk import InferenceHTTPClient + +PREDICTION_TYPE = "instance-segmentation" + +LONG_DESCRIPTION = """ +Run Segment Anything 3, a zero-shot instance segmentation model, on an image. + +You can use a text prompt (class names) for open-vocabulary segmentation. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "SAM 3", + "version": "v1", + "short_description": "Sam3", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["Sam3"], + "ui_manifest": { + "section": "model", + "icon": "fa-solid fa-eye", + "blockPriority": 9.49, + "needsGPU": True, + "inference": True, + }, + }, + protected_namespaces=(), + ) + + type: Literal["roboflow_core/sam3@v1"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), Optional[str]] = Field( + default="sam3/sam3_final", + description="model version. You only need to change this for fine tuned sam3 models.", + examples=["sam3/sam3_final", "$inputs.model_variant"], + ) + class_names: Optional[ + Union[List[str], str, Selector(kind=[LIST_OF_VALUES_KIND, STRING_KIND])] + ] = Field( + title="Class Names", + default=None, + description="List of classes to recognise", + examples=[["car", "person"], "$inputs.classes"], + ) + threshold: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + default=0.5, description="Threshold for predicted mask scores", examples=[0.3] + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + # INTENTIONAL CORRECTION vs the numpy sibling (v1.py), which returns + # ["images", "boxes"] while declaring no `boxes` field โ€” a manifest bug + # (it advertises a non-existent param as batched). SAM3 v1 has only the + # batched `images` input, so the correct value is ["images"]. The + # numpy-side discrepancy should be backported separately. + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restrictions = [ + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; run_locally() loads a model that needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + ] + if not CORE_MODEL_SAM3_ENABLED: + restrictions.append( + RuntimeRestriction( + severity=Severity.HARD, + note=( + "CORE_MODEL_SAM3_ENABLED=False on Roboflow Hosted " + "Serverless: the SAM3 endpoint is not registered, so " + "run_remotely() returns 404." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + ) + return restrictions + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + return ["sam3/sam3_final"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + if SAM3_EXEC_MODE == "remote": + # Proxy execution ignores the configured model id โ€” the proxy runs + # its own fixed SAM3 server-side; nothing to declare. + return [] + if self.model_id is None: + return [] + return [roboflow_platform_model(model_id=self.model_id)] + + +class SegmentAnything3BlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_names: Optional[Union[List[str], str]], + threshold: float, + ) -> BlockResult: + # RLE mask output is enforced ALWAYS on the tensor path. The numpy SAM3 v1 + # sibling exposes no mask-format field, so neither does this manifest; the + # shared internal helpers (with SAM2/v2/v3) still take a `mask_representation` + # argument, so it is pinned to "rle" here. + mask_representation = "rle" + class_names = _normalize_class_names(class_names) + if SAM3_EXEC_MODE == "remote": + return self.run_via_request( + images=images, + class_names=class_names, + threshold=threshold, + mask_representation=mask_representation, + ) + elif self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_names=class_names, + threshold=threshold, + mask_representation=mask_representation, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_names=class_names, + threshold=threshold, + mask_representation=mask_representation, + ) + raise ValueError(f"Unknown step execution mode: {self._step_execution_mode}") + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_names: List[Optional[str]], + threshold: float, + mask_representation: str, + ) -> BlockResult: + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + prompts = [{"text": class_name} for class_name in class_names] + + results: List[dict] = [] + for single_image in images: + # SAM3 converts its input to numpy internally and is colour-agnostic + # (expects RGB, like the materialised tensor). Pass the tensor when it is + # already on device, otherwise an RGB host frame โ€” avoiding a forced CHW + # transpose + H2D that SAM3 would immediately undo. + if single_image.is_tensor_materialised(): + model_image = single_image.tensor_image + else: + model_image = np.ascontiguousarray(single_image.numpy_image[:, :, ::-1]) + per_image = self._model_manager.run_tensor_native_inference( + model_id, + images=[model_image], + prompts=prompts, + output_prob_thresh=threshold, + ) + items = _collect_from_native( + per_prompt_results=per_image[0], + class_names=class_names, + threshold=threshold, + ) + results.append( + { + "predictions": _build_instance_detections( + items=items, + image=single_image, + mask_representation=mask_representation, + ) + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_names: List[Optional[str]], + threshold: float, + mask_representation: str, + ) -> BlockResult: + ensure_builtin_remote_execution_allowed("SAM3 remote execution") + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient(api_url=api_url, api_key=self._api_key) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + http_prompts = [{"type": "text", "text": cn} for cn in class_names] + + results: List[dict] = [] + for single_image in images: + resp_json = client.sam3_concept_segment( + inference_input=single_image.base64_image, + prompts=http_prompts, + model_id=model_id, + output_prob_thresh=threshold, + ) + results.append( + self._build_from_polygon_response( + resp_json=resp_json, + image=single_image, + class_names=class_names, + threshold=threshold, + mask_representation=mask_representation, + ) + ) + return results + + def run_via_request( + self, + images: Batch[WorkflowImageData], + class_names: List[Optional[str]], + threshold: float, + mask_representation: str, + ) -> BlockResult: + ensure_builtin_remote_execution_allowed("SAM3 inference proxy execution") + endpoint = f"{API_BASE_URL}/inferenceproxy/seg-preview" + http_prompts = [{"type": "text", "text": cn} for cn in class_names] + + results: List[dict] = [] + for single_image in images: + payload = { + "image": {"type": "base64", "value": single_image.base64_image}, + "prompts": http_prompts, + "output_prob_thresh": threshold, + } + headers = {"Content-Type": "application/json"} + if ROBOFLOW_INTERNAL_SERVICE_NAME: + headers["X-Roboflow-Internal-Service-Name"] = ( + ROBOFLOW_INTERNAL_SERVICE_NAME + ) + if ROBOFLOW_INTERNAL_SERVICE_SECRET: + headers["X-Roboflow-Internal-Service-Secret"] = ( + ROBOFLOW_INTERNAL_SERVICE_SECRET + ) + headers = build_roboflow_api_headers(explicit_headers=headers) + try: + response = requests.post( + wrap_url(f"{endpoint}?api_key={self._api_key}"), + json=payload, + headers=headers, + timeout=60, + ) + response.raise_for_status() + resp_json = response.json() + except Exception as exc: + raise Exception(f"SAM3 request failed: {exc}") + results.append( + self._build_from_polygon_response( + resp_json=resp_json, + image=single_image, + class_names=class_names, + threshold=threshold, + mask_representation=mask_representation, + ) + ) + return results + + def _build_from_polygon_response( + self, + resp_json: dict, + image: WorkflowImageData, + class_names: List[Optional[str]], + threshold: float, + mask_representation: str, + ) -> dict: + return { + "predictions": _build_instance_detections_from_polygons( + prompt_results=resp_json.get("prompt_results", []), + class_names=class_names, + image=image, + threshold=threshold, + mask_representation=mask_representation, + ) + } + + +def _normalize_class_names(class_names) -> List[Optional[str]]: + if isinstance(class_names, str): + class_names = class_names.split(",") + elif isinstance(class_names, list): + class_names = list(class_names) + elif class_names is None: + class_names = [] + else: + raise ValueError(f"Invalid class names type: {type(class_names)}") + if len(class_names) == 0: + # A single null prompt = unprompted (class_id 0 / "foreground"). + class_names = [None] + return class_names + + +# Each item: (mask (H,W) bool, score, class_id, class_name) +Item = Tuple[np.ndarray, float, int, str] + + +def _collect_from_native( + per_prompt_results: List[dict], + class_names: List[Optional[str]], + threshold: float, +) -> List[Item]: + items: List[Item] = [] + for prompt_result in per_prompt_results: + # Key class identity off the model's prompt_index field (not the loop + # order), matching the numpy block + the polygon path below. + idx = prompt_result.get("prompt_index", 0) + class_name = class_names[idx] if idx < len(class_names) else None + masks = prompt_result.get("masks") + if masks is None: + continue + masks = np.asarray(masks) + if masks.ndim == 4 and masks.shape[1] == 1: + # SAM3 PostProcessImage (non-RLE) returns (N, 1, H, W): the channel dim + # from mask.unsqueeze(1) is never squeezed on this path. Drop it so + # downstream np.where / RLE / stack see true (N, H, W) masks. + masks = masks[:, 0] + scores = list(prompt_result.get("scores", [])) + for i in range(masks.shape[0]): + score = float(scores[i]) if i < len(scores) else 1.0 + if score < threshold: + continue + mask = masks[i].astype(bool) + # Drop empty masks HERE (not in the packer) so _build_instance_detections + # is a pure packer: every emitted Item is guaranteed non-empty, keeping + # masks/xyxy/confidence/class rows in lockstep. `mask.any()` is the + # tensor-native equivalent of the previous `np.where(...).size == 0` + # skip that lived in the builder. + if not mask.any(): + continue + items.append((mask, score, idx, class_name or "foreground")) + return items + + +def _build_instance_detections( + items: List[Item], + image: WorkflowImageData, + mask_representation: str, +) -> InstanceDetections: + """Build InstanceDetections from items that already carry dense binary masks + (the LOCAL native path: SAM3 returns dense masks, so dense is inherent here). + + Pure packer: empty masks are dropped upstream in `_collect_from_native`, so + every Item here is non-empty and every list below grows by exactly one row per + item โ€” masks/xyxy/confidence/class can never desync. + """ + height, width = image._read_shape_without_materialization() + xyxy: List[List[float]] = [] + confidences: List[float] = [] + class_ids: List[int] = [] + class_names_map: Dict[int, str] = {} + bboxes_metadata: List[dict] = [] + kept_masks: List[np.ndarray] = [] + + for mask, score, class_id, class_name in items: + ys, xs = np.where(mask) + xyxy.append( + [float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max())] + ) + confidences.append(score) + class_ids.append(class_id) + class_names_map[class_id] = class_name + bboxes_metadata.append( + {DETECTION_ID_KEY: str(uuid.uuid4()), CLASS_NAME_KEY: class_name} + ) + kept_masks.append(mask) + + if mask_representation == "rle": + mask = InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), + masks=[torch_mask_to_coco_rle(torch.from_numpy(m)) for m in kept_masks], + ) + elif kept_masks: + mask = torch.from_numpy(np.stack(kept_masks, axis=0)).to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + else: + mask = torch.zeros( + (0, height, width), dtype=torch.bool, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + return _assemble_detections( + image=image, + xyxy=xyxy, + confidences=confidences, + class_ids=class_ids, + class_names_map=class_names_map, + bboxes_metadata=bboxes_metadata, + mask=mask, + ) + + +def _build_instance_detections_from_polygons( + prompt_results: List[dict], + class_names: List[Optional[str]], + image: WorkflowImageData, + threshold: float, + mask_representation: str, + class_mapping: Optional[Dict[str, str]] = None, +) -> InstanceDetections: + """Stream the proxy's polygon point-lists into masks one instance at a time. + + For the default ``rle`` carrier each polygon is converted straight to a compact + COCO RLE via ``pycocotools.frPyObjects`` (in C, no dense ``H x W`` array), so + peak memory is the sum of the run-length strings rather than ``N x H x W`` bools + โ€” the same fix applied to ``seg_preview/v1_tensor.py``. The ``dense`` carrier + inherently needs a full mask, so the RLE is decoded back per instance. Bbox is + taken from the polygon point min/max, matching the numpy SAM3 block. + """ + height, width = image._read_shape_without_materialization() + xyxy: List[List[float]] = [] + confidences: List[float] = [] + class_ids: List[int] = [] + class_names_map: Dict[int, str] = {} + bboxes_metadata: List[dict] = [] + rle_dicts: List[dict] = [] + dense_masks: List[np.ndarray] = [] + + for prompt_result in prompt_results: + idx = prompt_result.get("prompt_index", 0) + class_name = class_names[idx] if idx < len(class_names) else None + class_name = class_name or "foreground" + if class_mapping: + class_name = class_mapping.get(class_name, class_name) + for prediction in prompt_result.get("predictions", []): + confidence = float(prediction.get("confidence", 0.0)) + if confidence < threshold: + continue + for polygon in prediction.get("masks", []): + if polygon is None or len(polygon) < 3: + continue + xs = [float(point[0]) for point in polygon] + ys = [float(point[1]) for point in polygon] + x_min, y_min, x_max, y_max = min(xs), min(ys), max(xs), max(ys) + if x_max <= x_min or y_max <= y_min: + continue + flat_polygon = [ + coord + for point in polygon + for coord in (float(point[0]), float(point[1])) + ] + # polygon -> compact COCO RLE entirely in C; no dense mask is built. + rle = mask_utils.frPyObjects([flat_polygon], height, width)[0] + if mask_representation == "rle": + rle_dicts.append(rle) + else: + dense = mask_utils.decode(rle) + if dense.ndim == 3: + dense = dense[:, :, 0] + dense_masks.append(dense.astype(bool)) + xyxy.append([x_min, y_min, x_max, y_max]) + confidences.append(confidence) + class_ids.append(idx) + class_names_map[idx] = class_name + bboxes_metadata.append( + {DETECTION_ID_KEY: str(uuid.uuid4()), CLASS_NAME_KEY: class_name} + ) + + if mask_representation == "rle": + mask = InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), masks=rle_dicts + ) + elif dense_masks: + mask = torch.from_numpy(np.stack(dense_masks, axis=0)).to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + else: + mask = torch.zeros( + (0, height, width), dtype=torch.bool, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + return _assemble_detections( + image=image, + xyxy=xyxy, + confidences=confidences, + class_ids=class_ids, + class_names_map=class_names_map, + bboxes_metadata=bboxes_metadata, + mask=mask, + ) + + +def _assemble_detections( + image: WorkflowImageData, + xyxy: List[List[float]], + confidences: List[float], + class_ids: List[int], + class_names_map: Dict[int, str], + bboxes_metadata: List[dict], + mask, +) -> InstanceDetections: + """Pack the collected rows into an InstanceDetections, with every prediction + tensor allocated on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``.""" + n = len(xyxy) + if n == 0: + xyxy_t = torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + class_id_t = torch.zeros( + (0,), dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + confidence_t = torch.zeros( + (0,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + else: + xyxy_t = torch.tensor( + xyxy, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + class_id_t = torch.tensor( + class_ids, dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + confidence_t = torch.tensor( + confidences, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + detections = InstanceDetections( + xyxy=xyxy_t, class_id=class_id_t, confidence=confidence_t, mask=mask + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=class_names_map, + prediction_type=PREDICTION_TYPE, + inference_id=str(uuid.uuid4()), + ) + detections.bboxes_metadata = bboxes_metadata if n else None + return detections diff --git a/inference/core/workflows/core_steps/models/foundation/segment_anything3/v2_tensor.py b/inference/core/workflows/core_steps/models/foundation/segment_anything3/v2_tensor.py new file mode 100644 index 0000000000..09c885a857 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/segment_anything3/v2_tensor.py @@ -0,0 +1,785 @@ +"""Tensor-native sibling of `roboflow_core/sam3@v2`. + +Builds on the v1_tensor SAM3 conversion, adding per-class confidence thresholds + +cross-prompt NMS. Because `run_tensor_native_inference` is a THIN forward to +`segment_with_text_prompts` (it applies only a single global threshold floor), +this block must replicate the adapter's per-class-threshold + cross-prompt-NMS +orchestration itself. Rather than reach across the package boundary into the +private `inference.models.sam3` adapter internals (the previous, fragile design), +this block keeps LOCAL COPIES of those helpers (see the TODO(sam3-public-adapter) +note below); they run the per-class threshold + cross-prompt NMS on the raw +native masks, then build one `inference_models.InstanceDetections` per image +(RLE-default). + +When NMS runs against the RLE carrier, the COCO RLEs that NMS already encodes for +IoU are reused to build the output `InstancesRLEMasks` directly, instead of +re-encoding the kept dense masks a second time in `_build_instance_detections`. + +Output is tensor-native. REMOTE/proxy paths forward per-class thresholds + +nms_iou_threshold to the server (which applies them, matching the numpy block) and +rasterise the returned polygon response via the v1_tensor polygon path. +""" + +import uuid +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import requests +from pycocotools import mask as mask_utils +from pydantic import ConfigDict, Field, field_validator, model_validator + +from inference.core.entities.requests.sam3 import Sam3Prompt +from inference.core.env import ( + API_BASE_URL, + CORE_MODEL_SAM3_ENABLED, + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + ROBOFLOW_INTERNAL_SERVICE_NAME, + ROBOFLOW_INTERNAL_SERVICE_SECRET, + SAM3_EXEC_MODE, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import build_roboflow_api_headers +from inference.core.utils.url_utils import wrap_url +from inference.core.workflows.core_steps.common.entities import StepExecutionMode + +# Reuse the v1_tensor SAM3 conversion machinery verbatim. +from inference.core.workflows.core_steps.models.foundation.segment_anything3.v1_tensor import ( + Item, + _assemble_detections, + _build_instance_detections, + _build_instance_detections_from_polygons, + _normalize_class_names, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.offline import ensure_builtin_remote_execution_allowed +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_sdk import InferenceHTTPClient + +# TODO(sam3-public-adapter): these per-class-threshold + cross-prompt-NMS helpers +# are LOCAL COPIES of the private `inference.models.sam3. +# segment_anything3_inference_models` internals +# (`_to_numpy_masks`, `_collect_masks_with_per_prompt_threshold`, +# `_apply_nms_cross_prompt`, `_nms_greedy_pycocotools`). They were previously +# imported across the package boundary from those underscore-prefixed adapter +# internals, which carry no API-stability guarantee and forced this workflow +# block to track the adapter by hand. We copy them here (decoupling the block +# from adapter refactors) until inference_models exposes a PUBLIC, +# tested entrypoint that performs per-class threshold + cross-prompt NMS +# adapter-side (e.g. a `run_tensor_native_inference` mode that returns +# post-NMS, post-per-class-threshold native masks). When that lands, delete +# these copies and consume the public result directly. Keep this logic in +# byte-for-byte parity with the adapter so numpy/tensor behaviour stays defined +# in one place until the public API exists. + +LONG_DESCRIPTION = """ +Run Segment Anything 3 (zero-shot, text-prompted) with per-class confidence +thresholds and optional cross-prompt Non-Maximum Suppression. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "SAM 3", + "version": "v2", + "short_description": "Sam3", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["Sam3"], + "ui_manifest": { + "section": "model", + "icon": "fa-solid fa-eye", + "blockPriority": 9.49, + "needsGPU": True, + "inference": True, + }, + }, + protected_namespaces=(), + ) + + type: Literal["roboflow_core/sam3@v2"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), Optional[str]] = Field( + default="sam3/sam3_final", + description="model version. You only need to change this for fine tuned sam3 models.", + examples=["sam3/sam3_final", "$inputs.model_variant"], + ) + class_names: Optional[ + Union[List[str], str, Selector(kind=[LIST_OF_VALUES_KIND, STRING_KIND])] + ] = Field( + title="Class Names", + default=None, + description="List of classes to recognise", + examples=[["car", "person"], "$inputs.classes"], + ) + confidence: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + default=0.5, + title="Confidence Threshold", + description="Minimum confidence threshold for predicted masks", + examples=[0.3], + ) + per_class_confidence: Optional[ + Union[List[float], Selector(kind=[LIST_OF_VALUES_KIND])] + ] = Field( + default=None, + title="Per-Class Confidence", + description="List of confidence thresholds per class (must match class_names length)", + examples=[[0.3, 0.5, 0.7]], + ) + apply_nms: Union[Selector(kind=[BOOLEAN_KIND]), bool] = Field( + default=True, + title="Apply NMS", + description="Whether to apply Non-Maximum Suppression across prompts", + ) + nms_iou_threshold: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + default=0.9, + title="NMS IoU Threshold", + description="IoU threshold for cross-prompt NMS. Must be in [0.0, 1.0]", + examples=[0.5, 0.9], + ) + + @field_validator("nms_iou_threshold") + @classmethod + def _validate_nms_iou_threshold(cls, v): + if isinstance(v, (int, float)) and (v < 0.0 or v > 1.0): + raise ValueError("nms_iou_threshold must be between 0.0 and 1.0") + return v + + @model_validator(mode="after") + def _validate_per_class_confidence_length(self) -> "BlockManifest": + if not isinstance(self.per_class_confidence, list): + return self + if isinstance(self.class_names, list): + class_names_length = len(self.class_names) + elif isinstance(self.class_names, str): + class_names_length = len(self.class_names.split(",")) + else: + return self + if len(self.per_class_confidence) != class_names_length: + raise ValueError( + f"per_class_confidence length ({len(self.per_class_confidence)}) " + f"must match class_names length ({class_names_length})" + ) + return self + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restrictions = [ + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; run_locally() loads a model that needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + ] + if not CORE_MODEL_SAM3_ENABLED: + restrictions.append( + RuntimeRestriction( + severity=Severity.HARD, + note=( + "CORE_MODEL_SAM3_ENABLED=False on Roboflow Hosted " + "Serverless: the SAM3 endpoint is not registered, so " + "run_remotely() returns 404." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + ) + return restrictions + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + return ["sam3/sam3_final"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + if SAM3_EXEC_MODE == "remote": + # Proxy execution ignores the configured model id โ€” the proxy runs + # its own fixed SAM3 server-side; nothing to declare. + return [] + if self.model_id is None: + return [] + return [roboflow_platform_model(model_id=self.model_id)] + + +class SegmentAnything3BlockV2(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_names: Optional[Union[List[str], str]], + confidence: float, + per_class_confidence: Optional[List[float]], + apply_nms: bool, + nms_iou_threshold: float, + ) -> BlockResult: + # RLE mask output is enforced ALWAYS on the tensor path. The numpy SAM3 v2 + # sibling exposes no mask-format field, so neither does this manifest; the + # shared internal helpers (with SAM2/v1/v3) still take a `mask_representation` + # argument, so it is pinned to "rle" here. + mask_representation = "rle" + class_names = _normalize_class_names(class_names) + if SAM3_EXEC_MODE == "remote": + return self.run_via_request( + images=images, + class_names=class_names, + confidence=confidence, + per_class_confidence=per_class_confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + mask_representation=mask_representation, + ) + elif self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_names=class_names, + confidence=confidence, + per_class_confidence=per_class_confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + mask_representation=mask_representation, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_names=class_names, + confidence=confidence, + per_class_confidence=per_class_confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + mask_representation=mask_representation, + ) + raise ValueError(f"Unknown step execution mode: {self._step_execution_mode}") + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_names: List[Optional[str]], + confidence: float, + per_class_confidence: Optional[List[float]], + apply_nms: bool, + nms_iou_threshold: float, + mask_representation: str, + ) -> BlockResult: + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + + # Real Sam3Prompt objects carry .output_prob_thresh for the per-class helper. + sam3_prompts = [ + Sam3Prompt( + type="text", + text=class_name, + output_prob_thresh=_per_class_threshold(per_class_confidence, idx), + ) + for idx, class_name in enumerate(class_names) + ] + # The model forward takes plain {"text": cn} dicts. + native_prompts = [{"text": cn} for cn in class_names] + # MIN-floor so the backend floor never pre-clips a mask a LOWER per-class + # threshold would keep (mirrors adapter segment_image). + floor = _min_floor(confidence, per_class_confidence) + + results: List[dict] = [] + for single_image in images: + # SAM3 converts its input to numpy internally and is colour-agnostic + # (expects RGB, like the materialised tensor). Pass the tensor when it is + # already on device, otherwise an RGB host frame โ€” avoiding a forced CHW + # transpose + H2D that SAM3 would immediately undo. + if single_image.is_tensor_materialised(): + model_image = single_image.tensor_image + else: + model_image = np.ascontiguousarray(single_image.numpy_image[:, :, ::-1]) + per_image = self._model_manager.run_tensor_native_inference( + model_id, + images=[model_image], + prompts=native_prompts, + output_prob_thresh=float(floor), + ) + # When NMS runs against the RLE carrier, _collect_and_build reuses the + # COCO RLEs NMS already encoded for IoU (no second encode); otherwise it + # falls back to the dense packer (_build_instance_detections). + results.append( + { + "predictions": _collect_and_build( + per_prompt_results=per_image[0], + class_names=class_names, + prompts=sam3_prompts, + global_confidence=confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + image=single_image, + mask_representation=mask_representation, + ) + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_names: List[Optional[str]], + confidence: float, + per_class_confidence: Optional[List[float]], + apply_nms: bool, + nms_iou_threshold: float, + mask_representation: str, + ) -> BlockResult: + ensure_builtin_remote_execution_allowed("SAM3 remote execution") + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient(api_url=api_url, api_key=self._api_key) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + http_prompts = _build_http_prompts(class_names, per_class_confidence) + + results: List[dict] = [] + for single_image in images: + resp_json = client.sam3_concept_segment( + inference_input=single_image.base64_image, + prompts=http_prompts, + model_id=model_id, + output_prob_thresh=confidence, + nms_iou_threshold=nms_iou_threshold if apply_nms else None, + ) + results.append( + self._build_from_polygon_response( + resp_json=resp_json, + image=single_image, + class_names=class_names, + confidence=confidence, + mask_representation=mask_representation, + ) + ) + return results + + def run_via_request( + self, + images: Batch[WorkflowImageData], + class_names: List[Optional[str]], + confidence: float, + per_class_confidence: Optional[List[float]], + apply_nms: bool, + nms_iou_threshold: float, + mask_representation: str, + ) -> BlockResult: + ensure_builtin_remote_execution_allowed("SAM3 inference proxy execution") + endpoint = f"{API_BASE_URL}/inferenceproxy/seg-preview" + http_prompts = _build_http_prompts(class_names, per_class_confidence) + + results: List[dict] = [] + for single_image in images: + payload = { + "image": {"type": "base64", "value": single_image.base64_image}, + "prompts": http_prompts, + "output_prob_thresh": confidence, + "nms_iou_threshold": nms_iou_threshold if apply_nms else None, + } + headers = {"Content-Type": "application/json"} + if ROBOFLOW_INTERNAL_SERVICE_NAME: + headers["X-Roboflow-Internal-Service-Name"] = ( + ROBOFLOW_INTERNAL_SERVICE_NAME + ) + if ROBOFLOW_INTERNAL_SERVICE_SECRET: + headers["X-Roboflow-Internal-Service-Secret"] = ( + ROBOFLOW_INTERNAL_SERVICE_SECRET + ) + headers = build_roboflow_api_headers(explicit_headers=headers) + try: + response = requests.post( + wrap_url(f"{endpoint}?api_key={self._api_key}"), + json=payload, + headers=headers, + timeout=60, + ) + response.raise_for_status() + resp_json = response.json() + except Exception as exc: + raise Exception(f"SAM3 request failed: {exc}") + results.append( + self._build_from_polygon_response( + resp_json=resp_json, + image=single_image, + class_names=class_names, + confidence=confidence, + mask_representation=mask_representation, + ) + ) + return results + + def _build_from_polygon_response( + self, + resp_json: dict, + image: WorkflowImageData, + class_names: List[Optional[str]], + confidence: float, + mask_representation: str, + ) -> dict: + return { + "predictions": _build_instance_detections_from_polygons( + prompt_results=resp_json.get("prompt_results", []), + class_names=class_names, + image=image, + threshold=confidence, + mask_representation=mask_representation, + ) + } + + +def _build_http_prompts( + class_names: List[Optional[str]], + per_class_confidence: Optional[List[float]], +) -> List[dict]: + """Text prompt dicts for the REMOTE/proxy HTTP API, carrying per-class + output_prob_thresh when provided (matches numpy v2/v3). Unguarded index is + intentional parity โ€” the manifest validator enforces length for concrete lists.""" + prompts: List[dict] = [] + for idx, class_name in enumerate(class_names): + prompt: dict = {"type": "text", "text": class_name} + if per_class_confidence is not None: + prompt["output_prob_thresh"] = per_class_confidence[idx] + prompts.append(prompt) + return prompts + + +def _per_class_threshold( + per_class_confidence: Optional[List[float]], idx: int +) -> Optional[float]: + if per_class_confidence and idx < len(per_class_confidence): + return per_class_confidence[idx] + return None + + +def _min_floor(confidence: float, per_class_confidence: Optional[List[float]]) -> float: + if per_class_confidence: + return min([confidence, *[t for t in per_class_confidence if t is not None]]) + return confidence + + +# --------------------------------------------------------------------------- # +# LOCAL COPIES of the SAM3 adapter per-class-threshold + cross-prompt-NMS +# helpers. See TODO(sam3-public-adapter) at the top of this module. Keep these +# in byte-for-byte parity with +# inference.models.sam3.segment_anything3_inference_models until a public +# adapter entrypoint exists. +# --------------------------------------------------------------------------- # +def _to_numpy_masks(masks_any) -> np.ndarray: + if masks_any is None: + return np.zeros((0, 0, 0), dtype=np.uint8) + if hasattr(masks_any, "detach"): + masks_np = masks_any.detach().cpu().numpy().astype(np.uint8) + else: + arrs = [] + for m in masks_any: + if hasattr(m, "detach"): + arrs.append(m.detach().cpu().numpy().astype(np.uint8)) + else: + arrs.append(np.asarray(m, dtype=np.uint8)) + if not arrs: + return np.zeros((0, 0, 0), dtype=np.uint8) + masks_np = np.stack(arrs, axis=0) + if masks_np.ndim == 4 and masks_np.shape[1] == 1: + masks_np = masks_np[:, 0, ...] + elif masks_np.ndim == 2: + masks_np = masks_np[None, ...] + return masks_np + + +def _collect_masks_with_per_prompt_threshold( + processed: Dict[int, Dict[str, Any]], + prompts: List[Sam3Prompt], + default_threshold: float, +) -> List[Tuple[int, np.ndarray, float]]: + all_masks: List[Tuple[int, np.ndarray, float]] = [] + for idx, p in enumerate(prompts): + prompt_thresh = getattr(p, "output_prob_thresh", None) + if prompt_thresh is None: + prompt_thresh = default_threshold + masks_np = _to_numpy_masks(processed[idx]["masks"]) + scores = processed[idx]["scores"] + if masks_np.ndim != 3 or 0 in masks_np.shape: + continue + for mask, score in zip(masks_np, scores): + if score >= prompt_thresh: + all_masks.append((idx, mask, float(score))) + return all_masks + + +def _nms_greedy_pycocotools( + rles: List[dict], + confidences: np.ndarray, + iou_threshold: float = 0.5, +) -> np.ndarray: + num_detections = len(rles) + if num_detections == 0: + return np.array([], dtype=bool) + sort_index = np.argsort(confidences)[::-1] + sorted_rles = [rles[i] for i in sort_index] + ious = mask_utils.iou(sorted_rles, sorted_rles, [0] * num_detections) + keep = np.ones(num_detections, dtype=bool) + for i in range(num_detections): + if keep[i]: + condition = ious[i, :] > iou_threshold + keep[i + 1 :] = np.where(condition[i + 1 :], False, keep[i + 1 :]) + return keep[np.argsort(sort_index)] + + +def _encode_mask_to_rle(mask_np: np.ndarray) -> dict: + """COCO-RLE-encode a single (H,W) mask, returning the pycocotools dict + ({"size": [...], "counts": }) โ€” the exact shape + InstancesRLEMasks.from_coco_rle_masks consumes.""" + mb = (mask_np > 0).astype(np.uint8) + return mask_utils.encode(np.asfortranarray(mb)) + + +def _apply_nms_cross_prompt_with_rles( + all_masks: List[Tuple[int, np.ndarray, float]], + iou_threshold: float, +) -> Tuple[List[Tuple[int, np.ndarray, float]], List[dict]]: + """Cross-prompt greedy NMS (local copy of the adapter's + `_apply_nms_cross_prompt`) that ALSO returns the per-kept-mask COCO RLEs it + already computed for IoU. Reusing these downstream avoids encoding the kept + masks to RLE a second time when building the RLE output carrier. + """ + if not all_masks: + return all_masks, [] + rles = [_encode_mask_to_rle(mask_np) for _, mask_np, _ in all_masks] + confidences = np.array([score for _, _, score in all_masks]) + keep = _nms_greedy_pycocotools(rles, confidences, iou_threshold) + kept_masks = [all_masks[i] for i in range(len(all_masks)) if keep[i]] + kept_rles = [rles[i] for i in range(len(all_masks)) if keep[i]] + return kept_masks, kept_rles + + +def _threshold_and_nms( + per_prompt_results: List[dict], + prompts: List[Sam3Prompt], + global_confidence: float, + apply_nms: bool, + nms_iou_threshold: float, +) -> Tuple[List[Tuple[int, np.ndarray, float]], Optional[List[dict]]]: + """Run per-class threshold + (optional) cross-prompt NMS on the raw native + masks. Returns the kept (prompt_idx, mask(H,W), score) rows and, when NMS + actually ran, the matching COCO RLEs (else None). + + The adapter helpers key `processed` by enumerate(prompts) order, so we key the + native results by prompt_index (pre-seeding all idx) to keep them aligned, and + pass raw (N,1,H,W) masks straight through โ€” `_to_numpy_masks` squeezes/casts. + """ + processed: Dict[int, Dict[str, Any]] = { + idx: {"masks": None, "scores": []} for idx in range(len(prompts)) + } + for result in per_prompt_results: + idx = result.get("prompt_index", 0) + if idx not in processed: + continue + processed[idx] = { + "masks": result.get("masks"), + "scores": list(result.get("scores", [])), + } + + all_masks = _collect_masks_with_per_prompt_threshold( + processed=processed, + prompts=prompts, + default_threshold=global_confidence, + ) + kept_rles: Optional[List[dict]] = None + if apply_nms and nms_iou_threshold is not None and len(all_masks) > 0: + all_masks, kept_rles = _apply_nms_cross_prompt_with_rles( + all_masks, nms_iou_threshold + ) + return all_masks, kept_rles + + +def _collect_from_native_with_nms( + per_prompt_results: List[dict], + class_names: List[Optional[str]], + prompts: List[Sam3Prompt], + global_confidence: float, + apply_nms: bool, + nms_iou_threshold: float, +) -> List[Item]: + """Per-class-threshold + cross-prompt NMS on the raw native masks, flattened to + Items for `_build_instance_detections` (dense-mask packer). Used by v3_tensor + and by the v2 `dense` carrier path. + """ + all_masks, _ = _threshold_and_nms( + per_prompt_results=per_prompt_results, + prompts=prompts, + global_confidence=global_confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + ) + items: List[Item] = [] + for prompt_idx, mask, score in all_masks: + class_name = class_names[prompt_idx] if prompt_idx < len(class_names) else None + items.append( + (mask.astype(bool), float(score), prompt_idx, class_name or "foreground") + ) + return items + + +def _build_instance_detections_reusing_nms_rles( + all_masks: List[Tuple[int, np.ndarray, float]], + kept_rles: List[dict], + class_names: List[Optional[str]], + image: WorkflowImageData, +) -> InstanceDetections: + """RLE-carrier assembler that REUSES the COCO RLEs NMS already encoded, so the + kept masks are not re-encoded a second time. Bbox is derived from each kept + dense mask via np.where (cheap, no encode). Rows stay in lockstep because + `all_masks` and `kept_rles` are produced together by + `_apply_nms_cross_prompt_with_rles`. + """ + height, width = image._read_shape_without_materialization() + xyxy: List[List[float]] = [] + confidences: List[float] = [] + class_ids: List[int] = [] + class_names_map: Dict[int, str] = {} + bboxes_metadata: List[dict] = [] + rle_dicts: List[dict] = [] + + for (prompt_idx, mask_np, score), rle in zip(all_masks, kept_rles): + ys, xs = np.where(mask_np > 0) + if xs.size == 0: + continue + class_name = ( + class_names[prompt_idx] if prompt_idx < len(class_names) else None + ) or "foreground" + xyxy.append( + [float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max())] + ) + confidences.append(float(score)) + class_ids.append(prompt_idx) + class_names_map[prompt_idx] = class_name + bboxes_metadata.append( + {DETECTION_ID_KEY: str(uuid.uuid4()), CLASS_NAME_KEY: class_name} + ) + rle_dicts.append(rle) + + mask = InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), masks=rle_dicts + ) + return _assemble_detections( + image=image, + xyxy=xyxy, + confidences=confidences, + class_ids=class_ids, + class_names_map=class_names_map, + bboxes_metadata=bboxes_metadata, + mask=mask, + ) + + +def _collect_and_build( + per_prompt_results: List[dict], + class_names: List[Optional[str]], + prompts: List[Sam3Prompt], + global_confidence: float, + apply_nms: bool, + nms_iou_threshold: float, + image: WorkflowImageData, + mask_representation: str, +) -> InstanceDetections: + """LOCAL-path builder. When NMS runs against the RLE carrier, reuse the COCO + RLEs NMS already computed (avoids the double encode). Otherwise fall back to + the v1 dense packer (`_build_instance_detections`).""" + all_masks, kept_rles = _threshold_and_nms( + per_prompt_results=per_prompt_results, + prompts=prompts, + global_confidence=global_confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + ) + if mask_representation == "rle" and kept_rles is not None: + return _build_instance_detections_reusing_nms_rles( + all_masks=all_masks, + kept_rles=kept_rles, + class_names=class_names, + image=image, + ) + items: List[Item] = [] + for prompt_idx, mask, score in all_masks: + class_name = class_names[prompt_idx] if prompt_idx < len(class_names) else None + items.append( + (mask.astype(bool), float(score), prompt_idx, class_name or "foreground") + ) + return _build_instance_detections( + items=items, image=image, mask_representation=mask_representation + ) diff --git a/inference/core/workflows/core_steps/models/foundation/segment_anything3/v3_tensor.py b/inference/core/workflows/core_steps/models/foundation/segment_anything3/v3_tensor.py new file mode 100644 index 0000000000..3efdda319f --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/segment_anything3/v3_tensor.py @@ -0,0 +1,572 @@ +"""Tensor-native sibling of `roboflow_core/sam3@v3`. + +v3 = v2 (per-class thresholds + cross-prompt NMS) plus `class_mapping` (rename +predicted class names after inference). + +BREAKING CHANGE (tensor path) โ€” mask output format: this sibling KEEPS the numpy +v3 manifest field `output_format: Literal["rle", "polygons"]` verbatim (so existing +workflow definitions validate identically against the tensor block), but +`run_tensor_native_inference` always returns raw binary masks with no polygon +analog, so RLE output is ENFORCED ALWAYS on this tensor path (see run()). Selecting +`output_format="polygons"` is now a NO-OP โ€” it is silently downgraded to 'rle' with +a warning. That downgrade IS the breaking change: the field is honored for 'rle' +and ignored for 'polygons'. + +Output kinds: the numpy v3 declares DUAL kinds +(RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND + INSTANCE_SEGMENTATION_PREDICTION_KIND). +The tensor world has the matching pair +(TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND + +TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND), so `describe_outputs` declares +BOTH (RLE first) โ€” matching the roboflow instance_segmentation tensor producers and +the established convention. At runtime there is still ONE `InstanceDetections` +carrier (always RLE, since `output_format` is enforced to 'rle'); the dual kind is purely a +wiring/advertisement concern. Downstream tensor consumers (mask/polygon/halo/icon +visualizers, trackers, fusion blocks) accept both kinds, so no graph is broken by +the carrier choice. + +Private adapter coupling: the per-class-threshold + cross-prompt-NMS orchestration +is reused from `v2_tensor` (`_collect_from_native_with_nms`, `_min_floor`, +`_per_class_threshold`, `_build_http_prompts`). v2_tensor now holds LOCAL COPIES of +those adapter helpers (see its TODO(sam3-public-adapter) note) instead of importing +the private `inference.models.sam3` internals, so v3 no longer transitively depends +on that private cross-package surface. + +`class_mapping` is applied on the flat `items` list before building +`InstanceDetections`, so both image_metadata's class-names map and per-instance +bboxes_metadata inherit the mapped names (v3's sv-based `_apply_class_mapping` is a +no-op on InstanceDetections and cannot be reused). +""" + +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +import requests +from pydantic import ConfigDict, Field, field_validator, model_validator + +from inference.core import logger +from inference.core.entities.requests.sam3 import Sam3Prompt +from inference.core.env import ( + API_BASE_URL, + CORE_MODEL_SAM3_ENABLED, + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + ROBOFLOW_INTERNAL_SERVICE_NAME, + ROBOFLOW_INTERNAL_SERVICE_SECRET, + SAM3_EXEC_MODE, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import build_roboflow_api_headers +from inference.core.utils.url_utils import wrap_url +from inference.core.workflows.core_steps.common.entities import StepExecutionMode + +# Reuse the v1_tensor conversion machinery + the v2_tensor per-class/NMS collector. +from inference.core.workflows.core_steps.models.foundation.segment_anything3.v1_tensor import ( + Item, + _build_instance_detections, + _build_instance_detections_from_polygons, + _normalize_class_names, +) +from inference.core.workflows.core_steps.models.foundation.segment_anything3.v2_tensor import ( + _build_http_prompts, + _collect_from_native_with_nms, + _min_floor, + _per_class_threshold, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + DICTIONARY_KIND, + FLOAT_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.offline import ensure_builtin_remote_execution_allowed +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, +) +from inference_sdk import InferenceHTTPClient + +LONG_DESCRIPTION = """ +Run Segment Anything 3 (zero-shot, text-prompted) with per-class confidence +thresholds, optional cross-prompt NMS, and post-inference class renaming. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "SAM 3", + "version": "v3", + "short_description": "Sam3", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["Sam3"], + "ui_manifest": { + "section": "model", + "icon": "fa-solid fa-eye", + "blockPriority": 9.48, + "needsGPU": True, + "inference": True, + }, + }, + protected_namespaces=(), + ) + + type: Literal["roboflow_core/sam3@v3"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), Optional[str]] = Field( + default="sam3/sam3_final", + description="model version. You only need to change this for fine tuned sam3 models.", + examples=["sam3/sam3_final", "$inputs.model_variant"], + ) + class_names: Optional[ + Union[List[str], str, Selector(kind=[LIST_OF_VALUES_KIND, STRING_KIND])] + ] = Field( + title="Class Names", + default=None, + description="List of classes to recognise", + examples=[["car", "person"], "$inputs.classes"], + ) + class_mapping: Optional[Union[Dict[str, str], Selector(kind=[DICTIONARY_KIND])]] = ( + Field( + default=None, + title="Class Mapping", + description="Maps class names in predictions to different output names. Applied " + "after inference, e.g. {'cat': 'gato'} renames 'cat' predictions to 'gato'.", + examples=[{"cat": "gato", "dog": "perro"}], + ) + ) + confidence: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + default=0.5, + title="Confidence Threshold", + description="Minimum confidence threshold for predicted masks", + examples=[0.3], + ) + per_class_confidence: Optional[ + Union[List[float], Selector(kind=[LIST_OF_VALUES_KIND])] + ] = Field( + default=None, + title="Per-Class Confidence", + description="List of confidence thresholds per class (must match class_names length)", + examples=[[0.3, 0.5, 0.7]], + ) + apply_nms: Union[Selector(kind=[BOOLEAN_KIND]), bool] = Field( + default=True, + title="Apply NMS", + description="Whether to apply Non-Maximum Suppression across prompts", + ) + nms_iou_threshold: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + default=0.9, + title="NMS IoU Threshold", + description="IoU threshold for cross-prompt NMS. Must be in [0.0, 1.0]", + examples=[0.5, 0.9], + ) + # NOTE (BREAKING CHANGE, tensor path): this field mirrors the numpy v3 sibling's + # `output_format: Literal["rle", "polygons"]` verbatim so existing workflow + # definitions validate identically. At runtime 'polygons' is NOT honored โ€” RLE + # mask output is enforced ALWAYS (see run()); a 'polygons' selection is silently + # downgraded to 'rle' with a warning. That downgrade is the breaking change. + output_format: Literal["rle", "polygons"] = Field( + default="rle", + description="Mask output format. Mirrors the numpy v3 sibling. On the tensor " + "path only 'rle' is honored; 'polygons' is downgraded to 'rle' with a " + "warning (breaking change).", + ) + + @field_validator("nms_iou_threshold") + @classmethod + def _validate_nms_iou_threshold(cls, v): + if isinstance(v, (int, float)) and (v < 0.0 or v > 1.0): + raise ValueError("nms_iou_threshold must be between 0.0 and 1.0") + return v + + @model_validator(mode="after") + def _validate_per_class_confidence_length(self) -> "BlockManifest": + if not isinstance(self.per_class_confidence, list): + return self + if isinstance(self.class_names, list): + class_names_length = len(self.class_names) + elif isinstance(self.class_names, str): + class_names_length = len(self.class_names.split(",")) + else: + return self + if len(self.per_class_confidence) != class_names_length: + raise ValueError( + f"per_class_confidence length ({len(self.per_class_confidence)}) " + f"must match class_names length ({class_names_length})" + ) + return self + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + # Declare BOTH the RLE and the plain tensor instance-seg kinds (RLE first, + # mirroring the enforced `output_format="rle"`). This restores the + # dual-kind shape the numpy v3 sibling declares + # (RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND + + # INSTANCE_SEGMENTATION_PREDICTION_KIND) and matches the established tensor + # convention used by the roboflow instance_segmentation tensor producers + # (v1-v4) and accepted by every tensor instance-seg consumer (mask/polygon/ + # halo/icon visualizers, trackers, fusion blocks). The runtime carrier is + # still ONE InstanceDetections (always RLE, since `output_format` is enforced + # to 'rle'); the two kinds are advertised purely so downstream wiring that selects either + # the rle-tensor or the plain-tensor kind resolves against this producer. + # Audit: no tensor consumer requires the rle-tensor kind EXCLUSIVELY, so + # the previous single-kind output was wiring-compatible too, but it diverged + # from the dual-kind convention; this aligns it. + return [ + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restrictions = [ + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; run_locally() loads a model that needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + ] + if not CORE_MODEL_SAM3_ENABLED: + restrictions.append( + RuntimeRestriction( + severity=Severity.HARD, + note=( + "CORE_MODEL_SAM3_ENABLED=False on Roboflow Hosted " + "Serverless: the SAM3 endpoint is not registered, so " + "run_remotely() returns 404." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + ) + return restrictions + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + return ["sam3/sam3_final"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + if SAM3_EXEC_MODE == "remote": + # Proxy execution ignores the configured model id โ€” the proxy runs + # its own fixed SAM3 server-side; nothing to declare. + return [] + if self.model_id is None: + return [] + return [roboflow_platform_model(model_id=self.model_id)] + + +class SegmentAnything3BlockV3(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_names: Optional[Union[List[str], str]], + class_mapping: Optional[Dict[str, str]], + confidence: float, + per_class_confidence: Optional[List[float]], + apply_nms: bool, + nms_iou_threshold: float, + output_format: Literal["rle", "polygons"], + ) -> BlockResult: + # BREAKING CHANGE (tensor path): RLE mask output is enforced ALWAYS. The + # manifest keeps numpy v3's `output_format: Literal["rle", "polygons"]` + # verbatim, but `run_tensor_native_inference` has no polygon analog, so + # 'polygons' is silently downgraded to 'rle' with a warning. The shared + # internal helpers (with SAM2/v1/v2) take a `mask_representation` arg, so it + # is pinned to "rle" here. + if output_format != "rle": + logger.warning( + "SAM3 v3 (tensor) block: output_format=%r is not honored; RLE mask " + "output is enforced on the tensor path. Selecting 'polygons' is a " + "no-op (breaking change vs the numpy v3 sibling).", + output_format, + ) + mask_representation = "rle" + class_names = _normalize_class_names(class_names) + if SAM3_EXEC_MODE == "remote": + return self.run_via_request( + images=images, + class_names=class_names, + class_mapping=class_mapping, + confidence=confidence, + per_class_confidence=per_class_confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + mask_representation=mask_representation, + ) + elif self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_names=class_names, + class_mapping=class_mapping, + confidence=confidence, + per_class_confidence=per_class_confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + mask_representation=mask_representation, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_names=class_names, + class_mapping=class_mapping, + confidence=confidence, + per_class_confidence=per_class_confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + mask_representation=mask_representation, + ) + raise ValueError(f"Unknown step execution mode: {self._step_execution_mode}") + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_names: List[Optional[str]], + class_mapping: Optional[Dict[str, str]], + confidence: float, + per_class_confidence: Optional[List[float]], + apply_nms: bool, + nms_iou_threshold: float, + mask_representation: str, + ) -> BlockResult: + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + + sam3_prompts = [ + Sam3Prompt( + type="text", + text=class_name, + output_prob_thresh=_per_class_threshold(per_class_confidence, idx), + ) + for idx, class_name in enumerate(class_names) + ] + native_prompts = [{"text": cn} for cn in class_names] + floor = _min_floor(confidence, per_class_confidence) + + results: List[dict] = [] + for single_image in images: + # SAM3 converts its input to numpy internally and is colour-agnostic + # (expects RGB, like the materialised tensor). Pass the tensor when it is + # already on device, otherwise an RGB host frame โ€” avoiding a forced CHW + # transpose + H2D that SAM3 would immediately undo. + if single_image.is_tensor_materialised(): + model_image = single_image.tensor_image + else: + model_image = np.ascontiguousarray(single_image.numpy_image[:, :, ::-1]) + per_image = self._model_manager.run_tensor_native_inference( + model_id, + images=[model_image], + prompts=native_prompts, + output_prob_thresh=float(floor), + ) + items = _collect_from_native_with_nms( + per_prompt_results=per_image[0], + class_names=class_names, + prompts=sam3_prompts, + global_confidence=confidence, + apply_nms=apply_nms, + nms_iou_threshold=nms_iou_threshold, + ) + items = _apply_class_mapping_to_items(items, class_mapping) + results.append( + { + "predictions": _build_instance_detections( + items=items, + image=single_image, + mask_representation=mask_representation, + ) + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_names: List[Optional[str]], + class_mapping: Optional[Dict[str, str]], + confidence: float, + per_class_confidence: Optional[List[float]], + apply_nms: bool, + nms_iou_threshold: float, + mask_representation: str, + ) -> BlockResult: + ensure_builtin_remote_execution_allowed("SAM3 remote execution") + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient(api_url=api_url, api_key=self._api_key) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + http_prompts = _build_http_prompts(class_names, per_class_confidence) + + results: List[dict] = [] + for single_image in images: + resp_json = client.sam3_concept_segment( + inference_input=single_image.base64_image, + prompts=http_prompts, + model_id=model_id, + output_prob_thresh=confidence, + nms_iou_threshold=nms_iou_threshold if apply_nms else None, + ) + results.append( + self._build_from_polygon_response( + resp_json=resp_json, + image=single_image, + class_names=class_names, + class_mapping=class_mapping, + confidence=confidence, + mask_representation=mask_representation, + ) + ) + return results + + def run_via_request( + self, + images: Batch[WorkflowImageData], + class_names: List[Optional[str]], + class_mapping: Optional[Dict[str, str]], + confidence: float, + per_class_confidence: Optional[List[float]], + apply_nms: bool, + nms_iou_threshold: float, + mask_representation: str, + ) -> BlockResult: + ensure_builtin_remote_execution_allowed("SAM3 inference proxy execution") + endpoint = f"{API_BASE_URL}/inferenceproxy/seg-preview" + http_prompts = _build_http_prompts(class_names, per_class_confidence) + + results: List[dict] = [] + for single_image in images: + payload = { + "image": {"type": "base64", "value": single_image.base64_image}, + "prompts": http_prompts, + "output_prob_thresh": confidence, + "nms_iou_threshold": nms_iou_threshold if apply_nms else None, + } + headers = {"Content-Type": "application/json"} + if ROBOFLOW_INTERNAL_SERVICE_NAME: + headers["X-Roboflow-Internal-Service-Name"] = ( + ROBOFLOW_INTERNAL_SERVICE_NAME + ) + if ROBOFLOW_INTERNAL_SERVICE_SECRET: + headers["X-Roboflow-Internal-Service-Secret"] = ( + ROBOFLOW_INTERNAL_SERVICE_SECRET + ) + headers = build_roboflow_api_headers(explicit_headers=headers) + try: + response = requests.post( + wrap_url(f"{endpoint}?api_key={self._api_key}"), + json=payload, + headers=headers, + timeout=60, + ) + response.raise_for_status() + resp_json = response.json() + except Exception as exc: + raise Exception(f"SAM3 request failed: {exc}") + results.append( + self._build_from_polygon_response( + resp_json=resp_json, + image=single_image, + class_names=class_names, + class_mapping=class_mapping, + confidence=confidence, + mask_representation=mask_representation, + ) + ) + return results + + def _build_from_polygon_response( + self, + resp_json: dict, + image: WorkflowImageData, + class_names: List[Optional[str]], + class_mapping: Optional[Dict[str, str]], + confidence: float, + mask_representation: str, + ) -> dict: + return { + "predictions": _build_instance_detections_from_polygons( + prompt_results=resp_json.get("prompt_results", []), + class_names=class_names, + image=image, + threshold=confidence, + mask_representation=mask_representation, + class_mapping=class_mapping, + ) + } + + +def _apply_class_mapping_to_items( + items: List[Item], class_mapping: Optional[Dict[str, str]] +) -> List[Item]: + """Rename predicted class names on the flat items list before building + InstanceDetections, so image_metadata's class-names map and per-instance + bboxes_metadata both inherit the mapped names. (v3's sv-based + _apply_class_mapping is a no-op on InstanceDetections and is not reused.)""" + if not class_mapping: + return items + return [ + (mask, score, class_id, class_mapping.get(class_name, class_name)) + for (mask, score, class_id, class_name) in items + ] diff --git a/inference/core/workflows/core_steps/models/foundation/segment_anything3_3d/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/segment_anything3_3d/v1_tensor.py new file mode 100644 index 0000000000..526d6550bf --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/segment_anything3_3d/v1_tensor.py @@ -0,0 +1,305 @@ +"""Tensor-native sibling of `roboflow_core/segment_anything3_3d_objects@v1`. + +SCRATCH โ€” first pass for review. SAM3-3D has NO inference_models tensor path (it is +a request-based mesh / Gaussian-splat generator). So this sibling does NOT change +the model invocation โ€” it only adapts the DATA REPRESENTATION at the block +boundary so the tensor-native pipeline can drive it: + +- `mask_input` arrives as `inference_models.InstanceDetections` (the tensor-native + instance-segmentation kind) instead of `sv.Detections`; `extract_masks_from_input` + materialises a list of (H, W) binary numpy masks from it (RLE or dense), which is + exactly what the numpy block produced from `sv.Detections.mask`. +- The IMAGE is consumed via `base64_image` / `to_inference_format` exactly as the + numpy block โ€” `WorkflowImageData` materialises those from the CHW tensor_image + transparently; the request-based model needs base64/numpy, not a tensor, so there + is no tensor image path to take here. +- Outputs (mesh_glb / gaussian_ply / objects / inference_time) are 3D artifacts, + not predictions โ€” unchanged. +""" + +import base64 +from typing import Any, List, Literal, Optional, Type, Union + +import numpy as np +from pydantic import ConfigDict, Field + +from inference.core.entities.requests.sam3_3d import Sam3_3D_Objects_InferenceRequest +from inference.core.entities.responses.sam3_3d import Sam3_3D_Objects_Response +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + SAM3_3D_OBJECTS_ENABLED, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + instance_mask_to_numpy, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.offline import ensure_builtin_remote_execution_allowed +from inference.core.workflows.prototypes.block import ( + AirGappedAvailability, + BlockResult, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_sdk import InferenceHTTPClient + +LONG_DESCRIPTION = """ +Generate 3D meshes and Gaussian splatting from 2D images with mask prompts. + +Accepts masks as: tensor-native instance segmentation (InstanceDetections) or a +flat list of polygon coordinates in COCO format. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "SAM3D", + "version": "v1", + "short_description": "Generate 3D meshes and Gaussian splatting from 2D images with mask prompts.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "beta": True, + "search_keywords": ["SAM3_3D", "3D", "mesh", "gaussian splatting"], + "ui_manifest": { + "section": "model", + "icon": "far fa-cube", + "blockPriority": 9.0, + "needsGPU": True, + "inference": True, + }, + }, + protected_namespaces=(), + ) + + type: Literal["roboflow_core/segment_anything3_3d_objects@v1"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + mask_input: Selector( + kind=[LIST_OF_VALUES_KIND, TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND] + ) = Field( + description="Mask input - either instance segmentation predictions (e.g., from SAM2) or a flat list of polygon coordinates in COCO format [x1, y1, x2, y2, x3, y3, ...]", + examples=["$steps.sam2.predictions", "$steps.detections.mask_polygon"], + ) + + @classmethod + def get_air_gapped_availability(cls) -> AirGappedAvailability: + return AirGappedAvailability(available=False, reason="requires_internet") + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images", "mask_input"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="mesh_glb", + kind=[STRING_KIND], + description="Scene mesh in GLB format (base64 encoded)", + ), + OutputDefinition( + name="gaussian_ply", + kind=[STRING_KIND], + description="Combined Gaussian splatting in PLY format (base64 encoded)", + ), + OutputDefinition( + name="objects", + kind=[LIST_OF_VALUES_KIND], + description="List of individual objects, each with mesh_glb, gaussian_ply, and metadata (rotation, translation, scale)", + ), + OutputDefinition( + name="inference_time", + kind=[FLOAT_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restrictions = [ + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; run_locally() loads a model that needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + ] + if not SAM3_3D_OBJECTS_ENABLED: + restrictions.append( + RuntimeRestriction( + severity=Severity.HARD, + note=( + "SAM3_3D_OBJECTS_ENABLED=False on Roboflow Hosted " + "Serverless: the SAM3 3D endpoint is not registered, so " + "run_remotely() returns 404." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + ) + return restrictions + + +class SegmentAnything3_3D_ObjectsBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + mask_input: Batch[Union[InstanceDetections, List[float]]], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally(images=images, mask_input=mask_input) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely(images=images, mask_input=mask_input) + raise ValueError(f"Unknown step execution mode: {self._step_execution_mode}") + + def run_remotely( + self, + images: Batch[WorkflowImageData], + mask_input: Batch[Union[InstanceDetections, List[float]]], + ) -> BlockResult: + ensure_builtin_remote_execution_allowed("SAM3 3D remote execution") + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient(api_url=api_url, api_key=self._api_key) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + results = [] + model_id = "sam3-3d-objects" + for single_image, single_mask_input in zip(images, mask_input): + converted_mask = extract_masks_from_input(single_mask_input) + if isinstance(converted_mask, list): + converted_mask = [ + mask.tolist() if isinstance(mask, np.ndarray) else mask + for mask in converted_mask + ] + result = client.sam3_3d_infer( + inference_input=single_image.base64_image, + mask_input=converted_mask, + model_id=model_id, + ) + results.append( + { + "mesh_glb": result.get("mesh_glb"), + "gaussian_ply": result.get("gaussian_ply"), + "objects": result.get("objects", []), + "inference_time": result.get("time", 0.0), + } + ) + return results + + def run_locally( + self, + images: Batch[WorkflowImageData], + mask_input: Batch[Union[InstanceDetections, List[float]]], + ) -> BlockResult: + results = [] + model_id = "sam3-3d-objects" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + + for single_image, single_mask_input in zip(images, mask_input): + converted_mask = extract_masks_from_input(single_mask_input) + inference_request = Sam3_3D_Objects_InferenceRequest( + image=single_image.to_inference_format(numpy_preferred=True), + mask_input=converted_mask, + api_key=self._api_key, + model_id=model_id, + ) + response: Sam3_3D_Objects_Response = ( + self._model_manager.infer_from_request_sync(model_id, inference_request) + ) + results.append(_format_response(response)) + return results + + +def extract_masks_from_input(mask_input: Any) -> Any: + """Materialise a list of (H, W) binary numpy masks from a tensor-native + InstanceDetections (RLE or dense, via instance_mask_to_numpy); pass through + other formats (e.g. a flat polygon list) unchanged. + + Matches the numpy block, which returned `list(sv.Detections.mask)`. + """ + if isinstance(mask_input, InstanceDetections): + number_of_instances = len(mask_input) + if number_of_instances == 0: + raise ValueError("InstanceDetections contains no detections.") + return [ + instance_mask_to_numpy(mask_input, index) + for index in range(number_of_instances) + ] + return mask_input + + +def _format_response(response: Sam3_3D_Objects_Response) -> dict: + """Format response with base64 encoded outputs.""" + + def encode(data): + return base64.b64encode(data).decode("utf-8") if data else None + + objects_list = [ + { + "mesh_glb": encode(obj.mesh_glb), + "gaussian_ply": encode(obj.gaussian_ply), + "metadata": { + "rotation": obj.metadata.rotation, + "translation": obj.metadata.translation, + "scale": obj.metadata.scale, + }, + } + for obj in response.objects + ] + + return { + "mesh_glb": encode(response.mesh_glb), + "gaussian_ply": encode(response.gaussian_ply), + "objects": objects_list, + "inference_time": response.time, + } diff --git a/inference/core/workflows/core_steps/models/foundation/segment_anything3_interactive/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/segment_anything3_interactive/v1_tensor.py new file mode 100644 index 0000000000..1bb46d4594 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/segment_anything3_interactive/v1_tensor.py @@ -0,0 +1,872 @@ +"""Tensor-native sibling of ``segment_anything3_interactive/v1.py``, loaded when +``ENABLE_TENSOR_DATA_REPRESENTATION`` is on. + +The numpy source drives SAM3 PVS through SAM2-style requests against +``sam3/sam3_interactive`` and expands each polygon of the response into a +separate ``sv.Detections`` row. Tensor-native differences: + +- ``boxes`` prompts arrive as native ``Detections`` / ``InstanceDetections`` / + ``(KeyPoints, Detections)`` โ€” box coordinates are read from the ``xyxy`` + tensor, class names resolved via per-box ``CLASS_NAME_KEY`` override or the + ``image_metadata[CLASS_NAMES_KEY]`` map, ``detection_id`` from + ``bboxes_metadata`` (mirroring ``segment_anything2/v1_tensor.py``). +- LOCAL goes through ``ModelManager.run_tensor_native_inference`` (the + ``InferenceModelsSAM3InteractiveAdapter`` ``action="segment"`` bridge to + ``SAM3Torch.segment_with_visual_prompts``): tensor-resident images are handed + over directly (RGB; host-only images are flipped BGRโ†’RGB), prompts are + converted with the exact ``to_sam2_inputs`` + padding chain the legacy + adapter applied, and raw logits are binarised at the adapter's + ``MASK_THRESHOLD`` โ€” instances carry the model's RAW binary mask, one + instance per prompt, no polygon collapse (the family convention; numpy + splits a multi-region mask into one instance per polygon). REMOTE asks for + ``mask_input_format="rle"``; the proxy path (``SAM3_EXEC_MODE=remote``) + requests RLE too but falls back to rasterising polygon-format predictions if + the proxy strips the ``format`` field. +- Masks are carried as ``InstancesRLEMasks`` (compact, serverless-safe); + ``xyxy`` is the mask's tight bbox (inclusive min/max, no +1 โ€” family-wide + convention keeping numpy-faithful boxes). +""" + +import uuid +from dataclasses import dataclass +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import requests +import torch +from pydantic import ConfigDict, Field, model_validator + +from inference.core import logger +from inference.core.entities.requests.sam2 import Box, Point, Sam2Prompt, Sam2PromptSet +from inference.core.env import ( + API_BASE_URL, + CORE_MODEL_SAM3_ENABLED, + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + ROBOFLOW_INTERNAL_SERVICE_NAME, + ROBOFLOW_INTERNAL_SERVICE_SECRET, + SAM3_EXEC_MODE, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import build_roboflow_api_headers +from inference.core.utils.url_utils import wrap_url +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, + split_key_point_prediction, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + IMAGE_KIND, + LABELED_POINTS_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.offline import ensure_builtin_remote_execution_allowed +from inference.core.workflows.prototypes.block import ( + BlockResult, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import ( + coco_rle_masks_to_numpy_mask, + torch_mask_to_coco_rle, +) +from inference_sdk import InferenceHTTPClient + +DETECTIONS_CLASS_NAME_FIELD = "class_name" +DETECTION_ID_FIELD = "detection_id" + +SAM3_INTERACTIVE_MODEL_ID = "sam3/sam3_interactive" + +PREDICTION_TYPE = "instance-segmentation" + +# The legacy adapter's MASK_THRESHOLD: raw logits from the model are binarised +# at this cutoff (see `_build_rle_response` in +# inference/models/sam3/visual_segmentation_inference_models.py). +SAM3_MASK_LOGITS_THRESHOLD = 0.0 + +SHORT_DESCRIPTION = ( + "Segment a specific object with SAM3 using point and/or bounding box prompts." +) + +LONG_DESCRIPTION = """ +Run the interactive (promptable visual segmentation) head of Segment Anything 3 (SAM3) on an image. + +Unlike the SAM 3 concept segmentation block (which takes text or exemplar prompts and returns +ALL instances of a concept), this block performs SAM2-style interactive segmentation: each prompt +targets ONE object and the model returns a single mask for it. + +Two prompt inputs are supported (at least one must be provided): +- **points**: a list of labeled 2D points defining a single object. Positive points mark the + object to segment, negative points mark regions to exclude (useful to refine the mask). +- **boxes**: detections from another model. Each bounding box becomes a separate prompt and + the model segments the object inside it. Class names of the boxes are forwarded to the + predicted masks. +""" + + +def _as_sam2_points(points: List[Any]) -> List[Point]: + """Normalise points given as dicts or (x, y[, positive]) sequences into Sam2 Points.""" + result = [] + for raw_point in points: + if isinstance(raw_point, Point): + result.append(raw_point) + continue + if isinstance(raw_point, dict): + if "x" not in raw_point or "y" not in raw_point: + raise ValueError( + f"Each point prompt must define `x` and `y` coordinates - got: {raw_point}" + ) + x, y = raw_point["x"], raw_point["y"] + positive = raw_point.get("positive", True) + elif isinstance(raw_point, (list, tuple)) and len(raw_point) in {2, 3}: + x, y = raw_point[0], raw_point[1] + positive = raw_point[2] if len(raw_point) == 3 else True + else: + raise ValueError( + f"Invalid point prompt: {raw_point}. Expected dict with `x`, `y` and optional " + f"`positive` keys, or a sequence of (x, y) or (x, y, positive)." + ) + if isinstance(x, bool) or isinstance(y, bool): + raise ValueError(f"Point coordinates must be numbers - got: {raw_point}") + if not isinstance(x, (int, float)) or not isinstance(y, (int, float)): + raise ValueError(f"Point coordinates must be numbers - got: {raw_point}") + result.append(Point(x=float(x), y=float(y), positive=bool(positive))) + return result + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "SAM 3 Interactive", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": [ + "Sam", + "SAM3", + "segment anything", + "segment anything 3", + "point prompt", + "interactive segmentation", + "PVS", + ], + "ui_manifest": { + "section": "model", + "icon": "fa-solid fa-eye", + "blockPriority": 9.47, + "needsGPU": True, + "inference": True, + }, + }, + protected_namespaces=(), + ) + + type: Literal["roboflow_core/sam3_interactive@v1"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + points: Optional[Union[List[Any], Selector(kind=[LABELED_POINTS_KIND])]] = Field( + default=None, + title="Point Prompts", + description="Labeled points defining a single object to segment. " + "Each point is {'x': ..., 'y': ..., 'positive': ...} in absolute pixel coordinates - " + "positive points mark the object, negative points mark regions to exclude. " + "Plain (x, y) or (x, y, positive) sequences are also accepted.", + examples=[ + [{"x": 320, "y": 240, "positive": True}], + "$inputs.points", + ], + json_schema_extra={"always_visible": True}, + ) + boxes: Optional[ + Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) + ] = Field( # type: ignore + default=None, + description="Bounding boxes (from another model) to use as prompts - " + "the model segments the object inside each box", + examples=["$steps.object_detection_model.predictions"], + json_schema_extra={"always_visible": True}, + ) + threshold: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + default=0.0, + description="Minimum confidence threshold for predicted masks", + examples=[0.3], + ) + multimask_output: Union[Optional[bool], Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Flag to determine whether to use SAM3 internal multimask or single mask mode. " + "For ambiguous prompts (like a single point) setting to True is recommended.", + examples=[True, "$inputs.multimask_output"], + ) + + @model_validator(mode="after") + def _validate_points(self) -> "BlockManifest": + if isinstance(self.points, list): + _as_sam2_points(self.points) + return self + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images", "boxes"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restrictions = [ + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; run_locally() loads a model that needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + ] + if not CORE_MODEL_SAM3_ENABLED: + restrictions.append( + RuntimeRestriction( + severity=Severity.HARD, + note=( + "CORE_MODEL_SAM3_ENABLED=False on Roboflow Hosted " + "Serverless: the SAM3 endpoint is not registered, so " + "run_remotely() returns 404." + ), + applies_to_runtimes=[Runtime.HOSTED_SERVERLESS], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + ) + return restrictions + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + return [SAM3_INTERACTIVE_MODEL_ID] + + +class SegmentAnything3InteractiveBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + points: Optional[List[Any]], + boxes: Optional[Batch], + threshold: float, + multimask_output: bool, + ) -> BlockResult: + if SAM3_EXEC_MODE == "remote": + logger.debug( + "Running SAM3 Interactive via inference proxy (SAM3_EXEC_MODE=remote)" + ) + return self.run_via_request( + images=images, + points=points, + boxes=boxes, + threshold=threshold, + multimask_output=multimask_output, + ) + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + points=points, + boxes=boxes, + threshold=threshold, + multimask_output=multimask_output, + ) + if self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + points=points, + boxes=boxes, + threshold=threshold, + multimask_output=multimask_output, + ) + raise ValueError(f"Unknown step execution mode: {self._step_execution_mode}") + + def run_locally( + self, + images: Batch[WorkflowImageData], + points: Optional[List[Any]], + boxes: Optional[Batch], + threshold: float, + multimask_output: bool, + ) -> BlockResult: + if boxes is None: + boxes = [None] * len(images) + + self._model_manager.add_model( + model_id=SAM3_INTERACTIVE_MODEL_ID, + api_key=self._api_key, + ) + + results: List[dict] = [] + for single_image, boxes_for_image in zip(images, boxes): + groups = self._build_prompt_groups( + boxes_for_image=boxes_for_image, points=points + ) + if single_image.is_tensor_materialised(): + # SAM3Torch normalises CHW/HWC itself and expects RGB channel + # order - the tensor representation is already RGB. + model_image = single_image.tensor_image + else: + # The numpy representation is BGR; flip to the RGB that the + # legacy `load_image_rgb` preprocessing produced. + model_image = np.ascontiguousarray(single_image.numpy_image[:, :, ::-1]) + raw_masks, confidences, class_ids, class_names, detection_ids = ( + [], + [], + [], + [], + [], + ) + for group in groups: + model_inputs = _prompt_group_to_model_inputs(prompts=group.prompts) + sam3_predictions = self._model_manager.run_tensor_native_inference( + SAM3_INTERACTIVE_MODEL_ID, + action="segment", + images=[model_image], + point_coordinates=model_inputs["point_coords"], + point_labels=model_inputs["point_labels"], + boxes=model_inputs["box"], + multi_mask_output=multimask_output, + # Raw logits, binarised below at the legacy adapter's + # MASK_THRESHOLD - the family's explicit-binarisation + # convention (see segment_anything2/v1_tensor.py). + return_logits=True, + ) + prediction = sam3_predictions[0] + binary_masks = prediction.masks >= SAM3_MASK_LOGITS_THRESHOLD + scores = prediction.scores.detach().to("cpu").tolist() + for instance_mask, score in zip(binary_masks, scores): + raw_masks.append(instance_mask) + confidences.append(float(score)) + class_ids.extend(group.class_ids) + class_names.extend(group.class_names) + detection_ids.extend(group.detection_ids) + instance_detections = _segmentation_results_to_instance_detections( + raw_masks=raw_masks, + confidences=confidences, + class_ids=class_ids, + class_names=class_names, + detection_ids=detection_ids, + image=single_image, + threshold=threshold, + ) + results.append({"predictions": instance_detections}) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + points: Optional[List[Any]], + boxes: Optional[Batch], + threshold: float, + multimask_output: bool, + ) -> BlockResult: + ensure_builtin_remote_execution_allowed("SAM3 Interactive remote execution") + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + if boxes is None: + boxes = [None] * len(images) + + results: List[dict] = [] + for single_image, boxes_for_image in zip(images, boxes): + groups = self._build_prompt_groups( + boxes_for_image=boxes_for_image, points=points + ) + raw_masks, confidences, class_ids, class_names, detection_ids = ( + [], + [], + [], + [], + [], + ) + for group in groups: + result = client.sam3_visual_segment( + inference_input=single_image.base64_image, + prompts=[ + prompt.dict(exclude_none=True) for prompt in group.prompts + ], + multimask_output=multimask_output, + # Compact transfer + lossless masks (same as the SAM2 tensor + # sibling's remote path). + mask_input_format="rle", + ) + for masks_payload, confidence in _parse_segmentation_predictions( + result + ): + raw_masks.append(masks_payload) + confidences.append(confidence) + class_ids.extend(group.class_ids) + class_names.extend(group.class_names) + detection_ids.extend(group.detection_ids) + instance_detections = _segmentation_results_to_instance_detections( + raw_masks=raw_masks, + confidences=confidences, + class_ids=class_ids, + class_names=class_names, + detection_ids=detection_ids, + image=single_image, + threshold=threshold, + ) + results.append({"predictions": instance_detections}) + return results + + def run_via_request( + self, + images: Batch[WorkflowImageData], + points: Optional[List[Any]], + boxes: Optional[Batch], + threshold: float, + multimask_output: bool, + ) -> BlockResult: + ensure_builtin_remote_execution_allowed( + "SAM3 Interactive inference proxy execution" + ) + endpoint = f"{API_BASE_URL}/inferenceproxy/sam3-pvs" + + if boxes is None: + boxes = [None] * len(images) + + results: List[dict] = [] + for single_image, boxes_for_image in zip(images, boxes): + groups = self._build_prompt_groups( + boxes_for_image=boxes_for_image, points=points + ) + raw_masks, confidences, class_ids, class_names, detection_ids = ( + [], + [], + [], + [], + [], + ) + for group in groups: + payload = { + "image": {"type": "base64", "value": single_image.base64_image}, + "prompts": Sam2PromptSet(prompts=group.prompts).dict( + exclude_none=True + ), + "multimask_output": multimask_output, + # Ask for lossless RLE; if the proxy strips this field the + # polygon fallback in the converter rasterises the response. + "format": "rle", + } + try: + headers = {"Content-Type": "application/json"} + if ROBOFLOW_INTERNAL_SERVICE_NAME: + headers["X-Roboflow-Internal-Service-Name"] = ( + ROBOFLOW_INTERNAL_SERVICE_NAME + ) + if ROBOFLOW_INTERNAL_SERVICE_SECRET: + headers["X-Roboflow-Internal-Service-Secret"] = ( + ROBOFLOW_INTERNAL_SERVICE_SECRET + ) + headers = build_roboflow_api_headers(explicit_headers=headers) + response = requests.post( + wrap_url(f"{endpoint}?api_key={self._api_key}"), + json=payload, + headers=headers, + timeout=60, + ) + response.raise_for_status() + resp_json = response.json() + except Exception as e: + raise Exception(f"SAM3 interactive request failed: {e}") + + for masks_payload, confidence in _parse_segmentation_predictions( + resp_json + ): + raw_masks.append(masks_payload) + confidences.append(confidence) + class_ids.extend(group.class_ids) + class_names.extend(group.class_names) + detection_ids.extend(group.detection_ids) + instance_detections = _segmentation_results_to_instance_detections( + raw_masks=raw_masks, + confidences=confidences, + class_ids=class_ids, + class_names=class_names, + detection_ids=detection_ids, + image=single_image, + threshold=threshold, + ) + results.append({"predictions": instance_detections}) + return results + + @staticmethod + def _build_prompt_groups( + boxes_for_image, + points: Optional[List[Any]], + ) -> List["_PromptGroup"]: + # The SAM prompt encoder batches a request's prompts into a single + # tensor, so one request cannot mix box-carrying and box-less prompts - + # box and point prompts are therefore issued as separate requests. + groups: List[_PromptGroup] = [] + prompt_detections = _prompt_detections(boxes_for_image) + if prompt_detections is not None: + prompts: List[Sam2Prompt] = [] + class_ids: List[Optional[int]] = [] + class_names: List[str] = [] + detection_ids: List[Optional[str]] = [] + bboxes_metadata = prompt_detections.bboxes_metadata or [ + {} for _ in range(len(prompt_detections.xyxy)) + ] + for index in range(int(prompt_detections.xyxy.shape[0])): + x1, y1, x2, y2 = prompt_detections.xyxy[index].tolist() + width = x2 - x1 + height = y2 - y1 + prompts.append( + Sam2Prompt( + box=Box( + x=float(x1 + width / 2), + y=float(y1 + height / 2), + width=float(width), + height=float(height), + ) + ) + ) + class_ids.append(int(prompt_detections.class_id[index])) + class_names.append(_resolve_prompt_class_name(prompt_detections, index)) + detection_ids.append(bboxes_metadata[index].get(DETECTION_ID_KEY)) + groups.append( + _PromptGroup( + prompts=prompts, + class_ids=class_ids, + class_names=class_names, + detection_ids=detection_ids, + ) + ) + if points: + groups.append( + _PromptGroup( + prompts=[Sam2Prompt(points=_as_sam2_points(points))], + class_ids=[0], + class_names=["foreground"], + detection_ids=[None], + ) + ) + if boxes_for_image is None and not points: + raise ValueError( + "SAM 3 Interactive block requires at least one prompt - " + "provide `points` and/or `boxes` input." + ) + # boxes input connected but no detections for this image -> no prompts, + # runners emit empty predictions + return groups + + +@dataclass(frozen=True) +class _PromptGroup: + prompts: List[Sam2Prompt] + class_ids: List[Optional[int]] + class_names: List[str] + detection_ids: List[Optional[str]] + + +def _prompt_detections(boxes_for_image): + """Normalize the tensor-native ``boxes`` prompt into a bbox-bearing prediction. + + ``boxes`` arrives as native ``Detections`` / ``InstanceDetections`` or the + keypoint ``(KeyPoints, Optional[Detections])`` tuple; the tuple's instances + component is required (``split_key_point_prediction`` raises when missing). + Returns ``None`` when ``boxes`` is absent or carries zero instances. + """ + if boxes_for_image is None: + return None + _key_points, detections = split_key_point_prediction(boxes_for_image) + if int(detections.xyxy.shape[0]) == 0: + return None + return detections + + +def _resolve_prompt_class_name(prompt_detections, source_index: int) -> str: + """Forward the prompt box's class name, mirroring the numpy block. Native OD + producers carry names in ``image_metadata['class_names']`` keyed by ``class_id`` + (NOT a per-box field); a per-box ``class`` override (vlm/ocr prompts) wins if + present.""" + if prompt_detections.bboxes_metadata is not None and source_index < len( + prompt_detections.bboxes_metadata + ): + override = prompt_detections.bboxes_metadata[source_index].get(CLASS_NAME_KEY) + if override is not None: + return str(override) + class_names = (prompt_detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + return class_names.get(int(prompt_detections.class_id[source_index]), "foreground") + + +def _prompt_group_to_model_inputs( + prompts: List[Sam2Prompt], +) -> Dict[str, Optional[np.ndarray]]: + """Convert a prompt group into the array inputs that + ``SAM3Torch.segment_with_visual_prompts`` expects - the exact + ``to_sam2_inputs`` + point-padding + ``np.array`` chain the legacy adapter's + ``segment_image`` applies, so prompts hit the model identically on both + paths (padding is an identity for the groups this block builds today, but + keeps ragged point prompts safe).""" + args = Sam2PromptSet(prompts=prompts).to_sam2_inputs() + if args["point_coords"] is not None: + max_len = max(max(len(prompt) for prompt in args["point_coords"]), 1) + for prompt in args["point_coords"]: + for _ in range(max_len - len(prompt)): + prompt.append([0, 0]) + for label in args["point_labels"]: + for _ in range(max_len - len(label)): + label.append(-1) + elif args["point_labels"] is not None: + raise ValueError( + "Can't have point labels without corresponding point coordinates" + ) + return { + "point_coords": ( + np.array(args["point_coords"]) if args["point_coords"] is not None else None + ), + "point_labels": ( + np.array(args["point_labels"]) if args["point_labels"] is not None else None + ), + "box": np.array(args["box"]) if args["box"] is not None else None, + } + + +def _parse_segmentation_predictions( + response: Union[dict, list], +) -> List[Tuple[Any, float]]: + """Extract ``(masks_payload, confidence)`` pairs from a raw SAM3 PVS response. + + ``masks_payload`` is either a COCO RLE dict (``{"size": [h, w], "counts": ...}``, + when the server honored ``format="rle"``) or a list of polygons (default + polygon-format response) - the converter handles both. + """ + if isinstance(response, list): + raw_predictions = response + else: + raw_predictions = response.get("predictions", []) + return [ + ( + raw_prediction.get("masks", []), + float(raw_prediction.get("confidence", 0.0)), + ) + for raw_prediction in raw_predictions + ] + + +def _segmentation_results_to_instance_detections( + raw_masks: List[Any], + confidences: List[float], + class_ids: List[Optional[int]], + class_names: List[str], + detection_ids: List[Optional[str]], + image: WorkflowImageData, + threshold: float, +) -> InstanceDetections: + """Build a native ``InstanceDetections`` from per-prompt SAM3 PVS results. + + Mirrors ``convert_sam2_segmentation_response_to_inference_instances_seg_response`` + semantics (confidence filter, class/detection-id forwarding by prompt order) + with the family's raw-mask instance model: an RLE mask yields ONE instance per + prompt; a polygon-format payload is rasterised per polygon (numpy-parity + expansion for responses where RLE was unavailable). + """ + height, width = image._read_shape_without_materialization() + kept_masks: List[torch.Tensor] = [] + kept_confidences: List[float] = [] + kept_class_ids: List[int] = [] + kept_class_names: List[str] = [] + kept_detection_ids: List[Optional[str]] = [] + + def _keep( + mask: torch.Tensor, + confidence: float, + class_id: Optional[int], + class_name: str, + detection_id: Optional[str], + ) -> None: + kept_masks.append(mask) + kept_confidences.append(confidence) + kept_class_ids.append(int(class_id) if class_id is not None else 0) + kept_class_names.append(class_name) + kept_detection_ids.append(detection_id) + + for masks_payload, confidence, class_id, class_name, detection_id in zip( + raw_masks, confidences, class_ids, class_names, detection_ids + ): + if confidence < threshold: + # skipping masks below threshold - mirrors the numpy converter + continue + if isinstance(masks_payload, torch.Tensor): + # Native-bridge path: an already-binarised (H, W) mask straight from + # SAM3Torch - one instance per prompt; single DtoH copy here (the + # RLE carrier built below is CPU-side anyway). + _keep( + masks_payload.detach().to("cpu").to(torch.bool), + confidence, + class_id, + class_name, + detection_id, + ) + continue + if isinstance(masks_payload, dict): + dense = _coco_rle_to_torch_mask( + masks_payload, image_height=height, image_width=width + ) + _keep(dense, confidence, class_id, class_name, detection_id) + continue + # Polygon-format payload: one instance PER polygon (>= 3 points), exactly + # like the numpy converter's expansion. + for polygon in masks_payload: + if len(polygon) < 3: + # skipping empty masks + continue + dense = _polygon_to_torch_mask( + polygon, image_height=height, image_width=width + ) + _keep(dense, confidence, class_id, class_name, detection_id) + + n = len(kept_masks) + if n == 0: + xyxy = torch.zeros((0, 4), dtype=torch.float32) + class_id_tensor = torch.zeros((0,), dtype=torch.int64) + confidence_tensor = torch.zeros((0,), dtype=torch.float32) + else: + xyxy = torch.stack([_mask_to_xyxy(mask) for mask in kept_masks]) + class_id_tensor = torch.tensor(kept_class_ids, dtype=torch.int64) + confidence_tensor = torch.tensor(kept_confidences, dtype=torch.float32) + + rle_dicts = [torch_mask_to_coco_rle(mask) for mask in kept_masks] + mask_carrier = InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), masks=rle_dicts + ) + + detections = InstanceDetections( + xyxy=xyxy.to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + class_id=class_id_tensor.to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + confidence=confidence_tensor.to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + mask=mask_carrier, + ) + class_names_map: Dict[int, str] = {} + for instance_class_id, instance_class_name in zip(kept_class_ids, kept_class_names): + class_names_map[instance_class_id] = instance_class_name + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=class_names_map, + prediction_type=PREDICTION_TYPE, + inference_id=str(uuid.uuid4()), + ) + bboxes_metadata = [] + for instance_class_name, instance_detection_id in zip( + kept_class_names, kept_detection_ids + ): + entry = { + DETECTION_ID_KEY: ( + instance_detection_id + if instance_detection_id is not None + else str(uuid.uuid4()) + ), + # Per-box class override keeps numpy's per-row name fidelity even if + # two prompts share a class_id with different labels. + CLASS_NAME_KEY: instance_class_name, + } + bboxes_metadata.append(entry) + detections.bboxes_metadata = bboxes_metadata + return detections + + +def _coco_rle_to_torch_mask( + rle_payload: dict, image_height: int, image_width: int +) -> torch.Tensor: + counts = rle_payload["counts"] + if isinstance(counts, str): + counts = counts.encode("utf-8") + size = rle_payload.get("size") or [image_height, image_width] + dense = coco_rle_masks_to_numpy_mask( + InstancesRLEMasks(image_size=(int(size[0]), int(size[1])), masks=[counts]) + )[0] + return torch.from_numpy(dense).to(torch.bool) + + +def _polygon_to_torch_mask( + polygon: List[List[float]], image_height: int, image_width: int +) -> torch.Tensor: + import cv2 + + canvas = np.zeros((image_height, image_width), dtype=np.uint8) + contour = np.round(np.array(polygon, dtype=np.float32)).astype(np.int32) + cv2.fillPoly(canvas, [contour], 1) + return torch.from_numpy(canvas.astype(bool)) + + +def _mask_to_xyxy(mask: torch.Tensor) -> torch.Tensor: + """Tight xyxy bbox of a 2-D boolean mask (inclusive min/max, NO +1) - the + convention shared by the SAM tensor family, keeping boxes numpy-faithful.""" + nonzero = torch.nonzero(mask, as_tuple=False) + if nonzero.numel() == 0: + return torch.zeros((4,), dtype=torch.float32, device=mask.device) + ys, xs = nonzero[:, 0], nonzero[:, 1] + return torch.stack([xs.min(), ys.min(), xs.max(), ys.max()]).to(torch.float32) diff --git a/inference/core/workflows/core_steps/models/foundation/segment_anything3_video/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/segment_anything3_video/v1_tensor.py new file mode 100644 index 0000000000..f683444dd4 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/segment_anything3_video/v1_tensor.py @@ -0,0 +1,494 @@ +"""Tensor-native sibling of ``roboflow_core/sam3_video@v1``, loaded when +``ENABLE_TENSOR_DATA_REPRESENTATION`` is on. + +SAM3 Video Tracker is a STATEFUL, LOCAL-only streaming open-vocabulary +concept tracker (see the numpy source for the full behavioural notes). +The session bookkeeping, prompt-signature reseeding and prompt/track state +machine are verbatim copies of the numpy source; only the prediction +representation is rewritten: + +- INPUT frame: when the workflow image is tensor-materialised the block + forwards ``WorkflowImageData.tensor_image`` (CHW RGB) directly โ€” + ``SAM3Video``'s ``_ensure_numpy_image`` permutes CHW->HWC WITHOUT + reinterpreting channel order, so the HF processor receives HWC RGB (the + numpy sibling fed HWC BGR ``numpy_image``; the tensor path is the correct + colour order โ€” the same cross-repo contract as the SAM2 video sibling). + Non-materialised inputs are flipped BGR->RGB on host. +- OUTPUT: native ``inference_models.InstanceDetections`` instead of + ``sv.Detections``. Masks are carried as compact ``InstancesRLEMasks`` by + default; the carrier is an execution-level choice driven by the + ``WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION`` env variable ("rle"/"dense", + same convention as the SAM2 video sibling; GCP_SERVERLESS forces "rle") โ€” + NOT a manifest field, so the manifest stays identical to the numpy + sibling. Per-box + ``bboxes_metadata`` carries ``detection_id`` / ``class`` / ``tracker_id``; + ``image_metadata`` is built via ``build_native_image_metadata`` with the + full prompt-position ``class_id -> name`` map (``class_id`` is the + prompt's index in ``class_names``, exactly like the numpy sibling), which + replaces ``attach_prediction_type_info_to_sv_detections_batch`` + + ``attach_parents_coordinates_to_batch_of_sv_detections``. +- Objects whose id maps to no registered prompt keep the numpy fallback: + ``class_id=0`` with per-box class name ``"foreground"`` (the per-box name + wins at serialization, so output parity holds even on collision with + ``class_names[0]``). +""" + +import uuid +from dataclasses import dataclass, field +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import ( + GCP_SERVERLESS, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION, +) +from inference.core.managers.base import ModelManager +from inference.core.roboflow_api import get_extra_weights_provider_headers +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.core_steps.models.foundation._streaming_video_common import ( + normalise_class_names, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + DETECTION_ID_KEY, + TRACKER_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import torch_mask_to_coco_rle + +PREDICTION_TYPE = "instance-segmentation" + + +def _resolve_mask_representation() -> str: + """Execution-level selection of the instance-mask carrier ("rle"/"dense"). + + Deliberately NOT a manifest field: the numpy sibling has no such knob and + manifests must stay identical across the flag swap. Driven by the + ``WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION`` env variable ("rle" default); + ``GCP_SERVERLESS`` forces the compact RLE carrier regardless. + """ + if GCP_SERVERLESS: + return "rle" + return WORKFLOWS_SAM_VIDEO_MASK_REPRESENTATION + + +SHORT_DESCRIPTION = ( + "Segment and track objects across video frames from text prompts with " + "SAM3's streaming concept tracker." +) + +LONG_DESCRIPTION = """ +Run Segment Anything 3 on a live video stream frame by frame, keeping +per-video temporal memory so object identities are preserved across +frames. + +Provide the concepts to track as text in `class_names` (e.g. +`["person", "forklift"]`) โ€” no upstream detector is needed. SAM3 runs +fused detection and tracking on every frame, so objects matching a +concept that enter the scene mid-stream are picked up automatically and +assigned fresh `tracker_id`s. Each emitted mask carries the prompt it +matched as its class name and the model's detection score as its +confidence. + +The block multiplexes a single SAM3 streaming model across many video +streams by keying state on `video_metadata.video_identifier`; a session +is re-seeded only when the source stream restarts or `class_names` +changes. For detector-driven (box-prompted) video tracking, use the +SAM2 Video Tracker block instead. + +Intended for use with `InferencePipeline`, which delivers one frame at +a time and tags each frame with video metadata. +""" + + +@dataclass +class _ConceptSessionBookkeeping: + """Per-video session state for the concept tracker. + + ``prompt_signature`` records the exact concept set the live session + was seeded with, so a change in ``class_names`` (e.g. driven by a + workflow parameter) re-seeds instead of silently tracking stale + concepts. + """ + + state_dict: Optional[dict] = None + last_frame_number: int = -1 + prompt_signature: Tuple[str, ...] = field(default_factory=tuple) + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "SAM3 Video Tracker", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": [ + "SAM3", + "segment anything 3", + "video", + "tracking", + "open vocabulary", + "META", + ], + "ui_manifest": { + "section": "video", + "icon": "fa-brands fa-meta", + "blockPriority": 9.39, + "needsGPU": True, + "inference": True, + "trackers": True, + }, + }, + protected_namespaces=(), + ) + + type: Literal["roboflow_core/sam3_video@v1"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + class_names: Union[ + List[str], str, Selector(kind=[LIST_OF_VALUES_KIND, STRING_KIND]) + ] = Field( + description=( + "Concepts to segment and track, as a list of phrases (or a " + "single comma-separated string). Each emitted mask carries " + "the concept it matched as its class name." + ), + examples=[["person", "forklift"]], + json_schema_extra={"always_visible": True}, + ) + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = Field( + default="sam3video", + description="Streaming SAM3 model id resolved by `inference_models`.", + examples=["sam3video"], + ) + threshold: Union[ + Selector(kind=[FLOAT_KIND]), + float, + ] = Field( + default=0.5, + description=( + "Minimum detection score for emitted masks. Scores come " + "from SAM3's per-object concept detection head." + ), + examples=[0.5], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + RuntimeRestriction( + severity=Severity.HARD, + note="Requires a GPU; the streaming SAM3 video model needs CUDA.", + applies_to_runtimes=[Runtime.SELF_HOSTED_CPU], + applies_to_step_execution_modes=[StepExecutionMode.LOCAL], + ), + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + return ["sam3video"] + + +class SegmentAnything3VideoBlockV1(WorkflowBlock): + """Stateful SAM3 streaming concept tracking block (tensor-native output).""" + + _REMOTE_EXECUTION_NOT_SUPPORTED_MESSAGE = ( + "SAM3 Video Tracker only supports LOCAL workflow step " + "execution. Remote execution would ship each frame to a " + "separate process and break the per-video SAM3 session " + "that holds the temporal memory. Set " + "WORKFLOWS_STEP_EXECUTION_MODE=local (or run on a " + "dedicated deployment) to use this block." + ) + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + self._model = None # lazily loaded + self._current_model_id: Optional[str] = None + self._sessions: Dict[str, _ConceptSessionBookkeeping] = {} + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def _get_model(self, model_id: str): + if self._model is None or self._current_model_id != model_id: + from inference_models import AutoModel + + extra_weights_provider_headers = get_extra_weights_provider_headers() + self._model = AutoModel.from_pretrained( + model_id_or_path=model_id, + api_key=self._api_key, + weights_provider_extra_headers=extra_weights_provider_headers, + ) + self._current_model_id = model_id + # Switching model invalidates every session we held. + self._sessions.clear() + return self._model + + def run( + self, + images: Batch[WorkflowImageData], + class_names: Union[List[str], str], + model_id: str, + threshold: float, + ) -> BlockResult: + if self._step_execution_mode is not StepExecutionMode.LOCAL: + raise NotImplementedError(self._REMOTE_EXECUTION_NOT_SUPPORTED_MESSAGE) + model = self._get_model(model_id=model_id) + class_list = normalise_class_names(class_names) + prompt_signature = tuple(class_list) + + batch_predictions: List[InstanceDetections] = [] + for single_image in images: + metadata = single_image.video_metadata + video_id = metadata.video_identifier + frame_number = metadata.frame_number or 0 + + session = self._sessions.setdefault(video_id, _ConceptSessionBookkeeping()) + stream_restarted = ( + session.last_frame_number >= 0 + and frame_number < session.last_frame_number + ) + if stream_restarted or session.prompt_signature != prompt_signature: + session.state_dict = None + + # Same cross-repo contract as the SAM2 video sibling: SAM3Video's + # _ensure_numpy_image permutes a CHW tensor to HWC WITHOUT swapping + # channels and the HF processor expects RGB, so the materialised + # CHW RGB tensor is forwarded as-is, and a host-only frame is + # flipped BGR->RGB before the call. + if single_image.is_tensor_materialised(): + frame = single_image.tensor_image + if frame.dim() != 3: + raise ValueError( + "SAM3 video tracker expects a CHW (3-D) RGB frame tensor; got " + f"a tensor with {frame.dim()} dim(s). The model's " + "_ensure_numpy_image permutes CHW->HWC and assumes this layout." + ) + else: + frame = np.ascontiguousarray(single_image.numpy_image[:, :, ::-1]) + + if not class_list: + predictions = _empty_instance_detections( + image=single_image, class_names=class_list + ) + else: + if session.state_dict is None: + result = model.prompt( + image=frame, + text=class_list, + clear_old_prompts=True, + ) + session.prompt_signature = prompt_signature + else: + result = model.track(image=frame, state_dict=session.state_dict) + session.state_dict = result.state_dict + predictions = _concept_frame_to_instance_detections( + masks=result.masks, + object_ids=result.object_ids, + scores=result.scores, + boxes=result.boxes, + prompt_to_object_ids=result.prompt_to_object_ids, + class_names=class_list, + image=single_image, + threshold=threshold, + ) + + session.last_frame_number = frame_number + batch_predictions.append(predictions) + + return [{"predictions": predictions} for predictions in batch_predictions] + + +def _concept_frame_to_instance_detections( + masks: np.ndarray, + object_ids: np.ndarray, + scores: np.ndarray, + boxes: np.ndarray, + prompt_to_object_ids: Dict[str, List[int]], + class_names: List[str], + image: WorkflowImageData, + threshold: float, +) -> InstanceDetections: + """Tensor-native equivalent of ``concept_frame_to_sv_detections``: assemble one + ``InstanceDetections`` from one SAM3 concept-tracker frame. + + Class labels and confidences are rebuilt every frame from + ``prompt_to_object_ids`` / ``scores`` (mid-stream objects get labelled + correctly); ``class_id`` is the prompt's position in ``class_names``. Masks + with no positive pixels, or whose score < ``threshold``, are dropped. Kept + masks are carried as compact ``InstancesRLEMasks``. + """ + if masks.shape[0] == 0: + return _empty_instance_detections(image=image, class_names=class_names) + + object_id_to_prompt = { + int(obj_id): prompt + for prompt, obj_ids in prompt_to_object_ids.items() + for obj_id in obj_ids + } + + xyxy: List[List[float]] = [] + confidences: List[float] = [] + class_ids: List[int] = [] + bboxes_metadata: List[dict] = [] + kept_masks: List[np.ndarray] = [] + + for mask, obj_id, score, box in zip( + masks, object_ids.tolist(), scores.tolist(), boxes.tolist() + ): + if score < threshold: + continue + if not mask.any(): + continue + prompt = object_id_to_prompt.get(int(obj_id), "foreground") + try: + class_id = class_names.index(prompt) + except ValueError: + class_id = 0 + xyxy.append([float(v) for v in box[:4]]) + confidences.append(float(score)) + class_ids.append(class_id) + bboxes_metadata.append( + { + DETECTION_ID_KEY: str(uuid.uuid4()), + CLASS_NAME_KEY: prompt, + TRACKER_ID_KEY: int(obj_id), + } + ) + kept_masks.append(mask.astype(bool)) + + if not kept_masks: + return _empty_instance_detections(image=image, class_names=class_names) + + height, width = image._read_shape_without_materialization() + if _resolve_mask_representation() == "rle": + rle_dicts = [torch_mask_to_coco_rle(torch.from_numpy(m)) for m in kept_masks] + mask = InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), masks=rle_dicts + ) + else: + mask = torch.from_numpy(np.stack(kept_masks, axis=0)).to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + detections = InstanceDetections( + xyxy=torch.tensor(xyxy, dtype=torch.float32).to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + class_id=torch.tensor(class_ids, dtype=torch.int64).to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.tensor(confidences, dtype=torch.float32).to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + mask=mask, + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=_prompt_class_names_map(class_names), + prediction_type=PREDICTION_TYPE, + ) + detections.bboxes_metadata = bboxes_metadata + return detections + + +def _empty_instance_detections( + image: WorkflowImageData, + class_names: List[str], +) -> InstanceDetections: + height, width = image._read_shape_without_materialization() + if _resolve_mask_representation() == "rle": + mask = InstancesRLEMasks(image_size=(height, width), masks=[]) + else: + mask = torch.zeros((0, height, width), dtype=torch.bool).to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + detections = InstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32).to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + class_id=torch.zeros((0,), dtype=torch.int64).to(WORKFLOWS_IMAGE_TENSOR_DEVICE), + confidence=torch.zeros((0,), dtype=torch.float32).to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + mask=mask, + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=_prompt_class_names_map(class_names), + prediction_type=PREDICTION_TYPE, + ) + detections.bboxes_metadata = None + return detections + + +def _prompt_class_names_map(class_names: List[str]) -> Dict[int, str]: + """``class_id -> name`` over the FULL prompt list (ids are prompt positions, + stable for a given prompt set). Unmatched-prompt objects fall back to + ``class_id=0`` with a per-box ``"foreground"`` label, which wins over this map + at serialization (C1 convention).""" + return {class_id: class_name for class_id, class_name in enumerate(class_names)} diff --git a/inference/core/workflows/core_steps/models/foundation/smolvlm/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/smolvlm/v1_tensor.py new file mode 100644 index 0000000000..d6b3f7ca79 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/smolvlm/v1_tensor.py @@ -0,0 +1,148 @@ +"""Tensor-native sibling of `roboflow_core/smolvlm2@v1`. + +SCRATCH โ€” first pass for review. SmolVLM2's *output* is text wrapped in a dict +(`parsed_output`), NOT a prediction kind, so this block does not PRODUCE +tensor-native predictions and takes no tensor-native prediction INPUT โ€” the only +change is the local inference route. + +Manifest, class name, type literal and `describe_outputs` are identical to v1 and +imported verbatim. Local inference now routes through the inference_models adapter +(`run_tensor_native_inference`, CHW uint8 RGB tensor straight off +`image.tensor_image`) instead of `LMMInferenceRequest`; `run_remotely` is byte-for +-byte identical to v1 (the model output is text either way). +""" + +from typing import List, Optional, Type + +from inference.core.env import ( + HOSTED_CORE_MODEL_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_REMOTE_API_TARGET, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode + +# Unchanged from v1 โ€” verbatim manifest, class name and type literal. +from inference.core.workflows.core_steps.models.foundation.smolvlm.v1 import ( + BlockManifest, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_sdk import InferenceHTTPClient + + +class SmolVLM2BlockV1(WorkflowBlock): + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + ) -> BlockResult: + if self._step_execution_mode == StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_version=model_version, + prompt=prompt, + ) + elif self._step_execution_mode == StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_version=model_version, + prompt=prompt, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CORE_MODEL_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + + prompt = prompt or "Describe what's in this image." + predictions = [] + for image in images: + result = client.infer_lmm( + inference_input=image.base64_image, + model_id=model_version, + prompt=prompt, + model_id_in_path=True, + ) + response_text = result.get("response", result) + predictions.append({"parsed_output": response_text}) + + return predictions + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_version: str, + prompt: Optional[str], + ) -> BlockResult: + # Use the provided prompt or default to a generic image description request. + prompt = prompt or "Describe what's in this image." + + # Register SmolVLM2 with the model manager. + self._model_manager.add_model(model_id=model_version, api_key=self._api_key) + + predictions = [] + for image in images: + # Local exec goes through the inference_models adapter (CHW uint8 RGB + # tensor straight off image.tensor_image). The adapter's + # run_tensor_native_inference returns one generated string per image. + if image.is_tensor_materialised(): + model_image, image_color_format = image.tensor_image, "rgb" + else: + model_image, image_color_format = image.numpy_image, "bgr" + result = self._model_manager.run_tensor_native_inference( + model_version, + images=[model_image], + input_color_format=image_color_format, + prompt=prompt, + ) + response_text = result[0] + predictions.append( + { + "parsed_output": response_text, + } + ) + return predictions diff --git a/inference/core/workflows/core_steps/models/foundation/stability_ai/inpainting/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/stability_ai/inpainting/v1_tensor.py new file mode 100644 index 0000000000..fa6dba8adb --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/stability_ai/inpainting/v1_tensor.py @@ -0,0 +1,245 @@ +""" +Credits to: https://github.com/Fafruch for origin idea +""" + +from enum import Enum +from typing import List, Literal, Optional, Type, Union + +import cv2 +import numpy as np +import requests +import supervision as sv +from pydantic import ConfigDict, Field +from supervision import Color + +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + INTEGER_KIND, + SECRET_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + AirGappedAvailability, + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections + +LONG_DESCRIPTION = """ +The block wraps +[Stability AI inpainting API](https://platform.stability.ai/docs/legacy/grpc-api/features/inpainting#Python) and +let users use instance segmentation results to change the content of images in a creative way. +""" + +SHORT_DESCRIPTION = "Use segmentation masks to inpaint objects within an image." + +API_HOST = "https://api.stability.ai" +ENDPOINT = "/v2beta/stable-image/edit/inpaint" + + +class StabilityAIPresets(Enum): + THREE_D_MODEL = "3d-model" + ANALOG_FILM = "analog-film" + ANIME = "anime" + CINEMATIC = "cinematic" + COMIC_BOOK = "comic-book" + DIGITAL_ART = "digital-art" + ENHANCE = "enhance" + FANTASY_ART = "fantasy-art" + ISOMETRIC = "isometric" + LINE_ART = "line-art" + LOW_POLY = "low-poly" + MODELING_COMPOUND = "modeling-compound" + NEON_PUNK = "neon-punk" + ORIGAMI = "origami" + PHOTOGRAPHIC = "photographic" + PIXEL_ART = "pixel-art" + TILE_TEXTURE = "tile-texture" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Stability AI Inpainting", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": [ + "Stability AI", + "stability.ai", + "inpainting", + "image generation", + ], + "ui_manifest": { + "section": "model", + "icon": "far fa-palette", + "blockPriority": 14, + }, + } + ) + type: Literal["roboflow_core/stability_ai_inpainting@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="The image to inpaint.", + examples=["$inputs.image", "$steps.cropping.crops"], + ) + segmentation_mask: Selector( + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND] + ) = Field( + name="Segmentation Mask", + description="Model predictions from segmentation model.", + examples=["$steps.model.predictions"], + ) + prompt: Union[ + Selector(kind=[STRING_KIND]), + str, + ] = Field( + description="Prompt to inpainting model (what you wish to see).", + examples=["my prompt", "$inputs.prompt"], + json_schema_extra={ + "multiline": True, + }, + ) + negative_prompt: Optional[ + Union[ + Selector(kind=[STRING_KIND]), + str, + ] + ] = Field( + default=None, + description="Negative prompt to inpainting model (what you do not wish to see).", + examples=["my prompt", "$inputs.prompt"], + ) + api_key: Union[Selector(kind=[STRING_KIND, SECRET_KIND]), str] = Field( + description="Your Stability AI API key.", + examples=["xxx-xxx", "$inputs.stability_ai_api_key"], + private=True, + ) + invert_segmentation_mask: Union[ + Selector(kind=[BOOLEAN_KIND]), + bool, + ] = Field( + default=False, + description="Invert segmentation mask to inpaint background instead of foreground.", + ) + preset: Optional[StabilityAIPresets] = Field( + default=None, + description="Optional preset to apply when outpainting the image (what you wish to see)." + " If not provided, the image will be outpainted without any preset." + f" Avaliable presets: {', '.join(m.value for m in StabilityAIPresets)}", + examples=[StabilityAIPresets.THREE_D_MODEL], + ) + seed: Optional[ + Union[ + Selector(kind=[INTEGER_KIND]), + int, + ] + ] = Field( + default=None, + description="A specific value that is used to guide the 'randomness' of the generation." + " If not provided, a random seed is used." + " Must be a number between 0 and 4294967294", + examples=[200], + ) + + @classmethod + def get_air_gapped_availability(cls) -> AirGappedAvailability: + return AirGappedAvailability(available=False, reason="requires_internet") + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="image", kind=[IMAGE_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.4.0,<2.0.0" + + +class StabilityAIInpaintingBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + segmentation_mask: InstanceDetections, + prompt: str, + negative_prompt: str, + api_key: str, + invert_segmentation_mask: bool, + preset: Optional[StabilityAIPresets] = None, + seed: Optional[int] = None, + ) -> BlockResult: + segmentation_mask = segmentation_mask.to_supervision() + black_image = np.zeros_like(image.numpy_image) + mask_annotator = sv.MaskAnnotator(color=Color.WHITE, opacity=1.0) + mask = mask_annotator.annotate(black_image, segmentation_mask) + mask = cv2.GaussianBlur(mask, (15, 15), 0) + encoded_image = numpy_array_to_jpeg_bytes(image=image.numpy_image) + if invert_segmentation_mask: + mask = cv2.bitwise_not(mask) + encoded_mask = numpy_array_to_jpeg_bytes(image=mask) + request_data = { + "prompt": prompt, + "output_format": "jpeg", + } + preset = ( + preset.value if preset in set(e.value for e in StabilityAIPresets) else None + ) + if preset: + request_data["preset"] = preset + seed = max(0, min(4294967294, seed)) if seed else None + if seed: + request_data["seed"] = seed + response = requests.post( + f"{API_HOST}{ENDPOINT}", + headers={"authorization": f"Bearer {api_key}", "accept": "image/*"}, + files={ + "image": encoded_image, + "mask": encoded_mask, + }, + data=request_data, + ) + if response.status_code != 200: + raise RuntimeError( + f"Request to StabilityAI API failed: {str(response.json())}" + ) + result_image = bytes_to_opencv_image(payload=response.content) + return { + "image": WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=result_image, + ), + } + + +def numpy_array_to_jpeg_bytes( + image: np.ndarray, +) -> bytes: + _, img_encoded = cv2.imencode(".jpg", image) + return np.array(img_encoded).tobytes() + + +def bytes_to_opencv_image( + payload: bytes, array_type: np.number = np.uint8 +) -> np.ndarray: + bytes_array = np.frombuffer(payload, dtype=array_type) + decoding_result = cv2.imdecode(bytes_array, cv2.IMREAD_UNCHANGED) + if decoding_result is None: + raise ValueError("Could not encode bytes to OpenCV image.") + return decoding_result diff --git a/inference/core/workflows/core_steps/models/foundation/yolo_world/v1_tensor.py b/inference/core/workflows/core_steps/models/foundation/yolo_world/v1_tensor.py new file mode 100644 index 0000000000..7f4b56e632 --- /dev/null +++ b/inference/core/workflows/core_steps/models/foundation/yolo_world/v1_tensor.py @@ -0,0 +1,147 @@ +"""Tensor-native sibling of `roboflow_core/yolo_world_model@v1`. + +DEPRECATION STUB. YOLO-World has no inference_models tensor-native path, so rather +than port it, it is deprecated in the tensor-native Workflows pipeline (same pattern +as `gaze/v1.py`). The manifest is kept identical to the numpy block so existing +workflow JSON still loads/validates; invoking the block raises +`FeatureDeprecatedError` (HTTP 410 Gone) at runtime. Flag-off keeps the working +numpy `YoloWorldModelBlockV1`. +""" + +from typing import List, Literal, Optional, Type, Union + +from pydantic import ConfigDict, Field + +from inference.core.exceptions import FeatureDeprecatedError +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + OBJECT_DETECTION_PREDICTION_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) + +LONG_DESCRIPTION = """ +**DEPRECATED.** YOLO-World is deprecated in the tensor-native Workflows pipeline. +Invoking this block raises `FeatureDeprecatedError` (HTTP 410 Gone). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "YOLO-World Model", + "version": "v1", + "short_description": "Run a zero-shot object detection model (deprecated).", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "deprecated": True, + "ui_manifest": { + "section": "model", + "icon": "fal fa-atom", + "blockPriority": 8, + "inference": True, + }, + } + ) + type: Literal["roboflow_core/yolo_world_model@v1", "YoloWorldModel", "YoloWorld"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + class_names: Union[Selector(kind=[LIST_OF_VALUES_KIND]), List[str]] = Field( + description="One or more classes that you want YOLO-World to detect. The model accepts any string as an input, though does best with short descriptions of common objects.", + examples=[["person", "car", "license plate"], "$inputs.class_names"], + ) + version: Union[ + Literal["v2-s", "v2-m", "v2-l", "v2-x", "s", "m", "l", "x"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="v2-s", + description="Variant of YoloWorld model", + examples=["v2-s", "$inputs.variant"], + ) + confidence: Union[ + Optional[FloatZeroToOne], + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.005, + description="Confidence threshold for detections", + examples=[0.005, "$inputs.confidence"], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", kind=[OBJECT_DETECTION_PREDICTION_KIND] + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_supported_model_variants(cls) -> Optional[List[str]]: + """Return list of model_id variants that can satisfy this block.""" + return [ + "yolo_world/v2-s", + "yolo_world/v2-m", + "yolo_world/v2-l", + "yolo_world/v2-x", + "yolo_world/s", + "yolo_world/m", + "yolo_world/l", + "yolo_world/x", + ] + + +class YoloWorldModelBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + class_names: List[str], + version: str, + confidence: Optional[float], + ) -> BlockResult: + raise FeatureDeprecatedError( + feature="roboflow_core/yolo_world_model@v1", + reason="YOLO-World is deprecated in the tensor-native pipeline (no inference_models path)", + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v1_tensor.py b/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v1_tensor.py new file mode 100644 index 0000000000..d90fb1900b --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v1_tensor.py @@ -0,0 +1,699 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_instance_segmentation_model@v1`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.InstanceDetections`` (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``; masks dense or RLE) under +``TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND`` / +``TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND`` instead of +``sv.Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns + ``List[InstanceDetections]`` straight from the adapter. The mask carrier (dense + ``torch.Tensor`` vs ``InstancesRLEMasks``) is adapter-decided: the v1 manifest + exposes ``enforce_dense_masks_in_inference_models`` and the adapter consumes it + to choose dense vs RLE, so both carriers are possible and both are handled + downstream by the helpers and the tensor serialiser. The block applies + ``class_filter`` natively (the adapter/model does NOT read it on this path) and + attaches the producer contract (``image_metadata[class_names]`` + per-box + ``detection_id``) the tensor serialiser requires, via + ``attach_native_detection_metadata`` (it preserves masks + per-box metadata). +- REMOTE: standard inference instance-seg prediction dicts are rebuilt into a + native ``InstanceDetections`` via the in-file + ``_native_instance_detections_from_inference_predictions`` converter (never + ``sv.Detections``). The block requests ``response_mask_format="rle"`` and carries + RLE masks when the server honours it; it degrades gracefully to a dense mask + rasterised from polygon ``points`` when the server ignores that parameter + (older/alternative servers), matching the numpy flag-OFF polygon path. + +This block creates ONLY this file; it reuses the already-registered tensor +serialiser for ``instance_segmentation_prediction`` / +``rle_instance_segmentation_prediction`` (``serialise_sv_detections`` already +handles ``InstanceDetections``, dense or RLE). The numpy sibling lives in +``.../instance_segmentation/v1.py``; this manifest is identical except the output +kinds. +""" + +import uuid +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field, PositiveInt + +from inference.core.env import ( + HOSTED_INSTANCE_SEGMENTATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + build_native_image_metadata, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_ID_KEY, + CLASS_NAME_KEY, + CONFIDENCE_KEY, + DETECTION_ID_KEY, + HEIGHT_KEY, + INFERENCE_ID_KEY, + POLYGON_KEY, + WIDTH_KEY, + X_KEY, + Y_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "instance-segmentation" + +# Keys carrying the RLE mask in standard inference instance-seg prediction dicts. +# `InstanceSegmentationRLEPrediction` declares the mask under `rle`; some response +# paths use `rle_mask` (mirrors the numpy converter's `d.get("rle_mask") or d.get("rle")`). +RLE_MASK_KEYS = ("rle_mask", "rle") + +LONG_DESCRIPTION = """ +Run inference on an instance segmentation model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Instance Segmentation Model", + "version": "v1", + "short_description": "Predict the shape, size, and location of objects.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo", "rfdetr", "rf-detr"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 1, + "inference": True, + "popular": True, + }, + }, + protected_namespaces=(), + ) + type: Literal[ + "roboflow_core/roboflow_instance_segmentation_model@v1", + "RoboflowInstanceSegmentationModel", + "InstanceSegmentationModel", + ] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + mask_decode_mode: Union[ + Literal["accurate", "tradeoff", "fast"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="accurate", + description="Parameter of mask decoding in prediction post-processing.", + examples=["accurate", "$inputs.mask_decode_mode"], + ) + tradeoff_factor: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.0, + description="Post-processing parameter to dictate tradeoff between fast and accurate.", + examples=[0.3, "$inputs.tradeoff_factor"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + enforce_dense_masks_in_inference_models: Union[ + bool, Selector(kind=[BOOLEAN_KIND]) + ] = Field( + default=True, + description="Boolean flag to enforce dense masks when inference models backend is in use " + "(irrelevant in other cases). Dense masks are faster to process, but require more memory. " + "Users can't tweak this flag when running on Roboflow serverless platform.", + examples=[True, "$inputs.enforce_dense_masks_in_inference_models"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["instance-segmentation"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=INFERENCE_ID_KEY, kind=[STRING_KIND]), + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowInstanceSegmentationModelBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + enforce_dense_masks_in_inference_models: bool, + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + enforce_dense_masks_in_inference_models=enforce_dense_masks_in_inference_models, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + enforce_dense_masks_in_inference_models: bool, + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + detections_batch: List[InstanceDetections] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + enforce_dense_masks_in_inference_models=( + enforce_dense_masks_in_inference_models + or WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS + ), + ) + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, detections in zip(images, detections_batch): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). + detections = _filter_classes_native(detections, class_filter, class_names) + # Reuse the adapter-provided inference id when present (numpy parity: + # numpy v1 surfaces ``p.get(INFERENCE_ID_KEY)`` off the model dump); + # the tensor-native adapter normally attaches none, so fall back to a + # freshly minted uuid that is then shared with ``image_metadata``. + inference_id = getattr(detections, "inference_id", None) or str( + uuid.uuid4() + ) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_INSTANCE_SEGMENTATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + response_mask_format="rle", + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + model_id: str, + ) -> BlockResult: + # Fallback class_id -> name map from the model, used to name boxes whose + # remote prediction dict lacks a `class` key (otherwise the tensor + # serialiser hard-raises "class_id missing from mapping"). + model_class_names = _class_names_map( + self._model_manager.get_class_names(model_id) + ) + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) or [] + if class_filter: + accepted = set(class_filter) + detection_dicts = [ + d for d in detection_dicts if d.get(CLASS_NAME_KEY) in accepted + ] + detections = _native_instance_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + model_class_names=model_class_names, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + detections: InstanceDetections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> InstanceDetections: + if not class_filter: + return detections + accepted = set(class_filter) + accepted_ids = sorted( + class_id for class_id, name in class_names.items() if name in accepted + ) + if not accepted_ids: + return take_prediction_by_mask( + detections, + torch.zeros_like(detections.class_id, dtype=torch.bool), + ) + accepted_tensor = torch.as_tensor(accepted_ids, device=detections.class_id.device) + keep = torch.isin(detections.class_id, accepted_tensor) + return take_prediction_by_mask(detections, keep) + + +def _extract_rle_mask(prediction: dict) -> Optional[dict]: + """Pull the COCO-RLE mask ({"size": [H, W], "counts": ...}) from a standard + inference instance-seg prediction dict, trying both response key variants.""" + for key in RLE_MASK_KEYS: + mask = prediction.get(key) + if mask is not None: + return mask + return None + + +def _extract_polygon_points(prediction: dict) -> Optional[List[dict]]: + """Pull the polygon ``points`` ([{"x": .., "y": ..}, ...]) from a standard + inference instance-seg prediction dict that carries a polygon mask instead of + RLE. Returns ``None`` when the key is absent or the polygon is degenerate + (< 3 points) - mirroring supervision's ``Detections.from_inference`` / the numpy + ``filter_out_invalid_polygons``, which drop such instances entirely.""" + points = prediction.get(POLYGON_KEY) + if points is not None and len(points) >= 3: + return points + return None + + +def _polygon_points_to_dense_mask( + points: List[dict], height: int, width: int +) -> np.ndarray: + """Rasterise polygon ``points`` into a dense boolean ``(H, W)`` mask, + byte-identically to supervision's ``polygon_to_mask`` as used by + ``Detections.from_inference`` (the numpy REMOTE path): vertices are integer + *truncated* (``dtype=int``, NOT rounded), filled with ``cv2.fillPoly`` onto a + zero ``uint8`` canvas, then cast to ``bool``. Keeping the exact rasterisation + means the serialiser's mask->polygon re-contouring (``mask_to_polygon``, shared + verbatim with the numpy serialiser) reproduces the numpy ``points`` output.""" + polygon = np.array([[point[X_KEY], point[Y_KEY]] for point in points], dtype=int) + return sv.polygon_to_mask(polygon, resolution_wh=(width, height)).astype(bool) + + +def _native_instance_detections_from_inference_predictions( + image: WorkflowImageData, + predictions: List[dict], + prediction_type: str, + inference_id: Optional[str] = None, + device: Optional[torch.device] = None, + model_class_names: Optional[Dict[int, str]] = None, +) -> InstanceDetections: + """REMOTE-path converter: build a native ``InstanceDetections`` from standard + inference instance-segmentation prediction dicts. + + The block requests ``response_mask_format="rle"`` so a current inference server + returns an RLE-encoded COCO mask (``{"size": [H, W], "counts": ""}``) + under ``rle`` / ``rle_mask``; those are carried as ``InstancesRLEMasks`` (counts + normalised from a utf-8 string to bytes, what ``pycocotools`` / the serialiser's + RLE decode path expect). To stay compatible with servers that ignore that + request parameter and return polygon ``points`` instead (older/alternative + inference servers), the converter degrades gracefully: when the response is not + fully RLE-backed it rebuilds the masks from the polygon ``points`` into a dense + ``torch.bool`` ``(N, H, W)`` carrier - the same masks the numpy flag-OFF REMOTE + path builds via ``sv.Detections.from_inference`` - so the serialised ``points`` + output matches numpy's (the serialiser re-contours the dense mask with the same + ``mask_to_polygon``). Boxes are converted from center form to corner ``xyxy``. + The ``class_id -> name`` map (required by the tensor serialiser) and per-box + ``detection_id`` are built here too. + + Degenerate (< 3-point) polygons are dropped exactly like the numpy path + (``filter_out_invalid_polygons`` / supervision's ``from_inference``). When a + prediction omits its ``class`` key, the class name is backfilled from + ``model_class_names`` (the model's ``get_class_names`` map) so the tensor + serialiser does not hard-raise on an unmapped ``class_id``. + """ + height, width = image._read_shape_without_materialization() + # Prefer the RLE masks the block requests; fall back to dense polygon masks when + # the response is not fully RLE-backed (server ignored `response_mask_format`). + # `all([])` is True, so an empty response keeps the empty-RLE carrier unchanged. + use_rle = all( + _extract_rle_mask(prediction) is not None for prediction in predictions + ) + if use_rle: + kept_predictions = predictions + else: + # Drop degenerate polygons up front so the surviving box arrays and masks + # stay aligned (matches the numpy path's instance drop). + kept_predictions = [ + prediction + for prediction in predictions + if _extract_polygon_points(prediction) is not None + ] + xyxy: List[List[float]] = [] + class_id: List[int] = [] + confidence: List[float] = [] + rle_counts: List[bytes] = [] + dense_masks: List[np.ndarray] = [] + bboxes_metadata: List[dict] = [] + derived_class_names: Dict[int, str] = {} + for prediction in kept_predictions: + center_x = float(prediction[X_KEY]) + center_y = float(prediction[Y_KEY]) + box_width = float(prediction[WIDTH_KEY]) + box_height = float(prediction[HEIGHT_KEY]) + xyxy.append( + [ + center_x - box_width / 2, + center_y - box_height / 2, + center_x + box_width / 2, + center_y + box_height / 2, + ] + ) + prediction_class_id = int(prediction.get(CLASS_ID_KEY, 0)) + class_id.append(prediction_class_id) + confidence.append(float(prediction.get(CONFIDENCE_KEY, 1.0))) + if CLASS_NAME_KEY in prediction: + derived_class_names[prediction_class_id] = str(prediction[CLASS_NAME_KEY]) + elif model_class_names is not None and prediction_class_id in model_class_names: + derived_class_names[prediction_class_id] = model_class_names[ + prediction_class_id + ] + if use_rle: + mask = _extract_rle_mask(prediction) + raw_counts = mask["counts"] + # Normalise to bytes: pycocotools (used by the serialiser's RLE decode) + # expects byte counts; the remote rle response carries them as utf-8 strings. + rle_counts.append( + raw_counts.encode("utf-8") + if isinstance(raw_counts, str) + else raw_counts + ) + else: + dense_masks.append( + _polygon_points_to_dense_mask( + points=_extract_polygon_points(prediction), + height=height, + width=width, + ) + ) + bboxes_metadata.append( + {DETECTION_ID_KEY: str(prediction.get(DETECTION_ID_KEY) or uuid.uuid4())} + ) + number_of_detections = len(xyxy) + image_metadata = build_native_image_metadata( + image=image, + class_names=derived_class_names, + prediction_type=prediction_type, + inference_id=inference_id, + ) + if use_rle: + mask_carrier: Union[torch.Tensor, InstancesRLEMasks] = InstancesRLEMasks( + image_size=(height, width), masks=rle_counts + ) + elif dense_masks: + mask_carrier = torch.as_tensor( + np.stack(dense_masks), dtype=torch.bool, device=device + ) + else: + # No surviving polygons: an empty dense carrier the serialiser handles. + mask_carrier = torch.zeros((0, height, width), dtype=torch.bool, device=device) + return InstanceDetections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32, device=device).reshape(-1, 4), + class_id=torch.as_tensor(class_id, dtype=torch.long, device=device).reshape(-1), + confidence=torch.as_tensor( + confidence, dtype=torch.float32, device=device + ).reshape(-1), + mask=mask_carrier, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata if number_of_detections > 0 else None, + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v2_tensor.py b/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v2_tensor.py new file mode 100644 index 0000000000..924942a594 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v2_tensor.py @@ -0,0 +1,699 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_instance_segmentation_model@v2`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.InstanceDetections`` (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``; masks dense or RLE) under +``TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND`` / +``TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND`` instead of +``sv.Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns + ``List[InstanceDetections]`` straight from the adapter. The mask carrier (dense + ``torch.Tensor`` vs ``InstancesRLEMasks``) is adapter-decided: the v2 manifest + exposes ``enforce_dense_masks_in_inference_models`` and the adapter consumes it + to choose dense vs RLE, so both carriers are possible and both are handled + downstream by the helpers and the tensor serialiser. The block applies + ``class_filter`` natively (the adapter/model does NOT read it on this path) and + attaches the producer contract (``image_metadata[class_names]`` + per-box + ``detection_id``) the tensor serialiser requires, via + ``attach_native_detection_metadata`` (it preserves masks + per-box metadata). +- REMOTE: standard inference instance-seg prediction dicts are rebuilt into a + native ``InstanceDetections`` via the in-file + ``_native_instance_detections_from_inference_predictions`` converter (never + ``sv.Detections``). The block requests ``response_mask_format="rle"`` and carries + RLE masks when the server honours it; it degrades gracefully to a dense mask + rasterised from polygon ``points`` when the server ignores that parameter + (older/alternative servers), matching the numpy flag-OFF polygon path. + +This block creates ONLY this file; it reuses the already-registered tensor +serialiser for ``instance_segmentation_prediction`` / +``rle_instance_segmentation_prediction`` (``serialise_sv_detections`` already +handles ``InstanceDetections``, dense or RLE). The numpy sibling lives in +``.../instance_segmentation/v2.py``; this manifest is identical except the output +kinds. +""" + +import uuid +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field, PositiveInt + +from inference.core.env import ( + HOSTED_INSTANCE_SEGMENTATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + build_native_image_metadata, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_ID_KEY, + CLASS_NAME_KEY, + CONFIDENCE_KEY, + DETECTION_ID_KEY, + HEIGHT_KEY, + INFERENCE_ID_KEY, + POLYGON_KEY, + WIDTH_KEY, + X_KEY, + Y_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "instance-segmentation" + +# Keys carrying the RLE mask in standard inference instance-seg prediction dicts. +# `InstanceSegmentationRLEPrediction` declares the mask under `rle`; some response +# paths use `rle_mask` (mirrors the numpy converter's `d.get("rle_mask") or d.get("rle")`). +RLE_MASK_KEYS = ("rle_mask", "rle") + +LONG_DESCRIPTION = """ +Run inference on an instance segmentation model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Instance Segmentation Model", + "version": "v2", + "short_description": "Predict the shape, size, and location of objects.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo", "rfdetr", "rf-detr"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 1, + "inference": True, + "popular": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_instance_segmentation_model@v2"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + mask_decode_mode: Union[ + Literal["accurate", "tradeoff", "fast"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="accurate", + description="Parameter of mask decoding in prediction post-processing.", + examples=["accurate", "$inputs.mask_decode_mode"], + ) + tradeoff_factor: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.0, + description="Post-processing parameter to dictate tradeoff between fast and accurate.", + examples=[0.3, "$inputs.tradeoff_factor"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + enforce_dense_masks_in_inference_models: Union[ + bool, Selector(kind=[BOOLEAN_KIND]) + ] = Field( + default=True, + description="Boolean flag to enforce dense masks when inference models backend is in use " + "(irrelevant in other cases). Dense masks are faster to process, but require more memory. " + "Users can't tweak this flag when running on Roboflow serverless platform.", + examples=[True, "$inputs.enforce_dense_masks_in_inference_models"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["instance-segmentation"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowInstanceSegmentationModelBlockV2(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + enforce_dense_masks_in_inference_models: bool, + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + enforce_dense_masks_in_inference_models=enforce_dense_masks_in_inference_models, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + enforce_dense_masks_in_inference_models: bool, + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + detections_batch: List[InstanceDetections] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + enforce_dense_masks_in_inference_models=( + enforce_dense_masks_in_inference_models + or WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS + ), + ) + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, detections in zip(images, detections_batch): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). + detections = _filter_classes_native(detections, class_filter, class_names) + # Reuse the adapter-provided inference id when present (numpy parity: + # numpy v2 surfaces ``p.get(INFERENCE_ID_KEY)`` off the model dump); + # the tensor-native adapter normally attaches none, so fall back to a + # freshly minted uuid that is then shared with ``image_metadata``. + inference_id = getattr(detections, "inference_id", None) or str( + uuid.uuid4() + ) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_INSTANCE_SEGMENTATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + response_mask_format="rle", + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + model_id: str, + ) -> BlockResult: + # Fallback class_id -> name map from the model, used to name boxes whose + # remote prediction dict lacks a `class` key (otherwise the tensor + # serialiser hard-raises "class_id missing from mapping"). + model_class_names = _class_names_map( + self._model_manager.get_class_names(model_id) + ) + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) or [] + if class_filter: + accepted = set(class_filter) + detection_dicts = [ + d for d in detection_dicts if d.get(CLASS_NAME_KEY) in accepted + ] + detections = _native_instance_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + model_class_names=model_class_names, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + detections: InstanceDetections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> InstanceDetections: + if not class_filter: + return detections + accepted = set(class_filter) + accepted_ids = sorted( + class_id for class_id, name in class_names.items() if name in accepted + ) + if not accepted_ids: + return take_prediction_by_mask( + detections, + torch.zeros_like(detections.class_id, dtype=torch.bool), + ) + accepted_tensor = torch.as_tensor(accepted_ids, device=detections.class_id.device) + keep = torch.isin(detections.class_id, accepted_tensor) + return take_prediction_by_mask(detections, keep) + + +def _extract_rle_mask(prediction: dict) -> Optional[dict]: + """Pull the COCO-RLE mask ({"size": [H, W], "counts": ...}) from a standard + inference instance-seg prediction dict, trying both response key variants.""" + for key in RLE_MASK_KEYS: + mask = prediction.get(key) + if mask is not None: + return mask + return None + + +def _extract_polygon_points(prediction: dict) -> Optional[List[dict]]: + """Pull the polygon ``points`` ([{"x": .., "y": ..}, ...]) from a standard + inference instance-seg prediction dict that carries a polygon mask instead of + RLE. Returns ``None`` when the key is absent or the polygon is degenerate + (< 3 points) - mirroring supervision's ``Detections.from_inference`` / the numpy + ``filter_out_invalid_polygons``, which drop such instances entirely.""" + points = prediction.get(POLYGON_KEY) + if points is not None and len(points) >= 3: + return points + return None + + +def _polygon_points_to_dense_mask( + points: List[dict], height: int, width: int +) -> np.ndarray: + """Rasterise polygon ``points`` into a dense boolean ``(H, W)`` mask, + byte-identically to supervision's ``polygon_to_mask`` as used by + ``Detections.from_inference`` (the numpy REMOTE path): vertices are integer + *truncated* (``dtype=int``, NOT rounded), filled with ``cv2.fillPoly`` onto a + zero ``uint8`` canvas, then cast to ``bool``. Keeping the exact rasterisation + means the serialiser's mask->polygon re-contouring (``mask_to_polygon``, shared + verbatim with the numpy serialiser) reproduces the numpy ``points`` output.""" + polygon = np.array([[point[X_KEY], point[Y_KEY]] for point in points], dtype=int) + return sv.polygon_to_mask(polygon, resolution_wh=(width, height)).astype(bool) + + +def _native_instance_detections_from_inference_predictions( + image: WorkflowImageData, + predictions: List[dict], + prediction_type: str, + inference_id: Optional[str] = None, + device: Optional[torch.device] = None, + model_class_names: Optional[Dict[int, str]] = None, +) -> InstanceDetections: + """REMOTE-path converter: build a native ``InstanceDetections`` from standard + inference instance-segmentation prediction dicts. + + The block requests ``response_mask_format="rle"`` so a current inference server + returns an RLE-encoded COCO mask (``{"size": [H, W], "counts": ""}``) + under ``rle`` / ``rle_mask``; those are carried as ``InstancesRLEMasks`` (counts + normalised from a utf-8 string to bytes, what ``pycocotools`` / the serialiser's + RLE decode path expect). To stay compatible with servers that ignore that + request parameter and return polygon ``points`` instead (older/alternative + inference servers), the converter degrades gracefully: when the response is not + fully RLE-backed it rebuilds the masks from the polygon ``points`` into a dense + ``torch.bool`` ``(N, H, W)`` carrier - the same masks the numpy flag-OFF REMOTE + path builds via ``sv.Detections.from_inference`` - so the serialised ``points`` + output matches numpy's (the serialiser re-contours the dense mask with the same + ``mask_to_polygon``). Boxes are converted from center form to corner ``xyxy``. + The ``class_id -> name`` map (required by the tensor serialiser) and per-box + ``detection_id`` are built here too. + + Degenerate (< 3-point) polygons are dropped exactly like the numpy path + (``filter_out_invalid_polygons`` / supervision's ``from_inference``). When a + prediction omits its ``class`` key, the class name is backfilled from + ``model_class_names`` (the model's ``get_class_names`` map) so the tensor + serialiser does not hard-raise on an unmapped ``class_id``. + """ + height, width = image._read_shape_without_materialization() + # Prefer the RLE masks the block requests; fall back to dense polygon masks when + # the response is not fully RLE-backed (server ignored `response_mask_format`). + # `all([])` is True, so an empty response keeps the empty-RLE carrier unchanged. + use_rle = all( + _extract_rle_mask(prediction) is not None for prediction in predictions + ) + if use_rle: + kept_predictions = predictions + else: + # Drop degenerate polygons up front so the surviving box arrays and masks + # stay aligned (matches the numpy path's instance drop). + kept_predictions = [ + prediction + for prediction in predictions + if _extract_polygon_points(prediction) is not None + ] + xyxy: List[List[float]] = [] + class_id: List[int] = [] + confidence: List[float] = [] + rle_counts: List[bytes] = [] + dense_masks: List[np.ndarray] = [] + bboxes_metadata: List[dict] = [] + derived_class_names: Dict[int, str] = {} + for prediction in kept_predictions: + center_x = float(prediction[X_KEY]) + center_y = float(prediction[Y_KEY]) + box_width = float(prediction[WIDTH_KEY]) + box_height = float(prediction[HEIGHT_KEY]) + xyxy.append( + [ + center_x - box_width / 2, + center_y - box_height / 2, + center_x + box_width / 2, + center_y + box_height / 2, + ] + ) + prediction_class_id = int(prediction.get(CLASS_ID_KEY, 0)) + class_id.append(prediction_class_id) + confidence.append(float(prediction.get(CONFIDENCE_KEY, 1.0))) + if CLASS_NAME_KEY in prediction: + derived_class_names[prediction_class_id] = str(prediction[CLASS_NAME_KEY]) + elif model_class_names is not None and prediction_class_id in model_class_names: + derived_class_names[prediction_class_id] = model_class_names[ + prediction_class_id + ] + if use_rle: + mask = _extract_rle_mask(prediction) + raw_counts = mask["counts"] + # Normalise to bytes: pycocotools (used by the serialiser's RLE decode) + # expects byte counts; the remote rle response carries them as utf-8 strings. + rle_counts.append( + raw_counts.encode("utf-8") + if isinstance(raw_counts, str) + else raw_counts + ) + else: + dense_masks.append( + _polygon_points_to_dense_mask( + points=_extract_polygon_points(prediction), + height=height, + width=width, + ) + ) + bboxes_metadata.append( + {DETECTION_ID_KEY: str(prediction.get(DETECTION_ID_KEY) or uuid.uuid4())} + ) + number_of_detections = len(xyxy) + image_metadata = build_native_image_metadata( + image=image, + class_names=derived_class_names, + prediction_type=prediction_type, + inference_id=inference_id, + ) + if use_rle: + mask_carrier: Union[torch.Tensor, InstancesRLEMasks] = InstancesRLEMasks( + image_size=(height, width), masks=rle_counts + ) + elif dense_masks: + mask_carrier = torch.as_tensor( + np.stack(dense_masks), dtype=torch.bool, device=device + ) + else: + # No surviving polygons: an empty dense carrier the serialiser handles. + mask_carrier = torch.zeros((0, height, width), dtype=torch.bool, device=device) + return InstanceDetections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32, device=device).reshape(-1, 4), + class_id=torch.as_tensor(class_id, dtype=torch.long, device=device).reshape(-1), + confidence=torch.as_tensor( + confidence, dtype=torch.float32, device=device + ).reshape(-1), + mask=mask_carrier, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata if number_of_detections > 0 else None, + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v3_tensor.py b/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v3_tensor.py new file mode 100644 index 0000000000..badf1a2d80 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v3_tensor.py @@ -0,0 +1,740 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_instance_segmentation_model@v3`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.InstanceDetections`` (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``; masks dense or RLE) under +``TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND`` / +``TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND`` instead of +``sv.Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns + ``List[InstanceDetections]`` straight from the adapter. The mask carrier (dense + ``torch.Tensor`` vs ``InstancesRLEMasks``) is adapter-decided: the v3 manifest + exposes ``enforce_dense_masks_in_inference_models`` and the adapter consumes it + to choose dense vs RLE, so both carriers are possible and both are handled + downstream by the helpers and the tensor serialiser. The block applies + ``class_filter`` natively (the adapter/model does NOT read it on this path) and + attaches the producer contract (``image_metadata[class_names]`` + per-box + ``detection_id``) the tensor serialiser requires, via + ``attach_native_detection_metadata`` (it preserves masks + per-box metadata). +- REMOTE: standard inference instance-seg prediction dicts are rebuilt into a + native ``InstanceDetections`` via the in-file + ``_native_instance_detections_from_inference_predictions`` converter (never + ``sv.Detections``). The block requests ``response_mask_format="rle"`` and carries + RLE masks when the server honours it; it degrades gracefully to a dense mask + rasterised from polygon ``points`` when the server ignores that parameter + (older/alternative servers), matching the numpy flag-OFF polygon path. + +This block creates ONLY this file; it reuses the already-registered tensor +serialiser for ``instance_segmentation_prediction`` / +``rle_instance_segmentation_prediction`` (``serialise_sv_detections`` already +handles ``InstanceDetections``, dense or RLE). The numpy sibling lives in +``.../instance_segmentation/v3.py``; this manifest is identical except the output +kinds. +""" + +import uuid +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field, PositiveInt, model_validator + +from inference.core.env import ( + HOSTED_INSTANCE_SEGMENTATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + build_native_image_metadata, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_ID_KEY, + CLASS_NAME_KEY, + CONFIDENCE_KEY, + DETECTION_ID_KEY, + HEIGHT_KEY, + INFERENCE_ID_KEY, + POLYGON_KEY, + WIDTH_KEY, + X_KEY, + Y_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "instance-segmentation" + +# Keys carrying the RLE mask in standard inference instance-seg prediction dicts. +# `InstanceSegmentationRLEPrediction` declares the mask under `rle`; some response +# paths use `rle_mask` (mirrors the numpy converter's `d.get("rle_mask") or d.get("rle")`). +RLE_MASK_KEYS = ("rle_mask", "rle") + +LONG_DESCRIPTION = """ +Run inference on an instance segmentation model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Instance Segmentation Model", + "version": "v3", + "short_description": "Predict the shape, size, and location of objects.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo", "rfdetr", "rf-detr"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 1, + "inference": True, + "popular": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_instance_segmentation_model@v3"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence_mode: Union[ + Literal["best", "default", "custom"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="best", + description="How confidence thresholds are determined.", + json_schema_extra={ + "always_visible": True, + "values_metadata": { + "best": { + "name": "Best (Recommended)", + "description": "Use F1-optimal thresholds from model evaluation.", + }, + "default": { + "name": "Default", + "description": "Use the model's built-in default threshold.", + }, + "custom": { + "name": "Custom", + "description": "Specify a custom confidence threshold.", + }, + }, + }, + ) + custom_confidence: Union[ + Optional[FloatZeroToOne], + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Custom confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + json_schema_extra={ + "relevant_for": { + "confidence_mode": {"values": ["custom"], "required": True}, + }, + }, + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + mask_decode_mode: Union[ + Literal["accurate", "tradeoff", "fast"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="accurate", + description="Parameter of mask decoding in prediction post-processing.", + examples=["accurate", "$inputs.mask_decode_mode"], + ) + tradeoff_factor: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.0, + description="Post-processing parameter to dictate tradeoff between fast and accurate.", + examples=[0.3, "$inputs.tradeoff_factor"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + enforce_dense_masks_in_inference_models: Union[ + bool, Selector(kind=[BOOLEAN_KIND]) + ] = Field( + default=True, + description="Boolean flag to enforce dense masks when inference models backend is in use " + "(irrelevant in other cases). Dense masks are faster to process, but require more memory. " + "Users can't tweak this flag when running on Roboflow serverless platform.", + examples=[True, "$inputs.enforce_dense_masks_in_inference_models"], + ) + + @model_validator(mode="after") + def validate(self) -> "BlockManifest": + if self.confidence_mode == "custom" and self.custom_confidence is None: + raise ValueError( + "`custom_confidence` is required when `confidence_mode` is 'custom'" + ) + return self + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["instance-segmentation"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowInstanceSegmentationModelBlockV3(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence_mode: str, + custom_confidence: Optional[float], + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + enforce_dense_masks_in_inference_models: bool, + ) -> BlockResult: + confidence = ( + custom_confidence if confidence_mode == "custom" else confidence_mode + ) + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + enforce_dense_masks_in_inference_models=enforce_dense_masks_in_inference_models, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Union[None, float, Literal["best", "default"]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + enforce_dense_masks_in_inference_models: bool, + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + detections_batch: List[InstanceDetections] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + enforce_dense_masks_in_inference_models=( + enforce_dense_masks_in_inference_models + or WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS + ), + ) + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, detections in zip(images, detections_batch): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). + detections = _filter_classes_native(detections, class_filter, class_names) + # Reuse the adapter-provided inference id when present (numpy parity: + # numpy v3 surfaces ``p.get(INFERENCE_ID_KEY)`` off the model dump); + # the tensor-native adapter normally attaches none, so fall back to a + # freshly minted uuid that is then shared with ``image_metadata``. + inference_id = getattr(detections, "inference_id", None) or str( + uuid.uuid4() + ) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Union[None, float, Literal["best", "default"]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_INSTANCE_SEGMENTATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + response_mask_format="rle", + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + model_id: str, + ) -> BlockResult: + # Fallback class_id -> name map from the model, used to name boxes whose + # remote prediction dict lacks a `class` key (otherwise the tensor + # serialiser hard-raises "class_id missing from mapping"). + model_class_names = _class_names_map( + self._model_manager.get_class_names(model_id) + ) + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) or [] + if class_filter: + accepted = set(class_filter) + detection_dicts = [ + d for d in detection_dicts if d.get(CLASS_NAME_KEY) in accepted + ] + detections = _native_instance_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + model_class_names=model_class_names, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + detections: InstanceDetections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> InstanceDetections: + if not class_filter: + return detections + accepted = set(class_filter) + accepted_ids = sorted( + class_id for class_id, name in class_names.items() if name in accepted + ) + if not accepted_ids: + return take_prediction_by_mask( + detections, + torch.zeros_like(detections.class_id, dtype=torch.bool), + ) + accepted_tensor = torch.as_tensor(accepted_ids, device=detections.class_id.device) + keep = torch.isin(detections.class_id, accepted_tensor) + return take_prediction_by_mask(detections, keep) + + +def _extract_rle_mask(prediction: dict) -> Optional[dict]: + """Pull the COCO-RLE mask ({"size": [H, W], "counts": ...}) from a standard + inference instance-seg prediction dict, trying both response key variants.""" + for key in RLE_MASK_KEYS: + mask = prediction.get(key) + if mask is not None: + return mask + return None + + +def _extract_polygon_points(prediction: dict) -> Optional[List[dict]]: + """Pull the polygon ``points`` ([{"x": .., "y": ..}, ...]) from a standard + inference instance-seg prediction dict that carries a polygon mask instead of + RLE. Returns ``None`` when the key is absent or the polygon is degenerate + (< 3 points) - mirroring supervision's ``Detections.from_inference`` / the numpy + ``filter_out_invalid_polygons``, which drop such instances entirely.""" + points = prediction.get(POLYGON_KEY) + if points is not None and len(points) >= 3: + return points + return None + + +def _polygon_points_to_dense_mask( + points: List[dict], height: int, width: int +) -> np.ndarray: + """Rasterise polygon ``points`` into a dense boolean ``(H, W)`` mask, + byte-identically to supervision's ``polygon_to_mask`` as used by + ``Detections.from_inference`` (the numpy REMOTE path): vertices are integer + *truncated* (``dtype=int``, NOT rounded), filled with ``cv2.fillPoly`` onto a + zero ``uint8`` canvas, then cast to ``bool``. Keeping the exact rasterisation + means the serialiser's mask->polygon re-contouring (``mask_to_polygon``, shared + verbatim with the numpy serialiser) reproduces the numpy ``points`` output.""" + polygon = np.array([[point[X_KEY], point[Y_KEY]] for point in points], dtype=int) + return sv.polygon_to_mask(polygon, resolution_wh=(width, height)).astype(bool) + + +def _native_instance_detections_from_inference_predictions( + image: WorkflowImageData, + predictions: List[dict], + prediction_type: str, + inference_id: Optional[str] = None, + device: Optional[torch.device] = None, + model_class_names: Optional[Dict[int, str]] = None, +) -> InstanceDetections: + """REMOTE-path converter: build a native ``InstanceDetections`` from standard + inference instance-segmentation prediction dicts. + + The block requests ``response_mask_format="rle"`` so a current inference server + returns an RLE-encoded COCO mask (``{"size": [H, W], "counts": ""}``) + under ``rle`` / ``rle_mask``; those are carried as ``InstancesRLEMasks`` (counts + normalised from a utf-8 string to bytes, what ``pycocotools`` / the serialiser's + RLE decode path expect). To stay compatible with servers that ignore that + request parameter and return polygon ``points`` instead (older/alternative + inference servers), the converter degrades gracefully: when the response is not + fully RLE-backed it rebuilds the masks from the polygon ``points`` into a dense + ``torch.bool`` ``(N, H, W)`` carrier - the same masks the numpy flag-OFF REMOTE + path builds via ``sv.Detections.from_inference`` - so the serialised ``points`` + output matches numpy's (the serialiser re-contours the dense mask with the same + ``mask_to_polygon``). Boxes are converted from center form to corner ``xyxy``. + The ``class_id -> name`` map (required by the tensor serialiser) and per-box + ``detection_id`` are built here too. + + Degenerate (< 3-point) polygons are dropped exactly like the numpy path + (``filter_out_invalid_polygons`` / supervision's ``from_inference``). When a + prediction omits its ``class`` key, the class name is backfilled from + ``model_class_names`` (the model's ``get_class_names`` map) so the tensor + serialiser does not hard-raise on an unmapped ``class_id``. + """ + height, width = image._read_shape_without_materialization() + # Prefer the RLE masks the block requests; fall back to dense polygon masks when + # the response is not fully RLE-backed (server ignored `response_mask_format`). + # `all([])` is True, so an empty response keeps the empty-RLE carrier unchanged. + use_rle = all( + _extract_rle_mask(prediction) is not None for prediction in predictions + ) + if use_rle: + kept_predictions = predictions + else: + # Drop degenerate polygons up front so the surviving box arrays and masks + # stay aligned (matches the numpy path's instance drop). + kept_predictions = [ + prediction + for prediction in predictions + if _extract_polygon_points(prediction) is not None + ] + xyxy: List[List[float]] = [] + class_id: List[int] = [] + confidence: List[float] = [] + rle_counts: List[bytes] = [] + dense_masks: List[np.ndarray] = [] + bboxes_metadata: List[dict] = [] + derived_class_names: Dict[int, str] = {} + for prediction in kept_predictions: + center_x = float(prediction[X_KEY]) + center_y = float(prediction[Y_KEY]) + box_width = float(prediction[WIDTH_KEY]) + box_height = float(prediction[HEIGHT_KEY]) + xyxy.append( + [ + center_x - box_width / 2, + center_y - box_height / 2, + center_x + box_width / 2, + center_y + box_height / 2, + ] + ) + prediction_class_id = int(prediction.get(CLASS_ID_KEY, 0)) + class_id.append(prediction_class_id) + confidence.append(float(prediction.get(CONFIDENCE_KEY, 1.0))) + if CLASS_NAME_KEY in prediction: + derived_class_names[prediction_class_id] = str(prediction[CLASS_NAME_KEY]) + elif model_class_names is not None and prediction_class_id in model_class_names: + derived_class_names[prediction_class_id] = model_class_names[ + prediction_class_id + ] + if use_rle: + mask = _extract_rle_mask(prediction) + raw_counts = mask["counts"] + # Normalise to bytes: pycocotools (used by the serialiser's RLE decode) + # expects byte counts; the remote rle response carries them as utf-8 strings. + rle_counts.append( + raw_counts.encode("utf-8") + if isinstance(raw_counts, str) + else raw_counts + ) + else: + dense_masks.append( + _polygon_points_to_dense_mask( + points=_extract_polygon_points(prediction), + height=height, + width=width, + ) + ) + bboxes_metadata.append( + {DETECTION_ID_KEY: str(prediction.get(DETECTION_ID_KEY) or uuid.uuid4())} + ) + number_of_detections = len(xyxy) + image_metadata = build_native_image_metadata( + image=image, + class_names=derived_class_names, + prediction_type=prediction_type, + inference_id=inference_id, + ) + if use_rle: + mask_carrier: Union[torch.Tensor, InstancesRLEMasks] = InstancesRLEMasks( + image_size=(height, width), masks=rle_counts + ) + elif dense_masks: + mask_carrier = torch.as_tensor( + np.stack(dense_masks), dtype=torch.bool, device=device + ) + else: + # No surviving polygons: an empty dense carrier the serialiser handles. + mask_carrier = torch.zeros((0, height, width), dtype=torch.bool, device=device) + return InstanceDetections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32, device=device).reshape(-1, 4), + class_id=torch.as_tensor(class_id, dtype=torch.long, device=device).reshape(-1), + confidence=torch.as_tensor( + confidence, dtype=torch.float32, device=device + ).reshape(-1), + mask=mask_carrier, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata if number_of_detections > 0 else None, + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v4_tensor.py b/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v4_tensor.py new file mode 100644 index 0000000000..4a375a9b3a --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/instance_segmentation/v4_tensor.py @@ -0,0 +1,727 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_instance_segmentation_model@v4`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.InstanceDetections`` (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``; masks dense or RLE) under +``TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND`` / +``TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND`` instead of +``sv.Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns + ``List[InstanceDetections]`` straight from the adapter. The mask carrier (dense + ``torch.Tensor`` vs ``InstancesRLEMasks``) is adapter-decided; both are handled + downstream by the helpers and the tensor serialiser. The block applies + ``class_filter`` natively (the adapter/model does NOT read it on this path) and + attaches the producer contract (``image_metadata[class_names]`` + per-box + ``detection_id``) the tensor serialiser requires, via + ``attach_native_detection_metadata`` (it preserves masks + per-box metadata). +- REMOTE: standard inference instance-seg prediction dicts are rebuilt into a + native ``InstanceDetections`` via the in-file + ``_native_instance_detections_from_inference_predictions`` converter (never + ``sv.Detections``). The block requests ``response_mask_format="rle"`` and carries + RLE masks when the server honours it; it degrades gracefully to a dense mask + rasterised from polygon ``points`` when the server ignores that parameter + (older/alternative servers), matching the numpy flag-OFF polygon path. + +This block creates ONLY this file; it reuses the already-registered tensor +serialiser for ``instance_segmentation_prediction`` / +``rle_instance_segmentation_prediction`` (``serialise_sv_detections`` already +handles ``InstanceDetections``, dense or RLE). The numpy sibling lives in +``.../instance_segmentation/v4.py``; this manifest is identical except the output +kinds. +""" + +import uuid +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field, PositiveInt, model_validator + +from inference.core.env import ( + HOSTED_INSTANCE_SEGMENTATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + build_native_image_metadata, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_ID_KEY, + CLASS_NAME_KEY, + CONFIDENCE_KEY, + DETECTION_ID_KEY, + HEIGHT_KEY, + INFERENCE_ID_KEY, + POLYGON_KEY, + WIDTH_KEY, + X_KEY, + Y_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "instance-segmentation" + +# Keys carrying the RLE mask in standard inference instance-seg prediction dicts. +# `InstanceSegmentationRLEPrediction` declares the mask under `rle`; some response +# paths use `rle_mask` (mirrors the numpy converter's `d.get("rle_mask") or d.get("rle")`). +RLE_MASK_KEYS = ("rle_mask", "rle") + +LONG_DESCRIPTION = """ +Run inference on an instance segmentation model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). + +This version of block introduces breaking change in behaviour of mask construction - it uses +`rle` format instead `polygon` making it possible to retrieve +shapes of any kind from remote server. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Instance Segmentation Model", + "version": "v4", + "short_description": "Predict the shape, size, and location of objects.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo", "rfdetr", "rf-detr"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 1, + "inference": True, + "popular": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_instance_segmentation_model@v4"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence_mode: Union[ + Literal["best", "default", "custom"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="best", + description="How confidence thresholds are determined.", + json_schema_extra={ + "always_visible": True, + "values_metadata": { + "best": { + "name": "Best (Recommended)", + "description": "Use F1-optimal thresholds from model evaluation.", + }, + "default": { + "name": "Default", + "description": "Use the model's built-in default threshold.", + }, + "custom": { + "name": "Custom", + "description": "Specify a custom confidence threshold.", + }, + }, + }, + ) + custom_confidence: Union[ + Optional[FloatZeroToOne], + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Custom confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + json_schema_extra={ + "relevant_for": { + "confidence_mode": {"values": ["custom"], "required": True}, + }, + }, + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + mask_decode_mode: Union[ + Literal["accurate", "tradeoff", "fast"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="accurate", + description="Parameter of mask decoding in prediction post-processing.", + examples=["accurate", "$inputs.mask_decode_mode"], + ) + tradeoff_factor: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.0, + description="Post-processing parameter to dictate tradeoff between fast and accurate.", + examples=[0.3, "$inputs.tradeoff_factor"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @model_validator(mode="after") + def validate(self) -> "BlockManifest": + if self.confidence_mode == "custom" and self.custom_confidence is None: + raise ValueError( + "`custom_confidence` is required when `confidence_mode` is 'custom'" + ) + return self + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["instance-segmentation"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowInstanceSegmentationModelBlockV4(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence_mode: str, + custom_confidence: Optional[float], + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + confidence = ( + custom_confidence if confidence_mode == "custom" else confidence_mode + ) + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Union[None, float, Literal["best", "default"]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + detections_batch: List[InstanceDetections] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + enforce_dense_masks_in_inference_models=WORKFLOWS_ENFORCE_DENSE_INSTANCE_MASKS, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, detections in zip(images, detections_batch): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). + detections = _filter_classes_native(detections, class_filter, class_names) + # Reuse the adapter-provided inference id when present (numpy parity: + # numpy v4 surfaces ``p.get(INFERENCE_ID_KEY)`` off the model dump); + # the tensor-native adapter normally attaches none, so fall back to a + # freshly minted uuid that is then shared with ``image_metadata``. + inference_id = getattr(detections, "inference_id", None) or str( + uuid.uuid4() + ) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Union[None, float, Literal["best", "default"]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + mask_decode_mode: Literal["accurate", "tradeoff", "fast"], + tradeoff_factor: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_INSTANCE_SEGMENTATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + mask_decode_mode=mask_decode_mode, + tradeoff_factor=tradeoff_factor, + response_mask_format="rle", + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + model_id: str, + ) -> BlockResult: + # Fallback class_id -> name map from the model, used to name boxes whose + # remote prediction dict lacks a `class` key (otherwise the tensor + # serialiser hard-raises "class_id missing from mapping"). + model_class_names = _class_names_map( + self._model_manager.get_class_names(model_id) + ) + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) or [] + if class_filter: + accepted = set(class_filter) + detection_dicts = [ + d for d in detection_dicts if d.get(CLASS_NAME_KEY) in accepted + ] + detections = _native_instance_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + model_class_names=model_class_names, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + detections: InstanceDetections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> InstanceDetections: + if not class_filter: + return detections + accepted = set(class_filter) + accepted_ids = sorted( + class_id for class_id, name in class_names.items() if name in accepted + ) + if not accepted_ids: + return take_prediction_by_mask( + detections, + torch.zeros_like(detections.class_id, dtype=torch.bool), + ) + accepted_tensor = torch.as_tensor(accepted_ids, device=detections.class_id.device) + keep = torch.isin(detections.class_id, accepted_tensor) + return take_prediction_by_mask(detections, keep) + + +def _extract_rle_mask(prediction: dict) -> Optional[dict]: + """Pull the COCO-RLE mask ({"size": [H, W], "counts": ...}) from a standard + inference instance-seg prediction dict, trying both response key variants.""" + for key in RLE_MASK_KEYS: + mask = prediction.get(key) + if mask is not None: + return mask + return None + + +def _extract_polygon_points(prediction: dict) -> Optional[List[dict]]: + """Pull the polygon ``points`` ([{"x": .., "y": ..}, ...]) from a standard + inference instance-seg prediction dict that carries a polygon mask instead of + RLE. Returns ``None`` when the key is absent or the polygon is degenerate + (< 3 points) - mirroring supervision's ``Detections.from_inference`` / the numpy + ``filter_out_invalid_polygons``, which drop such instances entirely.""" + points = prediction.get(POLYGON_KEY) + if points is not None and len(points) >= 3: + return points + return None + + +def _polygon_points_to_dense_mask( + points: List[dict], height: int, width: int +) -> np.ndarray: + """Rasterise polygon ``points`` into a dense boolean ``(H, W)`` mask, + byte-identically to supervision's ``polygon_to_mask`` as used by + ``Detections.from_inference`` (the numpy REMOTE path): vertices are integer + *truncated* (``dtype=int``, NOT rounded), filled with ``cv2.fillPoly`` onto a + zero ``uint8`` canvas, then cast to ``bool``. Keeping the exact rasterisation + means the serialiser's mask->polygon re-contouring (``mask_to_polygon``, shared + verbatim with the numpy serialiser) reproduces the numpy ``points`` output.""" + polygon = np.array([[point[X_KEY], point[Y_KEY]] for point in points], dtype=int) + return sv.polygon_to_mask(polygon, resolution_wh=(width, height)).astype(bool) + + +def _native_instance_detections_from_inference_predictions( + image: WorkflowImageData, + predictions: List[dict], + prediction_type: str, + inference_id: Optional[str] = None, + device: Optional[torch.device] = None, + model_class_names: Optional[Dict[int, str]] = None, +) -> InstanceDetections: + """REMOTE-path converter: build a native ``InstanceDetections`` from standard + inference instance-segmentation prediction dicts. + + The block requests ``response_mask_format="rle"`` so a current inference server + returns an RLE-encoded COCO mask (``{"size": [H, W], "counts": ""}``) + under ``rle`` / ``rle_mask``; those are carried as ``InstancesRLEMasks`` (counts + normalised from a utf-8 string to bytes, what ``pycocotools`` / the serialiser's + RLE decode path expect). To stay compatible with servers that ignore that + request parameter and return polygon ``points`` instead (older/alternative + inference servers), the converter degrades gracefully: when the response is not + fully RLE-backed it rebuilds the masks from the polygon ``points`` into a dense + ``torch.bool`` ``(N, H, W)`` carrier - the same masks the numpy flag-OFF REMOTE + path builds via ``sv.Detections.from_inference`` - so the serialised ``points`` + output matches numpy's (the serialiser re-contours the dense mask with the same + ``mask_to_polygon``). Boxes are converted from center form to corner ``xyxy``. + The ``class_id -> name`` map (required by the tensor serialiser) and per-box + ``detection_id`` are built here too. + + Degenerate (< 3-point) polygons are dropped exactly like the numpy path + (``filter_out_invalid_polygons`` / supervision's ``from_inference``). When a + prediction omits its ``class`` key, the class name is backfilled from + ``model_class_names`` (the model's ``get_class_names`` map) so the tensor + serialiser does not hard-raise on an unmapped ``class_id``. + """ + height, width = image._read_shape_without_materialization() + # Prefer the RLE masks the block requests; fall back to dense polygon masks when + # the response is not fully RLE-backed (server ignored `response_mask_format`). + # `all([])` is True, so an empty response keeps the empty-RLE carrier unchanged. + use_rle = all( + _extract_rle_mask(prediction) is not None for prediction in predictions + ) + if use_rle: + kept_predictions = predictions + else: + # Drop degenerate polygons up front so the surviving box arrays and masks + # stay aligned (matches the numpy path's instance drop). + kept_predictions = [ + prediction + for prediction in predictions + if _extract_polygon_points(prediction) is not None + ] + xyxy: List[List[float]] = [] + class_id: List[int] = [] + confidence: List[float] = [] + rle_counts: List[bytes] = [] + dense_masks: List[np.ndarray] = [] + bboxes_metadata: List[dict] = [] + derived_class_names: Dict[int, str] = {} + for prediction in kept_predictions: + center_x = float(prediction[X_KEY]) + center_y = float(prediction[Y_KEY]) + box_width = float(prediction[WIDTH_KEY]) + box_height = float(prediction[HEIGHT_KEY]) + xyxy.append( + [ + center_x - box_width / 2, + center_y - box_height / 2, + center_x + box_width / 2, + center_y + box_height / 2, + ] + ) + prediction_class_id = int(prediction.get(CLASS_ID_KEY, 0)) + class_id.append(prediction_class_id) + confidence.append(float(prediction.get(CONFIDENCE_KEY, 1.0))) + if CLASS_NAME_KEY in prediction: + derived_class_names[prediction_class_id] = str(prediction[CLASS_NAME_KEY]) + elif model_class_names is not None and prediction_class_id in model_class_names: + derived_class_names[prediction_class_id] = model_class_names[ + prediction_class_id + ] + if use_rle: + mask = _extract_rle_mask(prediction) + raw_counts = mask["counts"] + # Normalise to bytes: pycocotools (used by the serialiser's RLE decode) + # expects byte counts; the remote rle response carries them as utf-8 strings. + rle_counts.append( + raw_counts.encode("utf-8") + if isinstance(raw_counts, str) + else raw_counts + ) + else: + dense_masks.append( + _polygon_points_to_dense_mask( + points=_extract_polygon_points(prediction), + height=height, + width=width, + ) + ) + bboxes_metadata.append( + {DETECTION_ID_KEY: str(prediction.get(DETECTION_ID_KEY) or uuid.uuid4())} + ) + number_of_detections = len(xyxy) + image_metadata = build_native_image_metadata( + image=image, + class_names=derived_class_names, + prediction_type=prediction_type, + inference_id=inference_id, + ) + if use_rle: + mask_carrier: Union[torch.Tensor, InstancesRLEMasks] = InstancesRLEMasks( + image_size=(height, width), masks=rle_counts + ) + elif dense_masks: + mask_carrier = torch.as_tensor( + np.stack(dense_masks), dtype=torch.bool, device=device + ) + else: + # No surviving polygons: an empty dense carrier the serialiser handles. + mask_carrier = torch.zeros((0, height, width), dtype=torch.bool, device=device) + return InstanceDetections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32, device=device).reshape(-1, 4), + class_id=torch.as_tensor(class_id, dtype=torch.long, device=device).reshape(-1), + confidence=torch.as_tensor( + confidence, dtype=torch.float32, device=device + ).reshape(-1), + mask=mask_carrier, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata if number_of_detections > 0 else None, + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v1_tensor.py b/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v1_tensor.py new file mode 100644 index 0000000000..7990d745b0 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v1_tensor.py @@ -0,0 +1,635 @@ +"""Tensor-native sibling of ``roboflow_core/roboflow_keypoint_detection_model@v1``. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits the native keypoint +prediction shape - a ``Tuple[inference_models.KeyPoints, inference_models.Detections]`` +(torch tensors on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``) - under +``TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND`` instead of ``sv.Detections``. + +This is the v1 manifest (three accepted ``type`` literals; ``inference_id`` output is +``STRING_KIND``; no ``model_id`` output) wired to the v3_tensor run bodies. v1 carries a +plain ``confidence`` field (there is no ``confidence_mode`` / ``custom_confidence`` split +- that was introduced in v3), so ``confidence`` is forwarded directly. + +Keypoint detection is unique among the tensor-native model blocks: the prediction +is a *2-tuple*. The producer contract that the tensor serialiser expects is carried +entirely on the **bbox** ``Detections`` component: + +- ``attach_native_detection_metadata`` first attaches ``image_metadata`` (with the + ``class_id -> name`` map) and a per-box ``detection_id`` to the bbox ``Detections``; +- then, for each instance ``i``, the per-instance keypoints are flattened into + ``bboxes_metadata[i]`` under the four ``KEYPOINTS_*_KEY_IN_SV_DETECTIONS`` keys the + serialiser reads. Keypoint class names are looked up via + ``key_points_classes[object_class_id]`` (variable ``K`` per object class); keypoints + with ``confidence <= 0.0`` are dropped (mirroring the numpy adapter's + ``model_keypoints_to_response``). + +The ``predictions`` output value is the full tuple ``(KeyPoints, Detections)`` so the +``KeyPoints`` component stays available to downstream tensor-native consumers; only the +serialiser unwraps the tuple back to the bbox ``Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns + ``Tuple[List[KeyPoints], List[Detections]]`` from the adapter. ``class_filter`` is + applied here natively (the adapter/model does NOT read it on this path) - the slice + is applied to the tuple so the ``KeyPoints`` and bbox ``Detections`` stay aligned. +- REMOTE: standard keypoint inference prediction dicts are rebuilt into a native bbox + ``Detections`` (never ``sv.Detections``); the per-detection ``keypoints`` lists are + flattened straight into ``bboxes_metadata``, and a parallel ``KeyPoints`` component is + reconstructed from the same dicts so the output keeps the tuple shape. +""" + +import uuid +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +import torch +from pydantic import ConfigDict, Field, PositiveInt + +from inference.core.env import ( + HOSTED_DETECT_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + native_detections_from_inference_predictions, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + CONFIDENCE_KEY, + INFERENCE_ID_KEY, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + X_KEY, + Y_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "keypoint-detection" + +# Keypoint-detection responses (both inference_models adapter and remote API) carry +# the per-detection keypoint list under this key. +KEYPOINTS_KEY = "keypoints" + +LONG_DESCRIPTION = """ +Run inference on a keypoint detection model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Keypoint Detection Model", + "version": "v1", + "short_description": "Predict skeletons on objects.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 4, + "inference": True, + }, + }, + protected_namespaces=(), + ) + type: Literal[ + "roboflow_core/roboflow_keypoint_detection_model@v1", + "RoboflowKeypointDetectionModel", + "KeypointsDetectionModel", + ] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + keypoint_confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.0, + description="Confidence threshold to predict a keypoint as visible.", + examples=[0.3, "$inputs.keypoint_confidence"], + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["keypoint-detection"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=INFERENCE_ID_KEY, kind=[STRING_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowKeypointDetectionModelBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + keypoint_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + keypoint_confidence=keypoint_confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + keypoint_confidence=keypoint_confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + keypoint_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + keypoints_batch: List[KeyPoints] + detections_batch: Optional[List[Detections]] + keypoints_batch, detections_batch = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + # The adapter's `map_inference_kwargs` only sets `key_points_threshold` + # from `kwargs["request"].keypoint_confidence`, which is absent on the + # tensor-native path - so pass it explicitly here. + key_points_threshold=keypoint_confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + if detections_batch is None: + raise RuntimeError( + "Keypoint detection model did not return the bounding-box component " + "required by the tensor-native keypoint prediction. Models adapted from " + "the `inference_models` package must provide instance detections." + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + # `key_points_classes` (List[List[str]], indexed by *object* class id) is only + # exposed on the inference_models adapter - reach the adapter directly through + # the manager's item access (the same handle that backs `get_class_names`). + key_points_classes = self._model_manager[model_id].key_points_classes + results: List[dict] = [] + for image, key_points, detections in zip( + images, keypoints_batch, detections_batch + ): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). The mask + # slices the tuple consistently across KeyPoints + bbox Detections. + key_points, detections = _filter_classes_native( + key_points=key_points, + detections=detections, + class_filter=class_filter, + class_names=class_names, + ) + inference_id = str(uuid.uuid4()) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + _attach_keypoints_to_bboxes_metadata( + detections=detections, + key_points=key_points, + key_points_classes=key_points_classes, + ) + results.append( + { + "inference_id": inference_id, + "predictions": (key_points, detections), + } + ) + return results + + def run_remotely( + self, + images: Batch[Optional[WorkflowImageData]], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + keypoint_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_DETECT_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + keypoint_confidence_threshold=keypoint_confidence, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + ) -> BlockResult: + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) + if class_filter: + detection_dicts = [ + d for d in detection_dicts if d.get("class") in class_filter + ] + detections = native_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + # The remote response already carries per-detection keypoint dicts + # (class_id / class / confidence / x / y) - flatten them straight onto the + # bbox Detections metadata and rebuild a parallel KeyPoints component so the + # output keeps the native tuple shape. + _attach_remote_keypoints_to_bboxes_metadata( + detections=detections, + detection_dicts=detection_dicts, + ) + key_points = _native_key_points_from_inference_predictions( + detection_dicts=detection_dicts, + image_metadata=detections.image_metadata, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + results.append( + { + "inference_id": inference_id, + "predictions": (key_points, detections), + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + key_points: KeyPoints, + detections: Detections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> Tuple[KeyPoints, Detections]: + if not class_filter: + return key_points, detections + accepted = set(class_filter) + keep = [ + class_names.get(int(class_id)) in accepted + for class_id in detections.class_id.tolist() + ] + # take_prediction_by_mask slices the (KeyPoints, Detections) tuple consistently. + return take_prediction_by_mask((key_points, detections), keep) + + +def _attach_keypoints_to_bboxes_metadata( + detections: Detections, + key_points: KeyPoints, + key_points_classes: List[List[str]], +) -> None: + """Flatten each instance's keypoints into the bbox ``Detections`` metadata under + the four keys the tensor serialiser reads. ``key_points_classes`` is indexed by + *object* class id (``K`` varies per class); keypoints with ``confidence <= 0.0`` + are dropped, mirroring the numpy adapter's ``model_keypoints_to_response``. + + Must run after ``attach_native_detection_metadata`` so ``bboxes_metadata`` exists. + """ + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + return + xy = key_points.xy.detach().cpu().tolist() + confidence = key_points.confidence.detach().cpu().tolist() + object_class_ids = [ + int(value) for value in key_points.class_id.detach().cpu().tolist() + ] + for index, entry in enumerate(bboxes_metadata): + instance_xy = xy[index] + instance_confidence = confidence[index] + object_class_id = object_class_ids[index] + keypoint_class_names = key_points_classes[object_class_id] + kept_class_id: List[int] = [] + kept_class_name: List[str] = [] + kept_confidence: List[float] = [] + kept_xy: List[List[float]] = [] + for keypoint_class_id, ((x, y), conf, keypoint_class_name) in enumerate( + zip(instance_xy, instance_confidence, keypoint_class_names) + ): + if conf <= 0.0: + continue + kept_class_id.append(int(keypoint_class_id)) + kept_class_name.append(str(keypoint_class_name)) + kept_confidence.append(float(conf)) + kept_xy.append([float(x), float(y)]) + entry[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS] = kept_class_id + entry[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS] = kept_class_name + entry[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS] = kept_confidence + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = kept_xy + + +def _attach_remote_keypoints_to_bboxes_metadata( + detections: Detections, + detection_dicts: List[dict], +) -> None: + """Flatten remote keypoint dicts (already keyed by class_id / class / confidence / + x / y per detection) into the bbox ``Detections`` metadata. The remote API has + already applied the keypoint confidence threshold, so no filtering is needed.""" + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + return + for entry, detection_dict in zip(bboxes_metadata, detection_dicts): + keypoints = detection_dict.get(KEYPOINTS_KEY, []) or [] + entry[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS] = [ + int(keypoint.get("class_id", keypoint_index)) + for keypoint_index, keypoint in enumerate(keypoints) + ] + entry[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS] = [ + str(keypoint.get("class", "")) for keypoint in keypoints + ] + entry[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS] = [ + float(keypoint.get(CONFIDENCE_KEY, 0.0)) for keypoint in keypoints + ] + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = [ + [float(keypoint[X_KEY]), float(keypoint[Y_KEY])] for keypoint in keypoints + ] + + +def _native_key_points_from_inference_predictions( + detection_dicts: List[dict], + image_metadata: Optional[dict], + device: Optional[Any] = None, +) -> KeyPoints: + """Rebuild a padded native ``KeyPoints`` from remote keypoint dicts so the REMOTE + output keeps the native tuple shape. Padded to a uniform ``K`` (ragged keypoint + counts across instances) with confidence 0.0 in the padding rows.""" + per_instance_xy: List[List[List[float]]] = [] + per_instance_confidence: List[List[float]] = [] + object_class_ids: List[int] = [] + for detection_dict in detection_dicts: + keypoints = detection_dict.get(KEYPOINTS_KEY, []) or [] + per_instance_xy.append( + [[float(keypoint[X_KEY]), float(keypoint[Y_KEY])] for keypoint in keypoints] + ) + per_instance_confidence.append( + [float(keypoint.get(CONFIDENCE_KEY, 0.0)) for keypoint in keypoints] + ) + object_class_ids.append(int(detection_dict.get("class_id", 0))) + number_of_instances = len(detection_dicts) + max_key_points = max((len(xy) for xy in per_instance_xy), default=0) + xy_tensor = torch.zeros( + (number_of_instances, max_key_points, 2), dtype=torch.float32, device=device + ) + confidence_tensor = torch.zeros( + (number_of_instances, max_key_points), dtype=torch.float32, device=device + ) + for index in range(number_of_instances): + count = len(per_instance_xy[index]) + if count > 0: + xy_tensor[index, :count] = torch.as_tensor( + per_instance_xy[index], dtype=torch.float32, device=device + ) + confidence_tensor[index, :count] = torch.as_tensor( + per_instance_confidence[index], dtype=torch.float32, device=device + ) + class_id_tensor = torch.as_tensor( + object_class_ids, dtype=torch.long, device=device + ).reshape(-1) + return KeyPoints( + xy=xy_tensor, + class_id=class_id_tensor, + confidence=confidence_tensor, + image_metadata=image_metadata, + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v2_tensor.py b/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v2_tensor.py new file mode 100644 index 0000000000..d148d75270 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v2_tensor.py @@ -0,0 +1,637 @@ +"""Tensor-native sibling of ``roboflow_core/roboflow_keypoint_detection_model@v2``. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits the native keypoint +prediction shape - a ``Tuple[inference_models.KeyPoints, inference_models.Detections]`` +(torch tensors on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``) - under +``TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND`` instead of ``sv.Detections``. + +This is the v2 manifest (single ``@v2`` ``type`` literal; ``inference_id`` output is +``INFERENCE_ID_KIND``; a ``model_id`` output is present, and ``model_id`` is echoed in +each result dict) wired to the v3_tensor run bodies. v2 carries a plain ``confidence`` +field (there is no ``confidence_mode`` / ``custom_confidence`` split - that was +introduced in v3), so ``confidence`` is forwarded directly. + +Keypoint detection is unique among the tensor-native model blocks: the prediction +is a *2-tuple*. The producer contract that the tensor serialiser expects is carried +entirely on the **bbox** ``Detections`` component: + +- ``attach_native_detection_metadata`` first attaches ``image_metadata`` (with the + ``class_id -> name`` map) and a per-box ``detection_id`` to the bbox ``Detections``; +- then, for each instance ``i``, the per-instance keypoints are flattened into + ``bboxes_metadata[i]`` under the four ``KEYPOINTS_*_KEY_IN_SV_DETECTIONS`` keys the + serialiser reads. Keypoint class names are looked up via + ``key_points_classes[object_class_id]`` (variable ``K`` per object class); keypoints + with ``confidence <= 0.0`` are dropped (mirroring the numpy adapter's + ``model_keypoints_to_response``). + +The ``predictions`` output value is the full tuple ``(KeyPoints, Detections)`` so the +``KeyPoints`` component stays available to downstream tensor-native consumers; only the +serialiser unwraps the tuple back to the bbox ``Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns + ``Tuple[List[KeyPoints], List[Detections]]`` from the adapter. ``class_filter`` is + applied here natively (the adapter/model does NOT read it on this path) - the slice + is applied to the tuple so the ``KeyPoints`` and bbox ``Detections`` stay aligned. +- REMOTE: standard keypoint inference prediction dicts are rebuilt into a native bbox + ``Detections`` (never ``sv.Detections``); the per-detection ``keypoints`` lists are + flattened straight into ``bboxes_metadata``, and a parallel ``KeyPoints`` component is + reconstructed from the same dicts so the output keeps the tuple shape. +""" + +import uuid +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +import torch +from pydantic import ConfigDict, Field, PositiveInt + +from inference.core.env import ( + HOSTED_DETECT_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + native_detections_from_inference_predictions, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + CONFIDENCE_KEY, + INFERENCE_ID_KEY, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + X_KEY, + Y_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "keypoint-detection" + +# Keypoint-detection responses (both inference_models adapter and remote API) carry +# the per-detection keypoint list under this key. +KEYPOINTS_KEY = "keypoints" + +LONG_DESCRIPTION = """ +Run inference on a keypoint detection model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Keypoint Detection Model", + "version": "v2", + "short_description": "Predict skeletons on objects.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 4, + "inference": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_keypoint_detection_model@v2"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + keypoint_confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.0, + description="Confidence threshold to predict a keypoint as visible.", + examples=[0.3, "$inputs.keypoint_confidence"], + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["keypoint-detection"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowKeypointDetectionModelBlockV2(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + keypoint_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + keypoint_confidence=keypoint_confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + keypoint_confidence=keypoint_confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + keypoint_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + keypoints_batch: List[KeyPoints] + detections_batch: Optional[List[Detections]] + keypoints_batch, detections_batch = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + # The adapter's `map_inference_kwargs` only sets `key_points_threshold` + # from `kwargs["request"].keypoint_confidence`, which is absent on the + # tensor-native path - so pass it explicitly here. + key_points_threshold=keypoint_confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + if detections_batch is None: + raise RuntimeError( + "Keypoint detection model did not return the bounding-box component " + "required by the tensor-native keypoint prediction. Models adapted from " + "the `inference_models` package must provide instance detections." + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + # `key_points_classes` (List[List[str]], indexed by *object* class id) is only + # exposed on the inference_models adapter - reach the adapter directly through + # the manager's item access (the same handle that backs `get_class_names`). + key_points_classes = self._model_manager[model_id].key_points_classes + results: List[dict] = [] + for image, key_points, detections in zip( + images, keypoints_batch, detections_batch + ): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). The mask + # slices the tuple consistently across KeyPoints + bbox Detections. + key_points, detections = _filter_classes_native( + key_points=key_points, + detections=detections, + class_filter=class_filter, + class_names=class_names, + ) + inference_id = str(uuid.uuid4()) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + _attach_keypoints_to_bboxes_metadata( + detections=detections, + key_points=key_points, + key_points_classes=key_points_classes, + ) + results.append( + { + "inference_id": inference_id, + "predictions": (key_points, detections), + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[Optional[WorkflowImageData]], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + keypoint_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_DETECT_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + keypoint_confidence_threshold=keypoint_confidence, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + model_id: str, + ) -> BlockResult: + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) + if class_filter: + detection_dicts = [ + d for d in detection_dicts if d.get("class") in class_filter + ] + detections = native_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + # The remote response already carries per-detection keypoint dicts + # (class_id / class / confidence / x / y) - flatten them straight onto the + # bbox Detections metadata and rebuild a parallel KeyPoints component so the + # output keeps the native tuple shape. + _attach_remote_keypoints_to_bboxes_metadata( + detections=detections, + detection_dicts=detection_dicts, + ) + key_points = _native_key_points_from_inference_predictions( + detection_dicts=detection_dicts, + image_metadata=detections.image_metadata, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + results.append( + { + "inference_id": inference_id, + "predictions": (key_points, detections), + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + key_points: KeyPoints, + detections: Detections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> Tuple[KeyPoints, Detections]: + if not class_filter: + return key_points, detections + accepted = set(class_filter) + keep = [ + class_names.get(int(class_id)) in accepted + for class_id in detections.class_id.tolist() + ] + # take_prediction_by_mask slices the (KeyPoints, Detections) tuple consistently. + return take_prediction_by_mask((key_points, detections), keep) + + +def _attach_keypoints_to_bboxes_metadata( + detections: Detections, + key_points: KeyPoints, + key_points_classes: List[List[str]], +) -> None: + """Flatten each instance's keypoints into the bbox ``Detections`` metadata under + the four keys the tensor serialiser reads. ``key_points_classes`` is indexed by + *object* class id (``K`` varies per class); keypoints with ``confidence <= 0.0`` + are dropped, mirroring the numpy adapter's ``model_keypoints_to_response``. + + Must run after ``attach_native_detection_metadata`` so ``bboxes_metadata`` exists. + """ + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + return + xy = key_points.xy.detach().cpu().tolist() + confidence = key_points.confidence.detach().cpu().tolist() + object_class_ids = [ + int(value) for value in key_points.class_id.detach().cpu().tolist() + ] + for index, entry in enumerate(bboxes_metadata): + instance_xy = xy[index] + instance_confidence = confidence[index] + object_class_id = object_class_ids[index] + keypoint_class_names = key_points_classes[object_class_id] + kept_class_id: List[int] = [] + kept_class_name: List[str] = [] + kept_confidence: List[float] = [] + kept_xy: List[List[float]] = [] + for keypoint_class_id, ((x, y), conf, keypoint_class_name) in enumerate( + zip(instance_xy, instance_confidence, keypoint_class_names) + ): + if conf <= 0.0: + continue + kept_class_id.append(int(keypoint_class_id)) + kept_class_name.append(str(keypoint_class_name)) + kept_confidence.append(float(conf)) + kept_xy.append([float(x), float(y)]) + entry[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS] = kept_class_id + entry[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS] = kept_class_name + entry[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS] = kept_confidence + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = kept_xy + + +def _attach_remote_keypoints_to_bboxes_metadata( + detections: Detections, + detection_dicts: List[dict], +) -> None: + """Flatten remote keypoint dicts (already keyed by class_id / class / confidence / + x / y per detection) into the bbox ``Detections`` metadata. The remote API has + already applied the keypoint confidence threshold, so no filtering is needed.""" + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + return + for entry, detection_dict in zip(bboxes_metadata, detection_dicts): + keypoints = detection_dict.get(KEYPOINTS_KEY, []) or [] + entry[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS] = [ + int(keypoint.get("class_id", keypoint_index)) + for keypoint_index, keypoint in enumerate(keypoints) + ] + entry[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS] = [ + str(keypoint.get("class", "")) for keypoint in keypoints + ] + entry[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS] = [ + float(keypoint.get(CONFIDENCE_KEY, 0.0)) for keypoint in keypoints + ] + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = [ + [float(keypoint[X_KEY]), float(keypoint[Y_KEY])] for keypoint in keypoints + ] + + +def _native_key_points_from_inference_predictions( + detection_dicts: List[dict], + image_metadata: Optional[dict], + device: Optional[Any] = None, +) -> KeyPoints: + """Rebuild a padded native ``KeyPoints`` from remote keypoint dicts so the REMOTE + output keeps the native tuple shape. Padded to a uniform ``K`` (ragged keypoint + counts across instances) with confidence 0.0 in the padding rows.""" + per_instance_xy: List[List[List[float]]] = [] + per_instance_confidence: List[List[float]] = [] + object_class_ids: List[int] = [] + for detection_dict in detection_dicts: + keypoints = detection_dict.get(KEYPOINTS_KEY, []) or [] + per_instance_xy.append( + [[float(keypoint[X_KEY]), float(keypoint[Y_KEY])] for keypoint in keypoints] + ) + per_instance_confidence.append( + [float(keypoint.get(CONFIDENCE_KEY, 0.0)) for keypoint in keypoints] + ) + object_class_ids.append(int(detection_dict.get("class_id", 0))) + number_of_instances = len(detection_dicts) + max_key_points = max((len(xy) for xy in per_instance_xy), default=0) + xy_tensor = torch.zeros( + (number_of_instances, max_key_points, 2), dtype=torch.float32, device=device + ) + confidence_tensor = torch.zeros( + (number_of_instances, max_key_points), dtype=torch.float32, device=device + ) + for index in range(number_of_instances): + count = len(per_instance_xy[index]) + if count > 0: + xy_tensor[index, :count] = torch.as_tensor( + per_instance_xy[index], dtype=torch.float32, device=device + ) + confidence_tensor[index, :count] = torch.as_tensor( + per_instance_confidence[index], dtype=torch.float32, device=device + ) + class_id_tensor = torch.as_tensor( + object_class_ids, dtype=torch.long, device=device + ).reshape(-1) + return KeyPoints( + xy=xy_tensor, + class_id=class_id_tensor, + confidence=confidence_tensor, + image_metadata=image_metadata, + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v3_tensor.py b/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v3_tensor.py new file mode 100644 index 0000000000..834f2b06bc --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/keypoint_detection/v3_tensor.py @@ -0,0 +1,670 @@ +"""Tensor-native sibling of ``roboflow_core/roboflow_keypoint_detection_model@v3``. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits the native keypoint +prediction shape - a ``Tuple[inference_models.KeyPoints, inference_models.Detections]`` +(torch tensors on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``) - under +``TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND`` instead of ``sv.Detections``. + +Keypoint detection is unique among the tensor-native model blocks: the prediction +is a *2-tuple*. The producer contract that the tensor serialiser expects is carried +entirely on the **bbox** ``Detections`` component: + +- ``attach_native_detection_metadata`` first attaches ``image_metadata`` (with the + ``class_id -> name`` map) and a per-box ``detection_id`` to the bbox ``Detections``; +- then, for each instance ``i``, the per-instance keypoints are flattened into + ``bboxes_metadata[i]`` under the four ``KEYPOINTS_*_KEY_IN_SV_DETECTIONS`` keys the + serialiser reads. Keypoint class names are looked up via + ``key_points_classes[object_class_id]`` (variable ``K`` per object class); keypoints + with ``confidence <= 0.0`` are dropped (mirroring the numpy adapter's + ``model_keypoints_to_response``). + +The ``predictions`` output value is the full tuple ``(KeyPoints, Detections)`` so the +``KeyPoints`` component stays available to downstream tensor-native consumers; only the +serialiser unwraps the tuple back to the bbox ``Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns + ``Tuple[List[KeyPoints], List[Detections]]`` from the adapter. ``class_filter`` is + applied here natively (the adapter/model does NOT read it on this path) - the slice + is applied to the tuple so the ``KeyPoints`` and bbox ``Detections`` stay aligned. +- REMOTE: standard keypoint inference prediction dicts are rebuilt into a native bbox + ``Detections`` (never ``sv.Detections``); the per-detection ``keypoints`` lists are + flattened straight into ``bboxes_metadata``, and a parallel ``KeyPoints`` component is + reconstructed from the same dicts so the output keeps the tuple shape. +""" + +import uuid +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +import torch +from pydantic import ConfigDict, Field, PositiveInt, model_validator + +from inference.core.env import ( + HOSTED_DETECT_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + native_detections_from_inference_predictions, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + CONFIDENCE_KEY, + INFERENCE_ID_KEY, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + X_KEY, + Y_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "keypoint-detection" + +# Keypoint-detection responses (both inference_models adapter and remote API) carry +# the per-detection keypoint list under this key. +KEYPOINTS_KEY = "keypoints" + +LONG_DESCRIPTION = """ +Run inference on a keypoint detection model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Keypoint Detection Model", + "version": "v3", + "short_description": "Predict skeletons on objects.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 4, + "inference": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_keypoint_detection_model@v3"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + # TODO: add "best" option once model eval supports keypoint detection. + confidence_mode: Union[ + Literal["default", "custom"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="default", + json_schema_extra={ + "always_visible": True, + "values_metadata": { + "default": { + "name": "Default", + "description": "Use the model's built-in default threshold.", + }, + "custom": { + "name": "Custom", + "description": "Specify a custom confidence threshold.", + }, + }, + }, + ) + custom_confidence: Union[ + Optional[FloatZeroToOne], + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + json_schema_extra={ + "relevant_for": { + "confidence_mode": { + "values": ["custom"], + "required": True, + }, + }, + }, + ) + keypoint_confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.0, + description="Confidence threshold to predict a keypoint as visible.", + examples=[0.3, "$inputs.keypoint_confidence"], + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @model_validator(mode="after") + def validate_custom_confidence(self) -> "BlockManifest": + if self.confidence_mode == "custom" and self.custom_confidence is None: + raise ValueError( + "custom_confidence must be provided when confidence_mode is 'custom'" + ) + return self + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["keypoint-detection"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowKeypointDetectionModelBlockV3(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence_mode: str, + custom_confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + keypoint_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + confidence = ( + custom_confidence if confidence_mode == "custom" else confidence_mode + ) + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + keypoint_confidence=keypoint_confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + keypoint_confidence=keypoint_confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Union[None, float, Literal["default"]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + keypoint_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + keypoints_batch: List[KeyPoints] + detections_batch: Optional[List[Detections]] + keypoints_batch, detections_batch = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + # The adapter's `map_inference_kwargs` only sets `key_points_threshold` + # from `kwargs["request"].keypoint_confidence`, which is absent on the + # tensor-native path - so pass it explicitly here. + key_points_threshold=keypoint_confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + if detections_batch is None: + raise RuntimeError( + "Keypoint detection model did not return the bounding-box component " + "required by the tensor-native keypoint prediction. Models adapted from " + "the `inference_models` package must provide instance detections." + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + # `key_points_classes` (List[List[str]], indexed by *object* class id) is only + # exposed on the inference_models adapter - reach the adapter directly through + # the manager's item access (the same handle that backs `get_class_names`). + key_points_classes = self._model_manager[model_id].key_points_classes + results: List[dict] = [] + for image, key_points, detections in zip( + images, keypoints_batch, detections_batch + ): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). The mask + # slices the tuple consistently across KeyPoints + bbox Detections. + key_points, detections = _filter_classes_native( + key_points=key_points, + detections=detections, + class_filter=class_filter, + class_names=class_names, + ) + inference_id = str(uuid.uuid4()) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + _attach_keypoints_to_bboxes_metadata( + detections=detections, + key_points=key_points, + key_points_classes=key_points_classes, + ) + results.append( + { + "inference_id": inference_id, + "predictions": (key_points, detections), + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[Optional[WorkflowImageData]], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Union[None, float, Literal["default"]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + keypoint_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_DETECT_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + keypoint_confidence_threshold=keypoint_confidence, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + model_id: str, + ) -> BlockResult: + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) + if class_filter: + detection_dicts = [ + d for d in detection_dicts if d.get("class") in class_filter + ] + detections = native_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + # The remote response already carries per-detection keypoint dicts + # (class_id / class / confidence / x / y) - flatten them straight onto the + # bbox Detections metadata and rebuild a parallel KeyPoints component so the + # output keeps the native tuple shape. + _attach_remote_keypoints_to_bboxes_metadata( + detections=detections, + detection_dicts=detection_dicts, + ) + key_points = _native_key_points_from_inference_predictions( + detection_dicts=detection_dicts, + image_metadata=detections.image_metadata, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + results.append( + { + "inference_id": inference_id, + "predictions": (key_points, detections), + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + key_points: KeyPoints, + detections: Detections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> Tuple[KeyPoints, Detections]: + if not class_filter: + return key_points, detections + accepted = set(class_filter) + keep = [ + class_names.get(int(class_id)) in accepted + for class_id in detections.class_id.tolist() + ] + # take_prediction_by_mask slices the (KeyPoints, Detections) tuple consistently. + return take_prediction_by_mask((key_points, detections), keep) + + +def _attach_keypoints_to_bboxes_metadata( + detections: Detections, + key_points: KeyPoints, + key_points_classes: List[List[str]], +) -> None: + """Flatten each instance's keypoints into the bbox ``Detections`` metadata under + the four keys the tensor serialiser reads. ``key_points_classes`` is indexed by + *object* class id (``K`` varies per class); keypoints with ``confidence <= 0.0`` + are dropped, mirroring the numpy adapter's ``model_keypoints_to_response``. + + Must run after ``attach_native_detection_metadata`` so ``bboxes_metadata`` exists. + """ + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + return + xy = key_points.xy.detach().cpu().tolist() + confidence = key_points.confidence.detach().cpu().tolist() + object_class_ids = [ + int(value) for value in key_points.class_id.detach().cpu().tolist() + ] + for index, entry in enumerate(bboxes_metadata): + instance_xy = xy[index] + instance_confidence = confidence[index] + object_class_id = object_class_ids[index] + keypoint_class_names = key_points_classes[object_class_id] + kept_class_id: List[int] = [] + kept_class_name: List[str] = [] + kept_confidence: List[float] = [] + kept_xy: List[List[float]] = [] + for keypoint_class_id, ((x, y), conf, keypoint_class_name) in enumerate( + zip(instance_xy, instance_confidence, keypoint_class_names) + ): + if conf <= 0.0: + continue + kept_class_id.append(int(keypoint_class_id)) + kept_class_name.append(str(keypoint_class_name)) + kept_confidence.append(float(conf)) + kept_xy.append([float(x), float(y)]) + entry[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS] = kept_class_id + entry[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS] = kept_class_name + entry[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS] = kept_confidence + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = kept_xy + + +def _attach_remote_keypoints_to_bboxes_metadata( + detections: Detections, + detection_dicts: List[dict], +) -> None: + """Flatten remote keypoint dicts (already keyed by class_id / class / confidence / + x / y per detection) into the bbox ``Detections`` metadata. The remote API has + already applied the keypoint confidence threshold, so no filtering is needed.""" + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + return + for entry, detection_dict in zip(bboxes_metadata, detection_dicts): + keypoints = detection_dict.get(KEYPOINTS_KEY, []) or [] + entry[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS] = [ + int(keypoint.get("class_id", keypoint_index)) + for keypoint_index, keypoint in enumerate(keypoints) + ] + entry[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS] = [ + str(keypoint.get("class", "")) for keypoint in keypoints + ] + entry[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS] = [ + float(keypoint.get(CONFIDENCE_KEY, 0.0)) for keypoint in keypoints + ] + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = [ + [float(keypoint[X_KEY]), float(keypoint[Y_KEY])] for keypoint in keypoints + ] + + +def _native_key_points_from_inference_predictions( + detection_dicts: List[dict], + image_metadata: Optional[dict], + device: Optional[Any] = None, +) -> KeyPoints: + """Rebuild a padded native ``KeyPoints`` from remote keypoint dicts so the REMOTE + output keeps the native tuple shape. Padded to a uniform ``K`` (ragged keypoint + counts across instances) with confidence 0.0 in the padding rows.""" + per_instance_xy: List[List[List[float]]] = [] + per_instance_confidence: List[List[float]] = [] + object_class_ids: List[int] = [] + for detection_dict in detection_dicts: + keypoints = detection_dict.get(KEYPOINTS_KEY, []) or [] + per_instance_xy.append( + [[float(keypoint[X_KEY]), float(keypoint[Y_KEY])] for keypoint in keypoints] + ) + per_instance_confidence.append( + [float(keypoint.get(CONFIDENCE_KEY, 0.0)) for keypoint in keypoints] + ) + object_class_ids.append(int(detection_dict.get("class_id", 0))) + number_of_instances = len(detection_dicts) + max_key_points = max((len(xy) for xy in per_instance_xy), default=0) + xy_tensor = torch.zeros( + (number_of_instances, max_key_points, 2), dtype=torch.float32, device=device + ) + confidence_tensor = torch.zeros( + (number_of_instances, max_key_points), dtype=torch.float32, device=device + ) + for index in range(number_of_instances): + count = len(per_instance_xy[index]) + if count > 0: + xy_tensor[index, :count] = torch.as_tensor( + per_instance_xy[index], dtype=torch.float32, device=device + ) + confidence_tensor[index, :count] = torch.as_tensor( + per_instance_confidence[index], dtype=torch.float32, device=device + ) + class_id_tensor = torch.as_tensor( + object_class_ids, dtype=torch.long, device=device + ).reshape(-1) + return KeyPoints( + xy=xy_tensor, + class_id=class_id_tensor, + confidence=confidence_tensor, + image_metadata=image_metadata, + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v1_tensor.py b/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v1_tensor.py new file mode 100644 index 0000000000..e020f0f92d --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v1_tensor.py @@ -0,0 +1,469 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_classification_model@v1`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits native +``inference_models.ClassificationPrediction`` objects (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``) under +``TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND`` instead of the legacy +classification response dict. + +The manifest mirrors the numpy ``v1`` sibling EXACTLY (3-way ``type`` Literal +with legacy aliases, plain ``confidence`` field, ``inference_id`` output kind +``STRING_KIND``, no ``model_id`` output) โ€” ONLY the ``predictions`` output kind +is changed to the tensor-native kind. The run bodies follow the ``v3_tensor`` +pattern: + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns ONE batched + ``ClassificationPrediction`` (``class_id`` shape ``(bs,)``, ``confidence`` shape + ``(bs, num_classes)`` full softmax). The consumer indexes per-image, so the + block fans the batched object out into ``bs`` single-row predictions, each + carrying the producer contract in ``images_metadata`` (PLURAL) that the + tensor classification serialiser requires (mirrors + ``formatters/vlm_as_classifier/v1_tensor.py``). +- REMOTE: standard inference classification response dicts are rebuilt into a + native ``ClassificationPrediction`` via an inline converter (never the legacy + response dict). +""" + +import uuid +from time import perf_counter +from typing import Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import ( + HOSTED_CLASSIFICATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_KEY, + CLASSIFICATION_STYLE_MODEL, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.classification import ClassificationPrediction +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "classification" + +# Numpy parity: the adapter `postprocess` (inference_models_adapters.py) drops +# classes whose score < confidence_threshold, falling back to 0.5 when the +# resolved confidence is non-numeric (e.g. the string "default"). The native +# LOCAL path bypasses `postprocess`, so we carry the resolved float in +# `image_metadata` for `serialise_native_classification` to apply the same cutoff. +_DEFAULT_CONFIDENCE_THRESHOLD = 0.5 +CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY = "classification_confidence_threshold" + + +def _resolve_confidence_threshold(confidence: Optional[Union[float, str]]) -> float: + if isinstance(confidence, (int, float)): + return float(confidence) + return _DEFAULT_CONFIDENCE_THRESHOLD + + +LONG_DESCRIPTION = """ +Run inference on a multi-class classification model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Single-Label Classification Model", + "version": "v1", + "short_description": "Apply a single tag to an image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + }, + protected_namespaces=(), + ) + type: Literal[ + "roboflow_core/roboflow_classification_model@v1", + "RoboflowClassificationModel", + "ClassificationModel", + ] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["classification"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND], + ), + OutputDefinition(name=INFERENCE_ID_KEY, kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowClassificationModelBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + # Single-label returns ONE batched ClassificationPrediction (NOT a list): + # class_id -> (bs,) + # confidence -> (bs, num_classes) full softmax + inference_start = perf_counter() + batched_prediction: ClassificationPrediction = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + # numpy parity: `Model.infer_from_request` stamps `time` (elapsed around + # the bare model call) on every response in the batch; mirror with one + # elapsed value for the whole batch (numpy's per-response deltas differ + # only by loop overhead). + inference_time = perf_counter() - inference_start + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + confidence_threshold = _resolve_confidence_threshold(confidence) + results: List[dict] = [] + for index, image in enumerate(images): + inference_id = str(uuid.uuid4()) + # Fan the batched object out into a single-row (bs=1) prediction the + # consumer can index per-image. confidence stays the FULL per-class + # softmax vector (do NOT collapse to a scalar). + prediction = _single_row_prediction( + batched_prediction=batched_prediction, + index=index, + image=image, + class_names=class_names, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=inference_time, + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + } + ) + return results + + def run_remotely( + self, + images: Batch[Optional[WorkflowImageData]], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CLASSIFICATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + confidence_threshold=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=non_empty_inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + confidence_threshold = _resolve_confidence_threshold(confidence) + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + prediction = _native_classification_from_inference_response( + image=image, + response=response, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=response.get("time"), + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _build_image_metadata( + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + confidence_threshold: float, + inference_time: Optional[float] = None, +) -> dict: + height, width = image._read_shape_without_materialization() + metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag the serialiser reads to pick the flag-OFF shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_MODEL, + PREDICTION_TYPE_KEY: PREDICTION_TYPE, + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + ROOT_PARENT_ID_KEY: image.workflow_root_ancestor_metadata.parent_id, + # C2: resolved float threshold so `serialise_native_classification` + # drops sub-threshold classes from `predictions`, matching numpy. + CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY: confidence_threshold, + } + # numpy parity: flag-off always dumps `time` โ€” stamped around the bare model + # call in `Model.infer_from_request` (LOCAL) or by the server (REMOTE). + if inference_time is not None: + metadata["time"] = inference_time + return metadata + + +def _single_row_prediction( + batched_prediction: ClassificationPrediction, + index: int, + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + confidence_threshold: float, + inference_time: Optional[float] = None, +) -> ClassificationPrediction: + image_metadata = _build_image_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=inference_time, + ) + # Slice (not index) to keep the leading batch dim: class_id -> (1,), + # confidence -> (1, num_classes). The confidence row stays the FULL softmax. + return ClassificationPrediction( + class_id=batched_prediction.class_id[index : index + 1].to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=batched_prediction.confidence[index : index + 1].to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + images_metadata=[image_metadata], + ) + + +def _native_classification_from_inference_response( + image: WorkflowImageData, + response: dict, + inference_id: str, + confidence_threshold: float, + inference_time: Optional[float] = None, +) -> ClassificationPrediction: + """Rebuild a native single-row ``ClassificationPrediction`` from a standard + inference classification response dict. + + The remote response carries ``predictions`` as a list of + ``{"class_id", "class", "confidence"}`` (already confidence-filtered/sorted) + plus a ``top`` class. We reconstruct a dense confidence vector indexed by + ``class_id`` and a ``class_id -> name`` map from those entries; classes not + in the response default to 0.0 confidence (the model package's full class + list is not available on this path). + """ + detection_dicts = response.get("predictions", []) or [] + class_names: Dict[int, str] = {} + confidence_by_id: Dict[int, float] = {} + for entry in detection_dicts: + class_id = int(entry["class_id"]) + class_names[class_id] = str(entry.get("class", class_id)) + confidence_by_id[class_id] = float(entry.get("confidence", 0.0)) + num_classes = (max(class_names.keys()) + 1) if class_names else 0 + # Fill gaps so the dense vector and the class_names map agree on every index. + for class_id in range(num_classes): + class_names.setdefault(class_id, str(class_id)) + confidence_vector = [confidence_by_id.get(i, 0.0) for i in range(num_classes)] + top_class = response.get("top") + top_class_id = next( + (cid for cid, name in class_names.items() if name == top_class), + int(max(confidence_by_id, key=confidence_by_id.get)) if confidence_by_id else 0, + ) + image_metadata = _build_image_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=inference_time, + ) + return ClassificationPrediction( + class_id=torch.tensor( + [top_class_id], dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.tensor( + [confidence_vector], + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + images_metadata=[image_metadata], + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v2_tensor.py b/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v2_tensor.py new file mode 100644 index 0000000000..71164f10f8 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v2_tensor.py @@ -0,0 +1,475 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_classification_model@v2`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits native +``inference_models.ClassificationPrediction`` objects (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``) under +``TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND`` instead of the legacy +classification response dict. + +The manifest mirrors the numpy ``v2`` sibling EXACTLY (single ``type`` Literal, +``ui_manifest``, plain ``confidence`` field, ``inference_id`` output kind +``INFERENCE_ID_KIND``, ``model_id`` output) โ€” ONLY the ``predictions`` output +kind is changed to the tensor-native kind. The run bodies follow the +``v3_tensor`` pattern: + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns ONE batched + ``ClassificationPrediction`` (``class_id`` shape ``(bs,)``, ``confidence`` shape + ``(bs, num_classes)`` full softmax). The consumer indexes per-image, so the + block fans the batched object out into ``bs`` single-row predictions, each + carrying the producer contract in ``images_metadata`` (PLURAL) that the + tensor classification serialiser requires (mirrors + ``formatters/vlm_as_classifier/v1_tensor.py``). +- REMOTE: standard inference classification response dicts are rebuilt into a + native ``ClassificationPrediction`` via an inline converter (never the legacy + response dict). +""" + +import uuid +from time import perf_counter +from typing import Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import ( + HOSTED_CLASSIFICATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_KEY, + CLASSIFICATION_STYLE_MODEL, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.classification import ClassificationPrediction +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "classification" + +# Numpy parity: the adapter `postprocess` (inference_models_adapters.py) drops +# classes whose score < confidence_threshold, falling back to 0.5 when the +# resolved confidence is non-numeric (e.g. the string "default"). The native +# LOCAL path bypasses `postprocess`, so we carry the resolved float in +# `image_metadata` for `serialise_native_classification` to apply the same cutoff. +_DEFAULT_CONFIDENCE_THRESHOLD = 0.5 +CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY = "classification_confidence_threshold" + + +def _resolve_confidence_threshold(confidence: Optional[Union[float, str]]) -> float: + if isinstance(confidence, (int, float)): + return float(confidence) + return _DEFAULT_CONFIDENCE_THRESHOLD + + +LONG_DESCRIPTION = """ +Run inference on a multi-class classification model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Single-Label Classification Model", + "version": "v2", + "short_description": "Apply a single tag to an image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 2, + "inference": True, + "popular": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_classification_model@v2"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["classification"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND], + ), + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowClassificationModelBlockV2(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + # Single-label returns ONE batched ClassificationPrediction (NOT a list): + # class_id -> (bs,) + # confidence -> (bs, num_classes) full softmax + inference_start = perf_counter() + batched_prediction: ClassificationPrediction = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + # numpy parity: `Model.infer_from_request` stamps `time` (elapsed around + # the bare model call) on every response in the batch; mirror with one + # elapsed value for the whole batch (numpy's per-response deltas differ + # only by loop overhead). + inference_time = perf_counter() - inference_start + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + confidence_threshold = _resolve_confidence_threshold(confidence) + results: List[dict] = [] + for index, image in enumerate(images): + inference_id = str(uuid.uuid4()) + # Fan the batched object out into a single-row (bs=1) prediction the + # consumer can index per-image. confidence stays the FULL per-class + # softmax vector (do NOT collapse to a scalar). + prediction = _single_row_prediction( + batched_prediction=batched_prediction, + index=index, + image=image, + class_names=class_names, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=inference_time, + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[Optional[WorkflowImageData]], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CLASSIFICATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + confidence_threshold=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=non_empty_inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + confidence_threshold = _resolve_confidence_threshold(confidence) + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + prediction = _native_classification_from_inference_response( + image=image, + response=response, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=response.get("time"), + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _build_image_metadata( + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + confidence_threshold: float, + inference_time: Optional[float] = None, +) -> dict: + height, width = image._read_shape_without_materialization() + metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag the serialiser reads to pick the flag-OFF shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_MODEL, + PREDICTION_TYPE_KEY: PREDICTION_TYPE, + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + ROOT_PARENT_ID_KEY: image.workflow_root_ancestor_metadata.parent_id, + # C2: resolved float threshold so `serialise_native_classification` + # drops sub-threshold classes from `predictions`, matching numpy. + CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY: confidence_threshold, + } + # numpy parity: flag-off always dumps `time` โ€” stamped around the bare model + # call in `Model.infer_from_request` (LOCAL) or by the server (REMOTE). + if inference_time is not None: + metadata["time"] = inference_time + return metadata + + +def _single_row_prediction( + batched_prediction: ClassificationPrediction, + index: int, + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + confidence_threshold: float, + inference_time: Optional[float] = None, +) -> ClassificationPrediction: + image_metadata = _build_image_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=inference_time, + ) + # Slice (not index) to keep the leading batch dim: class_id -> (1,), + # confidence -> (1, num_classes). The confidence row stays the FULL softmax. + return ClassificationPrediction( + class_id=batched_prediction.class_id[index : index + 1].to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=batched_prediction.confidence[index : index + 1].to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + images_metadata=[image_metadata], + ) + + +def _native_classification_from_inference_response( + image: WorkflowImageData, + response: dict, + inference_id: str, + confidence_threshold: float, + inference_time: Optional[float] = None, +) -> ClassificationPrediction: + """Rebuild a native single-row ``ClassificationPrediction`` from a standard + inference classification response dict. + + The remote response carries ``predictions`` as a list of + ``{"class_id", "class", "confidence"}`` (already confidence-filtered/sorted) + plus a ``top`` class. We reconstruct a dense confidence vector indexed by + ``class_id`` and a ``class_id -> name`` map from those entries; classes not + in the response default to 0.0 confidence (the model package's full class + list is not available on this path). + """ + detection_dicts = response.get("predictions", []) or [] + class_names: Dict[int, str] = {} + confidence_by_id: Dict[int, float] = {} + for entry in detection_dicts: + class_id = int(entry["class_id"]) + class_names[class_id] = str(entry.get("class", class_id)) + confidence_by_id[class_id] = float(entry.get("confidence", 0.0)) + num_classes = (max(class_names.keys()) + 1) if class_names else 0 + # Fill gaps so the dense vector and the class_names map agree on every index. + for class_id in range(num_classes): + class_names.setdefault(class_id, str(class_id)) + confidence_vector = [confidence_by_id.get(i, 0.0) for i in range(num_classes)] + top_class = response.get("top") + top_class_id = next( + (cid for cid, name in class_names.items() if name == top_class), + int(max(confidence_by_id, key=confidence_by_id.get)) if confidence_by_id else 0, + ) + image_metadata = _build_image_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=inference_time, + ) + return ClassificationPrediction( + class_id=torch.tensor( + [top_class_id], dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.tensor( + [confidence_vector], + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + images_metadata=[image_metadata], + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v3_tensor.py b/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v3_tensor.py new file mode 100644 index 0000000000..b62ffc2e37 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/multi_class_classification/v3_tensor.py @@ -0,0 +1,518 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_classification_model@v3`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits native +``inference_models.ClassificationPrediction`` objects (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``) under +``TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND`` instead of the legacy +classification response dict. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns ONE batched + ``ClassificationPrediction`` (``class_id`` shape ``(bs,)``, ``confidence`` shape + ``(bs, num_classes)`` full softmax). The consumer indexes per-image, so the + block fans the batched object out into ``bs`` single-row predictions, each + carrying the producer contract in ``images_metadata`` (PLURAL) that the + tensor classification serialiser requires (mirrors + ``formatters/vlm_as_classifier/v1_tensor.py``). +- REMOTE: standard inference classification response dicts are rebuilt into a + native ``ClassificationPrediction`` via an inline converter (never the legacy + response dict). +""" + +import uuid +from time import perf_counter +from typing import Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field, model_validator + +from inference.core.env import ( + HOSTED_CLASSIFICATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_KEY, + CLASSIFICATION_STYLE_MODEL, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.classification import ClassificationPrediction +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "classification" + +# Numpy parity: the adapter `postprocess` (inference_models_adapters.py) drops +# classes whose score < confidence_threshold, falling back to 0.5 when the +# resolved confidence is non-numeric (e.g. the string "default"). The native +# LOCAL path bypasses `postprocess`, so we (a) resolve the string mode to this +# numeric default BEFORE the native model call and (b) carry the resolved float +# in `image_metadata` for `serialise_native_classification` to apply the cutoff. +_DEFAULT_CONFIDENCE_THRESHOLD = 0.5 +CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY = "classification_confidence_threshold" + + +def _resolve_confidence_threshold(confidence: Optional[Union[float, str]]) -> float: + if isinstance(confidence, (int, float)): + return float(confidence) + return _DEFAULT_CONFIDENCE_THRESHOLD + + +LONG_DESCRIPTION = """ +Run inference on a multi-class classification model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Single-Label Classification Model", + "version": "v3", + "short_description": "Apply a single tag to an image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 2, + "inference": True, + "popular": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_classification_model@v3"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + # Single-label classification opts out of per-class refinement โ€” top-1 + # drives the response, so the "best" (F1-optimal) mode from model eval has + # no meaningful effect here and is omitted. + confidence_mode: Union[ + Literal["default", "custom"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="default", + description="How to determine the confidence threshold.", + json_schema_extra={ + "always_visible": True, + "values_metadata": { + "default": { + "name": "Default", + "description": "Use the model's built-in default threshold.", + }, + "custom": { + "name": "Custom", + "description": "Specify a custom confidence threshold.", + }, + }, + }, + ) + custom_confidence: Union[ + Optional[FloatZeroToOne], + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Custom confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + json_schema_extra={ + "relevant_for": { + "confidence_mode": { + "values": ["custom"], + "required": True, + }, + }, + }, + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @model_validator(mode="after") + def validate_custom_confidence(self) -> "BlockManifest": + if self.confidence_mode == "custom" and self.custom_confidence is None: + raise ValueError( + "custom_confidence must be provided when confidence_mode is 'custom'" + ) + return self + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["classification"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND], + ), + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowClassificationModelBlockV3(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence_mode: str, + custom_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + confidence = ( + custom_confidence if confidence_mode == "custom" else confidence_mode + ) + # Resolve the string mode ("default") to a concrete float BEFORE dispatch + # so the native model `__call__` (which bypasses the adapter `postprocess` + # string->0.5 fallback) never receives a non-numeric confidence. + confidence = _resolve_confidence_threshold(confidence) + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Union[None, float, Literal["default"]], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + # Single-label returns ONE batched ClassificationPrediction (NOT a list): + # class_id -> (bs,) + # confidence -> (bs, num_classes) full softmax + inference_start = perf_counter() + batched_prediction: ClassificationPrediction = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + # numpy parity: `Model.infer_from_request` stamps `time` (elapsed around + # the bare model call) on every response in the batch; mirror with one + # elapsed value for the whole batch (numpy's per-response deltas differ + # only by loop overhead). + inference_time = perf_counter() - inference_start + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + confidence_threshold = _resolve_confidence_threshold(confidence) + results: List[dict] = [] + for index, image in enumerate(images): + inference_id = str(uuid.uuid4()) + # Fan the batched object out into a single-row (bs=1) prediction the + # consumer can index per-image. confidence stays the FULL per-class + # softmax vector (do NOT collapse to a scalar). + prediction = _single_row_prediction( + batched_prediction=batched_prediction, + index=index, + image=image, + class_names=class_names, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=inference_time, + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[Optional[WorkflowImageData]], + model_id: str, + confidence: Union[None, float, Literal["default"]], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CLASSIFICATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + confidence_threshold=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=non_empty_inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + confidence_threshold = _resolve_confidence_threshold(confidence) + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + prediction = _native_classification_from_inference_response( + image=image, + response=response, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=response.get("time"), + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _build_image_metadata( + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + confidence_threshold: float, + inference_time: Optional[float] = None, +) -> dict: + height, width = image._read_shape_without_materialization() + metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag the serialiser reads to pick the flag-OFF shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_MODEL, + PREDICTION_TYPE_KEY: PREDICTION_TYPE, + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + ROOT_PARENT_ID_KEY: image.workflow_root_ancestor_metadata.parent_id, + # C2: resolved float threshold so `serialise_native_classification` + # drops sub-threshold classes from `predictions`, matching numpy. + CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY: confidence_threshold, + } + # numpy parity: flag-off always dumps `time` โ€” stamped around the bare model + # call in `Model.infer_from_request` (LOCAL) or by the server (REMOTE). + if inference_time is not None: + metadata["time"] = inference_time + return metadata + + +def _single_row_prediction( + batched_prediction: ClassificationPrediction, + index: int, + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + confidence_threshold: float, + inference_time: Optional[float] = None, +) -> ClassificationPrediction: + image_metadata = _build_image_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=inference_time, + ) + # Slice (not index) to keep the leading batch dim: class_id -> (1,), + # confidence -> (1, num_classes). The confidence row stays the FULL softmax. + return ClassificationPrediction( + class_id=batched_prediction.class_id[index : index + 1].to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=batched_prediction.confidence[index : index + 1].to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + images_metadata=[image_metadata], + ) + + +def _native_classification_from_inference_response( + image: WorkflowImageData, + response: dict, + inference_id: str, + confidence_threshold: float, + inference_time: Optional[float] = None, +) -> ClassificationPrediction: + """Rebuild a native single-row ``ClassificationPrediction`` from a standard + inference classification response dict. + + The remote response carries ``predictions`` as a list of + ``{"class_id", "class", "confidence"}`` (already confidence-filtered/sorted) + plus a ``top`` class. We reconstruct a dense confidence vector indexed by + ``class_id`` and a ``class_id -> name`` map from those entries; classes not + in the response default to 0.0 confidence (the model package's full class + list is not available on this path). + """ + detection_dicts = response.get("predictions", []) or [] + class_names: Dict[int, str] = {} + confidence_by_id: Dict[int, float] = {} + for entry in detection_dicts: + class_id = int(entry["class_id"]) + class_names[class_id] = str(entry.get("class", class_id)) + confidence_by_id[class_id] = float(entry.get("confidence", 0.0)) + num_classes = (max(class_names.keys()) + 1) if class_names else 0 + # Fill gaps so the dense vector and the class_names map agree on every index. + for class_id in range(num_classes): + class_names.setdefault(class_id, str(class_id)) + confidence_vector = [confidence_by_id.get(i, 0.0) for i in range(num_classes)] + top_class = response.get("top") + top_class_id = next( + (cid for cid, name in class_names.items() if name == top_class), + int(max(confidence_by_id, key=confidence_by_id.get)) if confidence_by_id else 0, + ) + image_metadata = _build_image_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + confidence_threshold=confidence_threshold, + inference_time=inference_time, + ) + return ClassificationPrediction( + class_id=torch.tensor( + [top_class_id], dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.tensor( + [confidence_vector], + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + images_metadata=[image_metadata], + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v1_tensor.py b/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v1_tensor.py new file mode 100644 index 0000000000..22cf4cdfe7 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v1_tensor.py @@ -0,0 +1,451 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_multi_label_classification_model@v1`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.MultiLabelClassificationPrediction`` (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``) under +``TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND`` instead of the standard +multi-label classification prediction dict. + +This mirrors ``v3_tensor`` but reproduces the v1 manifest exactly: the legacy +type aliases, a flat ``confidence`` float threshold (no ``confidence_mode`` / +``custom_confidence`` / validator), and v1's reduced outputs (``predictions`` + +``inference_id`` only โ€” NO ``model_id`` output, and ``inference_id`` is a plain +``STRING_KIND``). The result dicts likewise carry no ``model_id`` key. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns a + ``List[MultiLabelClassificationPrediction]`` (one per image) straight from the + adapter. Each carries ``class_ids`` (the already-threshold-filtered predicted + label ids โ€” the model's ``post_process`` applied the full priority chain, so we + do NOT re-threshold / re-build them) and ``confidence`` (the FULL sigmoid score + vector, shape ``(num_classes,)``). The block only ATTACHES the per-image + ``image_metadata`` (SINGULAR dict โ€” NOT plural; this is the key shape difference + from multi-class ``ClassificationPrediction``) that the tensor serialiser needs: + the ``class_id -> name`` map plus image dimensions and parent/root lineage. +- REMOTE: the standard inference multi-label response dict (a ``predictions`` dict + keyed by class name with per-class ``confidence``/``class_id`` plus the + ``predicted_classes`` name list) is rebuilt into a native + ``MultiLabelClassificationPrediction`` by the inline converter below. +""" + +import uuid +from time import perf_counter +from typing import Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import ( + HOSTED_CLASSIFICATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_KEY, + CLASSIFICATION_STYLE_MODEL, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.classification import ( + MultiLabelClassificationPrediction, +) +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "classification" + +LONG_DESCRIPTION = """ +Run inference on a multi-label classification model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Multi-Label Classification Model", + "version": "v1", + "short_description": "Apply multiple tags to an image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 3, + "inference": True, + }, + }, + protected_namespaces=(), + ) + type: Literal[ + "roboflow_core/roboflow_multi_label_classification_model@v1", + "RoboflowMultiLabelClassificationModel", + "MultiLabelClassificationModel", + ] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["multi-label-classification"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND], + ), + OutputDefinition(name=INFERENCE_ID_KEY, kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowMultiLabelClassificationModelBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + inference_start = perf_counter() + predictions: List[MultiLabelClassificationPrediction] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + # numpy parity: `Model.infer_from_request` stamps `time` (elapsed around + # the bare model call) on every response in the batch; mirror with one + # elapsed value for the whole batch (numpy's per-response deltas differ + # only by loop overhead). + inference_time = perf_counter() - inference_start + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, prediction in zip(images, predictions): + inference_id = str(uuid.uuid4()) + # The adapter/model already applied the full threshold chain when it + # built `class_ids`, and `confidence` is the FULL sigmoid vector. We do + # NOT re-threshold / re-softmax / rebuild `class_ids` โ€” only attach the + # (SINGULAR) image_metadata the tensor serialiser requires. + # Pin the prediction tensors to WORKFLOWS_IMAGE_TENSOR_DEVICE so LOCAL + # output matches the single-label sibling and this block's REMOTE path + # (which builds them on-device via torch.as_tensor). + prediction.class_ids = prediction.class_ids.to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + prediction.confidence = prediction.confidence.to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + prediction.image_metadata = _build_native_classification_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + inference_time=inference_time, + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + } + ) + return results + + def run_remotely( + self, + images: Batch[Optional[WorkflowImageData]], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CLASSIFICATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + confidence_threshold=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=non_empty_inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + ) -> BlockResult: + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + prediction = _native_multi_label_from_inference_response( + image=image, + response=response, + inference_id=inference_id, + inference_time=response.get("time"), + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _build_native_classification_metadata( + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + inference_time: Optional[float] = None, +) -> dict: + """Build the (SINGULAR) ``image_metadata`` dict carried by a tensor-native + ``MultiLabelClassificationPrediction``. + + Mirrors the ``build_native_image_metadata`` key conventions and the + ``vlm_as_classifier/v1_tensor`` classification pattern: the ``class_id -> + name`` map (required by the serialiser to resolve labels), the image + dimensions, and the parent/root lineage. The image shape is read without + forcing a device->host materialization so tensor-only inputs stay on device. + """ + height, width = image._read_shape_without_materialization() + metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag the serialiser reads to pick the flag-OFF shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_MODEL, + PREDICTION_TYPE_KEY: PREDICTION_TYPE, + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + ROOT_PARENT_ID_KEY: image.workflow_root_ancestor_metadata.parent_id, + } + # numpy parity: flag-off always dumps `time` โ€” stamped around the bare model + # call in `Model.infer_from_request` (LOCAL) or by the server (REMOTE). + if inference_time is not None: + metadata["time"] = inference_time + return metadata + + +def _native_multi_label_from_inference_response( + image: WorkflowImageData, + response: dict, + inference_id: str, + inference_time: Optional[float] = None, +) -> MultiLabelClassificationPrediction: + """Rebuild a native ``MultiLabelClassificationPrediction`` from the standard + inference multi-label response dict. + + The response ``predictions`` is a dict keyed by class name, each value holding + ``{"confidence": float, "class_id": int}``; ``predicted_classes`` is the list + of above-threshold class names. We reconstruct the dense (``num_classes``,) + sigmoid ``confidence`` vector indexed by ``class_id``, the ``class_ids`` of the + predicted (above-threshold) labels, and the ``class_id -> name`` map. The + server already applied the threshold chain, so ``predicted_classes`` is taken + as authoritative โ€” we do NOT re-threshold here. + """ + per_class = response.get("predictions", {}) or {} + class_names: Dict[int, str] = {} + name_to_id: Dict[str, int] = {} + confidence_by_id: Dict[int, float] = {} + for class_name, entry in per_class.items(): + class_id = int(entry["class_id"]) + class_names[class_id] = class_name + name_to_id[class_name] = class_id + confidence_by_id[class_id] = float(entry.get("confidence", 0.0)) + num_classes = (max(class_names) + 1) if class_names else 0 + confidence_vector = [confidence_by_id.get(i, 0.0) for i in range(num_classes)] + predicted_class_ids = [ + name_to_id[class_name] + for class_name in response.get("predicted_classes", []) + if class_name in name_to_id + ] + return MultiLabelClassificationPrediction( + class_ids=torch.as_tensor( + predicted_class_ids, + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.as_tensor( + confidence_vector, + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + image_metadata=_build_native_classification_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + inference_time=inference_time, + ), + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v2_tensor.py b/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v2_tensor.py new file mode 100644 index 0000000000..172f114949 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v2_tensor.py @@ -0,0 +1,452 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_multi_label_classification_model@v2`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.MultiLabelClassificationPrediction`` (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``) under +``TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND`` instead of the standard +multi-label classification prediction dict. + +This mirrors ``v3_tensor`` but reproduces the v2 manifest exactly: the single +``@v2`` type literal, a flat ``confidence`` float threshold (no +``confidence_mode`` / ``custom_confidence`` / validator), and v2's outputs +(``predictions`` + ``inference_id`` as ``INFERENCE_ID_KIND`` + ``model_id``). The +result dicts carry the ``model_id`` key as v2 does. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns a + ``List[MultiLabelClassificationPrediction]`` (one per image) straight from the + adapter. Each carries ``class_ids`` (the already-threshold-filtered predicted + label ids โ€” the model's ``post_process`` applied the full priority chain, so we + do NOT re-threshold / re-build them) and ``confidence`` (the FULL sigmoid score + vector, shape ``(num_classes,)``). The block only ATTACHES the per-image + ``image_metadata`` (SINGULAR dict โ€” NOT plural; this is the key shape difference + from multi-class ``ClassificationPrediction``) that the tensor serialiser needs: + the ``class_id -> name`` map plus image dimensions and parent/root lineage. +- REMOTE: the standard inference multi-label response dict (a ``predictions`` dict + keyed by class name with per-class ``confidence``/``class_id`` plus the + ``predicted_classes`` name list) is rebuilt into a native + ``MultiLabelClassificationPrediction`` by the inline converter below. +""" + +import uuid +from time import perf_counter +from typing import Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import ( + HOSTED_CLASSIFICATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_KEY, + CLASSIFICATION_STYLE_MODEL, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.classification import ( + MultiLabelClassificationPrediction, +) +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "classification" + +LONG_DESCRIPTION = """ +Run inference on a multi-label classification model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Multi-Label Classification Model", + "version": "v2", + "short_description": "Apply multiple tags to an image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 3, + "inference": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_multi_label_classification_model@v2"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["multi-label-classification"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND], + ), + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowMultiLabelClassificationModelBlockV2(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + inference_start = perf_counter() + predictions: List[MultiLabelClassificationPrediction] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + # numpy parity: `Model.infer_from_request` stamps `time` (elapsed around + # the bare model call) on every response in the batch; mirror with one + # elapsed value for the whole batch (numpy's per-response deltas differ + # only by loop overhead). + inference_time = perf_counter() - inference_start + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, prediction in zip(images, predictions): + inference_id = str(uuid.uuid4()) + # The adapter/model already applied the full threshold chain when it + # built `class_ids`, and `confidence` is the FULL sigmoid vector. We do + # NOT re-threshold / re-softmax / rebuild `class_ids` โ€” only attach the + # (SINGULAR) image_metadata the tensor serialiser requires. + # Pin the prediction tensors to WORKFLOWS_IMAGE_TENSOR_DEVICE so LOCAL + # output matches the single-label sibling and this block's REMOTE path + # (which builds them on-device via torch.as_tensor). + prediction.class_ids = prediction.class_ids.to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + prediction.confidence = prediction.confidence.to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + prediction.image_metadata = _build_native_classification_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + inference_time=inference_time, + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[Optional[WorkflowImageData]], + model_id: str, + confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CLASSIFICATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + confidence_threshold=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=non_empty_inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + model_id: str, + ) -> BlockResult: + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + prediction = _native_multi_label_from_inference_response( + image=image, + response=response, + inference_id=inference_id, + inference_time=response.get("time"), + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _build_native_classification_metadata( + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + inference_time: Optional[float] = None, +) -> dict: + """Build the (SINGULAR) ``image_metadata`` dict carried by a tensor-native + ``MultiLabelClassificationPrediction``. + + Mirrors the ``build_native_image_metadata`` key conventions and the + ``vlm_as_classifier/v1_tensor`` classification pattern: the ``class_id -> + name`` map (required by the serialiser to resolve labels), the image + dimensions, and the parent/root lineage. The image shape is read without + forcing a device->host materialization so tensor-only inputs stay on device. + """ + height, width = image._read_shape_without_materialization() + metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag the serialiser reads to pick the flag-OFF shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_MODEL, + PREDICTION_TYPE_KEY: PREDICTION_TYPE, + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + ROOT_PARENT_ID_KEY: image.workflow_root_ancestor_metadata.parent_id, + } + # numpy parity: flag-off always dumps `time` โ€” stamped around the bare model + # call in `Model.infer_from_request` (LOCAL) or by the server (REMOTE). + if inference_time is not None: + metadata["time"] = inference_time + return metadata + + +def _native_multi_label_from_inference_response( + image: WorkflowImageData, + response: dict, + inference_id: str, + inference_time: Optional[float] = None, +) -> MultiLabelClassificationPrediction: + """Rebuild a native ``MultiLabelClassificationPrediction`` from the standard + inference multi-label response dict. + + The response ``predictions`` is a dict keyed by class name, each value holding + ``{"confidence": float, "class_id": int}``; ``predicted_classes`` is the list + of above-threshold class names. We reconstruct the dense (``num_classes``,) + sigmoid ``confidence`` vector indexed by ``class_id``, the ``class_ids`` of the + predicted (above-threshold) labels, and the ``class_id -> name`` map. The + server already applied the threshold chain, so ``predicted_classes`` is taken + as authoritative โ€” we do NOT re-threshold here. + """ + per_class = response.get("predictions", {}) or {} + class_names: Dict[int, str] = {} + name_to_id: Dict[str, int] = {} + confidence_by_id: Dict[int, float] = {} + for class_name, entry in per_class.items(): + class_id = int(entry["class_id"]) + class_names[class_id] = class_name + name_to_id[class_name] = class_id + confidence_by_id[class_id] = float(entry.get("confidence", 0.0)) + num_classes = (max(class_names) + 1) if class_names else 0 + confidence_vector = [confidence_by_id.get(i, 0.0) for i in range(num_classes)] + predicted_class_ids = [ + name_to_id[class_name] + for class_name in response.get("predicted_classes", []) + if class_name in name_to_id + ] + return MultiLabelClassificationPrediction( + class_ids=torch.as_tensor( + predicted_class_ids, + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.as_tensor( + confidence_vector, + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + image_metadata=_build_native_classification_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + inference_time=inference_time, + ), + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v3_tensor.py b/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v3_tensor.py new file mode 100644 index 0000000000..2a6b6f0e9e --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/multi_label_classification/v3_tensor.py @@ -0,0 +1,491 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_multi_label_classification_model@v3`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.MultiLabelClassificationPrediction`` (torch tensors on +``WORKFLOWS_IMAGE_TENSOR_DEVICE``) under +``TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND`` instead of the standard +multi-label classification prediction dict. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns a + ``List[MultiLabelClassificationPrediction]`` (one per image) straight from the + adapter. Each carries ``class_ids`` (the already-threshold-filtered predicted + label ids โ€” the model's ``post_process`` applied the full priority chain, so we + do NOT re-threshold / re-build them) and ``confidence`` (the FULL sigmoid score + vector, shape ``(num_classes,)``). The block only ATTACHES the per-image + ``image_metadata`` (SINGULAR dict โ€” NOT plural; this is the key shape difference + from multi-class ``ClassificationPrediction``) that the tensor serialiser needs: + the ``class_id -> name`` map plus image dimensions and parent/root lineage. +- REMOTE: the standard inference multi-label response dict (a ``predictions`` dict + keyed by class name with per-class ``confidence``/``class_id`` plus the + ``predicted_classes`` name list) is rebuilt into a native + ``MultiLabelClassificationPrediction`` by the inline converter below. +""" + +import uuid +from time import perf_counter +from typing import Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field, model_validator + +from inference.core.env import ( + HOSTED_CLASSIFICATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + CLASSIFICATION_STYLE_KEY, + CLASSIFICATION_STYLE_MODEL, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.classification import ( + MultiLabelClassificationPrediction, +) +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "classification" + +LONG_DESCRIPTION = """ +Run inference on a multi-label classification model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Multi-Label Classification Model", + "version": "v3", + "short_description": "Apply multiple tags to an image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 3, + "inference": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_multi_label_classification_model@v3"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence_mode: Union[ + Literal["best", "default", "custom"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="best", + description="How to determine the confidence threshold.", + json_schema_extra={ + "always_visible": True, + "values_metadata": { + "best": { + "name": "Best (Recommended)", + "description": "Use F1-optimal thresholds from model evaluation.", + }, + "default": { + "name": "Default", + "description": "Use the model's built-in default threshold.", + }, + "custom": { + "name": "Custom", + "description": "Specify a custom confidence threshold.", + }, + }, + }, + ) + custom_confidence: Union[ + Optional[FloatZeroToOne], + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Custom confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + json_schema_extra={ + "relevant_for": { + "confidence_mode": { + "values": ["custom"], + "required": True, + }, + }, + }, + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @model_validator(mode="after") + def validate_custom_confidence(self) -> "BlockManifest": + if self.confidence_mode == "custom" and self.custom_confidence is None: + raise ValueError( + "custom_confidence must be provided when confidence_mode is 'custom'" + ) + return self + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["multi-label-classification"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND], + ), + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowMultiLabelClassificationModelBlockV3(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence_mode: str, + custom_confidence: Optional[float], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + confidence = ( + custom_confidence if confidence_mode == "custom" else confidence_mode + ) + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Union[None, float, Literal["best", "default"]], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + inference_start = perf_counter() + predictions: List[MultiLabelClassificationPrediction] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + # numpy parity: `Model.infer_from_request` stamps `time` (elapsed around + # the bare model call) on every response in the batch; mirror with one + # elapsed value for the whole batch (numpy's per-response deltas differ + # only by loop overhead). + inference_time = perf_counter() - inference_start + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, prediction in zip(images, predictions): + inference_id = str(uuid.uuid4()) + # The adapter/model already applied the full threshold chain when it + # built `class_ids`, and `confidence` is the FULL sigmoid vector. We do + # NOT re-threshold / re-softmax / rebuild `class_ids` โ€” only attach the + # (SINGULAR) image_metadata the tensor serialiser requires. + # Pin the prediction tensors to WORKFLOWS_IMAGE_TENSOR_DEVICE so LOCAL + # output matches the single-label sibling and this block's REMOTE path + # (which builds them on-device via torch.as_tensor). + prediction.class_ids = prediction.class_ids.to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + prediction.confidence = prediction.confidence.to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + prediction.image_metadata = _build_native_classification_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + inference_time=inference_time, + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[Optional[WorkflowImageData]], + model_id: str, + confidence: Union[None, float, Literal["best", "default"]], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_CLASSIFICATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + confidence_threshold=confidence, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=non_empty_inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + model_id: str, + ) -> BlockResult: + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + prediction = _native_multi_label_from_inference_response( + image=image, + response=response, + inference_id=inference_id, + inference_time=response.get("time"), + ) + results.append( + { + "inference_id": inference_id, + "predictions": prediction, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _build_native_classification_metadata( + image: WorkflowImageData, + class_names: Dict[int, str], + inference_id: str, + inference_time: Optional[float] = None, +) -> dict: + """Build the (SINGULAR) ``image_metadata`` dict carried by a tensor-native + ``MultiLabelClassificationPrediction``. + + Mirrors the ``build_native_image_metadata`` key conventions and the + ``vlm_as_classifier/v1_tensor`` classification pattern: the ``class_id -> + name`` map (required by the serialiser to resolve labels), the image + dimensions, and the parent/root lineage. The image shape is read without + forcing a device->host materialization so tensor-only inputs stay on device. + """ + height, width = image._read_shape_without_materialization() + metadata = { + CLASS_NAMES_KEY: class_names, + # Lane 1b: explicit style tag the serialiser reads to pick the flag-OFF shape. + CLASSIFICATION_STYLE_KEY: CLASSIFICATION_STYLE_MODEL, + PREDICTION_TYPE_KEY: PREDICTION_TYPE, + IMAGE_DIMENSIONS_KEY: [height, width], + INFERENCE_ID_KEY: inference_id, + PARENT_ID_KEY: image.parent_metadata.parent_id, + ROOT_PARENT_ID_KEY: image.workflow_root_ancestor_metadata.parent_id, + } + # numpy parity: flag-off always dumps `time` โ€” stamped around the bare model + # call in `Model.infer_from_request` (LOCAL) or by the server (REMOTE). + if inference_time is not None: + metadata["time"] = inference_time + return metadata + + +def _native_multi_label_from_inference_response( + image: WorkflowImageData, + response: dict, + inference_id: str, + inference_time: Optional[float] = None, +) -> MultiLabelClassificationPrediction: + """Rebuild a native ``MultiLabelClassificationPrediction`` from the standard + inference multi-label response dict. + + The response ``predictions`` is a dict keyed by class name, each value holding + ``{"confidence": float, "class_id": int}``; ``predicted_classes`` is the list + of above-threshold class names. We reconstruct the dense (``num_classes``,) + sigmoid ``confidence`` vector indexed by ``class_id``, the ``class_ids`` of the + predicted (above-threshold) labels, and the ``class_id -> name`` map. The + server already applied the threshold chain, so ``predicted_classes`` is taken + as authoritative โ€” we do NOT re-threshold here. + """ + per_class = response.get("predictions", {}) or {} + class_names: Dict[int, str] = {} + name_to_id: Dict[str, int] = {} + confidence_by_id: Dict[int, float] = {} + for class_name, entry in per_class.items(): + class_id = int(entry["class_id"]) + class_names[class_id] = class_name + name_to_id[class_name] = class_id + confidence_by_id[class_id] = float(entry.get("confidence", 0.0)) + num_classes = (max(class_names) + 1) if class_names else 0 + confidence_vector = [confidence_by_id.get(i, 0.0) for i in range(num_classes)] + predicted_class_ids = [ + name_to_id[class_name] + for class_name in response.get("predicted_classes", []) + if class_name in name_to_id + ] + return MultiLabelClassificationPrediction( + class_ids=torch.as_tensor( + predicted_class_ids, + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.as_tensor( + confidence_vector, + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + image_metadata=_build_native_classification_metadata( + image=image, + class_names=class_names, + inference_id=inference_id, + inference_time=inference_time, + ), + ) diff --git a/inference/core/workflows/core_steps/models/roboflow/object_detection/v1_tensor.py b/inference/core/workflows/core_steps/models/roboflow/object_detection/v1_tensor.py new file mode 100644 index 0000000000..d6dccf687d --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/object_detection/v1_tensor.py @@ -0,0 +1,438 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_object_detection_model@v1`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.Detections`` (torch tensors on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``) +under ``TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND`` instead of ``sv.Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns ``List[Detections]`` + straight from the adapter (xyxy / class_id / confidence only). The block applies + ``class_filter`` natively (the adapter/model does NOT read it on this path) and + attaches the producer contract (``image_metadata[class_names]`` + per-box + ``detection_id``) the tensor serialiser requires, via ``attach_native_detection_metadata``. +- REMOTE: standard inference prediction dicts are rebuilt into a native ``Detections`` + via ``native_detections_from_inference_predictions`` (never ``sv.Detections``). + +Manifest (type@version, class name, params, validators, compatibility) mirrors the +numpy v1 sibling exactly; only the ``predictions`` output kind and the run bodies +differ. v1 carries the legacy ``type`` aliases and declares only +``inference_id``/``predictions`` outputs (no ``model_id``). +""" + +import uuid +from typing import Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field, PositiveInt + +from inference.core.env import ( + HOSTED_DETECT_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + native_detections_from_inference_predictions, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import INFERENCE_ID_KEY +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.object_detection import Detections +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "object-detection" + +LONG_DESCRIPTION = """ +Run inference on a object-detection model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Object Detection Model", + "version": "v1", + "short_description": "Predict the location of objects with bounding boxes.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo", "rfdetr", "rf-detr"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 0, + "inference": True, + "popular": True, + }, + }, + protected_namespaces=(), + ) + type: Literal[ + "roboflow_core/roboflow_object_detection_model@v1", + "RoboflowObjectDetectionModel", + "ObjectDetectionModel", + ] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[Optional[bool], Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["object-detection"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="inference_id", kind=[STRING_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowObjectDetectionModelBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + detections_batch: List[Detections] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, detections in zip(images, detections_batch): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). + detections = _filter_classes_native(detections, class_filter, class_names) + # Reuse the adapter-provided inference id when present (numpy parity: + # numpy v1 surfaces ``p.get(INFERENCE_ID_KEY)`` off the model dump); + # the tensor-native adapter normally attaches none, so fall back to a + # freshly minted uuid that is then shared with ``image_metadata``. + inference_id = getattr(detections, "inference_id", None) or str( + uuid.uuid4() + ) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_DETECT_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=non_empty_inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + ) -> BlockResult: + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) + if class_filter: + detection_dicts = [ + d for d in detection_dicts if d.get("class") in class_filter + ] + detections = native_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + detections: Detections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> Detections: + if not class_filter: + return detections + accepted = set(class_filter) + accepted_ids = sorted( + class_id for class_id, name in class_names.items() if name in accepted + ) + if not accepted_ids: + return take_prediction_by_mask( + detections, + torch.zeros_like(detections.class_id, dtype=torch.bool), + ) + accepted_tensor = torch.as_tensor(accepted_ids, device=detections.class_id.device) + keep = torch.isin(detections.class_id, accepted_tensor) + return take_prediction_by_mask(detections, keep) diff --git a/inference/core/workflows/core_steps/models/roboflow/object_detection/v2_tensor.py b/inference/core/workflows/core_steps/models/roboflow/object_detection/v2_tensor.py new file mode 100644 index 0000000000..6f3bf2aef3 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/object_detection/v2_tensor.py @@ -0,0 +1,438 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_object_detection_model@v2`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.Detections`` (torch tensors on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``) +under ``TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND`` instead of ``sv.Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns ``List[Detections]`` + straight from the adapter (xyxy / class_id / confidence only). The block applies + ``class_filter`` natively (the adapter/model does NOT read it on this path) and + attaches the producer contract (``image_metadata[class_names]`` + per-box + ``detection_id``) the tensor serialiser requires, via ``attach_native_detection_metadata``. +- REMOTE: standard inference prediction dicts are rebuilt into a native ``Detections`` + via ``native_detections_from_inference_predictions`` (never ``sv.Detections``). + +Manifest (type@version, class name, params, validators, compatibility) mirrors the +numpy v2 sibling exactly; only the ``predictions`` output kind and the run bodies +differ. v2 declares ``inference_id``/``predictions``/``model_id`` outputs. +""" + +import uuid +from typing import Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field, PositiveInt + +from inference.core.env import ( + HOSTED_DETECT_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + native_detections_from_inference_predictions, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import INFERENCE_ID_KEY +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.object_detection import Detections +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "object-detection" + +LONG_DESCRIPTION = """ +Run inference on a object-detection model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Object Detection Model", + "version": "v2", + "short_description": "Predict the location of objects with bounding boxes.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo", "rfdetr", "rf-detr"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 0, + "inference": True, + "popular": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_object_detection_model@v2"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[Optional[bool], Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["object-detection"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="inference_id", kind=[INFERENCE_ID_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowObjectDetectionModelBlockV2(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + detections_batch: List[Detections] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, detections in zip(images, detections_batch): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). + detections = _filter_classes_native(detections, class_filter, class_names) + # Reuse the adapter-provided inference id when present (numpy parity: + # numpy v2 surfaces ``p.get(INFERENCE_ID_KEY)`` off the model dump); + # the tensor-native adapter normally attaches none, so fall back to a + # freshly minted uuid that is then shared with ``image_metadata``. + inference_id = getattr(detections, "inference_id", None) or str( + uuid.uuid4() + ) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Optional[float], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_DETECT_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=non_empty_inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + model_id: str, + ) -> BlockResult: + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) + if class_filter: + detection_dicts = [ + d for d in detection_dicts if d.get("class") in class_filter + ] + detections = native_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + detections: Detections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> Detections: + if not class_filter: + return detections + accepted = set(class_filter) + accepted_ids = sorted( + class_id for class_id, name in class_names.items() if name in accepted + ) + if not accepted_ids: + return take_prediction_by_mask( + detections, + torch.zeros_like(detections.class_id, dtype=torch.bool), + ) + accepted_tensor = torch.as_tensor(accepted_ids, device=detections.class_id.device) + keep = torch.isin(detections.class_id, accepted_tensor) + return take_prediction_by_mask(detections, keep) diff --git a/inference/core/workflows/core_steps/models/roboflow/object_detection/v3_tensor.py b/inference/core/workflows/core_steps/models/roboflow/object_detection/v3_tensor.py new file mode 100644 index 0000000000..a8a01fb748 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/object_detection/v3_tensor.py @@ -0,0 +1,477 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_object_detection_model@v3`. + +Under ENABLE_TENSOR_DATA_REPRESENTATION this block emits a native +``inference_models.Detections`` (torch tensors on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``) +under ``TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND`` instead of ``sv.Detections``. + +- LOCAL: ``ModelManager.run_tensor_native_inference`` returns ``List[Detections]`` + straight from the adapter (xyxy / class_id / confidence only). The block applies + ``class_filter`` natively (the adapter/model does NOT read it on this path) and + attaches the producer contract (``image_metadata[class_names]`` + per-box + ``detection_id``) the tensor serialiser requires, via ``attach_native_detection_metadata``. +- REMOTE: standard inference prediction dicts are rebuilt into a native ``Detections`` + via ``native_detections_from_inference_predictions`` (never ``sv.Detections``). +""" + +import uuid +from typing import Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field, PositiveInt, model_validator + +from inference.core.env import ( + HOSTED_DETECT_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + native_detections_from_inference_predictions, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import INFERENCE_ID_KEY +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_MODEL_ID_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, + roboflow_platform_project, +) +from inference_models.models.base.object_detection import Detections +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "object-detection" + +LONG_DESCRIPTION = """ +Run inference on a object-detection model hosted on or uploaded to Roboflow. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Object Detection Model", + "version": "v3", + "short_description": "Predict the location of objects with bounding boxes.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["yolo", "rfdetr", "rf-detr"], + "ui_manifest": { + "section": "model", + "icon": "far fa-chart-network", + "blockPriority": 0, + "inference": True, + "popular": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_object_detection_model@v3"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence_mode: Union[ + Literal["best", "default", "custom"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="best", + description="How confidence thresholds are determined.", + json_schema_extra={ + "always_visible": True, + "values_metadata": { + "best": { + "name": "Best (Recommended)", + "description": "Use F1-optimal thresholds from model evaluation.", + }, + "default": { + "name": "Default", + "description": "Use the model's built-in default threshold.", + }, + "custom": { + "name": "Custom", + "description": "Specify a custom confidence threshold.", + }, + }, + }, + ) + custom_confidence: Union[ + Optional[FloatZeroToOne], + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Custom confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + json_schema_extra={ + "relevant_for": { + "confidence_mode": {"values": ["custom"], "required": True}, + }, + }, + ) + class_filter: Union[Optional[List[str]], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( + default=None, + description="List of accepted classes. Classes must exist in the model's training set.", + examples=[["a", "b", "c"], "$inputs.class_filter"], + ) + ) + iou_threshold: Union[ + FloatZeroToOne, + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.3, + description="Minimum overlap threshold between boxes to combine them into a single detection, used in NMS. [Learn more](https://blog.roboflow.com/how-to-code-non-maximum-suppression-nms-in-plain-numpy/).", + examples=[0.4, "$inputs.iou_threshold"], + ) + max_detections: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=300, + description="Maximum number of detections to return.", + examples=[300, "$inputs.max_detections"], + ) + class_agnostic_nms: Union[Optional[bool], Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Boolean flag to specify if NMS is to be used in class-agnostic mode.", + examples=[True, "$inputs.class_agnostic_nms"], + ) + max_candidates: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + default=3000, + description="Maximum number of candidates as NMS input to be taken into account.", + examples=[3000, "$inputs.max_candidates"], + ) + disable_active_learning: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Boolean flag to disable project-level active learning for this block.", + examples=[True, "$inputs.disable_active_learning"], + ) + active_learning_target_dataset: Union[ + Selector(kind=[ROBOFLOW_PROJECT_KIND]), Optional[str] + ] = Field( + default=None, + description="Target dataset for active learning, if enabled.", + examples=["my_project", "$inputs.al_target_project"], + ) + + @model_validator(mode="after") + def validate(self) -> "BlockManifest": + if self.confidence_mode == "custom" and self.custom_confidence is None: + raise ValueError( + "`custom_confidence` is required when `confidence_mode` is 'custom'" + ) + return self + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["object-detection"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + resources = [roboflow_platform_model(model_id=self.model_id)] + if self.disable_active_learning is True: + # Active learning literally disabled (the default) โ€” the target + # project is dead configuration, not a dependency. + return resources + if self.active_learning_target_dataset is not None: + resources.append( + roboflow_platform_project( + project_url=self.active_learning_target_dataset + ) + ) + return resources + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="inference_id", kind=[INFERENCE_ID_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowObjectDetectionModelBlockV3(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence_mode: str, + custom_confidence: Optional[float], + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + confidence = ( + custom_confidence if confidence_mode == "custom" else confidence_mode + ) + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, + model_id=model_id, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Union[None, float, Literal["best", "default"]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching what the numpy + # block hands the model. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + detections_batch: List[Detections] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + iou_threshold=iou_threshold, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + max_detections=max_detections, + max_candidates=max_candidates, + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + ) + ) + class_names = _class_names_map(self._model_manager.get_class_names(model_id)) + results: List[dict] = [] + for image, detections in zip(images, detections_batch): + # The adapter/model does NOT honour class_filter on the native path, so + # filter here (the only place it is applied for LOCAL execution). + detections = _filter_classes_native(detections, class_filter, class_names) + # Reuse the adapter-provided inference id when present (numpy parity: + # numpy v3 surfaces ``p.get(INFERENCE_ID_KEY)`` off the model dump); + # the tensor-native adapter normally attaches none, so fall back to a + # freshly minted uuid that is then shared with ``image_metadata``. + inference_id = getattr(detections, "inference_id", None) or str( + uuid.uuid4() + ) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names=class_names, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + class_agnostic_nms: Optional[bool], + class_filter: Optional[List[str]], + confidence: Union[None, float, Literal["best", "default"]], + iou_threshold: Optional[float], + max_detections: Optional[int], + max_candidates: Optional[int], + disable_active_learning: Optional[bool], + active_learning_target_dataset: Optional[str], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_DETECT_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + disable_active_learning=disable_active_learning, + active_learning_target_dataset=active_learning_target_dataset, + class_agnostic_nms=class_agnostic_nms, + class_filter=class_filter, + confidence_threshold=confidence, + iou_threshold=iou_threshold, + max_detections=max_detections, + max_candidates=max_candidates, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + non_empty_inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=non_empty_inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, + predictions=predictions, + class_filter=class_filter, + model_id=model_id, + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + class_filter: Optional[List[str]], + model_id: str, + ) -> BlockResult: + results: List[dict] = [] + for image, response in zip(images, predictions): + inference_id = response.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detection_dicts = response.get("predictions", []) + if class_filter: + detection_dicts = [ + d for d in detection_dicts if d.get("class") in class_filter + ] + detections = native_detections_from_inference_predictions( + image=image, + predictions=detection_dicts, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + results.append( + { + "inference_id": inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _filter_classes_native( + detections: Detections, + class_filter: Optional[List[str]], + class_names: Dict[int, str], +) -> Detections: + if not class_filter: + return detections + accepted = set(class_filter) + accepted_ids = sorted( + class_id for class_id, name in class_names.items() if name in accepted + ) + if not accepted_ids: + return take_prediction_by_mask( + detections, + torch.zeros_like(detections.class_id, dtype=torch.bool), + ) + accepted_tensor = torch.as_tensor(accepted_ids, device=detections.class_id.device) + keep = torch.isin(detections.class_id, accepted_tensor) + return take_prediction_by_mask(detections, keep) diff --git a/inference/core/workflows/core_steps/models/roboflow/semantic_segmentation/v1_tensor.py b/inference/core/workflows/core_steps/models/roboflow/semantic_segmentation/v1_tensor.py new file mode 100644 index 0000000000..b885d31137 --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/semantic_segmentation/v1_tensor.py @@ -0,0 +1,616 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_semantic_segmentation_model@v1`. + +Numpy `semantic_segmentation/v1.py` runs the model (LOCAL: via +`SemanticSegmentationInferenceRequest` + `infer_from_request_sync`; REMOTE: via the +HTTP client) and converts each *dense* per-pixel response into an ``sv.Detections`` +carrying one COCO-RLE mask per class (under ``data['rle_mask']``). Under +ENABLE_TENSOR_DATA_REPRESENTATION this sibling must instead emit a native +``inference_models`` prediction so the shared tensor serialiser/consumers work +unchanged. + +DESIGN (Option B - per-class RLE detections, no dense carrier kind): +The block emits an ``inference_models.InstanceDetections`` (the same carrier the +instance-segmentation / seg_preview tensor blocks emit) with ONE instance per +non-background class id present in the segmentation map. Each instance carries: + * ``xyxy`` - tight bbox enclosing all pixels of the class, + * ``class_id`` - the class id, + * ``confidence`` - mean per-pixel confidence over that class's pixels, + * a COCO RLE in ``mask`` (``InstancesRLEMasks``), encoded one class at a time + (``torch_mask_to_coco_rle`` / ``pycocotools``) so a giant ``N x H x W`` dense + stack is never allocated. +This reuses the exact RLE approach of ``semantic_segmentation/v2_tensor.py``. + +KEY DIFFERENCE FROM v2_tensor (manifest only): +The v1 manifest has NO ``confidence_mode`` / ``custom_confidence`` params (and no +``validate`` validator); ``run`` takes only ``images`` + ``model_id``. Otherwise the +flow mirrors v2_tensor: LOCAL calls ``run_tensor_native_inference`` (-> +``SemanticSegmentationResult``) and builds per-class RLE detections via +``_build_instance_detections_from_segmentation``; REMOTE converts the HTTP +``segmentation_mask`` / ``confidence_mask`` / ``class_map`` response via +``_build_instance_detections_from_inference_response``. + +Because the produced object is an ``InstanceDetections``, the existing tensor +serialiser ``serializers_tensor.py::serialise_sv_detections`` handles it unchanged. + +BACKGROUND / IGNORE (numpy parity - D3: drop ONLY class id 0): +Numpy ``v1.py`` (line 231) filters the present class ids with ``if cid != 0`` and +keeps every other id as a detection - INCLUDING the ``255`` ignore/no-class +sentinel, which becomes a ``'255'`` detection named ``class_map.get('255', '255')``. +The flag-ON emitted-detection SET must match flag-OFF exactly (D3), so this sibling +replicates that filter verbatim: exclude id 0 only, keep everything else (255 +included), and name any id missing from the class map ``str(class_id)``. + +NOTE: this deliberately drops the earlier data-driven "background-by-name + 255" +exclusion. That rule was arguably more correct (it dropped ``background`` by name +even when not at id 0, and never emitted the 255 sentinel), but it diverged from +the numpy block's emitted set. D3 chooses byte-parity: drop id 0, keep 255. +Confidence VALUES may still differ (numpy quantises through a uint8/255 PNG mask, +the tensor path keeps full-precision float means) - that precision delta is accepted +under D3 and is NOT reconciled here. + +LOADER DELTAS (return-only - do NOT apply here): + +(a) serializer registration (``KINDS_SERIALIZERS``, loader.py ~1139) - the kind + ``semantic_segmentation_prediction`` is currently mapped to the NUMPY sv + serialiser ``serialise_rle_sv_detections``, which crashes on a native + ``InstanceDetections``. Under the tensor flag it must point at the tensor + serialiser ``serialise_sv_detections`` (already imported from + ``serializers_tensor`` when ``ENABLE_TENSOR_DATA_REPRESENTATION`` is set). This + is the SAME delta v2_tensor documents (the kind is shared by both versions): + + # loader.py KINDS_SERIALIZERS (inside the dict, ~line 1139) + SEMANTIC_SEGMENTATION_PREDICTION_KIND.name: ( + serialise_sv_detections + if ENABLE_TENSOR_DATA_REPRESENTATION + else serialise_rle_sv_detections + ), + +(b) block import-swap (loader.py, mirroring the object_detection@v3 swap) - import + the v1 tensor sibling under the flag: + + if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v1_tensor import ( + RoboflowSemanticSegmentationModelBlockV1, + ) + else: + from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v1 import ( + RoboflowSemanticSegmentationModelBlockV1, + ) +""" + +import base64 +import uuid +from typing import Dict, List, Literal, Optional, Type, Union + +import cv2 +import numpy as np +import torch +from pydantic import ConfigDict + +from inference.core.env import ( + HOSTED_SEMANTIC_SEGMENTATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + DETECTION_ID_KEY, + INFERENCE_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_SEMANTIC_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + INFERENCE_ID_KIND, + ROBOFLOW_MODEL_ID_KIND, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.semantic_segmentation import ( + SemanticSegmentationResult, +) +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import torch_mask_to_coco_rle +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "semantic-segmentation" + +# D3 (numpy parity): numpy `v1.py:231` / `v2.py:292` drop ONLY class id 0 +# (`if cid != 0`) and keep every other present id as a detection โ€” including the +# 255 ignore/no-class sentinel. The flag-ON emitted set must match numpy exactly, +# so id 0 is the single excluded id. (This intentionally does NOT do the +# data-driven background-by-name / 255 exclusion; see the module docstring.) +BACKGROUND_CLASS_ID = 0 + + +LONG_DESCRIPTION = """ +Run inference on a semantic segmentation model hosted on or uploaded to Roboflow. + +Semantic segmentation assigns a class label to every pixel in the image, producing a +dense segmentation mask rather than per-object bounding boxes or instance masks. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Semantic Segmentation Model", + "version": "v1", + "short_description": "Assign a class label to every pixel in the image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["semantic", "segmentation", "deeplab", "deep_lab"], + "ui_manifest": { + "section": "model", + "icon": "far fa-paint-brush", + "blockPriority": 3, + "inference": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_semantic_segmentation_model@v1"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["semantic-segmentation"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + return [roboflow_platform_model(model_id=self.model_id)] + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_SEMANTIC_SEGMENTATION_PREDICTION_KIND], + ), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowSemanticSegmentationModelBlockV1(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + ) -> BlockResult: + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally(images=images, model_id=model_id) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely(images=images, model_id=model_id) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + segmentation_results: List[SemanticSegmentationResult] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + ) + ) + class_names_map = _class_names_map( + list(self._model_manager.get_class_names(model_id)) + ) + excluded_ids = _excluded_class_ids() + results: List[dict] = [] + for image, segmentation in zip(images, segmentation_results): + inference_id = str(uuid.uuid4()) + detections = _build_instance_detections_from_segmentation( + segmentation_map=segmentation.segmentation_map, + confidence=segmentation.confidence, + image=image, + class_names_map=class_names_map, + excluded_ids=excluded_ids, + inference_id=inference_id, + ) + results.append( + { + INFERENCE_ID_KEY: inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_SEMANTIC_SEGMENTATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_result( + images=images, predictions=predictions, model_id=model_id + ) + + def _post_process_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + model_id: str, + ) -> BlockResult: + results: List[dict] = [] + for image, prediction in zip(images, predictions): + inference_id = prediction.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detections = _build_instance_detections_from_inference_response( + predictions_dict=prediction.get("predictions") or {}, + image=image, + inference_id=inference_id, + ) + results.append( + { + INFERENCE_ID_KEY: prediction.get(INFERENCE_ID_KEY), + "predictions": detections, + "model_id": model_id, + } + ) + return results + + +def _class_names_map_from_class_map(class_map: Dict[str, str]) -> Dict[int, str]: + return {int(k): v for k, v in class_map.items()} + + +def _excluded_class_ids() -> set: + """Class ids that must never become detections. To match numpy exactly (D3), + this is ONLY id 0: numpy ``v1.py``/``v2.py`` filter ``cid != 0`` and keep every + other present id โ€” the 255 ignore sentinel included.""" + return {BACKGROUND_CLASS_ID} + + +def _empty_instance_detections( + image: WorkflowImageData, + class_names_map: Dict[int, str], + height: int, + width: int, + inference_id: str, +) -> InstanceDetections: + detections = InstanceDetections( + xyxy=torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + class_id=torch.zeros( + (0,), dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.zeros( + (0,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + mask=InstancesRLEMasks(image_size=(height, width), masks=[]), + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=class_names_map, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + detections.bboxes_metadata = None + return detections + + +def _assemble_instance_detections( + xyxy: List[List[float]], + class_ids: List[int], + confidences: List[float], + rle_dicts: List[dict], + bboxes_metadata: List[dict], + image: WorkflowImageData, + class_names_map: Dict[int, str], + height: int, + width: int, + inference_id: str, +) -> InstanceDetections: + if len(rle_dicts) == 0: + return _empty_instance_detections( + image=image, + class_names_map=class_names_map, + height=height, + width=width, + inference_id=inference_id, + ) + detections = InstanceDetections( + xyxy=torch.tensor( + xyxy, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + class_id=torch.tensor( + class_ids, dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.tensor( + confidences, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + mask=InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), masks=rle_dicts + ), + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=class_names_map, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + detections.bboxes_metadata = bboxes_metadata + return detections + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _bbox_from_binary_mask(binary_mask: torch.Tensor) -> Optional[List[float]]: + """Tight ``[x_min, y_min, x_max, y_max]`` enclosing all True pixels of a + ``(H, W)`` torch bool mask, or ``None`` if empty (inclusive max, as numpy v1).""" + rows = torch.where(torch.any(binary_mask, dim=1))[0] + cols = torch.where(torch.any(binary_mask, dim=0))[0] + if rows.numel() == 0 or cols.numel() == 0: + return None + return [ + float(cols[0].item()), + float(rows[0].item()), + float(cols[-1].item()), + float(rows[-1].item()), + ] + + +def _build_instance_detections_from_segmentation( + segmentation_map: torch.Tensor, + confidence: torch.Tensor, + image: WorkflowImageData, + class_names_map: Dict[int, str], + excluded_ids: set, + inference_id: str, +) -> InstanceDetections: + """LOCAL path: dense ``(H, W)`` class grid + ``(H, W)`` confidence -> one RLE + instance per present non-background/non-ignore class. Masks are encoded one + class at a time (``torch_mask_to_coco_rle``); no ``N x H x W`` dense stack is + materialised. Mirrors ``semantic_segmentation/v2_tensor.py``.""" + height, width = int(segmentation_map.shape[0]), int(segmentation_map.shape[1]) + xyxy: List[List[float]] = [] + class_ids: List[int] = [] + confidences: List[float] = [] + bboxes_metadata: List[dict] = [] + rle_dicts: List[dict] = [] + present_ids = [int(value) for value in torch.unique(segmentation_map).tolist()] + for class_id in present_ids: + if class_id in excluded_ids: + continue + # numpy names any present id via `class_map.get(str(cid), str(cid))`; ids + # absent from the class map (e.g. the 255 sentinel) fall back to str(id) + # and are STILL emitted โ€” matching numpy, which never skips on name. + class_name = class_names_map.get(class_id, str(class_id)) + binary_mask = segmentation_map == class_id + bbox = _bbox_from_binary_mask(binary_mask) + if bbox is None: + continue + class_confidence = float(confidence[binary_mask].mean().item()) + rle = torch_mask_to_coco_rle(binary_mask.to(torch.uint8)) + rle = _normalise_rle_counts(rle) + xyxy.append(bbox) + class_ids.append(class_id) + confidences.append(class_confidence) + rle_dicts.append(rle) + bboxes_metadata.append( + {DETECTION_ID_KEY: str(uuid.uuid4()), CLASS_NAME_KEY: class_name} + ) + return _assemble_instance_detections( + xyxy=xyxy, + class_ids=class_ids, + confidences=confidences, + rle_dicts=rle_dicts, + bboxes_metadata=bboxes_metadata, + image=image, + class_names_map=class_names_map, + height=height, + width=width, + inference_id=inference_id, + ) + + +def _build_instance_detections_from_inference_response( + predictions_dict: Dict, + image: WorkflowImageData, + inference_id: str, +) -> InstanceDetections: + """Standard inference semantic-seg response (a single ``dict``, matching numpy + ``v1.py``'s ``_convert_to_sv_detections`` input - produced identically by v1's + LOCAL ``infer_from_request_sync`` dump and its REMOTE HTTP client) -> + one RLE instance per present non-background/non-ignore class. + + Response shape (see numpy ``v1.py``): + * ``segmentation_mask`` - base64 PNG, grayscale, pixel value == class id + * ``confidence_mask`` - base64 PNG, grayscale, 0..255 (== confidence*255) + * ``class_map`` - ``{str(class_id): class_name}`` + """ + seg_mask_b64 = predictions_dict.get("segmentation_mask", "") + conf_mask_b64 = predictions_dict.get("confidence_mask", "") + class_map: Dict[str, str] = predictions_dict.get("class_map", {}) + class_names_map = _class_names_map_from_class_map(class_map) + + height, width = image._read_shape_without_materialization() + + mask_array = _decode_b64_grayscale_png(seg_mask_b64) + if mask_array is None: + return _empty_instance_detections( + image=image, + class_names_map=class_names_map, + height=int(height), + width=int(width), + inference_id=inference_id, + ) + height, width = int(mask_array.shape[0]), int(mask_array.shape[1]) + conf_array = _decode_b64_grayscale_png(conf_mask_b64) + + excluded_ids = _excluded_class_ids() + xyxy: List[List[float]] = [] + class_ids: List[int] = [] + confidences: List[float] = [] + bboxes_metadata: List[dict] = [] + rle_dicts: List[dict] = [] + + present_ids = [int(value) for value in np.unique(mask_array).tolist()] + for class_id in present_ids: + if class_id in excluded_ids: + continue + class_name = class_names_map.get(class_id, str(class_id)) + binary_mask = mask_array == class_id + bbox = _bbox_from_numpy_mask(binary_mask) + if bbox is None: + continue + if conf_array is not None: + class_confidence = float(conf_array[binary_mask].mean()) / 255.0 + else: + class_confidence = 1.0 + mask_tensor = torch.from_numpy(binary_mask.astype(np.uint8)).to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + rle = torch_mask_to_coco_rle(mask_tensor) + rle = _normalise_rle_counts(rle) + xyxy.append(bbox) + class_ids.append(class_id) + confidences.append(class_confidence) + rle_dicts.append(rle) + bboxes_metadata.append( + {DETECTION_ID_KEY: str(uuid.uuid4()), CLASS_NAME_KEY: class_name} + ) + + return _assemble_instance_detections( + xyxy=xyxy, + class_ids=class_ids, + confidences=confidences, + rle_dicts=rle_dicts, + bboxes_metadata=bboxes_metadata, + image=image, + class_names_map=class_names_map, + height=height, + width=width, + inference_id=inference_id, + ) + + +def _bbox_from_numpy_mask(binary_mask: np.ndarray) -> Optional[List[float]]: + """Tight ``[x_min, y_min, x_max, y_max]`` enclosing all True pixels of a + ``(H, W)`` numpy bool mask, or ``None`` if empty (mirrors numpy ``v1.py``).""" + rows = np.where(np.any(binary_mask, axis=1))[0] + cols = np.where(np.any(binary_mask, axis=0))[0] + if rows.size == 0 or cols.size == 0: + return None + return [float(cols[0]), float(rows[0]), float(cols[-1]), float(rows[-1])] + + +def _decode_b64_grayscale_png(b64_value: str) -> Optional[np.ndarray]: + if not b64_value: + return None + mask_bytes = base64.b64decode(b64_value) + nparr = np.frombuffer(mask_bytes, np.uint8) + return cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE) + + +def _normalise_rle_counts(rle: dict) -> dict: + """``InstancesRLEMasks.from_coco_rle_masks`` keeps the raw ``counts`` payload + as-is; ``pycocotools`` returns it as ``bytes``. Decode to ``str`` so the + serialised RLE matches the rest of the tensor pipeline (and the numpy block, + which also stores ``counts`` as a utf-8 string).""" + counts = rle.get("counts") + if isinstance(counts, bytes): + rle = dict(rle) + rle["counts"] = counts.decode("utf-8") + return rle diff --git a/inference/core/workflows/core_steps/models/roboflow/semantic_segmentation/v2_tensor.py b/inference/core/workflows/core_steps/models/roboflow/semantic_segmentation/v2_tensor.py new file mode 100644 index 0000000000..7e8548686c --- /dev/null +++ b/inference/core/workflows/core_steps/models/roboflow/semantic_segmentation/v2_tensor.py @@ -0,0 +1,699 @@ +"""Tensor-native sibling of `roboflow_core/roboflow_semantic_segmentation_model@v2`. + +Numpy `semantic_segmentation/v2.py` produces a *dense* per-pixel class grid and +serialises it as an ``sv.Detections`` carrying one COCO-RLE mask per class (under +``data['rle_mask']``). Under ENABLE_TENSOR_DATA_REPRESENTATION this sibling must +instead emit a native ``inference_models`` prediction so the shared tensor +serialiser/consumers work unchanged. + +DESIGN (Option B - per-class RLE detections, no dense carrier kind): +The block emits an ``inference_models.InstanceDetections`` (the same carrier the +instance-segmentation / seg_preview tensor blocks emit) with ONE instance per +non-background class id present in the segmentation map. Each instance carries: + * ``xyxy`` - tight bbox enclosing all pixels of the class, + * ``class_id`` - the class id, + * ``confidence`` - mean per-pixel confidence over that class's pixels, + * a COCO RLE in ``mask`` (``InstancesRLEMasks``), encoded one class at a time + (``torch_mask_to_coco_rle`` / ``pycocotools``) so a giant ``N x H x W`` dense + stack is never allocated. +This reuses the exact RLE approach of ``foundation/seg_preview/v1_tensor.py``. + +Because the produced object is an ``InstanceDetections``, the existing tensor +serialiser ``serializers_tensor.py::serialise_sv_detections`` handles it +unchanged. (The serialiser turns each RLE instance into a polygon list - that is +the standard tensor-native instance-segmentation serialisation; see the loader +delta in this module's docstring tail.) + +BACKGROUND / IGNORE (numpy parity - D3: drop ONLY class id 0): +Numpy ``v2.py`` (line 292) filters the present class ids with ``if cid != 0`` and +keeps every other id as a detection - INCLUDING the ``255`` ignore/no-class +sentinel, which becomes a ``'255'`` detection named ``class_map.get('255', '255')``. +The flag-ON emitted-detection SET must match flag-OFF exactly (D3), so this sibling +replicates that filter verbatim: exclude id 0 only, keep everything else (255 +included), and name any id missing from the class map ``str(class_id)``. + +NOTE: this deliberately drops the earlier data-driven "background-by-name + 255" +exclusion. That rule was arguably more correct (it dropped ``background`` by name +even when not at id 0, and never emitted the 255 sentinel), but it diverged from +the numpy block's emitted set. D3 chooses byte-parity: drop id 0, keep 255. +Confidence VALUES may still differ (numpy quantises through a uint8/255 PNG mask, +the tensor path keeps full-precision float means) - that precision delta is accepted +under D3 and is NOT reconciled here. + +LOADER DELTAS (return-only - do NOT apply here): + +(a) serializer registration (``KINDS_SERIALIZERS``, loader.py ~1139) - the kind + ``semantic_segmentation_prediction`` is currently mapped to the NUMPY sv + serialiser ``serialise_rle_sv_detections``, which crashes on a native + ``InstanceDetections``. Under the tensor flag it must point at the tensor + serialiser ``serialise_sv_detections`` (already imported from + ``serializers_tensor`` when ``ENABLE_TENSOR_DATA_REPRESENTATION`` is set): + + # loader.py KINDS_SERIALIZERS (inside the dict, ~line 1139) + SEMANTIC_SEGMENTATION_PREDICTION_KIND.name: ( + serialise_sv_detections + if ENABLE_TENSOR_DATA_REPRESENTATION + else serialise_rle_sv_detections + ), + +(b) block import-swap (loader.py ~656, mirroring the object_detection@v3 swap at + ~645) - import the tensor sibling under the flag: + + from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v1 import ( + RoboflowSemanticSegmentationModelBlockV1, + ) + if ENABLE_TENSOR_DATA_REPRESENTATION: + from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v2_tensor import ( + RoboflowSemanticSegmentationModelBlockV2, + ) + else: + from inference.core.workflows.core_steps.models.roboflow.semantic_segmentation.v2 import ( + RoboflowSemanticSegmentationModelBlockV2, + ) +""" + +import base64 +import uuid +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +import cv2 +import numpy as np +import torch +from pydantic import ConfigDict, Field, model_validator + +from inference.core.env import ( + HOSTED_SEMANTIC_SEGMENTATION_URL, + LOCAL_INFERENCE_API_URL, + WORKFLOWS_IMAGE_TENSOR_DEVICE, + WORKFLOWS_REMOTE_API_TARGET, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, +) +from inference.core.managers.base import ModelManager +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + DETECTION_ID_KEY, + INFERENCE_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_SEMANTIC_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INFERENCE_ID_KIND, + ROBOFLOW_MODEL_ID_KIND, + STRING_KIND, + FloatZeroToOne, + ImageInputField, + RoboflowModelField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.semantic_segmentation import ( + SemanticSegmentationResult, +) +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import torch_mask_to_coco_rle +from inference_sdk import InferenceConfiguration, InferenceHTTPClient + +PREDICTION_TYPE = "semantic-segmentation" + +# D3 (numpy parity): numpy `v1.py:231` / `v2.py:292` drop ONLY class id 0 +# (`if cid != 0`) and keep every other present id as a detection โ€” including the +# 255 ignore/no-class sentinel. The flag-ON emitted set must match numpy exactly, +# so id 0 is the single excluded id. (This intentionally does NOT do the +# data-driven background-by-name / 255 exclusion; see the module docstring.) +BACKGROUND_CLASS_ID = 0 + +# `image_metadata` key under which the dense per-pixel confidence map is carried +# for numpy parity (numpy `v2.py` stores `conf_array` on the sv.Detections under +# `result["confidence_mask"]`). The serialiser never emits it into `predictions`, +# but it survives for consumers reading the prediction's `image_metadata`. +CONFIDENCE_MASK_KEY = "confidence_mask" + + +LONG_DESCRIPTION = """ +Run inference on a semantic segmentation model hosted on or uploaded to Roboflow. + +Semantic segmentation assigns a class label to every pixel in the image, producing a +dense segmentation mask rather than per-object bounding boxes or instance masks. + +You can query any model that is private to your account, or any public model available +on [Roboflow Universe](https://universe.roboflow.com). + +You will need to set your Roboflow API key in your Inference environment to use this +block. To learn more about setting your Roboflow API key, [refer to the Inference +documentation](https://inference.roboflow.com/quickstart/configure_api_key/). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Semantic Segmentation Model", + "version": "v2", + "short_description": "Assign a class label to every pixel in the image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "search_keywords": ["semantic", "segmentation", "deeplab", "deep_lab"], + "ui_manifest": { + "section": "model", + "icon": "far fa-paint-brush", + "blockPriority": 3, + "inference": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/roboflow_semantic_segmentation_model@v2"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + model_id: Union[Selector(kind=[ROBOFLOW_MODEL_ID_KIND]), str] = RoboflowModelField + confidence_mode: Union[ + Literal["best", "default", "custom"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="best", + description="How confidence thresholds are determined.", + json_schema_extra={ + "always_visible": True, + "values_metadata": { + "best": { + "name": "Best (Recommended)", + "description": "Use F1-optimal thresholds from model evaluation.", + }, + "default": { + "name": "Default", + "description": "Use the model's built-in default threshold.", + }, + "custom": { + "name": "Custom", + "description": "Specify a custom confidence threshold.", + }, + }, + }, + ) + custom_confidence: Union[ + Optional[FloatZeroToOne], + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + ] = Field( + default=0.4, + description="Custom confidence threshold for predictions.", + examples=[0.3, "$inputs.confidence_threshold"], + json_schema_extra={ + "relevant_for": { + "confidence_mode": {"values": ["custom"], "required": True}, + }, + }, + ) + + @model_validator(mode="after") + def validate(self) -> "BlockManifest": + if self.confidence_mode == "custom" and self.custom_confidence is None: + raise ValueError( + "`custom_confidence` is required when `confidence_mode` is 'custom'" + ) + return self + + @classmethod + def get_compatible_task_types(cls) -> Optional[List[str]]: + return ["semantic-segmentation"] + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + return [roboflow_platform_model(model_id=self.model_id)] + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=INFERENCE_ID_KEY, kind=[INFERENCE_ID_KIND]), + OutputDefinition( + name="predictions", + kind=[TENSOR_NATIVE_SEMANTIC_SEGMENTATION_PREDICTION_KIND], + ), + OutputDefinition(name="model_id", kind=[ROBOFLOW_MODEL_ID_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowSemanticSegmentationModelBlockV2(WorkflowBlock): + + def __init__( + self, + model_manager: ModelManager, + api_key: Optional[str], + step_execution_mode: StepExecutionMode, + ): + self._model_manager = model_manager + self._api_key = api_key + self._step_execution_mode = step_execution_mode + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["model_manager", "api_key", "step_execution_mode"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence_mode: str, + custom_confidence: Optional[float], + ) -> BlockResult: + confidence = ( + custom_confidence if confidence_mode == "custom" else confidence_mode + ) + if self._step_execution_mode is StepExecutionMode.LOCAL: + return self.run_locally( + images=images, model_id=model_id, confidence=confidence + ) + elif self._step_execution_mode is StepExecutionMode.REMOTE: + return self.run_remotely( + images=images, model_id=model_id, confidence=confidence + ) + else: + raise ValueError( + f"Unknown step execution mode: {self._step_execution_mode}" + ) + + def run_locally( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Union[None, float, Literal["best", "default"]], + ) -> BlockResult: + # Feed the representation already materialised on the images to avoid forcing a + # numpy->device conversion: GPU tensors (RGB) only when every image in the batch + # already has one, otherwise the numpy frames (BGR) โ€” matching the numpy block. + if all(image.is_tensor_materialised() for image in images): + model_inputs = [image.tensor_image for image in images] + image_color_format = "rgb" + else: + model_inputs = [image.numpy_image for image in images] + image_color_format = "bgr" + self._model_manager.add_model(model_id=model_id, api_key=self._api_key) + segmentation_results: List[SemanticSegmentationResult] = ( + self._model_manager.run_tensor_native_inference( + model_id=model_id, + images=model_inputs, + input_color_format=image_color_format, + confidence=confidence, + ) + ) + # class_names is index-ordered (id -> name); same source as the adapter's + # `class_map` ({str(idx): name}). + class_names = list(self._model_manager.get_class_names(model_id)) + class_names_map = _class_names_map(class_names) + excluded_ids = _excluded_class_ids() + results: List[dict] = [] + for image, segmentation in zip(images, segmentation_results): + inference_id = str(uuid.uuid4()) + detections = _build_instance_detections_from_segmentation( + segmentation_map=segmentation.segmentation_map, + confidence=segmentation.confidence, + image=image, + class_names_map=class_names_map, + excluded_ids=excluded_ids, + inference_id=inference_id, + ) + results.append( + { + INFERENCE_ID_KEY: inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + def run_remotely( + self, + images: Batch[WorkflowImageData], + model_id: str, + confidence: Union[None, float, Literal["best", "default"]], + ) -> BlockResult: + api_url = ( + LOCAL_INFERENCE_API_URL + if WORKFLOWS_REMOTE_API_TARGET != "hosted" + else HOSTED_SEMANTIC_SEGMENTATION_URL + ) + client = InferenceHTTPClient( + api_url=api_url, + api_key=self._api_key, + ) + if WORKFLOWS_REMOTE_API_TARGET == "hosted": + client.select_api_v0() + client_config = InferenceConfiguration( + confidence_threshold=confidence, + max_batch_size=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_BATCH_SIZE, + max_concurrent_requests=WORKFLOWS_REMOTE_EXECUTION_MAX_STEP_CONCURRENT_REQUESTS, + source="workflow-execution", + ) + client.configure(inference_configuration=client_config) + inference_images = [i.base64_image for i in images] + predictions = client.infer( + inference_input=inference_images, + model_id=model_id, + ) + if not isinstance(predictions, list): + predictions = [predictions] + return self._post_process_remote_result( + images=images, predictions=predictions, model_id=model_id + ) + + def _post_process_remote_result( + self, + images: Batch[WorkflowImageData], + predictions: List[dict], + model_id: str, + ) -> BlockResult: + results: List[dict] = [] + for image, prediction in zip(images, predictions): + inference_id = prediction.get(INFERENCE_ID_KEY) or str(uuid.uuid4()) + detections = _build_instance_detections_from_inference_response( + predictions_dict=prediction.get("predictions") or {}, + image=image, + inference_id=inference_id, + ) + results.append( + { + INFERENCE_ID_KEY: inference_id, + "predictions": detections, + "model_id": model_id, + } + ) + return results + + +def _class_names_map(class_names: List[str]) -> Dict[int, str]: + return {index: name for index, name in enumerate(class_names)} + + +def _excluded_class_ids() -> set: + """Class ids that must never become detections. To match numpy exactly (D3), + this is ONLY id 0: numpy ``v1.py``/``v2.py`` filter ``cid != 0`` and keep every + other present id โ€” the 255 ignore sentinel included.""" + return {BACKGROUND_CLASS_ID} + + +def _empty_instance_detections( + image: WorkflowImageData, + class_names_map: Dict[int, str], + height: int, + width: int, + inference_id: str, +) -> InstanceDetections: + detections = InstanceDetections( + xyxy=torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + class_id=torch.zeros( + (0,), dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.zeros( + (0,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + mask=InstancesRLEMasks(image_size=(height, width), masks=[]), + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=class_names_map, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + detections.bboxes_metadata = None + return detections + + +def _assemble_instance_detections( + xyxy: List[List[float]], + class_ids: List[int], + confidences: List[float], + rle_dicts: List[dict], + bboxes_metadata: List[dict], + image: WorkflowImageData, + class_names_map: Dict[int, str], + height: int, + width: int, + inference_id: str, + confidence_mask: Optional[torch.Tensor] = None, +) -> InstanceDetections: + if len(rle_dicts) == 0: + return _empty_instance_detections( + image=image, + class_names_map=class_names_map, + height=height, + width=width, + inference_id=inference_id, + ) + detections = InstanceDetections( + xyxy=torch.tensor( + xyxy, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + class_id=torch.tensor( + class_ids, dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.tensor( + confidences, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + mask=InstancesRLEMasks.from_coco_rle_masks( + image_size=(height, width), masks=rle_dicts + ), + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names=class_names_map, + prediction_type=PREDICTION_TYPE, + inference_id=inference_id, + ) + # Carry the dense per-pixel confidence map for numpy parity (numpy `v2.py` + # attaches `conf_array` to the sv.Detections under `confidence_mask`). The + # serialiser drops it from `predictions`, but consumers can still read it off + # the prediction's `image_metadata`. + if confidence_mask is not None: + detections.image_metadata[CONFIDENCE_MASK_KEY] = confidence_mask + detections.bboxes_metadata = bboxes_metadata + return detections + + +def _build_instance_detections_from_segmentation( + segmentation_map: torch.Tensor, + confidence: torch.Tensor, + image: WorkflowImageData, + class_names_map: Dict[int, str], + excluded_ids: set, + inference_id: str, +) -> InstanceDetections: + """LOCAL path: dense ``(H, W)`` class grid + ``(H, W)`` confidence -> + one RLE instance per present non-background/non-ignore class. + + Masks are encoded one class at a time via ``torch_mask_to_coco_rle`` (the same + compact COCO-RLE approach as ``seg_preview/v1_tensor.py``); no ``N x H x W`` + dense stack is ever materialised. + """ + height, width = int(segmentation_map.shape[0]), int(segmentation_map.shape[1]) + xyxy: List[List[float]] = [] + class_ids: List[int] = [] + confidences: List[float] = [] + bboxes_metadata: List[dict] = [] + rle_dicts: List[dict] = [] + + present_ids = [int(value) for value in torch.unique(segmentation_map).tolist()] + for class_id in present_ids: + if class_id in excluded_ids: + continue + # numpy names any present id via `class_map.get(str(cid), str(cid))`; ids + # absent from the class map (e.g. the 255 sentinel) fall back to str(id) + # and are STILL emitted โ€” matching numpy, which never skips on name. + class_name = class_names_map.get(class_id, str(class_id)) + binary_mask = segmentation_map == class_id + bbox = _bbox_from_binary_mask(binary_mask) + if bbox is None: + continue + # per-class mean confidence over that class's pixels + class_confidence = float(confidence[binary_mask].mean().item()) + # single (H, W) bool/byte mask -> compact COCO RLE in C; no dense stack. + rle = torch_mask_to_coco_rle(binary_mask.to(torch.uint8)) + rle = _normalise_rle_counts(rle) + xyxy.append(bbox) + class_ids.append(class_id) + confidences.append(class_confidence) + rle_dicts.append(rle) + bboxes_metadata.append( + {DETECTION_ID_KEY: str(uuid.uuid4()), CLASS_NAME_KEY: class_name} + ) + + return _assemble_instance_detections( + xyxy=xyxy, + class_ids=class_ids, + confidences=confidences, + rle_dicts=rle_dicts, + bboxes_metadata=bboxes_metadata, + image=image, + class_names_map=class_names_map, + height=height, + width=width, + inference_id=inference_id, + # Dense per-pixel confidence grid (kept on device) for numpy parity. + confidence_mask=confidence, + ) + + +def _build_instance_detections_from_inference_response( + predictions_dict: Dict, + image: WorkflowImageData, + inference_id: str, +) -> InstanceDetections: + """REMOTE path: standard inference semantic-seg response (a single ``dict``, + matching numpy ``v2.py``'s ``_convert_to_sv_detections`` input) -> + one RLE instance per present non-background/non-ignore class. + + Response shape (see numpy ``v2.py``): + * ``segmentation_mask`` - base64 PNG, grayscale, pixel value == class id + * ``confidence_mask`` - base64 PNG, grayscale, 0..255 (== confidence*255) + * ``class_map`` - ``{str(class_id): class_name}`` + """ + seg_mask_b64 = predictions_dict.get("segmentation_mask", "") + conf_mask_b64 = predictions_dict.get("confidence_mask", "") + class_map: Dict[str, str] = predictions_dict.get("class_map", {}) + class_names_map = {int(k): v for k, v in class_map.items()} + + height, width = image._read_shape_without_materialization() + + mask_array = _decode_b64_grayscale_png(seg_mask_b64) + if mask_array is None: + return _empty_instance_detections( + image=image, + class_names_map=class_names_map, + height=int(height), + width=int(width), + inference_id=inference_id, + ) + height, width = int(mask_array.shape[0]), int(mask_array.shape[1]) + conf_array = _decode_b64_grayscale_png(conf_mask_b64) + + excluded_ids = _excluded_class_ids() + xyxy: List[List[float]] = [] + class_ids: List[int] = [] + confidences: List[float] = [] + bboxes_metadata: List[dict] = [] + rle_dicts: List[dict] = [] + + present_ids = [int(value) for value in np.unique(mask_array).tolist()] + for class_id in present_ids: + if class_id in excluded_ids: + continue + class_name = class_names_map.get(class_id, str(class_id)) + binary_mask = mask_array == class_id + bbox = _bbox_from_numpy_mask(binary_mask) + if bbox is None: + continue + if conf_array is not None: + class_confidence = float(conf_array[binary_mask].mean()) / 255.0 + else: + class_confidence = 1.0 + mask_tensor = torch.from_numpy(binary_mask.astype(np.uint8)).to( + WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + rle = torch_mask_to_coco_rle(mask_tensor) + rle = _normalise_rle_counts(rle) + xyxy.append(bbox) + class_ids.append(class_id) + confidences.append(class_confidence) + rle_dicts.append(rle) + bboxes_metadata.append( + {DETECTION_ID_KEY: str(uuid.uuid4()), CLASS_NAME_KEY: class_name} + ) + + # Carry the dense per-pixel confidence map for numpy parity (numpy `v2.py` + # attaches the decoded `conf_array`). Mirror the native carrier convention by + # moving it to the workflow tensor device when present. + confidence_mask = ( + torch.from_numpy(conf_array).to(WORKFLOWS_IMAGE_TENSOR_DEVICE) + if conf_array is not None + else None + ) + return _assemble_instance_detections( + xyxy=xyxy, + class_ids=class_ids, + confidences=confidences, + rle_dicts=rle_dicts, + bboxes_metadata=bboxes_metadata, + image=image, + class_names_map=class_names_map, + height=height, + width=width, + inference_id=inference_id, + confidence_mask=confidence_mask, + ) + + +def _bbox_from_binary_mask(binary_mask: torch.Tensor) -> Optional[List[float]]: + """Tight ``[x_min, y_min, x_max, y_max]`` enclosing all True pixels of a + ``(H, W)`` torch bool mask, or ``None`` if empty.""" + rows = torch.any(binary_mask, dim=1) + cols = torch.any(binary_mask, dim=0) + row_indices = torch.nonzero(rows, as_tuple=False).flatten() + col_indices = torch.nonzero(cols, as_tuple=False).flatten() + if row_indices.numel() == 0 or col_indices.numel() == 0: + return None + y_min = float(row_indices[0].item()) + y_max = float(row_indices[-1].item()) + x_min = float(col_indices[0].item()) + x_max = float(col_indices[-1].item()) + return [x_min, y_min, x_max, y_max] + + +def _bbox_from_numpy_mask(binary_mask: np.ndarray) -> Optional[List[float]]: + """Tight ``[x_min, y_min, x_max, y_max]`` enclosing all True pixels of a + ``(H, W)`` numpy bool mask, or ``None`` if empty (mirrors numpy ``v2.py``).""" + rows = np.where(np.any(binary_mask, axis=1))[0] + cols = np.where(np.any(binary_mask, axis=0))[0] + if rows.size == 0 or cols.size == 0: + return None + return [float(cols[0]), float(rows[0]), float(cols[-1]), float(rows[-1])] + + +def _decode_b64_grayscale_png(b64_value: str) -> Optional[np.ndarray]: + if not b64_value: + return None + mask_bytes = base64.b64decode(b64_value) + nparr = np.frombuffer(mask_bytes, np.uint8) + return cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE) + + +def _normalise_rle_counts(rle: dict) -> dict: + """``InstancesRLEMasks.from_coco_rle_masks`` keeps the raw ``counts`` payload + as-is; ``pycocotools`` returns it as ``bytes``. Decode to ``str`` so the + serialised RLE matches the rest of the tensor pipeline (and the numpy block, + which also stores ``counts`` as a utf-8 string).""" + counts = rle.get("counts") + if isinstance(counts, bytes): + rle = dict(rle) + rle["counts"] = counts.decode("utf-8") + return rle diff --git a/inference/core/workflows/core_steps/models/third_party/barcode_detection/v1_tensor.py b/inference/core/workflows/core_steps/models/third_party/barcode_detection/v1_tensor.py new file mode 100644 index 0000000000..0329afa597 --- /dev/null +++ b/inference/core/workflows/core_steps/models/third_party/barcode_detection/v1_tensor.py @@ -0,0 +1,158 @@ +"""Tensor-native sibling of `roboflow_core/barcode_detector@v1`. + +Classical CV (zxingcpp) on numpy_image โ€” there is no inference_models path โ€” so the +ONLY change is the OUTPUT representation: instead of sv.Detections this builds +`inference_models.Detections` (xyxy / class_id / confidence) with `image_metadata` +(class_names map + prediction_type + lineage, via build_native_image_metadata) and +per-detection `bboxes_metadata` carrying the required detection_id plus the decoded +value under DETECTED_CODE_KEY. The output kind becomes the tensor-native barcode kind. +numpy_image is materialised transparently from the CHW tensor_image; zxingcpp needs +numpy, so the image path is unchanged. +""" + +from typing import List, Literal, Optional, Type +from uuid import uuid4 + +import torch +import zxingcpp +from pydantic import ConfigDict + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.execution_engine.constants import ( + DETECTED_CODE_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_BAR_CODE_DETECTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections + +PREDICTION_TYPE = "barcode-detection" +CLASS_NAME = "barcode" + +LONG_DESCRIPTION = """ +Detect the location of barcodes in an image. + +This block is useful for manufacturing and consumer packaged goods projects where you +need to detect a barcode region in an image. You can then apply Crop block to isolate +each barcode then apply further processing (i.e. OCR of the characters on a barcode). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Barcode Detection", + "version": "v1", + "short_description": "Detect and read barcodes in an image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "far fa-barcode", + "blockPriority": 12, + }, + } + ) + type: Literal[ + "roboflow_core/barcode_detector@v1", "BarcodeDetector", "BarcodeDetection" + ] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", kind=[TENSOR_NATIVE_BAR_CODE_DETECTION_KIND] + ) + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class BarcodeDetectorBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run(self, images: Batch[WorkflowImageData]) -> BlockResult: + return [{"predictions": detect_barcodes(image=image)} for image in images] + + +def detect_barcodes(image: WorkflowImageData) -> Detections: + barcodes = zxingcpp.read_barcodes(image.numpy_image) + xyxy: List[List[float]] = [] + codes: List[str] = [] + for barcode in barcodes: + x_min = barcode.position.top_left.x + y_min = barcode.position.top_left.y + x_max = barcode.position.bottom_right.x + y_max = barcode.position.bottom_right.y + xyxy.append([float(x_min), float(y_min), float(x_max), float(y_max)]) + codes.append(barcode.text) + return _build_code_detections( + image=image, xyxy=xyxy, codes=codes, class_name=CLASS_NAME + ) + + +def _build_code_detections( + image: WorkflowImageData, + xyxy: List[List[float]], + codes: List[str], + class_name: str, +) -> Detections: + n = len(xyxy) + detections = Detections( + xyxy=( + torch.tensor( + xyxy, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + if n + else torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + ), + class_id=torch.zeros( + (n,), dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.ones( + (n,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names={0: class_name}, + prediction_type=PREDICTION_TYPE, + inference_id=str(uuid4()), + ) + detections.bboxes_metadata = [ + {DETECTION_ID_KEY: str(uuid4()), DETECTED_CODE_KEY: codes[index]} + for index in range(n) + ] + return detections diff --git a/inference/core/workflows/core_steps/models/third_party/qr_code_detection/v1_tensor.py b/inference/core/workflows/core_steps/models/third_party/qr_code_detection/v1_tensor.py new file mode 100644 index 0000000000..6e1718b52b --- /dev/null +++ b/inference/core/workflows/core_steps/models/third_party/qr_code_detection/v1_tensor.py @@ -0,0 +1,160 @@ +"""Tensor-native sibling of `roboflow_core/qr_code_detector@v1`. + +Classical CV (cv2.QRCodeDetector) on numpy_image โ€” there is no inference_models path +โ€” so the ONLY change is the OUTPUT representation: instead of sv.Detections this +builds `inference_models.Detections` (xyxy / class_id / confidence) with +`image_metadata` (class_names map + prediction_type + lineage, via +build_native_image_metadata) and per-detection `bboxes_metadata` carrying the +required detection_id plus the decoded value under DETECTED_CODE_KEY. The output kind +becomes the tensor-native QR kind. numpy_image is materialised transparently from the +CHW tensor_image; cv2 needs numpy, so the image path is unchanged. +""" + +from typing import List, Literal, Optional, Type +from uuid import uuid4 + +import cv2 +import torch +from pydantic import ConfigDict + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.tensor_native import ( + build_native_image_metadata, +) +from inference.core.workflows.execution_engine.constants import ( + DETECTED_CODE_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_QR_CODE_DETECTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections + +PREDICTION_TYPE = "qrcode-detection" +CLASS_NAME = "qr_code" + +LONG_DESCRIPTION = """ +Detect the location of a QR code. + +This block is useful for manufacturing and consumer packaged goods projects where you +need to detect a QR code region in an image. You can then apply Crop block to isolate +each QR code then apply further processing (i.e. read a QR code with a custom block). +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "QR Code Detection", + "version": "v1", + "short_description": "Detect and read QR codes in an image.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "model", + "ui_manifest": { + "section": "model", + "icon": "fal fa-qrcode", + "blockPriority": 13, + "opencv": True, + }, + } + ) + type: Literal[ + "roboflow_core/qr_code_detector@v1", "QRCodeDetector", "QRCodeDetection" + ] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", kind=[TENSOR_NATIVE_QR_CODE_DETECTION_KIND] + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class QRCodeDetectorBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run(self, images: Batch[WorkflowImageData]) -> BlockResult: + return [{"predictions": detect_qr_codes(image=image)} for image in images] + + +def detect_qr_codes(image: WorkflowImageData) -> Detections: + detector = cv2.QRCodeDetector() + _retval, decoded, points_list, _ = detector.detectAndDecodeMulti(image.numpy_image) + xyxy: List[List[float]] = [] + codes: List[str] = [] + for data, points in zip(decoded, points_list): + x_min = points[0][0] + y_min = points[0][1] + x_max = x_min + (points[2][0] - points[0][0]) + y_max = y_min + (points[2][1] - points[0][1]) + xyxy.append([float(x_min), float(y_min), float(x_max), float(y_max)]) + codes.append(data) + return _build_code_detections( + image=image, xyxy=xyxy, codes=codes, class_name=CLASS_NAME + ) + + +def _build_code_detections( + image: WorkflowImageData, + xyxy: List[List[float]], + codes: List[str], + class_name: str, +) -> Detections: + n = len(xyxy) + detections = Detections( + xyxy=( + torch.tensor( + xyxy, dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + if n + else torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + ), + class_id=torch.zeros( + (n,), dtype=torch.int64, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.ones( + (n,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + ) + detections.image_metadata = build_native_image_metadata( + image=image, + class_names={0: class_name}, + prediction_type=PREDICTION_TYPE, + inference_id=str(uuid4()), + ) + detections.bboxes_metadata = [ + {DETECTION_ID_KEY: str(uuid4()), DETECTED_CODE_KEY: codes[index]} + for index in range(n) + ] + return detections diff --git a/inference/core/workflows/core_steps/sampling/identify_changes/v1_tensor.py b/inference/core/workflows/core_steps/sampling/identify_changes/v1_tensor.py new file mode 100644 index 0000000000..5504c973e0 --- /dev/null +++ b/inference/core/workflows/core_steps/sampling/identify_changes/v1_tensor.py @@ -0,0 +1,406 @@ +import math +from typing import List, Literal, Optional, Type, Union + +import numpy as np +import torch +from pydantic import ConfigDict, Field + +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_EMBEDDING_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) + +LONG_DESCRIPTION = """ +Identify changes and detect when data patterns change at unusual rates compared to historical norms by tracking embedding vectors over time, measuring cosine similarity changes, computing rate-of-change statistics, and flagging anomalies when changes occur faster or slower than expected for change detection, anomaly monitoring, rate-of-change analysis, and temporal pattern detection workflows. + +## How This Block Works + +This block detects changes by monitoring how quickly embeddings change over time and comparing the current rate of change against historical patterns. The block: + +1. Receives an embedding vector representing the current data point's features +2. Normalizes the embedding to unit length: + - Converts the embedding to a unit vector (length = 1) for cosine similarity calculations + - Enables comparison using angular similarity rather than distance-based metrics + - Handles zero vectors gracefully by skipping normalization +3. Tracks sample count and warmup status: + - Increments sample counter for each processed embedding + - Determines if still in warmup period (samples < warmup parameter) + - During warmup, no outliers are identified to allow baseline establishment +4. Maintains running statistics for embedding averages and standard deviations using one of three strategies: + + **For Exponential Moving Average (EMA) strategy:** + - Updates average and variance using exponential moving average with smoothing_factor + - More recent embeddings have greater weight (controlled by smoothing_factor) + - Adapts quickly to recent trends while maintaining historical context + - Smoothing factor determines responsiveness (higher = more responsive to recent changes) + + **For Simple Moving Average (SMA) strategy:** + - Uses Welford's method to calculate running mean and variance + - All historical samples contribute equally to statistics + - Provides stable, unbiased estimates over time + - Well-suited for consistent, long-term tracking + + **For Sliding Window strategy:** + - Maintains a fixed-size window of recent embeddings (window_size) + - Removes oldest embeddings when window exceeds size (FIFO) + - Calculates mean and standard deviation from window contents only + - Adapts quickly to recent trends, discarding older information + +5. Calculates cosine similarity between current embedding and running average: + - Measures how similar the current embedding is to the typical embedding pattern + - Cosine similarity ranges from -1 (opposite) to 1 (identical) + - Values close to 1 indicate the embedding is similar to the norm + - Values further from 1 indicate the embedding differs from the norm +6. Tracks rate of change by monitoring cosine similarity statistics: + - Maintains running average and standard deviation of cosine similarity values + - Uses the same strategy (EMA, SMA, or Sliding Window) for cosine similarity tracking + - Measures how quickly embeddings are changing compared to historical change rates + - Tracks both the average change rate and variability in change rates +7. Calculates z-score for current cosine similarity: + - Measures how many standard deviations the current cosine similarity is from the average + - Z-score = (current_cosine_similarity - average_cosine_similarity) / std_cosine_similarity + - Positive z-scores indicate faster-than-normal changes + - Negative z-scores indicate slower-than-normal changes +8. Converts z-score to percentile: + - Uses error function (erf) to convert z-score to percentile position + - Percentile represents where the current change rate ranks among historical rates + - Values near 0.0 indicate unusually slow changes (low percentiles) + - Values near 1.0 indicate unusually fast changes (high percentiles) +9. Determines outlier status based on percentile thresholds: + - Flags as outlier if percentile is below threshold_percentile/2 (unusually slow changes) + - Flags as outlier if percentile is above (1 - threshold_percentile/2) (unusually fast changes) + - Detects both abnormally fast and abnormally slow rates of change +10. Updates running statistics for next iteration: + - Updates embedding average and standard deviation using selected strategy + - Updates cosine similarity average and standard deviation using selected strategy + - Maintains state across workflow executions +11. Returns six outputs: + - **is_outlier**: Boolean flag indicating if the rate of change is anomalous + - **percentile**: Float value (0.0-1.0) representing where the change rate ranks historically + - **z_score**: Float value representing standard deviations from average change rate + - **average**: Current average embedding vector (running average of historical embeddings) + - **std**: Current standard deviation vector for embeddings (variability per dimension) + - **warming_up**: Boolean flag indicating if still in warmup period + +The block monitors the **rate of change** rather than just detecting outliers. It tracks how quickly embeddings are changing and compares this to historical change patterns. When embeddings change much faster or slower than they have in the past, the block flags this as anomalous. This makes it ideal for detecting sudden pattern shifts, unexpected changes in scenes, or unusual behavior patterns. The three strategies (EMA, SMA, Sliding Window) offer different trade-offs between responsiveness and stability. + +## Common Use Cases + +- **Change Detection**: Detect when scenes, environments, or patterns change unexpectedly (e.g., detect scene changes, identify sudden pattern shifts, flag unexpected environmental changes), enabling change detection workflows +- **Anomaly Monitoring**: Monitor for unusual changes in behavior or patterns (e.g., detect abnormal behavior changes, monitor unusual pattern variations, flag unexpected rate changes), enabling anomaly monitoring workflows +- **Rate-of-Change Analysis**: Analyze and detect unusual rates of change in data streams (e.g., detect unusually fast changes, identify unusually slow changes, monitor change rate patterns), enabling rate-of-change analysis workflows +- **Temporal Pattern Detection**: Identify when temporal patterns deviate from expected change rates (e.g., detect pattern disruptions, identify timeline anomalies, flag temporal inconsistencies), enabling temporal pattern detection workflows +- **Quality Monitoring**: Monitor for unexpected changes in quality or characteristics (e.g., detect quality degradation, identify unexpected quality changes, monitor characteristic variations), enabling quality monitoring workflows +- **Event Detection**: Detect significant events based on unusual change rates (e.g., detect significant events, identify important changes, flag notable pattern shifts), enabling event detection workflows + +## Connecting to Other Blocks + +This block receives embeddings and produces is_outlier, percentile, z_score, average, std, and warming_up outputs: + +- **After embedding model blocks** (CLIP, Perception Encoder, etc.) to analyze change rates from embeddings (e.g., detect changes from CLIP embeddings, analyze Perception Encoder change rates, monitor embedding-based changes), enabling embedding-to-change workflows +- **After classification or detection blocks** with embeddings to identify unusual change patterns (e.g., detect unusual detection changes, flag anomalous classification changes, monitor prediction pattern changes), enabling prediction-to-change workflows +- **Before logic blocks** like Continue If to make decisions based on change detection (e.g., continue if change detected, filter based on change rate, trigger actions on unusual changes), enabling change-based decision workflows +- **Before notification blocks** to alert on change detection (e.g., alert on significant changes, notify about pattern shifts, trigger alerts on rate anomalies), enabling change-based notification workflows +- **Before data storage blocks** to record change information (e.g., log change data, store change statistics, record rate-of-change metrics), enabling change data logging workflows +- **In monitoring pipelines** where change detection is part of continuous monitoring (e.g., monitor changes in observation systems, track pattern variations, detect anomalies in monitoring workflows), enabling change monitoring workflows + +## Requirements + +This block requires embeddings as input (typically from embedding model blocks like CLIP or Perception Encoder). The block maintains internal state across workflow executions, tracking running statistics for both embeddings and cosine similarity values. During the warmup period (first `warmup` samples), no outliers are identified and the block returns is_outlier=False, percentile=0.5, and warming_up=True. After warmup, the block uses the selected strategy (EMA, SMA, or Sliding Window) to track statistics and detect rate-of-change anomalies. The threshold_percentile parameter (0.0-1.0) controls sensitivity - lower values detect only extreme rate changes, while higher values detect more moderate rate deviations. The strategy choice affects responsiveness: EMA adapts quickly to recent trends, SMA provides stable long-term tracking, and Sliding Window adapts quickly but discards older information. The block works best with consistent embedding models and may need adjustment of threshold_percentile and strategy based on expected variation and change patterns in your data. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Identify Changes", + "version": "v1", + "short_description": "Identify changes compared to prior data via embeddings.", + "long_description": LONG_DESCRIPTION, + "license": "MIT", + "block_type": "video", + "ui_manifest": { + "section": "video", + "icon": "far fa-triangle", + }, + } + ) + type: Literal["roboflow_core/identify_changes@v1"] + name: str = Field(description="Unique name of step in workflows") + + strategy: Literal[ + "Exponential Moving Average (EMA)", + "Simple Moving Average (SMA)", + "Sliding Window", + ] = Field( + default="Exponential Moving Average (EMA)", + description="Statistical strategy for tracking embedding and change rate statistics. 'Exponential Moving Average (EMA)': Adapts quickly to recent trends, more weight on recent data. 'Simple Moving Average (SMA)': Stable long-term tracking, all data contributes equally. 'Sliding Window': Fast adaptation, uses only recent window_size samples. EMA is best for adaptive monitoring, SMA for stable tracking, Sliding Window for rapid adaptation.", + examples=[ + "Exponential Moving Average (EMA)", + "Simple Moving Average (SMA)", + "Sliding Window", + ], + json_schema_extra={ + "always_visible": True, + }, + ) + + embedding: Selector(kind=[TENSOR_NATIVE_EMBEDDING_KIND]) = Field( + description="Embedding vector representing the current data point's features. Typically from embedding models like CLIP or Perception Encoder. The embedding is normalized to unit length for cosine similarity calculations. The block compares current embedding to running average and tracks rate of change over time.", + examples=["$steps.clip.embedding", "$steps.perception_encoder.embedding"], + ) + + threshold_percentile: Union[Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), float] = Field( + default=0.2, + description="Percentile threshold for change rate anomaly detection, range 0.0-1.0. Change rates below threshold_percentile/2 or above (1 - threshold_percentile/2) are flagged as outliers. Lower values (e.g., 0.05) detect only extreme rate changes - very strict. Higher values (e.g., 0.3) detect more moderate rate deviations - more sensitive. Default 0.2 means bottom 10% and top 10% of change rates are outliers. Adjust based on expected variation in change rates.", + examples=[0.2, 0.05, 0.3, "$inputs.threshold_percentile"], + json_schema_extra={ + "always_visible": True, + }, + ) + + warmup: Union[Selector(kind=[INTEGER_KIND]), int] = Field( + default=3, + description="Number of initial data points required before change detection begins. During warmup, no outliers are identified (is_outlier=False) to allow baseline establishment for change rates. Must be at least 2 for statistical analysis. Typical range: 3-100 samples. Higher values provide more stable baselines but delay change detection. Lower values enable faster detection but may be less accurate initially.", + examples=[3, 10, 100], + ) + + smoothing_factor: Optional[ + Union[Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), float] + ] = Field( + default=0.1, + description="Smoothing factor (alpha) for Exponential Moving Average strategy, range 0.0-1.0. Controls responsiveness to recent data - higher values make statistics more responsive to recent changes, lower values maintain more historical context. Example: 0.1 means 10% weight on current value, 90% on historical average. Typical range: 0.05-0.3. Only used when strategy is 'Exponential Moving Average (EMA)'.", + examples=[0.1, 0.05, 0.25], + json_schema_extra={ + "relevant_for": { + "strategy": { + "values": {"Exponential Moving Average (EMA)"}, + }, + }, + }, + ) + + window_size: Optional[Union[Selector(kind=[INTEGER_KIND]), int]] = Field( + default=10, + description="Maximum number of recent embeddings to maintain in sliding window. When exceeded, oldest embeddings are removed (FIFO). Larger windows provide more stable statistics but adapt slower to changes. Smaller windows adapt faster but may be less stable. Only used when strategy is 'Sliding Window'. Must be at least 2. Typical range: 5-50 embeddings.", + examples=[10, 20, 50], + json_schema_extra={ + "relevant_for": { + "strategy": {"values": {"Sliding Window"}, "required": True}, + }, + }, + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="is_outlier", kind=[BOOLEAN_KIND]), + OutputDefinition(name="percentile", kind=[FLOAT_ZERO_TO_ONE_KIND]), + OutputDefinition(name="z_score", kind=[FLOAT_KIND]), + OutputDefinition(name="average", kind=[TENSOR_NATIVE_EMBEDDING_KIND]), + OutputDefinition(name="std", kind=[TENSOR_NATIVE_EMBEDDING_KIND]), + OutputDefinition(name="warming_up", kind=[BOOLEAN_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class IdentifyChangesBlockV1(WorkflowBlock): + def __init__(self): + self.average = None + self.std = None + self.var = None # For EMA variance tracking + self.M2 = None # For SMA variance tracking + self.sliding_window = [] + self.samples = 0 + + self.cosine_similarity_avg = None + self.cosine_similarity_std = None + self.cosine_similarity_var = None + self.cosine_similarity_m2 = None + self.cosine_similarity_sliding_window = [] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + strategy: str, + embedding: torch.Tensor, + threshold_percentile: float, + smoothing_factor: float, + window_size: int, + warmup: int, + ) -> BlockResult: + is_outlier = False + percentile = 0.5 + z_score = 0 + warming_up = False + + # Embeddings stay on-device: the running average/std are torch tensors; + # only the scalar cosine-similarity statistics use python/numpy scalars. + embedding = embedding.detach() + norm = torch.linalg.norm(embedding) + if norm != 0: + embedding = embedding / norm + + # determine if embedding is an outlier + if self.average is not None: + # np.float64 keeps the numpy divide-by-zero -> nan semantics in the + # downstream scalar statistics (no exception). + cs = np.float64( + ( + torch.dot(embedding, self.average) + / torch.sqrt( + torch.dot(embedding, embedding) + * torch.dot(self.average, self.average) + ) + ).item() + ) + + if self.cosine_similarity_avg is None: + self.cosine_similarity_avg = cs + self.cosine_similarity_std = 0 + self.cosine_similarity_var = 0 + self.cosine_similarity_m2 = 0 + else: + if strategy == "Exponential Moving Average (EMA)": + # Update EMA average: + self.cosine_similarity_avg = ( + 1 - smoothing_factor + ) * self.cosine_similarity_avg + smoothing_factor * cs + + # Update EMA variance: + # var_new = (1 - alpha)*var_old + alpha*(x - new_avg)^2 + diff = cs - self.cosine_similarity_avg + self.cosine_similarity_var = ( + 1 - smoothing_factor + ) * self.cosine_similarity_var + smoothing_factor * (diff**2) + self.cosine_similarity_std = np.sqrt(self.cosine_similarity_var) + elif strategy == "Simple Moving Average (SMA)": + count = self.samples + 1 + delta = cs - self.cosine_similarity_avg + + self.cosine_similarity_avg = ( + cs / count + self.cosine_similarity_avg * self.samples / count + ) + delta2 = cs - self.cosine_similarity_avg + + self.cosine_similarity_m2 = ( + self.cosine_similarity_m2 + delta * delta2 + ) + var = self.cosine_similarity_m2 / (count - 1) + self.cosine_similarity_std = np.sqrt(var) + elif strategy == "Sliding Window": + self.cosine_similarity_sliding_window.append(cs) + if len(self.cosine_similarity_sliding_window) > window_size: + self.cosine_similarity_sliding_window.pop(0) + + self.cosine_similarity_avg = np.mean( + self.cosine_similarity_sliding_window + ) + self.cosine_similarity_std = np.std( + self.cosine_similarity_sliding_window + ) + + if self.cosine_similarity_std == 0: + # No variance in cosine similarity (e.g. a static scene) means + # there is nothing to compare against; treat as non-outlier + # instead of dividing by zero and emitting NaN downstream. + z_score = 0 + percentile = 0.5 + else: + z_score = (cs - self.cosine_similarity_avg) / self.cosine_similarity_std + percentile = 1 - 0.5 * (1 + math.erf(z_score / np.sqrt(2))) + + # print(f"Z-score: {z_score}, Percentile: {percentile}, Cosine Similarity: {cs}, Average: {self.cosine_similarity_avg}, Std: {self.cosine_similarity_std}") + + if self.samples < warmup: + is_outlier = False + warming_up = True + else: + is_outlier = percentile <= threshold_percentile / 2 or percentile >= ( + 1 - threshold_percentile / 2 + ) + + # update average and std (kept as torch tensors, on the embedding's device) + if self.average is None: + self.average = embedding + self.std = torch.zeros_like(embedding) + self.var = torch.zeros_like(embedding) + self.M2 = torch.zeros_like(embedding) + else: + if strategy == "Exponential Moving Average (EMA)": + # Update EMA average: + self.average = ( + 1 - smoothing_factor + ) * self.average + smoothing_factor * embedding + + # Update EMA variance: + # var_new = (1 - alpha)*var_old + alpha*(x - new_avg)^2 + diff = embedding - self.average + self.var = (1 - smoothing_factor) * self.var + smoothing_factor * ( + diff**2 + ) + self.std = torch.sqrt(self.var) + + elif strategy == "Simple Moving Average (SMA)": + # Use Welford's method to update mean and variance + count = self.samples + 1 + delta = embedding - self.average + + # Update average: + self.average = self.average + delta / count + delta2 = embedding - self.average + + # Update M2: + self.M2 = self.M2 + delta * delta2 + var = self.M2 / (count - 1) + self.std = torch.sqrt(var) + + elif strategy == "Sliding Window": + self.sliding_window.append(embedding) + if len(self.sliding_window) > window_size: + self.sliding_window.pop(0) + + window = torch.stack(self.sliding_window) + self.average = window.mean(dim=0) + self.std = window.std(dim=0, unbiased=False) + + self.samples = self.samples + 1 + + return { + "is_outlier": is_outlier, + "percentile": percentile, + "z_score": z_score, + # `average`/`std` are embedding-kind outputs: torch.Tensor on the + # embedding's device. + "average": self.average, + "std": self.std, + "warming_up": warming_up, + } diff --git a/inference/core/workflows/core_steps/sampling/identify_outliers/v1_tensor.py b/inference/core/workflows/core_steps/sampling/identify_outliers/v1_tensor.py new file mode 100644 index 0000000000..47fdfc4bf6 --- /dev/null +++ b/inference/core/workflows/core_steps/sampling/identify_outliers/v1_tensor.py @@ -0,0 +1,266 @@ +from typing import List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field + +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_EMBEDDING_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) + +LONG_DESCRIPTION = """ +Identify outlier embeddings compared to prior data using von Mises-Fisher statistical distribution analysis to detect anomalies, unusual patterns, or deviations from normal behavior by comparing current embedding vectors against a sliding window of historical embeddings for quality control, anomaly detection, and data sampling workflows. + +## How This Block Works + +This block detects outliers by statistically comparing embedding vectors against historical data using directional statistics. The block: + +1. Receives an embedding vector representing the current data point's features +2. Normalizes the embedding to unit length: + - Converts the embedding to a unit vector (length = 1) for directional analysis + - Enables comparison using angular/directional statistics rather than distance-based metrics + - Handles zero vectors gracefully by skipping normalization +3. Tracks sample count and warmup status: + - Increments sample counter for each processed embedding + - Determines if still in warmup period (samples < warmup parameter) + - During warmup, no outliers are identified to allow baseline establishment +4. Maintains a sliding window of historical embeddings: + - Stores normalized embeddings in a buffer that grows up to window_size + - When buffer exceeds window_size, removes oldest embeddings (FIFO) + - Creates a rolling history of recent data for statistical comparison +5. Fits von Mises-Fisher (vMF) distribution parameters during warmup completion: + - **Mean Direction (mu)**: Calculates the average direction of all historical embeddings + - **Concentration Parameter (kappa)**: Measures how tightly clustered the embeddings are around the mean + - Uses statistical estimation to model the distribution of embedding directions + - vMF distribution is ideal for directional data on a hypersphere (unit vectors) +6. Computes alignment score for current embedding: + - Calculates dot product between current normalized embedding and mean direction vector + - Measures how well the current embedding aligns with the typical direction + - Higher values indicate closer alignment to the norm, lower values indicate deviation +7. Calculates empirical percentile of current embedding: + - Computes alignment scores for all historical embeddings against the mean direction + - Ranks the current embedding's alignment score among historical scores + - Determines percentile position (0.0 = lowest, 1.0 = highest) of current embedding +8. Determines outlier status based on percentile thresholds: + - Flags as outlier if percentile is below threshold_percentile (e.g., bottom 5%) + - Flags as outlier if percentile is above (1 - threshold_percentile) (e.g., top 5%) + - Detects both extreme low and extreme high deviations from the norm +9. Returns three outputs: + - **is_outlier**: Boolean flag indicating if the current embedding is an outlier + - **percentile**: Float value (0.0-1.0) representing where the embedding ranks among historical data + - **warming_up**: Boolean flag indicating if still in warmup period (always False after warmup) + +The block uses von Mises-Fisher distribution analysis, which is designed for directional data on a hypersphere (unit vectors). This makes it well-suited for high-dimensional embeddings where direction matters more than magnitude. The sliding window approach ensures the statistical model adapts to recent trends while the percentile-based detection identifies embeddings that are unusually different from the historical pattern. Lower percentiles indicate embeddings that are less aligned with typical patterns, while higher percentiles indicate embeddings that are unusually well-aligned or different in a positive direction. + +## Common Use Cases + +- **Anomaly Detection**: Detect unusual images, objects, or patterns that deviate from normal data (e.g., identify unusual product variations, detect anomalous behavior, flag unexpected patterns), enabling anomaly detection workflows +- **Quality Control**: Identify defective or unusual items in manufacturing or production (e.g., detect product defects, identify quality issues, flag manufacturing anomalies), enabling quality control workflows +- **Data Sampling**: Identify interesting or unusual data points for manual review or further analysis (e.g., sample unusual images for labeling, identify edge cases for model improvement, select interesting data for analysis), enabling intelligent data sampling workflows +- **Change Detection**: Detect when data patterns change significantly from historical norms (e.g., detect scene changes, identify pattern shifts, flag significant variations), enabling change detection workflows +- **Model Monitoring**: Monitor model performance by detecting when embeddings deviate from training distribution (e.g., detect distribution shift, identify out-of-distribution data, monitor model drift), enabling model monitoring workflows +- **Content Filtering**: Identify unusual or inappropriate content that differs from expected patterns (e.g., detect unusual content, flag inappropriate material, identify content anomalies), enabling content filtering workflows + +## Connecting to Other Blocks + +This block receives embeddings and produces is_outlier, percentile, and warming_up outputs: + +- **After embedding model blocks** (CLIP, Perception Encoder, etc.) to analyze embedding outliers (e.g., identify outliers from CLIP embeddings, analyze Perception Encoder outliers, detect anomalies from embeddings), enabling embedding-to-outlier workflows +- **After classification or detection blocks** with embeddings to identify unusual predictions (e.g., identify unusual detections, flag anomalous classifications, detect outlier predictions), enabling prediction-to-outlier workflows +- **Before logic blocks** like Continue If to make decisions based on outlier detection (e.g., continue if outlier detected, filter based on outlier status, trigger actions on anomalies), enabling outlier-based decision workflows +- **Before notification blocks** to alert on outlier detection (e.g., alert on anomalies, notify about unusual data, trigger alerts on outliers), enabling outlier-based notification workflows +- **Before data storage blocks** to record outlier information (e.g., log outlier data, store anomaly statistics, record unusual data points), enabling outlier data logging workflows +- **In quality control pipelines** where outlier detection is part of quality assurance (e.g., filter outliers in quality pipelines, identify issues in production workflows, detect problems in processing chains), enabling quality control workflows + +## Requirements + +This block requires embeddings as input (typically from embedding model blocks like CLIP or Perception Encoder). The block maintains internal state across workflow executions, accumulating a sliding window of historical embeddings. During the warmup period (first `warmup` samples), no outliers are identified and the block returns is_outlier=False and percentile=0.5. After warmup, the block uses at least `warmup` embeddings (up to `window_size` embeddings) to establish statistical baselines. The threshold_percentile parameter (0.0-1.0) controls sensitivity - lower values (e.g., 0.01) detect only extreme outliers, while higher values (e.g., 0.1) detect more moderate deviations. The block works best with consistent embedding models and may need adjustment of threshold_percentile based on expected variation in your data. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Identify Outliers", + "version": "v1", + "short_description": "Identify outlier embeddings compared to prior data.", + "long_description": LONG_DESCRIPTION, + "license": "MIT", + "block_type": "video", + "ui_manifest": { + "section": "video", + "icon": "far fa-chart-scatter-bubble", + }, + } + ) + type: Literal["roboflow_core/identify_outliers@v1"] + name: str = Field(description="Unique name of step in workflows") + + embedding: Selector(kind=[TENSOR_NATIVE_EMBEDDING_KIND]) = Field( + description="Embedding vector representing the current data point's features. Typically from embedding models like CLIP or Perception Encoder. The embedding is normalized to unit length for directional statistical analysis using von Mises-Fisher distribution. Must be a numerical vector of any dimension.", + examples=["$steps.clip.embedding", "$steps.perception_encoder.embedding"], + ) + + threshold_percentile: Union[Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), float] = Field( + default=0.05, + description="Percentile threshold for outlier detection, range 0.0-1.0. Embeddings below this percentile or above (1 - threshold_percentile) are flagged as outliers. Lower values (e.g., 0.01) detect only extreme outliers - very strict. Higher values (e.g., 0.1) detect more moderate deviations - more sensitive. Default 0.05 means bottom 5% and top 5% are outliers. Adjust based on expected variation in your data.", + examples=[0.05, 0.01, 0.1, "$inputs.threshold_percentile"], + json_schema_extra={ + "always_visible": True, + }, + ) + + warmup: Union[Selector(kind=[INTEGER_KIND]), int] = Field( + default=3, + description="Number of initial data points required before outlier detection begins. During warmup, no outliers are identified (is_outlier=False) to allow baseline establishment. Must be at least 2 for statistical analysis. Typical range: 3-100 samples. Higher values provide more stable baselines but delay outlier detection. Lower values enable faster detection but may be less accurate initially.", + examples=[3, 10, 100], + ) + + window_size: Optional[Union[Selector(kind=[INTEGER_KIND]), int]] = Field( + default=32, + description="Maximum number of historical embeddings to maintain in sliding window. The block keeps the most recent window_size embeddings for statistical comparison. When exceeded, oldest embeddings are removed (FIFO). Larger windows provide more stable statistics but adapt slower to distribution changes. Smaller windows adapt faster but may be less stable. Set to None for unlimited window (uses all historical data). Typical range: 10-100 embeddings.", + examples=[32, 64, 100, None], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="is_outlier", kind=[BOOLEAN_KIND]), + OutputDefinition(name="percentile", kind=[FLOAT_ZERO_TO_ONE_KIND]), + OutputDefinition(name="warming_up", kind=[BOOLEAN_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class IdentifyOutliersBlockV1(WorkflowBlock): + def __init__(self): + self.samples = 0 + + # Keep track of all embeddings for vMF parameter estimation: + self.all_embeddings = [] # Store normalized embeddings + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def _fit_vmf_parameters(self, embeddings: torch.Tensor): + """ + Fit a von Mises-Fisher distribution to the given set of unit-normalized embeddings. + Returns: + mu (torch.Tensor): Mean direction vector. + kappa (float): Concentration parameter. + """ + n, d = embeddings.shape + if n < 2: + # Not enough data to fit + return None, None + + # Sum all embeddings: + sum_vec = embeddings.sum(dim=0) + R = torch.linalg.norm(sum_vec) + if R == 0: + # All embeddings canceled out, no direction + return None, None + + mu = sum_vec / R + r_bar = R / n + + # Approximate kappa: + # For d>2, a known approximation: + # kappa โ‰ˆ r_bar*(d - r_bar^2)/(1 - r_bar^2) + d = float(d) + if r_bar < 1.0: + kappa = (r_bar * (d - r_bar**2)) / (1 - r_bar**2) + else: + # Degenerate case: all points are identical + kappa = float("inf") + + return mu, kappa + + def run( + self, + embedding: torch.Tensor, + threshold_percentile: float, + warmup: int, + window_size: int, + ) -> BlockResult: + # Tensor-native embeddings stay on-device: the vMF math runs in torch. + embedding = embedding.detach() + # Normalize embedding for vMF + norm = torch.linalg.norm(embedding) + if norm == 0: + # If zero vector, skip normalization + embedding_normed = embedding + else: + embedding_normed = embedding / norm + + self.samples += 1 + warming_up = self.samples < warmup + + # Store normalized embedding for vMF: + self.all_embeddings.append(embedding_normed) + if len(self.all_embeddings) > window_size: + # Remove oldest embedding + self.all_embeddings.pop(0) + + # If we're still in warmup, we cannot decide outliers yet + if warming_up: + return { + "is_outlier": False, + "percentile": 0.5, + "warming_up": True, + } + + # Fit vMF parameters based on all embeddings so far + all_emb_array = torch.stack(self.all_embeddings) + mu, kappa = self._fit_vmf_parameters(all_emb_array) + if mu is None or kappa is None: + # Fallback if we cannot fit (e.g. all embeddings identical) + mu = embedding_normed + kappa = 0.0 + + # Compute alignment score with the current mean direction + t_new = torch.dot(mu, embedding_normed) + + # Compute empirical percentile of t_new relative to historical t_i + # Using all previous embeddings (excluding the current one if desired) + t_values = all_emb_array @ mu + # Sort t-values to find percentile + sorted_t = torch.sort(t_values).values + rank = int(torch.searchsorted(sorted_t, t_new, right=False)) + percentile = rank / len(sorted_t) + + # Determine outlier based on percentile thresholds + is_outlier = (percentile < threshold_percentile) or ( + percentile > (1 - threshold_percentile) + ) + + return { + "is_outlier": bool(is_outlier), + "percentile": float(percentile), + "warming_up": warming_up, + } diff --git a/inference/core/workflows/core_steps/sinks/onvif_movement/v1_tensor.py b/inference/core/workflows/core_steps/sinks/onvif_movement/v1_tensor.py new file mode 100644 index 0000000000..0baf1feaba --- /dev/null +++ b/inference/core/workflows/core_steps/sinks/onvif_movement/v1_tensor.py @@ -0,0 +1,1049 @@ +import asyncio +import concurrent +import importlib +import os +import threading +import time +from threading import Thread +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import torch +import zeep +from onvif import ONVIFCamera, ONVIFService +from pydantic import ConfigDict, Field, PositiveInt +from simple_pid import PID + +from inference.core import logger +from inference.core.utils.function import experimental +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + take_prediction_by_indices, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + ROOT_PARENT_DIMENSIONS_KEY, + TRACKER_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + SECRET_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + Runtime, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +# max number of seconds to switch to zoom only (no xy movement) +ZOOM_MODE_SECONDS = 2 + +# After the first zoom mode, multiply pan/tilt speed by this much to +# help with control. Will revert once the camera goes back to the preset +# This could be improved in the future by more accurately measuring +# zoom level (note not all cameras can provide coordinates) +ZOOM_MODE_SPEED_REDUCER = 0.5 + +PREDICTIONS_OUTPUT_KEY: str = "predictions" +SEEKING_OUTPUT_KEY: str = "seeking" + +LONG_DESCRIPTION = """ +Control an ONVIF-compatible PTZ (Pan-Tilt-Zoom) camera to automatically follow detected objects, move to preset positions, and maintain smooth tracking using PID control for surveillance, security monitoring, and automated camera control workflows. + +## How This Block Works + +This block controls PTZ cameras through the ONVIF protocol to automatically track and follow objects in real-time. The block: + +1. Receives object detection or instance segmentation predictions from upstream blocks +2. Connects to an ONVIF-compatible PTZ camera using IP address, port, username, and password credentials +3. Determines the movement mode (Follow or Go To Preset): + + **For Follow Mode:** + - Selects the target object to track: either the highest confidence detection or a specific tracked object (if tracker IDs are present) + - Calculates the position error by comparing the object's bounding box center to the frame center + - Uses PID (Proportional-Integral-Derivative) control to calculate smooth movement commands: + - Proportional (Kp): Responds to current position error + - Integral (Ki): Corrects steady-state errors over time + - Derivative (Kd): Predicts future error and dampens oscillations + - Normalizes movement commands to the camera's velocity limits and applies rate limiting to prevent command flooding + - Sends continuous movement commands to the camera via ONVIF ContinuousMove service + - Monitors the dead zone (region around center where camera stops moving) to prevent hunting behavior + - Optionally zooms into the object when it's centered, adjusting zoom speed to fill the frame + - Maintains tracking of a specific object using tracker IDs until it disappears or tracking is reset + - Automatically moves to a preset position after a configurable idle period when no objects are detected + + **For Go To Preset Mode:** + - Moves the camera to a predefined preset position using the ONVIF GotoPreset service + - Requires the camera to support preset functionality and have presets configured + - Uses the preset name specified in the configuration + +4. Handles camera communication asynchronously using a separate event loop to prevent blocking workflow execution +5. Manages camera state including seeking status, tracked object IDs, zoom state, and movement history +6. Returns two outputs: + - **predictions**: The detection being tracked (empty if no object is being followed) + - **seeking**: Boolean indicating whether the camera is currently moving/seeking an object + +The block uses PID control to calculate smooth, proportional movement commands based on the distance between the object center and frame center. Movement is normalized as a percentage of the camera's maximum speed, and commands are rate-limited to prevent overwhelming the camera with updates. The dead zone prevents small movements when the object is near the center, reducing hunting behavior. When zoom is enabled, the camera first centers the object with pan/tilt, then zooms in to fill the frame while maintaining the object in view. + +## Common Use Cases + +- **Surveillance and Security**: Automatically track individuals or vehicles in surveillance scenarios (e.g., follow suspicious activity, track intruders, monitor security perimeters), enabling automated surveillance workflows +- **Sports and Event Coverage**: Track athletes or objects during sports events or performances (e.g., follow players on field, track ball movement, cover event action), enabling automated sports coverage workflows +- **Wildlife Monitoring**: Follow animals or wildlife in natural habitats (e.g., track bird movements, follow animals in reserves, monitor wildlife behavior), enabling wildlife observation workflows +- **Industrial Monitoring**: Automatically follow objects or personnel in industrial settings (e.g., track equipment movement, monitor worker activities, follow vehicles in facilities), enabling industrial automation workflows +- **Traffic Monitoring**: Track vehicles or objects in traffic scenarios (e.g., follow vehicles through intersections, track traffic violations, monitor road activity), enabling automated traffic monitoring workflows +- **Retail Analytics**: Track customers or products in retail environments (e.g., follow customer paths, track product interactions, monitor shopping behavior), enabling retail analytics workflows + +## Connecting to Other Blocks + +This block receives predictions and produces camera control commands and tracking status: + +- **After object detection or instance segmentation blocks** to track detected objects with the camera (e.g., follow detected people, track detected vehicles, monitor detected objects), enabling detection-to-camera-tracking workflows +- **After Byte Tracker blocks** to follow specific tracked objects with consistent IDs (e.g., follow tracked person across frames, maintain tracking of specific vehicle, monitor tracked object persistently), enabling tracking-to-camera workflows +- **After detection filter blocks** to track specific object classes or filtered detections (e.g., track only specific classes, follow filtered detections, monitor selected objects), enabling filtered-tracking workflows +- **Before visualization blocks** to display camera movement status and tracked objects (e.g., visualize tracking status, display seeking indicator, show camera control feedback), enabling camera control visualization workflows +- **Before notification blocks** to alert when camera starts or stops tracking (e.g., notify when tracking begins, alert on tracking loss, report camera status), enabling camera status notification workflows +- **In surveillance and monitoring pipelines** where automated camera control is part of a larger security or monitoring system (e.g., automated security systems, monitoring pipelines, camera control chains), enabling comprehensive surveillance workflows + +## Requirements + +This block requires an ONVIF-compatible PTZ camera with network access. The camera must support the ONVIF ContinuousMove service for Follow mode and GotoPreset service for preset movement. For optimal performance, use a camera with variable speed movement capability - cameras without variable speed can use the simulate_variable_speed option but may experience jerky movement. The block must run in local execution mode (not suitable for remote/cloud execution). PID tuning is recommended to achieve smooth tracking without overshooting or hunting - adjust pid_kp, pid_ki, and pid_kd parameters based on camera responsiveness and video latency. For accurate tracking, use an eager buffer consumption strategy to minimize lag between camera movement and video feedback. The camera must have presets configured if using preset movement or auto-reset functionality. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "PTZ Tracking (ONVIF)", + "version": "v1", + "short_description": "Control an ONVIF compatible PTZ camera to follow an object", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "sink", + "ui_manifest": { + "section": "video", + "icon": "fal fa-camera-cctv", + "blockPriority": 1, + "popular": False, + }, + } + ) + type: Literal["roboflow_core/onvif_sink@v1"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Object detection or instance segmentation predictions to track. In Follow mode, the block will follow the highest confidence prediction or a tracked object if tracker IDs are present. Predictions should include bounding box coordinates and optionally tracker IDs for persistent tracking.", + examples=["$steps.object_detection_model.predictions"], + ) + camera_ip: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="Camera IP address or hostname for ONVIF connection. Must be reachable from the workflow execution environment.", + ) + camera_port: Union[Selector(kind=[INTEGER_KIND]), PositiveInt] = Field( + description="Camera ONVIF service port (typically 80, 8080, or camera-specific port). Must match the camera's ONVIF configuration.", + ge=0, + le=65535, + ) + camera_username: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="Camera username for ONVIF authentication. Must have PTZ control permissions on the camera.", + ) + camera_password: Union[Selector(kind=[SECRET_KIND]), str] = Field( + description="Camera password for ONVIF authentication. Should be stored as a secret for security.", + ) + movement_type: Literal["Follow", "Go To Preset"] = Field( + default="Follow", + description="Movement mode for the camera. 'Follow' mode tracks detected objects using PID control. 'Go To Preset' mode moves the camera to a predefined preset position (requires default_position_preset to be configured).", + examples=["Follow", "Go To Preset", "$inputs.movement_type"], + ) + simulate_variable_speed: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Enable variable speed simulation for cameras without native variable speed support. When enabled, sends 100% speed commands followed by stop commands to approximate percentage speeds. May result in jerky movement - only use if camera lacks variable speed capability.", + examples=[True, False, "$inputs.simulate_variable_speed"], + ) + zoom_if_able: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="Enable automatic zoom to fill frame with the tracked object. When enabled, camera first centers the object with pan/tilt, then zooms in until the object fills the frame. Requires camera to support zoom functionality.", + examples=[True, False, "$inputs.zoom_if_able"], + ) + follow_tracker: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Enable persistent tracking using tracker IDs. When enabled, camera locks onto the tracker ID of the highest confidence detection and continues following that specific object until it disappears or tracking resets. Requires a Byte Tracker block in the workflow to assign tracker IDs.", + examples=[True, False, "$inputs.follow_tracker"], + ) + dead_zone: Union[Selector(kind=[INTEGER_KIND]), int] = Field( + default=50, + description="Dead zone size in pixels around frame center where camera movement stops. Prevents hunting behavior when object is near center. Larger values reduce pan/tilt hunting but may cause zoom hunting. Smaller values improve zoom stability but may cause pan/tilt oscillations. Typical range: 30-100 pixels.", + examples=[50, "$inputs.dead_zone"], + ) + default_position_preset: Union[Selector(kind=[STRING_KIND]), str] = Field( + description="Preset name for default/home position. Camera will return to this preset after idle period (if move_to_position_after_idle_seconds is set) or when using Go To Preset mode. Must match exactly a preset name configured on the camera. Required for preset movement functionality.", + default="", + examples=["", "$inputs.default_position_preset"], + ) + move_to_position_after_idle_seconds: Union[Selector(kind=[INTEGER_KIND]), int] = ( + Field( + default=0, + description="Auto-reset time in seconds. After camera stops seeking/moving for this duration, automatically moves to default_position_preset. Set to 0 to disable auto-reset. Requires default_position_preset to be configured.", + ) + ) + camera_update_rate_limit: Union[Selector(kind=[INTEGER_KIND]), int] = Field( + default=250, + description="Minimum time in milliseconds between camera movement commands. Rate limits ONVIF updates to prevent overwhelming the camera. Lower values provide more responsive tracking but may overload slower cameras. Higher values reduce camera load but may cause less smooth movement. Typical range: 100-500ms.", + ) + flip_x_movement: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + examples=[True, False], + description="Invert horizontal (pan) movement direction. Enable if the camera image is mirrored horizontally and movement appears reversed. Use to correct camera movement when image is flipped.", + ) + flip_y_movement: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + examples=[True, False], + description="Invert vertical (tilt) movement direction. Enabled by default as many cameras have inverted Y-axis. Disable if vertical movement appears reversed.", + ) + minimum_camera_speed: Union[float, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( + default=0.05, + description="Minimum movement speed as percentage (0.0-1.0). Movement commands below this threshold are boosted to this minimum. Some cameras ignore very low speeds - increase if camera doesn't respond to small movements. Typical range: 0.02-0.1 (2%-10%).", + ) + pid_kp: Union[float, Selector(kind=[FLOAT_KIND])] = Field( + default=0.25, + description="PID proportional gain (Kp). Controls response to position error - higher values make camera respond faster but may cause overshooting and hunting. Lower values reduce hunting but make tracking slower. Start with default and adjust based on camera responsiveness. Typical range: 0.1-0.5.", + ) + pid_ki: Union[float, Selector(kind=[FLOAT_KIND])] = Field( + default=0.0, + description="PID integral gain (Ki). Eliminates steady-state error by accumulating error over time. Usually kept at 0 as it can cause oscillations. Increase slightly (0.01-0.1) if camera consistently stops slightly off-center despite small errors.", + ) + pid_kd: Union[float, Selector(kind=[FLOAT_KIND])] = Field( + default=1, + description="PID derivative gain (Kd). Predicts future error and dampens oscillations. Higher values improve stability with video lag but excessive values can cause hunting. Increase (1-5) if there's significant delay between camera movement and video feedback. Decrease if tracking appears jerky.", + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=PREDICTIONS_OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + ], + ), + OutputDefinition( + name=SEEKING_OUTPUT_KEY, + kind=[ + BOOLEAN_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + no_remote_step_execution = RuntimeRestriction( + severity=Severity.HARD, + note=( + "Block requires step_execution_mode=local; raises ValueError " + "otherwise. ONVIF commands must be issued from the same " + "process that drives the workflow." + ), + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + ) + no_lan_from_hosted = RuntimeRestriction( + severity=Severity.HARD, + note=( + "Block requires LAN access to a PTZ camera. Hosted Serverless " + "and Roboflow Dedicated Deployments cannot reach customer LANs." + ), + applies_to_runtimes=[ + Runtime.HOSTED_SERVERLESS, + Runtime.DEDICATED_DEPLOYMENT, + ], + ) + return [no_remote_step_execution, no_lan_from_hosted] + + +# primarily used for rate limiting +def now() -> int: + return int(round(time.time() * 1000)) + + +# Builds a native-empty `Detections` carrying over the source `image_metadata` +# (the tensor-native analogue of `sv.Detections.empty()`). Shared-helper +# candidate โ€” the maintainer may consolidate this into `tensor_native`. +def _empty_detections_like( + predictions: Union[Detections, InstanceDetections], +) -> Detections: + device = predictions.xyxy.device + return Detections( + xyxy=torch.empty((0, 4), dtype=torch.float32, device=device), + class_id=torch.empty((0,), dtype=torch.long, device=device), + confidence=torch.empty((0,), dtype=torch.float32, device=device), + image_metadata=predictions.image_metadata, + bboxes_metadata=None, + ) + + +class Limits: + min: float + max: float + + def __init__(self, range): + self.min = range.Min + self.max = range.Max + + def __repr__(self): + return f"({self.min},{self.max})" + + def __eq__(self, value): + value.min == self.min and value.max == self.max + + +class VelocityLimits: + x: Limits + y: Limits + z: Union[Limits, None] + + def __init__(self, x: Limits, y: Limits, z: Limits): + self.x = x + self.y = y + self.z = z + + def __repr__(self): + return f"x:{self.x} y:{self.y} z:{self.z}" + + +def run_loop(loop): + asyncio.set_event_loop(loop) + loop.run_forever() + + +# camera wrapper is used to store config info so that we don't have +# to keep querying it from the camera on successive commands +class CameraWrapper: + camera: Optional[ONVIFCamera] = None + thread: Thread + run_loop: asyncio.AbstractEventLoop = None + media_profile = zeep.AnyObject + configuration_options = zeep.AnyObject + _media_profile_token: str = None + tracked_object: Optional[int] = None + _velocity_limits: Optional[VelocityLimits] + _presets: Optional[Dict[str, zeep.AnyObject]] = None + _last_update_ms: Optional[int] = None + _max_update_rate: int = 0 + _move_to_position_after_idle_seconds: int = 0 + _existing_preset_task: Optional[asyncio.Task] = None + _seeking: Optional[bool] = None + _prev_x: float = 100 + _prev_y: float = 100 + _prev_z: float = 100 + _start_zoom_time: Optional[int] = ( + None # don't allow xy movements when not None, zoom only + ) + _stop_preset: Optional[str] = None + _is_zoomed: bool = False + _has_config_error: bool = False + _x_count: int = 0 + _y_count: int = 0 + _z_count: int = 0 + + @experimental( + reason="Usage of CameraWrapper is an experimental feature. Please report any issues " + "here: https://github.com/roboflow/inference/issues" + ) + # create a new camera wrapper with an asyncio event loop + def __init__( + self, + max_update_rate: int, + move_to_position_after_idle_seconds: int, + run_loop: asyncio.AbstractEventLoop, + ): + self._max_update_rate = max_update_rate + self._move_to_position_after_idle_seconds = move_to_position_after_idle_seconds + self.run_loop = run_loop + + @classmethod + def create_event_loop(cls) -> asyncio.AbstractEventLoop: + async_run_loop = asyncio.new_event_loop() + thread = threading.Thread(target=run_loop, args=(async_run_loop,), daemon=True) + thread.start() + return async_run_loop + + def connect_camera( + self, camera_ip, camera_port, camera_username, camera_password + ) -> concurrent.futures._base.Future[None]: + return self.schedule( + self.connect_camera_async( + camera_ip, camera_port, camera_username, camera_password + ) + ) + + async def connect_camera_async( + self, camera_ip, camera_port, camera_username, camera_password + ): + if not self.camera: + # wsdls are in package directory, ex: "/usr/local/lib/python3.9/site-packages/onvif/wsdl" + spec = importlib.util.find_spec("onvif") + wdsl_path = f"{os.path.dirname(spec.origin)}/wsdl" + self.camera = ONVIFCamera( + camera_ip, camera_port, camera_username, camera_password, wdsl_path + ) + await self.configure_async() + else: + logger.debug("camera is already connected") + + def set_stop_preset(self, stop_preset: Optional[str]): + self._stop_preset = stop_preset + + # schedule a future inside the camera's event loop + def schedule(self, cor) -> concurrent.futures._base.Future[None]: + return asyncio.run_coroutine_threadsafe(cor, loop=self.run_loop) + + # pushes out the next scheduled move to preset (reset) + def schedule_next_reset(self): + if self._stop_preset: + self.schedule(self.next_reset()) + + def clear_next_reset(self): + if self._existing_preset_task: + self._existing_preset_task.cancel() + + async def next_reset(self): + self.clear_next_reset() + if self._stop_preset: + self._existing_preset_task = asyncio.create_task(self.reset_task()) + await self._existing_preset_task + + # sleep will expire after camera hasn't moved for idle seconds, and move to preset + async def reset_task(self): + await asyncio.sleep(self._move_to_position_after_idle_seconds) + logger.debug( + f'camera is idle for {self._move_to_position_after_idle_seconds}s: moving to preset "{self._stop_preset}"' + ) + # if we've been tracking an object, we want to clear it here + self.tracked_object = None + # note stop preset can be cleared even after task is scheduled + if self._stop_preset: + await self.go_to_preset_async(self._stop_preset) + + # true if movement update hasn't happened within max_update_rate + def _can_update(self) -> bool: + return ( + self._last_update_ms is None + or now() - self._last_update_ms > self._max_update_rate + ) + + # this is mainly used to allow stop commands through on new zero speeds + def save_last_speeds(self, x, y, z) -> Tuple[bool, bool, bool]: + x_changed = x != self._prev_x + y_changed = y != self._prev_y + z_changed = z != self._prev_z + self._prev_x = x + self._prev_y = y + self._prev_z = z + return (x_changed, y_changed, z_changed) + + async def ptz_service(self) -> Union[None, ONVIFService]: + """ + Creates the ONVIF PTZ service + This has to run on every command requiring the service - a service can't be awaited twice + """ + if not self.camera: + return None + return await self.camera.create_ptz_service() + + async def media_service(self) -> Union[None, ONVIFService]: + """ + Creates the ONVIF media service + This is primarily used to get the media token + """ + if not self.camera: + return None + return await self.camera.create_media_service() + + async def configure_async(self): + """ + Does initial configuration and gathers all camera info + Doesn't currently run in init since it needs to be awaited + """ + if not self.camera: + raise ValueError(f"Tried to configure camera, but camera was not created") + if self._has_config_error: + return + + try: + await self.camera.update_xaddrs() + # capabilities = await self.camera.get_capabilities() # <- could be useful in the future + media = await self.media_service() + self.media_profile = (await media.GetProfiles())[0] + ptz = await self.ptz_service() + config_request = ptz.create_type("GetConfigurationOptions") + config_request.ConfigurationToken = ( + self.media_profile.PTZConfiguration.token + ) + config_options = ptz.GetConfigurationOptions(config_request) + self.configuration_options = await config_options + pan_tilt_space = ( + self.configuration_options.Spaces.ContinuousPanTiltVelocitySpace[0] + if hasattr( + self.configuration_options.Spaces, "ContinuousPanTiltVelocitySpace" + ) + and len( + self.configuration_options.Spaces.ContinuousPanTiltVelocitySpace + ) + > 0 + else None + ) + # pan tilt space is necessary + if pan_tilt_space is None: + logger.error("Could not get pan tilt space for camera") + raise ValueError("Could not get pan tilt space for camera") + zoom_space = ( + self.configuration_options.Spaces.ContinuousZoomVelocitySpace[0] + if hasattr( + self.configuration_options.Spaces, "ContinuousZoomVelocitySpace" + ) + and len(self.configuration_options.Spaces.ContinuousZoomVelocitySpace) + > 0 + else None + ) + self._velocity_limits = VelocityLimits( + x=Limits(pan_tilt_space.XRange), + y=Limits(pan_tilt_space.YRange), + z=Limits(zoom_space.XRange) if zoom_space else None, + ) + self._media_profile_token = self.media_profile.token + presets = await ptz.GetPresets({"ProfileToken": self._media_profile_token}) + # reconfigure into a dict keyed by preset name + self._presets = {preset["Name"]: preset for preset in presets} + except Exception as e: + self._has_config_error = True + logger.warning(f"Config error: {e}") + + def seeking(self) -> Optional[bool]: + return self._seeking + + # stops the camera movement + def stop_camera(self): + # don't stop if already stopped as stop commands aren't rate limited + if self._seeking or self._seeking is None: + logger.debug("stop camera") + self.continuous_move(0, 0, 0) + self.schedule_next_reset() + + self._seeking = False + + def go_to_preset(self, preset_name: str, limit_rate: bool = False): + self.schedule(self.go_to_preset_async(preset_name, limit_rate)) + + async def go_to_preset_async(self, preset_name: str, limit_rate: bool = False): + """ + Tells the camera to move to a preset + This is not rate limited - all commands will be sent to the camera + + Args: + preset_name: The preset name to move to (this varies by camera) + """ + + # this is only used for blocks in go to preset mode + if limit_rate: + if not self._can_update(): + return + self._last_update_ms = now() + + if self._media_profile_token is None: + await self.configure_async() + + preset = self._presets.get(preset_name) + if not preset: + # TODO: since this is thrown in another thread, it isn't getting thrown, print to logs + logger.error( + f'Camera does not have preset "{preset_name}" - valid presets are {list(self._presets.keys())}' + ) + raise ValueError( + f'Camera does not have preset "{preset_name}" - valid presets are {list(self._presets.keys())}' + ) + + ptz = await self.ptz_service() + request = ptz.create_type("GotoPreset") + request.ProfileToken = self._media_profile_token + request.PresetToken = preset["token"] + request.Speed = {"PanTilt": {"x": 1.0, "y": 1.0}, "Zoom": {"x": 1.0}} + + self._seeking = False + self._is_zoomed = False + self.tracked_object = None + await ptz.GotoPreset(request) + + def zoom(self, z: float): + # start limited time zoom mode + if not self._start_zoom_time: + self._start_zoom_time = now() + # even though the zoom command should 0 out pan/tilt, sending an explicit stop on all axes seems to help + self.stop_camera() + self.schedule(self.continuous_move_async(0, 0, z)) + + def stop_zoom(self): + if self._start_zoom_time is not None: + self._start_zoom_time = None + + def zooming(self) -> bool: + global ZOOM_MODE_SECONDS + if ( + self._start_zoom_time + and now() > self._start_zoom_time + ZOOM_MODE_SECONDS * 1000 + ): + self.stop_zoom() + return self._start_zoom_time is not None + + # tells the camera to move at a continuous velocity + def continuous_move( + self, x: float, y: float, z: float, simulate_variable_speed: bool = False + ): + self.schedule(self.continuous_move_async(x, y, z, simulate_variable_speed)) + + def simulate_variable_speed(self, speed: float, count: int) -> Tuple[float, int]: + count = count + 1 + + if speed != 0 and count >= int(1.0 / speed): + speed = np.sign(speed) + if self._can_update(): + count = 0 + else: + speed = 0 + + # stop count until we let the next update through + return speed, count + + # x and y are velocities from -1 to 1 + async def continuous_move_async( + self, x: float, y: float, z: float, simulate_variable_speed: bool = False + ): + """ + Tells the camera to move at a continuous velocity, or 0 to stop + Note this is rate limited, some commands will be ignored + + Args: + x: The x velocity normalized as -1 to 1 + y: The y velocity normalized as -1 to 1 + z: The zoom velocity normalized as -1 to 1 + """ + + if self._media_profile_token is None: + await self.configure_async() + + # clear out any scheduled position resets + # they'll be rescheduled on the next stop + if x != 0 and y != 0 and z != 0: + self.clear_next_reset() + + # This option simulates a % speed by sending a number of move and stop + # commands that approximate the required % - so for 25% we'll send + # one 100% move command followed by one stop command. + if simulate_variable_speed: + x, self._x_count = self.simulate_variable_speed(x, self._x_count) + y, self._y_count = self.simulate_variable_speed(y, self._y_count) + z, self._z_count = self.simulate_variable_speed(z, self._z_count) + + # try to avoid hunting by allowing immediate stop + x_changed, y_changed, z_changed = self.save_last_speeds(x, y, z) + + # don't rate limit stop commands + if (x == 0 and x_changed) or (y == 0 and y_changed) or (z == 0 and z_changed): + pass + elif not self._can_update(): + return + + ptz = await self.ptz_service() + + # https://www.onvif.org/onvif/ver20/ptz/wsdl/ptz.wsdl#op.AbsoluteMove + + request = ptz.create_type("ContinuousMove") + request.ProfileToken = self._media_profile_token + + limits = self._velocity_limits + + # normalize to camera's velocity limits + x_limit = limits.x.min if x < 0 else limits.x.max + y_limit = limits.y.min if x < 0 else limits.y.max + if limits.z: + z_limit = limits.z.min if x < 0 else limits.z.max + else: + z_limit = 0 + + x = abs(x_limit) * x + y = abs(y_limit) * y + z = abs(z_limit) * z + + if self._is_zoomed: + x = x * ZOOM_MODE_SPEED_REDUCER + y = y * ZOOM_MODE_SPEED_REDUCER + + request.Velocity = {"PanTilt": {"x": x, "y": y}, "Zoom": {"x": z}} + + logger.debug( + f"ptz continuous move update: {x},{y},{z} in_zoom:{self._is_zoomed} tracker:{self.tracked_object}" + ) + + # Execute the movement + await ptz.ContinuousMove(request) + self._last_update_ms = now() + + if z > 0: + self._is_zoomed = True + + # prevent stops from re-flagging themselves as seeking + if x != 0 or y != 0 or z != 0: + self._seeking = True + + def stop_tracking(self): + self.stop_camera() + self.tracked_object = None + + # this just sends a stop command - should be ok if other blocks are controlling camera + def __del__(self): + self.stop_camera() + + +class ONVIFSinkBlockV1(WorkflowBlock): + + def __init__( + self, + step_execution_mode: StepExecutionMode, + disable_sinks: bool = False, + ): + self._step_execution_mode = step_execution_mode + self._disable_sinks = disable_sinks + # all commands will be send to the camera normalized to -1 to 1 + # all setpoints are 0, which represents the center of the frame + self.x_pid = PID(0, 0, 0, setpoint=0) + self.x_pid.output_limits = (-1, 1) + self.y_pid = PID(0, 0, 0, setpoint=0) + self.y_pid.output_limits = (-1, 1) + self.z_pid = PID(0, 0, 0, setpoint=0) + self.z_pid.output_limits = (-1, 1) + self.event_loop = CameraWrapper.create_event_loop() + # pool of camera services can can be used in block + self.cameras: Dict[Tuple[str, int], CameraWrapper] = {} + self._lock = threading.Lock() + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["step_execution_mode", "disable_sinks"] + + # gets the CameraWrapper from the static cameras collection + def get_camera( + self, + camera_ip: str, + camera_port: int, + camera_username: str, + camera_password: str, + max_update_rate: int, + move_to_position_after_idle_seconds: int, + event_loop: asyncio.AbstractEventLoop, + ) -> Optional[CameraWrapper]: + cameras = self.cameras + camera_key = (camera_ip, camera_port) + with self._lock: + if camera_key not in cameras: + try: + cameras[camera_key] = CameraWrapper( + max_update_rate, move_to_position_after_idle_seconds, event_loop + ) + cameras[camera_key].connect_camera( + camera_ip, camera_port, camera_username, camera_password + ) + except Exception as e: + if camera_key in cameras: + del cameras[camera_key] + raise ValueError( + f"Error connecting to camera at {camera_ip}:{camera_port}: {e}" + ) + return cameras.get(camera_key) + + def run( + self, + predictions: Union[Detections, InstanceDetections], + camera_ip: str, + camera_port: int, + camera_username: str, + camera_password: str, + movement_type: str, + default_position_preset: Union[str, None], + zoom_if_able: bool, + follow_tracker: bool, + dead_zone: int, + camera_update_rate_limit: int, + flip_y_movement: bool, + flip_x_movement: bool, + move_to_position_after_idle_seconds: int, + pid_kp: float, + pid_ki: float, + pid_kd: float, + minimum_camera_speed: float, + simulate_variable_speed: bool, + ) -> BlockResult: + if self._disable_sinks: + return { + PREDICTIONS_OUTPUT_KEY: predictions, + SEEKING_OUTPUT_KEY: False, + } + + # this is hard coded: if intermittent move signals are less + # than 10% then it's unlikely the camera will ever move + if simulate_variable_speed: + minimum_camera_speed = max(minimum_camera_speed, 0.1) + + if self._step_execution_mode != StepExecutionMode.LOCAL: + raise ValueError("Inference must be run locally for the ONVIF block") + + if move_to_position_after_idle_seconds and not default_position_preset: + raise ValueError( + "Move to position after idle is set, but no default position is set" + ) + + # disable the stop preset if necessary + stop_preset = ( + default_position_preset if move_to_position_after_idle_seconds else None + ) + + camera = self.get_camera( + camera_ip, + camera_port, + camera_username, + camera_password, + camera_update_rate_limit, + move_to_position_after_idle_seconds, + self.event_loop, + ) + camera.set_stop_preset(stop_preset) + + if movement_type == "Follow": + + # for v1 use the same constants for all axes + self.x_pid.tunings = (pid_kp, pid_ki, pid_kd) + self.y_pid.tunings = (pid_kp, pid_ki, pid_kd) + self.z_pid.tunings = (pid_kp, pid_ki, pid_kd) + + if len(predictions.xyxy) == 0: + # get/create the camera first so that we can move it to the preset + camera.stop_tracking() + + return { + PREDICTIONS_OUTPUT_KEY: _empty_detections_like(predictions), + SEEKING_OUTPUT_KEY: camera.seeking() if camera else False, + } + + tracked_object = camera.tracked_object + + max_confidence_prediction = None + + # tracker_id lives in per-detection bboxes_metadata for tensor-native + # predictions (numpy used predictions.tracker_id). + bboxes_metadata = predictions.bboxes_metadata or [ + {} for _ in range(len(predictions)) + ] + tracker_ids = [ + box_metadata.get(TRACKER_ID_KEY) for box_metadata in bboxes_metadata + ] + + # if there's a tracked object, continue to use it + if tracked_object: + tracked_mask = [ + tracker_id == tracked_object for tracker_id in tracker_ids + ] + tracked_predictions = take_prediction_by_mask(predictions, tracked_mask) + if len(tracked_predictions.xyxy) > 0: + max_confidence_prediction = take_prediction_by_indices( + tracked_predictions, [0] + ) + + # if there's no tracked object, use the max confidence prediction + if max_confidence_prediction is None: + max_confidence = predictions.confidence.max() + max_confidence_index = int( + torch.nonzero(predictions.confidence == max_confidence)[0] + ) + max_confidence_prediction = take_prediction_by_indices( + predictions, [max_confidence_index] + ) + # if we're not tracking at the moment, start tracking this one + max_confidence_tracker_id = ( + max_confidence_prediction.bboxes_metadata or [{}] + )[0].get(TRACKER_ID_KEY) + if follow_tracker and max_confidence_tracker_id is not None: + tracked_object = int(max_confidence_tracker_id) + + camera.tracked_object = tracked_object + + # adjust PID here as necessary, and send commands to wrapper + self.move_camera( + camera, + max_confidence_prediction, + zoom_if_able, + dead_zone, + flip_x_movement, + flip_y_movement, + minimum_camera_speed, + simulate_variable_speed, + ) + + return { + PREDICTIONS_OUTPUT_KEY: max_confidence_prediction, + SEEKING_OUTPUT_KEY: camera.seeking() if camera else False, + } + + elif movement_type == "Go To Preset": + camera.stop_camera() + camera.go_to_preset(default_position_preset, True) + + return { + PREDICTIONS_OUTPUT_KEY: _empty_detections_like(predictions), + SEEKING_OUTPUT_KEY: camera.seeking() if camera else False, + } + + def move_camera( + self, + camera: CameraWrapper, + prediction: Union[Detections, InstanceDetections], + zoom_if_able: bool, + dead_zone: int, + flip_x_movement: bool, + flip_y_movement: bool, + minimum_camera_speed: float, + simulate_variable_speed: bool = False, + ): + """ + This is where the PID changes are adjusted before the movement + commands are sent to CameraWrapper + + Args: + camera: CameraWrapper + prediction: The prediction containing a single object to point the camera to + zoom_if_able: True to try zooming in, False to ignore + dead_zone: The camera will send stop commands when the object is within this zone + flip_x_movement: Use to reverse the sign on the movement command if image is flipped + flip_y_movement: Use to reverse the sign on the movement command if image is flipped + minimum_camera_speed: Expressed as % from 0-1, movement commands are limited to this minimum + stop_preset: The name of the preset for the camera to go to after being idle + simulate_variable_speed: All speeds will be max % with start/stops used to simulate speed + """ + + # get dimensions of image and bounding box from prediction. For + # tensor-native predictions root_parent_dimensions is a single + # (height, width) in image_metadata (numpy stored an (N, 2) array + # in prediction.data and indexed the first row). + image_dimensions = (prediction.image_metadata or {})[ROOT_PARENT_DIMENSIONS_KEY] + xyxy = prediction.xyxy.detach().to("cpu").numpy() + + # calculate centers + x1, y1, x2, y2 = tuple(xyxy[0]) + image_height, image_width = tuple(image_dimensions) + center_point = (x1 + (x2 - x1) / 2, y1 + (y2 - y1) / 2) + + # calculate deltas from center and edge + zoom_delta = min([x1, image_width - x2, y1, image_height - y2]) + box_at_edge = zoom_delta <= 1 # allow 1px tolerance + # make the deltas x, y, zoom + delta = (image_width / 2 - center_point[0], image_height / 2 - center_point[1]) + + # if we're locked into zoom only mode, and the object goes to the edge, unlock it and go back to pan/tilt + if camera.zooming() and zoom_delta < dead_zone and not box_at_edge: + camera.stop_zoom() + + abs_delta = np.abs(delta) + + if np.all(abs_delta < dead_zone) or camera.zooming(): + # if within tolerance or currently zooming, then camera is in zoom mode if available + # this can continue for up to 5 seconds, or as long as the box is at the edge + if zoom_if_able: + if zoom_delta < dead_zone and not box_at_edge: + camera.stop_camera() + else: + # we want zoom_delta just past the dead zone, but not at the edge + # zoom delta normalized to % image like with pan/tilt + # this means constants mostly stay the same regardless of image size + normalized_zoom_delta = abs(zoom_delta - dead_zone / 2) / max( + image_width, image_height + ) + control_output_z = self.z_pid(normalized_zoom_delta) + + if abs(control_output_z) < minimum_camera_speed: + control_output_z = minimum_camera_speed * np.sign( + control_output_z + ) + + logger.debug( + f"zdelta:{normalized_zoom_delta} output:{control_output_z}" + ) + + # in the case where we've overshot, there's no signal to use for PID + # back off slowly if the box is at the edge as we likely just missed it + camera.zoom( + minimum_camera_speed * -1 + if box_at_edge + else control_output_z * -1 + ) + + else: + camera.stop_camera() + else: + # if not in zoom mode, then allow xy (pan/tilt) motion + + # normalize delta in terms of % for PID loop + # this means constants mostly stay the same regardless of image size + normalized_delta = delta / np.array([image_width, image_height]) + + control_output_x = self.x_pid(normalized_delta[0]) + control_output_y = self.y_pid(normalized_delta[1]) + + if abs(control_output_x) < minimum_camera_speed: + control_output_x = minimum_camera_speed * np.sign(control_output_x) + if abs(control_output_y) < minimum_camera_speed: + control_output_y = minimum_camera_speed * np.sign(control_output_y) + + logger.debug( + f"delta:{normalized_delta} output:{control_output_x} {control_output_y}" + ) + + # larger axis moves at max speed, so normalize to 100% + speeds = abs(delta / abs_delta.max()) + + # hard stop within deadzone seems to help, even though first pass will likely + # overshoot if there's a lot of lag + x = speeds[0] * control_output_x if abs_delta[0] > dead_zone else 0 + y = speeds[1] * control_output_y if abs_delta[1] > dead_zone else 0 + + # flip movement as necessary based on settings + # this is actually reversed (delta is backwards) + x_modifier = 1 if flip_x_movement else -1 + y_modifier = 1 if flip_y_movement else -1 + + camera.continuous_move( + x * x_modifier, y * y_modifier, 0, simulate_variable_speed + ) + + def __del__(self): + self.event_loop.stop() diff --git a/inference/core/workflows/core_steps/sinks/roboflow/custom_metadata/v1_tensor.py b/inference/core/workflows/core_steps/sinks/roboflow/custom_metadata/v1_tensor.py new file mode 100644 index 0000000000..d7cbeba724 --- /dev/null +++ b/inference/core/workflows/core_steps/sinks/roboflow/custom_metadata/v1_tensor.py @@ -0,0 +1,345 @@ +import hashlib +import logging +from concurrent.futures import ThreadPoolExecutor +from functools import partial +from typing import List, Literal, Optional, Tuple, Type, Union + +from fastapi import BackgroundTasks +from pydantic import ConfigDict, Field + +from inference.core.cache.base import BaseCache +from inference.core.roboflow_api import add_custom_metadata, get_roboflow_workspace +from inference.core.workflows.core_steps.common.tensor_native import ( + KeyPointPrediction, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.sinks.noop import disabled_sink_response +from inference.core.workflows.execution_engine.constants import INFERENCE_ID_KEY +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + AirGappedAvailability, + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +SHORT_DESCRIPTION = "Add custom metadata to the Roboflow Model Monitoring dashboard." + +LONG_DESCRIPTION = """ +Attach custom metadata fields to inference results in the Roboflow Model Monitoring dashboard by extracting inference IDs from predictions and adding name-value pairs that enable filtering, analysis, and organization of inference data for monitoring workflows, production analytics, and model performance tracking. + +## How This Block Works + +This block adds custom metadata to inference results stored in Roboflow Model Monitoring, allowing you to attach contextual information to predictions for filtering and analysis. The block: + +1. Receives model predictions and metadata configuration: + - Takes predictions from any supported model type (object detection, instance segmentation, keypoint detection, or classification) + - Receives field name and field value for the custom metadata to attach + - Accepts fire-and-forget flag for execution mode +2. Validates Roboflow API key: + - Checks that a valid Roboflow API key is available (required for API access) + - Raises an error if API key is missing with instructions on how to retrieve one +3. Extracts inference IDs from predictions: + - For supervision Detections objects: extracts inference IDs from the data dictionary + - For classification predictions: extracts inference ID from the prediction dictionary + - Collects all unique inference IDs that need metadata attached + - Handles cases where no inference IDs are found (returns error message) +4. Retrieves workspace information: + - Gets workspace ID from Roboflow API using the provided API key + - Uses caching (15-minute expiration) to avoid repeated API calls for workspace lookup + - Caches workspace name using MD5 hash of API key as cache key +5. Adds custom metadata via API: + - Calls Roboflow API to attach custom metadata field to each inference ID + - Associates the field name and field value with the inference results + - Metadata becomes available in the Model Monitoring dashboard for filtering and analysis +6. Executes synchronously or asynchronously: + - **Asynchronous mode (fire_and_forget=True)**: Submits task to background thread pool or FastAPI background tasks, allowing workflow to continue without waiting for API call to complete + - **Synchronous mode (fire_and_forget=False)**: Waits for API call to complete and returns immediate status, useful for debugging and error handling +7. Returns status information: + - Outputs error_status indicating success (False) or failure (True) + - Outputs message with upload status or error details + - Provides feedback on whether metadata was successfully attached + +The block enables attaching custom metadata to inference results, making it easier to filter and analyze predictions in the Model Monitoring dashboard. For example, you can attach location labels, quality scores, processing flags, or any other contextual information that helps organize and analyze your inference data. + +## Common Use Cases + +- **Location-Based Filtering**: Attach location metadata to inferences for geographic analysis and filtering (e.g., tag inferences with location labels like "toronto", "warehouse_a", "production_line_1"), enabling location-based monitoring workflows +- **Quality Control Tagging**: Attach quality or validation metadata to inferences for quality tracking (e.g., tag inferences as "pass", "fail", "requires_review", "approved"), enabling quality control workflows +- **Contextual Annotation**: Add contextual information to inferences for better organization and analysis (e.g., tag with camera ID, time period, batch number, operator ID, environmental conditions), enabling contextual analysis workflows +- **Classification Enhancement**: Attach custom labels or categories to inference results beyond model predictions (e.g., tag with business logic outcomes, workflow decisions, user feedback, manual corrections), enabling enhanced classification workflows +- **Production Analytics**: Track production metrics by attaching metadata that represents operational context (e.g., tag with shift information, production batch, equipment status, performance metrics), enabling production analytics workflows +- **Filtering and Segmentation**: Enable advanced filtering in Model Monitoring dashboard by attaching metadata that represents data segments (e.g., tag with customer segment, product category, use case type, deployment environment), enabling segmentation workflows + +## Connecting to Other Blocks + +This block receives predictions and outputs status information: + +- **After model blocks** (Object Detection Model, Instance Segmentation Model, Classification Model, Keypoint Detection Model) to attach metadata to inference results (e.g., add location tags to detections, attach quality labels to classifications, tag keypoint detections with context), enabling model-to-metadata workflows +- **After filtering or analytics blocks** (DetectionsFilter, ContinueIf, OverlapFilter) to tag filtered or analyzed results with metadata (e.g., tag filtered detections with filter criteria, attach analytics results as metadata, label processed results with workflow state), enabling analysis-to-metadata workflows +- **After conditional execution blocks** (ContinueIf, Expression) to attach metadata based on workflow decisions (e.g., tag with decision outcomes, attach conditional branch labels, mark results based on conditions), enabling conditional-to-metadata workflows +- **In parallel with other sink blocks** to combine metadata tagging with other data storage operations (e.g., tag while uploading to dataset, attach metadata while logging, combine with webhook notifications), enabling parallel sink workflows +- **Before or after visualization blocks** to ensure metadata is attached before or after visualization operations (e.g., tag visualizations with context, attach metadata to visualized results), enabling visualization workflows with metadata +- **At workflow endpoints** to ensure all inference results are tagged with metadata before workflow completion (e.g., final metadata attachment, comprehensive result tagging, complete metadata coverage), enabling end-to-end metadata workflows + +## Requirements + +This block requires a valid Roboflow API key configured in the environment or workflow configuration. The API key is required to authenticate with Roboflow API and access Model Monitoring features. Visit https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key to learn how to retrieve an API key. The block requires predictions that contain inference IDs (predictions must have been generated by models that include inference IDs). Supported prediction types: object detection, instance segmentation, keypoint detection, and classification. The block uses workspace caching (15-minute expiration) to optimize API calls. For more information on Model Monitoring at Roboflow, see https://docs.roboflow.com/deploy/model-monitoring. +""" + +WORKSPACE_NAME_CACHE_EXPIRE = 900 # 15 min + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Roboflow Custom Metadata", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "sink", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-analytics", + "blockPriority": 8, + "requires_rf_key": True, + }, + } + ) + type: Literal["roboflow_core/roboflow_custom_metadata@v1", "RoboflowCustomMetadata"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + ] + ) = Field( + description="Model predictions (object detection, instance segmentation, keypoint detection, or classification) to attach custom metadata to. The predictions must contain inference IDs that are used to associate metadata with specific inference results in Roboflow Model Monitoring. Inference IDs are automatically extracted from supervision Detections objects or classification prediction dictionaries. The metadata will be attached to all inference IDs found in the predictions.", + examples=[ + "$steps.object_detection.predictions", + "$steps.classification.predictions", + "$steps.instance_segmentation.predictions", + ], + ) + field_name: str = Field( + description="Name of the custom metadata field to create in Roboflow Model Monitoring. This becomes the field name that can be used for filtering and analysis in the Model Monitoring dashboard. Field names should be descriptive and represent the type of metadata being attached (e.g., 'location', 'quality', 'camera_id', 'batch_number'). The field name is used to organize and categorize metadata values.", + examples=[ + "location", + "quality", + "camera_id", + "batch_number", + "shift", + "operator", + ], + ) + field_value: Union[ + str, + Selector(kind=[STRING_KIND]), + Selector(kind=[STRING_KIND]), + ] = Field( + description="Value to assign to the custom metadata field. This is the actual data that will be attached to inference results and can be used for filtering and analysis in the Model Monitoring dashboard. Can be a string literal or a selector that references workflow outputs. Common values: location identifiers (e.g., 'toronto', 'warehouse_a'), quality labels (e.g., 'pass', 'fail', 'review'), identifiers (e.g., camera IDs, batch numbers), or any other contextual information relevant to your use case.", + examples=[ + "toronto", + "pass", + "fail", + "warehouse_a", + "camera_01", + "$steps.expression.output", + ], + ) + fire_and_forget: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Execution mode flag. When True (default), the block runs asynchronously in the background, allowing the workflow to continue processing without waiting for the API call to complete. This provides faster workflow execution but errors are not immediately available. When False, the block runs synchronously and waits for the API call to complete, returning immediate status and error information. Use False for debugging and error handling, True for production workflows where performance is prioritized.", + examples=[True, False], + ) + + @classmethod + def get_air_gapped_availability(cls) -> AirGappedAvailability: + return AirGappedAvailability(available=False, reason="requires_internet") + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="message", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class RoboflowCustomMetadataBlockV1(WorkflowBlock): + + def __init__( + self, + cache: BaseCache, + api_key: Optional[str], + background_tasks: Optional[BackgroundTasks], + thread_pool_executor: Optional[ThreadPoolExecutor], + disable_sinks: bool = False, + ): + self._api_key = api_key + self._cache = cache + self._background_tasks = background_tasks + self._thread_pool_executor = thread_pool_executor + self._disable_sinks = disable_sinks + + @classmethod + def get_init_parameters(cls) -> List[str]: + return [ + "api_key", + "cache", + "background_tasks", + "thread_pool_executor", + "disable_sinks", + ] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + fire_and_forget: bool, + field_name: str, + field_value: str, + predictions: Union[ + Detections, + InstanceDetections, + KeyPointPrediction, + ClassificationPrediction, + MultiLabelClassificationPrediction, + ], + ) -> BlockResult: + if self._disable_sinks: + return disabled_sink_response() + if self._api_key is None: + raise ValueError( + "RoboflowCustomMetadata block cannot run without Roboflow API key. " + "If you do not know how to get API key - visit " + "https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key to learn how to " + "retrieve one." + ) + inference_ids: List[str] = _extract_inference_ids(predictions) + if len(inference_ids) == 0: + return { + "error_status": True, + "message": "Custom metadata upload failed because no inference_ids were received. " + "This is known bug (https://github.com/roboflow/inference/issues/567). " + "Please provide a report for the problem under mentioned issue.", + } + inference_ids: List[str] = list(set(inference_ids)) + registration_task = partial( + add_custom_metadata_request, + cache=self._cache, + api_key=self._api_key, + inference_ids=inference_ids, + field_name=field_name, + field_value=field_value, + ) + error_status = False + message = "Registration happens in the background task" + if fire_and_forget and self._background_tasks: + self._background_tasks.add_task(registration_task) + elif fire_and_forget and self._thread_pool_executor: + self._thread_pool_executor.submit(registration_task) + else: + error_status, message = registration_task() + return { + "error_status": error_status, + "message": message, + } + + +def _extract_inference_ids( + predictions: Union[ + Detections, + InstanceDetections, + KeyPointPrediction, + ClassificationPrediction, + MultiLabelClassificationPrediction, + ], +) -> List[str]: + # Tensor-native predictions keep the inference_id as a single per-image value in + # `image_metadata` (detections / keypoints) or in the classification prediction's + # metadata, rather than the per-detection `sv.Detections.data[INFERENCE_ID_KEY]` + # array the numpy sibling read from. + if isinstance(predictions, ClassificationPrediction): + images_metadata = predictions.images_metadata or [{}] + image_metadata = images_metadata[0] if images_metadata else {} + elif isinstance(predictions, MultiLabelClassificationPrediction): + image_metadata = predictions.image_metadata or {} + else: + _key_points, detections = split_key_point_prediction(predictions) + image_metadata = detections.image_metadata or {} + inference_id = image_metadata.get(INFERENCE_ID_KEY) + if inference_id is None: + return [] + return [inference_id] + + +def get_workspace_name( + api_key: str, + cache: BaseCache, +) -> str: + # codeql[py/weak-sensitive-data-hashing]: MD5 cache fingerprint; not crypto storage. + api_key_hash = hashlib.md5(api_key.encode("utf-8")).hexdigest() + cache_key = f"workflows:api_key_to_workspace:{api_key_hash}" + cached_workspace_name = cache.get(cache_key) + if cached_workspace_name: + return cached_workspace_name + workspace_name_from_api = get_roboflow_workspace(api_key=api_key) + cache.set( + key=cache_key, value=workspace_name_from_api, expire=WORKSPACE_NAME_CACHE_EXPIRE + ) + return workspace_name_from_api + + +def add_custom_metadata_request( + cache: BaseCache, + api_key: str, + inference_ids: List[str], + field_name: str, + field_value: str, +) -> Tuple[bool, str]: + workspace_id = get_workspace_name(api_key=api_key, cache=cache) + try: + add_custom_metadata( + api_key=api_key, + workspace_id=workspace_id, + inference_ids=inference_ids, + field_name=field_name, + field_value=field_value, + ) + return ( + False, + "Custom metadata upload was successful", + ) + except Exception as error: + logging.warning( + f"Could not add custom metadata for inference IDs: {inference_ids}. Reason: {error}" + ) + return ( + True, + f"Error while custom metadata registration. Error type: {type(error)}. Details: {error}", + ) diff --git a/inference/core/workflows/core_steps/sinks/roboflow/dataset_upload/v1_tensor.py b/inference/core/workflows/core_steps/sinks/roboflow/dataset_upload/v1_tensor.py new file mode 100644 index 0000000000..5e318e7fed --- /dev/null +++ b/inference/core/workflows/core_steps/sinks/roboflow/dataset_upload/v1_tensor.py @@ -0,0 +1,886 @@ +""" +***************************************************************** +* WARNING! * +***************************************************************** +This module contains the utility functions used by +RoboflowDatasetUploadBlockV2. + +We do not recommend making multiple blocks dependent on the same code, +but the change between v1 and v2 was basically the default value of +some parameter - hence we decided not to replicate the code. + +If you need to modify this module beware that you may introduce +change to RoboflowDatasetUploadBlockV2! If that happens, +probably that's the time to disentangle those blocks and copy the +code. +""" + +import hashlib +import json +import logging +from collections import OrderedDict +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timedelta +from functools import partial +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import numpy as np +import supervision as sv +import torch +from fastapi import BackgroundTasks +from pydantic import AliasChoices, ConfigDict, Field + +from inference.core.active_learning.cache_operations import ( + return_strategy_credit, + use_credit_of_matching_strategy, +) +from inference.core.active_learning.core import prepare_image_to_registration +from inference.core.active_learning.entities import ( + ImageDimensions, + StrategyLimit, + StrategyLimitType, +) +from inference.core.cache.base import BaseCache +from inference.core.roboflow_api import ( + annotate_image_at_roboflow, + get_roboflow_workspace, + register_image_at_roboflow, +) +from inference.core.workflows.core_steps.common.query_language.operations.classification_results.base import ( + extract_top_class_tensor_native, +) +from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_sv_detections, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + KeyPointPrediction, + TensorNativeDetections, + instance_mask_to_numpy, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.sinks.noop import disabled_sink_message +from inference.core.workflows.execution_engine.constants import ( + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + POLYGON_KEY_IN_SV_DETECTIONS, + SCALING_RELATIVE_TO_PARENT_KEY, + SCALING_RELATIVE_TO_ROOT_PARENT_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + AirGappedAvailability, + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_project, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks + +# A tensor-native prediction may be a classification result, a detection-shaped +# prediction (object detection / instance segmentation), or a keypoint tuple. +TensorNativePrediction = Union[ + Detections, + InstanceDetections, + KeyPointPrediction, + ClassificationPrediction, + MultiLabelClassificationPrediction, +] + +SHORT_DESCRIPTION = "Save images and predictions to your Roboflow Dataset." + +LONG_DESCRIPTION = """ +Upload images and model predictions to a Roboflow dataset for active learning, model improvement, and data collection, with configurable usage quotas, batch organization, image compression, and optional annotation persistence. + +## How This Block Works + +This block uploads workflow images and predictions to your Roboflow dataset for storage, labeling, and model training. The block: + +1. Takes images and optional model predictions (object detection, instance segmentation, keypoint detection, or classification) as input +2. Validates the Roboflow API key is available (required for uploading) +3. Checks usage quotas (minutely, hourly, daily limits) to ensure uploads stay within configured rate limits for active learning strategies +4. Prepares images by resizing if they exceed maximum size (maintaining aspect ratio) and compressing to specified quality level +5. Generates labeling batch names based on the prefix and batch creation frequency (never, daily, weekly, or monthly), organizing uploaded data into batches +6. Optionally persists model predictions as annotations if `persist_predictions` is enabled, allowing predictions to serve as pre-labels for review and correction +7. Attaches registration tags to images for organization and filtering in the Roboflow platform +8. Registers the image (and annotations if enabled) to the specified Roboflow project via the Roboflow API +9. Executes synchronously or asynchronously based on `fire_and_forget` setting, allowing non-blocking uploads for faster workflow execution +10. Returns error status and messages indicating upload success or failure + +The block supports active learning workflows by implementing usage quotas that prevent excessive data collection, helping focus on collecting valuable training data within rate limits. Images are organized into labeling batches that can be automatically recreated on a schedule (daily, weekly, monthly), making it easier to manage and review collected data over time. The block can operate in fire-and-forget mode for asynchronous execution, allowing workflows to continue processing without waiting for uploads to complete, or synchronously for debugging and error handling. + +## Requirements + +**API Key Required**: This block requires a valid Roboflow API key to upload data. The API key must be configured in your environment or workflow configuration. Visit https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key to learn how to retrieve an API key. + +## Common Use Cases + +- **Active Learning Data Collection**: Collect images and predictions from production environments where models struggle or are uncertain (e.g., low-confidence detections, edge cases), enabling iterative model improvement by gathering challenging examples for retraining +- **Production Data Logging**: Continuously upload production inference data to Roboflow datasets for monitoring, analysis, and future model training, creating a growing dataset from real-world deployments +- **Pre-Labeled Data Collection**: Upload images with model predictions as pre-labels (when `persist_predictions` is enabled), accelerating annotation workflows by providing initial labels that can be reviewed and corrected rather than starting from scratch +- **Stratified Data Sampling**: Use rate limiting and quotas to selectively collect data based on specific criteria (e.g., combine with Rate Limiter or Continue If blocks), ensuring diverse and balanced dataset collection without overwhelming storage or annotation resources +- **Batch-Based Labeling Workflows**: Organize uploaded data into batches with automatic recreation schedules (daily, weekly, monthly), making it easier to manage labeling tasks, track progress, and organize data collection efforts over time +- **Tagged Data Organization**: Attach metadata tags to uploaded images (e.g., location, camera ID, time period, model version), enabling filtering and organization of collected data in Roboflow for better dataset management and analysis + +## Connecting to Other Blocks + +This block receives data from workflow steps and uploads it to Roboflow: + +- **After detection or analysis blocks** (e.g., Object Detection Model, Instance Segmentation Model, Classification Model, Keypoint Detection Model) to upload images along with their predictions, enabling active learning by collecting inference data with model outputs for annotation and retraining +- **After filtering or analytics blocks** (e.g., Detections Filter, Continue If, Overlap Filter) to selectively upload only specific types of data (e.g., low-confidence detections, overlapping objects, specific classes), focusing data collection on valuable edge cases or interesting scenarios +- **After rate limiter blocks** (e.g., Rate Limiter) to throttle upload frequency and stay within usage quotas, ensuring controlled data collection that respects rate limits and prevents excessive storage usage +- **Image inputs or preprocessing blocks** to upload raw images or processed images (e.g., crops, transformed images) without predictions, enabling collection of image data for future labeling or analysis +- **Conditional workflows** using flow control blocks (e.g., Continue If) to upload data only when certain conditions are met (e.g., upload only when detection count exceeds threshold, upload only errors or failures), enabling selective data collection based on workflow state +- **Batch processing workflows** where multiple images or predictions are generated, allowing bulk upload of workflow outputs to Roboflow datasets for organized data collection and management +""" + +WORKSPACE_NAME_CACHE_EXPIRE = 900 # 15 min +TIMESTAMP_FORMAT = "%Y_%m_%d" +DUPLICATED_STATUS = "Duplicated image" +BatchCreationFrequency = Literal["never", "daily", "weekly", "monthly"] + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Roboflow Dataset Upload", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "sink", + "ui_manifest": { + "section": "data_storage", + "icon": "fal fa-upload", + "blockPriority": 0, + "popular": True, + "requires_rf_key": True, + }, + } + ) + type: Literal["roboflow_core/roboflow_dataset_upload@v1", "RoboflowDatasetUpload"] + images: Selector(kind=[IMAGE_KIND]) = Field( + title="Input Image", + description="Image(s) to upload to the Roboflow dataset. Can be a single image or batch of images from workflow inputs or processing steps. Images are resized if they exceed max_image_size and compressed before uploading. Supports batch processing.", + examples=["$inputs.image", "$steps.cropping.crops"], + validation_alias=AliasChoices("image", "images"), + ) + predictions: Optional[ + Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + ] + ) + ] = Field( + default=None, + description="Optional model predictions to upload alongside images. Predictions are saved as annotations (pre-labels) in the Roboflow dataset when persist_predictions is enabled, allowing predictions to serve as starting points for annotation review and correction. Supports object detection, instance segmentation, keypoint detection, and classification predictions. If None, only images are uploaded.", + examples=["$steps.object_detection_model.predictions"], + ) + target_project: Union[Selector(kind=[ROBOFLOW_PROJECT_KIND]), str] = Field( + description="Roboflow project identifier where uploaded images and annotations will be saved. Must be a valid project in your Roboflow workspace. The project name can be specified directly or referenced from workflow inputs.", + examples=["my_project", "$inputs.target_project"], + ) + minutely_usage_limit: int = Field( + default=10, + description="Maximum number of image uploads allowed per minute for this quota. Part of the usage quota system that enforces rate limits for active learning data collection. Uploads exceeding this limit are skipped to prevent excessive data collection. Works together with hourly_usage_limit and daily_usage_limit to provide multi-level rate limiting.", + examples=[10, 60], + ) + hourly_usage_limit: int = Field( + default=100, + description="Maximum number of image uploads allowed per hour for this quota. Part of the usage quota system that enforces rate limits for active learning data collection. Uploads exceeding this limit are skipped to prevent excessive data collection. Works together with minutely_usage_limit and daily_usage_limit to provide multi-level rate limiting.", + examples=[10, 60], + ) + daily_usage_limit: int = Field( + default=1000, + description="Maximum number of image uploads allowed per day for this quota. Part of the usage quota system that enforces rate limits for active learning data collection. Uploads exceeding this limit are skipped to prevent excessive data collection. Works together with minutely_usage_limit and hourly_usage_limit to provide multi-level rate limiting.", + examples=[10, 60], + ) + usage_quota_name: str = Field( + description="Unique identifier for tracking usage quotas (minutely, hourly, daily limits). Used internally to manage rate limiting across multiple upload operations. Each unique quota name maintains separate counters, allowing different upload strategies or data collection workflows to have independent rate limits.", + examples=["quota-for-data-sampling-1"], + json_schema_extra={"hidden": True}, + ) + max_image_size: Tuple[int, int] = Field( + default=(512, 512), + description="Maximum dimensions (width, height) for uploaded images. Images exceeding these dimensions are automatically resized while preserving aspect ratio before uploading. Smaller sizes reduce storage and bandwidth but may lose image quality. Use larger sizes (e.g., (1920, 1080)) for high-resolution data collection, or smaller sizes (e.g., (512, 512)) for efficient storage and faster uploads.", + examples=[(512, 512), (1920, 1080)], + ) + compression_level: int = Field( + default=75, + gt=0, + le=100, + description="JPEG compression quality level for uploaded images, ranging from 1 (highest compression, smallest file size, lower quality) to 100 (no compression, largest file size, highest quality). Higher values preserve more image quality but increase storage and bandwidth usage. Typical values range from 70-90 for balanced quality and size. Default of 75 provides good quality with reasonable file sizes.", + examples=[75], + ) + registration_tags: Union[ + List[Union[Selector(kind=[STRING_KIND]), str]], + Selector(kind=[LIST_OF_VALUES_KIND]), + ] = Field( + default_factory=list, + description="List of tags to attach to uploaded images for organization and filtering in Roboflow. Tags can be static strings (e.g., 'location-florida', 'camera-1') or dynamic values from workflow inputs. Tags help organize collected data, filter images in Roboflow, and add metadata for dataset management. Can be an empty list if no tags are needed.", + examples=[ + ["location-florida", "factory-name", "$inputs.dynamic_tag"], + "$inputs.tags", + ], + ) + persist_predictions: bool = Field( + default=True, + description="If True, model predictions are saved as annotations (pre-labels) in the Roboflow dataset alongside images. This enables predictions to serve as starting points for annotation, allowing reviewers to correct or approve labels rather than creating them from scratch. If False, only images are uploaded without annotations. Enabling this accelerates annotation workflows by providing initial labels.", + examples=[True, False], + ) + disable_sink: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="If True, the block execution is disabled and no uploads occur. This allows temporarily disabling data collection without removing the block from workflows, useful for testing, debugging, or conditional data collection. When disabled, returns a message indicating the sink was disabled. Default is False (uploads enabled).", + examples=[True, "$inputs.disable_active_learning"], + ) + fire_and_forget: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="If True, uploads execute asynchronously (fire-and-forget mode), allowing the workflow to continue immediately without waiting for upload completion. This improves workflow performance but prevents error handling. If False, uploads execute synchronously, blocking workflow execution until completion and allowing proper error handling and status reporting. Use async mode (True) for production workflows where speed is prioritized, and sync mode (False) for debugging or when error handling is critical.", + examples=[True], + ) + labeling_batch_prefix: Union[str, Selector(kind=[STRING_KIND])] = Field( + default="workflows_data_collector", + description="Prefix used to generate labeling batch names for organizing uploaded images in Roboflow. Combined with the batch recreation frequency and timestamps to create batch names like 'workflows_data_collector_2024_01_15'. Batches help organize collected data for labeling, making it easier to manage and review uploaded images in groups. Can be customized to match your organization scheme.", + examples=["my_labeling_batch_name"], + ) + labeling_batches_recreation_frequency: BatchCreationFrequency = Field( + default="never", + description="Frequency at which new labeling batches are automatically created for uploaded images. Options: 'never' (all images go to the same batch), 'daily' (new batch each day), 'weekly' (new batch each week), 'monthly' (new batch each month). Batch timestamps are appended to the labeling_batch_prefix to create unique batch names. Automatically organizing uploads into time-based batches simplifies dataset management and makes it easier to track and review collected data over time.", + examples=["never", "daily"], + ) + image_name: Optional[Union[str, Selector(kind=[STRING_KIND])]] = Field( + default=None, + description="Optional custom name for the uploaded image. If provided, this name will be used instead of an auto-generated UUID. This is useful when you want to preserve the original filename or use a meaningful identifier (e.g., serial number, timestamp) for the image in the Roboflow dataset. The name should not include file extension. If not provided, a UUID will be generated automatically.", + examples=["serial_12345", "camera1_frame_001", "$inputs.filename"], + ) + + @classmethod + def get_air_gapped_availability(cls) -> AirGappedAvailability: + return AirGappedAvailability(available=False, reason="requires_internet") + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images", "predictions", "image_name"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="message", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + if self.disable_sink is True: + # Sink literally disabled โ€” execution returns before the target + # project is ever accessed. + return [] + return [roboflow_platform_project(project_url=self.target_project)] + + +class RoboflowDatasetUploadBlockV1(WorkflowBlock): + + def __init__( + self, + cache: BaseCache, + api_key: Optional[str], + background_tasks: Optional[BackgroundTasks], + thread_pool_executor: Optional[ThreadPoolExecutor], + disable_sinks: bool = False, + ): + self._cache = cache + self._api_key = api_key + self._background_tasks = background_tasks + self._thread_pool_executor = thread_pool_executor + self._disable_sinks = disable_sinks + + @classmethod + def get_init_parameters(cls) -> List[str]: + return [ + "cache", + "api_key", + "background_tasks", + "thread_pool_executor", + "disable_sinks", + ] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + predictions: Optional[Batch[TensorNativePrediction]], + target_project: str, + usage_quota_name: str, + minutely_usage_limit: int, + persist_predictions: bool, + hourly_usage_limit: int, + daily_usage_limit: int, + max_image_size: Tuple[int, int], + compression_level: int, + registration_tags: List[str], + disable_sink: bool, + fire_and_forget: bool, + labeling_batch_prefix: str, + labeling_batches_recreation_frequency: BatchCreationFrequency, + image_name: Optional[Batch[Optional[str]]] = None, + ) -> BlockResult: + if self._disable_sinks: + return [ + { + "error_status": False, + "message": disabled_sink_message(disabled_by_execution_policy=True), + } + for _ in range(len(images)) + ] + if self._api_key is None: + raise ValueError( + "RoboflowDataCollector block cannot run without Roboflow API key. " + "If you do not know how to get API key - visit " + "https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key to learn how to " + "retrieve one." + ) + if disable_sink: + return [ + { + "error_status": False, + "message": disabled_sink_message( + disabled_by_execution_policy=False + ), + } + for _ in range(len(images)) + ] + result = [] + predictions = [None] * len(images) if predictions is None else predictions + image_names = [None] * len(images) if image_name is None else image_name + for image, prediction, img_name in zip(images, predictions, image_names): + error_status, message = register_datapoint_at_roboflow( + image=image, + prediction=prediction, + target_project=target_project, + usage_quota_name=usage_quota_name, + persist_predictions=persist_predictions, + minutely_usage_limit=minutely_usage_limit, + hourly_usage_limit=hourly_usage_limit, + daily_usage_limit=daily_usage_limit, + max_image_size=max_image_size, + compression_level=compression_level, + registration_tags=registration_tags, + fire_and_forget=fire_and_forget, + labeling_batch_prefix=labeling_batch_prefix, + new_labeling_batch_frequency=labeling_batches_recreation_frequency, + cache=self._cache, + background_tasks=self._background_tasks, + thread_pool_executor=self._thread_pool_executor, + api_key=self._api_key, + image_name=img_name, + ) + result.append({"error_status": error_status, "message": message}) + return result + + +def register_datapoint_at_roboflow( + image: WorkflowImageData, + prediction: Optional[TensorNativePrediction], + target_project: str, + usage_quota_name: str, + persist_predictions: bool, + minutely_usage_limit: int, + hourly_usage_limit: int, + daily_usage_limit: int, + max_image_size: Tuple[int, int], + compression_level: int, + registration_tags: List[str], + fire_and_forget: bool, + labeling_batch_prefix: str, + new_labeling_batch_frequency: BatchCreationFrequency, + cache: BaseCache, + background_tasks: Optional[BackgroundTasks], + thread_pool_executor: Optional[ThreadPoolExecutor], + api_key: str, + image_name: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, +) -> Tuple[bool, str]: + registration_task = partial( + execute_registration, + image=image, + prediction=prediction, + target_project=target_project, + usage_quota_name=usage_quota_name, + persist_predictions=persist_predictions, + minutely_usage_limit=minutely_usage_limit, + hourly_usage_limit=hourly_usage_limit, + daily_usage_limit=daily_usage_limit, + max_image_size=max_image_size, + compression_level=compression_level, + registration_tags=registration_tags, + labeling_batch_prefix=labeling_batch_prefix, + new_labeling_batch_frequency=new_labeling_batch_frequency, + cache=cache, + api_key=api_key, + image_name=image_name, + metadata=metadata, + ) + if fire_and_forget and background_tasks: + background_tasks.add_task(registration_task) + return False, "Element registration happens in the background task" + if fire_and_forget and thread_pool_executor: + thread_pool_executor.submit(registration_task) + return False, "Element registration happens in the background task" + return registration_task() + + +def execute_registration( + image: WorkflowImageData, + prediction: Optional[TensorNativePrediction], + target_project: str, + persist_predictions: bool, + usage_quota_name: str, + minutely_usage_limit: int, + hourly_usage_limit: int, + daily_usage_limit: int, + max_image_size: Tuple[int, int], + compression_level: int, + registration_tags: List[str], + labeling_batch_prefix: str, + new_labeling_batch_frequency: BatchCreationFrequency, + cache: BaseCache, + api_key: str, + image_name: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, +) -> Tuple[bool, str]: + matching_strategies_limits = OrderedDict( + { + usage_quota_name: [ + StrategyLimit( + limit_type=StrategyLimitType.MINUTELY, value=minutely_usage_limit + ), + StrategyLimit( + limit_type=StrategyLimitType.HOURLY, value=hourly_usage_limit + ), + StrategyLimit( + limit_type=StrategyLimitType.DAILY, value=daily_usage_limit + ), + ] + } + ) + workspace_name = get_workspace_name(api_key=api_key, cache=cache) + strategy_with_spare_credit = use_credit_of_matching_strategy( + cache=cache, + workspace=workspace_name, + project=target_project, + matching_strategies_limits=matching_strategies_limits, + ) + if strategy_with_spare_credit is None: + return False, "Registration skipped due to usage quota exceeded" + credit_to_be_returned = False + try: + local_image_id = image_name if image_name else str(uuid4()) + encoded_image, scaling_factor = prepare_image_to_registration( + image=image.numpy_image, + desired_size=ImageDimensions( + width=max_image_size[0], height=max_image_size[1] + ), + jpeg_compression_level=compression_level, + ) + batch_name = generate_batch_name( + labeling_batch_prefix=labeling_batch_prefix, + new_labeling_batch_frequency=new_labeling_batch_frequency, + ) + if _is_tensor_native_detection_prediction(prediction): + prediction = scale_tensor_native_prediction( + prediction=prediction, scale=scaling_factor + ) + status = register_datapoint( + target_project=target_project, + encoded_image=encoded_image, + local_image_id=local_image_id, + prediction=prediction if persist_predictions else None, + api_key=api_key, + batch_name=batch_name, + tags=registration_tags, + metadata=metadata, + ) + if status == DUPLICATED_STATUS: + credit_to_be_returned = True + return False, status + except Exception as error: + credit_to_be_returned = True + logging.exception("Failed to register datapoint on the Roboflow platform") + return ( + True, + f"Error while registration. Error type: {type(error)}. Details: {error}", + ) + finally: + if credit_to_be_returned: + return_strategy_credit( + cache=cache, + workspace=workspace_name, + project=target_project, + strategy_name=strategy_with_spare_credit, + ) + + +def get_workspace_name( + api_key: str, + cache: BaseCache, +) -> str: + api_key_hash = hashlib.md5(api_key.encode("utf-8")).hexdigest() + cache_key = f"workflows:api_key_to_workspace:{api_key_hash}" + cached_workspace_name = cache.get(cache_key) + if cached_workspace_name: + return cached_workspace_name + workspace_name_from_api = get_roboflow_workspace(api_key=api_key) + cache.set( + key=cache_key, value=workspace_name_from_api, expire=WORKSPACE_NAME_CACHE_EXPIRE + ) + return workspace_name_from_api + + +def generate_batch_name( + labeling_batch_prefix: str, + new_labeling_batch_frequency: BatchCreationFrequency, +) -> str: + if new_labeling_batch_frequency == "never": + return labeling_batch_prefix + timestamp_generator = RECREATION_INTERVAL2TIMESTAMP_GENERATOR[ + new_labeling_batch_frequency + ] + timestamp = timestamp_generator() + return f"{labeling_batch_prefix}_{timestamp}" + + +def generate_today_timestamp() -> str: + return datetime.today().strftime(TIMESTAMP_FORMAT) + + +def generate_start_timestamp_for_this_week() -> str: + today = datetime.today() + return (today - timedelta(days=today.weekday())).strftime(TIMESTAMP_FORMAT) + + +def generate_start_timestamp_for_this_month() -> str: + return datetime.today().replace(day=1).strftime(TIMESTAMP_FORMAT) + + +RECREATION_INTERVAL2TIMESTAMP_GENERATOR = { + "daily": generate_today_timestamp, + "weekly": generate_start_timestamp_for_this_week, + "monthly": generate_start_timestamp_for_this_month, +} + + +def register_datapoint( + target_project: str, + encoded_image: bytes, + local_image_id: str, + prediction: Optional[TensorNativePrediction], + api_key: str, + batch_name: str, + tags: List[str], + metadata: Optional[Dict[str, Any]] = None, +) -> str: + inference_id = None + if _is_tensor_native_classification_prediction(prediction): + inference_id = _read_classification_inference_id(prediction=prediction) + if _is_tensor_native_detection_prediction(prediction): + # TODO: Lack of inference ID for empty prediction - + # dependent on https://github.com/roboflow/inference/issues/567 + _key_points, detections = split_key_point_prediction(prediction) + if len(detections) > 0: + image_metadata = detections.image_metadata or {} + inference_id = image_metadata.get(INFERENCE_ID_KEY) + roboflow_image_id = safe_register_image_at_roboflow( + target_project=target_project, + encoded_image=encoded_image, + local_image_id=local_image_id, + api_key=api_key, + batch_name=batch_name, + tags=tags, + inference_id=inference_id, + metadata=metadata, + ) + if roboflow_image_id is None: + return DUPLICATED_STATUS + if is_prediction_registration_forbidden(prediction=prediction): + return "Successfully registered image" + encoded_prediction, prediction_format = encode_prediction(prediction=prediction) + _ = annotate_image_at_roboflow( + api_key=api_key, + dataset_id=target_project, + local_image_id=local_image_id, + roboflow_image_id=roboflow_image_id, + annotation_content=encoded_prediction, + annotation_file_type=prediction_format, + is_prediction=True, + ) + return "Successfully registered image and annotation" + + +def safe_register_image_at_roboflow( + target_project: str, + encoded_image: bytes, + local_image_id: str, + api_key: str, + batch_name: str, + tags: List[str], + inference_id: Optional[str], + metadata: Optional[Dict[str, Any]] = None, +) -> Optional[str]: + registration_response = register_image_at_roboflow( + api_key=api_key, + dataset_id=target_project, + local_image_id=local_image_id, + image_bytes=encoded_image, + batch_name=batch_name, + tags=tags, + inference_id=inference_id, + metadata=metadata, + ) + image_duplicated = registration_response.get("duplicate", False) + if image_duplicated: + logging.warning(f"Image duplication detected: {registration_response}.") + return None + return registration_response["id"] + + +def is_prediction_registration_forbidden( + prediction: Optional[TensorNativePrediction], +) -> bool: + if prediction is None: + return True + if _is_tensor_native_detection_prediction(prediction): + _key_points, detections = split_key_point_prediction(prediction) + if len(detections) == 0: + return True + return False + + +def encode_prediction( + prediction: TensorNativePrediction, +) -> Tuple[str, str]: + if _is_tensor_native_classification_prediction(prediction): + top_classes = extract_top_class_tensor_native(prediction=prediction) + if isinstance(top_classes, list): + return ",".join(top_classes), "txt" + return top_classes, "txt" + _key_points, detections = split_key_point_prediction(prediction) + detections_in_inference_format = serialise_sv_detections(detections=detections) + return json.dumps(detections_in_inference_format), "json" + + +def _is_tensor_native_classification_prediction( + prediction: Optional[TensorNativePrediction], +) -> bool: + return isinstance( + prediction, (ClassificationPrediction, MultiLabelClassificationPrediction) + ) + + +def _is_tensor_native_detection_prediction( + prediction: Optional[TensorNativePrediction], +) -> bool: + if isinstance(prediction, (Detections, InstanceDetections)): + return True + # Keypoint workflow kind: Tuple[KeyPoints, Optional[Detections]]. + if isinstance(prediction, tuple): + return len(prediction) == 2 and isinstance(prediction[0], KeyPoints) + return False + + +def _read_classification_inference_id( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> Optional[str]: + # Tensor-native classification keeps the inference_id as a single per-image value + # in the prediction's metadata, rather than the per-detection + # `sv.Detections.data[INFERENCE_ID_KEY]` array the numpy sibling read from. + if isinstance(prediction, ClassificationPrediction): + images_metadata = prediction.images_metadata or [{}] + image_metadata = images_metadata[0] if images_metadata else {} + else: + image_metadata = prediction.image_metadata or {} + return image_metadata.get(INFERENCE_ID_KEY) + + +def scale_tensor_native_prediction( + prediction: Union[TensorNativeDetections, KeyPointPrediction], + scale: float, +) -> Union[TensorNativeDetections, KeyPointPrediction]: + """Scale a tensor-native detection-shaped prediction by ``scale`` to match a + resized uploaded image, building a new native prediction (no sv round-trip). + + Mirrors ``scale_sv_detections``: scales ``xyxy``, per-image + ``image_dimensions``, per-detection ``keypoints_xy`` / ``polygon`` / + ``scaling_relative_to_*`` metadata, and re-rasterises instance masks at the + scaled resolution (``sv.mask_to_polygons`` / ``sv.polygon_to_mask`` used as a + pure numpy algorithm). For the keypoint tuple, both the ``KeyPoints`` and the + bbox ``Detections`` components are scaled consistently. + """ + key_points, detections = split_key_point_prediction(prediction) + scaled_detections = _scale_tensor_native_detections( + detections=detections, scale=scale + ) + if key_points is None: + return scaled_detections + scaled_key_points = KeyPoints( + xy=(key_points.xy * scale).round(), + class_id=key_points.class_id, + confidence=key_points.confidence, + image_metadata=_scale_image_metadata( + image_metadata=key_points.image_metadata, scale=scale + ), + key_points_metadata=key_points.key_points_metadata, + ) + return scaled_key_points, scaled_detections + + +def _scale_tensor_native_detections( + detections: TensorNativeDetections, + scale: float, +) -> TensorNativeDetections: + if len(detections) == 0: + return detections + scaled_xyxy = (detections.xyxy * scale).round() + scaled_image_metadata = _scale_image_metadata( + image_metadata=detections.image_metadata, scale=scale + ) + scaled_bboxes_metadata = _scale_bboxes_metadata( + bboxes_metadata=detections.bboxes_metadata, + detections_number=len(detections), + scale=scale, + ) + if isinstance(detections, InstanceDetections): + scaled_mask = _scale_instance_masks(detections=detections, scale=scale) + return InstanceDetections( + xyxy=scaled_xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + mask=scaled_mask, + image_metadata=scaled_image_metadata, + bboxes_metadata=scaled_bboxes_metadata, + ) + return Detections( + xyxy=scaled_xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + image_metadata=scaled_image_metadata, + bboxes_metadata=scaled_bboxes_metadata, + ) + + +def _scale_image_metadata( + image_metadata: Optional[dict], + scale: float, +) -> Optional[dict]: + if image_metadata is None: + return None + scaled = dict(image_metadata) + image_dimensions = scaled.get(IMAGE_DIMENSIONS_KEY) + if image_dimensions is not None: + scaled[IMAGE_DIMENSIONS_KEY] = ( + np.asarray(image_dimensions).astype(float) * scale + ).round() + return scaled + + +def _scale_bboxes_metadata( + bboxes_metadata: Optional[List[dict]], + detections_number: int, + scale: float, +) -> List[dict]: + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(detections_number)] + scaled_bboxes_metadata = [] + for data in bboxes_metadata: + scaled_data = dict(data) + if KEYPOINTS_XY_KEY_IN_SV_DETECTIONS in scaled_data: + scaled_data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = ( + np.asarray(scaled_data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS]).astype( + np.float32 + ) + * scale + ).round() + if POLYGON_KEY_IN_SV_DETECTIONS in scaled_data: + scaled_data[POLYGON_KEY_IN_SV_DETECTIONS] = ( + (np.asarray(scaled_data[POLYGON_KEY_IN_SV_DETECTIONS]) * scale) + .round() + .astype(np.int32) + ) + if SCALING_RELATIVE_TO_PARENT_KEY in scaled_data: + scaled_data[SCALING_RELATIVE_TO_PARENT_KEY] = ( + scaled_data[SCALING_RELATIVE_TO_PARENT_KEY] * scale + ) + else: + scaled_data[SCALING_RELATIVE_TO_PARENT_KEY] = scale + if SCALING_RELATIVE_TO_ROOT_PARENT_KEY in scaled_data: + scaled_data[SCALING_RELATIVE_TO_ROOT_PARENT_KEY] = ( + scaled_data[SCALING_RELATIVE_TO_ROOT_PARENT_KEY] * scale + ) + else: + scaled_data[SCALING_RELATIVE_TO_ROOT_PARENT_KEY] = scale + scaled_bboxes_metadata.append(scaled_data) + return scaled_bboxes_metadata + + +def _scale_instance_masks( + detections: InstanceDetections, + scale: float, +) -> Union[torch.Tensor, InstancesRLEMasks]: + detections_number = len(detections) + sample_mask = instance_mask_to_numpy(detections, 0) + original_h, original_w = sample_mask.shape[:2] + scaled_mask_size_wh = ( + round(original_w * scale), + round(original_h * scale), + ) + scaled_masks = [] + for index in range(detections_number): + detection_mask = instance_mask_to_numpy(detections, index).astype(np.uint8) + polygons = sv.mask_to_polygons(mask=detection_mask) + polygon_masks = [] + for polygon in polygons: + scaled_polygon = (polygon * scale).round().astype(np.int32) + polygon_masks.append( + sv.polygon_to_mask( + polygon=scaled_polygon, resolution_wh=scaled_mask_size_wh + ) + ) + scaled_detection_mask = np.sum(polygon_masks, axis=0) > 0 + scaled_masks.append(scaled_detection_mask) + return torch.from_numpy(np.array(scaled_masks)).to(torch.bool) diff --git a/inference/core/workflows/core_steps/sinks/roboflow/dataset_upload/v2_tensor.py b/inference/core/workflows/core_steps/sinks/roboflow/dataset_upload/v2_tensor.py new file mode 100644 index 0000000000..e38170decf --- /dev/null +++ b/inference/core/workflows/core_steps/sinks/roboflow/dataset_upload/v2_tensor.py @@ -0,0 +1,439 @@ +import random +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +from fastapi import BackgroundTasks +from pydantic import AliasChoices, ConfigDict, Field +from typing_extensions import Annotated + +from inference.core.cache.base import BaseCache +from inference.core.workflows.core_steps.common.tensor_native import KeyPointPrediction +from inference.core.workflows.core_steps.sinks.noop import disabled_sink_message +from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1_tensor import ( + register_datapoint_at_roboflow, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + ROBOFLOW_PROJECT_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + AirGappedAvailability, + BlockResult, + DependentResource, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_project, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +# A tensor-native prediction may be a classification result, a detection-shaped +# prediction (object detection / instance segmentation), or a keypoint tuple. +TensorNativePrediction = Union[ + Detections, + InstanceDetections, + KeyPointPrediction, + ClassificationPrediction, + MultiLabelClassificationPrediction, +] + +FloatZeroToHundred = Annotated[float, Field(ge=0.0, le=100.0)] + +SHORT_DESCRIPTION = "Save images and predictions to your Roboflow Dataset." + +LONG_DESCRIPTION = """ +Upload images and model predictions to a Roboflow dataset for active learning, model improvement, and data collection, with configurable usage quotas, probabilistic sampling, batch organization, image compression, and optional annotation persistence. + +## How This Block Works + +This block uploads workflow images and predictions to your Roboflow dataset for storage, labeling, and model training. The block: + +1. Takes images and optional model predictions (object detection, instance segmentation, keypoint detection, or classification) as input +2. Validates the Roboflow API key is available (required for uploading) +3. Applies probabilistic sampling based on `data_percentage` setting, randomly selecting a percentage of inputs to upload (e.g., 50% uploads half the data, 100% uploads everything) +4. Checks usage quotas (minutely, hourly, daily limits) to ensure uploads stay within configured rate limits for active learning strategies +5. Prepares images by resizing if they exceed maximum size (maintaining aspect ratio) and compressing to specified quality level +6. Generates labeling batch names based on the prefix and batch creation frequency (never, daily, weekly, or monthly), organizing uploaded data into batches +7. Optionally persists model predictions as annotations if `persist_predictions` is enabled, allowing predictions to serve as pre-labels for review and correction +8. Attaches registration tags to images for organization and filtering in the Roboflow platform +9. Registers the image (and annotations if enabled) to the specified Roboflow project via the Roboflow API +10. Executes synchronously or asynchronously based on `fire_and_forget` setting, allowing non-blocking uploads for faster workflow execution +11. Returns error status and messages indicating upload success, failure, or sampling skip + +The block supports active learning workflows by implementing usage quotas that prevent excessive data collection, helping focus on collecting valuable training data within rate limits. The probabilistic sampling feature (new in v2) allows you to randomly sample a percentage of data for upload, enabling cost-effective data collection strategies where you want to collect representative samples rather than all data. Images are organized into labeling batches that can be automatically recreated on a schedule (daily, weekly, monthly), making it easier to manage and review collected data over time. The block can operate in fire-and-forget mode for asynchronous execution, allowing workflows to continue processing without waiting for uploads to complete, or synchronously for debugging and error handling. + +## Version Differences (v2 vs v1) + +**New Features in v2:** + +- **Probabilistic Data Sampling**: Added `data_percentage` parameter (0-100%) that enables random sampling of data for upload. This allows you to upload only a percentage of workflow inputs (e.g., 25% samples one in four images), reducing storage and annotation costs while still collecting representative data. When sampling skips an upload, the block returns a message indicating the skip. + +- **Improved Default Settings**: + - `max_image_size` default increased from (512, 512) to (1920, 1080) for higher resolution data collection + - `compression_level` default increased from 75 to 95 for better image quality preservation + +**Behavior Changes:** + +- By default, `data_percentage` is set to 100, so v2 behaves identically to v1 unless sampling is explicitly configured +- The block now uses probabilistic sampling before quota checking and image preparation, allowing efficient filtering before resource-intensive operations + +## Requirements + +**API Key Required**: This block requires a valid Roboflow API key to upload data. The API key must be configured in your environment or workflow configuration. Visit https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key to learn how to retrieve an API key. + +## Common Use Cases + +- **Active Learning Data Collection**: Collect images and predictions from production environments where models struggle or are uncertain (e.g., low-confidence detections, edge cases), enabling iterative model improvement by gathering challenging examples for retraining +- **Probabilistic Data Sampling**: Use `data_percentage` to randomly sample a subset of data for upload (e.g., upload 20% of all detections, 50% of low-confidence cases), enabling cost-effective data collection strategies that reduce storage and annotation overhead while maintaining dataset diversity +- **Production Data Logging**: Continuously upload production inference data to Roboflow datasets for monitoring, analysis, and future model training, creating a growing dataset from real-world deployments +- **Pre-Labeled Data Collection**: Upload images with model predictions as pre-labels (when `persist_predictions` is enabled), accelerating annotation workflows by providing initial labels that can be reviewed and corrected rather than starting from scratch +- **Stratified Data Sampling**: Combine probabilistic sampling with rate limiting and quotas to selectively collect data based on specific criteria (e.g., sample 30% of detections that pass filters), ensuring diverse and balanced dataset collection without overwhelming storage or annotation resources +- **Batch-Based Labeling Workflows**: Organize uploaded data into batches with automatic recreation schedules (daily, weekly, monthly), making it easier to manage labeling tasks, track progress, and organize data collection efforts over time + +## Connecting to Other Blocks + +This block receives data from workflow steps and uploads it to Roboflow: + +- **After detection or analysis blocks** (e.g., Object Detection Model, Instance Segmentation Model, Classification Model, Keypoint Detection Model) to upload images along with their predictions, enabling active learning by collecting inference data with model outputs for annotation and retraining +- **After filtering or analytics blocks** (e.g., Detections Filter, Continue If, Overlap Filter) to selectively upload only specific types of data (e.g., low-confidence detections, overlapping objects, specific classes), focusing data collection on valuable edge cases or interesting scenarios +- **After rate limiter blocks** (e.g., Rate Limiter) to throttle upload frequency and stay within usage quotas, ensuring controlled data collection that respects rate limits and prevents excessive storage usage +- **Image inputs or preprocessing blocks** to upload raw images or processed images (e.g., crops, transformed images) without predictions, enabling collection of image data for future labeling or analysis +- **Conditional workflows** using flow control blocks (e.g., Continue If) to upload data only when certain conditions are met (e.g., upload only when detection count exceeds threshold, upload only errors or failures), enabling selective data collection based on workflow state +- **Batch processing workflows** where multiple images or predictions are generated, allowing bulk upload of workflow outputs to Roboflow datasets with probabilistic sampling for organized and cost-effective data collection +""" + +WORKSPACE_NAME_CACHE_EXPIRE = 900 # 15 min +TIMESTAMP_FORMAT = "%Y_%m_%d" +DUPLICATED_STATUS = "Duplicated image" +BatchCreationFrequency = Literal["never", "daily", "weekly", "monthly"] + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Roboflow Dataset Upload", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "sink", + "ui_manifest": { + "section": "data_storage", + "icon": "fal fa-upload", + "blockPriority": 0, + "popular": True, + "requires_rf_key": True, + }, + } + ) + type: Literal["roboflow_core/roboflow_dataset_upload@v2"] + images: Selector(kind=[IMAGE_KIND]) = Field( + title="Image", + description="Image(s) to upload to the Roboflow dataset. Can be a single image or batch of images from workflow inputs or processing steps. Images are randomly sampled based on data_percentage, resized if they exceed max_image_size, and compressed before uploading. Supports batch processing.", + examples=["$inputs.image", "$steps.cropping.crops"], + validation_alias=AliasChoices("images", "image"), + ) + target_project: Union[Selector(kind=[ROBOFLOW_PROJECT_KIND]), str] = Field( + description="Roboflow project identifier where uploaded images and annotations will be saved. Must be a valid project in your Roboflow workspace. The project name can be specified directly or referenced from workflow inputs.", + examples=["my_dataset", "$inputs.target_al_dataset"], + ) + predictions: Optional[ + Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + ] + ) + ] = Field( + default=None, + description="Optional model predictions to upload alongside images. Predictions are saved as annotations (pre-labels) in the Roboflow dataset when persist_predictions is enabled, allowing predictions to serve as starting points for annotation review and correction. Supports object detection, instance segmentation, keypoint detection, and classification predictions. If None, only images are uploaded.", + examples=["$steps.object_detection_model.predictions"], + json_schema_extra={"always_visible": True}, + ) + data_percentage: Union[FloatZeroToHundred, Selector(kind=[FLOAT_KIND])] = Field( + default=100, + description="Percentage of input data (0.0 to 100.0) to randomly sample for upload. This enables probabilistic data collection where only a subset of inputs are uploaded, reducing storage and annotation costs. For example, 25.0 uploads approximately 25% of images (one in four on average), 50.0 uploads half, and 100.0 uploads everything (no sampling). Random sampling occurs before quota checking and image processing, making it efficient for large-scale data collection workflows.", + examples=[100, 25, 50, "$inputs.sampling_rate"], + ) + minutely_usage_limit: int = Field( + default=10, + description="Maximum number of image uploads allowed per minute for this quota. Part of the usage quota system that enforces rate limits for active learning data collection. Uploads exceeding this limit are skipped to prevent excessive data collection. Works together with hourly_usage_limit and daily_usage_limit to provide multi-level rate limiting. Note: This quota is checked after probabilistic sampling via data_percentage.", + examples=[10, 60], + ) + hourly_usage_limit: int = Field( + default=100, + description="Maximum number of image uploads allowed per hour for this quota. Part of the usage quota system that enforces rate limits for active learning data collection. Uploads exceeding this limit are skipped to prevent excessive data collection. Works together with minutely_usage_limit and daily_usage_limit to provide multi-level rate limiting. Note: This quota is checked after probabilistic sampling via data_percentage.", + examples=[10, 60], + ) + daily_usage_limit: int = Field( + default=1000, + description="Maximum number of image uploads allowed per day for this quota. Part of the usage quota system that enforces rate limits for active learning data collection. Uploads exceeding this limit are skipped to prevent excessive data collection. Works together with minutely_usage_limit and hourly_usage_limit to provide multi-level rate limiting. Note: This quota is checked after probabilistic sampling via data_percentage.", + examples=[10, 60], + ) + usage_quota_name: str = Field( + description="Unique identifier for tracking usage quotas (minutely, hourly, daily limits). Used internally to manage rate limiting across multiple upload operations. Each unique quota name maintains separate counters, allowing different upload strategies or data collection workflows to have independent rate limits.", + examples=["quota-for-data-sampling-1"], + json_schema_extra={"hidden": True}, + ) + max_image_size: Tuple[int, int] = Field( + default=(1920, 1080), + description="Maximum dimensions (width, height) for uploaded images. Images exceeding these dimensions are automatically resized while preserving aspect ratio before uploading. Default is (1920, 1080) for higher resolution data collection. Use smaller sizes (e.g., (512, 512)) for efficient storage and faster uploads, or keep the default for preserving image quality.", + examples=[(1920, 1080), (512, 512)], + ) + compression_level: int = Field( + default=95, + gt=0, + le=100, + description="JPEG compression quality level for uploaded images, ranging from 1 (highest compression, smallest file size, lower quality) to 100 (no compression, largest file size, highest quality). Default is 95 for better image quality preservation. Higher values preserve more image quality but increase storage and bandwidth usage. Typical values range from 70-95 for balanced quality and size.", + examples=[95, 75], + ) + registration_tags: Union[ + List[Union[Selector(kind=[STRING_KIND]), str]], + Selector(kind=[LIST_OF_VALUES_KIND]), + ] = Field( + default_factory=list, + description="List of tags to attach to uploaded images for organization and filtering in Roboflow. Tags can be static strings (e.g., 'location-florida', 'camera-1') or dynamic values from workflow inputs. Tags help organize collected data, filter images in Roboflow, and add metadata for dataset management. Can be an empty list if no tags are needed.", + examples=[ + ["location-florida", "factory-name", "$inputs.dynamic_tag"], + "$inputs.tags", + ], + ) + persist_predictions: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="If True, model predictions are saved as annotations (pre-labels) in the Roboflow dataset alongside images. This enables predictions to serve as starting points for annotation, allowing reviewers to correct or approve labels rather than creating them from scratch. If False, only images are uploaded without annotations. Enabling this accelerates annotation workflows by providing initial labels.", + examples=[True, False, "$inputs.persist_predictions"], + ) + disable_sink: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + description="If True, the block execution is disabled and no uploads occur. This allows temporarily disabling data collection without removing the block from workflows, useful for testing, debugging, or conditional data collection. When disabled, returns a message indicating the sink was disabled. Default is False (uploads enabled).", + examples=[True, "$inputs.disable_active_learning"], + ) + fire_and_forget: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="If True, uploads execute asynchronously (fire-and-forget mode), allowing the workflow to continue immediately without waiting for upload completion. This improves workflow performance but prevents error handling. If False, uploads execute synchronously, blocking workflow execution until completion and allowing proper error handling and status reporting. Use async mode (True) for production workflows where speed is prioritized, and sync mode (False) for debugging or when error handling is critical.", + ) + labeling_batch_prefix: Union[str, Selector(kind=[STRING_KIND])] = Field( + default="workflows_data_collector", + description="Prefix used to generate labeling batch names for organizing uploaded images in Roboflow. Combined with the batch recreation frequency and timestamps to create batch names like 'workflows_data_collector_2024_01_15'. Batches help organize collected data for labeling, making it easier to manage and review uploaded images in groups. Can be customized to match your organization scheme.", + examples=["my_labeling_batch_name"], + ) + labeling_batches_recreation_frequency: BatchCreationFrequency = Field( + default="never", + description="Frequency at which new labeling batches are automatically created for uploaded images. Options: 'never' (all images go to the same batch), 'daily' (new batch each day), 'weekly' (new batch each week), 'monthly' (new batch each month). Batch timestamps are appended to the labeling_batch_prefix to create unique batch names. Automatically organizing uploads into time-based batches simplifies dataset management and makes it easier to track and review collected data over time.", + examples=["never", "daily"], + ) + image_name: Optional[Union[str, Selector(kind=[STRING_KIND])]] = Field( + default=None, + description="Optional custom name for the uploaded image. This is useful when you want to preserve the original filename or use a meaningful identifier (e.g., serial number, timestamp) for the image in the Roboflow dataset. The name should not include file extension. If not provided, a UUID will be generated automatically.", + examples=["serial_12345", "camera1_frame_001", "$inputs.filename"], + ) + metadata: Dict[str, Union[str, int, float, bool, Selector()]] = Field( + default_factory=dict, + description="Optional key-value metadata to attach to uploaded images. Metadata is stored as user_metadata on the image in Roboflow and can be used for filtering and organization. Values can be static strings, numbers, booleans, or references to workflow inputs/steps.", + examples=[{"camera_id": "cam_01", "location": "$inputs.location"}, {}], + ) + + @classmethod + def get_air_gapped_availability(cls) -> AirGappedAvailability: + return AirGappedAvailability(available=False, reason="requires_internet") + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images", "predictions", "image_name"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="message", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + if self.disable_sink is True: + # Sink literally disabled โ€” execution returns before the target + # project is ever accessed. + return [] + return [roboflow_platform_project(project_url=self.target_project)] + + +class RoboflowDatasetUploadBlockV2(WorkflowBlock): + + def __init__( + self, + cache: BaseCache, + api_key: Optional[str], + background_tasks: Optional[BackgroundTasks], + thread_pool_executor: Optional[ThreadPoolExecutor], + disable_sinks: bool = False, + ): + self._cache = cache + self._api_key = api_key + self._background_tasks = background_tasks + self._thread_pool_executor = thread_pool_executor + self._disable_sinks = disable_sinks + + @classmethod + def get_init_parameters(cls) -> List[str]: + return [ + "cache", + "api_key", + "background_tasks", + "thread_pool_executor", + "disable_sinks", + ] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + predictions: Optional[Batch[TensorNativePrediction]], + target_project: str, + usage_quota_name: str, + data_percentage: float, + minutely_usage_limit: int, + persist_predictions: bool, + hourly_usage_limit: int, + daily_usage_limit: int, + max_image_size: Tuple[int, int], + compression_level: int, + registration_tags: List[str], + disable_sink: bool, + fire_and_forget: bool, + labeling_batch_prefix: str, + labeling_batches_recreation_frequency: BatchCreationFrequency, + image_name: Optional[Batch[Optional[str]]] = None, + metadata: Optional[Dict[str, Any]] = None, + ) -> BlockResult: + if self._disable_sinks: + return [ + { + "error_status": False, + "message": disabled_sink_message(disabled_by_execution_policy=True), + } + for _ in range(len(images)) + ] + if self._api_key is None: + raise ValueError( + "RoboflowDataCollector block cannot run without Roboflow API key. " + "If you do not know how to get API key - visit " + "https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key to learn how to " + "retrieve one." + ) + if disable_sink: + return [ + { + "error_status": False, + "message": disabled_sink_message( + disabled_by_execution_policy=False + ), + } + for _ in range(len(images)) + ] + result = [] + predictions = [None] * len(images) if predictions is None else predictions + image_names = [None] * len(images) if image_name is None else image_name + for image, prediction, img_name in zip(images, predictions, image_names): + error_status, message = maybe_register_datapoint_at_roboflow( + image=image, + prediction=prediction, + target_project=target_project, + usage_quota_name=usage_quota_name, + data_percentage=data_percentage, + persist_predictions=persist_predictions, + minutely_usage_limit=minutely_usage_limit, + hourly_usage_limit=hourly_usage_limit, + daily_usage_limit=daily_usage_limit, + max_image_size=max_image_size, + compression_level=compression_level, + registration_tags=registration_tags, + fire_and_forget=fire_and_forget, + labeling_batch_prefix=labeling_batch_prefix, + new_labeling_batch_frequency=labeling_batches_recreation_frequency, + cache=self._cache, + background_tasks=self._background_tasks, + thread_pool_executor=self._thread_pool_executor, + api_key=self._api_key, + image_name=img_name, + metadata=metadata, + ) + result.append({"error_status": error_status, "message": message}) + return result + + +def maybe_register_datapoint_at_roboflow( + image: WorkflowImageData, + prediction: Optional[TensorNativePrediction], + target_project: str, + usage_quota_name: str, + data_percentage: float, + persist_predictions: bool, + minutely_usage_limit: int, + hourly_usage_limit: int, + daily_usage_limit: int, + max_image_size: Tuple[int, int], + compression_level: int, + registration_tags: List[str], + fire_and_forget: bool, + labeling_batch_prefix: str, + new_labeling_batch_frequency: BatchCreationFrequency, + cache: BaseCache, + background_tasks: Optional[BackgroundTasks], + thread_pool_executor: Optional[ThreadPoolExecutor], + api_key: str, + image_name: Optional[str] = None, + metadata: Optional[Dict[str, Any]] = None, +) -> Tuple[bool, str]: + normalised_probability = data_percentage / 100 + if random.random() < normalised_probability: + return register_datapoint_at_roboflow( + image=image, + prediction=prediction, + target_project=target_project, + usage_quota_name=usage_quota_name, + persist_predictions=persist_predictions, + minutely_usage_limit=minutely_usage_limit, + hourly_usage_limit=hourly_usage_limit, + daily_usage_limit=daily_usage_limit, + max_image_size=max_image_size, + compression_level=compression_level, + registration_tags=registration_tags, + fire_and_forget=fire_and_forget, + labeling_batch_prefix=labeling_batch_prefix, + new_labeling_batch_frequency=new_labeling_batch_frequency, + cache=cache, + background_tasks=background_tasks, + thread_pool_executor=thread_pool_executor, + api_key=api_key, + image_name=image_name, + metadata=metadata, + ) + return False, "Registration skipped due to sampling settings" diff --git a/inference/core/workflows/core_steps/sinks/roboflow/model_monitoring_inference_aggregator/v1_tensor.py b/inference/core/workflows/core_steps/sinks/roboflow/model_monitoring_inference_aggregator/v1_tensor.py new file mode 100644 index 0000000000..35b561247e --- /dev/null +++ b/inference/core/workflows/core_steps/sinks/roboflow/model_monitoring_inference_aggregator/v1_tensor.py @@ -0,0 +1,604 @@ +import hashlib +import logging +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime +from functools import partial +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +from fastapi import BackgroundTasks +from pydantic import BaseModel, ConfigDict, Field, field_validator + +from inference.core.cache.base import BaseCache +from inference.core.env import DEVICE_ID +from inference.core.managers.metrics import get_system_info +from inference.core.roboflow_api import ( + get_roboflow_workspace, + send_inference_results_to_model_monitoring, +) +from inference.core.version import __version__ +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + KeyPointPrediction, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.sinks.noop import disabled_sink_response +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + INFERENCE_ID_KEY, + PREDICTION_TYPE_KEY, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + ROBOFLOW_MODEL_ID_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + AirGappedAvailability, + BlockResult, + DependentResource, + ModelRequiredAction, + Runtime, + RuntimeInputMode, + RuntimeRestriction, + Severity, + WorkflowBlock, + WorkflowBlockManifest, + roboflow_platform_model, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +TensorNativePrediction = Union[ + Detections, + InstanceDetections, + KeyPointPrediction, + ClassificationPrediction, + MultiLabelClassificationPrediction, +] + +# Key the classification model producers attach to `image_metadata` carrying the +# resolved confidence threshold. Defined locally (mirroring the other `_tensor` +# classification blocks) so this file stays self-contained. Flag-OFF the numpy +# classification response was already cut at this threshold before model monitoring +# ever saw it, so the tensor sink must apply the same cut here. +CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY = "classification_confidence_threshold" + +SHORT_DESCRIPTION = "Periodically report an aggregated sample of inference results to Roboflow Model Monitoring." + +LONG_DESCRIPTION = """ +Periodically aggregate and report a curated sample of inference predictions to Roboflow Model Monitoring by collecting predictions in memory, grouping by class, selecting the most confident prediction per class, and sending aggregated results at configurable intervals to enable efficient video processing monitoring, production analytics, and model performance tracking workflows with minimal performance overhead. + +## How This Block Works + +This block aggregates predictions over time and sends representative samples to Roboflow Model Monitoring at regular intervals, reducing API calls and maintaining video processing performance. The block: + +1. Receives predictions and configuration: + - Takes predictions from any supported model type (object detection, instance segmentation, keypoint detection, or classification) + - Receives model ID for identification in Model Monitoring + - Accepts frequency parameter specifying reporting interval in seconds + - Receives execution mode flag (fire-and-forget) +2. Validates Roboflow API key: + - Checks that a valid Roboflow API key is available (required for API access) + - Raises an error if API key is missing with instructions on how to retrieve one +3. Collects predictions in memory: + - Stores predictions in an in-memory aggregator organized by model ID + - Accumulates predictions between reporting intervals + - Maintains state for the duration of the workflow execution session +4. Checks reporting interval: + - Uses cache to track last report time based on unique aggregator key + - Calculates time elapsed since last report + - Compares elapsed time to configured frequency threshold + - Skips reporting if interval has not been reached (returns status message) +5. Consolidates predictions when reporting: + - Formats all collected predictions for Model Monitoring + - Groups predictions by class name across all collected data + - For each class, sorts predictions by confidence (highest first) + - Selects the most confident prediction per class as representative sample + - Creates a curated set of predictions (one per class with highest confidence) +6. Retrieves workspace information: + - Gets workspace ID from Roboflow API using the provided API key + - Uses caching (15-minute expiration) to avoid repeated API calls + - Caches workspace name using MD5 hash of API key as cache key +7. Sends aggregated data to Model Monitoring: + - Constructs inference data payload with timestamp, source info, device ID, and server version + - Includes system information (if available) for monitoring context + - Sends aggregated predictions (one per class) to Roboflow Model Monitoring API + - Flushes in-memory aggregator after sending (starts fresh collection) + - Updates last report time in cache +8. Executes synchronously or asynchronously: + - **Asynchronous mode (fire_and_forget=True)**: Submits task to background thread pool or FastAPI background tasks, allowing workflow to continue without waiting for API call to complete + - **Synchronous mode (fire_and_forget=False)**: Waits for API call to complete and returns immediate status, useful for debugging and error handling +9. Returns status information: + - Outputs error_status indicating success (False) or failure (True) + - Outputs message with reporting status or error details + - Provides feedback on whether aggregation was sent or skipped + +The block is optimized for video processing workflows where sending every prediction would create excessive API calls and impact performance. By aggregating predictions and selecting representative samples (most confident per class), the block provides meaningful monitoring data while minimizing overhead. The interval-based reporting ensures regular updates to Model Monitoring without constant API calls. + +## Common Use Cases + +#### ๐Ÿ” Why Use This Block? + +This block is a game-changer for projects relying on video processing in Workflows. +With its aggregation process, it identifies the most confident predictions across classes and sends +them at regular intervals in small messages to Roboflow backend - ensuring that video processing +performance is impacted to the least extent. + +Perfect for: + +* Monitoring production line performance in real-time ๐Ÿญ. + +* Debugging and validating your modelโ€™s performance over time โฑ๏ธ. + +* Providing actionable insights from inference workflows with minimal overhead ๐Ÿ”ง. + +#### ๐Ÿšจ Limitations + +* The block is should not be relied on when running Workflow in `inference` server or via HTTP request to Roboflow +hosted platform, as the internal state is not persisted in a memory that would be accessible for all requests to +the server, causing aggregation to **only have a scope of single request**. We will solve that problem in future +releases if proven to be serious limitation for clients. + +## Connecting to Other Blocks + +This block receives predictions and outputs status information: + +- **After model blocks** (Object Detection Model, Instance Segmentation Model, Classification Model, Keypoint Detection Model) to aggregate and report predictions to Model Monitoring (e.g., aggregate detection results, report classification outputs, monitor model predictions), enabling model-to-monitoring workflows +- **After filtering or analytics blocks** (DetectionsFilter, ContinueIf, OverlapFilter) to aggregate filtered or analyzed results for monitoring (e.g., aggregate filtered detections, report analytics results, monitor processed predictions), enabling analysis-to-monitoring workflows +- **In video processing workflows** to efficiently monitor video analysis with minimal performance impact (e.g., aggregate video frame detections, report video processing results, monitor video analysis performance), enabling video monitoring workflows +- **After preprocessing or transformation blocks** to monitor transformed predictions (e.g., aggregate transformed detections, report processed results, monitor transformation outputs), enabling transformation-to-monitoring workflows +- **In production deployment workflows** to track model performance in production environments (e.g., monitor production inference, track deployment performance, report production metrics), enabling production monitoring workflows +- **As a sink block** to send aggregated monitoring data without blocking workflow execution (e.g., background monitoring reporting, non-blocking analytics, efficient data collection), enabling sink-to-monitoring workflows + +## Requirements + +This block requires a valid Roboflow API key configured in the environment or workflow configuration. The API key is required to authenticate with Roboflow API and access Model Monitoring features. Visit https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key to learn how to retrieve an API key. The block maintains in-memory state for aggregation, which means it works best for long-running workflows (like video processing with InferencePipeline). The block should not be relied upon when running workflows in inference server or via HTTP requests to Roboflow hosted platform, as the internal state is only accessible for single requests and aggregation scope is limited to single request execution. The block aggregates data for all video feeds connected to a single InferencePipeline process (cannot separate aggregations per video feed). The frequency parameter must be at least 1 second. For more information on Model Monitoring at Roboflow, see https://docs.roboflow.com/deploy/model-monitoring. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Model Monitoring Inference Aggregator", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "sink", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-chart-line", + "blockPriority": 8.5, + "requires_rf_key": True, + }, + } + ) + type: Literal["roboflow_core/model_monitoring_inference_aggregator@v1"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + ] + ) = Field( + description="Model predictions (object detection, instance segmentation, keypoint detection, or classification) to aggregate and report to Roboflow Model Monitoring. Predictions are collected in memory, grouped by class name, and the most confident prediction per class is selected as a representative sample. Predictions accumulate between reporting intervals based on the frequency setting. Supported prediction types: supervision Detections objects or classification prediction dictionaries.", + examples=[ + "$steps.object_detection.predictions", + "$steps.classification.predictions", + "$steps.instance_segmentation.predictions", + ], + ) + model_id: Selector(kind=[ROBOFLOW_MODEL_ID_KIND]) = Field( + description="Roboflow model ID (format: 'project/version') to associate with the predictions in Model Monitoring. This identifies which model generated the predictions being reported. The model ID is included in the monitoring data sent to Roboflow, allowing you to track performance per model in the Model Monitoring dashboard.", + examples=["my_project/3", "production_model/1", "detection_model/5"], + ) + frequency: Union[ + int, + Selector(kind=[STRING_KIND]), + ] = Field( + default=5, + description="Reporting frequency in seconds. Specifies how often aggregated predictions are sent to Roboflow Model Monitoring. For example, if set to 5, the block collects predictions for 5 seconds, then sends the aggregated sample (one most confident prediction per class) to Model Monitoring. Must be at least 1 second. Lower values provide more frequent updates but increase API calls. Higher values reduce API calls but provide less frequent updates. Default: 5 seconds. Works well for video processing where you want regular but not excessive reporting.", + examples=[3, 5, 10, 30, 60], + ) + unique_aggregator_key: str = Field( + description="Unique key used internally to track the aggregation session and cache last report time. This key must be unique for each instance of this block in your workflow. The key is used to create cache entries that track when the last report was sent, enabling interval-based reporting. This field is automatically generated and hidden in the UI.", + examples=["session-1v73kdhfse"], + json_schema_extra={"hidden": True}, + ) + fire_and_forget: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + description="Execution mode flag. When True (default), the block runs asynchronously in the background, allowing the workflow to continue processing without waiting for the API call to complete. This provides faster workflow execution but errors are not immediately available. When False, the block runs synchronously and waits for the API call to complete, returning immediate status and error information. Use False for debugging and error handling, True for production workflows where performance is prioritized.", + examples=[True, False], + ) + + @classmethod + def get_air_gapped_availability(cls) -> AirGappedAvailability: + return AirGappedAvailability(available=False, reason="requires_internet") + + @field_validator("frequency") + @classmethod + def ensure_frequency_is_correct(cls, value: Any) -> Any: + if isinstance(value, int) and value < 1: + raise ValueError("`frequency` cannot be lower than 1.") + return value + + @field_validator("model_id") + @classmethod + def ensure_model_id_is_correct(cls, value: Any) -> Any: + if isinstance(value, str) and value == "": + raise ValueError("`model_id` cannot be empty.") + return value + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="message", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + def discover_dependent_resources(self) -> Optional[List[DependentResource]]: + # References the platform model entity for monitoring; no weights are pulled. + return [ + roboflow_platform_model( + model_id=self.model_id, + required_action=ModelRequiredAction.ACCESS, + ) + ] + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restriction = RuntimeRestriction( + severity=Severity.SOFT, + note=( + "Aggregation buffers are stored in process memory while the " + "reporting interval is tracked in cache. With remote step " + "execution on stateless or multi-replica HTTP runtimes, " + "predictions may be collected by different worker processes, " + "so reports can under-collect or flush partial aggregation " + "windows. Use local step execution in an InferencePipeline " + "for stable video aggregation." + ), + applies_to_runtimes=[ + Runtime.HOSTED_SERVERLESS, + Runtime.DEDICATED_DEPLOYMENT, + ], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + applies_to_input_modes=[RuntimeInputMode.VIDEO], + ) + return [restriction, STILL_IMAGE_INPUT_SOFT_RESTRICTION] + + +class ParsedPrediction(BaseModel): + model_config = ConfigDict( + protected_namespaces=(), + ) + + model_id: str + class_name: str + confidence: float + inference_id: str + model_type: str + + +class PredictionsAggregator(object): + + def __init__(self): + self._raw_predictions: dict[str, List[TensorNativePrediction]] = {} + + def collect(self, value: TensorNativePrediction, model_id: str) -> None: + # TODO: push into global state, otherwise for HTTP server use, + # state would at most have 1 prediction!!! + if model_id not in self._raw_predictions: + self._raw_predictions[model_id] = [] + self._raw_predictions[model_id].append(value) + + def get_and_flush(self) -> List[ParsedPrediction]: + predictions = self._consolidate() + self._raw_predictions = {} + return predictions + + def _consolidate(self) -> List[ParsedPrediction]: + formatted_predictions = [] + for model_id, predictions in self._raw_predictions.items(): + formatted_predictions.extend( + format_predictions_for_model_monitoring(predictions, model_id) + ) + class_groups: Dict[str, List[ParsedPrediction]] = defaultdict(list) + for prediction in formatted_predictions: + class_name = prediction.class_name + class_groups[class_name].append(prediction) + representative_predictions = [] + for class_name, predictions in class_groups.items(): + predictions.sort(key=lambda x: x.confidence, reverse=True) + representative_predictions.append(predictions[0]) + return representative_predictions + + +class ModelMonitoringInferenceAggregatorBlockV1(WorkflowBlock): + + def __init__( + self, + cache: BaseCache, + api_key: Optional[str], + background_tasks: Optional[BackgroundTasks], + thread_pool_executor: Optional[ThreadPoolExecutor], + disable_sinks: bool = False, + ): + if api_key is None and not disable_sinks: + raise ValueError( + "ModelMonitoringInferenceAggregator block cannot run without Roboflow API key. " + "If you do not know how to get API key - visit " + "https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key to learn how to " + "retrieve one." + ) + self._api_key = api_key + self._cache = cache + self._background_tasks = background_tasks + self._thread_pool_executor = thread_pool_executor + self._disable_sinks = disable_sinks + self._predictions_aggregator = PredictionsAggregator() + + @classmethod + def get_init_parameters(cls) -> List[str]: + return [ + "api_key", + "cache", + "background_tasks", + "thread_pool_executor", + "disable_sinks", + ] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + fire_and_forget: bool, + predictions: TensorNativePrediction, + frequency: int, + unique_aggregator_key: str, + model_id: str, + ) -> BlockResult: + if self._disable_sinks: + return disabled_sink_response() + self._last_report_time_cache_key = f"workflows:steps_cache:roboflow_core/model_monitoring_inference_aggregator@v1:{unique_aggregator_key}:last_report_time" + if predictions: + self._predictions_aggregator.collect(predictions, model_id) + if not self._is_in_reporting_range(frequency): + return { + "error_status": False, + "message": "Not in reporting range, skipping report. (Ok)", + } + preds = self._predictions_aggregator.get_and_flush() + registration_task = partial( + send_to_model_monitoring_request, + cache=self._cache, + last_report_time_cache_key=self._last_report_time_cache_key, + api_key=self._api_key, + predictions=preds, + ) + error_status = False + message = "Reporting happens in the background task" + if fire_and_forget and self._background_tasks: + self._background_tasks.add_task(registration_task) + elif fire_and_forget and self._thread_pool_executor: + self._thread_pool_executor.submit(registration_task) + else: + error_status, message = registration_task() + self._cache.set(self._last_report_time_cache_key, datetime.now().isoformat()) + return { + "error_status": error_status, + "message": message, + } + + def _is_in_reporting_range(self, frequency: int) -> bool: + now = datetime.now() + last_report_time_str = self._cache.get(self._last_report_time_cache_key) + if last_report_time_str is None: + self._cache.set(self._last_report_time_cache_key, now.isoformat()) + v = self._cache.get(self._last_report_time_cache_key) + last_report_time = now + else: + last_report_time = datetime.fromisoformat(last_report_time_str) + time_elapsed = int((now - last_report_time).total_seconds()) + return time_elapsed >= int(frequency) + + +# TODO: maybe make this a helper or decorator, it's used in multiple places +def get_workspace_name( + api_key: str, + cache: BaseCache, +) -> str: + # codeql[py/weak-sensitive-data-hashing]: MD5 cache fingerprint; not crypto storage. + api_key_hash = hashlib.md5(api_key.encode("utf-8")).hexdigest() + cache_key = f"workflows:api_key_to_workspace:{api_key_hash}" + cached_workspace_name = cache.get(cache_key) + if cached_workspace_name: + return cached_workspace_name + workspace_name_from_api = get_roboflow_workspace(api_key=api_key) + cache.set(key=cache_key, value=workspace_name_from_api, expire=900) + return workspace_name_from_api + + +def send_to_model_monitoring_request( + cache: BaseCache, + last_report_time_cache_key: str, + api_key: str, + predictions: List[ParsedPrediction], +) -> Tuple[bool, str]: + workspace_id = get_workspace_name(api_key=api_key, cache=cache) + try: + inference_data = { + "timestamp": datetime.now().isoformat(), + "source": "workflow", + "source_info": "ModelMonitoringInferenceAggregatorBlockV1", + "inference_results": [], + "device_id": DEVICE_ID, + "inference_server_version": __version__, + } + system_info = get_system_info() + if system_info: + for key, value in system_info.items(): + inference_data[key] = value + inference_data["inference_results"] = [p.model_dump() for p in predictions] + send_inference_results_to_model_monitoring( + api_key, workspace_id, inference_data + ) + cache.set(last_report_time_cache_key, datetime.now().isoformat()) + return ( + False, + "Data sent successfully", + ) + except Exception as error: + logging.warning(f"Could not upload inference data. Reason: {error}") + return ( + True, + f"Error while uploading inference data. Error type: {type(error)}. Details: {error}", + ) + + +def format_predictions_for_model_monitoring( + predictions_list: List[TensorNativePrediction], + model_id: str, +) -> List[ParsedPrediction]: + results = [] + for predictions in predictions_list: + if isinstance(predictions, ClassificationPrediction): + # Tensor-native single-label classification: `confidence` is the full + # softmax distribution `(bs, num_classes)` and class names live in + # `images_metadata[i][CLASS_NAMES_KEY]`. The numpy sibling iterated the + # `predictions` list (one entry per class) of the serialised dict, so we + # emit one `ParsedPrediction` per class with its distribution confidence. + images_metadata = predictions.images_metadata or [] + class_ids = predictions.class_id.detach().cpu().tolist() + for image_index in range(len(class_ids)): + image_metadata = ( + images_metadata[image_index] + if image_index < len(images_metadata) + else {} + ) + class_names = image_metadata.get(CLASS_NAMES_KEY) or {} + prediction_type = image_metadata.get(PREDICTION_TYPE_KEY, "") + inference_id = image_metadata.get(INFERENCE_ID_KEY, "") + image_confidence = predictions.confidence[image_index].detach().cpu() + confidence_threshold = image_metadata.get( + CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY + ) + # Mirror the numpy `prepare_classification_response` / + # `_serialise_classification_model_style` semantics that flag-OFF + # produced the serialised `predictions` list model monitoring then + # iterated: drop classes whose (raw) softmax score is below the + # confidence threshold when the producer attached one, round each + # surviving score to 4 dp, and sort descending by confidence. Without + # this the sink emitted one entry per class of the FULL class map + # (e.g. 1000 for a 1000-class model) with raw softmax scores, diverging + # from the few above-threshold, rounded entries flag-OFF reports. + per_class_predictions = [] + for class_id in sorted(class_names): + class_score = float(image_confidence[class_id]) + if ( + confidence_threshold is not None + and class_score < confidence_threshold + ): + continue + per_class_predictions.append( + (class_names[class_id], round(class_score, 4)) + ) + per_class_predictions.sort(key=lambda item: item[1], reverse=True) + for class_name, class_confidence in per_class_predictions: + pred_instance = ParsedPrediction( + class_name=class_name, + confidence=class_confidence, + inference_id=inference_id, + model_type=prediction_type, + model_id=model_id, + ) + results.append(pred_instance) + elif isinstance(predictions, MultiLabelClassificationPrediction): + # Tensor-native multi-label classification: `confidence` is the full + # sigmoid distribution `(num_classes,)` and class names live in + # `image_metadata[CLASS_NAMES_KEY]`. The numpy sibling iterated the + # `predictions` dict (one entry per class) of the serialised dict. + image_metadata = predictions.image_metadata or {} + class_names = image_metadata.get(CLASS_NAMES_KEY) or {} + prediction_type = image_metadata.get(PREDICTION_TYPE_KEY, "") + inference_id = image_metadata.get(INFERENCE_ID_KEY, "") + confidence = predictions.confidence.detach().cpu() + confidence_threshold = image_metadata.get( + CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY + ) + # Mirror the numpy multi-label serialisation model monitoring consumed + # flag-OFF (`prepare_multi_label_classification_response` / + # `_serialise_classification_model_style`): the `predictions` dict keeps + # every class whose (raw, UNROUNDED) sigmoid score clears the confidence + # threshold when the producer attached one (e.g. visual_search_classifier + # gap-fills a dense vector and attaches it); plain model predictions attach + # no threshold and keep every class. Unlike single-label, numpy does NOT + # round or sort the multi-label dict, so neither do we. + for class_id in sorted(class_names): + class_score = float(confidence[class_id]) + if ( + confidence_threshold is not None + and class_score < confidence_threshold + ): + continue + pred_instance = ParsedPrediction( + class_name=class_names[class_id], + confidence=class_score, + inference_id=inference_id, + model_type=prediction_type, + model_id=model_id, + ) + results.append(pred_instance) + else: + # Detections / InstanceDetections / keypoint-detection tuple. Iterate the + # native 7-tuple; the numpy sibling read per-detection `class_name`, + # `confidence`, `inference_id` and `prediction_type` from + # `sv.Detections.data[...]`. Tensor-native keeps the class-id -> name map, + # `inference_id` and `prediction_type` per-image in `image_metadata`, so we + # resolve `class_name` via that mapping and confidence from the tensor. + _key_points, detections = split_key_point_prediction(predictions) + image_metadata = detections.image_metadata or {} + class_names = image_metadata.get(CLASS_NAMES_KEY) or {} + inference_id = image_metadata.get(INFERENCE_ID_KEY, "") + prediction_type = image_metadata.get(PREDICTION_TYPE_KEY, "") + for ( + _xyxy, + _mask, + class_id, + confidence, + _tracker_id, + _data, + _meta, + ) in detections: + class_id_int = int(class_id) + class_name = class_names.get(class_id_int, "") + prediction = ParsedPrediction( + class_name=class_name, + confidence=(float(confidence) if confidence is not None else 0.0), + inference_id=inference_id, + model_type=prediction_type, + model_id=model_id, + ) + results.append(prediction) + return results diff --git a/inference/core/workflows/core_steps/sinks/roboflow/vision_events/v1_tensor.py b/inference/core/workflows/core_steps/sinks/roboflow/vision_events/v1_tensor.py new file mode 100644 index 0000000000..8b5377d380 --- /dev/null +++ b/inference/core/workflows/core_steps/sinks/roboflow/vision_events/v1_tensor.py @@ -0,0 +1,1114 @@ +import os +from concurrent.futures import ThreadPoolExecutor +from datetime import datetime, timezone +from functools import partial +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import numpy as np +import requests +from fastapi import BackgroundTasks +from pydantic import ConfigDict, Field, NonNegativeFloat, NonNegativeInt + +from inference.core.env import API_BASE_URL +from inference.core.logger import logger +from inference.core.utils.image_utils import encode_image_to_jpeg_bytes +from inference.core.utils.url_utils import wrap_url +from inference.core.workflows.core_steps.common.keypoints import real_keypoints_count +from inference.core.workflows.core_steps.common.serializers import mask_to_polygon +from inference.core.workflows.core_steps.common.tensor_native import ( + KeyPointPrediction, + instance_mask_to_numpy, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.sinks.noop import disabled_sink_message +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + POLYGON_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + IMAGE_KIND, + INTEGER_KIND, + ROBOFLOW_SOLUTION_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + COOLDOWN_HTTP_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +# Tensor-native prediction union the block accepts. Detection predictions arrive +# as `inference_models` dataclasses (or the keypoint `(KeyPoints, Detections)` +# tuple) and classification predictions as the native classification dataclasses, +# rather than the `sv.Detections` / dict the numpy sibling consumed. +TensorNativePrediction = Union[ + Detections, + InstanceDetections, + KeyPointPrediction, + ClassificationPrediction, + MultiLabelClassificationPrediction, +] + +VALID_EVENT_TYPES = [ + "quality_check", + "inventory_count", + "safety_alert", + "custom", + "operator_feedback", +] + +QUALITY_CHECK_RELEVANT = { + "event_type": {"values": ["quality_check"], "required": True}, +} +INVENTORY_COUNT_RELEVANT = { + "event_type": {"values": ["inventory_count"], "required": True}, +} +SAFETY_ALERT_RELEVANT = { + "event_type": {"values": ["safety_alert"], "required": True}, +} +CUSTOM_RELEVANT = { + "event_type": {"values": ["custom"], "required": True}, +} +OPERATOR_FEEDBACK_RELEVANT = { + "event_type": {"values": ["operator_feedback"], "required": True}, +} +ALL_DATA_SCHEMAS_RELEVANT = { + "event_type": { + "values": [ + "quality_check", + "inventory_count", + "safety_alert", + "custom", + ], + }, +} + +SHORT_DESCRIPTION = "Send vision events to the Roboflow Vision Events API." + +LONG_DESCRIPTION = """ +Send images, model predictions, and event metadata to the Roboflow Vision Events API for +monitoring, quality control, safety alerting, and custom event tracking. + +## How This Block Works + +This block uploads workflow images and model predictions to the Roboflow Vision Events API, +creating structured events that can be queried, filtered, and visualized in the Roboflow +dashboard. + +1. Optionally uploads an input image and/or output image (visualization) to the Vision Events + image storage via the public API +2. Converts model predictions (object detection, classification, instance segmentation, or + keypoint detection) into the Vision Events annotation format and attaches them to the + input image +3. Creates a vision event with the specified event type, use case, event data, + and custom metadata +4. Supports fire-and-forget mode for non-blocking execution + +## Deployment Modes + +By default this block sends events to the **Roboflow Vision Events API** (cloud / +Serverless API), uploading images and posting the event over the public API. + +For edge deployments, enable **Write to Local Event Store** to send events to a +local Event Ingestion Service instead. In this mode images are embedded directly +in the request (no upload step) and the event is posted to `/v2/events`. +The event store URL defaults to `http://localhost:8001` and can be overridden. No +Roboflow API key is required in this mode; if the local service requires +authentication, set the `EVENT_INGESTION_API_KEY` environment variable on the +inference server. + +## Event Types + +- **quality_check**: Manufacturing/inspection QA with pass/fail result and optional confidence +- **inventory_count**: Inventory tracking with location, item count, and item type +- **safety_alert**: Safety violations with alert type, severity (low/medium/high), and description +- **custom**: User-defined events with a free-form value string +- **operator_feedback**: Operator review/correction of previous events (correct/incorrect/inconclusive) + +## Requirements + +The default (cloud) mode requires a valid Roboflow API key with `vision-events:write` +scope, configured in your environment or workflow configuration. No Roboflow API key is +needed when **Write to Local Event Store** is enabled (see Deployment Modes above). + +## Common Use Cases + +- **Quality Control**: Automatically log inspection results with images and detection overlays +- **Safety Monitoring**: Send safety alerts when violations are detected in video streams +- **Production Analytics**: Track inventory counts and production metrics with visual evidence +- **Active Monitoring**: Fire-and-forget event logging from real-time video processing workflows +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Roboflow Vision Events", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "sink", + "ui_manifest": { + "section": "data_storage", + "icon": "fal fa-eye", + "blockPriority": 1, + "popular": False, + # Not unconditionally required: the local event store mode needs no + # Roboflow key. A True value here walls off the whole block config in + # the inference-embedded (edge) editor, which would defeat that mode. + # The cloud path still raises at runtime if no key is available. + "requires_rf_key": False, + }, + } + ) + type: Literal["roboflow_core/roboflow_vision_events@v1"] + input_image: Optional[Selector(kind=[IMAGE_KIND])] = Field( + default=None, + title="Input Image", + description="The original input image. Uploaded to the Vision Events API and " + "used as the base image for detection annotations.", + examples=["$inputs.image", "$steps.cropping.crops"], + json_schema_extra={"always_visible": True}, + ) + output_image: Optional[Selector(kind=[IMAGE_KIND])] = Field( + default=None, + title="Output Image", + description="An optional output/visualized image (e.g., from a visualization " + "block). Displayed as the primary image in the Vision Events dashboard.", + examples=["$steps.visualization.image"], + json_schema_extra={"always_visible": True}, + ) + predictions: Optional[ + Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, + ] + ) + ] = Field( + default=None, + title="Predictions", + description="Optional model predictions to include as detection annotations on " + "the input image. Supports object detection, instance segmentation, keypoint " + "detection, and classification predictions.", + examples=["$steps.object_detection_model.predictions"], + json_schema_extra={"always_visible": True}, + ) + event_type: Union[ + Literal[ + "quality_check", + "inventory_count", + "safety_alert", + "custom", + "operator_feedback", + ], + Selector(kind=[STRING_KIND]), + ] = Field( + title="Event Type", + description="The type of vision event to create.", + examples=["quality_check", "custom", "$inputs.event_type"], + json_schema_extra={ + "always_visible": True, + "values_metadata": { + "quality_check": { + "name": "Quality Check", + "description": "Manufacturing/inspection QA with pass/fail result and optional confidence", + }, + "inventory_count": { + "name": "Inventory Count", + "description": "Inventory tracking with location, item count, and item type", + }, + "safety_alert": { + "name": "Safety Alert", + "description": "Safety violations with alert type, severity, and description", + }, + "custom": { + "name": "Custom", + "description": "User-defined events with a free-form value string", + }, + "operator_feedback": { + "name": "Operator Feedback", + "description": "Operator review/correction of previous events", + }, + }, + }, + ) + solution: Union[str, Selector(kind=[ROBOFLOW_SOLUTION_KIND, STRING_KIND])] = Field( + title="Use Case", + description="The use case to associate the event with. Events are " + "namespaced by use case within a workspace.", + examples=["my-use-case", "$inputs.use_case"], + ) + # --- External ID (shared across schemas) --- + external_id: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + title="External ID", + description="External identifier for correlation with other systems (max 1000 chars).", + examples=["batch-2025-001", "$inputs.external_id"], + json_schema_extra={ + "relevant_for": ALL_DATA_SCHEMAS_RELEVANT, + }, + ) + + # --- Quality Check fields --- + qc_result: Optional[ + Union[Selector(kind=[STRING_KIND]), Literal["pass", "fail"]] + ] = Field( + default=None, + title="Result", + description="Quality check result: pass or fail.", + examples=["pass", "fail", "$steps.qc_logic.result"], + json_schema_extra={ + "relevant_for": QUALITY_CHECK_RELEVANT, + "always_visible": True, + }, + ) + + # --- Inventory Count fields --- + location: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + title="Location", + description="Location identifier for inventory count.", + examples=["warehouse-A", "$inputs.location"], + json_schema_extra={"relevant_for": INVENTORY_COUNT_RELEVANT}, + ) + item_count: Optional[Union[Selector(kind=[INTEGER_KIND]), int]] = Field( + default=None, + title="Item Count", + description="Number of items counted.", + examples=[42, "$steps.counter.count"], + json_schema_extra={ + "relevant_for": INVENTORY_COUNT_RELEVANT, + "always_visible": True, + }, + ) + item_type: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + title="Item Type", + description="Type of item being counted.", + examples=["widget", "$inputs.item_type"], + json_schema_extra={"relevant_for": INVENTORY_COUNT_RELEVANT}, + ) + + # --- Safety Alert fields --- + alert_type: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + title="Alert Type", + description="Alert type identifier (e.g. no_hardhat, spill_detected).", + examples=["no_hardhat", "$steps.classifier.top_class"], + json_schema_extra={ + "relevant_for": SAFETY_ALERT_RELEVANT, + "always_visible": True, + }, + ) + severity: Optional[ + Union[Selector(kind=[STRING_KIND]), Literal["low", "medium", "high"]] + ] = Field( + default=None, + title="Severity", + description="Severity level for the safety alert.", + examples=["high", "$inputs.severity"], + json_schema_extra={"relevant_for": SAFETY_ALERT_RELEVANT}, + ) + alert_description: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + title="Description", + description="Description of the safety alert.", + examples=["Worker detected without hardhat in zone B"], + json_schema_extra={"relevant_for": SAFETY_ALERT_RELEVANT}, + ) + + # --- Custom Event fields --- + custom_value: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + title="Value", + description="Arbitrary value for custom events.", + examples=["anomaly detected at 14:32"], + json_schema_extra={ + "relevant_for": CUSTOM_RELEVANT, + "always_visible": True, + }, + ) + + # --- Operator Feedback fields --- + related_event_id: Optional[Union[Selector(kind=[STRING_KIND]), str]] = Field( + default=None, + title="Related Event ID", + description="The event ID of the event being reviewed.", + examples=["evt_abc123", "$inputs.related_event_id"], + json_schema_extra={ + "relevant_for": OPERATOR_FEEDBACK_RELEVANT, + "always_visible": True, + }, + ) + feedback: Optional[ + Union[ + Selector(kind=[STRING_KIND]), + Literal["correct", "incorrect", "inconclusive"], + ] + ] = Field( + default=None, + title="Feedback", + description="Operator feedback on the related event.", + examples=["correct", "incorrect", "$inputs.feedback"], + json_schema_extra={"relevant_for": OPERATOR_FEEDBACK_RELEVANT}, + ) + + custom_metadata: Dict[str, Union[str, int, float, bool, Selector()]] = Field( + default_factory=dict, + title="Custom Metadata", + description="Flat key-value metadata to attach to the event. Keys must match " + "pattern [a-zA-Z0-9_ -]+ (max 100 chars). String values max 1000 chars.", + examples=[{"camera_id": "cam_01", "location": "$inputs.location"}], + json_schema_extra={"additional_section": True}, + ) + fire_and_forget: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=True, + title="Fire and Forget", + description="If True, the event is sent asynchronously and the workflow " + "continues without waiting. If False, the block waits for the API response.", + examples=[True, "$inputs.fire_and_forget"], + ) + disable_sink: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + title="Disable Sink", + description="If True, the block is disabled and no events are sent.", + examples=[False, "$inputs.disable_vision_events"], + ) + cooldown_seconds: Union[ + NonNegativeInt, NonNegativeFloat, Selector(kind=[INTEGER_KIND, FLOAT_KIND]) + ] = Field( + default=1, + title="Cooldown", + description="Minimum number of seconds between consecutive events sent by " + "this block. Events triggered during the cooldown period are dropped and " + "the `throttling_status` output is set to True. Defaults to 1 second (at " + "most 1 event per second) so high-frequency video workflows do not flood " + "the Vision Events API with an event per frame. Set to 0 to disable rate " + "limiting for intentionally bursty use cases.", + examples=[1, 0.5, "$inputs.cooldown_seconds"], + json_schema_extra={"always_visible": True}, + ) + write_to_event_store: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( + default=False, + title="Write to Local Event Store", + description="If True, send the event to a local Event Ingestion Service " + "(edge deployment) instead of the Roboflow Vision Events API (cloud). " + "Images are embedded in the request and the event is posted to " + "`/v2/events`. No Roboflow API key is required in this mode.", + examples=[False, True, "$inputs.write_to_event_store"], + ) + event_store_url: Union[Selector(kind=[STRING_KIND]), str] = Field( + default="http://localhost:8001", + title="Event Store URL", + description="Base URL of the local Event Ingestion Service. Only used when " + "`Write to Local Event Store` is enabled.", + examples=["http://localhost:8001", "$inputs.event_store_url"], + json_schema_extra={ + "relevant_for": { + "write_to_event_store": { + "values": [True], + "required": False, + }, + }, + }, + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="error_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="throttling_status", kind=[BOOLEAN_KIND]), + OutputDefinition(name="event_id", kind=[STRING_KIND]), + OutputDefinition(name="message", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [COOLDOWN_HTTP_SOFT_RESTRICTION] + + +class RoboflowVisionEventsBlockV1(WorkflowBlock): + + def __init__( + self, + api_key: Optional[str], + background_tasks: Optional[BackgroundTasks], + thread_pool_executor: Optional[ThreadPoolExecutor], + disable_sinks: bool = False, + ): + self._api_key = api_key + self._background_tasks = background_tasks + self._thread_pool_executor = thread_pool_executor + self._disable_sinks = disable_sinks + self._last_event_fired: Optional[datetime] = None + + @classmethod + def get_init_parameters(cls) -> List[str]: + return ["api_key", "background_tasks", "thread_pool_executor", "disable_sinks"] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + input_image: Optional[WorkflowImageData], + output_image: Optional[WorkflowImageData], + predictions: Optional[TensorNativePrediction], + event_type: str, + solution: str, + custom_metadata: Dict[str, Any], + fire_and_forget: bool, + disable_sink: bool, + cooldown_seconds: Union[int, float] = 1, + write_to_event_store: bool = False, + event_store_url: str = "http://localhost:8001", + external_id: Optional[str] = None, + qc_result: Optional[str] = None, + location: Optional[str] = None, + item_count: Optional[int] = None, + item_type: Optional[str] = None, + alert_type: Optional[str] = None, + severity: Optional[str] = None, + alert_description: Optional[str] = None, + custom_value: Optional[str] = None, + related_event_id: Optional[str] = None, + feedback: Optional[str] = None, + ) -> BlockResult: + if self._disable_sinks or disable_sink: + return { + "error_status": False, + "throttling_status": False, + "event_id": "", + "message": disabled_sink_message( + disabled_by_execution_policy=self._disable_sinks + ), + } + + # Selector-resolved values bypass manifest validation; a negative + # cooldown behaves as 0 (rate limiting disabled). + cooldown_seconds = max(cooldown_seconds, 0) + seconds_since_last_event = cooldown_seconds + if self._last_event_fired is not None: + seconds_since_last_event = ( + datetime.now() - self._last_event_fired + ).total_seconds() + if seconds_since_last_event < cooldown_seconds: + logger.info("Activated `roboflow_core/roboflow_vision_events@v1` cooldown.") + return { + "error_status": False, + "throttling_status": True, + "event_id": "", + "message": "Sink cooldown applies", + } + + event_data = _build_event_data( + event_type=event_type, + external_id=external_id, + qc_result=qc_result, + location=location, + item_count=item_count, + item_type=item_type, + alert_type=alert_type, + severity=severity, + alert_description=alert_description, + custom_value=custom_value, + related_event_id=related_event_id, + feedback=feedback, + ) + + if write_to_event_store: + task = partial( + _execute_local_event, + event_store_url=event_store_url, + input_image=input_image, + output_image=output_image, + prediction=predictions, + event_type=event_type, + solution=solution, + event_data=event_data, + custom_metadata=custom_metadata, + ) + else: + if self._api_key is None: + raise ValueError( + "VisionEvents block cannot run without Roboflow API key. " + "If you do not know how to get API key - visit " + "https://docs.roboflow.com/api-reference/authentication" + "#retrieve-an-api-key to learn how to retrieve one." + ) + task = partial( + _execute_vision_event, + api_base_url=API_BASE_URL, + api_key=self._api_key, + input_image=input_image, + output_image=output_image, + prediction=predictions, + event_type=event_type, + solution=solution, + event_data=event_data, + custom_metadata=custom_metadata, + ) + + self._last_event_fired = datetime.now() + + if fire_and_forget and self._background_tasks: + self._background_tasks.add_task(task) + return { + "error_status": False, + "throttling_status": False, + "event_id": "", + "message": "Vision event sent in background task", + } + elif fire_and_forget and self._thread_pool_executor: + self._thread_pool_executor.submit(task) + return { + "error_status": False, + "throttling_status": False, + "event_id": "", + "message": "Vision event sent in background task", + } + else: + error_status, message, event_id = task() + return { + "error_status": error_status, + "throttling_status": False, + "event_id": event_id, + "message": message, + } + + +def _build_event_data( + event_type: str, + external_id: Optional[str] = None, + qc_result: Optional[str] = None, + location: Optional[str] = None, + item_count: Optional[int] = None, + item_type: Optional[str] = None, + alert_type: Optional[str] = None, + severity: Optional[str] = None, + alert_description: Optional[str] = None, + custom_value: Optional[str] = None, + related_event_id: Optional[str] = None, + feedback: Optional[str] = None, +) -> Dict[str, Any]: + """Build schema-specific eventData dict with camelCase keys, stripping None values.""" + if event_type == "quality_check": + data = {"result": qc_result, "externalId": external_id} + elif event_type == "inventory_count": + data = { + "location": location, + "itemCount": item_count, + "itemType": item_type, + "externalId": external_id, + } + elif event_type == "safety_alert": + data = { + "alertType": alert_type, + "severity": severity, + "description": alert_description, + "externalId": external_id, + } + elif event_type == "custom": + data = {"value": custom_value, "externalId": external_id} + elif event_type == "operator_feedback": + data = { + "relatedEventId": related_event_id, + "feedback": feedback, + } + else: + logger.warning("Unknown event_type: %s", event_type) + data = {} + return {k: v for k, v in data.items() if v is not None} + + +def _convert_predictions_to_annotations( + prediction: Optional[TensorNativePrediction], +) -> Dict[str, List[dict]]: + """Convert predictions into vision events annotation lists. + + Returns a dict keyed by annotation type (objectDetections, classifications, + instanceSegmentations, keypoints), containing only the non-empty annotation + lists. Shared by the cloud and local event store code paths. + """ + object_detections: List[dict] = [] + classifications: List[dict] = [] + instance_segmentations: List[dict] = [] + keypoints_detections: List[dict] = [] + + if prediction is not None: + if isinstance( + prediction, + (ClassificationPrediction, MultiLabelClassificationPrediction), + ): + classifications = _convert_classification_to_vision_events_format( + prediction + ) + else: + # Detection-style prediction (object detection, instance segmentation, + # or the keypoint `(KeyPoints, Detections)` tuple). + ( + object_detections, + instance_segmentations, + keypoints_detections, + ) = _convert_native_detections_to_vision_events_format(prediction) + + annotations: Dict[str, List[dict]] = {} + if object_detections: + annotations["objectDetections"] = object_detections + if classifications: + annotations["classifications"] = classifications + if instance_segmentations: + annotations["instanceSegmentations"] = instance_segmentations + if keypoints_detections: + annotations["keypoints"] = keypoints_detections + return annotations + + +def _execute_vision_event( + api_base_url: str, + api_key: str, + input_image: Optional[WorkflowImageData], + output_image: Optional[WorkflowImageData], + prediction: Optional[TensorNativePrediction], + event_type: str, + solution: str, + event_data: Dict[str, Any], + custom_metadata: Dict[str, Any], +) -> Tuple[bool, str, str]: + try: + # Step 1: Convert predictions to vision events annotation format + annotations = _convert_predictions_to_annotations(prediction) + + # Step 2: Upload images and build a single image entry + # sourceId = output/display image, inputSourceId = original input image + image_entry: Dict[str, Any] = {} + + if output_image is not None: + output_source_id, _ = _upload_image(api_base_url, api_key, output_image) + image_entry["sourceId"] = output_source_id + + if input_image is not None: + input_source_id, _ = _upload_image( + api_base_url, api_key, input_image, jpeg_quality=95 + ) + image_entry["inputSourceId"] = input_source_id + + if image_entry: + image_entry["label"] = "workflow" + image_entry.update(annotations) + + images_payload: List[dict] = [image_entry] if image_entry else [] + + # Step 3: Build and send event + payload = _build_event_payload( + event_type=event_type, + solution=solution, + images=images_payload, + event_data=event_data, + custom_metadata=custom_metadata, + ) + + error_status, message = _send_event(api_base_url, api_key, payload) + # The eventId is generated client-side and sent in the payload, so it is the + # canonical id of the created event; surface it on success. + event_id = "" if error_status else payload.get("eventId", "") + return error_status, message, event_id + except Exception as error: + logger.warning("Failed to create vision event: %s", error) + return ( + True, + f"Error creating vision event: {type(error).__name__}: {error}", + "", + ) + + +def _execute_local_event( + event_store_url: str, + input_image: Optional[WorkflowImageData], + output_image: Optional[WorkflowImageData], + prediction: Optional[TensorNativePrediction], + event_type: str, + solution: str, + event_data: Dict[str, Any], + custom_metadata: Dict[str, Any], +) -> Tuple[bool, str, str]: + """Send an event to a local Event Ingestion Service (v2 API). + + Unlike the cloud path, images are embedded directly in the request as base64 + rather than uploaded first. The use case (`solution`) is forwarded so events are + namespaced consistently with the cloud path; the service requires it when cloud + upload is enabled. + """ + try: + # Convert predictions to vision events annotation format + annotations = _convert_predictions_to_annotations(prediction) + + # Build a single image entry with base64-embedded images + # base64Image = output/display image, inputBase64Image = original input image + image_entry: Dict[str, Any] = {} + if output_image is not None: + image_entry["base64Image"] = output_image.base64_image + if input_image is not None: + image_entry["inputBase64Image"] = input_image.base64_image + if image_entry: + image_entry["label"] = "workflow" + image_entry.update(annotations) + + images_payload: List[dict] = [image_entry] if image_entry else [] + + payload: Dict[str, Any] = { + "inference_timestamp": datetime.now(timezone.utc).isoformat(), + "solution": solution, + "event_schema": event_type, + "event_data": event_data, + "images": images_payload, + } + if custom_metadata: + payload["custom_metadata"] = custom_metadata + if images_payload: + payload["displayImagePosition"] = 0 + + url = f"{event_store_url.rstrip('/')}/v2/events" + return _send_local_event(url, payload) + except Exception as error: + logger.warning("Failed to write event to local event store: %s", error) + return ( + True, + f"Error writing event to event store: {type(error).__name__}: {error}", + "", + ) + + +def _send_local_event( + url: str, + payload: dict, +) -> Tuple[bool, str, str]: + """Send an event to the local Event Ingestion Service. + + Authenticates with the ``EVENT_INGESTION_API_KEY`` environment variable when set. + Mirrors the Event Writer sink's response handling for the shared ``/v2/events`` + endpoint: the service returns 201 with the server-assigned event ``id`` on success, + and 529 when the device is at capacity and applying backpressure while it waits for + cloud uploads to drain. + + Returns: + Tuple of (error_status, message, event_id) + """ + try: + headers = {"Content-Type": "application/json"} + api_key = os.environ.get("EVENT_INGESTION_API_KEY") + if api_key: + headers["X-API-Key"] = api_key + response = requests.post(url, headers=headers, json=payload, timeout=30) + + if response.status_code == 201: + event_id = str(response.json().get("id", "")) + return False, "Event written to local event store successfully", event_id + + if response.status_code == 529: + detail = _extract_detail(response) + return ( + True, + "Event Ingestion Service at capacity (529). The device is " + f"experiencing backpressure. Detail: {detail}", + "", + ) + + detail = _extract_detail(response) + logger.warning( + "Event Ingestion Service error (%s): %s", response.status_code, detail + ) + return ( + True, + f"Failed to write event to event store. HTTP {response.status_code}: {detail}", + "", + ) + except requests.exceptions.Timeout: + return True, "Request to event store timed out after 30s", "" + except Exception as e: + logger.warning("Failed to write event to local event store: %s", e) + return ( + True, + f"Failed to write event to event store. Error: {type(e).__name__}: {e}", + "", + ) + + +def _extract_detail(response: requests.Response) -> str: + """Extract the ``detail`` field from a JSON error response, falling back to text.""" + try: + return str(response.json().get("detail", response.text)) + except Exception: + return response.text + + +def _detect_prediction_type( + prediction: Union[Detections, InstanceDetections, KeyPointPrediction], +) -> str: + """Determine detection type from the tensor-native prediction shape. + + Keypoint predictions arrive as a `(KeyPoints, Detections)` tuple, instance + segmentation as `inference_models.InstanceDetections`, and plain object + detection as `inference_models.Detections`. + """ + if isinstance(prediction, tuple): + return "keypoint_detection" + if isinstance(prediction, InstanceDetections): + return "instance_segmentation" + return "object_detection" + + +def _convert_native_detections_to_vision_events_format( + prediction: Union[Detections, InstanceDetections, KeyPointPrediction], +) -> Tuple[List[dict], List[dict], List[dict]]: + """Convert tensor-native detections to vision events format. + + Returns: + Tuple of (object_detections, instance_segmentations, keypoints) + using center-based bounding boxes in absolute pixel coordinates. + """ + detection_type = _detect_prediction_type(prediction) + + # Keypoint predictions arrive as a `(KeyPoints, Detections)` tuple; the vision + # events payload is driven by the bbox component, whose `bboxes_metadata` + # carries the per-detection keypoint arrays. + _key_points, detections = split_key_point_prediction(prediction) + + object_detections: List[dict] = [] + instance_segmentations: List[dict] = [] + keypoints_list: List[dict] = [] + + # class_name of detection i is resolved through the per-image class_id -> name + # map (numpy stored it per-detection in `sv.Detections.data["class_name"]`). + class_names_map = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + + for index, ( + xyxy, + _mask, + class_id, + confidence, + _tracker_id, + data, + _meta, + ) in enumerate(detections): + xyxy = xyxy.detach().to("cpu").numpy().astype(float).tolist() + x1, y1, x2, y2 = xyxy + w = abs(x2 - x1) + h = abs(y2 - y1) + cx = x1 + w / 2 + cy = y1 + h / 2 + class_id_int = int(class_id) + class_name = str(class_names_map.get(class_id_int, "unknown")) + conf = float(confidence) if confidence is not None else 0.0 + # Clamp confidence to [0, 1] + conf = max(0.0, min(1.0, conf)) + + base = { + "class": class_name, + "x": cx, + "y": cy, + "width": w, + "height": h, + "confidence": conf, + } + + if detection_type == "keypoint_detection": + kp_entry = dict(base) + kp_xy = data.get(KEYPOINTS_XY_KEY_IN_SV_DETECTIONS) + kp_class_id = data.get(KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS) + kp_class_name = data.get(KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS) + if kp_xy is not None and len(kp_xy) > 0: + # Skip trailing padding slots (see common/keypoints.py). + kp_count = real_keypoints_count(kp_class_name, total=len(kp_xy)) + kp_entry["keypoints"] = [] + for i, (kx, ky) in enumerate(kp_xy[:kp_count]): + kp_id = ( + int(kp_class_id[i]) + if kp_class_id is not None and i < len(kp_class_id) + else i + ) + kp_entry["keypoints"].append( + {"id": kp_id, "x": float(kx), "y": float(ky)} + ) + keypoints_list.append(kp_entry) + + elif detection_type == "instance_segmentation": + seg_entry = dict(base) + polygon = data.get(POLYGON_KEY_IN_SV_DETECTIONS) + if polygon is None and detections.mask is not None: + polygon = mask_to_polygon( + mask=instance_mask_to_numpy(detections, index) + ) + if polygon is not None and len(polygon) >= 3: + if isinstance(polygon, np.ndarray): + polygon = polygon.astype(float).tolist() + seg_entry["points"] = [[float(pt[0]), float(pt[1])] for pt in polygon] + instance_segmentations.append(seg_entry) + else: + # Fall back to object detection if polygon is invalid + object_detections.append(base) + + else: + object_detections.append(base) + + return object_detections, instance_segmentations, keypoints_list + + +def _convert_classification_to_vision_events_format( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> List[dict]: + """Convert tensor-native classification predictions to vision events classification format. + + Handles both single-label (`ClassificationPrediction`, top-1 class) and + multi-label (`MultiLabelClassificationPrediction`, all above-threshold classes) + predictions. Class names are resolved through the per-image `class_names` map + carried in the prediction's image metadata; the `confidence` tensor is the full + distribution, so each emitted class is paired with its own per-class confidence. + """ + classifications: List[dict] = [] + + if isinstance(prediction, ClassificationPrediction): + images_metadata = prediction.images_metadata or [{}] + image_metadata = images_metadata[0] if images_metadata else {} + class_names_map = image_metadata.get(CLASS_NAMES_KEY) or {} + class_id = int(prediction.class_id[0]) + class_name = class_names_map.get(class_id, f"class_{class_id}") + conf = float(prediction.confidence[0, class_id]) + classifications.append({"class": str(class_name), "confidence": conf}) + return classifications + + if isinstance(prediction, MultiLabelClassificationPrediction): + image_metadata = prediction.image_metadata or {} + class_names_map = image_metadata.get(CLASS_NAMES_KEY) or {} + for class_id_scalar in prediction.class_ids.tolist(): + class_id = int(class_id_scalar) + class_name = class_names_map.get(class_id, f"class_{class_id}") + conf = float(prediction.confidence[class_id]) + classifications.append({"class": str(class_name), "confidence": conf}) + return classifications + + return classifications + + +def _upload_image( + api_base_url: str, + api_key: str, + image_data: WorkflowImageData, + jpeg_quality: int = 85, +) -> Tuple[str, str]: + """Upload an image to the Vision Events API. + + Returns: + Tuple of (sourceId, url) + """ + image_bytes = encode_image_to_jpeg_bytes( + image_data.numpy_image, jpeg_quality=jpeg_quality + ) + response = requests.post( + wrap_url(f"{api_base_url}/vision-events/upload"), + headers={"Authorization": f"Bearer {api_key}"}, + files={"file": ("image.jpg", image_bytes, "image/jpeg")}, + timeout=30, + ) + response.raise_for_status() + result = response.json() + return result["sourceId"], result.get("url", "") + + +def _build_event_payload( + event_type: str, + solution: str, + images: List[dict], + event_data: Dict[str, Any], + custom_metadata: Dict[str, Any], +) -> dict: + """Build the full event payload for the Vision Events API.""" + event_id = str(uuid4()) + timestamp = datetime.now(timezone.utc).isoformat() + + payload: Dict[str, Any] = { + "eventId": event_id, + "eventType": event_type, + "useCaseId": solution, + "timestamp": timestamp, + "images": images, + } + + if event_data: + payload["eventData"] = event_data + if custom_metadata: + payload["customMetadata"] = custom_metadata + # Output image is at position 0 (first in images array), always display it + if len(images) > 0: + payload["displayImagePosition"] = 0 + + return payload + + +def _send_event( + api_base_url: str, + api_key: str, + payload: dict, +) -> Tuple[bool, str]: + """Send a vision event to the API. + + Returns: + Tuple of (error_status, message) + """ + try: + response = requests.post( + wrap_url(f"{api_base_url}/vision-events"), + headers={ + "Authorization": f"Bearer {api_key}", + "Content-Type": "application/json", + }, + json=payload, + timeout=30, + ) + response.raise_for_status() + return False, "Vision event sent successfully" + except requests.exceptions.HTTPError as e: + status_code = e.response.status_code if e.response is not None else "unknown" + body = e.response.text if e.response is not None else "no response" + logger.warning("Vision Events API error (%s): %s", status_code, body) + return ( + True, + f"Failed to send vision event. Status: {status_code}. Details: {body}", + ) + except Exception as e: + logger.warning("Failed to send vision event: %s", e) + return True, f"Failed to send vision event. Error: {type(e).__name__}: {e}" diff --git a/inference/core/workflows/core_steps/trackers/_base_tensor.py b/inference/core/workflows/core_steps/trackers/_base_tensor.py new file mode 100644 index 0000000000..81572e4e8b --- /dev/null +++ b/inference/core/workflows/core_steps/trackers/_base_tensor.py @@ -0,0 +1,315 @@ +"""Shared base classes for tracker workflow blocks (tensor-native sibling). + +Tensor-native counterpart of ``_base.py``. Under +``ENABLE_TENSOR_DATA_REPRESENTATION`` the loader swaps each concrete tracker +block to its ``*_tensor.py`` sibling; those siblings import ``TrackerBlockBase`` +/ ``TRACKER_PREDICTION_KINDS`` / ``tracker_describe_outputs`` from here instead +of from ``_base.py``. + +The only representation-specific change versus ``_base.py`` is the native +input/output handling in ``_run_tracker``: detection predictions are native +``inference_models`` objects (``Detections`` / ``InstanceDetections`` / the +``(KeyPoints, Optional[Detections])`` keypoint tuple) rather than +``sv.Detections``. The third-party tracker libraries (``trackers`` package) are +``sv.Detections``-based, so ``_run_tracker`` materialises a minimal +``sv.Detections`` (bounding boxes only, with a stashed row index) as the +transport to/from the tracker, then maps the surviving rows back onto the +ORIGINAL native input โ€” preserving masks / keypoints / all native metadata โ€” +and writes the assigned ``tracker_id`` into ``bboxes_metadata``. Block outputs +are always native ``inference_models`` objects; ``sv.Detections`` is used only +as the algorithm boundary to the tracker library. + +Each concrete tracker block (ByteTrack, BoT-SORT, SORT, OC-SORT) inherits from +``TrackerBlockBase`` and implements ``_create_tracker`` and ``get_manifest``. +Sub-classes may override ``_tracker_update`` when the underlying tracker needs +extra per-frame context (e.g. a video frame for camera motion compensation). +``_tracker_update`` / ``_create_tracker`` stay ``sv.Detections``-based / +library-based and identical to ``_base.py`` (the third-party trackers are +``sv``-based) โ€” only ``_run_tracker`` does the nativeโ†”sv conversion. +""" + +from abc import abstractmethod +from collections import deque +from typing import Any, Dict, List, Optional, Tuple, Type, Union + +import numpy as np +import supervision as sv + +from inference.core import logger +from inference.core.workflows.core_steps.common.tensor_native import ( + split_key_point_prediction, + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "tracked_detections" + +#: Key under which the per-detection row index into the original native input is +#: stashed inside the transport ``sv.Detections.data`` dict. The third-party +#: tracker libraries index back into the input ``sv.Detections`` (preserving the +#: ``.data`` dict) so this row index travels through ``tracker.update`` and lets +#: us slice the ORIGINAL native input by the surviving rows afterwards. +_TRACKER_ROW_INDEX_KEY: str = "__tracker_row_index__" + +#: Detection kinds accepted as tracker input and declared on tracker output. +#: Trackers only use bounding boxes for association and preserve all other +#: fields (masks, keypoints, custom data) via native indexing back into the +#: original ``inference_models`` prediction. +TRACKER_PREDICTION_KINDS = [ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +] + + +class InstanceCache: + """FIFO cache that tracks which object track IDs have been seen before. + + Used to categorize tracked detections as new (first appearance) or + already seen (reappearance) across video frames. + """ + + def __init__(self, size: int): + size = max(1, size) + self._cache_inserts_track = deque(maxlen=size) + self._cache = set() + + def record_instance(self, tracker_id: int) -> bool: + """Record a tracker ID and return whether it was previously seen. + + Returns: + True if the tracker_id was already in the cache (seen before), + False if this is its first appearance. + """ + in_cache = tracker_id in self._cache + if not in_cache: + self._cache_new_tracker_id(tracker_id=tracker_id) + return in_cache + + def _cache_new_tracker_id(self, tracker_id: int) -> None: + while len(self._cache) >= self._cache_inserts_track.maxlen: + to_drop = self._cache_inserts_track.popleft() + self._cache.remove(to_drop) + self._cache_inserts_track.append(tracker_id) + self._cache.add(tracker_id) + + +# Native prediction shapes accepted by tracker blocks (object detection, +# instance segmentation, RLE instance segmentation, or the keypoint-detection +# tuple). Mirrors TRACKER_PREDICTION_KINDS. +TensorNativeTrackerPrediction = Union[ + Detections, + InstanceDetections, + Tuple[KeyPoints, Optional[Detections]], +] + + +class TrackerBlockBase(WorkflowBlock): + """Common run-loop shared by every tracker block. + + Sub-classes implement ``_create_tracker`` and ``get_manifest``. Override + ``_tracker_update`` only when the tracker API requires additional context + beyond ``sv.Detections`` (e.g. BoT-SORT with camera motion compensation). + """ + + def __init__(self) -> None: + self._trackers: Dict[str, Any] = {} + self._per_video_cache: Dict[str, InstanceCache] = {} + + @classmethod + @abstractmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: ... + + @abstractmethod + def _create_tracker(self, fps: int, **kwargs: Any) -> Any: + """Instantiate the concrete tracker with algorithm-specific params.""" + ... + + def _tracker_update( + self, + tracker: Any, + detections: sv.Detections, + image: WorkflowImageData, + ) -> sv.Detections: + """Invoke the tracker for one frame. + + Must call ``tracker.update`` only with arguments that library trackers + define for the per-frame step (typically detections, optionally a frame + tensor). Do **not** pass workflow/block kwargs used in ``_create_tracker``. + """ + return tracker.update(detections) + + def _run_tracker( + self, + image: WorkflowImageData, + detections: TensorNativeTrackerPrediction, + instances_cache_size: int, + **tracker_kwargs: Any, + ) -> BlockResult: + """Run one frame through the tracker. + + Note: tracker parameters (``tracker_kwargs``) are only used when the + tracker is **first created** for a given ``video_identifier``. + Changing parameter values on subsequent frames has no effect because + the tracker instance is cached for the lifetime of the video stream. + + Tensor-native note: ``detections`` is a native ``inference_models`` + prediction. The bounding-box-bearing component is materialised to a + minimal ``sv.Detections`` (with a stashed row index) as transport to the + ``sv``-based tracker library; the surviving rows are then mapped back + onto the ORIGINAL native input so masks / keypoints / native metadata are + preserved, and the assigned ``tracker_id`` is written into + ``bboxes_metadata``. All three outputs are native objects. + """ + metadata = image.video_metadata + fps = metadata.fps + if not fps: + fps = 30 + logger.warning( + f"fps not available in VideoMetadata for {self.__class__.__name__}, " + "defaulting to 30 fps for tracker initialisation" + ) + video_id = metadata.video_identifier + + if video_id not in self._trackers: + self._trackers[video_id] = self._create_tracker(fps=fps, **tracker_kwargs) + + tracker = self._trackers[video_id] + + # Materialise the bbox component to a minimal sv.Detections used purely as + # transport to/from the third-party (sv-based) tracker. For the keypoint + # tuple the bbox Detections (second element) drives association. + _, bbox = split_key_point_prediction(detections) + n = int(bbox.xyxy.shape[0]) + sv_input = sv.Detections( + xyxy=bbox.xyxy.detach().to("cpu").numpy().astype(float), + confidence=bbox.confidence.detach().to("cpu").numpy().astype(float), + class_id=bbox.class_id.detach().to("cpu").numpy().astype(int), + data={_TRACKER_ROW_INDEX_KEY: np.arange(n, dtype=np.int64)}, + ) + + tracked_sv = self._tracker_update(tracker, sv_input, image) + + # Filter out immature / unmatched tracks (tracker_id == -1). Mirror the + # numpy guard exactly: only filter when tracker_id is present and there + # is at least one tracked detection. + if tracked_sv.tracker_id is not None and len(tracked_sv) > 0: + valid_mask = tracked_sv.tracker_id != -1 + tracked_sv = tracked_sv[valid_mask] + + # Recover the surviving native-input row indices (stashed in .data and + # preserved by the tracker library's internal sv indexing) and the + # assigned tracker ids. The empty / library-emptied case yields no rows. + if ( + tracked_sv.data + and _TRACKER_ROW_INDEX_KEY in tracked_sv.data + and tracked_sv.tracker_id is not None + and len(tracked_sv) > 0 + ): + surviving = ( + np.asarray(tracked_sv.data[_TRACKER_ROW_INDEX_KEY]).astype(int).tolist() + ) + tracker_ids = [int(t) for t in tracked_sv.tracker_id.tolist()] + else: + surviving = [] + tracker_ids = [] + + # Slice the ORIGINAL native input by the surviving rows (handles + # Detections / InstanceDetections / the keypoint tuple, dense + RLE + # masks). This preserves every native field for the surviving rows. + tracked_detections = take_prediction_by_indices(detections, surviving) + + # Write the assigned tracker_id into the bbox component's + # bboxes_metadata (copy dicts so the caller's state is not mutated). + _patch_tracker_ids(tracked_detections, tracker_ids) + + if video_id not in self._per_video_cache: + self._per_video_cache[video_id] = InstanceCache(size=instances_cache_size) + cache = self._per_video_cache[video_id] + + not_seen_indices, seen_indices = [], [] + for position, tracker_id in enumerate(tracker_ids): + already_seen = cache.record_instance(tracker_id=tracker_id) + if already_seen: + seen_indices.append(position) + else: + not_seen_indices.append(position) + + return { + OUTPUT_KEY: tracked_detections, + "new_instances": take_prediction_by_indices( + tracked_detections, not_seen_indices + ), + "already_seen_instances": take_prediction_by_indices( + tracked_detections, seen_indices + ), + } + + +def _bbox_component( + prediction: TensorNativeTrackerPrediction, +) -> Union[Detections, InstanceDetections]: + """Return the bounding-box-bearing component of a tracker prediction (the + Detections/InstanceDetections itself, or the bbox element of the keypoint + tuple).""" + _, bbox = split_key_point_prediction(prediction) + return bbox + + +def _patch_tracker_ids( + prediction: TensorNativeTrackerPrediction, + tracker_ids: List[int], +) -> None: + """Write ``tracker_id`` into the bbox component's ``bboxes_metadata``. + + Builds ``bboxes_metadata`` as a list of dicts when it is ``None``, and + copies any existing per-detection dict so caller-owned state is not mutated. + Mutates the bbox component of ``prediction`` in place (the native objects + produced by ``take_prediction_by_indices`` are freshly allocated for this + block, so this is safe). + """ + bbox = _bbox_component(prediction) + n = int(bbox.xyxy.shape[0]) + existing = bbox.bboxes_metadata + new_meta: List[dict] = [] + for i in range(n): + base = dict(existing[i]) if existing is not None and existing[i] else {} + base["tracker_id"] = int(tracker_ids[i]) + new_meta.append(base) + bbox.bboxes_metadata = new_meta if new_meta else None + + +def tracker_describe_outputs() -> List[OutputDefinition]: + """Output definitions shared by all tracker blocks. + + Trackers preserve all detection fields (masks, keypoints, custom data) โ€” + they only use bounding boxes for association then index back into the + original native prediction. The output kinds therefore mirror the input + kinds accepted by every tracker manifest. + """ + return [ + OutputDefinition(name=OUTPUT_KEY, kind=TRACKER_PREDICTION_KINDS), + OutputDefinition(name="new_instances", kind=TRACKER_PREDICTION_KINDS), + OutputDefinition( + name="already_seen_instances", + kind=TRACKER_PREDICTION_KINDS, + ), + ] diff --git a/inference/core/workflows/core_steps/trackers/botsort/v1_tensor.py b/inference/core/workflows/core_steps/trackers/botsort/v1_tensor.py new file mode 100644 index 0000000000..23a608c5fa --- /dev/null +++ b/inference/core/workflows/core_steps/trackers/botsort/v1_tensor.py @@ -0,0 +1,325 @@ +from typing import Any, List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field +from trackers import BoTSORTTracker + +from inference.core import logger +from inference.core.workflows.core_steps.trackers._base_tensor import ( + TRACKER_PREDICTION_KINDS, + TensorNativeTrackerPrediction, + TrackerBlockBase, + tracker_describe_outputs, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +#: Camera motion compensation (CMC) backend for BoT-SORT. Valid string values: +#: +#: - ``orb``: ORB keypoints/descriptors with RANSAC affine estimation. +#: - ``sift``: SIFT keypoints/descriptors with RANSAC affine estimation (typically slower than ORB). +#: - ``sparseOptFlow``: ``goodFeaturesToTrack`` + Lucas-Kanade optical flow with RANSAC affine estimation. +#: - ``ecc``: Enhanced Correlation Coefficient alignment on image intensities. +CMCMethod = Literal["orb", "sift", "sparseOptFlow", "ecc"] + +DEFAULT_LOST_TRACK_BUFFER = 30 +DEFAULT_TRACK_ACTIVATION_THRESHOLD = 0.7 +DEFAULT_MINIMUM_CONSECUTIVE_FRAMES = 2 +DEFAULT_MINIMUM_IOU_THRESHOLD_FIRST_ASSOC = 0.2 +DEFAULT_MINIMUM_IOU_THRESHOLD_SECOND_ASSOC = 0.5 +DEFAULT_MINIMUM_IOU_THRESHOLD_UNCONFIRMED_ASSOC = 0.3 +DEFAULT_HIGH_CONF_DET_THRESHOLD = 0.6 +DEFAULT_ENABLE_CMC = False +DEFAULT_CMC_METHOD: CMCMethod = "sparseOptFlow" +DEFAULT_CMC_DOWNSCALE = 2 +DEFAULT_INSTANT_FIRST_FRAME_ACTIVATION = False +DEFAULT_INSTANCES_CACHE_SIZE = 16384 + +SHORT_DESCRIPTION = ( + "ByteTrack-style association with optional camera motion compensation (BoT-SORT)." +) +LONG_DESCRIPTION = """ +Track objects across video frames using the **BoT-SORT** algorithm from the +roboflow/trackers package. + +BoT-SORT follows a ByteTrack-style association pipeline (high- and low-confidence +detections, Kalman track states) and can apply **camera motion compensation (CMC)** +before association when enabled. CMC estimates a global affine motion between +frames so predicted boxes align better when the camera moves. + +**When to use BoT-SORT:** +- Scenes with **moving or shaking cameras** (enable **Camera motion compensation**). +- Dense detection noise where ByteTrack-style two-stage matching helps. +- When you want ByteTrack-like behaviour with an optional motion-compensation stage. + +**When to consider alternatives:** +- Fixed camera and you only need speed: **ByteTrack** or **SORT** may be simpler. +- Heavy occlusion and erratic object motion without camera motion: **OC-SORT**. +- Low-texture backgrounds where sparse-feature CMC is unreliable. + +**Camera motion compensation:** When enabled, the block passes the workflow image +pixels to the tracker each frame. If the image cannot be decoded to a numpy array, +the tracker runs without CMC for that frame (a warning is logged). + +**Instant first-frame activation** defaults to off so behaviour aligns with other +core tracker blocks for ``new_instances`` / ``already_seen_instances``. Enable it +if you want tracks on frame 1 to receive stable IDs immediately (original BoT-SORT +paper-style). + +Outputs three detection sets: +- **tracked_detections**: All confirmed tracked detections with assigned track IDs. +- **new_instances**: Detections whose track ID appears for the first time. +- **already_seen_instances**: Detections whose track ID has been seen in a prior frame. + +The block maintains separate tracker state and instance cache per `video_identifier`, +enabling multi-stream tracking within a single workflow. +""" + + +class BoTSORTManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "BoT-SORT Tracker", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "video", + "icon": "far fa-location-crosshairs", + "blockPriority": 0, + "trackers": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/trackers_botsort@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Input image with embedded video metadata (fps and video_identifier). " + "Used to initialise and retrieve per-video tracker state. When camera motion " + "compensation is enabled, frame pixels are read from this image.", + ) + detections: Selector( + kind=TRACKER_PREDICTION_KINDS, + ) = Field( + description="Detection predictions for the current frame to track.", + examples=["$steps.object_detection_model.predictions"], + ) + minimum_iou_threshold_first_assoc: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_MINIMUM_IOU_THRESHOLD_FIRST_ASSOC, + description="Minimum fused similarity (IoU ร— confidence) for the first " + f"(high-confidence) association step. Default: {DEFAULT_MINIMUM_IOU_THRESHOLD_FIRST_ASSOC}.", + examples=[ + DEFAULT_MINIMUM_IOU_THRESHOLD_FIRST_ASSOC, + "$inputs.minimum_iou_threshold_first_assoc", + ], + json_schema_extra={ + "always_visible": True, + }, + ) + minimum_iou_threshold_second_assoc: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_MINIMUM_IOU_THRESHOLD_SECOND_ASSOC, + description="Minimum IoU for the second (low-confidence) association step. " + f"Default: {DEFAULT_MINIMUM_IOU_THRESHOLD_SECOND_ASSOC}.", + examples=[ + DEFAULT_MINIMUM_IOU_THRESHOLD_SECOND_ASSOC, + "$inputs.minimum_iou_threshold_second_assoc", + ], + json_schema_extra={ + "always_visible": True, + }, + ) + minimum_iou_threshold_unconfirmed_assoc: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_MINIMUM_IOU_THRESHOLD_UNCONFIRMED_ASSOC, + description="Minimum fused similarity for matching unconfirmed tracks to " + f"remaining high-confidence detections. Default: {DEFAULT_MINIMUM_IOU_THRESHOLD_UNCONFIRMED_ASSOC}.", + examples=[ + DEFAULT_MINIMUM_IOU_THRESHOLD_UNCONFIRMED_ASSOC, + "$inputs.minimum_iou_threshold_unconfirmed_assoc", + ], + json_schema_extra={ + "always_visible": True, + }, + ) + minimum_consecutive_frames: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = ( + Field( + default=DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + description="Number of consecutive frames a track must be matched before it is " + f"emitted as a confirmed track (tracker_id != -1). Default: {DEFAULT_MINIMUM_CONSECUTIVE_FRAMES}.", + examples=[ + DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + "$inputs.minimum_consecutive_frames", + ], + json_schema_extra={ + "always_visible": True, + }, + ) + ) + lost_track_buffer: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( + default=DEFAULT_LOST_TRACK_BUFFER, + description="Number of frames to keep a track alive after it loses its matched " + f"detection. Higher values improve occlusion recovery. Default: {DEFAULT_LOST_TRACK_BUFFER}.", + examples=[DEFAULT_LOST_TRACK_BUFFER, "$inputs.lost_track_buffer"], + json_schema_extra={ + "always_visible": True, + }, + ) + track_activation_threshold: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_TRACK_ACTIVATION_THRESHOLD, + description="Minimum detection confidence required to spawn a new track. " + f"Detections below this threshold are not used to create new tracks. Default: {DEFAULT_TRACK_ACTIVATION_THRESHOLD}.", + examples=[ + DEFAULT_TRACK_ACTIVATION_THRESHOLD, + "$inputs.track_activation_threshold", + ], + ) + high_conf_det_threshold: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_HIGH_CONF_DET_THRESHOLD, + description="Confidence threshold for high-confidence detections used in " + f"association. Default: {DEFAULT_HIGH_CONF_DET_THRESHOLD}.", + examples=[DEFAULT_HIGH_CONF_DET_THRESHOLD, "$inputs.high_conf_det_threshold"], + ) + enable_cmc: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore[arg-type] + default=DEFAULT_ENABLE_CMC, + description="Enable camera motion compensation (uses per-frame image pixels). " + "Recommended for moving cameras.", + examples=[DEFAULT_ENABLE_CMC, "$inputs.enable_cmc"], + ) + cmc_method: CMCMethod = Field( + default=DEFAULT_CMC_METHOD, + description="Camera motion estimator. One of: orb, " + "sift, sparseOptFlow, ecc. Default: {DEFAULT_CMC_METHOD!r}.", + examples=[DEFAULT_CMC_METHOD, "$inputs.cmc_method"], + ) + cmc_downscale: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( + default=DEFAULT_CMC_DOWNSCALE, + description="Downscale factor applied inside CMC for speed and robustness. " + f"Default: {DEFAULT_CMC_DOWNSCALE}.", + examples=[DEFAULT_CMC_DOWNSCALE, "$inputs.cmc_downscale"], + ) + instant_first_frame_activation: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore[arg-type] + default=DEFAULT_INSTANT_FIRST_FRAME_ACTIVATION, + description="If true, tracks on the first frame receive IDs immediately (paper-style). " + "Default false so new/already-seen outputs match other core trackers.", + examples=[ + DEFAULT_INSTANT_FIRST_FRAME_ACTIVATION, + "$inputs.instant_first_frame_activation", + ], + ) + instances_cache_size: int = Field( + default=DEFAULT_INSTANCES_CACHE_SIZE, + description="Maximum number of track IDs retained in the instance cache for " + f"new/already-seen categorisation. Uses FIFO eviction. Default: {DEFAULT_INSTANCES_CACHE_SIZE}.", + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return tracker_describe_outputs() + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class BoTSORTBlockV1(TrackerBlockBase): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BoTSORTManifest + + def _create_tracker(self, fps: int, **kwargs: Any) -> Any: + return BoTSORTTracker( + lost_track_buffer=kwargs["lost_track_buffer"], + frame_rate=fps, + track_activation_threshold=kwargs["track_activation_threshold"], + minimum_consecutive_frames=kwargs["minimum_consecutive_frames"], + minimum_iou_threshold_first_assoc=kwargs[ + "minimum_iou_threshold_first_assoc" + ], + minimum_iou_threshold_second_assoc=kwargs[ + "minimum_iou_threshold_second_assoc" + ], + minimum_iou_threshold_unconfirmed_assoc=kwargs[ + "minimum_iou_threshold_unconfirmed_assoc" + ], + high_conf_det_threshold=kwargs["high_conf_det_threshold"], + enable_cmc=kwargs["enable_cmc"], + cmc_method=kwargs["cmc_method"], + cmc_downscale=kwargs["cmc_downscale"], + instant_first_frame_activation=kwargs["instant_first_frame_activation"], + ) + + def _tracker_update( + self, + tracker: Any, + detections: sv.Detections, + image: WorkflowImageData, + ) -> sv.Detections: + if not getattr(tracker, "enable_cmc", False): + return tracker.update(detections) + try: + frame = image.numpy_image + except Exception as exc: + logger.warning( + "%s: enable_cmc=True but frame unavailable (%s); running without CMC.", + self.__class__.__name__, + exc, + ) + return tracker.update(detections) + return tracker.update(detections, frame=frame) + + def run( + self, + image: WorkflowImageData, + detections: TensorNativeTrackerPrediction, + lost_track_buffer: int = DEFAULT_LOST_TRACK_BUFFER, + minimum_iou_threshold_first_assoc: float = DEFAULT_MINIMUM_IOU_THRESHOLD_FIRST_ASSOC, + minimum_iou_threshold_second_assoc: float = DEFAULT_MINIMUM_IOU_THRESHOLD_SECOND_ASSOC, + minimum_iou_threshold_unconfirmed_assoc: float = ( + DEFAULT_MINIMUM_IOU_THRESHOLD_UNCONFIRMED_ASSOC + ), + minimum_consecutive_frames: int = DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + instances_cache_size: int = DEFAULT_INSTANCES_CACHE_SIZE, + track_activation_threshold: float = DEFAULT_TRACK_ACTIVATION_THRESHOLD, + high_conf_det_threshold: float = DEFAULT_HIGH_CONF_DET_THRESHOLD, + enable_cmc: bool = DEFAULT_ENABLE_CMC, + cmc_method: CMCMethod = DEFAULT_CMC_METHOD, + cmc_downscale: int = DEFAULT_CMC_DOWNSCALE, + instant_first_frame_activation: bool = DEFAULT_INSTANT_FIRST_FRAME_ACTIVATION, + ) -> BlockResult: + return self._run_tracker( + image=image, + detections=detections, + instances_cache_size=instances_cache_size, + lost_track_buffer=lost_track_buffer, + minimum_iou_threshold_first_assoc=minimum_iou_threshold_first_assoc, + minimum_iou_threshold_second_assoc=minimum_iou_threshold_second_assoc, + minimum_iou_threshold_unconfirmed_assoc=minimum_iou_threshold_unconfirmed_assoc, + minimum_consecutive_frames=minimum_consecutive_frames, + track_activation_threshold=track_activation_threshold, + high_conf_det_threshold=high_conf_det_threshold, + enable_cmc=enable_cmc, + cmc_method=cmc_method, + cmc_downscale=cmc_downscale, + instant_first_frame_activation=instant_first_frame_activation, + ) diff --git a/inference/core/workflows/core_steps/trackers/bytetrack/v1_tensor.py b/inference/core/workflows/core_steps/trackers/bytetrack/v1_tensor.py new file mode 100644 index 0000000000..8f6e96ba8e --- /dev/null +++ b/inference/core/workflows/core_steps/trackers/bytetrack/v1_tensor.py @@ -0,0 +1,215 @@ +from typing import Any, List, Literal, Optional, Tuple, Type, Union + +from pydantic import ConfigDict, Field +from trackers import ByteTrackTracker + +from inference.core.workflows.core_steps.trackers._base_tensor import ( + TRACKER_PREDICTION_KINDS, + TrackerBlockBase, + tracker_describe_outputs, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +DEFAULT_LOST_TRACK_BUFFER = 30 +DEFAULT_TRACK_ACTIVATION_THRESHOLD = 0.7 +DEFAULT_MINIMUM_CONSECUTIVE_FRAMES = 2 +DEFAULT_MINIMUM_IOU_THRESHOLD = 0.1 +DEFAULT_HIGH_CONF_DET_THRESHOLD = 0.6 +DEFAULT_INSTANCES_CACHE_SIZE = 16384 + +SHORT_DESCRIPTION = "Tracks objects across frames. Best for most scenes." +LONG_DESCRIPTION = """ +Track objects across video frames using the **ByteTrack** algorithm from the +roboflow/trackers package. + +ByteTrack splits detections into high- and low-confidence pools and runs two +rounds of IoU-based association. The first round matches high-confidence +detections to existing tracks; the second recovers weak detections that overlap +unmatched tracks. This makes ByteTrack particularly effective in **dense +environments** where objects are frequently partially occluded and detector +confidence fluctuates. + +**When to use ByteTrack:** +- General-purpose tracking across diverse scenes. +- Dense or crowded environments with partial occlusions. +- Sports tracking and fast-moving objects (highest benchmark scores on SportsMOT). +- When your detector produces a mix of high- and low-confidence detections that + you want to retain. + +**When to consider alternatives:** +- For maximum simplicity and speed with a strong detector, use **SORT**. +- For scenes with heavy occlusion and non-linear motion, use **OC-SORT**. + +Outputs three detection sets: +- **tracked_detections**: All confirmed tracked detections with assigned track IDs. +- **new_instances**: Detections whose track ID appears for the first time. +- **already_seen_instances**: Detections whose track ID has been seen in a prior frame. + +The block maintains separate tracker state and instance cache per `video_identifier`, +enabling multi-stream tracking within a single workflow. +""" + + +class ByteTrackManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "ByteTrack Tracker", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "video", + "icon": "far fa-location-crosshairs", + "blockPriority": 0, + "trackers": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/trackers_bytetrack@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Input image with embedded video metadata (fps and video_identifier). " + "Used to initialise and retrieve per-video tracker state.", + ) + detections: Selector( + kind=TRACKER_PREDICTION_KINDS, + ) = Field( + description="Detection predictions for the current frame to track.", + examples=["$steps.object_detection_model.predictions"], + ) + minimum_iou_threshold: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_MINIMUM_IOU_THRESHOLD, + description="Minimum IoU required to associate a detection with an existing track. " + f"Default: {DEFAULT_MINIMUM_IOU_THRESHOLD}.", + examples=[DEFAULT_MINIMUM_IOU_THRESHOLD, "$inputs.minimum_iou_threshold"], + json_schema_extra={ + "always_visible": True, + }, + ) + minimum_consecutive_frames: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = ( + Field( + default=DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + description="Number of consecutive frames a track must be matched before it is " + f"emitted as a confirmed track (tracker_id != -1). Default: {DEFAULT_MINIMUM_CONSECUTIVE_FRAMES}.", + examples=[ + DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + "$inputs.minimum_consecutive_frames", + ], + json_schema_extra={ + "always_visible": True, + }, + ) + ) + lost_track_buffer: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( + default=DEFAULT_LOST_TRACK_BUFFER, + description="Number of frames to keep a track alive after it loses its matched " + f"detection. Higher values improve occlusion recovery. Default: {DEFAULT_LOST_TRACK_BUFFER}.", + examples=[DEFAULT_LOST_TRACK_BUFFER, "$inputs.lost_track_buffer"], + json_schema_extra={ + "always_visible": True, + }, + ) + track_activation_threshold: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_TRACK_ACTIVATION_THRESHOLD, + description="Minimum detection confidence required to spawn a new track. " + f"Detections below this threshold are not used to create new tracks. Default: {DEFAULT_TRACK_ACTIVATION_THRESHOLD}.", + examples=[ + DEFAULT_TRACK_ACTIVATION_THRESHOLD, + "$inputs.track_activation_threshold", + ], + ) + high_conf_det_threshold: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_HIGH_CONF_DET_THRESHOLD, + description="Confidence threshold for high-confidence detections used in " + f"association. Default: {DEFAULT_HIGH_CONF_DET_THRESHOLD}.", + examples=[DEFAULT_HIGH_CONF_DET_THRESHOLD, "$inputs.high_conf_det_threshold"], + ) + instances_cache_size: int = Field( + default=DEFAULT_INSTANCES_CACHE_SIZE, + description="Maximum number of track IDs retained in the instance cache for " + f"new/already-seen categorisation. Uses FIFO eviction. Default: {DEFAULT_INSTANCES_CACHE_SIZE}.", + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return tracker_describe_outputs() + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class ByteTrackBlockV1(TrackerBlockBase): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return ByteTrackManifest + + def _create_tracker(self, fps: int, **kwargs: Any) -> Any: + return ByteTrackTracker( + lost_track_buffer=kwargs["lost_track_buffer"], + frame_rate=fps, + track_activation_threshold=kwargs["track_activation_threshold"], + minimum_consecutive_frames=kwargs["minimum_consecutive_frames"], + minimum_iou_threshold=kwargs["minimum_iou_threshold"], + high_conf_det_threshold=kwargs["high_conf_det_threshold"], + ) + + def run( + self, + image: WorkflowImageData, + detections: Union[ + Detections, + InstanceDetections, + Tuple[KeyPoints, Optional[Detections]], + ], + lost_track_buffer: int = DEFAULT_LOST_TRACK_BUFFER, + minimum_iou_threshold: float = DEFAULT_MINIMUM_IOU_THRESHOLD, + minimum_consecutive_frames: int = DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + instances_cache_size: int = DEFAULT_INSTANCES_CACHE_SIZE, + track_activation_threshold: float = DEFAULT_TRACK_ACTIVATION_THRESHOLD, + high_conf_det_threshold: float = DEFAULT_HIGH_CONF_DET_THRESHOLD, + ) -> BlockResult: + return self._run_tracker( + image=image, + detections=detections, + instances_cache_size=instances_cache_size, + lost_track_buffer=lost_track_buffer, + minimum_iou_threshold=minimum_iou_threshold, + minimum_consecutive_frames=minimum_consecutive_frames, + track_activation_threshold=track_activation_threshold, + high_conf_det_threshold=high_conf_det_threshold, + ) diff --git a/inference/core/workflows/core_steps/trackers/ocsort/v1_tensor.py b/inference/core/workflows/core_steps/trackers/ocsort/v1_tensor.py new file mode 100644 index 0000000000..7d2869834b --- /dev/null +++ b/inference/core/workflows/core_steps/trackers/ocsort/v1_tensor.py @@ -0,0 +1,225 @@ +from typing import Any, List, Literal, Optional, Type, Union + +from pydantic import ConfigDict, Field +from trackers import OCSORTTracker + +from inference.core.workflows.core_steps.trackers._base_tensor import ( + TRACKER_PREDICTION_KINDS, + TensorNativeTrackerPrediction, + TrackerBlockBase, + tracker_describe_outputs, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlockManifest, +) + +DEFAULT_LOST_TRACK_BUFFER = 30 +DEFAULT_MINIMUM_CONSECUTIVE_FRAMES = 3 +DEFAULT_MINIMUM_IOU_THRESHOLD = 0.3 +DEFAULT_HIGH_CONF_DET_THRESHOLD = 0.6 +DEFAULT_DIRECTION_CONSISTENCY_WEIGHT = 0.2 +DEFAULT_DELTA_T = 3 +DEFAULT_INSTANCES_CACHE_SIZE = 16384 + +SHORT_DESCRIPTION = "Tracks objects through occlusion and unpredictable movement." +LONG_DESCRIPTION = """ +Track objects across video frames using the **OC-SORT** algorithm from the +roboflow/trackers package. + +OC-SORT extends SORT with two key mechanisms: + +1. **Observation-Centric Re-Update (OCR):** When a track reappears after + occlusion, OC-SORT retroactively corrects the Kalman filter using the real + observations before and after the gap, reducing accumulated drift. +2. **Observation-Centric Momentum (OCM):** A direction-consistency cost is + blended with IoU during association, penalising matches where the candidate + detection lies in a direction inconsistent with the track's recent motion. + +This makes OC-SORT significantly more robust than SORT in scenes with heavy +occlusion, erratic motion, and uniform appearance. + +**When to use OC-SORT:** +- Crowded scenes with frequent and prolonged occlusions (e.g. pedestrians, + warehouse workers). +- Non-linear or erratic motion patterns (e.g. dancing, sports with abrupt + direction changes). +- When identity consistency over long sequences is more important than raw speed. + +**When to consider alternatives:** +- For general-purpose tracking with mixed-confidence detections, try **ByteTrack**. +- For maximum simplicity and speed with a strong detector, try **SORT**. + +Outputs three detection sets: +- **tracked_detections**: All confirmed tracked detections with assigned track IDs. +- **new_instances**: Detections whose track ID appears for the first time. +- **already_seen_instances**: Detections whose track ID has been seen in a prior frame. + +The block maintains separate tracker state and instance cache per `video_identifier`, +enabling multi-stream tracking within a single workflow. +""" + + +class OCSORTManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "OC-SORT Tracker", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "video", + "icon": "far fa-location-crosshairs", + "blockPriority": 0, + "trackers": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/trackers_ocsort@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Input image with embedded video metadata (fps and video_identifier). " + "Used to initialise and retrieve per-video tracker state.", + ) + detections: Selector( + kind=TRACKER_PREDICTION_KINDS, + ) = Field( + description="Detection predictions for the current frame to track.", + examples=["$steps.object_detection_model.predictions"], + ) + minimum_iou_threshold: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_MINIMUM_IOU_THRESHOLD, + description="Minimum IoU required to associate a detection with an existing track. " + f"Default: {DEFAULT_MINIMUM_IOU_THRESHOLD}.", + examples=[DEFAULT_MINIMUM_IOU_THRESHOLD, "$inputs.minimum_iou_threshold"], + json_schema_extra={ + "always_visible": True, + }, + ) + minimum_consecutive_frames: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = ( + Field( + default=DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + description="Number of consecutive frames a track must be matched before it is " + f"emitted as a confirmed track (tracker_id != -1). Default: {DEFAULT_MINIMUM_CONSECUTIVE_FRAMES}.", + examples=[ + DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + "$inputs.minimum_consecutive_frames", + ], + json_schema_extra={ + "always_visible": True, + }, + ) + ) + lost_track_buffer: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( + default=DEFAULT_LOST_TRACK_BUFFER, + description="Number of frames to keep a track alive after it loses its matched " + f"detection. Higher values improve occlusion recovery. Default: {DEFAULT_LOST_TRACK_BUFFER}.", + examples=[DEFAULT_LOST_TRACK_BUFFER, "$inputs.lost_track_buffer"], + json_schema_extra={ + "always_visible": True, + }, + ) + high_conf_det_threshold: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_HIGH_CONF_DET_THRESHOLD, + description="Confidence threshold for high-confidence detections used in " + f"association. Default: {DEFAULT_HIGH_CONF_DET_THRESHOLD}.", + examples=[DEFAULT_HIGH_CONF_DET_THRESHOLD, "$inputs.high_conf_det_threshold"], + ) + direction_consistency_weight: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_DIRECTION_CONSISTENCY_WEIGHT, + description="Weight for the direction consistency term in the OC-SORT association " + "cost. Higher values prioritise alignment between historical motion direction and " + f"the direction to the candidate detection. Default: {DEFAULT_DIRECTION_CONSISTENCY_WEIGHT}.", + examples=[ + DEFAULT_DIRECTION_CONSISTENCY_WEIGHT, + "$inputs.direction_consistency_weight", + ], + ) + delta_t: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( + default=DEFAULT_DELTA_T, + description="Number of past frames used by OC-SORT to estimate per-track velocity " + f"for direction consistency momentum. Default: {DEFAULT_DELTA_T}.", + examples=[DEFAULT_DELTA_T, "$inputs.delta_t"], + ) + instances_cache_size: int = Field( + default=DEFAULT_INSTANCES_CACHE_SIZE, + description="Maximum number of track IDs retained in the instance cache for " + f"new/already-seen categorisation. Uses FIFO eviction. Default: {DEFAULT_INSTANCES_CACHE_SIZE}.", + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return tracker_describe_outputs() + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class OCSORTBlockV1(TrackerBlockBase): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return OCSORTManifest + + def _create_tracker(self, fps: int, **kwargs: Any) -> Any: + return OCSORTTracker( + lost_track_buffer=kwargs["lost_track_buffer"], + frame_rate=fps, + minimum_consecutive_frames=kwargs["minimum_consecutive_frames"], + minimum_iou_threshold=kwargs["minimum_iou_threshold"], + high_conf_det_threshold=kwargs["high_conf_det_threshold"], + direction_consistency_weight=kwargs["direction_consistency_weight"], + delta_t=kwargs["delta_t"], + ) + + def run( + self, + image: WorkflowImageData, + detections: TensorNativeTrackerPrediction, + lost_track_buffer: int = DEFAULT_LOST_TRACK_BUFFER, + minimum_iou_threshold: float = DEFAULT_MINIMUM_IOU_THRESHOLD, + minimum_consecutive_frames: int = DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + instances_cache_size: int = DEFAULT_INSTANCES_CACHE_SIZE, + high_conf_det_threshold: float = DEFAULT_HIGH_CONF_DET_THRESHOLD, + direction_consistency_weight: float = DEFAULT_DIRECTION_CONSISTENCY_WEIGHT, + delta_t: int = DEFAULT_DELTA_T, + ) -> BlockResult: + return self._run_tracker( + image=image, + detections=detections, + instances_cache_size=instances_cache_size, + lost_track_buffer=lost_track_buffer, + minimum_iou_threshold=minimum_iou_threshold, + minimum_consecutive_frames=minimum_consecutive_frames, + high_conf_det_threshold=high_conf_det_threshold, + direction_consistency_weight=direction_consistency_weight, + delta_t=delta_t, + ) diff --git a/inference/core/workflows/core_steps/trackers/sort/v1_tensor.py b/inference/core/workflows/core_steps/trackers/sort/v1_tensor.py new file mode 100644 index 0000000000..77e9299dde --- /dev/null +++ b/inference/core/workflows/core_steps/trackers/sort/v1_tensor.py @@ -0,0 +1,202 @@ +from typing import Any, List, Literal, Optional, Tuple, Type, Union + +from pydantic import ConfigDict, Field +from trackers import SORTTracker + +from inference.core.workflows.core_steps.trackers._base_tensor import ( + TRACKER_PREDICTION_KINDS, + TrackerBlockBase, + tracker_describe_outputs, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +DEFAULT_LOST_TRACK_BUFFER = 30 +DEFAULT_TRACK_ACTIVATION_THRESHOLD = 0.25 +DEFAULT_MINIMUM_CONSECUTIVE_FRAMES = 3 +DEFAULT_MINIMUM_IOU_THRESHOLD = 0.3 +DEFAULT_INSTANCES_CACHE_SIZE = 16384 + +SHORT_DESCRIPTION = ( + "Fast, lightweight object tracking. Works best when objects are clearly visible." +) +LONG_DESCRIPTION = """ +Track objects across video frames using the **SORT** algorithm from the +roboflow/trackers package. + +SORT pairs a Kalman filter motion model with single-stage IoU-based Hungarian +assignment. It has the fewest parameters and lowest overhead, processing +hundreds of frames per second. However, it lacks re-identification and +occlusion-recovery mechanisms, so tracks may fragment or switch IDs when objects +are temporarily hidden. + +**When to use SORT:** +- Controlled environments with reliable, high-confidence detections. +- Real-time pipelines where maximum throughput is critical. +- Simple scenes with minimal occlusion and predictable linear motion. + +**When to consider alternatives:** +- If you see fragmented tracks or missed weak detections, try **ByteTrack**. +- If objects undergo heavy occlusion or non-linear motion, try **OC-SORT**. + +Outputs three detection sets: +- **tracked_detections**: All confirmed tracked detections with assigned track IDs. +- **new_instances**: Detections whose track ID appears for the first time. +- **already_seen_instances**: Detections whose track ID has been seen in a prior frame. + +The block maintains separate tracker state and instance cache per `video_identifier`, +enabling multi-stream tracking within a single workflow. +""" + + +class SORTManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "SORT Tracker", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "video", + "icon": "far fa-location-crosshairs", + "blockPriority": 0, + "trackers": True, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/trackers_sort@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Input image with embedded video metadata (fps and video_identifier). " + "Used to initialise and retrieve per-video tracker state.", + ) + detections: Selector( + kind=TRACKER_PREDICTION_KINDS, + ) = Field( + description="Detection predictions for the current frame to track.", + examples=["$steps.object_detection_model.predictions"], + ) + minimum_iou_threshold: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_MINIMUM_IOU_THRESHOLD, + description="Minimum IoU required to associate a detection with an existing track. " + f"Default: {DEFAULT_MINIMUM_IOU_THRESHOLD}.", + examples=[DEFAULT_MINIMUM_IOU_THRESHOLD, "$inputs.minimum_iou_threshold"], + json_schema_extra={ + "always_visible": True, + }, + ) + minimum_consecutive_frames: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = ( + Field( + default=DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + description="Number of consecutive frames a track must be matched before it is " + f"emitted as a confirmed track (tracker_id != -1). Default: {DEFAULT_MINIMUM_CONSECUTIVE_FRAMES}.", + examples=[ + DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + "$inputs.minimum_consecutive_frames", + ], + json_schema_extra={ + "always_visible": True, + }, + ) + ) + lost_track_buffer: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( + default=DEFAULT_LOST_TRACK_BUFFER, + description="Number of frames to keep a track alive after it loses its matched " + f"detection. Higher values improve occlusion recovery. Default: {DEFAULT_LOST_TRACK_BUFFER}.", + examples=[DEFAULT_LOST_TRACK_BUFFER, "$inputs.lost_track_buffer"], + json_schema_extra={ + "always_visible": True, + }, + ) + track_activation_threshold: Union[ + Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]) + ] = Field( + default=DEFAULT_TRACK_ACTIVATION_THRESHOLD, + description="Minimum detection confidence required to spawn a new track. " + f"Detections below this threshold are not used to create new tracks. Default: {DEFAULT_TRACK_ACTIVATION_THRESHOLD}.", + examples=[ + DEFAULT_TRACK_ACTIVATION_THRESHOLD, + "$inputs.track_activation_threshold", + ], + ) + instances_cache_size: int = Field( + default=DEFAULT_INSTANCES_CACHE_SIZE, + description="Maximum number of track IDs retained in the instance cache for " + f"new/already-seen categorisation. Uses FIFO eviction. Default: {DEFAULT_INSTANCES_CACHE_SIZE}.", + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return tracker_describe_outputs() + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class SORTBlockV1(TrackerBlockBase): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return SORTManifest + + def _create_tracker(self, fps: int, **kwargs: Any) -> Any: + return SORTTracker( + lost_track_buffer=kwargs["lost_track_buffer"], + frame_rate=fps, + track_activation_threshold=kwargs["track_activation_threshold"], + minimum_consecutive_frames=kwargs["minimum_consecutive_frames"], + minimum_iou_threshold=kwargs["minimum_iou_threshold"], + ) + + def run( + self, + image: WorkflowImageData, + detections: Union[ + Detections, + InstanceDetections, + Tuple[KeyPoints, Optional[Detections]], + ], + lost_track_buffer: int = DEFAULT_LOST_TRACK_BUFFER, + minimum_iou_threshold: float = DEFAULT_MINIMUM_IOU_THRESHOLD, + minimum_consecutive_frames: int = DEFAULT_MINIMUM_CONSECUTIVE_FRAMES, + instances_cache_size: int = DEFAULT_INSTANCES_CACHE_SIZE, + track_activation_threshold: float = DEFAULT_TRACK_ACTIVATION_THRESHOLD, + ) -> BlockResult: + return self._run_tracker( + image=image, + detections=detections, + instances_cache_size=instances_cache_size, + lost_track_buffer=lost_track_buffer, + minimum_iou_threshold=minimum_iou_threshold, + minimum_consecutive_frames=minimum_consecutive_frames, + track_activation_threshold=track_activation_threshold, + ) diff --git a/inference/core/workflows/core_steps/transformations/absolute_static_crop/v1_tensor.py b/inference/core/workflows/core_steps/transformations/absolute_static_crop/v1_tensor.py new file mode 100644 index 0000000000..1528b089b0 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/absolute_static_crop/v1_tensor.py @@ -0,0 +1,187 @@ +from typing import List, Literal, Optional, Type, Union +from uuid import uuid4 + +from pydantic import ConfigDict, Field, PositiveInt + +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + INTEGER_KIND, + ImageInputField, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) + +LONG_DESCRIPTION = """ +Extract a fixed rectangular region from input images using absolute pixel coordinates specified by center point and dimensions, creating consistent crops from the same image location across all inputs for region-of-interest extraction and fixed-area analysis workflows. + +## How This Block Works + +This block crops a fixed rectangular region from input images using absolute pixel coordinates, unlike dynamic cropping which uses detection bounding boxes. The block: + +1. Receives input images and absolute coordinate specifications (x_center, y_center, width, height) +2. Calculates the crop boundaries from the center point and dimensions: + - Computes x_min and y_min by subtracting half the width/height from the center coordinates + - Computes x_max and y_max by adding the width/height to the minimum coordinates + - Rounds coordinate values to integer pixel positions +3. Extracts the rectangular region from the image using array slicing (from y_min to y_max, x_min to x_max) +4. Validates that the cropped region has content (returns None if the crop would be empty, such as when coordinates are outside image bounds) +5. Creates a cropped image object with metadata tracking the crop's origin (original image, offset coordinates, unique crop identifier) +6. Preserves video metadata if the input is from video (maintains frame information and temporal context) +7. Returns the cropped image for each input image + +The block uses fixed coordinates, so the same region is extracted from all images in a batch, making it suitable for extracting consistent regions across multiple images (e.g., always cropping the top-right corner, extracting a fixed area of interest, or focusing on a specific image section). The center-based coordinate system allows specifying crops by their center point rather than corner coordinates, which can be more intuitive for defining regions. The block handles edge cases gracefully by returning None for invalid crops (coordinates outside image bounds or resulting in empty regions). + +## Common Use Cases + +- **Fixed Region Extraction**: Extract the same image region from multiple images for consistent analysis (e.g., crop a specific area of interest like a logo zone, extract a fixed region for watermark detection, crop a consistent area for pattern matching), enabling standardized region analysis across image batches +- **Region-of-Interest Focus**: Isolate specific areas of images for detailed processing (e.g., crop a specific quadrant of surveillance frames, extract a fixed region for text recognition, focus on a known area of interest), enabling focused analysis of predetermined image regions +- **Multi-Stage Workflow Preparation**: Extract fixed regions for secondary processing steps (e.g., crop a specific area from full images, then run OCR or classification on the cropped region), enabling hierarchical workflows with fixed region focus +- **Standardized Crop Generation**: Create consistent crops from images for training or analysis (e.g., extract a fixed region from all images for dataset creation, crop a standard area for comparison, generate uniform crops for feature extraction), enabling standardized data preparation workflows +- **Fixed-Area Monitoring**: Monitor specific image regions across time or batches (e.g., crop the same area from video frames for change detection, extract a fixed region for tracking analysis, focus on a consistent monitoring zone), enabling temporal analysis of fixed regions +- **Pre-Processing for Specialized Blocks**: Extract fixed regions before processing with specialized models (e.g., crop a specific area before running OCR, extract a fixed region for fine-grained classification, isolate a region for specialized analysis), enabling optimized processing of known image regions + +## Connecting to Other Blocks + +This block receives images and produces cropped images from fixed regions: + +- **After image loading blocks** to extract a fixed region of interest before processing, enabling focused analysis of predetermined image areas without processing entire images +- **Before classification or analysis blocks** that need region-focused inputs (e.g., OCR for text in a fixed area, fine-grained classification for cropped regions, specialized models for specific image areas), enabling optimized processing of consistent regions +- **In video processing workflows** to extract the same region from multiple frames (e.g., crop a fixed area from each video frame for temporal analysis, extract a consistent monitoring zone for tracking, focus on a specific region across frames), enabling temporal analysis of fixed regions +- **After detection blocks** where you know the approximate location and want to extract a fixed-size region around it (e.g., detect objects in a general area, then crop a fixed region around that area for detailed analysis), enabling region-focused multi-stage workflows +- **Before visualization blocks** that display specific regions (e.g., display only the cropped region, visualize a fixed area of interest, show isolated region annotations), enabling focused visualization of extracted regions +- **In batch processing workflows** where the same region needs to be extracted from all images for consistent analysis or comparison, enabling standardized region extraction across image sets +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Absolute Static Crop", + "version": "v1", + "short_description": "Crop an image using fixed pixel coordinates.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "transformation", + "icon": "far fa-crop-alt", + "blockPriority": 1, + }, + } + ) + type: Literal["roboflow_core/absolute_static_crop@v1", "AbsoluteStaticCrop"] + images: Selector(kind=[IMAGE_KIND]) = ImageInputField + x_center: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + description="X coordinate of the center point of the crop region in absolute pixel coordinates. Must be a positive integer. The crop region is centered at this X coordinate. The actual crop boundaries are calculated as x_min = x_center - width/2 and x_max = x_min + width. If the calculated crop extends beyond image bounds, the crop will be clipped or may return None if the crop would be empty.", + examples=[40, "$inputs.center_x"], + ) + y_center: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + description="Y coordinate of the center point of the crop region in absolute pixel coordinates. Must be a positive integer. The crop region is centered at this Y coordinate. The actual crop boundaries are calculated as y_min = y_center - height/2 and y_max = y_min + height. If the calculated crop extends beyond image bounds, the crop will be clipped or may return None if the crop would be empty.", + examples=[40, "$inputs.center_y"], + ) + width: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + description="Width of the crop region in pixels. Must be a positive integer. Defines the horizontal extent of the crop. The crop extends width/2 pixels to the left and right of the x_center coordinate. Total crop width equals this value. If the calculated crop extends beyond the image's width, it will be clipped to image boundaries.", + examples=[40, "$inputs.width"], + ) + height: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + description="Height of the crop region in pixels. Must be a positive integer. Defines the vertical extent of the crop. The crop extends height/2 pixels above and below the y_center coordinate. Total crop height equals this value. If the calculated crop extends beyond the image's height, it will be clipped to image boundaries.", + examples=[40, "$inputs.height"], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="crops", kind=[IMAGE_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class AbsoluteStaticCropBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + x_center: int, + y_center: int, + width: int, + height: int, + ) -> BlockResult: + return [ + { + "crops": take_static_crop( + image=image, + x_center=x_center, + y_center=y_center, + width=width, + height=height, + ) + } + for image in images + ] + + +def take_static_crop( + image: WorkflowImageData, + x_center: int, + y_center: int, + width: int, + height: int, +) -> Optional[WorkflowImageData]: + x_min = round(x_center - width / 2) + y_min = round(y_center - height / 2) + x_max = round(x_min + width) + y_max = round(y_min + height) + # Clamp to image bounds before slicing โ€” negative indices would otherwise + # wrap on both the torch and numpy slicing paths. + image_height, image_width = image._read_shape_without_materialization() + x_min = max(0, x_min) + y_min = max(0, y_min) + x_max = min(image_width, x_max) + y_max = min(image_height, y_max) + if x_max <= x_min or y_max <= y_min: + return None + if image.is_tensor_materialised(): + cropped_tensor_image = image.tensor_image[:, y_min:y_max, x_min:x_max] + if cropped_tensor_image.numel() == 0: + return None + return WorkflowImageData.create_crop_from_tensor( + origin_image_data=image, + crop_identifier=f"absolute_static_crop.{uuid4()}", + cropped_tensor_image=cropped_tensor_image.contiguous(), + offset_x=x_min, + offset_y=y_min, + preserve_video_metadata=True, + ) + # Only numpy is materialised โ€” slice on the host and emit a numpy-backed + # crop; no full-image numpy->device conversion is forced. + cropped_image = image.numpy_image[y_min:y_max, x_min:x_max] + if cropped_image.size == 0: + return None + return WorkflowImageData.create_crop( + origin_image_data=image, + crop_identifier=f"absolute_static_crop.{uuid4()}", + cropped_image=cropped_image, + offset_x=x_min, + offset_y=y_min, + preserve_video_metadata=True, + ) diff --git a/inference/core/workflows/core_steps/transformations/bounding_rect/v1_tensor.py b/inference/core/workflows/core_steps/transformations/bounding_rect/v1_tensor.py new file mode 100644 index 0000000000..147b1fe222 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/bounding_rect/v1_tensor.py @@ -0,0 +1,243 @@ +from typing import List, Literal, Optional, Tuple, Type, Union + +import cv2 as cv +import numpy as np +import pycocotools.mask as mask_utils +import supervision as sv +import torch +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + instance_mask_to_numpy, +) +from inference.core.workflows.execution_engine.constants import ( + BOUNDING_RECT_ANGLE_KEY_IN_SV_DETECTIONS, + BOUNDING_RECT_HEIGHT_KEY_IN_SV_DETECTIONS, + BOUNDING_RECT_RECT_KEY_IN_SV_DETECTIONS, + BOUNDING_RECT_WIDTH_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import Selector +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks + +OUTPUT_KEY: str = "detections_with_rect" + +SHORT_DESCRIPTION = "Find the minimal bounding box surrounding the detected polygon." +LONG_DESCRIPTION = """ +Calculate the minimum rotated bounding rectangle around polygon segmentation masks, converting complex polygon shapes into simplified rectangular bounding boxes with orientation information for zone creation, region simplification, and rectangular area approximation based on detected object shapes. + +## How This Block Works + +This block processes instance segmentation predictions with polygon masks and calculates the minimum rotated bounding rectangle (smallest rectangle that can enclose the polygon) for each detection. The block: + +1. Receives instance segmentation predictions containing polygon masks (validates that masks are present) +2. Processes each detection's polygon mask individually +3. Extracts the largest contour from the polygon mask (handles multi-contour masks by selecting the largest) +4. Calculates the minimum rotated bounding rectangle using OpenCV's minAreaRect: + - Finds the smallest rotated rectangle that can completely enclose the polygon + - Determines the rectangle's center point, dimensions (width, height), and rotation angle + - Computes the four corner points of the rotated rectangle +5. Updates the detection's mask to be a rectangular mask matching the calculated bounding rectangle (converts the rotated rectangle polygon back to a mask format) +6. Updates the detection's axis-aligned bounding box (xyxy) to the bounding box of the rotated rectangle (fits the rotated rectangle into an axis-aligned box) +7. Stores additional rectangle metadata in the detection data: + - Rectangle corner coordinates (rotated rectangle points) + - Rectangle width and height (dimensions of the rotated rectangle) + - Rectangle angle (rotation angle in degrees) +8. Merges all processed detections and returns them with updated masks, bounding boxes, and rectangle metadata + +The block transforms complex polygon shapes into simplified rectangular representations, preserving orientation information through the rotation angle. This is particularly useful when you need to create zones or regions based on detected object shapes (e.g., sports fields, road segments, marked areas) and want to simplify them to rectangular approximations. The minimum rotated rectangle provides the most compact rectangular representation of the polygon, potentially at an angle to minimize area. + +## Common Use Cases + +- **Zone Creation from Object Shapes**: Convert detected polygon shapes into rectangular zones for area monitoring or analysis (e.g., create zones from basketball court detections, generate road segment zones from road markings, create rectangular regions from zebra crossing detections), enabling zone-based workflows from complex shapes +- **Region Simplification**: Simplify complex polygon shapes to rectangular approximations for easier processing (e.g., simplify irregular segmentation masks to rectangles, convert complex shapes to rectangular regions, approximate polygon areas with rectangles), enabling simplified region processing +- **Rotated Region Detection**: Detect and extract rotated rectangular regions from polygon detections (e.g., find rotated parking spaces from segmentation, detect angled road markings as rectangles, extract rotated objects as rectangular zones), enabling rotation-aware region extraction +- **Area Approximation**: Approximate polygon areas with compact rectangular bounding boxes (e.g., approximate sports field areas with minimal rectangles, estimate region sizes using rotated bounding boxes, calculate compact rectangular areas from complex shapes), enabling area estimation with rectangular approximations +- **Shape Normalization**: Normalize polygon shapes to rectangular representations for standardized processing (e.g., normalize detected shapes to rectangles for consistent analysis, standardize polygon regions to rectangular format, convert variable shapes to uniform rectangular regions), enabling shape normalization workflows +- **Multi-Object Zone Extraction**: Extract rectangular zones from multiple detected polygon objects (e.g., create zones from multiple road segment detections, generate rectangular regions from multiple field detections, extract zones from various marked area detections), enabling multi-object zone creation workflows + +## Connecting to Other Blocks + +This block receives instance segmentation predictions with polygon masks and produces detections with rectangular masks and bounding boxes: + +- **After instance segmentation blocks** to convert polygon masks to rectangular bounding boxes for zone creation or simplified processing, enabling rectangular zone generation from complex shapes +- **Before zone-based blocks** (e.g., Polygon Zone, Dynamic Zone) to prepare rectangular regions for zone-based workflows (e.g., create zones from simplified rectangles, use rectangular approximations for zone monitoring, enable zone workflows with rectangular regions), enabling zone-based workflows with simplified shapes +- **After filtering blocks** (e.g., Detections Filter) to process only specific polygon detections before converting to rectangles (e.g., filter detections by class before rectangular conversion, select specific polygon types for rectangle extraction, prepare filtered detections for zone creation), enabling selective rectangle extraction +- **Before crop blocks** to extract rectangular regions from polygon detections (e.g., crop rotated rectangular regions from polygon shapes, extract rectangular areas from complex detections, prepare rectangular crop regions from polygons), enabling rectangular region extraction +- **Before visualization blocks** to display simplified rectangular representations of complex polygons (e.g., visualize rectangular approximations of polygons, display rotated bounding rectangles, show simplified rectangular zones), enabling rectangular visualization outputs +- **Before analysis blocks** that work better with rectangular regions than complex polygons (e.g., analyze rectangular zones instead of polygons, process simplified rectangular regions, work with normalized rectangular shapes), enabling simplified region analysis workflows +""" + + +class BoundingRectManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Bounding Rectangle", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "transformation", + "icon": "fal fa-rectangles-mixed", + "blockPriority": 5, + }, + } + ) + type: Literal["roboflow_core/bounding_rect@v1"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Instance segmentation predictions containing polygon masks. The block requires masks to be present - it will raise an error if masks are missing. Each detection's polygon mask will be processed to calculate the minimum rotated bounding rectangle. The mask should contain polygon shapes that you want to convert to rectangular bounding boxes. Detections are processed individually, with the largest contour extracted from each mask. The block outputs detections with updated masks (rectangular), updated bounding boxes (axis-aligned boxes of the rotated rectangles), and additional rectangle metadata stored in `bboxes_metadata` (rectangle coordinates, width, height, angle).", + examples=["$segmentation.predictions"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +def calculate_minimum_bounding_rectangle( + mask: np.ndarray, +) -> Tuple[np.ndarray, float, float, float]: + contours = sv.mask_to_polygons(mask) + largest_contour = max(contours, key=len) + + rect = cv.minAreaRect(largest_contour) + box = cv.boxPoints(rect) + box = np.array(box, dtype=int) + + width, height = rect[1] + angle = rect[2] + return box, width, height, angle + + +def _mask_resolution( + mask: Union[torch.Tensor, InstancesRLEMasks], +) -> Tuple[int, int]: + """Return the mask `(H, W)` resolution without materialising the whole + stack โ€” `image_size` for RLE, the trailing dims for a dense tensor.""" + if isinstance(mask, InstancesRLEMasks): + return int(mask.image_size[0]), int(mask.image_size[1]) + return int(mask.shape[1]), int(mask.shape[2]) + + +def _repackage_masks( + new_dense_masks: List[np.ndarray], + original_mask: Union[torch.Tensor, InstancesRLEMasks], +) -> Union[torch.Tensor, InstancesRLEMasks]: + """Match the output mask representation to the input โ€” dense in, + dense out; RLE in, RLE out โ€” so the rest of the tensor pipeline + keeps the same storage shape it had upstream.""" + if isinstance(original_mask, torch.Tensor): + return torch.from_numpy(np.stack(new_dense_masks, axis=0)).to( + device=original_mask.device, dtype=original_mask.dtype + ) + rles = [ + mask_utils.encode(np.asfortranarray(m.astype(np.uint8))) + for m in new_dense_masks + ] + return InstancesRLEMasks.from_coco_rle_masks( + image_size=original_mask.image_size, masks=rles + ) + + +class BoundingRectBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BoundingRectManifest + + def run( + self, + predictions: InstanceDetections, + ) -> BlockResult: + if predictions.mask is None: + raise ValueError( + "Mask missing. This block operates on output from segmentation model." + ) + n = len(predictions) + if n == 0: + return { + OUTPUT_KEY: InstanceDetections( + xyxy=predictions.xyxy, + class_id=predictions.class_id, + confidence=predictions.confidence, + mask=predictions.mask, + image_metadata=predictions.image_metadata, + bboxes_metadata=None, + ) + } + mask_h, mask_w = _mask_resolution(predictions.mask) + + new_dense_masks: List[np.ndarray] = [] + new_xyxy_rows: List[np.ndarray] = [] + new_bboxes_metadata: List[dict] = [] + existing_meta = predictions.bboxes_metadata or [{} for _ in range(n)] + + for i in range(n): + # Decode one instance at a time (RLE row decoded on demand, dense + # row pulled to host individually) instead of holding two full + # (n, H, W) bool stacks โ€” the same streaming convention as the + # serialiser / instance_mask_to_numpy. + instance_mask = instance_mask_to_numpy(predictions, i).astype(bool) + rect, width, height, angle = calculate_minimum_bounding_rectangle( + instance_mask + ) + rect_polygon = np.around(rect).astype(np.int32) + new_dense_masks.append( + sv.polygon_to_mask( + polygon=rect_polygon, + resolution_wh=(mask_w, mask_h), + ).astype(bool) + ) + new_xyxy_rows.append(sv.polygon_to_xyxy(polygon=rect_polygon)) + + per_box_meta = { + **(existing_meta[i] or {}), + BOUNDING_RECT_RECT_KEY_IN_SV_DETECTIONS: rect.astype( + np.float16 + ).tolist(), + BOUNDING_RECT_WIDTH_KEY_IN_SV_DETECTIONS: float(width), + BOUNDING_RECT_HEIGHT_KEY_IN_SV_DETECTIONS: float(height), + BOUNDING_RECT_ANGLE_KEY_IN_SV_DETECTIONS: float(angle), + } + new_bboxes_metadata.append(per_box_meta) + + new_xyxy = torch.tensor( + np.stack(new_xyxy_rows, axis=0), + dtype=predictions.xyxy.dtype, + device=predictions.xyxy.device, + ) + new_mask = _repackage_masks(new_dense_masks, predictions.mask) + + return { + OUTPUT_KEY: InstanceDetections( + xyxy=new_xyxy, + class_id=predictions.class_id, + confidence=predictions.confidence, + mask=new_mask, + image_metadata=predictions.image_metadata, + bboxes_metadata=new_bboxes_metadata, + ) + } diff --git a/inference/core/workflows/core_steps/transformations/byte_tracker/v1_tensor.py b/inference/core/workflows/core_steps/transformations/byte_tracker/v1_tensor.py new file mode 100644 index 0000000000..feeee5b6c7 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/byte_tracker/v1_tensor.py @@ -0,0 +1,282 @@ +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + split_key_point_prediction, + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + VideoMetadata, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + VIDEO_METADATA_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "tracked_detections" +_INPUT_INDEX_KEY: str = "__byte_tracker_input_index__" + + +def _track_and_recover_indices( + tracker: sv.ByteTrack, + detections: Union[Detections, InstanceDetections], +) -> "tuple[List[int], np.ndarray]": + """Run one frame through the (numpy) ByteTrack tracker and recover, in sv + output order, the surviving input-row indices and their assigned tracker ids. + + Only ``xyxy`` / ``class_id`` / ``confidence`` are moved to the host (the + single, irreducible D2H) and tagged with a positional index; ByteTrack filters + and reorders detections, so the index is how the caller maps back onto the + original device tensors. + """ + n = int(detections.xyxy.shape[0]) + sv_input = sv.Detections( + xyxy=detections.xyxy.detach().to("cpu").numpy(), + class_id=detections.class_id.detach().to("cpu").numpy(), + confidence=detections.confidence.detach().to("cpu").numpy(), + data={_INPUT_INDEX_KEY: np.arange(n, dtype=np.int64)}, + ) + sv_tracked = tracker.update_with_detections(sv_input) + kept_indices = ( + sv_tracked.data.get(_INPUT_INDEX_KEY, np.empty((0,), dtype=np.int64)) + if sv_tracked.data + else np.empty((0,), dtype=np.int64) + ) + tracker_ids = ( + sv_tracked.tracker_id + if sv_tracked.tracker_id is not None + else np.full(len(sv_tracked), -1, dtype=np.int64) + ) + return [int(i) for i in np.asarray(kept_indices).tolist()], tracker_ids + + +def _write_tracker_ids( + detections: Union[Detections, InstanceDetections], + tracker_ids: np.ndarray, +) -> None: + """Write the assigned ``tracker_id`` into each surviving detection's + ``bboxes_metadata`` in place. + + ``tracker_ids`` is aligned with the sliced ``detections`` rows (both are in sv + output order). Existing per-detection dicts are copied so caller-owned state is + never mutated; when the input carried no metadata a fresh dict is created. + """ + n = int(detections.xyxy.shape[0]) + existing = detections.bboxes_metadata + new_meta: List[dict] = [] + for i in range(n): + base = dict(existing[i]) if existing is not None and existing[i] else {} + base["tracker_id"] = int(tracker_ids[i]) + new_meta.append(base) + detections.bboxes_metadata = new_meta if new_meta else None + + +SHORT_DESCRIPTION = ( + "Track and update object positions across video frames using ByteTrack." +) +LONG_DESCRIPTION = """ +Track objects across video frames using the ByteTrack algorithm to maintain consistent object identities, handle occlusions and temporary disappearances, associate detections with existing tracks, assign unique track IDs, and enable object behavior analysis, movement tracking, and video analytics workflows. + +## How This Block Works + +This block maintains object tracking across sequential video frames by associating detections from each frame with existing tracks and creating new tracks for new objects. The block: + +1. Receives detection predictions for the current frame and video metadata (including frame rate and video identifier) +2. Initializes or retrieves a ByteTrack tracker for the video: + - Creates a new tracker instance for each unique video (identified by video_identifier) + - Stores trackers in memory to maintain tracking state across frames + - Configures tracker with frame rate from metadata and user-specified parameters + - Reuses existing tracker for subsequent frames of the same video +3. Materialises the tensor-native predictions to sv.Detections at the input boundary (the ByteTrack implementation is numpy-based), runs ByteTrack, then re-packages the surviving detections back into `inference_models.Detections`. The tracker_id assigned by ByteTrack lands in `bboxes_metadata` per surviving detection. +4. Updates tracks using ByteTrack algorithm: + - **Track Association**: Matches current frame detections to existing tracks using IoU (Intersection over Union) matching + - **Track Activation**: Creates new tracks for detections with confidence above track_activation_threshold that don't match existing tracks + - **Track Matching**: Associates detections to tracks when IoU exceeds minimum_matching_threshold + - **Track Persistence**: Maintains tracks that don't have matches using lost_track_buffer to handle temporary occlusions + - **Track Validation**: Only outputs tracks that have been present for at least minimum_consecutive_frames consecutive frames +5. Handles tracking challenges: + - **Occlusions**: Maintains tracks when objects are temporarily hidden (using lost_track_buffer frames) + - **Missed Detections**: Keeps tracks alive through frames with missing detections + - **False Positives**: Filters out tracks that don't persist long enough (minimum_consecutive_frames) + - **Track Fragmentation**: Reduces track splits by maintaining buffer for lost objects +6. Assigns unique track IDs to each object: + - Each tracked object receives a consistent track_id that persists across frames + - Track IDs are assigned when tracks are activated and maintained throughout the video + - Enables tracking individual objects across the entire video sequence +7. Returns tracked detections with track IDs: + - Outputs detection predictions enhanced with track_id information + - Each detection includes its assigned track_id for identifying the same object across frames + - Maintains all original detection properties (bounding boxes, confidence, class names) plus tracking information + +ByteTrack is an efficient multi-object tracking algorithm that performs tracking-by-detection, associating detections across frames without requiring appearance features. It uses a two-stage association strategy: first matching high-confidence detections to tracks, then matching low-confidence detections to remaining tracks and lost tracks. The algorithm maintains a buffer for lost tracks, allowing it to recover tracks when objects temporarily disappear due to occlusions or detection failures. The configurable parameters allow fine-tuning tracking behavior: track_activation_threshold controls when new tracks are created (higher = more conservative), lost_track_buffer controls occlusion handling (higher = better occlusion recovery), minimum_matching_threshold controls association quality (higher = stricter matching), and minimum_consecutive_frames filters short-lived false tracks (higher = fewer false tracks). + +## Common Use Cases + +- **Video Analytics**: Track objects across video frames for behavior analysis and movement patterns (e.g., track people movement in videos, monitor vehicle paths, analyze object trajectories), enabling video analytics workflows +- **Traffic Monitoring**: Track vehicles and objects in traffic scenes for traffic analysis (e.g., track vehicles across frames, monitor vehicle paths, count vehicles with consistent IDs), enabling traffic monitoring workflows +- **Surveillance Systems**: Maintain object identities across video frames for security monitoring (e.g., track individuals in surveillance footage, monitor object movements, maintain object identities), enabling surveillance tracking workflows +- **Sports Analysis**: Track players and objects in sports videos for performance analysis (e.g., track player movements, analyze player trajectories, monitor ball positions), enabling sports analysis workflows +- **Retail Analytics**: Track customers and products across video frames for retail insights (e.g., track customer paths, monitor shopping behavior, analyze foot traffic patterns), enabling retail analytics workflows +- **Object Behavior Analysis**: Track objects to analyze their behavior and interactions over time (e.g., analyze object interactions, study movement patterns, track object relationships), enabling behavior analysis workflows + +## Connecting to Other Blocks + +This block receives detection predictions and video metadata, and produces tracked_detections with track IDs: + +- **After object detection or instance segmentation blocks** to track detected objects across video frames (e.g., track detected objects in video, add track IDs to detections, maintain object identities across frames), enabling detection-to-tracking workflows +- **Before video analysis blocks** that require consistent object identities (e.g., analyze tracked object behavior, process object trajectories, work with tracked object data), enabling tracking-to-analysis workflows +- **Before visualization blocks** to display tracked objects with consistent colors or labels (e.g., visualize tracked objects, display track IDs, show object paths), enabling tracking visualization workflows +- **Before logic blocks** like Continue If to make decisions based on track information (e.g., continue if object is tracked, filter based on track IDs, make decisions using tracking data), enabling tracking-based decision workflows +- **Before counting or aggregation blocks** to count tracked objects accurately (e.g., count unique tracked objects, aggregate track statistics, process track data), enabling tracking-to-counting workflows +- **In video processing pipelines** where object tracking is part of a larger video analysis workflow (e.g., track objects in video pipelines, maintain identities in processing chains, enable video analytics), enabling video tracking pipeline workflows + +## Requirements + +This block requires detection predictions (object detection or instance segmentation) and video metadata with frame rate (fps) information. The video metadata must include a valid fps value for ByteTrack initialization. The block maintains tracking state across frames for each video, so it should be used in video workflows where frames are processed sequentially. For optimal tracking performance, detections should be provided consistently across frames. The algorithm works best with stable detection performance and handles temporary detection gaps through the lost_track_buffer mechanism. + +Tensor-native note: the underlying `sv.ByteTrack` is numpy-based, so this block materialises only the bounding boxes / class ids / confidences to numpy at the boundary (tagged with a positional index) and runs the tracker. ByteTrack only filters and reorders detections โ€” it never mutates coordinates โ€” so the surviving rows are recovered by index and the ORIGINAL device tensors are sliced with `take_prediction_by_indices` (no GPU re-upload). Instance-segmentation masks follow the same index selection and are preserved (matching the numpy block, which keeps masks via `sv.Detections`). The assigned `tracker_id` is written into `bboxes_metadata` per surviving detection. +""" + + +class ByteTrackerBlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Byte Tracker", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "deprecated": True, + "ui_manifest": { + "section": "video", + "icon": "far fa-location-crosshairs", + "blockPriority": 0, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/byte_tracker@v1"] + metadata: Selector(kind=[VIDEO_METADATA_KIND]) = Field( + description="Video metadata containing frame rate (fps) and video identifier information required for ByteTrack initialization and tracking state management. The fps value is used to configure the tracker, and the video_identifier is used to maintain separate tracking state for different videos. The metadata must include valid fps information - ByteTrack requires frame rate to initialize. If processing multiple videos, each video's metadata should have a unique video_identifier to maintain separate tracking states. The block maintains persistent trackers across frames for each video using the video_identifier.", + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Detection predictions (object detection or instance segmentation) for the current video frame to be tracked. The block associates these detections with existing tracks or creates new tracks. Detections should be provided for each frame in sequence to maintain consistent tracking. The detections must include bounding boxes and class names. After tracking, the output will include the surviving detections enhanced with a `tracker_id` entry in `bboxes_metadata`, allowing identification of the same object across frames.", + examples=["$steps.object_detection_model.predictions"], + ) + track_activation_threshold: Union[Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + default=0.25, + description="Confidence threshold for activating new tracks from detections. Must be between 0.0 and 1.0. Default is 0.25. Only detections with confidence above this threshold can create new tracks. Increasing this threshold (e.g., 0.3-0.5) improves tracking accuracy and stability by only creating tracks from high-confidence detections, but might miss true detections with lower confidence. Decreasing this threshold (e.g., 0.15-0.2) increases tracking completeness by accepting lower-confidence detections, but risks introducing noise and instability from false positives. Adjust based on detection model performance: use lower values if detections are reliable, higher values if false positives are common.", + examples=[0.25, "$inputs.confidence"], + ) + lost_track_buffer: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + default=30, + description="Number of frames to maintain a track when it's lost (no matching detections). Must be a positive integer. Default is 30 frames. When an object temporarily disappears (due to occlusion, missed detection, or leaving frame), the track is maintained for this many frames before being considered lost. Increasing this value (e.g., 50-100) enhances occlusion handling and significantly reduces track fragmentation or disappearance caused by brief detection gaps, but increases memory usage. Decreasing this value (e.g., 10-20) reduces memory usage but may cause tracks to disappear during short occlusions. Adjust based on occlusion frequency: use higher values for frequent occlusions, lower values for stable tracking scenarios.", + examples=[30, "$inputs.lost_track_buffer"], + ) + minimum_matching_threshold: Union[Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + default=0.8, + description="IoU (Intersection over Union) threshold for matching detections to existing tracks. Must be between 0.0 and 1.0. Default is 0.8. Detections are associated with tracks when their bounding box IoU exceeds this threshold. Increasing this threshold (e.g., 0.85-0.95) improves tracking accuracy by requiring stronger spatial overlap, but risks track fragmentation when objects move quickly or detection boxes vary. Decreasing this threshold (e.g., 0.6-0.75) improves tracking completeness by accepting looser matches, but risks false positive associations and track drift. Adjust based on object movement speed and detection stability: use higher values for stable objects, lower values for fast-moving objects.", + examples=[0.8, "$inputs.min_matching_threshold"], + ) + minimum_consecutive_frames: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + default=1, + description="Minimum number of consecutive frames an object must be tracked before the track is considered valid and output. Must be a positive integer. Default is 1 (all tracks are immediately valid). Only tracks that persist for at least this many consecutive frames are included in the output. Increasing this value (e.g., 3-5) prevents the creation of accidental tracks from false detections or double detections, filtering out short-lived spurious tracks, but risks missing shorter legitimate tracks. Decreasing this value (e.g., 1) includes all tracks immediately, maximizing completeness but potentially including false tracks. Adjust based on false positive rate: use higher values if false detections are common, lower values if detections are reliable.", + examples=[1, "$inputs.min_consecutive_frames"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class ByteTrackerBlockV1(WorkflowBlock): + def __init__( + self, + ): + self._trackers: Dict[str, sv.ByteTrack] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return ByteTrackerBlockManifest + + def run( + self, + metadata: VideoMetadata, + detections: Union[Detections, InstanceDetections], + track_activation_threshold: float = 0.25, + lost_track_buffer: int = 30, + minimum_matching_threshold: float = 0.8, + minimum_consecutive_frames: int = 1, + ) -> BlockResult: + if not metadata.fps: + raise ValueError( + f"Malformed fps in VideoMetadata, {self.__class__.__name__} requires fps in order to initialize ByteTrack" + ) + if metadata.video_identifier not in self._trackers: + self._trackers[metadata.video_identifier] = sv.ByteTrack( + track_activation_threshold=track_activation_threshold, + lost_track_buffer=lost_track_buffer, + minimum_matching_threshold=minimum_matching_threshold, + minimum_consecutive_frames=minimum_consecutive_frames, + frame_rate=metadata.fps, + ) + tracker = self._trackers[metadata.video_identifier] + + # Materialise ONLY the bbox component to a minimal sv.Detections (the sole + # D2H) tagged with a positional index. ByteTrack drops untracked detections + # and reorders, so the index lets us recover the surviving input rows. + _, bbox = split_key_point_prediction(detections) + kept_indices, tracker_ids = _track_and_recover_indices( + tracker=tracker, detections=bbox + ) + + # ByteTrack only filters/reorders - coordinates are unchanged - so slice the + # ORIGINAL device tensors (masks included) by the surviving rows instead of + # re-uploading the numpy boxes. `take_prediction_by_indices` copies the + # surviving bboxes_metadata dicts, so the tracker_id write is leak-safe. + tracked_detections = take_prediction_by_indices(bbox, kept_indices) + _write_tracker_ids(tracked_detections, tracker_ids) + + return {OUTPUT_KEY: tracked_detections} diff --git a/inference/core/workflows/core_steps/transformations/byte_tracker/v2_tensor.py b/inference/core/workflows/core_steps/transformations/byte_tracker/v2_tensor.py new file mode 100644 index 0000000000..9f74bbef08 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/byte_tracker/v2_tensor.py @@ -0,0 +1,299 @@ +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core import logger +from inference.core.workflows.core_steps.common.tensor_native import ( + split_key_point_prediction, + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + Selector, + WorkflowImageSelector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "tracked_detections" +_INPUT_INDEX_KEY: str = "__byte_tracker_input_index__" + + +def _track_and_recover_indices( + tracker: sv.ByteTrack, + detections: Union[Detections, InstanceDetections], +) -> "tuple[List[int], np.ndarray]": + """Run one frame through the (numpy) ByteTrack tracker and recover, in sv + output order, the surviving input-row indices and their assigned tracker ids. + + Only ``xyxy`` / ``class_id`` / ``confidence`` are moved to the host (the + single, irreducible D2H) and tagged with a positional index; ByteTrack filters + and reorders detections, so the index is how the caller maps back onto the + original device tensors. + """ + n = int(detections.xyxy.shape[0]) + sv_input = sv.Detections( + xyxy=detections.xyxy.detach().to("cpu").numpy(), + class_id=detections.class_id.detach().to("cpu").numpy(), + confidence=detections.confidence.detach().to("cpu").numpy(), + data={_INPUT_INDEX_KEY: np.arange(n, dtype=np.int64)}, + ) + sv_tracked = tracker.update_with_detections(sv_input) + kept_indices = ( + sv_tracked.data.get(_INPUT_INDEX_KEY, np.empty((0,), dtype=np.int64)) + if sv_tracked.data + else np.empty((0,), dtype=np.int64) + ) + tracker_ids = ( + sv_tracked.tracker_id + if sv_tracked.tracker_id is not None + else np.full(len(sv_tracked), -1, dtype=np.int64) + ) + return [int(i) for i in np.asarray(kept_indices).tolist()], tracker_ids + + +def _write_tracker_ids( + detections: Union[Detections, InstanceDetections], + tracker_ids: np.ndarray, +) -> None: + """Write the assigned ``tracker_id`` into each surviving detection's + ``bboxes_metadata`` in place. + + ``tracker_ids`` is aligned with the sliced ``detections`` rows (both are in sv + output order). Existing per-detection dicts are copied so caller-owned state is + never mutated; when the input carried no metadata a fresh dict is created. + """ + n = int(detections.xyxy.shape[0]) + existing = detections.bboxes_metadata + new_meta: List[dict] = [] + for i in range(n): + base = dict(existing[i]) if existing is not None and existing[i] else {} + base["tracker_id"] = int(tracker_ids[i]) + new_meta.append(base) + detections.bboxes_metadata = new_meta if new_meta else None + + +SHORT_DESCRIPTION = ( + "Track and update object positions across video frames using ByteTrack." +) +LONG_DESCRIPTION = """ +Track objects across video frames using the ByteTrack algorithm to maintain consistent object identities, handle occlusions and temporary disappearances, associate detections with existing tracks, assign unique track IDs, and enable object behavior analysis, movement tracking, and video analytics workflows. + +## How This Block Works + +This block maintains object tracking across sequential video frames by associating detections from each frame with existing tracks and creating new tracks for new objects. The block: + +1. Receives detection predictions for the current frame and an image with embedded video metadata +2. Extracts video metadata from the image (including frame rate and video identifier): + - Accesses video_metadata from the WorkflowImageData object + - Extracts fps (frames per second) for tracker configuration + - Extracts video_identifier to maintain separate tracking state for different videos + - Handles missing fps gracefully (defaults to 0 and logs a warning instead of failing) +3. Initializes or retrieves a ByteTrack tracker for the video: + - Creates a new tracker instance for each unique video (identified by video_identifier) + - Stores trackers in memory to maintain tracking state across frames + - Configures tracker with frame rate from metadata and user-specified parameters + - Reuses existing tracker for subsequent frames of the same video +4. Materialises the tensor-native predictions to sv.Detections at the input boundary (the ByteTrack implementation is numpy-based), runs ByteTrack, then re-packages the surviving detections back into `inference_models.Detections`. The tracker_id assigned by ByteTrack lands in `bboxes_metadata` per surviving detection. +5. Updates tracks using ByteTrack algorithm: + - **Track Association**: Matches current frame detections to existing tracks using IoU (Intersection over Union) matching + - **Track Activation**: Creates new tracks for detections with confidence above track_activation_threshold that don't match existing tracks + - **Track Matching**: Associates detections to tracks when IoU exceeds minimum_matching_threshold + - **Track Persistence**: Maintains tracks that don't have matches using lost_track_buffer to handle temporary occlusions + - **Track Validation**: Only outputs tracks that have been present for at least minimum_consecutive_frames consecutive frames +6. Handles tracking challenges: + - **Occlusions**: Maintains tracks when objects are temporarily hidden (using lost_track_buffer frames) + - **Missed Detections**: Keeps tracks alive through frames with missing detections + - **False Positives**: Filters out tracks that don't persist long enough (minimum_consecutive_frames) + - **Track Fragmentation**: Reduces track splits by maintaining buffer for lost objects +7. Assigns unique track IDs to each object: + - Each tracked object receives a consistent track_id that persists across frames + - Track IDs are assigned when tracks are activated and maintained throughout the video + - Enables tracking individual objects across the entire video sequence +8. Returns tracked detections with track IDs: + - Outputs detection predictions enhanced with track_id information + - Each detection includes its assigned track_id for identifying the same object across frames + - Maintains all original detection properties (bounding boxes, confidence, class names) plus tracking information + +ByteTrack is an efficient multi-object tracking algorithm that performs tracking-by-detection, associating detections across frames without requiring appearance features. It uses a two-stage association strategy: first matching high-confidence detections to tracks, then matching low-confidence detections to remaining tracks and lost tracks. The algorithm maintains a buffer for lost tracks, allowing it to recover tracks when objects temporarily disappear due to occlusions or detection failures. The configurable parameters allow fine-tuning tracking behavior: track_activation_threshold controls when new tracks are created (higher = more conservative), lost_track_buffer controls occlusion handling (higher = better occlusion recovery), minimum_matching_threshold controls association quality (higher = stricter matching), and minimum_consecutive_frames filters short-lived false tracks (higher = fewer false tracks). + +## Common Use Cases + +- **Video Analytics**: Track objects across video frames for behavior analysis and movement patterns (e.g., track people movement in videos, monitor vehicle paths, analyze object trajectories), enabling video analytics workflows +- **Traffic Monitoring**: Track vehicles and objects in traffic scenes for traffic analysis (e.g., track vehicles across frames, monitor vehicle paths, count vehicles with consistent IDs), enabling traffic monitoring workflows +- **Surveillance Systems**: Maintain object identities across video frames for security monitoring (e.g., track individuals in surveillance footage, monitor object movements, maintain object identities), enabling surveillance tracking workflows +- **Sports Analysis**: Track players and objects in sports videos for performance analysis (e.g., track player movements, analyze player trajectories, monitor ball positions), enabling sports analysis workflows +- **Retail Analytics**: Track customers and products across video frames for retail insights (e.g., track customer paths, monitor shopping behavior, analyze foot traffic patterns), enabling retail analytics workflows +- **Object Behavior Analysis**: Track objects to analyze their behavior and interactions over time (e.g., analyze object interactions, study movement patterns, track object relationships), enabling behavior analysis workflows + +## Connecting to Other Blocks + +This block receives an image with video metadata and detection predictions, and produces tracked_detections with track IDs: + +- **After object detection or instance segmentation blocks** to track detected objects across video frames (e.g., track detected objects in video, add track IDs to detections, maintain object identities across frames), enabling detection-to-tracking workflows +- **Before video analysis blocks** that require consistent object identities (e.g., analyze tracked object behavior, process object trajectories, work with tracked object data), enabling tracking-to-analysis workflows +- **Before visualization blocks** to display tracked objects with consistent colors or labels (e.g., visualize tracked objects, display track IDs, show object paths), enabling tracking visualization workflows +- **Before logic blocks** like Continue If to make decisions based on track information (e.g., continue if object is tracked, filter based on track IDs, make decisions using tracking data), enabling tracking-based decision workflows +- **Before counting or aggregation blocks** to count tracked objects accurately (e.g., count unique tracked objects, aggregate track statistics, process track data), enabling tracking-to-counting workflows +- **In video processing pipelines** where object tracking is part of a larger video analysis workflow (e.g., track objects in video pipelines, maintain identities in processing chains, enable video analytics), enabling video tracking pipeline workflows + +## Version Differences + +**Enhanced from v1:** + +- **Simplified Input**: Uses `image` input that contains embedded video metadata instead of requiring a separate `metadata` field, simplifying workflow connections and reducing input complexity +- **Graceful FPS Handling**: Handles missing or invalid fps values gracefully by defaulting to 0 and logging a warning instead of raising an error, making the block more resilient to incomplete metadata +- **Improved Integration**: Better integration with image-based workflows since video metadata is accessed directly from the image object rather than requiring separate metadata input + +## Requirements + +This block requires detection predictions (object detection or instance segmentation) and an image with embedded video metadata containing frame rate (fps) and video identifier information. The image's video_metadata should include a valid fps value for optimal tracking performance, though the block will continue with fps=0 if missing. The block maintains tracking state across frames for each video, so it should be used in video workflows where frames are processed sequentially. For optimal tracking performance, detections should be provided consistently across frames. The algorithm works best with stable detection performance and handles temporary detection gaps through the lost_track_buffer mechanism. + +Tensor-native note: the underlying `sv.ByteTrack` is numpy-based, so this block materialises only the bounding boxes / class ids / confidences to numpy at the boundary (tagged with a positional index) and runs the tracker. ByteTrack only filters and reorders detections โ€” it never mutates coordinates โ€” so the surviving rows are recovered by index and the ORIGINAL device tensors are sliced with `take_prediction_by_indices` (no GPU re-upload). Instance-segmentation masks follow the same index selection and are preserved (matching the numpy block, which keeps masks via `sv.Detections`). The assigned `tracker_id` is written into `bboxes_metadata` per surviving detection. +""" + + +class ByteTrackerBlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Byte Tracker", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "deprecated": True, + "ui_manifest": { + "section": "video", + "icon": "far fa-location-crosshairs", + "blockPriority": 0, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/byte_tracker@v2"] + image: WorkflowImageSelector = Field( + description="Input image containing embedded video metadata (fps and video_identifier) required for ByteTrack initialization and tracking state management. The block extracts video_metadata from the WorkflowImageData object. The fps value is used to configure the tracker, and the video_identifier is used to maintain separate tracking state for different videos. If fps is missing or invalid, the block defaults to 0 and logs a warning but continues operation. If processing multiple videos, each video should have a unique video_identifier in its metadata to maintain separate tracking states. The block maintains persistent trackers across frames for each video using the video_identifier. This version simplifies input by embedding metadata in the image object rather than requiring a separate metadata field.", + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Detection predictions (object detection or instance segmentation) for the current video frame to be tracked. The block associates these detections with existing tracks or creates new tracks. Detections should be provided for each frame in sequence to maintain consistent tracking. The detections must include bounding boxes and class names. After tracking, the output will include the surviving detections enhanced with a `tracker_id` entry in `bboxes_metadata`, allowing identification of the same object across frames.", + examples=["$steps.object_detection_model.predictions"], + ) + track_activation_threshold: Union[Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + default=0.25, + description="Confidence threshold for activating new tracks from detections. Must be between 0.0 and 1.0. Default is 0.25. Only detections with confidence above this threshold can create new tracks. Increasing this threshold (e.g., 0.3-0.5) improves tracking accuracy and stability by only creating tracks from high-confidence detections, but might miss true detections with lower confidence. Decreasing this threshold (e.g., 0.15-0.2) increases tracking completeness by accepting lower-confidence detections, but risks introducing noise and instability from false positives. Adjust based on detection model performance: use lower values if detections are reliable, higher values if false positives are common.", + examples=[0.25, "$inputs.confidence"], + ) + lost_track_buffer: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + default=30, + description="Number of frames to maintain a track when it's lost (no matching detections). Must be a positive integer. Default is 30 frames. When an object temporarily disappears (due to occlusion, missed detection, or leaving frame), the track is maintained for this many frames before being considered lost. Increasing this value (e.g., 50-100) enhances occlusion handling and significantly reduces track fragmentation or disappearance caused by brief detection gaps, but increases memory usage. Decreasing this value (e.g., 10-20) reduces memory usage but may cause tracks to disappear during short occlusions. Adjust based on occlusion frequency: use higher values for frequent occlusions, lower values for stable tracking scenarios.", + examples=[30, "$inputs.lost_track_buffer"], + ) + minimum_matching_threshold: Union[Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + default=0.8, + description="IoU (Intersection over Union) threshold for matching detections to existing tracks. Must be between 0.0 and 1.0. Default is 0.8. Detections are associated with tracks when their bounding box IoU exceeds this threshold. Increasing this threshold (e.g., 0.85-0.95) improves tracking accuracy by requiring stronger spatial overlap, but risks track fragmentation when objects move quickly or detection boxes vary. Decreasing this threshold (e.g., 0.6-0.75) improves tracking completeness by accepting looser matches, but risks false positive associations and track drift. Adjust based on object movement speed and detection stability: use higher values for stable objects, lower values for fast-moving objects.", + examples=[0.8, "$inputs.min_matching_threshold"], + ) + minimum_consecutive_frames: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + default=1, + description="Minimum number of consecutive frames an object must be tracked before the track is considered valid and output. Must be a positive integer. Default is 1 (all tracks are immediately valid). Only tracks that persist for at least this many consecutive frames are included in the output. Increasing this value (e.g., 3-5) prevents the creation of accidental tracks from false detections or double detections, filtering out short-lived spurious tracks, but risks missing shorter legitimate tracks. Decreasing this value (e.g., 1) includes all tracks immediately, maximizing completeness but potentially including false tracks. Adjust based on false positive rate: use higher values if false detections are common, lower values if detections are reliable.", + examples=[1, "$inputs.min_consecutive_frames"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class ByteTrackerBlockV2(WorkflowBlock): + def __init__( + self, + ): + self._trackers: Dict[str, sv.ByteTrack] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return ByteTrackerBlockManifest + + def run( + self, + image: WorkflowImageData, + detections: Union[Detections, InstanceDetections], + track_activation_threshold: float = 0.25, + lost_track_buffer: int = 30, + minimum_matching_threshold: float = 0.8, + minimum_consecutive_frames: int = 1, + ) -> BlockResult: + metadata = image.video_metadata + fps = metadata.fps + if not fps: + fps = 0 + logger.warning( + f"Malformed fps in VideoMetadata, {self.__class__.__name__} requires fps in order to initialize ByteTrack" + ) + if metadata.video_identifier not in self._trackers: + self._trackers[metadata.video_identifier] = sv.ByteTrack( + track_activation_threshold=track_activation_threshold, + lost_track_buffer=lost_track_buffer, + minimum_matching_threshold=minimum_matching_threshold, + minimum_consecutive_frames=minimum_consecutive_frames, + frame_rate=fps, + ) + tracker = self._trackers[metadata.video_identifier] + + # Materialise ONLY the bbox component to a minimal sv.Detections (the sole + # D2H) tagged with a positional index. ByteTrack drops untracked detections + # and reorders, so the index lets us recover the surviving input rows. + _, bbox = split_key_point_prediction(detections) + kept_indices, tracker_ids = _track_and_recover_indices( + tracker=tracker, detections=bbox + ) + + # ByteTrack only filters/reorders - coordinates are unchanged - so slice the + # ORIGINAL device tensors (masks included) by the surviving rows instead of + # re-uploading the numpy boxes. `take_prediction_by_indices` copies the + # surviving bboxes_metadata dicts, so the tracker_id write is leak-safe. + tracked_detections = take_prediction_by_indices(bbox, kept_indices) + _write_tracker_ids(tracked_detections, tracker_ids) + + return {OUTPUT_KEY: tracked_detections} diff --git a/inference/core/workflows/core_steps/transformations/byte_tracker/v3_tensor.py b/inference/core/workflows/core_steps/transformations/byte_tracker/v3_tensor.py new file mode 100644 index 0000000000..ac36fca275 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/byte_tracker/v3_tensor.py @@ -0,0 +1,377 @@ +from collections import deque +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core import logger +from inference.core.workflows.core_steps.common.tensor_native import ( + split_key_point_prediction, + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "tracked_detections" +_INPUT_INDEX_KEY: str = "__byte_tracker_input_index__" + + +def _track_and_recover_indices( + tracker: sv.ByteTrack, + detections: Union[Detections, InstanceDetections], +) -> "tuple[List[int], np.ndarray]": + """Run one frame through the (numpy) ByteTrack tracker and recover, in sv + output order, the surviving input-row indices and their assigned tracker ids. + + Only ``xyxy`` / ``class_id`` / ``confidence`` are moved to the host (the + single, irreducible D2H) and tagged with a positional index; ByteTrack filters + and reorders detections, so the index is how the caller maps back onto the + original device tensors. + """ + n = int(detections.xyxy.shape[0]) + sv_input = sv.Detections( + xyxy=detections.xyxy.detach().to("cpu").numpy(), + class_id=detections.class_id.detach().to("cpu").numpy(), + confidence=detections.confidence.detach().to("cpu").numpy(), + data={_INPUT_INDEX_KEY: np.arange(n, dtype=np.int64)}, + ) + sv_tracked = tracker.update_with_detections(sv_input) + kept_indices = ( + sv_tracked.data.get(_INPUT_INDEX_KEY, np.empty((0,), dtype=np.int64)) + if sv_tracked.data + else np.empty((0,), dtype=np.int64) + ) + tracker_ids = ( + sv_tracked.tracker_id + if sv_tracked.tracker_id is not None + else np.full(len(sv_tracked), -1, dtype=np.int64) + ) + return [int(i) for i in np.asarray(kept_indices).tolist()], tracker_ids + + +def _write_tracker_ids( + detections: Union[Detections, InstanceDetections], + tracker_ids: np.ndarray, +) -> None: + """Write the assigned ``tracker_id`` into each surviving detection's + ``bboxes_metadata`` in place. + + ``tracker_ids`` is aligned with the sliced ``detections`` rows (both are in sv + output order). Existing per-detection dicts are copied so caller-owned state is + never mutated; when the input carried no metadata a fresh dict is created. + """ + n = int(detections.xyxy.shape[0]) + existing = detections.bboxes_metadata + new_meta: List[dict] = [] + for i in range(n): + base = dict(existing[i]) if existing is not None and existing[i] else {} + base["tracker_id"] = int(tracker_ids[i]) + new_meta.append(base) + detections.bboxes_metadata = new_meta if new_meta else None + + +SHORT_DESCRIPTION = ( + "Track and update object positions across video frames using ByteTrack." +) +LONG_DESCRIPTION = """ +Track objects across video frames using the ByteTrack algorithm to maintain consistent object identities, handle occlusions and temporary disappearances, associate detections with existing tracks, assign unique track IDs, categorize instances as new or previously seen, and enable object behavior analysis, movement tracking, first-appearance detection, and video analytics workflows. + +## How This Block Works + +This block maintains object tracking across sequential video frames by associating detections from each frame with existing tracks and creating new tracks for new objects, while also categorizing instances based on whether they've been seen before. The block: + +1. Receives detection predictions for the current frame and an image with embedded video metadata +2. Extracts video metadata from the image (including frame rate and video identifier): + - Accesses video_metadata from the WorkflowImageData object + - Extracts fps (frames per second) for tracker configuration + - Extracts video_identifier to maintain separate tracking state for different videos + - Handles missing fps gracefully (defaults to 0 and logs a warning instead of failing) +3. Initializes or retrieves a ByteTrack tracker for the video: + - Creates a new tracker instance for each unique video (identified by video_identifier) + - Stores trackers in memory to maintain tracking state across frames + - Configures tracker with frame rate from metadata and user-specified parameters + - Reuses existing tracker for subsequent frames of the same video +4. Initializes or retrieves an instance cache for the video: + - Creates a cache to track which track IDs have been seen before + - Maintains separate cache for each video using video_identifier + - Configures cache size using instances_cache_size parameter + - Uses FIFO (First-In-First-Out) strategy to manage cache capacity +5. Materialises the tensor-native predictions to sv.Detections at the input boundary, runs ByteTrack, then re-packages the surviving detections back into `inference_models.Detections`. Track IDs are stashed in `bboxes_metadata` per surviving detection. +6. Updates tracks using ByteTrack algorithm: + - **Track Association**: Matches current frame detections to existing tracks using IoU (Intersection over Union) matching + - **Track Activation**: Creates new tracks for detections with confidence above track_activation_threshold that don't match existing tracks + - **Track Matching**: Associates detections to tracks when IoU exceeds minimum_matching_threshold + - **Track Persistence**: Maintains tracks that don't have matches using lost_track_buffer to handle temporary occlusions + - **Track Validation**: Only outputs tracks that have been present for at least minimum_consecutive_frames consecutive frames +7. Categorizes tracked instances as new or already seen: + - For each tracked detection with a track_id, checks the instance cache + - **New Instances**: Track IDs not found in cache are marked as new (first appearance) + - **Already Seen Instances**: Track IDs found in cache are marked as already seen (reappearance) + - Updates cache with new track IDs, managing cache size with FIFO eviction +8. Handles tracking challenges: + - **Occlusions**: Maintains tracks when objects are temporarily hidden (using lost_track_buffer frames) + - **Missed Detections**: Keeps tracks alive through frames with missing detections + - **False Positives**: Filters out tracks that don't persist long enough (minimum_consecutive_frames) + - **Track Fragmentation**: Reduces track splits by maintaining buffer for lost objects +9. Assigns unique track IDs to each object: + - Each tracked object receives a consistent track_id that persists across frames + - Track IDs are assigned when tracks are activated and maintained throughout the video + - Enables tracking individual objects across the entire video sequence +10. Returns three sets of tracked detections: + - **tracked_detections**: All tracked detections with track IDs (same as v2) + - **new_instances**: Detections with track IDs that are appearing for the first time (each track ID appears only once when first generated) + - **already_seen_instances**: Detections with track IDs that have been seen before (track IDs appear each time the tracker associates them with detections) + +ByteTrack is an efficient multi-object tracking algorithm that performs tracking-by-detection, associating detections across frames without requiring appearance features. It uses a two-stage association strategy: first matching high-confidence detections to tracks, then matching low-confidence detections to remaining tracks and lost tracks. The algorithm maintains a buffer for lost tracks, allowing it to recover tracks when objects temporarily disappear due to occlusions or detection failures. The instance categorization feature enables detection of first appearances (new objects entering the scene) versus reappearances (objects returning after occlusion or leaving frame), which is useful for counting, behavior analysis, and event detection. The configurable parameters allow fine-tuning tracking behavior: track_activation_threshold controls when new tracks are created (higher = more conservative), lost_track_buffer controls occlusion handling (higher = better occlusion recovery), minimum_matching_threshold controls association quality (higher = stricter matching), minimum_consecutive_frames filters short-lived false tracks (higher = fewer false tracks), and instances_cache_size controls how many track IDs to remember for new/seen categorization (higher = longer memory). + +## Common Use Cases + +- **Video Analytics**: Track objects across video frames for behavior analysis and movement patterns (e.g., track people movement in videos, monitor vehicle paths, analyze object trajectories), enabling video analytics workflows +- **First Appearance Detection**: Identify new objects entering the scene for counting and event detection (e.g., detect new people entering area, identify new vehicles appearing, track first-time appearances), enabling new instance detection workflows +- **Traffic Monitoring**: Track vehicles and objects in traffic scenes with appearance tracking (e.g., track vehicles across frames, monitor vehicle paths, count unique vehicles with consistent IDs, detect new vehicles entering scene), enabling traffic monitoring workflows +- **Surveillance Systems**: Maintain object identities and detect new entries for security monitoring (e.g., track individuals in surveillance footage, detect new people entering area, monitor object movements, maintain object identities), enabling surveillance tracking workflows +- **Retail Analytics**: Track customers and products with entry detection for retail insights (e.g., track customer paths, detect new customers entering store, monitor shopping behavior, analyze foot traffic patterns), enabling retail analytics workflows +- **Object Counting**: Accurately count unique objects by tracking first appearances (e.g., count unique visitors by tracking new instances, count vehicles entering intersection, track unique object appearances), enabling accurate counting workflows + +## Connecting to Other Blocks + +This block receives an image with video metadata and detection predictions, and produces tracked_detections, new_instances, and already_seen_instances: + +- **After object detection, instance segmentation, or keypoint detection blocks** to track detected objects across video frames (e.g., track detected objects in video, add track IDs to detections, maintain object identities across frames), enabling detection-to-tracking workflows +- **Using new_instances output** to detect and process first appearances (e.g., count new objects, trigger actions on first appearance, detect new entries, initialize tracking for new objects), enabling new instance detection workflows +- **Using already_seen_instances output** to process reappearances and returning objects (e.g., handle returning objects, process reappearances, filter for existing objects), enabling reappearance handling workflows +- **Before video analysis blocks** that require consistent object identities (e.g., analyze tracked object behavior, process object trajectories, work with tracked object data), enabling tracking-to-analysis workflows +- **Before visualization blocks** to display tracked objects with consistent colors or labels (e.g., visualize tracked objects, display track IDs, show object paths, highlight new instances), enabling tracking visualization workflows +- **Before logic blocks** like Continue If to make decisions based on track information or instance status (e.g., continue if object is new, filter based on track IDs, make decisions using tracking data, handle new vs returning objects), enabling tracking-based decision workflows + +## Version Differences + +**Enhanced from v2:** + +- **Instance Categorization**: Adds two new outputs (`new_instances` and `already_seen_instances`) that categorize tracked objects based on whether their track IDs have been seen before, enabling first-appearance detection and reappearance tracking +- **Instance Cache**: Introduces an instance cache system that remembers previously seen track IDs across frames, allowing distinction between new objects entering the scene and objects reappearing after occlusion or leaving frame +- **Keypoint Detection Support**: Adds support for keypoint detection predictions in addition to object detection and instance segmentation, expanding tracking capabilities to keypoint-based detection models +- **Configurable Cache Size**: Adds `instances_cache_size` parameter to control how many track IDs are remembered in the cache, balancing memory usage with tracking history length +- **Enhanced Outputs**: Returns three outputs instead of one - `tracked_detections` (all tracked objects), `new_instances` (first appearances), and `already_seen_instances` (reappearances) + +## Requirements + +This block requires detection predictions (object detection, instance segmentation, or keypoint detection) and an image with embedded video metadata containing frame rate (fps) and video identifier information. The image's video_metadata should include a valid fps value for optimal tracking performance, though the block will continue with fps=0 if missing. The block maintains tracking state and instance cache across frames for each video, so it should be used in video workflows where frames are processed sequentially. For optimal tracking performance, detections should be provided consistently across frames. The algorithm works best with stable detection performance and handles temporary detection gaps through the lost_track_buffer mechanism. The instance cache maintains a history of seen track IDs with FIFO eviction when the cache size limit is reached. + +Tensor-native note: the underlying `sv.ByteTrack` is numpy-based, so this block materialises only the bounding boxes / class ids / confidences to numpy at the boundary (tagged with a positional index) and runs the tracker. ByteTrack only filters and reorders detections โ€” it never mutates coordinates โ€” so the surviving rows are recovered by index and the ORIGINAL device tensors are sliced with `take_prediction_by_indices` (no GPU re-upload); instance-segmentation masks follow the same index selection and are preserved. For keypoint input the bbox component of the `(KeyPoints, Detections)` tuple drives tracking and the output is that bbox `Detections` (keypoints are not carried onto the tracked output). The `new_instances` / `already_seen_instances` splits reuse the same index selection over the tracked prediction. +""" + + +class ByteTrackerBlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Byte Tracker", + "version": "v3", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "deprecated": True, + "deprecation_message": "This block is deprecated. Use the [ByteTrack Tracker](byte_track_tracker.md) block instead.", + "ui_manifest": { + "section": "video", + "icon": "far fa-location-crosshairs", + "blockPriority": 0, + }, + }, + protected_namespaces=(), + ) + type: Literal["roboflow_core/byte_tracker@v3"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Input image containing embedded video metadata (fps and video_identifier) required for ByteTrack initialization and tracking state management. The block extracts video_metadata from the WorkflowImageData object. The fps value is used to configure the tracker, and the video_identifier is used to maintain separate tracking state and instance cache for different videos. If fps is missing or invalid, the block defaults to 0 and logs a warning but continues operation. If processing multiple videos, each video should have a unique video_identifier in its metadata to maintain separate tracking states and caches. The block maintains persistent trackers and instance caches across frames for each video using the video_identifier.", + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Detection predictions (object detection, instance segmentation, or keypoint detection) for the current video frame to be tracked. The block associates these detections with existing tracks or creates new tracks. Supports object detection, instance segmentation, and keypoint detection predictions. Detections should be provided for each frame in sequence to maintain consistent tracking. The detections must include bounding boxes and class names (and keypoints if keypoint detection). After tracking, the output will include the surviving detections enhanced with a `tracker_id` entry in `bboxes_metadata`, allowing identification of the same object across frames. Masks and keypoints are dropped from the output.", + examples=["$steps.object_detection_model.predictions"], + ) + track_activation_threshold: Union[Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + default=0.25, + description="Confidence threshold for activating new tracks from detections. Must be between 0.0 and 1.0. Default is 0.25. Only detections with confidence above this threshold can create new tracks. Increasing this threshold (e.g., 0.3-0.5) improves tracking accuracy and stability by only creating tracks from high-confidence detections, but might miss true detections with lower confidence. Decreasing this threshold (e.g., 0.15-0.2) increases tracking completeness by accepting lower-confidence detections, but risks introducing noise and instability from false positives. Adjust based on detection model performance: use lower values if detections are reliable, higher values if false positives are common.", + examples=[0.25, "$inputs.confidence"], + ) + lost_track_buffer: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + default=30, + description="Number of frames to maintain a track when it's lost (no matching detections). Must be a positive integer. Default is 30 frames. When an object temporarily disappears (due to occlusion, missed detection, or leaving frame), the track is maintained for this many frames before being considered lost. Increasing this value (e.g., 50-100) enhances occlusion handling and significantly reduces track fragmentation or disappearance caused by brief detection gaps, but increases memory usage. Decreasing this value (e.g., 10-20) reduces memory usage but may cause tracks to disappear during short occlusions. Adjust based on occlusion frequency: use higher values for frequent occlusions, lower values for stable tracking scenarios.", + examples=[30, "$inputs.lost_track_buffer"], + ) + minimum_matching_threshold: Union[Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + default=0.8, + description="IoU (Intersection over Union) threshold for matching detections to existing tracks. Must be between 0.0 and 1.0. Default is 0.8. Detections are associated with tracks when their bounding box IoU exceeds this threshold. Increasing this threshold (e.g., 0.85-0.95) improves tracking accuracy by requiring stronger spatial overlap, but risks track fragmentation when objects move quickly or detection boxes vary. Decreasing this threshold (e.g., 0.6-0.75) improves tracking completeness by accepting looser matches, but risks false positive associations and track drift. Adjust based on object movement speed and detection stability: use higher values for stable objects, lower values for fast-moving objects.", + examples=[0.8, "$inputs.min_matching_threshold"], + ) + minimum_consecutive_frames: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + default=1, + description="Minimum number of consecutive frames an object must be tracked before the track is considered valid and output. Must be a positive integer. Default is 1 (all tracks are immediately valid). Only tracks that persist for at least this many consecutive frames are included in the output. Increasing this value (e.g., 3-5) prevents the creation of accidental tracks from false detections or double detections, filtering out short-lived spurious tracks, but risks missing shorter legitimate tracks. Decreasing this value (e.g., 1) includes all tracks immediately, maximizing completeness but potentially including false tracks. Adjust based on false positive rate: use higher values if false detections are common, lower values if detections are reliable.", + examples=[1, "$inputs.min_consecutive_frames"], + ) + instances_cache_size: int = Field( + default=16384, + description="Maximum number of track IDs to remember in the instance cache for determining if instances are new or already seen. Must be a positive integer. Default is 16384. The cache uses FIFO (First-In-First-Out) eviction - when the cache is full, the oldest track ID is removed to make room for new ones. Increasing this value (e.g., 32768-65536) maintains longer history of seen track IDs, allowing detection of reappearances after longer gaps, but uses more memory. Decreasing this value (e.g., 8192) reduces memory usage but may lose history of track IDs that appeared earlier, causing reappearing objects to be classified as new. Adjust based on video length and object reappearance patterns: use higher values for long videos or frequent reappearances, lower values for short videos or rare reappearances.", + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition( + name="new_instances", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + OutputDefinition( + name="already_seen_instances", + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +class ByteTrackerBlockV3(WorkflowBlock): + def __init__( + self, + ): + self._trackers: Dict[str, sv.ByteTrack] = {} + self._per_video_cache: Dict[str, "InstanceCache"] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return ByteTrackerBlockManifest + + def run( + self, + image: WorkflowImageData, + detections: Union[ + Detections, + InstanceDetections, + Tuple[KeyPoints, Optional[Detections]], + ], + track_activation_threshold: float = 0.25, + lost_track_buffer: int = 30, + minimum_matching_threshold: float = 0.8, + minimum_consecutive_frames: int = 1, + instances_cache_size: int = 16384, + ) -> BlockResult: + # For a keypoint tuple the bbox Detections drives tracking (and is the + # tracked output โ€” keypoints are not carried onto the result); for + # object-detection / instance-segmentation inputs the prediction is the + # bbox component itself (masks preserved through the index selection). + _, bbox = split_key_point_prediction(detections) + metadata = image.video_metadata + fps = metadata.fps + if not fps: + fps = 0 + logger.warning( + f"Malformed fps in VideoMetadata, {self.__class__.__name__} requires fps in order to initialize ByteTrack" + ) + if metadata.video_identifier not in self._trackers: + self._trackers[metadata.video_identifier] = sv.ByteTrack( + track_activation_threshold=track_activation_threshold, + lost_track_buffer=lost_track_buffer, + minimum_matching_threshold=minimum_matching_threshold, + minimum_consecutive_frames=minimum_consecutive_frames, + frame_rate=fps, + ) + tracker = self._trackers[metadata.video_identifier] + + kept_indices, tracker_ids = _track_and_recover_indices( + tracker=tracker, detections=bbox + ) + # ByteTrack only filters/reorders - coordinates are unchanged - so slice the + # ORIGINAL device tensors (masks included) by the surviving rows instead of + # re-uploading the numpy boxes. `take_prediction_by_indices` copies the + # surviving bboxes_metadata dicts, so the tracker_id write is leak-safe. + tracked_detections = take_prediction_by_indices(bbox, kept_indices) + _write_tracker_ids(tracked_detections, tracker_ids) + + if metadata.video_identifier not in self._per_video_cache: + self._per_video_cache[metadata.video_identifier] = InstanceCache( + size=instances_cache_size + ) + cache = self._per_video_cache[metadata.video_identifier] + not_seen_indices: List[int] = [] + seen_indices: List[int] = [] + for position, tid in enumerate(tracker_ids.tolist()): + already_seen = cache.record_instance(tracker_id=int(tid)) + if already_seen: + seen_indices.append(position) + else: + not_seen_indices.append(position) + + return { + OUTPUT_KEY: tracked_detections, + "new_instances": take_prediction_by_indices( + tracked_detections, not_seen_indices + ), + "already_seen_instances": take_prediction_by_indices( + tracked_detections, seen_indices + ), + } + + +class InstanceCache: + + def __init__(self, size: int): + size = max(1, size) + self._cache_inserts_track = deque(maxlen=size) + self._cache = set() + + def record_instance(self, tracker_id: int) -> bool: + in_cache = tracker_id in self._cache + if not in_cache: + self._cache_new_tracker_id(tracker_id=tracker_id) + return in_cache + + def _cache_new_tracker_id(self, tracker_id: int) -> None: + while len(self._cache) >= self._cache_inserts_track.maxlen: + to_drop = self._cache_inserts_track.popleft() + self._cache.remove(to_drop) + self._cache_inserts_track.append(tracker_id) + self._cache.add(tracker_id) diff --git a/inference/core/workflows/core_steps/transformations/detection_offset/v1_tensor.py b/inference/core/workflows/core_steps/transformations/detection_offset/v1_tensor.py new file mode 100644 index 0000000000..d1479c8827 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/detection_offset/v1_tensor.py @@ -0,0 +1,323 @@ +from typing import List, Literal, Optional, Type, Union +from uuid import uuid4 + +import torch +from pydantic import AliasChoices, ConfigDict, Field, PositiveInt + +from inference.core.workflows.core_steps.common.tensor_native import ( + HOST_MIRROR_KEYS, + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.execution_engine.constants import ( + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +LONG_DESCRIPTION = """ +Expand or contract detection bounding boxes by applying fixed offsets to their width and height, adding padding around detections to include more context, adjust bounding box sizes for downstream processing, or compensate for tight detections, supporting both pixel-based and percentage-based offset units for flexible bounding box adjustment. + +## How This Block Works + +This block adjusts the size of detection bounding boxes by adding offsets to their dimensions, effectively expanding or contracting the boxes to include more or less context around detected objects. The block: + +1. Receives detection predictions (object detection, instance segmentation, or keypoint detection) containing bounding boxes +2. Processes each detection's bounding box coordinates independently +3. Calculates offsets based on the selected unit type: + - **Pixel-based offsets**: Adds/subtracts a fixed number of pixels on each side (offset_width//2 pixels on left/right, offset_height//2 pixels on top/bottom) + - **Percentage-based offsets**: Calculates offsets as a percentage of the bounding box's width and height (offset_width% of box width, offset_height% of box height) +4. Applies the offsets to expand the bounding boxes: + - Subtracts half the width offset from x_min and adds half to x_max (expands horizontally) + - Subtracts half the height offset from y_min and adds half to y_max (expands vertically) +5. Clips the adjusted bounding boxes to image boundaries (ensures coordinates stay within image dimensions using min/max constraints) +6. Updates detection metadata: + - Sets parent_id_key to reference the original detection IDs (preserves traceability) + - Generates new detection IDs for the offset detections (tracks that these are modified versions) +7. Preserves all other detection properties (masks, keypoints, polygons, class labels, confidence scores) unchanged +8. Returns the modified detections with expanded or contracted bounding boxes + +The block applies offsets symmetrically around the center of each bounding box, expanding the box equally on all sides based on the width and height offsets. Positive offsets expand boxes (add padding), while the implementation always expands boxes outward. The pixel-based mode applies fixed pixel offsets regardless of box size, useful for consistent padding. The percentage-based mode applies offsets proportional to box size, useful when padding should scale with the detected object size. Boxes are automatically clipped to image boundaries to prevent invalid coordinates. + +## Common Use Cases + +- **Context Padding for Analysis**: Expand tight bounding boxes to include more surrounding context (e.g., add padding around detected objects for better classification, expand boxes to include object context for feature extraction, add margin around text detections for OCR), enabling improved analysis with additional context +- **Detection Size Adjustment**: Adjust bounding box sizes to match downstream processing requirements (e.g., expand boxes for models that need larger input regions, adjust box sizes to accommodate specific analysis needs, modify detections for compatibility with other blocks), enabling size customization for workflow compatibility +- **Tight Detection Compensation**: Expand overly tight bounding boxes that cut off parts of objects (e.g., add padding to tight object detections, expand boxes that miss object edges, compensate for models that produce undersized boxes), enabling better object coverage +- **Multi-Stage Workflow Preparation**: Prepare detections with adjusted sizes for secondary processing (e.g., expand initial detections before running secondary models, adjust box sizes for specialized analysis blocks, prepare detections with context for detailed processing), enabling optimized multi-stage workflows +- **Crop Region Optimization**: Adjust bounding boxes before cropping to include desired context (e.g., add padding before dynamic cropping to include surrounding context, expand boxes to capture more area for analysis, adjust crop regions for better feature extraction), enabling optimized region extraction +- **Visualization and Display**: Adjust bounding box sizes for better visualization or display purposes (e.g., expand boxes for clearer annotations, adjust box sizes for presentation, modify detections for visualization consistency), enabling improved visual outputs + +## Connecting to Other Blocks + +This block receives detection predictions and produces adjusted detections with modified bounding boxes: + +- **After detection blocks** (e.g., Object Detection, Instance Segmentation, Keypoint Detection) to expand or adjust bounding box sizes before further processing, enabling size-optimized detections for downstream analysis +- **Before dynamic crop blocks** to adjust bounding box sizes before cropping, enabling optimized crop regions with desired context or padding +- **Before classification or analysis blocks** that benefit from additional context around detections (e.g., classification with context, feature extraction from expanded regions, detailed analysis with padding), enabling improved analysis with context +- **In multi-stage detection workflows** where initial detections need size adjustments before secondary processing (e.g., expand initial detections before running specialized models, adjust box sizes for compatibility, prepare detections for optimized processing), enabling flexible multi-stage workflows +- **Before visualization blocks** to adjust bounding box sizes for display purposes (e.g., expand boxes for clearer annotations, adjust sizes for presentation, modify detections for visualization consistency), enabling optimized visual outputs +- **Before blocks that process detection regions** where bounding box size matters (e.g., OCR on text regions with padding, feature extraction from expanded regions, specialized models requiring specific box sizes), enabling size-optimized region processing +""" + +SHORT_DESCRIPTION = "Apply a padding around the width and height of detections." + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Detection Offset", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "transformation", + "icon": "fal fa-distribute-spacing-horizontal", + "blockPriority": 3, + }, + } + ) + type: Literal["roboflow_core/detection_offset@v1", "DetectionOffset"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + description="Detection predictions containing bounding boxes to adjust. Supports object detection, instance segmentation, or keypoint detection predictions. The bounding boxes in these predictions will be expanded or contracted based on the offset_width and offset_height values. All detection properties (masks, keypoints, polygons, classes, confidence) are preserved unchanged - only bounding box coordinates are modified.", + examples=["$steps.object_detection_model.predictions"], + ) + offset_width: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + description="Offset value to apply to bounding box width. Must be a positive integer. If units is 'Pixels', this is the number of pixels added to the box width (divided equally between left and right sides - offset_width//2 pixels on each side). If units is 'Percent (%)', this is the percentage of the bounding box width to add (calculated as percentage of the box's width, then divided between left and right). Positive values expand boxes horizontally. Boxes are clipped to image boundaries automatically.", + examples=[10, "$inputs.offset_x"], + validation_alias=AliasChoices("offset_width", "offset_x"), + ) + offset_height: Union[PositiveInt, Selector(kind=[INTEGER_KIND])] = Field( + description="Offset value to apply to bounding box height. Must be a positive integer. If units is 'Pixels', this is the number of pixels added to the box height (divided equally between top and bottom sides - offset_height//2 pixels on each side). If units is 'Percent (%)', this is the percentage of the bounding box height to add (calculated as percentage of the box's height, then divided between top and bottom). Positive values expand boxes vertically. Boxes are clipped to image boundaries automatically.", + examples=[10, "$inputs.offset_y"], + validation_alias=AliasChoices("offset_height", "offset_y"), + ) + units: Literal["Percent (%)", "Pixels"] = Field( + default="Pixels", + description="Unit type for offset values: 'Pixels' for fixed pixel offsets (same number of pixels for all boxes regardless of size) or 'Percent (%)' for percentage-based offsets (proportional to each bounding box's dimensions). Pixel offsets provide consistent padding in absolute terms. Percentage offsets scale with box size, providing proportional padding. Use pixels when you need consistent absolute padding. Use percentage when padding should scale with detected object size.", + examples=["Pixels", "Percent (%)"], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["predictions"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class DetectionOffsetBlockV1(WorkflowBlock): + # TODO: This block breaks parent coordinates. + # Issue report: https://github.com/roboflow/inference/issues/380 + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + predictions: Batch[TensorNativePrediction], + offset_width: int, + offset_height: int, + units: str = "Pixels", + ) -> BlockResult: + use_percentage = units == "Percent (%)" + return [ + { + "predictions": offset_detections( + prediction=prediction, + offset_width=offset_width, + offset_height=offset_height, + use_percentage=use_percentage, + ) + } + for prediction in predictions + ] + + +def _read_image_dimensions(detections: TensorNativeDetections) -> tuple: + """Return ``(height, width)`` from the prediction's ``image_metadata`` (one + ``[height, width]`` pair shared by every box of the prediction). Raises when + ``IMAGE_DIMENSIONS_KEY`` is absent rather than degrading to an un-clamped + offset.""" + image_metadata = detections.image_metadata or {} + image_dimensions = image_metadata.get(IMAGE_DIMENSIONS_KEY) + if image_dimensions is None: + raise ValueError( + "Detection Offset block requires image dimensions to clamp the " + f"offset boxes, but `{IMAGE_DIMENSIONS_KEY}` is missing from the " + "prediction's image_metadata." + ) + return int(image_dimensions[0]), int(image_dimensions[1]) + + +def _offset_xyxy( + xyxy: torch.Tensor, + offset_width: int, + offset_height: int, + image_height: Optional[int], + image_width: Optional[int], + use_percentage: bool, +) -> torch.Tensor: + """Expand each box outward, clamped to image bounds. Percentage mode pads by + a fraction of each box's own width/height; pixel mode pads by a fixed + (offset // 2) on each side.""" + x1 = xyxy[:, 0] + y1 = xyxy[:, 1] + x2 = xyxy[:, 2] + y2 = xyxy[:, 3] + if use_percentage: + box_width = x2 - x1 + box_height = y2 - y1 + pad_x = (box_width * offset_width / 200).to(torch.int64).to(xyxy.dtype) + pad_y = (box_height * offset_height / 200).to(torch.int64).to(xyxy.dtype) + else: + pad_x = torch.full_like(x1, float(offset_width // 2)) + pad_y = torch.full_like(y1, float(offset_height // 2)) + new_x1 = torch.clamp(x1 - pad_x, min=0) + new_y1 = torch.clamp(y1 - pad_y, min=0) + new_x2 = x2 + pad_x + new_y2 = y2 + pad_y + if image_width is not None: + new_x2 = torch.clamp(new_x2, max=image_width) + if image_height is not None: + new_y2 = torch.clamp(new_y2, max=image_height) + return torch.stack([new_x1, new_y1, new_x2, new_y2], dim=1) + + +def _rebuild_detections( + detections: TensorNativeDetections, + new_xyxy: torch.Tensor, + new_bboxes_metadata: Optional[List[dict]], +) -> TensorNativeDetections: + """Rebuild a native ``Detections`` / ``InstanceDetections`` with offset boxes; + class_id / confidence / mask / image_metadata are carried unchanged.""" + if isinstance(detections, InstanceDetections): + return InstanceDetections( + xyxy=new_xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + mask=detections.mask, + image_metadata=detections.image_metadata, + bboxes_metadata=new_bboxes_metadata, + ) + return Detections( + xyxy=new_xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + image_metadata=detections.image_metadata, + bboxes_metadata=new_bboxes_metadata, + ) + + +def _offset_bboxes_metadata( + detections: TensorNativeDetections, + number_of_detections: int, + parent_id_key: str, + detection_id_key: str, +) -> List[dict]: + """Set each box's parent_id to its prior detection_id and mint a fresh + detection_id. The per-box host mirror is dropped: the boxes are inflated by + the caller, so a carried mirror would be stale โ€” consumers fall back to + tensor reads.""" + existing = detections.bboxes_metadata or [{} for _ in range(number_of_detections)] + new_metadata = [] + for index in range(number_of_detections): + entry = dict(existing[index] or {}) if index < len(existing) else {} + entry[parent_id_key] = entry.get(detection_id_key) + entry[detection_id_key] = str(uuid4()) + for key in HOST_MIRROR_KEYS: + entry.pop(key, None) + new_metadata.append(entry) + return new_metadata + + +def offset_detections( + prediction: TensorNativePrediction, + offset_width: int, + offset_height: int, + parent_id_key: str = PARENT_ID_KEY, + detection_id_key: str = DETECTION_ID_KEY, + use_percentage: bool = False, +) -> TensorNativePrediction: + if prediction is None: + return prediction + # Only boxes are offset โ€” keypoints are left untouched and re-wrapped onto + # the output. A prediction with a missing/empty bbox component is returned + # unchanged. + if isinstance(prediction, tuple): + key_points, detections = prediction + else: + key_points, detections = None, prediction + if ( + detections is None + or detections.xyxy is None + or int(detections.xyxy.shape[0]) == 0 + ): + return prediction + number_of_detections = int(detections.xyxy.shape[0]) + image_height, image_width = _read_image_dimensions(detections) + new_xyxy = _offset_xyxy( + xyxy=detections.xyxy, + offset_width=offset_width, + offset_height=offset_height, + image_height=image_height, + image_width=image_width, + use_percentage=use_percentage, + ) + new_bboxes_metadata = _offset_bboxes_metadata( + detections=detections, + number_of_detections=number_of_detections, + parent_id_key=parent_id_key, + detection_id_key=detection_id_key, + ) + new_detections = _rebuild_detections( + detections=detections, + new_xyxy=new_xyxy, + new_bboxes_metadata=new_bboxes_metadata, + ) + if key_points is not None: + return key_points, new_detections + return new_detections diff --git a/inference/core/workflows/core_steps/transformations/detections_combine/v1_tensor.py b/inference/core/workflows/core_steps/transformations/detections_combine/v1_tensor.py new file mode 100644 index 0000000000..03e9ba709d --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/detections_combine/v1_tensor.py @@ -0,0 +1,343 @@ +from typing import Dict, List, Literal, Optional, Type, Union +from uuid import uuid4 + +import numpy as np +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.tensor_native import ( + instance_mask_to_numpy, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import Selector +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks + +LONG_DESCRIPTION = """ +Combine two sets of detection predictions into a single unified set of detections by merging both detection sets together, preserving all detections from both inputs for multi-source detection aggregation, combining results from multiple models, and consolidating detection sets from different processing stages into one workflow output. + +## How This Block Works + +This block combines two separate sets of detection predictions into a single unified detection set by merging all detections from both inputs. The block: + +1. Receives two separate detection prediction sets (prediction_one and prediction_two), each containing multiple detections from object detection or instance segmentation models +2. Processes both detection sets independently (each set maintains its own detections, properties, masks, and metadata) +3. Merges the two detection sets using supervision's Detections.merge() method: + - Combines all detections from prediction_one with all detections from prediction_two + - Preserves all detection properties from both sets (bounding boxes, masks, classes, confidence scores, metadata) + - Maintains detection order (typically prediction_one detections followed by prediction_two detections) + - Handles all detection attributes including masks (for instance segmentation), keypoints, class IDs, class names, confidence scores, and custom data fields +4. Returns a single unified detection set containing all detections from both inputs + +The block simply concatenates the two detection sets together, preserving all detections and their properties from both sources. Unlike the Detections Merge block (which creates a union bounding box from multiple detections), this block maintains all individual detections from both sets in the output. This is useful for combining detections from different models, different processing stages, or different detection sources into a single workflow stream for unified downstream processing. + +## Common Use Cases + +- **Multi-Model Detection Aggregation**: Combine detections from multiple detection models into a single unified set (e.g., combine detections from different object detection models, merge results from specialized models, aggregate detections from multiple model outputs), enabling multi-model detection workflows +- **Multi-Stage Detection Combination**: Combine detections from different processing stages or workflow branches (e.g., merge detections from different workflow paths, combine initial detections with refined detections, aggregate detections from multiple processing stages), enabling multi-stage detection aggregation +- **Detection Source Consolidation**: Consolidate detections from different sources or inputs into one set (e.g., combine detections from multiple images or frames, merge detections from different regions, aggregate detections from various sources), enabling detection source unification +- **Classification and Detection Combination**: Combine object detection results with classification results or other detection types (e.g., merge object detections with classification outputs, combine different detection types, aggregate complementary detection sets), enabling multi-type detection workflows +- **Filtered and Unfiltered Detection Combination**: Combine filtered detections with unfiltered detections or combine different filtered subsets (e.g., merge filtered detections by different criteria, combine specific class detections with general detections, aggregate different filtered detection sets), enabling flexible detection combination workflows +- **Workflow Branch Merging**: Merge detection results from different workflow branches back into a single detection stream (e.g., combine parallel processing branch results, merge conditional workflow paths, aggregate branch detection outputs), enabling workflow branch consolidation + +## Connecting to Other Blocks + +This block receives two detection prediction sets and produces a single combined detection set: + +- **After multiple detection blocks** to combine detections from different models into one unified set (e.g., combine detections from multiple object detection models, merge results from different segmentation models, aggregate detections from various model outputs), enabling multi-model detection aggregation workflows +- **After filtering blocks** to combine filtered detection subsets (e.g., merge detections filtered by different criteria, combine class-specific filtered detections, aggregate various filtered detection sets), enabling filtered detection combination workflows +- **At workflow merge points** where different workflow branches need to be combined (e.g., merge parallel processing branch results, combine conditional path outputs, aggregate branch detection streams), enabling workflow branch merging workflows +- **Before downstream processing blocks** that need unified detection sets (e.g., process combined detections together, visualize unified detection sets, analyze aggregated detections), enabling unified detection processing workflows +- **Before crop blocks** to process combined detections together (e.g., crop regions from combined detection sets, extract areas from aggregated detections, process unified detection regions), enabling combined detection region extraction +- **Before visualization blocks** to display unified detection sets (e.g., visualize combined detections from multiple sources, display aggregated detection results, show merged detection outputs), enabling unified detection visualization workflows +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Detections Combine", + "version": "v1", + "short_description": "Combines two sets of predictions into a single prediction.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "transformation", + "icon": "fal fa-object-union", + "blockPriority": 5, + }, + } + ) + type: Literal["roboflow_core/detections_combine@v1"] + prediction_one: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( + description="First set of detection predictions to combine. Supports object detection or instance segmentation predictions. All detections from this set will be included in the output. Detection properties (bounding boxes, masks, classes, confidence scores, metadata) are preserved as-is. This set is combined with prediction_two to create the unified output. Detections from this set typically appear first in the merged output.", + examples=["$steps.my_object_detection_model.predictions"], + ) + prediction_two: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( + description="Second set of detection predictions to combine. Supports object detection or instance segmentation predictions. All detections from this set will be included in the output. Detection properties (bounding boxes, masks, classes, confidence scores, metadata) are preserved as-is. This set is combined with prediction_one to create the unified output. Detections from this set are merged with detections from prediction_one to form a single combined detection set.", + examples=["$steps.my_object_detection_model.predictions"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +TensorNativeDetections = Union[Detections, InstanceDetections] + + +def _merged_class_names( + prediction_one: TensorNativeDetections, + prediction_two: TensorNativeDetections, +) -> Dict[int, str]: + """Merge the two ``image_metadata[CLASS_NAMES_KEY]`` maps into one combined + ``{class_id: name}`` map. + + Native predictions carry class names once per image (not per detection), so a + combine has to reconcile the two maps. ``sv.Detections.merge`` keeps each + detection's own ``class_name`` (it lives per-row in ``data``) and never re-maps + ``class_id`` โ€” it just ``np.hstack``-es the two id arrays. We mirror that by + keeping every id unchanged and unioning the two maps. On an id collision with + conflicting names the native single-map representation can only hold one name, + so ``prediction_one`` wins (its detections come first in the output and "own" + the id), matching the input ordering ``sv.Detections.merge`` preserves. + """ + names_two = (prediction_two.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + names_one = (prediction_one.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + return {**names_two, **names_one} + + +def _combine_dense_masks( + prediction_one: InstanceDetections, + prediction_two: InstanceDetections, + device: torch.device, +) -> torch.Tensor: + """Concatenate two dense / RLE mask stacks into one dense ``(n, H, W)`` torch + tensor on ``device``. Each instance is materialised to numpy (decoding RLE one + row at a time) and stacked, so dense and RLE inputs combine uniformly. The + common ``(H, W)`` is taken from whichever input carries instances.""" + mask_rows: List[np.ndarray] = [] + for prediction in (prediction_one, prediction_two): + for index in range(len(prediction)): + mask_rows.append(instance_mask_to_numpy(prediction, index)) + if not mask_rows: + return torch.zeros((0, 0, 0), dtype=torch.bool, device=device) + return torch.from_numpy(np.stack(mask_rows, axis=0)).to( + device=device, dtype=torch.bool + ) + + +def _combine_rle_masks( + prediction_one: InstanceDetections, + prediction_two: InstanceDetections, +) -> InstancesRLEMasks: + """Concatenate two ``InstancesRLEMasks`` lists into one. Both inputs must share + the same ``image_size`` (they are predictions over the same image); the encoded + RLE byte strings are simply concatenated, preserving detection order.""" + mask_one: InstancesRLEMasks = prediction_one.mask + mask_two: InstancesRLEMasks = prediction_two.mask + if tuple(mask_one.image_size) != tuple(mask_two.image_size): + raise ValueError( + "Cannot combine instance-segmentation predictions with RLE masks of " + f"different image sizes ({mask_one.image_size} vs " + f"{mask_two.image_size}); both predictions must be made against the " + "same input image." + ) + return InstancesRLEMasks( + image_size=mask_one.image_size, + masks=list(mask_one.masks) + list(mask_two.masks), + ) + + +def _combine_masks( + prediction_one: InstanceDetections, + prediction_two: InstanceDetections, + device: torch.device, +) -> Union[torch.Tensor, InstancesRLEMasks]: + """Combine the two instance masks, keeping the storage shape when it agrees: + RLE + RLE stays RLE (byte lists concatenated); dense + dense is concatenated + on-device with ``torch.cat`` (no host round-trip); any remaining RLE/dense mix + is materialised to a dense ``(n, H, W)`` torch tensor (mirroring bounding_rect's + RLE-in/RLE-out, dense-in/dense-out convention but for two inputs).""" + mask_one = prediction_one.mask + mask_two = prediction_two.mask + both_rle = isinstance(mask_one, InstancesRLEMasks) and isinstance( + mask_two, InstancesRLEMasks + ) + if both_rle: + return _combine_rle_masks(prediction_one, prediction_two) + both_dense = isinstance(mask_one, torch.Tensor) and isinstance( + mask_two, torch.Tensor + ) + if both_dense: + # Fast path: keep dense masks on-device, no device->host->device hop. + return torch.cat([mask_one.to(device), mask_two.to(device)], dim=0) + # RLE/dense mix: fall back to the per-row numpy materialise path. + return _combine_dense_masks(prediction_one, prediction_two, device=device) + + +def _combine_bboxes_metadata( + prediction_one: TensorNativeDetections, + prediction_two: TensorNativeDetections, +) -> Optional[List[dict]]: + """Concatenate the two ``bboxes_metadata`` lists, preserving each box's own + ``detection_id``. A missing list is padded with empty dicts so the result lines + up row-for-row with the combined tensors; a per-box ``detection_id`` is filled + in when absent (the tensor-native serialiser requires one per row).""" + if ( + prediction_one.bboxes_metadata is None + and prediction_two.bboxes_metadata is None + ): + return None + combined: List[dict] = [] + for prediction in (prediction_one, prediction_two): + per_box = prediction.bboxes_metadata + if per_box is None: + per_box = [{} for _ in range(len(prediction))] + for entry in per_box: + entry = dict(entry or {}) + entry.setdefault(DETECTION_ID_KEY, str(uuid4())) + combined.append(entry) + return combined + + +def _combine_image_metadata( + prediction_one: TensorNativeDetections, + prediction_two: TensorNativeDetections, +) -> Optional[dict]: + """Carry over the per-image lineage (parent/root coordinates, dimensions, ...) + from the first prediction that has it and overwrite its ``class_names`` with the + merged ``{class_id: name}`` map so every combined detection's id resolves.""" + base = prediction_one.image_metadata or prediction_two.image_metadata + if base is None: + return None + image_metadata = dict(base) + image_metadata[CLASS_NAMES_KEY] = _merged_class_names( + prediction_one=prediction_one, + prediction_two=prediction_two, + ) + return image_metadata + + +class DetectionsCombineBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + prediction_one: TensorNativeDetections, + prediction_two: TensorNativeDetections, + ) -> BlockResult: + # Native counterpart of sv.Detections.merge([p1, p2]): concatenate the two + # predictions' xyxy / class_id / confidence (torch.cat instead of + # np.vstack/np.hstack), reconcile the class-name maps, and concatenate the + # per-box metadata. The device is inherited from the inputs (first non-empty + # row), falling back to WORKFLOWS_IMAGE_TENSOR_DEVICE for the all-empty case + # โ€” matching bounding_rect/detections_merge. + device = WORKFLOWS_IMAGE_TENSOR_DEVICE + for prediction in (prediction_one, prediction_two): + if len(prediction) > 0: + device = prediction.xyxy.device + break + + xyxy = torch.cat( + [prediction_one.xyxy.to(device), prediction_two.xyxy.to(device)], dim=0 + ) + class_id = torch.cat( + [prediction_one.class_id.to(device), prediction_two.class_id.to(device)], + dim=0, + ) + confidence = torch.cat( + [ + prediction_one.confidence.to(device), + prediction_two.confidence.to(device), + ], + dim=0, + ) + image_metadata = _combine_image_metadata( + prediction_one=prediction_one, + prediction_two=prediction_two, + ) + bboxes_metadata = _combine_bboxes_metadata( + prediction_one=prediction_one, + prediction_two=prediction_two, + ) + + one_has_masks = isinstance(prediction_one, InstanceDetections) + two_has_masks = isinstance(prediction_two, InstanceDetections) + if one_has_masks != two_has_masks: + # Match numpy sv.Detections.merge: stack_or_none("mask") raises when + # some-but-not-all inputs carry masks ("All or none of the 'mask' + # fields must be None"). Mirror that here instead of silently dropping + # the instance-segmentation input's masks down the OD output path. + raise ValueError( + "Cannot combine an object-detection prediction with an " + "instance-segmentation prediction: all or none of the combined " + "predictions must carry masks (matching sv.Detections.merge)." + ) + has_masks = one_has_masks and two_has_masks + if has_masks: + mask = _combine_masks( + prediction_one=prediction_one, + prediction_two=prediction_two, + device=device, + ) + return { + "predictions": InstanceDetections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + } + return { + "predictions": Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + } diff --git a/inference/core/workflows/core_steps/transformations/detections_merge/v1_tensor.py b/inference/core/workflows/core_steps/transformations/detections_merge/v1_tensor.py new file mode 100644 index 0000000000..027ec777d8 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/detections_merge/v1_tensor.py @@ -0,0 +1,265 @@ +from typing import List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, + SCALING_RELATIVE_TO_PARENT_KEY, + SCALING_RELATIVE_TO_ROOT_PARENT_KEY, +) +from inference.core.workflows.execution_engine.entities.base import OutputDefinition +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import Selector +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "predictions" + +SHORT_DESCRIPTION = "Merge multiple detections into a single bounding box." +LONG_DESCRIPTION = """ +Combine multiple detection predictions into a single merged detection with a union bounding box that encompasses all input detections, simplifying multiple detections into one larger detection region for overlapping object consolidation, region creation from multiple objects, and detection simplification workflows. + +## How This Block Works + +This block merges multiple detections into a single detection by calculating a union bounding box that contains all input detections. The block: + +1. Receives detection predictions (object detection, instance segmentation, or keypoint detection) containing multiple detections +2. Validates input (handles empty detections by returning an empty detection result) +3. Calculates the union bounding box from all input detections: + - Extracts all bounding box coordinates (xyxy format) from input detections + - Finds the minimum x and y coordinates (leftmost and topmost points) across all boxes + - Finds the maximum x and y coordinates (rightmost and bottommost points) across all boxes + - Creates a single bounding box that completely encompasses all input detections +4. Determines the merged detection's confidence: + - Finds the detection with the lowest confidence score among all input detections + - Uses this lowest confidence as the merged detection's confidence (conservative approach) + - Handles cases where confidence scores may not be present +5. Creates a new merged detection with: + - The calculated union bounding box (encompasses all input detections) + - A customizable class name (default: "merged_detection", configurable via class_name parameter) + - The lowest confidence from input detections (conservative confidence assignment) + - A fixed class_id of 0 for the merged detection + - A newly generated detection ID (unique identifier for the merged detection) +6. Returns the single merged detection containing all input detections within its bounding box + +The block creates a unified bounding box representation of multiple detections, useful for consolidating overlapping or nearby detections into a single region. The union bounding box approach ensures all original detections are completely contained within the merged detection. By using the lowest confidence, the block adopts a conservative approach, ensuring the merged detection's confidence reflects the least certain input detection. The merged detection can be customized with a class name to indicate its merged nature or to represent a specific category. + +## Common Use Cases + +- **Overlapping Detection Consolidation**: Merge multiple overlapping detections of the same or related objects into a single unified detection (e.g., merge overlapping detections of the same person from multiple frames, consolidate duplicate detections from different models, combine overlapping object parts into one detection), enabling overlapping detection simplification +- **Multi-Object Region Creation**: Create a single bounding box region that encompasses multiple detected objects for area-based analysis (e.g., create a region containing multiple people for crowd analysis, merge detections of objects in a scene into one region, combine multiple detections into a single monitoring zone), enabling multi-object region workflows +- **Nearby Detection Grouping**: Group nearby detections together into a single merged detection (e.g., merge detections of objects close to each other, group nearby detections into clusters, combine adjacent detections for simplified processing), enabling spatial grouping workflows +- **Detection Simplification**: Simplify multiple detections into one larger detection for downstream processing (e.g., reduce multiple detections to one for simpler analysis, consolidate detections for easier visualization, merge detections for streamlined workflows), enabling detection simplification workflows +- **Zone Definition from Detections**: Create zone boundaries from multiple detection locations (e.g., define zones based on detection locations, create regions from detected object positions, establish boundaries from detection clusters), enabling zone creation from detections +- **Redundant Detection Removal**: Merge redundant or duplicate detections into a single representation (e.g., combine duplicate detections from different stages, merge redundant object detections, consolidate repeated detections), enabling redundant detection consolidation workflows + +## Connecting to Other Blocks + +This block receives multiple detection predictions and produces a single merged detection: + +- **After detection blocks** (e.g., Object Detection, Instance Segmentation, Keypoint Detection) to merge multiple detections into one unified detection for simplified processing, enabling detection consolidation workflows +- **After filtering blocks** (e.g., Detections Filter) to merge filtered detections that meet specific criteria into a single detection (e.g., merge filtered detections by class, combine detections after filtering, consolidate filtered results), enabling filtered detection consolidation +- **Before crop blocks** to create a single crop region from multiple detections (e.g., crop a region containing multiple objects, extract area encompassing multiple detections, create unified crop region), enabling multi-detection region extraction +- **Before zone-based blocks** (e.g., Polygon Zone, Dynamic Zone) to define zones based on merged detection regions (e.g., create zones from merged detection areas, establish monitoring zones from merged detections, define regions from consolidated detections), enabling zone creation from merged detections +- **Before visualization blocks** to display simplified merged detections instead of multiple individual detections (e.g., visualize consolidated detection regions, display merged bounding boxes, show simplified detection representation), enabling simplified visualization outputs +- **Before analysis blocks** that benefit from simplified detection representation (e.g., analyze merged detection regions, process consolidated detections, work with simplified detection data), enabling simplified detection analysis workflows +""" + +# Native tensor-data input/output shapes. The merge always collapses to a single +# object-detection bbox (no mask, no keypoints) regardless of the input shape, so +# the output kind is object-detection. The keypoint-detection input arrives as a +# Tuple[KeyPoints, Optional[Detections]]; its bounding-box component supplies xyxy. +TensorNativeDetections = Union[Detections, InstanceDetections] +KeyPointPrediction = Tuple[KeyPoints, Optional[Detections]] +TensorNativeMergeInput = Union[Detections, InstanceDetections, KeyPointPrediction] + + +class DetectionsMergeManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Detections Merge", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "transformation", + "icon": "fal fa-object-union", + "blockPriority": 5, + }, + } + ) + type: Literal["roboflow_core/detections_merge@v1"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + description="Detection predictions containing multiple detections to merge into a single detection. Supports object detection, instance segmentation, or keypoint detection predictions. All input detections will be combined into one merged detection with a union bounding box that encompasses all input detections. If empty detections are provided, the block returns an empty detection result. The merged detection will contain all input detections within its bounding box boundaries.", + examples=["$steps.object_detection_model.predictions"], + ) + class_name: str = Field( + default="merged_detection", + description="Class name to assign to the merged detection. The merged detection will use this class name in its data. Default is 'merged_detection' to indicate that this is a merged detection. You can customize this to represent a specific category or to indicate the purpose of the merged detection (e.g., 'crowd', 'group', 'region'). This class name will be stored in the detection's data dictionary.", + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +def _extract_bbox_detections( + predictions: TensorNativeMergeInput, +) -> Optional[TensorNativeDetections]: + """Return the bounding-box ``Detections`` carrier for any supported native + input. The merge only ever looks at xyxy / confidence, so the keypoint tuple + is reduced to its ``Detections`` component (which may be ``None`` / empty).""" + if isinstance(predictions, tuple): + _, detections = predictions + return detections + return predictions + + +def _build_merged_image_metadata( + source_metadata: Optional[dict], + class_id: int, + class_name: str, +) -> dict: + """Build the producer ``image_metadata`` for the merged single-row detection. + + The merged row is a fresh object-detection prediction, so ``class_names`` is + overridden to ``{class_id: class_name}`` and ``prediction_type`` becomes + ``object-detection``. Parent / root lineage and image dimensions are carried + over from the source prediction (shared across the same input image), mirroring + the convention used by the tensor-native detections-consensus block. + """ + source_metadata = source_metadata or {} + image_metadata = { + CLASS_NAMES_KEY: {int(class_id): class_name}, + PARENT_ID_KEY: source_metadata.get(PARENT_ID_KEY), + PREDICTION_TYPE_KEY: "object-detection", + PARENT_COORDINATES_KEY: source_metadata.get(PARENT_COORDINATES_KEY), + PARENT_DIMENSIONS_KEY: source_metadata.get(PARENT_DIMENSIONS_KEY), + ROOT_PARENT_ID_KEY: source_metadata.get(ROOT_PARENT_ID_KEY), + ROOT_PARENT_COORDINATES_KEY: source_metadata.get(ROOT_PARENT_COORDINATES_KEY), + ROOT_PARENT_DIMENSIONS_KEY: source_metadata.get(ROOT_PARENT_DIMENSIONS_KEY), + IMAGE_DIMENSIONS_KEY: source_metadata.get(IMAGE_DIMENSIONS_KEY), + } + image_metadata[SCALING_RELATIVE_TO_PARENT_KEY] = source_metadata.get( + SCALING_RELATIVE_TO_PARENT_KEY, 1.0 + ) + image_metadata[SCALING_RELATIVE_TO_ROOT_PARENT_KEY] = source_metadata.get( + SCALING_RELATIVE_TO_ROOT_PARENT_KEY, 1.0 + ) + return image_metadata + + +def _empty_detections(class_name: str) -> Detections: + """Empty merged result โ€” built on WORKFLOWS_IMAGE_TENSOR_DEVICE because there + is no source row whose device we could inherit.""" + return Detections( + xyxy=torch.zeros( + (0, 4), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + class_id=torch.zeros( + (0,), dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.zeros( + (0,), dtype=torch.float32, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + image_metadata=_build_merged_image_metadata( + source_metadata=None, class_id=0, class_name=class_name + ), + bboxes_metadata=None, + ) + + +class DetectionsMergeBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return DetectionsMergeManifest + + def run( + self, + predictions: TensorNativeMergeInput, + class_name: str = "merged_detection", + ) -> BlockResult: + detections = _extract_bbox_detections(predictions) + if ( + detections is None + or detections.xyxy is None + or int(detections.xyxy.shape[0]) == 0 + ): + return {OUTPUT_KEY: _empty_detections(class_name=class_name)} + + # Single device policy for the merged output: pin to + # WORKFLOWS_IMAGE_TENSOR_DEVICE so the non-empty path matches + # ``_empty_detections`` and the producer contract, instead of inheriting + # the (possibly divergent) input prediction device. + device = WORKFLOWS_IMAGE_TENSOR_DEVICE + + # Union bounding box: leftmost/topmost via torch.min over column 0/1, and + # rightmost/bottommost via torch.max over column 2/3 of the xyxy rows. + x1 = torch.min(detections.xyxy[:, 0]) + y1 = torch.min(detections.xyxy[:, 1]) + x2 = torch.max(detections.xyxy[:, 2]) + y2 = torch.max(detections.xyxy[:, 3]) + union_bbox = torch.stack([x1, y1, x2, y2]).reshape(1, 4).to(torch.float32) + + # Conservative confidence: the lowest-confidence input row's score. + if detections.confidence is not None: + lowest_conf_idx = int(torch.argmin(detections.confidence)) + confidence = ( + detections.confidence[lowest_conf_idx] + .reshape(1) + .to(device=device, dtype=torch.float32) + ) + else: + confidence = None + + merged_detection = Detections( + xyxy=union_bbox.to(device=device), + class_id=torch.zeros((1,), dtype=torch.long, device=device), + confidence=confidence, + image_metadata=_build_merged_image_metadata( + source_metadata=detections.image_metadata, + class_id=0, + class_name=class_name, + ), + bboxes_metadata=[{DETECTION_ID_KEY: str(uuid4())}], + ) + return {OUTPUT_KEY: merged_detection} diff --git a/inference/core/workflows/core_steps/transformations/detections_transformation/v1.py b/inference/core/workflows/core_steps/transformations/detections_transformation/v1.py index 7de7a6cd74..1d655bb87a 100644 --- a/inference/core/workflows/core_steps/transformations/detections_transformation/v1.py +++ b/inference/core/workflows/core_steps/transformations/detections_transformation/v1.py @@ -4,6 +4,7 @@ import supervision as sv from pydantic import ConfigDict, Field +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.common.query_language.entities.operations import ( DEFAULT_OPERAND_NAME, AllOperationsType, @@ -31,6 +32,9 @@ WorkflowBlock, WorkflowBlockManifest, ) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections LONG_DESCRIPTION = """ Apply customizable transformations to detection predictions using UQL (Query Language) operation chains, enabling flexible modification of bounding boxes, filtering detections, extracting properties, resizing boxes, and other detection manipulations through configurable operation sequences for advanced detection processing workflows. @@ -227,7 +231,7 @@ def execute_transformation( detections, global_parameters=single_evaluation_parameters, ) - if not isinstance(transformed_detections, sv.Detections): + if not _is_valid_transformation_output(transformed_detections): raise ValueError( "Definition of operation chain provided to `DetectionsTransformation` block " f"transforms sv.Detections into different type: {type(transformed_detections)} " @@ -235,3 +239,30 @@ def execute_transformation( ) results.append({"predictions": transformed_detections}) return results + + +def _is_native_key_point_prediction(value: Any) -> bool: + """True for the tensor-native keypoint-detection kind, carried as a 2-tuple + ``(KeyPoints, Optional[Detections])``.""" + return ( + isinstance(value, tuple) + and len(value) == 2 + and isinstance(value[0], KeyPoints) + and (value[1] is None or isinstance(value[1], (Detections, InstanceDetections))) + ) + + +def _is_valid_transformation_output(value: Any) -> bool: + """The operation chain must preserve the detection representation. With the + numpy representation that means ``sv.Detections``; under + ``ENABLE_TENSOR_DATA_REPRESENTATION`` the tensor-native dataclasses + (``Detections`` / ``InstanceDetections`` / a ``(KeyPoints, Detections)`` tuple) + are accepted as well. The numpy-only behaviour is preserved byte-for-byte when + the flag is off.""" + if isinstance(value, sv.Detections): + return True + if ENABLE_TENSOR_DATA_REPRESENTATION: + return isinstance( + value, (Detections, InstanceDetections) + ) or _is_native_key_point_prediction(value) + return False diff --git a/inference/core/workflows/core_steps/transformations/dynamic_crop/v1_tensor.py b/inference/core/workflows/core_steps/transformations/dynamic_crop/v1_tensor.py new file mode 100644 index 0000000000..ccab82bdf1 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/dynamic_crop/v1_tensor.py @@ -0,0 +1,603 @@ +"""Tensor-native sibling of ``roboflow_core/dynamic_crop@v1``. + +Consumes a native detection prediction (``Detections`` / ``InstanceDetections`` / +``Tuple[KeyPoints, Optional[Detections]]``) and produces one crop per detection +plus the same-shaped prediction translated into the crop's coordinate frame: +``xyxy``, masks (cropped to the box), ``KeyPoints.xy`` and the flattened +``keypoints_xy`` / ``polygon`` / OBB metadata entries are all offset by +``(-x_min, -y_min)``. Tensor-backed crops stay on-device via +``WorkflowImageData.create_crop_from_tensor``. +""" + +from copy import deepcopy +from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import torch +from pydantic import AliasChoices, ConfigDict, Field +from supervision.config import ORIENTED_BOX_COORDINATES + +from inference.core.workflows.core_steps.common.tensor_native import ( + HOST_MIRROR_KEYS, + take_prediction_by_indices, +) +from inference.core.workflows.execution_engine.constants import ( + DETECTION_ID_KEY, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + POLYGON_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + RGB_COLOR_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import ( + coco_rle_masks_to_numpy_mask, + torch_mask_to_coco_rle, +) + +TensorNativeDetections = Union[Detections, InstanceDetections] +KeyPointPrediction = Tuple[KeyPoints, Optional[Detections]] +TensorNativePrediction = Union[Detections, InstanceDetections, KeyPointPrediction] + +LONG_DESCRIPTION = """ +Extract cropped image regions from input images based on bounding boxes from detection model predictions, supporting object detection, instance segmentation, and keypoint detection models with optional background removal using segmentation masks for focused region extraction and multi-stage analysis workflows. + +## How This Block Works + +This block crops rectangular regions from input images using bounding boxes from detection model outputs, producing individual cropped images for each detected object. The block: + +1. Receives input images and detection predictions (object detection, instance segmentation, or keypoint detection) containing bounding boxes +2. Validates that predictions contain detection IDs required for crop tracking +3. Extracts each bounding box from the predictions and crops the corresponding rectangular region from the input image +4. For instance segmentation predictions with `mask_opacity > 0`: Applies background removal by overlaying the segmentation mask, replacing background pixels outside the detected instance with the specified `background_color` and blending with the original crop based on mask opacity +5. Creates cropped image objects with metadata tracking the crop's origin (original image, offset coordinates, detection ID) +6. Translates prediction coordinates from the original image space to the cropped region space (adjusts bounding boxes, masks, keypoints, and polygons to be relative to the crop origin) +7. Returns a list of results for each detection, each containing the cropped image and the translated predictions + +The block processes each detection's bounding box independently, creating separate crops for each detected object. For instance segmentation predictions, the optional background removal feature uses the segmentation mask to isolate the detected object from background pixels, useful for creating clean object-focused crops. All prediction coordinates (bounding boxes, keypoints, polygons, mask coordinates) are automatically translated to be relative to the cropped region's top-left corner, ensuring downstream blocks can process the crops correctly. The block increases output dimensionality by one (produces a list of crops per input image), enabling batch processing workflows where each crop can be processed independently. + +## Common Use Cases + +- **Multi-Stage Object Analysis**: Extract individual object crops from full images for detailed analysis (e.g., detect objects in a scene, crop each detected object, then run OCR or classification on individual crops), enabling focused analysis of specific regions without processing entire images +- **Background Removal for Object Focus**: Create clean object crops with background removed using segmentation masks (e.g., detect and segment objects, crop with background removal, create isolated object images for training or analysis), enabling focused object extraction and cleaner downstream processing +- **Region-Based Processing Pipelines**: Extract regions of interest for specialized processing (e.g., detect text regions, crop each text region, run OCR on crops; detect faces, crop each face, run face recognition), enabling efficient processing of specific image regions +- **Keypoint and Annotation Preservation**: Extract object crops while preserving detection annotations (e.g., detect objects with keypoints, crop objects maintaining keypoint coordinates, analyze keypoints in cropped context), enabling focused analysis with full annotation context +- **Batch Region Extraction**: Extract multiple regions from single images for parallel processing (e.g., detect all objects in image, crop each object separately, process crops in parallel for classification or analysis), enabling efficient batch processing of multiple regions +- **Training Data Preparation**: Create cropped object datasets from annotated images (e.g., detect objects with bounding boxes, crop each object individually, export crops for training data collection), enabling automated extraction of training samples from full images + +## Connecting to Other Blocks + +This block receives images and detection predictions, producing cropped images: + +- **After detection blocks** (e.g., Object Detection, Instance Segmentation, Keypoint Detection) to extract individual object regions based on detected bounding boxes, enabling focused analysis of detected objects in isolation +- **Before classification or analysis blocks** that need object-focused inputs (e.g., OCR for text regions, fine-grained classification for cropped objects, detailed feature extraction), enabling specialized processing of individual regions rather than full images +- **In multi-stage detection workflows** where initial detections are used to extract regions for secondary analysis (e.g., detect vehicles, crop each vehicle, detect license plates in crops), enabling hierarchical detection and analysis pipelines +- **Before visualization blocks** that display individual objects (e.g., display cropped objects separately, create galleries of detected objects, show isolated object annotations), enabling focused visualization of extracted regions +- **After detection blocks with instance segmentation** to create clean object crops with background removal, enabling isolated object images for analysis, training, or presentation +- **In keypoint detection workflows** where keypoint coordinates need to be preserved in cropped contexts (e.g., detect people with keypoints, crop each person, analyze pose in cropped images), enabling keypoint analysis in focused image regions +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Dynamic Crop", + "version": "v1", + "short_description": "Crop an image using bounding boxes from a detection model.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "transformation", + "icon": "far fa-crop-alt", + "blockPriority": 0, + "popular": True, + }, + } + ) + type: Literal["roboflow_core/dynamic_crop@v1", "DynamicCrop", "Crop"] + images: Selector(kind=[IMAGE_KIND]) = Field( + title="Image to Crop", + description="Input image(s) to extract cropped regions from. Can be a single image or batch of images. Each image will be processed with corresponding detection predictions to extract bounding box regions. Cropped regions are extracted based on bounding boxes in the predictions. Can also accept previously cropped images from another Dynamic Crop step for nested cropping workflows.", + examples=["$inputs.image", "$steps.cropping.crops"], + validation_alias=AliasChoices("images", "image"), + ) + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + title="Regions of Interest", + description="Detection model predictions containing bounding boxes that define regions to crop from the images. Supports object detection (bounding boxes), instance segmentation (bounding boxes with segmentation masks), or keypoint detection (bounding boxes with keypoints) predictions. Each bounding box in the predictions defines a rectangular region to extract. Predictions must include detection IDs for crop tracking. Multiple detections per image result in multiple crops per image.", + examples=["$steps.my_object_detection_model.predictions"], + validation_alias=AliasChoices("predictions", "detections"), + ) + mask_opacity: Union[ + Selector(kind=[FLOAT_ZERO_TO_ONE_KIND]), + float, + ] = Field( + default=0.0, + le=1.0, + ge=0.0, + description="Background removal opacity for instance segmentation crops (0.0 to 1.0). Only applies when predictions contain segmentation masks (instance segmentation predictions). Controls how aggressively background pixels outside the detected instance are removed: 0.0 leaves the crop unchanged (no background removal), 1.0 fully replaces background with background_color, values between blend the original crop with the background. Higher values create cleaner object-focused crops. Set to 0.0 to disable background removal. Requires instance segmentation predictions with masks.", + json_schema_extra={ + "relevant_for": { + "predictions": { + "kind": [TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND.name], + "required": True, + }, + } + }, + ) + background_color: Union[ + Selector(kind=[STRING_KIND]), + Selector(kind=[RGB_COLOR_KIND]), + str, + Tuple[int, int, int], + ] = Field( + default=(0, 0, 0), + description="Background color to use when removing background from instance segmentation crops. Only applies when mask_opacity > 0 and predictions contain segmentation masks. Background pixels outside the detected instance mask are replaced with this color. Can be specified as: hex string (e.g., '#431112' or '#fff'), RGB string in parentheses (e.g., '(128, 32, 64)'), or RGB tuple (e.g., (18, 17, 67)). Defaults to black (0, 0, 0). Use white (255, 255, 255) or '#ffffff' for white backgrounds, or match your use case's background requirements. Color values are interpreted as RGB and converted to BGR for image processing.", + examples=["#431112", "$inputs.bg_color", (18, 17, 67)], + json_schema_extra={ + "relevant_for": { + "predictions": { + "kind": [TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND.name], + "required": True, + }, + } + }, + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images", "predictions"] + + @classmethod + def get_output_dimensionality_offset(cls) -> int: + return 1 + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="crops", kind=[IMAGE_KIND]), + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class DynamicCropBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + predictions: Batch[TensorNativePrediction], + mask_opacity: float, + background_color: Union[str, Tuple[int, int, int]], + ) -> BlockResult: + return [ + crop_image( + image=image, + predictions=detections, + mask_opacity=mask_opacity, + background_color=background_color, + ) + for image, detections in zip(images, predictions) + ] + + +def crop_image( + image: WorkflowImageData, + predictions: TensorNativePrediction, + mask_opacity: float, + background_color: Union[str, Tuple[int, int, int]], + detection_id_key: str = DETECTION_ID_KEY, +) -> List[Dict[str, Any]]: + bbox_detections = _bbox_carrier(predictions) + if bbox_detections is None or len(bbox_detections) == 0: + return [] + bboxes_metadata = bbox_detections.bboxes_metadata + if bboxes_metadata is None or any( + detection_id_key not in (entry or {}) for entry in bboxes_metadata + ): + raise ValueError( + f"Detections object passed to crop step do not fulfill contract - lack of " + f"{detection_id_key} key in per-detection metadata." + ) + # Materialise all (n, 4) corner ints to host in one transfer. + xyxy_int = bbox_detections.xyxy.round().to(torch.int64).detach().to("cpu").numpy() + # Clamp corners to image bounds: a negative start index would slice from the + # end, and clamping keeps the (-x_min, -y_min) translation consistent with the + # actually-cropped region. Crop whichever representation is already + # materialised (CHW tensor or HWC numpy) to avoid forcing a numpy->device + # conversion. + use_tensor = image.is_tensor_materialised() + if use_tensor: + image_height = int(image.tensor_image.shape[1]) + image_width = int(image.tensor_image.shape[2]) + else: + image_height = int(image.numpy_image.shape[0]) + image_width = int(image.numpy_image.shape[1]) + want_overlay = ( + mask_opacity > 0 + and isinstance(bbox_detections, InstanceDetections) + and bbox_detections.mask is not None + ) + crops: List[Dict[str, Any]] = [] + for idx in range(len(bbox_detections)): + x_min, y_min, x_max, y_max = (int(v) for v in xyxy_int[idx]) + x_min = max(0, min(x_min, image_width)) + y_min = max(0, min(y_min, image_height)) + x_max = max(0, min(x_max, image_width)) + y_max = max(0, min(y_max, image_height)) + detection_id = bboxes_metadata[idx][detection_id_key] + cropped = _crop_region( + image=image, + use_tensor=use_tensor, + bbox_detections=bbox_detections, + index=idx, + detection_id=detection_id, + box=(x_min, y_min, x_max, y_max), + want_overlay=want_overlay, + mask_opacity=mask_opacity, + background_color=background_color, + ) + if cropped is None: + crops.append({"crops": None, "predictions": None}) + continue + translated_prediction = _translate_single_prediction( + prediction=predictions, + index=idx, + x_min=x_min, + y_min=y_min, + x_max=x_max, + y_max=y_max, + ) + crops.append( + { + "crops": cropped, + # preserve all masks, keypoints, and metadata if present + "predictions": translated_prediction, + } + ) + return crops + + +def _crop_region( + *, + image: WorkflowImageData, + use_tensor: bool, + bbox_detections, + index: int, + detection_id: str, + box: Tuple[int, int, int, int], + want_overlay: bool, + mask_opacity: float, + background_color: Union[str, Tuple[int, int, int]], +) -> Optional[WorkflowImageData]: + """Crop one detection's region from whichever image representation is + materialised (CHW tensor on-device, else HWC numpy on host โ€” no forced + numpy->device conversion). Returns None for an empty (out-of-bounds) box.""" + x_min, y_min, x_max, y_max = box + if use_tensor: + cropped_tensor_image = image.tensor_image[:, y_min:y_max, x_min:x_max] + if cropped_tensor_image.numel() == 0: + return None + cropped_tensor_image = cropped_tensor_image.contiguous() + if want_overlay: + cropped_tensor_image = _overlay_tensor_crop_with_mask( + crop=cropped_tensor_image, + detection_mask_2d=_instance_mask_bool_tensor( + detections=bbox_detections, + index=index, + device=cropped_tensor_image.device, + )[y_min:y_max, x_min:x_max], + mask_opacity=mask_opacity, + background_color=background_color, + ) + return WorkflowImageData.create_crop_from_tensor( + origin_image_data=image, + crop_identifier=detection_id, + cropped_tensor_image=cropped_tensor_image, + offset_x=x_min, + offset_y=y_min, + ) + cropped_image = image.numpy_image[y_min:y_max, x_min:x_max] + if cropped_image.size == 0: + return None + if want_overlay: + cropped_image = _overlay_numpy_crop_with_mask( + crop=cropped_image, + detection_mask_2d=_instance_mask_bool_numpy( + detections=bbox_detections, + index=index, + )[y_min:y_max, x_min:x_max], + mask_opacity=mask_opacity, + background_color=background_color, + ) + return WorkflowImageData.create_crop( + origin_image_data=image, + crop_identifier=detection_id, + cropped_image=cropped_image, + offset_x=x_min, + offset_y=y_min, + ) + + +def _bbox_carrier( + prediction: TensorNativePrediction, +) -> Optional[TensorNativeDetections]: + """Return the bounding-box ``Detections`` / ``InstanceDetections`` carrier for any + supported native input. For the keypoint tuple the bbox component supplies xyxy and + the per-box ``detection_id`` / metadata.""" + if isinstance(prediction, tuple): + _, detections = prediction + return detections + return prediction + + +def _translate_single_prediction( + prediction: TensorNativePrediction, + index: int, + x_min: int, + y_min: int, + x_max: int, + y_max: int, +) -> TensorNativePrediction: + """Slice the prediction down to a single detection ``index`` and translate + every geometry field into the crop's coordinate frame (offset + ``(-x_min, -y_min)``): xyxy is shifted, dense / RLE masks are cropped to the + box, and the flattened ``keypoints_xy`` / ``polygon`` / OBB metadata plus the + ``KeyPoints.xy`` tensor are offset. Returns the same shape as the input.""" + single = take_prediction_by_indices(prediction=prediction, indices=[index]) + is_tuple = isinstance(single, tuple) + key_points = None + if is_tuple: + key_points, detections = single + else: + detections = single + detections = deepcopy(detections) + offset = torch.tensor( + [-x_min, -y_min, -x_min, -y_min], + dtype=detections.xyxy.dtype, + device=detections.xyxy.device, + ) + detections.xyxy = detections.xyxy + offset + if isinstance(detections, InstanceDetections) and detections.mask is not None: + detections.mask = _crop_native_mask( + mask=detections.mask, x_min=x_min, y_min=y_min, x_max=x_max, y_max=y_max + ) + _offset_metadata_geometry( + bboxes_metadata=detections.bboxes_metadata, x_min=x_min, y_min=y_min + ) + if not is_tuple: + return detections + key_points = deepcopy(key_points) + if key_points.xy.numel() > 0: + kp_offset = torch.tensor( + [-x_min, -y_min], dtype=key_points.xy.dtype, device=key_points.xy.device + ) + key_points.xy = key_points.xy + kp_offset + return key_points, detections + + +def _offset_metadata_geometry( + bboxes_metadata: Optional[List[dict]], + x_min: int, + y_min: int, +) -> None: + """Translate the flattened geometry entries in per-box metadata โ€” + ``keypoints_xy``, ``polygon``, and the OBB ``xyxyxyxy`` corners โ€” writing + values back in their original list/array container so the serialiser keeps + working. The per-box host mirror is dropped: the caller shifted ``xyxy`` + into the crop frame, so a carried mirror would be stale โ€” consumers fall + back to tensor reads.""" + if not bboxes_metadata: + return + offset_xy = np.array([x_min, y_min]) + for entry in bboxes_metadata: + if not entry: + continue + for key in HOST_MIRROR_KEYS: + entry.pop(key, None) + if KEYPOINTS_XY_KEY_IN_SV_DETECTIONS in entry: + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = _subtract_offset( + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS], offset_xy + ) + if POLYGON_KEY_IN_SV_DETECTIONS in entry: + entry[POLYGON_KEY_IN_SV_DETECTIONS] = _subtract_offset( + entry[POLYGON_KEY_IN_SV_DETECTIONS], offset_xy + ) + if ORIENTED_BOX_COORDINATES in entry: + entry[ORIENTED_BOX_COORDINATES] = _subtract_offset( + entry[ORIENTED_BOX_COORDINATES], offset_xy + ) + + +def _subtract_offset( + value: Union[List, np.ndarray], offset_xy: np.ndarray +) -> Union[List, np.ndarray]: + """Subtract ``offset_xy`` ([x, y]) from a coordinate container, preserving the + container type (python list vs numpy array) and integer dtypes (integer + coordinates must not silently become floats).""" + if value is None: + return value + was_list = isinstance(value, list) + array = np.asarray(value) + if array.size == 0: + return value + is_integer = np.issubdtype(array.dtype, np.integer) + if is_integer: + shifted = array.astype(np.int64) - offset_xy.astype(np.int64) + else: + shifted = array.astype(float) - offset_xy + return shifted.tolist() if was_list else shifted + + +def _crop_native_mask( + mask: Union[torch.Tensor, InstancesRLEMasks], + x_min: int, + y_min: int, + x_max: int, + y_max: int, +) -> Union[torch.Tensor, InstancesRLEMasks]: + """Crop a single-instance native mask to the detection box, keeping the same + carrier (dense torch -> dense torch; RLE -> RLE).""" + if isinstance(mask, InstancesRLEMasks): + numpy_masks = coco_rle_masks_to_numpy_mask(mask) # (1, H, W) + cropped = numpy_masks[:, y_min:y_max, x_min:x_max] + target_height = cropped.shape[1] + target_width = cropped.shape[2] + rle_masks = [ + torch_mask_to_coco_rle(torch.as_tensor(single_mask, dtype=torch.bool))[ + "counts" + ] + for single_mask in cropped + ] + return InstancesRLEMasks( + image_size=(target_height, target_width), masks=rle_masks + ) + return mask[:, y_min:y_max, x_min:x_max].contiguous() + + +def _instance_mask_bool_tensor( + detections: InstanceDetections, + index: int, + device: torch.device, +) -> torch.Tensor: + """Materialise a single instance's full-image mask as a 2-D bool ``torch.Tensor`` + ``(H, W)`` on ``device`` for the background-removal overlay. A dense torch mask + stays on-device (no host round trip); RLE is decoded one instance at a time and + moved to ``device`` (the RLE codec is numpy-only).""" + mask = detections.mask + if isinstance(mask, InstancesRLEMasks): + single = coco_rle_masks_to_numpy_mask( + InstancesRLEMasks(image_size=mask.image_size, masks=[mask.masks[index]]) + )[0].astype(bool) + return torch.as_tensor(single, dtype=torch.bool, device=device) + return mask[index].to(device=device, dtype=torch.bool) + + +def _instance_mask_bool_numpy( + detections: InstanceDetections, + index: int, +) -> np.ndarray: + """Numpy counterpart of ``_instance_mask_bool_tensor``: a single instance's + full-image ``(H, W)`` bool mask on the host, for the numpy overlay path.""" + mask = detections.mask + if isinstance(mask, InstancesRLEMasks): + return coco_rle_masks_to_numpy_mask( + InstancesRLEMasks(image_size=mask.image_size, masks=[mask.masks[index]]) + )[0].astype(bool) + return mask[index].detach().to("cpu").numpy().astype(bool) + + +def _overlay_tensor_crop_with_mask( + crop: torch.Tensor, + detection_mask_2d: torch.Tensor, + mask_opacity: float, + background_color: Union[str, Tuple[int, int, int]], +) -> torch.Tensor: + """Background-removal overlay computed entirely on the crop's device (no + per-crop host round trip). Inside the instance mask the crop is kept verbatim; + outside it the pixel is faded toward ``background_color`` + (``mask_opacity * bg + (1 - mask_opacity) * crop``). ``crop`` is CHW RGB + uint8; ``detection_mask_2d`` is a (H, W) bool tensor already sliced to the + crop box, both on the same device.""" + device = crop.device + # convert_color_to_bgr_tuple yields BGR; the crop tensor is RGB, so reverse it. + bgr_color = convert_color_to_bgr_tuple(color=background_color) + rgb_color = tuple(bgr_color[::-1]) + bg = torch.tensor(rgb_color, dtype=torch.float32, device=device).reshape(3, 1, 1) + crop_float = crop.to(dtype=torch.float32) + faded = mask_opacity * bg + (1.0 - mask_opacity) * crop_float + mask_3c = detection_mask_2d.to(device=device, dtype=torch.bool).unsqueeze(0) + overlaid = torch.where(mask_3c, crop_float, faded) + return overlaid.round().clamp_(0, 255).to(dtype=crop.dtype).contiguous() + + +def _overlay_numpy_crop_with_mask( + crop: np.ndarray, + detection_mask_2d: np.ndarray, + mask_opacity: float, + background_color: Union[str, Tuple[int, int, int]], +) -> np.ndarray: + """Numpy counterpart of ``_overlay_tensor_crop_with_mask``: same blend for + host crops. ``crop`` is HWC BGR uint8 and ``convert_color_to_bgr_tuple`` + already yields BGR, so the colour is used directly (no channel reversal). + ``detection_mask_2d`` is a (H, W) bool array already sliced to the crop + box.""" + bgr_color = np.asarray( + convert_color_to_bgr_tuple(color=background_color), dtype=np.float32 + ) + crop_float = crop.astype(np.float32) + faded = mask_opacity * bgr_color + (1.0 - mask_opacity) * crop_float + mask_3c = np.asarray(detection_mask_2d, dtype=bool)[:, :, None] + overlaid = np.where(mask_3c, crop_float, faded) + return np.clip(np.round(overlaid), 0, 255).astype(crop.dtype) + + +def convert_color_to_bgr_tuple( + color: Union[str, Tuple[int, int, int]], +) -> Tuple[int, int, int]: + if isinstance(color, str): + return convert_string_color_to_bgr_tuple(color=color) + if isinstance(color, tuple) and len(color) == 3: + return color[::-1] + raise ValueError(f"Invalid color format: {color}") + + +def convert_string_color_to_bgr_tuple(color: str) -> Tuple[int, int, int]: + if color.startswith("#") and len(color) == 7: + try: + return tuple(int(color[i : i + 2], 16) for i in (5, 3, 1)) + except ValueError as e: + raise ValueError(f"Invalid hex color format: {color}") from e + if color.startswith("#") and len(color) == 4: + try: + return tuple(int(color[i] + color[i], 16) for i in (3, 2, 1)) + except ValueError as e: + raise ValueError(f"Invalid hex color format: {color}") from e + if color.startswith("(") and color.endswith(")"): + try: + return tuple(map(int, color[1:-1].split(",")))[::-1] + except ValueError as e: + raise ValueError(f"Invalid tuple color format: {color}") from e + raise ValueError(f"Invalid hex color format: {color}") diff --git a/inference/core/workflows/core_steps/transformations/dynamic_zones/v1_tensor.py b/inference/core/workflows/core_steps/transformations/dynamic_zones/v1_tensor.py new file mode 100644 index 0000000000..d7022769fa --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/dynamic_zones/v1_tensor.py @@ -0,0 +1,468 @@ +from typing import List, Literal, Optional, Tuple, Type, Union +from uuid import uuid4 + +import cv2 as cv +import numpy as np +import pycocotools.mask as mask_utils +import supervision as sv +import torch +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + instance_mask_to_numpy, +) +from inference.core.workflows.execution_engine.constants import ( + DETECTION_ID_KEY, + POLYGON_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks + +OUTPUT_KEY: str = "zones" +OUTPUT_KEY_DETECTIONS: str = "predictions" +OUTPUT_KEY_SIMPLIFICATION_CONVERGED: str = "simplification_converged" +TYPE: str = "roboflow_core/dynamic_zone@v1" +SHORT_DESCRIPTION = ( + "Simplify polygons so they are geometrically convex " + "and contain only the requested amount of vertices." +) +LONG_DESCRIPTION = """ +Generate simplified polygon zones from instance segmentation detections by converting masks to contours, computing convex hulls, reducing polygon vertices to a specified count using Douglas-Peucker approximation, optionally applying least squares edge fitting, and scaling polygons to create geometric zones based on detected object shapes for zone-based analytics, spatial filtering, and region-of-interest definition workflows. + +## How This Block Works + +This block creates simplified polygon zones from instance segmentation detections by converting complex mask shapes into geometrically convex polygons with a specified number of vertices. The block: + +1. Receives instance segmentation predictions containing masks (polygon representations) for detected objects +2. Converts masks to contours: + - Extracts contours from each detection mask using mask-to-polygon conversion + - Selects the largest contour from each detection (handles multiple contours per mask) +3. Computes convex hull: + - Calculates the convex hull of the largest contour using OpenCV's convex hull algorithm + - Ensures the resulting polygon is geometrically convex (no inward-facing angles) + - Creates a simplified outer boundary that encompasses all points in the contour +4. Simplifies polygon to required vertex count: + - Uses Douglas-Peucker polygon approximation algorithm to reduce vertices + - Iteratively adjusts epsilon parameter to achieve the target number of vertices + - Uses binary search to find the optimal epsilon value that produces the requested vertex count + - Handles convergence: if exact vertex count cannot be achieved, pads or truncates vertices +5. Optionally applies least squares edge fitting: + - If `apply_least_squares` is enabled, refines polygon edges by fitting lines to original contour points + - Selects contour points between polygon vertices + - Optionally filters to midpoint fraction (e.g., uses only central portion of each edge) to avoid edge effects + - Fits least squares lines to selected contour points for each edge + - Calculates intersections of fitted lines to create refined vertex positions + - Produces a polygon that better aligns with the original contour shape +6. Scales polygon (if `scale_ratio` != 1): + - Calculates polygon centroid (center of mass) + - Scales polygon relative to centroid by the specified scale ratio + - Expands or contracts polygon outward from center (scale > 1 expands, scale < 1 contracts) + - Useful for creating buffer zones or adjusting zone boundaries +7. Updates detections with simplified polygons: + - Stores simplified polygons in detection metadata under the polygon key + - Regenerates masks from simplified polygons for updated detection representation +8. Returns simplified zones and updated detections: + - `zones`: List of simplified polygons (one per detection) as coordinate lists + - `predictions`: Updated detections with simplified polygons and masks + - `simplification_converged`: Boolean indicating if all polygons converged to exact vertex count + +The block enables creation of geometric zones from complex object shapes detected by segmentation models. It's particularly useful when zones need to be created based on detected object shapes (e.g., basketball courts, road segments, parking lots, fields) where the zone should match the object's outline but be simplified for performance and ease of use. + +## Common Use Cases + +- **Zone Creation from Detections**: Create polygon zones based on detected object shapes (e.g., create basketball court zones from court detections, generate road segment zones from road detections, create field zones from sports field detections), enabling detection-based zone workflows +- **Geometric Zone Simplification**: Simplify complex object shapes into geometrically convex polygons with controlled vertex counts (e.g., simplify irregular shapes to rectangles/quadrilaterals, reduce complex polygons to manageable vertex counts, create geometric zones from masks), enabling zone simplification workflows +- **Dynamic Zone Definition**: Dynamically define zones based on detected objects in images (e.g., define zones from detected regions, create zones from object shapes, generate zones from segmentation results), enabling dynamic zone workflows +- **Zone-Based Analytics Setup**: Prepare zones for zone-based analytics and filtering (e.g., prepare zones for time-in-zone analytics, create zones for zone-based filtering, set up zones for spatial analytics), enabling zone-based analytics workflows +- **Region-of-Interest Definition**: Define regions of interest based on detected object boundaries (e.g., define ROIs from object detections, create ROI zones from segmentation, generate interest regions from masks), enabling ROI definition workflows +- **Spatial Filtering and Analysis**: Create zones for spatial filtering and analysis operations (e.g., create zones for spatial filtering, prepare zones for area calculations, generate zones for spatial queries), enabling spatial analysis workflows + +## Connecting to Other Blocks + +This block receives instance segmentation predictions and produces simplified polygon zones: + +- **After instance segmentation models** to create zones from detected object shapes (e.g., segmentation model to zones, masks to simplified polygons, detections to geometric zones), enabling segmentation-to-zone workflows +- **After detection filtering blocks** to create zones from filtered detections (e.g., filter detections then create zones, create zones from specific classes, generate zones from filtered results), enabling filter-to-zone workflows +- **Before zone-based analytics blocks** to provide simplified zones for analytics (e.g., zones for time-in-zone, zones for zone analytics, polygons for zone filtering), enabling zone-to-analytics workflows +- **Before visualization blocks** to display simplified zones (e.g., visualize zone polygons, display geometric zones, show simplified regions), enabling zone visualization workflows +- **Before spatial filtering blocks** to provide zones for spatial operations (e.g., zones for overlap filtering, polygons for spatial queries, regions for area calculations), enabling zone-to-filter workflows +- **In workflow outputs** to provide simplified zones as final output (e.g., zone generation workflows, polygon extraction workflows, geometric zone outputs), enabling zone output workflows + +## Requirements + +This block requires instance segmentation predictions with masks (polygon data). Input detections should be filtered to contain only the desired classes of interest before processing. The `required_number_of_vertices` parameter specifies the target vertex count for simplified polygons (e.g., 4 for rectangles/quadrilaterals, 3 for triangles). The block uses iterative Douglas-Peucker approximation with binary search to achieve the target vertex count, with a maximum of 1000 iterations. If convergence to exact vertex count fails, vertices are padded or truncated. The `scale_ratio` parameter (default 1) scales polygons relative to their centroid. The `apply_least_squares` parameter (default False) enables edge fitting to better align polygon edges with original contours. The `midpoint_fraction` parameter (0-1, default 1) controls which portion of contour points are used for least squares fitting (1 = all points, lower values use central portions of edges). The block outputs simplified polygons as lists of coordinate pairs, updated detections with simplified polygons, and a convergence flag. +""" + + +class DynamicZonesManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Dynamic Zone", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-square-dashed", + "blockPriority": 3, + "opencv": True, + }, + } + ) + type: Literal[f"{TYPE}", "DynamicZone"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Instance segmentation predictions containing masks (polygon data) for detected objects. Detections should be filtered to contain only desired classes of interest. Each detection's mask is converted to contours, and the largest contour is used to generate a simplified polygon zone. Supports instance segmentation format with mask data.", + examples=[ + "$steps.instance_segmentation_model.predictions", + "$segmentation.predictions", + ], + ) + required_number_of_vertices: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Target number of vertices for simplified polygons. The block uses Douglas-Peucker polygon approximation with iterative binary search to reduce polygon vertices to this count. Common values: 4 for rectangles/quadrilaterals, 3 for triangles, 6+ for more complex shapes. The algorithm attempts to converge to this exact count; if convergence fails (within iteration limit), vertices are padded or truncated to match the count.", + examples=[4, 3, 6, 8], + ) + scale_ratio: Union[float, Selector(kind=[FLOAT_KIND])] = Field( # type: ignore + default=1, + description="Scale factor to expand or contract resulting polygons relative to their centroid. Values > 1 expand polygons outward from center (create buffer zones), values < 1 contract polygons inward. Value of 1 (default) means no scaling. Scaling is applied after polygon simplification. Useful for creating buffer zones or adjusting zone boundaries.", + examples=[1.0, 1.05, 1.1, 0.95], + ) + apply_least_squares: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + default=False, + description="If True, applies least squares line fitting to refine polygon edges by aligning them with original contour points. For each edge of the simplified polygon, fits a line to contour points between vertices, then calculates intersections of fitted lines to create refined vertex positions. Produces polygons that better match the original contour shape, especially useful when simplified polygon vertices don't align well with contour edges.", + examples=[False, True], + ) + midpoint_fraction: Union[float, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + default=1, + description="Fraction (0-1) of contour points to use for least squares fitting on each edge. Value of 1 (default) uses all contour points between vertices. Lower values use only the central portion of each edge (e.g., 0.9 uses 90% of points, centered). Useful when convex polygon vertices are not well-aligned with edges, as it focuses fitting on the central portion of edges rather than edge effects near vertices. Only applies when apply_least_squares is True.", + examples=[1.0, 0.9, 0.8], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["predictions"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name=OUTPUT_KEY, kind=[LIST_OF_VALUES_KIND]), + OutputDefinition( + name=OUTPUT_KEY_DETECTIONS, + kind=[TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND], + ), + OutputDefinition( + name=OUTPUT_KEY_SIMPLIFICATION_CONVERGED, kind=[BOOLEAN_KIND] + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +def calculate_simplified_polygon( + contours: List[np.ndarray], required_number_of_vertices: int, max_steps: int = 1000 +) -> Tuple[np.ndarray, np.ndarray]: + largest_contour = max(contours, key=len) + + # https://docs.opencv.org/4.x/d3/dc0/group__imgproc__shape.html#ga014b28e56cb8854c0de4a211cb2be656 + convex_contour = cv.convexHull( + points=largest_contour, + returnPoints=True, + clockwise=True, + ) + # https://docs.opencv.org/4.9.0/d3/dc0/group__imgproc__shape.html#ga8d26483c636be6b35c3ec6335798a47c + perimeter = cv.arcLength(curve=convex_contour, closed=True) + upper_epsilon = perimeter + lower_epsilon = 0.0000001 + epsilon = lower_epsilon + upper_epsilon / 2 + # https://docs.opencv.org/4.9.0/d3/dc0/group__imgproc__shape.html#ga0012a5fdaea70b8a9970165d98722b4c + simplified_polygon = cv.approxPolyDP( + curve=convex_contour, epsilon=epsilon, closed=True + ) + for _ in range(max_steps): + if len(simplified_polygon) == required_number_of_vertices: + break + if len(simplified_polygon) > required_number_of_vertices: + lower_epsilon = epsilon + else: + upper_epsilon = epsilon + epsilon = lower_epsilon + (upper_epsilon - lower_epsilon) / 2 + simplified_polygon = cv.approxPolyDP( + curve=convex_contour, epsilon=epsilon, closed=True + ) + while len(simplified_polygon.shape) > 2: + simplified_polygon = np.concatenate(simplified_polygon) + return simplified_polygon, largest_contour + + +def calculate_least_squares_polygon( + contour: np.ndarray, polygon: np.ndarray, midpoint_fraction: float = 1 +) -> np.ndarray: + def find_closest_index(point: np.ndarray, contour: np.ndarray) -> int: + dists = np.linalg.norm(contour - point, axis=1) + return np.argmin(dists) + + def pick_contour_points_between_vertices( + point_1: np.ndarray, point_2: np.ndarray, contour: np.ndarray + ) -> np.ndarray: + i1 = find_closest_index(point_1, contour) + i2 = find_closest_index(point_2, contour) + + if i1 <= i2: + return contour[i1 : i2 + 1] + else: + return np.concatenate((contour[i1:], contour[: i2 + 1]), axis=0) + + def least_squares_line(points: np.ndarray) -> Optional[Tuple[float, float]]: + if len(points) < 2: + return None + x = points[:, 0] + y = points[:, 1] + A = np.vstack([x, np.ones_like(x)]).T + a, b = np.linalg.lstsq(A, y, rcond=None)[0] + return (a, b) + + def intersect_lines( + line_1: Optional[Tuple[float, float]], line_2: Optional[Tuple[float, float]] + ) -> Optional[np.ndarray]: + if line_1 is None or line_2 is None: + return None + a_1, b_1 = line_1 + a_2, b_2 = line_2 + if np.isclose(a_1, a_2): + return None + x = (b_2 - b_1) / (a_1 - a_2) + y = a_1 * x + b_1 + return np.array([x, y]) + + pairs = [[polygon[-1], polygon[0]]] + list(zip(polygon[:-1], polygon[1:])) + + lines = [] + for point_1, point_2 in pairs: + segment_points = pick_contour_points_between_vertices(point_1, point_2, contour) + if midpoint_fraction < 1: + number_of_points = int(round(len(segment_points) * midpoint_fraction)) + if number_of_points > 2: + number_of_points_to_discard = ( + len(segment_points) - number_of_points + ) // 2 + segment_points = segment_points[ + number_of_points_to_discard : len(segment_points) + - number_of_points_to_discard + ] + line_params = least_squares_line(segment_points) + lines.append(line_params) + + intersections = [] + for i in range(len(lines)): + line_1 = lines[i] + line_2 = lines[(i + 1) % len(lines)] + pt = intersect_lines(line_1, line_2) + intersections.append(pt) + + return np.array(intersections, dtype=float).round().astype(int) + + +def scale_polygon(polygon: np.ndarray, scale: float) -> np.ndarray: + if scale == 1: + return polygon + + M = cv.moments(polygon) + + if M["m00"] == 0: + return polygon + + centroid_x = M["m10"] / M["m00"] + centroid_y = M["m01"] / M["m00"] + + shifted = polygon - [centroid_x, centroid_y] + scaled = shifted * scale + result = scaled + [centroid_x, centroid_y] + + return result.round().astype(np.int32) + + +def _repackage_masks( + new_dense_masks: np.ndarray, + original_mask: Union[torch.Tensor, InstancesRLEMasks], +) -> Union[torch.Tensor, InstancesRLEMasks]: + """Match the output mask representation to the input โ€” dense in, dense out; + RLE in, RLE out โ€” so downstream keeps the upstream storage shape.""" + if isinstance(original_mask, torch.Tensor): + return torch.from_numpy(new_dense_masks).to( + device=original_mask.device, dtype=original_mask.dtype + ) + rles = [ + mask_utils.encode(np.asfortranarray(m.astype(np.uint8))) + for m in new_dense_masks + ] + return InstancesRLEMasks.from_coco_rle_masks( + image_size=original_mask.image_size, masks=rles + ) + + +class DynamicZonesBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return DynamicZonesManifest + + def run( + self, + predictions: Batch[InstanceDetections], + required_number_of_vertices: int, + scale_ratio: float = 1, + apply_least_squares: bool = False, + midpoint_fraction: float = 1, + ) -> BlockResult: + result = [] + for detections in predictions: + if detections is None: + result.append( + { + OUTPUT_KEY: None, + OUTPUT_KEY_DETECTIONS: None, + OUTPUT_KEY_SIMPLIFICATION_CONVERGED: False, + } + ) + continue + if detections.mask is None: + result.append( + { + OUTPUT_KEY: [], + OUTPUT_KEY_DETECTIONS: None, + OUTPUT_KEY_SIMPLIFICATION_CONVERGED: False, + } + ) + continue + number_of_detections = int(detections.xyxy.shape[0]) + if number_of_detections == 0: + # Empty-but-typed input (0 rows, empty mask carrier) is a + # legitimate upstream output: return an empty InstanceDetections + # preserving image_metadata + the original mask carrier, with + # convergence reported as True. + result.append( + { + OUTPUT_KEY: [], + OUTPUT_KEY_DETECTIONS: InstanceDetections( + xyxy=detections.xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + mask=detections.mask, + image_metadata=detections.image_metadata, + bboxes_metadata=None, + ), + OUTPUT_KEY_SIMPLIFICATION_CONVERGED: True, + } + ) + continue + simplified_polygons = [] + new_dense_masks: List[np.ndarray] = [] + existing_meta = detections.bboxes_metadata or [ + {} for _ in range(number_of_detections) + ] + new_bboxes_metadata: List[dict] = [] + all_converged = True + for i in range(number_of_detections): + mask = instance_mask_to_numpy(detections, i) + + contours = sv.mask_to_polygons(mask) + simplified_polygon, largest_contour = calculate_simplified_polygon( + contours=contours, + required_number_of_vertices=required_number_of_vertices, + ) + if apply_least_squares: + simplified_polygon = calculate_least_squares_polygon( + contour=largest_contour, + polygon=simplified_polygon, + midpoint_fraction=midpoint_fraction, + ) + vertices_count, _ = simplified_polygon.shape + if vertices_count < required_number_of_vertices: + all_converged = False + for _ in range(required_number_of_vertices - vertices_count): + simplified_polygon = np.append( + simplified_polygon, + [simplified_polygon[-1]], + axis=0, + ) + elif vertices_count > required_number_of_vertices: + all_converged = False + simplified_polygon = simplified_polygon[ + :required_number_of_vertices + ] + simplified_polygon = scale_polygon( + polygon=simplified_polygon, + scale=scale_ratio, + ) + # The stored per-box polygon carries the SCALED (original image + # resolution) coordinates, as the bare (V, 2) polygon โ€” an extra + # batch dim would nest the serialized field and bypass the + # declared-polygon fast path. + per_box_meta = { + **(existing_meta[i] or {}), + POLYGON_KEY_IN_SV_DETECTIONS: np.array(simplified_polygon), + } + # Backfill detection_id so output rows always satisfy the + # per-box detection_id producer contract even when upstream + # omitted it (preserves any existing id). + per_box_meta.setdefault(DETECTION_ID_KEY, str(uuid4())) + new_bboxes_metadata.append(per_box_meta) + simplified_polygons.append(simplified_polygon.tolist()) + new_dense_masks.append( + sv.polygon_to_mask( + polygon=simplified_polygon, + resolution_wh=mask.shape[::-1], + ).astype(bool) + ) + new_mask = _repackage_masks( + np.stack(new_dense_masks, axis=0), detections.mask + ) + result.append( + { + OUTPUT_KEY: simplified_polygons, + OUTPUT_KEY_DETECTIONS: InstanceDetections( + xyxy=detections.xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + mask=new_mask, + image_metadata=detections.image_metadata, + bboxes_metadata=new_bboxes_metadata, + ), + OUTPUT_KEY_SIMPLIFICATION_CONVERGED: all_converged, + } + ) + if not result: + result.append( + { + OUTPUT_KEY: [], + OUTPUT_KEY_DETECTIONS: None, + OUTPUT_KEY_SIMPLIFICATION_CONVERGED: False, + } + ) + return result diff --git a/inference/core/workflows/core_steps/transformations/geotag_detection/v1_tensor.py b/inference/core/workflows/core_steps/transformations/geotag_detection/v1_tensor.py new file mode 100644 index 0000000000..b3f8ecf04e --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/geotag_detection/v1_tensor.py @@ -0,0 +1,327 @@ +"""Tensor-native sibling of ``geotag_detection/v1.py``. + +Boxes / confidences / class ids are pulled to host in one DtoH copy per tensor. +The per-box class name resolves from ``bboxes_metadata[i][CLASS_NAME_KEY]`` when +present, else from the per-image ``image_metadata[CLASS_NAMES_KEY]`` +class_id -> name map, else degrades to ``"unknown"``. Image dimensions are read +via ``_read_shape_without_materialization()`` so a tensor-only image is not +copied to host just for its shape. +""" + +import math +from typing import List, Literal, Optional, Tuple, Type, Union + +import numpy as np +from pydantic import ConfigDict, Field + +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + DICTIONARY_KIND, + FLOAT_KIND, + IMAGE_KIND, + LIST_OF_VALUES_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +SHORT_DESCRIPTION = ( + "Project detections from pixel coordinates to real-world GPS positions." +) +LONG_DESCRIPTION = """ +Convert object detection bounding boxes to real-world GPS coordinates using camera position metadata. + +## How This Block Works + +This block takes object detection predictions and camera GPS metadata (latitude, longitude, altitude) and projects each detection's pixel position to a real-world ground coordinate. The projection uses the camera's field of view and altitude to compute a ground footprint, then maps pixel offsets from image center to geographic offsets from the camera position. + +1. **Receives detection predictions** from an upstream object detection block (any model that outputs bounding boxes) +2. **Takes camera GPS metadata** as inputs: latitude, longitude, altitude above ground, and optionally horizontal field of view +3. **Computes ground footprint** from altitude and FOV using basic trigonometry +4. **Projects each detection center** from pixel coordinates to lat/lon offset from camera position +5. **Outputs GeoJSON-ready records** with class, confidence, and geographic coordinates for each detection + +The projection assumes a nadir (straight-down) camera orientation, which is accurate for most drone survey flights. For oblique angles, accuracy decreases with distance from image center. + +## Common Use Cases + +- **Drone Survey Analysis**: Process drone footage to map detected objects (vehicles, people, animals, structures) at their real-world locations for survey, inspection, or monitoring applications +- **Agricultural Monitoring**: Map crop damage, equipment, or livestock positions from aerial imagery for precision agriculture workflows +- **Security and Surveillance**: Create geospatial awareness from aerial camera feeds, mapping detected activity to real-world coordinates for situational awareness +- **Wildlife Conservation**: Track and map animal detections from drone surveys to monitor populations, migration patterns, or habitat usage +- **Construction Site Monitoring**: Map equipment, materials, and personnel positions from aerial imagery for site management and safety compliance +- **Search and Rescue**: Rapidly map detected persons or objects across large areas from drone footage to coordinate response efforts + +## Connecting to Other Blocks + +This block receives detections and produces geospatial data: + +- **After object detection blocks** (YOLO, RF-DETR, etc.) to geotag their predictions with real-world coordinates +- **After tracking blocks** (ByteTrack, OC-SORT) to produce geotagged tracks with movement paths +- **Before data sink blocks** (CSV, JSON, Webhook) to export detection locations for GIS analysis +- **Before visualization blocks** to annotate frames with GPS coordinate labels +- **In video processing pipelines** where each frame's GPS comes from drone telemetry or EXIF metadata +""" + +OUTPUT_KEY_GEO_DETECTIONS = "geo_detections" +OUTPUT_KEY_GEOJSON = "geojson" + +METERS_PER_DEG_LAT = 111320.0 + + +class BlockManifest(WorkflowBlockManifest): + type: Literal["roboflow_core/geotag_detection@v1"] + model_config = ConfigDict( + json_schema_extra={ + "name": "GeoTag Detection", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "transformation", + "icon": "far fa-map-pin", + "blockPriority": 5, + }, + } + ) + + image: Selector(kind=[IMAGE_KIND]) = Field( + title="Input Image", + description="The image that detections were generated from. Used to determine image dimensions for coordinate projection.", + examples=["$inputs.image", "$steps.detection.image"], + ) + + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( + title="Detections", + description="Object detection predictions to geotag. Each detection's bounding box center will be projected to a GPS coordinate.", + examples=["$steps.detection.predictions"], + ) + + latitude: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + title="Camera Latitude", + description="GPS latitude of the camera position in decimal degrees. For drone imagery, this comes from the flight controller GPS. Positive values are North, negative are South.", + examples=[47.428681, "$inputs.latitude"], + ) + + longitude: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + title="Camera Longitude", + description="GPS longitude of the camera position in decimal degrees. For drone imagery, this comes from the flight controller GPS. Positive values are East, negative are West.", + examples=[-105.279125, "$inputs.longitude"], + ) + + altitude: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + title="Altitude (meters AGL)", + description="Camera altitude above ground level in meters. Used with field of view to compute the ground footprint. For drones, this is the relative altitude reported by the flight controller, not absolute altitude.", + examples=[69.0, "$inputs.altitude"], + ) + + horizontal_fov: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + title="Horizontal Field of View (degrees)", + description="Horizontal field of view of the camera in degrees. Default 73.7 covers most DJI consumer drones (Mini, Air, Mavic series). Adjust for other cameras. Wider FOV = larger ground footprint per frame.", + examples=[73.7, 84.0], + default=73.7, + ) + + heading: Union[Selector(kind=[FLOAT_KIND]), float] = Field( + title="Camera Heading (degrees clockwise from North)", + description="Compass bearing that the top of the image points toward, in degrees clockwise from true north. 0 means image-up is North (the default). When the gimbal does not report yaw, derive this from the flight course (bearing between successive GPS fixes) for nose-forward flight. Rotates the ground footprint so detections land on the correct real-world bearing instead of being pinned to North.", + examples=[0.0, 90.0, "$inputs.heading"], + default=0.0, + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY_GEO_DETECTIONS, + kind=[LIST_OF_VALUES_KIND], + description="List of GeoJSON-compatible detection records. Each record contains: class, confidence, lat, lon, pixel_x, pixel_y, width, height. Coordinates are in WGS-84 decimal degrees.", + ), + OutputDefinition( + name=OUTPUT_KEY_GEOJSON, + kind=[DICTIONARY_KIND], + description="Complete GeoJSON FeatureCollection with all detections as Point features. Ready for use with Mapbox, Leaflet, QGIS, or any GIS tool.", + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class GeoTagDetectionBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + predictions: Union[Detections, InstanceDetections], + latitude: float, + longitude: float, + altitude: float, + horizontal_fov: float = 73.7, + heading: float = 0.0, + ) -> BlockResult: + image_h, image_w = image._read_shape_without_materialization() + geo_detections, features = project_detections( + predictions, + image_w, + image_h, + latitude, + longitude, + altitude, + horizontal_fov, + heading, + ) + return { + OUTPUT_KEY_GEO_DETECTIONS: geo_detections, + OUTPUT_KEY_GEOJSON: {"type": "FeatureCollection", "features": features}, + } + + +def project_detections( + predictions: Union[Detections, InstanceDetections], + image_w: int, + image_h: int, + latitude: float, + longitude: float, + altitude: float, + horizontal_fov: float = 73.7, + heading: float = 0.0, +) -> Tuple[List[dict], List[dict]]: + """Project tensor-native detections to ground GPS coordinates. + + Returns (geo_detections records, GeoJSON features). + """ + geo_detections, features = [], [] + detections_number = int(predictions.xyxy.shape[0]) + if detections_number == 0: + return geo_detections, features + + image_metadata = predictions.image_metadata or {} + class_names_mapping = image_metadata.get(CLASS_NAMES_KEY) or {} + bboxes_metadata = predictions.bboxes_metadata or [ + {} for _ in range(detections_number) + ] + # Keep the boxes as float32 numpy for the projection arithmetic โ€” unpacking + # through Python floats changes the lat/lon low-order decimals. + boxes = predictions.xyxy.detach().to("cpu").numpy() + confidences = ( + predictions.confidence.detach().to("cpu").numpy() + if predictions.confidence is not None + else np.zeros(detections_number) + ) + class_ids = [ + int(value) for value in predictions.class_id.detach().to("cpu").tolist() + ] + + for i in range(detections_number): + x1, y1, x2, y2 = boxes[i] + cx, cy, w, h = (x1 + x2) / 2, (y1 + y2) / 2, x2 - x1, y2 - y1 + det_lat, det_lon = _pixel_to_gps( + cx, + cy, + image_w, + image_h, + latitude, + longitude, + altitude, + horizontal_fov, + heading, + ) + data = bboxes_metadata[i] + if CLASS_NAME_KEY in data: + class_name = str(data[CLASS_NAME_KEY]) + else: + class_name = str(class_names_mapping.get(class_ids[i], "unknown")) + record = { + "class": class_name, + "confidence": round(float(confidences[i]), 4), + "lat": round(det_lat, 7), + "lon": round(det_lon, 7), + "pixel_x": round(float(cx), 1), + "pixel_y": round(float(cy), 1), + "width": round(float(w), 1), + "height": round(float(h), 1), + } + geo_detections.append(record) + features.append( + { + "type": "Feature", + "geometry": { + "type": "Point", + "coordinates": [record["lon"], record["lat"]], + }, + "properties": { + "class": record["class"], + "confidence": record["confidence"], + }, + } + ) + return geo_detections, features + + +def _pixel_to_gps( + px: float, + py: float, + image_w: int, + image_h: int, + camera_lat: float, + camera_lon: float, + camera_alt: float, + fov_h: float = 73.7, + heading_deg: float = 0.0, +) -> Tuple[float, float]: + """Project a pixel to ground GPS for a nadir camera at a given heading. + + Similar triangles give the target's offset in the image's own right/up axes + (exact for a pinhole nadir camera over flat ground). The heading then rotates + those axes onto East/North before the lat/lon conversion. heading_deg = 0 + means image-up points true North, i.e. the original north-up behavior. + """ + if camera_alt <= 0: + return camera_lat, camera_lon + + fov_rad = math.radians(fov_h) + ground_width = 2 * camera_alt * math.tan(fov_rad / 2) + ground_height = ground_width * (image_h / image_w) + + # Offset in the image's own axes: +right toward frame-right, +up toward top. + right = (px - image_w / 2) / image_w * ground_width + up = -(py - image_h / 2) / image_h * ground_height + + # Rotate image axes onto compass axes by the camera heading (CW from North). + psi = math.radians(heading_deg) + east = right * math.cos(psi) + up * math.sin(psi) + north = -right * math.sin(psi) + up * math.cos(psi) + + meters_per_deg_lon = METERS_PER_DEG_LAT * math.cos(math.radians(camera_lat)) + det_lat = camera_lat + north / METERS_PER_DEG_LAT + det_lon = camera_lon + east / meters_per_deg_lon + return det_lat, det_lon diff --git a/inference/core/workflows/core_steps/transformations/image_slicer/v2.py b/inference/core/workflows/core_steps/transformations/image_slicer/v2.py index 5b942e8597..281ae475ff 100644 --- a/inference/core/workflows/core_steps/transformations/image_slicer/v2.py +++ b/inference/core/workflows/core_steps/transformations/image_slicer/v2.py @@ -177,6 +177,18 @@ def run( overlap_ratio_width: float, overlap_ratio_height: float, ) -> BlockResult: + # When a GPU/device tensor image is already materialised, slice the tensor + # (a zero-copy view) and emit tensor-backed crops so the per-slice + # tensor-native model reads them without a full-frame device->host round + # trip. Otherwise keep the numpy/cv2 path unchanged. + if image.is_tensor_materialised(): + return self._run_tensor( + image=image, + slice_width=slice_width, + slice_height=slice_height, + overlap_ratio_width=overlap_ratio_width, + overlap_ratio_height=overlap_ratio_height, + ) image_numpy = image.numpy_image resolution_wh = (image_numpy.shape[1], image_numpy.shape[0]) offsets = generate_offsets( @@ -201,6 +213,40 @@ def run( slices.append({"slices": None}) return slices + def _run_tensor( + self, + image: WorkflowImageData, + slice_width: int, + slice_height: int, + overlap_ratio_width: float, + overlap_ratio_height: float, + ) -> BlockResult: + # tensor_image is CHW; resolution is (W, H) read off the tensor shape so the + # whole frame is never transferred to host just to compute offsets. + tensor = image.tensor_image + resolution_wh = (int(tensor.shape[2]), int(tensor.shape[1])) + offsets = generate_offsets( + resolution_wh=resolution_wh, + slice_wh=(slice_width, slice_height), + overlap_ratio_wh=(overlap_ratio_width, overlap_ratio_height), + ) + slices = [] + for offset in offsets: + x_min, y_min, x_max, y_max = (int(v) for v in offset) + crop_tensor = tensor[:, y_min:y_max, x_min:x_max] + if crop_tensor.numel(): + cropped_image = WorkflowImageData.create_crop_from_tensor( + origin_image_data=image, + crop_identifier=f"image_slicer.{uuid4()}", + cropped_tensor_image=crop_tensor, + offset_x=x_min, + offset_y=y_min, + ) + slices.append({"slices": cropped_image}) + else: + slices.append({"slices": None}) + return slices + def generate_offsets( resolution_wh: Tuple[int, int], diff --git a/inference/core/workflows/core_steps/transformations/per_class_confidence_filter/v1_tensor.py b/inference/core/workflows/core_steps/transformations/per_class_confidence_filter/v1_tensor.py new file mode 100644 index 0000000000..a706d8939a --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/per_class_confidence_filter/v1_tensor.py @@ -0,0 +1,213 @@ +from typing import Any, Dict, List, Literal, Optional, Type, Union + +import torch +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, + take_prediction_by_mask, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + DICTIONARY_KIND, + FLOAT_ZERO_TO_ONE_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) + +SHORT_DESCRIPTION = "Filter detections by applying a per-class confidence threshold." + +LONG_DESCRIPTION = """ +Filter detection predictions by applying a different confidence threshold to each class, keeping only detections whose confidence meets or exceeds the threshold configured for their class (with a configurable fallback threshold for classes that are not listed). + +## How This Block Works + +This block applies class-aware confidence filtering to detection predictions, enabling precise control over which detections are retained based on per-class quality requirements. The block: + +1. Takes detection predictions (object detection, instance segmentation, or keypoint detection) and a dictionary mapping class names to confidence thresholds +2. Iterates through each detection, looking up the threshold associated with the detection's class name +3. If the class is not present in the dictionary, falls back to the configurable `default_threshold` value +4. Keeps only the detections whose confidence is greater than or equal to the resolved threshold +5. Returns the filtered detections while preserving all original metadata (class ids, masks, keypoints, tracker ids, etc.) + +Unlike a single global confidence threshold, this block lets you demand high-confidence predictions for classes that are prone to false positives while keeping a more permissive threshold for classes that are harder to detect. Unlike the generic detections filter, it exposes a purpose-built dictionary input that maps cleanly to a simple `{"class_name": threshold}` JSON object. + +## Common Use Cases + +- **Noise-prone classes**: Demand very high confidence (e.g. 0.9) for classes that frequently produce false positives, while accepting lower confidence for well-behaved classes +- **Hard-to-detect classes**: Lower the threshold for classes that the model rarely detects with high confidence so that they are not filtered out entirely +- **Production-grade filtering**: Apply domain-specific thresholds tuned during evaluation so that downstream analytics, alerts, or counting blocks only see detections that meet the project's quality bar +- **Multi-class pipelines**: Combine with object detection models that predict many classes at once when a single global confidence threshold is too coarse + +## Connecting to Other Blocks + +The filtered predictions from this block can be connected to: + +- **Visualization blocks** (Bounding Box Visualization, Label Visualization, Polygon Visualization) to render only detections that cleared their per-class threshold +- **Counting and analytics blocks** (Line Counter, Time in Zone, Velocity) so that metrics reflect only high-quality detections +- **Tracking blocks** (Byte Tracker) so that tracker associations are not polluted by low-confidence noise +- **Storage or sink blocks** (Roboflow Dataset Upload, Webhook Sink, CSV Formatter) so that only detections meeting the quality bar are persisted or transmitted +- **Downstream transformation blocks** (Dynamic Crop, Detection Offset) for subsequent processing on the filtered subset +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Per-Class Confidence Filter", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "flow_control", + "icon": "far fa-filter", + "blockPriority": 2, + }, + } + ) + type: Literal["roboflow_core/per_class_confidence_filter@v1"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + description="Detection predictions to filter. Each detection is kept only if its confidence is greater than or equal to the threshold configured for its class (with a fallback to default_threshold for classes that are not listed in class_thresholds).", + examples=["$steps.object_detection_model.predictions"], + ) + class_thresholds: Union[ + Dict[str, float], + Selector(kind=[DICTIONARY_KIND]), + ] = Field( + description="Mapping of class name to minimum confidence threshold. Detections whose class name is present in this dictionary are kept only if their confidence is at least the corresponding threshold. Classes not present fall back to default_threshold. Thresholds should be in the [0.0, 1.0] range.", + examples=[{"person": 0.98, "car": 0.5}, "$inputs.class_thresholds"], + ) + default_threshold: Union[float, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( + default=0.3, + description="Confidence threshold applied to detections whose class name is not listed in class_thresholds. Must be in the [0.0, 1.0] range.", + examples=[0.3, "$inputs.default_threshold"], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["predictions"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name="predictions", + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class PerClassConfidenceFilterBlockV1(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + predictions: Batch[TensorNativePrediction], + class_thresholds: Dict[str, Any], + default_threshold: float = 0.3, + ) -> BlockResult: + return [ + { + "predictions": filter_detections_by_class_confidence( + prediction=prediction, + class_thresholds=class_thresholds, + default_threshold=default_threshold, + ) + } + for prediction in predictions + ] + + +def filter_detections_by_class_confidence( + prediction: TensorNativePrediction, + class_thresholds: Dict[str, Any], + default_threshold: float = 0.3, +) -> TensorNativePrediction: + if prediction is None: + return prediction + # The keypoint kind is a (KeyPoints, Optional[Detections]) tuple; the bbox + # confidences live on the Detections component. + if isinstance(prediction, tuple): + _, detections = prediction + else: + detections = prediction + if detections is None: + return prediction + confidences = detections.confidence + if confidences is None or int(confidences.shape[0]) == 0: + return prediction + class_names = _resolve_class_names(detections) + thresholds = {str(k): float(v) for k, v in (class_thresholds or {}).items()} + default = float(default_threshold) + # Per-row thresholds are host data (class names live in image_metadata): + # built in one host pass, shipped in one H2D transfer, compared on-device. + per_row_thresholds = [ + thresholds.get(str(class_names[i] if i < len(class_names) else None), default) + for i in range(int(confidences.shape[0])) + ] + keep = confidences >= torch.as_tensor( + per_row_thresholds, dtype=confidences.dtype, device=confidences.device + ) + if bool(keep.all()): + return prediction + return take_prediction_by_mask(prediction, keep) + + +def _resolve_class_names( + detections: TensorNativeDetections, +) -> List[Optional[str]]: + """Per-box class name, preferring the ``image_metadata[CLASS_NAMES_KEY]`` + (class_id -> name) map and falling back to ``bboxes_metadata[i]["class"]`` + when the map is absent or lacks the id.""" + number_of_detections = int(detections.confidence.shape[0]) + name_by_class_id: Dict[int, str] = {} + if detections.image_metadata: + name_by_class_id = detections.image_metadata.get(CLASS_NAMES_KEY) or {} + bboxes_metadata = detections.bboxes_metadata or [] + class_ids = detections.class_id + resolved: List[Optional[str]] = [] + for i in range(number_of_detections): + class_name: Optional[str] = None + if class_ids is not None and i < int(class_ids.shape[0]): + class_name = name_by_class_id.get(int(class_ids[i])) + if class_name is None and i < len(bboxes_metadata): + entry = bboxes_metadata[i] or {} + class_name = entry.get(CLASS_NAME_KEY) + resolved.append(class_name) + return resolved diff --git a/inference/core/workflows/core_steps/transformations/perspective_correction/v1_tensor.py b/inference/core/workflows/core_steps/transformations/perspective_correction/v1_tensor.py new file mode 100644 index 0000000000..c681a99a65 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/perspective_correction/v1_tensor.py @@ -0,0 +1,1368 @@ +"""Tensor-native sibling of ``roboflow_core/perspective_correction@v1``. + +Under ENABLE_TENSOR_DATA_REPRESENTATION the ``predictions`` input/output kinds are +the native ``inference_models`` dataclasses (``Detections`` / ``InstanceDetections``, +or the keypoint-detection ``Tuple[KeyPoints, Optional[Detections]]``) rather than +``sv.Detections``. The block warps both the image and the predictions through a +``cv2`` perspective transform. + +The perspective-matrix math (``cv.getPerspectiveTransform`` / ``cv.perspectiveTransform`` +/ ``cv.warpPerspective``), polygon picking/sorting, and polygon extension logic are +re-used verbatim from the numpy sibling - only the bits that touched ``sv.Detections`` +are re-implemented natively: + +- ``get_anchors_coordinates`` (numpy ``extend_perspective_polygon``) -> + ``_native_anchor_coordinate`` computes the same anchor point from the box ``xyxy``. + ``CENTER_OF_MASS`` is the exception (it is the mask centroid, not derivable from + ``xyxy``): it is resolved from the decoded masks via ``_center_of_mass_coordinates``, + which delegates to supervision so the centroid matches numpy byte-for-byte and + object-detection input (no mask) raises the same ``ValueError``. +- ``.mask`` / ``.xyxy`` / ``detection[...]`` / ``sv.Detections.merge`` (numpy + ``correct_detections``) -> ``_correct_native_detections`` reads ``xyxy`` / decoded + masks / per-box keypoints off the native dataclasses, applies the same transform, and + rebuilds a native ``Detections`` / ``InstanceDetections`` via + ``attach_native_detection_metadata`` (class names + per-box ``detection_id``). + +Masks are transformed per-instance: decoded with ``instance_mask_to_numpy``, the mask +polygon is warped, and the result is re-encoded to match the input representation +(dense ``torch.Tensor`` in -> dense out; ``InstancesRLEMasks`` in -> RLE out). + +Keypoints: the canonical keypoint payload that the tensor serialiser reads lives in +``bboxes_metadata[i]["keypoints_xy"]`` (flattened by the keypoint-detection model +block), so those coordinates are warped there. When the input is the keypoint tuple, +the standalone ``KeyPoints.xy`` tensor is warped in parallel so downstream tensor-native +consumers stay consistent. + +There is no model call / ``run_remotely`` here (this is a pure transformation), so no +HTTP-response rebuild is needed - it mirrors the numpy sibling, which also has none. +""" + +import math +from typing import List, Optional, Tuple, Union + +import cv2 as cv +import numpy as np +import supervision as sv +import torch +from pydantic import AliasChoices, ConfigDict, Field +from supervision.config import ORIENTED_BOX_COORDINATES +from typing_extensions import Literal, Type + +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, + instance_mask_to_numpy, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + IMAGE_DIMENSIONS_KEY, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + PREDICTION_TYPE_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import torch_mask_to_coco_rle + +OUTPUT_DETECTIONS_KEY: str = "corrected_coordinates" +OUTPUT_IMAGE_KEY: str = "warped_image" +OUTPUT_EXTENDED_TRANSFORMED_RECT_WIDTH_KEY: str = "extended_transformed_rect_width" +OUTPUT_EXTENDED_TRANSFORMED_RECT_HEIGHT_KEY: str = "extended_transformed_rect_height" +TYPE: str = "PerspectiveCorrection" +SHORT_DESCRIPTION = ( + "Adjust detection coordinates from a polygon-defined plane " + "to a straight rectangular plane with specified width and height." +) +LONG_DESCRIPTION = """ +Transform detection coordinates and optionally images from a perspective view to a top-down orthographic view using perspective transformation, correcting camera angle distortions to enable accurate measurements, top-down analysis, and coordinate normalization for scenarios where objects are viewed at an angle (e.g., surveillance cameras, aerial imagery, or tilted camera setups). + +## How This Block Works + +This block corrects perspective distortion by transforming coordinates from a perspective view (where objects appear smaller when further away and angles are distorted) to a top-down orthographic view (as if the camera were directly above the scene). The block: + +1. Receives input data: images and/or detections (object detection or instance segmentation predictions), along with perspective polygons defining regions to transform +2. Processes perspective polygons: + - Selects the largest polygon from provided polygons (if multiple are provided) + - Sorts polygon vertices in clockwise order and orients them starting from the leftmost bottom vertex + - Ensures proper polygon ordering for transformation matrix calculation +3. Optionally extends perspective polygons to contain all detections: + - If `extend_perspective_polygon_by_detections_anchor` is set, extends the polygon to ensure all detection anchor points (or entire bounding boxes if "ALL" is specified) are contained within the polygon + - Calculates extension amounts needed to contain detections outside the original polygon + - Adjusts polygon vertices to create a larger region that encompasses all detections +4. Generates perspective transformation matrix: + - Maps the source polygon (4 vertices in the perspective view) to a destination rectangle (top-down view) with specified width and height + - Uses OpenCV's `getPerspectiveTransform` to compute the 3x3 transformation matrix + - Handles extended dimensions when polygon extension is enabled +5. Applies perspective transformation to detections (if provided): + - Transforms bounding box coordinates from perspective view to top-down coordinates + - Transforms instance segmentation masks by converting masks to polygons, transforming polygon vertices, and converting back to masks in the new coordinate space + - Transforms keypoint coordinates for keypoint detection predictions + - Updates all coordinate data to reflect the corrected perspective +6. Optionally warps images (if `warp_image` is True): + - Applies the perspective transformation to the entire image using OpenCV's `warpPerspective` + - Produces a top-down view of the image with corrected perspective + - Outputs the warped image at the specified transformed rectangle dimensions (plus any extensions) +7. Returns corrected outputs: + - `corrected_coordinates`: Detections with transformed coordinates in the top-down coordinate space + - `warped_image`: Perspective-corrected image (if image warping is enabled) + - `extended_transformed_rect_width` and `extended_transformed_rect_height`: Final dimensions including any polygon extensions + +The transformation effectively "unwarps" the perspective distortion, making coordinates and images appear as if viewed from directly above. This is useful for accurate measurements, area calculations, distance measurements, and spatial analysis where perspective distortion would otherwise introduce errors. + +## Common Use Cases + +- **Top-Down Analysis**: Correct perspective distortion for top-down analysis and measurement (e.g., surveillance camera analysis, overhead view generation, top-down coordinate normalization), enabling top-down analysis workflows +- **Accurate Measurements**: Enable accurate distance, area, and size measurements by removing perspective distortion (e.g., measure object sizes in real-world units, calculate areas accurately, measure distances without distortion), enabling measurement workflows +- **Spatial Analysis**: Perform spatial analysis and coordinate-based operations on corrected coordinates (e.g., zone-based analysis, spatial tracking, coordinate-based filtering), enabling spatial analysis workflows +- **Aerial and Overhead Imagery**: Process aerial imagery or overhead camera feeds with perspective correction (e.g., drone imagery analysis, overhead camera processing, satellite image analysis), enabling aerial analysis workflows +- **Quality Control and Inspection**: Correct perspective for quality control and inspection workflows (e.g., manufacturing inspection, product quality checks, defect detection with accurate measurements), enabling quality control workflows +- **Indoor Navigation and Mapping**: Correct perspective for indoor navigation and mapping applications (e.g., floor plan generation, indoor mapping, navigation systems), enabling mapping workflows + +## Connecting to Other Blocks + +This block receives images and/or detections and produces perspective-corrected outputs: + +- **After detection models** to correct coordinates for accurate analysis (e.g., object detection with perspective correction, instance segmentation with corrected coordinates), enabling detection-to-correction workflows +- **After zone or polygon definition blocks** to use defined regions as perspective polygons (e.g., use polygon zones as perspective regions, apply correction to specific regions), enabling zone-to-correction workflows +- **Before measurement blocks** to enable accurate measurements on corrected coordinates (e.g., distance measurement with corrected coordinates, size measurement on top-down view, area calculation on corrected coordinates), enabling correction-to-measurement workflows +- **Before analytics blocks** to perform analytics on corrected coordinates (e.g., zone analytics with corrected coordinates, tracking with top-down view, path analysis with corrected paths), enabling correction-to-analytics workflows +- **Before visualization blocks** to visualize corrected coordinates and warped images (e.g., display top-down view, visualize corrected detections, show perspective-corrected results), enabling correction-to-visualization workflows +- **In workflow outputs** to provide perspective-corrected final outputs (e.g., top-down coordinate outputs, corrected detection outputs, warped image outputs), enabling correction-to-output workflows + +## Requirements + +This block requires either images or predictions (detections) as input. The `perspective_polygons` parameter must contain at least one polygon with exactly 4 vertices defining the region to transform. Polygons can be provided as a list of 4 coordinate pairs `[[x1, y1], [x2, y2], [x3, y3], [x4, y4]]` or as NumPy arrays. If multiple polygons are provided, the largest polygon (by area) is selected for each batch element. The `transformed_rect_width` and `transformed_rect_height` parameters define the dimensions of the output top-down rectangle. The block uses OpenCV's perspective transformation functions, which require proper polygon ordering and valid coordinate data. If polygon extension is enabled, the output dimensions are automatically adjusted to include the extended regions. +""" +ALL_POSITIONS = "ALL" + +NativePrediction = Union[ + Detections, + InstanceDetections, + Tuple[KeyPoints, Optional[Detections]], +] + + +class PerspectiveCorrectionManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Perspective Correction", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-toolbox", + "blockPriority": 2, + "opencv": True, + }, + } + ) + type: Literal["roboflow_core/perspective_correction@v1", "PerspectiveCorrection"] + predictions: Optional[ + Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) + ] = Field( # type: ignore + description="Optional object detection or instance segmentation predictions to transform. If provided, bounding boxes, masks, and keypoints are transformed to the top-down coordinate space. If not provided, only image warping is performed (if enabled). Either predictions or images must be provided.", + default=None, + examples=[ + "$steps.object_detection_model.predictions", + "$steps.instance_segmentation_model.predictions", + ], + ) + images: Selector(kind=[IMAGE_KIND]) = Field( + title="Image to Crop", + description="Input images to optionally warp to top-down view. Required if warp_image is True. Images are transformed using the perspective transformation matrix to produce top-down views. If only images are provided (no predictions), only image warping is performed.", + examples=["$inputs.image", "$steps.cropping.crops"], + validation_alias=AliasChoices("images", "image"), + ) + perspective_polygons: Union[list, Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + description="Perspective polygons defining regions to transform from perspective view to top-down view. Each polygon must consist of exactly 4 vertices (coordinates). Format: list of 4 coordinate pairs [[x1, y1], [x2, y2], [x3, y3], [x4, y4]] or NumPy arrays. If multiple polygons are provided for a batch element, the largest polygon (by area) is selected. The polygon defines the source region in the perspective view that will be mapped to the destination rectangle.", + examples=[ + "$steps.perspective_wrap.zones", + [[100, 100], [500, 100], [500, 400], [100, 400]], + ], + ) + transformed_rect_width: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Width of the destination rectangle in the top-down view (in pixels). The perspective polygon is transformed to fit this width. Coordinates are scaled to match this dimension. If polygon extension is enabled, the actual output width may be larger to accommodate extended regions.", + default=1000, + examples=[1000, 1920], + ) + transformed_rect_height: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Height of the destination rectangle in the top-down view (in pixels). The perspective polygon is transformed to fit this height. Coordinates are scaled to match this dimension. If polygon extension is enabled, the actual output height may be larger to accommodate extended regions.", + default=1000, + examples=[1000, 1080], + ) + extend_perspective_polygon_by_detections_anchor: Union[str, Selector(kind=[STRING_KIND])] = Field( # type: ignore + description=f"Optional setting to extend the perspective polygon to contain all detection anchor points. If set to a Position value ({', '.join(sv.Position.list())}), extends the polygon to contain that anchor point from all detections. If set to '{ALL_POSITIONS}', extends to contain entire bounding boxes (all corners). Empty string (default) disables extension. Extension ensures all detections are within the transformed region, automatically adjusting polygon boundaries and output dimensions.", + default="", + examples=["CENTER", "BOTTOM_CENTER", "ALL"], + ) + warp_image: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + description="If True, applies perspective transformation to the input image, producing a warped image in the top-down view. The warped image shows the perspective-corrected view at the specified transformed rectangle dimensions (plus any extensions). If False (default), only detection coordinates are transformed, and the original image is returned unchanged. Images must be provided if this is True.", + default=False, + examples=[False, True], + ) + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["images", "predictions"] + + @classmethod + def get_parameters_accepting_batches_and_scalars(cls) -> List[str]: + return [ + "perspective_polygons", + "transformed_rect_width", + "transformed_rect_height", + ] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_DETECTIONS_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + OutputDefinition( + name=OUTPUT_IMAGE_KEY, + kind=[ + IMAGE_KIND, + ], + ), + OutputDefinition( + name=OUTPUT_EXTENDED_TRANSFORMED_RECT_WIDTH_KEY, + kind=[ + INTEGER_KIND, + ], + ), + OutputDefinition( + name=OUTPUT_EXTENDED_TRANSFORMED_RECT_HEIGHT_KEY, + kind=[ + INTEGER_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +def pick_largest_perspective_polygons( + perspective_polygons_batch: Union[ + List[np.ndarray], + List[List[np.ndarray]], + List[List[List[int]]], + List[List[List[List[int]]]], + ], +) -> List[np.ndarray]: + if not isinstance(perspective_polygons_batch, (list, Batch)): + raise ValueError("Unexpected type of input") + if not perspective_polygons_batch: + raise ValueError("Unexpected empty batch") + if len(perspective_polygons_batch) == 4 and all( + isinstance(p, list) and len(p) == 2 for p in perspective_polygons_batch + ): + perspective_polygons_batch = [perspective_polygons_batch] + + largest_perspective_polygons: List[np.ndarray] = [] + for polygons in perspective_polygons_batch: + if polygons is None: + continue + if not isinstance(polygons, list) and not isinstance(polygons, np.ndarray): + raise ValueError("Unexpected type of batch element") + if len(polygons) == 0: + raise ValueError("Unexpected empty batch element") + if isinstance(polygons, np.ndarray): + if polygons.shape != (4, 2): + raise ValueError("Unexpected shape of batch element") + largest_perspective_polygons.append(polygons) + continue + if len(polygons) == 4 and all( + isinstance(p, list) and len(p) == 2 for p in polygons + ): + largest_perspective_polygons.append(np.array(polygons)) + continue + polygons = [p if isinstance(p, np.ndarray) else np.array(p) for p in polygons] + polygons = [p for p in polygons if p.shape == (4, 2)] + if not polygons: + raise ValueError("No batch element consists of 4 vertices") + polygons = [np.around(p).astype(np.int32) for p in polygons] + largest_polygon = max(polygons, key=lambda p: cv.contourArea(p)) + largest_perspective_polygons.append(largest_polygon) + return largest_perspective_polygons + + +def sort_polygon_vertices_clockwise(polygon: np.ndarray) -> np.ndarray: + x_center = min(polygon[:, 0]) / 2 + max(polygon[:, 0]) / 2 + y_center = min(polygon[:, 1]) / 2 + max(polygon[:, 1]) / 2 + angle = lambda p: math.atan2(x_center - p[0], y_center - p[1]) + return np.array(sorted(polygon.tolist(), key=angle, reverse=True)) + + +def roll_polygon_vertices_to_start_from_leftmost_bottom( + polygon: np.ndarray, +) -> np.ndarray: + x_min = min(polygon[:, 0]) + x_max = max(polygon[:, 0]) + y_min = min(polygon[:, 1]) + y_max = max(polygon[:, 1]) + leftmost_bottom_rect = [ + [x_min, y_max], + [x_min, y_min], + [x_max, y_min], + [x_max, y_max], + ] + min_dist = sum( + ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5 + for (x1, y1), (x2, y2) in zip(leftmost_bottom_rect, polygon) + ) + closest = polygon + for shift in range(4): + rolled = np.roll(polygon, shift=(0, shift), axis=(0, 0)) + dist = sum( + ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5 + for (x1, y1), (x2, y2) in zip(leftmost_bottom_rect, rolled) + ) + if dist < min_dist or (dist == min_dist and rolled[0][0] < closest[0][0]): + min_dist = dist + closest = rolled + return closest + + +def ccw(x1, y1, x2, y2, x3, y3): + return (x2 - x1) * (y3 - y1) - (y2 - y1) * (x3 - x1) > 0 + + +def calculate_line_coeffs( + x1: int, y1: int, x2: int, y2: int +) -> Tuple[Optional[float], float]: + if x1 == x2: + return None, x1 + # Solved a and b for ax + b = y + return (y2 - y1) / (x2 - x1), (y1 * x2 - y2 * x1) / (x2 - x1) + + +def calculate_line_intercept_to_contain_point( + a: Optional[float], + x: int, + y: int, +) -> float: + if a is None: + return x + return y - a * x + + +def solve_line_intersection( + a1: Optional[float], + b1: float, + a2: Optional[float], + b2: float, +) -> Tuple[float, float]: + if a1 is None and a2 is None: + raise ValueError("Both lines are vertical") + if a1 is None: + x = b1 + y = a2 * x + b2 + elif a2 is None: + x = b2 + y = a1 * x + b1 + else: + x = (b2 - b1) / (a1 - a2) + y = a1 * x + b1 + return x, y + + +def calculate_vertices_to_contain_point( + vertex_1: np.ndarray, + vertex_2: np.ndarray, + vertex_3_from_1: np.ndarray, + vertex_4_from_2: np.ndarray, + x: int, + y: int, +) -> Tuple[np.ndarray, np.ndarray]: + a, _ = calculate_line_coeffs( + x1=vertex_1[0], + y1=vertex_1[1], + x2=vertex_2[0], + y2=vertex_2[1], + ) + b = calculate_line_intercept_to_contain_point( + a=a, + x=x, + y=y, + ) + a_3_1, b_3_1 = calculate_line_coeffs( + x1=vertex_1[0], + y1=vertex_1[1], + x2=vertex_3_from_1[0], + y2=vertex_3_from_1[1], + ) + vertex_1 = ( + np.array( + solve_line_intersection( + a1=a_3_1, + b1=b_3_1, + a2=a, + b2=b, + ) + ) + .round() + .astype(int) + ) + a_4_2, b_4_2 = calculate_line_coeffs( + x1=vertex_2[0], + y1=vertex_2[1], + x2=vertex_4_from_2[0], + y2=vertex_4_from_2[1], + ) + vertex_2 = ( + np.array( + solve_line_intersection( + a1=a_4_2, + b1=b_4_2, + a2=a, + b2=b, + ) + ) + .round() + .astype(int) + ) + return vertex_1, vertex_2 + + +def _native_anchor_coordinate( + box: np.ndarray, + anchor: sv.Position, +) -> Tuple[float, float]: + """Replicate ``sv.Detections.get_anchors_coordinates(anchor)[0]`` for a single + box read directly off the native ``xyxy`` (shape ``(4,)`` as ``[x1, y1, x2, y2]``). + Only the anchors that ``extend_perspective_polygon`` requests are required: + the four corners and the four edge mid-points used for the ``ALL`` mode and the + standard ``sv.Position`` extension positions. + + ``sv.Position.CENTER_OF_MASS`` is intentionally NOT handled here: it is the mask + centroid, not derivable from ``xyxy`` alone. It is resolved upstream from the + decoded masks (``_center_of_mass_coordinates``) and injected via + ``center_of_mass_xy`` in ``extend_perspective_polygon`` - so a CENTER_OF_MASS + request never reaches this xyxy-only mapping. Leaving it out makes an accidental + call fail loudly (KeyError) instead of silently substituting the box centre.""" + x1, y1, x2, y2 = float(box[0]), float(box[1]), float(box[2]), float(box[3]) + center_x = (x1 + x2) / 2 + center_y = (y1 + y2) / 2 + mapping = { + sv.Position.CENTER: (center_x, center_y), + sv.Position.TOP_LEFT: (x1, y1), + sv.Position.TOP_CENTER: (center_x, y1), + sv.Position.TOP_RIGHT: (x2, y1), + sv.Position.CENTER_LEFT: (x1, center_y), + sv.Position.CENTER_RIGHT: (x2, center_y), + sv.Position.BOTTOM_LEFT: (x1, y2), + sv.Position.BOTTOM_CENTER: (center_x, y2), + sv.Position.BOTTOM_RIGHT: (x2, y2), + } + return mapping[anchor] + + +def _center_of_mass_coordinates( + detections: Union[Detections, InstanceDetections], +) -> np.ndarray: + """Native replica of numpy v1's ``det.get_anchors_coordinates(CENTER_OF_MASS)[0]`` + used by ``extend_perspective_polygon``: the (integer-truncated) mask centroid per + instance, in source-frame pixel coordinates. Returns shape ``(N, 2)``. + + Parity: the centroid is computed by supervision's own + ``get_anchors_coordinates(CENTER_OF_MASS)`` (i.e. ``calculate_masks_centroids``) + on the decoded mask, so it matches numpy v1 byte-for-byte - the same ``+0.5`` + pixel offset and ``.astype(int)`` truncation. Object-detection input (no mask) + reproduces supervision's exact ``ValueError`` rather than silently falling back to + the box centre, matching flag-off (numpy v1 raises there). + + Materialisation cost: the box centre in the old code was free (xyxy only). The + mask centroid is not - the segmentation mask has to be decoded to a dense + ``(H, W)`` array per instance. Masks are decoded one instance at a time (via + ``instance_mask_to_numpy``, the same convention as the rest of this block) so the + full ``(N, H, W)`` stack is never materialised at once; peak extra cost is one + dense image-resolution mask. This only runs when + ``extend_perspective_polygon_by_detections_anchor == CENTER_OF_MASS`` (a rare + config), so it is off the default path. + """ + xyxy_numpy = _native_detections_xyxy_numpy(detections) + if not isinstance(detections, InstanceDetections) or detections.mask is None: + # Object detection (or masks stripped): numpy v1 delegates to supervision and + # raises here. Delegate too so the flag-on error is identical (message + # included) to flag-off - this call raises and never returns. + return sv.Detections(xyxy=xyxy_numpy).get_anchors_coordinates( + sv.Position.CENTER_OF_MASS + ) + number_of_detections = xyxy_numpy.shape[0] + centroids = np.empty((number_of_detections, 2), dtype=int) + for index in range(number_of_detections): + mask = instance_mask_to_numpy(detections, index) # (H, W) bool + centroids[index] = sv.Detections( + xyxy=xyxy_numpy[index][np.newaxis, :], + mask=mask[np.newaxis, :, :], + ).get_anchors_coordinates(sv.Position.CENTER_OF_MASS)[0] + return centroids + + +def _extension_anchor_point( + box: np.ndarray, + index: int, + bbox_position: sv.Position, + center_of_mass_xy: Optional[np.ndarray], +) -> Tuple[float, float]: + """Resolve the extension anchor point for detection ``index``. All positions are + derivable from ``xyxy`` except ``CENTER_OF_MASS``, whose per-instance mask centroid + is precomputed upstream (``_center_of_mass_coordinates``) and passed in via + ``center_of_mass_xy``.""" + if bbox_position == sv.Position.CENTER_OF_MASS: + # center_of_mass_xy is guaranteed non-None here: the caller only reaches + # CENTER_OF_MASS after computing it (and raising for OD input). + x, y = center_of_mass_xy[index] + return float(x), float(y) + return _native_anchor_coordinate(box, bbox_position) + + +def extend_perspective_polygon( + polygon: List[np.ndarray], + detections_xyxy: np.ndarray, + bbox_position: Union[sv.Position, Literal[ALL_POSITIONS]], + center_of_mass_xy: Optional[np.ndarray] = None, +) -> Tuple[np.ndarray, float, float, float, float]: + if not bbox_position: + return polygon + bottom_left, top_left, top_right, bottom_right = polygon + extended_width = 0 + extended_height = 0 + original_width = max( + ( + (bottom_left[0] - bottom_right[0]) ** 2 + + (bottom_left[1] - bottom_right[1]) ** 2 + ) + ** 0.5, + ((top_left[0] - top_right[0]) ** 2 + (top_left[1] - top_right[1]) ** 2) ** 0.5, + ) + original_height = max( + ((bottom_left[0] - top_left[0]) ** 2 + (bottom_left[1] - top_left[1]) ** 2) + ** 0.5, + ((bottom_right[0] - top_right[0]) ** 2 + (bottom_right[1] - top_right[1]) ** 2) + ** 0.5, + ) + for i in range(detections_xyxy.shape[0]): + box = detections_xyxy[i] + # extend to the left + points = [] + if bbox_position == ALL_POSITIONS: + points.append(_native_anchor_coordinate(box, sv.Position.BOTTOM_LEFT)) + points.append(_native_anchor_coordinate(box, sv.Position.TOP_LEFT)) + else: + points.append( + _extension_anchor_point(box, i, bbox_position, center_of_mass_xy) + ) + for x, y in points: + if ( + cv.pointPolygonTest( + np.array([bottom_left, top_left, top_right, bottom_right]), + (x, y), + False, + ) + >= 0 + ): + continue + if not ccw( + x1=bottom_left[0], + y1=bottom_left[1], + x2=top_left[0], + y2=top_left[1], + x3=x, + y3=y, + ): + original_bottom_left = bottom_left + original_top_left = top_left + bottom_left, top_left = calculate_vertices_to_contain_point( + vertex_1=original_bottom_left, + vertex_2=original_top_left, + vertex_3_from_1=bottom_right, + vertex_4_from_2=top_right, + x=x, + y=y, + ) + extended_width += ( + (bottom_left[0] - original_bottom_left[0]) ** 2 + + (bottom_left[1] - original_bottom_left[1]) ** 2 + ) ** 0.5 + # extend to the right + points = [] + if bbox_position == ALL_POSITIONS: + points.append(_native_anchor_coordinate(box, sv.Position.BOTTOM_RIGHT)) + points.append(_native_anchor_coordinate(box, sv.Position.TOP_RIGHT)) + else: + points.append( + _extension_anchor_point(box, i, bbox_position, center_of_mass_xy) + ) + for x, y in points: + if ( + cv.pointPolygonTest( + np.array([bottom_left, top_left, top_right, bottom_right]), + (x, y), + False, + ) + >= 0 + ): + continue + if not ccw( + x1=top_right[0], + y1=top_right[1], + x2=bottom_right[0], + y2=bottom_right[1], + x3=x, + y3=y, + ): + original_bottom_right = bottom_right + original_top_right = top_right + top_right, bottom_right = calculate_vertices_to_contain_point( + vertex_1=original_top_right, + vertex_2=original_bottom_right, + vertex_3_from_1=top_left, + vertex_4_from_2=bottom_left, + x=x, + y=y, + ) + extended_width += ( + (bottom_right[0] - original_bottom_right[0]) ** 2 + + (bottom_right[1] - original_bottom_right[1]) ** 2 + ) ** 0.5 + # extend to the bottom + points = [] + if bbox_position == ALL_POSITIONS: + points.append(_native_anchor_coordinate(box, sv.Position.BOTTOM_RIGHT)) + points.append(_native_anchor_coordinate(box, sv.Position.BOTTOM_LEFT)) + else: + points.append( + _extension_anchor_point(box, i, bbox_position, center_of_mass_xy) + ) + for x, y in points: + if ( + cv.pointPolygonTest( + np.array([bottom_left, top_left, top_right, bottom_right]), + (x, y), + False, + ) + >= 0 + ): + continue + if not ccw( + x1=bottom_right[0], + y1=bottom_right[1], + x2=bottom_left[0], + y2=bottom_left[1], + x3=x, + y3=y, + ): + original_bottom_right = bottom_right + original_bottom_left = bottom_left + bottom_right, bottom_left = calculate_vertices_to_contain_point( + vertex_1=original_bottom_right, + vertex_2=original_bottom_left, + vertex_3_from_1=top_right, + vertex_4_from_2=top_left, + x=x, + y=y, + ) + extended_height += ( + (bottom_left[0] - original_bottom_left[0]) ** 2 + + (bottom_left[1] - original_bottom_left[1]) ** 2 + ) ** 0.5 + # extend to the top + points = [] + if bbox_position == ALL_POSITIONS: + points.append(_native_anchor_coordinate(box, sv.Position.TOP_RIGHT)) + points.append(_native_anchor_coordinate(box, sv.Position.TOP_LEFT)) + else: + points.append( + _extension_anchor_point(box, i, bbox_position, center_of_mass_xy) + ) + for x, y in points: + if ( + cv.pointPolygonTest( + np.array([bottom_left, top_left, top_right, bottom_right]), + (x, y), + False, + ) + >= 0 + ): + continue + if not ccw( + x1=top_left[0], + y1=top_left[1], + x2=top_right[0], + y2=top_right[1], + x3=x, + y3=y, + ): + original_top_left = top_left + original_top_right = top_right + top_left, top_right = calculate_vertices_to_contain_point( + vertex_1=original_top_left, + vertex_2=original_top_right, + vertex_3_from_1=bottom_left, + vertex_4_from_2=bottom_right, + x=x, + y=y, + ) + extended_height += ( + (top_left[0] - original_top_left[0]) ** 2 + + (top_left[1] - original_top_left[1]) ** 2 + ) ** 0.5 + return ( + np.array( + [ + bottom_left, + top_left, + top_right, + bottom_right, + ] + ), + original_width, + original_height, + extended_width, + extended_height, + ) + + +def generate_transformation_matrix( + src_polygon: np.ndarray, + transformed_rect_width: int, + transformed_rect_height: int, + detections_xyxy: Optional[np.ndarray] = None, + detections_anchor: Optional[Union[sv.Position, Literal[ALL_POSITIONS]]] = None, + center_of_mass_xy: Optional[np.ndarray] = None, +) -> Tuple[np.ndarray, float, float]: + polygon_with_vertices_clockwise = sort_polygon_vertices_clockwise( + polygon=src_polygon + ) + src_polygon = roll_polygon_vertices_to_start_from_leftmost_bottom( + polygon=polygon_with_vertices_clockwise + ) + original_width = transformed_rect_width + original_height = transformed_rect_height + extended_width = 0 + extended_height = 0 + if ( + detections_xyxy is not None + and detections_xyxy.shape[0] > 0 + and detections_anchor + ): + ( + src_polygon, + original_width, + original_height, + extended_width, + extended_height, + ) = extend_perspective_polygon( + polygon=src_polygon, + detections_xyxy=detections_xyxy, + bbox_position=( + sv.Position(detections_anchor) + if detections_anchor != ALL_POSITIONS + else detections_anchor + ), + center_of_mass_xy=center_of_mass_xy, + ) + extended_width = extended_width * transformed_rect_width / max(original_width, 1) + extended_height = ( + extended_height * transformed_rect_height / max(original_height, 1) + ) + src_polygon = src_polygon.astype(np.float32) + dst_polygon = np.array( + [ + [0, transformed_rect_height + int(round(extended_height)) - 1], + [0, 0], + [transformed_rect_width + int(round(extended_width)) - 1, 0], + [ + transformed_rect_width + int(round(extended_width)) - 1, + transformed_rect_height + int(round(extended_height)) - 1, + ], + ] + ).astype(dtype=np.float32) + # https://docs.opencv.org/4.9.0/da/d54/group__imgproc__transform.html#ga20f62aa3235d869c9956436c870893ae + return ( + cv.getPerspectiveTransform( + src=src_polygon, + dst=dst_polygon, + ), + extended_width, + extended_height, + ) + + +def _native_detections_xyxy_numpy(detections: Detections) -> np.ndarray: + """Read the native ``xyxy`` tensor out as a host numpy array (the perspective + math below is pure numpy / cv2).""" + return detections.xyxy.detach().to("cpu").numpy() + + +def _class_names_for_output(detections: Detections) -> dict: + """The producer contract requires ``image_metadata[CLASS_NAMES_KEY]`` so the + serialiser can map every ``class_id`` to a name. Reuse the input map when present; + otherwise reconstruct it from per-box ``bboxes_metadata[i]["class"]`` (the OCR / + legacy convention) keyed by the box ``class_id``.""" + image_metadata = detections.image_metadata or {} + class_names = image_metadata.get(CLASS_NAMES_KEY) + if class_names: + return dict(class_names) + reconstructed: dict = {} + bboxes_metadata = detections.bboxes_metadata or [] + class_ids = detections.class_id.detach().to("cpu").tolist() + for index, class_id in enumerate(class_ids): + if index < len(bboxes_metadata): + entry = bboxes_metadata[index] or {} + if CLASS_NAME_KEY in entry: + reconstructed[int(class_id)] = str(entry[CLASS_NAME_KEY]) + return reconstructed + + +def _prediction_type_for_output(detections: Detections) -> str: + image_metadata = detections.image_metadata or {} + return image_metadata.get( + PREDICTION_TYPE_KEY, + ( + "instance-segmentation" + if isinstance(detections, InstanceDetections) + else "object-detection" + ), + ) + + +def _correct_box_xyxy( + box: np.ndarray, + perspective_transformer: np.ndarray, +) -> np.ndarray: + """Warp a single axis-aligned box (``[x1, y1, x2, y2]``) through the perspective + transform and return the axis-aligned box of the warped corners. Same math as the + numpy block's non-mask branch.""" + xmin, ymin, xmax, ymax = np.around(box).tolist() + polygon = np.array( + [[[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]]], + dtype=np.float32, + ) + # https://docs.opencv.org/4.9.0/d2/de8/group__core__array.html#gad327659ac03e5fd6894b90025e6900a7 + corrected_polygon = cv.perspectiveTransform( + src=polygon, m=perspective_transformer + ).reshape(-1, 2) + return np.around(sv.polygon_to_xyxy(polygon=corrected_polygon)).astype(np.int32) + + +def _correct_mask_and_xyxy( + mask: np.ndarray, + perspective_transformer: np.ndarray, + transformed_rect_width: float, + transformed_rect_height: float, +) -> Tuple[Optional[np.ndarray], np.ndarray]: + """Warp a single instance mask: convert to polygon, warp the polygon, rebuild the + mask in the destination resolution, and derive the new axis-aligned box. Mirrors + the numpy block's mask branch. If the mask has no contour, returns ``(None, ...)`` + so the caller falls back to box-only warping.""" + polygons = sv.mask_to_polygons(mask) + if not polygons: + return None, np.zeros((4,), dtype=np.int32) + # ``sv.mask_to_polygons`` returns a list of contours; an instance mask with + # multiple disconnected components yields contours of differing point counts, + # so ``np.array(polygons, dtype=np.float32)`` would build a ragged + # (inhomogeneous) array and raise. Concatenate all contour points into a + # single ``(1, total_pts, 2)`` array โ€” mirroring the single-contour shape + # the original ``np.array([single_poly])`` produced. + polygon = np.concatenate( + [np.asarray(p, dtype=np.float32) for p in polygons], axis=0 + ).reshape(1, -1, 2) + # https://docs.opencv.org/4.9.0/d2/de8/group__core__array.html#gad327659ac03e5fd6894b90025e6900a7 + corrected_polygon = cv.perspectiveTransform( + src=polygon, m=perspective_transformer + ).reshape(-1, 2) + corrected_mask = sv.polygon_to_mask( + polygon=np.around(corrected_polygon).astype(np.int32), + resolution_wh=( + int(round(transformed_rect_width)), + int(round(transformed_rect_height)), + ), + ).astype(bool) + corrected_xyxy = np.around(sv.polygon_to_xyxy(polygon=corrected_polygon)).astype( + np.int32 + ) + return corrected_mask, corrected_xyxy + + +def _correct_keypoints_in_metadata( + entry: dict, + perspective_transformer: np.ndarray, +) -> None: + """Warp the per-box keypoint coordinates that the tensor serialiser reads + (``bboxes_metadata[i]["keypoints_xy"]``). Mirrors the numpy block warping + ``detection.data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS]``. Mutates ``entry`` in place. + """ + keypoints = entry.get(KEYPOINTS_XY_KEY_IN_SV_DETECTIONS) + if keypoints is None: + return + keypoints_array = np.array(keypoints, dtype=np.float32) + if keypoints_array.size == 0: + return + corrected = cv.perspectiveTransform( + src=np.array([keypoints_array.reshape(-1, 2)], dtype=np.float32), + m=perspective_transformer, + ).reshape(-1, 2) + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = ( + np.around(corrected).astype(np.int32).tolist() + ) + + +def _correct_oriented_box_in_metadata( + entry: dict, + perspective_transformer: np.ndarray, +) -> None: + """Warp the per-box oriented-bounding-box corners carried in + ``bboxes_metadata[i]["xyxyxyxy"]`` (shape ``(4, 2)`` - the per-row equivalent of + the sv ``detection.data[ORIENTED_BOX_COORDINATES]`` column). Mirrors the numpy + block's OBB warp added in PR #2521. Mutates ``entry`` in place. + """ + oriented_box = entry.get(ORIENTED_BOX_COORDINATES) + if oriented_box is None: + return + corrected_obb: np.ndarray = cv.perspectiveTransform( + src=np.array([np.asarray(oriented_box)], dtype=np.float32), + m=perspective_transformer, + ).reshape(-1, 2) + entry[ORIENTED_BOX_COORDINATES] = np.around(corrected_obb).astype(np.float32) + + +def _correct_native_detections( + detections: Union[Detections, InstanceDetections], + image: WorkflowImageData, + perspective_transformer: np.ndarray, + transformed_rect_width: float, + transformed_rect_height: float, +) -> Union[Detections, InstanceDetections]: + """Native re-implementation of the numpy block's ``correct_detections``: warp each + box (and mask, if instance segmentation) through the perspective transform and + rebuild a native ``Detections`` / ``InstanceDetections``. Per-box keypoint + coordinates carried in ``bboxes_metadata`` are warped in place.""" + is_instance_segmentation = isinstance(detections, InstanceDetections) + number_of_detections = len(detections) + xyxy_numpy = _native_detections_xyxy_numpy(detections) + existing_bboxes_metadata = detections.bboxes_metadata + device = detections.xyxy.device + + corrected_xyxy_rows: List[np.ndarray] = [] + corrected_bboxes_metadata: List[dict] = [] + dense_masks: List[np.ndarray] = [] if is_instance_segmentation else None + input_mask_is_rle = is_instance_segmentation and isinstance( + detections.mask, InstancesRLEMasks + ) + + for index in range(number_of_detections): + entry = ( + dict(existing_bboxes_metadata[index]) + if existing_bboxes_metadata is not None + and index < len(existing_bboxes_metadata) + else {} + ) + corrected_mask = None + corrected_box = None + if is_instance_segmentation and detections.mask is not None: + mask = instance_mask_to_numpy(detections, index) + corrected_mask, corrected_box = _correct_mask_and_xyxy( + mask=mask, + perspective_transformer=perspective_transformer, + transformed_rect_width=transformed_rect_width, + transformed_rect_height=transformed_rect_height, + ) + if corrected_mask is None: + corrected_box = _correct_box_xyxy( + box=xyxy_numpy[index], + perspective_transformer=perspective_transformer, + ) + if is_instance_segmentation: + # No contour to warp: emit an all-false mask at the destination + # resolution so the mask stack stays aligned with the boxes. + corrected_mask = np.zeros( + ( + int(round(transformed_rect_height)), + int(round(transformed_rect_width)), + ), + dtype=bool, + ) + corrected_xyxy_rows.append(corrected_box) + if is_instance_segmentation: + dense_masks.append(corrected_mask) + _correct_keypoints_in_metadata( + entry=entry, perspective_transformer=perspective_transformer + ) + _correct_oriented_box_in_metadata( + entry=entry, perspective_transformer=perspective_transformer + ) + corrected_bboxes_metadata.append(entry) + + if number_of_detections > 0: + xyxy_tensor = torch.as_tensor( + np.stack(corrected_xyxy_rows, axis=0), + dtype=detections.xyxy.dtype, + device=device, + ) + else: + xyxy_tensor = torch.zeros((0, 4), dtype=detections.xyxy.dtype, device=device) + + class_names = _class_names_for_output(detections) + prediction_type = _prediction_type_for_output(detections) + + if is_instance_segmentation: + mask_value = _repackage_masks( + dense_masks=dense_masks, + input_is_rle=input_mask_is_rle, + device=device, + destination_height=int(round(transformed_rect_height)), + destination_width=int(round(transformed_rect_width)), + ) + corrected = InstanceDetections( + xyxy=xyxy_tensor, + class_id=detections.class_id, + confidence=detections.confidence, + mask=mask_value, + image_metadata=None, + bboxes_metadata=corrected_bboxes_metadata or None, + ) + else: + corrected = Detections( + xyxy=xyxy_tensor, + class_id=detections.class_id, + confidence=detections.confidence, + image_metadata=None, + bboxes_metadata=corrected_bboxes_metadata or None, + ) + # Producer contract: attach image_metadata (with class_id -> name map) and + # guarantee a detection_id per box (preserves keypoint / class / text keys already + # present in bboxes_metadata). + corrected = attach_native_detection_metadata( + detections=corrected, + image=image, + class_names=class_names, + prediction_type=prediction_type, + ) + # The corrected boxes/masks live in the destination (warped) frame, not the + # source image. ``attach_native_detection_metadata`` reports the source image + # dimensions, so override IMAGE_DIMENSIONS_KEY with the destination rect size + # [height, width] (parent/root lineage intentionally left as the source frame). + if corrected.image_metadata is not None: + corrected.image_metadata[IMAGE_DIMENSIONS_KEY] = [ + int(round(transformed_rect_height)), + int(round(transformed_rect_width)), + ] + return corrected + + +def _repackage_masks( + dense_masks: List[np.ndarray], + input_is_rle: bool, + device: torch.device, + destination_height: int, + destination_width: int, +) -> Union[torch.Tensor, InstancesRLEMasks]: + """Match the output mask representation to the input - dense torch in -> dense out, + RLE in -> RLE out - so the rest of the tensor pipeline keeps the same storage shape. + """ + if input_is_rle: + rle_masks = [ + torch_mask_to_coco_rle(torch.as_tensor(single_mask, dtype=torch.bool))[ + "counts" + ] + for single_mask in dense_masks + ] + return InstancesRLEMasks( + image_size=(destination_height, destination_width), masks=rle_masks + ) + if len(dense_masks) == 0: + return torch.zeros( + (0, destination_height, destination_width), + dtype=torch.bool, + device=device, + ) + return torch.as_tensor( + np.stack(dense_masks, axis=0), dtype=torch.bool, device=device + ) + + +def _warp_key_points_tensor( + key_points: KeyPoints, + perspective_transformer: np.ndarray, +) -> KeyPoints: + """Warp the standalone ``KeyPoints.xy`` tensor (shape ``(instances, K, 2)``) through + the perspective transform so downstream tensor-native consumers of the tuple stay + consistent with the bbox component. Only ``xy`` is warped; ``class_id`` / + ``confidence`` / ``covariance`` / ``detection_confidence`` are carried through + unchanged - numpy v1 warps only the keypoint coordinates and leaves the covariance + and per-instance detection confidence untouched, so this must too (byte parity).""" + if key_points.xy.numel() == 0: + return key_points + device = key_points.xy.device + dtype = key_points.xy.dtype + instances, num_key_points, _ = key_points.xy.shape + xy_numpy = ( + key_points.xy.detach().to("cpu").numpy().reshape(-1, 2).astype(np.float32) + ) + warped = cv.perspectiveTransform( + src=np.array([xy_numpy], dtype=np.float32), + m=perspective_transformer, + ).reshape(instances, num_key_points, 2) + return KeyPoints( + xy=torch.as_tensor(warped, dtype=dtype, device=device), + class_id=key_points.class_id, + confidence=key_points.confidence, + image_metadata=key_points.image_metadata, + key_points_metadata=key_points.key_points_metadata, + covariance=key_points.covariance, + detection_confidence=key_points.detection_confidence, + ) + + +class PerspectiveCorrectionBlockV1(WorkflowBlock): + def __init__(self): + self.perspective_transformers: List[Tuple[np.ndarray, float, float]] = [] + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return PerspectiveCorrectionManifest + + def run( + self, + images: Batch[WorkflowImageData], + predictions: Optional[Batch[NativePrediction]], + perspective_polygons: Union[ + List[np.ndarray], + List[List[np.ndarray]], + List[List[List[int]]], + List[List[List[List[int]]]], + ], + transformed_rect_width: Union[int, List[int], np.ndarray], + transformed_rect_height: Union[int, List[int], np.ndarray], + extend_perspective_polygon_by_detections_anchor: Union[ + sv.Position, Literal[ALL_POSITIONS] + ], + warp_image: Optional[bool], + ) -> BlockResult: + if not predictions and not images: + raise ValueError( + "Either predictions or images are required to apply perspective correction." + ) + if warp_image and not images: + raise ValueError( + "images are required to warp image into requested perspective." + ) + if not predictions: + predictions = [None] * len(images) + batch_size = len(predictions) if predictions else len(images) + if isinstance(transformed_rect_height, int): + transformed_rect_height = [transformed_rect_height] * batch_size + if isinstance(transformed_rect_width, int): + transformed_rect_width = [transformed_rect_width] * batch_size + + if ( + not self.perspective_transformers + or extend_perspective_polygon_by_detections_anchor + ): + self.perspective_transformers = [] + largest_perspective_polygons = pick_largest_perspective_polygons( + perspective_polygons + ) + + if len(largest_perspective_polygons) == 1 and batch_size > 1: + largest_perspective_polygons = largest_perspective_polygons * batch_size + + if len(largest_perspective_polygons) != batch_size: + raise ValueError( + f"Predictions batch size ({batch_size}) does not match number of perspective polygons ({largest_perspective_polygons})" + ) + for polygon, prediction, width, height in zip( + largest_perspective_polygons, + predictions, + list(transformed_rect_width), + list(transformed_rect_height), + ): + if polygon is None: + self.perspective_transformers.append(None) + continue + bbox_detections = _bbox_component(prediction) + detections_xyxy = ( + _native_detections_xyxy_numpy(bbox_detections) + if bbox_detections is not None + else None + ) + # CENTER_OF_MASS is the only extension anchor not derivable from + # ``xyxy``: it is the mask centroid. Resolve it from the (source-frame) + # native masks here, where the native detections are still in scope, and + # pass the (N, 2) centroids down. Only computed when the anchor is + # actually CENTER_OF_MASS and there are detections to extend by - mirrors + # numpy v1, which raises for object-detection input (no mask). + center_of_mass_xy = None + if ( + bbox_detections is not None + and detections_xyxy is not None + and detections_xyxy.shape[0] > 0 + and extend_perspective_polygon_by_detections_anchor + and extend_perspective_polygon_by_detections_anchor != ALL_POSITIONS + and sv.Position(extend_perspective_polygon_by_detections_anchor) + == sv.Position.CENTER_OF_MASS + ): + center_of_mass_xy = _center_of_mass_coordinates(bbox_detections) + self.perspective_transformers.append( + generate_transformation_matrix( + src_polygon=polygon, + detections_xyxy=detections_xyxy, + transformed_rect_width=width, + transformed_rect_height=height, + detections_anchor=extend_perspective_polygon_by_detections_anchor, + center_of_mass_xy=center_of_mass_xy, + ) + ) + + result = [] + for prediction, perspective_transformer_w_h, image, width, height in zip( + predictions, + self.perspective_transformers, + images, + transformed_rect_width, + transformed_rect_height, + ): + perspective_transformer, extended_width, extended_height = ( + perspective_transformer_w_h + ) + result_image = image + if warp_image: + # https://docs.opencv.org/4.9.0/da/d54/group__imgproc__transform.html#gaf73673a7e8e18ec6963e3774e6a94b87 + warped_image = cv.warpPerspective( + src=image.numpy_image, + M=perspective_transformer, + dsize=( + int(round(width)) + int(round(extended_width)), + int(round(height)) + int(round(extended_height)), + ), + ) + result_image = WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=warped_image, + ) + + if prediction is None: + result.append( + { + OUTPUT_DETECTIONS_KEY: None, + OUTPUT_IMAGE_KEY: result_image, + OUTPUT_EXTENDED_TRANSFORMED_RECT_WIDTH_KEY: width + + int(round(extended_width)), + OUTPUT_EXTENDED_TRANSFORMED_RECT_HEIGHT_KEY: height + + int(round(extended_height)), + } + ) + continue + + corrected_detections = self._correct_prediction( + prediction=prediction, + image=image, + perspective_transformer=perspective_transformer, + transformed_rect_width=width + int(round(extended_width)), + transformed_rect_height=height + int(round(extended_height)), + ) + + result.append( + { + OUTPUT_DETECTIONS_KEY: corrected_detections, + OUTPUT_IMAGE_KEY: result_image, + OUTPUT_EXTENDED_TRANSFORMED_RECT_WIDTH_KEY: width + + int(round(extended_width)), + OUTPUT_EXTENDED_TRANSFORMED_RECT_HEIGHT_KEY: height + + int(round(extended_height)), + } + ) + return result + + def _correct_prediction( + self, + prediction: NativePrediction, + image: WorkflowImageData, + perspective_transformer: np.ndarray, + transformed_rect_width: float, + transformed_rect_height: float, + ) -> NativePrediction: + # Keypoint-detection prediction: warp the bbox Detections (carrying the per-box + # keypoints the serialiser reads) and the standalone KeyPoints tensor in parallel. + if isinstance(prediction, tuple): + key_points, detections = prediction + corrected_detections = ( + _correct_native_detections( + detections=detections, + image=image, + perspective_transformer=perspective_transformer, + transformed_rect_width=transformed_rect_width, + transformed_rect_height=transformed_rect_height, + ) + if detections is not None + else None + ) + corrected_key_points = ( + _warp_key_points_tensor( + key_points=key_points, + perspective_transformer=perspective_transformer, + ) + if key_points is not None + else None + ) + return corrected_key_points, corrected_detections + return _correct_native_detections( + detections=prediction, + image=image, + perspective_transformer=perspective_transformer, + transformed_rect_width=transformed_rect_width, + transformed_rect_height=transformed_rect_height, + ) + + +def _bbox_component( + prediction: Optional[NativePrediction], +) -> Optional[Union[Detections, InstanceDetections]]: + """Return the bbox ``Detections`` component of a native prediction (unwrap the + keypoint tuple), or ``None`` when there is no bbox prediction to extend by.""" + if prediction is None: + return None + if isinstance(prediction, tuple): + _, detections = prediction + return detections + return prediction diff --git a/inference/core/workflows/core_steps/transformations/stabilize_detections/v1_tensor.py b/inference/core/workflows/core_steps/transformations/stabilize_detections/v1_tensor.py new file mode 100644 index 0000000000..0c238091ef --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/stabilize_detections/v1_tensor.py @@ -0,0 +1,524 @@ +from collections import deque +from typing import Deque, Dict, List, Literal, Optional, Set, Tuple, Type, Union + +import numpy as np +import torch +from pydantic import ConfigDict, Field + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE +from inference.core.workflows.core_steps.common.tensor_native import ( + strip_host_mirror_metadata, + take_detections_by_indices, +) +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + IMAGE_KIND, + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + RuntimeRestriction, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks + +OUTPUT_KEY: str = "tracked_detections" +LONG_DESCRIPTION = """ +Apply smoothing algorithms to reduce noise and flickering in tracked detections across video frames by using Kalman filtering to predict object velocities, exponential moving average to smooth bounding box positions, and gap filling to restore temporarily missing detections for improved tracking stability and smoother visualization workflows. + +## How This Block Works + +This block stabilizes tracked detections by reducing jitter, smoothing positions, and filling gaps when objects temporarily disappear from detection. The block: + +1. Receives tracked detection predictions with unique tracker IDs and an image with embedded video metadata +2. Extracts video metadata from the image: + - Accesses video_metadata to get video_identifier + - Uses video_identifier to maintain separate stabilization state for different videos +3. Validates that detections have tracker IDs (required for tracking object movement across frames) +4. Initializes or retrieves stabilization state for the video: + - Maintains a cache of last known detections for each tracker_id per video + - Creates or retrieves a Kalman filter for velocity prediction per video + - Stores separate state for each video using video_identifier +5. Measures object velocities for existing tracks: + - Calculates velocity by comparing current frame bounding box centers to previous frame centers + - Computes displacement (change in position) for objects present in both current and previous frames + - Velocity measurements are used to update the Kalman filter +6. Updates Kalman filter with velocity measurements: + - Uses Kalman filtering to predict smoothed velocities based on historical measurements + - Maintains a sliding window of velocity measurements (controlled by smoothing_window_size) + - Applies exponential moving average within the Kalman filter to smooth velocity estimates + - Filters out noise from detection inaccuracies and frame-to-frame variations +7. Smooths bounding boxes for objects present in current frame: + - Applies exponential moving average smoothing to bounding box coordinates + - Combines previous frame position with current frame position using bbox_smoothing_coefficient + - Formula: smoothed_bbox = alpha * current_bbox + (1 - alpha) * previous_bbox + - Reduces jitter and flickering from detection variations +8. Predicts positions for missing detections: + - Uses Kalman filter predicted velocities to estimate positions of objects that disappeared + - Applies predicted velocity to last known bounding box position + - Fills gaps by restoring detections that were temporarily missing from current frame + - Smooths predicted positions using exponential moving average +9. Manages tracking state: + - Updates cache with current frame detections for next frame calculations + - Removes tracking entries for objects that have been missing longer than smoothing_window_size frames + - Maintains separate state per video_identifier +10. Merges and returns stabilized detections: + - Combines smoothed detections (from current frame) and predicted detections (for missing objects) + - Outputs stabilized detection objects with reduced noise and filled gaps + - All detections maintain their tracker IDs for consistent tracking + +The block uses two complementary smoothing techniques: **Kalman filtering** for velocity prediction (estimating how fast objects are moving) and **exponential moving average** for position smoothing (reducing bounding box jitter). The Kalman filter maintains a history of velocity measurements and uses statistical estimation to predict future velocities while filtering out noise. The exponential moving average smooths bounding box coordinates by blending current and previous positions. Gap filling uses predicted velocities to restore detections that temporarily disappear, helping maintain track continuity. Note: This block may produce short-lived bounding boxes for unstable trackers, as it attempts to fill gaps even when objects are inconsistently detected. + +## Common Use Cases + +- **Video Visualization**: Reduce flickering and jitter in video annotations for smoother visualizations (e.g., smooth bounding box movements, reduce annotation noise, improve video visualization quality), enabling stable video visualization workflows +- **Tracking Stability**: Improve tracking stability when detections are noisy or inconsistent (e.g., stabilize noisy detections, reduce tracking jitter, improve tracking continuity), enabling stable tracking workflows +- **Temporary Occlusion Handling**: Fill gaps when objects are temporarily occluded or missing from detections (e.g., maintain tracks during brief occlusions, fill detection gaps, preserve tracking continuity), enabling occlusion handling workflows +- **Real-Time Monitoring**: Improve visual quality in real-time monitoring applications (e.g., smooth live video annotations, reduce flickering in monitoring displays, improve real-time visualization), enabling stable real-time monitoring workflows +- **Analytics Accuracy**: Reduce noise in analytics calculations that depend on stable detection positions (e.g., improve position-based analytics, reduce noise in measurements, stabilize movement calculations), enabling accurate analytics workflows +- **Quality Control**: Improve detection quality for downstream processing (e.g., smooth detections before analysis, reduce noise for better processing, stabilize inputs for other blocks), enabling quality improvement workflows + +## Connecting to Other Blocks + +This block receives tracked detections and an image, and produces stabilized tracked_detections: + +- **After Byte Tracker blocks** to stabilize tracked detections (e.g., smooth tracked object positions, reduce tracking jitter, fill tracking gaps), enabling tracking-stabilization workflows +- **After object detection or instance segmentation blocks** with tracking enabled to stabilize detections (e.g., smooth detection positions, reduce detection noise, improve tracking stability), enabling detection-stabilization workflows +- **Before visualization blocks** to display stabilized detections (e.g., visualize smooth bounding boxes, display stable annotations, show gap-filled detections), enabling stable visualization workflows +- **Before analytics blocks** to provide stable inputs for analysis (e.g., analyze stabilized positions, process smooth movement data, work with gap-filled detections), enabling stable analytics workflows +- **Before velocity or path analysis blocks** to improve measurement accuracy (e.g., calculate velocities from stable positions, analyze paths from smooth trajectories, measure from gap-filled detections), enabling accurate measurement workflows +- **In video processing pipelines** where detection stability is required for downstream processing (e.g., stabilize detections in processing chains, improve quality for analysis, reduce noise in pipelines), enabling stable video processing workflows + +## Requirements + +This block requires tracked detections with tracker_id information (detections must come from a tracking block like Byte Tracker). The image's video_metadata should include video_identifier to maintain separate stabilization state for different videos. The block maintains persistent stabilization state across frames for each video, so it should be used in video workflows where frames are processed sequentially. For optimal stabilization, detections should be provided consistently across frames with valid tracker IDs. The smoothing_window_size controls how many historical velocity measurements are used for Kalman filtering and how long missing detections are retained. The bbox_smoothing_coefficient (0-1) controls the balance between current and previous positions - lower values provide more smoothing but slower response to changes, higher values provide less smoothing but faster response. + +Tensor-native note: this block consumes and produces `inference_models` dataclasses. The per-box `tracker_id` is read from `bboxes_metadata[i]["tracker_id"]` (there is no `.tracker_id` attribute on the native types). Smoothing / Kalman prediction run on numpy scalars extracted from the `xyxy` tensors; the smoothed boxes are written back into per-tracker single-row native slices (preserving `bboxes_metadata` and masks for instance segmentation), then concatenated back into a single native prediction. +""" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Detections Stabilizer", + "version": "v1", + "short_description": "Apply a smoothing algorithm to reduce noise and flickering across video frames.", + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "video", + "icon": "fas fa-waveform-lines", + "blockPriority": 4, + }, + } + ) + type: Literal["roboflow_core/stabilize_detections@v1"] + image: Selector(kind=[IMAGE_KIND]) = Field( + description="Image with embedded video metadata. The video_metadata contains video_identifier to maintain separate stabilization state for different videos. Required for persistent state management across frames.", + ) + detections: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Tracked object detection or instance segmentation predictions. Must include tracker_id information from a tracking block. The block applies Kalman filtering for velocity prediction, exponential moving average for position smoothing, and gap filling for missing detections. Output detections are stabilized with reduced noise and jitter.", + examples=["$steps.object_detection_model.predictions"], + ) + smoothing_window_size: Union[Optional[int], Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + default=3, + description="Size of the sliding window for velocity smoothing in Kalman filter, controlling how many historical velocity measurements are used. Also determines how long missing detections are retained before removal. Larger values provide more smoothing but slower adaptation to changes. Smaller values provide less smoothing but faster adaptation. Detections missing for longer than this number of frames are removed from tracking state. Typical range: 3-10 frames.", + examples=[3, 5, 10, "$inputs.smoothing_window_size"], + ) + bbox_smoothing_coefficient: Union[Optional[float], Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + default=0.2, + description="Exponential moving average coefficient (alpha) for bounding box position smoothing, range 0.0-1.0. Controls the blend between current and previous bounding box positions: smoothed_bbox = alpha * current + (1-alpha) * previous. Lower values (closer to 0) provide more smoothing - slower response to changes, less jitter. Higher values (closer to 1) provide less smoothing - faster response to changes, more jitter. Default 0.2 balances smoothness and responsiveness. Typical range: 0.1-0.5.", + examples=[0.2, 0.1, 0.5, "$inputs.bbox_smoothing_coefficient"], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_KEY, + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ], + ), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + return [ + STATEFUL_VIDEO_HTTP_SOFT_RESTRICTION, + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + ] + + +NativeDetections = Union[Detections, InstanceDetections] + + +class StabilizeTrackedDetectionsBlockV1(WorkflowBlock): + def __init__(self): + # Cache of last-known single-row native detections per tracker_id per video. + self._batch_of_last_known_detections: Dict[ + str, Dict[Union[int, str], NativeDetections] + ] = {} + # Parallel cache of each cached slice's xyxy as a (4,) numpy array, kept + # in lockstep with ``_batch_of_last_known_detections``, so the Kalman / + # smoothing loops don't re-issue a device->host copy per cached + # detection every frame. + self._batch_of_last_known_xyxy: Dict[str, Dict[Union[int, str], np.ndarray]] = ( + {} + ) + self._batch_of_kalman_filters: Dict[Union[int, str], VelocityKalmanFilter] = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + detections: NativeDetections, + smoothing_window_size: int, + bbox_smoothing_coefficient: float, + ) -> BlockResult: + metadata = image.video_metadata + num_detections = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata + tracker_ids = [ + ( + (bboxes_metadata[i] or {}).get("tracker_id") + if bboxes_metadata is not None + else None + ) + for i in range(num_detections) + ] + if num_detections > 0 and any(tracker_id is None for tracker_id in tracker_ids): + raise ValueError( + f"tracker_id not initialized, {self.__class__.__name__} requires detections to be tracked" + ) + + device = detections.xyxy.device + # Per-row xyxy materialised to host once so the numpy Kalman / smoothing + # helpers stay device-agnostic. + xyxy_rows = detections.xyxy.detach().to("cpu").numpy().astype(float) + + cached_detections = self._batch_of_last_known_detections.setdefault( + metadata.video_identifier, {} + ) + cached_xyxy = self._batch_of_last_known_xyxy.setdefault( + metadata.video_identifier, {} + ) + kalman_filter = self._batch_of_kalman_filters.setdefault( + metadata.video_identifier, + VelocityKalmanFilter(smoothing_window_size=smoothing_window_size), + ) + + measured_velocities = {} + for i, tracker_id in enumerate(tracker_ids): + if tracker_id not in cached_detections: + continue + x1, y1, x2, y2 = xyxy_rows[i] + this_frame_center_xy = [x1 + abs(x2 - x1), y1 + abs(y2 - y1)] + px1, py1, px2, py2 = cached_xyxy[tracker_id] + prev_frame_center_xy = [px1 + abs(px2 - px1), py1 + abs(py2 - py1)] + measured_velocities[tracker_id] = ( + this_frame_center_xy[0] - prev_frame_center_xy[0], + this_frame_center_xy[1] - prev_frame_center_xy[1], + ) + predicted_velocities = kalman_filter.update(measurements=measured_velocities) + + # Single-row native slices keyed by tracker_id, in input order, with the + # current-frame box smoothed against the previous-frame box when known. + predicted_detections: Dict[Union[int, str], NativeDetections] = {} + for i, tracker_id in enumerate(tracker_ids): + curr_frame_detection = take_detections_by_indices(detections, [i]) + if tracker_id in cached_detections: + prev_frame_xyxy = np.asarray(cached_xyxy[tracker_id]) + curr_frame_xyxy = xyxy_rows[i] + smoothed = smooth_xyxy( + prev_xyxy=prev_frame_xyxy, + curr_xyxy=curr_frame_xyxy, + alpha=bbox_smoothing_coefficient, + ) + _set_row_xyxy(curr_frame_detection, smoothed, device=device) + predicted_detections[tracker_id] = curr_frame_detection + else: + predicted_detections[tracker_id] = curr_frame_detection + # Cache the unsmoothed current detection for next frame's velocity / + # gap-fill measurements. + cached_detections[tracker_id] = take_detections_by_indices(detections, [i]) + cached_xyxy[tracker_id] = np.asarray(xyxy_rows[i], dtype=float).reshape(-1)[ + :4 + ] + + for tracker_id, predicted_velocity in predicted_velocities.items(): + if tracker_id in predicted_detections: + continue + prev_frame_detection = cached_detections[tracker_id] + prev_frame_xyxy = np.asarray(cached_xyxy[tracker_id]) + curr_frame_xyxy = np.array( + [ + prev_frame_xyxy + + np.array([predicted_velocity, predicted_velocity]).flatten() + ] + ) + smoothed = smooth_xyxy( + prev_xyxy=prev_frame_xyxy, + curr_xyxy=curr_frame_xyxy, + alpha=bbox_smoothing_coefficient, + ) + _set_row_xyxy(prev_frame_detection, smoothed, device=device) + predicted_detections[tracker_id] = prev_frame_detection + # The cached slice was mutated in place to the smoothed/predicted box; + # keep the numpy xyxy cache in lockstep so next frame reads the same + # box without another device->host copy. + cached_xyxy[tracker_id] = np.asarray(smoothed, dtype=float).reshape(-1)[:4] + + for tracker_id in list(cached_detections.keys()): + if ( + tracker_id not in kalman_filter.tracked_vectors + and tracker_id not in predicted_detections + ): + del cached_detections[tracker_id] + cached_xyxy.pop(tracker_id, None) + + merged_detections = _merge_native_detections( + list(predicted_detections.values()), + empty_template=detections, + device=device, + ) + return {OUTPUT_KEY: merged_detections} + + +def _row_xyxy(detection: NativeDetections) -> Tuple[float, float, float, float]: + """Return the single (and only) box of a one-row native detection as four + python floats.""" + row = detection.xyxy[0].detach().to("cpu").numpy().astype(float) + return float(row[0]), float(row[1]), float(row[2]), float(row[3]) + + +def _set_row_xyxy( + detection: NativeDetections, + xyxy: np.ndarray, + device: torch.device, +) -> None: + """Overwrite the single box of a one-row native detection in place, keeping + the tensor on ``device`` with the existing dtype. The row's host mirror is + dropped โ€” the smoothed/predicted box no longer matches it, and consumers + fall back to tensor reads when any row lacks the mirror.""" + flat = np.asarray(xyxy, dtype=float).reshape(-1)[:4] + detection.xyxy = torch.as_tensor( + flat, dtype=detection.xyxy.dtype, device=device + ).reshape(1, 4) + detection.bboxes_metadata = strip_host_mirror_metadata(detection.bboxes_metadata) + + +def _merge_native_detections( + detections_list: List[NativeDetections], + empty_template: NativeDetections, + device: torch.device, +) -> NativeDetections: + """Concatenate single-row native detections back into one prediction: + xyxy / class_id / confidence are concatenated, ``bboxes_metadata`` lists are + joined (preserving each box's ``tracker_id`` and ``detection_id``), and + masks are concatenated (dense torch stack or RLE ``masks`` list). + ``image_metadata`` (carrying the ``class_names`` map needed by the + serialiser) is taken from the current input.""" + if not detections_list: + return _empty_like(empty_template, device=device) + + xyxy = torch.cat([d.xyxy for d in detections_list], dim=0) + class_id = torch.cat([d.class_id for d in detections_list], dim=0) + confidence = torch.cat([d.confidence for d in detections_list], dim=0) + + bboxes_metadata: List[dict] = [] + for d in detections_list: + if d.bboxes_metadata is not None: + bboxes_metadata.extend(d.bboxes_metadata) + else: + bboxes_metadata.append({}) + + image_metadata = empty_template.image_metadata + + if isinstance(empty_template, InstanceDetections): + mask = _merge_native_masks(detections_list) + return InstanceDetections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata if bboxes_metadata else None, + ) + return Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata if bboxes_metadata else None, + ) + + +def _merge_native_masks( + detections_list: List[NativeDetections], +) -> Union[torch.Tensor, InstancesRLEMasks]: + """Concatenate the per-instance masks of single-row InstanceDetections.""" + first_mask = detections_list[0].mask + if isinstance(first_mask, InstancesRLEMasks): + merged_masks: List[bytes] = [] + for d in detections_list: + merged_masks.extend(d.mask.masks) + return InstancesRLEMasks( + image_size=first_mask.image_size, + masks=merged_masks, + ) + return torch.cat([d.mask for d in detections_list], dim=0) + + +def _empty_like( + template: NativeDetections, + device: torch.device, +) -> NativeDetections: + """Empty merged result preserving the template's shape / image_metadata. + + Carries over the ``class_names`` map so the serialiser still resolves names. + """ + image_metadata = template.image_metadata + if image_metadata is None: + image_metadata = {CLASS_NAMES_KEY: {}} + xyxy = torch.zeros((0, 4), dtype=torch.float32, device=device) + class_id = torch.zeros((0,), dtype=torch.long, device=device) + confidence = torch.zeros((0,), dtype=torch.float32, device=device) + if isinstance(template, InstanceDetections): + if isinstance(template.mask, InstancesRLEMasks): + empty_mask: Union[torch.Tensor, InstancesRLEMasks] = InstancesRLEMasks( + image_size=template.mask.image_size, + masks=[], + ) + else: + empty_mask = torch.zeros( + (0,) + tuple(template.mask.shape[1:]), + dtype=template.mask.dtype, + device=device, + ) + return InstanceDetections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=empty_mask, + image_metadata=image_metadata, + bboxes_metadata=None, + ) + return Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + image_metadata=image_metadata, + bboxes_metadata=None, + ) + + +def smooth_xyxy(prev_xyxy: np.ndarray, curr_xyxy: np.ndarray, alpha=0.2) -> np.ndarray: + smoothed_xyxy = alpha * curr_xyxy + (1 - alpha) * prev_xyxy + + return smoothed_xyxy + + +class VelocityKalmanFilter: + def __init__(self, smoothing_window_size: int): + self.time_step = 1 + self.smoothing_window_size = smoothing_window_size + self.state_transition_matrix = np.array([[1, 0], [0, 1]]) + self.process_noise_covariance = np.eye(2) * 0.001 + self.measurement_noise_covariance = np.eye(2) * 0.01 + self.tracked_vectors: Dict[ + Union[int, str], + Dict[ + Literal["velocity", "error_covariance", "history"], + Union[np.ndarray, Deque[float, float]], + ], + ] = {} + + def predict(self) -> Dict[Union[int, str], np.ndarray]: + predictions: Dict[Union[int, str], np.ndarray] = {} + for tracker_id, data in self.tracked_vectors.items(): + data["velocity"] = np.dot(self.state_transition_matrix, data["velocity"]) + data["error_covariance"] = ( + np.dot( + np.dot(self.state_transition_matrix, data["error_covariance"]), + self.state_transition_matrix.T, + ) + + self.process_noise_covariance + ) + predictions[tracker_id] = data["velocity"] + return predictions + + def update( + self, measurements: Dict[Union[int, str], Tuple[float, float]] + ) -> Dict[Union[int, str], np.ndarray]: + updated_vector_ids: Set[Union[int, str]] = set() + for tracker_id, velocity in measurements.items(): + updated_vector_ids.add(tracker_id) + if tracker_id in self.tracked_vectors: + measurement = np.array(velocity).reshape(2, 1) + tracked_vector = self.tracked_vectors[tracker_id] + tracked_vector["history"].appendleft(measurement) + smoothed_measurement = np.mean(tracked_vector["history"], axis=0) + measurement_residual = smoothed_measurement - tracked_vector["velocity"] + residual_covariance = ( + tracked_vector["error_covariance"] + + self.measurement_noise_covariance + ) + kalman_gain = np.dot( + tracked_vector["error_covariance"], + np.linalg.inv(residual_covariance), + ) + tracked_vector["velocity"] = tracked_vector["velocity"] + np.dot( + kalman_gain, measurement_residual + ) + tracked_vector["error_covariance"] = tracked_vector[ + "error_covariance" + ] - np.dot(kalman_gain, tracked_vector["error_covariance"]) + else: + self.tracked_vectors[tracker_id] = { + "velocity": np.array([[velocity[0]], [velocity[1]]]), + "error_covariance": np.eye(2), + "history": deque( + [np.array([[velocity[0]], [velocity[1]]])], + maxlen=self.smoothing_window_size, + ), + } + + predicted_velocities = self.predict() + + for tracker_id in set(self.tracked_vectors.keys()) - updated_vector_ids: + if self.tracked_vectors[tracker_id]["history"]: + self.tracked_vectors[tracker_id]["history"].popleft() + if not self.tracked_vectors[tracker_id]["history"]: + del self.tracked_vectors[tracker_id] + + return predicted_velocities diff --git a/inference/core/workflows/core_steps/transformations/stitch_ocr_detections/v1_tensor.py b/inference/core/workflows/core_steps/transformations/stitch_ocr_detections/v1_tensor.py new file mode 100644 index 0000000000..4cbbf29b35 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/stitch_ocr_detections/v1_tensor.py @@ -0,0 +1,406 @@ +"""Tensor-native sibling of stitch_ocr_detections/v1.py. + +Loaded INSTEAD of v1.py when ENABLE_TENSOR_DATA_REPRESENTATION is set. Reuses the +exact same class name / block type@version (StitchOCRDetectionsBlockV1, +``roboflow_core/stitch_ocr_detections@v1``) so it is a drop-in swap in loader.py. + +This block is a CONSUMER only: it takes OCR object-detection predictions and emits +a stitched text string under ``ocr_text`` (a plain ``STRING_KIND`` output, NOT a +prediction kind), so there is nothing to rebuild natively on the produce side and +no run_remotely path on this block. + +Native-input deltas vs the numpy file: +- input is an ``inference_models.Detections`` (xyxy is a torch tensor; the legacy + numpy block received an ``sv.Detections``); +- the recognised text rides per-box on ``bboxes_metadata[i]["text"]`` under the + agreed easy_ocr / ocr convention (the numpy block read it from + ``detections.data["class_name"]``); +- box coordinates are read via ``.detach().to("cpu").numpy()`` then rounded to int, + matching ``detections.xyxy.round().astype(int)`` in the numpy file. + +Everything downstream of the two reads (line grouping, sorting, separators, +delimiter joining) is byte-for-byte identical to v1.py. +""" + +from enum import Enum +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +from pydantic import ConfigDict, Field, field_validator + +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections + +TEXT_KEY = "text" + +LONG_DESCRIPTION = """ +Combine individual OCR detection results (words, characters, or text regions) into coherent text strings by organizing detections spatially according to reading direction, grouping detections into lines, sorting them within lines, and concatenating text in proper reading order to reconstruct readable text from OCR model outputs. + +## How This Block Works + +This block reconstructs readable text from individual OCR detections by organizing them spatially and concatenating text in proper reading order. The block: + +1. Receives OCR detection predictions containing individual text detections with bounding boxes and class names (text content) +2. Prepares coordinates based on reading direction: + - For vertical reading directions, swaps x and y coordinates to enable vertical line processing + - For horizontal reading directions, uses coordinates as-is +3. Groups detections into lines: + - Groups detections based on vertical position (or horizontal position for vertical text) using the tolerance parameter + - Detections within the tolerance distance are considered part of the same line + - Higher tolerance values group detections that are further apart, useful for text with variable line spacing +4. Sorts lines based on reading direction: + - For left-to-right and vertical top-to-bottom: sorts lines from top to bottom + - For right-to-left and vertical bottom-to-top: sorts lines in reverse order (bottom to top) +5. Sorts detections within each line: + - For left-to-right and vertical top-to-bottom: sorts detections by horizontal position (left to right, or top to bottom for vertical) + - For right-to-left and vertical bottom-to-top: sorts detections in reverse order (right to left, or bottom to top for vertical) +6. Concatenates text in reading order: + - Extracts class names (text content) from detections in sorted order + - Adds line separators (newline for horizontal text, space for vertical text) between lines + - Optionally inserts a delimiter between each text element if specified + - Produces a single coherent text string with proper reading order +7. Handles automatic reading direction detection (if "auto" is selected): + - Analyzes average width and height of detection bounding boxes + - If average width > average height: detects horizontal text (left-to-right) + - If average height >= average width: detects vertical text (top-to-bottom) +8. Returns the stitched text string: + - Outputs a single text string under the `ocr_text` key + - Text is formatted with proper line breaks and spacing according to reading direction + +The block enables reconstruction of multi-line text from individual OCR detections, maintaining proper reading order for different languages and writing systems. It handles both horizontal (left-to-right, right-to-left) and vertical (top-to-bottom, bottom-to-top) text orientations, making it useful for processing text in various languages and formats. + +## Common Use Cases + +- **Text Reconstruction**: Convert individual word or character detections from OCR models into readable text blocks (e.g., reconstruct documents from word detections, combine character detections into words, stitch OCR results into paragraphs), enabling text reconstruction workflows +- **Multi-Line Text Processing**: Reconstruct multi-line text from OCR results with proper line breaks and formatting (e.g., extract paragraphs from OCR results, reconstruct formatted text, process multi-line documents), enabling multi-line text workflows +- **Multi-Language OCR**: Process OCR results from different languages and writing systems (e.g., process Arabic right-to-left text, handle vertical Chinese/Japanese text, support multiple reading directions), enabling multi-language OCR workflows +- **Document Processing**: Extract and reconstruct text from documents and images (e.g., extract text from scanned documents, process invoice text, extract text from forms), enabling document processing workflows +- **Text Extraction and Formatting**: Extract text from images and format it for downstream use (e.g., extract text for database storage, format text for API responses, prepare text for analysis), enabling text extraction workflows +- **OCR Result Post-Processing**: Post-process OCR model outputs to produce usable text strings (e.g., format OCR outputs, organize OCR results, prepare text for downstream blocks), enabling OCR post-processing workflows + +## Connecting to Other Blocks + +This block receives OCR detection predictions and produces stitched text strings: + +- **After OCR model blocks** to convert detection results into readable text (e.g., OCR model to text string, OCR detections to formatted text, OCR results to text output), enabling OCR-to-text workflows +- **Before data storage blocks** to store extracted text (e.g., store OCR text in databases, save extracted text, log OCR results), enabling text storage workflows +- **Before notification blocks** to send extracted text in notifications (e.g., send OCR text in alerts, include extracted text in messages, notify with OCR results), enabling text notification workflows +- **Before text processing blocks** to process stitched text (e.g., process text with NLP models, analyze extracted text, apply text transformations), enabling text processing workflows +- **Before API output blocks** to provide text in API responses (e.g., return OCR text in API, format text for responses, provide extracted text output), enabling text output workflows +- **In workflow outputs** to provide stitched text as final output (e.g., text extraction workflows, OCR output workflows, document processing workflows), enabling text output workflows + +## Requirements + +This block requires OCR detection predictions (object detection format) with bounding boxes and class names containing text content. The `tolerance` parameter must be greater than zero and controls the vertical (or horizontal for vertical text) distance threshold for grouping detections into lines. The `reading_direction` parameter supports five modes: "left_to_right" (standard horizontal), "right_to_left" (Arabic-style), "vertical_top_to_bottom" (vertical), "vertical_bottom_to_top" (vertical reversed), and "auto" (automatic detection based on bounding box dimensions). The `delimiter` parameter is optional and inserts a delimiter between each text element (empty string by default, meaning no delimiter). The block outputs a single text string under the `ocr_text` key. +""" + +SHORT_DESCRIPTION = "Combines OCR detection results into a coherent text string by organizing detections spatially." + + +class ReadingDirection(str, Enum): + LEFT_TO_RIGHT = "left_to_right" + RIGHT_TO_LEFT = "right_to_left" + VERTICAL_TOP_TO_BOTTOM = "vertical_top_to_bottom" + VERTICAL_BOTTOM_TO_TOP = "vertical_bottom_to_top" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Stitch OCR Detections", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-reel", + "blockPriority": 2, + }, + } + ) + type: Literal["roboflow_core/stitch_ocr_detections@v1"] + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + title="OCR Detections", + description="OCR detection predictions from an OCR model. Should contain bounding boxes and class names with text content. Each detection represents a word, character, or text region that will be stitched together into coherent text. Supports object detection format with bounding boxes (xyxy) and class names in the data dictionary.", + examples=[ + "$steps.ocr_model.predictions", + "$steps.my_ocr_detection_model.predictions", + ], + ) + reading_direction: Literal[ + "left_to_right", + "right_to_left", + "vertical_top_to_bottom", + "vertical_bottom_to_top", + "auto", + ] = Field( + title="Reading Direction", + description="Direction to read and organize text detections. 'left_to_right': Standard horizontal reading (English, most languages). 'right_to_left': Right-to-left reading (Arabic, Hebrew). 'vertical_top_to_bottom': Vertical reading from top to bottom (Traditional Chinese, Japanese). 'vertical_bottom_to_top': Vertical reading from bottom to top (rare vertical formats). 'auto': Automatically detects reading direction based on average bounding box dimensions (width > height = horizontal, height >= width = vertical). Determines how detections are grouped into lines and sorted within lines.", + examples=["left_to_right", "right_to_left", "auto"], + json_schema_extra={ + "values_metadata": { + "left_to_right": { + "name": "Left To Right", + "description": "Standard left-to-right reading (e.g., English language)", + }, + "right_to_left": { + "name": "Right To Left", + "description": "Right-to-left reading (e.g., Arabic)", + }, + "vertical_top_to_bottom": { + "name": "Top To Bottom (Vertical)", + "description": "Vertical reading from top to bottom", + }, + "vertical_bottom_to_top": { + "name": "Bottom To Top (Vertical)", + "description": "Vertical reading from bottom to top", + }, + "auto": { + "name": "Auto", + "description": "Automatically detect the reading direction based on text arrangement.", + }, + } + }, + ) + tolerance: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + title="Tolerance", + description="Vertical (or horizontal for vertical text) distance threshold in pixels for grouping detections into the same line. Detections within this tolerance distance are grouped into the same line. Higher values group detections that are further apart (useful for text with variable line spacing or slanted text). Lower values create more lines (useful for tightly spaced text). Must be greater than zero.", + default=10, + examples=[10, 20, 5], + ) + delimiter: Union[str, Selector(kind=[STRING_KIND])] = Field( + title="Delimiter", + description="Optional delimiter string to insert between each text element (word/character) when stitching. Empty string (default) means no delimiter - text elements are concatenated directly. Useful for adding spaces between words, commas between elements, or custom separators. Example: use ' ' (space) to add spaces between words, or ',' to add commas.", + default="", + examples=["", " ", ",", "-"], + ) + + @field_validator("tolerance") + @classmethod + def ensure_tolerance_greater_than_zero( + cls, value: Union[int, str] + ) -> Union[int, str]: + if isinstance(value, int) and value <= 0: + raise ValueError( + "Stitch OCR detections block expects `tolerance` to be greater than zero." + ) + return value + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["predictions"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="ocr_text", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.0.0,<2.0.0" + + +def _native_xyxy_int(detections: Detections) -> np.ndarray: + """Native equivalent of ``detections.xyxy.round().astype(int)`` for a torch + xyxy tensor: pull to CPU numpy, round, cast to int.""" + return detections.xyxy.detach().to("cpu").numpy().round().astype(dtype=int) + + +def _native_class_names(detections: Detections) -> np.ndarray: + """Native equivalent of ``detections.data["class_name"]`` for OCR stitching. + + Resolves each detection's label the way the tensor serialiser does: a per-box + override wins (OCR-as-detector rides the recognised text on ``class``; the legacy + ``text`` key is also honoured), otherwise the label is resolved from + ``image_metadata['class_names']`` keyed by ``class_id`` โ€” which is where a plain + character/word *detection* model carries its label. Returned as an object array so + the numpy fancy-indexing in ``stitch_ocr_detections`` works unchanged.""" + n = int(detections.xyxy.shape[0]) + class_names_map = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + class_ids = detections.class_id.detach().cpu().tolist() + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + names = [] + for i in range(n): + meta = bboxes_metadata[i] or {} + if CLASS_NAME_KEY in meta: + names.append(str(meta[CLASS_NAME_KEY])) + elif TEXT_KEY in meta: + names.append(str(meta[TEXT_KEY])) + else: + names.append(str(class_names_map.get(int(class_ids[i]), ""))) + return np.array(names, dtype=object) + + +def detect_reading_direction(detections: Detections) -> str: + if int(detections.xyxy.shape[0]) == 0: + return "left_to_right" + + xyxy = detections.xyxy.detach().to("cpu").numpy() + widths = xyxy[:, 2] - xyxy[:, 0] + heights = xyxy[:, 3] - xyxy[:, 1] + + avg_width = np.mean(widths) + avg_height = np.mean(heights) + + if avg_width > avg_height: + return "left_to_right" + else: + return "vertical_top_to_bottom" + + +class StitchOCRDetectionsBlockV1(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + predictions: Batch[Detections], + reading_direction: str, + tolerance: int, + delimiter: str = "", + ) -> BlockResult: + if reading_direction == "auto": + reading_direction = detect_reading_direction(predictions[0]) + return [ + stitch_ocr_detections( + detections=detections, + reading_direction=reading_direction, + tolerance=tolerance, + delimiter=delimiter, + ) + for detections in predictions + ] + + +def stitch_ocr_detections( + detections: Detections, + reading_direction: str = "left_to_right", + tolerance: int = 10, + delimiter: str = "", +) -> Dict[str, str]: + """ + Stitch OCR detections into coherent text based on spatial arrangement. + + Args: + detections: Native inference_models Detections object containing OCR results + reading_direction: Direction to read text ("left_to_right", "right_to_left", + "vertical_top_to_bottom", "vertical_bottom_to_top") + tolerance: Vertical tolerance for grouping text into lines + + Returns: + Dict containing stitched OCR text under 'ocr_text' key + """ + if int(detections.xyxy.shape[0]) == 0: + return {"ocr_text": ""} + + xyxy = _native_xyxy_int(detections) + class_names = _native_class_names(detections) + + # Prepare coordinates based on reading direction + xyxy = prepare_coordinates(xyxy, reading_direction) + + # Group detections into lines + boxes_by_line = group_detections_by_line(xyxy, reading_direction, tolerance) + # Sort lines based on reading direction + lines = sorted( + boxes_by_line.keys(), reverse=reading_direction in ["vertical_bottom_to_top"] + ) + + # Build final text + ordered_class_names = [] + for i, key in enumerate(lines): + line_data = boxes_by_line[key] + line_xyxy = np.array(line_data["xyxy"]) + line_idx = np.array(line_data["idx"]) + + # Sort detections within line + sort_idx = sort_line_detections(line_xyxy, reading_direction) + + # Add sorted class names for this line + ordered_class_names.extend(class_names[line_idx[sort_idx]]) + + # Add line separator if not last line + if i < len(lines) - 1: + ordered_class_names.append(get_line_separator(reading_direction)) + + return {"ocr_text": delimiter.join(ordered_class_names)} + + +def prepare_coordinates( + xyxy: np.ndarray, + reading_direction: str, +) -> np.ndarray: + """Prepare coordinates based on reading direction.""" + if reading_direction in ["vertical_top_to_bottom", "vertical_bottom_to_top"]: + # Swap x and y coordinates: [x1,y1,x2,y2] -> [y1,x1,y2,x2] + return xyxy[:, [1, 0, 3, 2]] + return xyxy + + +def group_detections_by_line( + xyxy: np.ndarray, + reading_direction: str, + tolerance: int, +) -> Dict[float, Dict[str, List]]: + """Group detections into lines based on primary coordinate.""" + # After prepare_coordinates swap, we always group by y ([:, 1]) + primary_coord = xyxy[:, 1] # This is y for horizontal, swapped x for vertical + + # Round primary coordinate to group into lines + rounded_primary = np.round(primary_coord / tolerance) * tolerance + + boxes_by_line = {} + # Group bounding boxes and associated indices by line + for i, (bbox, line_pos) in enumerate(zip(xyxy, rounded_primary)): + if line_pos not in boxes_by_line: + boxes_by_line[line_pos] = {"xyxy": [bbox], "idx": [i]} + else: + boxes_by_line[line_pos]["xyxy"].append(bbox) + boxes_by_line[line_pos]["idx"].append(i) + + return boxes_by_line + + +def sort_line_detections( + line_xyxy: np.ndarray, + reading_direction: str, +) -> np.ndarray: + """Sort detections within a line based on reading direction.""" + # After prepare_coordinates swap, we always sort by x ([:, 0]) + if reading_direction in ["left_to_right", "vertical_top_to_bottom"]: + return line_xyxy[:, 0].argsort() # Sort by x1 (original x or swapped y) + else: # right_to_left or vertical_bottom_to_top + return (-line_xyxy[:, 0]).argsort() # Sort by -x1 (original -x or swapped -y) + + +def get_line_separator(reading_direction: str) -> str: + """Get the appropriate separator based on reading direction.""" + return "\n" if reading_direction in ["left_to_right", "right_to_left"] else " " diff --git a/inference/core/workflows/core_steps/transformations/stitch_ocr_detections/v2_tensor.py b/inference/core/workflows/core_steps/transformations/stitch_ocr_detections/v2_tensor.py new file mode 100644 index 0000000000..28c3c55d2e --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/stitch_ocr_detections/v2_tensor.py @@ -0,0 +1,958 @@ +"""Tensor-native sibling of stitch_ocr_detections/v2.py. + +Loaded INSTEAD of v2.py when ENABLE_TENSOR_DATA_REPRESENTATION is set. Reuses the +exact same class name / block type@version (StitchOCRDetectionsBlockV2, +``roboflow_core/stitch_ocr_detections@v2``) so it is a drop-in swap in loader.py. + +This block is a CONSUMER only: it takes OCR object-detection predictions and emits +a stitched text string under ``ocr_text`` (a plain ``STRING_KIND`` output, NOT a +prediction kind), so there is nothing to rebuild natively on the produce side and +no run_remotely path on this block. + +Native-input deltas vs the numpy file (applied identically to all three stitching +algorithms - tolerance, otsu, collimate): +- input is an ``inference_models.Detections`` (xyxy is a torch tensor; the legacy + numpy block received an ``sv.Detections``); +- the recognised text rides per-box on ``bboxes_metadata[i]["text"]`` under the + agreed easy_ocr / ocr convention (the numpy block read it from + ``detections.data["class_name"]``); +- box coordinates are read via ``.detach().to("cpu").numpy()`` - the tolerance + algorithm rounds-to-int (matching ``.round().astype(int)``), otsu / collimate + keep float coordinates exactly as the numpy file did with the raw ``.xyxy``. + +Everything downstream of those two reads (Otsu thresholding, collimate traversal, +line grouping, sorting, separators, delimiter joining) is byte-for-byte identical +to v2.py. +""" + +from enum import Enum +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +from pydantic import ConfigDict, Field, field_validator + +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + OutputDefinition, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + INTEGER_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections + +TEXT_KEY = "text" + +LONG_DESCRIPTION = """ +Combine individual OCR detection results (words, characters, or text regions) into coherent text strings by organizing detections spatially, grouping them into lines, and concatenating text in proper reading order. + +## Stitching Algorithms + +This block supports three algorithms for reconstructing text from OCR detections: + +### Tolerance-based (default) +Groups detections into lines using a fixed pixel tolerance. Detections within the tolerance distance vertically (or horizontally for vertical text) are grouped into the same line, then sorted by position within each line. + +- **Best for**: Consistent font sizes and well-aligned horizontal/vertical text +- **Parameters**: `tolerance` (pixel threshold for line grouping) + +### Otsu Thresholding +Uses Otsu's method on normalized gap distances to automatically find the optimal threshold separating character gaps from word gaps. Gaps are normalized by local character width, making it resolution-invariant. + +- **Best for**: Variable font sizes, automatic word boundary detection +- **Parameters**: `otsu_threshold_multiplier` (adjust threshold sensitivity) +- **Key feature**: Detects bimodal distributions to distinguish single words from multi-word text + +### Collimate (Skewed Text) +Uses greedy parent-child traversal to follow text flow. Starting from the first detection, it finds subsequent detections that "follow" in reading order (similar alignment + correct direction), building lines through traversal rather than bucketing. + +- **Best for**: Skewed, curved, or non-axis-aligned text +- **Parameters**: `collimate_tolerance` (alignment tolerance in pixels) +- **Note**: Does not detect word boundaries - use `delimiter` parameter if spacing is needed + +## Reading Directions + +All algorithms support multiple reading directions: +- `left_to_right`: Standard horizontal (English, most languages) +- `right_to_left`: Right-to-left (Arabic, Hebrew) +- `vertical_top_to_bottom`: Vertical top-to-bottom (Traditional Chinese, Japanese) +- `vertical_bottom_to_top`: Vertical bottom-to-top +- `auto`: Automatically detect based on bounding box dimensions + +## Common Use Cases + +- **Document OCR**: Reconstruct paragraphs and lines from character/word detections +- **Multi-language support**: Handle different reading directions and writing systems +- **Skewed text processing**: Use collimate algorithm for tilted or curved text +- **Word detection**: Use Otsu algorithm to automatically insert spaces between words +""" + +SHORT_DESCRIPTION = "Combines OCR detection results into a coherent text string by organizing detections spatially." + + +class ReadingDirection(str, Enum): + LEFT_TO_RIGHT = "left_to_right" + RIGHT_TO_LEFT = "right_to_left" + VERTICAL_TOP_TO_BOTTOM = "vertical_top_to_bottom" + VERTICAL_BOTTOM_TO_TOP = "vertical_bottom_to_top" + + +class StitchingAlgorithm(str, Enum): + """Algorithm for grouping detections into words/lines. + + TOLERANCE: Uses fixed pixel tolerance for line grouping (original algorithm). + Good for consistent font sizes and line spacing. + + OTSU: Uses Otsu's method on normalized gaps to find natural breaks. + Resolution-invariant and works well with bimodal distributions + (e.g., character-level vs word-level spacing). + + COLLIMATE: Uses greedy parent-child traversal to group detections. + Good for skewed or curved text where bucket-based approaches fail. + """ + + TOLERANCE = "tolerance" + OTSU = "otsu" + COLLIMATE = "collimate" + + +class BlockManifest(WorkflowBlockManifest): + model_config = ConfigDict( + json_schema_extra={ + "name": "Stitch OCR Detections", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "transformation", + "ui_manifest": { + "section": "advanced", + "icon": "fal fa-reel", + "blockPriority": 2, + }, + } + ) + type: Literal["roboflow_core/stitch_ocr_detections@v2"] + stitching_algorithm: Literal["tolerance", "otsu", "collimate"] = Field( + title="Stitching Algorithm", + description="Algorithm for grouping detections into words/lines. 'tolerance': Uses fixed pixel tolerance for line grouping (original algorithm). Good for consistent font sizes and line spacing. 'otsu': Uses Otsu's method on normalized gaps to find natural breaks between words. Resolution-invariant and works well with bimodal gap distributions. 'collimate': Uses greedy parent-child traversal to group detections. Good for skewed or curved text where bucket-based approaches fail.", + examples=["tolerance", "otsu", "collimate"], + json_schema_extra={ + "values_metadata": { + "tolerance": { + "name": "Tolerance-based", + "description": "Uses fixed pixel tolerance for line grouping. Good for consistent font sizes.", + }, + "otsu": { + "name": "Otsu Thresholding", + "description": "Resolution-invariant algorithm using normalized gaps. Works well with varying font sizes.", + }, + "collimate": { + "name": "Collimate (Skewed Text)", + "description": "Greedy parent-child traversal for grouping. Best for skewed or curved text.", + }, + } + }, + ) + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + ] + ) = Field( + title="OCR Detections", + description="OCR detection predictions from an OCR model. Should contain bounding boxes and class names with text content. Each detection represents a word, character, or text region that will be stitched together into coherent text. Supports object detection format with bounding boxes (xyxy) and class names in the data dictionary.", + examples=[ + "$steps.ocr_model.predictions", + "$steps.my_ocr_detection_model.predictions", + ], + ) + reading_direction: Literal[ + "left_to_right", + "right_to_left", + "vertical_top_to_bottom", + "vertical_bottom_to_top", + "auto", + ] = Field( + title="Reading Direction", + description="Direction to read and organize text detections. 'left_to_right': Standard horizontal reading (English, most languages). 'right_to_left': Right-to-left reading (Arabic, Hebrew). 'vertical_top_to_bottom': Vertical reading from top to bottom (Traditional Chinese, Japanese). 'vertical_bottom_to_top': Vertical reading from bottom to top (rare vertical formats). 'auto': Automatically detects reading direction based on average bounding box dimensions (width > height = horizontal, height >= width = vertical). Determines how detections are grouped into lines and sorted within lines.", + examples=["left_to_right", "right_to_left", "auto"], + json_schema_extra={ + "values_metadata": { + "left_to_right": { + "name": "Left To Right", + "description": "Standard left-to-right reading (e.g., English language)", + }, + "right_to_left": { + "name": "Right To Left", + "description": "Right-to-left reading (e.g., Arabic)", + }, + "vertical_top_to_bottom": { + "name": "Top To Bottom (Vertical)", + "description": "Vertical reading from top to bottom", + }, + "vertical_bottom_to_top": { + "name": "Bottom To Top (Vertical)", + "description": "Vertical reading from bottom to top", + }, + "auto": { + "name": "Auto", + "description": "Automatically detect the reading direction based on text arrangement.", + }, + } + }, + ) + tolerance: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + title="Tolerance", + description="Vertical (or horizontal for vertical text) distance threshold in pixels for grouping detections into the same line. Detections within this tolerance distance are grouped into the same line. Higher values group detections that are further apart (useful for text with variable line spacing or slanted text). Lower values create more lines (useful for tightly spaced text). Must be greater than zero.", + default=10, + examples=[10, 20, 5], + json_schema_extra={ + "relevant_for": { + "stitching_algorithm": { + "values": ["tolerance"], + "required": True, + }, + }, + }, + ) + delimiter: Union[str, Selector(kind=[STRING_KIND])] = Field( + title="Delimiter", + description="Optional delimiter string to insert between each text element (word/character) when stitching. Empty string (default) means no delimiter - text elements are concatenated directly. Useful for adding spaces between words, commas between elements, or custom separators. Example: use ' ' (space) to add spaces between words, or ',' to add commas.", + default="", + examples=["", " ", ",", "-"], + ) + otsu_threshold_multiplier: Union[float, Selector(kind=[FLOAT_KIND])] = Field( + title="Otsu Threshold Multiplier", + description="Multiplier applied to the Otsu-computed threshold when using the 'otsu' stitching algorithm. Values > 1.0 make word breaks less frequent (more conservative, fewer splits), values < 1.0 make word breaks more frequent (more aggressive, more splits). Default is 1.0 (use Otsu threshold as-is). Try 1.3-1.5 if words are being incorrectly split, or 0.7-0.9 if words are being incorrectly merged.", + default=1.0, + examples=[1.0, 1.3, 1.5, 0.8], + json_schema_extra={ + "relevant_for": { + "stitching_algorithm": { + "values": ["otsu"], + "required": True, + }, + }, + }, + ) + collimate_tolerance: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + title="Collimate Tolerance", + description="Pixel tolerance for the 'collimate' stitching algorithm. Controls how much vertical (for horizontal text) or horizontal (for vertical text) deviation is allowed when determining if a detection follows another in reading order. Higher values handle more skewed text but may incorrectly merge separate lines. Default is 10 pixels.", + default=10, + examples=[5, 10, 15, 20], + json_schema_extra={ + "relevant_for": { + "stitching_algorithm": { + "values": ["collimate"], + "required": True, + }, + }, + }, + ) + + @field_validator("tolerance") + @classmethod + def ensure_tolerance_greater_than_zero( + cls, value: Union[int, str] + ) -> Union[int, str]: + if isinstance(value, int) and value <= 0: + raise ValueError( + "Stitch OCR detections block expects `tolerance` to be greater than zero." + ) + return value + + @classmethod + def get_parameters_accepting_batches(cls) -> List[str]: + return ["predictions"] + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition(name="ocr_text", kind=[STRING_KIND]), + ] + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.0.0,<2.0.0" + + +def _native_xyxy_float(detections: Detections) -> np.ndarray: + """Native equivalent of the raw ``detections.xyxy`` numpy array used by the + otsu / collimate algorithms: pull the torch xyxy tensor to CPU numpy, keeping + float coordinates.""" + return detections.xyxy.detach().to("cpu").numpy() + + +def _native_xyxy_int(detections: Detections) -> np.ndarray: + """Native equivalent of ``detections.xyxy.round().astype(int)`` used by the + tolerance algorithm: pull to CPU numpy, round, cast to int.""" + return detections.xyxy.detach().to("cpu").numpy().round().astype(dtype=int) + + +def _native_class_names(detections: Detections) -> np.ndarray: + """Native equivalent of ``detections.data["class_name"]`` for OCR stitching. + + Resolves each detection's label the way the tensor serialiser does: a per-box + override wins (OCR-as-detector rides the recognised text on ``class``; the legacy + ``text`` key is also honoured), otherwise the label is resolved from + ``image_metadata['class_names']`` keyed by ``class_id`` โ€” which is where a plain + character/word *detection* model carries its label. Returned as an object array so + both fancy-indexing (tolerance) and scalar indexing (otsu / collimate) work + exactly as in the numpy file.""" + n = int(detections.xyxy.shape[0]) + class_names_map = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + class_ids = detections.class_id.detach().cpu().tolist() + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + names = [] + for i in range(n): + meta = bboxes_metadata[i] or {} + if CLASS_NAME_KEY in meta: + names.append(str(meta[CLASS_NAME_KEY])) + elif TEXT_KEY in meta: + names.append(str(meta[TEXT_KEY])) + else: + names.append(str(class_names_map.get(int(class_ids[i]), ""))) + return np.array(names, dtype=object) + + +def detect_reading_direction(detections: Detections) -> str: + if int(detections.xyxy.shape[0]) == 0: + return "left_to_right" + + xyxy = detections.xyxy.detach().to("cpu").numpy() + widths = xyxy[:, 2] - xyxy[:, 0] + heights = xyxy[:, 3] - xyxy[:, 1] + + avg_width = np.mean(widths) + avg_height = np.mean(heights) + + if avg_width > avg_height: + return "left_to_right" + else: + return "vertical_top_to_bottom" + + +class StitchOCRDetectionsBlockV2(WorkflowBlock): + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + predictions: Batch[Detections], + stitching_algorithm: str = "tolerance", + reading_direction: str = "auto", + tolerance: int = 10, + delimiter: str = "", + otsu_threshold_multiplier: float = 1.0, + collimate_tolerance: int = 10, + ) -> BlockResult: + if reading_direction == "auto": + reading_direction = detect_reading_direction(predictions[0]) + + if stitching_algorithm == "otsu": + return [ + adaptive_word_grouping( + detections=detections, + reading_direction=reading_direction, + delimiter=delimiter, + threshold_multiplier=otsu_threshold_multiplier, + ) + for detections in predictions + ] + elif stitching_algorithm == "collimate": + return [ + collimate_word_grouping( + detections=detections, + reading_direction=reading_direction, + delimiter=delimiter, + tolerance=collimate_tolerance, + ) + for detections in predictions + ] + else: + return [ + stitch_ocr_detections( + detections=detections, + reading_direction=reading_direction, + tolerance=tolerance, + delimiter=delimiter, + ) + for detections in predictions + ] + + +def stitch_ocr_detections( + detections: Detections, + reading_direction: str = "left_to_right", + tolerance: int = 10, + delimiter: str = "", +) -> Dict[str, str]: + """ + Stitch OCR detections into coherent text based on spatial arrangement. + + Args: + detections: Native inference_models Detections object containing OCR results + reading_direction: Direction to read text ("left_to_right", "right_to_left", + "vertical_top_to_bottom", "vertical_bottom_to_top") + tolerance: Vertical tolerance for grouping text into lines + + Returns: + Dict containing stitched OCR text under 'ocr_text' key + """ + if int(detections.xyxy.shape[0]) == 0: + return {"ocr_text": ""} + + xyxy = _native_xyxy_int(detections) + class_names = _native_class_names(detections) + + # Prepare coordinates based on reading direction + xyxy = prepare_coordinates(xyxy, reading_direction) + + # Group detections into lines + boxes_by_line = group_detections_by_line(xyxy, reading_direction, tolerance) + # Sort lines based on reading direction + lines = sorted( + boxes_by_line.keys(), reverse=reading_direction in ["vertical_bottom_to_top"] + ) + + # Build final text + ordered_class_names = [] + for i, key in enumerate(lines): + line_data = boxes_by_line[key] + line_xyxy = np.array(line_data["xyxy"]) + line_idx = np.array(line_data["idx"]) + + # Sort detections within line + sort_idx = sort_line_detections(line_xyxy, reading_direction) + + # Add sorted class names for this line + ordered_class_names.extend(class_names[line_idx[sort_idx]]) + + # Add line separator if not last line + if i < len(lines) - 1: + ordered_class_names.append(get_line_separator(reading_direction)) + + return {"ocr_text": delimiter.join(ordered_class_names)} + + +def prepare_coordinates( + xyxy: np.ndarray, + reading_direction: str, +) -> np.ndarray: + """Prepare coordinates based on reading direction.""" + if reading_direction in ["vertical_top_to_bottom", "vertical_bottom_to_top"]: + # Swap x and y coordinates: [x1,y1,x2,y2] -> [y1,x1,y2,x2] + return xyxy[:, [1, 0, 3, 2]] + return xyxy + + +def group_detections_by_line( + xyxy: np.ndarray, + reading_direction: str, + tolerance: int, +) -> Dict[float, Dict[str, List]]: + """Group detections into lines based on primary coordinate.""" + # After prepare_coordinates swap, we always group by y ([:, 1]) + primary_coord = xyxy[:, 1] # This is y for horizontal, swapped x for vertical + + # Round primary coordinate to group into lines + rounded_primary = np.round(primary_coord / tolerance) * tolerance + + boxes_by_line = {} + # Group bounding boxes and associated indices by line + for i, (bbox, line_pos) in enumerate(zip(xyxy, rounded_primary)): + if line_pos not in boxes_by_line: + boxes_by_line[line_pos] = {"xyxy": [bbox], "idx": [i]} + else: + boxes_by_line[line_pos]["xyxy"].append(bbox) + boxes_by_line[line_pos]["idx"].append(i) + + return boxes_by_line + + +def sort_line_detections( + line_xyxy: np.ndarray, + reading_direction: str, +) -> np.ndarray: + """Sort detections within a line based on reading direction.""" + # After prepare_coordinates swap, we always sort by x ([:, 0]) + if reading_direction in ["left_to_right", "vertical_top_to_bottom"]: + return line_xyxy[:, 0].argsort() # Sort by x1 (original x or swapped y) + else: # right_to_left or vertical_bottom_to_top + return (-line_xyxy[:, 0]).argsort() # Sort by -x1 (original -x or swapped -y) + + +def get_line_separator(reading_direction: str) -> str: + """Get the appropriate separator based on reading direction.""" + return "\n" if reading_direction in ["left_to_right", "right_to_left"] else " " + + +def find_otsu_threshold(gaps: np.ndarray) -> tuple[float, bool]: + """Find natural break between intra-word and inter-word gaps using Otsu's method. + + This is a resolution-invariant approach that finds the optimal threshold + to separate two classes of gaps (e.g., gaps within words vs gaps between words). + + Also detects whether the distribution is bimodal (two distinct groups) or + unimodal (single group, suggesting single word or uniform spacing). + + Args: + gaps: Array of normalized gap values + + Returns: + Tuple of (threshold, is_bimodal): + - threshold: Optimal threshold value that maximizes between-class variance + - is_bimodal: True if distribution appears bimodal, False if unimodal + """ + if len(gaps) < 2: + return 0.0, False + + # Create histogram of gaps + hist, bin_edges = np.histogram(gaps, bins=min(50, len(gaps))) + bin_centers = (bin_edges[:-1] + bin_edges[1:]) / 2 + + best_thresh = 0.0 + best_variance = 0.0 + best_below_mean = 0.0 + best_above_mean = 0.0 + + for t in bin_centers: + below = gaps[gaps <= t] + above = gaps[gaps > t] + + if len(below) == 0 or len(above) == 0: + continue + + # Between-class variance (Otsu's criterion) + variance = len(below) * len(above) * (below.mean() - above.mean()) ** 2 + + if variance > best_variance: + best_variance = variance + best_thresh = t + best_below_mean = below.mean() + best_above_mean = above.mean() + + # Check if distribution is bimodal using several heuristics: + # 1. The gap between class means should be significant relative to overall spread + # 2. There should be meaningful absolute separation between classes + + overall_std = gaps.std() + overall_mean = gaps.mean() + + # Separation ratio: how far apart are the two class means relative to overall std + mean_separation = abs(best_above_mean - best_below_mean) + separation_ratio = mean_separation / overall_std if overall_std > 0 else 0 + + # Bimodality criteria - MUST have meaningful word gaps (not just outliers): + # The key insight is that real word gaps are typically 0.5+ in normalized units. + # A distribution with all gaps < 0.3 is unimodal (single word), even if there + # are outliers (like overlapping characters with negative gaps) that inflate + # the mean separation. + # + # Primary criterion: above-class mean must indicate actual word gaps exist + has_positive_word_gaps = ( + best_above_mean > 0.3 + ) # Word gaps should be clearly positive + + # Secondary criterion: if we have good separation AND positive gaps + has_good_relative_separation = separation_ratio > 1.5 and mean_separation > 0.3 + + # Must have positive word gaps to be considered bimodal + is_bimodal = has_positive_word_gaps and ( + mean_separation > 0.3 or has_good_relative_separation + ) + + return best_thresh, is_bimodal + + +def adaptive_word_grouping( + detections: Detections, + reading_direction: str, + delimiter: str = "", + threshold_multiplier: float = 1.0, +) -> Dict[str, str]: + """Stitch OCR detections using adaptive gap analysis with Otsu thresholding. + + This approach is resolution-invariant because it normalizes gaps by local + character dimensions. It works well with bimodal gap distributions + (e.g., character-level vs word-level spacing). + + The algorithm computes a global threshold across all lines to leverage + the full dataset of gaps, which provides more robust Otsu thresholding + than per-line computation. + + Args: + detections: Native inference_models Detections object containing OCR results + reading_direction: Direction to read text + delimiter: String to insert between text elements + threshold_multiplier: Multiplier applied to Otsu threshold (>1.0 = fewer word breaks, <1.0 = more word breaks) + + Returns: + Dict containing stitched OCR text under 'ocr_text' key + """ + if int(detections.xyxy.shape[0]) == 0: + return {"ocr_text": ""} + + xyxy = _native_xyxy_float(detections) + class_names = _native_class_names(detections) + + # Determine if we're working with vertical text + is_vertical = reading_direction in [ + "vertical_top_to_bottom", + "vertical_bottom_to_top", + ] + + # For vertical text, swap x/y for processing + if is_vertical: + # Swap coordinates: treat y as x for sorting + x_centers = (xyxy[:, 1] + xyxy[:, 3]) / 2 # y becomes primary axis + y_centers = (xyxy[:, 0] + xyxy[:, 2]) / 2 # x becomes secondary axis + widths = xyxy[:, 3] - xyxy[:, 1] # height becomes "width" + heights = xyxy[:, 2] - xyxy[:, 0] # width becomes "height" + else: + x_centers = (xyxy[:, 0] + xyxy[:, 2]) / 2 + y_centers = (xyxy[:, 1] + xyxy[:, 3]) / 2 + widths = xyxy[:, 2] - xyxy[:, 0] + heights = xyxy[:, 3] - xyxy[:, 1] + + # First, group detections into lines based on y-coordinate clustering + # Use adaptive threshold based on median height + median_height = np.median(heights) + line_tolerance = median_height * 0.5 + + # Sort by y to group into lines + y_sorted_indices = np.argsort(y_centers) + + lines = [] + current_line = [y_sorted_indices[0]] + current_line_y = y_centers[y_sorted_indices[0]] + + for idx in y_sorted_indices[1:]: + if abs(y_centers[idx] - current_line_y) <= line_tolerance: + current_line.append(idx) + # Update line y as running average + current_line_y = np.mean([y_centers[i] for i in current_line]) + else: + lines.append(current_line) + current_line = [idx] + current_line_y = y_centers[idx] + lines.append(current_line) + + # Sort lines by y position + line_y_positions = [np.mean([y_centers[i] for i in line]) for line in lines] + if reading_direction in ["vertical_bottom_to_top"]: + sorted_line_indices = np.argsort(line_y_positions)[::-1] + else: + sorted_line_indices = np.argsort(line_y_positions) + + # First pass: compute normalized gaps for ALL lines to get global threshold + all_normalized_gaps = [] + line_data = [] # Store sorted line info for second pass + + for line_idx in sorted_line_indices: + line = lines[line_idx] + + if len(line) == 1: + line_data.append((line, None, None, None)) + continue + + # Sort detections in line by x position + line_x_centers = x_centers[line] + line_widths = widths[line] + + if reading_direction in ["right_to_left", "vertical_bottom_to_top"]: + x_sorted_order = np.argsort(line_x_centers)[::-1] + else: + x_sorted_order = np.argsort(line_x_centers) + + sorted_line = [line[i] for i in x_sorted_order] + sorted_x_centers = line_x_centers[x_sorted_order] + sorted_widths = line_widths[x_sorted_order] + + # Compute normalized gaps for this line + normalized_gaps = [] + for i in range(1, len(sorted_line)): + prev_idx, curr_idx = i - 1, i + # Raw gap between detection edges + if reading_direction in ["right_to_left", "vertical_bottom_to_top"]: + raw_gap = ( + sorted_x_centers[prev_idx] + - sorted_x_centers[curr_idx] + - (sorted_widths[prev_idx] + sorted_widths[curr_idx]) / 2 + ) + else: + raw_gap = ( + sorted_x_centers[curr_idx] + - sorted_x_centers[prev_idx] + - (sorted_widths[prev_idx] + sorted_widths[curr_idx]) / 2 + ) + + # Normalize by local character scale + local_scale = (sorted_widths[prev_idx] + sorted_widths[curr_idx]) / 2 + if local_scale > 0: + normalized_gaps.append(raw_gap / local_scale) + else: + normalized_gaps.append(0.0) + + normalized_gaps = np.array(normalized_gaps) + all_normalized_gaps.extend(normalized_gaps.tolist()) + line_data.append( + (sorted_line, sorted_x_centers, sorted_widths, normalized_gaps) + ) + + # Compute global threshold using all gaps, then apply multiplier + all_normalized_gaps = np.array(all_normalized_gaps) + global_threshold, is_bimodal = find_otsu_threshold(all_normalized_gaps) + global_threshold *= threshold_multiplier + + # Second pass: use global threshold to group words + all_text_parts = [] + + for sorted_line, sorted_x_centers, sorted_widths, normalized_gaps in line_data: + if normalized_gaps is None: + # Single detection in line + all_text_parts.append(class_names[sorted_line[0]]) + continue + + # If distribution is not bimodal (likely single word or uniform spacing), + # treat all detections as a single word to avoid incorrect splitting + if not is_bimodal: + word_text = delimiter.join([class_names[idx] for idx in sorted_line]) + all_text_parts.append(word_text) + continue + + # Group into words based on global threshold + words = [[sorted_line[0]]] + for i, det_idx in enumerate(sorted_line[1:]): + if normalized_gaps[i] > global_threshold: + words.append([det_idx]) + else: + words[-1].append(det_idx) + + # Build text for this line + line_text_parts = [] + for word in words: + word_text = delimiter.join([class_names[idx] for idx in word]) + line_text_parts.append(word_text) + + # Join words with space (or delimiter if specified and non-empty) + word_separator = " " if delimiter == "" else delimiter + all_text_parts.append(word_separator.join(line_text_parts)) + + # Join lines with appropriate separator + line_separator = get_line_separator(reading_direction) + return {"ocr_text": line_separator.join(all_text_parts)} + + +class CollimateDetection: + """Helper class for collimate algorithm to store detection properties.""" + + def __init__(self, xyxy: np.ndarray, class_name: str, idx: int): + self.x = (xyxy[0] + xyxy[2]) / 2 + self.y = (xyxy[1] + xyxy[3]) / 2 + self.width = xyxy[2] - xyxy[0] + self.height = xyxy[3] - xyxy[1] + self.class_name = class_name + self.idx = idx # Original index for tracking + + def __repr__(self) -> str: + return f"{self.class_name}" + + +def _detection_follows( + parent: CollimateDetection, + child: CollimateDetection, + reading_direction: str, + tolerance: int, +) -> bool: + """Check if child detection follows parent in reading order within tolerance. + + For horizontal text: child should be roughly on the same line (similar y) + and to the right of parent. + + For vertical text: child should be roughly in the same column (similar x) + and below parent. + + Args: + parent: The reference detection + child: The detection to check + reading_direction: Reading direction + tolerance: Pixel tolerance for alignment + + Returns: + True if child follows parent in reading order + """ + is_vertical = reading_direction in [ + "vertical_top_to_bottom", + "vertical_bottom_to_top", + ] + + if is_vertical: + # For vertical text, check if x-coordinates align and child is below/above + # Use max width to handle narrow letters + width_tolerance = max(parent.width, child.width) / 2 + x_aligned = abs(parent.x - child.x) < width_tolerance + tolerance + + if reading_direction == "vertical_top_to_bottom": + return x_aligned and parent.y <= child.y + else: # vertical_bottom_to_top + return x_aligned and parent.y >= child.y + else: + # For horizontal text, check if y-coordinates align and child is to the right/left + # Use max height to handle varying letter sizes + height_tolerance = max(parent.height, child.height) / 2 + y_aligned = abs(parent.y - child.y) < height_tolerance + tolerance + + if reading_direction == "left_to_right": + return y_aligned and parent.x <= child.x + else: # right_to_left + return y_aligned and parent.x >= child.x + + +def _sort_detections_for_collimate( + detections: List[CollimateDetection], + reading_direction: str, +) -> List[CollimateDetection]: + """Sort detections by primary reading coordinate.""" + is_vertical = reading_direction in [ + "vertical_top_to_bottom", + "vertical_bottom_to_top", + ] + + if is_vertical: + # Sort by y for vertical text + reverse = reading_direction == "vertical_bottom_to_top" + return sorted(detections, key=lambda d: d.y, reverse=reverse) + else: + # Sort by x for horizontal text + reverse = reading_direction == "right_to_left" + return sorted(detections, key=lambda d: d.x, reverse=reverse) + + +def _get_line_avg_coord( + line: List[CollimateDetection], + reading_direction: str, +) -> float: + """Get average coordinate for sorting lines.""" + if len(line) == 0: + return 0.0 + + is_vertical = reading_direction in [ + "vertical_top_to_bottom", + "vertical_bottom_to_top", + ] + + if is_vertical: + # For vertical text, lines are columns - sort by x + return sum(d.x for d in line) / len(line) + else: + # For horizontal text, lines are rows - sort by y + return sum(d.y for d in line) / len(line) + + +def collimate_word_grouping( + detections: Detections, + reading_direction: str, + delimiter: str = "", + tolerance: int = 10, +) -> Dict[str, str]: + """Stitch OCR detections using greedy parent-child traversal (collimate algorithm). + + This algorithm is good for skewed or curved text where traditional bucket-based + line grouping may fail. It works by: + 1. Sorting detections by primary reading coordinate + 2. Starting with the first detection as a "parent" + 3. Finding all detections that "follow" the parent (within tolerance) + 4. Building lines/columns through greedy traversal + + Args: + detections: Native inference_models Detections object containing OCR results + reading_direction: Direction to read text + delimiter: String to insert between characters within words + tolerance: Pixel tolerance for alignment + + Returns: + Dict containing stitched OCR text under 'ocr_text' key + """ + n = int(detections.xyxy.shape[0]) + if n == 0: + return {"ocr_text": ""} + + xyxy = _native_xyxy_float(detections) + class_names = _native_class_names(detections) + + # Convert to CollimateDetection objects + coll_detections = [CollimateDetection(xyxy[i], class_names[i], i) for i in range(n)] + + # Sort by primary reading coordinate + coll_detections = _sort_detections_for_collimate(coll_detections, reading_direction) + + if len(coll_detections) == 0: + return {"ocr_text": ""} + + # Build lines through greedy parent-child traversal + remaining = list(coll_detections) + lines: List[List[CollimateDetection]] = [[remaining.pop(0)]] + + while len(remaining) > 0: + found_child = False + + # Try to extend existing lines + for line in lines: + parent = line[-1] + + # Find children that follow parent + for det in remaining.copy(): + if _detection_follows(parent, det, reading_direction, tolerance): + found_child = True + line.append(det) + parent = det # New parent for next iteration + remaining.remove(det) + + # If no children found for any line, start a new line + if not found_child and len(remaining) > 0: + lines.append([remaining.pop(0)]) + + # Sort lines by their average secondary coordinate + is_vertical = reading_direction in [ + "vertical_top_to_bottom", + "vertical_bottom_to_top", + ] + if is_vertical: + # For vertical text, sort columns left-to-right (or right-to-left) + reverse = reading_direction == "vertical_bottom_to_top" + else: + # For horizontal text, sort rows top-to-bottom + reverse = False + + lines = sorted( + lines, + key=lambda line: _get_line_avg_coord(line, reading_direction), + reverse=reverse, + ) + + # Build output text + line_texts = [] + for line in lines: + # Characters within a line are concatenated with delimiter + line_text = delimiter.join(d.class_name for d in line) + line_texts.append(line_text) + + # Join lines with appropriate separator + line_separator = get_line_separator(reading_direction) + return {"ocr_text": line_separator.join(line_texts)} diff --git a/inference/core/workflows/core_steps/transformations/track_class_lock/v1_tensor.py b/inference/core/workflows/core_steps/transformations/track_class_lock/v1_tensor.py new file mode 100644 index 0000000000..b98cb11696 --- /dev/null +++ b/inference/core/workflows/core_steps/transformations/track_class_lock/v1_tensor.py @@ -0,0 +1,400 @@ +"""Tensor-native sibling of track_class_lock/v1.py. + +The voting/lock logic is numpy/``sv.Detections``-based and stateful across +engine runs, so the native input is materialised to ``sv.Detections`` at the +boundary (tracker_id / class_name / detection_id carried from +``bboxes_metadata`` / ``image_metadata[CLASS_NAMES_KEY]`` into ``.data``), the +lock/vote state machine runs on it (``self._per_video_state`` persists across +runs, keyed by video identifier), and the result is repacked natively: the +locked ``class``, ``class_id``, ``confidence`` and a boolean ``class_locked`` +flag land in ``bboxes_metadata[i]``, and ``image_metadata[CLASS_NAMES_KEY]`` is +rebuilt so a relabelled class id resolves to a name downstream. + +Keypoint input arrives as a ``(KeyPoints, Detections)`` tuple; the bbox +component drives the voting and the relabelled tuple is returned so keypoints +survive. Instance-segmentation masks are carried through unchanged. +""" + +from collections import OrderedDict, defaultdict +from copy import deepcopy +from typing import List, Optional, Set, Tuple, Type, Union + +import numpy as np +import supervision as sv + +from inference.core.workflows.core_steps.common.tensor_native import ( + HOST_MIRROR_KEYS, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.transformations.track_class_lock.v1 import ( + MAX_TRACKED_VIDEOS, + BlockManifest, + _find_lock_to_inherit, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + TRACKER_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +OUTPUT_KEY: str = "tracked_detections" +CLASS_LOCKED_KEY: str = "class_locked" + +_CLASS_NAME_DATA_KEY = "class_name" + +TensorNativeDetectionsLike = Union[Detections, InstanceDetections] +TensorNativeTrackInput = Union[ + Detections, + InstanceDetections, + Tuple[KeyPoints, Optional[Detections]], +] + + +def _resolve_class_names(detections: TensorNativeDetectionsLike) -> List[Optional[str]]: + """Per-box class name: prefer ``bboxes_metadata[i]["class"]`` (the + ``CLASS_NAME_KEY`` convention), fall back to the + ``image_metadata[CLASS_NAMES_KEY]`` class_id -> name map, else ``None`` + (voting then uses ``str(class_id)``).""" + n = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + class_names_map = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + class_id = detections.class_id.detach().to("cpu").numpy() + resolved: List[Optional[str]] = [] + for index in range(n): + per_box = (bboxes_metadata[index] or {}).get(CLASS_NAME_KEY) + if per_box is not None: + resolved.append(str(per_box)) + continue + mapped = class_names_map.get(int(class_id[index])) + resolved.append(str(mapped) if mapped is not None else None) + return resolved + + +def _to_supervision_boundary(detections: TensorNativeDetectionsLike) -> sv.Detections: + """Materialise a native ``Detections``/``InstanceDetections`` to + ``sv.Detections`` carrying tracker_id + class_name + detection_id in + ``.data``. Masks are dropped โ€” the voting/lock logic only uses + xyxy/class_id/confidence/tracker_id/class_name.""" + n = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata or [{} for _ in range(n)] + xyxy = detections.xyxy.detach().to("cpu").numpy().astype(np.float64) + class_id = detections.class_id.detach().to("cpu").numpy() + confidence = detections.confidence.detach().to("cpu").numpy() + tracker_ids = np.array( + [(bboxes_metadata[i] or {}).get(TRACKER_ID_KEY, -1) for i in range(n)], + dtype=np.int64, + ) + detection_ids = np.array( + [(bboxes_metadata[i] or {}).get(DETECTION_ID_KEY, "") for i in range(n)], + dtype=object, + ) + class_names = np.array( + [name if name is not None else "" for name in _resolve_class_names(detections)], + dtype=object, + ) + data = { + _CLASS_NAME_DATA_KEY: class_names, + DETECTION_ID_KEY: detection_ids, + } + return sv.Detections( + xyxy=xyxy if n else np.empty((0, 4), dtype=np.float64), + class_id=class_id if n else np.empty((0,), dtype=int), + confidence=confidence if n else np.empty((0,), dtype=np.float32), + tracker_id=tracker_ids if n else np.empty((0,), dtype=np.int64), + data=data, + ) + + +def _vote_and_lock( + self: "TrackClassLockBlockV1", + image: WorkflowImageData, + dets: sv.Detections, + min_votes: int, + vote_confidence: float, + lead_margin: int, + switch_after: int, + state_ttl: int, + reattach_window: int, + reattach_iou: float, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Run the majority-vote / lock state machine over a boundary + ``sv.Detections``. Mutates ``dets.class_id``/``dets.confidence`` and the + ``class_name`` column in place and returns + ``(class_names, class_id, locked_flags)`` for repacking into the native + output.""" + video_id = image.video_metadata.video_identifier + video_state = self._per_video_state.setdefault(video_id, {"tracks": {}, "frame": 0}) + self._per_video_state.move_to_end(video_id) + while len(self._per_video_state) > MAX_TRACKED_VIDEOS: + self._per_video_state.popitem(last=False) + video_state["frame"] += 1 + frame = video_state["frame"] + tracks = video_state["tracks"] + + n = len(dets) + locked_flags = np.zeros(n, dtype=bool) + + if dets.confidence is None and n > 0: + dets.confidence = np.ones(n, dtype=np.float32) + + class_names = dets.data.get(_CLASS_NAME_DATA_KEY) + if class_names is not None: + class_names = np.asarray(class_names).astype(object) + dets.data[_CLASS_NAME_DATA_KEY] = class_names + active_tids: Set[int] = set() + if dets.tracker_id is not None: + active_tids = {int(t) for t in dets.tracker_id if t is not None and int(t) >= 0} + for i in range(n): + tid = dets.tracker_id[i] + if tid is None or int(tid) < 0: + continue + tid = int(tid) + if tid not in tracks: + inherited = _find_lock_to_inherit( + tracks=tracks, + xyxy=dets.xyxy[i], + frame=frame, + reattach_window=reattach_window, + reattach_iou=reattach_iou, + active_tids=active_tids, + ) + if inherited is not None: + tracks[tid] = tracks.pop(inherited) + st = tracks.setdefault( + tid, + { + "votes": defaultdict(int), + "conf_sum": defaultdict(float), + "class_ids": {}, + "locked": None, + "last_seen": frame, + "challenger": None, + "streak": 0, + "streak_conf": 0.0, + "last_xyxy": None, + }, + ) + st["last_seen"] = frame + st["last_xyxy"] = np.array(dets.xyxy[i], copy=True) + + if class_names is not None and str(class_names[i]) != "": + cname = str(class_names[i]) + elif dets.class_id is not None: + cname = str(dets.class_id[i]) + else: + continue + conf = float(dets.confidence[i]) if dets.confidence is not None else 1.0 + qualifying = conf >= vote_confidence + + if st["locked"] is None: + if qualifying: + st["votes"][cname] += 1 + st["conf_sum"][cname] += conf + if dets.class_id is not None: + st["class_ids"][cname] = int(dets.class_id[i]) + if st["votes"]: + ranked = sorted(st["votes"].items(), key=lambda kv: kv[1], reverse=True) + top_c, top_v = ranked[0] + runner_v = ranked[1][1] if len(ranked) > 1 else 0 + if top_v >= min_votes and top_v - runner_v >= lead_margin: + st["locked"] = top_c + else: + if qualifying and cname != st["locked"]: + if cname == st["challenger"]: + st["streak"] += 1 + st["streak_conf"] += conf + else: + st["challenger"] = cname + st["streak"] = 1 + st["streak_conf"] = conf + if dets.class_id is not None: + st["class_ids"][cname] = int(dets.class_id[i]) + if st["streak"] >= switch_after: + new = st["challenger"] + st["locked"] = new + st["votes"] = defaultdict(int, {new: st["streak"]}) + st["conf_sum"] = defaultdict(float, {new: st["streak_conf"]}) + st["challenger"], st["streak"], st["streak_conf"] = None, 0, 0.0 + else: + st["challenger"], st["streak"], st["streak_conf"] = None, 0, 0.0 + + if st["locked"] is not None: + lc = st["locked"] + if class_names is not None: + class_names[i] = lc + if dets.class_id is not None and lc in st["class_ids"]: + dets.class_id[i] = st["class_ids"][lc] + denom = max(st["votes"][lc], 1) + dets.confidence[i] = min(1.0, st["conf_sum"][lc] / denom) + locked_flags[i] = True + + stale = [t for t, st in tracks.items() if frame - st["last_seen"] > state_ttl] + for t in stale: + del tracks[t] + + out_class_names = ( + np.asarray(class_names) + if class_names is not None + else np.empty((n,), dtype=object) + ) + out_class_id = ( + dets.class_id.copy() if dets.class_id is not None else np.empty((n,), dtype=int) + ) + return out_class_names, out_class_id, locked_flags + + +def _repack_native( + detections: TensorNativeDetectionsLike, + class_names: np.ndarray, + class_id: np.ndarray, + confidence: np.ndarray, + locked_flags: np.ndarray, +) -> TensorNativeDetectionsLike: + """Write the locked/relabelled class, confidence and ``class_locked`` flag back + into a COPY of the native prediction. ``bboxes_metadata[i]`` receives the new + ``class`` name, ``class_locked`` flag (and ``class_id``/``confidence`` mirrors); + ``image_metadata[CLASS_NAMES_KEY]`` is rebuilt so any relabelled class id resolves + to a name. The xyxy/mask tensors are untouched; class_id/confidence tensors are + rebuilt from the (possibly mutated) numpy values on the original device/dtype. + """ + result = deepcopy(detections) + n = int(result.xyxy.shape[0]) + new_class_id = ( + result.class_id.new_tensor(class_id.tolist()) if n else result.class_id + ) + new_confidence = ( + result.confidence.new_tensor(confidence.tolist()) if n else result.confidence + ) + result.class_id = new_class_id + result.confidence = new_confidence + + class_names_map = dict((result.image_metadata or {}).get(CLASS_NAMES_KEY) or {}) + bboxes_metadata = result.bboxes_metadata or [{} for _ in range(n)] + new_bboxes_metadata: List[dict] = [] + for index in range(n): + entry = dict(bboxes_metadata[index] or {}) + cid = int(class_id[index]) + name = class_names[index] if index < len(class_names) else None + if name is not None and str(name) != "": + entry[CLASS_NAME_KEY] = str(name) + class_names_map[cid] = str(name) + entry[CLASS_LOCKED_KEY] = bool(locked_flags[index]) + # Drop the per-box host mirror: class_id / confidence were rebuilt + # above, so a carried mirror would be stale โ€” consumers fall back to + # tensor reads. + for mirror_key in HOST_MIRROR_KEYS: + entry.pop(mirror_key, None) + new_bboxes_metadata.append(entry) + result.bboxes_metadata = new_bboxes_metadata if n else None + + image_metadata = dict(result.image_metadata or {}) + image_metadata[CLASS_NAMES_KEY] = class_names_map + result.image_metadata = image_metadata + return result + + +class TrackClassLockBlockV1(WorkflowBlock): + def __init__(self): + self._per_video_state: "OrderedDict[str, dict]" = OrderedDict() + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + detections: TensorNativeTrackInput, + min_votes: int, + vote_confidence: float, + lead_margin: int, + switch_after: int, + state_ttl: int, + reattach_window: int, + reattach_iou: float, + ) -> BlockResult: + # selector-provided params bypass the manifest's pydantic Field bounds, + # so the same constraints are re-checked here at runtime + if min_votes < 1: + raise ValueError(f"`min_votes` must be >= 1, got {min_votes}") + if not 0.0 <= vote_confidence <= 1.0: + raise ValueError( + f"`vote_confidence` must be within [0.0, 1.0], got {vote_confidence}" + ) + if lead_margin < 0: + raise ValueError(f"`lead_margin` must be >= 0, got {lead_margin}") + if switch_after < 1: + raise ValueError(f"`switch_after` must be >= 1, got {switch_after}") + if state_ttl < 1: + raise ValueError(f"`state_ttl` must be >= 1, got {state_ttl}") + if reattach_window < 0: + raise ValueError(f"`reattach_window` must be >= 0, got {reattach_window}") + if not 0.0 <= reattach_iou <= 1.0: + raise ValueError( + f"`reattach_iou` must be within [0.0, 1.0], got {reattach_iou}" + ) + if reattach_window > state_ttl: + raise ValueError( + f"`reattach_window` ({reattach_window}) must not exceed `state_ttl` " + f"({state_ttl}) - lost tracks are purged after `state_ttl` frames, so " + "re-attachment beyond that point can never happen." + ) + + key_points, bbox_detections = split_key_point_prediction(detections) + + n = int(bbox_detections.xyxy.shape[0]) + if n > 0 and not _has_tracker_ids(bbox_detections): + raise ValueError( + f"tracker_id not initialized, {self.__class__.__name__} requires " + "detections to be tracked" + ) + + sv_dets = _to_supervision_boundary(bbox_detections) + class_names, class_id, locked_flags = _vote_and_lock( + self, + image=image, + dets=sv_dets, + min_votes=min_votes, + vote_confidence=vote_confidence, + lead_margin=lead_margin, + switch_after=switch_after, + state_ttl=state_ttl, + reattach_window=reattach_window, + reattach_iou=reattach_iou, + ) + confidence = ( + sv_dets.confidence + if sv_dets.confidence is not None + else np.ones(n, dtype=np.float32) + ) + native_output = _repack_native( + detections=bbox_detections, + class_names=class_names, + class_id=class_id, + confidence=confidence, + locked_flags=locked_flags, + ) + if key_points is not None: + return {OUTPUT_KEY: (key_points, native_output)} + return {OUTPUT_KEY: native_output} + + +def _has_tracker_ids(detections: TensorNativeDetectionsLike) -> bool: + """True when at least one box carries a ``tracker_id`` in + ``bboxes_metadata``.""" + if detections.bboxes_metadata is None: + return False + return any( + (entry or {}).get(TRACKER_ID_KEY) is not None + for entry in detections.bboxes_metadata + ) diff --git a/inference/core/workflows/core_steps/visualizations/background_color/v1.py b/inference/core/workflows/core_steps/visualizations/background_color/v1.py index 30c50ce32f..b9160b3e29 100644 --- a/inference/core/workflows/core_steps/visualizations/background_color/v1.py +++ b/inference/core/workflows/core_steps/visualizations/background_color/v1.py @@ -148,8 +148,13 @@ def run( opacity: Optional[float], ) -> BlockResult: annotator = self.getAnnotator(color, opacity) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/background_color/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/background_color/v1_tensor.py new file mode 100644 index 0000000000..76d42540ed --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/background_color/v1_tensor.py @@ -0,0 +1,170 @@ +from typing import Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.annotators.background_color import ( + BackgroundColorAnnotator, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + PredictionsVisualizationBlock, + PredictionsVisualizationManifest, + to_supervision_for_annotation, +) +from inference.core.workflows.core_steps.visualizations.common.utils import str_to_color +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + STRING_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/background_color_visualization@v1" +SHORT_DESCRIPTION = ( + "Apply a mask to cover all areas outside the detected regions in an image." +) +LONG_DESCRIPTION = """ +Apply a colored overlay to areas outside detected regions, effectively masking the background while preserving detected objects at their original appearance. + +## How This Block Works + +This block takes an image and detection predictions and applies a colored overlay to all areas outside of the detected objects, leaving the detected regions unchanged. The block: + +1. Takes an image and predictions as input +2. Creates a colored mask layer with the specified background color +3. Identifies detected regions from bounding boxes or segmentation masks (preserves detected objects) +4. Applies the colored overlay to all areas outside the detected regions with the specified opacity +5. Blends the colored overlay with the original image based on the opacity setting +6. Returns an annotated image where detected objects appear unchanged, while the background is filled with the specified color + +The block works with both object detection predictions (using bounding boxes) and instance segmentation predictions (using masks). When masks are available, it preserves the exact shape of detected objects; otherwise, it uses bounding box regions. The opacity parameter controls how transparent or opaque the background overlay is, allowing you to create effects ranging from subtle background dimming (low opacity) to complete background replacement (high opacity). This creates a visual focus effect that highlights the detected objects by de-emphasizing or completely hiding the background. + +## Common Use Cases + +- **Object Focus and Highlighting**: Highlight detected objects by dimming or replacing the background, making objects stand out for presentations, documentation, or user interfaces +- **Background Removal Effects**: Create images where backgrounds are replaced with solid colors or semi-transparent overlays for product photography, content creation, or design workflows +- **Privacy and Anonymization**: Mask backgrounds while preserving detected objects (e.g., people, vehicles) to anonymize images, protect privacy, or comply with data protection requirements +- **Visual Debugging and Validation**: Dim backgrounds to focus attention on detected regions when validating model performance, checking detection accuracy, or debugging detection results +- **Presentation and Documentation**: Create clean, professional visualizations for reports, presentations, or documentation where you want to emphasize detected objects without distracting backgrounds +- **Content Creation and Editing**: Prepare images for further processing, compositing, or editing by isolating detected objects with colored backgrounds for easier manipulation or integration into other workflows + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Bounding Box Visualization, Polygon Visualization) to add additional annotations on top of the background-colored image for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save images with background coloring for documentation, reporting, or archiving +- **Webhook blocks** to send visualized results with background coloring to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with background coloring as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with background coloring for live monitoring, tracking visualization, or post-processing analysis +""" + + +class BackgroundColorManifest(PredictionsVisualizationManifest): + type: Literal[f"{TYPE}", "BackgroundColorVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Background Color Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-fill-drip", + "blockPriority": 3, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + color: Union[str, Selector(kind=[STRING_KIND])] = Field( # type: ignore + description="Color to use for the background overlay. Areas outside detected regions will be filled with this color. Can be a color name (e.g., 'BLACK', 'WHITE') or color code in HEX format (e.g., '#000000') or RGB format (e.g., 'rgb(0, 0, 0)').", + default="BLACK", + examples=["WHITE", "#FFFFFF", "rgb(255, 255, 255)" "$inputs.background_color"], + ) + + opacity: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + description="Opacity of the background overlay, ranging from 0.0 (fully transparent, original background visible) to 1.0 (fully opaque, complete background replacement). Values between 0.0 and 1.0 create a blend between the original image and the background color. Lower values create subtle background dimming, while higher values create stronger background replacement effects.", + default=0.5, + examples=[0.5, "$inputs.opacity"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class BackgroundColorVisualizationBlockV1(PredictionsVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BackgroundColorManifest + + def getAnnotator( + self, + color: str, + opacity: float, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color, + opacity, + ], + ) + ) + + if key not in self.annotatorCache: + background_color = str_to_color(color) + self.annotatorCache[key] = BackgroundColorAnnotator( + color=background_color, + opacity=opacity, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color: str, + opacity: Optional[float], + ) -> BlockResult: + predictions = to_supervision_for_annotation(predictions) + annotator = self.getAnnotator(color, opacity) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/blur/v1.py b/inference/core/workflows/core_steps/visualizations/blur/v1.py index 24dfc2156c..79f6f88b75 100644 --- a/inference/core/workflows/core_steps/visualizations/blur/v1.py +++ b/inference/core/workflows/core_steps/visualizations/blur/v1.py @@ -118,8 +118,13 @@ def run( kernel_size: Optional[int], ) -> BlockResult: annotator = self.getAnnotator(kernel_size) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/blur/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/blur/v1_tensor.py new file mode 100644 index 0000000000..fea68ba1ab --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/blur/v1_tensor.py @@ -0,0 +1,144 @@ +from typing import Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + PredictionsVisualizationBlock, + PredictionsVisualizationManifest, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/blur_visualization@v1" +SHORT_DESCRIPTION = "Blur detected objects in an image." +LONG_DESCRIPTION = """ +Apply blur effects to detected objects in an image, obscuring their details while preserving the background, useful for privacy protection, content filtering, or visual emphasis. + +## How This Block Works + +This block takes an image and detection predictions and applies a blur effect to the detected objects, leaving the background unchanged. The block: + +1. Takes an image and predictions as input +2. Identifies detected regions from bounding boxes or segmentation masks +3. Applies a blur effect (using average pooling) to the detected object regions +4. Preserves the background and areas outside detected objects unchanged +5. Returns an annotated image where detected objects are blurred, while the rest of the image remains sharp + +The block works with both object detection predictions (using bounding boxes) and instance segmentation predictions (using masks). When masks are available, it blurs the exact shape of detected objects; otherwise, it blurs rectangular bounding box regions. The blur intensity is controlled by the kernel size parameter, where larger kernel sizes create stronger blur effects. This creates a visual effect that obscures or anonymizes detected objects while maintaining context from the surrounding image, making it ideal for privacy protection, content filtering, or focusing attention on the background. + +## Common Use Cases + +- **Privacy Protection and Anonymization**: Blur faces, people, license plates, or other sensitive information in images or videos to protect privacy, comply with data protection regulations, or anonymize content before sharing or publishing +- **Content Filtering and Moderation**: Obscure inappropriate or sensitive content in images or videos for content moderation workflows, safe content previews, or user-generated content filtering +- **Visual Emphasis and Focus**: Blur detected objects to draw attention to other parts of the image, create visual contrast between blurred foreground objects and sharp backgrounds, or emphasize specific elements in composition +- **Product Photography and E-commerce**: Blur detected distracting elements or secondary products in images to keep the main subject sharp and prominent for product photography, catalog creation, or e-commerce image preparation +- **Security and Surveillance**: Anonymize people, vehicles, or other identifiable elements in security footage or surveillance images while preserving scene context for analysis, reporting, or public sharing +- **Documentation and Reporting**: Create anonymized or censored versions of images for reports, documentation, or case studies where sensitive information needs to be obscured but overall context should remain visible + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Bounding Box Visualization, Polygon Visualization) to add additional annotations on top of blurred objects for comprehensive visualization or to indicate what was blurred +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save blurred images for documentation, reporting, or archiving privacy-protected content +- **Webhook blocks** to send blurred images to external systems, APIs, or web applications for content moderation, privacy-compliant sharing, or anonymized analysis +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send blurred images as privacy-protected visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with blurred objects for live monitoring, privacy-compliant video processing, or post-processing analysis +""" + + +class BlurManifest(PredictionsVisualizationManifest): + type: Literal[f"{TYPE}", "BlurVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Blur Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "fad fa-glasses", + "blockPriority": 4, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + kernel_size: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Size of the blur kernel used for average pooling. Larger values create stronger blur effects, making objects more obscured. Smaller values create subtle blur effects. Typical values range from 5 (light blur) to 51 (strong blur). Must be an odd number for optimal blurring performance.", + default=15, + examples=[15, "$inputs.kernel_size"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class BlurVisualizationBlockV1(PredictionsVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlurManifest + + def getAnnotator( + self, + kernel_size: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join(map(str, [kernel_size])) + + if key not in self.annotatorCache: + self.annotatorCache[key] = sv.BlurAnnotator(kernel_size=kernel_size) + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + kernel_size: Optional[int], + ) -> BlockResult: + # sv.BlurAnnotator blurs the `xyxy` box region and never reads `.mask`; + # skip the device->host dense-mask materialisation. + predictions = to_supervision_for_annotation( + predictions, materialise_masks=False + ) + annotator = self.getAnnotator(kernel_size) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/bounding_box/v1.py b/inference/core/workflows/core_steps/visualizations/bounding_box/v1.py index 507db0df29..c8af7bdd47 100644 --- a/inference/core/workflows/core_steps/visualizations/bounding_box/v1.py +++ b/inference/core/workflows/core_steps/visualizations/bounding_box/v1.py @@ -162,8 +162,13 @@ def run( thickness, roundness, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/bounding_box/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/bounding_box/v1_tensor.py new file mode 100644 index 0000000000..98bea1835c --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/bounding_box/v1_tensor.py @@ -0,0 +1,490 @@ +from functools import lru_cache +from typing import Dict, List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field + +from inference.core.logger import logger +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, + read_host_mirror, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + empty_predictions_passthrough, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +_EMPTY_I64 = np.zeros(0, dtype=np.int64) + +TYPE: str = "roboflow_core/bounding_box_visualization@v1" +SHORT_DESCRIPTION = "Draw a box around detected objects in an image." +LONG_DESCRIPTION = """ +Draw bounding boxes around detected objects in an image, with customizable colors, thickness, and corner roundness. + +## How This Block Works + +This block takes an image and detection predictions (from object detection, instance segmentation, or keypoint detection models) and draws rectangular bounding boxes around each detected object. The block: + +1. Takes an image and predictions as input +2. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +3. Draws bounding boxes using Supervision's BoxAnnotator (for square corners) or RoundBoxAnnotator (for rounded corners) based on the roundness setting +4. Applies the specified box thickness to control the line width of the bounding boxes +5. Returns an annotated image with bounding boxes overlaid on the original image + +The block supports various color palettes (default, Roboflow, Matplotlib palettes, or custom colors) and can color boxes based on detection class, index, or tracker ID. When roundness is set to 0, square corners are used; when roundness is greater than 0, rounded corners are applied for a softer visual appearance. You can choose whether to modify the original image or create a copy for visualization, which is useful when stacking multiple visualization blocks. + +## Common Use Cases + +- **Model Validation and Debugging**: Visualize detection results to verify model performance, check bounding box accuracy, identify false positives or false negatives, and debug model outputs +- **Results Presentation**: Create annotated images for reports, dashboards, or presentations showing what objects were detected in images or video frames +- **Quality Control**: Overlay bounding boxes on production line images to visualize detected defects, products, or components for quality assurance workflows +- **Monitoring and Alerting**: Generate visual outputs showing detected objects for security monitoring, surveillance systems, or compliance tracking with annotated evidence +- **Training Data Review**: Review and validate training datasets by visualizing annotations and bounding boxes to ensure labeling accuracy and consistency +- **Interactive Applications**: Create user interfaces that display real-time detection results with bounding boxes for object tracking, counting, or identification applications + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Polygon Visualization, Mask Visualization) to stack multiple annotations on the same image for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images for documentation, archiving, or training data preparation +- **Webhook blocks** to send visualized results to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with bounding boxes for live monitoring or post-processing analysis +""" + + +@lru_cache(maxsize=256) +def _quarter_arc_offsets(radius: int, thickness: int) -> Tuple[np.ndarray, np.ndarray]: + """(dy, dx) pixel offsets of a top-left quarter ring around an arc + center โ€” a `thickness`-wide annulus at `radius`. Mirrored by sign for + the other three corners. Cached per (radius, thickness) across frames.""" + outer = thickness // 2 + span = radius + outer + yy, xx = np.mgrid[-span:1, -span:1] + dist = np.sqrt(yy**2 + xx**2) + ring = (dist >= radius - (thickness - 1 - outer) - 0.5) & ( + dist <= radius + outer + 0.5 + ) + return yy[ring].astype(np.int64), xx[ring].astype(np.int64) + + +def gpu_draw_boxes( + scene_chw: torch.Tensor, + xyxy: np.ndarray, + colors_rgb: np.ndarray, + thickness: int, + roundness: float = 0.0, +) -> torch.Tensor: + """Draw box borders on a CHW RGB uint8 device tensor, in place. + + Approximate rendering: each border is a plain ``thickness``-wide band + centered on the box edge. ``roundness > 0`` rounds the corners with + analytic quarter-ring arcs at sv's corner radius + (``int(min_side // 2 * roundness)``). cv2 (the sv path) rasterises + joins/caps/arcs slightly differently, so output is visually equivalent, + not bit-identical. + + Painted with a fixed number of torch ops regardless of box count โ€” a + per-box loop is dispatch-bound on Jetson (measured 43 ms @ 50 boxes): + + 1. Host: 4 band rectangles per box, clamped to the frame. + 2. One packed int32 H2D upload (per-transfer fixed cost is ~0.36 ms on + AGX Orin), then two-level ragged expansion to flat pixel indices via + ``repeat_interleave`` with host-computed ``output_size`` โ€” no + GPU->CPU syncs. + 3. One indexed store. When boxes overlap, sv's later-box-wins paint + order is reproduced with a deterministic ``scatter_reduce_(amax)`` + of box indices; disjoint boxes (the common case) skip it. + """ + device = scene_chw.device + height, width = int(scene_chw.shape[1]), int(scene_chw.shape[2]) + if height * width >= 2**31: + raise ValueError("frame too large for int32 pixel indexing") + n = int(xyxy.shape[0]) + outer = thickness // 2 # band extends `outer` outward, rest inward + + xy = np.asarray(xyxy, dtype=np.int64) + bx1 = np.minimum(xy[:, 0], xy[:, 2]) + bx2 = np.maximum(xy[:, 0], xy[:, 2]) + by1 = np.minimum(xy[:, 1], xy[:, 3]) + by2 = np.maximum(xy[:, 1], xy[:, 3]) + if roundness > 0: + # sv.RoundBoxAnnotator's corner radius, from the smaller box side. + radii = (np.minimum(bx2 - bx1, by2 - by1) // 2 * roundness).astype(np.int64) + else: + radii = np.zeros(n, dtype=np.int64) + x1, x2, y1, y2 = bx1 - outer, bx2 + outer, by1 - outer, by2 + outer + + # 4 bands per box (inclusive coords). Square corners: top/bottom span + # the full width, left/right fill between them. Rounded: every band is + # inset by the corner radius; quarter-ring arcs (below) fill the joins. + # Degenerate boxes just overlap bands of the same color โ€” harmless. + t = thickness + inset = radii - outer # 0 boxes: -outer == square full-width bands + rect_r1 = np.concatenate([y1, y2 - t + 1, by1 + radii, by1 + radii]) + rect_r2 = np.concatenate([y1 + t - 1, y2, by2 - radii, by2 - radii]) + rect_c1 = np.concatenate([x1 + inset + outer, x1 + inset + outer, x1, x2 - t + 1]) + rect_c2 = np.concatenate([x2 - inset - outer, x2 - inset - outer, x1 + t - 1, x2]) + rect_box = np.tile(np.arange(n), 4) + np.clip(rect_r1, 0, None, out=rect_r1) + np.clip(rect_c1, 0, None, out=rect_c1) + np.clip(rect_r2, None, height - 1, out=rect_r2) + np.clip(rect_c2, None, width - 1, out=rect_c2) + heights = np.maximum(rect_r2 - rect_r1 + 1, 0) + widths = np.maximum(rect_c2 - rect_c1 + 1, 0) + total_rows = int(heights.sum()) + total_px = int((heights * widths).sum()) + if total_px == 0: + return scene_chw + + # Pairwise overlap test on the expanded bounds: disjoint borders make + # every pixel's winner its own box, so owner resolution can be skipped. + inter_x = np.maximum(x1[:, None], x1[None, :]) <= np.minimum( + x2[:, None], x2[None, :] + ) + inter_y = np.maximum(y1[:, None], y1[None, :]) <= np.minimum( + y2[:, None], y2[None, :] + ) + boxes_overlap = int((inter_x & inter_y).sum()) > n # diagonal always True + + # Rounded corners: quarter-ring arc pixels per box, grouped by radius so + # the ring offsets are computed (and lru-cached) once per radius. Coords + # are pre-resolved to flat indices on host and ride the packed upload. + arc_flat = arc_box = _EMPTY_I64 + if roundness > 0: + flat_parts, box_parts = [], [] + for radius in np.unique(radii): + idx = np.nonzero(radii == radius)[0] + dy, dx = _quarter_arc_offsets(int(radius), thickness) + centers_y = ( + by1[idx] + radius, + by1[idx] + radius, + by2[idx] - radius, + by2[idx] - radius, + ) + centers_x = ( + bx1[idx] + radius, + bx2[idx] - radius, + bx1[idx] + radius, + bx2[idx] - radius, + ) + signs = ((1, 1), (1, -1), (-1, 1), (-1, -1)) + for cy, cx, (sy, sx) in zip(centers_y, centers_x, signs): + rows = (cy[:, None] + sy * dy[None, :]).ravel() + cols = (cx[:, None] + sx * dx[None, :]).ravel() + keep = (rows >= 0) & (rows < height) & (cols >= 0) & (cols < width) + flat_parts.append(rows[keep] * width + cols[keep]) + box_parts.append(np.repeat(idx, len(dy))[keep]) + arc_flat = np.concatenate(flat_parts) + arc_box = np.concatenate(box_parts) + + # Single packed upload; int32 halves the transfer and every device-side + # memory pass (torch index ops demand int64 โ€” `flat` is cast once). + segments = [ + rect_r1, + heights, + rect_c1, + widths, + rect_box, + arc_flat, + arc_box, + colors_rgb.astype(np.int64).ravel(), + ] + packed = torch.from_numpy(np.concatenate(segments).astype(np.int32)).to(device) + views = [] + offset = 0 + for segment in segments: + views.append(packed[offset : offset + len(segment)]) + offset += len(segment) + r1_t, heights_t, c1_t, widths_t, box_t, arc_flat_t, arc_box_t, colors_flat_t = views + + # Two-level ragged expansion: rect -> row -> pixel. + rect_of_row = torch.repeat_interleave( + torch.arange(4 * n, device=device), heights_t.long(), output_size=total_rows + ) + row_starts = heights_t.cumsum(0) - heights_t + row_intra = ( + torch.arange(total_rows, device=device, dtype=torch.int32) + - row_starts[rect_of_row] + ) + row_base = (r1_t[rect_of_row] + row_intra) * width + c1_t[rect_of_row] + row_width = widths_t[rect_of_row] + row_box = box_t[rect_of_row] + px_of_row = torch.repeat_interleave( + torch.arange(total_rows, device=device), row_width.long(), output_size=total_px + ) + col_starts = row_width.cumsum(0) - row_width + px_intra = ( + torch.arange(total_px, device=device, dtype=torch.int32) - col_starts[px_of_row] + ) + flat = row_base[px_of_row] + px_intra + pixel_box = row_box[px_of_row] + if len(arc_flat): + flat = torch.cat([flat, arc_flat_t]) + pixel_box = torch.cat([pixel_box, arc_box_t]) + flat = flat.long() + + colors_dev = colors_flat_t.view(n, 3).to(torch.uint8) + if boxes_overlap: + # include_self=False: uninitialized cells never participate, and + # every gathered position below was scattered to. + owner = torch.empty(height * width, dtype=torch.int32, device=device) + owner.scatter_reduce_(0, flat, pixel_box, reduce="amax", include_self=False) + winner_colors = colors_dev[owner[flat].long()] # (P, 3) uint8 + else: + winner_colors = colors_dev[pixel_box.long()] + # .view (not .reshape): guarantees the write lands in the caller's storage + # (raises on a non-contiguous scene -> caught by the block's sv fallback). + scene_chw.view(3, -1)[:, flat] = winner_colors.t() + return scene_chw + + +def _gpu_box_draw_eligible( + detections, color_axis: str, image: WorkflowImageData +) -> bool: + """True when the torch painter can replace the sv path.""" + if color_axis not in ("CLASS", "INDEX", "TRACK"): + return False + if not image.is_tensor_materialised(): + # Forcing tensor_image on a numpy-sourced image is a costly host-side + # conversion โ€” the sv path is faster there. + return False + xyxy = getattr(detections, "xyxy", None) + if not isinstance(xyxy, torch.Tensor) or int(xyxy.shape[0]) == 0: + # Nothing to paint; the sv path is a trivial no-op. + return False + return True + + +class BoundingBoxManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "BoundingBoxVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Bounding Box Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-object-group", + "blockPriority": 0, + "supervision": True, + "popular": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the bounding box edges in pixels. Higher values create thicker, more visible box outlines.", + default=2, + examples=[2, "$inputs.thickness"], + ) + + roundness: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + description="Roundness of the bounding box corners, ranging from 0.0 (square corners) to 1.0 (fully rounded corners). When set to 0.0, square-cornered boxes are used; higher values create progressively more rounded corners.", + default=0.0, + examples=[0.0, "$inputs.roundness"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class BoundingBoxVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BoundingBoxManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + thickness: int, + roundness: float, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map(str, [color_palette, palette_size, color_axis, thickness, roundness]) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + if roundness == 0: + self.annotatorCache[key] = sv.BoxAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + thickness=thickness, + ) + else: + self.annotatorCache[key] = sv.RoundBoxAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + thickness=thickness, + roundness=roundness, + ) + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + thickness: Optional[int], + roundness: Optional[float], + ) -> BlockResult: + detections = ( + split_key_point_prediction(predictions)[1] + if isinstance(predictions, tuple) + else predictions + ) + passthrough = empty_predictions_passthrough( + image=image, detections=detections, copy_image=copy_image + ) + if passthrough is not None: + return passthrough + if _gpu_box_draw_eligible(detections, color_axis, image): + try: + palette = self.getPalette(color_palette, palette_size, custom_colors) + if not isinstance(palette, sv.ColorPalette): + raise TypeError("expected sv.ColorPalette") + # The per-box host mirror (attach_native_detection_metadata) + # serves class ids and box coordinates without a device read โ€” + # a `.cpu()` here queues behind unrelated kernels on the + # default CUDA stream. Mirror-less predictions keep the + # tensor-read path; both produce bit-identical arrays. + host_mirror = read_host_mirror( + detections.bboxes_metadata, int(detections.xyxy.shape[0]) + ) + # Same colors sv's resolve_color would pick: CLASS -> + # by_idx(class_id), INDEX -> by_idx(det index), TRACK -> + # by_idx(tracker_id) with sv's gray for pending tracks (-1). + # Missing tracker ids raise here -> the sv path raises the + # same ValueError sv.resolve_color would. + if color_axis == "CLASS": + ids = ( + host_mirror[1] + if host_mirror is not None + else detections.class_id.detach().cpu().numpy().astype(int) + ) + elif color_axis == "TRACK": + ids = np.asarray( + [ + int(per_box["tracker_id"]) + for per_box in detections.bboxes_metadata + ] + ) + else: + ids = np.arange(int(detections.xyxy.shape[0])) + pending_gray = (128, 128, 128) # sv PENDING_TRACK_COLOR + colors_rgb = np.asarray( + [ + ( + pending_gray + if color_axis == "TRACK" and idx == -1 + else palette.by_idx(int(idx)).as_rgb() + ) + for idx in ids + ], + dtype=np.uint8, + ) + # Same int conversion as the sv annotator loop (truncation + # toward zero, incl. negative coords). The mirror is float32 + # exactly like the tensor read, so the truncation matches. + xyxy = ( + host_mirror[0].astype(int) + if host_mirror is not None + else detections.xyxy.detach().cpu().numpy().astype(int) + ) + # Tensor pipeline contract: the image is a CHW RGB device + # tensor โ€” zero-copy in, tensor out (downstream materialises + # numpy lazily only if something asks for it). + scene_t = image.tensor_image + if int(scene_t.shape[0]) != 3: + raise ValueError("GPU box painter requires a 3-channel image") + if copy_image: + scene_t = scene_t.clone() + annotated_tensor = gpu_draw_boxes( + scene_t, xyxy, colors_rgb, int(thickness), float(roundness) + ) + if not copy_image: + # The painter mutated `image.tensor_image` storage in + # place (the sv-path contract for copy_image=False); + # invalidate the derived numpy/base64 caches. + image.declare_tensor_image_mutated() + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, tensor_image=annotated_tensor + ) + } + except Exception as gpu_error: + logger.debug( + "GPU box painter failed (%s); falling back to " + "sv.BoxAnnotator path.", + gpu_error, + ) + # sv.BoxAnnotator / sv.RoundBoxAnnotator draw from `xyxy` only and never + # read `.mask`; skip the device->host dense-mask materialisation. + predictions = to_supervision_for_annotation( + predictions, materialise_masks=False + ) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + thickness, + roundness, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/circle/v1.py b/inference/core/workflows/core_steps/visualizations/circle/v1.py index 46890d7f29..f2ed06220f 100644 --- a/inference/core/workflows/core_steps/visualizations/circle/v1.py +++ b/inference/core/workflows/core_steps/visualizations/circle/v1.py @@ -152,8 +152,13 @@ def run( color_axis, thickness, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/circle/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/circle/v1_tensor.py new file mode 100644 index 0000000000..ff805a2438 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/circle/v1_tensor.py @@ -0,0 +1,178 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/circle_visualization@v1" +SHORT_DESCRIPTION = "Draw a circle around detected objects in an image." +LONG_DESCRIPTION = """ +Draw circular outlines around detected objects, providing an alternative to rectangular bounding boxes with a softer, more rounded visualization style. + +## How This Block Works + +This block takes an image and detection predictions and draws circular outlines around each detected object. The block: + +1. Takes an image and predictions as input +2. Calculates the center point and size for each detection based on its bounding box +3. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +4. Draws circular outlines around each detected object using Supervision's CircleAnnotator +5. Applies the specified circle thickness to control the line width of the circular outlines +6. Returns an annotated image with circular outlines overlaid on the original image + +The block draws circles that are typically centered on each detection's bounding box, with the circle size determined by the detection dimensions. Circles provide a softer, more organic visual style compared to rectangular bounding boxes, while still clearly marking the location and extent of detected objects. Unlike dot visualization (which marks specific points), circle visualization draws full circular outlines that encompass the detected objects, making it useful when you want a rounded geometric shape that's less angular than bounding boxes but more prominent than small dot markers. + +## Common Use Cases + +- **Soft Geometric Visualization**: Use circular outlines instead of rectangular bounding boxes for a softer, more organic visual style in presentations, dashboards, or user interfaces where rounded shapes are preferred +- **Object Highlighting with Rounded Shapes**: Highlight detected objects with circular outlines when working with circular or spherical objects (e.g., balls, coins, circular logos, round products) where circles naturally fit the object shape +- **Aesthetic Visualization Alternatives**: Create visually distinct annotations compared to standard bounding boxes for design purposes, artistic visualizations, or when circular shapes better match the overall design aesthetic +- **Detection Visualization with Variation**: Provide an alternative visualization style to bounding boxes for comparison, experimentation, or when multiple visualization types are used together to distinguish different detection sets +- **User Interface Design**: Use circular outlines in user interfaces, mobile apps, or interactive displays where rounded shapes are more visually appealing or match design guidelines +- **Scientific and Medical Imaging**: Visualize detections with circular outlines in scientific or medical imaging contexts where rounded shapes may be more appropriate than angular bounding boxes + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Dot Visualization, Bounding Box Visualization) to combine circular outlines with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images with circular outlines for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with circular outlines to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with circular outlines as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with circular outlines for live monitoring, tracking visualization, or post-processing analysis +""" + + +class CircleManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "CircleVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Circle Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-circle", + "blockPriority": 5, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the circle outline in pixels. Higher values create thicker, more visible circular outlines.", + default=2, + examples=[2, "$inputs.thickness"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class CircleVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return CircleManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + thickness: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + thickness, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = sv.CircleAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + thickness=thickness, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + thickness: Optional[int], + ) -> BlockResult: + # sv.CircleAnnotator draws from `xyxy` only and never reads `.mask`; + # skip the device->host dense-mask materialisation. + predictions = to_supervision_for_annotation( + predictions, materialise_masks=False + ) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + thickness, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/classification_label/v1.py b/inference/core/workflows/core_steps/visualizations/classification_label/v1.py index 92bcae115b..1032d1537f 100644 --- a/inference/core/workflows/core_steps/visualizations/classification_label/v1.py +++ b/inference/core/workflows/core_steps/visualizations/classification_label/v1.py @@ -327,8 +327,13 @@ def run( tracker_id=np.array([0 for _ in predictions_to_use]), ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=pseudo_detections, labels=labels, ) diff --git a/inference/core/workflows/core_steps/visualizations/classification_label/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/classification_label/v1_tensor.py new file mode 100644 index 0000000000..72a9e63be0 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/classification_label/v1_tensor.py @@ -0,0 +1,661 @@ +from typing import List, Literal, Optional, Tuple, Type, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, +) +from inference.core.workflows.core_steps.visualizations.common.utils import str_to_color +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + IMAGE_DIMENSIONS_KEY, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + INTEGER_KIND, + STRING_KIND, + StepOutputSelector, + WorkflowParameterSelector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) + +SHORT_DESCRIPTION = "Visualize both single-label and multi-label classification predictions with customizable display options." + +LONG_DESCRIPTION = """ +Visualize classification predictions as text labels positioned on images, automatically handling both single-label and multi-label classification formats with customizable styling and positioning. + +## How This Block Works + +This block takes an image and classification predictions (for entire image classification, not object detection) and displays text labels showing the predicted class names and confidence scores. The block: + +1. Takes an image and classification predictions as input +2. Automatically detects whether predictions are single-label (one class per image) or multi-label (multiple classes per image) +3. For single-label predictions: selects the highest confidence prediction to display +4. For multi-label predictions: formats and sorts all predicted classes by confidence score (highest first) +5. Extracts label text based on the selected text option (class name, confidence score, or both) +6. Positions labels on the image at the specified location (top, center, or bottom edges, with left/center/right alignment) +7. Applies background color styling based on the selected color palette, with colors assigned by class +8. Renders text labels with customizable text color, scale, thickness, padding, and border radius +9. Returns an annotated image with classification labels overlaid on the original image + +Unlike the regular Label Visualization block (which labels detected objects with bounding boxes), this block is designed for image-level classification where the entire image is classified into one or more categories. Labels are positioned at the edges or center of the image itself, not relative to object locations. For multi-label predictions, multiple labels are stacked vertically at the chosen position, making it easy to see all predicted classes and their confidence scores. + +## Common Use Cases + +- **Image Classification Results Display**: Visualize the predicted class and confidence score for classified images in applications like content moderation, product categorization, or medical image analysis +- **Multi-Class Probability Visualization**: Display multiple predicted classes with their confidence scores for multi-label classification tasks, such as tagging images with multiple attributes, detecting multiple defects, or identifying multiple objects in scene classification +- **Model Performance Validation**: Show classification predictions directly on images to validate model performance, verify correct classifications, and identify misclassifications during model development or testing +- **User Interface Integration**: Create clean, professional displays of classification results for applications, dashboards, or mobile apps where users need to see what an image was classified as +- **Documentation and Reporting**: Generate annotated images showing classification results for reports, documentation, or training data review to demonstrate model predictions +- **Quality Control Workflows**: Display classification results on production images for quality control, content filtering, or automated categorization workflows where visual confirmation of predictions is needed + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images with classification labels for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with classification labels to external systems, APIs, or web applications for display in dashboards or classification monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with classification labels as visual evidence in alerts or reports when specific classes are detected +- **Video output blocks** to create annotated video streams or recordings with classification labels for live monitoring, real-time classification display, or post-processing analysis +- **Conditional logic blocks** (e.g., Continue If) to route workflow execution based on classification results or confidence scores displayed in the labels +""" + + +class ClassificationLabelManifest(ColorableVisualizationManifest): + type: Literal["roboflow_core/classification_label_visualization@v1"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Classification Label Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-tags", + "blockPriority": 2.5, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + text: Union[ + Literal["Class", "Confidence", "Class and Confidence"], + WorkflowParameterSelector(kind=[STRING_KIND]), + ] = Field( # type: ignore + default="Class", + description="Content to display in text labels. Options: 'Class' (class name only), 'Confidence' (confidence score only, formatted as decimal), or 'Class and Confidence' (both class name and confidence score).", + examples=["LABEL", "$inputs.text"], + json_schema_extra={ + "always_visible": True, + }, + ) + + text_position: Union[ + Literal[ + "CENTER", + "CENTER_LEFT", + "CENTER_RIGHT", + "TOP_CENTER", + "TOP_LEFT", + "TOP_RIGHT", + "BOTTOM_LEFT", + "BOTTOM_CENTER", + "BOTTOM_RIGHT", + ], + WorkflowParameterSelector(kind=[STRING_KIND]), + ] = Field( # type: ignore + default="TOP_LEFT", + description="Position for placing labels on the image. Options include: TOP (TOP_LEFT, TOP_CENTER, TOP_RIGHT), CENTER (CENTER_LEFT, CENTER, CENTER_RIGHT), or BOTTOM (BOTTOM_LEFT, BOTTOM_CENTER, BOTTOM_RIGHT). For multi-label predictions, labels are stacked vertically at the chosen position.", + examples=["CENTER", "$inputs.text_position"], + ) + + predictions: StepOutputSelector(kind=[TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND]) = Field( # type: ignore + description="Classification predictions from a single-label or multi-label classification model. The block automatically detects the prediction format and handles both types accordingly.", + examples=["$steps.classification_model.predictions"], + ) + + text_color: Union[str, WorkflowParameterSelector(kind=[STRING_KIND])] = Field( # type: ignore + description="Color of the label text. Can be a color name (e.g., 'WHITE', 'BLACK') or color code in HEX format (e.g., '#FFFFFF') or RGB format (e.g., 'rgb(255, 255, 255)').", + default="WHITE", + examples=["WHITE", "#FFFFFF", "rgb(255, 255, 255)" "$inputs.text_color"], + ) + + text_scale: Union[float, WorkflowParameterSelector(kind=[FLOAT_KIND])] = Field( # type: ignore + description="Scale factor for text size. Higher values create larger text. Default is 1.0.", + default=1.0, + examples=[1.0, "$inputs.text_scale"], + ) + + text_thickness: Union[int, WorkflowParameterSelector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of text characters in pixels. Higher values create bolder, thicker text for better visibility.", + default=1, + examples=[1, "$inputs.text_thickness"], + ) + + text_padding: Union[int, WorkflowParameterSelector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Padding around the text in pixels. Controls the spacing between the text and the label background border, and the spacing between multiple labels in multi-label predictions.", + default=10, + examples=[10, "$inputs.text_padding"], + ) + + border_radius: Union[int, WorkflowParameterSelector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Border radius of the label background in pixels. Set to 0 for square corners. Higher values create more rounded corners for a softer appearance.", + default=0, + examples=[0, "$inputs.border_radius"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class ClassificationLabelVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return ClassificationLabelManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + text_position: str, + text_color: str, + text_scale: float, + text_thickness: int, + text_padding: int, + border_radius: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + text_position, + text_color, + text_scale, + text_thickness, + text_padding, + border_radius, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + text_color = str_to_color(text_color) + + self.annotatorCache[key] = sv.LabelAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + text_position=getattr(sv.Position, text_position), + text_color=text_color, + text_scale=text_scale, + text_thickness=text_thickness, + text_padding=text_padding, + border_radius=border_radius, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[ + ClassificationPrediction, MultiLabelClassificationPrediction + ], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + text: Optional[str], + text_position: Optional[str], + text_color: Optional[str], + text_scale: Optional[float], + text_thickness: Optional[int], + text_padding: Optional[int], + border_radius: Optional[int], + ) -> BlockResult: + predictions = to_legacy_classification_prediction(predictions) + try: + # Detect task type from predictions + task_type = detect_prediction_type(predictions) + + # Get sorted predictions based on detected task type + if task_type == "single-label": + sorted_predictions = sorted( + predictions["predictions"], + key=lambda x: x["confidence"], + reverse=True, + )[:1] + else: # multi-label + formatted_predictions = format_multi_label_predictions(predictions) + sorted_predictions = sorted( + formatted_predictions, key=lambda x: x["confidence"], reverse=True + ) + + # Early return if no predictions + if not sorted_predictions: + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=( + image.numpy_image.copy() + if copy_image + else image.numpy_image + ), + ) + } + + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + text_position, + text_color, + text_scale, + text_thickness, + text_padding, + border_radius, + ) + + # Calculate spacing based on all relevant parameters + base_text_height = 20 # OpenCV's base text height in pixels + scaled_text_height = base_text_height * text_scale + padding_space = ( + text_padding * 2 + ) # multiply by 2 since padding is applied top and bottom + thickness_space = text_thickness * 2 # account for text thickness + total_spacing = ( + scaled_text_height + padding_space + thickness_space + 5 + ) # added small buffer + initial_offset = total_spacing + w = predictions["image"]["width"] + h = predictions["image"]["height"] + + # Both single and multi-label use the same visualization creation + xyxy, labels, predictions_to_use = create_label_visualization( + sorted_predictions=sorted_predictions, + text_position=text_position, + text=text, + w=w, + h=h, + initial_offset=initial_offset, + total_spacing=total_spacing, + text_scale=text_scale, + text_padding=text_padding, + ) + + if not predictions_to_use: + # If no predictions, return the original image + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=( + image.numpy_image.copy() + if copy_image + else image.numpy_image + ), + ) + } + + pseudo_detections = sv.Detections( + xyxy=xyxy, + class_id=np.array([p["class_id"] for p in predictions_to_use]), + confidence=np.array([p["confidence"] for p in predictions_to_use]), + tracker_id=np.array([0 for _ in predictions_to_use]), + ) + + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=pseudo_detections, + labels=labels, + ) + + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } + + except ValueError as e: + raise ValueError( + f"Invalid prediction format: {str(e)}. Please check if the task_type matches your model's output format." + ) from e + + +def to_legacy_classification_prediction( + prediction: Union[ClassificationPrediction, MultiLabelClassificationPrediction], +) -> dict: + """Materialise a tensor-native classification prediction into the legacy + ``dict`` shape the visualisation layout helpers consume. + + Reconstructed shapes: + + * single-label (``ClassificationPrediction``) -> + ``{"image": {"width", "height"}, "predictions": [{"class", "class_id", + "confidence"}, ...]}`` over the full softmax distribution; class names come + from ``images_metadata[0]["class_names"]`` (``{int class_id: str name}``), + * multi-label (``MultiLabelClassificationPrediction``) -> + ``{"image": {"width", "height"}, "predictions": {class_name: + {"confidence", "class_id"}, ...}, "predicted_classes": [...]}``; class + names / dimensions come from ``image_metadata``, and ``predicted_classes`` + is the above-threshold ``class_ids``. + """ + if isinstance(prediction, ClassificationPrediction): + image_metadata = (prediction.images_metadata or [{}])[0] or {} + height, width = _resolve_image_dimensions(image_metadata) + class_names_mapping = image_metadata.get(CLASS_NAMES_KEY) or {} + confidence = prediction.confidence.detach().cpu().reshape(-1).numpy() + predictions = [ + { + "class": _resolve_class_name(class_id, class_names_mapping), + "class_id": int(class_id), + "confidence": float(confidence[class_id]), + } + for class_id in range(int(confidence.shape[-1])) + ] + return { + "image": {"width": width, "height": height}, + "predictions": predictions, + } + image_metadata = prediction.image_metadata or {} + height, width = _resolve_image_dimensions(image_metadata) + class_names_mapping = image_metadata.get(CLASS_NAMES_KEY) or {} + confidence = prediction.confidence.detach().cpu().numpy() + class_ids = ( + prediction.class_ids.detach().cpu().numpy().tolist() + if prediction.class_ids is not None + else [] + ) + num_classes = int(confidence.shape[-1]) + predictions = { + _resolve_class_name(class_id, class_names_mapping): { + "confidence": float(confidence[class_id]), + "class_id": int(class_id), + } + for class_id in range(num_classes) + } + predicted_classes = [ + _resolve_class_name(int(class_id), class_names_mapping) + for class_id in class_ids + ] + return { + "image": {"width": width, "height": height}, + "predictions": predictions, + "predicted_classes": predicted_classes, + } + + +def _resolve_image_dimensions(image_metadata: dict) -> Tuple[int, int]: + image_dimensions = image_metadata.get(IMAGE_DIMENSIONS_KEY) + if image_dimensions is None: + raise ValueError( + "Classification prediction is missing image dimensions required to " + "place labels." + ) + height, width = int(image_dimensions[0]), int(image_dimensions[1]) + return height, width + + +def _resolve_class_name(class_id: int, class_names_mapping: dict) -> str: + class_name = class_names_mapping.get(class_id) + if class_name is None: + return f"class_{class_id}" + return str(class_name) + + +def handle_bottom_position( + sorted_predictions: List[dict], + text: str, + w: int, + h: int, + initial_offset: float, + total_spacing: float, +) -> Tuple[np.ndarray, List[str], List[dict]]: + """Handle visualization layout for bottom positions.""" + reversed_predictions = sorted_predictions[::-1] + xyxy = np.array( + [ + [0, 0, w, h - (initial_offset + i * total_spacing)] + for i in range(len(reversed_predictions)) + ] + ) + labels = format_labels(reversed_predictions, text) + return xyxy, labels, reversed_predictions + + +def handle_center_position( + sorted_predictions: List[dict], + text: str, + text_position: str, + w: int, + h: int, + total_spacing: float, + text_scale: float, + text_padding: int, +) -> Tuple[np.ndarray, List[str], List[dict]]: + """Handle visualization layout for center positions.""" + labels = format_labels(sorted_predictions, text) + n_predictions = len(sorted_predictions) + total_height = total_spacing * n_predictions + start_y = max(0, min((h - total_height) / 2, h - total_height)) + + max_label_length = max(len(label) for label in labels) + char_width = 15 + label_width = (max_label_length * char_width * text_scale) + (text_padding * 2) + extra_padding = 20 + max(0, 10 - text_padding) * 3 + + if text_position == "CENTER_LEFT": + x_start = label_width + extra_padding + xyxy = np.array( + [ + [ + x_start, + start_y + i * total_spacing, + w, + start_y + (i + 1) * total_spacing, + ] + for i in range(n_predictions) + ] + ) + elif text_position == "CENTER_RIGHT": + x_end = w - (label_width + extra_padding) + xyxy = np.array( + [ + [ + 0, + start_y + i * total_spacing, + x_end, + start_y + (i + 1) * total_spacing, + ] + for i in range(n_predictions) + ] + ) + else: # CENTER + xyxy = np.array( + [ + [0, start_y + i * total_spacing, w, start_y + (i + 1) * total_spacing] + for i in range(n_predictions) + ] + ) + + return xyxy, labels, sorted_predictions + + +def handle_top_position( + sorted_predictions: List[dict], + text: str, + w: int, + h: int, + initial_offset: float, + total_spacing: float, +) -> Tuple[np.ndarray, List[str], List[dict]]: + """Handle visualization layout for top positions.""" + xyxy = np.array( + [ + [0, initial_offset + i * total_spacing, w, h] + for i in range(len(sorted_predictions)) + ] + ) + labels = format_labels(sorted_predictions, text) + return xyxy, labels, sorted_predictions + + +def create_label_visualization( + sorted_predictions: List[dict], + text_position: str, + text: str, + w: int, + h: int, + initial_offset: float, + total_spacing: float, + text_scale: float, + text_padding: int, +) -> Tuple[np.ndarray, List[str], List[dict]]: + """Create visualization layout for classification labels.""" + if text_position in ["BOTTOM_LEFT", "BOTTOM_CENTER", "BOTTOM_RIGHT"]: + return handle_bottom_position( + sorted_predictions, text, w, h, initial_offset, total_spacing + ) + elif text_position in ["CENTER", "CENTER_LEFT", "CENTER_RIGHT"]: + return handle_center_position( + sorted_predictions, + text, + text_position, + w, + h, + total_spacing, + text_scale, + text_padding, + ) + else: # Top positions + return handle_top_position( + sorted_predictions, text, w, h, initial_offset, total_spacing + ) + + +def detect_prediction_type(predictions: dict) -> str: + """ + Detect whether predictions are single-label or multi-label based on structure. + + Args: + predictions (dict): The predictions dictionary + + Returns: + str: 'single-label' or 'multi-label' + """ + if isinstance(predictions.get("predictions"), list): + return "single-label" + return "multi-label" + + +def validate_prediction_format(predictions: dict, task_type: str) -> None: + """ + Validate that the predictions format matches the specified task type. + + Args: + predictions (dict): The predictions dictionary + task_type (str): The specified task type ('single-label' or 'multi-label') + + Raises: + ValueError: If prediction format doesn't match task type + """ + actual_type = detect_prediction_type(predictions) + + if actual_type != task_type: + if actual_type == "single-label": + raise ValueError( + "Received single-label predictions but task_type is set to 'multi-label'. Please correct the task_type setting." + ) + else: + raise ValueError( + "Received multi-label predictions but task_type is set to 'single-label'. Please correct the task_type setting." + ) + + +def format_multi_label_predictions(predictions: dict) -> List[dict]: + """ + Transform multi-label predictions from predicted_classes into standard format. + + Args: + predictions (dict): The predictions dictionary + + Returns: + List[dict]: Formatted predictions list + """ + formatted_predictions = [] + for class_name in predictions["predicted_classes"]: + pred_info = predictions["predictions"][class_name] + formatted_predictions.append( + { + "class": class_name, + "class_id": pred_info["class_id"], + "confidence": pred_info["confidence"], + } + ) + return formatted_predictions + + +def format_labels(predictions, text="Class and Confidence"): + """ + Format labels based on specified text option. + + Args: + predictions (list): List of prediction dictionaries containing 'class' and 'confidence' + text (str): One of "class", "confidence", or "class and confidence" + + Returns: + list: Formatted label strings + """ + if text == "Class": + labels = [f"{p['class']}" for p in predictions] + elif text == "Confidence": + labels = [f"{p['confidence']:.2f}" for p in predictions] + elif text == "Class and Confidence": + labels = [f"{p['class']} {p['confidence']:.2f}" for p in predictions] + else: + raise ValueError( + "text must be one of: 'class', 'confidence', or 'class and confidence'" + ) + + return labels diff --git a/inference/core/workflows/core_steps/visualizations/color/v1.py b/inference/core/workflows/core_steps/visualizations/color/v1.py index 73ee056038..a81116da64 100644 --- a/inference/core/workflows/core_steps/visualizations/color/v1.py +++ b/inference/core/workflows/core_steps/visualizations/color/v1.py @@ -153,8 +153,13 @@ def run( color_axis, opacity, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/color/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/color/v1_tensor.py new file mode 100644 index 0000000000..9a18e901b6 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/color/v1_tensor.py @@ -0,0 +1,179 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/color_visualization@v1" +SHORT_DESCRIPTION = "Paint a solid color on detected objects in an image." +LONG_DESCRIPTION = """ +Fill detected objects with solid colors using customizable color palettes, creating color-coded overlays that distinguish different objects or classes while preserving image details through opacity blending. + +## How This Block Works + +This block takes an image and detection predictions and fills the detected object regions with solid colors. The block: + +1. Takes an image and predictions as input +2. Identifies detected regions from bounding boxes or segmentation masks +3. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +4. Fills detected object regions with solid colors using Supervision's ColorAnnotator +5. Blends the colored overlay with the original image based on the opacity setting +6. Returns an annotated image where detected objects are filled with colors, while the rest of the image remains unchanged + +The block works with both object detection predictions (using bounding boxes) and instance segmentation predictions (using masks). When masks are available, it fills the exact shape of detected objects; otherwise, it fills rectangular bounding box regions. Colors are assigned from the selected palette based on the color axis setting (class, index, or track ID), allowing different objects or classes to be distinguished by color. The opacity parameter controls how transparent the color overlay is, allowing you to create effects ranging from subtle color tinting (low opacity) where original image details remain visible, to solid color fills (high opacity) that completely replace object appearance. + +## Common Use Cases + +- **Color-Coded Object Classification**: Fill detected objects with different colors based on their class, category, or classification results to create intuitive color-coded visualizations for quick object identification and categorization +- **Multi-Object Tracking Visualization**: Color-code tracked objects with distinct colors based on their tracking IDs to visualize object trajectories, track persistence, or distinguish multiple tracked objects across frames +- **Visual Category Distinction**: Use different colors for different object categories or types (e.g., vehicles, people, products) to create clear visual distinctions in monitoring, surveillance, or inventory management workflows +- **Mask-Based Segmentation Display**: Fill segmented regions with colors to visualize instance segmentation results, highlight segmented objects, or create colored mask overlays for analysis or presentation +- **Interactive Visualization and UI**: Create color-coded visualizations for user interfaces, dashboards, or interactive applications where color-coding provides intuitive visual feedback or object grouping +- **Presentation and Reporting**: Generate color-filled visualizations for reports, documentation, or presentations where color-coding helps distinguish object types, highlight specific categories, or create visually appealing detection displays + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Bounding Box Visualization, Polygon Visualization) to combine color fills with additional annotations (labels, outlines) for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save color-coded images for documentation, reporting, or analysis +- **Webhook blocks** to send color-coded visualizations to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send color-coded images as visual evidence in alerts or reports +- **Video output blocks** to create color-coded video streams or recordings for live monitoring, tracking visualization, or post-processing analysis +""" + + +class ColorManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "ColorVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Color Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-palette", + "blockPriority": 6, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + opacity: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + description="Opacity of the color overlay, ranging from 0.0 (fully transparent, original object appearance visible) to 1.0 (fully opaque, solid color fill). Values between 0.0 and 1.0 create a blend between the original image and the color overlay. Lower values create subtle color tinting where object details remain visible, while higher values create stronger color fills that obscure original object appearance.", + default=0.5, + examples=[0.5, "$inputs.opacity"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class ColorVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return ColorManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + opacity: float, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + opacity, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = sv.ColorAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + opacity=opacity, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + opacity: Optional[float], + ) -> BlockResult: + # sv.ColorAnnotator fills the `xyxy` box rectangle and never reads + # `.mask`; skip the device->host dense-mask materialisation. + predictions = to_supervision_for_annotation( + predictions, materialise_masks=False + ) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + opacity, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/common/base_colorable_tensor.py b/inference/core/workflows/core_steps/visualizations/common/base_colorable_tensor.py new file mode 100644 index 0000000000..6d0fbaba0b --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/common/base_colorable_tensor.py @@ -0,0 +1,162 @@ +from abc import ABC, abstractmethod +from typing import List, Literal, Optional, Union + +import supervision as sv +from pydantic import Field + +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + PredictionsVisualizationBlock, + PredictionsVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.utils import str_to_color +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult + + +class ColorableVisualizationManifest(PredictionsVisualizationManifest, ABC): + color_palette: Union[ + Literal[ + "DEFAULT", + "CUSTOM", + "ROBOFLOW", + "Matplotlib Viridis", + "Matplotlib Plasma", + "Matplotlib Inferno", + "Matplotlib Magma", + "Matplotlib Cividis", + # TODO: Re-enable once supervision 0.23 is released with a fix + # "Matplotlib Twilight", + # "Matplotlib Twilight_Shifted", + # "Matplotlib HSV", + # "Matplotlib Jet", + # "Matplotlib Turbo", + # "Matplotlib Rainbow", + # "Matplotlib gist_rainbow", + # "Matplotlib nipy_spectral", + # "Matplotlib gist_ncar", + "Matplotlib Pastel1", + "Matplotlib Pastel2", + "Matplotlib Paired", + "Matplotlib Accent", + "Matplotlib Dark2", + "Matplotlib Set1", + "Matplotlib Set2", + "Matplotlib Set3", + "Matplotlib Tab10", + "Matplotlib Tab20", + "Matplotlib Tab20b", + "Matplotlib Tab20c", + # TODO: Re-enable once supervision 0.23 is released with a fix + # "Matplotlib Ocean", + # "Matplotlib Gist_Earth", + # "Matplotlib Terrain", + # "Matplotlib Stern", + # "Matplotlib gnuplot", + # "Matplotlib gnuplot2", + # "Matplotlib Spring", + # "Matplotlib Summer", + # "Matplotlib Autumn", + # "Matplotlib Winter", + # "Matplotlib Cool", + # "Matplotlib Hot", + # "Matplotlib Copper", + # "Matplotlib Bone", + # "Matplotlib Greys_R", + # "Matplotlib Purples_R", + # "Matplotlib Blues_R", + # "Matplotlib Greens_R", + # "Matplotlib Oranges_R", + # "Matplotlib Reds_R", + ], + Selector(kind=[STRING_KIND]), + ] = Field( # type: ignore + default="DEFAULT", + description="Select a color palette for the visualised elements.", + examples=["DEFAULT", "$inputs.color_palette"], + ) + + palette_size: Union[ + int, + Selector(kind=[INTEGER_KIND]), + ] = Field( # type: ignore + default=10, + description="Specify the number of colors in the palette. This applies when using custom or Matplotlib palettes.", + examples=[10, "$inputs.palette_size"], + ) + + custom_colors: Union[List[str], Selector(kind=[LIST_OF_VALUES_KIND])] = ( + Field( # type: ignore + default=[], + description="Define a list of custom colors for bounding boxes in HEX format.", + examples=[["#FF0000", "#00FF00", "#0000FF"], "$inputs.custom_colors"], + ) + ) + + color_axis: Union[ + Literal["INDEX", "CLASS", "TRACK"], + Selector(kind=[STRING_KIND]), + ] = Field( # type: ignore + default="CLASS", + description="Choose how bounding box colors are assigned.", + examples=["CLASS", "$inputs.color_axis"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class ColorableVisualizationBlock(PredictionsVisualizationBlock, ABC): + @classmethod + def getPalette(self, color_palette, palette_size, custom_colors): + if color_palette == "CUSTOM": + return sv.ColorPalette( + colors=[str_to_color(color) for color in custom_colors] + ) + elif hasattr(sv.ColorPalette, color_palette): + return getattr(sv.ColorPalette, color_palette) + else: + palette_name = color_palette.replace("Matplotlib ", "") + + if palette_name in [ + "Greys_R", + "Purples_R", + "Blues_R", + "Greens_R", + "Oranges_R", + "Reds_R", + "Wistia", + "Pastel1", + "Pastel2", + "Paired", + "Accent", + "Dark2", + "Set1", + "Set2", + "Set3", + ]: + palette_name = palette_name.capitalize() + else: + palette_name = palette_name.lower() + + return sv.ColorPalette.from_matplotlib(palette_name, int(palette_size)) + + @abstractmethod + def run( + self, + image: WorkflowImageData, + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + *args, + **kwargs + ) -> BlockResult: + pass diff --git a/inference/core/workflows/core_steps/visualizations/common/base_tensor.py b/inference/core/workflows/core_steps/visualizations/common/base_tensor.py new file mode 100644 index 0000000000..e54f62baa4 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/common/base_tensor.py @@ -0,0 +1,364 @@ +from abc import ABC, abstractmethod +from typing import List, Optional, Type, Union + +import numpy as np +import supervision as sv +from pydantic import AliasChoices, ConfigDict, Field +from supervision.detection.compact_mask import CompactMask + +from inference.core.workflows.core_steps.common.rle_compact import ( + instances_rle_to_compact_mask, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + HOST_MIRROR_KEYS, + TensorNativeDetections, + TensorNativePrediction, + read_host_mirror, + split_key_point_prediction, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_ID_KEY, + TRACKER_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + OutputDefinition, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + IMAGE_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.types import InstancesRLEMasks + +OUTPUT_IMAGE_KEY: str = "image" + + +def predictions_are_empty(detections) -> bool: + """True when there is provably nothing to draw (no boxes at all).""" + if detections is None: + return True + xyxy = getattr(detections, "xyxy", None) + if xyxy is None: + return True + try: + return int(xyxy.shape[0]) == 0 + except (AttributeError, TypeError): + try: + return len(xyxy) == 0 + except TypeError: + return False + + +def empty_predictions_passthrough( + image: WorkflowImageData, detections, copy_image: bool +) -> Optional[dict]: + """``{OUTPUT_IMAGE_KEY: ...}`` when there is nothing to draw, else ``None``. + + With no detections every annotator is a no-op, but the sv fallback still + pays ``image.numpy_image`` first - on the tensor pipeline that is a + full-resolution device->host materialisation (~30 ms per 2K frame under + load on Orin NX, measured on the ~60% of live camera frames with no + detections - the dominant cost of visualization blocks in that regime). + Passing the input representation through keeps the exact output + semantics of an empty annotate (an independent copy when + ``copy_image=True``, shared backing otherwise) without ever leaving the + device. + """ + if not predictions_are_empty(detections): + return None + if image.is_tensor_materialised(): + tensor = image.tensor_image + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, + tensor_image=tensor.clone() if copy_image else tensor, + ) + } + numpy_image = image.numpy_image + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=numpy_image.copy() if copy_image else numpy_image, + ) + } + + +#: ``sv.Detections.data`` key the supervision annotators read class names from. +CLASS_NAME_DATA_FIELD: str = "class_name" + + +def to_supervision_for_annotation( + prediction: Union[TensorNativePrediction, TensorNativeDetections], + materialise_masks: bool = True, +) -> sv.Detections: + """Materialise a tensor-native prediction into an ``sv.Detections`` carrying + everything the supervision annotators read. + + ``materialise_masks=False`` leaves ``mask`` as ``None``, skipping the + device->host mask transfer/decode for annotators that never read it. + + The reconstructed ``sv.Detections`` carries: + + * ``xyxy`` / ``class_id`` / ``confidence`` (plus ``mask`` for instance + segmentation), + * ``tracker_id`` (from ``bboxes_metadata[i]["tracker_id"]`` when present), + * ``data["class_name"]`` resolved from ``image_metadata["class_names"]`` + (``{int class_id: str name}``), falling back to ``f"class_{id}"``, + * ``data[DETECTION_ID_KEY]`` (from ``bboxes_metadata``), + * ``data[IMAGE_DIMENSIONS_KEY]`` (broadcast from ``image_metadata``), and + * any extra per-box ``bboxes_metadata`` keys (``time_in_zone``, + ``area``-derived keys, etc.) that specific annotators consume. + + For the keypoint-detection tuple input, only the bounding-box component is + converted. + """ + if isinstance(prediction, tuple): + _, detections = split_key_point_prediction(prediction) + elif isinstance(prediction, KeyPoints): + raise ValueError( + "A bare `KeyPoints` prediction (without its bounding-box component) " + "cannot be visualised by this block: the supervision annotators " + "require the bounding-box `Detections`. Provide the keypoint-detection " + "tuple `(KeyPoints, Detections)` instead." + ) + else: + detections = prediction + image_metadata = detections.image_metadata or {} + detections_number = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(detections_number)] + class_names_mapping = image_metadata.get(CLASS_NAMES_KEY) or {} + # Prefer the per-box host mirror written by + # ``attach_native_detection_metadata``: when EVERY box carries it, the sv + # view's xyxy/class_id/confidence are assembled on the host with ZERO device + # reads โ€” a ``.cpu()`` here queues behind unrelated kernels on the default + # CUDA stream (up to ~65 ms per batch measured on Jetson under load). + # Mirror-less predictions (remote/deserialized/transformed) keep the + # tensor-read path unchanged; both paths produce bit-identical arrays. + host_mirror = read_host_mirror(bboxes_metadata, detections_number) + if host_mirror is not None: + xyxy, class_id, confidence = host_mirror + else: + xyxy = detections.xyxy.detach().cpu().numpy().astype(np.float32) + class_id = detections.class_id.detach().cpu().numpy().astype(int) + confidence = detections.confidence.detach().cpu().numpy().astype(np.float32) + mask = ( + _materialise_mask(detections, detections_number, xyxy) + if materialise_masks + else None + ) + tracker_id = _materialise_tracker_id(bboxes_metadata) + data = _materialise_data( + bboxes_metadata=bboxes_metadata, + class_id=class_id, + class_names_mapping=class_names_mapping, + image_metadata=image_metadata, + detections_number=detections_number, + ) + return sv.Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask, + tracker_id=tracker_id, + data=data, + ) + + +def _materialise_mask( + detections: TensorNativeDetections, + detections_number: int, + xyxy: np.ndarray, +) -> Optional[Union[np.ndarray, CompactMask]]: + if not isinstance(detections, InstanceDetections): + return None + if detections_number == 0: + return None + mask = detections.mask + if isinstance(mask, InstancesRLEMasks): + # RLE is transcoded without decoding the full-frame (N, H, W) boolean + # stack; the boxes provide the per-crop bounds. + return instances_rle_to_compact_mask(mask, xyxy) + # Dense (N, H, W) masks: one bulk device->host transfer instead of N + # per-instance round-trips (each a blocking CUDA sync). + return mask.detach().cpu().numpy().astype(bool) + + +def _materialise_tracker_id( + bboxes_metadata: List[dict], +) -> Optional[np.ndarray]: + tracker_ids = [data.get(TRACKER_ID_KEY) for data in bboxes_metadata] + if any(tracker_id is None for tracker_id in tracker_ids): + return None + return np.asarray([int(tracker_id) for tracker_id in tracker_ids]) + + +def _materialise_data( + bboxes_metadata: List[dict], + class_id: np.ndarray, + class_names_mapping: dict, + image_metadata: dict, + detections_number: int, +) -> dict: + class_names = [ + _resolve_class_name(int(value), class_names_mapping) for value in class_id + ] + data: dict = {CLASS_NAME_DATA_FIELD: np.asarray(class_names, dtype=object)} + detection_ids = [ + str(per_box.get(DETECTION_ID_KEY, "")) for per_box in bboxes_metadata + ] + data[DETECTION_ID_KEY] = np.asarray(detection_ids, dtype=object) + image_dimensions = image_metadata.get(IMAGE_DIMENSIONS_KEY) + if image_dimensions is not None: + data[IMAGE_DIMENSIONS_KEY] = np.asarray( + [list(image_dimensions) for _ in range(detections_number)] + ) + # Per-image lineage that flag-off carries per-box in ``sv.Detections.data`` + # (the single per-image value broadcast to every row). The tensor-native + # path stores these once in ``image_metadata``; re-broadcast them here so a + # custom-text Label lookup (e.g. ``predictions["parent_id"]``) resolves + # identically to flag-off instead of raising ``KeyError``. Only keys present + # on this image are emitted, matching flag-off (which omits ``inference_id`` + # when the model did not supply one). Fixed-anchor annotators never read + # these, so adding them leaves their rendered output unchanged. + for lineage_key in ( + PARENT_ID_KEY, + ROOT_PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + INFERENCE_ID_KEY, + ): + if lineage_key in image_metadata: + data[lineage_key] = np.asarray( + [image_metadata[lineage_key]] * detections_number, dtype=object + ) + extra_keys = set() + for per_box in bboxes_metadata: + extra_keys.update(per_box.keys()) + extra_keys.discard(DETECTION_ID_KEY) + extra_keys.discard(TRACKER_ID_KEY) + # The private host mirror of the box tensors is an internal transport + # channel, not per-box user data โ€” it must never surface in ``.data``. + extra_keys.difference_update(HOST_MIRROR_KEYS) + for key in extra_keys: + data[key] = np.asarray( + [per_box.get(key) for per_box in bboxes_metadata], dtype=object + ) + return data + + +def _resolve_class_name(class_id: int, class_names_mapping: dict) -> str: + class_name = class_names_mapping.get(class_id) + if class_name is None: + return f"class_{class_id}" + return str(class_name) + + +class VisualizationManifest(WorkflowBlockManifest, ABC): + model_config = ConfigDict( + json_schema_extra={ + "license": "Apache-2.0", + "block_type": "visualization", + } + ) + image: Selector(kind=[IMAGE_KIND]) = Field( + title="Input Image", + description="The image to visualize on.", + examples=["$inputs.image", "$steps.cropping.crops"], + validation_alias=AliasChoices("image", "images"), + ) + copy_image: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + description="Enable this option to create a copy of the input image for visualization, preserving the original. Use this when stacking multiple visualizations.", + default=True, + examples=[True, False], + ) + + @classmethod + def describe_outputs(cls) -> List[OutputDefinition]: + return [ + OutputDefinition( + name=OUTPUT_IMAGE_KEY, + kind=[ + IMAGE_KIND, + ], + ), + ] + + +class VisualizationBlock(WorkflowBlock, ABC): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @classmethod + @abstractmethod + def get_manifest(cls) -> Type[VisualizationManifest]: + pass + + @abstractmethod + def getAnnotator(self, *args, **kwargs) -> sv.annotators.base.BaseAnnotator: + pass + + @abstractmethod + def run( + self, image: WorkflowImageData, copy_image: bool, *args, **kwargs + ) -> BlockResult: + pass + + +class PredictionsVisualizationManifest(VisualizationManifest, ABC): + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Model predictions to visualize.", + examples=["$steps.object_detection_model.predictions"], + ) + + +class PredictionsVisualizationBlock(VisualizationBlock, ABC): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + @classmethod + @abstractmethod + def get_manifest(cls) -> Type[VisualizationManifest]: + pass + + @abstractmethod + def getAnnotator(self, *args, **kwargs) -> sv.annotators.base.BaseAnnotator: + pass + + @abstractmethod + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + *args, + **kwargs, + ) -> BlockResult: + pass diff --git a/inference/core/workflows/core_steps/visualizations/corner/v1.py b/inference/core/workflows/core_steps/visualizations/corner/v1.py index 2196377a4d..7def44a333 100644 --- a/inference/core/workflows/core_steps/visualizations/corner/v1.py +++ b/inference/core/workflows/core_steps/visualizations/corner/v1.py @@ -164,8 +164,13 @@ def run( thickness, corner_length, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/corner/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/corner/v1_tensor.py new file mode 100644 index 0000000000..f17418a91f --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/corner/v1_tensor.py @@ -0,0 +1,190 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/corner_visualization@v1" +SHORT_DESCRIPTION = "Draw the corners of detected objects in an image." +LONG_DESCRIPTION = """ +Draw corner markers at the four corners of detected object bounding boxes, providing a minimal, clean visualization style that marks object locations without full bounding box outlines. + +## How This Block Works + +This block takes an image and detection predictions and draws corner markers at the four corners of each detected object's bounding box. The block: + +1. Takes an image and predictions as input +2. Identifies bounding box coordinates for each detected object +3. Calculates the four corner positions (top-left, top-right, bottom-left, bottom-right) of each bounding box +4. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +5. Draws corner markers (typically L-shaped lines or corner indicators) at each corner position using Supervision's BoxCornerAnnotator +6. Applies the specified thickness and corner length to control the appearance of the corner markers +7. Returns an annotated image with corner markers overlaid on the original image + +The block draws minimal corner markers instead of full bounding boxes, creating a clean, unobtrusive visualization style. This approach marks object locations clearly while maintaining a minimal aesthetic that doesn't overwhelm the image with full rectangular outlines. The corner markers can be customized with different thickness and length values, and colors can be assigned based on object class, index, or tracking ID, making it easy to distinguish between different objects or object types. + +## Common Use Cases + +- **Minimal Object Marking**: Mark detected objects with corner indicators instead of full bounding boxes for a clean, unobtrusive visualization style that preserves image clarity while still indicating object locations +- **Aesthetic Visualization Design**: Create visually minimal annotations for presentations, dashboards, or user interfaces where full bounding boxes would be too visually intrusive but corner markers provide sufficient location indication +- **Dense Scene Visualization**: Use corner markers when working with many detected objects in dense scenes where full bounding boxes would overlap excessively and create visual clutter +- **Design-Oriented Applications**: Apply corner markers in design workflows, artistic visualizations, or creative applications where a minimal, modern aesthetic is preferred over traditional bounding box outlines +- **Subtle Object Highlighting**: Mark object locations subtly without drawing attention away from the main image content, useful for background annotations or when object location indication is needed without visual prominence +- **UI and Dashboard Integration**: Integrate corner markers into user interfaces, dashboards, or interactive applications where minimal visual indicators are preferred for better user experience and reduced visual noise + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Dot Visualization, Bounding Box Visualization) to combine corner markers with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images with corner markers for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with corner markers to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with corner markers as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with corner markers for live monitoring, tracking visualization, or post-processing analysis +""" + + +class CornerManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "CornerVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Corner Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-expand", + "blockPriority": 7, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the corner marker lines in pixels. Higher values create thicker, more visible corner markers.", + default=4, + examples=[4, "$inputs.thickness"], + ) + + corner_length: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Length of each corner marker line segment in pixels. This controls how long the corner indicators extend from each corner point. Higher values create longer, more prominent corner markers.", + default=15, + examples=[15, "$inputs.corner_length"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class CornerVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return CornerManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + thickness: int, + corner_length: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + thickness, + corner_length, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = sv.BoxCornerAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + thickness=thickness, + corner_length=corner_length, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + thickness: Optional[int], + corner_length: Optional[int], + ) -> BlockResult: + # sv.BoxCornerAnnotator draws from `xyxy` only and never reads `.mask`; + # skip the device->host dense-mask materialisation. + predictions = to_supervision_for_annotation( + predictions, materialise_masks=False + ) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + thickness, + corner_length, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/crop/v1.py b/inference/core/workflows/core_steps/visualizations/crop/v1.py index a78e0441b3..a51c68739f 100644 --- a/inference/core/workflows/core_steps/visualizations/crop/v1.py +++ b/inference/core/workflows/core_steps/visualizations/crop/v1.py @@ -192,8 +192,13 @@ def run( scale_factor, border_thickness, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/crop/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/crop/v1_tensor.py new file mode 100644 index 0000000000..2777f0e511 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/crop/v1_tensor.py @@ -0,0 +1,214 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + INTEGER_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/crop_visualization@v1" +SHORT_DESCRIPTION = "Draw scaled up crops of detections on the scene." +LONG_DESCRIPTION = """ +Display scaled-up, zoomed-in views of detected objects overlaid on the original image, allowing detailed inspection of small or distant objects while maintaining context with the full scene. + +## How This Block Works + +This block takes an image and detection predictions and creates scaled-up, zoomed-in crops of each detected object, then displays these enlarged crops on the original image. The block: + +1. Takes an image and predictions as input +2. Identifies detected regions from bounding boxes or segmentation masks +3. Extracts the image region for each detected object (crops the object from the original image) +4. Scales up each crop by the specified scale factor (e.g., 2x makes objects twice as large) +5. Applies color styling to the crop border based on the selected color palette, with colors assigned by class, index, or track ID +6. Positions the scaled crop on the image at the specified anchor point relative to the original detection location using Supervision's CropAnnotator +7. Draws a colored border around the scaled crop with the specified thickness +8. Returns an annotated image with scaled-up object crops overlaid on the original image + +The block works with both object detection predictions (using bounding boxes) and instance segmentation predictions (using masks). When masks are available, it crops the exact shape of detected objects; otherwise, it crops rectangular bounding box regions. The scale factor allows you to zoom in on objects, making small or distant objects more visible and easier to inspect. The scaled crops are positioned relative to their original detection locations, allowing you to see both the zoomed-in detail and the object's position in the full scene context. + +## Common Use Cases + +- **Small Object Inspection**: Zoom in on small detected objects (e.g., defects, small products, distant objects) to make them more visible and easier to inspect while maintaining scene context +- **Detail Visualization**: Display enlarged views of detected objects for detailed analysis, quality control, or inspection workflows where fine details need to be visible +- **Multi-Scale Object Display**: Show both the full scene and zoomed-in object details simultaneously, useful for applications where context and detail are both important +- **Quality Control and Inspection**: Inspect detected defects, products, or components at higher magnification while keeping the original detection location visible for reference +- **Presentation and Reporting**: Create visualizations that highlight detected objects with zoomed-in views for reports, documentation, or presentations where both overview and detail are needed +- **User Interface Enhancement**: Provide zoomed-in object views in user interfaces, dashboards, or interactive applications where users need to see object details without losing scene context + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Bounding Box Visualization, Polygon Visualization) to combine scaled crops with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save images with scaled crops for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with scaled crops to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with scaled crops as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with scaled crops for live monitoring, detailed inspection, or post-processing analysis +""" + + +class CropManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "CropVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Crop Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-crop-alt", + "blockPriority": 8, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + position: Union[ + Literal[ + "CENTER", + "CENTER_LEFT", + "CENTER_RIGHT", + "TOP_CENTER", + "TOP_LEFT", + "TOP_RIGHT", + "BOTTOM_LEFT", + "BOTTOM_CENTER", + "BOTTOM_RIGHT", + "CENTER_OF_MASS", + ], + Selector(kind=[STRING_KIND]), + ] = Field( # type: ignore + description="Anchor position for placing the scaled crop relative to the original detection's bounding box. Options include: CENTER (center of box), corners (TOP_LEFT, TOP_RIGHT, BOTTOM_LEFT, BOTTOM_RIGHT), edge midpoints (TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, BOTTOM_CENTER), or CENTER_OF_MASS (center of mass of the object). The scaled crop will be positioned at this anchor point relative to the original detection location.", + default="TOP_CENTER", + examples=["CENTER", "$inputs.position"], + ) + + scale_factor: Union[float, Selector(kind=[FLOAT_KIND])] = Field( # type: ignore + description="Factor by which to scale (zoom) the cropped object region. A factor of 2.0 doubles the size of the crop, making objects twice as large. A factor of 1.0 shows the crop at original size. Higher values (e.g., 3.0, 4.0) create more zoomed-in views, useful for inspecting small or distant objects. Lower values (e.g., 1.5) provide subtle magnification.", + default=2.0, + examples=[2.0, "$inputs.scale_factor"], + ) + + border_thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the border outline around the scaled crop in pixels. Higher values create thicker, more visible borders that help distinguish the scaled crop from the background.", + default=2, + examples=[2, "$inputs.border_thickness"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class CropVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return CropManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + position: str, + scale_factor: float, + border_thickness: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + position, + scale_factor, + border_thickness, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = sv.CropAnnotator( + border_color=palette, + border_color_lookup=getattr(sv.ColorLookup, color_axis), + position=getattr(sv.Position, position), + scale_factor=scale_factor, + border_thickness=border_thickness, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + position: Optional[str], + scale_factor: Optional[float], + border_thickness: Optional[int], + ) -> BlockResult: + predictions = to_supervision_for_annotation(predictions) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + position, + scale_factor, + border_thickness, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/dot/v1.py b/inference/core/workflows/core_steps/visualizations/dot/v1.py index 8bd83af6d7..1d80d0536f 100644 --- a/inference/core/workflows/core_steps/visualizations/dot/v1.py +++ b/inference/core/workflows/core_steps/visualizations/dot/v1.py @@ -183,8 +183,13 @@ def run( radius, outline_thickness, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/dot/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/dot/v1_tensor.py new file mode 100644 index 0000000000..d0a94d3ed6 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/dot/v1_tensor.py @@ -0,0 +1,205 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/dot_visualization@v1" +SHORT_DESCRIPTION = ( + "Draw dots on an image at specific coordinates based on provided detections." +) +LONG_DESCRIPTION = """ +Draw circular dots on an image to mark specific points on detected objects, with customizable position, size, color, and outline styling. + +## How This Block Works + +This block takes an image and detection predictions and draws circular dot markers at specified anchor positions on each detected object. The block: + +1. Takes an image and predictions as input +2. Determines the dot position for each detection based on the selected anchor point (center, corners, edges, or center of mass) +3. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +4. Draws circular dots with the specified radius and optional outline thickness using Supervision's DotAnnotator +5. Returns an annotated image with dots overlaid on the original image + +The block supports various position options including the center of the bounding box, any of the four corners, edge midpoints, or the center of mass (useful for objects with irregular shapes). Dots can be customized with different sizes (radius), optional outlines for better visibility, and various color palettes. This provides a minimal, clean visualization style that marks detection locations without the visual clutter of full bounding boxes, making it ideal for dense scenes or when you need to highlight specific points of interest. + +## Common Use Cases + +- **Minimal Object Marking**: Mark detected objects with small dots instead of bounding boxes for cleaner, less cluttered visualizations when working with dense scenes or many detections +- **Point of Interest Highlighting**: Mark specific anchor points (corners, center, center of mass) on detected objects for applications like object tracking, pose estimation, or spatial analysis +- **Tracking Visualization**: Use dots to visualize object trajectories or tracking IDs over time, creating a cleaner alternative to bounding boxes for tracking workflows +- **Crowd Counting and Density Analysis**: Mark people or objects with dots to visualize density patterns, crowd distribution, or object counts without overlapping bounding boxes +- **Keypoint and Landmark Marking**: Mark specific points on objects (such as the center of mass for irregular shapes) for physics simulations, measurement workflows, or spatial relationship analysis +- **Minimal UI Overlays**: Create clean, unobtrusive visual overlays for user interfaces, dashboards, or mobile applications where full bounding boxes would be too visually intrusive + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Bounding Box Visualization, Label Visualization, Trace Visualization) to combine dot markers with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images with dot markers for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with dot markers to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with dot markers as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with dot markers for live monitoring, tracking visualization, or post-processing analysis +""" + + +class DotManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "DotVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Dot Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-palette", + "blockPriority": 1, + "opencv": True, + }, + } + ) + + position: Union[ + Literal[ + "CENTER", + "CENTER_LEFT", + "CENTER_RIGHT", + "TOP_CENTER", + "TOP_LEFT", + "TOP_RIGHT", + "BOTTOM_LEFT", + "BOTTOM_CENTER", + "BOTTOM_RIGHT", + "CENTER_OF_MASS", + ], + Selector(kind=[STRING_KIND]), + ] = Field( # type: ignore + description="Anchor position for placing the dot relative to each detection's bounding box. Options include: CENTER (center of box), corners (TOP_LEFT, TOP_RIGHT, BOTTOM_LEFT, BOTTOM_RIGHT), edge midpoints (TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, BOTTOM_CENTER), or CENTER_OF_MASS (center of mass of the object, useful for irregular shapes).", + default="CENTER", + examples=["CENTER", "$inputs.position"], + ) + + radius: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Radius of the dot in pixels. Higher values create larger, more visible dots.", + default=4, + examples=[4, "$inputs.radius"], + ) + + outline_thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the dot outline in pixels. Set to 0 for no outline (filled dots only). Higher values create thicker outlines around the dot for better visibility against varying backgrounds.", + default=0, + examples=[2, "$inputs.outline_thickness"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class DotVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return DotManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + position: str, + radius: int, + outline_thickness: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + position, + radius, + outline_thickness, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = sv.DotAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + position=getattr(sv.Position, position), + radius=radius, + outline_thickness=outline_thickness, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + position: Optional[str], + radius: Optional[int], + outline_thickness: Optional[int], + ) -> BlockResult: + predictions = to_supervision_for_annotation(predictions) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + position, + radius, + outline_thickness, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/ellipse/v1.py b/inference/core/workflows/core_steps/visualizations/ellipse/v1.py index 888dab6ab6..144c77d093 100644 --- a/inference/core/workflows/core_steps/visualizations/ellipse/v1.py +++ b/inference/core/workflows/core_steps/visualizations/ellipse/v1.py @@ -176,8 +176,13 @@ def run( start_angle, end_angle, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/ellipse/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/ellipse/v1_tensor.py new file mode 100644 index 0000000000..3d92303959 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/ellipse/v1_tensor.py @@ -0,0 +1,202 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/ellipse_visualization@v1" +SHORT_DESCRIPTION = "Draw ellipses that highlight detected objects in an image." +LONG_DESCRIPTION = """ +Draw elliptical outlines around detected objects, providing oval-shaped annotations that can be customized to full ellipses or partial arcs, offering more flexible geometry than circular visualizations. + +## How This Block Works + +This block takes an image and detection predictions and draws elliptical (oval-shaped) outlines around each detected object. The block: + +1. Takes an image and predictions as input +2. Identifies bounding box coordinates for each detected object +3. Calculates the center point and dimensions (width and height) for each detection based on its bounding box +4. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +5. Draws elliptical outlines around each detected object using Supervision's EllipseAnnotator +6. Applies the specified thickness to control the line width of the elliptical outlines +7. Uses start and end angle parameters to draw either full ellipses or partial elliptical arcs +8. Returns an annotated image with elliptical outlines overlaid on the original image + +The block draws ellipses that are typically fitted to each detection's bounding box, creating oval-shaped outlines that adapt to the object's aspect ratio (unlike circles which are always round). Ellipses provide more flexible geometry than circles, as they can represent both round and elongated objects more accurately. The start and end angle parameters allow you to draw full ellipses (360 degrees) or partial arcs, providing additional visual style options. Unlike circle visualization (which always draws complete circles), ellipse visualization can draw partial arcs, creating distinctive visual markers while still clearly indicating object locations. + +## Common Use Cases + +- **Oval and Elongated Object Highlighting**: Highlight detected objects with elliptical outlines when objects are oval-shaped or elongated (e.g., vehicles, elongated products, elliptical objects) where ellipses provide a better fit than circular or rectangular shapes +- **Flexible Geometric Visualization**: Use elliptical shapes as a more geometrically flexible alternative to circles or bounding boxes, adapting to object aspect ratios for more accurate visual representation +- **Partial Arc Annotations**: Draw partial elliptical arcs (using start and end angles) to create distinctive, stylized visual markers that indicate object locations with a unique visual style +- **Aesthetic Visualization Alternatives**: Create visually distinct annotations compared to standard bounding boxes or circles for design purposes, artistic visualizations, or when elliptical shapes better match design aesthetics +- **Object Shape Adaptation**: Use ellipses when objects have non-square aspect ratios where circular outlines would be less accurate, providing better visual fit for rectangular or elongated objects +- **Design and UI Applications**: Integrate elliptical outlines into user interfaces, dashboards, or interactive applications where oval shapes provide a softer, more organic visual style than angular rectangles + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Dot Visualization, Bounding Box Visualization) to combine elliptical outlines with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images with elliptical outlines for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with elliptical outlines to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with elliptical outlines as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with elliptical outlines for live monitoring, tracking visualization, or post-processing analysis +""" + + +class EllipseManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "EllipseVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Ellipse Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "fad fa-dot-circle", + "blockPriority": 10, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the ellipse outline in pixels. Higher values create thicker, more visible elliptical outlines.", + default=2, + examples=[2, "$inputs.thickness"], + ) + + start_angle: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Starting angle for drawing the ellipse arc in degrees. Used together with end_angle to control whether a full ellipse (360 degrees) or partial arc is drawn. Angles are measured from the positive x-axis (0 degrees = right, 90 degrees = down, 180 degrees = left, 270 degrees = up).", + default=-45, + examples=[-45, "$inputs.start_angle"], + ) + + end_angle: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Ending angle for drawing the ellipse arc in degrees. Used together with start_angle to control whether a full ellipse (360 degrees) or partial arc is drawn. To draw a full ellipse, set end_angle = start_angle + 360 (or equivalent). Angles are measured from the positive x-axis (0 degrees = right, 90 degrees = down, 180 degrees = left, 270 degrees = up).", + default=235, + examples=[235, "$inputs.end_angle"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class EllipseVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return EllipseManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + thickness: int, + start_angle: int, + end_angle: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + thickness, + start_angle, + end_angle, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = sv.EllipseAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + thickness=thickness, + start_angle=start_angle, + end_angle=end_angle, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + thickness: Optional[int], + start_angle: Optional[int], + end_angle: Optional[int], + ) -> BlockResult: + # sv.EllipseAnnotator draws from `xyxy` only and never reads `.mask`; + # skip the device->host dense-mask materialisation. + predictions = to_supervision_for_annotation( + predictions, materialise_masks=False + ) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + thickness, + start_angle, + end_angle, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/halo/v1.py b/inference/core/workflows/core_steps/visualizations/halo/v1.py index 58a47204ee..9167bb00a8 100644 --- a/inference/core/workflows/core_steps/visualizations/halo/v1.py +++ b/inference/core/workflows/core_steps/visualizations/halo/v1.py @@ -180,8 +180,13 @@ def run( opacity, kernel_size, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/halo/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/halo/v1_tensor.py new file mode 100644 index 0000000000..7515bed75b --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/halo/v1_tensor.py @@ -0,0 +1,204 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.annotators.halo import ( + HaloAnnotator, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/halo_visualization@v1" +SHORT_DESCRIPTION = "Paint a halo around detected objects in an image." +LONG_DESCRIPTION = """ +Create a soft, glowing halo effect around detected objects by blurring and overlaying colored masks, providing a distinctive visual style that highlights object boundaries with a smooth, illuminated appearance. + +## How This Block Works + +This block takes an image and instance segmentation predictions (with masks) and creates a glowing halo effect around each detected object. The block: + +1. Takes an image and instance segmentation predictions (with masks) as input +2. Extracts segmentation masks for each detected object (uses masks from predictions, or creates bounding box masks if masks are not available) +3. Applies color styling to each mask based on the selected color palette, with colors assigned by class, index, or track ID +4. Creates colored mask overlays for each detection, combining masks from largest to smallest area (to handle overlapping objects correctly) +5. Applies a blur filter (average pooling with specified kernel size) to the colored masks, creating a soft, diffused halo effect around object edges +6. Blends the blurred halo overlay with the original image using the specified opacity level, creating a glowing appearance around detected objects +7. Returns an annotated image with soft halo effects overlaid around each detected object + +The block creates halos by blurring the colored masks, which produces a soft, glowing effect that extends beyond the object boundaries. Unlike hard-edged visualizations (like bounding boxes or polygons), halos provide a smooth, illuminated appearance that makes objects stand out while maintaining a visually appealing aesthetic. The blur kernel size controls how far the halo extends beyond the object (larger kernel = wider halo), and the opacity controls the intensity of the glow effect. This block requires instance segmentation predictions with masks, as it uses mask shapes to create the halo effect around object perimeters. + +## Common Use Cases + +- **Artistic and Aesthetic Visualizations**: Create visually appealing, glowing effects around detected objects for artistic presentations, design applications, or user interfaces where soft, illuminated halos provide a modern, polished appearance +- **Soft Object Highlighting**: Highlight detected objects with gentle, diffused halos when hard edges would be too harsh or distracting, useful for presentations, marketing materials, or consumer-facing applications +- **Overlapping Object Visualization**: Use halos to visualize overlapping or closely-spaced objects where hard boundaries would create visual clutter, allowing multiple objects to be distinguished while maintaining visual clarity +- **Brand and Design Applications**: Integrate halo effects into brand visuals, promotional materials, or design systems where soft, glowing annotations match design aesthetics better than angular bounding boxes +- **Visual Emphasis and Focus**: Draw attention to detected objects with glowing halos that create a natural visual focus point, useful in dashboards, monitoring interfaces, or interactive applications +- **Mask-Based Object Highlighting**: Visualize instance segmentation results with soft halo effects, providing an alternative to solid mask overlays when you want to show object boundaries without obscuring image details + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Dot Visualization, Bounding Box Visualization) to combine halo effects with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save images with halo effects for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with halo effects to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with halo effects as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with halo effects for live monitoring, artistic visualizations, or post-processing analysis +""" + + +class HaloManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "HaloVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Halo Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-lightbulb-on", + "blockPriority": 11, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + predictions: Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Instance segmentation predictions containing masks for detected objects. The block uses segmentation masks to create halo effects around object boundaries. If masks are not available, it will create masks from bounding boxes. Requires instance segmentation model outputs with mask data.", + examples=["$steps.instance_segmentation_model.predictions"], + ) + + opacity: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + description="Opacity of the halo overlay, ranging from 0.0 (fully transparent) to 1.0 (fully opaque). Controls the intensity of the glowing halo effect. Lower values create more subtle, softer halos that blend with the background, while higher values create more intense, visible glows. Typical values range from 0.5 to 0.9 for balanced visual effects.", + default=0.8, + examples=[0.8, "$inputs.opacity"], + ) + + kernel_size: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Size of the blur kernel (in pixels) used for creating the halo effect. This controls how far the halo extends beyond the object boundaries and how soft/diffused the glow appears. Larger values create wider, more spread-out halos with smoother gradients, while smaller values create tighter, more concentrated glows. Values typically range from 20 to 80 pixels, with 40 being a good default for most use cases.", + default=40, + examples=[40, "$inputs.kernel_size"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class HaloVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return HaloManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + opacity: float, + kernel_size: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + opacity, + kernel_size, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = HaloAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + opacity=opacity, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + opacity: Optional[float], + kernel_size: Optional[int], + ) -> BlockResult: + predictions = to_supervision_for_annotation(predictions) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + opacity, + kernel_size, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/halo/v2.py b/inference/core/workflows/core_steps/visualizations/halo/v2.py index c2a232e4a9..a516700c9a 100644 --- a/inference/core/workflows/core_steps/visualizations/halo/v2.py +++ b/inference/core/workflows/core_steps/visualizations/halo/v2.py @@ -180,8 +180,13 @@ def run( opacity, kernel_size, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/halo/v2_tensor.py b/inference/core/workflows/core_steps/visualizations/halo/v2_tensor.py new file mode 100644 index 0000000000..1e9c7b44b6 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/halo/v2_tensor.py @@ -0,0 +1,204 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field +from supervision import HaloAnnotator + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + INTEGER_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/halo_visualization@v2" +SHORT_DESCRIPTION = "Paint a halo around detected objects in an image." +LONG_DESCRIPTION = """ +Create a soft, glowing halo effect around detected objects by blurring and overlaying colored masks, providing a distinctive visual style that highlights object boundaries with a smooth, illuminated appearance. + +## How This Block Works + +This block takes an image and instance segmentation predictions (with masks) and creates a glowing halo effect around each detected object. The block: + +1. Takes an image and instance segmentation predictions (with masks) as input +2. Extracts segmentation masks for each detected object (uses masks from predictions, or creates bounding box masks if masks are not available) +3. Applies color styling to each mask based on the selected color palette, with colors assigned by class, index, or track ID +4. Creates colored mask overlays for each detection, combining masks from largest to smallest area (to handle overlapping objects correctly) +5. Applies a blur filter (average pooling with specified kernel size) to the colored masks, creating a soft, diffused halo effect around object edges +6. Blends the blurred halo overlay with the original image using the specified opacity level, creating a glowing appearance around detected objects +7. Returns an annotated image with soft halo effects overlaid around each detected object + +The block creates halos by blurring the colored masks, which produces a soft, glowing effect that extends beyond the object boundaries. Unlike hard-edged visualizations (like bounding boxes or polygons), halos provide a smooth, illuminated appearance that makes objects stand out while maintaining a visually appealing aesthetic. The blur kernel size controls how far the halo extends beyond the object (larger kernel = wider halo), and the opacity controls the intensity of the glow effect. This block requires instance segmentation predictions with masks, as it uses mask shapes to create the halo effect around object perimeters. + +## Common Use Cases + +- **Artistic and Aesthetic Visualizations**: Create visually appealing, glowing effects around detected objects for artistic presentations, design applications, or user interfaces where soft, illuminated halos provide a modern, polished appearance +- **Soft Object Highlighting**: Highlight detected objects with gentle, diffused halos when hard edges would be too harsh or distracting, useful for presentations, marketing materials, or consumer-facing applications +- **Overlapping Object Visualization**: Use halos to visualize overlapping or closely-spaced objects where hard boundaries would create visual clutter, allowing multiple objects to be distinguished while maintaining visual clarity +- **Brand and Design Applications**: Integrate halo effects into brand visuals, promotional materials, or design systems where soft, glowing annotations match design aesthetics better than angular bounding boxes +- **Visual Emphasis and Focus**: Draw attention to detected objects with glowing halos that create a natural visual focus point, useful in dashboards, monitoring interfaces, or interactive applications +- **Mask-Based Object Highlighting**: Visualize instance segmentation results with soft halo effects, providing an alternative to solid mask overlays when you want to show object boundaries without obscuring image details + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Dot Visualization, Bounding Box Visualization) to combine halo effects with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save images with halo effects for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with halo effects to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with halo effects as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with halo effects for live monitoring, artistic visualizations, or post-processing analysis +""" + + +class HaloManifest(ColorableVisualizationManifest): + type: Literal[TYPE] + model_config = ConfigDict( + json_schema_extra={ + "name": "Halo Visualization", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-lightbulb-on", + "blockPriority": 11, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + predictions: Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Instance segmentation predictions containing masks for detected objects. The block uses segmentation masks to create halo effects around object boundaries. If masks are not available, it will create masks from bounding boxes. Requires instance segmentation model outputs with mask data.", + examples=["$steps.instance_segmentation_model.predictions"], + ) + + opacity: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + description="Opacity of the halo overlay, ranging from 0.0 (fully transparent) to 1.0 (fully opaque). Controls the intensity of the glowing halo effect. Lower values create more subtle, softer halos that blend with the background, while higher values create more intense, visible glows. Typical values range from 0.5 to 0.9 for balanced visual effects.", + default=0.8, + examples=[0.8, "$inputs.opacity"], + ) + + kernel_size: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Size of the blur kernel (in pixels) used for creating the halo effect. This controls how far the halo extends beyond the object boundaries and how soft/diffused the glow appears. Larger values create wider, more spread-out halos with smoother gradients, while smaller values create tighter, more concentrated glows. Values typically range from 20 to 80 pixels, with 40 being a good default for most use cases.", + default=40, + examples=[40, "$inputs.kernel_size"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class HaloVisualizationBlockV2(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return HaloManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + opacity: float, + kernel_size: int, + ) -> HaloAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + opacity, + kernel_size, + "_".join(custom_colors) if custom_colors else "", + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = HaloAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + opacity=opacity, + kernel_size=kernel_size, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + opacity: Optional[float], + kernel_size: Optional[int], + ) -> BlockResult: + predictions = to_supervision_for_annotation(predictions) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + opacity, + kernel_size, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/heatmap/v1.py b/inference/core/workflows/core_steps/visualizations/heatmap/v1.py index f058d10a2f..3b8bbdd2a2 100644 --- a/inference/core/workflows/core_steps/visualizations/heatmap/v1.py +++ b/inference/core/workflows/core_steps/visualizations/heatmap/v1.py @@ -324,8 +324,13 @@ def run( top_hue, low_hue, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=detections_to_plot, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/heatmap/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/heatmap/v1_tensor.py new file mode 100644 index 0000000000..4050df8c80 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/heatmap/v1_tensor.py @@ -0,0 +1,352 @@ +import time +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + PredictionsVisualizationBlock, + PredictionsVisualizationManifest, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import ( + VideoMetadata, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.entities.types import ( + BOOLEAN_KIND, + FLOAT_KIND, + INTEGER_KIND, + STRING_KIND, + VIDEO_METADATA_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + Runtime, + RuntimeInputMode, + RuntimeRestriction, + Severity, + WorkflowBlockManifest, +) + +TYPE: str = "roboflow_core/heatmap_visualization@v1" +SHORT_DESCRIPTION = "Draw a heatmap based on detections in an image." +LONG_DESCRIPTION = """ +Draw heatmaps on an image based on provided detections. Heat accumulates over time and is drawn as a semi-transparent overlay of blurred circles. + +## How This Block Works + +This block takes an image and detection predictions and draws a heatmap. The block: + +1. Takes an image and predictions as input. +2. Accumulates heat based on the position of detections. +3. Draws a semi-transparent overlay of blurred circles representing the heat. + +## Common Use Cases + +- **Density Analysis**: Visualize the density of objects in a scene. +- **Traffic Monitoring**: Identify high-traffic areas. +- **Retail Analytics**: Analyze foot traffic in stores. +""" + + +class HeatmapManifest(PredictionsVisualizationManifest): + type: Literal[f"{TYPE}", "HeatmapVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Heatmap Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator", "heatmap"], + "ui_manifest": { + "section": "visualization", + "icon": "fas fa-fire", + "blockPriority": 4, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + metadata: Selector(kind=[VIDEO_METADATA_KIND]) = Field( + description="Video metadata containing video_identifier to maintain separate state for different videos.", + default=None, + examples=["$inputs.video_metadata"], + ) + + position: Union[ + Literal[ + "CENTER", + "CENTER_LEFT", + "CENTER_RIGHT", + "TOP_CENTER", + "TOP_LEFT", + "TOP_RIGHT", + "BOTTOM_CENTER", + "BOTTOM_LEFT", + "BOTTOM_RIGHT", + ], + Selector(kind=[STRING_KIND]), + ] = Field( # type: ignore + default="BOTTOM_CENTER", + description="The position of the heatmap relative to the detection.", + examples=["BOTTOM_CENTER", "$inputs.position"], + ) + + opacity: Union[float, Selector(kind=[FLOAT_KIND])] = Field( # type: ignore + description="Opacity of the overlay mask, between 0 and 1.", + default=0.2, + examples=[0.2, "$inputs.opacity"], + ) + + radius: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Radius of the heat circle.", + default=40, + examples=[40, "$inputs.radius"], + ) + + kernel_size: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Kernel size for blurring the heatmap.", + default=25, + examples=[25, "$inputs.kernel_size"], + ) + + top_hue: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Hue at the top of the heatmap. Defaults to 0 (red).", + default=0, + examples=[0, "$inputs.top_hue"], + ) + + low_hue: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Hue at the bottom of the heatmap. Defaults to 125 (blue).", + default=125, + examples=[125, "$inputs.low_hue"], + ) + + ignore_stationary: Union[bool, Selector(kind=[BOOLEAN_KIND])] = Field( # type: ignore + description="If True, only moving objects (based on tracker ID) will contribute to the heatmap.", + default=True, + examples=[True, "$inputs.ignore_stationary"], + ) + + motion_threshold: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Minimum movement in pixels required to consider an object as moving.", + default=25, + examples=[25, "$inputs.motion_threshold"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restriction = RuntimeRestriction( + severity=Severity.SOFT, + note=( + "Heatmap accumulation and stationary-object filtering keep " + "per-video tracking state in process memory. With remote step " + "execution on stateless or multi-replica HTTP runtimes, " + "successive frames may be served by different worker " + "processes, so heat history resets or splits across workers. " + "Use local step execution in an InferencePipeline for stable " + "cross-frame visualizations." + ), + applies_to_runtimes=[ + Runtime.HOSTED_SERVERLESS, + Runtime.DEDICATED_DEPLOYMENT, + ], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + applies_to_input_modes=[RuntimeInputMode.VIDEO], + ) + return [restriction, STILL_IMAGE_INPUT_SOFT_RESTRICTION] + + +class HeatmapVisualizationBlockV1(PredictionsVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + # Dictionary to store track history: {video_id: {tracker_id: (x, y, timestamp)}} + self._track_history: Dict[str, Dict[int, tuple]] = {} + self._last_cleanup_time = time.time() + self._cleanup_interval = 10.0 # seconds + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return HeatmapManifest + + def _cleanup_history(self): + current_time = time.time() + if current_time - self._last_cleanup_time < self._cleanup_interval: + return + + # Clean up stale trackers (e.g., older than 60s) + # Using 60s as a conservative estimate for ~1800 frames at 30fps + stale_threshold = 60.0 + empty_videos = [] + + for video_id, history in self._track_history.items(): + expired_trackers = [ + tid + for tid, data in history.items() + if current_time - data[2] > stale_threshold + ] + for tid in expired_trackers: + del history[tid] + + if not history: + empty_videos.append(video_id) + + # Clean up empty video histories + for video_id in empty_videos: + del self._track_history[video_id] + + self._last_cleanup_time = current_time + + def getAnnotator( + self, + video_id: str, + position: str, + opacity: float, + radius: int, + kernel_size: int, + top_hue: int, + low_hue: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + video_id, + position, + opacity, + radius, + kernel_size, + top_hue, + low_hue, + ], + ) + ) + + if key not in self.annotatorCache: + position_enum = getattr(sv.Position, position) + self.annotatorCache[key] = sv.HeatMapAnnotator( + position=position_enum, + opacity=opacity, + radius=radius, + kernel_size=kernel_size, + top_hue=top_hue, + low_hue=low_hue, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + position: Optional[str], + opacity: Optional[float], + radius: Optional[int], + kernel_size: Optional[int], + top_hue: Optional[int], + low_hue: Optional[int], + metadata: Optional[VideoMetadata] = None, + ignore_stationary: bool = True, + motion_threshold: int = 25, + ) -> BlockResult: + # sv.HeatMapAnnotator accumulates heat at box anchor points + # (get_anchors_coordinates); this block's `position` excludes + # CENTER_OF_MASS, so no path reads `.mask`. Skip the device->host + # dense-mask materialisation. + predictions = to_supervision_for_annotation( + predictions, materialise_masks=False + ) + self._cleanup_history() + detections_to_plot = predictions + video_id = metadata.video_identifier if metadata else "default_video" + + if ignore_stationary and predictions.tracker_id is not None: + if video_id not in self._track_history: + self._track_history[video_id] = {} + + current_history = self._track_history[video_id] + moving_indices = [] + current_time = time.time() + + # Calculate centers for current detections + # Use the specified position anchor for tracking consistency + anchor_position = ( + getattr(sv.Position, position) + if position + else sv.Position.BOTTOM_CENTER + ) + anchors = predictions.get_anchors_coordinates(anchor=anchor_position) + + for i, (tracker_id, point) in enumerate( + zip(predictions.tracker_id, anchors) + ): + tracker_id = int(tracker_id) + x, y = point + + if tracker_id in current_history: + # Check for movement + prev_x, prev_y, _ = current_history[tracker_id] + dist = np.sqrt((x - prev_x) ** 2 + (y - prev_y) ** 2) + + if dist >= motion_threshold: + moving_indices.append(i) + # Update history with new position and timestamp + current_history[tracker_id] = (x, y, current_time) + else: + # New track, initialize history + current_history[tracker_id] = (x, y, current_time) + + # Filter detections + if len(moving_indices) > 0: + detections_to_plot = predictions[np.array(moving_indices)] + else: + detections_to_plot = sv.Detections.empty() + + annotator = self.getAnnotator( + video_id, + position, + opacity, + radius, + kernel_size, + top_hue, + low_hue, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=detections_to_plot, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/icon/v1.py b/inference/core/workflows/core_steps/visualizations/icon/v1.py index d3462916b6..87f1d9ecff 100644 --- a/inference/core/workflows/core_steps/visualizations/icon/v1.py +++ b/inference/core/workflows/core_steps/visualizations/icon/v1.py @@ -254,7 +254,11 @@ def run( x_position: Optional[int], y_position: Optional[int], ) -> BlockResult: - annotated_image = image.numpy_image.copy() if copy_image else image.numpy_image + annotated_image = image.numpy_image + if copy_image: + annotated_image = annotated_image.copy() + else: + image.declare_numpy_image_mutated() icon_np = icon.numpy_image.copy() import os diff --git a/inference/core/workflows/core_steps/visualizations/icon/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/icon/v1_tensor.py new file mode 100644 index 0000000000..26a1204c00 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/icon/v1_tensor.py @@ -0,0 +1,398 @@ +from typing import Dict, List, Literal, Optional, Type, Union + +import numpy as np +import supervision as sv +from pydantic import AliasChoices, ConfigDict, Field, model_validator + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + VisualizationBlock, + VisualizationManifest, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + IMAGE_KIND, + INTEGER_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/icon_visualization@v1" +SHORT_DESCRIPTION = "Draw icons on an image either at specific static coordinates or dynamically based on detections." +LONG_DESCRIPTION = """ +Place custom icon images on images either at fixed positions (static mode) or dynamically positioned on detected objects (dynamic mode), useful for watermarks, labels, badges, or visual markers. + +## How This Block Works + +This block takes an image and optionally detection predictions, then places a custom icon image on the image. The block supports two modes: + +**Static Mode** (for watermarks and fixed positioning): +1. Takes an image and an icon image as input +2. Places the icon at fixed x and y coordinates on the image +3. Supports negative coordinates for positioning from the right or bottom edges +4. Returns an annotated image with the icon at the specified static location + +**Dynamic Mode** (for detection-based positioning): +1. Takes an image, an icon image, and detection predictions as input +2. Positions the icon on each detected object based on the selected anchor point (center, corners, edges, or center of mass) +3. Places the icon at the same position relative to each detection +4. Returns an annotated image with icons overlaid on detected objects + +The block supports PNG images with transparency (alpha channel), allowing icons to blend naturally with the background. Icons can be resized to any width and height, making them suitable for various use cases from small badges to large watermarks. In static mode, icons are placed at fixed coordinates, making it ideal for watermarks or branding. In dynamic mode, icons automatically follow detected objects, making it useful for labeling, categorizing, or marking detected items with custom visual indicators. + +## Common Use Cases + +- **Watermarks and Branding**: Place logos, watermarks, or branding elements at fixed positions (static mode) on images or videos for content protection, copyright marking, or brand identification +- **Object Labeling with Icons**: Place custom icons on detected objects (dynamic mode) to categorize, label, or mark objects with visual indicators (e.g., warning icons on unsafe objects, category icons for products, status badges) +- **Visual Status Indicators**: Display status icons (e.g., checkmarks, warning signs, information badges) on detected objects based on classification results, confidence levels, or custom logic for quick visual feedback +- **Product Marking and Categorization**: Place category icons, product type indicators, or custom markers on detected products in retail, e-commerce, or inventory management workflows +- **Custom Annotation Systems**: Create custom annotation workflows with specialized icons for quality control, defect marking, or compliance tracking in manufacturing or inspection workflows +- **Interactive UI Elements**: Add icon-based visual elements to images or videos for user interfaces, dashboards, or interactive applications where custom icons provide intuitive visual cues + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Bounding Box Visualization, Polygon Visualization) to combine icon placement with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save images with icons for documentation, reporting, or archiving +- **Webhook blocks** to send visualized results with icons to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with icons as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with icons for live monitoring, tracking visualization, or post-processing analysis +""" + + +class IconManifest(VisualizationManifest): + type: Literal[f"{TYPE}", "IconVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Icon Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator", "icon", "watermark"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-image", + "blockPriority": 5, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + icon: Selector(kind=[IMAGE_KIND]) = Field( + title="Icon Image", + description="The icon image to place on the input image. PNG format with transparency (alpha channel) is recommended for best results, as it allows the icon to blend naturally with the background. The icon will be resized to the specified width and height.", + examples=["$inputs.icon", "$steps.image_loader.image"], + json_schema_extra={ + "always_visible": True, + "order": 3, + }, + ) + + mode: Union[ + Literal["static", "dynamic"], + Selector(kind=[STRING_KIND]), + ] = Field( + default="dynamic", + description="Mode for placing icons. 'static' mode places the icon at fixed x,y coordinates (useful for watermarks or fixed-position elements). 'dynamic' mode places icons on detected objects based on their positions (useful for object labeling or categorization).", + examples=["static", "dynamic", "$inputs.mode"], + json_schema_extra={ + "always_visible": True, + "order": 1, + }, + ) + + predictions: Optional[ + Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) + ] = Field( + default=None, + description="Model predictions to place icons on (required for dynamic mode). Icons will be positioned on each detected object based on the selected position anchor point.", + examples=["$steps.object_detection_model.predictions"], + json_schema_extra={ + "relevant_for": { + "mode": {"values": ["dynamic"], "required": True}, + }, + "order": 4, + }, + ) + + icon_width: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=64, + description="Width of the icon in pixels. The icon image will be resized to this width while maintaining aspect ratio if height is also specified.", + examples=[64, "$inputs.icon_width"], + json_schema_extra={ + "always_visible": True, + }, + ) + + icon_height: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=64, + description="Height of the icon in pixels. The icon image will be resized to this height while maintaining aspect ratio if width is also specified.", + examples=[64, "$inputs.icon_height"], + json_schema_extra={ + "always_visible": True, + }, + ) + + position: Optional[ + Union[ + Literal[ + "CENTER", + "CENTER_LEFT", + "CENTER_RIGHT", + "TOP_CENTER", + "TOP_LEFT", + "TOP_RIGHT", + "BOTTOM_LEFT", + "BOTTOM_CENTER", + "BOTTOM_RIGHT", + "CENTER_OF_MASS", + ], + Selector(kind=[STRING_KIND]), + ] + ] = Field( + default="TOP_CENTER", + description="Anchor position for placing icons relative to each detection's bounding box (dynamic mode only). Options include: CENTER (center of box), corners (TOP_LEFT, TOP_RIGHT, BOTTOM_LEFT, BOTTOM_RIGHT), edge midpoints (TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, BOTTOM_CENTER), or CENTER_OF_MASS (center of mass of the object).", + examples=["TOP_CENTER", "$inputs.position"], + json_schema_extra={ + "relevant_for": { + "mode": {"values": ["dynamic"], "required": False}, + }, + }, + ) + + x_position: Optional[Union[int, Selector(kind=[INTEGER_KIND])]] = Field( + default=10, + description="X coordinate for static mode positioning. Positive values position from the left edge of the image. Negative values position from the right edge (e.g., -10 places the icon 10 pixels from the right edge).", + examples=[10, -10, "$inputs.x_position"], + json_schema_extra={ + "relevant_for": { + "mode": {"values": ["static"], "required": True}, + }, + }, + ) + + y_position: Optional[Union[int, Selector(kind=[INTEGER_KIND])]] = Field( + default=10, + description="Y coordinate for static mode positioning. Positive values position from the top edge of the image. Negative values position from the bottom edge (e.g., -10 places the icon 10 pixels from the bottom edge).", + examples=[10, -10, "$inputs.y_position"], + json_schema_extra={ + "relevant_for": { + "mode": {"values": ["static"], "required": True}, + }, + }, + ) + + @model_validator(mode="after") + def validate_mode_parameters(self) -> "IconManifest": + if self.mode == "dynamic": + if self.predictions is None: + raise ValueError("The 'predictions' field is required for dynamic mode") + return self + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class IconVisualizationBlockV1(VisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return IconManifest + + def getAnnotator( + self, + icon_width: int, + icon_height: int, + position: Optional[str] = None, + ) -> Optional[sv.annotators.base.BaseAnnotator]: + if position is not None: + key = f"dynamic_{icon_width}_{icon_height}_{position}" + if key not in self.annotatorCache: + self.annotatorCache[key] = sv.IconAnnotator( + icon_resolution_wh=(icon_width, icon_height), + icon_position=getattr(sv.Position, position), + ) + return self.annotatorCache[key] + return None + + def run( + self, + image: WorkflowImageData, + copy_image: bool, + mode: str, + icon: WorkflowImageData, + predictions: Optional[Union[TensorNativePrediction, TensorNativeDetections]], + icon_width: int, + icon_height: int, + position: Optional[str], + x_position: Optional[int], + y_position: Optional[int], + ) -> BlockResult: + if predictions is not None: + predictions = to_supervision_for_annotation(predictions) + annotated_image = image.numpy_image + if copy_image: + annotated_image = annotated_image.copy() + else: + image.declare_numpy_image_mutated() + icon_np = icon.numpy_image.copy() + + import os + import tempfile + + import cv2 + + # WorkflowImageData loses alpha channels when loading images. + # Try to recover them from the original source. + if icon_np.shape[2] == 3: + # Try reloading from file with IMREAD_UNCHANGED + if ( + hasattr(icon, "_image_reference") + and icon._image_reference + and not icon._image_reference.startswith("http") + ): + try: + icon_with_alpha = cv2.imread( + icon._image_reference, cv2.IMREAD_UNCHANGED + ) + if icon_with_alpha is not None and icon_with_alpha.shape[2] == 4: + icon_np = icon_with_alpha + except: + pass + + # Try decoding base64 with alpha preserved + if ( + icon_np.shape[2] == 3 + and hasattr(icon, "_base64_image") + and icon._base64_image + ): + try: + import base64 + + image_bytes = base64.b64decode(icon._base64_image) + nparr = np.frombuffer(image_bytes, np.uint8) + decoded = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED) + if decoded is not None and len(decoded.shape) >= 2: + if len(decoded.shape) == 2: + decoded = cv2.cvtColor(decoded, cv2.COLOR_GRAY2BGR) + if decoded.shape[2] == 4: + icon_np = decoded + except: + pass + + with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f: + # Ensure proper format for IconAnnotator + if len(icon_np.shape) == 2: + icon_np = cv2.cvtColor(icon_np, cv2.COLOR_GRAY2BGR) + alpha = ( + np.ones( + (icon_np.shape[0], icon_np.shape[1], 1), dtype=icon_np.dtype + ) + * 255 + ) + icon_np = np.concatenate([icon_np, alpha], axis=2) + elif icon_np.shape[2] == 3: + alpha = ( + np.ones( + (icon_np.shape[0], icon_np.shape[1], 1), dtype=icon_np.dtype + ) + * 255 + ) + icon_np = np.concatenate([icon_np, alpha], axis=2) + + cv2.imwrite(f.name, icon_np) + icon_path = f.name + + try: + if mode == "static": + img_height, img_width = annotated_image.shape[:2] + + # Handle negative positioning (from right/bottom edges) + if x_position < 0: + actual_x = img_width + x_position - icon_width + else: + actual_x = x_position + + if y_position < 0: + actual_y = img_height + y_position - icon_height + else: + actual_y = y_position + + # IconAnnotator expects a detection, so create one at the desired position + center_x = actual_x + icon_width // 2 + center_y = actual_y + icon_height // 2 + + static_detections = sv.Detections( + xyxy=np.array( + [[center_x - 1, center_y - 1, center_x + 1, center_y + 1]], + dtype=np.float64, + ), + class_id=np.array([0]), + confidence=np.array([1.0]), + ) + + annotator = sv.IconAnnotator( + icon_resolution_wh=(icon_width, icon_height), + icon_position=sv.Position.CENTER, + ) + + annotated_image = annotator.annotate( + scene=annotated_image, + detections=static_detections, + icon_path=icon_path, + ) + + elif mode == "dynamic" and predictions is not None and len(predictions) > 0: + annotator = self.getAnnotator( + icon_width=icon_width, + icon_height=icon_height, + position=position, + ) + + if annotator is not None: + annotated_image = annotator.annotate( + scene=annotated_image, + detections=predictions, + icon_path=icon_path, + ) + finally: + os.unlink(icon_path) + + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/keypoint/v1.py b/inference/core/workflows/core_steps/visualizations/keypoint/v1.py index 1def2aea6c..7b6821c3ef 100644 --- a/inference/core/workflows/core_steps/visualizations/keypoint/v1.py +++ b/inference/core/workflows/core_steps/visualizations/keypoint/v1.py @@ -312,8 +312,13 @@ def run( keypoints = self.convert_detections_to_keypoints(predictions) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, key_points=keypoints, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/keypoint/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/keypoint/v1_tensor.py new file mode 100644 index 0000000000..52e3d192e6 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/keypoint/v1_tensor.py @@ -0,0 +1,304 @@ +from typing import List, Literal, Optional, Tuple, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + VisualizationBlock, + VisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.utils import str_to_color +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + INTEGER_KIND, + LIST_OF_VALUES_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +TYPE: str = "roboflow_core/keypoint_visualization@v1" +SHORT_DESCRIPTION = "Draw keypoints on detected objects in an image." +LONG_DESCRIPTION = """ +Visualize keypoints (landmark points) detected on objects by drawing point markers, connecting edges, or labeled vertices, providing pose estimation visualization for anatomical points, structural landmarks, or object key features. + +## How This Block Works + +This block takes an image and keypoint detection predictions and visualizes the detected keypoints using one of three visualization modes. The block: + +1. Takes an image and keypoint detection predictions as input (predictions must include keypoint coordinates, confidence scores, and class names) +2. Extracts keypoint data (coordinates, confidence values, and class names) from the predictions +3. Converts the detection data into a KeyPoints format suitable for visualization +4. Applies one of three visualization modes based on the annotator_type setting: + - **Edge mode**: Draws connecting lines (edges) between keypoints using specified edge pairs to show keypoint relationships (e.g., skeleton connections in pose estimation) + - **Vertex mode**: Draws circular markers at each keypoint location without connections, showing individual keypoint positions + - **Vertex label mode**: Draws circular markers with text labels identifying each keypoint class name, providing labeled keypoint visualization +5. Applies color styling, sizing, and optional text labeling based on the selected parameters +6. Returns an annotated image with keypoints visualized according to the selected mode + +The block supports three visualization styles to suit different use cases. Edge mode connects related keypoints with lines (useful for pose estimation skeletons or structural relationships), vertex mode shows individual keypoint locations as circular markers, and vertex label mode adds text labels to identify each keypoint type. This visualization is essential for pose estimation workflows, anatomical point detection, or any application where specific landmark points on objects need to be identified and visualized. + +## Common Use Cases + +- **Human Pose Estimation**: Visualize human body keypoints (joints, body parts) for pose estimation, activity recognition, or motion analysis applications where anatomical points need to be displayed with skeleton connections or labeled markers +- **Animal Pose Estimation**: Display animal keypoints for behavior analysis, veterinary applications, or wildlife monitoring where anatomical landmarks need to be visualized for pose analysis or movement tracking +- **Structural Landmark Detection**: Visualize keypoints on objects, structures, or machinery for structural analysis, quality control, or measurement workflows where specific landmark points need to be identified and displayed +- **Facial Landmark Detection**: Display facial keypoints (eye corners, nose tip, mouth corners, etc.) for facial recognition, expression analysis, or face alignment applications where facial features need to be visualized +- **Sports and Movement Analysis**: Visualize keypoints for sports analysis, biomechanics, or movement studies where body positions, joint angles, or movement patterns need to be analyzed and displayed +- **Quality Control and Inspection**: Display keypoints for manufacturing, quality assurance, or inspection workflows where specific points on products or components need to be identified, measured, or validated + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Keypoint Detection Model blocks** to receive keypoint predictions that are visualized with point markers, edges, or labeled vertices +- **Other visualization blocks** (e.g., Bounding Box Visualization, Label Visualization, Polygon Visualization) to combine keypoint visualization with additional annotations for comprehensive pose or structure visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save images with keypoint visualizations for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with keypoints to external systems, APIs, or web applications for display in dashboards, pose analysis tools, or monitoring interfaces +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with keypoints as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with keypoint visualizations for live pose estimation, movement analysis, or post-processing workflows +""" + + +class KeypointManifest(VisualizationManifest): + type: Literal[f"{TYPE}", "KeypointVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Keypoint Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-braille", + "blockPriority": 20, + }, + } + ) + + predictions: Selector( + kind=[ + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Keypoint detection predictions containing keypoint coordinates, confidence scores, and class names. Predictions must include keypoints_xy (keypoint coordinates), keypoints_confidence (confidence values), and keypoints_class_name (keypoint class/type names). Requires outputs from a keypoint detection model block.", + examples=["$steps.keypoint_detection_model.predictions"], + ) + + annotator_type: Literal["edge", "vertex", "vertex_label"] = Field( + description="Type of keypoint visualization mode. Options: 'edge' (draws connecting lines between keypoints using edge pairs, useful for skeleton/pose visualization), 'vertex' (draws circular markers at keypoint locations without connections), 'vertex_label' (draws circular markers with text labels identifying each keypoint class name).", + default="edge", + json_schema_extra={"always_visible": True}, + ) + + color: Union[str, Selector(kind=[STRING_KIND])] = Field( # type: ignore + description="Color of the keypoint markers, edges, or labels. Can be specified as a color name (e.g., 'green', 'red', 'blue'), hex color code (e.g., '#A351FB', '#FF0000'), or RGB format. Used for keypoint circles (vertex/vertex_label modes) or edge lines (edge mode).", + default="#A351FB", + examples=["#A351FB", "green", "$inputs.color"], + ) + + text_color: Union[str, Selector(kind=[STRING_KIND])] = Field( # type: ignore + description="Color of the text labels displayed on keypoints (vertex_label mode only). Can be specified as a color name (e.g., 'black', 'white'), hex color code, or RGB format. Only applies when annotator_type is 'vertex_label'.", + default="black", + examples=["black", "$inputs.text_color"], + json_schema_extra={ + "relevant_for": { + "annotator_type": { + "values": ["vertex_label"], + }, + }, + }, + ) + text_scale: Union[float, Selector(kind=[FLOAT_KIND])] = Field( # type: ignore + description="Scale factor for keypoint label text size (vertex_label mode only). Controls the size of text labels displayed on keypoints. Values greater than 1.0 make text larger, values less than 1.0 make text smaller. Only applies when annotator_type is 'vertex_label'. Typical values range from 0.3 to 1.0.", + default=0.5, + examples=[0.5, "$inputs.text_scale"], + json_schema_extra={ + "relevant_for": { + "annotator_type": { + "values": ["vertex_label"], + }, + }, + }, + ) + + text_thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the keypoint label text characters in pixels (vertex_label mode only). Controls how bold the text labels appear. Higher values create thicker, bolder text. Only applies when annotator_type is 'vertex_label'. Typical values range from 1 to 3.", + default=1, + examples=[1, "$inputs.text_thickness"], + json_schema_extra={ + "relevant_for": { + "annotator_type": { + "values": ["vertex_label"], + }, + }, + }, + ) + + text_padding: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Padding around keypoint label text in pixels (vertex_label mode only). Controls the spacing between the text label and its background border. Higher values create more space around text. Only applies when annotator_type is 'vertex_label'. Typical values range from 5 to 20 pixels.", + default=10, + examples=[10, "$inputs.text_padding"], + json_schema_extra={ + "relevant_for": { + "annotator_type": { + "values": ["vertex_label"], + }, + }, + }, + ) + + thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the edge lines connecting keypoints in pixels (edge mode only). Controls how thick the connecting lines between keypoints appear. Higher values create thicker, more visible edges. Only applies when annotator_type is 'edge'. Typical values range from 1 to 5 pixels.", + default=2, + examples=[2, "$inputs.thickness"], + json_schema_extra={ + "relevant_for": { + "annotator_type": { + "values": ["edge"], + }, + }, + }, + ) + + radius: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Radius of the circular keypoint markers in pixels (vertex and vertex_label modes only). Controls the size of circular markers drawn at keypoint locations. Higher values create larger, more visible markers. Only applies when annotator_type is 'vertex' or 'vertex_label'. Typical values range from 5 to 20 pixels.", + default=10, + examples=[10, "$inputs.radius"], + json_schema_extra={ + "relevant_for": { + "annotator_type": { + "values": ["vertex", "vertex_label"], + }, + }, + }, + ) + edges: Union[list, Selector(kind=[LIST_OF_VALUES_KIND])] = Field( # type: ignore + description="Edge connections between keypoints (edge mode only). List of pairs of keypoint indices (e.g., [(0, 1), (1, 2), ...]) defining which keypoints should be connected with lines. For pose estimation, this typically represents skeleton connections (e.g., connecting joints). Only applies when annotator_type is 'edge'. Required for edge visualization.", + default=None, + examples=["$inputs.edges"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.2.0,<2.0.0" + + +class KeypointVisualizationBlockV1(VisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return KeypointManifest + + def getAnnotator( + self, + color: str, + text_color: str, + text_scale: float, + text_thickness: int, + text_padding: int, + thickness: int, + radius: int, + annotator_type: str, + edges: List[Tuple[int, int]], + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color, + text_color, + text_scale, + text_thickness, + text_padding, + thickness, + radius, + annotator_type, + ], + ) + ) + + if key not in self.annotatorCache: + color = str_to_color(color) + text_color = str_to_color(text_color) + + if annotator_type == "edge": + self.annotatorCache[key] = sv.EdgeAnnotator( + color=color, + thickness=thickness, + edges=edges, + ) + elif annotator_type == "vertex": + self.annotatorCache[key] = sv.VertexAnnotator( + color=color, + radius=radius, + ) + elif annotator_type == "vertex_label": + self.annotatorCache[key] = sv.VertexLabelAnnotator( + color=color, + text_color=text_color, + text_scale=text_scale, + text_thickness=text_thickness, + text_padding=text_padding, + border_radius=radius, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Tuple[KeyPoints, Optional[Detections]], + copy_image: bool, + annotator_type: Optional[str], + color: Optional[str], + text_color: Optional[str], + text_scale: Optional[float], + text_thickness: Optional[int], + text_padding: Optional[int], + thickness: Optional[int], + radius: Optional[int], + edges: Optional[List[Tuple[int, int]]] = None, + ) -> BlockResult: + annotator: sv.EdgeAnnotator = self.getAnnotator( + color, + text_color, + text_scale, + text_thickness, + text_padding, + thickness, + radius, + annotator_type, + edges, + ) + + key_points = predictions[0] if isinstance(predictions, tuple) else predictions + keypoints = key_points.to_supervision() + + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + key_points=keypoints, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/label/v1.py b/inference/core/workflows/core_steps/visualizations/label/v1.py index b6ba3d6f21..391629475a 100644 --- a/inference/core/workflows/core_steps/visualizations/label/v1.py +++ b/inference/core/workflows/core_steps/visualizations/label/v1.py @@ -266,8 +266,13 @@ def run( border_radius, ) labels = build_detection_labels(predictions, text) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, labels=labels, ) diff --git a/inference/core/workflows/core_steps/visualizations/label/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/label/v1_tensor.py new file mode 100644 index 0000000000..8ad3896438 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/label/v1_tensor.py @@ -0,0 +1,1166 @@ +from collections import OrderedDict +from typing import List, Literal, Optional, Tuple, Type, Union + +import cv2 +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field +from supervision.annotators.utils import resolve_text_background_xyxy, wrap_text + +from inference.core.logger import logger +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + empty_predictions_passthrough, + to_supervision_for_annotation, +) +from inference.core.workflows.core_steps.visualizations.common.label_text import ( + build_detection_labels, +) +from inference.core.workflows.core_steps.visualizations.common.utils import str_to_color +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_KIND, + INTEGER_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/label_visualization@v1" +SHORT_DESCRIPTION = ( + "Draw labels on an image at specific coordinates based on provided detections." +) +LONG_DESCRIPTION = """ +Draw text labels on detected objects with customizable content, position, styling, and background colors to display information like class names, confidence scores, tracking IDs, or other detection metadata. + +## How This Block Works + +This block takes an image and detection predictions and draws text labels on each detected object. The block: + +1. Takes an image and predictions as input +2. Extracts label text for each detection based on the selected text option (class name, confidence, tracker ID, dimensions, area, time in zone, or index) +3. Determines label position based on the selected anchor point (center, corners, edges, or center of mass) +4. Applies background color styling based on the selected color palette, with colors assigned by class, index, or track ID +5. Renders text labels with customizable text color, scale, thickness, padding, and border radius using Supervision's LabelAnnotator +6. Returns an annotated image with text labels overlaid on the original image + +The block supports various text content options including class names, confidence scores, combination of class and confidence, tracker IDs (for tracked objects), time in zone (for zone analysis), object dimensions (center coordinates and width/height), area, or detection index. Labels are rendered with colored backgrounds that match the object's assigned color from the palette, and text styling (color, size, thickness) can be customized for optimal visibility. The labels can be positioned at any anchor point relative to each detection, allowing flexible placement for different visualization needs. + +## Common Use Cases + +- **Information Display on Detections**: Add informative text labels showing class names, confidence scores, or other metadata directly on detected objects for quick identification and validation +- **Model Performance Visualization**: Display confidence scores or class predictions on detected objects to visualize model certainty, identify low-confidence detections, and validate model performance +- **Object Tracking Visualization**: Show tracker IDs on tracked objects to visualize object tracking across frames, monitor persistent object identities, or debug tracking algorithms +- **Zone Analysis and Monitoring**: Display "Time In Zone" labels on objects to visualize how long objects have been in specific zones for occupancy monitoring, dwell time analysis, or compliance tracking +- **Spatial Information Display**: Show object dimensions (center coordinates, width, height) or area measurements directly on detections for spatial analysis, measurement workflows, or quality control +- **Professional Presentation and Reporting**: Create clean, informative visualizations with labeled detections for reports, dashboards, or presentations that combine visual results with textual information + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Bounding Box Visualization, Polygon Visualization, Dot Visualization) to combine text labels with geometric annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images with labels for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with labels to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with labels as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with labels for live monitoring, tracking visualization, or post-processing analysis +""" + +#: The font every supervision label draw call uses +#: (``supervision/annotators/core.py`` ``CV2_FONT``). +_CV2_FONT = cv2.FONT_HERSHEY_SIMPLEX + +#: ``supervision.annotators.utils.PENDING_TRACK_ID`` / ``PENDING_TRACK_COLOR``: +#: with ``ColorLookup.TRACK``, ``resolve_color`` (utils.py:139) returns the +#: pending gray for BOTH the background and the text color whenever the +#: resolved id is -1 โ€” before it ever consults the palette or the configured +#: text color. +_PENDING_TRACK_ID = -1 +_PENDING_TRACK_COLOR_BGR = (128, 128, 128) + +#: LRU bound on cached label sprites. The steady-state label vocabulary is +#: tiny (class names, `class 0.87`-style strings), but `Dimensions`-like texts +#: change every frame โ€” the bound turns a pathological stream into an LRU +#: churn (a cache miss costs one small CPU patch render + one small H2D) +#: instead of an unbounded device-memory leak. +_SPRITE_CACHE_MAX_ENTRIES = 512 + +#: Slots in the per-block pinned staging ring for the per-frame paste tables. +#: Must be >= 2x the pipeline batch size: on slot reuse the previous copy from +#: that slot was enqueued a whole batch (one model step) earlier, so its event +#: has fired long before the slot comes around again โ€” the reuse-wait below is +#: a pathological-backpressure guard, never the steady state. +_TABLE_RING_DEPTH = 8 + +#: Initial per-slot slab capacity (int64 elements). Covers 32 labels per frame +#: (2 table entries per label); grown geometrically on capacity misses. +_TABLE_SLAB_MIN_CAPACITY = 64 + + +def _staged_upload( + host_values: np.ndarray, + device: torch.device, + pending: Optional[List[Tuple[torch.Tensor, "torch.cuda.Event"]]] = None, +) -> torch.Tensor: + """One-shot hostโ†’device upload that never drains the device queue. + + A pageable ``torch.from_numpy(...).to(device)`` follows cudaMemcpy + semantics: the driver first waits for every kernel already queued on the + stream before the bytes move (measured 55-65 ms per first label upload of + a batch on Jetson, queued behind the mask-composite kernels). Staging + through pinned memory with ``copy_(non_blocking=True)`` enqueues the + transfer asynchronously instead. + + On CUDA the pinned staging tensor plus a recorded event is appended to + ``pending`` so the caller provably keeps the host memory intact until the + copy has executed (entries may be dropped only once ``event.query()`` is + True). Callers that pass no ``pending`` (the ring-less direct-call path of + ``gpu_paste_label_sprites``) fall back on torch's caching host allocator, + which event-guards a freed pinned block internally before reusing it โ€” + safe, just not explicit. + + On CPU devices there is nothing to pin and the copy is a plain memcpy, + but the identical ``copy_(..., non_blocking=True)`` is still dispatched + (a same-device ``.to`` would short-circuit without dispatching) so the op + trace carries the flag on every platform โ€” the viz-phase transfer audit + asserts it on the CPU test host too. + """ + source = torch.from_numpy(np.ascontiguousarray(host_values)) + if device.type == "cuda": + source = source.pin_memory() + staged = torch.empty(tuple(source.shape), dtype=source.dtype, device=device) + staged.copy_(source, non_blocking=True) + if device.type == "cuda" and pending is not None: + event = torch.cuda.Event() + event.record(torch.cuda.current_stream(device)) + pending.append((source, event)) + return staged + + +class _TableRingSlot: + __slots__ = ("host_slab", "host_view", "device_slab", "event", "capacity") + + def __init__(self): + self.host_slab: Optional[torch.Tensor] = None + self.host_view: Optional[np.ndarray] = None + self.device_slab: Optional[torch.Tensor] = None + self.event = None + self.capacity = 0 + + +class _PinnedSlabRing: + """Reusable pinned staging ring for the per-frame packed paste tables. + + Every warm-path ``gpu_paste_label_sprites`` call uploads one small int64 + table (the per-label paste offsets + lengths). Uploading it as a pageable + copy would synchronize the stream (see ``_staged_upload``); allocating and + pinning a fresh buffer per frame would thrash the pinned allocator. This + ring keeps ``depth`` pre-pinned host slabs, each paired with a same-size + device slab and a reuse event. ``upload_int64``: + + 1. takes the next slot round-robin; + 2. waits on the slot's recorded event ONLY when the copy issued on the + slot's previous use has not executed yet (``event.query()`` False) โ€” + with ``depth`` >= 2x the pipeline batch size that copy ran during an + earlier model step, so the wait is a pathological-backpressure guard, + never the steady state; + 3. grows the slab geometrically on a capacity miss (rare realloc, then + re-pin โ€” safe because step 2 just proved no copy is in flight from the + old slab); + 4. writes the table into the pinned host slab (a plain numpy write, no + torch dispatch) and enqueues ONE full-slab ``copy_(non_blocking=True)`` + into the device slab, then records the slot's event behind it. + + The returned device slab is safe for consumers enqueued later on the same + stream (stream order puts the copy before them); the recorded event only + guards the HOST slab against being overwritten while its copy is pending. + On CPU devices pinning and events are skipped entirely โ€” plain tensors and + the same single ``copy_`` (still flagged ``non_blocking=True``), so the + code runs identically on the CPU test host. + """ + + __slots__ = ("_depth", "_min_capacity", "_slots", "_cursor", "_initialised") + + def __init__( + self, + depth: int = _TABLE_RING_DEPTH, + min_capacity: int = _TABLE_SLAB_MIN_CAPACITY, + ): + self._depth = int(depth) + self._min_capacity = int(min_capacity) + self._slots = [_TableRingSlot() for _ in range(self._depth)] + self._cursor = 0 + self._initialised = False + + def _ensure_slot( + self, slot: _TableRingSlot, needed: int, device: torch.device + ) -> None: + if ( + slot.device_slab is not None + and slot.capacity >= needed + and slot.device_slab.device == device + ): + return + capacity = max(self._min_capacity, slot.capacity) + while capacity < needed: + capacity *= 2 + pin = device.type == "cuda" + slot.host_slab = torch.empty(capacity, dtype=torch.int64, pin_memory=pin) + slot.host_view = slot.host_slab.numpy() + slot.device_slab = torch.empty(capacity, dtype=torch.int64, device=device) + slot.capacity = capacity + slot.event = None + + def upload_int64(self, values: np.ndarray, device: torch.device) -> torch.Tensor: + """Stage 1-D int64 ``values`` and return the slot's device slab (its + first ``len(values)`` elements hold the table; callers slice it).""" + slot = self._slots[self._cursor] + self._cursor = (self._cursor + 1) % self._depth + if slot.event is not None and not slot.event.query(): + # Pathological backpressure only (the device is more than a whole + # ring โ€” >= 2 batches โ€” behind the host); steady state never waits. + slot.event.synchronize() + needed = int(values.shape[0]) + if not self._initialised: + # First touch sizes EVERY slot so steady-state frames (any ring + # phase) never dispatch an allocation. + for other in self._slots: + self._ensure_slot(other, needed, device) + self._initialised = True + else: + self._ensure_slot(slot, needed, device) + slot.host_view[:needed] = values + slot.device_slab.copy_(slot.host_slab, non_blocking=True) + if device.type == "cuda": + if slot.event is None: + slot.event = torch.cuda.Event() + slot.event.record(torch.cuda.current_stream(device)) + return slot.device_slab + + +class _SceneDependentLabelError(ValueError): + """Raised when a label patch cannot be pre-rendered scene-independently. + + ``sv.LabelAnnotator`` draws the anti-aliased text directly on the scene, so + any text ink that escapes the opaque label background (descenders when + ``text_padding`` is smaller than the font's baseline extent, glyphs leaking + into a rounded-off corner when ``border_radius`` is large relative to the + padding) is blended with the *scene* pixels underneath. A cached sprite + cannot reproduce that blend, so such configurations are refused here and + the block falls back to the bit-identical sv path instead of approximating. + """ + + +class _LabelMeasurement: + """The size math of ``LabelAnnotator._get_label_properties`` (supervision + 0.29.0, ``annotators/core.py:1301``) for one label: sv's own ``wrap_text`` + (v1 exposes no ``max_line_length``, so ``None`` โ€” newline splitting only), + per-line ``cv2.getTextSize`` with ``CV2_FONT``, ``max_width`` / + ``total_height`` aggregation and ``2 * text_padding`` padding on both + axes. ``margin`` bounds how far any ink (descenders below the last + baseline, stroke thickness, the 1-px AA fringe) can reach beyond the + background box, so a canvas with a ``margin`` border provably captures + every touched pixel.""" + + __slots__ = ("lines", "width_padded", "height_padded", "margin") + + def __init__(self, lines, width_padded, height_padded, margin): + self.lines = lines + self.width_padded = width_padded + self.height_padded = height_padded + self.margin = margin + + +def _measure_label( + label, text_scale: float, text_thickness: int, text_padding: int +) -> _LabelMeasurement: + lines = wrap_text(label, None) + line_heights: List[int] = [] + line_widths: List[int] = [] + baselines: List[int] = [] + for line in lines: + (line_w, line_h), baseline = cv2.getTextSize( + line, _CV2_FONT, text_scale, text_thickness + ) + line_widths.append(line_w) + line_heights.append(line_h) + baselines.append(baseline) + max_width = max(line_widths) if line_widths else 0 + total_height = sum(line_heights) + (len(line_heights) - 1) * text_padding + width_padded = max_width + 2 * text_padding + height_padded = total_height + 2 * text_padding + tg_baseline = cv2.getTextSize("Tg", _CV2_FONT, text_scale, text_thickness)[1] + margin = max(baselines + [tg_baseline]) + int(text_thickness) + 3 + return _LabelMeasurement(lines, width_padded, height_padded, margin) + + +class _LabelSprite: + """A single label patch pre-rendered with sv's exact draw calls. + + The patch is rendered once on CPU (see ``_render_label_sprite``) and kept + as a sparse pixel list: host-side ``(rows, cols)`` canvas coordinates of + the opaque (scene-independent) pixels plus their RGB values as a device + tensor. Per frame the only work is offsetting cached flat-index templates + โ€” no re-render, no per-frame H2D of pixel payloads. + + ``(offset_x, offset_y)`` is where the sv background-box origin sits inside + the sprite canvas, so pasting the canvas origin at ``(box_x1 - offset_x, + box_y1 - offset_y)`` reproduces sv's drawing at ``box_xyxy`` exactly. For + an interior sprite the offsets both equal the measurement margin; for a + frame-clipped variant the canvas edges coincide with the frame edges on + the clipped sides (cv2's anti-aliased rasterisation treats strokes cut by + the canvas border differently from uncut ones, so the border must sit + exactly where sv's frame border sits). + """ + + __slots__ = ( + "offset_x", + "offset_y", + "rows", + "cols", + "row_min", + "row_max", + "col_min", + "col_max", + "colors_dev", + "_flat_by_width", + "_pending", + ) + + def __init__( + self, + offset_x: int, + offset_y: int, + rows: np.ndarray, + cols: np.ndarray, + colors_dev: torch.Tensor, + ): + self.offset_x = offset_x + self.offset_y = offset_y + self.rows = rows # (K,) int32, canvas-relative + self.cols = cols # (K,) int32, canvas-relative + self.row_min = int(rows.min()) if rows.size else 0 + self.row_max = int(rows.max()) if rows.size else -1 + self.col_min = int(cols.min()) if cols.size else 0 + self.col_max = int(cols.max()) if cols.size else -1 + self.colors_dev = colors_dev # (K, 3) uint8 RGB on the pipeline device + self._flat_by_width: dict = {} + # (pinned staging tensor, recorded event) pairs for this sprite's + # in-flight non-blocking uploads: the pinned source provably outlives + # the async copy โ€” dropped only once the event has fired. + self._pending: List = [] + + @property + def pixel_count(self) -> int: + return int(self.rows.shape[0]) + + def _release_completed_staging(self) -> None: + """Drop pinned staging buffers whose uploads provably executed (their + recorded events query True). Never blocks; events recorded on one + stream fire in order, so popping from the front suffices.""" + while self._pending and self._pending[0][1].query(): + self._pending.pop(0) + + def flat_template(self, frame_width: int) -> torch.Tensor: + """Device-resident ``rows * frame_width + cols`` template, cached per + frame width (streams keep a constant width, so this is a one-time + staged non-blocking upload โ€” a pageable ``.to(device)`` here would + drain the queued viz kernels). Callers must add the paste offset out + of place โ€” the cached tensor is shared across frames.""" + self._release_completed_staging() + flat = self._flat_by_width.get(frame_width) + if flat is None: + flat_np = self.rows.astype(np.int64) * frame_width + self.cols.astype( + np.int64 + ) + flat = _staged_upload(flat_np, self.colors_dev.device, self._pending) + self._flat_by_width[frame_width] = flat + return flat + + +def _draw_label_patch( + canvas: np.ndarray, + box_xyxy: Tuple[int, int, int, int], + lines: List[str], + text_color_bgr: Tuple[int, int, int], + background_color_bgr: Tuple[int, int, int], + text_scale: float, + text_thickness: int, + text_padding: int, + border_radius: int, +) -> None: + """The exact per-label draw sequence of ``sv.LabelAnnotator`` (supervision + 0.29.0, ``annotators/core.py``): ``draw_rounded_rectangle`` (line 1437 โ€” + two filled rectangles + four filled circles, radius clipped to + ``min(width, height) // 2``) followed by the multiline ``cv2.putText`` + loop of ``_draw_labels`` (line 1355 โ€” empty lines advance by the height of + ``"Tg"`` without drawing). cv2 rasterisation is invariant under integer + translation, so drawing at the canvas-local box and pasting equals drawing + at the scene-global box.""" + x1, y1, x2, y2 = box_xyxy + width = x2 - x1 + height = y2 - y1 + radius = min(border_radius, min(width, height) // 2) + cv2.rectangle( + canvas, (x1 + radius, y1), (x2 - radius, y2), background_color_bgr, -1 + ) + cv2.rectangle( + canvas, (x1, y1 + radius), (x2, y2 - radius), background_color_bgr, -1 + ) + for center in ( + (x1 + radius, y1 + radius), + (x2 - radius, y1 + radius), + (x1 + radius, y2 - radius), + (x2 - radius, y2 - radius), + ): + cv2.circle(canvas, center, radius, background_color_bgr, -1) + current_y = y1 + text_padding + for line in lines: + # sv measures "Tg" (a string with ascender + descender) to advance + # over empty lines, and the line itself otherwise. + text_h = cv2.getTextSize( + line if line else "Tg", _CV2_FONT, text_scale, text_thickness + )[0][1] + if line: + cv2.putText( + img=canvas, + text=line, + org=(x1 + text_padding, current_y + text_h), + fontFace=_CV2_FONT, + fontScale=text_scale, + color=text_color_bgr, + thickness=text_thickness, + lineType=cv2.LINE_AA, + ) + current_y += text_h + text_padding + + +def _render_label_sprite( + measurement: _LabelMeasurement, + text_color_bgr: Tuple[int, int, int], + background_color_bgr: Tuple[int, int, int], + text_scale: float, + text_thickness: int, + text_padding: int, + border_radius: int, + device: torch.device, + box_in_canvas: Tuple[int, int, int, int], + canvas_hw: Tuple[int, int], + frame_edge_sides: Tuple[bool, bool, bool, bool], +) -> _LabelSprite: + """Render one label patch with sv's exact draw calls and extract its + scene-independent pixels. + + The patch is drawn twice with identical calls โ€” once over a black canvas, + once over a white one. Pixels where the two renders agree are fully + determined by the drawing (opaque); pixels where they differ depend on + what is underneath: either never touched (transparent โ€” sv leaves the + scene) or touched by anti-aliased text ink *outside* the opaque background + (sv blends that ink with the scene โ€” not representable by a cached sprite, + so ``_SceneDependentLabelError`` sends the block to the sv path). + + ``box_in_canvas`` places the sv background box inside the canvas. On sides + marked in ``frame_edge_sides`` (left, top, right, bottom) the caller + aligned the canvas edge with the *frame* edge, so cv2 clips exactly where + sv would โ€” cv2's AA rasterisation of a stroke cut by the canvas border + differs from an uncut stroke, so an edge-crossing label must be rendered + against the identical border. The remaining sides carry a ``margin``-wide + border that provably contains all reachable ink; if ink still lands on + such a border ring the sprite is refused rather than silently clipped + where sv would have drawn. + """ + canvas_h, canvas_w = canvas_hw + on_black = np.zeros((canvas_h, canvas_w, 3), dtype=np.uint8) + on_white = np.full((canvas_h, canvas_w, 3), 255, dtype=np.uint8) + for canvas in (on_black, on_white): + _draw_label_patch( + canvas=canvas, + box_xyxy=box_in_canvas, + lines=measurement.lines, + text_color_bgr=text_color_bgr, + background_color_bgr=background_color_bgr, + text_scale=text_scale, + text_thickness=text_thickness, + text_padding=text_padding, + border_radius=border_radius, + ) + opaque = (on_black == on_white).all(axis=2) + touched = opaque | (on_black != 0).any(axis=2) | (on_white != 255).any(axis=2) + left_edge, top_edge, right_edge, bottom_edge = frame_edge_sides + if ( + (not top_edge and touched[0].any()) + or (not bottom_edge and touched[-1].any()) + or (not left_edge and touched[:, 0].any()) + or (not right_edge and touched[:, -1].any()) + ): + raise _SceneDependentLabelError( + "label ink reaches the sprite canvas border on a side that is not " + "a frame edge" + ) + if (touched & ~opaque).any(): + raise _SceneDependentLabelError( + "anti-aliased label text escapes the opaque background patch " + "(text_padding too small for the font extent, or border_radius " + "cutting under the text); sv blends such ink with the scene, " + "which a cached sprite cannot reproduce" + ) + rows, cols = np.nonzero(opaque) + colors_rgb = np.ascontiguousarray(on_black[rows, cols][:, ::-1]) # BGR -> RGB + # Cache-miss pixel payload: staged through pinned memory + a non-blocking + # copy (a pageable upload would synchronize the stream and drain queued + # viz kernels โ€” measured 55-62 ms per miss on Jetson under load); the + # staging buffer is parked on the sprite until its event fires. + pending: List = [] + colors_dev = _staged_upload(colors_rgb, device, pending) + sprite = _LabelSprite( + offset_x=int(box_in_canvas[0]), + offset_y=int(box_in_canvas[1]), + rows=rows.astype(np.int32), + cols=cols.astype(np.int32), + colors_dev=colors_dev, + ) + sprite._pending.extend(pending) + return sprite + + +def gpu_paste_label_sprites( + scene_chw: torch.Tensor, + sprites: List[_LabelSprite], + canvas_origins: List[Tuple[int, int]], + table_ring: Optional[_PinnedSlabRing] = None, +) -> torch.Tensor: + """Composite pre-rendered label sprites into a CHW RGB uint8 tensor, in + place, with a device-op budget that is flat in the label count. + + A per-label paste loop is dispatch-bound on Jetson (the bbox painter in + this package measured 43 ms @ 50 boxes for a per-box loop), so all labels + land in ONE indexed store: + + 1. Host: per label, the cached flat-index template (device-resident, see + ``_LabelSprite.flat_template``) is selected and its scalar paste offset + computed from the sprite's canvas origin โ€” the cached color tensor is + reused untouched, so a cache-hit frame does no per-pixel host work and + no pixel-payload H2D at all. + 2. One packed upload of the per-label (offset, length) table, staged + through ``table_ring`` (the block's pinned slab ring) as a single + NON-BLOCKING copy โ€” a pageable upload here follows cudaMemcpy + semantics and drains every queued kernel first (measured 62-65 ms per + first label of a batch on Jetson behind the mask compositor). Then one + ``repeat_interleave`` expansion, one ``cat`` + offset add โ€” flat pixel + indices for every sprite pixel of every label. Ring-less direct calls + stage through a one-shot pinned buffer instead + (``_staged_upload``). + 3. sv draws labels sequentially, so a later label's patch overwrites an + earlier one. When the pasted rectangles overlap, that order is + reproduced deterministically with a ``scatter_reduce_(amax)`` of the + global pixel position (ascending in label order) followed by a winner + gather โ€” the same later-wins resolution the bbox painter uses. + Disjoint labels (the common case) skip it. + 4. One final indexed store. This is the ONLY write to ``scene_chw``: any + failure before it leaves the scene untouched, so the block's sv + fallback can never double-draw a partially annotated frame (and the + paste is an opaque overwrite, hence idempotent, unlike a blend). + + No ``.contiguous()`` is taken: ``.view`` raises on a non-contiguous scene + (routing to the sv fallback) rather than silently writing into a copy โ€” + ``copy_image=False`` must mutate the caller's storage. + + Every sprite pixel must land inside the frame at its origin (the caller + picked an interior sprite whose pixel bounds fit, or a frame-clipped + variant whose canvas lies inside the frame); a violation raises before + anything is written. + """ + device = scene_chw.device + height, width = int(scene_chw.shape[1]), int(scene_chw.shape[2]) + flat_parts: List[torch.Tensor] = [] + color_parts: List[torch.Tensor] = [] + piece_offsets: List[int] = [] + piece_lengths: List[int] = [] + piece_rects: List[Tuple[int, int, int, int]] = [] + for sprite, (origin_x, origin_y) in zip(sprites, canvas_origins): + if sprite.pixel_count == 0: + continue + if ( + origin_x + sprite.col_min < 0 + or origin_y + sprite.row_min < 0 + or origin_x + sprite.col_max >= width + or origin_y + sprite.row_max >= height + ): + raise ValueError("label sprite pixels fall outside the frame") + flat_parts.append(sprite.flat_template(width)) + color_parts.append(sprite.colors_dev) + piece_offsets.append(origin_y * width + origin_x) + piece_lengths.append(sprite.pixel_count) + piece_rects.append( + ( + origin_x + sprite.col_min, + origin_y + sprite.row_min, + origin_x + sprite.col_max + 1, + origin_y + sprite.row_max + 1, + ) + ) + if not flat_parts: + return scene_chw + pieces = len(flat_parts) + total = int(sum(piece_lengths)) + packed_host = np.empty(2 * pieces, dtype=np.int64) + packed_host[:pieces] = piece_offsets + packed_host[pieces:] = piece_lengths + if table_ring is not None: + packed = table_ring.upload_int64(packed_host, device) + else: + packed = _staged_upload(packed_host, device) + offsets_t, lengths_t = packed[:pieces], packed[pieces : 2 * pieces] + piece_of_px = torch.repeat_interleave( + torch.arange(pieces, device=device), lengths_t, output_size=total + ) + flat = torch.cat(flat_parts) if pieces > 1 else flat_parts[0] + flat = flat + offsets_t[piece_of_px] + colors = torch.cat(color_parts) if pieces > 1 else color_parts[0] + # Pairwise overlap test of the pasted (frame-clipped) canvas rectangles: + # disjoint labels make every pixel its own winner, so owner resolution is + # skipped. Conservative (canvas margins are mostly transparent), which + # only costs the extra pass, never correctness. + rect = np.asarray(piece_rects, dtype=np.int64) + inter_x = np.maximum(rect[:, 0][:, None], rect[:, 0][None, :]) < np.minimum( + rect[:, 2][:, None], rect[:, 2][None, :] + ) + inter_y = np.maximum(rect[:, 1][:, None], rect[:, 1][None, :]) < np.minimum( + rect[:, 3][:, None], rect[:, 3][None, :] + ) + labels_overlap = int((inter_x & inter_y).sum()) > pieces # diagonal always True + if labels_overlap: + # include_self=False: uninitialized cells never participate, and every + # gathered position below was scattered to. + order = torch.arange(total, device=device, dtype=torch.int32) + owner = torch.empty(height * width, dtype=torch.int32, device=device) + owner.scatter_reduce_(0, flat, order, reduce="amax", include_self=False) + colors = colors[owner[flat].long()] + scene_chw.view(3, -1)[:, flat] = colors.t() + return scene_chw + + +def _gpu_label_paste_eligible( + detections, color_axis: str, image: WorkflowImageData +) -> bool: + """True when the sprite compositor can replace the sv path.""" + if color_axis not in ("CLASS", "INDEX", "TRACK"): + # Custom lookups keep the battle-tested sv path. + return False + if not image.is_tensor_materialised(): + # Numpy-sourced images must behave EXACTLY as before: forcing + # tensor_image would be a costly host-side conversion and the sv path + # is faster there. + return False + xyxy = getattr(detections, "xyxy", None) + if not isinstance(xyxy, torch.Tensor) or int(xyxy.shape[0]) == 0: + # Nothing to draw; the sv path is a trivial no-op. + return False + return True + + +def _resolve_color_ids_for_labels( + detections: sv.Detections, color_axis: str +) -> np.ndarray: + """The palette indices sv's ``resolve_color_idx`` (supervision 0.29.0, + ``annotators/utils.py:40``) would use on this materialised view, raising + its exact ``ValueError``s when ids are missing (the sv fallback would then + raise the very same error from the annotator).""" + n = len(detections) + if color_axis == "INDEX": + return np.arange(n) + if color_axis == "CLASS": + if detections.class_id is None: + raise ValueError( + "Could not resolve color by class because " + "Detections do not have class_id. If using an annotator, " + "try setting color_lookup to sv.ColorLookup.INDEX or " + "sv.ColorLookup.TRACK." + ) + return detections.class_id.astype(int) + if detections.tracker_id is None: + raise ValueError( + "Could not resolve color by track because " + "Detections do not have tracker_id. Did you call " + "tracker.update_with_detections(...) before annotating?" + ) + return detections.tracker_id.astype(int) + + +class LabelManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "LabelVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Label Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-tag", + "blockPriority": 2, + "popular": True, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + text: Union[ + Literal[ + "Class", + "Confidence", + "Class and Confidence", + "Index", + "Dimensions", + "Area", + "Area (mask)", + "Area (converted)", + "Tracker Id", + "Time In Zone", + ], + Selector(kind=[STRING_KIND]), + ] = Field( # type: ignore + default="Class", + description="Content to display in text labels. Options: 'Class' (class name), 'Confidence' (confidence score), 'Class and Confidence' (both), 'Tracker Id' (tracking ID for tracked objects), 'Time In Zone' (time spent in zone), 'Dimensions' (center coordinates and width x height), 'Area' (bounding box area in pixels), 'Area (mask)' (mask area in pixels from Mask Area Measurement block), 'Area (converted)' (mask area in converted units from Mask Area Measurement block), or 'Index' (detection index).", + examples=["LABEL", "$inputs.text"], + json_schema_extra={ + "always_visible": True, + }, + ) + + text_position: Union[ + Literal[ + "CENTER", + "CENTER_LEFT", + "CENTER_RIGHT", + "TOP_CENTER", + "TOP_LEFT", + "TOP_RIGHT", + "BOTTOM_LEFT", + "BOTTOM_CENTER", + "BOTTOM_RIGHT", + "CENTER_OF_MASS", + ], + Selector(kind=[STRING_KIND]), + ] = Field( # type: ignore + default="TOP_LEFT", + description="Anchor position for placing labels relative to each detection's bounding box. Options include: CENTER (center of box), corners (TOP_LEFT, TOP_RIGHT, BOTTOM_LEFT, BOTTOM_RIGHT), edge midpoints (TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, BOTTOM_CENTER), or CENTER_OF_MASS (center of mass of the object).", + examples=["CENTER", "$inputs.text_position"], + ) + + text_color: Union[str, Selector(kind=[STRING_KIND])] = Field( # type: ignore + description="Color of the label text. Can be a color name (e.g., 'WHITE', 'BLACK') or color code in HEX format (e.g., '#FFFFFF') or RGB format (e.g., 'rgb(255, 255, 255)').", + default="WHITE", + examples=["WHITE", "#FFFFFF", "rgb(255, 255, 255)" "$inputs.text_color"], + ) + + text_scale: Union[float, Selector(kind=[FLOAT_KIND])] = Field( # type: ignore + description="Scale factor for text size. Higher values create larger text. Default is 1.0.", + default=1.0, + examples=[1.0, "$inputs.text_scale"], + ) + + text_thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of text characters in pixels. Higher values create bolder, thicker text for better visibility.", + default=1, + examples=[1, "$inputs.text_thickness"], + ) + + text_padding: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Padding around the text in pixels. Controls the spacing between the text and the label background border.", + default=10, + examples=[10, "$inputs.text_padding"], + ) + + border_radius: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Border radius of the label background in pixels. Set to 0 for square corners. Higher values create more rounded corners for a softer appearance.", + default=0, + examples=[0, "$inputs.border_radius"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class LabelVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + self._sprite_cache: "OrderedDict" = OrderedDict() + # Pinned staging ring for the per-frame packed paste tables โ€” the + # steady-state path allocates/pins nothing per frame and its single + # upload never blocks the stream (see _PinnedSlabRing). + self._table_ring = _PinnedSlabRing() + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return LabelManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + text_position: str, + text_color: str, + text_scale: float, + text_thickness: int, + text_padding: int, + border_radius: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + text_position, + text_color, + text_scale, + text_thickness, + text_padding, + border_radius, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + text_color = str_to_color(text_color) + + self.annotatorCache[key] = sv.LabelAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + text_position=getattr(sv.Position, text_position), + text_color=text_color, + text_scale=text_scale, + text_thickness=text_thickness, + text_padding=text_padding, + border_radius=border_radius, + ) + + return self.annotatorCache[key] + + def _get_sprite( + self, + label, + measurement: _LabelMeasurement, + text_color_bgr: Tuple[int, int, int], + background_color_bgr: Tuple[int, int, int], + text_scale: float, + text_thickness: int, + text_padding: int, + border_radius: int, + device: torch.device, + box_in_canvas: Tuple[int, int, int, int], + canvas_hw: Tuple[int, int], + frame_edge_sides: Tuple[bool, bool, bool, bool], + ) -> _LabelSprite: + """LRU-cached sprite lookup. The key covers everything that affects + the patch's pixels: the label + typography + colors, plus the canvas + geometry (interior sprites share one canonical geometry; frame-clipped + variants are keyed by their clip window, which changes only when the + label's overhang does). The raw label object is used as the key + component โ€” ``wrap_text`` special-cases falsy labels before + stringifying, so pre-stringifying here could change sv's rendering; + str-like labels hash consistently across ``str``/``np.str_``. In + practice the label vocabulary is tiny (class names x rounded + confidences), so steady-state hit rate is ~100% and no H2D happens per + frame.""" + key = ( + label, + text_color_bgr, + background_color_bgr, + float(text_scale), + int(text_thickness), + int(text_padding), + int(border_radius), + str(device), + box_in_canvas, + canvas_hw, + frame_edge_sides, + ) + sprite = self._sprite_cache.get(key) + if sprite is not None: + self._sprite_cache.move_to_end(key) + return sprite + sprite = _render_label_sprite( + measurement=measurement, + text_color_bgr=text_color_bgr, + background_color_bgr=background_color_bgr, + text_scale=float(text_scale), + text_thickness=int(text_thickness), + text_padding=int(text_padding), + border_radius=int(border_radius), + device=device, + box_in_canvas=box_in_canvas, + canvas_hw=canvas_hw, + frame_edge_sides=frame_edge_sides, + ) + self._sprite_cache[key] = sprite + if len(self._sprite_cache) > _SPRITE_CACHE_MAX_ENTRIES: + self._sprite_cache.popitem(last=False) + return sprite + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + text: Optional[str], + text_position: Optional[str], + text_color: Optional[str], + text_scale: Optional[float], + text_thickness: Optional[int], + text_padding: Optional[int], + border_radius: Optional[int], + ) -> BlockResult: + detections = ( + split_key_point_prediction(predictions)[1] + if isinstance(predictions, tuple) + else predictions + ) + passthrough = empty_predictions_passthrough( + image=image, detections=detections, copy_image=copy_image + ) + if passthrough is not None: + return passthrough + # The Label annotator reads `.mask` for exactly two configurations, and + # only for instance-segmentation input (there is no mask to materialise + # otherwise, so `materialise_masks=True` is a no-op for OD input): + # * text == "Area": `sv.Detections.area` returns MASK area when a mask + # is present and BOX area when it is None โ€” flag-off shows mask area + # on IS input, so the mask must be materialised to match. + # * text_position == "CENTER_OF_MASS": `sv.LabelAnnotator` anchors on + # the mask centroid; `get_anchors_coordinates` RAISES without a mask. + # Every other label reads xyxy / confidence / per-box metadata, so the + # device->host dense-mask copy is skipped for them. Both configurations + # also keep the sv path below (the sprite compositor never materialises + # masks). + needs_masks = text == "Area" or text_position == "CENTER_OF_MASS" + if not needs_masks and _gpu_label_paste_eligible(detections, color_axis, image): + # GPU sprite path: labels are rendered once on CPU with sv's exact + # draw calls, cached as device sprites, and pasted with a fixed + # number of torch ops โ€” the full-frame device->host materialisation + # the sv path pays (~30 ms/2K frame on Orin NX) never happens. + try: + palette = self.getPalette(color_palette, palette_size, custom_colors) + if not isinstance(palette, sv.ColorPalette): + raise TypeError("expected sv.ColorPalette") + scene_t = image.tensor_image + if int(scene_t.shape[0]) != 3: + raise ValueError("GPU label compositor requires a 3-channel image") + # Mask-free sv view: one tiny xyxy/class/confidence D2H plus + # per-box metadata โ€” labels never read masks on this path. + sv_view = to_supervision_for_annotation( + predictions, materialise_masks=False + ) + labels = build_detection_labels(sv_view, text) + if len(labels) != len(sv_view): + # sv's _validate_labels contract. + raise ValueError( + f"The number of labels ({len(labels)}) does not match " + f"the number of detections ({len(sv_view)}). Each " + "detection should have exactly 1 label." + ) + text_color_bgr = str_to_color(text_color).as_bgr() + ids = _resolve_color_ids_for_labels(sv_view, color_axis) + # Same anchor math as LabelAnnotator._get_label_properties + # (supervision 0.29.0, annotators/core.py:1301): sv's own + # get_anchors_coordinates on the same materialised view, + # truncated to int; the label box comes from sv's + # resolve_text_background_xyxy with the float32 round-trip the + # annotator applies to it. + anchors = sv_view.get_anchors_coordinates( + anchor=getattr(sv.Position, text_position) + ).astype(int) + frame_h = int(scene_t.shape[1]) + frame_w = int(scene_t.shape[2]) + sprites: List[_LabelSprite] = [] + canvas_origins: List[Tuple[int, int]] = [] + for idx, label in enumerate(labels): + if color_axis == "TRACK" and int(ids[idx]) == _PENDING_TRACK_ID: + # sv resolve_color returns the pending-track gray for + # BOTH background and text color (utils.py:139). + background_bgr = _PENDING_TRACK_COLOR_BGR + line_color_bgr = _PENDING_TRACK_COLOR_BGR + else: + background_bgr = palette.by_idx(int(ids[idx])).as_bgr() + line_color_bgr = text_color_bgr + measurement = _measure_label( + label, + float(text_scale), + int(text_thickness), + int(text_padding), + ) + background_xyxy = resolve_text_background_xyxy( + center_coordinates=(int(anchors[idx][0]), int(anchors[idx][1])), + text_wh=( + measurement.width_padded, + measurement.height_padded, + ), + position=getattr(sv.Position, text_position), + ) + background_xyxy = np.asarray( + background_xyxy, dtype=np.float32 + ).astype(int) + bx1, by1, bx2, by2 = (int(value) for value in background_xyxy) + margin = measurement.margin + sprite = self._get_sprite( + label=label, + measurement=measurement, + text_color_bgr=line_color_bgr, + background_color_bgr=background_bgr, + text_scale=text_scale, + text_thickness=text_thickness, + text_padding=text_padding, + border_radius=border_radius, + device=scene_t.device, + box_in_canvas=( + margin, + margin, + margin + measurement.width_padded, + margin + measurement.height_padded, + ), + canvas_hw=( + measurement.height_padded + 1 + 2 * margin, + measurement.width_padded + 1 + 2 * margin, + ), + frame_edge_sides=(False, False, False, False), + ) + origin_x = bx1 - sprite.offset_x + origin_y = by1 - sprite.offset_y + if not ( + origin_x + sprite.col_min >= 0 + and origin_y + sprite.row_min >= 0 + and origin_x + sprite.col_max < frame_w + and origin_y + sprite.row_max < frame_h + ): + # The label crosses the frame boundary: cv2 clips at + # the frame edge and its AA rasterisation of a clipped + # stroke differs from an unclipped one, so a variant + # is rendered in a canvas whose edges coincide with + # the frame edges on the crossing sides (cached per + # clip window). + window_x1 = max(bx1 - margin, 0) + window_y1 = max(by1 - margin, 0) + window_x2 = min(bx2 + 1 + margin, frame_w) + window_y2 = min(by2 + 1 + margin, frame_h) + if window_x2 <= window_x1 or window_y2 <= window_y1: + # Entirely outside the frame: sv draws nothing + # visible. + continue + sprite = self._get_sprite( + label=label, + measurement=measurement, + text_color_bgr=line_color_bgr, + background_color_bgr=background_bgr, + text_scale=text_scale, + text_thickness=text_thickness, + text_padding=text_padding, + border_radius=border_radius, + device=scene_t.device, + box_in_canvas=( + bx1 - window_x1, + by1 - window_y1, + bx2 - window_x1, + by2 - window_y1, + ), + canvas_hw=( + window_y2 - window_y1, + window_x2 - window_x1, + ), + frame_edge_sides=( + bx1 - margin < 0, + by1 - margin < 0, + bx2 + 1 + margin > frame_w, + by2 + 1 + margin > frame_h, + ), + ) + origin_x = bx1 - sprite.offset_x + origin_y = by1 - sprite.offset_y + sprites.append(sprite) + canvas_origins.append((origin_x, origin_y)) + # All validation and sprite rendering succeeded โ€” only now may + # the scene be touched (clone, or the single in-place store + # inside the paste). + if copy_image: + scene_t = scene_t.clone() + annotated_tensor = gpu_paste_label_sprites( + scene_t, sprites, canvas_origins, table_ring=self._table_ring + ) + if not copy_image: + # The paste mutated `image.tensor_image` storage in place + # (the sv-path contract for copy_image=False); invalidate + # the derived numpy/base64 caches. + image.declare_tensor_image_mutated() + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, tensor_image=annotated_tensor + ) + } + except Exception as gpu_error: + logger.debug( + "GPU label compositor failed (%s); falling back to " + "sv.LabelAnnotator path.", + gpu_error, + ) + predictions = to_supervision_for_annotation( + predictions, materialise_masks=needs_masks + ) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + text_position, + text_color, + text_scale, + text_thickness, + text_padding, + border_radius, + ) + labels = build_detection_labels(predictions, text) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + labels=labels, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/label/v2_tensor.py b/inference/core/workflows/core_steps/visualizations/label/v2_tensor.py new file mode 100644 index 0000000000..b4fa6d20fa --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/label/v2_tensor.py @@ -0,0 +1,148 @@ +from typing import List, Optional, Type, Union + +from pydantic import Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + empty_predictions_passthrough, + to_supervision_for_annotation, +) +from inference.core.workflows.core_steps.visualizations.common.label_text import ( + TEXT_SIZE_MODE_MANUAL, + build_detection_labels, + compute_adaptive_label_text_scale, +) +from inference.core.workflows.core_steps.visualizations.label.v2 import TYPE +from inference.core.workflows.core_steps.visualizations.label.v2 import ( + LabelManifestV2 as _NumpyLabelManifestV2, +) +from inference.core.workflows.core_steps.visualizations.label.v2 import ( + LabelVisualizationBlockV2 as _NumpyLabelVisualizationBlockV2, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import Selector +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + + +class LabelManifestV2(_NumpyLabelManifestV2): + """The numpy manifest reused verbatim (same ``type`` literal, fields and + I/O contract) - only the ``predictions`` selector is re-declared with the + tensor-native kinds, mirroring ``base_tensor.PredictionsVisualizationManifest``.""" + + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Model predictions to visualize.", + examples=["$steps.object_detection_model.predictions"], + ) + + +class LabelVisualizationBlockV2(_NumpyLabelVisualizationBlockV2): + """Tensor-native sibling of Label Visualization v2. + + All drawing internals are reused from the numpy block: ``getAnnotator`` + (annotator construction + caching) is inherited, label text comes from the + shared ``label_text.build_detection_labels`` and adaptive sizing from + ``label_text.compute_adaptive_label_text_scale``. This class only adds the + tensor-side glue: device-resident empty passthrough and native->sv + materialisation before the inherently-CPU cv2 text rasterisation. + """ + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return LabelManifestV2 + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + text: Optional[str], + text_position: Optional[str], + text_color: Optional[str], + text_size_mode: Optional[str], + text_scale: Optional[float], + text_thickness: Optional[int], + text_padding: Optional[int], + border_radius: Optional[int], + ) -> BlockResult: + detections = ( + split_key_point_prediction(predictions)[1] + if isinstance(predictions, tuple) + else predictions + ) + passthrough = empty_predictions_passthrough( + image=image, detections=detections, copy_image=copy_image + ) + if passthrough is not None: + return passthrough + # `.mask` is read for exactly two configurations, and only for + # instance-segmentation input (see label/v1_tensor.py for details): + # * text == "Area": `sv.Detections.area` reports MASK area when a + # mask is present and BOX area when it is None - flag-off shows + # mask area on IS input, so the mask must be materialised to match; + # * text_position == "CENTER_OF_MASS": the annotator anchors on the + # mask centroid; `get_anchors_coordinates` RAISES without a mask. + # Every other label reads xyxy / confidence / per-box metadata, so the + # device->host dense-mask copy is skipped for them. + needs_masks = text == "Area" or text_position == "CENTER_OF_MASS" + sv_detections = to_supervision_for_annotation( + predictions, materialise_masks=needs_masks + ) + # Non-empty frame: label rendering is inherently CPU work (cv2 text + # rasterisation), so the numpy image is materialised here. + height, width = image.numpy_image.shape[:2] + effective_text_scale = compute_adaptive_label_text_scale( + height, + width, + manual_text_scale=text_scale, + text_size_mode=text_size_mode or TEXT_SIZE_MODE_MANUAL, + ) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + text_position, + text_color, + effective_text_scale, + text_thickness, + text_padding, + border_radius, + ) + labels = build_detection_labels(sv_detections, text) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=sv_detections, + labels=labels, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/mask/v1.py b/inference/core/workflows/core_steps/visualizations/mask/v1.py index f550e3f804..0efc319b21 100644 --- a/inference/core/workflows/core_steps/visualizations/mask/v1.py +++ b/inference/core/workflows/core_steps/visualizations/mask/v1.py @@ -186,8 +186,13 @@ def run( ] ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/mask/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/mask/v1_tensor.py new file mode 100644 index 0000000000..62c1aa0cc8 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/mask/v1_tensor.py @@ -0,0 +1,639 @@ +from typing import List, Literal, Optional, Type, Union + +import numpy as np +import supervision as sv +import torch +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + empty_predictions_passthrough, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + SEMANTIC_SEGMENTATION_PREDICTION_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks + +TYPE: str = "roboflow_core/mask_visualization@v1" +SHORT_DESCRIPTION = "Apply a mask over detected objects in an image." +LONG_DESCRIPTION = """ +Fill segmentation masks with semi-transparent color overlays, creating solid color fills that precisely follow the shape of detected objects from instance segmentation predictions. + +## How This Block Works + +This block takes an image and instance segmentation predictions (with masks) and fills the mask regions with colored overlays. The block: + +1. Takes an image and instance segmentation predictions (with masks) as input +2. Extracts segmentation masks for each detected object from the predictions +3. Applies color styling to each mask based on the selected color palette, with colors assigned by class, index, or track ID +4. Fills the mask regions with solid colors using Supervision's MaskAnnotator +5. Blends the colored mask overlays with the original image using the specified opacity level +6. Returns an annotated image where mask regions are filled with semi-transparent colors, while non-masked areas remain unchanged + +The block fills the exact shape of each object's segmentation mask with colored overlays, creating solid color fills that precisely follow object boundaries. Unlike polygon visualization (which draws outlines) or bounding box visualizations (which use rectangular regions), mask visualization fills the entire mask area with color, providing clear visual indication of the segmented regions. The opacity parameter controls how transparent the mask overlay is, allowing you to see the original image details through the colored mask (lower opacity) or create more opaque fills (higher opacity) that better obscure background details. This block requires instance segmentation predictions with mask data, as it specifically works with segmentation masks to create precise, shape-following color fills. + +## Common Use Cases + +- **Instance Segmentation Visualization**: Visualize instance segmentation results by filling mask regions with colors to clearly show segmented objects, validate segmentation quality, or highlight detected regions in analysis workflows +- **Precise Shape-Following Overlays**: Fill objects with colors that exactly match their segmented shapes, useful for applications requiring accurate region visualization such as medical imaging, quality control, or precise object identification +- **Mask-Based Object Highlighting**: Highlight segmented objects with colored overlays that follow exact object boundaries, providing clear visual distinction between different objects or object classes +- **Segmentation Model Validation**: Visualize segmentation predictions with colored mask fills to verify model performance, identify segmentation errors, or validate mask accuracy in model development and debugging workflows +- **Medical and Scientific Imaging**: Display segmented regions in medical imaging, microscopy, or scientific analysis applications where colored mask overlays help visualize tissue boundaries, cell regions, or measured areas +- **Mask Quality Inspection**: Use colored mask fills to inspect segmentation quality, verify mask boundaries, or identify areas where segmentation may need improvement in training data or model outputs + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Polygon Visualization, Bounding Box Visualization) to combine mask fills with additional annotations (labels, outlines) for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save images with mask overlays for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with mask fills to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with mask overlays as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with mask fills for live monitoring, segmentation visualization, or post-processing analysis +""" + +#: supervision's pending-track sentinel and its color (``PENDING_TRACK_ID`` / +#: ``PENDING_TRACK_COLOR = Color.GREY``, supervision 0.29.0 +#: ``annotators/utils.py``). Values are copied so the runtime path never +#: touches supervision. +_PENDING_TRACK_ID = -1 +_PENDING_TRACK_COLOR_RGB = (128, 128, 128) + +_SUPPORTED_COLOR_AXES = ("CLASS", "INDEX", "TRACK") + + +def _coco_rle_counts_to_runs(counts) -> np.ndarray: + """Decode a COCO compressed-RLE ``counts`` payload into run lengths. + + Vectorised numpy port of pycocotools' ``rleFrString``: each value is a + varint of base-48 chars carrying 5 data bits + a continuation bit (0x20), + with sign bit 0x10 in the final char, and from the 4th value on each count + stored as a delta against the value two back. Uncompressed payloads (a + list of ints) pass through. Returns int64 run lengths alternating + background/foreground over the column-major (Fortran) pixel order, + background first. + """ + if isinstance(counts, (list, tuple, np.ndarray)): + return np.asarray(counts, dtype=np.int64) + if isinstance(counts, str): + counts = counts.encode("ascii") + chars = np.frombuffer(counts, dtype=np.uint8).astype(np.int64) - 48 + if chars.size == 0: + return np.zeros(0, dtype=np.int64) + ends = (chars & 0x20) == 0 # final char of each varint + ends_idx = np.flatnonzero(ends) + starts_idx = np.concatenate(([0], ends_idx[:-1] + 1)) + value_id = np.cumsum(np.concatenate(([False], ends[:-1]))) + bit_shift = 5 * (np.arange(chars.size) - starts_idx[value_id]) + values = np.zeros(ends_idx.size, dtype=np.int64) + np.add.at(values, value_id, (chars & 0x1F) << bit_shift) + negative = (chars[ends_idx] & 0x10) != 0 + values[negative] -= np.int64(1) << (5 * (ends_idx - starts_idx + 1)[negative]) + # Undo the delta coding: values 3+ each add the decoded value two back, + # which is a running sum over the odd and even positions independently. + if values.size > 3: + values[3::2] = values[1] + np.cumsum(values[3::2]) + if values.size > 4: + values[4::2] = values[2] + np.cumsum(values[4::2]) + return values + + +def _rle_foreground_pixels_in_roi( + masks: "InstancesRLEMasks", + roi: tuple, + device: torch.device, +): + """Turn COCO-RLE masks into ``(flat ROI pixel index, detection id)`` device + tensors without materialising any dense mask. + + Host work is proportional to the encoded byte size (the varint decode is + vectorised numpy); only the per-run tables cross the bus โ€” one packed H2D + upload โ€” and the runโ†’pixel expansion happens on the device with a + host-known ``output_size`` (no deviceโ†’host sync). + """ + uy1, ux1, uy2, ux2 = roi + height = int(masks.image_size[0]) # runs are column-major over (h, w) + starts_l, lens_l, dets_l = [], [], [] + for det_idx, payload in enumerate(masks.masks): + runs = _coco_rle_counts_to_runs(payload) + bounds = np.concatenate(([0], np.cumsum(runs))) + fg_starts, fg_lens = bounds[1::2], runs[1::2] + keep = fg_lens > 0 + starts_l.append(fg_starts[: fg_lens.size][keep]) + lens_l.append(fg_lens[keep]) + dets_l.append(np.full(int(keep.sum()), det_idx, dtype=np.int64)) + starts = np.concatenate(starts_l) if starts_l else np.zeros(0, dtype=np.int64) + lens = np.concatenate(lens_l) if lens_l else np.zeros(0, dtype=np.int64) + total = int(lens.sum()) + empty = torch.zeros(0, dtype=torch.int64, device=device) + if total == 0: + return empty, empty + offsets = np.concatenate(([0], np.cumsum(lens)[:-1])) + packed = torch.from_numpy( + np.stack([starts, lens, np.concatenate(dets_l), offsets]) + ).to(device) + run_starts, run_lens, run_dets, run_offsets = packed + run_ids = torch.repeat_interleave( + torch.arange(run_lens.shape[0], device=device), + run_lens, + output_size=total, + ) + pix_f = run_starts[run_ids] + ( + torch.arange(total, device=device, dtype=torch.int64) - run_offsets[run_ids] + ) + rows, cols = pix_f % height, pix_f // height + inside = (rows >= uy1) & (rows < uy2) & (cols >= ux1) & (cols < ux2) + rows, cols = rows[inside], cols[inside] + return (rows - uy1) * (ux2 - ux1) + (cols - ux1), run_dets[run_ids][inside] + + +def _rle_to_dense_masks( + masks: "InstancesRLEMasks", device: torch.device +) -> torch.Tensor: + """Decode COCO-RLE payloads into a dense bool ``(N, H, W)`` mask stack on + ``device``, feeding the same fixed-shape composite the dense carrier uses. + + The RLE bytes live on the host, so the varint decode is inherently host + work; the per-run tables cross the bus once (H2D upload โ€” it does not + drain the device queue) and the runโ†’pixel expansion plus the scatter into + the dense stack happen on the device with host-known sizes: no + deviceโ†’host sync anywhere. + """ + height, width = (int(size) for size in masks.image_size) + n = len(masks.masks) + flat_idx, det_ids = _rle_foreground_pixels_in_roi( + masks, (0, 0, height, width), device + ) + dense = torch.zeros(n * height * width, dtype=torch.bool, device=device) + dense[det_ids * (height * width) + flat_idx] = True + return dense.view(n, height, width) + + +def gpu_mask_composite( + scene: torch.Tensor, + mask: Union[torch.Tensor, "InstancesRLEMasks", np.ndarray], + colors_rgb: Union[torch.Tensor, np.ndarray], + opacity: float, +) -> torch.Tensor: + """Torch-native mask compositor replicating ``sv.MaskAnnotator.annotate`` + with a ZERO-SYNC dense hot path. + + ``scene`` is a CHW RGB uint8 torch tensor (``WorkflowImageData``'s + ``tensor_image`` layout), mutated IN PLACE and returned. The function is + device-agnostic: the same code runs on CUDA tensors (the tensor pipeline) + and CPU tensors (the numpy-image path of the block). + + Sync-freedom contract (the whole point of this block): for a dense bool + ``(N, H, W)`` mask on the scene's device, the composite enqueues a FIXED, + N-independent number of device ops and never reads a device value back to + the host โ€” no ``.cpu()``, no ``.item()``, no ``nonzero``, no data-dependent + shapes. A sync-free viz phase queues a whole batch back-to-back instead of + serialising each image against the previous one's kernels (measured ~35 ms + of pure queue-wait per batch on Jetson NX before this rewrite). There is + deliberately NO ROI narrowing (the old union-of-boxes crop needed an + ``xyxy`` deviceโ†’host read) and NO large-N scatter branch: full-frame fp32 + bandwidth at realistic N (a few dozen instances) is far cheaper than one + default-stream sync, and the GEMM formulation below keeps the traffic at + one bool read of the mask stack plus a fixed number of full-frame fp32 + passes. + + Compositing math โ€” identical to the previous revision, kept on purpose: + + count = ฮฃ_i mask_i # (HยทW,) + color_sum = colorsแต€ยทopacity @ masks # (3, HยทW) โ€” one GEMM + out = round(color_sum / count + (1-opacity) ยท scene) where count > 0 + + Overlap semantics are intentionally simpler than supervision's: every mask + covering a pixel contributes equally (MEAN of the covering colors, + alpha-composited once), which is order-independent and diverges from + supervision's smallest-area-on-top painter's algorithm on OVERLAPPING + pixels only. Pixels covered by exactly one mask match ``sv.MaskAnnotator`` + bit-for-bit: same premultiplied blend, and ``torch.round_`` is + round-half-to-even like ``cv2.addWeighted``'s ``cvRound``. A convex + combination of uint8 values needs no clamp. fp32 accumulation is exact for + any realistic N (counts and premultiplied sums stay far below 2^24) and + the GEMM is deterministic โ€” note it assumes the default + ``torch.backends.cuda.matmul.allow_tf32 = False``; enabling TF32 globally + may shift overlap pixels by ยฑ1. + + Accepted mask carriers: + + * dense bool ``torch.Tensor`` ``(N, H, W)`` โ€” the zero-sync hot path + (non-bool dtypes are cast on device; a carrier on another device is + moved to the scene's device, which is only ever paid on the host-bound + numpy-image path); + * ``np.ndarray`` ``(N, H, W)`` โ€” uploaded once, then the dense path; + * ``InstancesRLEMasks`` (COCO compressed RLE, column-major โ€” the SAM3 / + semantic-segmentation carrier) โ€” decoded host-side (inherent: the bytes + live on the host) and scattered into a dense stack on device, then the + identical fixed-shape composite. No deviceโ†’host sync either. + + All validation happens before any write; the single in-place write to + ``scene`` is staged last, so a raising call never leaves a partially + painted image. + + Args: + scene: CHW RGB uint8 torch tensor, any device. Mutated in place โ€” + callers that need the original must pass a clone. + mask: one of the carriers above; the mask canvas must match the scene. + colors_rgb: ``(N, 3)`` uint8 per-detection RGB colors (device tensor + preferred; numpy is uploaded). + opacity: overlay opacity, matches ``sv.MaskAnnotator(opacity=...)``. + + Returns: + ``scene`` (same tensor, annotated in place). + """ + if scene.ndim != 3 or int(scene.shape[0]) != 3: + raise ValueError( + "mask visualization requires a 3-channel CHW RGB scene tensor, " + f"got shape {tuple(scene.shape)}" + ) + height, width = int(scene.shape[1]), int(scene.shape[2]) + if isinstance(mask, InstancesRLEMasks): + canvas = tuple(int(size) for size in mask.image_size) + if canvas != (height, width): + raise ValueError( + f"mask canvas {canvas} does not match scene {(height, width)}" + ) + mask = _rle_to_dense_masks(mask, scene.device) + else: + if isinstance(mask, np.ndarray): + mask = torch.from_numpy(np.ascontiguousarray(mask)) + if not isinstance(mask, torch.Tensor) or mask.ndim != 3: + raise ValueError( + "mask visualization requires a dense (N, H, W) mask tensor, " + "a (N, H, W) numpy array or InstancesRLEMasks, got " + f"{type(mask).__name__}" + ) + canvas = (int(mask.shape[1]), int(mask.shape[2])) + if canvas != (height, width): + raise ValueError( + f"mask canvas {canvas} does not match scene {(height, width)}" + ) + if mask.device != scene.device: + # Only reachable on the host-bound numpy-image path (device masks + # with a CPU scene); the tensor pipeline keeps everything on one + # device and never dispatches a device change here. + mask = mask.to(scene.device) + if mask.dtype != torch.bool: + mask = mask != 0 + if isinstance(colors_rgb, np.ndarray): + colors_rgb = torch.from_numpy(np.ascontiguousarray(colors_rgb)).to(scene.device) + n = int(mask.shape[0]) + if int(colors_rgb.shape[0]) != n: + raise ValueError( + f"got {int(colors_rgb.shape[0])} colors for {n} masks โ€” one RGB " + "color per detection is required" + ) + pixels = height * width + masks_flat = mask.reshape(n, pixels).to(torch.float32) # (N, P) + count = masks_flat.sum(dim=0) # (P,) + hit = (count > 0).unsqueeze(0) # (1, P) + colors_premul = colors_rgb.to(dtype=torch.float32).mul_(float(opacity)) # (N, 3) + # One deterministic GEMM instead of an (N, 3, H, W) broadcast or an + # index_add scatter: (3, N) @ (N, P) โ†’ the premultiplied color sum per + # pixel, with no giant intermediate and no data-dependent indexing. + color_sum = colors_premul.t() @ masks_flat # (3, P) + scene_flat = scene.reshape(3, pixels) + blended = color_sum.div_(count.clamp_(min=1.0)) + blended = blended.add_(scene_flat.to(torch.float32), alpha=1.0 - float(opacity)) + blended_u8 = blended.round_().to(torch.uint8) + # Single staged write: `where` materialises the full output first, then + # one in-place copy lands it in the caller's storage. + scene.copy_(torch.where(hit, blended_u8, scene_flat).view(3, height, width)) + return scene + + +def _resolve_color_ids( + predictions: "InstanceDetections", + color_axis: str, + device: torch.device, +) -> torch.Tensor: + """Palette ids as an ``(N,)`` int64 tensor on ``device`` โ€” the same ids + supervision's ``resolve_color_idx`` would produce, without any + deviceโ†’host sync: + + * ``INDEX`` โ†’ ``arange(n)`` built on device (``n`` comes from + ``xyxy.shape[0]`` โ€” shape metadata, no sync); + * ``CLASS`` โ†’ ``class_id`` stays a device tensor (dtype cast only, never a + ``.cpu()``); + * ``TRACK`` โ†’ tracker ids read from the host-side ``bboxes_metadata`` + dicts and uploaded H2D (uploads do not drain the device queue). The + pending-track sentinel ``-1`` is passed through for the caller to map to + supervision's gray. + + Missing ids raise ``ValueError`` BEFORE any mask work, so a doomed run + never pays a mask decode first. + """ + n = int(predictions.xyxy.shape[0]) + if color_axis == "INDEX": + return torch.arange(n, dtype=torch.int64, device=device) + if color_axis == "CLASS": + class_id = predictions.class_id + if class_id is None: + raise ValueError( + "Could not resolve color by class because " + "Detections do not have class_id. If using an annotator, " + "try setting color_lookup to sv.ColorLookup.INDEX or " + "sv.ColorLookup.TRACK." + ) + return class_id.detach().to(device=device, dtype=torch.int64) + if color_axis == "TRACK": + metadata = predictions.bboxes_metadata or [] + tracker_ids = [box.get("tracker_id") for box in metadata] + if len(tracker_ids) != n or any( + tracker_id is None for tracker_id in tracker_ids + ): + raise ValueError( + "Could not resolve color by track because " + "Detections do not have tracker_id. Did you call " + "tracker.update_with_detections(...) before annotating?" + ) + return torch.tensor( + [int(tracker_id) for tracker_id in tracker_ids], + dtype=torch.int64, + device=device, + ) + raise ValueError( + f"mask visualization supports color_axis in {_SUPPORTED_COLOR_AXES}, " + f"got {color_axis!r}" + ) + + +def _validate_inputs(predictions, color_axis: str) -> None: + """Raise a clear ``ValueError`` for inputs the torch compositor cannot + paint โ€” the previous silent supervision fallback is gone, so invalid + inputs now fail loudly instead of quietly taking a different code path.""" + if color_axis not in _SUPPORTED_COLOR_AXES: + raise ValueError( + f"mask visualization supports color_axis in {_SUPPORTED_COLOR_AXES}, " + f"got {color_axis!r}" + ) + if not isinstance(predictions, InstanceDetections): + raise ValueError( + "mask visualization requires instance segmentation predictions " + f"(InstanceDetections with masks), got {type(predictions).__name__}" + ) + n = int(predictions.xyxy.shape[0]) + mask = predictions.mask + if isinstance(mask, InstancesRLEMasks): + if len(mask.masks) != n: + raise ValueError( + f"predictions carry {len(mask.masks)} RLE masks for {n} boxes โ€” " + "one mask per detection is required" + ) + elif isinstance(mask, (torch.Tensor, np.ndarray)): + if mask.ndim != 3 or int(mask.shape[0]) != n: + raise ValueError( + "predictions must carry a dense (N, H, W) mask stack with one " + f"mask per detection; got shape {tuple(mask.shape)} for {n} boxes" + ) + else: + raise ValueError( + "predictions carry no usable mask (expected a dense (N, H, W) " + "tensor/array or InstancesRLEMasks, got " + f"{type(mask).__name__}) โ€” mask visualization requires " + "segmentation masks" + ) + + +class MaskManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "MaskVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Mask Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-mask", + "blockPriority": 12, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + predictions: Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + SEMANTIC_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Segmentation predictions containing masks for detected objects. The block uses segmentation masks to create colored fills that precisely follow object or class boundaries. Requires segmentation model outputs with mask data, which may be RLE-encoded.", + examples=["$steps.instance_segmentation_model.predictions"], + ) + + opacity: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + description="Opacity of the mask overlay, ranging from 0.0 (fully transparent) to 1.0 (fully opaque). Controls the transparency of the colored mask fill. Lower values (e.g., 0.3-0.5) create semi-transparent overlays that allow original image details to show through, while higher values (e.g., 0.7-1.0) create more opaque fills that better obscure background details. Typical values range from 0.4 to 0.7 for balanced visualization where both the mask and underlying image are visible.", + default=0.5, + examples=[0.5, "$inputs.opacity"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class MaskVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # (color_palette, palette_size, custom_colors, device) โ†’ (P, 3) uint8 + # RGB LUT. supervision's ColorPalette is consulted ONCE here, at + # cache-build time (pure configuration); the per-frame tensor path is + # supervision-free. + self._palette_lut_cache = {} + # sv.MaskAnnotator cache for the numpy-sourced-image path. + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return MaskManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + opacity: float, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + opacity, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = sv.MaskAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + opacity=opacity, + ) + + return self.annotatorCache[key] + + def _get_palette_lut( + self, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + device: torch.device, + ) -> torch.Tensor: + key = ( + color_palette, + int(palette_size) if palette_size is not None else None, + tuple(custom_colors or ()), + str(device), + ) + lut = self._palette_lut_cache.get(key) + if lut is None: + palette = self.getPalette(color_palette, palette_size, custom_colors) + palette_colors = getattr(palette, "colors", None) + if not palette_colors: + raise ValueError( + f"color palette {color_palette!r} did not resolve to a " + "color palette with at least one color" + ) + lut = torch.tensor( + [color.as_rgb() for color in palette_colors], + dtype=torch.uint8, + device=device, + ) + self._palette_lut_cache[key] = lut + return lut + + def _annotate_scene( + self, + scene: torch.Tensor, + predictions: "InstanceDetections", + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: str, + opacity: float, + ) -> torch.Tensor: + device = scene.device + ids = _resolve_color_ids(predictions, color_axis, device) + lut = self._get_palette_lut(color_palette, palette_size, custom_colors, device) + # Same color sv's `by_idx` picks: idx % palette size, on device. (For + # negative ids torch's remainder wraps instead of raising like sv's + # by_idx โ€” checking would require a deviceโ†’host read.) + colors_rgb = lut[ids.remainder(int(lut.shape[0]))] + if color_axis == "TRACK": + # sv resolve_color: the pending-track id (-1) gets Color.GREY. + pending = torch.tensor( + _PENDING_TRACK_COLOR_RGB, dtype=torch.uint8, device=device + ) + colors_rgb = torch.where( + (ids == _PENDING_TRACK_ID).unsqueeze(1), pending, colors_rgb + ) + return gpu_mask_composite(scene, predictions.mask, colors_rgb, float(opacity)) + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + opacity: Optional[float], + ) -> BlockResult: + passthrough = empty_predictions_passthrough( + image=image, detections=predictions, copy_image=copy_image + ) + if passthrough is not None: + return passthrough + if image.is_tensor_materialised(): + _validate_inputs(predictions, color_axis) + # Tensor pipeline contract: CHW RGB device tensor in, tensor out โ€” + # zero deviceโ†’host syncs on the dense path (downstream materialises + # numpy lazily only if something asks for it). + scene = image.tensor_image + if copy_image: + scene = scene.clone() + annotated = self._annotate_scene( + scene, + predictions, + color_palette, + palette_size, + custom_colors, + color_axis, + opacity, + ) + if not copy_image: + # The compositor mutated `image.tensor_image` storage in + # place; invalidate the derived numpy/base64 caches. + image.declare_tensor_image_mutated() + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, tensor_image=annotated + ) + } + # Numpy-sourced image (flag-on cv2-fallback frames): behave EXACTLY as + # before โ€” the battle-tested sv.MaskAnnotator path, byte-identical to + # the numpy v1 block. Forcing a CHW tensor out of such a frame would be + # pure host-side conversion overhead. + predictions = to_supervision_for_annotation(predictions) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + opacity, + ) + + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/model_comparison/v1.py b/inference/core/workflows/core_steps/visualizations/model_comparison/v1.py index dc0c850e41..c41eedf93c 100644 --- a/inference/core/workflows/core_steps/visualizations/model_comparison/v1.py +++ b/inference/core/workflows/core_steps/visualizations/model_comparison/v1.py @@ -206,8 +206,13 @@ def run( opacity=opacity, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections_a=predictions_a, detections_b=predictions_b, ) diff --git a/inference/core/workflows/core_steps/visualizations/model_comparison/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/model_comparison/v1_tensor.py new file mode 100644 index 0000000000..b717bc6793 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/model_comparison/v1_tensor.py @@ -0,0 +1,234 @@ +from typing import Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.annotators.model_comparison import ( + ModelComparisonAnnotator, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + PredictionsVisualizationBlock, + VisualizationManifest, + to_supervision_for_annotation, +) +from inference.core.workflows.core_steps.visualizations.common.utils import str_to_color +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + FLOAT_ZERO_TO_ONE_KIND, + STRING_KIND, + FloatZeroToOne, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/model_comparison_visualization@v1" +SHORT_DESCRIPTION = "Visualize the difference between two models' detections." +LONG_DESCRIPTION = """ +Compare predictions from two different models by color-coding areas where only one model detected objects, highlighting model differences while leaving overlapping predictions unchanged to visualize model agreement and disagreement. + +## How This Block Works + +This block takes an image and predictions from two models (Model A and Model B) and creates a visual comparison overlay that highlights differences between the models. The block: + +1. Takes an image and two sets of predictions (predictions_a and predictions_b) as input +2. Creates masks for areas predicted by each model (using bounding boxes or segmentation masks if available) +3. Identifies four distinct regions: + - Areas predicted only by Model A (colored with color_a, default green) + - Areas predicted only by Model B (colored with color_b, default red) + - Areas predicted by both models (left unchanged, allowing the original image to show through) + - Areas predicted by neither model (colored with background_color, default black) +4. Applies colored overlays to the identified regions using the specified opacity +5. Returns an annotated image where model differences are visually distinguished with color coding + +The block creates a side-by-side comparison visualization that makes it easy to see where models agree (unchanged areas) and where they disagree (color-coded areas). Areas where both models made predictions are left unchanged, allowing the original image to "shine through" and clearly showing model consensus. This visualization helps identify model strengths, weaknesses, and differences in detection behavior. The block works with object detection predictions (using bounding boxes) or instance segmentation predictions (using masks), making it versatile for comparing different model types. + +## Common Use Cases + +- **Model Evaluation and Comparison**: Compare two models' detection performance side-by-side to identify where models agree, disagree, or have different detection behaviors for model evaluation, benchmarking, or selection workflows +- **Model Development and Debugging**: Visualize differences between model versions, architectures, or configurations to understand how changes affect detection behavior, identify improvement opportunities, or debug model performance issues +- **Ensemble Model Analysis**: Compare predictions from different models in ensemble workflows to understand model agreement patterns, identify complementary strengths, or analyze consensus areas for ensemble decision-making +- **Training Data Analysis**: Compare model predictions to ground truth annotations or between training runs to identify patterns in detection differences, validate training improvements, or analyze model behavior across datasets +- **A/B Testing and Model Selection**: Visually compare candidate models to evaluate relative performance, identify detection differences, or make informed model selection decisions for deployment +- **Quality Assurance and Validation**: Validate model consistency, compare model performance on edge cases, or identify systematic differences between models for quality assurance, validation, or compliance workflows + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Model blocks** (e.g., Object Detection Model, Instance Segmentation Model) to receive predictions_a and predictions_b from different models for comparison +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save comparison visualizations for documentation, reporting, or analysis +- **Webhook blocks** to send comparison visualizations to external systems, APIs, or web applications for display in dashboards, model monitoring tools, or evaluation interfaces +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send comparison visualizations as visual evidence in alerts or reports for model performance monitoring +- **Video output blocks** to create annotated video streams or recordings with model comparison visualizations for live model evaluation, performance monitoring, or post-processing analysis +""" + + +class ModelComparisonManifest(VisualizationManifest): + type: Literal[f"{TYPE}"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Model Comparison Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-not-equal", + "blockPriority": 16, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + predictions_a: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Predictions from Model A (the first model being compared). Can be object detection, instance segmentation, or keypoint detection predictions. Areas predicted only by Model A (and not by Model B) will be colored with color_a. Works with bounding boxes or masks depending on prediction type.", + examples=["$steps.object_detection_model.predictions"], + ) + + color_a: Union[str, Selector(kind=[STRING_KIND])] = Field( # type: ignore + description="Color used to highlight areas predicted only by Model A (that Model B did not predict). Can be specified as a color name (e.g., 'GREEN', 'BLUE'), hex color code (e.g., '#00FF00', '#FFFFFF'), or RGB format (e.g., 'rgb(0, 255, 0)'). Default is GREEN to indicate Model A's unique predictions.", + default="GREEN", + examples=["GREEN", "#FFFFFF", "rgb(255, 255, 255)", "$inputs.color_a"], + ) + + predictions_b: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Predictions from Model B (the second model being compared). Can be object detection, instance segmentation, or keypoint detection predictions. Areas predicted only by Model B (and not by Model A) will be colored with color_b. Works with bounding boxes or masks depending on prediction type.", + examples=["$steps.object_detection_model.predictions"], + ) + + color_b: Union[str, Selector(kind=[STRING_KIND])] = Field( # type: ignore + description="Color used to highlight areas predicted only by Model B (that Model A did not predict). Can be specified as a color name (e.g., 'RED', 'BLUE'), hex color code (e.g., '#FF0000', '#FFFFFF'), or RGB format (e.g., 'rgb(255, 0, 0)'). Default is RED to indicate Model B's unique predictions.", + default="RED", + examples=["RED", "#FFFFFF", "rgb(255, 255, 255)", "$inputs.color_b"], + ) + + background_color: Union[str, Selector(kind=[STRING_KIND])] = Field( # type: ignore + description="Color used for areas predicted by neither model. Can be specified as a color name (e.g., 'BLACK', 'GRAY'), hex color code (e.g., '#000000', '#808080'), or RGB format (e.g., 'rgb(0, 0, 0)'). Default is BLACK to indicate areas where both models missed detections.", + default="BLACK", + examples=["BLACK", "#FFFFFF", "rgb(255, 255, 255)", "$inputs.background_color"], + ) + + opacity: Union[FloatZeroToOne, Selector(kind=[FLOAT_ZERO_TO_ONE_KIND])] = Field( # type: ignore + description="Opacity of the comparison overlay, ranging from 0.0 (fully transparent) to 1.0 (fully opaque). Controls how transparent the color-coded overlays appear over the original image. Lower values create more transparent overlays where original image details remain more visible, while higher values create more opaque overlays with stronger color emphasis. Typical values range from 0.5 to 0.8 for balanced visibility.", + default=0.7, + examples=[0.7, "$inputs.opacity"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class ModelComparisonVisualizationBlockV1(PredictionsVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return ModelComparisonManifest + + def getAnnotator( + self, + color_a: str, + color_b: str, + background_color: str, + opacity: float, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_a, + color_b, + background_color, + opacity, + ], + ) + ) + + if key not in self.annotatorCache: + color_a = str_to_color(color_a) + color_b = str_to_color(color_b) + background_color = str_to_color(background_color) + self.annotatorCache[key] = ModelComparisonAnnotator( + color_a=color_a, + color_b=color_b, + background_color=background_color, + opacity=opacity, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions_a: Union[TensorNativePrediction, TensorNativeDetections], + color_a: str, + predictions_b: Union[TensorNativePrediction, TensorNativeDetections], + color_b: str, + background_color: str, + opacity: Optional[float], + copy_image: bool, + ) -> BlockResult: + predictions_a = to_supervision_for_annotation(predictions_a) + predictions_b = to_supervision_for_annotation(predictions_b) + annotator = self.getAnnotator( + color_a=color_a, + color_b=color_b, + background_color=background_color, + opacity=opacity, + ) + + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections_a=predictions_a, + detections_b=predictions_b, + ) + + output = WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=annotated_image, + ) + + return {OUTPUT_IMAGE_KEY: output} diff --git a/inference/core/workflows/core_steps/visualizations/pixelate/v1.py b/inference/core/workflows/core_steps/visualizations/pixelate/v1.py index 0babd0f521..1929021148 100644 --- a/inference/core/workflows/core_steps/visualizations/pixelate/v1.py +++ b/inference/core/workflows/core_steps/visualizations/pixelate/v1.py @@ -122,8 +122,13 @@ def run( pixel_size, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/pixelate/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/pixelate/v1_tensor.py new file mode 100644 index 0000000000..0d97757980 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/pixelate/v1_tensor.py @@ -0,0 +1,148 @@ +from typing import Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + PredictionsVisualizationBlock, + PredictionsVisualizationManifest, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/pixelate_visualization@v1" +SHORT_DESCRIPTION = "Pixelate detected objects in an image." +LONG_DESCRIPTION = """ +Apply a pixelated mosaic effect to detected objects, creating a blocky, pixelated appearance that obscures object details while maintaining recognizable shapes, useful for privacy protection and content anonymization. + +## How This Block Works + +This block takes an image and detection predictions and applies a pixelation (mosaic) effect to the detected objects, leaving the background unchanged. The block: + +1. Takes an image and predictions as input +2. Identifies detected regions from bounding boxes or segmentation masks +3. Divides the detected object regions into square blocks (pixels) of the specified size +4. Replaces each block with a single color value (typically the average color of that block), creating a mosaic-like pixelated effect +5. Preserves the background and areas outside detected objects unchanged +6. Returns an annotated image where detected objects are pixelated with blocky, mosaic appearance, while the rest of the image remains sharp + +The block works with both object detection predictions (using bounding boxes) and instance segmentation predictions (using masks). When masks are available, it pixelates the exact shape of detected objects; otherwise, it pixelates rectangular bounding box regions. The pixel size parameter controls how large each square block is, where larger pixel sizes create more pronounced pixelation with fewer, larger blocks (more obscured), while smaller pixel sizes create finer pixelation with more, smaller blocks (less obscured). Unlike blur visualization (which creates smooth, gradient-like obscuration), pixelation creates distinct, blocky squares that maintain a more stylized, mosaic appearance while still effectively obscuring details. + +## Common Use Cases + +- **Privacy Protection and Anonymization**: Pixelate faces, people, license plates, or other sensitive information in images or videos to protect privacy, comply with data protection regulations, or anonymize content before sharing or publishing, using the distinctive pixelated mosaic effect +- **Content Filtering and Censorship**: Obscure inappropriate or sensitive content in images or videos for content moderation workflows, safe content previews, or user-generated content filtering, where the pixelated effect provides clear visual indication that content has been processed +- **Stylized Content Anonymization**: Create pixelated effects for artistic or stylized content anonymization where the mosaic appearance is preferred over smooth blur, useful for creative projects, stylized presentations, or distinctive visual effects +- **Visual Emphasis and Focus**: Pixelate detected objects to draw attention to other parts of the image, create visual contrast between pixelated foreground objects and sharp backgrounds, or emphasize specific elements in composition with a distinctive visual style +- **Security and Surveillance**: Anonymize people, vehicles, or other identifiable elements in security footage or surveillance images while preserving scene context, using pixelation as an alternative to blur for a more stylized anonymization effect +- **Documentation and Reporting**: Create pixelated, anonymized versions of images for reports, documentation, or case studies where sensitive information needs to be obscured but overall context should remain visible, with a distinctive mosaic aesthetic + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Bounding Box Visualization, Polygon Visualization) to add additional annotations on top of pixelated objects for comprehensive visualization or to indicate what was pixelated +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save pixelated images for documentation, reporting, or archiving privacy-protected content +- **Webhook blocks** to send pixelated images to external systems, APIs, or web applications for content moderation, privacy-compliant sharing, or anonymized analysis +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send pixelated images as privacy-protected visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with pixelated objects for live monitoring, privacy-compliant video processing, or post-processing analysis +""" + + +class PixelateManifest(PredictionsVisualizationManifest): + type: Literal[f"{TYPE}", "PixelateVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Pixelate Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "fad fa-grid", + "blockPriority": 13, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + pixel_size: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Size of each square pixel block in the pixelated effect, measured in pixels. Controls the granularity of the pixelation: larger values create bigger, more blocky pixels with stronger obscuration (fewer blocks, more abstract appearance), while smaller values create finer, more detailed pixelation (more blocks, less obscured). Typical values range from 10 to 50 pixels, with 20 being a good default that balances obscuration with recognizable object shape.", + default=20, + examples=[20, "$inputs.pixel_size"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class PixelateVisualizationBlockV1(PredictionsVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return PixelateManifest + + def getAnnotator( + self, + pixel_size: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join(map(str, [pixel_size])) + + if key not in self.annotatorCache: + self.annotatorCache[key] = sv.PixelateAnnotator(pixel_size=pixel_size) + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + pixel_size: Optional[int], + ) -> BlockResult: + # sv.PixelateAnnotator pixelates the `xyxy` box region and never reads + # `.mask`; skip the device->host dense-mask materialisation. + predictions = to_supervision_for_annotation( + predictions, materialise_masks=False + ) + annotator = self.getAnnotator( + pixel_size, + ) + + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/polygon/v1.py b/inference/core/workflows/core_steps/visualizations/polygon/v1.py index a62f4cdd39..2f004c6fd4 100644 --- a/inference/core/workflows/core_steps/visualizations/polygon/v1.py +++ b/inference/core/workflows/core_steps/visualizations/polygon/v1.py @@ -168,8 +168,13 @@ def run( thickness, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/polygon/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/polygon/v1_tensor.py new file mode 100644 index 0000000000..2126dcb5ce --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/polygon/v1_tensor.py @@ -0,0 +1,192 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.annotators.polygon import ( + PolygonAnnotator, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/polygon_visualization@v1" +SHORT_DESCRIPTION = "Draw a polygon around detected objects in an image." +LONG_DESCRIPTION = """ +Draw polygon outlines around detected objects that follow the exact shape of object masks, providing precise boundary visualization for instance segmentation results. + +## How This Block Works + +This block takes an image and instance segmentation predictions (which include segmentation masks) and draws polygon outlines that precisely follow the shape of each detected object. The block: + +1. Takes an image and instance segmentation predictions as input (predictions must include mask data) +2. Converts segmentation masks to polygon coordinates that trace the object boundaries +3. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +4. Draws polygon outlines with the specified thickness using the PolygonAnnotator +5. Returns an annotated image with polygon outlines overlaid on the original image + +The block extracts the exact shape of each object from its segmentation mask and draws polygon outlines that follow these precise boundaries. This provides much more accurate visualization than bounding boxes, as polygons conform to the actual object shape rather than enclosing them in rectangles. If mask data is not available, the block falls back to drawing bounding boxes. The polygon outlines can be customized with different thickness values and color palettes, allowing you to clearly distinguish between different objects or object classes. + +## Common Use Cases + +- **Precise Object Boundary Visualization**: Visualize the exact shape and boundaries of segmented objects for applications requiring accurate object outlines, such as medical imaging, manufacturing quality control, or precise measurement workflows +- **Instance Segmentation Model Validation**: Verify and debug instance segmentation model performance by visualizing how well polygon predictions match object boundaries, identify segmentation errors, and validate mask quality +- **Irregular Shape Analysis**: Visualize objects with irregular or non-rectangular shapes (e.g., people, animals, complex machinery parts) where bounding boxes would be inaccurate or misleading +- **Overlapping Object Visualization**: Clearly show object boundaries when multiple objects overlap, as polygons accurately represent each object's shape without the ambiguity of overlapping bounding boxes +- **Shape-Based Quality Control**: Inspect object shapes and boundaries in manufacturing, agriculture, or quality assurance workflows where precise object contours are critical for defect detection or classification +- **Scientific and Medical Imaging**: Visualize segmented regions in medical imaging, microscopy, or scientific analysis where accurate boundary representation is essential for measurement, analysis, or diagnosis + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Mask Visualization, Bounding Box Visualization) to combine polygon outlines with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images with polygon outlines for documentation, reporting, or training data validation +- **Webhook blocks** to send visualized results with polygon outlines to external systems, APIs, or web applications for display in dashboards or analysis tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with polygon outlines as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with polygon outlines for live monitoring, tracking, or post-processing analysis +""" + + +class PolygonManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "PolygonVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Polygon Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-shapes", + "blockPriority": 1, + "popular": True, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + predictions: Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Instance segmentation predictions containing mask data. The block converts masks to polygon outlines that follow the exact shape of each detected object.", + examples=["$steps.instance_segmentation_model.predictions"], + ) + + thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the polygon outline in pixels. Higher values create thicker, more visible outlines.", + default=2, + examples=[2, "$inputs.thickness"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class PolygonVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return PolygonManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + thickness: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + thickness, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = PolygonAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + thickness=thickness, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + thickness: Optional[int], + ) -> BlockResult: + predictions = to_supervision_for_annotation(predictions) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + thickness, + ) + + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/polygon/v2.py b/inference/core/workflows/core_steps/visualizations/polygon/v2.py index b3cd8bb108..e476f4d676 100644 --- a/inference/core/workflows/core_steps/visualizations/polygon/v2.py +++ b/inference/core/workflows/core_steps/visualizations/polygon/v2.py @@ -167,8 +167,13 @@ def run( thickness, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/polygon/v2_tensor.py b/inference/core/workflows/core_steps/visualizations/polygon/v2_tensor.py new file mode 100644 index 0000000000..f41f7bac6b --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/polygon/v2_tensor.py @@ -0,0 +1,191 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field +from supervision import PolygonAnnotator + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/polygon_visualization@v2" +SHORT_DESCRIPTION = "Draw a polygon around detected objects in an image." +LONG_DESCRIPTION = """ +Draw polygon outlines around detected objects that follow the exact shape of object masks, providing precise boundary visualization for instance segmentation results. + +## How This Block Works + +This block takes an image and instance segmentation predictions (which include segmentation masks) and draws polygon outlines that precisely follow the shape of each detected object. The block: + +1. Takes an image and instance segmentation predictions as input (predictions must include mask data) +2. Converts segmentation masks to polygon coordinates that trace the object boundaries +3. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +4. Draws polygon outlines with the specified thickness using the PolygonAnnotator +5. Returns an annotated image with polygon outlines overlaid on the original image + +The block extracts the exact shape of each object from its segmentation mask and draws polygon outlines that follow these precise boundaries. This provides much more accurate visualization than bounding boxes, as polygons conform to the actual object shape rather than enclosing them in rectangles. If mask data is not available, the block falls back to drawing bounding boxes. The polygon outlines can be customized with different thickness values and color palettes, allowing you to clearly distinguish between different objects or object classes. + +## Common Use Cases + +- **Precise Object Boundary Visualization**: Visualize the exact shape and boundaries of segmented objects for applications requiring accurate object outlines, such as medical imaging, manufacturing quality control, or precise measurement workflows +- **Instance Segmentation Model Validation**: Verify and debug instance segmentation model performance by visualizing how well polygon predictions match object boundaries, identify segmentation errors, and validate mask quality +- **Irregular Shape Analysis**: Visualize objects with irregular or non-rectangular shapes (e.g., people, animals, complex machinery parts) where bounding boxes would be inaccurate or misleading +- **Overlapping Object Visualization**: Clearly show object boundaries when multiple objects overlap, as polygons accurately represent each object's shape without the ambiguity of overlapping bounding boxes +- **Shape-Based Quality Control**: Inspect object shapes and boundaries in manufacturing, agriculture, or quality assurance workflows where precise object contours are critical for defect detection or classification +- **Scientific and Medical Imaging**: Visualize segmented regions in medical imaging, microscopy, or scientific analysis where accurate boundary representation is essential for measurement, analysis, or diagnosis + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Label Visualization, Mask Visualization, Bounding Box Visualization) to combine polygon outlines with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images with polygon outlines for documentation, reporting, or training data validation +- **Webhook blocks** to send visualized results with polygon outlines to external systems, APIs, or web applications for display in dashboards or analysis tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with polygon outlines as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with polygon outlines for live monitoring, tracking, or post-processing analysis +""" + + +class PolygonManifest(ColorableVisualizationManifest): + type: Literal[TYPE] + model_config = ConfigDict( + json_schema_extra={ + "name": "Polygon Visualization", + "version": "v2", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-shapes", + "blockPriority": 1, + "popular": True, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + predictions: Selector( + kind=[ + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Instance segmentation predictions containing mask data. The block converts masks to polygon outlines that follow the exact shape of each detected object.", + examples=["$steps.instance_segmentation_model.predictions"], + ) + + thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the polygon outline in pixels. Higher values create thicker, more visible outlines.", + default=2, + examples=[2, "$inputs.thickness"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class PolygonVisualizationBlockV2(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return PolygonManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + thickness: int, + ) -> PolygonAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + thickness, + "_".join(custom_colors) if custom_colors else "", + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = PolygonAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + thickness=thickness, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + thickness: Optional[int], + ) -> BlockResult: + predictions = to_supervision_for_annotation(predictions) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + thickness, + ) + + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/reference_path/v1.py b/inference/core/workflows/core_steps/visualizations/reference_path/v1.py index ca6a1bf43c..6e76d9314e 100644 --- a/inference/core/workflows/core_steps/visualizations/reference_path/v1.py +++ b/inference/core/workflows/core_steps/visualizations/reference_path/v1.py @@ -137,9 +137,13 @@ def run( thickness: int, ) -> BlockResult: reference_path_array = np.array(reference_path)[:, :2].astype(np.int32) - numpy_image = image.numpy_image + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() result_image = cv2.polylines( - numpy_image if not copy_image else numpy_image.copy(), + scene, [reference_path_array], False, str_to_color(color).as_bgr(), diff --git a/inference/core/workflows/core_steps/visualizations/rich_label/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/rich_label/v1_tensor.py new file mode 100644 index 0000000000..8f871e8f08 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/rich_label/v1_tensor.py @@ -0,0 +1,152 @@ +from typing import List, Optional, Type, Union + +from pydantic import Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, + split_key_point_prediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + empty_predictions_passthrough, + to_supervision_for_annotation, +) +from inference.core.workflows.core_steps.visualizations.common.label_text import ( + TEXT_SIZE_MODE_MANUAL, + build_detection_labels, + compute_adaptive_rich_font_size, +) +from inference.core.workflows.core_steps.visualizations.rich_label.v1 import TYPE +from inference.core.workflows.core_steps.visualizations.rich_label.v1 import ( + RichLabelManifest as _NumpyRichLabelManifest, +) +from inference.core.workflows.core_steps.visualizations.rich_label.v1 import ( + RichLabelVisualizationBlockV1 as _NumpyRichLabelVisualizationBlockV1, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.tensor_native_types import ( + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, +) +from inference.core.workflows.execution_engine.entities.types import Selector +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + + +class RichLabelManifest(_NumpyRichLabelManifest): + """The numpy manifest reused verbatim (same ``type`` literal, fields, + font-family validator and I/O contract) - only the ``predictions`` selector + is re-declared with the tensor-native kinds, mirroring + ``base_tensor.PredictionsVisualizationManifest``.""" + + predictions: Selector( + kind=[ + TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND, + TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND, + ] + ) = Field( # type: ignore + description="Model predictions to visualize.", + examples=["$steps.object_detection_model.predictions"], + ) + + +class RichLabelVisualizationBlockV1(_NumpyRichLabelVisualizationBlockV1): + """Tensor-native sibling of Rich Label Visualization v1. + + All drawing internals are reused from the numpy block: ``getAnnotator`` + (approved-font resolution via ``fonts.resolve_font_path`` + + ``sv.RichLabelAnnotator`` construction/caching) is inherited, label text + comes from the shared ``label_text.build_detection_labels`` and adaptive + sizing from ``label_text.compute_adaptive_rich_font_size``. This class only + adds the tensor-side glue: device-resident empty passthrough and native->sv + materialisation before the inherently-CPU Pillow text rendering. + """ + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return RichLabelManifest + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + text: Optional[str], + text_position: Optional[str], + text_color: Optional[str], + font_family: Optional[str], + text_size_mode: Optional[str], + font_size: Optional[int], + text_padding: Optional[int], + border_radius: Optional[int], + max_line_length: Optional[int], + ) -> BlockResult: + detections = ( + split_key_point_prediction(predictions)[1] + if isinstance(predictions, tuple) + else predictions + ) + passthrough = empty_predictions_passthrough( + image=image, detections=detections, copy_image=copy_image + ) + if passthrough is not None: + return passthrough + # `.mask` is read for exactly two configurations, and only for + # instance-segmentation input (see label/v1_tensor.py for details): + # * text == "Area": `sv.Detections.area` reports MASK area when a + # mask is present and BOX area when it is None - flag-off shows + # mask area on IS input, so the mask must be materialised to match; + # * text_position == "CENTER_OF_MASS": the annotator anchors on the + # mask centroid; `get_anchors_coordinates` RAISES without a mask. + # Every other label reads xyxy / confidence / per-box metadata, so the + # device->host dense-mask copy is skipped for them. + needs_masks = text == "Area" or text_position == "CENTER_OF_MASS" + sv_detections = to_supervision_for_annotation( + predictions, materialise_masks=needs_masks + ) + # Non-empty frame: rich-label rendering is inherently CPU work (Pillow + # TrueType rasterisation), so the numpy image is materialised here. + height, width = image.numpy_image.shape[:2] + effective_font_size = compute_adaptive_rich_font_size( + height, + width, + manual_font_size=font_size, + text_size_mode=text_size_mode or TEXT_SIZE_MODE_MANUAL, + ) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + text_position, + text_color, + font_family, + effective_font_size, + text_padding, + border_radius, + max_line_length, + ) + labels = build_detection_labels(sv_detections, text) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=sv_detections, + labels=labels, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/core_steps/visualizations/text_display/v1.py b/inference/core/workflows/core_steps/visualizations/text_display/v1.py index 049dbfdcab..cc9b4bb71e 100644 --- a/inference/core/workflows/core_steps/visualizations/text_display/v1.py +++ b/inference/core/workflows/core_steps/visualizations/text_display/v1.py @@ -388,7 +388,11 @@ def run( text_parameters=text_parameters, text_parameters_operations=text_parameters_operations, ) - output_image = image.numpy_image.copy() if copy_image else image.numpy_image + output_image = image.numpy_image + if copy_image: + output_image = output_image.copy() + else: + image.declare_numpy_image_mutated() text_color_bgr = str_to_color(text_color).as_bgr() bg_color_bgr = ( diff --git a/inference/core/workflows/core_steps/visualizations/trace/v1.py b/inference/core/workflows/core_steps/visualizations/trace/v1.py index 0a276c9313..c75c9b846b 100644 --- a/inference/core/workflows/core_steps/visualizations/trace/v1.py +++ b/inference/core/workflows/core_steps/visualizations/trace/v1.py @@ -193,8 +193,13 @@ def run( trace_length=trace_length, thickness=thickness, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/trace/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/trace/v1_tensor.py new file mode 100644 index 0000000000..8ee2d0e384 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/trace/v1_tensor.py @@ -0,0 +1,249 @@ +from typing import Any, List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field, field_validator +from supervision.annotators.base import BaseAnnotator + +from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import ( + STILL_IMAGE_INPUT_SOFT_RESTRICTION, + BlockResult, + Runtime, + RuntimeInputMode, + RuntimeRestriction, + Severity, + WorkflowBlockManifest, +) + +SHORT_DESCRIPTION = "Draw traces based on detections tracking results." +LONG_DESCRIPTION = """ +Draw trajectory paths for tracked objects, visualizing their movement history by connecting recent positions with colored lines to show object movement patterns, paths, and tracking behavior over time. + +## How This Block Works + +This block takes an image and tracked predictions (with tracker IDs) and draws trajectory paths showing the recent movement history of each tracked object. The block: + +1. Takes an image and tracked predictions as input (predictions must include tracker_id data from a tracking block) +2. Extracts tracking IDs and position history for each tracked object +3. Determines the reference point for drawing traces based on the selected position anchor (center, corners, edges, or center of mass) +4. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +5. Draws trajectory lines connecting the recent positions (up to trace_length positions) for each tracked object using Supervision's TraceAnnotator +6. Connects historical positions sequentially, creating path traces that show object movement direction and patterns +7. Returns an annotated image with trajectory paths overlaid on the original image + +The block visualizes object tracking by drawing the path that each tracked object has taken over recent frames. Each tracked object gets a unique trace line (colored by track ID, class, or index) that connects its recent positions, creating a visual trail that shows movement direction, speed, and trajectory patterns. The trace_length parameter controls how many historical positions are included in each trace (longer traces show more movement history, shorter traces show recent movement only). This visualization requires predictions with tracker IDs from tracking blocks (like Byte Tracker), as it needs the tracking information to connect positions across frames. The traces help visualize object movement, identify tracking patterns, and understand object behavior over time. + +## Common Use Cases + +- **Object Trajectory Visualization**: Visualize movement paths and trajectories of tracked objects to understand object behavior, movement patterns, or navigation routes for applications like vehicle tracking, pedestrian flow analysis, or object movement monitoring +- **Tracking Performance Validation**: Validate tracking performance by visualizing object paths to ensure tracking consistency, identify tracking errors or ID switches, or verify that objects maintain consistent trajectories +- **Movement Pattern Analysis**: Analyze movement patterns, speeds, or direction changes by visualizing trajectory traces to understand object behavior, detect anomalies, or identify movement trends in surveillance, security, or traffic monitoring workflows +- **Path Deviation Detection**: Visualize object paths to detect deviations from expected routes, identify unusual movement patterns, or monitor object trajectories for safety, security, or compliance workflows +- **Real-Time Tracking Monitoring**: Display trajectory traces in real-time monitoring interfaces, dashboards, or live video feeds to visualize object movement and tracking behavior as it happens +- **Video Analysis and Post-Processing**: Create trajectory visualizations for video analysis, post-processing workflows, or forensic analysis where understanding object movement paths and patterns is critical + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Tracking blocks** (e.g., Byte Tracker) to receive tracked predictions with tracker IDs that are required for trace visualization +- **Other visualization blocks** (e.g., Bounding Box Visualization, Label Visualization, Dot Visualization) to combine trajectory traces with additional annotations for comprehensive tracking visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save images with trajectory traces for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with trajectory traces to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with trajectory traces as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with trajectory traces for live monitoring, tracking visualization, or post-processing analysis +""" + + +class TraceManifest(ColorableVisualizationManifest): + type: Literal["roboflow_core/trace_visualization@v1"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Trace Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-scribble", + "blockPriority": 17, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + position: Union[ + Literal[ + "CENTER", + "CENTER_LEFT", + "CENTER_RIGHT", + "TOP_CENTER", + "TOP_LEFT", + "TOP_RIGHT", + "BOTTOM_LEFT", + "BOTTOM_CENTER", + "BOTTOM_RIGHT", + "CENTER_OF_MASS", + ], + Selector(kind=[STRING_KIND]), + ] = Field( # type: ignore + default="CENTER", + description="Anchor position for drawing trajectory traces relative to each detection's bounding box. Options include: CENTER (center of box), corners (TOP_LEFT, TOP_RIGHT, BOTTOM_LEFT, BOTTOM_RIGHT), edge midpoints (TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, BOTTOM_CENTER), or CENTER_OF_MASS (center of mass of the object). The trace path is drawn connecting positions at this anchor point across recent frames.", + examples=["CENTER", "$inputs.text_position"], + ) + trace_length: Union[int, Selector(kind=[INTEGER_KIND])] = Field( + default=30, + description="Maximum number of historical tracked object positions to include in each trajectory trace. Controls how long the movement trail appears. Higher values create longer traces showing more movement history, while lower values create shorter traces showing only recent movement. Must be at least 1. Typical values range from 10 to 50 frames depending on the desired trail length and frame rate.", + examples=[30, "$inputs.trace_length"], + ) + thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the trajectory trace lines in pixels. Controls how thick the path lines appear. Higher values create thicker, more visible traces, while lower values create thinner, more subtle traces. Must be at least 1. Typical values range from 1 to 5 pixels.", + default=1, + examples=[1, "$inputs.track_thickness"], + ) + + @field_validator("trace_length", "thickness") + @classmethod + def ensure_max_entries_per_file_is_correct(cls, value: Any) -> Any: + if isinstance(value, int) and value < 1: + raise ValueError("`trace_length` and `thickness` cannot be lower than 1.") + return value + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + @classmethod + def get_restrictions(cls) -> List[RuntimeRestriction]: + restriction = RuntimeRestriction( + severity=Severity.SOFT, + note=( + "Trajectory history is stored inside a cached TraceAnnotator " + "in process memory. With remote step execution on stateless " + "or multi-replica HTTP runtimes, successive frames may be " + "served by different worker processes, so traces reset or " + "split across workers. Use local step execution in an " + "InferencePipeline for stable cross-frame visualizations." + ), + applies_to_runtimes=[ + Runtime.HOSTED_SERVERLESS, + Runtime.DEDICATED_DEPLOYMENT, + ], + applies_to_step_execution_modes=[StepExecutionMode.REMOTE], + applies_to_input_modes=[RuntimeInputMode.VIDEO], + ) + return [restriction, STILL_IMAGE_INPUT_SOFT_RESTRICTION] + + +class TraceVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return TraceManifest + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + position: str, + trace_length: int, + thickness: int, + ) -> BlockResult: + predictions = to_supervision_for_annotation(predictions) + if predictions.tracker_id is None: + raise ValueError( + "Expected tracked predictions in `roboflow_core/trace_visualization@v1` block." + ) + annotator = self.getAnnotator( + color_palette=color_palette, + palette_size=palette_size, + custom_colors=custom_colors, + color_axis=color_axis, + position=position, + trace_length=trace_length, + thickness=thickness, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + position: str, + trace_length: int, + thickness: int, + ) -> BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + position, + trace_length, + thickness, + ], + ) + ) + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + self.annotatorCache[key] = sv.TraceAnnotator( + color=palette, + position=getattr(sv.Position, position), + trace_length=trace_length, + thickness=thickness, + color_lookup=getattr(sv.ColorLookup, color_axis), + ) + return self.annotatorCache[key] diff --git a/inference/core/workflows/core_steps/visualizations/triangle/v1.py b/inference/core/workflows/core_steps/visualizations/triangle/v1.py index f8ae75534f..521a2833be 100644 --- a/inference/core/workflows/core_steps/visualizations/triangle/v1.py +++ b/inference/core/workflows/core_steps/visualizations/triangle/v1.py @@ -199,8 +199,13 @@ def run( height, outline_thickness, ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() annotated_image = annotator.annotate( - scene=image.numpy_image.copy() if copy_image else image.numpy_image, + scene=scene, detections=predictions, ) return { diff --git a/inference/core/workflows/core_steps/visualizations/triangle/v1_tensor.py b/inference/core/workflows/core_steps/visualizations/triangle/v1_tensor.py new file mode 100644 index 0000000000..eec08d8113 --- /dev/null +++ b/inference/core/workflows/core_steps/visualizations/triangle/v1_tensor.py @@ -0,0 +1,221 @@ +from typing import List, Literal, Optional, Type, Union + +import supervision as sv +from pydantic import ConfigDict, Field + +from inference.core.workflows.core_steps.common.tensor_native import ( + TensorNativeDetections, + TensorNativePrediction, +) +from inference.core.workflows.core_steps.visualizations.common.base_colorable_tensor import ( + ColorableVisualizationBlock, + ColorableVisualizationManifest, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + OUTPUT_IMAGE_KEY, + to_supervision_for_annotation, +) +from inference.core.workflows.execution_engine.entities.base import WorkflowImageData +from inference.core.workflows.execution_engine.entities.types import ( + INTEGER_KIND, + STRING_KIND, + Selector, +) +from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest + +TYPE: str = "roboflow_core/triangle_visualization@v1" +SHORT_DESCRIPTION = "Draw triangle markers on an image at specific coordinates based on provided detections." +LONG_DESCRIPTION = """ +Draw triangular markers on an image to mark specific points on detected objects, with customizable position, size, color, and outline styling, providing directional indicators and geometric markers for visual annotations. + +## How This Block Works + +This block takes an image and detection predictions and draws triangular markers at specified anchor positions on each detected object. The block: + +1. Takes an image and predictions as input +2. Determines the triangle position for each detection based on the selected anchor point (center, corners, edges, or center of mass) +3. Applies color styling based on the selected color palette, with colors assigned by class, index, or track ID +4. Draws triangular markers with the specified base width and height dimensions, with optional outline thickness using Supervision's TriangleAnnotator +5. Returns an annotated image with triangular markers overlaid on the original image + +The block supports various position options including the center of the bounding box, any of the four corners, edge midpoints, or the center of mass (useful for objects with irregular shapes). Triangles can be customized with different sizes (base width and height), optional outlines for better visibility, and various color palettes. Triangular markers provide a distinctive geometric shape that can serve as directional indicators (pointing in a specific direction) or geometric markers, offering an alternative to circular dots or square markers. This provides a clean visualization style that marks detection locations with directional or geometric emphasis, making it ideal for applications requiring directional indicators, geometric markers, or distinctive point markers. + +## Common Use Cases + +- **Directional Object Marking**: Mark detected objects with triangular markers that can indicate direction or orientation, useful for tracking applications, motion analysis, or directional workflows where the triangle's pointed shape provides directional information +- **Geometric Marker Visualization**: Use triangular markers as distinctive geometric shapes to mark detection locations, providing visual variety compared to circular dots or rectangular bounding boxes for design purposes or geometric emphasis +- **Minimal Object Marking**: Mark detected objects with small triangular markers instead of bounding boxes for cleaner, less cluttered visualizations when working with dense scenes or many detections, with the triangular shape providing visual distinction +- **Tracking Visualization**: Use triangular markers to visualize object trajectories or tracking IDs over time, creating a cleaner alternative to bounding boxes for tracking workflows, with triangles potentially indicating movement direction +- **Point of Interest Highlighting**: Mark specific anchor points (corners, center, center of mass) on detected objects with triangular markers for applications like object tracking, spatial analysis, or geometric annotation workflows +- **Design and Aesthetic Applications**: Create triangular markers for design purposes, user interfaces, dashboards, or artistic visualizations where geometric shapes provide distinctive visual style or aesthetic appeal + +## Connecting to Other Blocks + +The annotated image from this block can be connected to: + +- **Other visualization blocks** (e.g., Bounding Box Visualization, Label Visualization, Trace Visualization) to combine triangular markers with additional annotations for comprehensive visualization +- **Data storage blocks** (e.g., Local File Sink, CSV Formatter, Roboflow Dataset Upload) to save annotated images with triangular markers for documentation, reporting, or analysis +- **Webhook blocks** to send visualized results with triangular markers to external systems, APIs, or web applications for display in dashboards or monitoring tools +- **Notification blocks** (e.g., Email Notification, Slack Notification) to send annotated images with triangular markers as visual evidence in alerts or reports +- **Video output blocks** to create annotated video streams or recordings with triangular markers for live monitoring, tracking visualization, or post-processing analysis +""" + + +class TriangleManifest(ColorableVisualizationManifest): + type: Literal[f"{TYPE}", "TriangleVisualization"] + model_config = ConfigDict( + json_schema_extra={ + "name": "Triangle Visualization", + "version": "v1", + "short_description": SHORT_DESCRIPTION, + "long_description": LONG_DESCRIPTION, + "license": "Apache-2.0", + "block_type": "visualization", + "search_keywords": ["annotator"], + "ui_manifest": { + "section": "visualization", + "icon": "far fa-triangle", + "blockPriority": 14, + "supervision": True, + "warnings": [ + { + "property": "copy_image", + "value": False, + "message": "This setting will mutate its input image. If the input is used by other blocks, it may cause unexpected behavior.", + } + ], + }, + } + ) + + position: Union[ + Literal[ + "CENTER", + "CENTER_LEFT", + "CENTER_RIGHT", + "TOP_CENTER", + "TOP_LEFT", + "TOP_RIGHT", + "BOTTOM_LEFT", + "BOTTOM_CENTER", + "BOTTOM_RIGHT", + "CENTER_OF_MASS", + ], + Selector(kind=[STRING_KIND]), + ] = Field( # type: ignore + default="TOP_CENTER", + description="Anchor position for placing the triangle marker relative to the detection's bounding box. Options include: CENTER (center of box), corners (TOP_LEFT, TOP_RIGHT, BOTTOM_LEFT, BOTTOM_RIGHT), edge midpoints (TOP_CENTER, CENTER_LEFT, CENTER_RIGHT, BOTTOM_CENTER), or CENTER_OF_MASS (center of mass of the object, useful for irregular shapes).", + examples=["CENTER", "$inputs.position"], + ) + + base: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Base width of the triangle in pixels. Controls the horizontal width of the triangular marker at its base. Larger values create wider triangles, while smaller values create narrower triangles. Works together with height to control the overall triangle size and shape.", + default=10, + examples=[10, "$inputs.base"], + ) + + height: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Height of the triangle in pixels. Controls the vertical height of the triangular marker from base to tip. Larger values create taller triangles, while smaller values create shorter triangles. Works together with base to control the overall triangle size and shape.", + default=10, + examples=[10, "$inputs.height"], + ) + + outline_thickness: Union[int, Selector(kind=[INTEGER_KIND])] = Field( # type: ignore + description="Thickness of the triangle outline in pixels. A value of 0 creates a filled triangle with no outline. Higher values create thicker outlines around the triangle border, improving visibility and contrast. Useful for making triangular markers more visible against complex backgrounds.", + default=0, + examples=[2, "$inputs.outline_thickness"], + ) + + @classmethod + def get_execution_engine_compatibility(cls) -> Optional[str]: + return ">=1.3.0,<2.0.0" + + +class TriangleVisualizationBlockV1(ColorableVisualizationBlock): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.annotatorCache = {} + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return TriangleManifest + + def getAnnotator( + self, + color_palette: str, + palette_size: int, + custom_colors: List[str], + color_axis: str, + position: str, + base: int, + height: int, + outline_thickness: int, + ) -> sv.annotators.base.BaseAnnotator: + key = "_".join( + map( + str, + [ + color_palette, + palette_size, + color_axis, + position, + base, + height, + outline_thickness, + ], + ) + ) + + if key not in self.annotatorCache: + palette = self.getPalette(color_palette, palette_size, custom_colors) + + self.annotatorCache[key] = sv.TriangleAnnotator( + color=palette, + color_lookup=getattr(sv.ColorLookup, color_axis), + position=getattr(sv.Position, position), + base=base, + height=height, + outline_thickness=outline_thickness, + ) + + return self.annotatorCache[key] + + def run( + self, + image: WorkflowImageData, + predictions: Union[TensorNativePrediction, TensorNativeDetections], + copy_image: bool, + color_palette: Optional[str], + palette_size: Optional[int], + custom_colors: Optional[List[str]], + color_axis: Optional[str], + position: Optional[str], + base: Optional[int], + height: Optional[int], + outline_thickness: Optional[int], + ) -> BlockResult: + predictions = to_supervision_for_annotation(predictions) + annotator = self.getAnnotator( + color_palette, + palette_size, + custom_colors, + color_axis, + position, + base, + height, + outline_thickness, + ) + scene = image.numpy_image + if copy_image: + scene = scene.copy() + else: + image.declare_numpy_image_mutated() + annotated_image = annotator.annotate( + scene=scene, + detections=predictions, + ) + return { + OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace( + origin_image_data=image, numpy_image=annotated_image + ) + } diff --git a/inference/core/workflows/execution_engine/constants.py b/inference/core/workflows/execution_engine/constants.py index ba07835836..4eb6f714f4 100644 --- a/inference/core/workflows/execution_engine/constants.py +++ b/inference/core/workflows/execution_engine/constants.py @@ -65,4 +65,23 @@ AREA_CONVERTED_KEY_IN_INFERENCE_RESPONSE = "area_converted" RLE_MASK_KEY_IN_SV_DETECTIONS = "rle_mask" RLE_MASK_KEY_IN_INFERENCE_RESPONSE = "rle_mask" +CLASS_NAMES_KEY = "class_names" +# D1: optional boolean carried in a tensor-native detection's ``image_metadata``. +# Meaning: "emit per-instance mask polygons (the ``points`` field) during +# serialisation?". When the key is ABSENT the serialiser behaves exactly as before +# (polygons ARE computed and emitted). A producer that sets this to ``False`` +# instructs the serialiser to skip the per-instance RLE-decode + ``mask_to_polygon`` +# / ``findContours`` work. Wave-2 InstanceDetections producers set this value. +SERIALISE_POLYGONS_KEY = "serialise_polygons" +# Lane 1b: explicit style tag carried in a tensor-native classification prediction's +# ``image_metadata``, telling ``serialise_native_classification`` which flag-OFF shape +# to reproduce. ``"model"`` -> numpy ``*InferenceResponse``-derived shape (roboflow +# classification model blocks + visual_search_classifier + clip_comparison); +# ``"formatter"`` -> hand-built ``vlm_as_classifier`` shape. Replaces the interim +# ``threshold|root_parent_id|time`` heuristic; when the key is ABSENT (or an +# unrecognised value) the serialiser falls back to that heuristic, so this is +# non-breaking for any producer that has not yet been wired. +CLASSIFICATION_STYLE_KEY = "classification_style" +CLASSIFICATION_STYLE_MODEL = "model" +CLASSIFICATION_STYLE_FORMATTER = "formatter" NEAREST_TARGET_DISTANCE_KEY = "nearest_target_distance" diff --git a/inference/core/workflows/execution_engine/entities/base.py b/inference/core/workflows/execution_engine/entities/base.py index b3d4d8643b..0f5ea65a26 100644 --- a/inference/core/workflows/execution_engine/entities/base.py +++ b/inference/core/workflows/execution_engine/entities/base.py @@ -18,9 +18,15 @@ import cv2 import numpy as np +import torch from pydantic import BaseModel, Field +from torchvision.io import ImageReadMode, decode_image, read_file from typing_extensions import Annotated, Literal +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_IMAGE_TENSOR_DEVICE, +) from inference.core.utils.image_utils import ( attempt_loading_image_from_string, encode_image_to_jpeg_bytes, @@ -264,6 +270,31 @@ class ImageParentMetadata: class WorkflowImageData: + """Container hosting the image in one of several representations, materialised + lazily on access. + + Layout contract: + * ``numpy_image``: HWC uint8 BGR (cv2 native); 2-D ``(H, W)`` for + single-channel images (grayscale / threshold outputs). + * ``tensor_image``: CHW uint8 RGB on ``WORKFLOWS_IMAGE_TENSOR_DEVICE``; + ``(1, H, W)`` for single-channel images. Single-channel data carries no + BGR/RGB semantics, so it is never channel-reversed in either direction. + + Mutation contract: representations may be cached simultaneously (fan-out + readers alternate between them for free) and are exposed as the raw mutable + buffers - legacy blocks mutate ``numpy_image`` in place (e.g. visualizations + with ``copy_image=False``) and the class does not police that. A client + that mutates a representation in place MUST fetch the buffer via the + property, mutate it, and then declare the mutation via + ``declare_numpy_image_mutated()`` / ``declare_tensor_image_mutated()``. + Declaring makes the mutated representation the SOLE source of truth: the + derived sibling caches are removed so later readers re-derive from the + mutated pixels, and the original source (base64 string / file / URL + reference) is cut off as well - it no longer describes the pixels, so + serialization must never hand it out. A ``base64_image`` access AFTER the + declare re-encodes from the mutated pixels - that re-derived value is + valid and cached again. Undeclared in-place mutation leaves sibling caches + stale.""" def __init__( self, @@ -273,8 +304,14 @@ def __init__( base64_image: Optional[str] = None, numpy_image: Optional[np.ndarray] = None, video_metadata: Optional[VideoMetadata] = None, + tensor_image: Optional[torch.Tensor] = None, ): - if not base64_image and numpy_image is None and not image_reference: + if ( + not base64_image + and numpy_image is None + and not image_reference + and tensor_image is None + ): raise ValueError("Could not initialise empty `WorkflowImageData`.") self._parent_metadata = parent_metadata self._workflow_root_ancestor_metadata = ( @@ -285,7 +322,17 @@ def __init__( self._image_reference = image_reference self._base64_image = base64_image self._numpy_image = numpy_image + self._tensor_image = ( + tensor_image.to(WORKFLOWS_IMAGE_TENSOR_DEVICE) + if tensor_image is not None + else None + ) self._video_metadata = video_metadata + # Flipped by declare_*_image_mutated(): once an in-place mutation is + # declared, the original base64/reference sources no longer describe + # the pixels and must never be decoded again (see the guard in the + # image properties). + self._sources_invalidated_by_mutation = False @classmethod def copy_and_replace( @@ -300,10 +347,11 @@ def copy_and_replace( * image_reference * base64_image * numpy_image + * tensor_image * video_metadata - When more than one from ["numpy_image", "base64_image", "image_reference"] args are - given, they MUST be compliant. + When more than one from ["numpy_image", "base64_image", "image_reference", "tensor_image"] + args are given, they MUST be compliant. """ parent_metadata = origin_image_data._parent_metadata workflow_root_ancestor_metadata = ( @@ -312,11 +360,16 @@ def copy_and_replace( image_reference = origin_image_data._image_reference base64_image = origin_image_data._base64_image numpy_image = origin_image_data._numpy_image + tensor_image = origin_image_data._tensor_image video_metadata = origin_image_data._video_metadata - if any(k in kwargs for k in ["numpy_image", "base64_image", "image_reference"]): + if any( + k in kwargs + for k in ["numpy_image", "base64_image", "image_reference", "tensor_image"] + ): numpy_image = kwargs.get("numpy_image") base64_image = kwargs.get("base64_image") image_reference = kwargs.get("image_reference") + tensor_image = kwargs.get("tensor_image") if "parent_metadata" in kwargs: if workflow_root_ancestor_metadata is parent_metadata: workflow_root_ancestor_metadata = kwargs["parent_metadata"] @@ -333,6 +386,7 @@ def copy_and_replace( image_reference=image_reference, base64_image=base64_image, numpy_image=numpy_image, + tensor_image=tensor_image, video_metadata=video_metadata, ) @@ -383,15 +437,63 @@ def create_crop( video_metadata=video_metadata, ) + @classmethod + def create_crop_from_tensor( + cls, + origin_image_data: "WorkflowImageData", + crop_identifier: str, + cropped_tensor_image: torch.Tensor, + offset_x: int, + offset_y: int, + preserve_video_metadata: bool = False, + ) -> "WorkflowImageData": + """ + Tensor-native mirror of `create_crop`. Identical metadata math; + the child carries `tensor_image` instead of `numpy_image`. + """ + origin_h, origin_w = origin_image_data._read_shape_without_materialization() + parent_metadata = ImageParentMetadata( + parent_id=crop_identifier, + origin_coordinates=OriginCoordinatesSystem( + left_top_x=offset_x, + left_top_y=offset_y, + origin_width=origin_w, + origin_height=origin_h, + ), + ) + workflow_root_ancestor_coordinates = replace( + origin_image_data.workflow_root_ancestor_metadata.origin_coordinates, + left_top_x=origin_image_data.workflow_root_ancestor_metadata.origin_coordinates.left_top_x + + offset_x, + left_top_y=origin_image_data.workflow_root_ancestor_metadata.origin_coordinates.left_top_y + + offset_y, + ) + workflow_root_ancestor_metadata = ImageParentMetadata( + parent_id=origin_image_data.workflow_root_ancestor_metadata.parent_id, + origin_coordinates=workflow_root_ancestor_coordinates, + ) + video_metadata = None + if preserve_video_metadata and origin_image_data._video_metadata is not None: + video_metadata = copy(origin_image_data._video_metadata) + video_metadata.video_identifier = ( + f"{video_metadata.video_identifier} | crop: {crop_identifier}" + ) + return WorkflowImageData( + parent_metadata=parent_metadata, + workflow_root_ancestor_metadata=workflow_root_ancestor_metadata, + tensor_image=cropped_tensor_image, + video_metadata=video_metadata, + ) + @property def parent_metadata(self) -> ImageParentMetadata: if self._parent_metadata.origin_coordinates is None: - numpy_image = self.numpy_image + h, w = self._read_shape_without_materialization() origin_coordinates = OriginCoordinatesSystem( left_top_y=0, left_top_x=0, - origin_width=numpy_image.shape[1], - origin_height=numpy_image.shape[0], + origin_width=w, + origin_height=h, ) self._parent_metadata = replace( self._parent_metadata, origin_coordinates=origin_coordinates @@ -401,12 +503,12 @@ def parent_metadata(self) -> ImageParentMetadata: @property def workflow_root_ancestor_metadata(self) -> ImageParentMetadata: if self._workflow_root_ancestor_metadata.origin_coordinates is None: - numpy_image = self.numpy_image + h, w = self._read_shape_without_materialization() origin_coordinates = OriginCoordinatesSystem( left_top_y=0, left_top_x=0, - origin_width=numpy_image.shape[1], - origin_height=numpy_image.shape[0], + origin_width=w, + origin_height=h, ) self._workflow_root_ancestor_metadata = replace( self._workflow_root_ancestor_metadata, @@ -414,13 +516,54 @@ def workflow_root_ancestor_metadata(self) -> ImageParentMetadata: ) return self._workflow_root_ancestor_metadata + def _read_shape_without_materialization(self) -> Tuple[int, int]: + """Returns (height, width). Prefers whichever representation is already + set, so no numpy<->tensor conversion is ever triggered. When neither is + materialised yet (a base64/reference-born image), the source is decoded + into the representation the current mode works with - tensor under + ENABLE_TENSOR_DATA_REPRESENTATION (every caller of this helper is a + tensor-mode block that will need the tensor anyway), numpy otherwise - + so the shape read does not leave behind a representation the run has no + use for.""" + if self._numpy_image is not None: + return self._numpy_image.shape[0], self._numpy_image.shape[1] + if self._tensor_image is not None: + # tensor_image is CHW -> H=shape[1], W=shape[2] + return ( + int(self._tensor_image.shape[1]), + int(self._tensor_image.shape[2]), + ) + if ENABLE_TENSOR_DATA_REPRESENTATION: + tensor = self.tensor_image + return int(tensor.shape[1]), int(tensor.shape[2]) + np_img = self.numpy_image + return np_img.shape[0], np_img.shape[1] + @property def numpy_image(self) -> np.ndarray: + # Layout + mutation contract: see the class docstring. In-place mutators + # of the returned buffer must call declare_numpy_image_mutated(). if self._numpy_image is not None: return self._numpy_image + if self._tensor_image is not None: + if int(self._tensor_image.shape[0]) == 1: + # Single-channel: (1, H, W) -> (H, W), no channel reversal. + self._numpy_image = ( + self._tensor_image.detach().squeeze(0).to("cpu").numpy().copy() + ) + else: + # CHW RGB -> HWC (permute on-device before host transfer) -> BGR + hwc_rgb = self._tensor_image.detach().permute(1, 2, 0).to("cpu").numpy() + self._numpy_image = hwc_rgb[:, :, ::-1].copy() + return self._numpy_image if self._base64_image: + # Post-mutation this can only be a RE-DERIVED base64 (the declare + # nulls the original and base64_image re-encodes from the mutated + # pixels), so decoding it is always valid - the staleness guard + # only covers the reference branch below. self._numpy_image = attempt_loading_image_from_string(self._base64_image)[0] return self._numpy_image + self._ensure_original_sources_usable() if self._image_reference.startswith( "http://" ) or self._image_reference.startswith("https://"): @@ -429,6 +572,139 @@ def numpy_image(self) -> np.ndarray: self._numpy_image = cv2.imread(self._image_reference) return self._numpy_image + @property + def tensor_image(self) -> torch.Tensor: + # Layout + mutation contract: see the class docstring. In-place mutators + # of the returned tensor must call declare_tensor_image_mutated(). + if self._tensor_image is not None: + return self._tensor_image + if self._numpy_image is not None: + bgr_np = self._numpy_image + if bgr_np.ndim == 2: + # Single-channel (grayscale / threshold outputs): (H, W) -> + # (1, H, W), no channel reversal - there is no BGR/RGB + # semantics to convert. `.copy()` so the tensor owns its buffer: + # without it a CPU grayscale tensor aliases `self._numpy_image` + # (the trailing `.to(...)` is a no-op on CPU), breaking mutation + # isolation between the two representations - the 3-channel branch + # below already copies. + chw = torch.from_numpy(np.ascontiguousarray(bgr_np).copy()).unsqueeze(0) + else: + # HWC BGR -> HWC RGB -> CHW RGB; contiguous so model ingestion + # gets a dense buffer. + chw = torch.from_numpy(bgr_np[:, :, ::-1].copy()).permute(2, 0, 1) + else: + # A base64/reference-born image asked for the tensor first: the + # source decodes DIRECTLY into a CHW RGB tensor - no numpy hop and + # nothing cached besides the tensor itself. + chw = self._decode_source_to_tensor() + self._tensor_image = chw.contiguous().to(WORKFLOWS_IMAGE_TENSOR_DEVICE) + return self._tensor_image + + def _decode_source_to_tensor(self) -> torch.Tensor: + """Decode the base64 / file / URL source straight into a CHW RGB uint8 + tensor via ``torchvision.io`` - no numpy intermediate, no cv2. Only + reached when neither in-memory representation exists (the constructor + guarantees a source). EXIF handling mirrors the numpy path: cv2.imdecode + (the base64 route) does not apply EXIF orientation, cv2.imread (local + files) does. URLs are the one exception: their bytes come through the + SSRF-guarded numpy loader today, so they keep a single numpy hop until a + bytes-level guarded fetcher exists.""" + if self._base64_image: + # Post-mutation this can only be a RE-DERIVED base64 (see the + # matching note in numpy_image) - always valid to decode. + payload = torch.frombuffer( + bytearray(base64.b64decode(self._base64_image)), dtype=torch.uint8 + ) + return decode_image(payload, mode=ImageReadMode.RGB) + self._ensure_original_sources_usable() + if self._image_reference.startswith( + "http://" + ) or self._image_reference.startswith("https://"): + hwc_bgr = load_image_from_url(value=self._image_reference) + return torch.from_numpy(hwc_bgr[:, :, ::-1].copy()).permute(2, 0, 1) + return decode_image( + read_file(self._image_reference), + mode=ImageReadMode.RGB, + apply_exif_orientation=True, + ) + + def _ensure_original_sources_usable(self) -> None: + """Guard for the properties' reference-decoding fallback: once an + in-place mutation was declared, the original file/URL reference no + longer describes the pixels, so loading from it would silently + resurrect the pre-mutation image. A cached base64 is deliberately NOT + covered: the declare nulls the original one, so any base64 present + afterwards was re-derived from the mutated pixels by ``base64_image`` + and is valid. Reaching this with the flag set means every valid + channel (mutated representation, re-derived base64) is gone - an + internal-consistency bug worth failing loudly on, with the cause + named.""" + if self._sources_invalidated_by_mutation: + raise ValueError( + "Cannot load this WorkflowImageData from its original source " + "reference: an in-place mutation was declared " + "(declare_numpy_image_mutated / declare_tensor_image_mutated), " + "which invalidated the original source, and neither the mutated " + "representation nor a re-derived base64 is materialised any " + "more. This state is induced by the image mutation - the " + "pre-mutation source must not be decoded as it no longer " + "describes the image." + ) + + def declare_numpy_image_mutated(self) -> None: + """A client that mutated the ``numpy_image`` buffer in place declares it + here: the mutated numpy becomes the SOLE source of truth. Every other + channel is cut off - the derived siblings (tensor / base64) are removed + so later readers re-derive from the mutated pixels, and the original + source reference is removed too, because it no longer describes these + pixels (serialization such as ``to_inference_format`` must never hand + out the pre-mutation source). Declaring a mutation of a representation + that was never materialised is a caller bug and raises.""" + if self._numpy_image is None: + raise ValueError( + "declare_numpy_image_mutated() called, but no numpy representation " + "is materialised - nothing could have been mutated. Access " + "`numpy_image` before declaring." + ) + self._tensor_image = None + self._base64_image = None + self._image_reference = None + self._sources_invalidated_by_mutation = True + + def declare_tensor_image_mutated(self) -> None: + """A client that mutated the ``tensor_image`` in place declares it here: + the mutated tensor becomes the SOLE source of truth. Every other channel + is cut off - the derived siblings (numpy / base64) are removed so later + readers re-derive from the mutated pixels, and the original source + reference is removed too, because it no longer describes these pixels. + Declaring a mutation of a representation that was never materialised is + a caller bug and raises.""" + if self._tensor_image is None: + raise ValueError( + "declare_tensor_image_mutated() called, but no tensor representation " + "is materialised - nothing could have been mutated. Access " + "`tensor_image` before declaring." + ) + self._numpy_image = None + self._base64_image = None + self._image_reference = None + self._sources_invalidated_by_mutation = True + + def is_tensor_materialised(self) -> bool: + """Whether the CHW RGB tensor image already exists on device. + + ``True`` only when a tensor representation is already present (the caller fed + one in, or something downstream already built it) โ€” reading ``tensor_image`` is + then free. ``False`` means only the numpy (HWC BGR) representation is available, + so accessing ``tensor_image`` would trigger an eager numpy->device conversion. + + Blocks use this to pick the representation that is already materialised instead + of forcing a conversion: ``tensor_image`` (RGB) when ``True``, ``numpy_image`` + (BGR) otherwise. The check itself does no I/O and never materialises anything. + """ + return self._tensor_image is not None + @property def base64_image(self) -> str: if self._base64_image is not None: diff --git a/inference/core/workflows/execution_engine/entities/tensor_native_types.py b/inference/core/workflows/execution_engine/entities/tensor_native_types.py new file mode 100644 index 0000000000..ac79fc489e --- /dev/null +++ b/inference/core/workflows/execution_engine/entities/tensor_native_types.py @@ -0,0 +1,742 @@ +from inference.core.workflows.execution_engine.entities.types import Kind + +TENSOR_NATIVE_EMBEDDING_KIND_DOCS = """ +This kind represents a vector embedding. It is a list of floating point numbers. + +Embeddings are used in various machine learning tasks like clustering, classification, +and similarity search. They are used to represent data in a continuous, low-dimensional space. + +Typically, vectors that are close to each other in the embedding space are considered similar. +""" +TENSOR_NATIVE_EMBEDDING_KIND = Kind( + name="embedding", + description="A list of floating point numbers representing a vector embedding.", + docs=TENSOR_NATIVE_EMBEDDING_KIND_DOCS, + serialised_data_type="List[float]", + internal_data_type="torch.Tensor", +) + + +TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND_DOCS = """ +This kind represent predictions from Classification Models. Internal representation is +based on `inference_models.ClassificationPrediction` and `inference_models.MultiLabelClassificationPrediction` +respectivelly. Below one can find external - serialised representation. + +Examples: +``` +# in case of multi-class classification +{ + "image": {"height": 128, "width": 256}, + "predictions": [{"class_name": "A", "class_id": 0, "confidence": 0.3}], + "top": "A", + "confidence": 0.3, + "parent_id": "some", + "prediction_type": "classification", + "inference_id": "some", + "root_parent_id": "some", +} + +# in case of multi-label classification +{ + "image": {"height": 128, "width": 256}, + "predictions": { + "a": {"confidence": 0.3, "class_id": 0}, + "b": {"confidence": 0.3, "class_id": 1}, + } + "predicted_classes": ["a", "b"], + "parent_id": "some", + "prediction_type": "classification", + "inference_id": "some", + "root_parent_id": "some", +} +``` +""" +TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND = Kind( + name="classification_prediction", + description="Predictions from classifier", + docs=TENSOR_NATIVE_CLASSIFICATION_PREDICTION_KIND_DOCS, + serialised_data_type="dict", + internal_data_type="Union[inference_models.ClassificationPrediction, inference_models.MultiLabelClassificationPrediction]", +) + +TENSOR_NATIVE_DETECTION_KIND_DOCS = """ +This kind represents single detection in prediction from a model that detects multiple elements +(like object detection or instance segmentation model). It is represented as a tuple +that is created from `inference_models.Detections(...)` object while iterating over its content. `workflows` +utilises `bboxes_metadata` as well as `image_metadata` property of `inference_models.Detections(...)` to keep +additional metadata which will be available in the tuple. Some properties may not always be present. Take a look at +documentation of `object_detection_prediction`, `instance_segmentation_prediction`, `keypoint_detection_prediction` +kinds to discover which additional metadata are available. + +**Per-box metadata keys** (read from `bboxes_metadata[i]`, the entry for the i-th box): + +* `class` - per-box class label string. When present it overrides the class name resolved from +`image_metadata[CLASS_NAMES_KEY]` keyed by `class_id` (so producers may carry an arbitrary +per-box label, e.g. VLM-as-detector labels or recognised OCR text). When absent, the class name +is resolved from `image_metadata[CLASS_NAMES_KEY][class_id]`. + +* `tracker_id` - per-box tracker identifier (present once a tracking block has run). + +* `detection_id` - per-box unique identifier (see `object_detection_prediction` for details). + +* `text` - per-box recognised text (used by OCR producers); note that the serialised +representation surfaces recognised text under the `class` key, not `text`. +""" + +TENSOR_NATIVE_DETECTION_KIND = Kind( + name="detection", + description="Single element of detections-based prediction (like `object_detection_prediction`)", + docs=TENSOR_NATIVE_DETECTION_KIND_DOCS, + serialised_data_type="Tuple[list, Optional[list], Optional[float], Optional[float], Optional[int], dict, dict]", + internal_data_type="Tuple[torch.Tensor, Optional[torch.Tensor], Optional[float], Optional[float], Optional[int], dict, dict]", +) + + +TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND_DOCS = """ +This kind represents single object detection prediction in form of `inference_models.Detections` object. + +Example: +``` +inference_models.Detections( + xyxy=torch.Tensor([ + [ 865, 153.5, 1189, 422.5], + [ 192.5, 77.5, 995.5, 722.5], + [ 194, 82, 996, 726], + [ 460, 333, 704, 389]] + ), + confidence=torch.Tensor([ 0.84955, 0.74344, 0.45636, 0.86537]), + class_id=torch.Tensor([2, 7, 2, 0]), + tracker_id=None, + image_metadata={ + "class_names": {0: "car", 1: "truck"}, + "image_dimensions": [425, 640], + "parent_id": "image.[0]", + "inference_id": "51dfa8d5-261c-4dcb-ab30-9aafe9b52379", + "prediction_type": "object-detection", + "root_parent_id": "image.[0]", + "root_parent_coordinates": [0, 0], + "root_parent_dimensions": [425, 640], + "parent_coordinates": [0, 0], + "parent_dimensions": [425, 640], + "scaling_relative_to_parent": 1.0, + "scaling_relative_to_root_parent": 1.0, + } + bboxes_metadata=[ + {'detection_id': '51dfa8d5-261c-4dcb-ab30-9aafe9b52379'}, + {'detection_id': 'c0c684d1-1e30-4880-aedd-29e67e417264'}, + {'detection_id': '8cfc543b-9cfe-493b-b5ad-77afed7bee83'}, + {'detection_id': 'c0c684d1-1e30-4880-aedd-38e67e441454'}, + ], +) +``` +**Details of additional fields:** + +* `detection_id` - unique identifier for each detection, to be used for when dependent elements +are created based on specific detection (example: Dynamic Crop takes this value as parent id for new image) + +* `parent_id` - identifier of image that generated prediction (to be fetched from `WorkflowImageData` object) + +* `image_dimensions` - dimensions of image that was basis for prediction - format: `(height, width)` + +* `inference_id` - identifier of inference request (optional, relevant for Roboflow models) + +* `prediction_type` - type of prediction + +* `root_parent_id` - identifier of primary Workflow input that was responsible for downstream prediction +(to be fetched from `WorkflowImageData` object) - usually identifier of Workflow input placeholder + +* `root_parent_coordinates` - offset regarding origin input - format (`offset_x`, `offset_y`) + +* `root_parent_dimensions` - dimensions of origin input image `(height, width)` + +* `parent_coordinates` - offset regarding parent - format (`offset_x`, `offset_y`) + +* `parent_dimensions` - dimensions of parent image `(height, width)` + +* `scaling_relative_to_parent` - scaling factor regarding parent image + +* `scaling_relative_to_root_parent` - scaling factor regarding origin input image + + +**SERIALISATION:** + +Execution Engine behind API will serialise underlying data once selector of this kind is declared as +Workflow output - serialisation will be executed such that `sv.Detections.from_inference(...)` +can decode the output. Entity details: [ObjectDetectionInferenceResponse](https://detect.roboflow.com/docs) +""" +TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND = Kind( + name="object_detection_prediction", + description="Prediction with detected bounding boxes in form of inference_models.Detections(...) object", + docs=TENSOR_NATIVE_OBJECT_DETECTION_PREDICTION_KIND_DOCS, + serialised_data_type="dict", + internal_data_type="inference_models.Detections", +) + + +TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND_DOCS = """ +This kind represents single object detection prediction in form of `inference_models.InstanceDetections` object. + +Example: +``` +inference_models.InstanceDetections( + xyxy=torch.Tensor([ + [ 865, 153.5, 1189, 422.5], + [ 192.5, 77.5, 995.5, 722.5], + [ 194, 82, 996, 726], + [ 460, 333, 704, 389]] + ), + mask=InstancesRLEMasks(...), + confidence=torch.Tensor([ 0.84955, 0.74344, 0.45636, 0.86537]), + class_id=torch.Tensor([2, 7, 2, 0]), + tracker_id=None, + image_metadata={ + "class_names": {0: "car", 1: "truck"}, + "image_dimensions": [425, 640], + "parent_id": "image.[0]", + "inference_id": "51dfa8d5-261c-4dcb-ab30-9aafe9b52379", + "prediction_type": "instance-segmentation", + "root_parent_id": "image.[0]", + "root_parent_coordinates": [0, 0], + "root_parent_dimensions": [425, 640], + "parent_coordinates": [0, 0], + "parent_dimensions": [425, 640], + "scaling_relative_to_parent": 1.0, + "scaling_relative_to_root_parent": 1.0, + } + bboxes_metadata=[ + {'detection_id': '51dfa8d5-261c-4dcb-ab30-9aafe9b52379'}, + {'detection_id': 'c0c684d1-1e30-4880-aedd-29e67e417264'}, + {'detection_id': '8cfc543b-9cfe-493b-b5ad-77afed7bee83'}, + {'detection_id': 'c0c684d1-1e30-4880-aedd-38e67e441454'}, + ], +) +``` +**Details of additional fields:** + +* `detection_id` - unique identifier for each detection, to be used for when dependent elements +are created based on specific detection (example: Dynamic Crop takes this value as parent id for new image) + +* `parent_id` - identifier of image that generated prediction (to be fetched from `WorkflowImageData` object) + +* `image_dimensions` - dimensions of image that was basis for prediction - format: `(height, width)` + +* `inference_id` - identifier of inference request (optional, relevant for Roboflow models) + +* `prediction_type` - type of prediction + +* `root_parent_id` - identifier of primary Workflow input that was responsible for downstream prediction +(to be fetched from `WorkflowImageData` object) - usually identifier of Workflow input placeholder + +* `root_parent_coordinates` - offset regarding origin input - format (`offset_x`, `offset_y`) + +* `root_parent_dimensions` - dimensions of origin input image `(height, width)` + +* `parent_coordinates` - offset regarding parent - format (`offset_x`, `offset_y`) + +* `parent_dimensions` - dimensions of parent image `(height, width)` + +* `scaling_relative_to_parent` - scaling factor regarding parent image + +* `scaling_relative_to_root_parent` - scaling factor regarding origin input image + +**SERIALISATION:** + +Execution Engine behind API will serialise underlying data once selector of this kind is declared as +Workflow output - serialisation will be executed such that `sv.Detections.from_inference(...)` +can decode the output (up to RLE representation). Entity details: [InstanceSegmentationInferenceResponse](https://detect.roboflow.com/docs) +""" +TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND = Kind( + name="instance_segmentation_prediction", + description="Prediction with detected bounding boxes and segmentation masks in form of inference_models.InstanceDetections(...) object", + docs=TENSOR_NATIVE_INSTANCE_SEGMENTATION_PREDICTION_KIND_DOCS, + serialised_data_type="dict", + internal_data_type="inference_models.InstanceDetections", +) + + +TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND_DOCS = """ +This kind represents single instance segmentation prediction in form of an +`inference_models.InstanceDetections` object whose masks are RLE-encoded +(`mask` is an `inference_models.models.base.types.InstancesRLEMasks`, rather +than a dense `torch.Tensor`). It is otherwise identical to +`instance_segmentation_prediction` - the same `image_metadata` and +`bboxes_metadata` conventions apply. + +Example: +``` +inference_models.InstanceDetections( + xyxy=torch.Tensor([ + [ 865, 153.5, 1189, 422.5], + [ 192.5, 77.5, 995.5, 722.5]] + ), + mask=InstancesRLEMasks(image_size=(425, 640), masks=[...]), + confidence=torch.Tensor([ 0.84955, 0.74344]), + class_id=torch.Tensor([2, 7]), + image_metadata={ + "class_names": {0: "car", 1: "truck"}, + "image_dimensions": [425, 640], + "parent_id": "image.[0]", + "inference_id": "51dfa8d5-261c-4dcb-ab30-9aafe9b52379", + "prediction_type": "rle-instance-segmentation", + "root_parent_id": "image.[0]", + "root_parent_coordinates": [0, 0], + "root_parent_dimensions": [425, 640], + "parent_coordinates": [0, 0], + "parent_dimensions": [425, 640], + "scaling_relative_to_parent": 1.0, + "scaling_relative_to_root_parent": 1.0, + } + bboxes_metadata={ + 'detection_id': [ + '51dfa8d5-261c-4dcb-ab30-9aafe9b52379', 'c0c684d1-1e30-4880-aedd-29e67e417264' + ], + } +) +``` + +The additional `image_metadata` / `bboxes_metadata` fields carry the same +meaning as documented for `instance_segmentation_prediction`. + +**SERIALISATION:** + +Execution Engine behind API will serialise underlying data once selector of this kind is declared as +Workflow output - serialisation preserves the RLE mask representation. Entity +details: [InstanceSegmentationInferenceResponse](https://detect.roboflow.com/docs) +""" +TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND = Kind( + name="rle_instance_segmentation_prediction", + description="Prediction with detected bounding boxes and RLE-encoded segmentation masks in form of inference_models.InstanceDetections(...) object", + docs=TENSOR_NATIVE_RLE_INSTANCE_SEGMENTATION_PREDICTION_KIND_DOCS, + serialised_data_type="dict", + internal_data_type="inference_models.InstanceDetections", +) + + +TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND_DOCS = """ +This kind represents single object detection prediction in form of Tuple[`inference_models.KeyPoints`, Optional[`inference_models.Detections`]] object. + + +Example: +``` +( + inference_models.KeyPoints( + xy=torch.Tensor([ + [[10, 20], [10, 20], [10, 20], [10, 20]], + [[10, 20], [10, 20], [10, 20], [10, 20]], + [[10, 20], [10, 20], [10, 20], [10, 20]], + [[10, 20], [10, 20], [10, 20], [10, 20]], + ]) + confidence=torch.Tensor([ + [ 0.84955, 0.74344, 0.45636, 0.86537], + [ 0.84955, 0.74344, 0.45636, 0.86537], + [ 0.84955, 0.74344, 0.45636, 0.86537], + [ 0.84955, 0.74344, 0.45636, 0.86537] + ]), + class_id=torch.Tensor([2, 7, 2, 0]), + image_metadata={ + "keypoints_class_names": {0: "car", 1: "truck"}, + "image_dimensions": [425, 640], + "parent_id": "image.[0]", + "inference_id": "51dfa8d5-261c-4dcb-ab30-9aafe9b52379", + "prediction_type": "object-detection", + "root_parent_id": "image.[0]", + "root_parent_coordinates": [0, 0], + "root_parent_dimensions": [425, 640], + "parent_coordinates": [0, 0], + "parent_dimensions": [425, 640], + "scaling_relative_to_parent": 1.0, + "scaling_relative_to_root_parent": 1.0, + }, + ) + inference_models.Detections( + xyxy=torch.Tensor([ + [ 865, 153.5, 1189, 422.5], + [ 192.5, 77.5, 995.5, 722.5], + [ 194, 82, 996, 726], + [ 460, 333, 704, 389]] + ), + confidence=torch.Tensor([ 0.84955, 0.74344, 0.45636, 0.86537]), + class_id=torch.Tensor([2, 7, 2, 0]), + tracker_id=None, + image_metadata={ + "class_names": {0: "car", 1: "truck"}, + "image_dimensions": [425, 640], + "parent_id": "image.[0]", + "inference_id": "51dfa8d5-261c-4dcb-ab30-9aafe9b52379", + "prediction_type": "object-detection", + "root_parent_id": "image.[0]", + "root_parent_coordinates": [0, 0], + "root_parent_dimensions": [425, 640], + "parent_coordinates": [0, 0], + "parent_dimensions": [425, 640], + "scaling_relative_to_parent": 1.0, + "scaling_relative_to_root_parent": 1.0, + } + key_points_metadata={ + 'detection_id': [ + '51dfa8d5-261c-4dcb-ab30-9aafe9b52379', 'c0c684d1-1e30-4880-aedd-29e67e417264' + '8cfc543b-9cfe-493b-b5ad-77afed7bee83', 'c0c684d1-1e30-4880-aedd-38e67e441454' + ], + } + ) +) + +``` + +Prior to [sv.Keypoints(...)](https://supervision.roboflow.com/0.21.0/keypoint/core/) we introduced +keypoints detection based on [`sv.Detections(...)`](https://supervision.roboflow.com/latest/detection/core/) object. +The decision was suboptimal so we would need to revert in the future, but for now this is the format of +data for keypoints detection. + +The design of metadata is also suboptimal (as metadata regarding whole image is duplicated across all +bounding boxes and there is no way on how to save metadata for empty predictions). We +have [GH issue](https://github.com/roboflow/inference/issues/567) to communicate around this +problem. + +**Details of additional fields:** + +* `detection_id` - unique identifier for each detection, to be used for when dependent elements +are created based on specific detection (example: Dynamic Crop takes this value as parent id for new image) + +* `parent_id` - identifier of image that generated prediction (to be fetched from `WorkflowImageData` object) + +* `image_dimensions` - dimensions of image that was basis for prediction - format: `(height, width)` + +* `inference_id` - identifier of inference request (optional, relevant for Roboflow models) + +* `prediction_type` - type of prediction + +* `root_parent_id` - identifier of primary Workflow input that was responsible for downstream prediction +(to be fetched from `WorkflowImageData` object) - usually identifier of Workflow input placeholder + +* `root_parent_coordinates` - offset regarding origin input - format (`offset_x`, `offset_y`) + +* `root_parent_dimensions` - dimensions of origin input image `(height, width)` + +* `parent_coordinates` - offset regarding parent - format (`offset_x`, `offset_y`) + +* `parent_dimensions` - dimensions of parent image `(height, width)` + +* `scaling_relative_to_parent` - scaling factor regarding parent image + +* `scaling_relative_to_root_parent` - scaling factor regarding origin input image + +* `keypoints_class_name` array of variable size 1D arrays of string with key points class names + +* `keypoints_class_id` array of variable size 1D arrays of int with key points class ids + +* `keypoints_confidence` array of variable size 1D arrays of float with key points confidence + +* `keypoints_xy` array of variable size 2D arrays of coordinates of keypoints in `(x, y)` format + +**SERIALISATION:** + +Execution Engine behind API will serialise underlying data once selector of this kind is declared as +Workflow output - serialisation will be executed such that `sv.Detections.from_inference(...)` +can decode the output, but **loosing keypoints details** - which can be recovered if output +JSON field is parsed. Entity details: [KeypointsDetectionInferenceResponse](https://detect.roboflow.com/docs) +""" +TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND = Kind( + name="keypoint_detection_prediction", + description="Prediction with detected bounding boxes and detected keypoints in form of a `(inference_models.KeyPoints, Optional[inference_models.Detections])` tuple", + docs=TENSOR_NATIVE_KEYPOINT_DETECTION_PREDICTION_KIND_DOCS, + serialised_data_type="dict", + internal_data_type="Tuple[inference_models.KeyPoints, Optional[inference_models.Detections]]", +) + + +TENSOR_NATIVE_QR_CODE_DETECTION_KIND_DOCS = """ +This kind represents batch of predictions regarding QR codes location and data their provide. + +Example: +``` +inference_models.Detections( + xyxy=torch.Tensor([ + [ 865, 153.5, 1189, 422.5], + [ 192.5, 77.5, 995.5, 722.5], + [ 194, 82, 996, 726], + [ 460, 333, 704, 389]] + ), + confidence=torch.Tensor([ 0.84955, 0.74344, 0.45636, 0.86537]), + class_id=torch.Tensor([2, 7, 2, 0]), + tracker_id=None, + image_metadata={ + "class_names": {0: "qr-code"}, + "image_dimensions": [425, 640], + "parent_id": "image.[0]", + "inference_id": "51dfa8d5-261c-4dcb-ab30-9aafe9b52379", + "prediction_type": "qrcode-detection", + "root_parent_id": "image.[0]", + "root_parent_coordinates": [0, 0], + "root_parent_dimensions": [425, 640], + "parent_coordinates": [0, 0], + "parent_dimensions": [425, 640], + "scaling_relative_to_parent": 1.0, + "scaling_relative_to_root_parent": 1.0, + } + bboxes_metadata={ + 'detection_id': [ + '51dfa8d5-261c-4dcb-ab30-9aafe9b52379', 'c0c684d1-1e30-4880-aedd-29e67e417264' + '8cfc543b-9cfe-493b-b5ad-77afed7bee83', 'c0c684d1-1e30-4880-aedd-38e67e441454' + ], + "qr_codes": ["a", "b", "c", "d"] + } +) +``` + +As you can see, we have extended the standard set of metadata for predictions maintained by `supervision`. +Adding this metadata is needed to ensure compatibility with blocks from `roboflow_core` plugin. + +The design of metadata is suboptimal (as metadata regarding whole image is duplicated across all +bounding boxes and there is no way on how to save metadata for empty predictions). We +have [GH issue](https://github.com/roboflow/inference/issues/567) to communicate around this +problem. + +**Details of additional fields:** + +* `detection_id` - unique identifier for each detection, to be used for when dependent elements +are created based on specific detection (example: Dynamic Crop takes this value as parent id for new image) + +* `parent_id` - identifier of image that generated prediction (to be fetched from `WorkflowImageData` object) + +* `image_dimensions` - dimensions of image that was basis for prediction - format: `(height, width)` + +* `inference_id` - identifier of inference request (optional, relevant for Roboflow models) + +* `prediction_type` - type of prediction + +* `root_parent_id` - identifier of primary Workflow input that was responsible for downstream prediction +(to be fetched from `WorkflowImageData` object) - usually identifier of Workflow input placeholder + +* `root_parent_coordinates` - offset regarding origin input - format (`offset_x`, `offset_y`) + +* `root_parent_dimensions` - dimensions of origin input image `(height, width)` + +* `parent_coordinates` - offset regarding parent - format (`offset_x`, `offset_y`) + +* `parent_dimensions` - dimensions of parent image `(height, width)` + +* `scaling_relative_to_parent` - scaling factor regarding parent image + +* `scaling_relative_to_root_parent` - scaling factor regarding origin input image + +* `qr_codes` - extracted QR code + +**SERIALISATION:** +Execution Engine behind API will serialise underlying data once selector of this kind is declared as +Workflow output - serialisation will be executed such that `sv.Detections.from_inference(...)` +can decode the output. Entity details: [ObjectDetectionInferenceResponse](https://detect.roboflow.com/docs) +""" +TENSOR_NATIVE_QR_CODE_DETECTION_KIND = Kind( + name="qr_code_detection", + description="Prediction with QR code detection", + docs=TENSOR_NATIVE_QR_CODE_DETECTION_KIND_DOCS, + serialised_data_type="dict", + internal_data_type="inference_models.Detections", +) + + +TENSOR_NATIVE_BAR_CODE_DETECTION_KIND_DOCS = """ +This kind represents batch of predictions regarding barcodes location and data their provide. + +Example: +``` +inference_models.Detections( + xyxy=torch.Tensor([ + [ 865, 153.5, 1189, 422.5], + [ 192.5, 77.5, 995.5, 722.5], + [ 194, 82, 996, 726], + [ 460, 333, 704, 389]] + ), + confidence=torch.Tensor([ 0.84955, 0.74344, 0.45636, 0.86537]), + class_id=torch.Tensor([2, 7, 2, 0]), + tracker_id=None, + image_metadata={ + "class_names": {0: "barcode"}, + "image_dimensions": [425, 640], + "parent_id": "image.[0]", + "inference_id": "51dfa8d5-261c-4dcb-ab30-9aafe9b52379", + "prediction_type": "barcode-detection", + "root_parent_id": "image.[0]", + "root_parent_coordinates": [0, 0], + "root_parent_dimensions": [425, 640], + "parent_coordinates": [0, 0], + "parent_dimensions": [425, 640], + "scaling_relative_to_parent": 1.0, + "scaling_relative_to_root_parent": 1.0, + } + bboxes_metadata={ + 'detection_id': [ + '51dfa8d5-261c-4dcb-ab30-9aafe9b52379', 'c0c684d1-1e30-4880-aedd-29e67e417264' + '8cfc543b-9cfe-493b-b5ad-77afed7bee83', 'c0c684d1-1e30-4880-aedd-38e67e441454' + ], + "barcodes": ["a", "b", "c", "d"] + } +) +``` + +As you can see, we have extended the standard set of metadata for predictions maintained by `supervision`. +Adding this metadata is needed to ensure compatibility with blocks from `roboflow_core` plugin. + +The design of metadata is suboptimal (as metadata regarding whole image is duplicated across all +bounding boxes and there is no way on how to save metadata for empty predictions). We +have [GH issue](https://github.com/roboflow/inference/issues/567) to communicate around this +problem. + +**Details of additional fields:** + +* `detection_id` - unique identifier for each detection, to be used for when dependent elements +are created based on specific detection (example: Dynamic Crop takes this value as parent id for new image) + +* `parent_id` - identifier of image that generated prediction (to be fetched from `WorkflowImageData` object) + +* `image_dimensions` - dimensions of image that was basis for prediction - format: `(height, width)` + +* `inference_id` - identifier of inference request (optional, relevant for Roboflow models) + +* `prediction_type` - type of prediction + +* `root_parent_id` - identifier of primary Workflow input that was responsible for downstream prediction +(to be fetched from `WorkflowImageData` object) - usually identifier of Workflow input placeholder + +* `root_parent_coordinates` - offset regarding origin input - format (`offset_x`, `offset_y`) + +* `root_parent_dimensions` - dimensions of origin input image `(height, width)` + +* `parent_coordinates` - offset regarding parent - format (`offset_x`, `offset_y`) + +* `parent_dimensions` - dimensions of parent image `(height, width)` + +* `scaling_relative_to_parent` - scaling factor regarding parent image + +* `scaling_relative_to_root_parent` - scaling factor regarding origin input image + +* `barcodes` - extracted barcode + +**SERIALISATION:** +Execution Engine behind API will serialise underlying data once selector of this kind is declared as +Workflow output - serialisation will be executed such that `sv.Detections.from_inference(...)` +can decode the output. Entity details: [ObjectDetectionInferenceResponse](https://detect.roboflow.com/docs) +""" +TENSOR_NATIVE_BAR_CODE_DETECTION_KIND = Kind( + name="bar_code_detection", + description="Prediction with barcode detection", + docs=TENSOR_NATIVE_BAR_CODE_DETECTION_KIND_DOCS, + serialised_data_type="dict", + internal_data_type="inference_models.Detections", +) + + +TENSOR_NATIVE_SEMANTIC_SEGMENTATION_PREDICTION_KIND_DOCS = """ +This kind represents a single semantic segmentation prediction in form of an +`inference_models.InstanceDetections` object with one detection per predicted class. Each +detection carries an RLE-encoded mask covering all pixels assigned to that class. + +**Why RLE and not polygons:** + +Semantic segmentation assigns a class label to every pixel in the image. A single class +can appear in multiple spatially disconnected regions (e.g., two separate "person" regions +on opposite sides of the frame). Polygon-based serialization uses `cv2.findContours()`, +which only retains the first contiguous contour and silently discards all others โ€” causing +irreversible data loss for non-contiguous masks. RLE (Run-Length Encoding, COCO standard) +is a pixel-level encoding that represents the complete mask regardless of spatial topology, +making it the only correct serialization format for semantic segmentation masks. + +**Internal representation:** an `inference_models.InstanceDetections` carrier (the same +per-class RLE carrier used by `rle_instance_segmentation_prediction`), with one row per +predicted class: +- `xyxy` โ€” `torch.Tensor` of tight bounding boxes enclosing all pixels of each class +- `mask` โ€” `inference_models.models.base.types.InstancesRLEMasks` holding one COCO RLE + entry per class (NOT a dense `torch.Tensor`, NO polygon collapse) +- `class_id` โ€” `torch.Tensor` of integer class IDs +- `confidence` โ€” `torch.Tensor` of mean confidence over all pixels of each class +- `image_metadata[CLASS_NAMES_KEY]` โ€” class-id -> class-name mapping for the prediction +- `bboxes_metadata` โ€” same per-box metadata conventions as `instance_segmentation_prediction` + (e.g. `detection_id`); see that kind for the full key list. + +Example: +``` +inference_models.InstanceDetections( + xyxy=torch.Tensor([ + [ 0, 0, 640, 480], + [ 100, 120, 400, 360]] + ), + mask=InstancesRLEMasks(image_size=(480, 640), masks=[...]), + confidence=torch.Tensor([ 0.97, 0.92]), + class_id=torch.Tensor([0, 1]), + image_metadata={ + "class_names": {0: "background", 1: "person"}, + "image_dimensions": [480, 640], + "parent_id": "image.[0]", + "inference_id": "51dfa8d5-261c-4dcb-ab30-9aafe9b52379", + "prediction_type": "semantic-segmentation", + "root_parent_id": "image.[0]", + "root_parent_coordinates": [0, 0], + "root_parent_dimensions": [480, 640], + "parent_coordinates": [0, 0], + "parent_dimensions": [480, 640], + "scaling_relative_to_parent": 1.0, + "scaling_relative_to_root_parent": 1.0, + } + bboxes_metadata={ + 'detection_id': [ + '51dfa8d5-261c-4dcb-ab30-9aafe9b52379', 'c0c684d1-1e30-4880-aedd-29e67e417264' + ], + } +) +``` + +**Serialised format** (one entry per class in `predictions`): + +```json +{ + "image": {"width": 640, "height": 480}, + "predictions": [ + { + "x": 320.0, "y": 240.0, "width": 200.0, "height": 180.0, + "confidence": 0.92, + "class_id": 1, + "class": "person", + "detection_id": "a1b2c3d4-...", + "rle_mask": {"size": [480, 640], "counts": "XYZ..."} + } + ] +} +``` + +**Decoding RLE masks:** + +```python +import pycocotools.mask as mask_utils +import numpy as np + +rle = prediction["rle_mask"] +binary_mask = mask_utils.decode(rle).astype(bool) # shape: (H, W) +``` +""" +TENSOR_NATIVE_SEMANTIC_SEGMENTATION_PREDICTION_KIND = Kind( + name="semantic_segmentation_prediction", + description="Prediction with per-pixel class label and confidence for semantic segmentation", + docs=TENSOR_NATIVE_SEMANTIC_SEGMENTATION_PREDICTION_KIND_DOCS, + serialised_data_type="dict", + internal_data_type="inference_models.InstanceDetections", +) + + +TENSOR_KIND_DOCS = """ +This kind represents a raw, multi-dimensional numeric tensor kept on-device as a +`torch.Tensor`. It is the tensor-native counterpart of `numpy_array`, but it is a +standalone kind with its own name (`tensor`) - it does NOT replace or alias +`numpy_array`; the two coexist. + +It is used by blocks that emit dense numeric maps (for example, the depth map +produced by Depth Estimation) and want to keep the data on the accelerator, +avoiding an eager GPU -> CPU / numpy round-trip. +""" +TENSOR_KIND = Kind( + name="tensor", + description="A raw multi-dimensional numeric tensor (torch.Tensor)", + docs=TENSOR_KIND_DOCS, + serialised_data_type="list", + internal_data_type="torch.Tensor", +) diff --git a/inference/core/workflows/execution_engine/v1/dynamic_blocks/block_assembler.py b/inference/core/workflows/execution_engine/v1/dynamic_blocks/block_assembler.py index a336efe229..6afd02854f 100644 --- a/inference/core/workflows/execution_engine/v1/dynamic_blocks/block_assembler.py +++ b/inference/core/workflows/execution_engine/v1/dynamic_blocks/block_assembler.py @@ -6,6 +6,7 @@ from inference.core.env import ( ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS, + ENABLE_TENSOR_DATA_REPRESENTATION, WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE, ) from inference.core.workflows.errors import ( @@ -46,6 +47,7 @@ DynamicOutputDefinition, ManifestDescription, SelectorType, + TensorCompatibility, ValueType, ) from inference.core.workflows.prototypes.block import WorkflowBlockManifest @@ -114,6 +116,9 @@ def create_dynamic_block_specification( api_key: Optional[str] = None, skip_class_eval: Optional[bool] = False, ) -> BlockSpecification: + ensure_tensor_compatibility_supported( + manifest_description=dynamic_block_definition.manifest, + ) unique_identifier = str(uuid4()) block_manifest = assembly_dynamic_block_manifest( unique_identifier=unique_identifier, @@ -127,6 +132,7 @@ def create_dynamic_block_specification( python_code=dynamic_block_definition.code, api_key=api_key, skip_class_eval=skip_class_eval, + manifest_description=dynamic_block_definition.manifest, ) return BlockSpecification( block_source=BLOCK_SOURCE, @@ -136,6 +142,38 @@ def create_dynamic_block_specification( ) +def ensure_tensor_compatibility_supported( + manifest_description: ManifestDescription, +) -> None: + """D4 of the tensor_compatibility plan: `tensor_native` blocks fail fast at + compile time when the server cannot honor the declared contract โ€” their + tensor-expecting user code would break mid-run anyway.""" + if ( + manifest_description.tensor_compatibility + is not TensorCompatibility.TENSOR_NATIVE + ): + return + # Deliberate precedence: with flag-off AND modal both misconfigured, the flag + # error fires alone โ€” enabling the flag is the prerequisite that makes the + # modal limitation relevant at all. + if not ENABLE_TENSOR_DATA_REPRESENTATION: + raise DynamicBlockError( + public_message=f"Dynamic block `{manifest_description.block_type}` declares " + f"`tensor_compatibility=tensor_native`, but this server runs the numpy data " + f"representation. Use `legacy_compatibility` (the default) or enable " + f"`ENABLE_TENSOR_DATA_REPRESENTATION` on the server.", + context="workflow_compilation | dynamic_blocks_compilation", + ) + if WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE == "modal": + raise DynamicBlockError( + public_message=f"Dynamic block `{manifest_description.block_type}` declares " + f"`tensor_compatibility=tensor_native`, which is not yet supported for remote " + f"(modal) custom Python execution. Use `legacy_compatibility` or run custom " + f"Python code locally.", + context="workflow_compilation | dynamic_blocks_compilation", + ) + + def assembly_dynamic_block_manifest( unique_identifier: str, manifest_description: ManifestDescription, diff --git a/inference/core/workflows/execution_engine/v1/dynamic_blocks/block_scaffolding.py b/inference/core/workflows/execution_engine/v1/dynamic_blocks/block_scaffolding.py index b1413a7984..a043836968 100644 --- a/inference/core/workflows/execution_engine/v1/dynamic_blocks/block_scaffolding.py +++ b/inference/core/workflows/execution_engine/v1/dynamic_blocks/block_scaffolding.py @@ -7,6 +7,7 @@ from inference.core.env import ( ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS, + ENABLE_TENSOR_DATA_REPRESENTATION, MODAL_ANONYMOUS_WORKSPACE_NAME, WEBEXEC_MODAL_EXECUTOR_IDLE_TTL_SECONDS, WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE, @@ -23,6 +24,7 @@ get_active_collector, ) from inference.core.workflows.execution_engine.v1.dynamic_blocks.entities import ( + ManifestDescription, PythonCode, ) from inference.core.workflows.execution_engine.v1.dynamic_blocks.error_utils import ( @@ -30,6 +32,12 @@ create_dynamic_block_code_error, extract_code_snippet, ) +from inference.core.workflows.execution_engine.v1.dynamic_blocks.representation_boundary import ( + collect_declared_input_kind_names, + collect_declared_output_kind_names, + convert_block_result_to_native, + convert_kwargs_to_legacy, +) from inference.core.workflows.prototypes.block import ( BlockResult, WorkflowBlock, @@ -57,6 +65,28 @@ "from inference.core.workflows.execution_engine.v1.dynamic_blocks.workflow_debug import debug_traces", ] +# Authoring surface for `tensor_compatibility=tensor_native` dynamic blocks: the +# native prediction types plus the metadata helpers that mint serializer-compliant +# outputs (`image_metadata['class_names']` + per-box `detection_id`), and the +# device native tensors pin to. Appended to IMPORTS_LINES ONLY under the flag +# (load-time, mirroring the pivot's swap philosophy) so the flag-off generated +# module source โ€” including error-line offsets derived from it โ€” stays +# byte-identical to the legacy list above. The numpy imports stay available in +# tensor mode too (user code may mix representations for its own math). +# NOTE: `modal/modal_app.py` mirrors this list into the remote sandbox namespace +# via a guarded import of this constant โ€” keep it importable and self-contained. +TENSOR_NATIVE_IMPORTS_LINES = [ + "import torch", + "from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE", + "from inference_models.models.base.object_detection import Detections", + "from inference_models.models.base.instance_segmentation import InstanceDetections", + "from inference_models.models.base.keypoints_detection import KeyPoints", + "from inference_models.models.base.classification import ClassificationPrediction, MultiLabelClassificationPrediction", + "from inference.core.workflows.core_steps.common.tensor_native import build_native_image_metadata, attach_native_detection_metadata", +] +if ENABLE_TENSOR_DATA_REPRESENTATION: + IMPORTS_LINES = IMPORTS_LINES + TENSOR_NATIVE_IMPORTS_LINES + # Shared globals dict for all custom python blocks in local mode _LOCAL_SHARED_GLOBALS = {} @@ -169,6 +199,7 @@ def assembly_custom_python_block( python_code: PythonCode, api_key: Optional[str] = None, skip_class_eval: Optional[bool] = False, + manifest_description: Optional[ManifestDescription] = None, ) -> Type[WorkflowBlock]: code_module = create_dynamic_module( @@ -187,9 +218,40 @@ def assembly_custom_python_block( ) run_function = getattr(code_module, python_code.run_function_name) + # Declared-kind lookups are static per block โ€” computed once at assembly so + # the boundary does not re-derive them per kwarg per run(). + declared_input_kinds = collect_declared_input_kind_names(manifest_description) + declared_output_kinds = collect_declared_output_kind_names(manifest_description) + def run(self, *args, **kwargs) -> BlockResult: + step_name = getattr(self, "_workflow_step_name", None) or block_type_name + # Representation boundary: under ENABLE_TENSOR_DATA_REPRESENTATION, + # `legacy_compatibility` blocks receive the documented sv/numpy + # representations instead of native tensor objects, and their returned + # legacy objects are converted back. Identity when the flag is off or + # the block declared `tensor_native`. D2-REVISED (Option A, symmetric + # sv-hop): BOTH execution arms convert on both sides โ€” Modal inputs + # ride the modal serializer's existing sv arms, and Modal results + # (rebuilt as sv by the numpy deserializers on the return leg) are + # converted to native here. Boundary calls sit OUTSIDE + # capture_output/run_function on purpose: a boundary failure is an + # engine-boundary error, not a user-code crash. Each arm converts + # AFTER its own execution gate so misconfiguration errors keep + # precedence over conversion errors. if WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE == "modal": - try: + # Remote execution via Modal - allowed even if local execution is disabled + from inference.core.workflows.execution_engine.v1.dynamic_blocks.modal_executor import ( + ModalExecutor, + ) + + kwargs = convert_kwargs_to_legacy( + kwargs=kwargs, + manifest_description=self._manifest_description, + block_name=step_name, + declared_input_kinds=declared_input_kinds, + ) + + try: # Get workspace_id from context if available workspace_id = get_roboflow_workspace(self._api_key) except WorkspaceLoadError: workspace_id = None @@ -198,13 +260,19 @@ def run(self, *args, **kwargs) -> BlockResult: workspace_id = MODAL_ANONYMOUS_WORKSPACE_NAME with _acquire_modal_executor(workspace_id) as executor: - return executor.execute_remote( + remote_result = executor.execute_remote( block_type_name=block_type_name, python_code=python_code, inputs=kwargs, workspace_id=workspace_id, workflow_context=self.get_workflow_context(), ) + return convert_block_result_to_native( + result=remote_result, + manifest_description=self._manifest_description, + block_name=step_name, + declared_output_kinds=declared_output_kinds, + ) else: # Local execution - check if allowed if not ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS: @@ -214,8 +282,13 @@ def run(self, *args, **kwargs) -> BlockResult: "`ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS=True`", context="workflow_execution | step_execution | dynamic_step", ) + kwargs = convert_kwargs_to_legacy( + kwargs=kwargs, + manifest_description=self._manifest_description, + block_name=step_name, + declared_input_kinds=declared_input_kinds, + ) import_lines_count = len(_get_python_code_imports(python_code).splitlines()) - step_name = getattr(self, "_workflow_step_name", None) or block_type_name try: with capture_output() as (stdout_buf, stderr_buf): # stdout/stderr already reach the process streams in real time via the @@ -236,7 +309,12 @@ def run(self, *args, **kwargs) -> BlockResult: block_type_name=block_type_name, ) from error _record_logs_to_active_collector(step_name, stdout_buf, stderr_buf) - return result + return convert_block_result_to_native( + result=result, + manifest_description=self._manifest_description, + block_name=step_name, + declared_output_kinds=declared_output_kinds, + ) if python_code.init_function_code is not None and not hasattr( code_module, python_code.init_function_name @@ -278,6 +356,12 @@ def get_manifest(cls) -> Type[WorkflowBlockManifest]: "get_init_parameters": get_init_parameters, "get_manifest": get_manifest, "run": run, + # AUTHORITATIVE source of the raw dynamic-block manifest description + # (carries `tensor_compatibility`): run() reads self._manifest_description, + # it is introspectable, and the Step-1 assembler tests pin it. The + # constructor parameter of assembly_custom_python_block is only the + # input that seeds this attribute (None for callers predating the knob). + "_manifest_description": manifest_description, }, ) diff --git a/inference/core/workflows/execution_engine/v1/dynamic_blocks/entities.py b/inference/core/workflows/execution_engine/v1/dynamic_blocks/entities.py index b34a45b6cb..fffa589fbb 100644 --- a/inference/core/workflows/execution_engine/v1/dynamic_blocks/entities.py +++ b/inference/core/workflows/execution_engine/v1/dynamic_blocks/entities.py @@ -22,6 +22,11 @@ class ValueType(Enum): STRING = "string" +class TensorCompatibility(str, Enum): + LEGACY_COMPATIBILITY = "legacy_compatibility" + TENSOR_NATIVE = "tensor_native" + + class DynamicInputDefinition(BaseModel): type: Literal["DynamicInputDefinition"] has_default_value: bool = Field( @@ -127,6 +132,16 @@ class ManifestDescription(BaseModel): "blocks decreasing output dimensionality which do not define neither `batch_oriented_parameters` nor " "`parameters_with_scalars_and_batches`.", ) + tensor_compatibility: TensorCompatibility = Field( + default=TensorCompatibility.LEGACY_COMPATIBILITY, + description="Representation contract of the run() function when the server runs the tensor " + "data representation (ENABLE_TENSOR_DATA_REPRESENTATION). `legacy_compatibility` (default): " + "the execution engine converts native tensor predictions into the documented " + "sv.Detections / numpy representations at the block's input boundary and converts returned " + "legacy objects back at the output boundary. `tensor_native`: run() receives and must " + "return native `inference_models` objects. The knob is a no-op when the server runs the " + "numpy representation (`tensor_native` then fails at compile time).", + ) class PythonCode(BaseModel): diff --git a/inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py b/inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py index 4b888fd1de..c195df5108 100644 --- a/inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py +++ b/inference/core/workflows/execution_engine/v1/dynamic_blocks/modal_executor.py @@ -179,6 +179,67 @@ def _serialise_image_for_webexec(image: Any) -> dict: return result +def _raise_on_unconverted_tensor_native_value(obj: Any) -> None: + """Defense-in-depth for the Modal wire (tensor pivot, D2-REVISED / Step 7). + + Under ``ENABLE_TENSOR_DATA_REPRESENTATION`` the dynamic-block representation + boundary converts ``legacy_compatibility`` inputs BEFORE remote execution and + ``tensor_native`` is compile-blocked on modal, so no native + ``inference_models`` object (nor a bare ``torch.Tensor``) should ever reach + this serializer. If one does, the generic ``__dict__`` fallback would + silently stringify it on the wire โ€” raise loudly instead. Provably inert + when the flag is off (guard first; natives cannot exist flag-off anyway), + keeping flag-off behavior byte-identical. Lazy imports mirror this file's + style; the modules are already in ``sys.modules`` on any code path that + reaches Modal serialization (the block scaffolding imports the boundary). + """ + # The boundary module's import-time constant is the established patch point + # for tests; reading it through the module keeps the two arms in lockstep. + from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( + representation_boundary, + ) + + if not representation_boundary._TENSOR_REPRESENTATION_ACTIVE: + return + import torch + + from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, + ) + from inference_models.models.base.instance_segmentation import InstanceDetections + from inference_models.models.base.keypoints_detection import KeyPoints + from inference_models.models.base.object_detection import Detections + + if not isinstance( + obj, + ( + Detections, + InstanceDetections, + KeyPoints, + ClassificationPrediction, + MultiLabelClassificationPrediction, + torch.Tensor, + ), + ): + return + offending_type = type(obj) + raise representation_boundary.RepresentationBoundaryError( + public_message=( + f"A native tensor object of type " + f"`{offending_type.__module__}.{offending_type.__qualname__}` reached the " + f"Modal wire serializer unconverted. The representation boundary converts " + f"`legacy_compatibility` inputs before remote execution (and " + f"`tensor_native` is not executable over modal), so this indicates the " + f"value bypassed the boundary. Declare the input's kind on the dynamic " + f"block so the boundary can convert it โ€” refusing to fall back to generic " + f"object serialization, which would silently corrupt the value on the wire." + ), + context="workflow_execution | modal_executor | input_serialization", + offending_type=offending_type, + ) + + def serialize_for_modal_remote_execution(inputs: Dict[str, Any]) -> str: from datetime import datetime @@ -201,6 +262,10 @@ def default(self, obj): "shape": obj.shape, } elif hasattr(obj, "__dict__"): + # Native tensor objects also carry __dict__ โ€” refuse to + # stringify them (see the helper's docstring); arbitrary + # non-native objects keep the pre-existing generic contract. + _raise_on_unconverted_tensor_native_value(obj) return { "_type": "object", "class": obj.__class__.__name__, diff --git a/inference/core/workflows/execution_engine/v1/dynamic_blocks/representation_boundary.py b/inference/core/workflows/execution_engine/v1/dynamic_blocks/representation_boundary.py new file mode 100644 index 0000000000..d2be5cd767 --- /dev/null +++ b/inference/core/workflows/execution_engine/v1/dynamic_blocks/representation_boundary.py @@ -0,0 +1,1455 @@ +"""Representation boundary for custom Python (dynamic) blocks. + +Dynamic blocks are compiled per-workflow from user code the engine cannot +rewrite, so the sibling-file swap pattern of the tensor pivot cannot apply to +them. Instead, blocks declaring ``tensor_compatibility=legacy_compatibility`` +(the default) get a conversion boundary around their ``run()``: native +``inference_models`` objects are converted into the documented legacy +``sv.Detections`` / numpy representations at the input boundary, and the +returned legacy objects are converted back at the output boundary. + +This module hosts both directions of the boundary: the IN direction +(``convert_kwargs_to_legacy`` โ€” native -> legacy at the block's input) and the +OUT direction (``convert_block_result_to_native`` โ€” legacy -> native over the +returned ``BlockResult``, ``FlowControl`` passing through untouched). Both are +kind-driven first (declared kinds pick the converter) with best-effort +sniffing for wildcard values. + +The whole boundary is an IDENTITY when ``ENABLE_TENSOR_DATA_REPRESENTATION`` +is off โ€” resolved once at import time into ``_TENSOR_REPRESENTATION_ACTIVE`` +(tests patch that constant) โ€” which is the flag-off byte-parity guarantee of +the tensor pivot. +""" + +import dataclasses +from typing import Any, Dict, List, Optional, Tuple, Union +from uuid import uuid4 + +import numpy as np +import supervision as sv +import torch + +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_IMAGE_TENSOR_DEVICE, +) +from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_native_classification, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + HOST_MIRROR_KEYS, + build_native_key_points, +) +from inference.core.workflows.errors import DynamicBlockError +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, + TRACKER_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.v1.dynamic_blocks.entities import ( + DynamicOutputDefinition, + ManifestDescription, + TensorCompatibility, +) +from inference.core.workflows.execution_engine.v1.entities import FlowControl +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import ( + coco_rle_masks_to_numpy_mask, + torch_mask_to_coco_rle, +) + +# Resolved once at import time; the boundary is a strict identity when the flag +# is off. +_TENSOR_REPRESENTATION_ACTIVE: bool = ENABLE_TENSOR_DATA_REPRESENTATION + +# The sv.Detections.data column carrying the class-name string (supervision's own +# convention; the numpy serializer reads data["class_name"] literally). +CLASS_NAME_DATA_COLUMN = "class_name" + +_BOUNDARY_CONTEXT = ( + "workflow_execution | step_execution | dynamic_block_representation_boundary" +) + +# Kind names whose runtime representation differs between the numpy and tensor +# paths โ€” exactly the kinds whose serializers/deserializers are flag-swapped in +# `core_steps/loader.py`. The kind NAME is the stable conversion key (same-name +# kinds by design). All other kinds are representation-invariant. +_DETECTIONS_FAMILY_KIND_NAMES = { + "object_detection_prediction", + "instance_segmentation_prediction", + "rle_instance_segmentation_prediction", + "semantic_segmentation_prediction", + "qr_code_detection", + "bar_code_detection", +} +_KEYPOINT_KIND_NAME = "keypoint_detection_prediction" +_CLASSIFICATION_KIND_NAME = "classification_prediction" +_EMBEDDING_KIND_NAME = "embedding" +_TENSOR_KIND_NAME = "tensor" +_IMAGE_KIND_NAME = "image" +CONVERTIBLE_KIND_NAMES = _DETECTIONS_FAMILY_KIND_NAMES | { + _KEYPOINT_KIND_NAME, + _CLASSIFICATION_KIND_NAME, + _EMBEDDING_KIND_NAME, + _TENSOR_KIND_NAME, + _IMAGE_KIND_NAME, +} + +_NativeDetections = (Detections, InstanceDetections) + + +class RepresentationBoundaryError(DynamicBlockError): + """Raised when a value cannot cross the dynamic-block representation + boundary โ€” carries the block name, the input/output name, the offending + type and a remediation hint, all baked into the public message.""" + + def __init__( + self, + public_message: str, + context: str, + inner_error: Optional[Exception] = None, + block_name: Optional[str] = None, + value_name: Optional[str] = None, + offending_type: Optional[type] = None, + ): + super().__init__( + public_message=public_message, + context=context, + inner_error=inner_error, + ) + self.block_name = block_name + self.value_name = value_name + self.offending_type = offending_type + + +def _raise_boundary_error( + block_name: str, + value_name: str, + value: Any, + problem: str, + remediation: str, +) -> None: + offending_type = type(value) + raise RepresentationBoundaryError( + public_message=( + f"Dynamic block `{block_name}`, input/output `{value_name}`: {problem} " + f"(offending type: `{offending_type.__module__}.{offending_type.__qualname__}`). " + f"{remediation}" + ), + context=_BOUNDARY_CONTEXT, + block_name=block_name, + value_name=value_name, + offending_type=offending_type, + ) + + +def convert_kwargs_to_legacy( + kwargs: Dict[str, Any], + manifest_description: Optional[ManifestDescription], + block_name: str, + declared_input_kinds: Optional[Dict[str, frozenset]] = None, +) -> Dict[str, Any]: + """IN boundary: convert native tensor objects in a dynamic block's assembled + kwargs into the documented legacy representations, driven by declared kinds + with best-effort sniffing for wildcard inputs. + + Identity function (the very same ``kwargs`` object) when the tensor + representation is off or the block declared ``tensor_native``. + + ``declared_input_kinds`` is the optional assembly-time precomputed lookup + (see ``collect_declared_input_kind_names``); when omitted it is derived from + the manifest here โ€” same behavior, one extra computation per call. + """ + if not _TENSOR_REPRESENTATION_ACTIVE: + return kwargs + if ( + manifest_description is not None + and manifest_description.tensor_compatibility + is TensorCompatibility.TENSOR_NATIVE + ): + return kwargs + if declared_input_kinds is None: + declared_input_kinds = ( + collect_declared_input_kind_names(manifest_description) or {} + ) + converted = {} + for name, value in kwargs.items(): + converted[name] = _walk_value_to_legacy( + value=value, + declared_kinds=declared_input_kinds.get(name, frozenset()), + block_name=block_name, + value_name=name, + ) + return converted + + +def collect_declared_input_kind_names( + manifest_description: Optional[ManifestDescription], +) -> Optional[Dict[str, frozenset]]: + """Precompute the per-input declared-kind lookup once per block โ€” the + manifest is static, so callers (block assembly) hoist this out of run().""" + if manifest_description is None: + return None + return { + name: frozenset(_collect_declared_kind_names(definition)) + for name, definition in manifest_description.inputs.items() + } + + +def collect_declared_output_kind_names( + manifest_description: Optional[ManifestDescription], +) -> Optional[Dict[str, frozenset]]: + """Per-output twin of ``collect_declared_input_kind_names``.""" + if manifest_description is None: + return None + return { + name: frozenset(_collect_declared_output_kind_names(definition)) + for name, definition in manifest_description.outputs.items() + } + + +def _collect_declared_kind_names(input_definition: Optional[Any]) -> set: + """Union of kind names across the input's selector types. Empty set means + wildcard (no declaration โ‡’ sniffing).""" + if input_definition is None: + return set() + selector_data_kind = getattr(input_definition, "selector_data_kind", None) or {} + declared = set() + for kind_names in selector_data_kind.values(): + declared.update(kind_names) + declared.discard("*") + return declared + + +def _walk_value_to_legacy( + value: Any, + declared_kinds: set, + block_name: str, + value_name: str, +) -> Any: + """Walk containers (Batch preserving indices, dict/list/plain tuple) down to + leaf values and convert each leaf. The keypoint-detection tuple is a LEAF, + not a container โ€” checked before generic tuple recursion.""" + if value is None: + return None + if isinstance(value, Batch): + converted_content = [ + _walk_value_to_legacy( + value=element, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + for element in value + ] + return Batch.init(content=converted_content, indices=value.indices) + if _is_key_point_prediction_tuple(value): + return _convert_leaf_to_legacy( + value=value, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + if isinstance(value, dict): + return { + key: _walk_value_to_legacy( + value=element, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + for key, element in value.items() + } + if isinstance(value, (list, tuple)): + converted = [ + _walk_value_to_legacy( + value=element, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + for element in value + ] + if isinstance(value, tuple): + # Namedtuples take one positional argument PER FIELD - splat the + # converted elements; a plain tuple takes the iterable whole. + if hasattr(value, "_fields"): + return type(value)(*converted) + return type(value)(converted) + return converted + return _convert_leaf_to_legacy( + value=value, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + + +def _convert_leaf_to_legacy( + value: Any, + declared_kinds: set, + block_name: str, + value_name: str, +) -> Any: + convertible_declared = declared_kinds & CONVERTIBLE_KIND_NAMES + if convertible_declared: + return _convert_by_declared_kinds( + value=value, + declared_kinds=convertible_declared, + block_name=block_name, + value_name=value_name, + ) + return _sniff_convert_to_legacy( + value=value, + block_name=block_name, + value_name=value_name, + ) + + +def _convert_by_declared_kinds( + value: Any, + declared_kinds: set, + block_name: str, + value_name: str, +) -> Any: + """Kind-driven conversion: the declared kind names pick the converter; a value + whose type does not match any declared kind's native representation fails + loudly (the producer/consumer contract is broken upstream).""" + if _IMAGE_KIND_NAME in declared_kinds and isinstance(value, WorkflowImageData): + # Images need no conversion in either direction: WorkflowImageData is + # lazily dual-representation. + return value + if declared_kinds & _DETECTIONS_FAMILY_KIND_NAMES and isinstance( + value, _NativeDetections + ): + return native_detections_to_sv(detections=value) + if _KEYPOINT_KIND_NAME in declared_kinds and _is_key_point_prediction_tuple(value): + return native_key_point_prediction_to_sv( + prediction=value, block_name=block_name, value_name=value_name + ) + if _CLASSIFICATION_KIND_NAME in declared_kinds and isinstance( + value, (ClassificationPrediction, MultiLabelClassificationPrediction) + ): + return serialise_native_classification(prediction=value) + if _EMBEDDING_KIND_NAME in declared_kinds and isinstance(value, torch.Tensor): + # numpy embedding blocks (clip / perception_encoder) emit List[float] + # in-memory (`predictions.embeddings[0]`). + return value.detach().cpu().reshape(-1).tolist() + if _TENSOR_KIND_NAME in declared_kinds and isinstance(value, torch.Tensor): + # The `tensor` kind has no numpy-path producer; ndarray is the legacy + # in-memory equivalent numpy user code can operate on. + return value.detach().cpu().numpy() + if _is_representation_invariant(value): + # Declared kinds also admit static values / already-legacy payloads + # (e.g. a plain list fed to an embedding input) โ€” those pass through. + return value + _raise_boundary_error( + block_name=block_name, + value_name=value_name, + value=value, + problem=( + f"value does not match the native representation of the declared " + f"kind(s) {sorted(declared_kinds)}" + ), + remediation=( + "Check the upstream step's output kind, or remove the kind " + "declaration to enable best-effort conversion." + ), + ) + + +def _sniff_convert_to_legacy( + value: Any, + block_name: str, + value_name: str, +) -> Any: + """Wildcard best-effort: convert the known native prediction types, pass + representation-invariant values through, and fail loudly on tensor-only + values with no legacy equivalent.""" + if isinstance(value, _NativeDetections): + return native_detections_to_sv(detections=value) + if _is_key_point_prediction_tuple(value): + return native_key_point_prediction_to_sv( + prediction=value, block_name=block_name, value_name=value_name + ) + if isinstance( + value, (ClassificationPrediction, MultiLabelClassificationPrediction) + ): + return serialise_native_classification(prediction=value) + if isinstance(value, KeyPoints): + _raise_boundary_error( + block_name=block_name, + value_name=value_name, + value=value, + problem="a bare `KeyPoints` prediction (without its bounding-box " + "component) has no legacy sv.Detections equivalent", + remediation="Pass the full keypoint-detection prediction, or switch " + "the block to `tensor_compatibility=tensor_native`.", + ) + if isinstance(value, torch.Tensor): + _raise_boundary_error( + block_name=block_name, + value_name=value_name, + value=value, + problem="a bare torch.Tensor has no unambiguous legacy equivalent", + remediation="Declare the input's kind (e.g. `embedding` or `tensor`) " + "or switch the block to `tensor_compatibility=tensor_native`.", + ) + if _is_representation_invariant(value): + return value + if dataclasses.is_dataclass(value) and not isinstance(value, type): + _raise_boundary_error( + block_name=block_name, + value_name=value_name, + value=value, + problem="unrecognised tensor-native dataclass cannot be converted " + "to a legacy representation", + remediation="Declare the input's kind or switch the block to " + "`tensor_compatibility=tensor_native`.", + ) + # INTENTIONAL representation-invariant passthrough. Anything not caught above + # (arbitrary objects, pydantic models, NamedTuples, ...) is assumed + # representation-agnostic and shipped as-is; the loud-error net covers only + # the tensor-only shapes we can positively identify (native dataclasses, + # bare tensors, bare KeyPoints). + return value + + +def _is_representation_invariant(value: Any) -> bool: + return isinstance( + value, (WorkflowImageData, str, bytes, bool, int, float, np.ndarray) + ) + + +def _is_key_point_prediction_tuple(value: Any) -> bool: + return ( + isinstance(value, tuple) + and len(value) == 2 + and isinstance(value[0], KeyPoints) + and (value[1] is None or isinstance(value[1], _NativeDetections)) + ) + + +def native_key_point_prediction_to_sv( + prediction: Tuple[KeyPoints, Optional[Detections]], + block_name: str, + value_name: str, +) -> sv.Detections: + """Convert the native keypoint-detection tuple via its bounding-box component + (per-instance keypoint payloads already ride in ``bboxes_metadata`` under the + sv-data key names).""" + _, detections = prediction + if detections is None: + _raise_boundary_error( + block_name=block_name, + value_name=value_name, + value=prediction, + problem="keypoint prediction is missing the bounding-box component " + "required to build the legacy sv.Detections", + remediation="Use a keypoint model output that carries bounding boxes, " + "or switch the block to `tensor_compatibility=tensor_native`.", + ) + return native_detections_to_sv(detections=detections) + + +def native_detections_to_sv( + detections: Union[Detections, InstanceDetections], +) -> sv.Detections: + """Materialise a native ``Detections`` / ``InstanceDetections`` into the full + legacy ``sv.Detections`` the numpy blocks operate on. + + Extends ``to_supervision_for_annotation`` semantics (visualizations + ``base_tensor.py``) with what legacy USER CODE additionally relies on: + + * ``data['class_name']`` resolved effective-name / override-first โ€” per-box + ``bboxes_metadata['class']`` when present, else the + ``image_metadata['class_names']`` map (the tensor serializer's rule), + * per-box ``detection_id`` minted (uuid4) when missing โ€” mirroring the numpy + deserializer, never an empty string, + * lineage broadcast into ``.data`` exactly as + ``attach_parents_coordinates_to_sv_detections`` shapes it, plus + ``prediction_type`` / ``image_dimensions`` / ``inference_id``, + * masks ALWAYS dense ``(N, H, W)`` bool ndarray โ€” RLE carriers are decoded + (legacy user code indexes ``detections.mask``; the device->host + densify + cost is inherent to running numpy user code), + * per-box keypoint payloads padded into proper N-d arrays following the + numpy ``add_inference_keypoints_to_sv_detections`` convention (ragged + object arrays break supervision's ``is_data_equal``). + """ + image_metadata = detections.image_metadata or {} + detections_number = int(detections.xyxy.shape[0]) + bboxes_metadata = detections.bboxes_metadata + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(detections_number)] + class_names_mapping = image_metadata.get(CLASS_NAMES_KEY) or {} + # Single deviceโ†’host sync for the three per-box tensors (class_id round-trips + # through float32, which is exact for any realistic class count). + packed = ( + torch.cat( + [ + detections.xyxy.reshape(detections_number, 4).to(torch.float32), + detections.confidence.reshape(detections_number, 1).to(torch.float32), + detections.class_id.reshape(detections_number, 1).to(torch.float32), + ], + dim=1, + ) + .detach() + .cpu() + .numpy() + ) + xyxy = packed[:, :4].astype(np.float32) + confidence = packed[:, 4].astype(np.float32) + class_id = packed[:, 5].astype(int) + mask = _materialise_dense_mask(detections=detections) + tracker_id = _materialise_tracker_id(bboxes_metadata=bboxes_metadata) + data: Dict[str, np.ndarray] = {} + class_names = [ + _resolve_effective_class_name( + per_box=bboxes_metadata[index], + class_id=int(class_id[index]), + class_names_mapping=class_names_mapping, + ) + for index in range(detections_number) + ] + data[CLASS_NAME_DATA_COLUMN] = np.asarray(class_names, dtype=object) + detection_ids = [ + str(per_box.get(DETECTION_ID_KEY) or uuid4()) for per_box in bboxes_metadata + ] + data[DETECTION_ID_KEY] = np.array(detection_ids) + _broadcast_image_level_columns( + data=data, + image_metadata=image_metadata, + detections_number=detections_number, + ) + _attach_per_box_columns( + data=data, + bboxes_metadata=bboxes_metadata, + detections_number=detections_number, + ) + return sv.Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask, + tracker_id=tracker_id, + data=data, + ) + + +def _resolve_effective_class_name( + per_box: dict, + class_id: int, + class_names_mapping: dict, +) -> str: + if CLASS_NAME_KEY in per_box: + return str(per_box[CLASS_NAME_KEY]) + class_name = class_names_mapping.get(class_id) + if class_name is None: + return f"class_{class_id}" + return str(class_name) + + +def _materialise_dense_mask( + detections: Union[Detections, InstanceDetections], +) -> Optional[np.ndarray]: + if not isinstance(detections, InstanceDetections): + return None + mask = detections.mask + if mask is None: + return None + if isinstance(mask, InstancesRLEMasks): + # Iterative densification: decode one instance at a time into a single + # preallocated (N, H, W) output. The bulk pycocotools decode holds up to + # three full-stack transients at peak ((H, W, N) uint8 + bool astype + + # contiguous copy); per-instance decoding caps the transient at one mask. + height, width = (int(value) for value in mask.image_size) + dense = np.zeros((len(mask.masks), height, width), dtype=bool) + for index in range(len(mask.masks)): + dense[index] = coco_rle_masks_to_numpy_mask( + InstancesRLEMasks( + image_size=mask.image_size, + masks=[mask.masks[index]], + ) + )[0] + return dense + # single bulk device->host transfer for the whole stack + return mask.detach().cpu().numpy().astype(bool) + + +def _materialise_tracker_id(bboxes_metadata: List[dict]) -> Optional[np.ndarray]: + tracker_ids = [per_box.get(TRACKER_ID_KEY) for per_box in bboxes_metadata] + if not tracker_ids or any(tracker_id is None for tracker_id in tracker_ids): + return None + return np.asarray([int(tracker_id) for tracker_id in tracker_ids]) + + +_IMAGE_LEVEL_ID_KEYS = (PARENT_ID_KEY, ROOT_PARENT_ID_KEY) +_IMAGE_LEVEL_PAIR_KEYS = ( + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + IMAGE_DIMENSIONS_KEY, +) +_IMAGE_LEVEL_SCALAR_KEYS = (PREDICTION_TYPE_KEY, INFERENCE_ID_KEY) + + +def _broadcast_image_level_columns( + data: Dict[str, np.ndarray], + image_metadata: dict, + detections_number: int, +) -> None: + """Broadcast per-image metadata into per-row ``.data`` columns with the exact + key names and shapes the numpy path produces: id/scalar keys as string arrays, + coordinate/dimension keys as ``(n, 2)`` arrays (``[x, y]`` for coordinates, + ``[h, w]`` for dimensions โ€” the native metadata already stores them in those + orders).""" + for key in _IMAGE_LEVEL_ID_KEYS + _IMAGE_LEVEL_SCALAR_KEYS: + value = image_metadata.get(key) + if value is not None: + data[key] = np.array([value] * detections_number) + for key in _IMAGE_LEVEL_PAIR_KEYS: + value = image_metadata.get(key) + if value is not None: + data[key] = np.array([list(value)] * detections_number).reshape( + detections_number, 2 + ) + + +_KEYPOINT_PAYLOAD_KEYS = ( + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, +) + + +def _attach_per_box_columns( + data: Dict[str, np.ndarray], + bboxes_metadata: List[dict], + detections_number: int, +) -> None: + extra_keys = set() + for per_box in bboxes_metadata: + extra_keys.update(per_box.keys()) + extra_keys.discard(DETECTION_ID_KEY) + extra_keys.discard(TRACKER_ID_KEY) + extra_keys.discard(CLASS_NAME_KEY) # resolved into data['class_name'] already + # The private per-box host mirror of the detection tensors (tensor_native) + # is an internal transport channel โ€” never surfaced to legacy user code. + # Excluding it here also keeps the sv->native round-trip from re-attaching + # a mirror the user code may have invalidated by editing xyxy/class_id. + extra_keys.difference_update(HOST_MIRROR_KEYS) + keypoint_keys_present = extra_keys & set(_KEYPOINT_PAYLOAD_KEYS) + if keypoint_keys_present == set(_KEYPOINT_PAYLOAD_KEYS): + _attach_padded_keypoint_columns( + data=data, + bboxes_metadata=bboxes_metadata, + detections_number=detections_number, + ) + extra_keys -= keypoint_keys_present + for key in extra_keys: + data[key] = np.asarray( + [per_box.get(key) for per_box in bboxes_metadata], dtype=object + ) + + +def _attach_padded_keypoint_columns( + data: Dict[str, np.ndarray], + bboxes_metadata: List[dict], + detections_number: int, +) -> None: + """Pad per-box keypoint payloads into proper N-d arrays โ€” the exact + convention of ``add_inference_keypoints_to_sv_detections`` (utils.py): ragged + object arrays break supervision's ``is_data_equal``.""" + keypoints_xy = [ + list(per_box.get(KEYPOINTS_XY_KEY_IN_SV_DETECTIONS) or []) + for per_box in bboxes_metadata + ] + keypoints_confidence = [ + list(per_box.get(KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS) or []) + for per_box in bboxes_metadata + ] + keypoints_class_id = [ + list(per_box.get(KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS) or []) + for per_box in bboxes_metadata + ] + keypoints_class_name = [ + list(per_box.get(KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS) or []) + for per_box in bboxes_metadata + ] + max_kps = max((len(kp) for kp in keypoints_xy), default=0) + padded_xy = np.zeros((detections_number, max_kps, 2), dtype=np.float32) + padded_conf = np.zeros((detections_number, max_kps), dtype=np.float32) + padded_class_id = np.zeros((detections_number, max_kps), dtype=int) + padded_class_name = np.full((detections_number, max_kps), "", dtype=object) + for index in range(detections_number): + kps_in_instance = len(keypoints_xy[index]) + if kps_in_instance > 0: + padded_xy[index, :kps_in_instance] = keypoints_xy[index] + padded_conf[index, :kps_in_instance] = keypoints_confidence[index] + padded_class_id[index, :kps_in_instance] = keypoints_class_id[index] + padded_class_name[index, :kps_in_instance] = keypoints_class_name[index] + data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = padded_xy + data[KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS] = padded_conf + data[KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS] = padded_class_id + data[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS] = padded_class_name + + +# --------------------------------------------------------------------------- # +# OUT boundary: legacy -> native over the returned BlockResult # +# --------------------------------------------------------------------------- # + +# Same literal the classification tensor blocks attach; the tensor serializer +# reads it from image metadata to filter sub-threshold classes. +CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY = "classification_confidence_threshold" + +RLE_MASK_DATA_COLUMN = "rle_mask" + +# Kinds whose native instance-segmentation carrier should be RLE rather than a +# dense torch stack (the tensor path's compact-mask convention). +_RLE_CARRIER_KIND_NAMES = { + "rle_instance_segmentation_prediction", + "semantic_segmentation_prediction", +} + +# Declared kinds whose native representation is `InstanceDetections` (mask +# carrier required) โ€” an empty output under these must stay instance-shaped. +_MASK_CARRIER_KIND_NAMES = { + "instance_segmentation_prediction" +} | _RLE_CARRIER_KIND_NAMES + +_NATIVE_CLASSIFICATION_TYPES = ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) + + +def convert_block_result_to_native( + result: Any, + manifest_description: Optional[ManifestDescription], + block_name: str, + declared_output_kinds: Optional[Dict[str, frozenset]] = None, +) -> Any: + """OUT boundary: convert legacy objects in a dynamic block's ``BlockResult`` + back into native tensor representations, driven by declared output kinds + with best-effort sniffing for wildcard outputs. + + ``BlockResult`` is ``dict | FlowControl | List[...] | List[List[...]]``; + ``FlowControl`` elements pass through untouched (they carry no + representation-dependent payload). Identity function (the very same object) + when the tensor representation is off or the block declared + ``tensor_native``. + + ``declared_output_kinds`` is the optional assembly-time precomputed lookup + (see ``collect_declared_output_kind_names``); when omitted it is derived + from the manifest here โ€” same behavior, one extra computation per call. + """ + if not _TENSOR_REPRESENTATION_ACTIVE: + return result + if ( + manifest_description is not None + and manifest_description.tensor_compatibility + is TensorCompatibility.TENSOR_NATIVE + ): + return result + if declared_output_kinds is None: + declared_output_kinds = ( + collect_declared_output_kind_names(manifest_description) or {} + ) + return _walk_result_to_native( + result=result, + declared_output_kinds=declared_output_kinds, + block_name=block_name, + ) + + +def _walk_result_to_native( + result: Any, + declared_output_kinds: Dict[str, frozenset], + block_name: str, +) -> Any: + if isinstance(result, FlowControl): + return result + if isinstance(result, list): + return [ + _walk_result_to_native( + result=element, + declared_output_kinds=declared_output_kinds, + block_name=block_name, + ) + for element in result + ] + if isinstance(result, dict): + return { + output_name: _walk_output_value_to_native( + value=value, + declared_kinds=declared_output_kinds.get(output_name, frozenset()), + block_name=block_name, + value_name=output_name, + ) + for output_name, value in result.items() + } + # Anything else is not a valid BlockResult shape โ€” the engine rejects it + # downstream with its own diagnostics; representation is not the problem. + return result + + +def _collect_declared_output_kind_names( + output_definition: Optional[DynamicOutputDefinition], +) -> set: + if output_definition is None: + return set() + declared = set(getattr(output_definition, "kind", None) or []) + declared.discard("*") + return declared + + +def _walk_output_value_to_native( + value: Any, + declared_kinds: set, + block_name: str, + value_name: str, +) -> Any: + """Walk containers down to leaf values (mirroring the IN walker; results + never carry ``Batch``). Legacy predictions (sv.Detections, classification + dicts under a declared kind) are leaves.""" + if value is None: + return None + if _is_key_point_prediction_tuple(value): + # Already-native keypoint tuple returned under legacy mode: idempotent + # passthrough (it is the target representation already). + return value + if isinstance(value, dict): + if declared_kinds & {_CLASSIFICATION_KIND_NAME}: + return _convert_output_leaf_to_native( + value=value, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + return { + key: _walk_output_value_to_native( + value=element, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + for key, element in value.items() + } + if isinstance(value, (list, tuple)): + if declared_kinds & {_EMBEDDING_KIND_NAME, _TENSOR_KIND_NAME}: + # A declared embedding/tensor output returning a plain list IS the + # legacy leaf value - convert it, don't recurse into the numbers. + return _convert_output_leaf_to_native( + value=value, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + converted = [ + _walk_output_value_to_native( + value=element, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + for element in value + ] + if isinstance(value, tuple): + # Namedtuples take one positional argument PER FIELD - splat the + # converted elements; a plain tuple takes the iterable whole. + if hasattr(value, "_fields"): + return type(value)(*converted) + return type(value)(converted) + return converted + return _convert_output_leaf_to_native( + value=value, + declared_kinds=declared_kinds, + block_name=block_name, + value_name=value_name, + ) + + +def _convert_output_leaf_to_native( + value: Any, + declared_kinds: set, + block_name: str, + value_name: str, +) -> Any: + convertible_declared = declared_kinds & CONVERTIBLE_KIND_NAMES + if convertible_declared: + return _convert_output_by_declared_kinds( + value=value, + declared_kinds=convertible_declared, + block_name=block_name, + value_name=value_name, + ) + return _sniff_convert_to_native( + value=value, block_name=block_name, value_name=value_name + ) + + +def _convert_output_by_declared_kinds( + value: Any, + declared_kinds: set, + block_name: str, + value_name: str, +) -> Any: + if _IMAGE_KIND_NAME in declared_kinds and isinstance(value, WorkflowImageData): + return value + if _KEYPOINT_KIND_NAME in declared_kinds: + if _is_key_point_prediction_tuple(value): + return value + # Kind-UNION disambiguation: outputs commonly declare + # [object_detection, instance_segmentation, keypoint_detection] together + # (mirroring model-agnostic inputs). Rebuilding the KP tuple from an sv + # that carries no keypoint payload would fabricate a keypoint prediction + # out of a plain detection - so within a union the choice is data-driven + # (payload columns present -> tuple), and only a SOLE keypoint + # declaration keeps the strict always-tuple contract. + if isinstance(value, sv.Detections) and ( + _sv_detections_carry_keypoint_payload(value) + or not declared_kinds & _DETECTIONS_FAMILY_KIND_NAMES + ): + return sv_detections_to_native_key_point_prediction(sv_detections=value) + if declared_kinds & _DETECTIONS_FAMILY_KIND_NAMES: + if isinstance(value, _NativeDetections): + return value + if isinstance(value, sv.Detections): + prefer_rle = bool(declared_kinds & _RLE_CARRIER_KIND_NAMES) + converted = sv_detections_to_native( + sv_detections=value, + prefer_rle=prefer_rle, + ) + # An empty sv carries no mask, so the converter degrades to plain + # `Detections`; under a declared mask-carrying kind the native + # convention is an EMPTY `InstanceDetections` (mirrors + # `_empty_instance_detections` in sam3_video / the module's own + # empty prefer_rle branch). Image dims are unrecoverable from zero + # rows, hence the (0, 0)-sized empty carrier. Non-empty box-only sv + # stays plain `Detections` โ€” masks cannot be invented, and the + # numpy path serializes box-only rows the same way. + if ( + len(value) == 0 + and declared_kinds & _MASK_CARRIER_KIND_NAMES + and not isinstance(converted, InstanceDetections) + ): + converted = _as_empty_instance_detections( + empty_detections=converted, prefer_rle=prefer_rle + ) + return converted + if _CLASSIFICATION_KIND_NAME in declared_kinds: + if isinstance(value, _NATIVE_CLASSIFICATION_TYPES): + return value + if isinstance(value, dict): + return classification_dict_to_native( + prediction=value, + block_name=block_name, + value_name=value_name, + ) + if declared_kinds & {_EMBEDDING_KIND_NAME, _TENSOR_KIND_NAME}: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (list, tuple, np.ndarray)): + return _legacy_array_to_tensor(value=value, declared_kinds=declared_kinds) + if _is_representation_invariant(value): + return value + _raise_boundary_error( + block_name=block_name, + value_name=value_name, + value=value, + problem=( + f"returned value does not match the legacy representation of the " + f"declared output kind(s) {sorted(declared_kinds)}" + ), + remediation=( + "Return the documented legacy type for the declared kind, or remove " + "the kind declaration to enable best-effort conversion." + ), + ) + + +def _sniff_convert_to_native(value: Any, block_name: str, value_name: str) -> Any: + """Wildcard best-effort for outputs: ``sv.Detections`` converts data-driven + (keypoint payload columns -> the KP tuple; masks -> ``InstanceDetections``; + else ``Detections``); already-native objects pass through; a bare + ``sv.KeyPoints`` fails loudly (mirror of the IN-side bare-``KeyPoints`` + rule); EVERYTHING else - dicts included (a classification dict through a + wildcard output stays a dict) - is representation-invariant.""" + if isinstance(value, sv.Detections): + if _sv_detections_carry_keypoint_payload(value): + return sv_detections_to_native_key_point_prediction(sv_detections=value) + return sv_detections_to_native(sv_detections=value, prefer_rle=False) + if isinstance(value, sv.KeyPoints): + _raise_boundary_error( + block_name=block_name, + value_name=value_name, + value=value, + problem="a bare `sv.KeyPoints` (without its detections component) " + "has no native keypoint-prediction equivalent", + remediation="Return the full keypoint payload on an sv.Detections " + "(the numpy keypoint blocks' shape), declare the output kind, or " + "switch the block to `tensor_compatibility=tensor_native`.", + ) + return value + + +def _sv_detections_carry_keypoint_payload(sv_detections: sv.Detections) -> bool: + return all(key in sv_detections.data for key in _KEYPOINT_PAYLOAD_KEYS) + + +def _legacy_array_to_tensor(value: Any, declared_kinds: set) -> torch.Tensor: + if _EMBEDDING_KIND_NAME in declared_kinds and not isinstance(value, np.ndarray): + # IN emits embeddings as List[float]; the exact inverse. + return torch.as_tensor( + np.asarray(value, dtype=np.float32), device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + array = np.asarray(value) + if _EMBEDDING_KIND_NAME in declared_kinds and array.dtype != np.float32: + array = array.astype(np.float32) + return torch.as_tensor(array.copy(), device=WORKFLOWS_IMAGE_TENSOR_DEVICE) + + +def sv_detections_to_native_key_point_prediction( + sv_detections: sv.Detections, +) -> Tuple[KeyPoints, Union[Detections, InstanceDetections]]: + """Rebuild the native ``(KeyPoints, Detections)`` keypoint prediction from an + ``sv.Detections`` whose rows carry the keypoint payload columns (the shape + the IN converter / numpy keypoint blocks produce). The bbox component keeps + the payload in ``bboxes_metadata`` (the serializer's convention); the + ``KeyPoints`` component is rebuilt via the shared + ``build_native_key_points`` helper.""" + bbox_component = sv_detections_to_native(sv_detections=sv_detections) + detections_number = int(len(sv_detections)) + per_instance_xy: List[List[List[float]]] = [] + per_instance_confidence: List[List[float]] = [] + xy_column = sv_detections.data.get(KEYPOINTS_XY_KEY_IN_SV_DETECTIONS) + confidence_column = sv_detections.data.get( + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS + ) + for index in range(detections_number): + xy_row = xy_column[index] if xy_column is not None else [] + confidence_row = ( + confidence_column[index] if confidence_column is not None else [] + ) + per_instance_xy.append(np.asarray(xy_row).reshape(-1, 2).tolist()) + per_instance_confidence.append(np.asarray(confidence_row).reshape(-1).tolist()) + class_ids = ( + bbox_component.class_id.detach().cpu().tolist() if detections_number > 0 else [] + ) + key_points = build_native_key_points( + per_instance_xy=per_instance_xy, + per_instance_confidence=per_instance_confidence, + object_class_ids=class_ids, + image_metadata=bbox_component.image_metadata or {}, + ) + return key_points, bbox_component + + +def sv_detections_to_native( + sv_detections: sv.Detections, + *, + prefer_rle: bool = False, +) -> Union[Detections, InstanceDetections]: + """Direct ``sv.Detections`` -> native converter (no dict round-trip, dense + masks preserved). + + Inverts ``native_detections_to_sv`` exactly: + + * xyxy/class_id/confidence become torch tensors on + ``WORKFLOWS_IMAGE_TENSOR_DEVICE``; + * ``image_metadata`` is rebuilt from the broadcast ``.data`` lineage columns + (row 0 of each; rows share them by construction) โ€” parent/root ids, + coordinates and dimensions (including ``ROOT_PARENT_DIMENSIONS_KEY``, the + root-coordinates precondition for masked instances), prediction_type, + image_dimensions, inference_id; + * ``CLASS_NAMES_KEY`` maps class_id -> first-seen ``data['class_name']`` + (fallback ``f"class_{id}"``); rows whose name differs from the map entry + (id-sharing overrides) keep a per-box ``class`` override โ€” the exact + inverse of the IN converter's effective-name resolution; + * ``bboxes_metadata`` carries per-box ``detection_id`` (minted uuid4 when + missing โ€” the tensor serializer hard-requires it), ``tracker_id`` from the + sv field, keypoint payload rows and every remaining ``.data`` column; + * mask carrier: a ``data['rle_mask']`` column wins (COCO dicts -> + ``InstancesRLEMasks``, no re-encode); otherwise a dense sv mask becomes a + torch bool stack, or โ€” under ``prefer_rle=True`` (declared + RLE/semantic-seg output kinds) โ€” is encoded to ``InstancesRLEMasks``. + """ + detections_number = int(len(sv_detections)) + data = sv_detections.data or {} + image_metadata = _rebuild_image_metadata_from_sv( + data=data, detections_number=detections_number + ) + class_names_column = data.get(CLASS_NAME_DATA_COLUMN) + class_id_values = ( + [int(value) for value in sv_detections.class_id] + if sv_detections.class_id is not None + else [0] * detections_number + ) + class_names_mapping: Dict[int, str] = {} + per_box_class_overrides: Dict[int, str] = {} + for index in range(detections_number): + class_id = class_id_values[index] + class_name = ( + str(class_names_column[index]) + if class_names_column is not None + else f"class_{class_id}" + ) + if class_id not in class_names_mapping: + # First-seen name wins the map entry. KNOWN DRIFT: if row 0 of a + # shared class_id was itself a per-box override, the override + # string lands in CLASS_NAMES_KEY and the base name becomes the + # per-box override โ€” the base-map/override split is NOT recoverable + # from sv.Detections. Effective names and serialized output are + # preserved (the collision round-trip test proves it); only + # consumers reading image_metadata[CLASS_NAMES_KEY][id] directly + # may observe the swapped roles. + class_names_mapping[class_id] = class_name + elif class_names_mapping[class_id] != class_name: + per_box_class_overrides[index] = class_name + image_metadata[CLASS_NAMES_KEY] = class_names_mapping + bboxes_metadata = _rebuild_bboxes_metadata_from_sv( + sv_detections=sv_detections, + data=data, + detections_number=detections_number, + per_box_class_overrides=per_box_class_overrides, + ) + xyxy = torch.as_tensor( + np.asarray(sv_detections.xyxy, dtype=np.float32).reshape(-1, 4), + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ) + class_id_tensor = torch.as_tensor( + class_id_values, dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ).reshape(-1) + confidence_values = ( + np.asarray(sv_detections.confidence, dtype=np.float32).reshape(-1) + if sv_detections.confidence is not None + else np.zeros(detections_number, dtype=np.float32) + ) + confidence_tensor = torch.as_tensor( + confidence_values, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + mask_carrier = _rebuild_mask_carrier_from_sv( + sv_detections=sv_detections, + data=data, + detections_number=detections_number, + prefer_rle=prefer_rle, + ) + if mask_carrier is not None: + return InstanceDetections( + xyxy=xyxy, + class_id=class_id_tensor, + confidence=confidence_tensor, + mask=mask_carrier, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + return Detections( + xyxy=xyxy, + class_id=class_id_tensor, + confidence=confidence_tensor, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def _rebuild_image_metadata_from_sv( + data: Dict[str, Any], + detections_number: int, +) -> dict: + image_metadata: dict = {} + if detections_number == 0: + return image_metadata + for key in _IMAGE_LEVEL_ID_KEYS + _IMAGE_LEVEL_SCALAR_KEYS: + column = data.get(key) + if column is not None and len(column) > 0: + image_metadata[key] = str(column[0]) + for key in _IMAGE_LEVEL_PAIR_KEYS: + column = data.get(key) + if column is not None and len(column) > 0: + image_metadata[key] = [int(value) for value in np.asarray(column[0])] + return image_metadata + + +_IMAGE_LEVEL_COLUMN_KEYS = frozenset( + _IMAGE_LEVEL_ID_KEYS + _IMAGE_LEVEL_PAIR_KEYS + _IMAGE_LEVEL_SCALAR_KEYS +) + + +def _as_empty_instance_detections( + empty_detections: Detections, + prefer_rle: bool, +) -> InstanceDetections: + """Re-shape an empty plain ``Detections`` into the empty + ``InstanceDetections`` convention for declared mask-carrying kinds.""" + if prefer_rle: + mask_carrier: Union[torch.Tensor, InstancesRLEMasks] = InstancesRLEMasks( + image_size=(0, 0), masks=[] + ) + else: + mask_carrier = torch.zeros( + (0, 0, 0), dtype=torch.bool, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + return InstanceDetections( + xyxy=empty_detections.xyxy, + class_id=empty_detections.class_id, + confidence=empty_detections.confidence, + mask=mask_carrier, + image_metadata=empty_detections.image_metadata, + bboxes_metadata=empty_detections.bboxes_metadata, + ) + + +def _rebuild_bboxes_metadata_from_sv( + sv_detections: sv.Detections, + data: Dict[str, Any], + detections_number: int, + per_box_class_overrides: Dict[int, str], +) -> Optional[List[dict]]: + # Zero rows -> None, matching the native convention + # (`attach_native_detection_metadata` / the fusion empty builders); the + # root-shift helper and the serializer both special-case None already. + if detections_number == 0: + return None + detection_id_column = data.get(DETECTION_ID_KEY) + per_box_keys = [ + key + for key in data.keys() + if key + not in _IMAGE_LEVEL_COLUMN_KEYS + | {CLASS_NAME_DATA_COLUMN, DETECTION_ID_KEY, RLE_MASK_DATA_COLUMN} + ] + keypoint_keys = set(_KEYPOINT_PAYLOAD_KEYS) + bboxes_metadata: List[dict] = [] + for index in range(detections_number): + per_box: dict = {} + carried_id = ( + detection_id_column[index] if detection_id_column is not None else None + ) + per_box[DETECTION_ID_KEY] = str(carried_id) if carried_id else str(uuid4()) + if sv_detections.tracker_id is not None: + per_box[TRACKER_ID_KEY] = int(sv_detections.tracker_id[index]) + if index in per_box_class_overrides: + per_box[CLASS_NAME_KEY] = per_box_class_overrides[index] + for key in per_box_keys: + row_value = data[key][index] + if key in keypoint_keys: + row_value = np.asarray(row_value).tolist() + per_box[key] = row_value + bboxes_metadata.append(per_box) + return bboxes_metadata + + +def _rebuild_mask_carrier_from_sv( + sv_detections: sv.Detections, + data: Dict[str, Any], + detections_number: int, + prefer_rle: bool, +) -> Optional[Union[torch.Tensor, InstancesRLEMasks]]: + rle_column = data.get(RLE_MASK_DATA_COLUMN) + if rle_column is not None and len(rle_column) == detections_number > 0: + coco_masks = [] + for index, entry in enumerate(rle_column): + if ( + not isinstance(entry, dict) + or "size" not in entry + or "counts" not in entry + ): + raise ValueError( + f"sv.Detections carries a `{RLE_MASK_DATA_COLUMN}` column whose " + f"entry at index {index} is not a COCO-RLE dict with `size` and " + f"`counts` keys (got `{type(entry).__name__}`) - cannot rebuild " + f"the native RLE mask carrier from it." + ) + coco_masks.append(dict(entry)) + image_size = tuple(int(value) for value in coco_masks[0]["size"]) + return InstancesRLEMasks.from_coco_rle_masks( + image_size=image_size, masks=coco_masks + ) + if sv_detections.mask is None: + return None + dense = np.asarray(sv_detections.mask) + if dense.dtype != np.bool_: + # Both carriers require bool; copy ONLY when the dtype actually differs + # (sv masks are typically bool already, and an unconditional astype would + # duplicate the full (N, H, W) stack). + dense = dense.astype(bool) + if prefer_rle: + if detections_number == 0: + return InstancesRLEMasks(image_size=(0, 0), masks=[]) + # Iterative encode straight off the existing rows โ€” rows of a contiguous + # stack share memory with torch.as_tensor, so the only per-instance + # transient is the fortran-order flatten inside torch_mask_to_coco_rle. + coco_masks = [ + torch_mask_to_coco_rle(torch.as_tensor(np.ascontiguousarray(instance_mask))) + for instance_mask in dense + ] + image_size = tuple(int(value) for value in coco_masks[0]["size"]) + return InstancesRLEMasks.from_coco_rle_masks( + image_size=image_size, masks=coco_masks + ) + return torch.as_tensor(dense, device=WORKFLOWS_IMAGE_TENSOR_DEVICE) + + +def classification_dict_to_native( + prediction: dict, + block_name: str, + value_name: str, +) -> Union[ClassificationPrediction, MultiLabelClassificationPrediction]: + """Rebuild a native classification prediction from the legacy in-memory dict. + + The confidence vector is reconstructed SPARSE โ€” zeros for classes absent + from the (already thresholded) dict. Single-label vs multi-label is + discriminated by the legacy dict shape (``predictions`` list + ``top`` vs + ``predictions`` dict + ``predicted_classes``).""" + predictions_field = prediction.get("predictions") + if "predicted_classes" in prediction and isinstance(predictions_field, dict): + return _multi_label_dict_to_native(prediction=prediction) + if isinstance(predictions_field, list): + return _single_label_dict_to_native(prediction=prediction) + _raise_boundary_error( + block_name=block_name, + value_name=value_name, + value=prediction, + problem="dict is not a recognised legacy classification prediction " + "(expected a `predictions` list with `top`, or a `predictions` dict " + "with `predicted_classes`)", + remediation="Return the documented classification dict shape, or remove " + "the kind declaration.", + ) + + +def _classification_metadata_from_dict(prediction: dict, class_names: dict) -> dict: + metadata: dict = { + CLASS_NAMES_KEY: class_names, + PREDICTION_TYPE_KEY: str( + prediction.get(PREDICTION_TYPE_KEY) or "classification" + ), + } + image_field = prediction.get("image") or {} + height, width = image_field.get("height"), image_field.get("width") + if height is not None and width is not None: + metadata[IMAGE_DIMENSIONS_KEY] = [int(height), int(width)] + for key in (INFERENCE_ID_KEY, PARENT_ID_KEY, ROOT_PARENT_ID_KEY): + value = prediction.get(key) + if value is not None: + metadata[key] = str(value) + # `time` round-trips when the dict carries it (the tensor serializer emits it). + if prediction.get("time") is not None: + metadata["time"] = float(prediction["time"]) + return metadata + + +def _single_label_dict_to_native(prediction: dict) -> ClassificationPrediction: + entries = prediction.get("predictions") or [] + entry_names: Dict[int, str] = {} + entry_confidences: Dict[int, float] = {} + for entry in entries: + class_id = int(entry["class_id"]) + entry_names[class_id] = str(entry.get("class", class_id)) + entry_confidences[class_id] = float(entry.get("confidence", 0.0)) + # Sparse-reconstruction approximation, part 2: the vector length is + # max(SURVIVING class_id) + 1 โ€” when the highest-id classes were thresholded + # out of the legacy dict, the rebuilt vector is SHORTER than the model's + # original class count (serialization-stable: the surviving set round-trips + # exactly; only the trailing all-zero tail is unrecoverable). + num_classes = (max(entry_names.keys()) + 1) if entry_names else 0 + class_names = { + class_id: entry_names.get(class_id, str(class_id)) + for class_id in range(num_classes) + } + confidence_vector = [entry_confidences.get(i, 0.0) for i in range(num_classes)] + top_class = prediction.get("top") + top_class_id = next( + (cid for cid, name in entry_names.items() if name == top_class), + ( + int(max(entry_confidences, key=entry_confidences.get)) + if entry_confidences + else 0 + ), + ) + metadata = _classification_metadata_from_dict( + prediction=prediction, class_names=class_names + ) + # Re-attaching the smallest listed confidence as the serializer threshold + # keeps the gap-filled zero classes out of re-serialization, so the + # round-tripped dict matches the input (edge: a listed 0.0-confidence entry + # means no threshold can be attached and gap classes reappear). + positive_confidences = [c for c in entry_confidences.values() if c > 0.0] + if entry_confidences and len(positive_confidences) == len(entry_confidences): + metadata[CLASSIFICATION_CONFIDENCE_THRESHOLD_KEY] = min(positive_confidences) + return ClassificationPrediction( + class_id=torch.tensor( + [top_class_id], dtype=torch.long, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + confidence=torch.tensor( + [confidence_vector], + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + images_metadata=[metadata], + ) + + +def _multi_label_dict_to_native( + prediction: dict, +) -> MultiLabelClassificationPrediction: + entries = prediction.get("predictions") or {} + name_to_id = { + str(name): int(payload["class_id"]) for name, payload in entries.items() + } + num_classes = (max(name_to_id.values()) + 1) if name_to_id else 0 + class_names = {class_id: str(class_id) for class_id in range(num_classes)} + confidence_vector = [0.0] * num_classes + for name, payload in entries.items(): + class_id = int(payload["class_id"]) + class_names[class_id] = str(name) + confidence_vector[class_id] = float(payload.get("confidence", 0.0)) + predicted_class_ids = [ + name_to_id[str(name)] + for name in prediction.get("predicted_classes") or [] + if str(name) in name_to_id + ] + metadata = _classification_metadata_from_dict( + prediction=prediction, class_names=class_names + ) + return MultiLabelClassificationPrediction( + class_ids=torch.tensor( + predicted_class_ids, + dtype=torch.long, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + confidence=torch.tensor( + confidence_vector, + dtype=torch.float32, + device=WORKFLOWS_IMAGE_TENSOR_DEVICE, + ), + # MultiLabel carries SINGULAR image_metadata (dict). + image_metadata=metadata, + ) diff --git a/inference/core/workflows/execution_engine/v1/executor/output_constructor.py b/inference/core/workflows/execution_engine/v1/executor/output_constructor.py index 265102b5cf..fa7bddebb2 100644 --- a/inference/core/workflows/execution_engine/v1/executor/output_constructor.py +++ b/inference/core/workflows/execution_engine/v1/executor/output_constructor.py @@ -9,6 +9,10 @@ from networkx import DiGraph from inference.core import logger +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION +from inference.core.workflows.core_steps.common.tensor_native import ( + native_detections_to_root_coordinates, +) from inference.core.workflows.core_steps.common.utils import ( sv_detections_to_root_coordinates, ) @@ -42,6 +46,13 @@ maybe_resolve_futures, resolve_futures, ) +from inference_models.models.base.instance_segmentation import ( + InstanceDetections as NativeInstanceDetections, +) +from inference_models.models.base.keypoints_detection import ( + KeyPoints as NativeKeyPoints, +) +from inference_models.models.base.object_detection import Detections as NativeDetections def construct_workflow_output( @@ -383,6 +394,12 @@ def _convert_sv_detections_coordinates_if_present(data: Any) -> Any: def data_needs_sv_detections_coordinate_conversion(data: Any) -> bool: if isinstance(data, sv.Detections): return _sv_detections_need_root_coordinate_conversion(detections=data) + if ENABLE_TENSOR_DATA_REPRESENTATION and _is_native_prediction(data): + # native_detections_to_root_coordinates() no-ops on predictions that + # carry no root offset, so natives are always admitted here; this + # branch must precede the generic tuple recursion below, which would + # otherwise swallow the (KeyPoints, Detections) prediction tuple. + return True if isinstance(data, (list, tuple)): return any(data_needs_sv_detections_coordinate_conversion(data=e) for e in data) if isinstance(data, dict): @@ -555,9 +572,26 @@ def place_data_in_array(array: list, index: DynamicBatchIndex, data: Any) -> Non place_data_in_array(array=array[first_chunk], index=remaining_index, data=data) +def _is_native_prediction(data: Any) -> bool: + """Recognise a tensor-native prediction (the inference_models dataclasses, or the + ``(KeyPoints, Detections)`` keypoint-detection tuple) emitted by ``_tensor`` + blocks under ``ENABLE_TENSOR_DATA_REPRESENTATION``. These are not ``sv.Detections``, + so the existing sv-only coordinate-conversion gate skips them.""" + if isinstance(data, (NativeDetections, NativeInstanceDetections, NativeKeyPoints)): + return True + if isinstance(data, tuple) and len(data) == 2: + key_points, detections = data + return isinstance(key_points, NativeKeyPoints) and ( + detections is None or isinstance(detections, NativeDetections) + ) + return False + + def data_contains_sv_detections(data: Any) -> bool: if isinstance(data, sv.Detections): return True + if ENABLE_TENSOR_DATA_REPRESENTATION and _is_native_prediction(data): + return True if isinstance(data, dict): result = set() for value in data.values(): @@ -574,6 +608,8 @@ def data_contains_sv_detections(data: Any) -> bool: def convert_sv_detections_coordinates(data: Any) -> Any: if isinstance(data, sv.Detections): return sv_detections_to_root_coordinates(detections=data) + if ENABLE_TENSOR_DATA_REPRESENTATION and _is_native_prediction(data): + return native_detections_to_root_coordinates(prediction=data) if isinstance(data, dict): return {k: convert_sv_detections_coordinates(data=v) for k, v in data.items()} if isinstance(data, list): diff --git a/inference/models/clip/clip_inference_models.py b/inference/models/clip/clip_inference_models.py index 4cb9bf2d3e..85f892efc5 100644 --- a/inference/models/clip/clip_inference_models.py +++ b/inference/models/clip/clip_inference_models.py @@ -1,8 +1,10 @@ from time import perf_counter -from typing import Any, Dict, List, Tuple, Union +from typing import Any, Dict, List, Literal, Tuple, Union import numpy as np import onnxruntime +import torch +import torch.nn.functional as F from inference.core.entities.requests.clip import ( ClipCompareRequest, @@ -79,6 +81,33 @@ def __init__( **kwargs, ) + def run_tensor_native_inference( + self, action: Literal["compare", "embed-image", "embed-text"], **kwargs + ) -> torch.Tensor: + if action == "embed-image": + return self._model.embed_images(**kwargs) + elif action == "embed-text": + return self._model.embed_text(**kwargs) + subject_type = kwargs.get("subject_type", "image") + prompt_type = kwargs.get("prompt_type", "text") + if subject_type == "image": + subject_embeddings = self._model.embed_images( + images=kwargs["subject"], **kwargs + ) + else: + subject_embeddings = self._model.embed_text( + text=kwargs["subject"], **kwargs + ) + if prompt_type == "image": + prompt_embeddings = self._model.embed_images( + images=kwargs["prompt"], **kwargs + ) + else: + prompt_embeddings = self._model.embed_text(text=kwargs["prompt"], **kwargs) + subject_embeddings_norm = F.normalize(subject_embeddings, dim=1) + prompt_embeddings_norm = F.normalize(prompt_embeddings, dim=1) + return subject_embeddings_norm @ prompt_embeddings_norm.T + def compare( self, subject: Any, diff --git a/inference/models/depth_anything_v2/depth_anything_v2_inference_models.py b/inference/models/depth_anything_v2/depth_anything_v2_inference_models.py index f0478d5a1c..477ca002ac 100644 --- a/inference/models/depth_anything_v2/depth_anything_v2_inference_models.py +++ b/inference/models/depth_anything_v2/depth_anything_v2_inference_models.py @@ -1,4 +1,4 @@ -from typing import Any, List, Tuple +from typing import Any, List, Tuple, Union from uuid import uuid4 import cv2 @@ -60,6 +60,13 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[torch.Tensor]: + return self._model(images, **kwargs) + def preprocess(self, image: Any, **kwargs): if isinstance(image, list): raise ValueError("DepthAnythingV2 does not support batched inference.") diff --git a/inference/models/depth_anything_v3/depth_anything_v3_inference_models.py b/inference/models/depth_anything_v3/depth_anything_v3_inference_models.py index 2fd4fb1eee..48d4e25e56 100644 --- a/inference/models/depth_anything_v3/depth_anything_v3_inference_models.py +++ b/inference/models/depth_anything_v3/depth_anything_v3_inference_models.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, List, Tuple +from typing import Any, List, Tuple, Union from uuid import uuid4 import matplotlib.pyplot as plt @@ -73,6 +73,13 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[torch.Tensor]: + return self._model(images, **kwargs) + def preprocess(self, image: Any, **kwargs): if isinstance(image, list): raise ValueError("DepthAnythingV3 does not support batched inference.") diff --git a/inference/models/doctr/doctr_model_inference_models.py b/inference/models/doctr/doctr_model_inference_models.py index 75a6de5c0c..ec7dd03858 100644 --- a/inference/models/doctr/doctr_model_inference_models.py +++ b/inference/models/doctr/doctr_model_inference_models.py @@ -2,6 +2,8 @@ from time import perf_counter from typing import Any, List, Tuple, Union +import numpy as np +import torch from PIL import Image from inference.core.entities.requests.doctr import DoctrOCRInferenceRequest @@ -20,7 +22,7 @@ from inference.core.models.base import Model from inference.core.roboflow_api import get_extra_weights_provider_headers from inference.core.utils.image_utils import load_image_bgr -from inference_models import AutoModel +from inference_models import AutoModel, Detections from inference_models.models.doctr.doctr_torch import DocTR @@ -56,6 +58,13 @@ def __init__( **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs + ) -> Tuple[List[str], List[Detections]]: + return self._model.infer(images=images, **kwargs) + def clear_cache(self, delete_from_disk: bool = True) -> None: pass diff --git a/inference/models/easy_ocr/easy_ocr_inference_models.py b/inference/models/easy_ocr/easy_ocr_inference_models.py index efaf057799..14269988c1 100644 --- a/inference/models/easy_ocr/easy_ocr_inference_models.py +++ b/inference/models/easy_ocr/easy_ocr_inference_models.py @@ -4,6 +4,7 @@ from typing import Any, List, Tuple, Union import numpy as np +import torch from inference.core.entities.requests.easy_ocr import EasyOCRInferenceRequest from inference.core.entities.responses.inference import ( @@ -63,6 +64,13 @@ def __init__( **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> Tuple[List[str], List[Detections]]: + return self._model.infer(images=images, **kwargs) + def predict(self, image_in: np.ndarray, **kwargs) -> Tuple[str, Detections]: parsed_texts, parsed_structures = self._model.infer(images=image_in, **kwargs) parsed_text = parsed_texts[0] @@ -105,6 +113,7 @@ def single_request(self, request: EasyOCRInferenceRequest) -> OCRInferenceRespon predictions_for_image = [] for instance_id in range(prediction_result[1].xyxy.shape[0]): x_min, y_min, x_max, y_max = prediction_result[1].xyxy[instance_id].tolist() + instance_confidence = prediction_result[1].confidence[instance_id].item() width = x_max - x_min height = y_max - y_min center_x = (x_min + x_max) / 2 @@ -117,7 +126,7 @@ def single_request(self, request: EasyOCRInferenceRequest) -> OCRInferenceRespon "y": center_y, "width": width, "height": height, - "confidence": 1.0, # confidence is not returned by the model + "confidence": instance_confidence, "class": prediction_result[1].bboxes_metadata[instance_id][ "text" ], diff --git a/inference/models/florence2/florence2_inference_models.py b/inference/models/florence2/florence2_inference_models.py index a51d5cbaf9..2225614b6f 100644 --- a/inference/models/florence2/florence2_inference_models.py +++ b/inference/models/florence2/florence2_inference_models.py @@ -1,5 +1,6 @@ -from typing import Any, List +from typing import Any, List, Union +import numpy as np import torch from inference.core.entities.responses import ( @@ -49,6 +50,23 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[str]: + kwargs = self.map_inference_kwargs(kwargs) + # Derive the Florence task token from the leading `` marker of the + # prompt. NOTE: `"...".split(">")[0] + ">"` always ends with ">", so the + # previous `if not task: task = None` guard was dead code and is removed. + # For a plain-text prompt (no `` token) this yields a non-Florence + # key, but `post_process_generation` still wraps the answer under that key + # and the caller (Florence2 block) unwraps the single key, so the output + # matches the numpy sibling (whose `postprocess` derives the key the same + # way). This mirrors `postprocess` below; keep them in lockstep. + task = kwargs.get("prompt", "").split(">")[0] + ">" + return self._model.prompt(images=images, task=task, **kwargs) + def map_inference_kwargs(self, kwargs: dict) -> dict: pre_processing_overrides = PreProcessingOverrides( disable_contrast_enhancement=kwargs.get("disable_preproc_contrast", False), diff --git a/inference/models/glm_ocr/glm_ocr_inference_models.py b/inference/models/glm_ocr/glm_ocr_inference_models.py index ae178394b1..c63a26616a 100644 --- a/inference/models/glm_ocr/glm_ocr_inference_models.py +++ b/inference/models/glm_ocr/glm_ocr_inference_models.py @@ -1,5 +1,6 @@ -from typing import Any, List, Optional +from typing import Any, List, Optional, Union +import numpy as np import torch from inference.core.entities.responses import ( @@ -50,6 +51,13 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[str]: + return self._model.prompt(images=images, **kwargs) + def preprocess(self, image: Any, prompt: Optional[str] = None, **kwargs): is_batch = isinstance(image, list) if is_batch: diff --git a/inference/models/grounding_dino/grounding_dino_inference_models.py b/inference/models/grounding_dino/grounding_dino_inference_models.py index b9920fd24a..2a63dace33 100644 --- a/inference/models/grounding_dino/grounding_dino_inference_models.py +++ b/inference/models/grounding_dino/grounding_dino_inference_models.py @@ -1,5 +1,8 @@ from time import perf_counter -from typing import Any, List +from typing import Any, List, Union + +import numpy as np +import torch from inference.core.entities.requests.groundingdino import GroundingDINOInferenceRequest from inference.core.entities.requests.inference import InferenceRequestImage @@ -19,7 +22,7 @@ from inference.core.models.base import Model from inference.core.roboflow_api import get_extra_weights_provider_headers from inference.core.utils.image_utils import load_image_bgr, xyxy_to_xywh -from inference_models import AutoModel +from inference_models import AutoModel, Detections from inference_models.models.grounding_dino.grounding_dino_torch import ( GroundingDinoForObjectDetectionTorch, ) @@ -65,6 +68,13 @@ def __init__( **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs + ) -> List[Detections]: + return self._model(images=images, **kwargs) + def preproc_image(self, image: Any): """Preprocesses an image. diff --git a/inference/models/moondream2/moondream2_inference_models.py b/inference/models/moondream2/moondream2_inference_models.py index 0b2a28d1b6..cd6137031f 100644 --- a/inference/models/moondream2/moondream2_inference_models.py +++ b/inference/models/moondream2/moondream2_inference_models.py @@ -2,6 +2,7 @@ import cv2 import numpy as np +import torch from PIL import Image from inference.core.entities.responses.inference import ( @@ -52,6 +53,13 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[Detections]: + return self._model.detect(images=images, **kwargs) + def preprocess(self, image: Any, **kwargs): is_batch = isinstance(image, list) if is_batch: diff --git a/inference/models/owlv2/rf_instant_inference_models.py b/inference/models/owlv2/rf_instant_inference_models.py index e42437a9e6..5dafc95691 100644 --- a/inference/models/owlv2/rf_instant_inference_models.py +++ b/inference/models/owlv2/rf_instant_inference_models.py @@ -1,4 +1,7 @@ -from typing import List, Optional, Tuple +from typing import List, Optional, Tuple, Union + +import numpy as np +import torch from inference.core.entities.requests import ObjectDetectionInferenceRequest from inference.core.entities.responses import ( @@ -86,6 +89,13 @@ def __init__( **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[Detections]: + return self._model.infer(images=images, **kwargs) + def infer( self, image, diff --git a/inference/models/paligemma/paligemma_inference_models.py b/inference/models/paligemma/paligemma_inference_models.py index 25a094c21a..f329d1f00a 100644 --- a/inference/models/paligemma/paligemma_inference_models.py +++ b/inference/models/paligemma/paligemma_inference_models.py @@ -1,5 +1,6 @@ -from typing import Any, List +from typing import Any, List, Union +import numpy as np import torch from inference.core.entities.responses import ( @@ -52,6 +53,14 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[str]: + kwargs = self.map_inference_kwargs(kwargs) + return self._model.prompt(images=images, **kwargs) + def map_inference_kwargs(self, kwargs: dict) -> dict: pre_processing_overrides = PreProcessingOverrides( disable_contrast_enhancement=kwargs.get("disable_preproc_contrast", False), diff --git a/inference/models/perception_encoder/perception_encoder_inference_models.py b/inference/models/perception_encoder/perception_encoder_inference_models.py index 67865868dd..d4367b959b 100644 --- a/inference/models/perception_encoder/perception_encoder_inference_models.py +++ b/inference/models/perception_encoder/perception_encoder_inference_models.py @@ -1,8 +1,9 @@ from time import perf_counter -from typing import Any, Dict, List, Tuple, Union +from typing import Any, Dict, List, Literal, Tuple, Union import numpy as np import torch +import torch.nn.functional as F from inference.core.entities.requests.inference import InferenceRequestImage from inference.core.entities.requests.perception_encoder import ( @@ -232,6 +233,33 @@ def make_embed_text_response( response = PerceptionEncoderEmbeddingResponse(embeddings=embeddings.tolist()) return response + def run_tensor_native_inference( + self, action: Literal["compare", "embed-image", "embed-text"], **kwargs + ) -> torch.Tensor: + if action == "embed-image": + return self._model.embed_images(**kwargs) + elif action == "embed-text": + return self._model.embed_text(**kwargs) + subject_type = kwargs.get("subject_type", "image") + prompt_type = kwargs.get("prompt_type", "text") + if subject_type == "image": + subject_embeddings = self._model.embed_images( + images=kwargs["subject"], **kwargs + ) + else: + subject_embeddings = self._model.embed_text( + text=kwargs["subject"], **kwargs + ) + if prompt_type == "image": + prompt_embeddings = self._model.embed_images( + images=kwargs["prompt"], **kwargs + ) + else: + prompt_embeddings = self._model.embed_text(text=kwargs["prompt"], **kwargs) + subject_embeddings_norm = F.normalize(subject_embeddings, dim=1) + prompt_embeddings_norm = F.normalize(prompt_embeddings, dim=1) + return subject_embeddings_norm @ prompt_embeddings_norm.T + def infer_from_request( self, request: PerceptionEncoderInferenceRequest ) -> PerceptionEncoderEmbeddingResponse: diff --git a/inference/models/qwen25vl/qwen25vl_inference_models.py b/inference/models/qwen25vl/qwen25vl_inference_models.py index e6a0850ae9..6a83eb3d76 100644 --- a/inference/models/qwen25vl/qwen25vl_inference_models.py +++ b/inference/models/qwen25vl/qwen25vl_inference_models.py @@ -1,5 +1,6 @@ -from typing import Any, List +from typing import Any, List, Union +import numpy as np import torch from inference.core.entities.responses import ( @@ -51,6 +52,14 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[str]: + kwargs = self.map_inference_kwargs(kwargs) + return self._model.prompt(images=images, **kwargs) + def map_inference_kwargs(self, kwargs: dict) -> dict: pre_processing_overrides = PreProcessingOverrides( disable_contrast_enhancement=kwargs.get("disable_preproc_contrast", False), diff --git a/inference/models/qwen3_5vl/qwen3_5vl_inference_models.py b/inference/models/qwen3_5vl/qwen3_5vl_inference_models.py index 64ab0285d1..8e69668d2c 100644 --- a/inference/models/qwen3_5vl/qwen3_5vl_inference_models.py +++ b/inference/models/qwen3_5vl/qwen3_5vl_inference_models.py @@ -1,5 +1,6 @@ -from typing import Any, List +from typing import Any, List, Union +import numpy as np import torch from inference.core.entities.responses import ( @@ -46,6 +47,14 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): def map_inference_kwargs(self, kwargs: dict) -> dict: return kwargs + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[str]: + kwargs = self.map_inference_kwargs(kwargs) + return self._model.prompt(images=images, **kwargs) + def preprocess(self, image: Any, prompt: str = "", **kwargs): is_batch = isinstance(image, list) if is_batch: diff --git a/inference/models/qwen3vl/qwen3vl_inference_models.py b/inference/models/qwen3vl/qwen3vl_inference_models.py index 06b7c6b57d..10b52b62e5 100644 --- a/inference/models/qwen3vl/qwen3vl_inference_models.py +++ b/inference/models/qwen3vl/qwen3vl_inference_models.py @@ -1,5 +1,6 @@ -from typing import Any, List +from typing import Any, List, Union +import numpy as np import torch from inference.core.entities.responses import ( @@ -50,6 +51,14 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[str]: + kwargs = self.map_inference_kwargs(kwargs) + return self._model.prompt(images=images, **kwargs) + def map_inference_kwargs(self, kwargs: dict) -> dict: pre_processing_overrides = PreProcessingOverrides( disable_contrast_enhancement=kwargs.get("disable_preproc_contrast", False), diff --git a/inference/models/sam/segment_anything_inference_models.py b/inference/models/sam/segment_anything_inference_models.py index 8d37379d69..da8d60a5d2 100644 --- a/inference/models/sam/segment_anything_inference_models.py +++ b/inference/models/sam/segment_anything_inference_models.py @@ -1,7 +1,7 @@ import base64 from io import BytesIO from time import perf_counter -from typing import Any, List, Optional, Tuple, Union +from typing import Any, List, Literal, Optional, Tuple, Union import numpy as np import rasterio @@ -35,7 +35,7 @@ SamImageEmbeddingsInMemoryCache, SamLowResolutionMasksInMemoryCache, ) -from inference_models.models.sam.entities import SAMImageEmbeddings +from inference_models.models.sam.entities import SAMImageEmbeddings, SAMPrediction from inference_models.models.sam.sam_torch import SAMTorch, compute_image_hash MASK_THRESHOLD = 0.0 @@ -89,6 +89,14 @@ def __init__( def map_inference_kwargs(self, kwargs: dict) -> dict: return kwargs + def run_tensor_native_inference( + self, action: Literal["embed", "segment"], **kwargs + ) -> List[Union[SAMImageEmbeddings, SAMPrediction]]: + kwargs = self.map_inference_kwargs(kwargs) + if action == "embed": + return self._model.embed_images(**kwargs) + return self._model.segment_images(**kwargs) + def infer_from_request(self, request: SamInferenceRequest): t1 = perf_counter() if isinstance(request, SamEmbeddingRequest): diff --git a/inference/models/sam2/segment_anything2_inference_models.py b/inference/models/sam2/segment_anything2_inference_models.py index c79c3ed69d..3dab74ec0b 100644 --- a/inference/models/sam2/segment_anything2_inference_models.py +++ b/inference/models/sam2/segment_anything2_inference_models.py @@ -3,7 +3,7 @@ from io import BytesIO from threading import RLock from time import perf_counter -from typing import Any, Dict, List, Optional, Tuple, TypedDict, TypeVar, Union +from typing import Any, Dict, List, Literal, Optional, Tuple, TypedDict, TypeVar, Union import numpy as np import sam2.utils.misc @@ -18,6 +18,7 @@ Sam2ImageEmbeddingsInMemoryCache, Sam2LowResolutionMasksInMemoryCache, ) +from inference_models.models.sam2.entities import SAM2ImageEmbeddings, SAM2Prediction from inference_models.models.sam2.sam2_torch import SAM2Torch sam2.utils.misc.get_sdp_backends = lambda z: [ @@ -125,6 +126,13 @@ def __init__( **kwargs, ) + def run_tensor_native_inference( + self, action: Literal["embed", "segment"], **kwargs + ) -> List[Union[SAM2ImageEmbeddings, SAM2Prediction]]: + if action == "embed": + return self._model.embed_images(**kwargs) + return self._model.segment_images(**kwargs) + @usage_collector("model") def infer_from_request(self, request: Sam2InferenceRequest): """Performs inference based on the request type. diff --git a/inference/models/sam3/segment_anything3_inference_models.py b/inference/models/sam3/segment_anything3_inference_models.py index 4798c73e3c..8516eda6a5 100644 --- a/inference/models/sam3/segment_anything3_inference_models.py +++ b/inference/models/sam3/segment_anything3_inference_models.py @@ -91,6 +91,20 @@ def infer_from_request(self, request: Sam3InferenceRequest): ) raise ValueError(f"Invalid request type {type(request)}") + def run_tensor_native_inference(self, **kwargs) -> List[List[Dict]]: + """Minimal tensor-native bridge to the inference_models SAM3 model function. + + Forwards straight to ``SAM3Torch.segment_with_text_prompts`` โ€” the library + model function โ€” which accepts CHW image tensors directly (its + ``_normalize_to_hwc_uint8`` transposes channels-first to HWC and rescales). + The workflow block passes ``images`` (tensor) / ``prompts`` / + ``output_prob_thresh`` as kwargs and owns all downstream shaping (per-prompt + threshold, cross-prompt NMS, InstanceDetections build); this method performs + none of the ``load_image_rgb`` / polygon-RLE response work that + ``segment_image`` does. + """ + return self._model.segment_with_text_prompts(**kwargs) + def segment_image( self, image: InferenceRequestImage, diff --git a/inference/models/sam3/visual_segmentation_inference_models.py b/inference/models/sam3/visual_segmentation_inference_models.py index b8715efb02..f3028f004f 100644 --- a/inference/models/sam3/visual_segmentation_inference_models.py +++ b/inference/models/sam3/visual_segmentation_inference_models.py @@ -1,7 +1,7 @@ import copy from io import BytesIO from time import perf_counter -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union import numpy as np import torch @@ -42,6 +42,7 @@ Sam3ImageEmbeddingsInMemoryCache, Sam3LowResolutionMasksInMemoryCache, ) +from inference_models.models.sam3.entities import SAM3ImageEmbeddings, SAM3Prediction from inference_models.models.sam3.sam3_torch import SAM3Torch if DEVICE is None: @@ -103,6 +104,13 @@ def __init__( **kwargs, ) + def run_tensor_native_inference( + self, action: Literal["embed", "segment"], **kwargs + ) -> List[Union[SAM3ImageEmbeddings, SAM3Prediction]]: + if action == "embed": + return self._model.embed_images(**kwargs) + return self._model.segment_with_visual_prompts(**kwargs) + @usage_collector("model") def infer_from_request(self, request: Sam2InferenceRequest): t1 = perf_counter() diff --git a/inference/models/smolvlm/smolvlm_inference_models.py b/inference/models/smolvlm/smolvlm_inference_models.py index a2ad3fd22f..f31a390343 100644 --- a/inference/models/smolvlm/smolvlm_inference_models.py +++ b/inference/models/smolvlm/smolvlm_inference_models.py @@ -1,5 +1,6 @@ -from typing import Any, List +from typing import Any, List, Union +import numpy as np import torch from inference.core.entities.responses import ( @@ -55,6 +56,14 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[str]: + kwargs = self.map_inference_kwargs(kwargs) + return self._model.prompt(images=images, **kwargs) + def map_inference_kwargs(self, kwargs: dict) -> dict: pre_processing_overrides = PreProcessingOverrides( disable_contrast_enhancement=kwargs.get("disable_preproc_contrast", False), diff --git a/inference/models/trocr/trocr_inference_models.py b/inference/models/trocr/trocr_inference_models.py index 5158e51069..3513d88d9c 100644 --- a/inference/models/trocr/trocr_inference_models.py +++ b/inference/models/trocr/trocr_inference_models.py @@ -1,7 +1,8 @@ from time import perf_counter -from typing import Any, Tuple +from typing import Any, List, Tuple, Union import numpy as np +import torch from inference.core.entities.requests.trocr import TrOCRInferenceRequest from inference.core.entities.responses.ocr import OCRInferenceResponse @@ -48,6 +49,13 @@ def __init__(self, model_id: str, api_key: str = None, **kwargs): **kwargs, ) + def run_tensor_native_inference( + self, + images: Union[torch.Tensor, List[torch.Tensor], np.ndarray, List[np.ndarray]], + **kwargs, + ) -> List[str]: + return self._model(images=images, **kwargs) + def preprocess( self, image: Any, **kwargs ) -> Tuple[np.ndarray, PreprocessReturnMetadata]: diff --git a/inference_cli/lib/container_adapter.py b/inference_cli/lib/container_adapter.py index 49d2562488..18b8d0421b 100644 --- a/inference_cli/lib/container_adapter.py +++ b/inference_cli/lib/container_adapter.py @@ -98,7 +98,12 @@ class _JetsonImage(NamedTuple): # Ordered by (l4t_major DESC, l4t_minor_min DESC). First match wins. # See https://developer.nvidia.com/embedded/jetpack-archive for the full mapping. _JETSON_IMAGES: List[_JetsonImage] = [ - _JetsonImage(38, 0, "7", "roboflow/roboflow-inference-server-jetson-7.1.0:latest"), + _JetsonImage( + 39, 2, "7.2", "roboflow/roboflow-inference-server-jetson-7.2.0:latest" + ), + _JetsonImage( + 38, 4, "7.1", "roboflow/roboflow-inference-server-jetson-7.1.0:latest" + ), _JetsonImage( 36, 4, "6.2", "roboflow/roboflow-inference-server-jetson-6.2.0:latest" ), diff --git a/inference_models/dockerfiles/jp51.cu114.core.dockerfile b/inference_models/dockerfiles/jp51.cu114.core.dockerfile index 3102805015..b9eb9a96dd 100644 --- a/inference_models/dockerfiles/jp51.cu114.core.dockerfile +++ b/inference_models/dockerfiles/jp51.cu114.core.dockerfile @@ -1,5 +1,14 @@ FROM nvcr.io/nvidia/l4t-ml:r35.2.1-py3 AS builder +# Tell apt that libopencv-dev is already satisfied (the CUDA OpenCV in +# opencv-dev 4.5.0 from the base image covers it). This prevents apt from +# pulling in NVIDIA's libopencv-dev which conflicts on file ownership. +RUN mkdir -p /tmp/dummy/DEBIAN \ + && printf 'Package: libopencv-dev\nVersion: 4.5.0\nArchitecture: arm64\nMaintainer: local \nDescription: dummy satisfying libopencv-dev (real CUDA OpenCV lives in opencv-dev)\n' > /tmp/dummy/DEBIAN/control \ + && dpkg-deb --build /tmp/dummy /tmp/libopencv-dev-dummy.deb \ + && dpkg -i /tmp/libopencv-dev-dummy.deb \ + && rm -rf /tmp/dummy /tmp/libopencv-dev-dummy.deb + # install Python 3.12 RUN apt-get update -y && apt-get install -y \ libssl-dev \ @@ -16,13 +25,16 @@ RUN apt-get update -y && apt-get install -y \ libswscale-dev \ libavutil-dev \ libgstreamer1.0-dev \ - libgstreamer-plugins-base1.0-dev + libgstreamer-plugins-base1.0-dev \ + libgstreamer-plugins-bad1.0-dev \ + libjson-glib-dev \ + libnice-dev RUN mkdir -p /build/python-3.12 WORKDIR /build/python-3.12 RUN wget https://www.python.org/ftp/python/3.12.12/Python-3.12.12.tgz && tar -xzf Python-3.12.12.tgz WORKDIR /build/python-3.12/Python-3.12.12 -RUN ./configure --enable-optimizations +RUN ./configure --enable-optimizations --enable-shared --prefix=/usr/local LDFLAGS="-Wl,-rpath,/usr/local/lib" RUN make -j$(nproc) && make altinstall RUN update-alternatives --install /usr/bin/python python /usr/local/bin/python3.12 1 @@ -158,6 +170,24 @@ RUN CC=/root/GCC-11/bin/gcc CXX=/root/GCC-11/bin/g++ FORCE_CUDA=1 PATH=/build/cm RUN python3.12 -m pip install dist/torchvision-*.whl RUN cp dist/torchvision-*.whl /build/out/wheels/ +RUN apt-get update -y && apt-get install -y \ + libglew-dev \ + libgstrtspserver-1.0-dev \ + && mkdir -p /build/jetson-utils \ + && cd /build/jetson-utils \ + && git clone https://github.com/dusty-nv/jetson-utils \ + && cd /build/jetson-utils/jetson-utils \ + && git checkout fae4b4250f985ab92180b6ec955b79995b7a34ff \ + && sed -i 's/2\.7 3\.6 3\.7 3\.8 3\.10 3\.12/3.12/' python/CMakeLists.txt \ + && mkdir build \ + && cd /build/jetson-utils/jetson-utils/build \ + && cmake ../ \ + && make -j$(nproc) \ + && make install \ + && ldconfig \ + && python3.12 -m pip install termcolor tabulate docker \ + && echo "/usr/lib/python3.12/dist-packages" > /usr/local/lib/python3.12/site-packages/jetson_utils.pth + FROM nvcr.io/nvidia/l4t-ml:r35.2.1-py3 AS target RUN apt-get update -y && apt-get install -y \ @@ -195,6 +225,15 @@ ENV LD_LIBRARY_PATH="/opt/gcc-11/lib64:$$LD_LIBRARY_PATH" RUN update-alternatives --install /usr/bin/python python /usr/local/bin/python3.12 1 RUN update-alternatives --install /usr/bin/python3 python3 /usr/local/bin/python3.12 1 +# Tell apt that libopencv-dev is already satisfied (the CUDA OpenCV in +# opencv-dev 4.5.0 from the base image covers it). This prevents apt from +# pulling in NVIDIA's libopencv-dev which conflicts on file ownership. +RUN mkdir -p /tmp/dummy/DEBIAN \ + && printf 'Package: libopencv-dev\nVersion: 4.5.0\nArchitecture: arm64\nMaintainer: local \nDescription: dummy satisfying libopencv-dev (real CUDA OpenCV lives in opencv-dev)\n' > /tmp/dummy/DEBIAN/control \ + && dpkg-deb --build /tmp/dummy /tmp/libopencv-dev-dummy.deb \ + && dpkg -i /tmp/libopencv-dev-dummy.deb \ + && rm -rf /tmp/dummy /tmp/libopencv-dev-dummy.deb + # Install ffmpeg/gstreamer dev packages for OpenCV build (must be after COPY from builder) RUN apt-get update -y && apt-get install -y --no-install-recommends \ pkg-config \ @@ -202,8 +241,13 @@ RUN apt-get update -y && apt-get install -y --no-install-recommends \ libavformat-dev \ libswscale-dev \ libavutil-dev \ - libgstreamer1.0-dev \ - libgstreamer-plugins-base1.0-dev \ + libgstreamer1.0 \ + libgstreamer-plugins-base1.0 \ + libgstreamer-plugins-bad1.0 \ + libjson-glib-1.0-0 \ + libnice10 \ + libglew2.1 \ + libgstrtspserver-1.0 \ && rm -rf /var/lib/apt/lists/* # Install OpenCV diff --git a/inference_models/docs/changelog.md b/inference_models/docs/changelog.md index 2b4362a35c..7364a879e5 100644 --- a/inference_models/docs/changelog.md +++ b/inference_models/docs/changelog.md @@ -13,6 +13,9 @@ have explicit `base` implementations with typed metadata and runtime selection records, while preserving the existing tensor ownership, CUDA stream/event behavior, and protected TensorRT forward path. +- Adjustment to `KeyPoints` interface to expose `__len__(...)` method. +- Adjustment to `InstanceDetections` interface to expose `__len__(...)` and `__iter__(...)` method. +- Adjustment to `Detections` interface to expose `__len__(...)` and `__iter__(...)` method. ### Fixed @@ -301,6 +304,7 @@ otherwise float16). - Transitive dependency vulnerability patched - `idna>=3.15` required by the package --- + ## `0.29.1` ### Fixed diff --git a/inference_models/inference_models/models/base/classification.py b/inference_models/inference_models/models/base/classification.py index 37ca50c799..8f41d4ed21 100644 --- a/inference_models/inference_models/models/base/classification.py +++ b/inference_models/inference_models/models/base/classification.py @@ -11,7 +11,7 @@ @dataclass class ClassificationPrediction: class_id: torch.Tensor # (bs, ) - confidence: torch.Tensor # (bs, ) + confidence: torch.Tensor # (bs, num_classes) โ€” full softmax distribution images_metadata: Optional[List[dict]] = None # if given, list of size equal to bs @@ -73,8 +73,8 @@ def __call__( @dataclass class MultiLabelClassificationPrediction: - class_ids: torch.Tensor # (predicted_labels_ids, ) - confidence: torch.Tensor # (predicted_labels_confidence, ) + class_ids: torch.Tensor # (predicted_labels_ids, ) โ€” indices of above-threshold classes + confidence: torch.Tensor # (num_classes, ) โ€” full sigmoid distribution image_metadata: Optional[dict] = None diff --git a/inference_models/inference_models/models/base/instance_segmentation.py b/inference_models/inference_models/models/base/instance_segmentation.py index 2bc522d5a1..123f563d99 100644 --- a/inference_models/inference_models/models/base/instance_segmentation.py +++ b/inference_models/inference_models/models/base/instance_segmentation.py @@ -10,6 +10,7 @@ Set, Tuple, Union, + Iterator, runtime_checkable, ) @@ -144,6 +145,47 @@ class InstanceDetections: None # if given, list of size equal to # of bboxes ) + def __len__(self) -> int: + return int(self.xyxy.shape[0]) + + def __iter__(self) -> Iterator[Tuple]: + """Iterates detections yielding 7-tuples: + (xyxy, mask, class_id, confidence, tracker_id, data, metadata) + + - xyxy: torch.Tensor of shape (4,) - single bbox [x1, y1, x2, y2] + - mask: per-instance dense torch.Tensor of shape (H, W) or coco-representation + of RLE mask: {"size": (h, w), "counts": count} + - class_id: scalar tensor (0-dim) + - confidence: scalar tensor (0-dim) + - tracker_id: value of `bboxes_metadata[i]["tracker_id"]` or None + - data: per-detection dict (`bboxes_metadata[i]`, `{}` if not set) + - metadata: per-image dict (`image_metadata`, `{}` if not set) + """ + bboxes_metadata = self.bboxes_metadata + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(len(self))] + image_metadata = self.image_metadata or {} + for index in range(len(self)): + data = bboxes_metadata[index] + if self.mask is None: + selected_mask = None + elif isinstance(self.mask, InstancesRLEMasks): + selected_mask = { + "size": list(self.mask.image_size), + "counts": self.mask.masks[index], + } + else: + selected_mask = self.mask[index] + yield ( + self.xyxy[index], + selected_mask, + self.class_id[index], + self.confidence[index], + data.get("tracker_id"), + data, + image_metadata, + ) + def to_supervision(self) -> sv.Detections: """Convert instance segmentation detections to Supervision Detections format. diff --git a/inference_models/inference_models/models/base/keypoints_detection.py b/inference_models/inference_models/models/base/keypoints_detection.py index 66cff35c2e..9b6e14791f 100644 --- a/inference_models/inference_models/models/base/keypoints_detection.py +++ b/inference_models/inference_models/models/base/keypoints_detection.py @@ -30,6 +30,11 @@ class KeyPoints: None # if given, per-instance object confidence (instances, ) ) + def __len__(self) -> int: + # number of skeleton instances (xy is (instances, instance_key_points, 2)); + # NOT the per-instance keypoint count (that would be xy.shape[1]). + return int(self.xy.shape[0]) + def to_supervision(self) -> sv.KeyPoints: """Convert keypoints to Supervision KeyPoints format. diff --git a/inference_models/inference_models/models/base/object_detection.py b/inference_models/inference_models/models/base/object_detection.py index 177bd6d62d..813b9d6b7b 100644 --- a/inference_models/inference_models/models/base/object_detection.py +++ b/inference_models/inference_models/models/base/object_detection.py @@ -1,6 +1,6 @@ from abc import ABC, abstractmethod from dataclasses import dataclass -from typing import Generic, List, Optional, Tuple, Union +from typing import Generic, Iterator, List, Optional, Tuple, Union import numpy as np import supervision as sv @@ -23,6 +23,37 @@ class Detections: None # if given, list of size equal to # of bboxes ) + def __len__(self) -> int: + return int(self.xyxy.shape[0]) + + def __iter__(self) -> Iterator[Tuple]: + """Iterates detections yielding 7-tuples: + (xyxy, mask, class_id, confidence, tracker_id, data, metadata) + + - xyxy: torch.Tensor of shape (4,) - single bbox [x1, y1, x2, y2] + - mask: always None for object detection + - class_id: scalar tensor (0-dim) + - confidence: scalar tensor (0-dim) + - tracker_id: value of `bboxes_metadata[i]["tracker_id"]` or None + - data: per-detection dict (`bboxes_metadata[i]`, `{}` if not set) + - metadata: per-image dict (`image_metadata`, `{}` if not set) + """ + bboxes_metadata = self.bboxes_metadata + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(len(self))] + image_metadata = self.image_metadata or {} + for index in range(len(self)): + data = bboxes_metadata[index] + yield ( + self.xyxy[index], + None, + self.class_id[index], + self.confidence[index], + data.get("tracker_id"), + data, + image_metadata, + ) + def to_supervision(self) -> sv.Detections: """Convert detections to Supervision Detections format. diff --git a/inference_models/inference_models/models/common/hf_streaming_video.py b/inference_models/inference_models/models/common/hf_streaming_video.py index 0dfdc5334d..7885058fa4 100644 --- a/inference_models/inference_models/models/common/hf_streaming_video.py +++ b/inference_models/inference_models/models/common/hf_streaming_video.py @@ -313,6 +313,11 @@ def _extract_masks_and_ids( def _ensure_numpy_image(image: Union[np.ndarray, torch.Tensor]) -> np.ndarray: if isinstance(image, torch.Tensor): + # Workflow tensor images are channels-first (CHW); the HF processor expects + # channels-last (HWC). Permute before the host transfer so the processor + # receives a correctly-shaped image. + if image.ndim == 3 and image.shape[0] in (1, 3, 4): + image = image.permute(1, 2, 0) return image.detach().cpu().numpy() return image diff --git a/inference_models/inference_models/models/common/rle_utils.py b/inference_models/inference_models/models/common/rle_utils.py index 36525ebeba..1e62075747 100644 --- a/inference_models/inference_models/models/common/rle_utils.py +++ b/inference_models/inference_models/models/common/rle_utils.py @@ -1,4 +1,4 @@ -from typing import Optional +from typing import List, Optional, Tuple import numpy as np import torch diff --git a/inference_models/inference_models/models/easy_ocr/easy_ocr_torch.py b/inference_models/inference_models/models/easy_ocr/easy_ocr_torch.py index 32eb04e6d0..e7b38e9750 100644 --- a/inference_models/inference_models/models/easy_ocr/easy_ocr_torch.py +++ b/inference_models/inference_models/models/easy_ocr/easy_ocr_torch.py @@ -208,9 +208,13 @@ def post_process( while_image_text_joined = text_regions_separator.join(whole_image_text) rendered_texts.append(while_image_text_joined) data = [{"text": text} for text in whole_image_text] + if xyxy: + xyxy_tensor = torch.tensor(xyxy, device=self._device) + else: + xyxy_tensor = torch.empty((0, 4), device=self._device) all_detections.append( Detections( - xyxy=torch.tensor(xyxy, device=self._device), + xyxy=xyxy_tensor, class_id=torch.tensor(class_id, device=self._device), confidence=torch.tensor( predictions_confidence, device=self._device diff --git a/inference_models/pyproject.toml b/inference_models/pyproject.toml index f2d0bcd1be..d854c9289e 100644 --- a/inference_models/pyproject.toml +++ b/inference_models/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "inference-models" -version = "0.34.4" +version = "0.35.0rc3" description = "The new inference engine for Computer Vision models" readme = "README.md" requires-python = ">=3.10,<3.13" diff --git a/inference_models/tests/unit_tests/models/base/test_instance_segmentation.py b/inference_models/tests/unit_tests/models/base/test_instance_segmentation.py new file mode 100644 index 0000000000..4cc9904700 --- /dev/null +++ b/inference_models/tests/unit_tests/models/base/test_instance_segmentation.py @@ -0,0 +1,45 @@ +import torch + +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks + + +def test_len_when_no_instances() -> None: + # given + detections = InstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, 8, 8), dtype=torch.bool), + ) + + # when / then + assert len(detections) == 0 + + +def test_len_when_multiple_instances_with_dense_mask() -> None: + # given + detections = InstanceDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + mask=torch.zeros((2, 8, 8), dtype=torch.bool), + ) + + # when / then + assert len(detections) == 2 + + +def test_len_counts_boxes_regardless_of_mask_representation() -> None: + # given - len() reads xyxy.shape[0], so the RLE mask representation must not matter + detections = InstanceDetections( + xyxy=torch.tensor( + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], dtype=torch.float32 + ), + class_id=torch.tensor([0, 1, 0], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4, 0.5], dtype=torch.float32), + mask=InstancesRLEMasks(image_size=(8, 8), masks=[b"", b"", b""]), + ) + + # when / then + assert len(detections) == 3 diff --git a/inference_models/tests/unit_tests/models/base/test_keypoints_detection.py b/inference_models/tests/unit_tests/models/base/test_keypoints_detection.py new file mode 100644 index 0000000000..a556185d6f --- /dev/null +++ b/inference_models/tests/unit_tests/models/base/test_keypoints_detection.py @@ -0,0 +1,59 @@ +import pytest +import torch + +from inference_models.models.base.keypoints_detection import KeyPoints + + +def test_len_when_no_instances() -> None: + # given + key_points = KeyPoints( + xy=torch.zeros((0, 17, 2), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0, 17), dtype=torch.float32), + ) + + # when / then + assert len(key_points) == 0 + + +def test_len_when_single_instance() -> None: + # given + key_points = KeyPoints( + xy=torch.tensor([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]), # (1, 3, 2) + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.9, 0.8, 0.7]]), # (1, 3) + ) + + # when / then + assert len(key_points) == 1 + + +def test_len_counts_instances_not_keypoints_per_instance() -> None: + # given - 2 instances, each with 17 keypoints; len() must be 2 (instance axis), + # NOT 17 (keypoints-per-instance axis, xy.shape[1]). + key_points = KeyPoints( + xy=torch.zeros((2, 17, 2), dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.zeros((2, 17), dtype=torch.float32), + ) + + # when / then + assert len(key_points) == 2 + + +def test_key_points_is_not_iterable() -> None: + # KeyPoints intentionally has no __iter__ (and no __getitem__): unlike Detections it + # does NOT support positional tuple-destructuring iteration - only __len__ and + # to_supervision(). This guard documents that contract so that adding iteration later + # (with some field order) is a conscious, explicitly-tested change rather than a + # silent one that downstream positional destructures could come to depend on. + key_points = KeyPoints( + xy=torch.tensor([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]), # (1, 3, 2) + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.9, 0.8, 0.7]]), # (1, 3) + ) + + # when / then + assert not hasattr(type(key_points), "__iter__") + with pytest.raises(TypeError): + iter(key_points) diff --git a/inference_models/tests/unit_tests/models/base/test_object_detection.py b/inference_models/tests/unit_tests/models/base/test_object_detection.py new file mode 100644 index 0000000000..94d9874462 --- /dev/null +++ b/inference_models/tests/unit_tests/models/base/test_object_detection.py @@ -0,0 +1,102 @@ +import pytest +import torch + +from inference_models.models.base.object_detection import Detections + + +def test_len_when_no_detections() -> None: + # given + detections = Detections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + ) + + # when / then + assert len(detections) == 0 + + +def test_len_when_single_detection() -> None: + # given + detections = Detections( + xyxy=torch.tensor([[0, 1, 2, 3]], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([0.9], dtype=torch.float32), + ) + + # when / then + assert len(detections) == 1 + + +def test_len_when_multiple_detections() -> None: + # given + detections = Detections( + xyxy=torch.tensor( + [[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], dtype=torch.float32 + ), + class_id=torch.tensor([0, 1, 2], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4, 0.5], dtype=torch.float32), + ) + + # when / then + assert len(detections) == 3 + + +def test_iter_yields_seven_tuple_in_documented_field_order() -> None: + # given - class_id and confidence are deliberately distinct in dtype and value so + # that a silent confidence/class_id swap is caught: the sv-mode 6-tuple yields these + # two fields in the OPPOSITE order, and positional destructures downstream + # (detection_event_log/v1_tensor, UQL detection/base.py, detections/base.py) depend + # on this exact native order. + detections = Detections( + xyxy=torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.float32), + class_id=torch.tensor([3, 7], dtype=torch.long), + confidence=torch.tensor([0.9, 0.1], dtype=torch.float32), + image_metadata={"image": "meta"}, + bboxes_metadata=[ + {"tracker_id": 10, "label": "a"}, + {"tracker_id": 20, "label": "b"}, + ], + ) + + # when + rows = list(detections) + + # then - order is (xyxy, mask, class_id, confidence, tracker_id, data, metadata) + assert len(rows) == 2 + for i, row in enumerate(rows): + assert len(row) == 7, "native __iter__ must yield a 7-tuple" + xyxy, mask, class_id, confidence, tracker_id, data, metadata = row + assert torch.equal(xyxy, detections.xyxy[i]) # 0: xyxy + assert mask is None # 1: mask (always None for object detection) + assert torch.equal(class_id, detections.class_id[i]) # 2: class_id + assert torch.equal(confidence, detections.confidence[i]) # 3: confidence + assert tracker_id == detections.bboxes_metadata[i]["tracker_id"] # 4: tracker_id + assert data == detections.bboxes_metadata[i] # 5: per-detection data dict + assert metadata == detections.image_metadata # 6: per-image metadata dict + + # class_id is the integer label and confidence is the float score - guards directly + # against a class_id<->confidence positional swap. + assert rows[0][2].item() == 3 + assert rows[0][3].item() == pytest.approx(0.9) + assert rows[1][2].item() == 7 + assert rows[1][3].item() == pytest.approx(0.1) + + +def test_iter_defaults_when_no_per_box_or_image_metadata() -> None: + # given - neither bboxes_metadata nor image_metadata provided + detections = Detections( + xyxy=torch.tensor([[0, 1, 2, 3]], dtype=torch.float32), + class_id=torch.tensor([5], dtype=torch.long), + confidence=torch.tensor([0.42], dtype=torch.float32), + ) + + # when + (row,) = list(detections) + + # then + xyxy, mask, class_id, confidence, tracker_id, data, metadata = row + assert mask is None + assert tracker_id is None # tracker_id absent -> None + assert data == {} # per-detection data defaults to empty dict + assert metadata == {} # per-image metadata defaults to empty dict diff --git a/inference_models/uv.lock b/inference_models/uv.lock index b20db65dd5..342a826e9c 100644 --- a/inference_models/uv.lock +++ b/inference_models/uv.lock @@ -913,7 +913,7 @@ wheels = [ [[package]] name = "inference-models" -version = "0.34.4" +version = "0.35.0rc3" source = { virtual = "." } dependencies = [ { name = "accelerate" }, diff --git a/modal/modal_app.py b/modal/modal_app.py index 6fb28f006d..7085630cc5 100644 --- a/modal/modal_app.py +++ b/modal/modal_app.py @@ -68,6 +68,10 @@ def append(self, *args, **kwargs) -> None: app = modal.App(WEBEXEC_MODAL_APP_NAME) +INFERENCE_VERSION = os.getenv("INFERENCE_VERSION") +WEBEXEC_INFERENCE_DOCKER_IMAGE = os.getenv( + "WEBEXEC_INFERENCE_DOCKER_IMAGE", "roboflow/roboflow-inference-server-cpu" +) WEBEXEC_INFERENCE_DOCKER_IMAGE = os.getenv("WEBEXEC_INFERENCE_DOCKER_IMAGE", "roboflow/roboflow-inference-server-cpu") @@ -85,7 +89,9 @@ def get_inference_image(): inference_version = "latest" image = ( - modal.Image.from_registry(f"{WEBEXEC_INFERENCE_DOCKER_IMAGE}:{inference_version}") + modal.Image.from_registry( + f"{WEBEXEC_INFERENCE_DOCKER_IMAGE}:{INFERENCE_VERSION}" + ) .apt_install( "libgl1-mesa-glx", "libglib2.0-0", @@ -96,7 +102,7 @@ def get_inference_image(): "ffmpeg", "wget", ) - .pip_install("fastapi[standard]", "msgpack") # Add FastAPI for web endpoints + .pip_install("fastapi[standard]") # Add FastAPI for web endpoints .entrypoint([]) ) return image @@ -184,6 +190,26 @@ def _get_or_initialize_namespace( "debug_traces": _NoopDebugTraces(), } import_code = "\n".join(imports) if imports else "" + # Mirror of block_scaffolding's tensor-native IMPORTS_LINES extension. + # Guarded import: this runs inside the sandbox on the PINNED inference + # image โ€” releases predating the tensor pivot lack the constant (and + # releases that carry it keep the extension a no-op unless the sandbox + # env enables the flag, which today it never does; tensor_native+modal + # is additionally blocked at compile time). Importing the constant + # instead of copy-pasting the lines keeps the two lists drift-free. + try: + from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION + from inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding import ( + TENSOR_NATIVE_IMPORTS_LINES, + ) + + tensor_native_imports = ( + "\n".join(TENSOR_NATIVE_IMPORTS_LINES) + if ENABLE_TENSOR_DATA_REPRESENTATION + else "" + ) + except ImportError: + tensor_native_imports = "" full_imports = f""" from typing import Any, List, Dict, Set, Optional import supervision as sv @@ -197,6 +223,7 @@ def _get_or_initialize_namespace( import shapely from inference.core.workflows.execution_engine.entities.base import Batch, WorkflowImageData from inference.core.workflows.prototypes.block import BlockResult +{tensor_native_imports} {import_code} diff --git a/requirements/_requirements.txt b/requirements/_requirements.txt index 5f3a2764b8..f11051f8d0 100644 --- a/requirements/_requirements.txt +++ b/requirements/_requirements.txt @@ -8,15 +8,15 @@ python-dotenv>=1.2.2,<2.0.0 tzdata>=2024.1 fastapi>=0.100,<0.116 # be careful with upper pin - fastapi might remove support for on_event numpy>=2.0.0,<2.4.0 -opencv-python>=4.8.1.78,<=4.10.0.84 -opencv-contrib-python>=4.8.1.78,<=4.10.0.84 # Note: opencv-python considers this as a bad practice, but since our dependencies rely on both we pin both here +opencv-python>=4.8.1.78,<4.13.0 +opencv-contrib-python>=4.8.1.78,<4.13.0 # Note: opencv-python considers this as a bad practice, but since our dependencies rely on both we pin both here pillow>=12.3.0,<13.0.0 prometheus-fastapi-instrumentator<=6.0.0 psutil>=7.0.0 redis~=5.0.0 requests>=2.33.0,<3.0.0 rich>=13.0.0,<15.0.0 -supervision>=0.27.0 +supervision>=0.29.0 pybase64~=1.0.0 scikit-image>=0.19.0,<=0.25.2 requests-toolbelt~=1.0.0 diff --git a/requirements/requirements.cli.txt b/requirements/requirements.cli.txt index 02688aabf0..ef946ce4a3 100644 --- a/requirements/requirements.cli.txt +++ b/requirements/requirements.cli.txt @@ -5,7 +5,7 @@ typer>=0.9.0,<=0.16.0 rich>=13.0.0,<15.0.0 PyYAML~=6.0.0 supervision>=0.26 -opencv-python>=4.8.1.78,<=4.10.0.84 +opencv-python>=4.8.1.78,<=4.13.0 tqdm>=4.0.0,<5.0.0 nvidia-ml-py<13.0.0 py-cpuinfo~=9.0.0 diff --git a/requirements/requirements.cpu.txt b/requirements/requirements.cpu.txt index 0391be3dfc..e15696285e 100644 --- a/requirements/requirements.cpu.txt +++ b/requirements/requirements.cpu.txt @@ -1,3 +1,3 @@ onnxruntime>=1.15.1,<1.22.0 nvidia-ml-py<13.0.0 -inference-models[torch-cpu,onnx-cpu]~=0.34.4 # keep in sync between requirements.gpu.txt, requirements.cpu.txt, requirements.vino.txt +inference-models[torch-cpu,onnx-cpu]~=0.35.0rc3 # keep in sync between requirements.gpu.txt, requirements.cpu.txt, requirements.vino.txt diff --git a/requirements/requirements.gpu.txt b/requirements/requirements.gpu.txt index 9b086366a5..2c3c231546 100644 --- a/requirements/requirements.gpu.txt +++ b/requirements/requirements.gpu.txt @@ -1,2 +1,3 @@ onnxruntime-gpu>=1.15.1,<1.22.0 -inference-models[torch-cu124,onnx-cu12]~=0.34.4 # keep in sync between requirements.jetson requirements.gpu.txt, requirements.cpu.txt, requirements.vino.txt +inference-models[torch-cu124,onnx-cu12]~=0.35.0rc3 # keep in sync between requirements.jetson requirements.gpu.txt, requirements.cpu.txt, requirements.vino.txt +pynvvideocodec~=2.1.0; platform_machine == "x86_64" diff --git a/requirements/requirements.jetson.txt b/requirements/requirements.jetson.txt index 7373c2100a..6417a00e1e 100644 --- a/requirements/requirements.jetson.txt +++ b/requirements/requirements.jetson.txt @@ -1,3 +1,3 @@ pypdfium2>=4.11.0,<5.0.0 PyYAML~=6.0.0 -inference-models~=0.34.4 # keep in sync between requirements.jetson requirements.gpu.txt, requirements.cpu.txt, requirements.vino.txt +inference-models~=0.35.0rc3 # keep in sync between requirements.jetson requirements.gpu.txt, requirements.cpu.txt, requirements.vino.txt diff --git a/requirements/requirements.sdk.http.txt b/requirements/requirements.sdk.http.txt index 9db5aabc75..6f655cebbd 100644 --- a/requirements/requirements.sdk.http.txt +++ b/requirements/requirements.sdk.http.txt @@ -2,7 +2,7 @@ requests>=2.33.0,<3.0.0 urllib3>=1.26,<3.0.0 tldextract~=5.1.2 dataclasses-json~=0.6.0 -opencv-python>=4.8.1.78,<=4.10.0.84 +opencv-python>=4.8.1.78,<=4.13.0 pillow>=12.3.0,<13.0.0 supervision>=0.26 numpy>=2.0.0,<2.4.0 diff --git a/requirements/requirements.vino.txt b/requirements/requirements.vino.txt index a655fc60d1..65a3ad6904 100644 --- a/requirements/requirements.vino.txt +++ b/requirements/requirements.vino.txt @@ -1,2 +1,2 @@ onnxruntime-openvino>=1.15.0,<1.22.0 -inference-models[torch-cpu]~=0.34.4 # keep in sync between requirements.jetson requirements.gpu.txt, requirements.cpu.txt, requirements.vino.txt +inference-models[torch-cpu]~=0.35.0rc3 # keep in sync between requirements.jetson requirements.gpu.txt, requirements.cpu.txt, requirements.vino.txt diff --git a/tests/inference/unit_tests/core/interfaces/camera/test_collection_policy.py b/tests/inference/unit_tests/core/interfaces/camera/test_collection_policy.py new file mode 100644 index 0000000000..a1f1f95eee --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/camera/test_collection_policy.py @@ -0,0 +1,498 @@ +"""Unit tests for AUTO/EVERY_FRAME collection policies and their dispatch. + +Covers: mode resolution (explicit > tensor-flag default > legacy), the +EMA-driven adaptive window, bounded-staleness FIFO reads with the file +exemption, and the `_multiplex_videos` dispatch guarantee that policy=None +uses the legacy retrieval path untouched. +""" + +from datetime import datetime, timedelta +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from inference.core import env as core_env +from inference.core.interfaces.camera import collection_policy as cp +from inference.core.interfaces.camera import utils as camera_utils +from inference.core.interfaces.camera.collection_policy import ( + AdaptiveWindowController, + CollectionPolicy, + VideoProcessingMode, + resolve_video_processing_mode, +) +from inference.core.interfaces.camera.exceptions import EndOfStreamError +from inference.core.interfaces.camera.utils import VideoSources, _multiplex_videos + + +def test_resolve_mode_explicit_argument_wins(monkeypatch) -> None: + monkeypatch.setattr(core_env, "ENABLE_TENSOR_DATA_REPRESENTATION", True) + + resolved = resolve_video_processing_mode(explicit_mode="every_frame") + + assert resolved is VideoProcessingMode.EVERY_FRAME + + +def test_resolve_mode_defaults_to_auto_for_tensor_cohort(monkeypatch) -> None: + monkeypatch.setattr(core_env, "ENABLE_TENSOR_DATA_REPRESENTATION", True) + + resolved = resolve_video_processing_mode(explicit_mode=None) + + assert resolved is VideoProcessingMode.AUTO + + +def test_resolve_mode_preserves_legacy_behavior_outside_tensor_cohort( + monkeypatch, +) -> None: + monkeypatch.setattr(core_env, "ENABLE_TENSOR_DATA_REPRESENTATION", False) + + resolved = resolve_video_processing_mode(explicit_mode=None) + + assert resolved is None + + +def test_resolve_mode_rejects_unknown_mode() -> None: + with pytest.raises(ValueError): + resolve_video_processing_mode(explicit_mode="turbo") + + +@pytest.mark.parametrize("alias", ["legacy", "none", "LEGACY", "None"]) +def test_resolve_mode_legacy_alias_overrides_tensor_cohort_default( + monkeypatch, alias: str +) -> None: + # given - the tensor cohort, where the implicit default is AUTO + monkeypatch.setattr(core_env, "ENABLE_TENSOR_DATA_REPRESENTATION", True) + + resolved = resolve_video_processing_mode(explicit_mode=alias) + + # then - the escape hatch forces the legacy path anyway + assert resolved is None + + +def test_adaptive_window_starts_at_initial_value() -> None: + controller = AdaptiveWindowController(initial_window=0.005, clock=lambda: 0.0) + + assert controller.on_collection_start() == 0.005 + + +def test_adaptive_window_tracks_execution_gap() -> None: + # given - a fake clock advanced manually + now = {"value": 0.0} + controller = AdaptiveWindowController(clock=lambda: now["value"]) + + # when - first round collects frames, execution takes 100 ms + controller.on_collection_start() + controller.on_collection_end(collected_any_frame=True) + now["value"] += 0.1 + window = controller.on_collection_start() + + # then - first gap sample seeds the EMA directly: 0.2 * 100 ms = 20 ms + assert controller.execution_gap_ema == pytest.approx(0.1) + assert window == pytest.approx(0.02) + + +def test_adaptive_window_is_clamped_to_bounds() -> None: + now = {"value": 0.0} + controller = AdaptiveWindowController(clock=lambda: now["value"]) + + controller.on_collection_end(collected_any_frame=True) + now["value"] += 10.0 # absurdly slow execution + upper = controller.on_collection_start() + controller.on_collection_end(collected_any_frame=True) + now["value"] += 0.000001 # near-instant execution + controller.on_collection_end(collected_any_frame=True) + # drive the EMA down with repeated near-zero gaps + for _ in range(64): + controller.on_collection_start() + controller.on_collection_end(collected_any_frame=True) + now["value"] += 0.000001 + lower = controller.on_collection_start() + + assert upper == pytest.approx(cp.MAX_COLLECTION_WINDOW_SECONDS) + assert lower == pytest.approx(cp.MIN_COLLECTION_WINDOW_SECONDS) + + +def test_adaptive_window_ignores_empty_rounds() -> None: + now = {"value": 0.0} + controller = AdaptiveWindowController(clock=lambda: now["value"]) + + controller.on_collection_start() + controller.on_collection_end(collected_any_frame=False) + now["value"] += 5.0 # long idle gap after an EMPTY round + controller.on_collection_start() + + # then - idle time never contaminates the execution-time estimate + assert controller.execution_gap_ema is None + + +def test_arrival_estimator_uniform_arrivals() -> None: + estimator = cp._SourceArrivalEstimator() + start = datetime(2026, 1, 1, 12, 0, 0) + + for index in range(32): + estimator.observe( + frame_timestamp=start + timedelta(seconds=index / 15.0), now=100.0 + ) + + assert estimator.period(now=100.0) == pytest.approx(1 / 15.0, rel=1e-6) + + +def test_arrival_estimator_bursty_arrivals_converge_on_mean_rate() -> None: + # given - the consumer-camera pattern: clusters of 50 ms spacing with a + # 450 ms encoder pause every 30 frames (true mean rate 15 fps-ish) + estimator = cp._SourceArrivalEstimator() + start = datetime(2026, 1, 1, 12, 0, 0) + timestamp = start + for index in range(64): + gap = 0.450 if index % 30 == 29 else 0.050 + timestamp = timestamp + timedelta(seconds=gap) + estimator.observe(frame_timestamp=timestamp, now=100.0) + + period = estimator.period(now=100.0) + + # then - close to the mean period, NOT the intra-burst 50 ms spacing + mean_gap = (0.050 * 29 + 0.450) / 30.0 + assert period == pytest.approx(mean_gap, rel=0.15) + assert period > 0.055 + + +def test_arrival_estimator_resets_after_reconnect_gap() -> None: + estimator = cp._SourceArrivalEstimator() + start = datetime(2026, 1, 1, 12, 0, 0) + for index in range(32): + estimator.observe( + frame_timestamp=start + timedelta(seconds=index / 15.0), now=100.0 + ) + + # when - a 30 s reconnect gap, then only a few fresh samples + rejoined = start + timedelta(seconds=40) + for index in range(4): + estimator.observe( + frame_timestamp=rejoined + timedelta(seconds=index / 15.0), now=141.0 + ) + + # then - the pre-gap history is discarded, too few samples to trust + assert estimator.period(now=141.0) is None + + +def test_arrival_estimator_dormant_source_reports_no_period() -> None: + estimator = cp._SourceArrivalEstimator() + start = datetime(2026, 1, 1, 12, 0, 0) + for index in range(32): + estimator.observe( + frame_timestamp=start + timedelta(seconds=index / 15.0), now=100.0 + ) + + assert estimator.period(now=100.0) is not None + # then - nothing seen for longer than the activity horizon -> excluded + assert estimator.period(now=103.0) is None + + +def test_controller_rate_matches_when_arrival_period_known() -> None: + now = {"value": 0.0} + controller = AdaptiveWindowController(clock=lambda: now["value"]) + + controller.on_collection_start() + controller.on_collection_end(collected_any_frame=True) + now["value"] += 0.042 # execution takes 42 ms + window = controller.on_collection_start(minimum_arrival_period=1 / 15.0) + + # then - window = frame period - exec: rounds lock to the arrival rate + assert window == pytest.approx(1 / 15.0 - 0.042, rel=1e-6) + + +def test_controller_floors_under_saturation_with_known_period() -> None: + now = {"value": 0.0} + controller = AdaptiveWindowController(clock=lambda: now["value"]) + + controller.on_collection_end(collected_any_frame=True) + now["value"] += 0.129 # x-model regime: exec far above the frame period + window = controller.on_collection_start(minimum_arrival_period=1 / 15.0) + + assert window == pytest.approx(cp.MIN_COLLECTION_WINDOW_SECONDS) + + +def test_controller_caps_rate_matched_window() -> None: + now = {"value": 0.0} + controller = AdaptiveWindowController(clock=lambda: now["value"]) + + controller.on_collection_end(collected_any_frame=True) + now["value"] += 0.005 # near-instant execution, very slow source (2 fps) + window = controller.on_collection_start(minimum_arrival_period=0.5) + + assert window == pytest.approx(cp.RATE_MATCHED_WINDOW_CAP_SECONDS) + + +def test_policy_minimum_period_tracks_fastest_source_and_skips_files() -> None: + policy = CollectionPolicy(mode=VideoProcessingMode.AUTO, max_staleness=0.4) + start = datetime.now() + fast_frames = [ + SimpleNamespace( + frame_id=index, + source_id=0, + frame_timestamp=start + timedelta(seconds=index / 30.0), + ) + for index in range(32) + ] + slow_frames = [ + SimpleNamespace( + frame_id=index, + source_id=1, + frame_timestamp=start + timedelta(seconds=index / 5.0), + ) + for index in range(32) + ] + file_frames = [ + SimpleNamespace( + frame_id=index, + source_id=2, + frame_timestamp=start + timedelta(seconds=index / 120.0), + ) + for index in range(32) + ] + fast = _FakeSource(frames=fast_frames, is_file=False) + slow = _FakeSource(frames=slow_frames, is_file=False) + file_source = _FakeSource(frames=file_frames, is_file=True) + + for _ in range(32): + policy.read_frame(source_ord=0, source=fast, timeout=0.1) + policy.read_frame(source_ord=1, source=slow, timeout=0.1) + policy.read_frame(source_ord=2, source=file_source, timeout=0.1) + + # then - the fastest LIVE source binds; the 120 fps file never counts + assert policy.minimum_live_arrival_period() == pytest.approx( + 1 / 30.0, rel=0.05 + ) + + +def test_policy_feeds_estimator_with_staleness_drained_frames() -> None: + policy = CollectionPolicy(mode=VideoProcessingMode.AUTO, max_staleness=0.4) + start = datetime.now() - timedelta(seconds=10) + stale_frames = [ + SimpleNamespace( + frame_id=index, + source_id=0, + frame_timestamp=start + timedelta(seconds=index / 15.0), + ) + for index in range(32) + ] + source = _FakeSource(frames=stale_frames, is_file=False) + + frame = policy.read_frame(source_ord=0, source=source, timeout=0.1) + + # then - every drained arrival counted: period known despite 0 returns + assert frame is None + assert policy.minimum_live_arrival_period() == pytest.approx( + 1 / 15.0, rel=0.05 + ) + + +def _fake_frame(age_seconds: float, frame_id: int = 1, source_id: int = 0): + return SimpleNamespace( + frame_id=frame_id, + source_id=source_id, + frame_timestamp=datetime.now() - timedelta(seconds=age_seconds), + ) + + +class _FakeSource: + def __init__(self, frames, is_file): + self._frames = list(frames) + self._is_file = is_file + + def read_frame(self, timeout=None): + if not self._frames: + return None + item = self._frames.pop(0) + if isinstance(item, Exception): + raise item + return item + + def describe_source(self): + properties = ( + None if self._is_file is None else SimpleNamespace(is_file=self._is_file) + ) + return SimpleNamespace(source_properties=properties) + + +def test_auto_policy_drops_stale_frames_for_live_sources() -> None: + # given - two stale frames queued ahead of a fresh one + dropped = [] + policy = CollectionPolicy( + mode=VideoProcessingMode.AUTO, + max_staleness=0.4, + on_frame_dropped=dropped.append, + ) + source = _FakeSource( + frames=[ + _fake_frame(1.2, frame_id=1), + _fake_frame(0.9, frame_id=2), + _fake_frame(0.05, frame_id=3), + ], + is_file=False, + ) + + frame = policy.read_frame(source_ord=0, source=source, timeout=0.1) + + assert frame.frame_id == 3 + assert [f.frame_id for f in dropped] == [1, 2] + assert policy.frames_dropped_on_staleness == {0: 2} + + +def test_auto_policy_never_drops_file_frames() -> None: + policy = CollectionPolicy(mode=VideoProcessingMode.AUTO, max_staleness=0.4) + source = _FakeSource(frames=[_fake_frame(10.0, frame_id=1)], is_file=True) + + frame = policy.read_frame(source_ord=0, source=source, timeout=0.1) + + assert frame.frame_id == 1 + assert policy.frames_dropped_on_staleness == {} + + +def test_auto_policy_treats_unknown_source_properties_as_file() -> None: + # given - source metadata not resolved yet (initialising source) + policy = CollectionPolicy(mode=VideoProcessingMode.AUTO, max_staleness=0.4) + source = _FakeSource(frames=[_fake_frame(10.0, frame_id=1)], is_file=None) + + frame = policy.read_frame(source_ord=0, source=source, timeout=0.1) + + # then - never drop while liveness is unknown + assert frame.frame_id == 1 + assert policy.frames_dropped_on_staleness == {} + + +def test_every_frame_policy_disables_staleness_budget() -> None: + policy = CollectionPolicy(mode=VideoProcessingMode.EVERY_FRAME, max_staleness=0.4) + source = _FakeSource(frames=[_fake_frame(10.0, frame_id=1)], is_file=False) + + frame = policy.read_frame(source_ord=0, source=source, timeout=0.1) + + assert frame.frame_id == 1 + assert policy.max_staleness is None + + +def test_policy_rejects_freshest_mode() -> None: + with pytest.raises(ValueError): + CollectionPolicy(mode=VideoProcessingMode.FRESHEST) + + +def test_policy_propagates_end_of_stream() -> None: + policy = CollectionPolicy(mode=VideoProcessingMode.AUTO) + source = _FakeSource(frames=[EndOfStreamError()], is_file=False) + + with pytest.raises(EndOfStreamError): + policy.read_frame(source_ord=0, source=source, timeout=0.1) + + +def test_dropped_frame_callback_errors_never_break_reads() -> None: + def raising_callback(frame): + raise RuntimeError("sink exploded") + + policy = CollectionPolicy( + mode=VideoProcessingMode.AUTO, + max_staleness=0.4, + on_frame_dropped=raising_callback, + ) + source = _FakeSource( + frames=[_fake_frame(1.0, frame_id=1), _fake_frame(0.0, frame_id=2)], + is_file=False, + ) + + frame = policy.read_frame(source_ord=0, source=source, timeout=0.1) + + assert frame.frame_id == 2 + + +def _ended_video_sources() -> VideoSources: + # single fake source that ends immediately - drives the multiplex loop to + # exactly one retrieval round + source = _FakeSource(frames=[EndOfStreamError()], is_file=False) + return VideoSources( + all_sources=[source], allow_reconnection=[False], managed_sources=[] + ) + + +def test_multiplex_dispatch_uses_legacy_path_without_policy(monkeypatch) -> None: + # given - spies on both retrieval methods + legacy_spy = MagicMock(return_value=None) + policy_spy = MagicMock(return_value=None) + monkeypatch.setattr( + camera_utils.VideoSourcesManager, + "retrieve_frames_from_sources", + legacy_spy, + ) + monkeypatch.setattr( + camera_utils.VideoSourcesManager, + "retrieve_frames_from_sources_with_policy", + policy_spy, + ) + + # when + list( + _multiplex_videos( + video_sources=_ended_video_sources(), + batch_collection_timeout=0.123, + should_stop=lambda: False, + on_reconnection_error=lambda *_: None, + ) + ) + + # then - policy machinery is not even touched, timeout forwarded verbatim + legacy_spy.assert_called_with(batch_collection_timeout=0.123) + policy_spy.assert_not_called() + + +def test_multiplex_dispatch_uses_policy_path_when_policy_given(monkeypatch) -> None: + legacy_spy = MagicMock(return_value=None) + policy_spy = MagicMock(return_value=None) + monkeypatch.setattr( + camera_utils.VideoSourcesManager, + "retrieve_frames_from_sources", + legacy_spy, + ) + monkeypatch.setattr( + camera_utils.VideoSourcesManager, + "retrieve_frames_from_sources_with_policy", + policy_spy, + ) + policy = CollectionPolicy(mode=VideoProcessingMode.AUTO) + + list( + _multiplex_videos( + video_sources=_ended_video_sources(), + batch_collection_timeout=None, + should_stop=lambda: False, + on_reconnection_error=lambda *_: None, + collection_policy=policy, + ) + ) + + policy_spy.assert_called_with(collection_policy=policy) + legacy_spy.assert_not_called() + + +def test_policy_retrieval_round_collects_and_registers_eos() -> None: + # given - one live source with a fresh frame, one source at EOS + fresh = _fake_frame(0.01, frame_id=7, source_id=0) + source_with_frame = _FakeSource(frames=[fresh], is_file=False) + ended_source = _FakeSource(frames=[EndOfStreamError()], is_file=False) + video_sources = VideoSources( + all_sources=[source_with_frame, ended_source], + allow_reconnection=[False, False], + managed_sources=[], + ) + manager = camera_utils.VideoSourcesManager.init( + video_sources=video_sources, + should_stop=lambda: False, + on_reconnection_error=lambda *_: None, + ) + policy = CollectionPolicy(mode=VideoProcessingMode.AUTO, max_staleness=0.4) + + batch = manager.retrieve_frames_from_sources_with_policy(collection_policy=policy) + + assert batch == [fresh] + assert manager.all_sources_ended() is False # only one of two sources ended + second_batch = manager.retrieve_frames_from_sources_with_policy( + collection_policy=policy + ) + assert second_batch == [] # ended source inactive, first source drained diff --git a/tests/inference/unit_tests/core/interfaces/camera/test_discoverability.py b/tests/inference/unit_tests/core/interfaces/camera/test_discoverability.py new file mode 100644 index 0000000000..ac3b8c18f5 --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/camera/test_discoverability.py @@ -0,0 +1,160 @@ +from unittest.mock import patch + +from inference.core.interfaces.camera.discoverability import ( + DGPU, + GSTREAMER_CUDA, + JETSON, + ProducerAvailability, + _resolution_order, + available_producers, + build_hw_producer, + check_gstreamer_cuda, + check_jetson_gstreamer, +) + + +@patch("platform.system", return_value="Linux") +@patch("platform.machine", return_value="aarch64") +def test_jetson_routes_gstreamer_for_streams_and_files( + machine_mock, + system_mock, +) -> None: + # Streams / cameras -> Jetson HW GStreamer (latest-wins slot). + assert _resolution_order(prefer=None, video="rtsp://cam/stream") == [JETSON] + # Local files -> the same producer via the bridge's lossless handoff + # (bridge v6+); cv2 remains only the construction-failure fallback. + assert _resolution_order(prefer=None, video="sample.mp4") == [JETSON] + + +@patch("platform.system", return_value="Linux") +@patch("platform.machine", return_value="x86_64") +def test_dgpu_routes_gstreamer_for_streams_and_pynvdec_for_files( + machine_mock, + system_mock, +) -> None: + # Streams -> GStreamer CUDA. + assert _resolution_order(prefer=None, video="rtsp://cam/stream") == [GSTREAMER_CUDA] + # Local files -> PyNvVideoCodec (dGPU). + assert _resolution_order(prefer=None, video="sample.mp4") == [DGPU] + + +def test_explicit_backend_preference_takes_priority() -> None: + assert _resolution_order(prefer=GSTREAMER_CUDA) == [ + GSTREAMER_CUDA, + JETSON, + DGPU, + ] + + +def test_generic_file_decoders_reject_v4l2_device_paths() -> None: + availability = check_gstreamer_cuda("/dev/video0") + + assert not availability.available + + +@patch( + "inference.core.interfaces.camera.discoverability.check_pynvvideocodec", + return_value=ProducerAvailability(DGPU, True, "ok"), +) +@patch( + "inference.core.interfaces.camera.discoverability.check_jetson_gstreamer", + return_value=ProducerAvailability(JETSON, True, "ok"), +) +@patch( + "inference.core.interfaces.camera.discoverability.check_gstreamer_cuda", + return_value=ProducerAvailability(GSTREAMER_CUDA, True, "ok"), +) +def test_generic_cuda_backend_remains_available_for_numpy_consumers( + check_gstreamer_cuda_mock, + check_jetson_gstreamer_mock, + check_pynvvideocodec_mock, +) -> None: + availability = available_producers(video="sample.mp4", require_cuda_tensor=False) + + assert availability[GSTREAMER_CUDA].available + check_gstreamer_cuda_mock.assert_called_once_with("sample.mp4") + check_jetson_gstreamer_mock.assert_called_once_with( + video="sample.mp4", require_cuda_tensor=False + ) + check_pynvvideocodec_mock.assert_not_called() + + +@patch("torch.cuda.is_available", return_value=True) +@patch( + "inference.core.interfaces.camera.jetson_tensor_bridge.jetson_tensor_bridge_available", + return_value=(True, "ok"), +) +@patch( + "inference.core.interfaces.camera.jetson_producer.probe_gstreamer_elements", + return_value=(True, "ok"), +) +def test_jetson_numpy_probe_requires_the_native_cuda_bridge( + probe_gstreamer_elements_mock, + bridge_available_mock, + cuda_available_mock, +) -> None: + availability = check_jetson_gstreamer(video="sample.mp4", require_cuda_tensor=False) + + assert availability.available + bridge_available_mock.assert_called_once_with() + cuda_available_mock.assert_called_once_with() + required_elements = probe_gstreamer_elements_mock.call_args.args[0] + assert "nvvidconv" in required_elements + assert "videoconvert" not in required_elements + + +@patch("platform.system", return_value="Linux") +@patch("platform.machine", return_value="aarch64") +@patch( + "inference.core.interfaces.camera.jetson_producer.JetsonVideoFrameProducer", + side_effect=RuntimeError("preroll failed"), +) +@patch( + "inference.core.interfaces.camera.discoverability.available_producers", + return_value={ + GSTREAMER_CUDA: ProducerAvailability(GSTREAMER_CUDA, False, "unavailable"), + JETSON: ProducerAvailability(JETSON, True, "ok"), + DGPU: ProducerAvailability(DGPU, False, "unavailable"), + }, +) +def test_factory_logs_construction_failures_before_falling_back( + available_producers_mock, + producer_class_mock, + machine_mock, + system_mock, + caplog, +) -> None: + # given - the probe passes but the producer constructor raises; without a + # log line the caller-side fallback to cv2 is undiagnosable + with caplog.at_level("WARNING"): + producer = build_hw_producer("sample.mp4") + + assert producer is None + assert "Constructing the 'jetson' hardware decoder" in caplog.text + assert "preroll failed" in caplog.text + + +@patch( + "inference.core.interfaces.camera.gstreamer_cuda_producer.GstreamerCudaVideoFrameProducer" +) +@patch( + "inference.core.interfaces.camera.discoverability.available_producers", + return_value={ + GSTREAMER_CUDA: ProducerAvailability(GSTREAMER_CUDA, True, "ok"), + JETSON: ProducerAvailability(JETSON, False, "unavailable"), + DGPU: ProducerAvailability(DGPU, False, "unavailable"), + }, +) +def test_factory_requests_numpy_from_the_native_generic_cuda_producer( + available_producers_mock, + producer_class_mock, +) -> None: + producer = build_hw_producer( + "sample.mp4", prefer=GSTREAMER_CUDA, output_tensor=False + ) + + assert producer is producer_class_mock.return_value + producer_class_mock.assert_called_once_with("sample.mp4", output_tensor=False) + available_producers_mock.assert_called_once_with( + video="sample.mp4", require_cuda_tensor=False + ) diff --git a/tests/inference/unit_tests/core/interfaces/camera/test_gstreamer_cuda_producer.py b/tests/inference/unit_tests/core/interfaces/camera/test_gstreamer_cuda_producer.py new file mode 100644 index 0000000000..953788ce30 --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/camera/test_gstreamer_cuda_producer.py @@ -0,0 +1,173 @@ +from types import SimpleNamespace + +import numpy as np +import torch + +from inference.core.interfaces.camera.gstreamer_cuda_producer import ( + GstreamerCudaVideoFrameProducer, + build_gstreamer_cuda_pipeline, + required_gstreamer_cuda_elements, +) + + +class _NativePipeline: + def __init__(self, frame=None, factories=()) -> None: + self.grab_calls = 0 + self.retrieve_calls = 0 + self.frame = frame if frame is not None else object() + self.factories = set(factories) + self.factory_queries = [] + self.last_grab_timeout_ns = None + + def grab(self, timeout_ns=None) -> bool: + self.grab_calls += 1 + self.last_grab_timeout_ns = timeout_ns + return True + + def retrieve(self): + self.retrieve_calls += 1 + return self.frame + + def has_factory(self, factory: str) -> bool: + self.factory_queries.append(factory) + return factory in self.factories + + def frame_info(self): + return SimpleNamespace( + width=320, + height=180, + fps_numerator=30, + fps_denominator=1, + duration_ns=0, + ) + + +def _producer( + native_pipeline: _NativePipeline, *, output_tensor: bool = True +) -> GstreamerCudaVideoFrameProducer: + producer = GstreamerCudaVideoFrameProducer.__new__(GstreamerCudaVideoFrameProducer) + producer._source_ref = "rtsp://camera.example.test/live" + producer._native_pipeline = native_pipeline + producer._output_tensor = output_tensor + producer._decoder_validated = True + producer._prerolled_frame_pending = False + producer._cached_source_properties = None + producer._grab_timeout_ns = 5_000_000_000 + producer._closed = False + producer._eos = False + return producer + + +def test_metadata_preroll_is_consumable_once_and_later_grabs_advance() -> None: + native_pipeline = _NativePipeline() + producer = _producer(native_pipeline) + + first_properties = producer.discover_source_properties() + second_properties = producer.discover_source_properties() + + assert first_properties is second_properties + assert native_pipeline.grab_calls == 1 + assert producer.grab() + assert native_pipeline.grab_calls == 1 + assert producer.grab() + assert native_pipeline.grab_calls == 2 + # The producer must hand the native pull a finite deadline so a stalled + # source raises instead of blocking VideoSource forever (FQ-1). + assert native_pipeline.last_grab_timeout_ns == producer._grab_timeout_ns + + +def test_retrieving_preroll_requires_next_grab_to_advance_native_pipeline() -> None: + native_pipeline = _NativePipeline() + producer = _producer(native_pipeline) + + producer.discover_source_properties() + success, _ = producer.retrieve() + + assert success + assert native_pipeline.retrieve_calls == 1 + assert producer.grab() + assert native_pipeline.grab_calls == 2 + + +def test_numpy_retrieve_materializes_bgr_hwc_from_native_rgb_tensor() -> None: + rgb_tensor = torch.tensor( + [ + [[10, 20], [30, 40]], + [[50, 60], [70, 80]], + [[90, 100], [110, 120]], + ], + dtype=torch.uint8, + ) + producer = _producer(_NativePipeline(frame=rgb_tensor), output_tensor=False) + + success, image = producer.retrieve() + + assert success + assert isinstance(image, np.ndarray) + assert image.dtype == np.uint8 + assert image.flags.c_contiguous + np.testing.assert_array_equal( + image, + np.array( + [ + [[90, 50, 10], [100, 60, 20]], + [[110, 70, 30], [120, 80, 40]], + ], + dtype=np.uint8, + ), + ) + + +def test_numpy_mode_validates_cuda_conversion_and_hardware_decoder() -> None: + native_pipeline = _NativePipeline(factories={"cudaconvertscale", "nvh264dec"}) + producer = _producer(native_pipeline, output_tensor=False) + producer._decoder_validated = False + + assert producer.grab() + + assert producer._decoder_validated + assert "cudaconvertscale" in native_pipeline.factory_queries + assert "nvh264dec" in native_pipeline.factory_queries + + +def test_tensor_pipeline_keeps_frames_in_cuda_memory() -> None: + pipeline = build_gstreamer_cuda_pipeline( + "rtsps://camera.example.test/live", device_id=2 + ) + + assert 'caps="video/x-raw(memory:CUDAMemory)"' in pipeline + assert "cudaconvertscale cuda-device-id=2" in pipeline + assert "video/x-raw(memory:CUDAMemory),format=RGBP" in pipeline + assert "appsink name=rf_tensor_sink" in pipeline + assert "cudaupload" not in pipeline + assert "cudadownload" not in pipeline + assert "videoconvert" not in pipeline + + +def test_rtsps_element_contract_includes_tls_capable_rtsp_source() -> None: + elements = set(required_gstreamer_cuda_elements("rtsps://camera.example.test/live")) + + assert { + "cudaconvertscale", + "h264parse", + "h265parse", + "rtph264depay", + "rtph265depay", + "rtspsrc", + "uridecodebin", + }.issubset(elements) + + +def test_local_mp4_contract_includes_demuxer() -> None: + elements = set(required_gstreamer_cuda_elements("sample.mp4")) + + assert "qtdemux" in elements + + +def test_v4l2_device_is_not_treated_as_a_regular_file() -> None: + try: + GstreamerCudaVideoFrameProducer("/dev/video0") + except TypeError: + pass + else: + raise AssertionError("V4L2 device path must not use the URI producer") diff --git a/tests/inference/unit_tests/core/interfaces/camera/test_jetson_producer.py b/tests/inference/unit_tests/core/interfaces/camera/test_jetson_producer.py new file mode 100644 index 0000000000..12009837ee --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/camera/test_jetson_producer.py @@ -0,0 +1,390 @@ +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from inference.core.interfaces.camera.jetson_producer import ( + JetsonVideoFrameProducer, + build_gstreamer_pipeline, + required_gstreamer_elements, +) + + +class _NativePipeline: + def __init__(self, frame=None, factories=()) -> None: + self.grab_calls = 0 + self.retrieve_calls = 0 + self.interrupt_calls = 0 + self.frame = frame if frame is not None else object() + self.factories = set(factories) + self.factory_queries = [] + self.last_grab_timeout_ns = None + + def grab(self, timeout_ns=None) -> bool: + self.grab_calls += 1 + self.last_grab_timeout_ns = timeout_ns + return True + + def retrieve(self): + self.retrieve_calls += 1 + return self.frame + + def has_factory(self, factory: str) -> bool: + self.factory_queries.append(factory) + return factory in self.factories + + def frame_info(self): + return SimpleNamespace( + width=320, + height=180, + fps_numerator=30, + fps_denominator=1, + duration_ns=0, + ) + + def interrupt(self) -> None: + self.interrupt_calls += 1 + + +def _native_producer( + native_pipeline: _NativePipeline, + *, + output_tensor: bool = True, + source_ref: str = "rtsp://camera.example.test/live", +) -> JetsonVideoFrameProducer: + producer = JetsonVideoFrameProducer.__new__(JetsonVideoFrameProducer) + producer._source_ref = source_ref + producer._native_pipeline = native_pipeline + producer._capture = None + producer._output_tensor = output_tensor + producer._decoder_validated = True + producer._prerolled_frame_pending = False + producer._cached_source_properties = None + producer._grab_timeout_ns = 5_000_000_000 + producer._closed = False + producer._eos = False + return producer + + +def test_metadata_preroll_is_consumable_once_and_later_grabs_advance() -> None: + native_pipeline = _NativePipeline() + producer = _native_producer(native_pipeline) + + first_properties = producer.discover_source_properties() + second_properties = producer.discover_source_properties() + + assert first_properties is second_properties + assert native_pipeline.grab_calls == 1 + assert producer.grab() + assert native_pipeline.grab_calls == 1 + assert producer.grab() + assert native_pipeline.grab_calls == 2 + # The producer must hand the native pull a finite deadline so a stalled + # source raises instead of blocking VideoSource forever (FQ-1). + assert native_pipeline.last_grab_timeout_ns == producer._grab_timeout_ns + + +def test_retrieving_preroll_clears_pending_grab_and_interrupts_native_wait() -> None: + native_pipeline = _NativePipeline() + producer = _native_producer(native_pipeline) + + producer.discover_source_properties() + success, _ = producer.retrieve() + producer.interrupt() + + assert success + assert native_pipeline.retrieve_calls == 1 + assert native_pipeline.interrupt_calls == 1 + assert not producer.isOpened() + + +def test_numpy_retrieve_materializes_bgr_hwc_from_native_rgb_tensor() -> None: + rgb_tensor = torch.tensor( + [ + [[10, 20], [30, 40]], + [[50, 60], [70, 80]], + [[90, 100], [110, 120]], + ], + dtype=torch.uint8, + ) + producer = _native_producer(_NativePipeline(frame=rgb_tensor), output_tensor=False) + + success, image = producer.retrieve() + + assert success + assert isinstance(image, np.ndarray) + assert image.dtype == np.uint8 + assert image.flags.c_contiguous + np.testing.assert_array_equal( + image, + np.array( + [ + [[90, 50, 10], [100, 60, 20]], + [[110, 70, 30], [120, 80, 40]], + ], + dtype=np.uint8, + ), + ) + + +def test_numpy_mode_validates_the_instantiated_hardware_decoder() -> None: + native_pipeline = _NativePipeline(factories={"nvv4l2decoder"}) + producer = _native_producer(native_pipeline, output_tensor=False) + producer._decoder_validated = False + + assert producer.grab() + + assert producer._decoder_validated + assert "nvv4l2decoder" in native_pipeline.factory_queries + + +def test_raw_v4l2_source_does_not_require_a_decoder_element() -> None: + producer = _native_producer( + _NativePipeline(), output_tensor=False, source_ref="/dev/video0" + ) + producer._decoder_validated = False + + assert producer.grab() + + assert producer._decoder_validated + + +def test_compressed_v4l2_source_rejects_a_software_decoder() -> None: + producer = _native_producer( + _NativePipeline(factories={"jpegdec"}), + output_tensor=False, + source_ref="/dev/video0", + ) + producer._decoder_validated = False + + with pytest.raises(RuntimeError, match="software decoders: jpegdec"): + producer.grab() + + +def test_rtsps_source_uses_live_rtsp_pipeline() -> None: + pipeline = build_gstreamer_pipeline( + "rtsps://camera.example.test:7441/live?token=secret" + ) + + # Explicit video-only chain: a codec-specific depayloader never links the + # audio track, so an audio-muxing camera cannot poison the pipeline. The + # decoder's NV12 NVMM output feeds the appsink directly (the bridge + # converts NV12->RGB in CUDA) โ€” no nvvidconv VIC pass, and the queue + # buffers compressed data before the depayloader instead of leaking + # decoded frames. + assert pipeline.startswith( + 'rtspsrc location="rtsps://camera.example.test:7441/live?token=secret" ' + "protocols=tcp latency=200 drop-on-latency=true teardown-timeout=0 ! " + "application/x-rtp,media=video ! queue ! " + ) + assert ( + "rtph264depay request-keyframe=true wait-for-keyframe=true ! " + "h264parse ! h264timestamper ! " + "nvv4l2decoder enable-max-performance=1" + ) in pipeline + assert "uridecodebin" not in pipeline + assert "nvvidconv" not in pipeline + assert "video/x-raw(memory:NVMM),format=NV12" in pipeline + assert "appsink name=rf_tensor_sink" in pipeline + assert "max-buffers=4 drop=false sync=false" in pipeline + assert "leaky" not in pipeline + + +def test_rtsps_sdes_source_decrypts_before_codec_autoplugging() -> None: + """Route explicitly requested SRTP through native SDES key handling.""" + + video = "rtsps://camera.example.test/live?enableSrtp" + + pipeline = build_gstreamer_pipeline(video) + elements = set(required_gstreamer_elements(video)) + + assert ( + "capssetter name=rf_srtp_caps caps=application/x-srtp " + "join=false replace=false ! srtpdec ! " + "rtph264depay request-keyframe=true wait-for-keyframe=true ! " + "h264parse ! h264timestamper ! " + "nvv4l2decoder enable-max-performance=1 !" + ) in pipeline + assert {"capssetter", "srtpdec"} <= elements + + +@pytest.mark.parametrize("value", ("0", "false", "no", "off")) +def test_rtsps_sdes_source_can_explicitly_disable_srtp(value: str) -> None: + """Treat false-like enableSrtp query values as clear RTP media.""" + + video = f"rtsps://camera.example.test/live?enableSrtp={value}" + + pipeline = build_gstreamer_pipeline(video) + elements = set(required_gstreamer_elements(video)) + + assert "rf_srtp_caps" not in pipeline + assert "srtpdec" not in elements + + +def test_rtspt_source_is_recognised_as_rtsp() -> None: + pipeline = build_gstreamer_pipeline("rtspt://camera.example.test/live") + + assert pipeline.startswith('rtspsrc location="rtspt://camera.example.test/live"') + + +def test_rtsp_codec_env_selects_h265_chain(monkeypatch) -> None: + monkeypatch.setenv("ROBOFLOW_RTSP_VIDEO_CODEC", "h265") + + pipeline = build_gstreamer_pipeline("rtsp://camera.example.test/live") + + assert ("rtph265depay ! h265parse ! h265timestamper ! nvv4l2decoder") in pipeline + assert "request-keyframe" not in pipeline + assert "wait-for-keyframe" not in pipeline + + +def test_rtsp_codec_override_requires_only_selected_parser(monkeypatch) -> None: + """Keep a forced codec usable in slim GStreamer images.""" + + monkeypatch.setenv("ROBOFLOW_RTSP_VIDEO_CODEC", "h265") + + elements = set(required_gstreamer_elements("rtsp://camera.example.test/live")) + + assert { + "rtspsrc", + "rtph265depay", + "h265parse", + "h265timestamper", + "nvv4l2decoder", + } <= elements + assert "parsebin" not in elements + assert "rtph264depay" not in elements + + +def test_rtsp_codec_env_rejects_unsupported_codec(monkeypatch) -> None: + monkeypatch.setenv("ROBOFLOW_RTSP_VIDEO_CODEC", "mjpeg") + + with pytest.raises(ValueError, match="Unsupported RTSP video codec"): + build_gstreamer_pipeline("rtsp://camera.example.test/live") + + +def test_rtsp_transport_env_overrides_protocols_and_latency(monkeypatch) -> None: + monkeypatch.setenv("ROBOFLOW_RTSP_PROTOCOLS", "tcp+udp") + monkeypatch.setenv("ROBOFLOW_RTSP_LATENCY_MS", "1000") + + pipeline = build_gstreamer_pipeline("rtsp://camera.example.test/live") + + expected_transport = ( + "protocols=tcp+udp latency=1000 drop-on-latency=true teardown-timeout=0 ! " + ) + assert expected_transport in pipeline + + +def test_rtsp_tls_validation_flags_are_opt_in(monkeypatch) -> None: + secure_pipeline = build_gstreamer_pipeline("rtsps://camera.example.test/live") + assert "tls-validation-flags" not in secure_pipeline + + monkeypatch.setenv("ROBOFLOW_RTSP_TLS_VALIDATION_FLAGS", "0") + self_signed_pipeline = build_gstreamer_pipeline("rtsps://camera.example.test/live") + assert "tls-validation-flags=0 ! " in self_signed_pipeline + + +def test_rtsp_tls_validation_flags_can_be_scoped_to_one_source( + monkeypatch, +) -> None: + monkeypatch.setenv("ROBOFLOW_RTSP_TLS_VALIDATION_FLAGS", "7") + + self_signed_pipeline = build_gstreamer_pipeline( + "rtsps://camera.example.test/self-signed", + rtsp_tls_validation_flags=0, + ) + default_pipeline = build_gstreamer_pipeline( + "rtsps://camera.example.test/default" + ) + + assert "tls-validation-flags=0 ! " in self_signed_pipeline + assert "tls-validation-flags=7 ! " in default_pipeline + + +@pytest.mark.parametrize("value", ("nope", "-1")) +def test_rtsp_tls_validation_flags_reject_invalid_values(monkeypatch, value) -> None: + monkeypatch.setenv("ROBOFLOW_RTSP_TLS_VALIDATION_FLAGS", value) + + with pytest.raises(ValueError, match="TLS_VALIDATION_FLAGS"): + build_gstreamer_pipeline("rtsps://camera.example.test/live") + + +def test_tensor_rtsps_pipeline_keeps_nvmm_at_named_appsink() -> None: + pipeline = build_gstreamer_pipeline( + "rtsps://camera.example.test:7441/live", + output_tensor=True, + ) + + assert "video/x-raw(memory:NVMM),format=NV12" in pipeline + assert "appsink name=rf_tensor_sink" in pipeline + assert "videoconvert" not in pipeline + assert "video/x-raw,format=BGR" not in pipeline + + +def test_rtsps_source_requires_rtsp_and_nvidia_decode_elements() -> None: + elements = set(required_gstreamer_elements("rtsps://camera.example.test/live")) + + assert { + "appsink", + "h264parse", + "h264timestamper", + "nvvidconv", + "nvv4l2decoder", + "rtph264depay", + "rtspsrc", + } <= elements + assert {"h265parse", "parsebin", "rtph265depay"}.isdisjoint(elements) + # The explicit rtspsrc chain does not autoplug, so the uridecodebin stack + # is no longer part of the RTSP requirements. + assert "uridecodebin" not in elements + assert "videoconvert" not in elements + + +def test_tensor_source_requires_hardware_jpeg_without_cpu_converter() -> None: + elements = set(required_gstreamer_elements("/tmp/sample.mp4", output_tensor=True)) + + assert "nvjpegdec" in elements + assert "nvv4l2decoder" in elements + assert "videoconvert" not in elements + + +def test_file_source_uses_non_dropping_sink_and_matching_demuxer() -> None: + pipeline = build_gstreamer_pipeline("/tmp/sample.mp4") + elements = set(required_gstreamer_elements("/tmp/sample.mp4")) + + assert pipeline.startswith( + f'uridecodebin uri="{Path("/tmp/sample.mp4").resolve().as_uri()}"' + ) + assert "max-buffers=4 drop=false sync=false" in pipeline + assert "qtdemux" in elements + + +def test_csi_source_uses_nvargus_camera() -> None: + pipeline = build_gstreamer_pipeline("csi://2") + + assert pipeline.startswith("nvarguscamerasrc sensor-id=2") + assert "video/x-raw(memory:NVMM),format=NV12" in pipeline + + +def test_v4l2_device_path_uses_live_sink() -> None: + pipeline = build_gstreamer_pipeline("/dev/video3") + + assert pipeline.startswith('v4l2src device="/dev/video3" ! decodebin !') + assert "max-buffers=1 drop=true sync=false" in pipeline + + +def test_v4l2_decodebin_can_negotiate_raw_mjpeg_and_h264_sources() -> None: + pipeline = build_gstreamer_pipeline("/dev/video3", output_tensor=True) + elements = set(required_gstreamer_elements("/dev/video3", output_tensor=True)) + + assert 'v4l2src device="/dev/video3" ! decodebin ! queue' in pipeline + assert "video/x-raw(memory:NVMM),format=RGBA" in pipeline + assert { + "decodebin", + "h264parse", + "jpegparse", + "nvjpegdec", + "nvv4l2decoder", + "v4l2src", + } <= elements diff --git a/tests/inference/unit_tests/core/interfaces/camera/test_rtsp_opencv_tls.py b/tests/inference/unit_tests/core/interfaces/camera/test_rtsp_opencv_tls.py index bf59712c57..282d8071ef 100644 --- a/tests/inference/unit_tests/core/interfaces/camera/test_rtsp_opencv_tls.py +++ b/tests/inference/unit_tests/core/interfaces/camera/test_rtsp_opencv_tls.py @@ -33,7 +33,9 @@ def test_build_options_rtsps_default_strict_verify( def test_build_options_rtsps_with_ca(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setenv(GST_SSL_CA_CERTIFICATE_ENV_VAR, "/etc/ssl/certs/ca-certificates.crt") + monkeypatch.setenv( + GST_SSL_CA_CERTIFICATE_ENV_VAR, "/etc/ssl/certs/ca-certificates.crt" + ) got = build_opencv_ffmpeg_capture_options("rtsps://cam/stream") assert got is not None assert "cafile;/etc/ssl/certs/ca-certificates.crt" in got diff --git a/tests/inference/unit_tests/core/interfaces/camera/test_video_source.py b/tests/inference/unit_tests/core/interfaces/camera/test_video_source.py index 1594dbf747..93b53b940d 100644 --- a/tests/inference/unit_tests/core/interfaces/camera/test_video_source.py +++ b/tests/inference/unit_tests/core/interfaces/camera/test_video_source.py @@ -1,7 +1,8 @@ import time from datetime import datetime +from inspect import signature from queue import Queue -from threading import Thread +from threading import Event, Thread from unittest import mock from unittest.mock import MagicMock, call, patch @@ -10,6 +11,7 @@ import pytest import supervision as sv +from inference.core.env import DEFAULT_BUFFER_SIZE, ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.interfaces.camera import video_source from inference.core.interfaces.camera.entities import ( StatusUpdate, @@ -30,12 +32,299 @@ StreamState, VideoConsumer, VideoSource, + _build_default_producer, decode_video_frame_to_buffer, drop_single_frame_from_buffer, get_fps_if_tick_happens_now, get_from_queue, ) +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="Exercises legacy (flag-off) producer selection; under " + "ENABLE_TENSOR_DATA_REPRESENTATION sources route through the hardware " + "decoder discovery covered by test_discoverability.py", +) + + +@patch("inference.core.interfaces.camera.discoverability.build_hw_producer") +@patch.object(video_source, "ENABLE_TENSOR_DATA_REPRESENTATION", True) +def test_default_producer_requests_numpy_frames_for_standard_consumers( + build_hw_producer: MagicMock, +) -> None: + producer = MagicMock() + build_hw_producer.return_value = producer + + result = _build_default_producer("rtsps://camera.example.test/live") + + assert result is producer + build_hw_producer.assert_called_once_with( + "rtsps://camera.example.test/live", + output_tensor=False, + ) + + +@patch("inference.core.interfaces.camera.discoverability.build_hw_producer") +@patch.object(video_source, "ENABLE_TENSOR_DATA_REPRESENTATION", True) +def test_default_producer_requests_tensor_frames_for_tensor_consumers( + build_hw_producer: MagicMock, +) -> None: + producer = MagicMock() + build_hw_producer.return_value = producer + + result = _build_default_producer( + "rtsps://camera.example.test/live", output_tensor=True + ) + + assert result is producer + build_hw_producer.assert_called_once_with( + "rtsps://camera.example.test/live", + output_tensor=True, + ) + + +@patch("inference.core.interfaces.camera.discoverability.build_hw_producer") +@patch.object(video_source, "ENABLE_TENSOR_DATA_REPRESENTATION", True) +def test_default_producer_forwards_per_source_options( + build_hw_producer: MagicMock, +) -> None: + producer = MagicMock() + build_hw_producer.return_value = producer + + result = _build_default_producer( + "rtsps://camera.example.test/live", + output_tensor=True, + producer_options={"rtsp_tls_validation_flags": 0}, + ) + + assert result is producer + build_hw_producer.assert_called_once_with( + "rtsps://camera.example.test/live", + output_tensor=True, + rtsp_tls_validation_flags=0, + ) + + +@pytest.mark.timeout(90) +def test_tensor_enabled_source_requests_tensor_producer() -> None: + # Covers the allow_tensor_frames kwarg hop from VideoSource.init through + # _start to _build_default_producer; a dropped kwarg silently degrades + # tensor pipelines to the numpy path. + properties = SourceProperties( + width=320, + height=180, + total_frames=0, + is_file=False, + fps=30.0, + ) + + class ImmediateEosProducer: + def __init__(self) -> None: + self.opened = True + + def isOpened(self) -> bool: + return self.opened + + def initialize_source_properties(self, properties) -> None: + return None + + def discover_source_properties(self) -> SourceProperties: + return properties + + def grab(self) -> bool: + self.opened = False + return False + + def release(self) -> None: + self.opened = False + + producer = ImmediateEosProducer() + with patch.object( + video_source, "_build_default_producer", return_value=producer + ) as build_mock: + source = VideoSource.init( + video_reference="rtsp://camera.example.test/live", + allow_tensor_frames=True, + ) + source.start() + source._stream_consumption_thread.join(timeout=5.0) + + build_mock.assert_called_once_with( + "rtsp://camera.example.test/live", + output_tensor=True, + ) + assert not source._stream_consumption_thread.is_alive() + + +@pytest.mark.timeout(90) +def test_video_source_forwards_options_to_default_producer() -> None: + properties = SourceProperties( + width=320, + height=180, + total_frames=0, + is_file=False, + fps=30.0, + ) + + class ImmediateEosProducer: + def isOpened(self) -> bool: + return True + + def initialize_source_properties(self, properties) -> None: + return None + + def discover_source_properties(self) -> SourceProperties: + return properties + + def grab(self) -> bool: + return False + + def release(self) -> None: + return None + + with patch.object( + video_source, + "_build_default_producer", + return_value=ImmediateEosProducer(), + ) as build_mock: + source = VideoSource.init( + video_reference="rtsps://camera.example.test/live", + allow_tensor_frames=True, + video_source_options={"rtsp_tls_validation_flags": 0}, + ) + source.start() + source._stream_consumption_thread.join(timeout=5.0) + + build_mock.assert_called_once_with( + "rtsps://camera.example.test/live", + output_tensor=True, + producer_options={"rtsp_tls_validation_flags": 0}, + ) + + +def test_video_source_init_preserves_legacy_positional_parameter_order() -> None: + legacy_parameters = ( + "video_reference", + "buffer_size", + "status_update_handlers", + "buffer_filling_strategy", + "buffer_consumption_strategy", + "adaptive_mode_stream_pace_tolerance", + "adaptive_mode_reader_pace_tolerance", + "minimum_adaptive_mode_samples", + "maximum_adaptive_frames_dropped_in_row", + "video_source_properties", + "source_id", + "desired_fps", + "allow_tensor_frames", + ) + + assert tuple(signature(VideoSource.init).parameters) == ( + *legacy_parameters, + "video_source_options", + ) + + +@pytest.mark.timeout(90) +def test_async_hardware_initialisation_failure_releases_and_uses_cv2() -> None: + properties = SourceProperties( + width=320, + height=180, + total_frames=0, + is_file=False, + fps=30.0, + ) + hardware_producer = MagicMock() + hardware_producer.isOpened.return_value = True + hardware_producer.discover_source_properties.side_effect = RuntimeError( + "decoder negotiation failed" + ) + + class FallbackProducer: + def __init__(self, video) -> None: + self.opened = True + + def isOpened(self) -> bool: + return self.opened + + def initialize_source_properties(self, properties) -> None: + return None + + def discover_source_properties(self) -> SourceProperties: + return properties + + def grab(self) -> bool: + self.opened = False + return False + + def release(self) -> None: + self.opened = False + + with patch.object( + video_source, "_build_default_producer", return_value=hardware_producer + ), patch.object(video_source, "CV2VideoFrameProducer", FallbackProducer): + source = VideoSource.init(video_reference="rtsp://camera.example.test/live") + source.start() + source._stream_consumption_thread.join(timeout=1.0) + + hardware_producer.release.assert_called_once_with() + assert isinstance(source._video, FallbackProducer) + assert not source._stream_consumption_thread.is_alive() + + +def test_termination_interrupts_a_producer_blocked_waiting_for_a_frame() -> None: + entered_grab = Event() + interrupted = Event() + properties = SourceProperties( + width=320, + height=180, + total_frames=0, + is_file=False, + fps=30.0, + ) + + class BlockingProducer: + def __init__(self) -> None: + self.opened = True + self.interrupt_calls = 0 + + def isOpened(self) -> bool: + return self.opened + + def initialize_source_properties(self, properties) -> None: + return None + + def discover_source_properties(self) -> SourceProperties: + return properties + + def grab(self) -> bool: + entered_grab.set() + interrupted.wait() + return False + + def retrieve(self): + raise AssertionError("retrieve must not run after interruption") + + def interrupt(self) -> None: + self.interrupt_calls += 1 + interrupted.set() + + def release(self) -> None: + self.opened = False + + producer = BlockingProducer() + source = VideoSource.init(video_reference=lambda: producer) + source.start() + assert entered_grab.wait(timeout=1.0) + + started = time.monotonic() + source.terminate(wait_on_frames_consumption=False) + elapsed = time.monotonic() - started + + assert elapsed < 1.0 + assert producer.interrupt_calls == 1 + assert not source._stream_consumption_thread.is_alive() + def tear_down_source(source: VideoSource) -> None: source.terminate(wait_on_frames_consumption=False) @@ -171,7 +460,7 @@ def test_video_source_describe_source_when_stream_consumption_not_yet_started() assert result == SourceMetadata( source_properties=None, source_reference="invalid", - buffer_size=64, + buffer_size=DEFAULT_BUFFER_SIZE, state=StreamState.NOT_STARTED, buffer_filling_strategy=None, buffer_consumption_strategy=None, @@ -179,6 +468,7 @@ def test_video_source_describe_source_when_stream_consumption_not_yet_started() ), "Source description must denote NOT_STARTED state and invalid source reference" +@_NUMPY_ONLY def test_video_source_selects_gstreamer_producer_for_rtsps_on_jetson() -> None: credentialed_url = "rtsps://user:secret@192.168.1.1:554/stream" with patch("inference.core.interfaces.camera.video_source.RUNS_ON_JETSON", True): @@ -190,14 +480,12 @@ def test_video_source_selects_gstreamer_producer_for_rtsps_on_jetson() -> None: "inference.core.interfaces.camera.gstreamer_rtsp_producer.GStreamerRtspVideoFrameProducer" ) as mock_producer_cls: mock_producer_cls.return_value.isOpened.return_value = True - mock_producer_cls.return_value.discover_source_properties.return_value = ( - SourceProperties( - width=640, - height=480, - fps=30.0, - total_frames=0, - is_file=False, - ) + mock_producer_cls.return_value.discover_source_properties.return_value = SourceProperties( + width=640, + height=480, + fps=30.0, + total_frames=0, + is_file=False, ) source = VideoSource.init(video_reference=credentialed_url) source.start() @@ -242,7 +530,7 @@ def test_video_source_describe_source_when_invalid_video_reference_consumption_s assert result == SourceMetadata( source_properties=None, source_reference="invalid", - buffer_size=64, + buffer_size=DEFAULT_BUFFER_SIZE, state=StreamState.ERROR, buffer_filling_strategy=None, buffer_consumption_strategy=None, diff --git a/tests/inference/unit_tests/core/interfaces/stream/test_interface_pipeline.py b/tests/inference/unit_tests/core/interfaces/stream/test_interface_pipeline.py index 020c7eb039..35d44806f9 100644 --- a/tests/inference/unit_tests/core/interfaces/stream/test_interface_pipeline.py +++ b/tests/inference/unit_tests/core/interfaces/stream/test_interface_pipeline.py @@ -42,6 +42,7 @@ default_process_frame, ) from inference.core.interfaces.stream.sinks import active_learning_sink, multi_sink +from inference.core.interfaces.stream.utils import VideoSourceOptions from inference.core.interfaces.stream.watchdog import BasePipelineWatchDog @@ -812,6 +813,177 @@ def test_inference_pipeline_factories_expose_optional_exec_session_id( assert parameter.default is None +@pytest.mark.parametrize( + "factory_name", + ["init", "init_with_yolo_world", "init_with_workflow", "init_with_custom_logic"], +) +def test_inference_pipeline_factories_expose_per_source_options( + factory_name: str, +) -> None: + parameter = signature(getattr(InferencePipeline, factory_name)).parameters[ + "video_source_options" + ] + + assert parameter.annotation == Optional[VideoSourceOptions] + assert parameter.default is None + + +_LEGACY_FACTORY_PARAMETERS = { + "init": ( + "video_reference", + "model_id", + "on_prediction", + "api_key", + "max_fps", + "watchdog", + "status_update_handlers", + "source_buffer_filling_strategy", + "source_buffer_consumption_strategy", + "class_agnostic_nms", + "confidence", + "iou_threshold", + "max_candidates", + "max_detections", + "mask_decode_mode", + "tradeoff_factor", + "active_learning_enabled", + "video_source_properties", + "active_learning_target_dataset", + "batch_collection_timeout", + "video_processing_mode", + "max_staleness", + "sink_mode", + "predictions_queue_size", + "decoding_buffer_size", + "exec_session_id", + ), + "init_with_yolo_world": ( + "video_reference", + "classes", + "model_size", + "on_prediction", + "max_fps", + "watchdog", + "status_update_handlers", + "source_buffer_filling_strategy", + "source_buffer_consumption_strategy", + "class_agnostic_nms", + "confidence", + "iou_threshold", + "max_candidates", + "max_detections", + "video_source_properties", + "batch_collection_timeout", + "video_processing_mode", + "max_staleness", + "sink_mode", + "predictions_queue_size", + "decoding_buffer_size", + "exec_session_id", + ), + "init_with_workflow": ( + "video_reference", + "workflow_specification", + "workspace_name", + "workflow_id", + "api_key", + "image_input_name", + "workflows_parameters", + "on_prediction", + "max_fps", + "watchdog", + "status_update_handlers", + "source_buffer_filling_strategy", + "source_buffer_consumption_strategy", + "video_source_properties", + "workflow_init_parameters", + "disable_sinks", + "workflows_thread_pool_workers", + "cancel_thread_pool_tasks_on_exit", + "video_metadata_input_name", + "batch_collection_timeout", + "video_processing_mode", + "max_staleness", + "profiling_directory", + "use_workflow_definition_cache", + "serialize_results", + "predictions_queue_size", + "decoding_buffer_size", + "model_manager", + "_is_preview", + "workflow_version_id", + "exec_session_id", + "workflows_dependencies_pre_init", + ), + "init_with_custom_logic": ( + "video_reference", + "on_video_frame", + "on_prediction", + "on_pipeline_start", + "on_pipeline_end", + "max_fps", + "watchdog", + "status_update_handlers", + "source_buffer_filling_strategy", + "source_buffer_consumption_strategy", + "video_source_properties", + "batch_collection_timeout", + "video_processing_mode", + "max_staleness", + "sink_mode", + "predictions_queue_size", + "decoding_buffer_size", + "exec_session_id", + "allow_tensor_frames", + ), +} + + +@pytest.mark.parametrize("factory_name", _LEGACY_FACTORY_PARAMETERS) +def test_inference_pipeline_factories_preserve_legacy_positional_parameter_order( + factory_name: str, +) -> None: + parameter_names = tuple( + signature(getattr(InferencePipeline, factory_name)).parameters + ) + + assert parameter_names == ( + *_LEGACY_FACTORY_PARAMETERS[factory_name], + "video_source_options", + ) + + +def test_inference_pipeline_init_propagates_per_source_options(monkeypatch) -> None: + pipeline = MagicMock() + init_with_custom_logic = MagicMock(return_value=pipeline) + monkeypatch.setattr( + "inference.core.interfaces.stream.inference_pipeline.get_model", + MagicMock(), + ) + monkeypatch.setattr( + InferencePipeline, + "init_with_custom_logic", + init_with_custom_logic, + ) + video_source_options = [ + {"rtsp_tls_validation_flags": 0}, + None, + ] + + result = InferencePipeline.init( + video_reference=["rtsps://camera-1", "rtsp://camera-2"], + model_id="example/1", + active_learning_enabled=False, + video_source_options=video_source_options, + ) + + assert result is pipeline + assert ( + init_with_custom_logic.call_args.kwargs["video_source_options"] + is video_source_options + ) + + @pytest.mark.parametrize("disable_sinks", [False, True]) def test_init_with_workflow_injects_sink_execution_policy( disable_sinks: bool, diff --git a/tests/inference/unit_tests/core/interfaces/stream/test_tensor_sink_boundary.py b/tests/inference/unit_tests/core/interfaces/stream/test_tensor_sink_boundary.py new file mode 100644 index 0000000000..6701edee70 --- /dev/null +++ b/tests/inference/unit_tests/core/interfaces/stream/test_tensor_sink_boundary.py @@ -0,0 +1,134 @@ +from datetime import datetime + +import numpy as np +import pytest + +from inference.core.interfaces.camera.entities import VideoFrame +from inference.core.interfaces.stream.inference_pipeline import InferencePipeline +from inference.core.interfaces.stream.utils import materialise_video_frame_for_sink + +torch = pytest.importorskip("torch") + + +def test_tensor_frame_reaches_sink_unmaterialised() -> None: + # Under ENABLE_TENSOR_DATA_REPRESENTATION nothing materialises in dispatch: + # the sink receives the ORIGINAL on-device tensor frame. Pixel-consuming + # sinks convert at their own boundary via materialise_video_frame_for_sink. + tensor_image = torch.tensor( + [ + [[10, 11], [12, 13]], + [[20, 21], [22, 23]], + [[30, 31], [32, 33]], + ], + dtype=torch.uint8, + ) + video_frame = VideoFrame( + image=tensor_image, + frame_id=1, + frame_timestamp=datetime.now(), + source_id=0, + ) + received = {} + + def sink(predictions, frame): + received["predictions"] = predictions + received["frame"] = frame + + pipeline = InferencePipeline.__new__(InferencePipeline) + pipeline._on_prediction = sink + pipeline._status_update_handlers = [] + + pipeline._use_sink({"result": "ok"}, video_frame) + + assert received["predictions"] == {"result": "ok"} + assert received["frame"] is video_frame + assert received["frame"].image is tensor_image + + +def test_tensor_frames_reach_batch_sink_unmaterialised() -> None: + video_frame = VideoFrame( + image=torch.arange(4, dtype=torch.uint8).reshape(1, 2, 2), + frame_id=1, + frame_timestamp=datetime.now(), + source_id=0, + ) + received = {} + + def sink(predictions, frames): + received["predictions"] = predictions + received["frames"] = frames + + pipeline = InferencePipeline.__new__(InferencePipeline) + pipeline._on_prediction = sink + pipeline._status_update_handlers = [] + + pipeline._use_sink([{"result": "ok"}, None], [video_frame, None]) + + assert received["predictions"] == [{"result": "ok"}, None] + assert received["frames"][0] is video_frame + assert received["frames"][1] is None + + +def test_numpy_frame_is_passed_to_sink_without_copying() -> None: + video_frame = VideoFrame( + image=np.zeros((2, 2, 3), dtype=np.uint8), + frame_id=1, + frame_timestamp=datetime.now(), + source_id=0, + ) + received = {} + + def sink(predictions, frame): + received["frame"] = frame + + pipeline = InferencePipeline.__new__(InferencePipeline) + pipeline._on_prediction = sink + pipeline._status_update_handlers = [] + + pipeline._use_sink({}, video_frame) + + assert received["frame"] is video_frame + + +def test_materialise_helper_converts_3chw_tensor_to_bgr_hwc_numpy() -> None: + tensor_image = torch.tensor( + [ + [[10, 11], [12, 13]], + [[20, 21], [22, 23]], + [[30, 31], [32, 33]], + ], + dtype=torch.uint8, + ) + video_frame = VideoFrame( + image=tensor_image, + frame_id=1, + frame_timestamp=datetime.now(), + source_id=0, + ) + + materialised = materialise_video_frame_for_sink(video_frame) + + expected_bgr = np.array( + [ + [[30, 20, 10], [31, 21, 11]], + [[32, 22, 12], [33, 23, 13]], + ], + dtype=np.uint8, + ) + assert materialised is not video_frame + assert materialised.image.flags.c_contiguous + np.testing.assert_array_equal(materialised.image, expected_bgr) + # The original frame stays untouched (frozen-dataclass replace semantics). + assert video_frame.image is tensor_image + + +def test_materialise_helper_passes_numpy_frame_through() -> None: + video_frame = VideoFrame( + image=np.zeros((2, 2, 3), dtype=np.uint8), + frame_id=1, + frame_timestamp=datetime.now(), + source_id=0, + ) + + assert materialise_video_frame_for_sink(video_frame) is video_frame + assert materialise_video_frame_for_sink(None) is None diff --git a/tests/inference/unit_tests/core/interfaces/stream/test_utils.py b/tests/inference/unit_tests/core/interfaces/stream/test_utils.py index 4843f5b513..94061f9d40 100644 --- a/tests/inference/unit_tests/core/interfaces/stream/test_utils.py +++ b/tests/inference/unit_tests/core/interfaces/stream/test_utils.py @@ -2,6 +2,7 @@ import os from concurrent.futures import ThreadPoolExecutor from glob import glob +from inspect import signature from unittest import mock from unittest.mock import MagicMock @@ -10,7 +11,9 @@ from inference.core.interfaces.stream import utils from inference.core.interfaces.stream.utils import ( broadcast_elements, + initialise_video_sources, on_pipeline_end, + prepare_video_sources, save_workflows_profiler_trace, wrap_in_list, ) @@ -90,6 +93,91 @@ def test_broadcast_elements_when_input_is_empty() -> None: ) +@pytest.mark.parametrize( + "callable_under_test", + [prepare_video_sources, initialise_video_sources], +) +def test_video_source_helpers_preserve_legacy_positional_parameter_order( + callable_under_test, +) -> None: + legacy_parameters = ( + "video_reference", + "video_source_properties", + "status_update_handlers", + "source_buffer_filling_strategy", + "source_buffer_consumption_strategy", + "desired_source_fps", + "decoding_buffer_size", + "allow_tensor_frames", + ) + + assert tuple(signature(callable_under_test).parameters) == ( + *legacy_parameters, + "video_source_options", + ) + + +@mock.patch.object(utils.VideoSource, "init") +def test_prepare_video_sources_accepts_legacy_positional_arguments( + video_source_init: MagicMock, +) -> None: + prepare_video_sources(["a"], None, None, None, None, 12, 3, True) + + video_source_init.assert_called_once_with( + video_reference="a", + status_update_handlers=None, + buffer_filling_strategy=None, + buffer_consumption_strategy=None, + video_source_properties=None, + video_source_options=None, + source_id=0, + desired_fps=12, + buffer_size=3, + allow_tensor_frames=True, + ) + + +@mock.patch.object(utils.VideoSource, "init") +def test_prepare_video_sources_broadcasts_per_source_options( + video_source_init: MagicMock, +) -> None: + prepare_video_sources( + video_reference=["a", "b"], + video_source_properties=None, + video_source_options={"rtsp_tls_validation_flags": 0}, + status_update_handlers=None, + source_buffer_filling_strategy=None, + source_buffer_consumption_strategy=None, + ) + + assert video_source_init.call_count == 2 + for call in video_source_init.call_args_list: + assert call.kwargs["video_source_options"] == {"rtsp_tls_validation_flags": 0} + + +@mock.patch.object(utils.VideoSource, "init") +def test_prepare_video_sources_applies_aligned_per_source_options( + video_source_init: MagicMock, +) -> None: + prepare_video_sources( + video_reference=["a", "b"], + video_source_properties=None, + video_source_options=[ + {"rtsp_tls_validation_flags": 0}, + None, + ], + status_update_handlers=None, + source_buffer_filling_strategy=None, + source_buffer_consumption_strategy=None, + ) + + assert video_source_init.call_count == 2 + assert video_source_init.call_args_list[0].kwargs["video_source_options"] == { + "rtsp_tls_validation_flags": 0 + } + assert video_source_init.call_args_list[1].kwargs["video_source_options"] is None + + def test_save_workflows_profiler_trace(empty_directory: str) -> None: # when save_workflows_profiler_trace( @@ -182,5 +270,6 @@ def mock_makedirs(name, *args, **kwargs): assert thread_pool_executor._shutdown is True, "Expected pool executor to be closed" # No profiling files should have been created json_files_in_directory = glob(os.path.join(read_only_dir, "*.json")) - assert len(json_files_in_directory) == 0, "Expected no profiler trace saved on read-only FS" - + assert ( + len(json_files_in_directory) == 0 + ), "Expected no profiler trace saved on read-only FS" diff --git a/tests/inference/unit_tests/core/managers/test_base.py b/tests/inference/unit_tests/core/managers/test_base.py index 78b62284a7..f3c9dbc266 100644 --- a/tests/inference/unit_tests/core/managers/test_base.py +++ b/tests/inference/unit_tests/core/managers/test_base.py @@ -258,6 +258,60 @@ async def test_infer_from_request_skips_model_monitoring_cache_offline( cache_mock.zadd.assert_not_called() +def test_run_tensor_native_inference_records_telemetry( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model = MagicMock() + model.run_tensor_native_inference.return_value = "result" + model_manager = ModelManager(model_registry=MagicMock(), models={"some/1": model}) + start_span_mock = MagicMock() + record_inference_mock = MagicMock() + record_error_mock = MagicMock() + monkeypatch.setattr(base_module, "start_span", start_span_mock) + monkeypatch.setattr(base_module, "record_inference", record_inference_mock) + monkeypatch.setattr(base_module, "record_error", record_error_mock) + + result = model_manager.run_tensor_native_inference( + model_id="some/1", images="tensor" + ) + + assert result == "result" + model.run_tensor_native_inference.assert_called_once_with(images="tensor") + start_span_mock.assert_called_once_with( + "model.infer", + { + "model.id": "some/1", + "model.infer.caller": "run_tensor_native_inference", + }, + ) + record_inference_mock.assert_called_once() + assert record_inference_mock.call_args.args[0] == "some/1" + assert record_inference_mock.call_args.args[1] >= 0 + record_error_mock.assert_not_called() + + +def test_run_tensor_native_inference_records_errors( + monkeypatch: pytest.MonkeyPatch, +) -> None: + error = RuntimeError("inference failed") + model = MagicMock() + model.run_tensor_native_inference.side_effect = error + model_manager = ModelManager(model_registry=MagicMock(), models={"some/1": model}) + start_span_mock = MagicMock() + record_inference_mock = MagicMock() + record_error_mock = MagicMock() + monkeypatch.setattr(base_module, "start_span", start_span_mock) + monkeypatch.setattr(base_module, "record_inference", record_inference_mock) + monkeypatch.setattr(base_module, "record_error", record_error_mock) + + with pytest.raises(RuntimeError) as raised_error: + model_manager.run_tensor_native_inference(model_id="some/1", images="tensor") + + assert raised_error.value is error + record_inference_mock.assert_not_called() + record_error_mock.assert_called_once_with(error) + + def test_make_response_when_model_available() -> None: # given model_registry = MagicMock() diff --git a/tests/inference/unit_tests/core/managers/test_cuda_memory_watchdog.py b/tests/inference/unit_tests/core/managers/test_cuda_memory_watchdog.py index d123ceb973..90004a0b54 100644 --- a/tests/inference/unit_tests/core/managers/test_cuda_memory_watchdog.py +++ b/tests/inference/unit_tests/core/managers/test_cuda_memory_watchdog.py @@ -127,7 +127,9 @@ def test_start_does_not_spin_daemon_when_cuda_unavailable() -> None: def test_start_launches_daemon_and_reclaims_then_stop_joins() -> None: # given reclaimed = threading.Event() - with mock.patch.object(mod, "cuda_is_available", return_value=True), mock.patch.object( + with mock.patch.object( + mod, "cuda_is_available", return_value=True + ), mock.patch.object( mod, "reclaim_cuda_memory", side_effect=lambda: reclaimed.set() ): wd = CudaMemoryReclamationWatchdog(interval_seconds=5.0) @@ -145,9 +147,9 @@ def test_start_launches_daemon_and_reclaims_then_stop_joins() -> None: def test_start_is_idempotent_while_running() -> None: - with mock.patch.object(mod, "cuda_is_available", return_value=True), mock.patch.object( - mod, "reclaim_cuda_memory", return_value=None - ): + with mock.patch.object( + mod, "cuda_is_available", return_value=True + ), mock.patch.object(mod, "reclaim_cuda_memory", return_value=None): wd = CudaMemoryReclamationWatchdog(interval_seconds=5.0) wd.start() try: diff --git a/tests/inference/unit_tests/core/models/test_inference_models_adapters.py b/tests/inference/unit_tests/core/models/test_inference_models_adapters.py index 3b981cf6f9..cdc35bdd28 100644 --- a/tests/inference/unit_tests/core/models/test_inference_models_adapters.py +++ b/tests/inference/unit_tests/core/models/test_inference_models_adapters.py @@ -4,6 +4,7 @@ from concurrent.futures import Future from types import SimpleNamespace +import numpy as np import pytest import torch @@ -25,6 +26,7 @@ InstanceDetections, MultiLabelClassificationPrediction, ) +from inference_models.models.auto_loaders.entities import PreProcessingOverrides from inference_models.models.base.async_handoff import attach_adapter_mapped_kwargs @@ -552,3 +554,71 @@ def test_depth_estimation_adapter_normalization_matches_depth_anything_conventio expected = np.array([[1.0, 2 / 3], [1 / 3, 0.0]], dtype=np.float32) assert np.allclose(result["normalized_depth"], expected, atol=1e-6) + + +def _make_depth_estimation_adapter(model) -> InferenceModelsDepthEstimationAdapter: + adapter = object.__new__(InferenceModelsDepthEstimationAdapter) + adapter._model = model + return adapter + + +def test_depth_estimation_adapter_tensor_native_negates_metric_depth() -> None: + """The tensor-native depth contract mirrors the DepthAnything adapters: + raw per-image maps in which larger means closer, normalized by the caller. + YOLO26-depth emits metric depth (larger == farther), so the adapter must + negate it while keeping kwargs mapping consistent with the other five + tensor-native adapter overrides.""" + captured = {} + + def fake_model(images, **kwargs): + captured["images"] = images + captured["kwargs"] = kwargs + return [torch.tensor([[1.0, 2.0], [3.0, 4.0]])] + + adapter = _make_depth_estimation_adapter(fake_model) + images = [torch.zeros((3, 2, 2), dtype=torch.uint8)] + + result = adapter.run_tensor_native_inference(images, input_color_format="rgb") + + assert captured["images"] is images + assert captured["kwargs"]["input_color_format"] == "rgb" + assert isinstance( + captured["kwargs"]["pre_processing_overrides"], PreProcessingOverrides + ) + assert len(result) == 1 + assert torch.equal(result[0], torch.tensor([[-1.0, -2.0], [-3.0, -4.0]])) + + +def test_depth_estimation_adapter_tensor_native_passes_missing_color_format_as_none() -> ( + None +): + captured = {} + + def fake_model(images, **kwargs): + captured["kwargs"] = kwargs + return [torch.tensor([[0.0, 1.0]])] + + adapter = _make_depth_estimation_adapter(fake_model) + + adapter.run_tensor_native_inference([torch.zeros((3, 1, 2), dtype=torch.uint8)]) + + assert captured["kwargs"]["input_color_format"] is None + + +def test_depth_estimation_adapter_tensor_native_composes_to_numpy_normalization() -> ( + None +): + """min-max normalization of the tensor-native output (what the flag-on + depth-estimation block computes) must reproduce the numpy path's + `(max - map) / (max - min)` proximity map exactly.""" + adapter = _make_depth_estimation_adapter( + lambda images, **kwargs: [torch.tensor([[1.0, 2.0], [3.0, 4.0]])] + ) + + (depth_map,) = adapter.run_tensor_native_inference( + [torch.zeros((3, 2, 2), dtype=torch.uint8)], input_color_format="bgr" + ) + normalized = (depth_map - depth_map.min()) / (depth_map.max() - depth_map.min()) + + expected = np.array([[1.0, 2 / 3], [1 / 3, 0.0]], dtype=np.float32) + assert np.allclose(normalized.numpy(), expected, atol=1e-6) diff --git a/tests/inference_cli/unit_tests/lib/test_container_adapter.py b/tests/inference_cli/unit_tests/lib/test_container_adapter.py index a099f4a832..6d630f25f0 100644 --- a/tests/inference_cli/unit_tests/lib/test_container_adapter.py +++ b/tests/inference_cli/unit_tests/lib/test_container_adapter.py @@ -281,12 +281,18 @@ def test_jetson_images_table_is_sorted_descending() -> None: ) +def test_jetson_images_have_unique_jetpack_prefixes() -> None: + prefixes = [entry.jetpack_prefix for entry in _JETSON_IMAGES] + assert len(prefixes) == len(set(prefixes)) + + JETSON_450 = "roboflow/roboflow-inference-server-jetson-4.5.0:latest" JETSON_461 = "roboflow/roboflow-inference-server-jetson-4.6.1:latest" JETSON_511 = "roboflow/roboflow-inference-server-jetson-5.1.1:latest" JETSON_600 = "roboflow/roboflow-inference-server-jetson-6.0.0:latest" JETSON_620 = "roboflow/roboflow-inference-server-jetson-6.2.0:latest" JETSON_710 = "roboflow/roboflow-inference-server-jetson-7.1.0:latest" +JETSON_720 = "roboflow/roboflow-inference-server-jetson-7.2.0:latest" class TestParseTegraRelease: @@ -330,8 +336,8 @@ class TestImageForL4t: (36, 3, JETSON_600), (36, 4, JETSON_620), (36, 5, JETSON_620), - (38, 0, JETSON_710), (38, 4, JETSON_710), + (39, 2, JETSON_720), ], ) def test_l4t_to_image(self, l4t_major: int, l4t_minor: int, expected: str) -> None: @@ -340,6 +346,14 @@ def test_l4t_to_image(self, l4t_major: int, l4t_minor: int, expected: str) -> No def test_returns_none_for_unknown_l4t_major(self) -> None: assert _image_for_l4t(99, 1) is None + @pytest.mark.parametrize( + "l4t_major, l4t_minor", [(38, 0), (38, 2), (39, 0), (39, 1)] + ) + def test_returns_none_for_unsupported_jetpack_7_release( + self, l4t_major: int, l4t_minor: int + ) -> None: + assert _image_for_l4t(l4t_major, l4t_minor) is None + class TestGetJetpackImage: @pytest.mark.parametrize( @@ -358,37 +372,68 @@ class TestGetJetpackImage: ("6.2.0", JETSON_620), ("7.1", JETSON_710), ("7.1.0", JETSON_710), + ("7.1-b123", JETSON_710), + ("7.2", JETSON_720), + ("7.2.0", JETSON_720), + ("7.2-b187", JETSON_720), ], ) def test_returns_correct_image(self, version: str, expected_image: str) -> None: assert _get_jetpack_image(version) == expected_image - def test_raises_for_unsupported_version(self) -> None: + @pytest.mark.parametrize("version", ["3.0", "7"]) + def test_raises_for_unsupported_version(self, version: str) -> None: with pytest.raises(RuntimeError, match="not supported"): - _get_jetpack_image("3.0") + _get_jetpack_image(version) class TestDetectJetson: - def test_detects_from_tegra_release(self) -> None: - content = "# R36 (release), REVISION: 4.0, GCID: 12345, BOARD: generic" + @pytest.mark.parametrize( + "l4t_major, l4t_minor, expected_image", + [ + (36, 4, JETSON_620), + (38, 4, JETSON_710), + (39, 2, JETSON_720), + ], + ) + def test_detects_from_tegra_release( + self, l4t_major: int, l4t_minor: int, expected_image: str + ) -> None: + content = ( + f"# R{l4t_major} (release), REVISION: {l4t_minor}.0, " + "GCID: 12345, BOARD: generic" + ) with patch("builtins.open", mock_open(read_data=content)): result = _detect_jetson() assert result is not None image, source = result - assert image == JETSON_620 + assert image == expected_image assert "/etc/nv_tegra_release" in source @patch.object(container_adapter, "_parse_tegra_release", return_value=None) - @patch.object( - container_adapter, "_get_jetpack_version_from_dpkg", return_value="6.0" + @pytest.mark.parametrize( + "jetpack_version, expected_image", + [ + ("6.0", JETSON_600), + ("7.1", JETSON_710), + ("7.2", JETSON_720), + ], ) def test_falls_back_to_dpkg( - self, _dpkg_mock: MagicMock, _tegra_mock: MagicMock + self, + _tegra_mock: MagicMock, + jetpack_version: str, + expected_image: str, ) -> None: - result = _detect_jetson() + with patch.object( + container_adapter, + "_get_jetpack_version_from_dpkg", + return_value=jetpack_version, + ): + result = _detect_jetson() assert result is not None image, source = result - assert image == JETSON_600 + assert image == expected_image assert "dpkg" in source @patch.object(container_adapter, "_parse_tegra_release", return_value=None) diff --git a/tests/workflows/integration_tests/execution/control_flow_with_side_effects/test_workflow_with_control_flow_with_side_effects.py b/tests/workflows/integration_tests/execution/control_flow_with_side_effects/test_workflow_with_control_flow_with_side_effects.py index 0fa536a5f5..739a996853 100644 --- a/tests/workflows/integration_tests/execution/control_flow_with_side_effects/test_workflow_with_control_flow_with_side_effects.py +++ b/tests/workflows/integration_tests/execution/control_flow_with_side_effects/test_workflow_with_control_flow_with_side_effects.py @@ -12,6 +12,7 @@ import pandas as pd import pytest import supervision as sv +import torch from inference.core.entities.requests.inference import ObjectDetectionInferenceRequest from inference.core.entities.responses.inference import ( @@ -19,17 +20,38 @@ ObjectDetectionInferenceResponse, ObjectDetectionPrediction, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.core_steps.fusion.detections_stitch.v1 import ( DetectionsStitchBlockV1, ) +from inference.core.workflows.core_steps.fusion.detections_stitch.v1_tensor import ( + DetectionsStitchBlockV1 as DetectionsStitchBlockV1Tensor, +) from inference.core.workflows.errors import ( ControlFlowDefinitionError, StepInputLineageError, ) from inference.core.workflows.execution_engine.core import ExecutionEngine from inference.core.workflows.execution_engine.introspection import blocks_loader +from inference_models.models.base.object_detection import Detections as NativeDetections + +# Model-based workflow tests mock inference at the ModelManager seam. Under +# ENABLE_TENSOR_DATA_REPRESENTATION the OD block routes through +# `run_tensor_native_inference` / `get_class_names` (not `infer_from_request_sync`), +# so the numpy-mocked tests skip when the flag is on and a `*_tensor_native` parity +# test (skipped when the flag is off) drives the same scenario through the native +# mock harness (`_run_workflow_tensor_native`). +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="numpy ModelManager mock (infer_from_request_sync); the OD block is " + "native under ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) _WORKFLOW_DEFINITIONS_DIR = Path(__file__).resolve().parent / "workflow_definitions" @@ -229,493 +251,575 @@ def _run_workflow( ) -@patch( - "inference.core.workflows.core_steps.sinks.email_notification.v2.send_email_via_roboflow_proxy" +def _make_person_detections_native(k: int) -> NativeDetections: + """Native equivalent of `_make_person_prediction` x k: a `Detections` with k + person boxes (class_id 0, confidence 0.9). The OD tensor block attaches the + workflow image lineage / detection ids itself, so this is left raw.""" + if k == 0: + xyxy = torch.zeros((0, 4), dtype=torch.float32) + else: + xyxy = torch.tensor([[75.0, 75.0, 125.0, 125.0]] * k, dtype=torch.float32) + return NativeDetections( + xyxy=xyxy, + class_id=torch.zeros((k,), dtype=torch.long), + confidence=torch.full((k,), 0.9, dtype=torch.float32), + ) + + +def make_mock_tensor_native_detections_per_model( + model_id_to_counts: Dict[str, List[int]], +): + """Tensor-native analogue of `make_mock_detection_responses_per_model`: a side + effect for `ModelManager.run_tensor_native_inference` returning one native + `Detections` per image, with the configured number of person boxes.""" + model_id_index: Dict[str, int] = {mid: 0 for mid in model_id_to_counts} + + def mock_fn(model_id: str, images, **kwargs) -> List[NativeDetections]: + if model_id not in model_id_to_counts: + raise ValueError( + f"Mock received unknown model_id={model_id!r}. " + f"Known: {list(model_id_to_counts)}." + ) + counts = model_id_to_counts[model_id] + n = len(images) + start = model_id_index[model_id] + end = start + n + if end > len(counts): + raise ValueError( + f"Mock for model_id={model_id!r} expected at least {end} counts " + f"but only {len(counts)} defined. Check detection_counts for this scenario." + ) + chunk = counts[start:end] + model_id_index[model_id] = end + return [_make_person_detections_native(k) for k in chunk] + + return mock_fn + + +def _run_workflow_tensor_native( + workflow_definition: dict, + runtime_parameters: dict, + model_manager, + detection_counts_per_model: Dict[str, List[int]], +): + """Mirror of `_run_workflow` for ENABLE_TENSOR_DATA_REPRESENTATION. The tensor OD + block runs inference via `ModelManager.run_tensor_native_inference` and reads class + names via `get_class_names` (instead of `infer_from_request_sync`), so those are the + seams mocked here; `add_model` stays a no-op as in the numpy harness.""" + init_params = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + + engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + init_parameters=init_params, + max_concurrent_steps=1, + ) + + mock_fn = make_mock_tensor_native_detections_per_model(detection_counts_per_model) + + with patch.object(ModelManager, "add_model"): + with patch.object(ModelManager, "get_class_names", return_value=["person"]): + with patch.object( + ModelManager, + "run_tensor_native_inference", + side_effect=mock_fn, + ): + return engine.run(runtime_parameters=runtime_parameters) + + +_SIDE_EFFECT_ARGNAMES = ( + "image_gen_fn, names, detection_counts_per_model, enable_email, workflow_name, " + "expected_call_count, expected_receiver_email, expected_subject, " + "expected_message_parameters, expected_result" ) -@pytest.mark.parametrize( - "image_gen_fn,\ - names,\ - detection_counts_per_model,\ - enable_email,\ - workflow_name,\ - expected_call_count,\ - expected_receiver_email,\ - expected_subject,\ - expected_message_parameters,\ - expected_result", - [ + +# Shared scenario table for the side-effect control-flow test, consumed by both the +# @_NUMPY_ONLY and @_TENSOR_ONLY variants so they stay in lockstep. +_SIDE_EFFECT_SCENARIOS = [ + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + True, + "with_email_message_params", + 2, + "noreply@example.com", + "Detections found", ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - True, - "with_email_message_params", - 2, - "noreply@example.com", - "Detections found", - ( - {"num_detections": 2}, - {"num_detections": 1}, - ), - ( - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {"num_detections": 2}, + {"num_detections": 1}, ), ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - True, - "without_email_message_params", - 2, - "noreply@example.com", - "Detections found", - ( - {}, - {}, - ), - ( - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, ), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + True, + "without_email_message_params", + 2, + "noreply@example.com", + "Detections found", ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - True, - "with_image_names_and_email_message_params", - 2, - "noreply@example.com", - "Detections found", - ( - {"num_detections": 2, "name": "img1"}, - {"num_detections": 1, "name": "img3"}, - ), - ( - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {}, + {}, ), ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - True, - "with_image_names_and_without_email_message_params", - 2, - "noreply@example.com", - "Detections found", - ( - {"name": "img1"}, - {"name": "img3"}, - ), - ( - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, ), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + True, + "with_image_names_and_email_message_params", + 2, + "noreply@example.com", + "Detections found", ( - _sliced_4_images, - SLICED_NAMES, - {"yolov8n-640": SLICED_DETECTION_COUNTS}, - True, - "sliced_image_with_email_message_params", - 3, - "noreply@example.com", - "Detections found", - ( - {"num_detections": 1}, - {"num_detections": 1}, - {"num_detections": 1}, - ), - ( - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - ] - }, - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - None, - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - None, - None, - ] - }, - ), + {"num_detections": 2, "name": "img1"}, + {"num_detections": 1, "name": "img3"}, ), ( - _sliced_4_images, - SLICED_NAMES, - {"yolov8n-640": SLICED_DETECTION_COUNTS}, - True, - "sliced_image_without_email_message_params", - 3, - "noreply@example.com", - "Detections found", - ( - {}, - {}, - {}, - ), - ( - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - ] - }, - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - None, - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - None, - None, - ] - }, - ), + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, ), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + True, + "with_image_names_and_without_email_message_params", + 2, + "noreply@example.com", + "Detections found", ( - _sliced_4_images, - SLICED_NAMES, - {"yolov8n-640": SLICED_DETECTION_COUNTS}, - True, - "sliced_image_with_email_message_params_and_area_size_step", - 3, - "noreply@example.com", - "Detections found", - ( - {"num_detections": 1, "area_converted": 2500}, - {"num_detections": 1, "area_converted": 2500}, - {"num_detections": 1, "area_converted": 2500}, - ), - ( - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - ] - }, - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - None, - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - None, - None, - ] - }, - ), + {"name": "img1"}, + {"name": "img3"}, ), ( - _sliced_4_images, - SLICED_NAMES, - {"yolov8n-640": SLICED_DETECTION_COUNTS}, - True, - "sliced_image_without_email_message_params_and_area_size_step", - 3, - "noreply@example.com", - "Detections found", - ( - {"area_converted": 2500}, - {"area_converted": 2500}, - {"area_converted": 2500}, - ), - ( - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - ] - }, - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - None, - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - None, - None, - ] - }, - ), + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, ), + ), + ( + _sliced_4_images, + SLICED_NAMES, + {"yolov8n-640": SLICED_DETECTION_COUNTS}, + True, + "sliced_image_with_email_message_params", + 3, + "noreply@example.com", + "Detections found", ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - True, - "with_email_gate_and_with_email_message_params", - 2, - "noreply@example.com", - "Detections found", - ( - {"num_detections": 2}, - {"num_detections": 1}, - ), - ( - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {"num_detections": 1}, + {"num_detections": 1}, + {"num_detections": 1}, ), ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - True, - "with_email_gate_and_without_email_message_params", - 2, - "noreply@example.com", - "Detections found", - ( - {}, - {}, - ), - ( - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + ] + }, + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + None, + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + None, + None, + ] + }, ), + ), + ( + _sliced_4_images, + SLICED_NAMES, + {"yolov8n-640": SLICED_DETECTION_COUNTS}, + True, + "sliced_image_without_email_message_params", + 3, + "noreply@example.com", + "Detections found", ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - False, - "with_email_gate_and_without_email_message_params", - 0, - "noreply@example.com", - "Detections found", - ({},), - ( - {"email_message": None}, - {"email_message": None}, - {"email_message": None}, - {"email_message": None}, - ), + {}, + {}, + {}, ), ( - _sliced_4_images, - SLICED_NAMES, - {"yolov8n-640": SLICED_DETECTION_COUNTS}, - True, - "with_detection_collapse_right_after_slice", # after the dim collapse we are dim=1 - 4, # In this scenario the continue_if step counts the number of slices for each image, so 4 calls to the email step - "noreply@example.com", - "Detections found", - ( - {"num_slices": 4}, - {"num_slices": 4}, - {"num_slices": 4}, - {"num_slices": 8}, - ), - ( # In this scenario the email step is called 4 times, once for each image, as each image has at least one slice - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + ] + }, + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + None, + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + None, + None, + ] + }, ), + ), + ( + _sliced_4_images, + SLICED_NAMES, + {"yolov8n-640": SLICED_DETECTION_COUNTS}, + True, + "sliced_image_with_email_message_params_and_area_size_step", + 3, + "noreply@example.com", + "Detections found", ( - _sliced_4_images, - SLICED_NAMES, - {"yolov8n-640": SLICED_DETECTION_COUNTS}, - True, - "with_detection_collapse_right_after_slice_with_agg_operation", - 2, # The continue-if correctly checks the number of detections for each image (given slices of that image) - "noreply@example.com", - "Detections found", - ( # The operations of counting the detection are done after receiving the params, so here we get the slices - {"num_slices": 4}, - {"num_slices": 8}, - ), - ( - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {"num_detections": 1, "area_converted": 2500}, + {"num_detections": 1, "area_converted": 2500}, + {"num_detections": 1, "area_converted": 2500}, ), ( - _sliced_4_images, - SLICED_NAMES, - {"yolov8n-640": SLICED_DETECTION_COUNTS}, - True, - "with_detection_collapse_right_after_slice_with_agg_operation_without_message_params", - 2, # The continue-if correctly checks the number of detections for each image (given slices of that image) - "noreply@example.com", - "Detections found", - ( - {}, - {}, - ), - ( - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + ] + }, + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + None, + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + None, + None, + ] + }, ), + ), + ( + _sliced_4_images, + SLICED_NAMES, + {"yolov8n-640": SLICED_DETECTION_COUNTS}, + True, + "sliced_image_without_email_message_params_and_area_size_step", + 3, + "noreply@example.com", + "Detections found", ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - True, - "with_detection_collapse_right_after_detect_with_agg_operation", - 1, # The continue-if correctly checks the total number of detections in the batch - "noreply@example.com", - "Detections found", - ( # The operations of counting the detection are done after receiving the params, so here we get the size of the batch - {"num_batch_detections": 4}, - ), - ({"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK},), + {"area_converted": 2500}, + {"area_converted": 2500}, + {"area_converted": 2500}, ), ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - True, - "with_detection_collapse_right_after_detect_with_agg_operation_without_message_params", - 1, # The continue-if correctly checks the total number of detections in the batch - "noreply@example.com", - "Detections found", - ({},), - ({"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK},), + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + ] + }, + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + None, + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + None, + None, + ] + }, ), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + True, + "with_email_gate_and_with_email_message_params", + 2, + "noreply@example.com", + "Detections found", ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - True, - "with_detection_collapse_after_continue_if", - 1, # We aggregated the detection lists after continue_if - "noreply@example.com", - "Detections found", - ({"num_batch_filtered_detections": 2},), # Only two images had detections - ({"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK},), + {"num_detections": 2}, + {"num_detections": 1}, ), ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": [3, 0, 1, 2]}, - True, - "with_two_continue_if", - 1, - "noreply@example.com", - "Detections found", - ({},), - ( - {"email_message": None}, - {"email_message": None}, - {"email_message": None}, - {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, - ), + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, ), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + True, + "with_email_gate_and_without_email_message_params", + 2, + "noreply@example.com", + "Detections found", ( - _sliced_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": SLICED_DETECTION_COUNTS}, - True, - "with_two_continue_if_different_control_flow_lineage", - 3, # The deepest control-flow-lineage is used, thus we are masking up to the slice level - "noreply@example.com", - "Detections found", - ( - {}, - {}, - {}, - ), - ( - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - ] - }, - {"email_message": [None, None, None, None]}, - { - "email_message": [ - None, - None, - None, - None, - None, - SUCCESSFUL_EMAIL_MESSAGE_MOCK, - None, - None, - ] - }, - ), + {}, + {}, ), - ], - ids=[ - "with_email_message_params", - "without_email_message_params", - "with_image_names_and_email_message_params", - "without_image_names_and_email_message_params", - "sliced_image_with_email_message_params", - "sliced_image_without_email_message_params", - "sliced_image_with_email_message_params_and_area_size_step", - "sliced_image_without_email_message_params_and_area_size_step", - "with_email_gate_and_with_email_message_params", + ( + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + ), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + False, "with_email_gate_and_without_email_message_params", - "with_email_gate_and_without_email_message_params_and_email_disabled", - "with_detection_collapse_right_after_slice", + 0, + "noreply@example.com", + "Detections found", + ({},), + ( + {"email_message": None}, + {"email_message": None}, + {"email_message": None}, + {"email_message": None}, + ), + ), + ( + _sliced_4_images, + SLICED_NAMES, + {"yolov8n-640": SLICED_DETECTION_COUNTS}, + True, + "with_detection_collapse_right_after_slice", # after the dim collapse we are dim=1 + 4, # In this scenario the continue_if step counts the number of slices for each image, so 4 calls to the email step + "noreply@example.com", + "Detections found", + ( + {"num_slices": 4}, + {"num_slices": 4}, + {"num_slices": 4}, + {"num_slices": 8}, + ), + ( # In this scenario the email step is called 4 times, once for each image, as each image has at least one slice + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + ), + ), + ( + _sliced_4_images, + SLICED_NAMES, + {"yolov8n-640": SLICED_DETECTION_COUNTS}, + True, "with_detection_collapse_right_after_slice_with_agg_operation", + 2, # The continue-if correctly checks the number of detections for each image (given slices of that image) + "noreply@example.com", + "Detections found", + ( # The operations of counting the detection are done after receiving the params, so here we get the slices + {"num_slices": 4}, + {"num_slices": 8}, + ), + ( + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + ), + ), + ( + _sliced_4_images, + SLICED_NAMES, + {"yolov8n-640": SLICED_DETECTION_COUNTS}, + True, "with_detection_collapse_right_after_slice_with_agg_operation_without_message_params", + 2, # The continue-if correctly checks the number of detections for each image (given slices of that image) + "noreply@example.com", + "Detections found", + ( + {}, + {}, + ), + ( + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + ), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + True, "with_detection_collapse_right_after_detect_with_agg_operation", + 1, # The continue-if correctly checks the total number of detections in the batch + "noreply@example.com", + "Detections found", + ( # The operations of counting the detection are done after receiving the params, so here we get the size of the batch + {"num_batch_detections": 4}, + ), + ({"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK},), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + True, "with_detection_collapse_right_after_detect_with_agg_operation_without_message_params", + 1, # The continue-if correctly checks the total number of detections in the batch + "noreply@example.com", + "Detections found", + ({},), + ({"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK},), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, + True, "with_detection_collapse_after_continue_if", + 1, # We aggregated the detection lists after continue_if + "noreply@example.com", + "Detections found", + ({"num_batch_filtered_detections": 2},), # Only two images had detections + ({"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK},), + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": [3, 0, 1, 2]}, + True, "with_two_continue_if", + 1, + "noreply@example.com", + "Detections found", + ({},), + ( + {"email_message": None}, + {"email_message": None}, + {"email_message": None}, + {"email_message": SUCCESSFUL_EMAIL_MESSAGE_MOCK}, + ), + ), + ( + _sliced_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": SLICED_DETECTION_COUNTS}, + True, "with_two_continue_if_different_control_flow_lineage", - ], + 3, # The deepest control-flow-lineage is used, thus we are masking up to the slice level + "noreply@example.com", + "Detections found", + ( + {}, + {}, + {}, + ), + ( + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + ] + }, + {"email_message": [None, None, None, None]}, + { + "email_message": [ + None, + None, + None, + None, + None, + SUCCESSFUL_EMAIL_MESSAGE_MOCK, + None, + None, + ] + }, + ), + ), +] + +_SIDE_EFFECT_IDS = [ + "with_email_message_params", + "without_email_message_params", + "with_image_names_and_email_message_params", + "without_image_names_and_email_message_params", + "sliced_image_with_email_message_params", + "sliced_image_without_email_message_params", + "sliced_image_with_email_message_params_and_area_size_step", + "sliced_image_without_email_message_params_and_area_size_step", + "with_email_gate_and_with_email_message_params", + "with_email_gate_and_without_email_message_params", + "with_email_gate_and_without_email_message_params_and_email_disabled", + "with_detection_collapse_right_after_slice", + "with_detection_collapse_right_after_slice_with_agg_operation", + "with_detection_collapse_right_after_slice_with_agg_operation_without_message_params", + "with_detection_collapse_right_after_detect_with_agg_operation", + "with_detection_collapse_right_after_detect_with_agg_operation_without_message_params", + "with_detection_collapse_after_continue_if", + "with_two_continue_if", + "with_two_continue_if_different_control_flow_lineage", +] + + +@patch( + "inference.core.workflows.core_steps.sinks.email_notification.v2.send_email_via_roboflow_proxy" +) +@pytest.mark.parametrize( + _SIDE_EFFECT_ARGNAMES, _SIDE_EFFECT_SCENARIOS, ids=_SIDE_EFFECT_IDS ) +@_NUMPY_ONLY def test_properly_running_side_effect_step_and_returning_results_in_different_data_lineage_control_lineage_scenarios( send_email_mock, image_gen_fn: callable, @@ -783,6 +887,86 @@ def test_properly_running_side_effect_step_and_returning_results_in_different_da assert result.get("email_message") == expected_result[i].get("email_message") +@patch( + "inference.core.workflows.core_steps.sinks.email_notification.v2.send_email_via_roboflow_proxy" +) +@pytest.mark.parametrize( + _SIDE_EFFECT_ARGNAMES, _SIDE_EFFECT_SCENARIOS, ids=_SIDE_EFFECT_IDS +) +@_TENSOR_ONLY +def test_properly_running_side_effect_step_and_returning_results_in_different_data_lineage_control_lineage_scenarios_tensor_native( + send_email_mock, + image_gen_fn: callable, + names: List[str], + detection_counts_per_model: Dict[str, List[int]], + enable_email: bool, + workflow_name: str, + expected_call_count: int, + expected_receiver_email: str, + expected_subject: str, + expected_message_parameters: tuple, + expected_result: tuple, + model_manager, +) -> None: + send_email_mock.return_value = (False, "Notification sent successfully") + workflow_definition = _load_workflow_definition(workflow_name) + + runtime_parameters = {"image": image_gen_fn()} + inputs = {inp["name"] for inp in workflow_definition.get("inputs", [])} + if "names" in inputs: + runtime_parameters["names"] = names + if "enable_email" in inputs: + runtime_parameters["enable_email"] = enable_email + + result = _run_workflow_tensor_native( + workflow_definition, + runtime_parameters, + model_manager, + detection_counts_per_model=detection_counts_per_model, + ) + + assert send_email_mock.call_count == expected_call_count + for i, call in enumerate(send_email_mock.call_args_list): + assert call.kwargs["receiver_email"] == [expected_receiver_email] + assert call.kwargs["subject"] == expected_subject + + expected_params = expected_message_parameters[i] + assert len(call.kwargs["message_parameters"]) == len(expected_params) + + for param_name, param_value in expected_params.items(): + actual = call.kwargs["message_parameters"][param_name] + + if param_name == "num_detections": + # native parity: intermediate detections are inference_models.Detections + assert isinstance(actual, NativeDetections) + assert len(actual) == param_value + continue + + if param_name in [ + "num_slices", + "num_batch_detections", + "num_batch_filtered_detections", + ]: + assert isinstance(actual, list) + assert len(actual) == param_value + continue + + if param_name == "area_converted": + # native parity: area lands per-box in bboxes_metadata, not in a + # subscriptable sv.Detections `.data` field + assert isinstance(actual, NativeDetections) + assert [m["area_converted"] for m in actual.bboxes_metadata] == [ + param_value + ] * len(actual) + continue + + assert actual == param_value + + assert len(result) == len(expected_result) + for i, result in enumerate(result): + assert result.get("email_message") == expected_result[i].get("email_message") + + @patch( "inference.core.workflows.core_steps.sinks.email_notification.v2.send_email_via_roboflow_proxy" ) @@ -1007,58 +1191,59 @@ def test_control_flow_lineage_using_workflow_with_batch_only_block_that_gets_bat assert result[i]["result"] == expect_result[i] -@pytest.mark.parametrize( - "image_gen_fn,\ - names,\ - detection_counts_per_model,\ - workflow_name,\ - expected_results,\ - expected_num_files,\ - expected_columns,\ - expected_num_rows,\ - expected_names,\ - expected_num_detections", - [ - ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - "with_csv_sink_and_with_detection_input", - [ - {"save_message": None}, - {"save_message": None}, - {"save_message": None}, - {"save_message": "Data saved successfully"}, - ], - 1, - ["num_detections", "name", "timestamp"], - 2, - ["img1", "img3"], - [2, 1], - ), - ( - _batch_4_images, - BATCH_4_IMAGE_NAMES, - {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, - "with_csv_sink_and_without_detection_input", - [ - {"save_message": None}, - {"save_message": None}, - {"save_message": None}, - {"save_message": "Data saved successfully"}, - ], - 1, - ["name", "timestamp"], - 2, - ["img1", "img3"], - [2, 1], - ), - ], - ids=[ +_CSV_SINK_ARGNAMES = ( + "image_gen_fn, names, detection_counts_per_model, workflow_name, expected_results, " + "expected_num_files, expected_columns, expected_num_rows, expected_names, " + "expected_num_detections" +) + +# Shared scenario table for the csv-sink control-flow test, consumed by both the +# @_NUMPY_ONLY and @_TENSOR_ONLY variants. +_CSV_SINK_SCENARIOS = [ + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, "with_csv_sink_and_with_detection_input", + [ + {"save_message": None}, + {"save_message": None}, + {"save_message": None}, + {"save_message": "Data saved successfully"}, + ], + 1, + ["num_detections", "name", "timestamp"], + 2, + ["img1", "img3"], + [2, 1], + ), + ( + _batch_4_images, + BATCH_4_IMAGE_NAMES, + {"yolov8n-640": BATCH_4_DETECTION_COUNTS}, "with_csv_sink_and_without_detection_input", - ], -) + [ + {"save_message": None}, + {"save_message": None}, + {"save_message": None}, + {"save_message": "Data saved successfully"}, + ], + 1, + ["name", "timestamp"], + 2, + ["img1", "img3"], + [2, 1], + ), +] + +_CSV_SINK_IDS = [ + "with_csv_sink_and_with_detection_input", + "with_csv_sink_and_without_detection_input", +] + + +@pytest.mark.parametrize(_CSV_SINK_ARGNAMES, _CSV_SINK_SCENARIOS, ids=_CSV_SINK_IDS) +@_NUMPY_ONLY def test_control_flow_lineage_using_workflow_with_csv_sink_and_detection_input( image_gen_fn: callable, names: List[str], @@ -1105,6 +1290,54 @@ def test_control_flow_lineage_using_workflow_with_csv_sink_and_detection_input( assert df["name"].tolist() == expected_names +@pytest.mark.parametrize(_CSV_SINK_ARGNAMES, _CSV_SINK_SCENARIOS, ids=_CSV_SINK_IDS) +@_TENSOR_ONLY +def test_control_flow_lineage_using_workflow_with_csv_sink_and_detection_input_tensor_native( + image_gen_fn: callable, + names: List[str], + detection_counts_per_model: Dict[str, List[int]], + workflow_name: str, + expected_results: List[dict], + expected_num_files: int, + expected_columns: List[str], + expected_num_rows: int, + expected_names: List[str], + expected_num_detections: List[int], + model_manager, + empty_directory, +) -> None: + workflow_definition = _load_workflow_definition(workflow_name) + + runtime_parameters = { + "image": image_gen_fn(), + "output_directory": empty_directory, + } + + inputs = {inp["name"] for inp in workflow_definition.get("inputs", [])} + if "names" in inputs: + runtime_parameters["names"] = names + + result = _run_workflow_tensor_native( + workflow_definition, + runtime_parameters, + model_manager, + detection_counts_per_model=detection_counts_per_model, + ) + assert result == expected_results + + csv_files = glob(os.path.join(empty_directory, "detection_log_*.csv")) + assert len(csv_files) == expected_num_files + df = pd.read_csv(csv_files[0]) + + assert set(df.columns) == set(expected_columns) + assert len(df) == expected_num_rows + + if "num_detections" in expected_columns: + assert df["num_detections"].tolist() == expected_num_detections + if "name" in expected_columns: + assert df["name"].tolist() == expected_names + + @pytest.mark.parametrize( "image_gen_fn,\ detection_counts_per_model,\ @@ -1132,6 +1365,7 @@ def test_control_flow_lineage_using_workflow_with_csv_sink_and_detection_input( "with_two_continue_if_data_lineage_present", ], ) +@_NUMPY_ONLY def test_side_effect_step_with_data_lineage_and_continue_if_zero_calls( image_gen_fn: callable, detection_counts_per_model: Dict[str, List[int]], @@ -1174,3 +1408,77 @@ def counting_run(self, *args, **kwargs): len(item["stitched_predictions"]) == expected_results[i]["stitched_predictions"][1] ) + + +@pytest.mark.parametrize( + "image_gen_fn,\ + detection_counts_per_model,\ + workflow_name,\ + expected_call_count, \ + expected_results", + [ + ( + _sliced_4_images, + { + "yolov8n-640": SLICED_DETECTION_COUNTS, + "yolov8s-640": BATCH_4_DETECTION_COUNTS, + }, + "with_two_continue_if_data_lineage_present", + 2, + [ + {"stitched_predictions": (type(None),)}, + {"stitched_predictions": (NativeDetections, 2)}, + {"stitched_predictions": (type(None),)}, + {"stitched_predictions": (NativeDetections, 1)}, + ], + ), + ], + ids=[ + "with_two_continue_if_data_lineage_present", + ], +) +@_TENSOR_ONLY +def test_side_effect_step_with_data_lineage_and_continue_if_zero_calls_tensor_native( + image_gen_fn: callable, + detection_counts_per_model: Dict[str, List[int]], + workflow_name: str, + expected_call_count: int, + expected_results: List[dict], + model_manager, +) -> None: + workflow_definition = _load_workflow_definition(workflow_name) + runtime_parameters = {"image": image_gen_fn()} + + stitch_run_call_count = [] + # Under the flag the workflow loads the tensor-native stitch block, so the + # call-count probe must patch that class, not the numpy one. + real_run = DetectionsStitchBlockV1Tensor.run + + def counting_run(self, *args, **kwargs): + stitch_run_call_count.append(1) + return real_run(self, *args, **kwargs) + + with patch.object(DetectionsStitchBlockV1Tensor, "run", counting_run): + result = _run_workflow_tensor_native( + workflow_definition, + runtime_parameters, + model_manager, + detection_counts_per_model=detection_counts_per_model, + ) + + assert len(stitch_run_call_count) == expected_call_count, ( + f"DetectionsStitchBlockV1.run should be called {expected_call_count} times, " + f"was {len(stitch_run_call_count)}" + ) + assert len(result) == len(expected_results) + for i, item in enumerate(result): + assert isinstance( + item["stitched_predictions"], + expected_results[i]["stitched_predictions"][0], + ) + + if isinstance(item["stitched_predictions"], NativeDetections): + assert ( + len(item["stitched_predictions"]) + == expected_results[i]["stitched_predictions"][1] + ) diff --git a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_after_continue_if_with_crop_batch_lineage.py b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_after_continue_if_with_crop_batch_lineage.py index 161d4aaec4..0c66ef33fe 100644 --- a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_after_continue_if_with_crop_batch_lineage.py +++ b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_after_continue_if_with_crop_batch_lineage.py @@ -5,6 +5,8 @@ from unittest import mock import numpy as np +import pytest +import torch from inference.core.entities.requests.inference import ( ClassificationInferenceRequest, @@ -17,12 +19,35 @@ ObjectDetectionInferenceResponse, ObjectDetectionPrediction, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.managers.base import ModelManager +from inference_models.models.base.classification import ( + ClassificationPrediction as NativeClassificationPrediction, +) +from inference_models.models.base.object_detection import Detections as NativeDetections from tests.workflows.integration_tests.execution.inner_workflow_inlining._common import ( echo_child_workflow, execution_engine, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION both the OD block and the classification +# block run inference through `run_tensor_native_inference` (+ `get_class_names`), +# not `infer_from_request_sync`. The OD block emits native `inference_models.Detections` +# and the classification block emits native `inference_models.ClassificationPrediction`. +# The numpy test skips when the flag is on; the `*_tensor_native` parity test (skipped +# when off) drives the same scenario natively. The continue_if gate, the crop-batch +# lineage and the dimension-collapse echo output are representation-independent, so the +# final `["passed", None]` is asserted identically in both variants. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="numpy ModelManager mock; OD + classification blocks are native under the " + "flag โ€” see *_tensor_native", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + _CLASSIFICATION_REQUEST_CONFIDENCE_THRESHOLD = 0.2 _CONTINUE_IF_CONFIDENCE_THRESHOLD = 0.5 @@ -258,6 +283,7 @@ def _flat_workflow() -> dict: } +@_NUMPY_ONLY def test_inlined_continue_if_echo_matches_inner_workflow( model_manager: ModelManager, dogs_image: np.ndarray, @@ -307,3 +333,136 @@ def test_inlined_continue_if_echo_matches_inner_workflow( assert nested_result == flat_result assert len(nested_result) == 1 assert nested_result[0]["from_child"] == ["passed", None] + + +# Class names for the native path. The OD `class_filter=["dog"]` is applied natively by +# the block against get_class_names("yolov8n-640"), so "dog" must sit at the OD +# class_id (0); the classification class_id-to-name map needs an entry at index 1 too. +_OD_CLASS_NAMES = ["dog"] +_CLS_CLASS_NAMES = ["dog", "cat"] + + +def _run_tensor_native_inference_factory(h: int, w: int): + # Native equivalent of _infer_from_request_sync_factory, dispatching on model_id + # because under the flag BOTH the OD block and the classification block call + # `run_tensor_native_inference`. + # + # OD (yolov8n-640): list of ONE raw `Detections` for the single input image, two + # "dog" boxes (xyxy from the numpy x/y/w/h centres: (80,80,100,100) -> + # [30,30,130,130]; (300,80,100,100) -> [250,30,350,130]). detection_id rides in + # bboxes_metadata; the OD block attaches the image lineage + class names itself. + # classification (dog-breed/1): ONE batched `ClassificationPrediction` over the two + # crops. class_id is (bs,), confidence is the FULL (bs, num_classes) softmax. The + # continue_if gate reads top_class_confidence = confidence[i, class_id[i]]: + # crop 0 -> class_id 0, confidence[0,0]=0.9 >= 0.5 -> PASS -> echo "passed" + # crop 1 -> class_id 1, confidence[1,1]=0.4 < 0.5 -> GATED -> None + # (mirrors the numpy responses top=dog @0.9 and @0.4). The block fans this batched + # object out per-crop and attaches images_metadata itself. + def run_tensor_native_inference(model_id: str, images, **kwargs): + imgs = images if isinstance(images, list) else [images] + if model_id == "yolov8n-640": + assert len(imgs) == 1, f"Mock Expected 1 OD image, got {len(imgs)}" + return [ + NativeDetections( + xyxy=torch.tensor( + [[30, 30, 130, 130], [250, 30, 350, 130]], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 0], dtype=torch.long), + confidence=torch.tensor([0.99, 0.99], dtype=torch.float32), + bboxes_metadata=[ + {"detection_id": "mock-d0"}, + {"detection_id": "mock-d1"}, + ], + ) + ] + if model_id == "dog-breed/1": + assert len(imgs) == 2, f"Mock Expected 2 crop images, got {len(imgs)}" + assert ( + kwargs.get("confidence") == _CLASSIFICATION_REQUEST_CONFIDENCE_THRESHOLD + ) + return NativeClassificationPrediction( + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([[0.9, 0.1], [0.6, 0.4]], dtype=torch.float32), + ) + raise AssertionError(f"Unexpected model_id: {model_id!r}") + + return run_tensor_native_inference + + +def _get_class_names_factory(): + def get_class_names(model_id: str): + if model_id == "yolov8n-640": + return list(_OD_CLASS_NAMES) + if model_id == "dog-breed/1": + return list(_CLS_CLASS_NAMES) + raise AssertionError(f"Unexpected model_id: {model_id!r}") + + return get_class_names + + +@_TENSOR_ONLY +def test_inlined_continue_if_echo_matches_inner_workflow_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + h, w = dogs_image.shape[:2] + run_mock = mock.MagicMock( + side_effect=_run_tensor_native_inference_factory(h, w), + ) + inner = echo_child_workflow() + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, + "get_class_names", + new=mock.MagicMock(side_effect=_get_class_names_factory()), + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine(model_manager, _nested_workflow(inner)) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run( + runtime_parameters={ + "image": dogs_image, + "crop_label": "passed", + }, + ) + flat_result = flat_engine.run( + runtime_parameters={ + "image": dogs_image, + "crop_label": "passed", + }, + ) + + # Same dispatch count as the numpy mock (4): OD + classification per engine run. + # model_id/images may be positional OR kwarg โ€” read tolerantly. + assert run_mock.call_count == 4 + od_calls = [] + cls_calls = [] + for call in run_mock.call_args_list: + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + images = call.kwargs.get("images", call.args[1] if len(call.args) > 1 else None) + assert images is not None + if model_id == "yolov8n-640": + assert len(images) == 1 + od_calls.append(call) + elif model_id == "dog-breed/1": + assert len(images) == 2 + cls_calls.append(call) + else: + raise AssertionError(f"Unexpected model_id: {model_id!r}") + assert len(od_calls) == 2 + assert len(cls_calls) == 2 + + # The continue_if gate, crop-batch lineage and dimension-collapse echo are + # representation-independent: crop 0 passes (top conf 0.9 >= 0.5) -> "passed", + # crop 1 is gated (top conf 0.4 < 0.5) -> None. Native Detections / + # ClassificationPrediction never surface in `from_child` (the echo emits scalars), + # so the output is identical to the numpy variant. Assert it on BOTH runs rather + # than comparing nested_result == flat_result (native carriers have no value __eq__). + assert len(nested_result) == 1 + assert len(flat_result) == 1 + assert nested_result[0]["from_child"] == ["passed", None] + assert flat_result[0]["from_child"] == ["passed", None] diff --git a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_child_runs_dimension_collapse_on_parent_detections.py b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_child_runs_dimension_collapse_on_parent_detections.py index 4ad28b2c6a..e565be3fe4 100644 --- a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_child_runs_dimension_collapse_on_parent_detections.py +++ b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_child_runs_dimension_collapse_on_parent_detections.py @@ -13,7 +13,9 @@ from unittest import mock import numpy as np +import pytest import supervision as sv +import torch from inference.core.entities.requests.inference import ObjectDetectionInferenceRequest from inference.core.entities.responses.inference import ( @@ -21,12 +23,28 @@ ObjectDetectionInferenceResponse, ObjectDetectionPrediction, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.managers.base import ModelManager +from inference_models.models.base.object_detection import Detections as NativeDetections from tests.workflows.integration_tests.execution.inner_workflow_inlining._common import ( child_dimension_collapse_from_parent_detections, execution_engine, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the OD block runs inference through +# `run_tensor_native_inference` (+ `get_class_names`), not `infer_from_request_sync`, +# and the collapsed predictions are native `inference_models.Detections` (no +# `__getitem__` / value `__eq__`). The numpy test skips when the flag is on; the +# `*_tensor_native` parity test (skipped when off) drives the same scenario natively. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="numpy ModelManager mock; OD block is native under the flag โ€” see *_tensor_native", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + def _mock_od_response(h: int, w: int) -> ObjectDetectionInferenceResponse: return ObjectDetectionInferenceResponse( @@ -151,6 +169,7 @@ def _flat_workflow() -> dict: } +@_NUMPY_ONLY def test_inlined_child_dimension_collapse_matches_flat_workflow( model_manager: ModelManager, dogs_image: np.ndarray, @@ -186,6 +205,7 @@ def test_inlined_child_dimension_collapse_matches_flat_workflow( assert len(nested_result) == 1 +@_NUMPY_ONLY def test_inlined_child_dimension_collapse_matches_flat_workflow_runtime_image_list( model_manager: ModelManager, dogs_image: np.ndarray, @@ -221,3 +241,137 @@ def test_inlined_child_dimension_collapse_matches_flat_workflow_runtime_image_li assert nested_result == flat_result assert len(nested_result) == 1 _assert_collapsed_two_image_batch(nested_result[0]["collapsed"]) + + +# --- tensor-native parity -------------------------------------------------- + + +def _one_native_detections() -> NativeDetections: + # Native equivalent of _mock_od_response: same three boxes as raw `Detections` + # (xyxy from the numpy x/y/w/h centres). The OD block attaches the image lineage; + # detection_id rides in bboxes_metadata, class names via get_class_names. + return NativeDetections( + xyxy=torch.tensor( + [[30, 30, 90, 90], [210, 20, 290, 100], [400, 10, 500, 110]], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 1, 2], dtype=torch.long), + confidence=torch.tensor([0.99, 0.95, 0.90], dtype=torch.float32), + bboxes_metadata=[ + {"detection_id": "mock-d0"}, + {"detection_id": "mock-d1"}, + {"detection_id": "mock-d2"}, + ], + ) + + +def _run_tensor_native_inference(model_id: str, images, **kwargs): + # Native equivalent of _mock_infer_object_detection_from_request_sync: one raw + # `Detections` per input image (the OD block zips images with this list). A fresh + # object per image keeps the two-image batch from aliasing one tensor. + imgs = images if isinstance(images, list) else [images] + return [_one_native_detections() for _ in imgs] + + +_EXPECTED_BOXES = [ + ("mock-d0", "x", 0, 0.99, [30, 30, 90, 90]), + ("mock-d1", "y", 1, 0.95, [210, 20, 290, 100]), + ("mock-d2", "z", 2, 0.90, [400, 10, 500, 110]), +] + + +def _assert_native_three_box_detections(det: NativeDetections) -> None: + # Native parity of the per-detection value checks: Detections has no __getitem__, + # so detection_id reads from bboxes_metadata and class_name from image_metadata. + assert isinstance(det, NativeDetections) + assert len(det) == 3 + for i, (det_id, class_name, class_id, conf, expected_xyxy) in enumerate( + _EXPECTED_BOXES + ): + assert det.bboxes_metadata[i]["detection_id"] == det_id + assert int(det.class_id[i]) == class_id + assert float(det.confidence[i]) == pytest.approx(conf) + assert det.image_metadata["class_names"][int(det.class_id[i])] == class_name + np.testing.assert_allclose( + det.xyxy[i].cpu().numpy(), expected_xyxy, rtol=0, atol=1e-3 + ) + + +def _assert_collapsed_native(collapsed: list, expected_images: int) -> None: + # Native parity of _assert_collapsed_two_image_batch (generalised to N images): + # the collapsed list holds one native Detections per image, each with 3 boxes. + assert isinstance(collapsed, list) + assert len(collapsed) == expected_images + for det in collapsed: + _assert_native_three_box_detections(det) + + +@_TENSOR_ONLY +def test_inlined_child_dimension_collapse_matches_flat_workflow_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + run_mock = mock.MagicMock(side_effect=_run_tensor_native_inference) + inner = child_dimension_collapse_from_parent_detections() + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", return_value=["x", "y", "z"] + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine(model_manager, _nested_workflow(inner)) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run(runtime_parameters={"image": dogs_image}) + flat_result = flat_engine.run(runtime_parameters={"image": dogs_image}) + + assert run_mock.call_count == 2 + for call in run_mock.call_args_list: + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + assert model_id == "yolov8n-640" + imgs = call.kwargs.get("images", call.args[1] if len(call.args) > 1 else None) + assert imgs is not None and len(imgs) == 1 + + # Native Detections has no value __eq__; assert the values on BOTH runs instead of + # comparing nested_result == flat_result. + assert len(nested_result) == 1 + assert len(flat_result) == 1 + _assert_collapsed_native(nested_result[0]["collapsed"], expected_images=1) + _assert_collapsed_native(flat_result[0]["collapsed"], expected_images=1) + + +@_TENSOR_ONLY +def test_inlined_child_dimension_collapse_matches_flat_workflow_runtime_image_list_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + run_mock = mock.MagicMock(side_effect=_run_tensor_native_inference) + inner = child_dimension_collapse_from_parent_detections() + images = [dogs_image, dogs_image.copy()] + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", return_value=["x", "y", "z"] + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine(model_manager, _nested_workflow(inner)) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run(runtime_parameters={"image": images}) + flat_result = flat_engine.run(runtime_parameters={"image": images}) + + assert run_mock.call_count == 2 + for call in run_mock.call_args_list: + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + assert model_id == "yolov8n-640" + imgs = call.kwargs.get("images", call.args[1] if len(call.args) > 1 else None) + assert imgs is not None and len(imgs) == 2 + + # Native Detections has no value __eq__; assert the values on BOTH runs instead of + # comparing nested_result == flat_result. + assert len(nested_result) == 1 + assert len(flat_result) == 1 + _assert_collapsed_native(nested_result[0]["collapsed"], expected_images=2) + _assert_collapsed_native(flat_result[0]["collapsed"], expected_images=2) diff --git a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_child_runs_dynamic_crop_on_parent_detections.py b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_child_runs_dynamic_crop_on_parent_detections.py index f7c5215758..cc15f04205 100644 --- a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_child_runs_dynamic_crop_on_parent_detections.py +++ b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_child_runs_dynamic_crop_on_parent_detections.py @@ -5,7 +5,9 @@ from unittest import mock import numpy as np +import pytest import supervision as sv +import torch from inference.core.entities.requests.inference import ObjectDetectionInferenceRequest from inference.core.entities.responses.inference import ( @@ -13,12 +15,28 @@ ObjectDetectionInferenceResponse, ObjectDetectionPrediction, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.managers.base import ModelManager +from inference_models.models.base.object_detection import Detections as NativeDetections from tests.workflows.integration_tests.execution.inner_workflow_inlining._common import ( child_dynamic_crop_from_parent_detections, execution_engine, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the OD block runs inference through +# `run_tensor_native_inference` (+ `get_class_names`), not `infer_from_request_sync`, +# and downstream crop predictions are native `inference_models.Detections` (no +# `__getitem__` / value `__eq__`). The numpy test skips when the flag is on; the +# `*_tensor_native` parity test (skipped when off) drives the same scenario natively. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="numpy ModelManager mock; OD block is native under the flag โ€” see *_tensor_native", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + def _infer_from_request_sync_factory(h: int, w: int): def infer_from_request_sync( @@ -146,6 +164,7 @@ def _assert_crop_predictions_equal(crop_preds: list) -> None: np.testing.assert_allclose(det.xyxy[0], expected_xyxy, rtol=0, atol=1e-3) +@_NUMPY_ONLY def test_inlined_dynamic_crop_matches_inner_workflow( model_manager: ModelManager, dogs_image: np.ndarray, @@ -181,3 +200,91 @@ def test_inlined_dynamic_crop_matches_inner_workflow( assert nested_result == flat_result assert len(nested_result) == 1 _assert_crop_predictions_equal(nested_result[0]["from_child"]) + + +def _run_tensor_native_inference_factory(h: int, w: int): + # Native equivalent of _infer_from_request_sync_factory: same three boxes as raw + # `Detections` (xyxy from the numpy x/y/w/h centres). The OD block attaches the + # image lineage; detection_id rides in bboxes_metadata, class names via get_class_names. + def run_tensor_native_inference(model_id: str, images, **kwargs): + imgs = images if isinstance(images, list) else [images] + assert len(imgs) == 1, f"Mock Expected 1 image, got {len(imgs)}" + return [ + NativeDetections( + xyxy=torch.tensor( + [[30, 30, 90, 90], [210, 20, 290, 100], [400, 10, 500, 110]], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 1, 2], dtype=torch.long), + confidence=torch.tensor([0.99, 0.95, 0.90], dtype=torch.float32), + bboxes_metadata=[ + {"detection_id": "mock-d0"}, + {"detection_id": "mock-d1"}, + {"detection_id": "mock-d2"}, + ], + ) + ] + + return run_tensor_native_inference + + +def _assert_crop_predictions_equal_native(crop_preds: list) -> None: + # Native parity of _assert_crop_predictions_equal: Detections has no __getitem__, + # so detection_id reads from bboxes_metadata and class_name from image_metadata. + expected = [ + ("mock-d0", "x", 0, 0.99, [0, 0, 60, 60]), + ("mock-d1", "y", 1, 0.95, [0, 0, 80, 80]), + ("mock-d2", "z", 2, 0.90, [0, 0, 100, 100]), + ] + assert isinstance(crop_preds, list) + assert len(crop_preds) == 3 + for det, (det_id, class_name, class_id, conf, expected_xyxy) in zip( + crop_preds, expected + ): + assert isinstance(det, NativeDetections) + assert len(det) == 1 + assert det.bboxes_metadata[0]["detection_id"] == det_id + assert int(det.class_id[0]) == class_id + assert float(det.confidence[0]) == pytest.approx(conf) + assert det.image_metadata["class_names"][int(det.class_id[0])] == class_name + np.testing.assert_allclose( + det.xyxy[0].cpu().numpy(), expected_xyxy, rtol=0, atol=1e-3 + ) + + +@_TENSOR_ONLY +def test_inlined_dynamic_crop_matches_inner_workflow_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + h, w = dogs_image.shape[:2] + run_mock = mock.MagicMock( + side_effect=_run_tensor_native_inference_factory(h, w), + ) + inner = child_dynamic_crop_from_parent_detections() + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", return_value=["x", "y", "z"] + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine(model_manager, _nested_workflow(inner)) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run(runtime_parameters={"image": dogs_image}) + flat_result = flat_engine.run(runtime_parameters={"image": dogs_image}) + + assert run_mock.call_count == 2 + for call in run_mock.call_args_list: + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + assert model_id == "yolov8n-640" + imgs = call.kwargs.get("images", call.args[1] if len(call.args) > 1 else None) + assert imgs is not None and len(imgs) == 1 + + # Native Detections has no value __eq__; assert the values on BOTH runs instead of + # comparing nested_result == flat_result. + assert len(nested_result) == 1 + assert len(flat_result) == 1 + _assert_crop_predictions_equal_native(nested_result[0]["from_child"]) + _assert_crop_predictions_equal_native(flat_result[0]["from_child"]) diff --git a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_continue_if_inside_inner_with_crop_batch_lineage.py b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_continue_if_inside_inner_with_crop_batch_lineage.py index 5606e8a349..db568afde9 100644 --- a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_continue_if_inside_inner_with_crop_batch_lineage.py +++ b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_continue_if_inside_inner_with_crop_batch_lineage.py @@ -11,6 +11,8 @@ from unittest import mock import numpy as np +import pytest +import torch from inference.core.entities.requests.inference import ( ClassificationInferenceRequest, @@ -23,11 +25,35 @@ ObjectDetectionInferenceResponse, ObjectDetectionPrediction, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.managers.base import ModelManager +from inference_models.models.base.classification import ( + ClassificationPrediction as NativeClassificationPrediction, +) +from inference_models.models.base.object_detection import Detections as NativeDetections from tests.workflows.integration_tests.execution.inner_workflow_inlining._common import ( execution_engine, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION both the OD block AND the classification +# block run inference through `run_tensor_native_inference` (+ `get_class_names`), +# not `infer_from_request_sync`. The numpy tests skip when the flag is on; the +# `*_tensor_native` parity tests (skipped when off) drive the same scenario with a +# single `run_tensor_native_inference` mock that dispatches on `model_id` and a +# `get_class_names` that dispatches likewise. The workflow output (`from_child`) is +# a scalar echo list (`["passed", None]`) gated by the `continue_if` on the +# classification top-class confidence, so the value-level expectations are identical +# to numpy โ€” only the mock seams change. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="numpy ModelManager mock; OD/classification blocks are native under the " + "flag โ€” see *_tensor_native", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + _CLASSIFICATION_REQUEST_CONFIDENCE_THRESHOLD = 0.2 _CONTINUE_IF_CONFIDENCE_THRESHOLD = 0.5 @@ -403,6 +429,7 @@ def _flat_workflow() -> dict: } +@_NUMPY_ONLY def test_inlined_continue_if_inside_inner_matches_flat_workflow( model_manager: ModelManager, dogs_image: np.ndarray, @@ -454,6 +481,124 @@ def test_inlined_continue_if_inside_inner_matches_flat_workflow( assert nested_result[0]["from_child"] == ["passed", None] +def _run_tensor_native_inference_factory(h: int, w: int): + # Native equivalent of _infer_from_request_sync_factory. Under the flag BOTH the + # OD block and the classification block dispatch through `run_tensor_native_inference`, + # keyed by `model_id` (the numpy mock keyed on request type instead): + # * "yolov8n-640" -> List[Detections] (one per input image). Two boxes, both + # class_id=0 -> "dog" via get_class_names, so both survive the native + # class_filter=["dog"]. xyxy from the numpy centre/size boxes + # (80,80,100,100) and (300,80,100,100). detection_id rides in bboxes_metadata. + # * "dog-breed/1" -> ONE batched ClassificationPrediction over the 2 crops: + # class_id (bs,) and confidence (bs, num_classes) full softmax. top_class_confidence + # reads confidence[0, class_id] per single-row slice, so row 0 -> 0.9 (>=0.5, passes) + # and row 1 -> 0.4 (<0.5, blocked), reproducing the numpy ["passed", None] gating. + def run_tensor_native_inference(model_id: str, images, **kwargs): + imgs = images if isinstance(images, list) else [images] + if model_id == "yolov8n-640": + assert len(imgs) == 1, f"Mock Expected 1 OD image, got {len(imgs)}" + return [ + NativeDetections( + xyxy=torch.tensor( + [[30, 30, 130, 130], [250, 30, 350, 130]], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 0], dtype=torch.long), + confidence=torch.tensor([0.99, 0.99], dtype=torch.float32), + bboxes_metadata=[ + {"detection_id": "mock-d0"}, + {"detection_id": "mock-d1"}, + ], + ) + ] + if model_id == "dog-breed/1": + assert len(imgs) == 2, f"Mock Expected 2 crop images, got {len(imgs)}" + return NativeClassificationPrediction( + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([[0.9, 0.1], [0.6, 0.4]], dtype=torch.float32), + ) + raise AssertionError(f"Unexpected model_id: {model_id!r}") + + return run_tensor_native_inference + + +def _get_class_names_native(model_id: str): + # Single get_class_names mock dispatching on model_id (mirrors the request-type + # dispatch in the numpy infer mock). OD class 0 -> "dog" so both boxes survive the + # native class_filter=["dog"]. The classification names don't affect the continue_if + # gate (it reads confidence[0, class_id]), but the list must cover the class ids. + if model_id == "yolov8n-640": + return ["dog"] + if model_id == "dog-breed/1": + return ["dog", "other"] + raise AssertionError(f"Unexpected model_id: {model_id!r}") + + +def _assert_native_call_order(run_mock: mock.MagicMock) -> None: + # Both engines run OD then classification, so the 4 calls alternate + # yolov8n-640, dog-breed/1. model_id/images may arrive positional OR kwarg. + assert run_mock.call_count == 4 + for idx in (0, 2): + call = run_mock.call_args_list[idx] + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + assert model_id == "yolov8n-640" + od_imgs = call.kwargs.get( + "images", call.args[1] if len(call.args) > 1 else None + ) + assert od_imgs is not None and len(od_imgs) == 1 + for idx in (1, 3): + call = run_mock.call_args_list[idx] + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + assert model_id == "dog-breed/1" + cls_imgs = call.kwargs.get( + "images", call.args[1] if len(call.args) > 1 else None + ) + assert cls_imgs is not None and len(cls_imgs) == 2 + + +@_TENSOR_ONLY +def test_inlined_continue_if_inside_inner_matches_flat_workflow_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + h, w = dogs_image.shape[:2] + run_mock = mock.MagicMock( + side_effect=_run_tensor_native_inference_factory(h, w), + ) + inner = _inner_continue_if_then_pick() + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", side_effect=_get_class_names_native + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine(model_manager, _nested_workflow(inner)) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run( + runtime_parameters={ + "image": dogs_image, + "crop_label": "passed", + }, + ) + flat_result = flat_engine.run( + runtime_parameters={ + "image": dogs_image, + "crop_label": "passed", + }, + ) + + _assert_native_call_order(run_mock) + + # `from_child` is a plain scalar echo list (no native Detections), so the value-level + # expectation is identical to numpy and safe to assert directly on both runs. + assert nested_result == flat_result + assert len(nested_result) == 1 + assert nested_result[0]["from_child"] == ["passed", None] + + +@_NUMPY_ONLY def test_inlined_continue_if_last_in_inner_echo_on_outer_matches_flat_workflow( model_manager: ModelManager, dogs_image: np.ndarray, @@ -506,3 +651,48 @@ def test_inlined_continue_if_last_in_inner_echo_on_outer_matches_flat_workflow( assert nested_result == flat_result assert len(nested_result) == 1 assert nested_result[0]["from_child"] == ["passed", None] + + +@_TENSOR_ONLY +def test_inlined_continue_if_last_in_inner_echo_on_outer_matches_flat_workflow_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + h, w = dogs_image.shape[:2] + run_mock = mock.MagicMock( + side_effect=_run_tensor_native_inference_factory(h, w), + ) + inner = _inner_continue_if_only_outer_echo_name_matches_parent() + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", side_effect=_get_class_names_native + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine( + model_manager, + _nested_workflow_continue_if_inner_echo_outer(inner), + ) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run( + runtime_parameters={ + "image": dogs_image, + "crop_label": "passed", + }, + ) + flat_result = flat_engine.run( + runtime_parameters={ + "image": dogs_image, + "crop_label": "passed", + }, + ) + + _assert_native_call_order(run_mock) + + # `from_child` is a plain scalar echo list (no native Detections), so the value-level + # expectation is identical to numpy and safe to assert directly on both runs. + assert nested_result == flat_result + assert len(nested_result) == 1 + assert nested_result[0]["from_child"] == ["passed", None] diff --git a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_detection_offset_after_nested_crop.py b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_detection_offset_after_nested_crop.py index f80d83cc93..7002bcb9b9 100644 --- a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_detection_offset_after_nested_crop.py +++ b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_detection_offset_after_nested_crop.py @@ -6,7 +6,9 @@ from unittest import mock import numpy as np +import pytest import supervision as sv +import torch from inference.core.entities.requests.inference import ObjectDetectionInferenceRequest from inference.core.entities.responses.inference import ( @@ -14,12 +16,28 @@ ObjectDetectionInferenceResponse, ObjectDetectionPrediction, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.managers.base import ModelManager +from inference_models.models.base.object_detection import Detections as NativeDetections from tests.workflows.integration_tests.execution.inner_workflow_inlining._common import ( child_dynamic_crop_from_parent_detections, execution_engine, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the OD block runs inference through +# `run_tensor_native_inference` (+ `get_class_names`), not `infer_from_request_sync`, +# and downstream crop/offset predictions are native `inference_models.Detections` (no +# `__getitem__` / value `__eq__`). The numpy test skips when the flag is on; the +# `*_tensor_native` parity test (skipped when off) drives the same scenario natively. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="numpy ModelManager mock; OD block is native under the flag โ€” see *_tensor_native", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + def _infer_from_request_sync_factory(h: int, w: int): def infer_from_request_sync( @@ -231,6 +249,7 @@ def _assert_offset_detection_lists_equal( ) +@_NUMPY_ONLY def test_inlined_crop_and_offset_match_nested_inner_workflow( model_manager: ModelManager, dogs_image: np.ndarray, @@ -282,3 +301,151 @@ def test_inlined_crop_and_offset_match_nested_inner_workflow( infer_mock.assert_called() assert infer_mock.call_count == 4 + + +def _run_tensor_native_inference_factory(h: int, w: int): + # Native equivalent of _infer_from_request_sync_factory: same three boxes as raw + # `Detections` (xyxy from the numpy x/y/w/h centres). The OD block attaches the + # image lineage; detection_id rides in bboxes_metadata, class names via get_class_names. + def run_tensor_native_inference(model_id: str, images, **kwargs): + imgs = images if isinstance(images, list) else [images] + assert len(imgs) == 1 + return [ + NativeDetections( + xyxy=torch.tensor( + [[30, 30, 90, 90], [210, 20, 290, 100], [400, 10, 500, 110]], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 1, 2], dtype=torch.long), + confidence=torch.tensor([0.99, 0.95, 0.90], dtype=torch.float32), + bboxes_metadata=[ + {"detection_id": "mock-d0"}, + {"detection_id": "mock-d1"}, + {"detection_id": "mock-d2"}, + ], + ) + ] + + return run_tensor_native_inference + + +def _native_class_name(det: NativeDetections, index: int) -> str: + return det.image_metadata["class_names"][int(det.class_id[index])] + + +def _assert_crop_predictions_equal_native(crop_preds: list) -> None: + # Native parity of the numpy crop-only comparison: each crop yields one detection at + # the crop's local origin. Detections has no __getitem__, so detection_id reads from + # bboxes_metadata and class_name from image_metadata. + expected = [ + ("mock-d0", "x", 0, 0.99, [0, 0, 60, 60]), + ("mock-d1", "y", 1, 0.95, [0, 0, 80, 80]), + ("mock-d2", "z", 2, 0.90, [0, 0, 100, 100]), + ] + assert isinstance(crop_preds, list) + assert len(crop_preds) == 3 + for det, (det_id, class_name, class_id, conf, expected_xyxy) in zip( + crop_preds, expected + ): + assert isinstance(det, NativeDetections) + assert len(det) == 1 + assert det.bboxes_metadata[0]["detection_id"] == det_id + assert int(det.class_id[0]) == class_id + assert float(det.confidence[0]) == pytest.approx(conf) + assert _native_class_name(det, 0) == class_name + np.testing.assert_allclose( + det.xyxy[0].cpu().numpy(), expected_xyxy, rtol=0, atol=1e-3 + ) + + +def _assert_offset_detection_lists_equal_native( + a_list: list, + b_list: list, +) -> None: + """Native parity of _assert_offset_detection_lists_equal. + + ``detection_offset`` issues new UUIDs, so geometry and class metadata are compared + nested-vs-flat only. Native ``Detections`` has no value ``__eq__`` / ``.data`` โ€” + xyxy is a torch tensor and class names resolve through ``image_metadata``. + """ + assert len(a_list) == len(b_list) + for det_a, det_b in zip(a_list, b_list): + assert isinstance(det_a, NativeDetections) and isinstance( + det_b, NativeDetections + ) + assert len(det_a) == len(det_b) + np.testing.assert_allclose( + det_a.xyxy.cpu().numpy().astype(np.float32), + det_b.xyxy.cpu().numpy().astype(np.float32), + rtol=0, + atol=1e-3, + ) + assert _native_class_name(det_a, 0) == _native_class_name(det_b, 0) + assert int(det_a.class_id[0]) == int(det_b.class_id[0]) + np.testing.assert_allclose( + float(det_a.confidence[0]), + float(det_b.confidence[0]), + rtol=0, + atol=1e-6, + ) + + +@_TENSOR_ONLY +def test_inlined_crop_and_offset_match_nested_inner_workflow_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + h, w = dogs_image.shape[:2] + run_mock = mock.MagicMock( + side_effect=_run_tensor_native_inference_factory(h, w), + ) + inner = child_dynamic_crop_from_parent_detections() + offset_x = 10 + offset_y = 10 + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", return_value=["x", "y", "z"] + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_crop = execution_engine( + model_manager, + _nested_parent_workflow_crop_only(inner), + ) + flat_crop = execution_engine( + model_manager, + _flat_parent_workflow_crop_only(), + ) + nested_off = execution_engine( + model_manager, + _nested_parent_workflow_with_detection_offset( + inner, + offset_x, + offset_y, + ), + ) + flat_off = execution_engine( + model_manager, + _flat_parent_workflow_with_detection_offset(offset_x, offset_y), + ) + + res_nested_crop = nested_crop.run(runtime_parameters={"image": dogs_image}) + res_flat_crop = flat_crop.run(runtime_parameters={"image": dogs_image}) + res_nested_off = nested_off.run(runtime_parameters={"image": dogs_image}) + res_flat_off = flat_off.run(runtime_parameters={"image": dogs_image}) + + # Native Detections has no value __eq__; assert the crop-only values on BOTH runs + # instead of comparing res_nested_crop == res_flat_crop. + assert len(res_nested_crop) == 1 + assert len(res_flat_crop) == 1 + _assert_crop_predictions_equal_native(res_nested_crop[0]["crop_preds"]) + _assert_crop_predictions_equal_native(res_flat_crop[0]["crop_preds"]) + + nested_off_preds = res_nested_off[0]["padded_preds"] + flat_off_preds = res_flat_off[0]["padded_preds"] + _assert_offset_detection_lists_equal_native(nested_off_preds, flat_off_preds) + + run_mock.assert_called() + assert run_mock.call_count == 4 diff --git a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_runs_dimension_collapse_on_inner_detections.py b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_runs_dimension_collapse_on_inner_detections.py index 32faf2c670..f98a5a2ebc 100644 --- a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_runs_dimension_collapse_on_inner_detections.py +++ b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_runs_dimension_collapse_on_inner_detections.py @@ -14,7 +14,9 @@ from unittest import mock import numpy as np +import pytest import supervision as sv +import torch from inference.core.entities.requests.inference import ObjectDetectionInferenceRequest from inference.core.entities.responses.inference import ( @@ -22,12 +24,28 @@ ObjectDetectionInferenceResponse, ObjectDetectionPrediction, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.managers.base import ModelManager +from inference_models.models.base.object_detection import Detections as NativeDetections from tests.workflows.integration_tests.execution.inner_workflow_inlining._common import ( child_detection_only_for_parent_dynamic_crop, execution_engine, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the OD block runs inference through +# `run_tensor_native_inference` (+ `get_class_names`), not `infer_from_request_sync`, +# and the collapsed predictions are native `inference_models.Detections` (no +# `__getitem__` / value `__eq__`). The numpy test skips when the flag is on; the +# `*_tensor_native` parity test (skipped when off) drives the same scenario natively. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="numpy ModelManager mock; OD block is native under the flag โ€” see *_tensor_native", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + def _mock_od_response(h: int, w: int) -> ObjectDetectionInferenceResponse: return ObjectDetectionInferenceResponse( @@ -152,6 +170,7 @@ def _flat_workflow() -> dict: } +@_NUMPY_ONLY def test_inlined_parent_dimension_collapse_matches_inner_workflow_detection( model_manager: ModelManager, dogs_image: np.ndarray, @@ -187,6 +206,7 @@ def test_inlined_parent_dimension_collapse_matches_inner_workflow_detection( assert len(nested_result) == 1 +@_NUMPY_ONLY def test_inlined_parent_dimension_collapse_matches_inner_workflow_detection_runtime_image_list( model_manager: ModelManager, dogs_image: np.ndarray, @@ -223,3 +243,137 @@ def test_inlined_parent_dimension_collapse_matches_inner_workflow_detection_runt assert nested_result == flat_result assert len(nested_result) == 1 _assert_collapsed_two_image_batch(nested_result[0]["collapsed"]) + + +# --- tensor-native parity ------------------------------------------------------------- + + +def _native_od_detections() -> NativeDetections: + """Native equivalent of ``_mock_od_response``: the same three full-frame boxes as a + raw ``Detections`` (xyxy from the numpy x/y/w/h centres). The OD block attaches the + image lineage; ``detection_id`` rides in ``bboxes_metadata``, class names via + ``get_class_names``. A fresh instance is built per image so the block-side metadata + attach never aliases tensors across the batch.""" + return NativeDetections( + xyxy=torch.tensor( + [[30, 30, 90, 90], [210, 20, 290, 100], [400, 10, 500, 110]], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 1, 2], dtype=torch.long), + confidence=torch.tensor([0.99, 0.95, 0.90], dtype=torch.float32), + bboxes_metadata=[ + {"detection_id": "mock-d0"}, + {"detection_id": "mock-d1"}, + {"detection_id": "mock-d2"}, + ], + ) + + +def _run_tensor_native_inference(model_id: str, images, **kwargs) -> list: + # Native equivalent of _mock_infer_object_detection_from_request_sync: one raw + # `Detections` per input image (the block zips images with this list). + imgs = images if isinstance(images, list) else [images] + return [_native_od_detections() for _ in imgs] + + +def _assert_native_detection(det: NativeDetections) -> None: + # Value parity for the three full-frame OD boxes. Native Detections has no + # __getitem__: detection_id reads from bboxes_metadata, class_name from + # image_metadata, and xyxy entries are torch tensors. + expected = [ + ("mock-d0", "x", 0, 0.99, [30, 30, 90, 90]), + ("mock-d1", "y", 1, 0.95, [210, 20, 290, 100]), + ("mock-d2", "z", 2, 0.90, [400, 10, 500, 110]), + ] + assert isinstance(det, NativeDetections) + assert len(det) == 3 + for i, (det_id, class_name, class_id, conf, expected_xyxy) in enumerate(expected): + assert det.bboxes_metadata[i]["detection_id"] == det_id + assert int(det.class_id[i]) == class_id + assert float(det.confidence[i]) == pytest.approx(conf) + assert det.image_metadata["class_names"][int(det.class_id[i])] == class_name + np.testing.assert_allclose( + det.xyxy[i].cpu().numpy(), expected_xyxy, rtol=0, atol=1e-3 + ) + + +def _assert_collapsed_native_batch(collapsed: list, expected_images: int) -> None: + # Native parity of _assert_collapsed_two_image_batch: collapse flattens the OD batch + # to one native `Detections` per image (each a full frame with the three boxes). + assert isinstance(collapsed, list) + assert len(collapsed) == expected_images + for det in collapsed: + _assert_native_detection(det) + + +@_TENSOR_ONLY +def test_inlined_parent_dimension_collapse_matches_inner_workflow_detection_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + run_mock = mock.MagicMock(side_effect=_run_tensor_native_inference) + inner = child_detection_only_for_parent_dynamic_crop() + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", return_value=["x", "y", "z"] + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine(model_manager, _nested_workflow(inner)) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run(runtime_parameters={"image": dogs_image}) + flat_result = flat_engine.run(runtime_parameters={"image": dogs_image}) + + assert run_mock.call_count == 2 + for call in run_mock.call_args_list: + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + assert model_id == "yolov8n-640" + imgs = call.kwargs.get("images", call.args[1] if len(call.args) > 1 else None) + assert imgs is not None and len(imgs) == 1 + + # Native Detections has no value __eq__; assert the values on BOTH runs instead of + # comparing nested_result == flat_result. Single image -> collapse yields one frame. + assert len(nested_result) == 1 + assert len(flat_result) == 1 + _assert_collapsed_native_batch(nested_result[0]["collapsed"], expected_images=1) + _assert_collapsed_native_batch(flat_result[0]["collapsed"], expected_images=1) + + +@_TENSOR_ONLY +def test_inlined_parent_dimension_collapse_matches_inner_workflow_detection_runtime_image_list_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Two runtime images are still one workflow row; ``dimension_collapse`` flattens the OD batch.""" + run_mock = mock.MagicMock(side_effect=_run_tensor_native_inference) + inner = child_detection_only_for_parent_dynamic_crop() + images = [dogs_image, dogs_image.copy()] + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", return_value=["x", "y", "z"] + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine(model_manager, _nested_workflow(inner)) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run(runtime_parameters={"image": images}) + flat_result = flat_engine.run(runtime_parameters={"image": images}) + + assert run_mock.call_count == 2 + for call in run_mock.call_args_list: + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + assert model_id == "yolov8n-640" + imgs = call.kwargs.get("images", call.args[1] if len(call.args) > 1 else None) + assert imgs is not None and len(imgs) == 2 + + # Native Detections has no value __eq__; assert the values on BOTH runs instead of + # comparing nested_result == flat_result. Two images -> collapse flattens to two + # full-frame Detections. + assert len(nested_result) == 1 + assert len(flat_result) == 1 + _assert_collapsed_native_batch(nested_result[0]["collapsed"], expected_images=2) + _assert_collapsed_native_batch(flat_result[0]["collapsed"], expected_images=2) diff --git a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_runs_dynamic_crop_on_inner_detections.py b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_runs_dynamic_crop_on_inner_detections.py index 6d44e49d01..da814d30c4 100644 --- a/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_runs_dynamic_crop_on_inner_detections.py +++ b/tests/workflows/integration_tests/execution/inner_workflow_inlining/test_inner_workflow_parent_runs_dynamic_crop_on_inner_detections.py @@ -10,7 +10,9 @@ from unittest import mock import numpy as np +import pytest import supervision as sv +import torch from inference.core.entities.requests.inference import ObjectDetectionInferenceRequest from inference.core.entities.responses.inference import ( @@ -18,12 +20,28 @@ ObjectDetectionInferenceResponse, ObjectDetectionPrediction, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.managers.base import ModelManager +from inference_models.models.base.object_detection import Detections as NativeDetections from tests.workflows.integration_tests.execution.inner_workflow_inlining._common import ( child_detection_only_for_parent_dynamic_crop, execution_engine, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the OD block runs inference through +# `run_tensor_native_inference` (+ `get_class_names`), not `infer_from_request_sync`, +# and downstream crop predictions are native `inference_models.Detections` (no +# `__getitem__` / value `__eq__`). The numpy test skips when the flag is on; the +# `*_tensor_native` parity test (skipped when off) drives the same scenario natively. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="numpy ModelManager mock; OD block is native under the flag โ€” see *_tensor_native", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + def _mock_od_response(h: int, w: int) -> ObjectDetectionInferenceResponse: return ObjectDetectionInferenceResponse( @@ -162,6 +180,7 @@ def _assert_crop_predictions_equal(crop_preds: list) -> None: np.testing.assert_allclose(det.xyxy[0], expected_xyxy, rtol=0, atol=1e-3) +@_NUMPY_ONLY def test_inlined_parent_crop_matches_inner_workflow_detection( model_manager: ModelManager, dogs_image: np.ndarray, @@ -198,6 +217,7 @@ def test_inlined_parent_crop_matches_inner_workflow_detection( _assert_crop_predictions_equal(nested_result[0]["from_child"]) +@_NUMPY_ONLY def test_inlined_parent_crop_matches_inner_workflow_detection_runtime_image_list( model_manager: ModelManager, dogs_image: np.ndarray, @@ -235,3 +255,135 @@ def test_inlined_parent_crop_matches_inner_workflow_detection_runtime_image_list assert len(nested_result) == 2 for row in nested_result: _assert_crop_predictions_equal(row["from_child"]) + + +def _native_od_detections() -> NativeDetections: + # Native equivalent of _mock_od_response: same three boxes as a raw `Detections` + # (xyxy from the numpy x/y/w/h centres). The OD block attaches the image lineage; + # detection_id rides in bboxes_metadata, class names via get_class_names. + return NativeDetections( + xyxy=torch.tensor( + [[30, 30, 90, 90], [210, 20, 290, 100], [400, 10, 500, 110]], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 1, 2], dtype=torch.long), + confidence=torch.tensor([0.99, 0.95, 0.90], dtype=torch.float32), + bboxes_metadata=[ + {"detection_id": "mock-d0"}, + {"detection_id": "mock-d1"}, + {"detection_id": "mock-d2"}, + ], + ) + + +def _run_tensor_native_inference_factory(): + # Native equivalent of _mock_infer_object_detection_from_request_sync: one raw + # `Detections` per input image (the native OD block zips images with the returned + # batch), so a list of N is returned for N images. + def run_tensor_native_inference(model_id: str, images, **kwargs): + imgs = images if isinstance(images, list) else [images] + return [_native_od_detections() for _ in imgs] + + return run_tensor_native_inference + + +def _assert_crop_predictions_equal_native(crop_preds: list) -> None: + # Native parity of _assert_crop_predictions_equal: Detections has no __getitem__, + # so detection_id reads from bboxes_metadata and class_name from image_metadata. + expected = [ + ("mock-d0", "x", 0, 0.99, [0, 0, 60, 60]), + ("mock-d1", "y", 1, 0.95, [0, 0, 80, 80]), + ("mock-d2", "z", 2, 0.90, [0, 0, 100, 100]), + ] + assert isinstance(crop_preds, list) + assert len(crop_preds) == 3 + for det, (det_id, class_name, class_id, conf, expected_xyxy) in zip( + crop_preds, expected + ): + assert isinstance(det, NativeDetections) + assert len(det) == 1 + assert det.bboxes_metadata[0]["detection_id"] == det_id + assert int(det.class_id[0]) == class_id + assert float(det.confidence[0]) == pytest.approx(conf) + assert det.image_metadata["class_names"][int(det.class_id[0])] == class_name + np.testing.assert_allclose( + det.xyxy[0].cpu().numpy(), expected_xyxy, rtol=0, atol=1e-3 + ) + + +@_TENSOR_ONLY +def test_inlined_parent_crop_matches_inner_workflow_detection_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + run_mock = mock.MagicMock( + side_effect=_run_tensor_native_inference_factory(), + ) + inner = child_detection_only_for_parent_dynamic_crop() + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", return_value=["x", "y", "z"] + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine(model_manager, _nested_workflow(inner)) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run(runtime_parameters={"image": dogs_image}) + flat_result = flat_engine.run(runtime_parameters={"image": dogs_image}) + + assert run_mock.call_count == 2 + for call in run_mock.call_args_list: + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + assert model_id == "yolov8n-640" + imgs = call.kwargs.get("images", call.args[1] if len(call.args) > 1 else None) + assert imgs is not None and len(imgs) == 1 + + # Native Detections has no value __eq__; assert the values on BOTH runs instead of + # comparing nested_result == flat_result. + assert len(nested_result) == 1 + assert len(flat_result) == 1 + _assert_crop_predictions_equal_native(nested_result[0]["from_child"]) + _assert_crop_predictions_equal_native(flat_result[0]["from_child"]) + + +@_TENSOR_ONLY +def test_inlined_parent_crop_matches_inner_workflow_detection_runtime_image_list_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Same equivalence as the single-image case, but ``image`` is a list at runtime (batch).""" + run_mock = mock.MagicMock( + side_effect=_run_tensor_native_inference_factory(), + ) + inner = child_detection_only_for_parent_dynamic_crop() + images = [dogs_image, dogs_image.copy()] + + with mock.patch.object(ModelManager, "add_model"), mock.patch.object( + ModelManager, "get_class_names", return_value=["x", "y", "z"] + ), mock.patch.object( + ModelManager, + "run_tensor_native_inference", + new=run_mock, + ): + nested_engine = execution_engine(model_manager, _nested_workflow(inner)) + flat_engine = execution_engine(model_manager, _flat_workflow()) + nested_result = nested_engine.run(runtime_parameters={"image": images}) + flat_result = flat_engine.run(runtime_parameters={"image": images}) + + assert run_mock.call_count == 2 + for call in run_mock.call_args_list: + model_id = call.kwargs.get("model_id", call.args[0] if call.args else None) + assert model_id == "yolov8n-640" + imgs = call.kwargs.get("images", call.args[1] if len(call.args) > 1 else None) + assert imgs is not None and len(imgs) == 2 + + # Native Detections has no value __eq__; assert the values on BOTH runs instead of + # comparing nested_result == flat_result. + assert len(nested_result) == 2 + assert len(flat_result) == 2 + for row in nested_result: + _assert_crop_predictions_equal_native(row["from_child"]) + for row in flat_result: + _assert_crop_predictions_equal_native(row["from_child"]) diff --git a/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/__init__.py b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/__init__.py index 4be91a3310..c33a6c9474 100644 --- a/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/__init__.py +++ b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/__init__.py @@ -1,27 +1,59 @@ from typing import List, Type +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.prototypes.block import WorkflowBlock + +# numpy / sv-shaped blocks from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.detections_to_parent_coordinates_batch import ( DetectionsToParentCoordinatesBatchBlock, ) + +# tensor-native siblings +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.detections_to_parent_coordinates_batch_tensor import ( + DetectionsToParentCoordinatesBatchBlock as DetectionsToParentCoordinatesBatchBlockTensor, +) from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.detections_to_parent_coordinates_non_batch import ( DetectionsToParentCoordinatesNonBatchBlock, ) +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.detections_to_parent_coordinates_non_batch_tensor import ( + DetectionsToParentCoordinatesNonBatchBlock as DetectionsToParentCoordinatesNonBatchBlockTensor, +) from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.stitch_detections_batch import ( StitchDetectionsBatchBlock, ) +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.stitch_detections_batch_tensor import ( + StitchDetectionsBatchBlock as StitchDetectionsBatchBlockTensor, +) from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.stitch_detections_non_batch import ( StitchDetectionsNonBatchBlock, ) +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.stitch_detections_non_batch_tensor import ( + StitchDetectionsNonBatchBlock as StitchDetectionsNonBatchBlockTensor, +) from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.tile_detections_batch import ( TileDetectionsBatchBlock, ) +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.tile_detections_batch_tensor import ( + TileDetectionsBatchBlock as TileDetectionsBatchBlockTensor, +) from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.tile_detections_non_batch import ( TileDetectionsNonBatchBlock, ) +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.tile_detections_non_batch_tensor import ( + TileDetectionsNonBatchBlock as TileDetectionsNonBatchBlockTensor, +) def load_blocks() -> List[Type[WorkflowBlock]]: + if ENABLE_TENSOR_DATA_REPRESENTATION: + return [ + DetectionsToParentCoordinatesBatchBlockTensor, + DetectionsToParentCoordinatesNonBatchBlockTensor, + StitchDetectionsBatchBlockTensor, + StitchDetectionsNonBatchBlockTensor, + TileDetectionsBatchBlockTensor, + TileDetectionsNonBatchBlockTensor, + ] return [ DetectionsToParentCoordinatesBatchBlock, DetectionsToParentCoordinatesNonBatchBlock, diff --git a/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/detections_to_parent_coordinates_batch_tensor.py b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/detections_to_parent_coordinates_batch_tensor.py new file mode 100644 index 0000000000..f8d1f51276 --- /dev/null +++ b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/detections_to_parent_coordinates_batch_tensor.py @@ -0,0 +1,91 @@ +""" +Tensor-native sibling of ``detections_to_parent_coordinates_batch.py``: writes +``parent_id`` / ``parent_coordinates`` / ``parent_dimensions`` per-box into +``bboxes_metadata`` (the native counterpart of sv ``.data``). + +This is just example, test implementation, please do not assume it being fully functional. +""" + +from typing import Type + +from inference.core.workflows.core_steps.common.query_language.operations.detections.base import ( + _copy_detections, +) +from inference.core.workflows.execution_engine.constants import ( + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.detections_to_parent_coordinates_batch import ( + BlockManifest, +) + + +def _write_parent_metadata_native( + prediction: Detections, + parent_id: str, + parent_coordinates, +) -> Detections: + """Return a copy of the prediction with parent lineage written per-box + into ``bboxes_metadata``.""" + prediction_copy = _copy_detections(prediction) + number_of_boxes = len(prediction_copy) + bboxes_metadata = prediction_copy.bboxes_metadata + if bboxes_metadata is None: + bboxes_metadata = [{} for _ in range(number_of_boxes)] + else: + bboxes_metadata = [dict(entry) for entry in bboxes_metadata] + offset = [0, 0] + dimensions = ( + [parent_coordinates.origin_height, parent_coordinates.origin_width] + if parent_coordinates + else None + ) + for entry in bboxes_metadata: + entry[PARENT_ID_KEY] = parent_id + if parent_coordinates: + entry[PARENT_COORDINATES_KEY] = list(offset) + entry[PARENT_DIMENSIONS_KEY] = list(dimensions) + prediction_copy.bboxes_metadata = bboxes_metadata + return prediction_copy + + +class DetectionsToParentCoordinatesBatchBlock(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + @classmethod + def accepts_batch_input(cls) -> bool: + return True + + def run( + self, + images: Batch[WorkflowImageData], + images_predictions: Batch[Batch[Detections]], + ) -> BlockResult: + result = [] + for image, image_predictions in zip(images, images_predictions): + parent_id = image.parent_metadata.parent_id + parent_coordinates = image.parent_metadata.origin_coordinates + transformed_predictions = [] + for prediction in image_predictions: + prediction_copy = _write_parent_metadata_native( + prediction=prediction, + parent_id=parent_id, + parent_coordinates=parent_coordinates, + ) + transformed_predictions.append({"predictions": prediction_copy}) + result.append(transformed_predictions) + return result diff --git a/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/detections_to_parent_coordinates_non_batch_tensor.py b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/detections_to_parent_coordinates_non_batch_tensor.py new file mode 100644 index 0000000000..bd21a72515 --- /dev/null +++ b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/detections_to_parent_coordinates_non_batch_tensor.py @@ -0,0 +1,50 @@ +""" +Tensor-native sibling of ``detections_to_parent_coordinates_non_batch.py``: writes +``parent_id`` / ``parent_coordinates`` / ``parent_dimensions`` per-box into +``bboxes_metadata`` (the native counterpart of sv ``.data``). + +This is just example, test implementation, please do not assume it being fully functional. +""" + +from typing import Type + +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.detections_to_parent_coordinates_batch_tensor import ( + _write_parent_metadata_native, +) +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.detections_to_parent_coordinates_non_batch import ( + BlockManifest, +) + + +class DetectionsToParentCoordinatesNonBatchBlock(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + image_predictions: Batch[Detections], + ) -> BlockResult: + parent_id = image.parent_metadata.parent_id + parent_coordinates = image.parent_metadata.origin_coordinates + transformed_predictions = [] + for prediction in image_predictions: + prediction_copy = _write_parent_metadata_native( + prediction=prediction, + parent_id=parent_id, + parent_coordinates=parent_coordinates, + ) + transformed_predictions.append({"predictions": prediction_copy}) + return transformed_predictions diff --git a/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/stitch_detections_batch_tensor.py b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/stitch_detections_batch_tensor.py new file mode 100644 index 0000000000..bbf311f662 --- /dev/null +++ b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/stitch_detections_batch_tensor.py @@ -0,0 +1,90 @@ +""" +Tensor-native sibling of ``stitch_detections_batch.py``: shifts each crop's ``xyxy`` +by the origin the native OD producer stored per-image in +``image_metadata[PARENT_COORDINATES_KEY]``, then concatenates the per-crop +``Detections``. + +This is just example, test implementation, please do not assume it being fully functional. +""" + +from copy import deepcopy +from typing import List, Type + +import torch + +from inference.core.workflows.core_steps.common.query_language.operations.detections.base import ( + _concatenate_detections, + _copy_detections, +) +from inference.core.workflows.execution_engine.constants import PARENT_COORDINATES_KEY +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.stitch_detections_batch import ( + BlockManifest, +) + + +def _shift_native_to_parent(prediction: Detections) -> Detections: + """Return a copy of the prediction with ``xyxy`` shifted by the crop origin + stored in ``image_metadata[PARENT_COORDINATES_KEY]`` as ``[x, y]``.""" + prediction_copy = _copy_detections(prediction) + image_metadata = prediction_copy.image_metadata or {} + coords = image_metadata.get(PARENT_COORDINATES_KEY, [0, 0]) + shift_x, shift_y = float(coords[0]), float(coords[1]) + shift = torch.as_tensor( + [shift_x, shift_y, shift_x, shift_y], + dtype=prediction_copy.xyxy.dtype, + device=prediction_copy.xyxy.device, + ) + prediction_copy.xyxy = prediction_copy.xyxy + shift + return prediction_copy + + +def _empty_like(prediction: Detections) -> Detections: + """Build an empty native ``Detections`` on the same device/dtype as ``prediction``.""" + xyxy = prediction.xyxy + class_id = prediction.class_id + confidence = prediction.confidence + return Detections( + xyxy=xyxy.new_zeros((0, 4)), + class_id=class_id.new_zeros((0,)), + confidence=confidence.new_zeros((0,)), + image_metadata=deepcopy(prediction.image_metadata), + bboxes_metadata=[], + ) + + +def merge_native_predictions(image_predictions: List[Detections]) -> Detections: + non_empty = [_shift_native_to_parent(p) for p in image_predictions if len(p)] + if not non_empty: + return _empty_like(image_predictions[0]) + merged = non_empty[0] + for prediction in non_empty[1:]: + merged = _concatenate_detections(merged, prediction) + return merged + + +class StitchDetectionsBatchBlock(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images: Batch[WorkflowImageData], + images_predictions: Batch[Batch[Detections]], + ) -> BlockResult: + result = [] + for image, image_predictions in zip(images, images_predictions): + merged_prediction = merge_native_predictions(list(image_predictions)) + result.append({"predictions": merged_prediction}) + return result diff --git a/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/stitch_detections_non_batch_tensor.py b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/stitch_detections_non_batch_tensor.py new file mode 100644 index 0000000000..8b1140787a --- /dev/null +++ b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/stitch_detections_non_batch_tensor.py @@ -0,0 +1,42 @@ +""" +Tensor-native sibling of ``stitch_detections_non_batch.py``: shifts each crop's +``xyxy`` by the origin stored per-image in +``image_metadata[PARENT_COORDINATES_KEY]``, then concatenates the per-crop +``Detections``. + +This is just example, test implementation, please do not assume it being fully functional. +""" + +from typing import Type + +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.stitch_detections_batch_tensor import ( + merge_native_predictions, +) +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.stitch_detections_non_batch import ( + BlockManifest, +) + + +class StitchDetectionsNonBatchBlock(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + image: WorkflowImageData, + image_predictions: Batch[Detections], + ) -> BlockResult: + merged_prediction = merge_native_predictions(list(image_predictions)) + return {"predictions": merged_prediction} diff --git a/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/tile_detections_batch_tensor.py b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/tile_detections_batch_tensor.py new file mode 100644 index 0000000000..ddb4ea0f47 --- /dev/null +++ b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/tile_detections_batch_tensor.py @@ -0,0 +1,52 @@ +""" +Tensor-native sibling of ``tile_detections_batch.py``: native ``Detections`` is +converted via ``.to_supervision()`` only to feed ``sv.BoxAnnotator``; the block +outputs image tiles, not detections. + +This is just example, test implementation, please do not assume it being fully functional. +""" + +from typing import Type + +import supervision as sv + +from inference.core.utils.drawing import create_tiles +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.tile_detections_batch import ( + BlockManifest, +) + + +class TileDetectionsBatchBlock(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + images_crops: Batch[Batch[WorkflowImageData]], + crops_predictions: Batch[Batch[Detections]], + ) -> BlockResult: + annotator = sv.BoxAnnotator() + visualisations = [] + for image_crops, crop_predictions in zip(images_crops, crops_predictions): + visualisations_batch_element = [] + for image, prediction in zip(image_crops, crop_predictions): + annotated_image = annotator.annotate( + image.numpy_image.copy(), + prediction.to_supervision(), + ) + visualisations_batch_element.append(annotated_image) + tile = create_tiles(visualisations_batch_element) + visualisations.append({"visualisations": tile}) + return visualisations diff --git a/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/tile_detections_non_batch_tensor.py b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/tile_detections_non_batch_tensor.py new file mode 100644 index 0000000000..e497d65fb9 --- /dev/null +++ b/tests/workflows/integration_tests/execution/stub_plugins/dimensionality_manipulation_plugin/tile_detections_non_batch_tensor.py @@ -0,0 +1,49 @@ +""" +Tensor-native sibling of ``tile_detections_non_batch.py``: native ``Detections`` is +converted via ``.to_supervision()`` only to feed ``sv.BoxAnnotator``; the block +outputs image tiles, not detections. + +This is just example, test implementation, please do not assume it being fully functional. +""" + +from typing import Type + +import supervision as sv + +from inference.core.utils.drawing import create_tiles +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + WorkflowImageData, +) +from inference.core.workflows.prototypes.block import ( + BlockResult, + WorkflowBlock, + WorkflowBlockManifest, +) +from inference_models.models.base.object_detection import Detections +from tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin.tile_detections_non_batch import ( + BlockManifest, +) + + +class TileDetectionsNonBatchBlock(WorkflowBlock): + + @classmethod + def get_manifest(cls) -> Type[WorkflowBlockManifest]: + return BlockManifest + + def run( + self, + crops: Batch[WorkflowImageData], + crops_predictions: Batch[Detections], + ) -> BlockResult: + annotator = sv.BoxAnnotator() + visualisations = [] + for image, prediction in zip(crops, crops_predictions): + annotated_image = annotator.annotate( + image.numpy_image.copy(), + prediction.to_supervision(), + ) + visualisations.append(annotated_image) + tile = create_tiles(visualisations) + return {"visualisations": tile} diff --git a/tests/workflows/integration_tests/execution/stub_plugins/plugin_handling_video_metadata/__init__.py b/tests/workflows/integration_tests/execution/stub_plugins/plugin_handling_video_metadata/__init__.py index 60e6b4d048..e004f4b519 100644 --- a/tests/workflows/integration_tests/execution/stub_plugins/plugin_handling_video_metadata/__init__.py +++ b/tests/workflows/integration_tests/execution/stub_plugins/plugin_handling_video_metadata/__init__.py @@ -88,9 +88,19 @@ def run( ) -> BlockResult: if metadata.video_identifier not in self._trackers: self._trackers[metadata.video_identifier] = sv.ByteTrack() + if isinstance(predictions, sv.Detections): + tracker_input = deepcopy(predictions) + else: + # Under ENABLE_TENSOR_DATA_REPRESENTATION predictions is a native + # inference_models.Detections; materialise it for sv.ByteTrack. + tracker_input = sv.Detections( + xyxy=predictions.xyxy.detach().cpu().numpy(), + class_id=predictions.class_id.detach().cpu().numpy(), + confidence=predictions.confidence.detach().cpu().numpy(), + ) tracked_detections = self._trackers[ metadata.video_identifier - ].update_with_detections(detections=deepcopy(predictions)) + ].update_with_detections(detections=tracker_input) result = ( tracked_detections.tracker_id.tolist() if tracked_detections.tracker_id is not None diff --git a/tests/workflows/integration_tests/execution/tensor_input_utils.py b/tests/workflows/integration_tests/execution/tensor_input_utils.py new file mode 100644 index 0000000000..f5f0a4a97e --- /dev/null +++ b/tests/workflows/integration_tests/execution/tensor_input_utils.py @@ -0,0 +1,28 @@ +"""Helpers for tensor-input parity tests. + +Under ENABLE_TENSOR_DATA_REPRESENTATION the `_tensor_native` integration tests still feed +the workflow a *numpy* fixture, so the model blocks hit their numpy fallback. To exercise +the OTHER branch (`WorkflowImageData.is_tensor_materialised()` -> True, blocks use the +on-device tensor + "rgb"), feed the image already materialised as a tensor. + +`numpy_image_as_tensor` converts a BGR HWC uint8 fixture into exactly the CHW RGB uint8 +tensor that `WorkflowImageData.tensor_image` builds, on WORKFLOWS_IMAGE_TENSOR_DEVICE. The +image deserializer wraps a raw torch.Tensor as `WorkflowImageData(tensor_image=...)`, so +passing the result as the `image` runtime parameter makes the input arrive pre-materialised. +""" + +import numpy as np +import torch + +from inference.core.env import WORKFLOWS_IMAGE_TENSOR_DEVICE + + +def numpy_image_as_tensor(bgr_hwc_image: np.ndarray) -> torch.Tensor: + """BGR HWC uint8 numpy fixture -> CHW RGB uint8 tensor on WORKFLOWS_IMAGE_TENSOR_DEVICE.""" + rgb = bgr_hwc_image[:, :, ::-1].copy() + return ( + torch.from_numpy(rgb) + .permute(2, 0, 1) + .contiguous() + .to(WORKFLOWS_IMAGE_TENSOR_DEVICE) + ) diff --git a/tests/workflows/integration_tests/execution/test_edge_snap_cv_pipeline.py b/tests/workflows/integration_tests/execution/test_edge_snap_cv_pipeline.py index 011c2858e4..832b058332 100644 --- a/tests/workflows/integration_tests/execution/test_edge_snap_cv_pipeline.py +++ b/tests/workflows/integration_tests/execution/test_edge_snap_cv_pipeline.py @@ -3,11 +3,37 @@ import numpy as np import pytest import supervision as sv +import torch -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.instance_segmentation import ( + InstanceDetections as NativeInstanceDetections, +) +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# The edge-snap pipeline pipes a `segmentation` runtime input into `mask_edge_snap`. +# Under ENABLE_TENSOR_DATA_REPRESENTATION that block accepts a tensor-native +# `InstanceDetections`, not sv.Detections โ€” so the two segmentation-feeding tests below +# are split into `@_NUMPY_ONLY` (sv input) + `@_TENSOR_ONLY` (native input) parity pairs +# asserting the same outputs. The preprocessing-only tests in this file feed no +# segmentation and run unchanged in both modes. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections segmentation input; mask_edge_snap is native-only under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) WORKFLOW_WITH_CLASSICAL_CV_PREPROCESSING = { "version": "1.0", @@ -359,6 +385,7 @@ def test_classical_cv_preprocessing_pipeline_all_intermediate_outputs( assert preprocessed.shape == original.shape +@_NUMPY_ONLY def test_edge_snap_cv_pipeline_with_segmentation( model_manager: ModelManager, ) -> None: @@ -421,6 +448,7 @@ def test_edge_snap_cv_pipeline_with_segmentation( assert refined_seg.mask is not None or len(refined_seg.mask) == 0 +@_NUMPY_ONLY def test_edge_snap_cv_pipeline_with_empty_segmentation( model_manager: ModelManager, ) -> None: @@ -459,3 +487,227 @@ def test_edge_snap_cv_pipeline_with_empty_segmentation( # Verify preprocessing still works assert preprocessed.shape == test_image.shape assert preprocessed.dtype == np.uint8 + + +# --------------------------------------------------------------------------- +# Tensor-native parity variants (run only under ENABLE_TENSOR_DATA_REPRESENTATION): +# same pipeline, but the segmentation runtime input is a native InstanceDetections. +# --------------------------------------------------------------------------- + + +@_TENSOR_ONLY +def test_edge_snap_cv_pipeline_with_segmentation_tensor_native( + model_manager: ModelManager, +) -> None: + """Edge snap pipeline with preprocessing + native InstanceDetections segmentation.""" + # given - Create a test image with a synthetic object + test_image = np.ones((200, 200, 3), dtype=np.uint8) * 100 + test_image[50:150, 50:150] = 150 + np.random.seed(42) + noise = np.random.randint(-20, 20, (200, 200, 3), dtype=np.int16) + noisy_image = np.clip(test_image.astype(np.int16) + noise, 0, 255).astype(np.uint8) + + # Synthetic native segmentation with a single rectangular mask + mask = np.zeros((200, 200), dtype=np.uint8) + mask[50:150, 50:150] = 1 + + detections = NativeInstanceDetections( + xyxy=torch.tensor([[50, 50, 150, 150]], dtype=torch.float32), + confidence=torch.tensor([0.95], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + mask=torch.as_tensor(mask[np.newaxis, :, :], dtype=torch.bool), + image_metadata=None, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_EDGE_SNAP, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": noisy_image, "segmentation": detections}, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + result_dict = result[0] + assert "refined_segmentation" in result_dict + assert "preprocessed_image" in result_dict + + refined_seg = result_dict["refined_segmentation"] + preprocessed = result_dict["preprocessed_image"].numpy_image + + # Verify outputs + assert refined_seg is not None + assert preprocessed.shape == noisy_image.shape + assert preprocessed.dtype == np.uint8 + + # Verify that refined_segmentation has masks + assert hasattr(refined_seg, "mask") + assert refined_seg.mask is not None or len(refined_seg.mask) == 0 + + +@_TENSOR_ONLY +def test_edge_snap_cv_pipeline_with_segmentation_with_tensor_input( + model_manager: ModelManager, +) -> None: + """Same as test_edge_snap_cv_pipeline_with_segmentation_tensor_native, but the image + arrives ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() == + True), so the preprocessing blocks run their on-device tensor path. The segmentation is + a native InstanceDetections; results must match the numpy-input variant.""" + # given - Create a test image with a synthetic object + test_image = np.ones((200, 200, 3), dtype=np.uint8) * 100 + test_image[50:150, 50:150] = 150 + np.random.seed(42) + noise = np.random.randint(-20, 20, (200, 200, 3), dtype=np.int16) + noisy_image = np.clip(test_image.astype(np.int16) + noise, 0, 255).astype(np.uint8) + + # Synthetic native segmentation with a single rectangular mask + mask = np.zeros((200, 200), dtype=np.uint8) + mask[50:150, 50:150] = 1 + + detections = NativeInstanceDetections( + xyxy=torch.tensor([[50, 50, 150, 150]], dtype=torch.float32), + confidence=torch.tensor([0.95], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + mask=torch.as_tensor(mask[np.newaxis, :, :], dtype=torch.bool), + image_metadata=None, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_EDGE_SNAP, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the image as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(noisy_image), + "segmentation": detections, + }, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + result_dict = result[0] + assert "refined_segmentation" in result_dict + assert "preprocessed_image" in result_dict + + refined_seg = result_dict["refined_segmentation"] + preprocessed = result_dict["preprocessed_image"].numpy_image + + # Verify outputs + assert refined_seg is not None + assert preprocessed.shape == noisy_image.shape + assert preprocessed.dtype == np.uint8 + + # Verify that refined_segmentation has masks + assert hasattr(refined_seg, "mask") + assert refined_seg.mask is not None or len(refined_seg.mask) == 0 + + +@_TENSOR_ONLY +def test_edge_snap_cv_pipeline_with_empty_segmentation_tensor_native( + model_manager: ModelManager, +) -> None: + """Edge snap pipeline with empty native segmentation.""" + # given - Create a simple test image + test_image = np.ones((150, 150, 3), dtype=np.uint8) * 120 + + # Empty native detections + detections = NativeInstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, 150, 150), dtype=torch.bool), + image_metadata=None, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_EDGE_SNAP, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": test_image, "segmentation": detections}, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + result_dict = result[0] + assert "refined_segmentation" in result_dict + preprocessed = result_dict["preprocessed_image"].numpy_image + + # Verify preprocessing still works + assert preprocessed.shape == test_image.shape + assert preprocessed.dtype == np.uint8 + + +@_TENSOR_ONLY +def test_edge_snap_cv_pipeline_with_empty_segmentation_with_tensor_input( + model_manager: ModelManager, +) -> None: + """Same as test_edge_snap_cv_pipeline_with_empty_segmentation_tensor_native, but the + image arrives ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() + == True), so the preprocessing blocks run their on-device tensor path. Results must match + the numpy-input variant.""" + # given - Create a simple test image + test_image = np.ones((150, 150, 3), dtype=np.uint8) * 120 + + # Empty native detections + detections = NativeInstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, 150, 150), dtype=torch.bool), + image_metadata=None, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_EDGE_SNAP, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the image as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(test_image), + "segmentation": detections, + }, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + result_dict = result[0] + assert "refined_segmentation" in result_dict + preprocessed = result_dict["preprocessed_image"].numpy_image + + # Verify preprocessing still works + assert preprocessed.shape == test_image.shape + assert preprocessed.dtype == np.uint8 diff --git a/tests/workflows/integration_tests/execution/test_mask_edge_snap.py b/tests/workflows/integration_tests/execution/test_mask_edge_snap.py index 39d4983305..6838491160 100644 --- a/tests/workflows/integration_tests/execution/test_mask_edge_snap.py +++ b/tests/workflows/integration_tests/execution/test_mask_edge_snap.py @@ -1,11 +1,56 @@ import numpy as np import pytest import supervision as sv +import torch -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.instance_segmentation import ( + InstanceDetections as NativeInstanceDetections, +) +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the `mask_edge_snap` block's `segmentation` +# input is a tensor-native `InstanceDetections` (not sv.Detections), and its outputs +# (`refined_segmentation`, `edges`) are native `InstanceDetections`. The sv-based tests +# below are skipped when the flag is on; each has a `*_tensor_native` parity test +# (skipped when the flag is off) that feeds the native equivalent and asserts the same +# refined shapes. `InstanceDetections` exposes `__len__` and a dense bool `mask` +# tensor of shape (n, H, W), so the assertions mirror the sv ones. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections segmentation input; mask_edge_snap is native-only under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + + +def _native_segmentation( + xyxy: np.ndarray, + masks: np.ndarray, + confidence: np.ndarray, + class_id: np.ndarray, +) -> NativeInstanceDetections: + """Tensor-native `InstanceDetections` equivalent of the sv.Detections fixtures + used by the numpy tests: dense bool masks of shape (n, H, W) carried as a torch + tensor โ€” the same representation a segmentation-model tensor sibling produces.""" + return NativeInstanceDetections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32), + class_id=torch.as_tensor(class_id, dtype=torch.long), + confidence=torch.as_tensor(confidence, dtype=torch.float32), + mask=torch.as_tensor(masks, dtype=torch.bool), + image_metadata=None, + ) def _build_mask_edge_snap_workflow() -> dict: @@ -67,6 +112,7 @@ def _build_mask_edge_snap_workflow() -> dict: @pytest.mark.slow +@_NUMPY_ONLY def test_mask_edge_snap_workflow_with_empty_segmentation( model_manager: ModelManager, dogs_image: np.ndarray, @@ -111,6 +157,7 @@ def test_mask_edge_snap_workflow_with_empty_segmentation( @pytest.mark.slow +@_NUMPY_ONLY def test_mask_edge_snap_workflow_with_single_mask( model_manager: ModelManager, dogs_image: np.ndarray, @@ -167,6 +214,7 @@ def test_mask_edge_snap_workflow_with_single_mask( @pytest.mark.slow +@_NUMPY_ONLY def test_mask_edge_snap_workflow_with_multiple_masks( model_manager: ModelManager, dogs_image: np.ndarray, @@ -225,6 +273,7 @@ def test_mask_edge_snap_workflow_with_multiple_masks( @pytest.mark.slow +@_NUMPY_ONLY def test_mask_edge_snap_workflow_with_permissive_parameters( model_manager: ModelManager, dogs_image: np.ndarray, @@ -277,6 +326,7 @@ def test_mask_edge_snap_workflow_with_permissive_parameters( @pytest.mark.slow +@_NUMPY_ONLY def test_mask_edge_snap_workflow_with_strict_parameters( model_manager: ModelManager, dogs_image: np.ndarray, @@ -329,6 +379,7 @@ def test_mask_edge_snap_workflow_with_strict_parameters( @pytest.mark.slow +@_NUMPY_ONLY def test_mask_edge_snap_workflow_with_different_mask_sizes( model_manager: ModelManager, dogs_image: np.ndarray, @@ -429,6 +480,7 @@ def _build_mask_edge_snap_with_morphological_preprocessing_workflow() -> dict: @pytest.mark.slow +@_NUMPY_ONLY def test_mask_edge_snap_workflow_with_morphological_preprocessing( model_manager: ModelManager, dogs_image: np.ndarray, @@ -481,3 +533,788 @@ def test_mask_edge_snap_workflow_with_morphological_preprocessing( assert len(refined) == 1 assert refined.mask is not None assert refined.mask[0].shape == (h, w) + + +# --------------------------------------------------------------------------- +# Tensor-native parity variants (run only under ENABLE_TENSOR_DATA_REPRESENTATION). +# Same scenarios as the sv.Detections tests above, but feeding the segmentation as a +# native `inference_models.InstanceDetections` and asserting the same refined shapes +# on the native `InstanceDetections` outputs. +# --------------------------------------------------------------------------- + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_empty_segmentation_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Test Mask Edge Snap workflow with empty native segmentation.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + empty_segmentation = NativeInstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, h, w), dtype=torch.bool), + image_metadata=None, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "segmentation": empty_segmentation, + "pixel_tolerance": 15, + "sigma": 1.0, + "min_contour_area": 50.0, + "dilation_iterations": 2, + "boundary_band_width": 15, + "adaptive_window_size": 41, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + assert len(output["refined_segmentation"]) == 0 + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_empty_segmentation_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Same as test_mask_edge_snap_workflow_with_empty_segmentation_tensor_native, but the + image arrives ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() + == True), so the block runs its on-device tensor path. Results must match the numpy-input + variant.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + empty_segmentation = NativeInstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, h, w), dtype=torch.bool), + image_metadata=None, + ) + + # when โ€” feed the image as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "segmentation": empty_segmentation, + "pixel_tolerance": 15, + "sigma": 1.0, + "min_contour_area": 50.0, + "dilation_iterations": 2, + "boundary_band_width": 15, + "adaptive_window_size": 41, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + assert len(output["refined_segmentation"]) == 0 + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_single_mask_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Test Mask Edge Snap workflow with a single native mask.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask = np.zeros((h, w), dtype=bool) + mask[50:150, 100:200] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0]]), + masks=np.array([mask]), + confidence=np.array([0.9]), + class_id=np.array([0]), + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "segmentation": segmentation, + "pixel_tolerance": 15, + "sigma": 1.0, + "min_contour_area": 50.0, + "dilation_iterations": 2, + "boundary_band_width": 15, + "adaptive_window_size": 41, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + refined = output["refined_segmentation"] + assert len(refined) == 1 + assert refined.mask is not None + assert tuple(refined.mask[0].shape) == (h, w) + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_single_mask_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Same as test_mask_edge_snap_workflow_with_single_mask_tensor_native, but the image + arrives ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() == + True), so the block runs its on-device tensor path. Results must match the numpy-input + variant.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask = np.zeros((h, w), dtype=bool) + mask[50:150, 100:200] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0]]), + masks=np.array([mask]), + confidence=np.array([0.9]), + class_id=np.array([0]), + ) + + # when โ€” feed the image as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "segmentation": segmentation, + "pixel_tolerance": 15, + "sigma": 1.0, + "min_contour_area": 50.0, + "dilation_iterations": 2, + "boundary_band_width": 15, + "adaptive_window_size": 41, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + refined = output["refined_segmentation"] + assert len(refined) == 1 + assert refined.mask is not None + assert tuple(refined.mask[0].shape) == (h, w) + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_multiple_masks_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Test Mask Edge Snap workflow with multiple native masks.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask1 = np.zeros((h, w), dtype=bool) + mask1[50:150, 100:200] = True + + mask2 = np.zeros((h, w), dtype=bool) + mask2[200:300, 300:400] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0], [300.0, 200.0, 400.0, 300.0]]), + masks=np.array([mask1, mask2]), + confidence=np.array([0.9, 0.85]), + class_id=np.array([0, 1]), + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "segmentation": segmentation, + "pixel_tolerance": 15, + "sigma": 1.0, + "min_contour_area": 50.0, + "dilation_iterations": 2, + "boundary_band_width": 15, + "adaptive_window_size": 41, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + refined = output["refined_segmentation"] + assert len(refined) == 2 + assert refined.mask.shape[0] == 2 + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_multiple_masks_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Same as test_mask_edge_snap_workflow_with_multiple_masks_tensor_native, but the image + arrives ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() == + True), so the block runs its on-device tensor path. Results must match the numpy-input + variant.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask1 = np.zeros((h, w), dtype=bool) + mask1[50:150, 100:200] = True + + mask2 = np.zeros((h, w), dtype=bool) + mask2[200:300, 300:400] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0], [300.0, 200.0, 400.0, 300.0]]), + masks=np.array([mask1, mask2]), + confidence=np.array([0.9, 0.85]), + class_id=np.array([0, 1]), + ) + + # when โ€” feed the image as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "segmentation": segmentation, + "pixel_tolerance": 15, + "sigma": 1.0, + "min_contour_area": 50.0, + "dilation_iterations": 2, + "boundary_band_width": 15, + "adaptive_window_size": 41, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + refined = output["refined_segmentation"] + assert len(refined) == 2 + assert refined.mask.shape[0] == 2 + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_permissive_parameters_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Test Mask Edge Snap workflow with permissive parameters (native input).""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask = np.zeros((h, w), dtype=bool) + mask[50:150, 100:200] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0]]), + masks=np.array([mask]), + confidence=np.array([0.9]), + class_id=np.array([0]), + ) + + # when - using permissive parameters + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "segmentation": segmentation, + "pixel_tolerance": 5, + "sigma": 0.3, + "min_contour_area": 10.0, + "dilation_iterations": 1, + "boundary_band_width": 10, + "adaptive_window_size": 21, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1 + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_permissive_parameters_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Same as test_mask_edge_snap_workflow_with_permissive_parameters_tensor_native, but the + image arrives ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() + == True), so the block runs its on-device tensor path. Results must match the numpy-input + variant.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask = np.zeros((h, w), dtype=bool) + mask[50:150, 100:200] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0]]), + masks=np.array([mask]), + confidence=np.array([0.9]), + class_id=np.array([0]), + ) + + # when - using permissive parameters, image fed as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "segmentation": segmentation, + "pixel_tolerance": 5, + "sigma": 0.3, + "min_contour_area": 10.0, + "dilation_iterations": 1, + "boundary_band_width": 10, + "adaptive_window_size": 21, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1 + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_strict_parameters_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Test Mask Edge Snap workflow with strict parameters (native input).""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask = np.zeros((h, w), dtype=bool) + mask[50:150, 100:200] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0]]), + masks=np.array([mask]), + confidence=np.array([0.9]), + class_id=np.array([0]), + ) + + # when - using strict parameters + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "segmentation": segmentation, + "pixel_tolerance": 50, + "sigma": 2.0, + "min_contour_area": 200.0, + "dilation_iterations": 5, + "boundary_band_width": 50, + "adaptive_window_size": 81, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1 + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_strict_parameters_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Same as test_mask_edge_snap_workflow_with_strict_parameters_tensor_native, but the + image arrives ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() + == True), so the block runs its on-device tensor path. Results must match the numpy-input + variant.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask = np.zeros((h, w), dtype=bool) + mask[50:150, 100:200] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0]]), + masks=np.array([mask]), + confidence=np.array([0.9]), + class_id=np.array([0]), + ) + + # when - using strict parameters, image fed as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "segmentation": segmentation, + "pixel_tolerance": 50, + "sigma": 2.0, + "min_contour_area": 200.0, + "dilation_iterations": 5, + "boundary_band_width": 50, + "adaptive_window_size": 81, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1 + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_different_mask_sizes_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Test Mask Edge Snap workflow with native masks of different sizes.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + + small_mask = np.zeros((h, w), dtype=bool) + small_mask[100:120, 150:170] = True + + medium_mask = np.zeros((h, w), dtype=bool) + medium_mask[50:200, 100:300] = True + + segmentation = _native_segmentation( + xyxy=np.array([[150.0, 100.0, 170.0, 120.0], [100.0, 50.0, 300.0, 200.0]]), + masks=np.array([small_mask, medium_mask]), + confidence=np.array([0.9, 0.85]), + class_id=np.array([0, 1]), + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "segmentation": segmentation, + "pixel_tolerance": 15, + "sigma": 1.0, + "min_contour_area": 50.0, + "dilation_iterations": 2, + "boundary_band_width": 15, + "adaptive_window_size": 41, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + refined = output["refined_segmentation"] + assert len(refined) == 2, "Both masks should be refined" + assert refined.mask.shape[0] == 2 + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_different_mask_sizes_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Same as test_mask_edge_snap_workflow_with_different_mask_sizes_tensor_native, but the + image arrives ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() + == True), so the block runs its on-device tensor path. Results must match the numpy-input + variant.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + + small_mask = np.zeros((h, w), dtype=bool) + small_mask[100:120, 150:170] = True + + medium_mask = np.zeros((h, w), dtype=bool) + medium_mask[50:200, 100:300] = True + + segmentation = _native_segmentation( + xyxy=np.array([[150.0, 100.0, 170.0, 120.0], [100.0, 50.0, 300.0, 200.0]]), + masks=np.array([small_mask, medium_mask]), + confidence=np.array([0.9, 0.85]), + class_id=np.array([0, 1]), + ) + + # when โ€” feed the image as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "segmentation": segmentation, + "pixel_tolerance": 15, + "sigma": 1.0, + "min_contour_area": 50.0, + "dilation_iterations": 2, + "boundary_band_width": 15, + "adaptive_window_size": 41, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + refined = output["refined_segmentation"] + assert len(refined) == 2, "Both masks should be refined" + assert refined.mask.shape[0] == 2 + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_morphological_preprocessing_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Test Mask Edge Snap workflow with morphological preprocessing (native input).""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_with_morphological_preprocessing_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask = np.zeros((h, w), dtype=bool) + mask[50:150, 100:200] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0]]), + masks=np.array([mask]), + confidence=np.array([0.9]), + class_id=np.array([0]), + ) + + # when - run workflow with morphological preprocessing + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "segmentation": segmentation, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + refined = output["refined_segmentation"] + assert len(refined) == 1 + assert refined.mask is not None + assert tuple(refined.mask[0].shape) == (h, w) + + +@pytest.mark.slow +@_TENSOR_ONLY +def test_mask_edge_snap_workflow_with_morphological_preprocessing_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Same as test_mask_edge_snap_workflow_with_morphological_preprocessing_tensor_native, + but the image arrives ALREADY materialised as a CHW RGB device tensor + (is_tensor_materialised() == True), so the preprocessing + edge-snap blocks run their + on-device tensor path. Results must match the numpy-input variant.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_mask_edge_snap_with_morphological_preprocessing_workflow(), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + h, w = dogs_image.shape[:2] + mask = np.zeros((h, w), dtype=bool) + mask[50:150, 100:200] = True + + segmentation = _native_segmentation( + xyxy=np.array([[100.0, 50.0, 200.0, 150.0]]), + masks=np.array([mask]), + confidence=np.array([0.9]), + class_id=np.array([0]), + ) + + # when - run workflow with morphological preprocessing, image fed as a tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "segmentation": segmentation, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + output = result[0]["result"] + assert "refined_segmentation" in output + assert "edges" in output + refined = output["refined_segmentation"] + assert len(refined) == 1 + assert refined.mask is not None + assert tuple(refined.mask[0].shape) == (h, w) diff --git a/tests/workflows/integration_tests/execution/test_plugins_enforcing_scalars_to_fit_into_batch_parameters.py b/tests/workflows/integration_tests/execution/test_plugins_enforcing_scalars_to_fit_into_batch_parameters.py index 2eebf419dd..51ca8f6baf 100644 --- a/tests/workflows/integration_tests/execution/test_plugins_enforcing_scalars_to_fit_into_batch_parameters.py +++ b/tests/workflows/integration_tests/execution/test_plugins_enforcing_scalars_to_fit_into_batch_parameters.py @@ -4,12 +4,34 @@ import numpy as np import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.errors import AssumptionError +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.core import ExecutionEngine from inference.core.workflows.execution_engine.introspection import blocks_loader +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the classification block emits native +# inference_models.ClassificationPrediction objects instead of sv-shaped dicts, so the +# numpy-shaped `e["top"]` subscript below fails. The numpy assertions run only with the +# flag off; an equivalent *_tensor_native test asserts the same facts against the native +# classification carrier with the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="classification output; native under ENABLE_TENSOR_DATA_REPRESENTATION " + "โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) WORKFLOW_IMAGE_PRODUCER_SINGLE_IMAGE_SIMD_CONSUMER = { "version": "1.1", @@ -4436,6 +4458,7 @@ def test_workflow_with_batch_inputs_feeding_simd_consumer_raising_dim( } +@_NUMPY_ONLY @mock.patch.object(blocks_loader, "get_plugin_modules") def test_workflow_with_input_derived_dims_and_emergent_dims( get_plugin_modules_mock: MagicMock, @@ -4490,3 +4513,141 @@ def test_workflow_with_input_derived_dims_and_emergent_dims( assert result[1]["collapsed_output"] == ["[192, 168, 3][292, 168, 3]"] assert result[1]["collapsed_output_2"] == [["[192, 168, 3][292, 168, 3]"]] assert result[1]["breds_classification"] == [] + + +@_TENSOR_ONLY +@mock.patch.object(blocks_loader, "get_plugin_modules") +def test_workflow_with_input_derived_dims_and_emergent_dims_tensor_native( + get_plugin_modules_mock: MagicMock, + model_manager: ModelManager, + dogs_image: np.ndarray, + crowd_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + get_plugin_modules_mock.return_value = [ + "tests.workflows.integration_tests.execution.stub_plugins.plugin_image_producer" + ] + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + + # then + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_INPUTS_DERIVED_NESTED_DIMS_AND_EMERGED_NESTED_DIMS, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [dogs_image, crowd_image], + } + ) + + # then + assert ( + len(result) == 2 + ), "Two inputs provided, their dimensions survived to the output, hence 2 outputs expected" + assert len(result[0]["collapsed_input"]) == 2 + assert np.allclose(result[0]["collapsed_input"][0].numpy_image, dogs_image) + assert np.allclose(result[0]["collapsed_input"][1].numpy_image, crowd_image) + assert np.allclose(result[0]["input_image"].numpy_image, dogs_image) + assert result[0]["shapes"] == ["[192, 168, 3][292, 168, 3]"] + assert result[0]["collapsed_output"] == ["[192, 168, 3][292, 168, 3]"] + assert result[0]["collapsed_output_2"] == [["[192, 168, 3][292, 168, 3]"]] + # Under the tensor flag, breds_classification entries are native + # ClassificationPrediction objects (with None slots where EachSecondPass gated out + # the crop). The top class name lives at images_metadata[0][CLASS_NAMES_KEY] indexed + # by the top class id; the native subscript reproduces the same expected literals. + assert [ + ( + p.images_metadata[0][CLASS_NAMES_KEY][int(p.class_id.reshape(-1)[0])] + if p is not None + else None + ) + for p in result[0]["breds_classification"] + ] == ["116.Parson_russell_terrier", None] + assert len(result[1]["collapsed_input"]) == 2 + assert np.allclose(result[1]["collapsed_input"][0].numpy_image, dogs_image) + assert np.allclose(result[1]["collapsed_input"][1].numpy_image, crowd_image) + assert np.allclose(result[1]["input_image"].numpy_image, crowd_image) + assert result[1]["shapes"] == ["[192, 168, 3][292, 168, 3]"] + assert result[1]["collapsed_output"] == ["[192, 168, 3][292, 168, 3]"] + assert result[1]["collapsed_output_2"] == [["[192, 168, 3][292, 168, 3]"]] + assert result[1]["breds_classification"] == [] + + +@_TENSOR_ONLY +@mock.patch.object(blocks_loader, "get_plugin_modules") +def test_workflow_with_input_derived_dims_and_emergent_dims_with_tensor_input( + get_plugin_modules_mock: MagicMock, + model_manager: ModelManager, + dogs_image: np.ndarray, + crowd_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as test_workflow_with_input_derived_dims_and_emergent_dims_tensor_native, but the + # images arrive ALREADY materialised as CHW RGB device tensors (is_tensor_materialised() + # == True), so the model blocks run their on-device tensor path. + # given + get_plugin_modules_mock.return_value = [ + "tests.workflows.integration_tests.execution.stub_plugins.plugin_image_producer" + ] + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + + # then + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_INPUTS_DERIVED_NESTED_DIMS_AND_EMERGED_NESTED_DIMS, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(crowd_image), + ], + } + ) + + # then + assert ( + len(result) == 2 + ), "Two inputs provided, their dimensions survived to the output, hence 2 outputs expected" + assert len(result[0]["collapsed_input"]) == 2 + assert np.allclose(result[0]["collapsed_input"][0].numpy_image, dogs_image) + assert np.allclose(result[0]["collapsed_input"][1].numpy_image, crowd_image) + assert np.allclose(result[0]["input_image"].numpy_image, dogs_image) + assert result[0]["shapes"] == ["[192, 168, 3][292, 168, 3]"] + assert result[0]["collapsed_output"] == ["[192, 168, 3][292, 168, 3]"] + assert result[0]["collapsed_output_2"] == [["[192, 168, 3][292, 168, 3]"]] + # Under the tensor flag, breds_classification entries are native + # ClassificationPrediction objects (with None slots where EachSecondPass gated out + # the crop). The top class name lives at images_metadata[0][CLASS_NAMES_KEY] indexed + # by the top class id; the native subscript reproduces the same expected literals. + assert [ + ( + p.images_metadata[0][CLASS_NAMES_KEY][int(p.class_id.reshape(-1)[0])] + if p is not None + else None + ) + for p in result[0]["breds_classification"] + ] == ["116.Parson_russell_terrier", None] + assert len(result[1]["collapsed_input"]) == 2 + assert np.allclose(result[1]["collapsed_input"][0].numpy_image, dogs_image) + assert np.allclose(result[1]["collapsed_input"][1].numpy_image, crowd_image) + assert np.allclose(result[1]["input_image"].numpy_image, crowd_image) + assert result[1]["shapes"] == ["[192, 168, 3][292, 168, 3]"] + assert result[1]["collapsed_output"] == ["[192, 168, 3][292, 168, 3]"] + assert result[1]["collapsed_output_2"] == [["[192, 168, 3][292, 168, 3]"]] + assert result[1]["breds_classification"] == [] diff --git a/tests/workflows/integration_tests/execution/test_workflow_detection_plus_classification.py b/tests/workflows/integration_tests/execution/test_workflow_detection_plus_classification.py index 5efee690f4..26b0533277 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_detection_plus_classification.py +++ b/tests/workflows/integration_tests/execution/test_workflow_detection_plus_classification.py @@ -1,14 +1,36 @@ import numpy as np import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.classification import ClassificationPrediction +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the classification block emits a list of native +# inference_models.ClassificationPrediction objects instead of sv-shaped dicts. The +# numpy-shaped assertions below run only with the flag off; an equivalent *_tensor_native +# test asserts the same facts against the native carrier with the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="classification dict output; native ClassificationPrediction under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + LEGACY_DETECTION_PLUS_CLASSIFICATION_WORKFLOW = { "version": "1.0", "inputs": [{"type": "WorkflowImage", "name": "image"}], @@ -44,6 +66,7 @@ } +@_NUMPY_ONLY def test_legacy_detection_plus_classification_workflow_when_minimal_valid_input_provided( model_manager: ModelManager, dogs_image: np.ndarray, @@ -82,6 +105,113 @@ def test_legacy_detection_plus_classification_workflow_when_minimal_valid_input_ ], "Expected predictions to be as measured in reference run" +@_TENSOR_ONLY +def test_legacy_detection_plus_classification_workflow_when_minimal_valid_input_provided_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=LEGACY_DETECTION_PLUS_CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "predictions", + }, "Expected all declared outputs to be delivered" + # Under the flag, predictions is a Python list of native ClassificationPrediction + # objects (one per crop); list length is unchanged. + predictions = result[0]["predictions"] + assert ( + len(predictions) == 2 + ), "Expected 2 dogs crops on input image, hence 2 nested classification results" + for pred in predictions: + assert isinstance( + pred, ClassificationPrediction + ), "Each crop prediction must be a native ClassificationPrediction" + # Native equivalent of ['top']: class name at the top class id. Note images_metadata + # is PLURAL and indexed [0]; top class id comes from class_id. + top_names = [ + pred.images_metadata[0][CLASS_NAMES_KEY][int(pred.class_id.reshape(-1)[0])] + for pred in predictions + ] + assert top_names == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected predictions to be as measured in reference run" + + +@_TENSOR_ONLY +def test_legacy_detection_plus_classification_workflow_when_minimal_valid_input_provided_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as the *_tensor_native test, but the image arrives ALREADY materialised as a CHW + # RGB device tensor (is_tensor_materialised() == True), so the OD block runs its on-device + # tensor path. Results must match the numpy-input variant. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=LEGACY_DETECTION_PLUS_CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "predictions", + }, "Expected all declared outputs to be delivered" + # Under the flag, predictions is a Python list of native ClassificationPrediction + # objects (one per crop); list length is unchanged. + predictions = result[0]["predictions"] + assert ( + len(predictions) == 2 + ), "Expected 2 dogs crops on input image, hence 2 nested classification results" + for pred in predictions: + assert isinstance( + pred, ClassificationPrediction + ), "Each crop prediction must be a native ClassificationPrediction" + # Native equivalent of ['top']: class name at the top class id. Note images_metadata + # is PLURAL and indexed [0]; top class id comes from class_id. + top_names = [ + pred.images_metadata[0][CLASS_NAMES_KEY][int(pred.class_id.reshape(-1)[0])] + for pred in predictions + ] + assert top_names == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected predictions to be as measured in reference run" + + # Matched pairs only โ€” OD and classification blocks are bumped together. DETECTION_PLUS_CLASSIFICATION_BLOCK_PAIRS = [ ( @@ -148,6 +278,7 @@ def _build_detection_plus_classification_workflow( ) +@_NUMPY_ONLY @add_to_workflows_gallery( category="Workflows with multiple models", use_case_title="Workflow detection model followed by classifier", @@ -213,6 +344,127 @@ def test_detection_plus_classification_workflow_when_minimal_valid_input_provide ], "Expected predictions to be as measured in reference run" +@_TENSOR_ONLY +@pytest.mark.parametrize( + "od_block_type, cls_block_type", DETECTION_PLUS_CLASSIFICATION_BLOCK_PAIRS +) +def test_detection_plus_classification_workflow_when_minimal_valid_input_provided_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, + od_block_type: str, + cls_block_type: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_detection_plus_classification_workflow( + od_block_type, cls_block_type + ), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "predictions", + }, "Expected all declared outputs to be delivered" + # Under the flag, predictions is a Python list of native ClassificationPrediction + # objects (one per crop); list length is unchanged. + predictions = result[0]["predictions"] + assert ( + len(predictions) == 2 + ), "Expected 2 dogs crops on input image, hence 2 nested classification results" + for pred in predictions: + assert isinstance( + pred, ClassificationPrediction + ), "Each crop prediction must be a native ClassificationPrediction" + # Native equivalent of ['top']: class name at the top class id. Note images_metadata + # is PLURAL and indexed [0]; top class id comes from class_id. + top_names = [ + pred.images_metadata[0][CLASS_NAMES_KEY][int(pred.class_id.reshape(-1)[0])] + for pred in predictions + ] + assert top_names == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected predictions to be as measured in reference run" + + +@_TENSOR_ONLY +@pytest.mark.parametrize( + "od_block_type, cls_block_type", DETECTION_PLUS_CLASSIFICATION_BLOCK_PAIRS +) +def test_detection_plus_classification_workflow_when_minimal_valid_input_provided_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, + od_block_type: str, + cls_block_type: str, +) -> None: + # Same as the *_tensor_native test, but the image arrives ALREADY materialised as a CHW + # RGB device tensor (is_tensor_materialised() == True), so the OD block runs its on-device + # tensor path. Results must match the numpy-input variant. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_detection_plus_classification_workflow( + od_block_type, cls_block_type + ), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "predictions", + }, "Expected all declared outputs to be delivered" + # Under the flag, predictions is a Python list of native ClassificationPrediction + # objects (one per crop); list length is unchanged. + predictions = result[0]["predictions"] + assert ( + len(predictions) == 2 + ), "Expected 2 dogs crops on input image, hence 2 nested classification results" + for pred in predictions: + assert isinstance( + pred, ClassificationPrediction + ), "Each crop prediction must be a native ClassificationPrediction" + # Native equivalent of ['top']: class name at the top class id. Note images_metadata + # is PLURAL and indexed [0]; top class id comes from class_id. + top_names = [ + pred.images_metadata[0][CLASS_NAMES_KEY][int(pred.class_id.reshape(-1)[0])] + for pred in predictions + ] + assert top_names == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected predictions to be as measured in reference run" + + @pytest.mark.parametrize( "od_block_type, cls_block_type", DETECTION_PLUS_CLASSIFICATION_BLOCK_PAIRS ) diff --git a/tests/workflows/integration_tests/execution/test_workflow_inference_id_response.py b/tests/workflows/integration_tests/execution/test_workflow_inference_id_response.py index de0838cbc9..8d0384a63e 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_inference_id_response.py +++ b/tests/workflows/integration_tests/execution/test_workflow_inference_id_response.py @@ -1,10 +1,37 @@ import numpy as np import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + INFERENCE_ID_KEY, +) from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.classification import ClassificationPrediction +from inference_models.models.base.object_detection import Detections as NativeDetections +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the detection/classification blocks emit native +# inference_models dataclasses (Detections / ClassificationPrediction) instead of +# sv.Detections / sv-shaped dicts, so inference_id moves to image-level metadata. The +# numpy-shaped assertions below run only with the flag off; an equivalent *_tensor_native +# test asserts the same facts against the native carrier with the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections / sv-shaped output; native under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) DETECTION_PLUS_CLASSIFICATION_WORKFLOW = { "version": "1.0", @@ -83,6 +110,7 @@ } +@_NUMPY_ONLY @pytest.mark.workflows def test_detection_plus_classification_workflow_with_inference_id( model_manager: ModelManager, @@ -123,6 +151,114 @@ def test_detection_plus_classification_workflow_with_inference_id( ], "Expected predictions to be as measured in reference run" +@_TENSOR_ONLY +@pytest.mark.workflows +def test_detection_plus_classification_workflow_with_inference_id_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTION_PLUS_CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + } + ) + + # then + assert ( + len(result[0]["predictions"]) == 2 + ), "Expected 2 dogs crops on input image, hence 2 nested classification results" + + # Under the flag each nested classification result is a native + # ClassificationPrediction dataclass (not iterable). inference_id lives on + # image-level metadata at images_metadata[0][INFERENCE_ID_KEY]; the top class is + # resolved via images_metadata[0][CLASS_NAMES_KEY][top class_id]. + tops = [] + for prediction in result[0]["predictions"]: + assert isinstance( + prediction, ClassificationPrediction + ), "Expected native ClassificationPrediction under the flag" + meta = prediction.images_metadata[0] + assert INFERENCE_ID_KEY in meta, "Expected inference_id in image metadata" + assert meta[INFERENCE_ID_KEY] is not None, "Expected non-null inference_id" + top_class_id = int(prediction.class_id.reshape(-1)[0]) + tops.append(meta[CLASS_NAMES_KEY][top_class_id]) + + assert tops == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected predictions to be as measured in reference run" + + +@_TENSOR_ONLY +@pytest.mark.workflows +def test_detection_plus_classification_workflow_with_inference_id_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as test_detection_plus_classification_workflow_with_inference_id_tensor_native, + # but the image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the OD block runs its on-device tensor path. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTION_PLUS_CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + } + ) + + # then + assert ( + len(result[0]["predictions"]) == 2 + ), "Expected 2 dogs crops on input image, hence 2 nested classification results" + + # Under the flag each nested classification result is a native + # ClassificationPrediction dataclass (not iterable). inference_id lives on + # image-level metadata at images_metadata[0][INFERENCE_ID_KEY]; the top class is + # resolved via images_metadata[0][CLASS_NAMES_KEY][top class_id]. + tops = [] + for prediction in result[0]["predictions"]: + assert isinstance( + prediction, ClassificationPrediction + ), "Expected native ClassificationPrediction under the flag" + meta = prediction.images_metadata[0] + assert INFERENCE_ID_KEY in meta, "Expected inference_id in image metadata" + assert meta[INFERENCE_ID_KEY] is not None, "Expected non-null inference_id" + top_class_id = int(prediction.class_id.reshape(-1)[0]) + tops.append(meta[CLASS_NAMES_KEY][top_class_id]) + + assert tops == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected predictions to be as measured in reference run" + + +@_NUMPY_ONLY @pytest.mark.workflows def test_object_detection_workflow_with_inference_id( model_manager: ModelManager, @@ -158,6 +294,87 @@ def test_object_detection_workflow_with_inference_id( ), "Expected non-null inference_id" +@_TENSOR_ONLY +@pytest.mark.workflows +def test_object_detection_workflow_with_inference_id_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + } + ) + + # then + # Under the flag predictions is a native Detections carrier (no __getitem__). + # inference_id is image-level metadata shared by both boxes, not per-box. + dets = result[0]["predictions"] + assert isinstance( + dets, NativeDetections + ), "Expected native inference_models.Detections under the flag" + assert dets.xyxy.shape[0] == 2, "Expected 2 predictions" + assert ( + dets.image_metadata[INFERENCE_ID_KEY] is not None + ), "Expected non-null inference_id" + + +@_TENSOR_ONLY +@pytest.mark.workflows +def test_object_detection_workflow_with_inference_id_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as test_object_detection_workflow_with_inference_id_tensor_native, but the image + # arrives ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() == + # True), so the OD block runs its on-device tensor path. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + } + ) + + # then + # Under the flag predictions is a native Detections carrier (no __getitem__). + # inference_id is image-level metadata shared by both boxes, not per-box. + dets = result[0]["predictions"] + assert isinstance( + dets, NativeDetections + ), "Expected native inference_models.Detections under the flag" + assert dets.xyxy.shape[0] == 2, "Expected 2 predictions" + assert ( + dets.image_metadata[INFERENCE_ID_KEY] is not None + ), "Expected non-null inference_id" + + @pytest.mark.workflows def test_instance_segmentation_workflow_with_inference_id( model_manager: ModelManager, diff --git a/tests/workflows/integration_tests/execution/test_workflow_json_parser_config.py b/tests/workflows/integration_tests/execution/test_workflow_json_parser_config.py index cfed8003e5..51b788f27d 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_json_parser_config.py +++ b/tests/workflows/integration_tests/execution/test_workflow_json_parser_config.py @@ -1,10 +1,32 @@ import numpy as np +import pytest import supervision as sv -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.object_detection import Detections as NativeDetections +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the object detection model block emits a +# native inference_models.Detections instead of sv.Detections. The sv-typed test below +# is skipped when the flag is on; the `*_tensor_native` parity test (skipped when the +# flag is off) asserts the same workflow result expressed as the native dataclass. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output; model block emits native Detections under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) JSON_PARSER_WORKFLOW = { "version": "1.0", @@ -45,6 +67,7 @@ } +@_NUMPY_ONLY def test_workflow_with_json_parameter( model_manager: ModelManager, dogs_image: np.ndarray, @@ -71,3 +94,77 @@ def test_workflow_with_json_parameter( assert set(result[0].keys()) == {"json_parser", "model_predictions"} assert result[0]["json_parser"] == "yolov8n-640" assert isinstance(result[0]["model_predictions"], sv.Detections) + + +@_TENSOR_ONLY +def test_workflow_with_json_parameter_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=JSON_PARSER_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "config": '{"model_id": "yolov8n-640"}', + } + ) + + assert len(result) == 1 + assert set(result[0].keys()) == {"json_parser", "model_predictions"} + assert result[0]["json_parser"] == "yolov8n-640" + # Under ENABLE_TENSOR_DATA_REPRESENTATION the model block emits a native + # inference_models.Detections (torch-backed) rather than sv.Detections. + predictions = result[0]["model_predictions"] + assert isinstance(predictions, NativeDetections) + # Same semantic result as the numpy path: two dogs detected (COCO class_id 16). + assert len(predictions.xyxy) == 2 + assert predictions.class_id.tolist() == [16, 16] + + +@_TENSOR_ONLY +def test_workflow_with_json_parameter_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + # Same as test_workflow_with_json_parameter_tensor_native, but the image arrives + # ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() == True), + # so the OD model block runs its on-device tensor path. Results must match the + # numpy-input variant. + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=JSON_PARSER_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "config": '{"model_id": "yolov8n-640"}', + } + ) + + assert len(result) == 1 + assert set(result[0].keys()) == {"json_parser", "model_predictions"} + assert result[0]["json_parser"] == "yolov8n-640" + # Under ENABLE_TENSOR_DATA_REPRESENTATION the model block emits a native + # inference_models.Detections (torch-backed) rather than sv.Detections. + predictions = result[0]["model_predictions"] + assert isinstance(predictions, NativeDetections) + # Same semantic result as the numpy path: two dogs detected (COCO class_id 16). + assert len(predictions.xyxy) == 2 + assert predictions.class_id.tolist() == [16, 16] diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_active_learning_sink.py b/tests/workflows/integration_tests/execution/test_workflow_with_active_learning_sink.py index d5ce4607ac..77daeb3939 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_active_learning_sink.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_active_learning_sink.py @@ -1,14 +1,37 @@ import numpy as np import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.classification import ClassificationPrediction +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the classification block emits native +# inference_models.ClassificationPrediction objects instead of the sv-shaped +# classification dicts (with a string "top" key). The numpy-shaped assertions below +# run only with the flag off; an equivalent *_tensor_native test asserts the same +# facts against the native carrier with the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="classification dict output; native ClassificationPrediction under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + ACTIVE_LEARNING_WORKFLOW = { "version": "1.0", "inputs": [ @@ -86,6 +109,7 @@ } +@_NUMPY_ONLY @add_to_workflows_gallery( category="Workflows enhanced by Roboflow Platform", use_case_title="Data Collection for Active Learning", @@ -142,6 +166,124 @@ def test_detection_plus_classification_workflow_when_minimal_valid_input_provide ), "Expected data not registered due to sampling" +@_TENSOR_ONLY +def test_detection_plus_classification_workflow_when_minimal_valid_input_provided_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=ACTIVE_LEARNING_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "data_percentage": 0.0, + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "predictions", + "registration_message", + }, "Expected all declared outputs to be delivered" + assert ( + len(result[0]["predictions"]) == 2 + ), "Expected 2 dogs crops on input image, hence 2 nested classification results" + # Under ENABLE_TENSOR_DATA_REPRESENTATION each per-crop prediction is a native + # inference_models.ClassificationPrediction. The top-class STRING the numpy test + # reads at predictions[i]["top"] is recovered here from the native carrier: + # top class id = int(class_id.reshape(-1)[0]); the class_id -> name mapping lives + # in images_metadata[0][CLASS_NAMES_KEY] (PLURAL, indexed [0]). + top_class_names = [] + for prediction in result[0]["predictions"]: + assert isinstance( + prediction, ClassificationPrediction + ), "Expected native inference_models.ClassificationPrediction" + top_id = int(prediction.class_id.reshape(-1)[0]) + names = prediction.images_metadata[0][CLASS_NAMES_KEY] + top_class_names.append(names[top_id]) + assert top_class_names == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected predictions to be as measured in reference run" + assert ( + result[0]["registration_message"] + == ["Registration skipped due to sampling settings"] * 2 + ), "Expected data not registered due to sampling" + + +@_TENSOR_ONLY +def test_detection_plus_classification_workflow_when_minimal_valid_input_provided_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as + # test_detection_plus_classification_workflow_when_minimal_valid_input_provided_tensor_native, + # but the image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the OD block runs its on-device tensor path. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=ACTIVE_LEARNING_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "data_percentage": 0.0, + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "predictions", + "registration_message", + }, "Expected all declared outputs to be delivered" + assert ( + len(result[0]["predictions"]) == 2 + ), "Expected 2 dogs crops on input image, hence 2 nested classification results" + # Under ENABLE_TENSOR_DATA_REPRESENTATION each per-crop prediction is a native + # inference_models.ClassificationPrediction. The top-class STRING the numpy test + # reads at predictions[i]["top"] is recovered here from the native carrier: + # top class id = int(class_id.reshape(-1)[0]); the class_id -> name mapping lives + # in images_metadata[0][CLASS_NAMES_KEY] (PLURAL, indexed [0]). + top_class_names = [] + for prediction in result[0]["predictions"]: + assert isinstance( + prediction, ClassificationPrediction + ), "Expected native inference_models.ClassificationPrediction" + top_id = int(prediction.class_id.reshape(-1)[0]) + names = prediction.images_metadata[0][CLASS_NAMES_KEY] + top_class_names.append(names[top_id]) + assert top_class_names == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected predictions to be as measured in reference run" + assert ( + result[0]["registration_message"] + == ["Registration skipped due to sampling settings"] * 2 + ), "Expected data not registered due to sampling" + + def test_detection_plus_classification_workflow_when_nothing_to_be_registered( model_manager: ModelManager, crowd_image: np.ndarray, diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_arbitrary_batch_inputs.py b/tests/workflows/integration_tests/execution/test_workflow_with_arbitrary_batch_inputs.py index 388631197c..f334a71102 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_arbitrary_batch_inputs.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_arbitrary_batch_inputs.py @@ -5,7 +5,10 @@ import pytest import supervision as sv -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.utils.image_utils import load_image from inference.core.workflows.core_steps.common.entities import StepExecutionMode @@ -16,6 +19,41 @@ ) from inference.core.workflows.execution_engine.core import ExecutionEngine from inference.core.workflows.execution_engine.introspection import blocks_loader +from inference_models.models.base.object_detection import Detections as NativeDetections +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the ObjectDetectionModel emits a native +# `inference_models.Detections` and the ClassificationModel emits native +# `ClassificationPrediction` objects (per-image, single-row: class_id -> (1,), +# confidence -> (1, num_classes) full softmax) instead of sv.Detections / inference +# classification dicts. The sv/dict-shaped tests below are skipped when the flag is +# on; each has a `*_tensor_native` parity test (skipped when the flag is off) that +# asserts the same semantic result expressed for the native objects. Serialized +# branches still yield dicts and are left unchanged. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections / classification-dict reads; models are native under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + + +def _native_top_class_id(prediction) -> int: + # Native per-image ClassificationPrediction: class_id -> (1,) holding the top + # class id (argmax over the softmax row). Equivalent to sv's `p["top"]` class. + return int(prediction.class_id[0]) + + +def _native_top_confidence(prediction) -> float: + # confidence -> (1, num_classes) full softmax; the top-class confidence is the + # row max, equivalent to sv's `p["confidence"]`. + return float(prediction.confidence[0].max()) + TWO_STAGE_WORKFLOW = { "version": "1.3.0", @@ -128,6 +166,7 @@ } +@_NUMPY_ONLY def test_debug_execution_of_workflow_for_single_image_without_conditional_evaluation( model_manager: ModelManager, dogs_image: np.ndarray, @@ -198,6 +237,161 @@ def test_debug_execution_of_workflow_for_single_image_without_conditional_evalua ), "Expected confidences from step-by-step execution to match e2e execution" +@_TENSOR_ONLY +def test_debug_execution_of_workflow_for_single_image_without_conditional_evaluation_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + end_to_end_execution_engine = ExecutionEngine.init( + workflow_definition=TWO_STAGE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + first_step_execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + second_step_execution_engine = ExecutionEngine.init( + workflow_definition=CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + third_step_execution_engine = ExecutionEngine.init( + workflow_definition=CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + e2e_results = end_to_end_execution_engine.run( + runtime_parameters={ + "image": dogs_image, + } + ) + detection_results = first_step_execution_engine.run( + runtime_parameters={ + "image": dogs_image, + } + ) + cropping_results = second_step_execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "predictions": detection_results[0]["result"]["predictions"], + } + ) + classification_results = third_step_execution_engine.run( + runtime_parameters={ + "crops": [[e["crops"] for e in cropping_results[0]["result"]]], + } + ) + + # then + # Native per-image ClassificationPrediction objects: top class via class_id, + # top confidence via the per-image max over the full softmax row. + e2e_top_classes = [_native_top_class_id(p) for p in e2e_results[0]["predictions"]] + debug_top_classes = [ + _native_top_class_id(p) for p in classification_results[0]["predictions"] + ] + assert ( + e2e_top_classes == debug_top_classes + ), "Expected top class prediction from step-by-step execution to match e2e execution" + e2e_confidence = [_native_top_confidence(p) for p in e2e_results[0]["predictions"]] + debug_confidence = [ + _native_top_confidence(p) for p in classification_results[0]["predictions"] + ] + assert np.allclose( + e2e_confidence, debug_confidence, atol=1e-4 + ), "Expected confidences from step-by-step execution to match e2e execution" + + +@_TENSOR_ONLY +def test_debug_execution_of_workflow_for_single_image_without_conditional_evaluation_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as + # test_debug_execution_of_workflow_for_single_image_without_conditional_evaluation_tensor_native, + # but each image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the model blocks run their on-device tensor path. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + end_to_end_execution_engine = ExecutionEngine.init( + workflow_definition=TWO_STAGE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + first_step_execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + second_step_execution_engine = ExecutionEngine.init( + workflow_definition=CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + third_step_execution_engine = ExecutionEngine.init( + workflow_definition=CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + e2e_results = end_to_end_execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + } + ) + detection_results = first_step_execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + } + ) + cropping_results = second_step_execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "predictions": detection_results[0]["result"]["predictions"], + } + ) + classification_results = third_step_execution_engine.run( + runtime_parameters={ + "crops": [[e["crops"] for e in cropping_results[0]["result"]]], + } + ) + + # then + # Native per-image ClassificationPrediction objects: top class via class_id, + # top confidence via the per-image max over the full softmax row. + e2e_top_classes = [_native_top_class_id(p) for p in e2e_results[0]["predictions"]] + debug_top_classes = [ + _native_top_class_id(p) for p in classification_results[0]["predictions"] + ] + assert ( + e2e_top_classes == debug_top_classes + ), "Expected top class prediction from step-by-step execution to match e2e execution" + e2e_confidence = [_native_top_confidence(p) for p in e2e_results[0]["predictions"]] + debug_confidence = [ + _native_top_confidence(p) for p in classification_results[0]["predictions"] + ] + assert np.allclose( + e2e_confidence, debug_confidence, atol=1e-4 + ), "Expected confidences from step-by-step execution to match e2e execution" + + +@_NUMPY_ONLY def test_debug_execution_of_workflow_for_single_image_without_conditional_evaluation_when_serialization_is_requested( model_manager: ModelManager, dogs_image: np.ndarray, @@ -322,7 +516,8 @@ def test_debug_execution_of_workflow_for_single_image_without_conditional_evalua ), "Expected confidences from step-by-step execution to match e2e execution" -def test_debug_execution_of_workflow_for_batch_of_images_without_conditional_evaluation( +@_TENSOR_ONLY +def test_debug_execution_of_workflow_for_single_image_without_conditional_evaluation_when_serialization_is_requested_tensor_native( model_manager: ModelManager, dogs_image: np.ndarray, roboflow_api_key: str, @@ -357,76 +552,531 @@ def test_debug_execution_of_workflow_for_batch_of_images_without_conditional_eva # when e2e_results = end_to_end_execution_engine.run( runtime_parameters={ - "image": [dogs_image, dogs_image], - } + "image": dogs_image, + }, + serialize_results=True, ) detection_results = first_step_execution_engine.run( runtime_parameters={ - "image": [dogs_image, dogs_image], - } + "image": dogs_image, + }, + serialize_results=True, + ) + detection_results_not_serialized = first_step_execution_engine.run( + runtime_parameters={ + "image": dogs_image, + }, ) cropping_results = second_step_execution_engine.run( runtime_parameters={ - "image": [dogs_image, dogs_image], - "predictions": [ - detection_results[0]["result"]["predictions"], - detection_results[1]["result"]["predictions"], - ], - } + "image": dogs_image, + "predictions": detection_results[0]["result"]["predictions"], + }, + serialize_results=True, + ) + cropping_results_not_serialized = second_step_execution_engine.run( + runtime_parameters={ + "image": dogs_image, + "predictions": detection_results_not_serialized[0]["result"]["predictions"], + }, + serialize_results=False, ) classification_results = third_step_execution_engine.run( runtime_parameters={ - "crops": [ - [e["crops"] for e in cropping_results[0]["result"]], - [e["crops"] for e in cropping_results[1]["result"]], - ], - } + "crops": [[e["crops"] for e in cropping_results[0]["result"]]], + }, + serialize_results=True, ) # then - e2e_top_classes = [p["top"] for r in e2e_results for p in r["predictions"]] - debug_top_classes = [ - p["top"] for r in classification_results for p in r["predictions"] - ] + # Serialized branch still yields a dict (unchanged); the NOT-serialized branch + # yields a native inference_models.Detections under the flag. + assert isinstance( + detection_results[0]["result"]["predictions"], dict + ), "Expected sv.Detections to be serialized" + assert isinstance( + detection_results_not_serialized[0]["result"]["predictions"], NativeDetections + ), "Expected native Detections not to be serialized" + deserialized_detections = sv.Detections.from_inference( + detection_results[0]["result"]["predictions"] + ) + assert np.allclose( + deserialized_detections.confidence, + detection_results_not_serialized[0]["result"]["predictions"] + .confidence.cpu() + .numpy(), + atol=1e-4, + ), "Expected confidence match when serialized detections are deserialized" + intermediate_crop = cropping_results[0]["result"][0]["crops"] + assert ( + intermediate_crop["type"] == "base64" + ), "Expected crop to be serialized to base64" + decoded_image, _ = load_image(intermediate_crop) + number_of_pixels = ( + decoded_image.shape[0] * decoded_image.shape[1] * decoded_image.shape[2] + ) + assert ( + decoded_image.shape + == cropping_results_not_serialized[0]["result"][0]["crops"].numpy_image.shape + ), "Expected deserialized crop to match in size with not serialized one" + assert ( + abs( + (decoded_image.sum() / number_of_pixels) + - ( + cropping_results_not_serialized[0]["result"][0][ + "crops" + ].numpy_image.sum() + / number_of_pixels + ) + ) + < 1e-1 + ), "Content of serialized and not serialized crop should roughly match (up to compression)" + # Both e2e and debug classification ran with serialize_results=True, so these + # remain inference classification dicts (unchanged from the sv-path test). + e2e_top_classes = [p["top"] for p in e2e_results[0]["predictions"]] + debug_top_classes = [p["top"] for p in classification_results[0]["predictions"]] assert ( e2e_top_classes == debug_top_classes ), "Expected top class prediction from step-by-step execution to match e2e execution" - e2e_confidence = [p["confidence"] for r in e2e_results for p in r["predictions"]] + e2e_confidence = [p["confidence"] for p in e2e_results[0]["predictions"]] debug_confidence = [ - p["confidence"] for r in classification_results for p in r["predictions"] + p["confidence"] for p in classification_results[0]["predictions"] ] assert np.allclose( - e2e_confidence, debug_confidence, atol=1e-4 + e2e_confidence, debug_confidence, atol=1e-1 ), "Expected confidences from step-by-step execution to match e2e execution" -TWO_STAGE_WORKFLOW_WITH_FLOW_CONTROL = { - "version": "1.3.0", - "inputs": [{"type": "WorkflowImage", "name": "image"}], - "steps": [ - { - "type": "ObjectDetectionModel", - "name": "general_detection", - "image": "$inputs.image", - "model_id": "yolov8n-640", - "class_filter": ["dog"], - }, - { - "type": "Crop", - "name": "cropping", - "image": "$inputs.image", - "predictions": "$steps.general_detection.predictions", - }, - { - "type": "roboflow_core/continue_if@v1", - "name": "verify_crop_size", - "condition_statement": { - "type": "StatementGroup", - "statements": [ - { - "type": "BinaryStatement", - "left_operand": { - "type": "DynamicOperand", +@_TENSOR_ONLY +def test_debug_execution_of_workflow_for_single_image_without_conditional_evaluation_when_serialization_is_requested_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as + # test_debug_execution_of_workflow_for_single_image_without_conditional_evaluation_when_serialization_is_requested_tensor_native, + # but each image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the model blocks run their on-device tensor path. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + end_to_end_execution_engine = ExecutionEngine.init( + workflow_definition=TWO_STAGE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + first_step_execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + second_step_execution_engine = ExecutionEngine.init( + workflow_definition=CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + third_step_execution_engine = ExecutionEngine.init( + workflow_definition=CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + e2e_results = end_to_end_execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + }, + serialize_results=True, + ) + detection_results = first_step_execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + }, + serialize_results=True, + ) + detection_results_not_serialized = first_step_execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + }, + ) + cropping_results = second_step_execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "predictions": detection_results[0]["result"]["predictions"], + }, + serialize_results=True, + ) + cropping_results_not_serialized = second_step_execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + "predictions": detection_results_not_serialized[0]["result"]["predictions"], + }, + serialize_results=False, + ) + classification_results = third_step_execution_engine.run( + runtime_parameters={ + "crops": [[e["crops"] for e in cropping_results[0]["result"]]], + }, + serialize_results=True, + ) + + # then + # Serialized branch still yields a dict (unchanged); the NOT-serialized branch + # yields a native inference_models.Detections under the flag. + assert isinstance( + detection_results[0]["result"]["predictions"], dict + ), "Expected sv.Detections to be serialized" + assert isinstance( + detection_results_not_serialized[0]["result"]["predictions"], NativeDetections + ), "Expected native Detections not to be serialized" + deserialized_detections = sv.Detections.from_inference( + detection_results[0]["result"]["predictions"] + ) + assert np.allclose( + deserialized_detections.confidence, + detection_results_not_serialized[0]["result"]["predictions"] + .confidence.cpu() + .numpy(), + atol=1e-4, + ), "Expected confidence match when serialized detections are deserialized" + intermediate_crop = cropping_results[0]["result"][0]["crops"] + assert ( + intermediate_crop["type"] == "base64" + ), "Expected crop to be serialized to base64" + decoded_image, _ = load_image(intermediate_crop) + number_of_pixels = ( + decoded_image.shape[0] * decoded_image.shape[1] * decoded_image.shape[2] + ) + assert ( + decoded_image.shape + == cropping_results_not_serialized[0]["result"][0]["crops"].numpy_image.shape + ), "Expected deserialized crop to match in size with not serialized one" + assert ( + abs( + (decoded_image.sum() / number_of_pixels) + - ( + cropping_results_not_serialized[0]["result"][0][ + "crops" + ].numpy_image.sum() + / number_of_pixels + ) + ) + < 1e-1 + ), "Content of serialized and not serialized crop should roughly match (up to compression)" + # Both e2e and debug classification ran with serialize_results=True, so these + # remain inference classification dicts (unchanged from the sv-path test). + e2e_top_classes = [p["top"] for p in e2e_results[0]["predictions"]] + debug_top_classes = [p["top"] for p in classification_results[0]["predictions"]] + assert ( + e2e_top_classes == debug_top_classes + ), "Expected top class prediction from step-by-step execution to match e2e execution" + e2e_confidence = [p["confidence"] for p in e2e_results[0]["predictions"]] + debug_confidence = [ + p["confidence"] for p in classification_results[0]["predictions"] + ] + assert np.allclose( + e2e_confidence, debug_confidence, atol=1e-1 + ), "Expected confidences from step-by-step execution to match e2e execution" + + +@_NUMPY_ONLY +def test_debug_execution_of_workflow_for_batch_of_images_without_conditional_evaluation( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + end_to_end_execution_engine = ExecutionEngine.init( + workflow_definition=TWO_STAGE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + first_step_execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + second_step_execution_engine = ExecutionEngine.init( + workflow_definition=CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + third_step_execution_engine = ExecutionEngine.init( + workflow_definition=CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + e2e_results = end_to_end_execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + } + ) + detection_results = first_step_execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + } + ) + cropping_results = second_step_execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + "predictions": [ + detection_results[0]["result"]["predictions"], + detection_results[1]["result"]["predictions"], + ], + } + ) + classification_results = third_step_execution_engine.run( + runtime_parameters={ + "crops": [ + [e["crops"] for e in cropping_results[0]["result"]], + [e["crops"] for e in cropping_results[1]["result"]], + ], + } + ) + + # then + e2e_top_classes = [p["top"] for r in e2e_results for p in r["predictions"]] + debug_top_classes = [ + p["top"] for r in classification_results for p in r["predictions"] + ] + assert ( + e2e_top_classes == debug_top_classes + ), "Expected top class prediction from step-by-step execution to match e2e execution" + e2e_confidence = [p["confidence"] for r in e2e_results for p in r["predictions"]] + debug_confidence = [ + p["confidence"] for r in classification_results for p in r["predictions"] + ] + assert np.allclose( + e2e_confidence, debug_confidence, atol=1e-4 + ), "Expected confidences from step-by-step execution to match e2e execution" + + +@_TENSOR_ONLY +def test_debug_execution_of_workflow_for_batch_of_images_without_conditional_evaluation_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + end_to_end_execution_engine = ExecutionEngine.init( + workflow_definition=TWO_STAGE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + first_step_execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + second_step_execution_engine = ExecutionEngine.init( + workflow_definition=CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + third_step_execution_engine = ExecutionEngine.init( + workflow_definition=CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + e2e_results = end_to_end_execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + } + ) + detection_results = first_step_execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + } + ) + cropping_results = second_step_execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + "predictions": [ + detection_results[0]["result"]["predictions"], + detection_results[1]["result"]["predictions"], + ], + } + ) + classification_results = third_step_execution_engine.run( + runtime_parameters={ + "crops": [ + [e["crops"] for e in cropping_results[0]["result"]], + [e["crops"] for e in cropping_results[1]["result"]], + ], + } + ) + + # then + # Native per-image ClassificationPrediction objects across the flattened batch. + e2e_top_classes = [ + _native_top_class_id(p) for r in e2e_results for p in r["predictions"] + ] + debug_top_classes = [ + _native_top_class_id(p) + for r in classification_results + for p in r["predictions"] + ] + assert ( + e2e_top_classes == debug_top_classes + ), "Expected top class prediction from step-by-step execution to match e2e execution" + e2e_confidence = [ + _native_top_confidence(p) for r in e2e_results for p in r["predictions"] + ] + debug_confidence = [ + _native_top_confidence(p) + for r in classification_results + for p in r["predictions"] + ] + assert np.allclose( + e2e_confidence, debug_confidence, atol=1e-4 + ), "Expected confidences from step-by-step execution to match e2e execution" + + +@_TENSOR_ONLY +def test_debug_execution_of_workflow_for_batch_of_images_without_conditional_evaluation_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as + # test_debug_execution_of_workflow_for_batch_of_images_without_conditional_evaluation_tensor_native, + # but each image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the model blocks run their on-device tensor path. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + end_to_end_execution_engine = ExecutionEngine.init( + workflow_definition=TWO_STAGE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + first_step_execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + second_step_execution_engine = ExecutionEngine.init( + workflow_definition=CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + third_step_execution_engine = ExecutionEngine.init( + workflow_definition=CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + e2e_results = end_to_end_execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(dogs_image), + ], + } + ) + detection_results = first_step_execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(dogs_image), + ], + } + ) + cropping_results = second_step_execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(dogs_image), + ], + "predictions": [ + detection_results[0]["result"]["predictions"], + detection_results[1]["result"]["predictions"], + ], + } + ) + classification_results = third_step_execution_engine.run( + runtime_parameters={ + "crops": [ + [e["crops"] for e in cropping_results[0]["result"]], + [e["crops"] for e in cropping_results[1]["result"]], + ], + } + ) + + # then + # Native per-image ClassificationPrediction objects across the flattened batch. + e2e_top_classes = [ + _native_top_class_id(p) for r in e2e_results for p in r["predictions"] + ] + debug_top_classes = [ + _native_top_class_id(p) + for r in classification_results + for p in r["predictions"] + ] + assert ( + e2e_top_classes == debug_top_classes + ), "Expected top class prediction from step-by-step execution to match e2e execution" + e2e_confidence = [ + _native_top_confidence(p) for r in e2e_results for p in r["predictions"] + ] + debug_confidence = [ + _native_top_confidence(p) + for r in classification_results + for p in r["predictions"] + ] + assert np.allclose( + e2e_confidence, debug_confidence, atol=1e-4 + ), "Expected confidences from step-by-step execution to match e2e execution" + + +TWO_STAGE_WORKFLOW_WITH_FLOW_CONTROL = { + "version": "1.3.0", + "inputs": [{"type": "WorkflowImage", "name": "image"}], + "steps": [ + { + "type": "ObjectDetectionModel", + "name": "general_detection", + "image": "$inputs.image", + "model_id": "yolov8n-640", + "class_filter": ["dog"], + }, + { + "type": "Crop", + "name": "cropping", + "image": "$inputs.image", + "predictions": "$steps.general_detection.predictions", + }, + { + "type": "roboflow_core/continue_if@v1", + "name": "verify_crop_size", + "condition_statement": { + "type": "StatementGroup", + "statements": [ + { + "type": "BinaryStatement", + "left_operand": { + "type": "DynamicOperand", "operand_name": "crops", "operations": [ { @@ -463,6 +1113,7 @@ def test_debug_execution_of_workflow_for_batch_of_images_without_conditional_eva } +@_NUMPY_ONLY def test_debug_execution_of_workflow_for_batch_of_images_with_conditional_evaluation( model_manager: ModelManager, dogs_image: np.ndarray, @@ -561,6 +1212,227 @@ def test_debug_execution_of_workflow_for_batch_of_images_with_conditional_evalua ), "Expected confidences from step-by-step execution to match e2e execution" +@_TENSOR_ONLY +def test_debug_execution_of_workflow_for_batch_of_images_with_conditional_evaluation_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + end_to_end_execution_engine = ExecutionEngine.init( + workflow_definition=TWO_STAGE_WORKFLOW_WITH_FLOW_CONTROL, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + first_step_execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + second_step_execution_engine = ExecutionEngine.init( + workflow_definition=CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + third_step_execution_engine = ExecutionEngine.init( + workflow_definition=CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + e2e_results = end_to_end_execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + } + ) + detection_results = first_step_execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + } + ) + cropping_results = second_step_execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + "predictions": [ + detection_results[0]["result"]["predictions"], + detection_results[1]["result"]["predictions"], + ], + } + ) + classification_results = third_step_execution_engine.run( + runtime_parameters={ + "crops": [ + [cropping_results[0]["result"][0]["crops"], None], + [cropping_results[1]["result"][0]["crops"], None], + ], + } + ) + + # then + assert ( + e2e_results[0]["predictions"][0] is not None + ), "Expected first dog crop not to be excluded by conditional eval" + assert ( + e2e_results[0]["predictions"][1] is None + ), "Expected second dog crop to be excluded by conditional eval" + assert ( + e2e_results[1]["predictions"][0] is not None + ), "Expected first dog crop not to be excluded by conditional eval" + assert ( + e2e_results[1]["predictions"][1] is None + ), "Expected second dog crop to be excluded by conditional eval" + # Native per-image ClassificationPrediction objects; excluded entries stay None. + e2e_top_classes = [ + _native_top_class_id(p) if p else None + for r in e2e_results + for p in r["predictions"] + ] + debug_top_classes = [ + _native_top_class_id(p) if p else None + for r in classification_results + for p in r["predictions"] + ] + assert ( + e2e_top_classes == debug_top_classes + ), "Expected top class prediction from step-by-step execution to match e2e execution" + e2e_confidence = [ + _native_top_confidence(p) if p else -1000.0 + for r in e2e_results + for p in r["predictions"] + ] + debug_confidence = [ + _native_top_confidence(p) if p else -1000.0 + for r in classification_results + for p in r["predictions"] + ] + assert np.allclose( + e2e_confidence, debug_confidence, atol=1e-4 + ), "Expected confidences from step-by-step execution to match e2e execution" + + +@_TENSOR_ONLY +def test_debug_execution_of_workflow_for_batch_of_images_with_conditional_evaluation_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as + # test_debug_execution_of_workflow_for_batch_of_images_with_conditional_evaluation_tensor_native, + # but each image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the model blocks run their on-device tensor path. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + end_to_end_execution_engine = ExecutionEngine.init( + workflow_definition=TWO_STAGE_WORKFLOW_WITH_FLOW_CONTROL, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + first_step_execution_engine = ExecutionEngine.init( + workflow_definition=OBJECT_DETECTION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + second_step_execution_engine = ExecutionEngine.init( + workflow_definition=CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + third_step_execution_engine = ExecutionEngine.init( + workflow_definition=CLASSIFICATION_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + e2e_results = end_to_end_execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(dogs_image), + ], + } + ) + detection_results = first_step_execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(dogs_image), + ], + } + ) + cropping_results = second_step_execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(dogs_image), + ], + "predictions": [ + detection_results[0]["result"]["predictions"], + detection_results[1]["result"]["predictions"], + ], + } + ) + classification_results = third_step_execution_engine.run( + runtime_parameters={ + "crops": [ + [cropping_results[0]["result"][0]["crops"], None], + [cropping_results[1]["result"][0]["crops"], None], + ], + } + ) + + # then + assert ( + e2e_results[0]["predictions"][0] is not None + ), "Expected first dog crop not to be excluded by conditional eval" + assert ( + e2e_results[0]["predictions"][1] is None + ), "Expected second dog crop to be excluded by conditional eval" + assert ( + e2e_results[1]["predictions"][0] is not None + ), "Expected first dog crop not to be excluded by conditional eval" + assert ( + e2e_results[1]["predictions"][1] is None + ), "Expected second dog crop to be excluded by conditional eval" + # Native per-image ClassificationPrediction objects; excluded entries stay None. + e2e_top_classes = [ + _native_top_class_id(p) if p else None + for r in e2e_results + for p in r["predictions"] + ] + debug_top_classes = [ + _native_top_class_id(p) if p else None + for r in classification_results + for p in r["predictions"] + ] + assert ( + e2e_top_classes == debug_top_classes + ), "Expected top class prediction from step-by-step execution to match e2e execution" + e2e_confidence = [ + _native_top_confidence(p) if p else -1000.0 + for r in e2e_results + for p in r["predictions"] + ] + debug_confidence = [ + _native_top_confidence(p) if p else -1000.0 + for r in classification_results + for p in r["predictions"] + ] + assert np.allclose( + e2e_confidence, debug_confidence, atol=1e-4 + ), "Expected confidences from step-by-step execution to match e2e execution" + + def test_debug_execution_when_empty_batch_oriented_input_provided( model_manager: ModelManager, dogs_image: np.ndarray, diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_bounding_rectangle.py b/tests/workflows/integration_tests/execution/test_workflow_with_bounding_rectangle.py index d1ee443c9a..5c667f88c1 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_bounding_rectangle.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_bounding_rectangle.py @@ -2,15 +2,39 @@ import pytest import supervision as sv -from inference.core.env import USE_INFERENCE_MODELS, WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + USE_INFERENCE_MODELS, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.errors import RuntimeInputError, StepExecutionError from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.instance_segmentation import ( + InstanceDetections as NativeInstanceDetections, +) +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the bounding_rect block emits a native +# inference_models.InstanceDetections instead of sv.Detections. The numpy-shaped +# assertions below run only with the flag off; an equivalent *_tensor_native test +# asserts the same facts against the native carrier with the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output; native under ENABLE_TENSOR_DATA_REPRESENTATION " + "โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + BOUNDNG_RECTANGLE_WORKFLOW = { "version": "1.0", "inputs": [ @@ -39,6 +63,7 @@ } +@_NUMPY_ONLY @add_to_workflows_gallery( category="Basic Workflows", use_case_title="Workflow with bounding rect", @@ -108,3 +133,165 @@ def test_rectangle_bounding_workflow( result[0]["result"]["height"], np.array([178.4, 135.2]), atol=6.0 ) assert np.allclose(result[0]["result"]["angle"], np.array([0.826, 79.5]), atol=0.5) + + +@_TENSOR_ONLY +def test_rectangle_bounding_workflow_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=BOUNDNG_RECTANGLE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [dogs_image], + } + ) + + # then + assert len(result) == 1, "One set ot outputs expected" + assert "result" in result[0], "Output must contain key 'result'" + detections = result[0]["result"] + assert isinstance( + detections, NativeInstanceDetections + ), "Output must be instance of native inference_models.InstanceDetections" + assert detections.xyxy.shape[0] == 2, "Two dogs on the image" + + # On the native carrier the rect/width/height/angle geometry that lived in the + # vectorized sv.Detections.data columns is stored per-box in bboxes_metadata[i]. + assert detections.bboxes_metadata is not None, "Should have per-box metadata" + for i in range(2): + meta = detections.bboxes_metadata[i] + assert "rect" in meta, "'rect' geometry must be found in per-box metadata" + assert "width" in meta, "'width' geometry must be found in per-box metadata" + assert "height" in meta, "'height' geometry must be found in per-box metadata" + assert "angle" in meta, "'angle' geometry must be found in per-box metadata" + + assert np.allclose( + np.array(detections.bboxes_metadata[0]["rect"]), + np.array([[322.0, 402.0], [325.0, 224.0], [586.0, 228.0], [583.0, 406.0]]), + atol=5.0, + ) + assert np.allclose( + np.array(detections.bboxes_metadata[1]["rect"]), + np.array([[219.0, 82.0], [352.0, 57.0], [409.0, 363.0], [276.0, 388.0]]), + atol=6.0, + ) + assert np.allclose( + [ + detections.bboxes_metadata[0]["width"], + detections.bboxes_metadata[1]["width"], + ], + np.array([261.5, 311.25]), + atol=5.0, + ) + assert np.allclose( + [ + detections.bboxes_metadata[0]["height"], + detections.bboxes_metadata[1]["height"], + ], + np.array([178.4, 135.2]), + atol=6.0, + ) + assert np.allclose( + [ + detections.bboxes_metadata[0]["angle"], + detections.bboxes_metadata[1]["angle"], + ], + np.array([0.826, 79.5]), + atol=0.5, + ) + + +@_TENSOR_ONLY +def test_rectangle_bounding_workflow_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + # Same as test_rectangle_bounding_workflow_tensor_native, but the image arrives ALREADY + # materialised as a CHW RGB device tensor (is_tensor_materialised() == True), so the + # instance-seg block runs its on-device tensor path. Results must match the numpy-input + # variant. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=BOUNDNG_RECTANGLE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the fixture as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": [numpy_image_as_tensor(dogs_image)], + } + ) + + # then + assert len(result) == 1, "One set ot outputs expected" + assert "result" in result[0], "Output must contain key 'result'" + detections = result[0]["result"] + assert isinstance( + detections, NativeInstanceDetections + ), "Output must be instance of native inference_models.InstanceDetections" + assert detections.xyxy.shape[0] == 2, "Two dogs on the image" + + # On the native carrier the rect/width/height/angle geometry that lived in the + # vectorized sv.Detections.data columns is stored per-box in bboxes_metadata[i]. + assert detections.bboxes_metadata is not None, "Should have per-box metadata" + for i in range(2): + meta = detections.bboxes_metadata[i] + assert "rect" in meta, "'rect' geometry must be found in per-box metadata" + assert "width" in meta, "'width' geometry must be found in per-box metadata" + assert "height" in meta, "'height' geometry must be found in per-box metadata" + assert "angle" in meta, "'angle' geometry must be found in per-box metadata" + + assert np.allclose( + np.array(detections.bboxes_metadata[0]["rect"]), + np.array([[322.0, 402.0], [325.0, 224.0], [586.0, 228.0], [583.0, 406.0]]), + atol=5.0, + ) + assert np.allclose( + np.array(detections.bboxes_metadata[1]["rect"]), + np.array([[219.0, 82.0], [352.0, 57.0], [409.0, 363.0], [276.0, 388.0]]), + atol=6.0, + ) + assert np.allclose( + [ + detections.bboxes_metadata[0]["width"], + detections.bboxes_metadata[1]["width"], + ], + np.array([261.5, 311.25]), + atol=5.0, + ) + assert np.allclose( + [ + detections.bboxes_metadata[0]["height"], + detections.bboxes_metadata[1]["height"], + ], + np.array([178.4, 135.2]), + atol=6.0, + ) + assert np.allclose( + [ + detections.bboxes_metadata[0]["angle"], + detections.bboxes_metadata[1]["angle"], + ], + np.array([0.826, 79.5]), + atol=0.5, + ) diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_classical_pattern_matching.py b/tests/workflows/integration_tests/execution/test_workflow_with_classical_pattern_matching.py index f876afa5ac..19b37a1880 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_classical_pattern_matching.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_classical_pattern_matching.py @@ -1,9 +1,33 @@ import numpy as np +import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.object_detection import Detections as NativeDetections +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the template_matching block emits a native +# inference_models.Detections instead of sv.Detections. The numpy-shaped assertion that +# reaches for the sv-style ["class_name"] subscript runs only with the flag off; an +# equivalent *_tensor_native test asserts the same facts against the native carrier with +# the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output; native under ENABLE_TENSOR_DATA_REPRESENTATION " + "โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) WORKFLOW_WITH_CLASSICAL_PATTERN_MATCHING = { "version": "1.0", @@ -49,6 +73,7 @@ } +@_NUMPY_ONLY def test_workflow_with_classical_pattern_matching( model_manager: ModelManager, dogs_image: np.ndarray, @@ -113,6 +138,153 @@ def test_workflow_with_classical_pattern_matching( ), "Expected fixed confidence" +@_TENSOR_ONLY +def test_workflow_with_classical_pattern_matching_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + crowd_image: np.ndarray, +) -> None: + """ + In this test set we check how classical pattern matching block integrates with + other blocks that accept sv.Detections. + + Please point out that a single image is passed as template, and batch of images + are passed as images to look for template. This workflow does also validate + Execution Engine capabilities to broadcast batch-oriented inputs properly. + """ + # given + template = dogs_image[220:280, 310:410] # dog's head + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_CLASSICAL_PATTERN_MATCHING, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [crowd_image, dogs_image], + "template": template, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 2, "Two images provided, so two outputs expected" + assert set(result[0].keys()) == { + "predictions", + "visualization", + }, "Expected all declared outputs to be delivered" + assert set(result[1].keys()) == { + "predictions", + "visualization", + }, "Expected all declared outputs to be delivered" + assert ( + len(result[0]["predictions"]) == 0 + ), "Expected no patterns matched in first image" + assert ( + len(result[1]["predictions"]) == 1 + ), "Expected single pattern match in second image" + assert np.allclose( + result[1]["predictions"].xyxy, np.array([312, 222, 412, 282]) + ), "Expected to find single template match" + assert result[1]["predictions"].class_id.tolist() == [ + 0 + ], "Expected fixed class id 0" + # Native Detections carry class names in image_metadata[CLASS_NAMES_KEY] (a + # {id: name} mapping) rather than via the sv-style ["class_name"] subscript. + preds = result[1]["predictions"] + assert isinstance( + preds, NativeDetections + ), "Output must be instance of native inference_models.Detections" + assert [ + preds.image_metadata[CLASS_NAMES_KEY][int(c)] for c in preds.class_id.tolist() + ] == ["template_match"], "Expected fixed class name" + assert np.allclose( + result[1]["predictions"].confidence, np.array([1.0]) + ), "Expected fixed confidence" + + +@_TENSOR_ONLY +def test_workflow_with_classical_pattern_matching_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + crowd_image: np.ndarray, +) -> None: + """ + Same as test_workflow_with_classical_pattern_matching_tensor_native, but every + image arrives ALREADY materialised as a CHW RGB device tensor + (is_tensor_materialised() == True), so the blocks run their on-device tensor path. + Results must match the numpy-input variant. + + Please point out that a single image is passed as template, and batch of images + are passed as images to look for template. This workflow does also validate + Execution Engine capabilities to broadcast batch-oriented inputs properly. + """ + # given + template = dogs_image[220:280, 310:410] # dog's head + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_CLASSICAL_PATTERN_MATCHING, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the fixtures as pre-materialised tensors + result = execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(crowd_image), + numpy_image_as_tensor(dogs_image), + ], + "template": numpy_image_as_tensor(template), + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 2, "Two images provided, so two outputs expected" + assert set(result[0].keys()) == { + "predictions", + "visualization", + }, "Expected all declared outputs to be delivered" + assert set(result[1].keys()) == { + "predictions", + "visualization", + }, "Expected all declared outputs to be delivered" + assert ( + len(result[0]["predictions"]) == 0 + ), "Expected no patterns matched in first image" + assert ( + len(result[1]["predictions"]) == 1 + ), "Expected single pattern match in second image" + assert np.allclose( + result[1]["predictions"].xyxy, np.array([312, 222, 412, 282]) + ), "Expected to find single template match" + assert result[1]["predictions"].class_id.tolist() == [ + 0 + ], "Expected fixed class id 0" + # Native Detections carry class names in image_metadata[CLASS_NAMES_KEY] (a + # {id: name} mapping) rather than via the sv-style ["class_name"] subscript. + preds = result[1]["predictions"] + assert isinstance( + preds, NativeDetections + ), "Output must be instance of native inference_models.Detections" + assert [ + preds.image_metadata[CLASS_NAMES_KEY][int(c)] for c in preds.class_id.tolist() + ] == ["template_match"], "Expected fixed class name" + assert np.allclose( + result[1]["predictions"].confidence, np.array([1.0]) + ), "Expected fixed confidence" + + WORKFLOW_WITH_CLASSICAL_PATTERN_MATCHING_REFERRING_THE_SAME_AS_IMAGE_AND_TEMPLATE = { "version": "1.0", "inputs": [ diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_clip.py b/tests/workflows/integration_tests/execution/test_workflow_with_clip.py index 228820e2ab..e3a41d0f82 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_clip.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_clip.py @@ -1,15 +1,40 @@ import numpy as np import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.errors import RuntimeInputError, StepExecutionError +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.classification import ClassificationPrediction +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the CLIP block emits the embedding as the +# tensor-native EMBEDDING kind (a torch.Tensor) instead of a numpy array / list. +# `ExecutionEngine.run` does not serialise outputs (serialize_results defaults to +# False), so the raw torch.Tensor reaches the test. Calling np.mean/np.max/np.min/ +# np.std directly on a torch.Tensor raises TypeError. The numeric content is +# identical to the numpy baseline; the *_tensor_native sibling asserts the same +# statistics after converting the tensor to numpy. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="embedding output is a torch.Tensor under ENABLE_TENSOR_DATA_REPRESENTATION " + "โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + CLIP_WORKFLOW = { "version": "1.0", "inputs": [ @@ -296,6 +321,7 @@ def test_clip_embedding_model_on_batches_of_cross_type_data_with_different_embed } +@_NUMPY_ONLY def test_clip_text_embedding_model( model_manager: ModelManager, license_plate_image: np.ndarray, @@ -339,6 +365,54 @@ def test_clip_text_embedding_model( ), "Expected embedding to have a value similar to during testing" +@_TENSOR_ONLY +def test_clip_text_embedding_model_tensor_native( + model_manager: ModelManager, + license_plate_image: np.ndarray, + crowd_image: np.ndarray, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=CLIP_TEXT_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run(runtime_parameters={"prompt": "Foo Bar"}) + + # then + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output" + assert set(result[0].keys()) == { + "text_embeddings", + }, "Expected all declared outputs to be delivered" + # Under ENABLE_TENSOR_DATA_REPRESENTATION the embedding is a tensor-native + # EMBEDDING (a torch.Tensor). It serialises to List[float] at the API boundary; + # here we convert to numpy and assert the same statistics as the numpy baseline. + text_embeddings = np.asarray(result[0]["text_embeddings"].detach().cpu()) + assert ( + len(text_embeddings) == 1024 + ), "Expected text embedding to be of dimension 1024 for RN50 model" + assert ( + pytest.approx(np.mean(text_embeddings), 0.0001) == -0.016772 + ), "Expected embedding to have a value similar to during testing" + assert ( + pytest.approx(np.max(text_embeddings), 0.0001) == 1.65736556 + ), "Expected embedding to have a value similar to during testing" + assert ( + pytest.approx(np.min(text_embeddings), 0.0001) == -10.109556 + ), "Expected embedding to have a value similar to during testing" + assert ( + pytest.approx(np.std(text_embeddings), 0.0001) == 0.39733439 + ), "Expected embedding to have a value similar to during testing" + + CLIP_COMPARISON_WORKFLOW = { "version": "1.0", "inputs": [ @@ -464,6 +538,7 @@ def test_clip_workflow_when_minimal_valid_input_provided( } +@_NUMPY_ONLY def test_workflow_with_clip_comparison_v2_and_property_definition_with_valid_input( model_manager: ModelManager, license_plate_image: np.ndarray, @@ -564,6 +639,234 @@ def test_workflow_with_clip_comparison_v2_and_property_definition_with_valid_inp ), "Expected property definition step to cooperate nicely with clip output" +@_TENSOR_ONLY +def test_workflow_with_clip_comparison_v2_and_property_definition_with_valid_input_tensor_native( + model_manager: ModelManager, + license_plate_image: np.ndarray, + crowd_image: np.ndarray, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_CLIP_COMPARISON_V2, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [license_plate_image, crowd_image], + "reference": ["car", "crowd"], + } + ) + + # then + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 2, "Expected 2 elements in the output for two input images" + assert set(result[0].keys()) == { + "clip_output", + "class_name", + }, "Expected all declared outputs to be delivered" + assert set(result[1].keys()) == { + "clip_output", + "class_name", + }, "Expected all declared outputs to be delivered" + assert np.allclose( + result[0]["clip_output"]["similarities"], + [0.23334351181983948, 0.17259158194065094], + atol=1e-2, + ), "Expected predicted similarities to match values verified at test creation" + assert ( + abs( + result[0]["clip_output"]["similarities"][0] + - result[0]["clip_output"]["max_similarity"] + ) + < 1e-2 + ), "Expected max similarity to be correct" + assert ( + abs( + result[0]["clip_output"]["similarities"][1] + - result[0]["clip_output"]["min_similarity"] + ) + < 1e-2 + ), "Expected max similarity to be correct" + assert ( + result[0]["clip_output"]["most_similar_class"] == "car" + ), "Expected most similar class to be extracted properly" + assert ( + result[0]["clip_output"]["least_similar_class"] == "crowd" + ), "Expected least similar class to be extracted properly" + # Under ENABLE_TENSOR_DATA_REPRESENTATION clip_comparison@v2_tensor places a native + # ClassificationPrediction at clip_output["classification_predictions"]; the top class + # name lives at images_metadata[0][CLASS_NAMES_KEY] (plural, indexed) keyed by the top + # class id (pred.class_id) rather than under a ["top"] subscript. + pred_0 = result[0]["clip_output"]["classification_predictions"] + assert isinstance(pred_0, ClassificationPrediction) + assert ( + pred_0.images_metadata[0][CLASS_NAMES_KEY][int(pred_0.class_id.reshape(-1)[0])] + == "car" + ), "Expected classifier output to be shaped correctly" + assert ( + result[0]["class_name"] == "car" + ), "Expected property definition step to cooperate nicely with clip output" + assert np.allclose( + result[1]["clip_output"]["similarities"], + [0.18426208198070526, 0.207647442817688], + atol=1e-2, + ), "Expected predicted similarities to match values verified at test creation" + assert ( + abs( + result[1]["clip_output"]["similarities"][1] + - result[1]["clip_output"]["max_similarity"] + ) + < 1e-5 + ), "Expected max similarity to be correct" + assert ( + abs( + result[1]["clip_output"]["similarities"][0] + - result[1]["clip_output"]["min_similarity"] + ) + < 1e-5 + ), "Expected max similarity to be correct" + assert ( + result[1]["clip_output"]["most_similar_class"] == "crowd" + ), "Expected most similar class to be extracted properly" + assert ( + result[1]["clip_output"]["least_similar_class"] == "car" + ), "Expected least similar class to be extracted properly" + pred_1 = result[1]["clip_output"]["classification_predictions"] + assert isinstance(pred_1, ClassificationPrediction) + assert ( + pred_1.images_metadata[0][CLASS_NAMES_KEY][int(pred_1.class_id.reshape(-1)[0])] + == "crowd" + ), "Expected classifier output to be shaped correctly" + assert ( + result[1]["class_name"] == "crowd" + ), "Expected property definition step to cooperate nicely with clip output" + + +@_TENSOR_ONLY +def test_workflow_with_clip_comparison_v2_and_property_definition_with_valid_input_with_tensor_input( + model_manager: ModelManager, + license_plate_image: np.ndarray, + crowd_image: np.ndarray, +) -> None: + # Same as the *_tensor_native variant, but both images arrive ALREADY materialised as + # CHW RGB device tensors (is_tensor_materialised() == True), so the CLIP block runs its + # on-device tensor path. Results must match the numpy-input variant. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_CLIP_COMPARISON_V2, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the fixtures as pre-materialised tensors + result = execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(license_plate_image), + numpy_image_as_tensor(crowd_image), + ], + "reference": ["car", "crowd"], + } + ) + + # then + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 2, "Expected 2 elements in the output for two input images" + assert set(result[0].keys()) == { + "clip_output", + "class_name", + }, "Expected all declared outputs to be delivered" + assert set(result[1].keys()) == { + "clip_output", + "class_name", + }, "Expected all declared outputs to be delivered" + assert np.allclose( + result[0]["clip_output"]["similarities"], + [0.23334351181983948, 0.17259158194065094], + atol=1e-2, + ), "Expected predicted similarities to match values verified at test creation" + assert ( + abs( + result[0]["clip_output"]["similarities"][0] + - result[0]["clip_output"]["max_similarity"] + ) + < 1e-2 + ), "Expected max similarity to be correct" + assert ( + abs( + result[0]["clip_output"]["similarities"][1] + - result[0]["clip_output"]["min_similarity"] + ) + < 1e-2 + ), "Expected max similarity to be correct" + assert ( + result[0]["clip_output"]["most_similar_class"] == "car" + ), "Expected most similar class to be extracted properly" + assert ( + result[0]["clip_output"]["least_similar_class"] == "crowd" + ), "Expected least similar class to be extracted properly" + # Under ENABLE_TENSOR_DATA_REPRESENTATION clip_comparison@v2_tensor places a native + # ClassificationPrediction at clip_output["classification_predictions"]; the top class + # name lives at images_metadata[0][CLASS_NAMES_KEY] (plural, indexed) keyed by the top + # class id (pred.class_id) rather than under a ["top"] subscript. + pred_0 = result[0]["clip_output"]["classification_predictions"] + assert isinstance(pred_0, ClassificationPrediction) + assert ( + pred_0.images_metadata[0][CLASS_NAMES_KEY][int(pred_0.class_id.reshape(-1)[0])] + == "car" + ), "Expected classifier output to be shaped correctly" + assert ( + result[0]["class_name"] == "car" + ), "Expected property definition step to cooperate nicely with clip output" + assert np.allclose( + result[1]["clip_output"]["similarities"], + [0.18426208198070526, 0.207647442817688], + atol=1e-2, + ), "Expected predicted similarities to match values verified at test creation" + assert ( + abs( + result[1]["clip_output"]["similarities"][1] + - result[1]["clip_output"]["max_similarity"] + ) + < 1e-5 + ), "Expected max similarity to be correct" + assert ( + abs( + result[1]["clip_output"]["similarities"][0] + - result[1]["clip_output"]["min_similarity"] + ) + < 1e-5 + ), "Expected max similarity to be correct" + assert ( + result[1]["clip_output"]["most_similar_class"] == "crowd" + ), "Expected most similar class to be extracted properly" + assert ( + result[1]["clip_output"]["least_similar_class"] == "car" + ), "Expected least similar class to be extracted properly" + pred_1 = result[1]["clip_output"]["classification_predictions"] + assert isinstance(pred_1, ClassificationPrediction) + assert ( + pred_1.images_metadata[0][CLASS_NAMES_KEY][int(pred_1.class_id.reshape(-1)[0])] + == "crowd" + ), "Expected classifier output to be shaped correctly" + assert ( + result[1]["class_name"] == "crowd" + ), "Expected property definition step to cooperate nicely with clip output" + + def test_workflow_with_clip_comparison_v2_and_property_definition_with_empty_class_list( model_manager: ModelManager, license_plate_image: np.ndarray, diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_custom_python_block.py b/tests/workflows/integration_tests/execution/test_workflow_with_custom_python_block.py index a128f77cde..98f0e01deb 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_custom_python_block.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_custom_python_block.py @@ -4,7 +4,11 @@ import numpy as np import pytest -from inference.core.env import USE_INFERENCE_MODELS, WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + USE_INFERENCE_MODELS, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.errors import ( @@ -540,6 +544,22 @@ def my_function(self, prediction: sv.Detections, crops: Batch[WorkflowImageData] } +def _class_names_of_prediction(prediction) -> list: + """Representation-agnostic class-name extraction: raw engine outputs are + sv.Detections flag-off and native inference_models Detections flag-on (the + declared-kind OUT boundary converts them by design).""" + import supervision as sv + + if isinstance(prediction, sv.Detections): + return prediction["class_name"].tolist() + per_box = prediction.bboxes_metadata or [{} for _ in range(len(prediction))] + mapping = (prediction.image_metadata or {}).get("class_names", {}) + return [ + str(box.get("class", mapping.get(int(class_id)))) + for box, class_id in zip(per_box, prediction.class_id.detach().cpu().tolist()) + ] + + def test_workflow_with_custom_python_block_operating_cross_dimensions( model_manager: ModelManager, dogs_image: np.ndarray, @@ -575,13 +595,13 @@ def test_workflow_with_custom_python_block_operating_cross_dimensions( }, "Expected all declared outputs to be delivered" assert len(result[1]["associated_detections"]) == 12 class_names_first_image_crops = [ - e["class_name"].tolist() for e in result[0]["associated_detections"] + _class_names_of_prediction(e) for e in result[0]["associated_detections"] ] for class_names in class_names_first_image_crops: assert len(class_names) == 1, "Expected single bbox to be associated" assert len(class_names_first_image_crops) == 2, "Expected 2 crops for first image" class_names_second_image_crops = [ - e["class_name"].tolist() for e in result[1]["associated_detections"] + _class_names_of_prediction(e) for e in result[1]["associated_detections"] ] for class_names in class_names_second_image_crops: assert len(class_names) == 1, "Expected single bbox to be associated" @@ -1050,3 +1070,581 @@ def test_workflow_with_custom_python_block_when_code_does_not_define_declared_in init_parameters=workflow_init_parameters, max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, ) + + +FUNCTION_TO_FILTER_AND_PROBE = """ +def run(self, predictions: sv.Detections) -> BlockResult: + got_sv = isinstance(predictions, sv.Detections) + filtered = predictions[predictions.confidence > 0.6] + _ = filtered.xyxy.copy() # numpy-style access, the documented legacy contract + return {"filtered": filtered, "got_sv": got_sv} +""" + +WORKFLOW_WITH_BOTH_BOUNDARIES_AND_TENSOR_CONSUMER = { + "version": "1.0", + "inputs": [ + {"type": "WorkflowImage", "name": "image"}, + ], + "dynamic_blocks_definitions": [ + { + "type": "DynamicBlockDefinition", + "manifest": { + "type": "ManifestDescription", + "block_type": "FilterProbe", + "inputs": { + "predictions": { + "type": "DynamicInputDefinition", + "selector_types": ["step_output"], + "selector_data_kind": { + "step_output": ["object_detection_prediction"] + }, + }, + }, + "outputs": { + "filtered": { + "type": "DynamicOutputDefinition", + "kind": ["object_detection_prediction"], + }, + "got_sv": {"type": "DynamicOutputDefinition", "kind": []}, + }, + }, + "code": { + "type": "PythonCode", + "run_function_code": FUNCTION_TO_FILTER_AND_PROBE, + }, + }, + ], + "steps": [ + { + "type": "RoboflowObjectDetectionModel", + "name": "model", + "image": "$inputs.image", + "model_id": "yolov8n-640", + }, + { + "type": "FilterProbe", + "name": "filter_probe", + "predictions": "$steps.model.predictions", + }, + { + "type": "roboflow_core/property_definition@v1", + "name": "count", + "data": "$steps.filter_probe.filtered", + "operations": [{"type": "SequenceLength"}], + }, + ], + "outputs": [ + { + "type": "JsonField", + "name": "filtered", + "selector": "$steps.filter_probe.filtered", + }, + {"type": "JsonField", "name": "count", "selector": "$steps.count.output"}, + { + "type": "JsonField", + "name": "got_sv", + "selector": "$steps.filter_probe.got_sv", + }, + ], +} + + +def test_workflow_with_custom_python_block_between_model_and_tensor_consumer( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Both boundaries in one workflow: model producer -> legacy custom block + (default knob; must see sv.Detections and use the numpy contract) -> + downstream UQL consumer (native ops under the flag) -> serialized output. + The hardcoded expectations hold in BOTH flag directions, which is the + integration-level parity statement (byte-parity is proven at the unit + layer's round-trip tests).""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_BOTH_BOUNDARIES_AND_TENSOR_CONSUMER, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": [dogs_image]}, + serialize_results=True, + ) + + # then + assert isinstance(result, list) and len(result) == 1 + assert result[0]["got_sv"] is True, "User code must receive sv.Detections" + assert result[0]["count"] == 1, "Expected exactly one detection above 0.6" + serialized = result[0]["filtered"]["predictions"] + assert len(serialized) == 1 + assert serialized[0]["class"] == "dog" + assert abs(serialized[0]["confidence"] - 0.856) < 0.1 + assert { + "x", + "y", + "width", + "height", + "confidence", + "class", + "class_id", + "detection_id", + } <= set(serialized[0].keys()) + + +FUNCTION_RETURNING_BARE_TENSOR = """ +def run(self, predictions) -> BlockResult: + import torch + return {"weird": torch.zeros(3)} +""" + +WORKFLOW_WITH_BARE_TENSOR_WILDCARD_OUTPUT = { + "version": "1.0", + "inputs": [ + {"type": "WorkflowImage", "name": "image"}, + ], + "dynamic_blocks_definitions": [ + { + "type": "DynamicBlockDefinition", + "manifest": { + "type": "ManifestDescription", + "block_type": "BareTensorBlock", + "inputs": { + "predictions": { + "type": "DynamicInputDefinition", + "selector_types": ["step_output"], + }, + }, + "outputs": {"weird": {"type": "DynamicOutputDefinition", "kind": []}}, + }, + "code": { + "type": "PythonCode", + "run_function_code": FUNCTION_RETURNING_BARE_TENSOR, + }, + }, + ], + "steps": [ + { + "type": "RoboflowObjectDetectionModel", + "name": "model", + "image": "$inputs.image", + "model_id": "yolov8n-640", + }, + { + "type": "BareTensorBlock", + "name": "bad_block", + "predictions": "$steps.model.predictions", + }, + ], + "outputs": [ + {"type": "JsonField", "name": "weird", "selector": "$steps.bad_block.weird"}, + ], +} + + +def test_workflow_with_custom_python_block_returning_bare_tensor_via_wildcard( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """OUT-direction wildcard semantics: a bare torch.Tensor returned through a + wildcard output PASSES THROUGH in both flag directions โ€” on the OUT side a + tensor is a native-representation value (the `tensor` kind's runtime type), + so the boundary forwards it untouched. The loud-error contract for bare + tensors applies to the IN direction (legacy user code cannot operate on + them) and is covered by the representation-boundary unit tests.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_BARE_TENSOR_WILDCARD_OUTPUT, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run(runtime_parameters={"image": [dogs_image]}) + + # then + import torch + + assert isinstance(result[0]["weird"], torch.Tensor) + assert result[0]["weird"].shape == (3,) + + +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason=( + "tensor_native dynamic blocks compile only under " + "ENABLE_TENSOR_DATA_REPRESENTATION (the flag-off compile-time error is " + "unit-tested in test_block_assembler.py)" + ), +) + +FUNCTION_MINTING_NATIVE_DETECTIONS = """ +def run(self, image) -> BlockResult: + detections = Detections( + xyxy=torch.tensor( + [[10.0, 10.0, 60.0, 60.0]], device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + class_id=torch.tensor([0], device=WORKFLOWS_IMAGE_TENSOR_DEVICE), + confidence=torch.tensor([0.9], device=WORKFLOWS_IMAGE_TENSOR_DEVICE), + ) + detections = attach_native_detection_metadata( + detections=detections, + image=image, + class_names={0: "widget"}, + prediction_type="object-detection", + ) + return {"predictions": detections} +""" + +WORKFLOW_WITH_TENSOR_NATIVE_MINTING_BLOCK = { + "version": "1.0", + "inputs": [ + {"type": "WorkflowImage", "name": "image"}, + ], + "dynamic_blocks_definitions": [ + { + "type": "DynamicBlockDefinition", + "manifest": { + "type": "ManifestDescription", + "block_type": "NativeMinter", + "tensor_compatibility": "tensor_native", + "inputs": { + "image": { + "type": "DynamicInputDefinition", + "selector_types": ["input_image"], + }, + }, + "outputs": { + "predictions": { + "type": "DynamicOutputDefinition", + "kind": ["object_detection_prediction"], + }, + }, + }, + "code": { + "type": "PythonCode", + "run_function_code": FUNCTION_MINTING_NATIVE_DETECTIONS, + }, + }, + ], + "steps": [ + { + "type": "NativeMinter", + "name": "minter", + "image": "$inputs.image", + }, + ], + "outputs": [ + { + "type": "JsonField", + "name": "predictions", + "selector": "$steps.minter.predictions", + }, + ], +} + + +@_TENSOR_ONLY +def test_workflow_with_tensor_native_custom_block_minting_native_detections( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """tensor_native authoring surface: the extended IMPORTS_LINES must resolve in + the generated module namespace (torch, Detections, + WORKFLOWS_IMAGE_TENSOR_DEVICE, attach_native_detection_metadata), and the + minted output must satisfy the tensor serializer's hard requirements + (class_names map + per-box detection_id) end-to-end.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_TENSOR_NATIVE_MINTING_BLOCK, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": [dogs_image]}, + serialize_results=True, + ) + + # then + assert isinstance(result, list) and len(result) == 1 + serialized = result[0]["predictions"]["predictions"] + assert len(serialized) == 1 + assert serialized[0]["class"] == "widget" + assert serialized[0]["class_id"] == 0 + assert abs(serialized[0]["confidence"] - 0.9) < 1e-6 + assert serialized[0]["x"] == 35.0 and serialized[0]["y"] == 35.0 + assert serialized[0]["width"] == 50.0 and serialized[0]["height"] == 50.0 + assert len(serialized[0]["detection_id"]) > 0 + + +FUNCTION_PROBING_NATIVE_PASSTHROUGH = """ +def run(self, predictions) -> BlockResult: + got_native = isinstance(predictions, Detections) + xyxy_is_tensor = isinstance(predictions.xyxy, torch.Tensor) + return { + "got_native": got_native, + "xyxy_is_tensor": xyxy_is_tensor, + "passthrough": predictions, + } +""" + +WORKFLOW_WITH_TENSOR_NATIVE_PROBING_BLOCK = { + "version": "1.0", + "inputs": [ + {"type": "WorkflowImage", "name": "image"}, + ], + "dynamic_blocks_definitions": [ + { + "type": "DynamicBlockDefinition", + "manifest": { + "type": "ManifestDescription", + "block_type": "NativeProbe", + "tensor_compatibility": "tensor_native", + "inputs": { + "predictions": { + "type": "DynamicInputDefinition", + "selector_types": ["step_output"], + "selector_data_kind": { + "step_output": ["object_detection_prediction"] + }, + }, + }, + "outputs": { + "got_native": {"type": "DynamicOutputDefinition", "kind": []}, + "xyxy_is_tensor": {"type": "DynamicOutputDefinition", "kind": []}, + "passthrough": { + "type": "DynamicOutputDefinition", + "kind": ["object_detection_prediction"], + }, + }, + }, + "code": { + "type": "PythonCode", + "run_function_code": FUNCTION_PROBING_NATIVE_PASSTHROUGH, + }, + }, + ], + "steps": [ + { + "type": "RoboflowObjectDetectionModel", + "name": "model", + "image": "$inputs.image", + "model_id": "yolov8n-640", + }, + { + "type": "NativeProbe", + "name": "probe", + "predictions": "$steps.model.predictions", + }, + ], + "outputs": [ + { + "type": "JsonField", + "name": "got_native", + "selector": "$steps.probe.got_native", + }, + { + "type": "JsonField", + "name": "xyxy_is_tensor", + "selector": "$steps.probe.xyxy_is_tensor", + }, + { + "type": "JsonField", + "name": "passthrough", + "selector": "$steps.probe.passthrough", + }, + ], +} + + +@_TENSOR_ONLY +def test_workflow_with_tensor_native_custom_block_receiving_native_predictions( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """tensor_native pass-through: user code receives the NATIVE prediction from a + model step (no boundary conversion in either direction) and can return it + for downstream serialization.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_TENSOR_NATIVE_PROBING_BLOCK, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": [dogs_image]}, + serialize_results=True, + ) + + # then + assert isinstance(result, list) and len(result) == 1 + assert result[0]["got_native"] is True, "User code must receive native Detections" + assert result[0]["xyxy_is_tensor"] is True, "Native xyxy must stay a torch.Tensor" + serialized = result[0]["passthrough"]["predictions"] + assert len(serialized) >= 1 + assert {"x", "y", "width", "height", "confidence", "class", "detection_id"} <= set( + serialized[0].keys() + ) + + +FUNCTION_PASSING_PREDICTIONS_THROUGH = """ +def run(self, predictions) -> BlockResult: + return {"passthrough": predictions} +""" + +WORKFLOW_WITH_WILDCARD_PREDICTIONS_OUTPUT = { + "version": "1.0", + "inputs": [ + {"type": "WorkflowImage", "name": "image"}, + ], + "dynamic_blocks_definitions": [ + { + "type": "DynamicBlockDefinition", + "manifest": { + "type": "ManifestDescription", + "block_type": "WildcardPassthrough", + "inputs": { + "predictions": { + "type": "DynamicInputDefinition", + "selector_types": ["step_output"], + }, + }, + "outputs": { + "passthrough": {"type": "DynamicOutputDefinition", "kind": []} + }, + }, + "code": { + "type": "PythonCode", + "run_function_code": FUNCTION_PASSING_PREDICTIONS_THROUGH, + }, + }, + ], + "steps": [ + { + "type": "RoboflowObjectDetectionModel", + "name": "model", + "image": "$inputs.image", + "model_id": "yolov8n-640", + }, + { + "type": "WildcardPassthrough", + "name": "passer", + "predictions": "$steps.model.predictions", + }, + ], + "outputs": [ + { + "type": "JsonField", + "name": "forwarded", + "selector": "$steps.passer.passthrough", + }, + ], +} + + +def test_workflow_with_custom_python_block_wildcard_predictions_output_serialization( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Step-6 closure: predictions routed through a custom block's WILDCARD + output must serialize to the standard predictions dict at workflow outputs. + Flag-on: the OUT boundary converts the returned sv to native and the + flag-swapped tensor `serialize_wildcard_kind` emits the standard dict; + flag-off: the numpy wildcard sv arm does โ€” identical assertions in both + directions are the parity statement.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_WILDCARD_PREDICTIONS_OUTPUT, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": [dogs_image]}, + serialize_results=True, + ) + + # then + assert isinstance(result, list) and len(result) == 1 + forwarded = result[0]["forwarded"] + assert isinstance(forwarded, dict), "Wildcard output must serialize to a dict" + serialized = forwarded["predictions"] + assert len(serialized) == 2, "Expected the two dogs the model detects" + assert {prediction["class"] for prediction in serialized} == {"dog"} + assert { + "x", + "y", + "width", + "height", + "confidence", + "class", + "class_id", + "detection_id", + } <= set(serialized[0].keys()) + + +def test_workflow_with_custom_python_block_bare_tensor_wildcard_output_serialization( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + """Step-6 closure of the gap the engine-level bare-tensor test left open: + with serialize_results=True, a bare torch.Tensor through a wildcard output + now serializes to the JSON-safe nested list under the flag (the tensor + wildcard serialiser mirrors the `tensor` kind) instead of crashing HTTP + encoding. Flag-off keeps the numpy wildcard contract verbatim: raw + passthrough (pre-existing behavior, deliberately unchanged).""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_BARE_TENSOR_WILDCARD_OUTPUT, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": [dogs_image]}, + serialize_results=True, + ) + + # then + import torch + + if ENABLE_TENSOR_DATA_REPRESENTATION: + assert result[0]["weird"] == [0.0, 0.0, 0.0] + else: + assert isinstance(result[0]["weird"], torch.Tensor) diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_dataset_upload_metadata.py b/tests/workflows/integration_tests/execution/test_workflow_with_dataset_upload_metadata.py index f3abf04b65..966e29f1a1 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_dataset_upload_metadata.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_dataset_upload_metadata.py @@ -2,12 +2,39 @@ from unittest.mock import MagicMock import numpy as np +import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode -from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload import v2 +from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload import ( + v2, + v2_tensor, +) from inference.core.workflows.execution_engine.core import ExecutionEngine +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the loader registers the tensor-native +# RoboflowDatasetUploadBlockV2 from `v2_tensor`, whose run() calls +# `v2_tensor.maybe_register_datapoint_at_roboflow`. The numpy test below patches the +# seam on the `v2` module, which is no longer the executing module when the flag is +# on (so the real, network-hitting code path runs). The flag-on parity test patches +# the equivalent seam on `v2_tensor` and asserts the same semantic result. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="patches v2.maybe_register_datapoint_at_roboflow; under " + "ENABLE_TENSOR_DATA_REPRESENTATION the executing block is v2_tensor โ€” see the " + "*_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) WORKFLOW_WITH_DATASET_UPLOAD_METADATA_SELECTOR = { "version": "1.0", @@ -48,6 +75,7 @@ } +@_NUMPY_ONLY @mock.patch.object(v2, "maybe_register_datapoint_at_roboflow") def test_workflow_with_dataset_upload_metadata_selector_inside_dict_for_batch( maybe_register_datapoint_at_roboflow_mock: MagicMock, @@ -89,3 +117,98 @@ def test_workflow_with_dataset_upload_metadata_selector_inside_dict_for_batch( "location": "warehouse_a", "source": "edge_camera", } + + +@_TENSOR_ONLY +@mock.patch.object(v2_tensor, "maybe_register_datapoint_at_roboflow") +def test_workflow_with_dataset_upload_metadata_selector_inside_dict_for_batch_with_tensor_input( + maybe_register_datapoint_at_roboflow_mock: MagicMock, + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + # Same as + # test_workflow_with_dataset_upload_metadata_selector_inside_dict_for_batch_tensor_native, + # but each image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the block runs its on-device tensor path. + # given + maybe_register_datapoint_at_roboflow_mock.return_value = False, "OK" + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": "my_api_key", + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_DATASET_UPLOAD_METADATA_SELECTOR, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(dogs_image), + ], + "location": "warehouse_a", + } + ) + + # then + assert len(result) == 2, "Expected one output per input image" + assert result[0]["registration_message"] == "OK" + assert result[1]["registration_message"] == "OK" + assert maybe_register_datapoint_at_roboflow_mock.call_count == 2 + calls = maybe_register_datapoint_at_roboflow_mock.call_args_list + assert calls[0].kwargs["metadata"] == { + "location": "warehouse_a", + "source": "edge_camera", + } + assert calls[1].kwargs["metadata"] == { + "location": "warehouse_a", + "source": "edge_camera", + } + + +@_TENSOR_ONLY +@mock.patch.object(v2_tensor, "maybe_register_datapoint_at_roboflow") +def test_workflow_with_dataset_upload_metadata_selector_inside_dict_for_batch_tensor_native( + maybe_register_datapoint_at_roboflow_mock: MagicMock, + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + # given + maybe_register_datapoint_at_roboflow_mock.return_value = False, "OK" + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": "my_api_key", + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_DATASET_UPLOAD_METADATA_SELECTOR, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image], + "location": "warehouse_a", + } + ) + + # then + assert len(result) == 2, "Expected one output per input image" + assert result[0]["registration_message"] == "OK" + assert result[1]["registration_message"] == "OK" + assert maybe_register_datapoint_at_roboflow_mock.call_count == 2 + calls = maybe_register_datapoint_at_roboflow_mock.call_args_list + assert calls[0].kwargs["metadata"] == { + "location": "warehouse_a", + "source": "edge_camera", + } + assert calls[1].kwargs["metadata"] == { + "location": "warehouse_a", + "source": "edge_camera", + } diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_detection_renaming.py b/tests/workflows/integration_tests/execution/test_workflow_with_detection_renaming.py index ef7be64378..42b13bba60 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_detection_renaming.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_detection_renaming.py @@ -1,10 +1,15 @@ +from collections import Counter from copy import deepcopy from typing import Dict import numpy as np import pytest -from inference.core.env import USE_INFERENCE_MODELS, WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + USE_INFERENCE_MODELS, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.core_steps.common.query_language.errors import ( @@ -12,11 +17,36 @@ OperationError, UndeclaredSymbolError, ) +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.object_detection import Detections as NativeDetections +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the DetectionsRename UQL op emits a native +# inference_models.Detections instead of sv.Detections. Native Detections is a plain +# dataclass: it has no __getitem__ (so `det["class_name"]` raises TypeError) and class +# names live in image_metadata[CLASS_NAMES_KEY] rather than data["class_name"]. The +# upstream tensor ObjectDetectionModel/NMS also returns detections in a different +# positional order than the numpy path, so the sv-shaped tests below (which assert an +# exact ordering via subscript access) run only with the flag off; an equivalent +# *_tensor_native test asserts the same order-independent semantic result against the +# native carrier with the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output + exact ordering; native carrier with different " + "upstream ordering under ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the " + "*_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + def build_class_remapping_workflow_definition( class_map: Dict[str, str], @@ -84,6 +114,7 @@ class name changes to `fruit` and class id is 1024. ), workflow_name_in_app="detections-class-remapping", ) +@_NUMPY_ONLY def test_class_rename_workflow_with_non_strict_mapping( model_manager: ModelManager, fruit_image: np.ndarray, @@ -153,6 +184,7 @@ def test_class_rename_workflow_with_non_strict_mapping( ), "Expected length of predictions no to change" +@_NUMPY_ONLY def test_class_rename_workflow_with_strict_mapping_when_all_classes_are_remapped( model_manager: ModelManager, fruit_image: np.ndarray, @@ -300,6 +332,7 @@ def test_class_rename_workflow_with_strict_mapping_when_not_all_classes_are_rema } +@_NUMPY_ONLY def test_class_rename_workflow_when_mapping_is_parametrised( model_manager: ModelManager, fruit_image: np.ndarray, @@ -421,3 +454,316 @@ def test_class_rename_workflow_when_mapping_is_not_passed_as_operation_parameter "strict": False, }, ) + + +def _native_class_names(detections: NativeDetections) -> Counter: + """Resolve per-detection class names from the native carrier. + + On native inference_models.Detections the class names live in + image_metadata[CLASS_NAMES_KEY] keyed by class id (there is no data["class_name"] + and no __getitem__). The renamed detections come back in a different positional + order than the numpy path, so callers compare the resulting multiset rather than + an exact ordering. + """ + class_names_map = detections.image_metadata[CLASS_NAMES_KEY] + return Counter( + class_names_map[int(class_id)] for class_id in detections.class_id.tolist() + ) + + +@_TENSOR_ONLY +def test_class_rename_workflow_with_non_strict_mapping_tensor_native( + model_manager: ModelManager, + fruit_image: np.ndarray, +) -> None: + workflow_definition = build_class_remapping_workflow_definition( + class_map={"apple": "fruit", "banana": "fruit"}, + strict=False, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": fruit_image, + "model_id": "yolov8n-640", + }, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + renamed = result[0]["renamed_predictions"] + assert isinstance( + renamed, NativeDetections + ), "Output must be native inference_models.Detections under tensor mode" + # Same semantic result as the numpy path (order-independent): 5 boxes renamed to + # `fruit` with new non-strict id 1024, the single `orange` (id 49) kept as-is. + assert _native_class_names(renamed) == Counter( + {"fruit": 5, "orange": 1} + ), "Expected renamed class names multiset to match the numpy baseline" + assert Counter(renamed.class_id.tolist()) == Counter( + {1024: 5, 49: 1} + ), "Expected renamed class id multiset to match the numpy baseline" + assert len(renamed) == len( + result[0]["original_predictions"] + ), "Expected length of predictions not to change" + + +@_TENSOR_ONLY +def test_class_rename_workflow_with_strict_mapping_when_all_classes_are_remapped_tensor_native( + model_manager: ModelManager, + fruit_image: np.ndarray, +) -> None: + workflow_definition = build_class_remapping_workflow_definition( + class_map={"apple": "fruit", "banana": "fruit", "orange": "my-orange"}, + strict=True, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": fruit_image, + "model_id": "yolov8n-640", + }, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + renamed = result[0]["renamed_predictions"] + assert isinstance( + renamed, NativeDetections + ), "Output must be native inference_models.Detections under tensor mode" + # Strict mapping reindexes the resulting classes from 0 (sorted): fruit -> 0, + # my-orange -> 1. Same semantic result as the numpy path (order-independent). + assert _native_class_names(renamed) == Counter( + {"fruit": 5, "my-orange": 1} + ), "Expected renamed class names multiset to match the numpy baseline" + assert Counter(renamed.class_id.tolist()) == Counter( + {0: 5, 1: 1} + ), "Expected renamed class id multiset to match the numpy baseline" + assert len(renamed) == len( + result[0]["original_predictions"] + ), "Expected length of predictions not to change" + + +@_TENSOR_ONLY +def test_class_rename_workflow_when_mapping_is_parametrised_tensor_native( + model_manager: ModelManager, + fruit_image: np.ndarray, +) -> None: + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_PARAMETRISED_DETECTIONS_RENAME, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": fruit_image, + "model_id": "yolov8n-640", + "class_map": {"apple": "fruit", "banana": "fruit"}, + "strict": False, + }, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + renamed = result[0]["renamed_predictions"] + assert isinstance( + renamed, NativeDetections + ), "Output must be native inference_models.Detections under tensor mode" + # Same semantic result as the numpy path (order-independent): 5 boxes renamed to + # `fruit` with new non-strict id 1024, the single `orange` (id 49) kept as-is. + assert _native_class_names(renamed) == Counter( + {"fruit": 5, "orange": 1} + ), "Expected renamed class names multiset to match the numpy baseline" + assert Counter(renamed.class_id.tolist()) == Counter( + {1024: 5, 49: 1} + ), "Expected renamed class id multiset to match the numpy baseline" + assert len(renamed) == len( + result[0]["original_predictions"] + ), "Expected length of predictions not to change" + + +@_TENSOR_ONLY +def test_class_rename_workflow_with_non_strict_mapping_with_tensor_input( + model_manager: ModelManager, + fruit_image: np.ndarray, +) -> None: + # Same as test_class_rename_workflow_with_non_strict_mapping_tensor_native, but the + # image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the OD block runs its on-device tensor path. + workflow_definition = build_class_remapping_workflow_definition( + class_map={"apple": "fruit", "banana": "fruit"}, + strict=False, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(fruit_image), + "model_id": "yolov8n-640", + }, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + renamed = result[0]["renamed_predictions"] + assert isinstance( + renamed, NativeDetections + ), "Output must be native inference_models.Detections under tensor mode" + # Same semantic result as the numpy path (order-independent): 5 boxes renamed to + # `fruit` with new non-strict id 1024, the single `orange` (id 49) kept as-is. + assert _native_class_names(renamed) == Counter( + {"fruit": 5, "orange": 1} + ), "Expected renamed class names multiset to match the numpy baseline" + assert Counter(renamed.class_id.tolist()) == Counter( + {1024: 5, 49: 1} + ), "Expected renamed class id multiset to match the numpy baseline" + assert len(renamed) == len( + result[0]["original_predictions"] + ), "Expected length of predictions not to change" + + +@_TENSOR_ONLY +def test_class_rename_workflow_with_strict_mapping_when_all_classes_are_remapped_with_tensor_input( + model_manager: ModelManager, + fruit_image: np.ndarray, +) -> None: + # Same as + # test_class_rename_workflow_with_strict_mapping_when_all_classes_are_remapped_tensor_native, + # but the image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the OD block runs its on-device tensor path. + workflow_definition = build_class_remapping_workflow_definition( + class_map={"apple": "fruit", "banana": "fruit", "orange": "my-orange"}, + strict=True, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(fruit_image), + "model_id": "yolov8n-640", + }, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + renamed = result[0]["renamed_predictions"] + assert isinstance( + renamed, NativeDetections + ), "Output must be native inference_models.Detections under tensor mode" + # Strict mapping reindexes the resulting classes from 0 (sorted): fruit -> 0, + # my-orange -> 1. Same semantic result as the numpy path (order-independent). + assert _native_class_names(renamed) == Counter( + {"fruit": 5, "my-orange": 1} + ), "Expected renamed class names multiset to match the numpy baseline" + assert Counter(renamed.class_id.tolist()) == Counter( + {0: 5, 1: 1} + ), "Expected renamed class id multiset to match the numpy baseline" + assert len(renamed) == len( + result[0]["original_predictions"] + ), "Expected length of predictions not to change" + + +@_TENSOR_ONLY +def test_class_rename_workflow_when_mapping_is_parametrised_with_tensor_input( + model_manager: ModelManager, + fruit_image: np.ndarray, +) -> None: + # Same as test_class_rename_workflow_when_mapping_is_parametrised_tensor_native, but + # the image arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the OD block runs its on-device tensor path. + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_PARAMETRISED_DETECTIONS_RENAME, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(fruit_image), + "model_id": "yolov8n-640", + "class_map": {"apple": "fruit", "banana": "fruit"}, + "strict": False, + }, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + renamed = result[0]["renamed_predictions"] + assert isinstance( + renamed, NativeDetections + ), "Output must be native inference_models.Detections under tensor mode" + # Same semantic result as the numpy path (order-independent): 5 boxes renamed to + # `fruit` with new non-strict id 1024, the single `orange` (id 49) kept as-is. + assert _native_class_names(renamed) == Counter( + {"fruit": 5, "orange": 1} + ), "Expected renamed class names multiset to match the numpy baseline" + assert Counter(renamed.class_id.tolist()) == Counter( + {1024: 5, 49: 1} + ), "Expected renamed class id multiset to match the numpy baseline" + assert len(renamed) == len( + result[0]["original_predictions"] + ), "Expected length of predictions not to change" diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_detections_classes_replacement_by_specialised_classification.py b/tests/workflows/integration_tests/execution/test_workflow_with_detections_classes_replacement_by_specialised_classification.py index f6cd3532d0..ceecf66eb6 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_detections_classes_replacement_by_specialised_classification.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_detections_classes_replacement_by_specialised_classification.py @@ -1,14 +1,48 @@ import numpy as np import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.core import ExecutionEngine +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the DetectionsClassesReplacement block and +# the ObjectDetectionModel producer both emit a native inference_models.Detections (a +# plain @dataclass of torch tensors) which has no sv `.data` / `__getitem__`, so the +# sv-only `det["class_name"]` reads raise `'Detections' object is not subscriptable`. +# The sv-shaped test below is skipped when the flag is on; the `*_tensor_native` parity +# test (skipped when the flag is off) asserts the SAME semantic result by resolving +# class names from `image_metadata[class_names]` keyed by each `class_id`. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason='det["class_name"] subscripting; output is native Detections under ' + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + + +def _native_class_names(predictions) -> list: + """Resolve a native Detections' per-detection class names from its + `image_metadata[class_names]` map keyed by each `class_id` (the same resolution + the block performs on serialisation), mirroring sv `.data["class_name"]`.""" + class_names_map = (predictions.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + class_id = predictions.class_id.detach().to("cpu").numpy() + return [class_names_map.get(int(cid), f"class_{int(cid)}") for cid in class_id] + + DETECTION_CLASSES_REPLACEMENT_WORKFLOW = { "version": "1.0", "inputs": [{"type": "WorkflowImage", "name": "image"}], @@ -77,6 +111,7 @@ workflow_definition=DETECTION_CLASSES_REPLACEMENT_WORKFLOW, workflow_name_in_app="detections-classes-replacement", ) +@_NUMPY_ONLY def test_detection_plus_classification_workflow_when_minimal_valid_input_provided( model_manager: ModelManager, dogs_image: np.ndarray, @@ -171,3 +206,216 @@ def test_detection_plus_classification_workflow_when_minimal_valid_input_provide result[1]["predictions_with_replaced_classes"].xyxy, result[1]["original_predictions"].xyxy, ), "Expected values of other fields in detections to be untouched" + + +@_TENSOR_ONLY +def test_detection_plus_classification_workflow_when_minimal_valid_input_provided_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + crowd_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as ..._tensor_native, but each image arrives ALREADY materialised as a CHW RGB + # device tensor (is_tensor_materialised() == True), so the OD producer runs its + # on-device tensor path. Results must match the numpy-input variant. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTION_CLASSES_REPLACEMENT_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed each fixture as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(crowd_image), + ], + } + ) + + # then + # Under ENABLE_TENSOR_DATA_REPRESENTATION both `original_predictions` and + # `predictions_with_replaced_classes` are native inference_models.Detections + # (torch tensors, no sv `.data` / `__getitem__`). Per-box class names are + # resolved from `image_metadata[class_names]` keyed by `class_id`. Same + # semantic result as the sv `test_...minimal_valid_input_provided` above. + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 3, "Expected 3 element in the output for three input images" + assert set(result[0].keys()) == { + "original_predictions", + "predictions_with_replaced_classes", + }, "Expected all declared outputs to be delivered for first output" + assert set(result[1].keys()) == { + "original_predictions", + "predictions_with_replaced_classes", + }, "Expected all declared outputs to be delivered for second output" + assert set(result[2].keys()) == { + "original_predictions", + "predictions_with_replaced_classes", + }, "Expected all declared outputs to be delivered for third output" + assert ( + len(result[0]["original_predictions"]) + == len(result[0]["predictions_with_replaced_classes"]) + == 2 + ), "Expected 2 dogs detected for first image" + assert ( + len(result[1]["original_predictions"]) + == len(result[1]["predictions_with_replaced_classes"]) + == 2 + ), "Expected 2 dogs detected for second image" + assert ( + len(result[2]["original_predictions"]) + == len(result[2]["predictions_with_replaced_classes"]) + == 0 + ), "Expected 0 dogs detected for third image" + assert ( + result[0]["predictions_with_replaced_classes"].confidence.tolist() + != result[0]["original_predictions"].confidence.tolist() + ), "Expected confidences to be altered" + assert ( + result[0]["predictions_with_replaced_classes"].class_id.tolist() + != result[0]["original_predictions"].class_id.tolist() + ), "Expected class_id to be altered" + assert _native_class_names(result[0]["predictions_with_replaced_classes"]) == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected classes to be changed" + assert _native_class_names(result[0]["original_predictions"]) == [ + "dog", + "dog", + ], "Expected classes not to be changed" + assert _native_class_names(result[1]["predictions_with_replaced_classes"]) == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected classes to be changed" + assert _native_class_names(result[1]["original_predictions"]) == [ + "dog", + "dog", + ], "Expected classes not to be changed" + assert ( + result[0]["predictions_with_replaced_classes"].xyxy + is not result[0]["original_predictions"].xyxy + ), "Expected copy of data to be created by step" + assert ( + result[1]["predictions_with_replaced_classes"].xyxy + is not result[1]["original_predictions"].xyxy + ), "Expected copy of data to be created by step" + assert np.allclose( + result[0]["predictions_with_replaced_classes"].xyxy, + result[0]["original_predictions"].xyxy, + ), "Expected values of other fields in detections to be untouched" + assert np.allclose( + result[1]["predictions_with_replaced_classes"].xyxy, + result[1]["original_predictions"].xyxy, + ), "Expected values of other fields in detections to be untouched" + + +@_TENSOR_ONLY +def test_detection_plus_classification_workflow_when_minimal_valid_input_provided_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + crowd_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTION_CLASSES_REPLACEMENT_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [dogs_image, dogs_image, crowd_image], + } + ) + + # then + # Under ENABLE_TENSOR_DATA_REPRESENTATION both `original_predictions` and + # `predictions_with_replaced_classes` are native inference_models.Detections + # (torch tensors, no sv `.data` / `__getitem__`). Per-box class names are + # resolved from `image_metadata[class_names]` keyed by `class_id`. Same + # semantic result as the sv `test_...minimal_valid_input_provided` above. + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 3, "Expected 3 element in the output for three input images" + assert set(result[0].keys()) == { + "original_predictions", + "predictions_with_replaced_classes", + }, "Expected all declared outputs to be delivered for first output" + assert set(result[1].keys()) == { + "original_predictions", + "predictions_with_replaced_classes", + }, "Expected all declared outputs to be delivered for second output" + assert set(result[2].keys()) == { + "original_predictions", + "predictions_with_replaced_classes", + }, "Expected all declared outputs to be delivered for third output" + assert ( + len(result[0]["original_predictions"]) + == len(result[0]["predictions_with_replaced_classes"]) + == 2 + ), "Expected 2 dogs detected for first image" + assert ( + len(result[1]["original_predictions"]) + == len(result[1]["predictions_with_replaced_classes"]) + == 2 + ), "Expected 2 dogs detected for second image" + assert ( + len(result[2]["original_predictions"]) + == len(result[2]["predictions_with_replaced_classes"]) + == 0 + ), "Expected 0 dogs detected for third image" + assert ( + result[0]["predictions_with_replaced_classes"].confidence.tolist() + != result[0]["original_predictions"].confidence.tolist() + ), "Expected confidences to be altered" + assert ( + result[0]["predictions_with_replaced_classes"].class_id.tolist() + != result[0]["original_predictions"].class_id.tolist() + ), "Expected class_id to be altered" + assert _native_class_names(result[0]["predictions_with_replaced_classes"]) == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected classes to be changed" + assert _native_class_names(result[0]["original_predictions"]) == [ + "dog", + "dog", + ], "Expected classes not to be changed" + assert _native_class_names(result[1]["predictions_with_replaced_classes"]) == [ + "116.Parson_russell_terrier", + "131.Wirehaired_pointing_griffon", + ], "Expected classes to be changed" + assert _native_class_names(result[1]["original_predictions"]) == [ + "dog", + "dog", + ], "Expected classes not to be changed" + assert ( + result[0]["predictions_with_replaced_classes"].xyxy + is not result[0]["original_predictions"].xyxy + ), "Expected copy of data to be created by step" + assert ( + result[1]["predictions_with_replaced_classes"].xyxy + is not result[1]["original_predictions"].xyxy + ), "Expected copy of data to be created by step" + assert np.allclose( + result[0]["predictions_with_replaced_classes"].xyxy, + result[0]["original_predictions"].xyxy, + ), "Expected values of other fields in detections to be untouched" + assert np.allclose( + result[1]["predictions_with_replaced_classes"].xyxy, + result[1]["original_predictions"].xyxy, + ), "Expected values of other fields in detections to be untouched" diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_detections_merge.py b/tests/workflows/integration_tests/execution/test_workflow_with_detections_merge.py index 92417573fd..5abd2cffb6 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_detections_merge.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_detections_merge.py @@ -2,14 +2,39 @@ import pytest import supervision as sv -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, +) from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.object_detection import Detections as NativeDetections +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the detections_merge block emits a native +# inference_models.Detections instead of sv.Detections. The numpy-shaped assertions +# below run only with the flag off; an equivalent *_tensor_native test asserts the same +# facts against the native carrier with the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output; native under ENABLE_TENSOR_DATA_REPRESENTATION " + "โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + DETECTIONS_MERGE_WORKFLOW = { "version": "1.0", "inputs": [ @@ -38,6 +63,7 @@ } +@_NUMPY_ONLY @add_to_workflows_gallery( category="Basic Workflows", use_case_title="Workflow with detections merge", @@ -109,3 +135,131 @@ def test_detections_merge_workflow( box_height = merged_bbox[3] - merged_bbox[1] assert box_width > 100, "Merged box should be reasonably wide" assert box_height > 100, "Merged box should be reasonably tall" + + +@_TENSOR_ONLY +def test_detections_merge_workflow_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTIONS_MERGE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [dogs_image], + } + ) + + # then + assert len(result) == 1, "One set of outputs expected" + assert "result" in result[0], "Output must contain key 'result'" + merged = result[0]["result"] + assert isinstance( + merged, NativeDetections + ), "Output must be instance of native inference_models.Detections" + + # Check that we have exactly one merged detection + assert merged.xyxy.shape[0] == 1, "Should have exactly one merged detection" + + # Check that the merged detection carries the required fields on the native + # carrier: class name lives in image_metadata[CLASS_NAMES_KEY]; detection_id is a + # per-box entry in bboxes_metadata (the native equivalents of sv.Detections.data). + assert ( + CLASS_NAMES_KEY in merged.image_metadata + ), "Should have class_names in metadata" + assert ( + "merged_detection" in merged.image_metadata[CLASS_NAMES_KEY].values() + ), "Merged detection should carry the merged class name" + assert merged.bboxes_metadata is not None, "Should have per-box metadata" + assert ( + DETECTION_ID_KEY in merged.bboxes_metadata[0] + ), "Should have detection_id in per-box metadata" + + # Check that the bounding box has reasonable dimensions + merged_bbox = merged.xyxy[0].tolist() + image_height, image_width = dogs_image.shape[:2] + + # Check that coordinates are within image bounds + assert 0 <= merged_bbox[0] <= image_width, "x1 should be within image bounds" + assert 0 <= merged_bbox[1] <= image_height, "y1 should be within image bounds" + assert 0 <= merged_bbox[2] <= image_width, "x2 should be within image bounds" + assert 0 <= merged_bbox[3] <= image_height, "y2 should be within image bounds" + + # Check that the box has reasonable dimensions + assert merged_bbox[2] > merged_bbox[0], "x2 should be greater than x1" + assert merged_bbox[3] > merged_bbox[1], "y2 should be greater than y1" + + # Check that the box is large enough to likely contain the dogs + box_width = merged_bbox[2] - merged_bbox[0] + box_height = merged_bbox[3] - merged_bbox[1] + assert box_width > 100, "Merged box should be reasonably wide" + assert box_height > 100, "Merged box should be reasonably tall" + + +@_TENSOR_ONLY +def test_detections_merge_workflow_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, +) -> None: + # Same as test_detections_merge_workflow_tensor_native, but the image arrives ALREADY + # materialised as a CHW RGB device tensor (is_tensor_materialised() == True), so the OD + # block runs its on-device tensor path. Results must match the numpy-input variant. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTIONS_MERGE_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the fixture as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": [numpy_image_as_tensor(dogs_image)], + } + ) + + # then + assert len(result) == 1, "One set of outputs expected" + assert "result" in result[0], "Output must contain key 'result'" + merged = result[0]["result"] + assert isinstance( + merged, NativeDetections + ), "Output must be instance of native inference_models.Detections" + assert merged.xyxy.shape[0] == 1, "Should have exactly one merged detection" + assert ( + CLASS_NAMES_KEY in merged.image_metadata + ), "Should have class_names in metadata" + assert ( + "merged_detection" in merged.image_metadata[CLASS_NAMES_KEY].values() + ), "Merged detection should carry the merged class name" + assert merged.bboxes_metadata is not None, "Should have per-box metadata" + assert ( + DETECTION_ID_KEY in merged.bboxes_metadata[0] + ), "Should have detection_id in per-box metadata" + + merged_bbox = merged.xyxy[0].tolist() + image_height, image_width = dogs_image.shape[:2] + assert 0 <= merged_bbox[0] <= image_width, "x1 should be within image bounds" + assert 0 <= merged_bbox[1] <= image_height, "y1 should be within image bounds" + assert 0 <= merged_bbox[2] <= image_width, "x2 should be within image bounds" + assert 0 <= merged_bbox[3] <= image_height, "y2 should be within image bounds" + assert merged_bbox[2] > merged_bbox[0], "x2 should be greater than x1" + assert merged_bbox[3] > merged_bbox[1], "y2 should be greater than y1" + assert merged_bbox[2] - merged_bbox[0] > 100, "Merged box should be reasonably wide" + assert merged_bbox[3] - merged_bbox[1] > 100, "Merged box should be reasonably tall" diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_detections_nearest_neighbor_block.py b/tests/workflows/integration_tests/execution/test_workflow_with_detections_nearest_neighbor_block.py index f1315e95b8..42ac03cd66 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_detections_nearest_neighbor_block.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_detections_nearest_neighbor_block.py @@ -9,8 +9,15 @@ """ import numpy as np +import pytest +from inference_models.models.base.object_detection import ( + Detections as NativeDetections, +) -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.constants import ( @@ -21,6 +28,16 @@ add_to_workflows_gallery, ) +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="Asserts on sv.Detections data columns produced by the numpy path", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="Asserts on native Detections bboxes_metadata produced under " + "ENABLE_TENSOR_DATA_REPRESENTATION", +) + DETECTIONS_NEAREST_NEIGHBOR_WORKFLOW = { "version": "1.0", "inputs": [ @@ -95,6 +112,7 @@ workflow_definition=DETECTIONS_NEAREST_NEIGHBOR_WORKFLOW, workflow_name_in_app="detections-nearest-neighbor", ) +@_NUMPY_ONLY def test_detections_nearest_neighbor_workflow( model_manager: ModelManager, crowd_image: np.ndarray, @@ -159,6 +177,7 @@ def test_detections_nearest_neighbor_workflow( ) +@_NUMPY_ONLY def test_detections_nearest_neighbor_workflow_when_target_set_is_empty( model_manager: ModelManager, crowd_image: np.ndarray, @@ -203,3 +222,118 @@ def test_detections_nearest_neighbor_workflow_when_target_set_is_empty( "With an empty target set every query detection must be unmatched " "(nearest_target_distance=None)." ) + + +@_TENSOR_ONLY +def test_detections_nearest_neighbor_workflow_tensor_native( + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTIONS_NEAREST_NEIGHBOR_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": crowd_image, + "model_id": "yolov8n-640", + "target_confidence": 0.5, + } + ) + + # then + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected one output element for one input image" + assert set(result[0].keys()) == { + "query_predictions", + "matched_query_detections", + "matched_target_detections", + } + + # Under ENABLE_TENSOR_DATA_REPRESENTATION detections-kind outputs are + # native inference_models Detections; per-box scalars ride in + # bboxes_metadata instead of sv data columns. + query_predictions = result[0]["query_predictions"] + matched_query_detections = result[0]["matched_query_detections"] + matched_target_detections = result[0]["matched_target_detections"] + assert isinstance(query_predictions, NativeDetections) + + assert ( + len(query_predictions) > 0 + ), "Expected at least one person detection at confidence=0.3 on the crowd image." + assert len(matched_query_detections) == len(matched_target_detections), ( + "matched_query_detections and matched_target_detections must stay the " + "same length and index-aligned." + ) + assert len(matched_query_detections) > 0, ( + "Expected at least one query detection to find a nearest target " + "detection on the crowd image." + ) + + for entry in query_predictions.bboxes_metadata: + assert NEAREST_TARGET_DISTANCE_KEY in entry + distance = entry[NEAREST_TARGET_DISTANCE_KEY] + assert distance is None or isinstance(distance, float) + + for entry in matched_query_detections.bboxes_metadata: + assert entry[NEAREST_TARGET_DISTANCE_KEY] is not None, ( + "Every detection appearing in matched_query_detections must carry " + "a real (non-None) nearest_target_distance." + ) + + +@_TENSOR_ONLY +def test_detections_nearest_neighbor_workflow_when_target_set_is_empty_tensor_native( + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + # given: target_confidence=0.999 is unreachable on this model/image, so + # target_predictions is empty and no query detection can find a match. + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTIONS_NEAREST_NEIGHBOR_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": crowd_image, + "model_id": "yolov8n-640", + "target_confidence": 0.999, + } + ) + + # then + assert isinstance(result, list) + assert len(result) == 1 + query_predictions = result[0]["query_predictions"] + matched_query_detections = result[0]["matched_query_detections"] + matched_target_detections = result[0]["matched_target_detections"] + assert isinstance(query_predictions, NativeDetections) + + assert len(query_predictions) > 0, ( + "Expected at least one person detection at confidence=0.3 on the crowd " + "image, even though the target set is empty." + ) + assert len(matched_query_detections) == 0 + assert len(matched_target_detections) == 0 + for entry in query_predictions.bboxes_metadata: + assert entry[NEAREST_TARGET_DISTANCE_KEY] is None, ( + "With an empty target set every query detection must be unmatched " + "(nearest_target_distance=None)." + ) diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_detections_rollup.py b/tests/workflows/integration_tests/execution/test_workflow_with_detections_rollup.py index 1eacce4fa2..86d6186e49 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_detections_rollup.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_detections_rollup.py @@ -6,14 +6,35 @@ import numpy as np import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.keypoints_detection import KeyPoints +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the keypoint rollup emits a native +# (KeyPoints, Detections) tuple (the keypoint kind's runtime shape) instead of an +# sv.Detections carrying keypoints in .data. The numpy-shaped assertions run only with +# the flag off; the *_tensor_native twin asserts the same facts on the native tuple. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output; native under ENABLE_TENSOR_DATA_REPRESENTATION " + "โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + # Full workflow with object detection, keypoint, and segmentation rollup FULL_DIMENSION_ROLLUP_WORKFLOW = { "version": "1.0", @@ -598,6 +619,7 @@ def test_dimension_rollup_with_object_detection_only( ), "Visualization should match image dimensions" +@_NUMPY_ONLY @add_to_workflows_gallery( category="Fusion Workflows", use_case_title="Dimension Rollup with Keypoint Detection", @@ -656,6 +678,119 @@ def test_dimension_rollup_with_keypoint_detection_only( ), "Visualization should match image dimensions" +@_TENSOR_ONLY +def test_dimension_rollup_with_keypoint_detection_only_tensor_native( + model_manager: ModelManager, + crowd_image: np.ndarray, + roboflow_api_key: str, +) -> None: + """Tensor-native parity of test_dimension_rollup_with_keypoint_detection_only.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=KEYPOINT_ROLLUP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": crowd_image, + } + ) + + # then + assert isinstance(result, list) + assert len(result) == 1 + assert set(result[0].keys()) == { + "rolled_up_detections", + "keypoint_visualization", + } + + # The keypoint kind's runtime payload is a (KeyPoints, Detections) tuple; the + # KeyPoints component carries the rolled-up keypoint data that lived in + # sv.Detections.data["keypoints_xy"] on the numpy path. + rolled_up = result[0]["rolled_up_detections"] + assert isinstance(rolled_up, tuple), "Keypoint rollup output is a tuple" + key_points, detections = rolled_up + assert isinstance(key_points, KeyPoints) + assert hasattr(detections, "xyxy"), "Bbox component should have xyxy" + # One skeleton instance per bounding box; each instance carries (K, 2) keypoints + # (or there are no instances at all, mirroring the numpy len==0 fallback). + assert key_points.xy.shape[0] == detections.xyxy.shape[0] + assert key_points.xy.shape[0] == 0 or key_points.xy.shape[-1] == 2 + + # Validate visualization + assert ( + result[0]["keypoint_visualization"].numpy_image.shape[:2] + == crowd_image.shape[:2] + ), "Visualization should match image dimensions" + + +@_TENSOR_ONLY +def test_dimension_rollup_with_keypoint_detection_only_with_tensor_input( + model_manager: ModelManager, + crowd_image: np.ndarray, + roboflow_api_key: str, +) -> None: + """Tensor-input parity of test_dimension_rollup_with_keypoint_detection_only. + + Same as the *_tensor_native variant, but the image arrives ALREADY materialised as + a CHW RGB device tensor (is_tensor_materialised() == True), so the OD producer runs + its on-device tensor path. Results must match the numpy-input variant. + """ + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=KEYPOINT_ROLLUP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the fixture as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(crowd_image), + } + ) + + # then + assert isinstance(result, list) + assert len(result) == 1 + assert set(result[0].keys()) == { + "rolled_up_detections", + "keypoint_visualization", + } + + # The keypoint kind's runtime payload is a (KeyPoints, Detections) tuple; the + # KeyPoints component carries the rolled-up keypoint data that lived in + # sv.Detections.data["keypoints_xy"] on the numpy path. + rolled_up = result[0]["rolled_up_detections"] + assert isinstance(rolled_up, tuple), "Keypoint rollup output is a tuple" + key_points, detections = rolled_up + assert isinstance(key_points, KeyPoints) + assert hasattr(detections, "xyxy"), "Bbox component should have xyxy" + # One skeleton instance per bounding box; each instance carries (K, 2) keypoints + # (or there are no instances at all, mirroring the numpy len==0 fallback). + assert key_points.xy.shape[0] == detections.xyxy.shape[0] + assert key_points.xy.shape[0] == 0 or key_points.xy.shape[-1] == 2 + + # Validate visualization + assert ( + result[0]["keypoint_visualization"].numpy_image.shape[:2] + == crowd_image.shape[:2] + ), "Visualization should match image dimensions" + + @add_to_workflows_gallery( category="Fusion Workflows", use_case_title="Dimension Rollup with Instance Segmentation", diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_dimensionality_change.py b/tests/workflows/integration_tests/execution/test_workflow_with_dimensionality_change.py index bb0d550af8..5ef3e5bb30 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_dimensionality_change.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_dimensionality_change.py @@ -4,11 +4,50 @@ import numpy as np import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import PARENT_ID_KEY from inference.core.workflows.execution_engine.core import ExecutionEngine from inference.core.workflows.execution_engine.introspection import blocks_loader +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# The dimensionality stub plugin follows the v1/v1_tensor parity pattern: under +# ENABLE_TENSOR_DATA_REPRESENTATION the loader swaps in tensor-native siblings that +# emit `inference_models.Detections` (per-box state in `bboxes_metadata`, no sv +# `.data`/`__getitem__`). Tests whose assertions read sv-specific things get the +# `_NUMPY_ONLY` original + a `_TENSOR_ONLY` `*_tensor_native` sibling asserting the +# same semantics natively. Tests whose assertions are representation-agnostic +# (image/tile shapes, `len()` which native supports) run under BOTH flags unsplit. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output; the dimensionality stub emits native " + "`inference_models.Detections` under ENABLE_TENSOR_DATA_REPRESENTATION โ€” see " + "the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + + +def _native_parent_ids(detections) -> list: + """Per-box `parent_id` for native detections, mirroring sv `.data["parent_id"]`. + + The tensor-native coordinates block writes `parent_id` per-box into + `bboxes_metadata`; the raw native producer instead carries the crop's parent_id + once in per-image `image_metadata` (per-box holds only `detection_id`). So fall + back to the per-image value broadcast across boxes โ€” this makes the own-side read + the real crop id (matching the numpy path) instead of a vacuous `None`.""" + bboxes_metadata = detections.bboxes_metadata or [] + image_parent_id = (detections.image_metadata or {}).get(PARENT_ID_KEY) + return [entry.get(PARENT_ID_KEY) or image_parent_id for entry in bboxes_metadata] + DETECTIONS_TO_PARENT_COORDINATES_BATCH_VARIANT_WORKFLOW = { "version": "1.0", @@ -64,6 +103,7 @@ } +@_NUMPY_ONLY @mock.patch.object(blocks_loader, "get_plugin_modules") def test_workflow_with_detections_coordinates_transformation_in_batch_variant( get_plugin_modules_mock: MagicMock, @@ -156,6 +196,181 @@ def test_workflow_with_detections_coordinates_transformation_in_batch_variant( ), "Expected parent of each bounding box to point into original input image instead of crop ID" +@_TENSOR_ONLY +@mock.patch.object(blocks_loader, "get_plugin_modules") +def test_workflow_with_detections_coordinates_transformation_in_batch_variant_tensor_native( + get_plugin_modules_mock: MagicMock, + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + """Tensor-native parity of + `test_workflow_with_detections_coordinates_transformation_in_batch_variant`. + + Same scenario, but the stub emits native `inference_models.Detections`. The + `parent_id` the coordinates block writes lives per-box in `bboxes_metadata` + (the native equivalent of sv `.data["parent_id"]`), and `len()` works on the + native object via `__len__`. Asserts the SAME semantics natively. + """ + # given + get_plugin_modules_mock.return_value = [ + "tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin" + ] + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTIONS_TO_PARENT_COORDINATES_BATCH_VARIANT_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": [crowd_image, crowd_image]} + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 2, "Two images provided, so two output elements expected" + assert set(result[0].keys()) == { + "predictions_in_own_coordinates", + "predictions_in_original_coordinates", + }, "Expected all declared outputs to be delivered" + assert set(result[1].keys()) == { + "predictions_in_own_coordinates", + "predictions_in_original_coordinates", + }, "Expected all declared outputs to be delivered" + assert len(result[0]["predictions_in_own_coordinates"]) == len( + result[0]["predictions_in_original_coordinates"] + ), "Expected the same number of nested detections in both outputs for input image 0" + assert len(result[1]["predictions_in_own_coordinates"]) == len( + result[1]["predictions_in_original_coordinates"] + ), "Expected the same number of nested detections in both outputs for input image 1" + for own_coords_detection, original_coords_detection in zip( + result[0]["predictions_in_own_coordinates"], + result[0]["predictions_in_original_coordinates"], + ): + assert len(own_coords_detection) == len( + original_coords_detection + ), "Expected number of bounding boxes in nested Detections not to change" + own_parent_ids = _native_parent_ids(own_coords_detection) + original_parent_ids = _native_parent_ids(original_coords_detection) + assert all( + own is not None and own != original + for own, original in zip(own_parent_ids, original_parent_ids) + ), "Expected parent_id to be modified (own-side crop id present and differing)" + assert original_parent_ids == ["image.[0]"] * len( + original_coords_detection + ), "Expected parent of each bounding box to point into original input image instead of crop ID" + for own_coords_detection, original_coords_detection in zip( + result[1]["predictions_in_own_coordinates"], + result[1]["predictions_in_original_coordinates"], + ): + assert len(own_coords_detection) == len( + original_coords_detection + ), "Expected number of bounding boxes in nested Detections not to change" + own_parent_ids = _native_parent_ids(own_coords_detection) + original_parent_ids = _native_parent_ids(original_coords_detection) + assert all( + own is not None and own != original + for own, original in zip(own_parent_ids, original_parent_ids) + ), "Expected parent_id to be modified (own-side crop id present and differing)" + assert original_parent_ids == ["image.[1]"] * len( + original_coords_detection + ), "Expected parent of each bounding box to point into original input image instead of crop ID" + + +@_TENSOR_ONLY +@mock.patch.object(blocks_loader, "get_plugin_modules") +def test_workflow_with_detections_coordinates_transformation_in_batch_variant_with_tensor_input( + get_plugin_modules_mock: MagicMock, + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + """Same as + test_workflow_with_detections_coordinates_transformation_in_batch_variant_tensor_native, + but each input image arrives ALREADY materialised as a CHW RGB device tensor + (is_tensor_materialised() == True), so the OD block runs its on-device tensor path. + Asserts the SAME native semantics as the numpy-input variant. + """ + # given + get_plugin_modules_mock.return_value = [ + "tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin" + ] + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTIONS_TO_PARENT_COORDINATES_BATCH_VARIANT_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the fixtures as pre-materialised tensors + result = execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(crowd_image), + numpy_image_as_tensor(crowd_image), + ] + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 2, "Two images provided, so two output elements expected" + assert set(result[0].keys()) == { + "predictions_in_own_coordinates", + "predictions_in_original_coordinates", + }, "Expected all declared outputs to be delivered" + assert set(result[1].keys()) == { + "predictions_in_own_coordinates", + "predictions_in_original_coordinates", + }, "Expected all declared outputs to be delivered" + assert len(result[0]["predictions_in_own_coordinates"]) == len( + result[0]["predictions_in_original_coordinates"] + ), "Expected the same number of nested detections in both outputs for input image 0" + assert len(result[1]["predictions_in_own_coordinates"]) == len( + result[1]["predictions_in_original_coordinates"] + ), "Expected the same number of nested detections in both outputs for input image 1" + for own_coords_detection, original_coords_detection in zip( + result[0]["predictions_in_own_coordinates"], + result[0]["predictions_in_original_coordinates"], + ): + assert len(own_coords_detection) == len( + original_coords_detection + ), "Expected number of bounding boxes in nested Detections not to change" + own_parent_ids = _native_parent_ids(own_coords_detection) + original_parent_ids = _native_parent_ids(original_coords_detection) + assert all( + own is not None and own != original + for own, original in zip(own_parent_ids, original_parent_ids) + ), "Expected parent_id to be modified (own-side crop id present and differing)" + assert original_parent_ids == ["image.[0]"] * len( + original_coords_detection + ), "Expected parent of each bounding box to point into original input image instead of crop ID" + for own_coords_detection, original_coords_detection in zip( + result[1]["predictions_in_own_coordinates"], + result[1]["predictions_in_original_coordinates"], + ): + assert len(own_coords_detection) == len( + original_coords_detection + ), "Expected number of bounding boxes in nested Detections not to change" + own_parent_ids = _native_parent_ids(own_coords_detection) + original_parent_ids = _native_parent_ids(original_coords_detection) + assert all( + own is not None and own != original + for own, original in zip(own_parent_ids, original_parent_ids) + ), "Expected parent_id to be modified (own-side crop id present and differing)" + assert original_parent_ids == ["image.[1]"] * len( + original_coords_detection + ), "Expected parent of each bounding box to point into original input image instead of crop ID" + + DETECTIONS_TO_PARENT_COORDINATES_NON_BATCH_VARIANT_WORKFLOW = { "version": "1.0", "inputs": [{"type": "WorkflowImage", "name": "image"}], @@ -210,6 +425,7 @@ def test_workflow_with_detections_coordinates_transformation_in_batch_variant( } +@_NUMPY_ONLY @mock.patch.object(blocks_loader, "get_plugin_modules") def test_workflow_with_detections_coordinates_transformation_in_non_batch_variant( get_plugin_modules_mock: MagicMock, @@ -302,6 +518,179 @@ def test_workflow_with_detections_coordinates_transformation_in_non_batch_varian ), "Expected parent of each bounding box to point into original input image instead of crop ID" +@_TENSOR_ONLY +@mock.patch.object(blocks_loader, "get_plugin_modules") +def test_workflow_with_detections_coordinates_transformation_in_non_batch_variant_tensor_native( + get_plugin_modules_mock: MagicMock, + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + """Tensor-native parity of + `test_workflow_with_detections_coordinates_transformation_in_non_batch_variant`. + + Same scenario on native `inference_models.Detections`: `parent_id` is read + per-box from `bboxes_metadata`, `len()` via native `__len__`. + """ + # given + get_plugin_modules_mock.return_value = [ + "tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin" + ] + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTIONS_TO_PARENT_COORDINATES_NON_BATCH_VARIANT_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": [crowd_image, crowd_image]} + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 2, "Two images provided, so two output elements expected" + assert set(result[0].keys()) == { + "predictions_in_own_coordinates", + "predictions_in_original_coordinates", + }, "Expected all declared outputs to be delivered" + assert set(result[1].keys()) == { + "predictions_in_own_coordinates", + "predictions_in_original_coordinates", + }, "Expected all declared outputs to be delivered" + assert len(result[0]["predictions_in_own_coordinates"]) == len( + result[0]["predictions_in_original_coordinates"] + ), "Expected the same number of nested detections in both outputs for input image 0" + assert len(result[1]["predictions_in_own_coordinates"]) == len( + result[1]["predictions_in_original_coordinates"] + ), "Expected the same number of nested detections in both outputs for input image 1" + for own_coords_detection, original_coords_detection in zip( + result[0]["predictions_in_own_coordinates"], + result[0]["predictions_in_original_coordinates"], + ): + assert len(own_coords_detection) == len( + original_coords_detection + ), "Expected number of bounding boxes in nested Detections not to change" + own_parent_ids = _native_parent_ids(own_coords_detection) + original_parent_ids = _native_parent_ids(original_coords_detection) + assert all( + own is not None and own != original + for own, original in zip(own_parent_ids, original_parent_ids) + ), "Expected parent_id to be modified (own-side crop id present and differing)" + assert original_parent_ids == ["image.[0]"] * len( + original_coords_detection + ), "Expected parent of each bounding box to point into original input image instead of crop ID" + for own_coords_detection, original_coords_detection in zip( + result[1]["predictions_in_own_coordinates"], + result[1]["predictions_in_original_coordinates"], + ): + assert len(own_coords_detection) == len( + original_coords_detection + ), "Expected number of bounding boxes in nested Detections not to change" + own_parent_ids = _native_parent_ids(own_coords_detection) + original_parent_ids = _native_parent_ids(original_coords_detection) + assert all( + own is not None and own != original + for own, original in zip(own_parent_ids, original_parent_ids) + ), "Expected parent_id to be modified (own-side crop id present and differing)" + assert original_parent_ids == ["image.[1]"] * len( + original_coords_detection + ), "Expected parent of each bounding box to point into original input image instead of crop ID" + + +@_TENSOR_ONLY +@mock.patch.object(blocks_loader, "get_plugin_modules") +def test_workflow_with_detections_coordinates_transformation_in_non_batch_variant_with_tensor_input( + get_plugin_modules_mock: MagicMock, + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + """Same as + test_workflow_with_detections_coordinates_transformation_in_non_batch_variant_tensor_native, + but each input image arrives ALREADY materialised as a CHW RGB device tensor + (is_tensor_materialised() == True), so the OD block runs its on-device tensor path. + Asserts the SAME native semantics as the numpy-input variant. + """ + # given + get_plugin_modules_mock.return_value = [ + "tests.workflows.integration_tests.execution.stub_plugins.dimensionality_manipulation_plugin" + ] + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=DETECTIONS_TO_PARENT_COORDINATES_NON_BATCH_VARIANT_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the fixtures as pre-materialised tensors + result = execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(crowd_image), + numpy_image_as_tensor(crowd_image), + ] + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 2, "Two images provided, so two output elements expected" + assert set(result[0].keys()) == { + "predictions_in_own_coordinates", + "predictions_in_original_coordinates", + }, "Expected all declared outputs to be delivered" + assert set(result[1].keys()) == { + "predictions_in_own_coordinates", + "predictions_in_original_coordinates", + }, "Expected all declared outputs to be delivered" + assert len(result[0]["predictions_in_own_coordinates"]) == len( + result[0]["predictions_in_original_coordinates"] + ), "Expected the same number of nested detections in both outputs for input image 0" + assert len(result[1]["predictions_in_own_coordinates"]) == len( + result[1]["predictions_in_original_coordinates"] + ), "Expected the same number of nested detections in both outputs for input image 1" + for own_coords_detection, original_coords_detection in zip( + result[0]["predictions_in_own_coordinates"], + result[0]["predictions_in_original_coordinates"], + ): + assert len(own_coords_detection) == len( + original_coords_detection + ), "Expected number of bounding boxes in nested Detections not to change" + own_parent_ids = _native_parent_ids(own_coords_detection) + original_parent_ids = _native_parent_ids(original_coords_detection) + assert all( + own is not None and own != original + for own, original in zip(own_parent_ids, original_parent_ids) + ), "Expected parent_id to be modified (own-side crop id present and differing)" + assert original_parent_ids == ["image.[0]"] * len( + original_coords_detection + ), "Expected parent of each bounding box to point into original input image instead of crop ID" + for own_coords_detection, original_coords_detection in zip( + result[1]["predictions_in_own_coordinates"], + result[1]["predictions_in_original_coordinates"], + ): + assert len(own_coords_detection) == len( + original_coords_detection + ), "Expected number of bounding boxes in nested Detections not to change" + own_parent_ids = _native_parent_ids(own_coords_detection) + original_parent_ids = _native_parent_ids(original_coords_detection) + assert all( + own is not None and own != original + for own, original in zip(own_parent_ids, original_parent_ids) + ), "Expected parent_id to be modified (own-side crop id present and differing)" + assert original_parent_ids == ["image.[1]"] * len( + original_coords_detection + ), "Expected parent of each bounding box to point into original input image instead of crop ID" + + DETECTIONS_STITCHING_BATCH_VARIANT_WORKFLOW = { "version": "1.0", "inputs": [{"type": "WorkflowImage", "name": "image"}], @@ -478,7 +867,7 @@ def test_workflow_with_detections_stitching_in_batch_variant( @mock.patch.object(blocks_loader, "get_plugin_modules") -def test_workflow_with_detections_stitching_in_batch_variant( +def test_workflow_with_detections_stitching_in_non_batch_variant( get_plugin_modules_mock: MagicMock, model_manager: ModelManager, crowd_image: np.ndarray, diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_masked_crop.py b/tests/workflows/integration_tests/execution/test_workflow_with_masked_crop.py index 32e06fac47..724f3c5fb6 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_masked_crop.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_masked_crop.py @@ -1,13 +1,38 @@ import numpy as np +import pytest -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.core import ExecutionEngine +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION, `$steps.segmentation.predictions` is a +# native `inference_models.InstanceDetections` (torch-backed @dataclass) rather than +# an `sv.Detections`. The sv/numpy-typed tests below assert numpy semantics directly +# on it (`predictions.xyxy[0].round().astype(int)`, `predictions.mask[...]`), which +# fails on torch tensors. Those tests are skipped when the flag is on; each has a +# `*_tensor_native` parity test (skipped when the flag is off) that recovers numpy via +# `predictions.to_supervision()` and asserts the same semantic result. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output asserted via .xyxy.astype/.mask numpy ops; " + "native-only under ENABLE_TENSOR_DATA_REPRESENTATION โ€” see *_tensor_native " + "parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + MASKED_CROP_LEGACY_WORKFLOW = { "version": "1.0", "inputs": [ @@ -55,6 +80,7 @@ } +@_NUMPY_ONLY def test_legacy_workflow_with_masked_crop( model_manager: ModelManager, dogs_image: np.ndarray, @@ -100,6 +126,104 @@ def test_legacy_workflow_with_masked_crop( assert pixels_sum == 0, "Expected everything black outside mask" +@_TENSOR_ONLY +def test_legacy_workflow_with_masked_crop_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=MASKED_CROP_LEGACY_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "crops", + "predictions", + }, "Expected all declared outputs to be delivered" + assert len(result[0]["crops"]) == 2, "Expected 2 crops for two dogs detected" + crop_image = result[0]["crops"][0].numpy_image + # Native InstanceDetections -> sv.Detections recovers numpy xyxy (N, 4) and a + # (N, H, W) bool mask, matching the numpy baseline the sv test asserts. + sv_predictions = result[0]["predictions"].to_supervision() + x_min, y_min, x_max, y_max = sv_predictions.xyxy[0].round().astype(dtype=int) + crop_mask = sv_predictions.mask[0][y_min:y_max, x_min:x_max] + pixels_outside_mask = np.where( + np.stack([crop_mask] * 3, axis=-1) == 0, + crop_image, + np.zeros_like(crop_image), + ) + pixels_sum = pixels_outside_mask.sum() + assert pixels_sum == 0, "Expected everything black outside mask" + + +@_TENSOR_ONLY +def test_legacy_workflow_with_masked_crop_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as test_legacy_workflow_with_masked_crop_tensor_native, but the image arrives + # ALREADY materialised as a CHW RGB device tensor (is_tensor_materialised() == True), + # so the segmentation block runs its on-device tensor path. Results must match the + # numpy-input variant. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=MASKED_CROP_LEGACY_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the fixture as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "crops", + "predictions", + }, "Expected all declared outputs to be delivered" + assert len(result[0]["crops"]) == 2, "Expected 2 crops for two dogs detected" + crop_image = result[0]["crops"][0].numpy_image + # Native InstanceDetections -> sv.Detections recovers numpy xyxy (N, 4) and a + # (N, H, W) bool mask, matching the numpy baseline the sv test asserts. + sv_predictions = result[0]["predictions"].to_supervision() + x_min, y_min, x_max, y_max = sv_predictions.xyxy[0].round().astype(dtype=int) + crop_mask = sv_predictions.mask[0][y_min:y_max, x_min:x_max] + pixels_outside_mask = np.where( + np.stack([crop_mask] * 3, axis=-1) == 0, + crop_image, + np.zeros_like(crop_image), + ) + pixels_sum = pixels_outside_mask.sum() + assert pixels_sum == 0, "Expected everything black outside mask" + + MASKED_CROP_WORKFLOW = { "version": "1.0", "inputs": [ @@ -147,6 +271,7 @@ def test_legacy_workflow_with_masked_crop( } +@_NUMPY_ONLY @add_to_workflows_gallery( category="Workflows with data transformations", use_case_title="Instance Segmentation results with background subtracted", @@ -202,6 +327,104 @@ def test_workflow_with_masked_crop( assert pixels_sum == 0, "Expected everything black outside mask" +@_TENSOR_ONLY +def test_workflow_with_masked_crop_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=MASKED_CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": dogs_image, + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "crops", + "predictions", + }, "Expected all declared outputs to be delivered" + assert len(result[0]["crops"]) == 2, "Expected 2 crops for two dogs detected" + crop_image = result[0]["crops"][0].numpy_image + # Native InstanceDetections -> sv.Detections recovers numpy xyxy (N, 4) and a + # (N, H, W) bool mask, matching the numpy baseline the sv test asserts. + sv_predictions = result[0]["predictions"].to_supervision() + x_min, y_min, x_max, y_max = sv_predictions.xyxy[0].round().astype(dtype=int) + crop_mask = sv_predictions.mask[0][y_min:y_max, x_min:x_max] + pixels_outside_mask = np.where( + np.stack([crop_mask] * 3, axis=-1) == 0, + crop_image, + np.zeros_like(crop_image), + ) + pixels_sum = pixels_outside_mask.sum() + assert pixels_sum == 0, "Expected everything black outside mask" + + +@_TENSOR_ONLY +def test_workflow_with_masked_crop_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as test_workflow_with_masked_crop_tensor_native, but the image arrives ALREADY + # materialised as a CHW RGB device tensor (is_tensor_materialised() == True), so the + # segmentation block runs its on-device tensor path. Results must match the numpy-input + # variant. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=MASKED_CROP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed the fixture as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(dogs_image), + } + ) + + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "crops", + "predictions", + }, "Expected all declared outputs to be delivered" + assert len(result[0]["crops"]) == 2, "Expected 2 crops for two dogs detected" + crop_image = result[0]["crops"][0].numpy_image + # Native InstanceDetections -> sv.Detections recovers numpy xyxy (N, 4) and a + # (N, H, W) bool mask, matching the numpy baseline the sv test asserts. + sv_predictions = result[0]["predictions"].to_supervision() + x_min, y_min, x_max, y_max = sv_predictions.xyxy[0].round().astype(dtype=int) + crop_mask = sv_predictions.mask[0][y_min:y_max, x_min:x_max] + pixels_outside_mask = np.where( + np.stack([crop_mask] * 3, axis=-1) == 0, + crop_image, + np.zeros_like(crop_image), + ) + pixels_sum = pixels_outside_mask.sum() + assert pixels_sum == 0, "Expected everything black outside mask" + + def test_workflow_with_masked_crop_when_nothing_gets_predicted( model_manager: ModelManager, dogs_image: np.ndarray, diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_ocr_detections_stitching.py b/tests/workflows/integration_tests/execution/test_workflow_with_ocr_detections_stitching.py index 7e3dbc19bd..c4f4a03e1b 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_ocr_detections_stitching.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_ocr_detections_stitching.py @@ -89,7 +89,6 @@ def test_detection_plus_classification_workflow_when_minimal_valid_input_provide "confidence": 0.6, } ) - assert isinstance(result, list), "Expected list to be delivered" assert len(result) == 1, "Expected 1 element in the output for one input image" assert set(result[0].keys()) == { diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_opc_writer.py b/tests/workflows/integration_tests/execution/test_workflow_with_opc_writer.py index 6f0057fba4..de78c9e6d4 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_opc_writer.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_opc_writer.py @@ -1,4 +1,5 @@ import asyncio +import os import threading import time from typing import Optional, Union @@ -24,6 +25,7 @@ get_connection_manager, opc_connect_and_write_value, ) +from tests.workflows.integration_tests.execution.conftest import bool_env from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) @@ -395,6 +397,10 @@ def test_opc_server(): workflow_definition=WORKFLOW_OPC_WRITER, workflow_name_in_app="opc_writer", ) +@pytest.mark.skipif( + bool_env(os.getenv("SKIP_OPC_WRITER_TEST", False)), + reason="Skipping OPC writer test due to flag `SKIP_OPC_WRITER_TEST`", +) @pytest.mark.timeout(10) @pytest.mark.parametrize( "value_type,variable_name,test_value,expected_value", @@ -475,6 +481,10 @@ def test_workflow_with_opc_writer_sink( assert result_value == expected_value +@pytest.mark.skipif( + bool_env(os.getenv("SKIP_OPC_WRITER_TEST", False)), + reason="Skipping OPC writer test due to flag `SKIP_OPC_WRITER_TEST`", +) @pytest.mark.timeout(5) def test_workflow_with_opc_writer_sink_direct_mode(test_opc_server) -> None: """Test OPC writer with direct NodeId lookup mode using the shared server""" @@ -534,6 +544,10 @@ def reset_connection_manager(): manager.close_all() +@pytest.mark.skipif( + bool_env(os.getenv("SKIP_OPC_WRITER_TEST", False)), + reason="Skipping OPC writer test due to flag `SKIP_OPC_WRITER_TEST`", +) @pytest.mark.timeout(10) def test_connection_pooling_reuses_connections( test_opc_server, reset_connection_manager @@ -594,6 +608,10 @@ def test_connection_pooling_reuses_connections( assert result_value == 200 +@pytest.mark.skipif( + bool_env(os.getenv("SKIP_OPC_WRITER_TEST", False)), + reason="Skipping OPC writer test due to flag `SKIP_OPC_WRITER_TEST`", +) @pytest.mark.timeout(10) def test_connection_invalidation_creates_new_connection( test_opc_server, reset_connection_manager @@ -643,6 +661,10 @@ def test_connection_invalidation_creates_new_connection( assert stats["total_connections"] == 1 +@pytest.mark.skipif( + bool_env(os.getenv("SKIP_OPC_WRITER_TEST", False)), + reason="Skipping OPC writer test due to flag `SKIP_OPC_WRITER_TEST`", +) @pytest.mark.timeout(10) def test_connection_failure_fails_fast(reset_connection_manager) -> None: """Test that connection failures fail fast without blocking.""" @@ -667,6 +689,10 @@ def test_connection_failure_fails_fast(reset_connection_manager) -> None: assert "Failed to write" in message +@pytest.mark.skipif( + bool_env(os.getenv("SKIP_OPC_WRITER_TEST", False)), + reason="Skipping OPC writer test due to flag `SKIP_OPC_WRITER_TEST`", +) @pytest.mark.timeout(10) def test_different_servers_get_different_connections( test_opc_server, reset_connection_manager @@ -711,6 +737,10 @@ def test_different_servers_get_different_connections( assert "AUTH ERROR" in message2 +@pytest.mark.skipif( + bool_env(os.getenv("SKIP_OPC_WRITER_TEST", False)), + reason="Skipping OPC writer test due to flag `SKIP_OPC_WRITER_TEST`", +) @pytest.mark.timeout(10) def test_close_all_connections(test_opc_server, reset_connection_manager) -> None: """Test that close_all properly closes all pooled connections.""" diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_overlap.py b/tests/workflows/integration_tests/execution/test_workflow_with_overlap.py index b34074f90f..c359e15ee8 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_overlap.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_overlap.py @@ -1,13 +1,47 @@ import numpy as np +import pytest -from inference.core.env import USE_INFERENCE_MODELS, WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + USE_INFERENCE_MODELS, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.core import ExecutionEngine +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the overlap block returns a native +# inference_models.Detections (a plain @dataclass of torch tensors) which has no +# sv `.data["class_name"]`. The sv-shaped test below is skipped when the flag is +# on; the `*_tensor_native` parity test (skipped when the flag is off) asserts the +# SAME semantic result by resolving class names from `image_metadata[class_names]`. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="overlap block returns native Detections (no .data) under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + + +def _native_class_names(predictions) -> list: + """Resolve a native Detections' per-detection class names from its + `image_metadata[class_names]` map keyed by each `class_id` (the same + resolution the overlap block performs), mirroring sv `.data["class_name"]`.""" + class_names_map = (predictions.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + class_id = predictions.class_id.detach().to("cpu").numpy() + return [class_names_map.get(int(cid), f"class_{int(cid)}") for cid in class_id] + + # check both overlap types OVERLAP_WORKFLOW = { "version": "1.0", @@ -64,6 +98,7 @@ ''' +@_NUMPY_ONLY def test_workflow_with_overlap_all( model_manager: ModelManager, fruit_image: np.ndarray, @@ -121,3 +156,102 @@ def test_workflow_with_overlap_all( assert "banana" not in class_names assert "apple" in class_names assert "orange" in class_names + + +@_TENSOR_ONLY +def test_workflow_with_overlap_all_tensor_native( + model_manager: ModelManager, + fruit_image: np.ndarray, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=OVERLAP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": fruit_image, + } + ) + + # then + assert len(result) == 1, "One set of images provided, so one output expected" + + # The overlap block returns a native inference_models.Detections (torch tensors, + # no sv `.data`); class names are resolved from `image_metadata[class_names]`. + # Same semantic result as the sv `test_workflow_with_overlap_all` else-branch. + + # if overlap_type is "Any Overlap", both the apples and orange will overlap the banana + any_redictions = result[0]["any_predictions"] + assert len(any_redictions.class_id) == 5 + class_names = _native_class_names(any_redictions) + assert "banana" not in class_names + assert "apple" in class_names + assert "orange" in class_names + + # if overlap_type is "Center Overlap" the apple and orange will overlap + any_redictions = result[0]["center_predictions"] + assert len(any_redictions.class_id) == 2 + class_names = _native_class_names(any_redictions) + assert "banana" not in class_names + assert "apple" in class_names + assert "orange" in class_names + + +@_TENSOR_ONLY +def test_workflow_with_overlap_all_with_tensor_input( + model_manager: ModelManager, + fruit_image: np.ndarray, +) -> None: + # Same as test_workflow_with_overlap_all_tensor_native, but the image arrives ALREADY + # materialised as a CHW RGB device tensor (is_tensor_materialised() == True), so the + # OD block runs its on-device tensor path. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=OVERLAP_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(fruit_image), + } + ) + + # then + assert len(result) == 1, "One set of images provided, so one output expected" + + # The overlap block returns a native inference_models.Detections (torch tensors, + # no sv `.data`); class names are resolved from `image_metadata[class_names]`. + # Same semantic result as the sv `test_workflow_with_overlap_all` else-branch. + + # if overlap_type is "Any Overlap", both the apples and orange will overlap the banana + any_redictions = result[0]["any_predictions"] + assert len(any_redictions.class_id) == 5 + class_names = _native_class_names(any_redictions) + assert "banana" not in class_names + assert "apple" in class_names + assert "orange" in class_names + + # if overlap_type is "Center Overlap" the apple and orange will overlap + any_redictions = result[0]["center_predictions"] + assert len(any_redictions.class_id) == 2 + class_names = _native_class_names(any_redictions) + assert "banana" not in class_names + assert "apple" in class_names + assert "orange" in class_names diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_per_class_confidence_filter.py b/tests/workflows/integration_tests/execution/test_workflow_with_per_class_confidence_filter.py index 1f8e1d12aa..1f7c40c049 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_per_class_confidence_filter.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_per_class_confidence_filter.py @@ -1,10 +1,38 @@ import numpy as np +import pytest import supervision as sv +import torch -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + DETECTION_ID_KEY, +) from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.object_detection import Detections as NativeDetections + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the object_detection_prediction deserializer +# accepts only a native inference_models.Detections (or a serialised dict), not a raw +# sv.Detections, and the per_class_confidence_filter block emits a native Detections +# instead of sv.Detections. The numpy-shaped tests below (which feed sv.Detections and +# assert sv attributes like .data / .confidence) run only with the flag off; an +# equivalent *_tensor_native test asserts the same semantic filtering result against the +# native carrier with the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections input/output; native under ENABLE_TENSOR_DATA_REPRESENTATION " + "โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) PER_CLASS_CONFIDENCE_FILTER_WORKFLOW = { "version": "1.3.0", @@ -53,6 +81,45 @@ def _make_detections(class_names: list[str], confidences: list[float]) -> sv.Det ) +def _make_native_detections( + class_names: list[str], confidences: list[float] +) -> NativeDetections: + """Native inference_models.Detections equivalent of ``_make_detections``. + + Per-box class name is carried both via the canonical + ``image_metadata[CLASS_NAMES_KEY]`` (class_id -> name) map and via + ``bboxes_metadata[i][CLASS_NAME_KEY]`` โ€” the two sources the block's + ``_resolve_class_names`` reads (native equivalents of sv.Detections.data). + """ + n = len(class_names) + name_to_id: dict[str, int] = {} + for name in class_names: + name_to_id.setdefault(name, len(name_to_id)) + class_id = torch.tensor( + [name_to_id[name] for name in class_names], dtype=torch.int64 + ) + name_by_class_id = {class_id: name for name, class_id in name_to_id.items()} + return NativeDetections( + xyxy=torch.tensor([[0, 0, 10, 10]] * n, dtype=torch.float32), + class_id=class_id, + confidence=torch.tensor(confidences, dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: name_by_class_id}, + bboxes_metadata=[ + {CLASS_NAME_KEY: class_names[i], DETECTION_ID_KEY: f"d{i}"} + for i in range(n) + ], + ) + + +def _native_class_names(detections: NativeDetections) -> list[str]: + name_by_class_id = detections.image_metadata.get(CLASS_NAMES_KEY) or {} + return [ + name_by_class_id.get(int(detections.class_id[i])) + for i in range(int(detections.confidence.shape[0])) + ] + + +@_NUMPY_ONLY def test_per_class_confidence_filter_end_to_end( model_manager: ModelManager, ) -> None: @@ -89,6 +156,7 @@ def test_per_class_confidence_filter_end_to_end( assert list(filtered.confidence) == [0.99, 0.6] +@_NUMPY_ONLY def test_per_class_confidence_filter_default_threshold_filters_unknown_class( model_manager: ModelManager, ) -> None: @@ -123,6 +191,7 @@ def test_per_class_confidence_filter_default_threshold_filters_unknown_class( assert list(filtered.confidence) == [0.8] +@_NUMPY_ONLY def test_per_class_confidence_filter_handles_batch_of_images( model_manager: ModelManager, ) -> None: @@ -155,3 +224,114 @@ def test_per_class_confidence_filter_handles_batch_of_images( assert len(result) == 2 assert list(result[0]["filtered"].data["class_name"]) == ["person"] assert list(result[1]["filtered"].data["class_name"]) == ["car"] + + +@_TENSOR_ONLY +def test_per_class_confidence_filter_end_to_end_tensor_native( + model_manager: ModelManager, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=PER_CLASS_CONFIDENCE_FILTER_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + predictions = _make_native_detections( + class_names=["person", "person", "car", "dog"], + confidences=[0.99, 0.7, 0.6, 0.4], + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "predictions": [predictions], + "class_thresholds": {"person": 0.98, "car": 0.5}, + "default_threshold": 0.5, + } + ) + + # then + assert isinstance(result, list) + assert len(result) == 1 + filtered = result[0]["filtered"] + assert isinstance(filtered, NativeDetections) + assert _native_class_names(filtered) == ["person", "car"] + assert filtered.confidence.tolist() == pytest.approx([0.99, 0.6]) + + +@_TENSOR_ONLY +def test_per_class_confidence_filter_default_threshold_filters_unknown_class_tensor_native( + model_manager: ModelManager, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=PER_CLASS_CONFIDENCE_FILTER_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + predictions = _make_native_detections( + class_names=["cat", "cat"], + confidences=[0.2, 0.8], + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "predictions": [predictions], + "class_thresholds": {"person": 0.98}, + "default_threshold": 0.5, + } + ) + + # then + filtered = result[0]["filtered"] + assert isinstance(filtered, NativeDetections) + assert _native_class_names(filtered) == ["cat"] + assert filtered.confidence.tolist() == pytest.approx([0.8]) + + +@_TENSOR_ONLY +def test_per_class_confidence_filter_handles_batch_of_images_tensor_native( + model_manager: ModelManager, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=PER_CLASS_CONFIDENCE_FILTER_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + predictions_image_1 = _make_native_detections( + class_names=["person", "car"], confidences=[0.99, 0.4] + ) + predictions_image_2 = _make_native_detections( + class_names=["car"], confidences=[0.55] + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "predictions": [predictions_image_1, predictions_image_2], + "class_thresholds": {"person": 0.98, "car": 0.5}, + "default_threshold": 0.3, + } + ) + + # then + assert len(result) == 2 + assert _native_class_names(result[0]["filtered"]) == ["person"] + assert _native_class_names(result[1]["filtered"]) == ["car"] diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_pp_ocr.py b/tests/workflows/integration_tests/execution/test_workflow_with_pp_ocr.py index 1a13a32abf..206ffa56e9 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_pp_ocr.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_pp_ocr.py @@ -2,10 +2,30 @@ import pytest import supervision as sv -from inference.core.env import USE_INFERENCE_MODELS, WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + USE_INFERENCE_MODELS, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.object_detection import Detections as NativeDetections + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the pp_ocr block emits its `predictions` +# output as a native inference_models.Detections instead of sv.Detections. The +# isinstance check below is numpy-shaped, so it runs only with the flag off; an +# equivalent *_tensor_native test asserts the same facts against the native carrier +# with the flag on. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output; native under ENABLE_TENSOR_DATA_REPRESENTATION " + "โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) PP_OCR_WORKFLOW = { "version": "1.0", @@ -36,6 +56,7 @@ } +@_NUMPY_ONLY @pytest.mark.skipif( not USE_INFERENCE_MODELS, reason="PP-OCR is backed by the inference-models package", @@ -76,6 +97,52 @@ def test_pp_ocr_workflow_extracts_text( assert len(detections) > 0, "Expected text lines to be detected" +@_TENSOR_ONLY +@pytest.mark.skipif( + not USE_INFERENCE_MODELS, + reason="PP-OCR is backed by the inference-models package", +) +def test_pp_ocr_workflow_extracts_text_tensor_native( + model_manager: ModelManager, + multi_line_text_image: np.ndarray, +) -> None: + # Parity sibling of test_pp_ocr_workflow_extracts_text: under + # ENABLE_TENSOR_DATA_REPRESENTATION the pp_ocr `predictions` output is a native + # inference_models.Detections rather than sv.Detections. Same transcription facts. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=PP_OCR_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": multi_line_text_image, + } + ) + + # then + assert isinstance(result, list), "Expected list to be delivered" + assert len(result) == 1, "Expected 1 element in the output for one input image" + assert set(result[0].keys()) == { + "extracted_text", + "text_detections", + }, "Expected all declared outputs to be delivered" + normalized_text = result[0]["extracted_text"].upper().replace(" ", "") + assert "GREAT" in normalized_text, "Expected the text to be transcribed" + detections = result[0]["text_detections"] + assert isinstance( + detections, NativeDetections + ), "Output must be instance of native inference_models.Detections" + assert len(detections) > 0, "Expected text lines to be detected" + + @pytest.mark.skipif( not USE_INFERENCE_MODELS, reason="PP-OCR is backed by the inference-models package", diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_property_extraction.py b/tests/workflows/integration_tests/execution/test_workflow_with_property_extraction.py index 9e5781452b..d86f0976c6 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_property_extraction.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_property_extraction.py @@ -6,7 +6,10 @@ import pytest import supervision as sv -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.interfaces.camera.video_source import VideoSource from inference.core.interfaces.stream.entities import VideoFrame from inference.core.interfaces.stream.inference_pipeline import InferencePipeline @@ -14,10 +17,29 @@ from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.object_detection import Detections as NativeDetections +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# DimensionCollapse is type-agnostic; under ENABLE_TENSOR_DATA_REPRESENTATION the +# RoboflowObjectDetectionModel producer emits native inference_models.Detections, +# so the collapsed elements are native, not sv.Detections. The sv-typed assertion +# test is skipped when the flag is on; the *_tensor_native sibling asserts the same +# collapse with native element types. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="asserts sv.Detections elements; producer is native under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + WORKFLOW_WITH_EXTRACTION_OF_CLASSES_FOR_DETECTIONS = { "version": "1.0", "inputs": [ @@ -713,6 +735,7 @@ def test_workflow_with_pass_fail_applied_for_each_ocr_result( } +@_NUMPY_ONLY def test_workflow_when_there_is_faulty_application_of_aggregation_step_at_batch_with_dimension_1( model_manager: ModelManager, license_plate_image: np.ndarray, @@ -749,6 +772,87 @@ def test_workflow_when_there_is_faulty_application_of_aggregation_step_at_batch_ assert isinstance(result[0]["result"][1], sv.Detections) +@_TENSOR_ONLY +def test_workflow_when_there_is_faulty_application_of_aggregation_step_at_batch_with_dimension_1_tensor_native( + model_manager: ModelManager, + license_plate_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_INVALID_AGGREGATION, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [ + np.zeros((192, 168, 3), dtype=np.uint8), + np.zeros((200, 168, 3), dtype=np.uint8), + ] + } + ) + + # then + assert len(result) == 1, "Expected result to collapse" + assert ( + len(result[0]["result"]) == 2 + ), "Expected both predictions to be placed in the list" + # DimensionCollapse preserves element type; under the flag the producer emits + # native inference_models.Detections, so collapsed elements are native. + assert isinstance(result[0]["result"][0], NativeDetections) + assert isinstance(result[0]["result"][1], NativeDetections) + + +@_TENSOR_ONLY +def test_workflow_when_there_is_faulty_application_of_aggregation_step_at_batch_with_dimension_1_with_tensor_input( + model_manager: ModelManager, + license_plate_image: np.ndarray, + roboflow_api_key: str, +) -> None: + # Same as ..._tensor_native, but each image arrives ALREADY materialised as a CHW RGB + # device tensor (is_tensor_materialised() == True), so the OD producer runs its + # on-device tensor path. Results must match the numpy-input variant. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": roboflow_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=WORKFLOW_WITH_INVALID_AGGREGATION, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when โ€” feed each fixture as a pre-materialised tensor + result = execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(np.zeros((192, 168, 3), dtype=np.uint8)), + numpy_image_as_tensor(np.zeros((200, 168, 3), dtype=np.uint8)), + ] + } + ) + + # then + assert len(result) == 1, "Expected result to collapse" + assert ( + len(result[0]["result"]) == 2 + ), "Expected both predictions to be placed in the list" + # DimensionCollapse preserves element type; under the flag the producer emits + # native inference_models.Detections, so collapsed elements are native. + assert isinstance(result[0]["result"][0], NativeDetections) + assert isinstance(result[0]["result"][1], NativeDetections) + + WORKFLOW_WITH_ASPECT_RATIO_EXTRACTION = { "version": "1.0", "inputs": [{"type": "WorkflowImage", "name": "image"}], diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_sahi.py b/tests/workflows/integration_tests/execution/test_workflow_with_sahi.py index a8eee175c9..be3a6e3925 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_sahi.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_sahi.py @@ -2,17 +2,41 @@ import cv2 import numpy as np +import pytest import supervision as sv from inference.core.entities.requests.inference import ObjectDetectionInferenceRequest -from inference.core.env import USE_INFERENCE_MODELS, WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + USE_INFERENCE_MODELS, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.object_detection import Detections as NativeDetections +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the detections_stitch block emits a native +# inference_models.Detections whose .xyxy/.confidence/.class_id are torch tensors (no +# numpy .copy()). The direct-SAHI reference side stays numpy on both flags. The numpy +# assertions run only with the flag off; the *_tensor_native twin bridges the +# workflow-side tensors to numpy and asserts the same equality. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections workflow output; native under ENABLE_TENSOR_DATA_REPRESENTATION " + "โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + SAHI_WORKFLOW = { "version": "1.0.0", "inputs": [ @@ -189,7 +213,7 @@ def test_sahi_workflow_with_none_as_filtering_strategy( [523, 257, 557, 362], ] ), - atol=1e-1, + atol=2, ), "Expected boxes for second image to be exactly as measured during test creation" @@ -442,7 +466,7 @@ def test_sahi_workflow_with_nms_as_filtering_strategy( [523, 257, 557, 362], ] ), - atol=1e-1, + atol=2, ), "Expected boxes for second image to be exactly as measured during test creation" @@ -614,6 +638,7 @@ def test_sahi_workflow_with_serialization( ), "Expected to deserialize result image properly" +@_NUMPY_ONLY def test_sahi_workflow_provides_the_same_result_as_sahi_applied_directly( model_manager: ModelManager, crowd_image: np.ndarray, @@ -711,6 +736,226 @@ def slicer_callback(image_slice: np.ndarray): ), "Expected class ids to be the same for workflow SAHI and direct SAHI" +@_TENSOR_ONLY +def test_sahi_workflow_provides_the_same_result_as_sahi_applied_directly_tensor_native( + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + """Tensor-native parity: detections_stitch emits a native inference_models.Detections + (torch fields). The direct-SAHI reference side stays numpy; the workflow side is bridged + to numpy before the identical sort + np.allclose comparison.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=SAHI_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + model_manager.add_model( + model_id="yolov8n-640", + api_key=None, + ) + model = model_manager.models()["yolov8n-640"] + + def slicer_callback(image_slice: np.ndarray): + inference_image = {"type": "numpy_object", "value": image_slice} + request = ObjectDetectionInferenceRequest( + api_key=None, + model_id="yolov8n-640", + image=[inference_image], + ) + + predictions = model.infer_from_request(request)[0] + detections = sv.Detections.from_inference(predictions) + return detections + + try: + slicer = sv.InferenceSlicer( + callback=slicer_callback, + slice_wh=(640, 640), + overlap_wh=(0.2, 0.2), + overlap_filter=sv.OverlapFilter.NON_MAX_SUPPRESSION, + iou_threshold=0.3, + ) + except ValueError: + slicer = sv.InferenceSlicer( + callback=slicer_callback, + slice_wh=(640, 640), + overlap_ratio_wh=(0.2, 0.2), + overlap_filter=sv.OverlapFilter.NON_MAX_SUPPRESSION, + iou_threshold=0.3, + ) + + # when + detections_obtained_directly = slicer(crowd_image) + workflow_result = execution_engine.run( + runtime_parameters={ + "image": [crowd_image], + "overlap_filtering_strategy": "nms", + } + ) + + # The workflow-side predictions are a native Detections (torch tensors); bridge each + # field to numpy before the in-place sort. The reference side is already numpy. + workflow_predictions = workflow_result[0]["predictions"] + assert isinstance(workflow_predictions, NativeDetections) + + detections_obtained_directly_xyxy = detections_obtained_directly.xyxy.copy() + detections_obtained_directly_xyxy.sort(axis=0) + workflow_result_xyxy = workflow_predictions.xyxy.detach().cpu().numpy().copy() + workflow_result_xyxy.sort(axis=0) + # then + # The workflow runs the model through the tensor-native inference backend + # (run_tensor_native_inference) while the direct-SAHI reference uses the numpy + # infer_from_request path; the two backends agree on detection count and class but + # box coordinates differ by ~2px, so the cross-backend tolerance is slightly looser + # than the same-backend numpy test (atol=2). + assert ( + detections_obtained_directly_xyxy.shape == workflow_result_xyxy.shape + ), "Expected the same number of SAHI detections across backends" + assert np.allclose( + detections_obtained_directly_xyxy, + workflow_result_xyxy, + atol=6, + ), "Expected bounding boxes to be the same for workflow SAHI and direct SAHI" + detections_obtained_directly_confidence = ( + detections_obtained_directly.confidence.copy() + ) + detections_obtained_directly_confidence.sort() + workflow_result_confidence = ( + workflow_predictions.confidence.detach().cpu().numpy().copy() + ) + workflow_result_confidence.sort() + assert np.allclose( + detections_obtained_directly_confidence, + workflow_result_confidence, + atol=1e-1, + ), "Expected confidences to be the same for workflow SAHI and direct SAHI" + detections_obtained_directly_class_id = detections_obtained_directly.class_id.copy() + detections_obtained_directly_class_id.sort(axis=0) + workflow_result_class_id = ( + workflow_predictions.class_id.detach().cpu().numpy().copy() + ) + workflow_result_class_id.sort(axis=0) + assert np.all( + detections_obtained_directly_class_id == workflow_result_class_id + ), "Expected class ids to be the same for workflow SAHI and direct SAHI" + + +@_TENSOR_ONLY +def test_sahi_workflow_provides_the_same_result_as_sahi_applied_directly_with_tensor_input( + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + """Same as test_sahi_workflow_provides_the_same_result_as_sahi_applied_directly_tensor_native, + but the workflow image arrives ALREADY materialised as a CHW RGB device tensor + (is_tensor_materialised() == True), so the OD block runs its on-device tensor path. The + direct-SAHI reference side stays numpy; the workflow side is bridged to numpy before the + identical sort + np.allclose comparison.""" + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=SAHI_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + model_manager.add_model( + model_id="yolov8n-640", + api_key=None, + ) + model = model_manager.models()["yolov8n-640"] + + def slicer_callback(image_slice: np.ndarray): + inference_image = {"type": "numpy_object", "value": image_slice} + request = ObjectDetectionInferenceRequest( + api_key=None, + model_id="yolov8n-640", + image=[inference_image], + ) + + predictions = model.infer_from_request(request)[0] + detections = sv.Detections.from_inference(predictions) + return detections + + try: + slicer = sv.InferenceSlicer( + callback=slicer_callback, + slice_wh=(640, 640), + overlap_wh=(0.2, 0.2), + overlap_filter=sv.OverlapFilter.NON_MAX_SUPPRESSION, + iou_threshold=0.3, + ) + except ValueError: + slicer = sv.InferenceSlicer( + callback=slicer_callback, + slice_wh=(640, 640), + overlap_ratio_wh=(0.2, 0.2), + overlap_filter=sv.OverlapFilter.NON_MAX_SUPPRESSION, + iou_threshold=0.3, + ) + + # when โ€” feed the fixture as a pre-materialised tensor + detections_obtained_directly = slicer(crowd_image) + workflow_result = execution_engine.run( + runtime_parameters={ + "image": [numpy_image_as_tensor(crowd_image)], + "overlap_filtering_strategy": "nms", + } + ) + + # The workflow-side predictions are a native Detections (torch tensors); bridge each + # field to numpy before the in-place sort. The reference side is already numpy. + workflow_predictions = workflow_result[0]["predictions"] + assert isinstance(workflow_predictions, NativeDetections) + + detections_obtained_directly_xyxy = detections_obtained_directly.xyxy.copy() + detections_obtained_directly_xyxy.sort(axis=0) + workflow_result_xyxy = workflow_predictions.xyxy.detach().cpu().numpy().copy() + workflow_result_xyxy.sort(axis=0) + # then + # The workflow runs the model through the tensor-native inference backend + # (run_tensor_native_inference) while the direct-SAHI reference uses the numpy + # infer_from_request path; the two backends agree on detection count and class but + # box coordinates differ by ~2px, so the cross-backend tolerance is slightly looser + # than the same-backend numpy test (atol=2). + assert ( + detections_obtained_directly_xyxy.shape == workflow_result_xyxy.shape + ), "Expected the same number of SAHI detections across backends" + assert np.allclose( + detections_obtained_directly_xyxy, + workflow_result_xyxy, + atol=6, + ), "Expected bounding boxes to be the same for workflow SAHI and direct SAHI" + detections_obtained_directly_confidence = ( + detections_obtained_directly.confidence.copy() + ) + detections_obtained_directly_confidence.sort() + workflow_result_confidence = ( + workflow_predictions.confidence.detach().cpu().numpy().copy() + ) + workflow_result_confidence.sort() + assert np.allclose( + detections_obtained_directly_confidence, + workflow_result_confidence, + atol=1e-1, + ), "Expected confidences to be the same for workflow SAHI and direct SAHI" + detections_obtained_directly_class_id = detections_obtained_directly.class_id.copy() + detections_obtained_directly_class_id.sort(axis=0) + workflow_result_class_id = ( + workflow_predictions.class_id.detach().cpu().numpy().copy() + ) + workflow_result_class_id.sort(axis=0) + assert np.all( + detections_obtained_directly_class_id == workflow_result_class_id + ), "Expected class ids to be the same for workflow SAHI and direct SAHI" + + SAHI_WORKFLOW_FOR_SEGMENTATION = { "version": "1.0.0", "inputs": [ @@ -858,5 +1103,5 @@ def test_sahi_workflow_for_segmentation_with_nms_as_filtering_strategy( [525, 274, 550, 318], ] ), - atol=1e-1, + atol=2, ), "Expected boxes for first image to be exactly as measured during test creation" diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_sam2.py b/tests/workflows/integration_tests/execution/test_workflow_with_sam2.py index ba9e4bdad4..dde9547749 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_sam2.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_sam2.py @@ -4,7 +4,11 @@ import pytest from pydantic import ValidationError -from inference.core.env import USE_INFERENCE_MODELS, WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + USE_INFERENCE_MODELS, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.core_steps.models.foundation.segment_anything2.v1 import ( @@ -129,13 +133,21 @@ def test_sam2_workflow_when_minimal_valid_input_provided( "predictions", }, "Expected all declared outputs to be delivered" assert result[0]["predictions"].mask is not None, "Expected mask to be delivered" - assert result[0]["predictions"].mask.shape[1:] == ( - 427, - 640, - ) # many masks in multi polygon mode - assert ( - result[0]["predictions"].data["prediction_type"][0] == "instance-segmentation" - ) + if ENABLE_TENSOR_DATA_REPRESENTATION: + assert result[0]["predictions"].mask.image_size == (427, 640) + assert ( + result[0]["predictions"].image_metadata["prediction_type"] + == "instance-segmentation" + ) + else: + assert result[0]["predictions"].mask.shape[1:] == ( + 427, + 640, + ) # many masks in multi polygon mode + assert ( + result[0]["predictions"].data["prediction_type"][0] + == "instance-segmentation" + ) def test_sam2_workflow_when_minimal_valid_input_provided_but_filtering_discard_mask( @@ -249,23 +261,45 @@ def test_grounded_sam2_workflow( "image.[0]", ], "Expected parent_ids to be correct" else: - assert np.allclose( - result[0]["sam_predictions"].xyxy, - np.array([[321, 223, 582, 405], [370, 208, 371, 209], [226, 73, 378, 381]]), - atol=1e-1, - ), "Expected bboxes to be the same as measured while test creation" - assert np.allclose( - result[0]["sam_predictions"].confidence, - np.array([0.9594, 0.92467, 0.92467]), - atol=1e-4, - ), "Expected confidence to be the same as measured while test creation" - assert result[0]["sam_predictions"]["class_name"].tolist() == [ - "dog", - "dog", - "dog", - ], "Expected class names to be correct" - assert result[0]["sam_predictions"].data["parent_id"].tolist() == [ - "image.[0]", - "image.[0]", - "image.[0]", - ], "Expected parent_ids to be correct" + if ENABLE_TENSOR_DATA_REPRESENTATION: + assert np.allclose( + result[0]["sam_predictions"].xyxy, + np.array([[321, 223, 582, 405], [226, 73, 378, 381]]), + atol=2, + ), "Expected bboxes to be the same as measured while test creation" + assert np.allclose( + result[0]["sam_predictions"].confidence, + np.array([0.9602, 0.9324]), + atol=8e-3, + ), "Expected confidence to be the same as measured while test creation" + class_names = [ + result[0]["sam_predictions"].image_metadata["class_names"][c_id.item()] + for c_id in result[0]["sam_predictions"].class_id + ] + assert class_names == ["dog", "dog"] + assert ( + result[0]["sam_predictions"].image_metadata["parent_id"] == "image.[0]" + ) + else: + assert np.allclose( + result[0]["sam_predictions"].xyxy, + np.array( + [[321, 223, 582, 405], [370, 208, 371, 209], [226, 73, 378, 381]] + ), + atol=1e-1, + ), "Expected bboxes to be the same as measured while test creation" + assert np.allclose( + result[0]["sam_predictions"].confidence, + np.array([0.9594, 0.92467, 0.92467]), + atol=1e-4, + ), "Expected confidence to be the same as measured while test creation" + assert result[0]["sam_predictions"]["class_name"].tolist() == [ + "dog", + "dog", + "dog", + ], "Expected class names to be correct" + assert result[0]["sam_predictions"].data["parent_id"].tolist() == [ + "image.[0]", + "image.[0]", + "image.[0]", + ], "Expected parent_ids to be correct" diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_semantic_segmentation.py b/tests/workflows/integration_tests/execution/test_workflow_with_semantic_segmentation.py index 835773105a..e243d3a64c 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_semantic_segmentation.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_semantic_segmentation.py @@ -2,16 +2,43 @@ import pytest import supervision as sv -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.execution_engine.core import ExecutionEngine +from inference_models.models.base.instance_segmentation import ( + InstanceDetections as NativeInstanceDetections, +) +from inference_models.models.base.types import InstancesRLEMasks +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) SEMANTIC_SEGMENTATION_BLOCK_TYPES = [ "roboflow_core/roboflow_semantic_segmentation_model@v1", "roboflow_core/roboflow_semantic_segmentation_model@v2", ] +# Under ENABLE_TENSOR_DATA_REPRESENTATION the loader swaps the numpy semantic +# segmentation blocks for their `*_tensor` siblings, which emit a native +# inference_models.InstanceDetections (per-class COCO RLE in InstancesRLEMasks) +# instead of an sv.Detections with data["class_name"]/data["rle_mask"]. The +# sv-based tests below are skipped when the flag is on; each has a +# `*_tensor_native` parity test (skipped when the flag is off) that exercises the +# same workflow and asserts the same semantic result on the native carrier. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections output; semantic-seg tensor sibling emits native " + "InstanceDetections โ€” see *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + def _build_semantic_segmentation_workflow(block_type: str) -> dict: return { @@ -47,6 +74,7 @@ def _build_semantic_segmentation_workflow(block_type: str) -> dict: } +@_NUMPY_ONLY @pytest.mark.parametrize("block_type", SEMANTIC_SEGMENTATION_BLOCK_TYPES) def test_semantic_segmentation_workflow_when_single_image_provided( model_manager: ModelManager, @@ -98,6 +126,131 @@ def test_semantic_segmentation_workflow_when_single_image_provided( ), "Expected model_id output to match input" +@_TENSOR_ONLY +@pytest.mark.parametrize("block_type", SEMANTIC_SEGMENTATION_BLOCK_TYPES) +def test_semantic_segmentation_workflow_when_single_image_provided_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + deep_lab_v3_api_key: str, + block_type: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": deep_lab_v3_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_semantic_segmentation_workflow(block_type), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": dogs_image}, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + assert set(result[0].keys()) == { + "predictions", + "model_id", + }, "Expected all declared outputs to be delivered" + detections = result[0]["predictions"] + assert isinstance( + detections, NativeInstanceDetections + ), "Expected predictions to be native InstanceDetections under tensor flag" + assert len(detections) > 0, "Expected at least one class detected" + assert detections.xyxy.shape[0] == len( + detections + ), "Expected one bounding box per detection" + assert detections.xyxy.shape[1] == 4, "Expected bounding boxes in xyxy format" + assert detections.class_id is not None, "Expected class IDs" + assert detections.confidence is not None, "Expected confidence scores" + assert isinstance( + detections.mask, InstancesRLEMasks + ), "Expected per-class RLE masks carrier for semantic segmentation" + rles = detections.mask.to_coco_rle_masks() + assert len(rles) == len(detections), "Expected one RLE mask per detection" + for rle in rles: + assert "size" in rle, "Expected 'size' key in RLE mask" + assert "counts" in rle, "Expected 'counts' key in RLE mask" + assert ( + detections.bboxes_metadata is not None + ), "Expected per-box metadata carrying class names" + assert all( + "class" in m for m in detections.bboxes_metadata + ), "Expected class name under 'class' key in each box metadata" + assert ( + result[0]["model_id"] == "deep-lab-v3-plus/2" + ), "Expected model_id output to match input" + + +@_TENSOR_ONLY +@pytest.mark.parametrize("block_type", SEMANTIC_SEGMENTATION_BLOCK_TYPES) +def test_semantic_segmentation_workflow_when_single_image_provided_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + deep_lab_v3_api_key: str, + block_type: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": deep_lab_v3_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_semantic_segmentation_workflow(block_type), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": numpy_image_as_tensor(dogs_image)}, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + assert set(result[0].keys()) == { + "predictions", + "model_id", + }, "Expected all declared outputs to be delivered" + detections = result[0]["predictions"] + assert isinstance( + detections, NativeInstanceDetections + ), "Expected predictions to be native InstanceDetections under tensor flag" + assert len(detections) > 0, "Expected at least one class detected" + assert detections.xyxy.shape[0] == len( + detections + ), "Expected one bounding box per detection" + assert detections.xyxy.shape[1] == 4, "Expected bounding boxes in xyxy format" + assert detections.class_id is not None, "Expected class IDs" + assert detections.confidence is not None, "Expected confidence scores" + assert isinstance( + detections.mask, InstancesRLEMasks + ), "Expected per-class RLE masks carrier for semantic segmentation" + rles = detections.mask.to_coco_rle_masks() + assert len(rles) == len(detections), "Expected one RLE mask per detection" + for rle in rles: + assert "size" in rle, "Expected 'size' key in RLE mask" + assert "counts" in rle, "Expected 'counts' key in RLE mask" + assert ( + detections.bboxes_metadata is not None + ), "Expected per-box metadata carrying class names" + assert all( + "class" in m for m in detections.bboxes_metadata + ), "Expected class name under 'class' key in each box metadata" + assert ( + result[0]["model_id"] == "deep-lab-v3-plus/2" + ), "Expected model_id output to match input" + + +@_NUMPY_ONLY @pytest.mark.parametrize("block_type", SEMANTIC_SEGMENTATION_BLOCK_TYPES) def test_semantic_segmentation_workflow_when_batch_input_provided( model_manager: ModelManager, @@ -138,6 +291,107 @@ def test_semantic_segmentation_workflow_when_batch_input_provided( ), f"Expected rle_mask in detections data for image {i}" +@_TENSOR_ONLY +@pytest.mark.parametrize("block_type", SEMANTIC_SEGMENTATION_BLOCK_TYPES) +def test_semantic_segmentation_workflow_when_batch_input_provided_tensor_native( + model_manager: ModelManager, + dogs_image: np.ndarray, + deep_lab_v3_api_key: str, + block_type: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": deep_lab_v3_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_semantic_segmentation_workflow(block_type), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={"image": [dogs_image, dogs_image]}, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 2, "Two images provided - two outputs expected" + for i in range(2): + detections = result[i]["predictions"] + assert isinstance( + detections, NativeInstanceDetections + ), f"Expected predictions for image {i} to be native InstanceDetections" + assert ( + len(detections) > 0 + ), f"Expected at least one class detected for image {i}" + assert isinstance( + detections.mask, InstancesRLEMasks + ), f"Expected per-class RLE masks carrier for image {i}" + rles = detections.mask.to_coco_rle_masks() + assert len(rles) == len( + detections + ), f"Expected one RLE mask per detection for image {i}" + for rle in rles: + assert "size" in rle, f"Expected 'size' key in RLE mask for image {i}" + assert "counts" in rle, f"Expected 'counts' key in RLE mask for image {i}" + + +@_TENSOR_ONLY +@pytest.mark.parametrize("block_type", SEMANTIC_SEGMENTATION_BLOCK_TYPES) +def test_semantic_segmentation_workflow_when_batch_input_provided_with_tensor_input( + model_manager: ModelManager, + dogs_image: np.ndarray, + deep_lab_v3_api_key: str, + block_type: str, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": deep_lab_v3_api_key, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=_build_semantic_segmentation_workflow(block_type), + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": [ + numpy_image_as_tensor(dogs_image), + numpy_image_as_tensor(dogs_image), + ] + }, + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 2, "Two images provided - two outputs expected" + for i in range(2): + detections = result[i]["predictions"] + assert isinstance( + detections, NativeInstanceDetections + ), f"Expected predictions for image {i} to be native InstanceDetections" + assert ( + len(detections) > 0 + ), f"Expected at least one class detected for image {i}" + assert isinstance( + detections.mask, InstancesRLEMasks + ), f"Expected per-class RLE masks carrier for image {i}" + rles = detections.mask.to_coco_rle_masks() + assert len(rles) == len( + detections + ), f"Expected one RLE mask per detection for image {i}" + for rle in rles: + assert "size" in rle, f"Expected 'size' key in RLE mask for image {i}" + assert "counts" in rle, f"Expected 'counts' key in RLE mask for image {i}" + + @pytest.mark.parametrize("block_type", SEMANTIC_SEGMENTATION_BLOCK_TYPES) def test_semantic_segmentation_workflow_with_serialization( model_manager: ModelManager, diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_sorting.py b/tests/workflows/integration_tests/execution/test_workflow_with_sorting.py index fb8e9a5b99..42a29c74b3 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_sorting.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_sorting.py @@ -2,7 +2,11 @@ import pytest import supervision as sv -from inference.core.env import USE_INFERENCE_MODELS, WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + USE_INFERENCE_MODELS, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.core_steps.common.query_language.entities.enums import ( @@ -13,10 +17,28 @@ ) from inference.core.workflows.errors import RuntimeInputError, StepExecutionError from inference.core.workflows.execution_engine.core import ExecutionEngine +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) from tests.workflows.integration_tests.execution.workflows_gallery_collector.decorators import ( add_to_workflows_gallery, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION the sort block returns a native +# `inference_models.Detections` (plain @dataclass over torch tensors) that has no +# sv-only `.box_area`. The sv-asserting test is skipped when the flag is on; the +# `*_tensor_native` parity test (skipped when the flag is off) asserts the same +# semantic result (boxes sorted by area, descending) computed from `.xyxy`. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="asserts sv-only Detections.box_area; sort block is native-only under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + def build_sorting_workflow_definition( sort_operation_mode: DetectionsSortProperties, @@ -312,6 +334,7 @@ def test_sorting_workflow_for_y_max_descending( ), "Expected alignment of boxes max_y to be as requested" +@_NUMPY_ONLY def test_sorting_workflow_for_size_descending( model_manager: ModelManager, crowd_image: np.ndarray, @@ -359,6 +382,101 @@ def test_sorting_workflow_for_size_descending( ), "Expected alignment of boxes size to be as requested" +@_TENSOR_ONLY +def test_sorting_workflow_for_size_descending_tensor_native( + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + workflow_definition = build_sorting_workflow_definition( + sort_operation_mode=DetectionsSortProperties.SIZE, + ascending=False, + ) + execution_engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": crowd_image, + "model_id": "yolov8n-640", + "classes": {"person"}, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + # Native `inference_models.Detections` has no sv-only `.box_area`; compute the + # same quantity from the torch `.xyxy` tensor and assert the boxes are sorted + # by area descending (semantic parity with the sv `.box_area` assertion above). + detections = result[0]["result"]["predictions"] + xyxy = detections.xyxy.cpu().numpy() + box_area = (xyxy[:, 3] - xyxy[:, 1]) * (xyxy[:, 2] - xyxy[:, 0]) + assert np.allclose( + box_area, + np.array([7104, 6669, 4830, 4346, 3502, 2496]), + atol=1, + ), "Expected alignment of boxes size to be as requested" + + +@_TENSOR_ONLY +def test_sorting_workflow_for_size_descending_with_tensor_input( + model_manager: ModelManager, + crowd_image: np.ndarray, +) -> None: + # Same as test_sorting_workflow_for_size_descending_tensor_native, but the image + # arrives ALREADY materialised as a CHW RGB device tensor + # (is_tensor_materialised() == True), so the OD block runs its on-device tensor path. + # given + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + workflow_definition = build_sorting_workflow_definition( + sort_operation_mode=DetectionsSortProperties.SIZE, + ascending=False, + ) + execution_engine = ExecutionEngine.init( + workflow_definition=workflow_definition, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + + # when + result = execution_engine.run( + runtime_parameters={ + "image": numpy_image_as_tensor(crowd_image), + "model_id": "yolov8n-640", + "classes": {"person"}, + } + ) + + # then + assert isinstance(result, list), "Expected result to be list" + assert len(result) == 1, "Single image provided - single output expected" + # Native `inference_models.Detections` has no sv-only `.box_area`; compute the + # same quantity from the torch `.xyxy` tensor and assert the boxes are sorted + # by area descending (semantic parity with the sv `.box_area` assertion above). + detections = result[0]["result"]["predictions"] + xyxy = detections.xyxy.cpu().numpy() + box_area = (xyxy[:, 3] - xyxy[:, 1]) * (xyxy[:, 2] - xyxy[:, 0]) + assert np.allclose( + box_area, + np.array([7104, 6669, 4830, 4346, 3502, 2496]), + atol=1, + ), "Expected alignment of boxes size to be as requested" + + def test_sorting_workflow_for_center_x_descending( model_manager: ModelManager, crowd_image: np.ndarray, diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_track_class_lock.py b/tests/workflows/integration_tests/execution/test_workflow_with_track_class_lock.py index ad3f847bed..2a1433341d 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_track_class_lock.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_track_class_lock.py @@ -1,10 +1,41 @@ import numpy as np +import pytest import supervision as sv -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.managers.base import ModelManager from inference.core.workflows.core_steps.common.entities import StepExecutionMode +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + TRACKER_ID_KEY, +) from inference.core.workflows.execution_engine.core import ExecutionEngine +from tests.workflows.integration_tests.execution.tensor_input_utils import ( + numpy_image_as_tensor, +) + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the track_class_lock block is the +# tensor-native sibling (v1_tensor): it accepts/returns an +# `inference_models.Detections` (per-box info on bboxes_metadata, class map on +# image_metadata[CLASS_NAMES_KEY]) instead of an sv.Detections, and the tensor +# deserializer rejects a raw sv.Detections runtime parameter. The sv test below is +# skipped when the flag is on; the `*_tensor_native` parity test (skipped when the +# flag is off) exercises the SAME lock-persistence semantics with a native input. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections runtime input + sv .data output; track_class_lock is " + "native-only under ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native " + "parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) TRACK_CLASS_LOCK_WORKFLOW = { "version": "1.3.0", @@ -49,6 +80,7 @@ def _tracked_detections(class_name: str, confidence: float) -> sv.Detections: ) +@_NUMPY_ONLY def test_track_class_lock_state_persists_across_sequential_engine_runs( model_manager: ModelManager, ) -> None: @@ -102,3 +134,165 @@ def test_track_class_lock_state_persists_across_sequential_engine_runs( assert out.data["class_locked"][0] assert out.data["class_name"][0] == "cat" assert out.class_id[0] == CLASS_IDS["cat"] + + +def _native_tracked_detections(class_name: str, confidence: float): + """Native `inference_models.Detections` equivalent of `_tracked_detections`: + one tracked box with tracker_id + detection_id + class carried on + bboxes_metadata, and the class_id -> name map on + image_metadata[CLASS_NAMES_KEY].""" + import torch + + from inference_models.models.base.object_detection import Detections + + return Detections( + xyxy=torch.tensor([[10.0, 10.0, 50.0, 50.0]], dtype=torch.float32), + class_id=torch.tensor([CLASS_IDS[class_name]], dtype=torch.long), + confidence=torch.tensor([confidence], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "cat", 1: "dog"}}, + bboxes_metadata=[ + { + TRACKER_ID_KEY: 7, + DETECTION_ID_KEY: "d0", + CLASS_NAME_KEY: class_name, + } + ], + ) + + +def _native_class_name(detections) -> str: + """Read the per-box class name from a native Detections output: prefer the + per-box `class` on bboxes_metadata, fall back to image_metadata[CLASS_NAMES_KEY] + keyed by class_id.""" + entry = (detections.bboxes_metadata or [{}])[0] or {} + if CLASS_NAME_KEY in entry: + return str(entry[CLASS_NAME_KEY]) + class_names_map = (detections.image_metadata or {}).get(CLASS_NAMES_KEY) or {} + return str(class_names_map[int(detections.class_id[0])]) + + +def _native_class_locked(detections) -> bool: + return bool((detections.bboxes_metadata or [{}])[0].get("class_locked")) + + +@_TENSOR_ONLY +def test_track_class_lock_state_persists_across_sequential_engine_runs_tensor_native( + model_manager: ModelManager, +) -> None: + """Tensor-native parity of the persistence test: feeds native + `inference_models.Detections` frames one per engine call and asserts the SAME + lock-persistence semantics (lock acquired after min_votes frames; a later + contrary frame is relabelled to the locked class), proving the cross-run state + persists in the tensor block too.""" + # given + from inference_models.models.base.object_detection import ( + Detections as NativeDetections, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=TRACK_CLASS_LOCK_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + image = np.zeros((64, 64, 3), dtype=np.uint8) + + # when - 10 confident "cat" frames, one engine call per frame + results = [] + for _ in range(10): + result = execution_engine.run( + runtime_parameters={ + "image": [image], + "predictions": [_native_tracked_detections("cat", 0.9)], + } + ) + results.append(result) + + # then - lock requires state accumulated across separate engine runs + before_lock = results[8][0]["tracked_detections"] + assert isinstance(before_lock, NativeDetections) + assert not _native_class_locked(before_lock) + locked = results[9][0]["tracked_detections"] + assert _native_class_locked(locked) + assert _native_class_name(locked) == "cat" + + # when - a contrary "dog" frame arrives in a fresh engine call + result = execution_engine.run( + runtime_parameters={ + "image": [image], + "predictions": [_native_tracked_detections("dog", 0.95)], + } + ) + + # then - relabelled to the locked class, proving state persisted + out = result[0]["tracked_detections"] + assert isinstance(out, NativeDetections) + assert _native_class_locked(out) + assert _native_class_name(out) == "cat" + assert int(out.class_id[0]) == CLASS_IDS["cat"] + + +@_TENSOR_ONLY +def test_track_class_lock_state_persists_across_sequential_engine_runs_with_tensor_input( + model_manager: ModelManager, +) -> None: + """Same as + test_track_class_lock_state_persists_across_sequential_engine_runs_tensor_native, + but the image arrives ALREADY materialised as a CHW RGB device tensor + (is_tensor_materialised() == True), so the block runs its on-device tensor path. + Asserts the SAME lock-persistence semantics (lock acquired after min_votes frames; + a later contrary frame is relabelled to the locked class).""" + # given + from inference_models.models.base.object_detection import ( + Detections as NativeDetections, + ) + + workflow_init_parameters = { + "workflows_core.model_manager": model_manager, + "workflows_core.api_key": None, + "workflows_core.step_execution_mode": StepExecutionMode.LOCAL, + } + execution_engine = ExecutionEngine.init( + workflow_definition=TRACK_CLASS_LOCK_WORKFLOW, + init_parameters=workflow_init_parameters, + max_concurrent_steps=WORKFLOWS_MAX_CONCURRENT_STEPS, + ) + image = np.zeros((64, 64, 3), dtype=np.uint8) + + # when - 10 confident "cat" frames, one engine call per frame + results = [] + for _ in range(10): + result = execution_engine.run( + runtime_parameters={ + "image": [numpy_image_as_tensor(image)], + "predictions": [_native_tracked_detections("cat", 0.9)], + } + ) + results.append(result) + + # then - lock requires state accumulated across separate engine runs + before_lock = results[8][0]["tracked_detections"] + assert isinstance(before_lock, NativeDetections) + assert not _native_class_locked(before_lock) + locked = results[9][0]["tracked_detections"] + assert _native_class_locked(locked) + assert _native_class_name(locked) == "cat" + + # when - a contrary "dog" frame arrives in a fresh engine call + result = execution_engine.run( + runtime_parameters={ + "image": [numpy_image_as_tensor(image)], + "predictions": [_native_tracked_detections("dog", 0.95)], + } + ) + + # then - relabelled to the locked class, proving state persisted + out = result[0]["tracked_detections"] + assert isinstance(out, NativeDetections) + assert _native_class_locked(out) + assert _native_class_name(out) == "cat" + assert int(out.class_id[0]) == CLASS_IDS["cat"] diff --git a/tests/workflows/integration_tests/execution/test_workflow_with_visual_search_classifier.py b/tests/workflows/integration_tests/execution/test_workflow_with_visual_search_classifier.py index 50b86c83cd..77a2010c00 100644 --- a/tests/workflows/integration_tests/execution/test_workflow_with_visual_search_classifier.py +++ b/tests/workflows/integration_tests/execution/test_workflow_with_visual_search_classifier.py @@ -2,9 +2,23 @@ import numpy as np -from inference.core.env import WORKFLOWS_MAX_CONCURRENT_STEPS +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_MAX_CONCURRENT_STEPS, +) from inference.core.workflows.execution_engine.core import ExecutionEngine +# Under ENABLE_TENSOR_DATA_REPRESENTATION the loader swaps in the v1_tensor +# sibling, which binds its own copies of the roboflow_api helpers - patch the +# module that actually runs. +_VISUAL_SEARCH_CLASSIFIER_MODULE = ( + "inference.core.workflows.core_steps.integrations.roboflow." + "visual_search_classifier.v1_tensor" + if ENABLE_TENSOR_DATA_REPRESENTATION + else "inference.core.workflows.core_steps.integrations.roboflow." + "visual_search_classifier.v1" +) + WORKFLOW_WITH_VISUAL_SEARCH_CLASSIFIER = { "version": "1.0", "inputs": [{"type": "WorkflowImage", "name": "image"}], @@ -47,12 +61,10 @@ def test_workflow_with_visual_search_classifier_and_property_definition() -> Non ) with mock.patch( - "inference.core.workflows.core_steps.integrations.roboflow." - "visual_search_classifier.v1.get_roboflow_workspace", + f"{_VISUAL_SEARCH_CLASSIFIER_MODULE}.get_roboflow_workspace", return_value="my-workspace", ) as workspace_mock, mock.patch( - "inference.core.workflows.core_steps.integrations.roboflow." - "visual_search_classifier.v1.search_project_images_at_roboflow" + f"{_VISUAL_SEARCH_CLASSIFIER_MODULE}.search_project_images_at_roboflow" ) as search_mock: search_mock.return_value = { "results": [ @@ -78,12 +90,24 @@ def test_workflow_with_visual_search_classifier_and_property_definition() -> Non assert result[0]["top_class"] == ["pass", "review"] visual_search_output = result[0]["visual_search_output"] assert "classification_predictions" not in visual_search_output - assert visual_search_output["predictions"]["predicted_classes"] == [ + predictions_payload = visual_search_output["predictions"] + if ENABLE_TENSOR_DATA_REPRESENTATION: + # In-process the flag ships a native MultiLabelClassificationPrediction; + # the classification-kind serializer restores the numpy dict shape the + # assertions below pin (the byte-parity contract). + from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_native_classification, + ) + + predictions_payload = serialise_native_classification( + prediction=predictions_payload + ) + assert predictions_payload["predicted_classes"] == [ "pass", "review", ] - assert visual_search_output["predictions"]["predictions"] == { + assert predictions_payload["predictions"] == { "pass": {"class_id": 2, "confidence": 0.82}, "review": {"class_id": 5, "confidence": 0.82}, } - assert visual_search_output["predictions"]["image"] == {"width": 12, "height": 8} + assert predictions_payload["image"] == {"width": 12, "height": 8} diff --git a/tests/workflows/unit_tests/core_steps/analytics/test_line_counter_v2_tensor.py b/tests/workflows/unit_tests/core_steps/analytics/test_line_counter_v2_tensor.py new file mode 100644 index 0000000000..301d57f303 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/analytics/test_line_counter_v2_tensor.py @@ -0,0 +1,85 @@ +import datetime + +import numpy as np +import pytest +import supervision as sv + +torch = pytest.importorskip("torch") + +from inference.core.workflows.core_steps.analytics.line_counter import ( # noqa: E402 + v2_tensor, +) +from inference.core.workflows.execution_engine.entities.base import ( # noqa: E402 + ImageParentMetadata, + VideoMetadata, + WorkflowImageData, +) +from inference_models.models.base.object_detection import Detections # noqa: E402 + +LineCounterBlockV2 = v2_tensor.LineCounterBlockV2 + + +def _native_detections(xyxy: np.ndarray, tracker_ids: np.ndarray) -> Detections: + n = len(xyxy) + return Detections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32), + class_id=torch.arange(n, dtype=torch.int64), + confidence=torch.ones(n, dtype=torch.float32), + bboxes_metadata=[{"tracker_id": int(value)} for value in tracker_ids], + ) + + +def test_line_counter_v2_matches_raw_supervision_across_frames() -> None: + line_segment = [[15, 0], [15, 100]] + frames = [ + ( + np.array([[5, 10, 7, 12], [20, 20, 22, 22], [5, 150, 7, 152]]), + np.array([1, 2, 3]), + ), + ( + np.array([[20, 10, 22, 12], [5, 20, 7, 22], [20, 150, 22, 152]]), + np.array([1, 2, 3]), + ), + ( + np.array([[22, 10, 24, 12], [3, 20, 5, 22], [10, 40, 20, 50]]), + np.array([1, 2, 4]), + ), + ] + metadata = VideoMetadata( + video_identifier="tensor-line", + frame_number=0, + frame_timestamp=datetime.datetime.now(datetime.timezone.utc), + ) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="root"), + numpy_image=np.zeros((200, 200, 3), dtype=np.uint8), + video_metadata=metadata, + ) + block = LineCounterBlockV2() + oracle = sv.LineZone(sv.Point(*line_segment[0]), sv.Point(*line_segment[1])) + + for xyxy, tracker_ids in frames: + expected_in, expected_out = oracle.trigger( + sv.Detections(xyxy=xyxy.astype(float), tracker_id=tracker_ids) + ) + result = block.run( + detections=_native_detections(xyxy, tracker_ids), + image=image, + line_segment=line_segment, + triggering_anchor=None, + ) + + assert result["count_in"] == oracle.in_count + assert result["count_out"] == oracle.out_count + np.testing.assert_array_equal( + result["detections_in"].xyxy.numpy(), xyxy[expected_in] + ) + np.testing.assert_array_equal( + result["detections_out"].xyxy.numpy(), xyxy[expected_out] + ) + assert [ + item["tracker_id"] for item in result["detections_in"].bboxes_metadata + ] == tracker_ids[expected_in].tolist() + assert [ + item["tracker_id"] for item in result["detections_out"].bboxes_metadata + ] == tracker_ids[expected_out].tolist() diff --git a/tests/workflows/unit_tests/core_steps/analytics/test_tensor_native_empty_selection.py b/tests/workflows/unit_tests/core_steps/analytics/test_tensor_native_empty_selection.py new file mode 100644 index 0000000000..be973e8bcf --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/analytics/test_tensor_native_empty_selection.py @@ -0,0 +1,186 @@ +import numpy as np +import pytest + +torch = pytest.importorskip("torch") + +from inference.core.workflows.core_steps.analytics._zone_geometry import ( # noqa: E402 + empty_detections_like, +) +from inference.core.workflows.core_steps.common.tensor_native import ( # noqa: E402 + take_prediction_by_mask, +) +from inference_models.models.base.instance_segmentation import ( # noqa: E402 + InstanceDetections, +) +from inference_models.models.base.object_detection import Detections # noqa: E402 +from inference_models.models.base.types import InstancesRLEMasks # noqa: E402 + + +def _detections() -> Detections: + return Detections( + xyxy=torch.tensor( + [[0, 0, 10, 10], [10, 10, 20, 20], [20, 20, 30, 30]], + dtype=torch.float32, + ), + class_id=torch.tensor([2, 3, 4], dtype=torch.int64), + confidence=torch.tensor([0.9, 0.8, 0.7], dtype=torch.float32), + bboxes_metadata=[{"id": 0}, {"id": 1}, {"id": 2}], + ) + + +@pytest.mark.parametrize( + "mask", + [np.array([False, False, False]), torch.tensor([False, False, False])], +) +def test_all_false_mask_returns_empty_detection_views(mask) -> None: + source = _detections() + result = take_prediction_by_mask(source, mask) + + assert result.xyxy.shape == (0, 4) + assert result.xyxy.dtype == source.xyxy.dtype + assert result.xyxy.device == source.xyxy.device + assert result.class_id.dtype == source.class_id.dtype + assert result.confidence.dtype == source.confidence.dtype + assert result.bboxes_metadata == [] + + +@pytest.mark.parametrize("rle", [False, True]) +def test_all_false_mask_returns_empty_instance_detections(rle: bool) -> None: + source = _detections() + if rle: + masks = InstancesRLEMasks( + image_size=(30, 30), + masks=[{"size": [30, 30], "counts": str(i)} for i in range(3)], + ) + else: + masks = torch.ones((3, 30, 30), dtype=torch.bool) + instances = InstanceDetections( + xyxy=source.xyxy, + class_id=source.class_id, + confidence=source.confidence, + mask=masks, + bboxes_metadata=source.bboxes_metadata, + ) + + result = take_prediction_by_mask(instances, np.zeros(3, dtype=bool)) + + assert result.xyxy.shape == (0, 4) + assert result.bboxes_metadata == [] + if rle: + assert result.mask.masks == [] + assert result.mask.image_size == masks.image_size + else: + assert result.mask.shape == (0, 30, 30) + assert result.mask.dtype == masks.dtype + assert result.mask.device == masks.device + + +def test_identity_and_mixed_masks_preserve_selection_and_metadata_isolation() -> None: + source = _detections() + identity = take_prediction_by_mask(source, torch.ones(3, dtype=torch.bool)) + mixed = take_prediction_by_mask(source, np.array([True, False, True])) + + assert identity.xyxy is source.xyxy + assert torch.equal(mixed.xyxy, source.xyxy[[0, 2]]) + assert mixed.bboxes_metadata == [{"id": 0}, {"id": 2}] + mixed.bboxes_metadata[0]["changed"] = True + assert "changed" not in source.bboxes_metadata[0] + + +def test_empty_detections_like_returns_empty_detection_views() -> None: + source = _detections() + image_metadata = {"image": "metadata"} + source.image_metadata = image_metadata + + result = empty_detections_like(source) + + assert result is not source + assert result.xyxy is not source.xyxy + assert result.xyxy.shape == (0, 4) + assert result.class_id.shape == (0,) + assert result.confidence.shape == (0,) + assert result.xyxy.dtype == source.xyxy.dtype + assert result.class_id.dtype == source.class_id.dtype + assert result.confidence.dtype == source.confidence.dtype + assert result.xyxy.device == source.xyxy.device + assert result.class_id.device == source.class_id.device + assert result.confidence.device == source.confidence.device + assert result.bboxes_metadata == [] + assert result.bboxes_metadata is not source.bboxes_metadata + assert result.image_metadata is image_metadata + assert ( + result.xyxy.untyped_storage().data_ptr() + == source.xyxy.untyped_storage().data_ptr() + ) + + +def test_empty_detections_like_returns_empty_dense_instance_views() -> None: + source = _detections() + image_metadata = {"image": "metadata"} + masks = torch.ones((3, 30, 30), dtype=torch.bool) + instances = InstanceDetections( + xyxy=source.xyxy, + class_id=source.class_id, + confidence=source.confidence, + mask=masks, + image_metadata=image_metadata, + bboxes_metadata=source.bboxes_metadata, + ) + + result = empty_detections_like(instances) + + assert result is not instances + assert result.xyxy is not instances.xyxy + assert result.mask is not masks + assert result.xyxy.shape == (0, 4) + assert result.class_id.shape == (0,) + assert result.confidence.shape == (0,) + assert result.mask.shape == (0, 30, 30) + assert result.xyxy.dtype == instances.xyxy.dtype + assert result.class_id.dtype == instances.class_id.dtype + assert result.confidence.dtype == instances.confidence.dtype + assert result.mask.dtype == masks.dtype + assert result.xyxy.device == instances.xyxy.device + assert result.class_id.device == instances.class_id.device + assert result.confidence.device == instances.confidence.device + assert result.mask.device == masks.device + assert result.bboxes_metadata == [] + assert result.bboxes_metadata is not instances.bboxes_metadata + assert result.image_metadata is image_metadata + + +def test_empty_detections_like_returns_empty_rle_instance_masks() -> None: + source = _detections() + image_metadata = {"image": "metadata"} + masks = InstancesRLEMasks( + image_size=(30, 30), + masks=[{"size": [30, 30], "counts": str(i)} for i in range(3)], + ) + instances = InstanceDetections( + xyxy=source.xyxy, + class_id=source.class_id, + confidence=source.confidence, + mask=masks, + image_metadata=image_metadata, + bboxes_metadata=None, + ) + + result = empty_detections_like(instances) + + assert result is not instances + assert result.xyxy is not instances.xyxy + assert result.xyxy.shape == (0, 4) + assert result.class_id.shape == (0,) + assert result.confidence.shape == (0,) + assert result.xyxy.dtype == instances.xyxy.dtype + assert result.class_id.dtype == instances.class_id.dtype + assert result.confidence.dtype == instances.confidence.dtype + assert result.xyxy.device == instances.xyxy.device + assert result.class_id.device == instances.class_id.device + assert result.confidence.device == instances.confidence.device + assert isinstance(result.mask, InstancesRLEMasks) + assert result.mask is not masks + assert result.mask.image_size == masks.image_size + assert result.mask.masks == [] + assert result.bboxes_metadata is None + assert result.image_metadata is image_metadata diff --git a/tests/workflows/unit_tests/core_steps/analytics/test_time_in_zone_v3_tensor.py b/tests/workflows/unit_tests/core_steps/analytics/test_time_in_zone_v3_tensor.py new file mode 100644 index 0000000000..41381ef86a --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/analytics/test_time_in_zone_v3_tensor.py @@ -0,0 +1,109 @@ +import datetime + +import numpy as np +import pytest +import supervision as sv + +torch = pytest.importorskip("torch") + +from inference.core.workflows.core_steps.analytics.time_in_zone import ( # noqa: E402 + v3_tensor, +) +from inference.core.workflows.execution_engine.constants import ( # noqa: E402 + TIME_IN_ZONE_KEY_IN_SV_DETECTIONS, +) +from inference.core.workflows.execution_engine.entities.base import ( # noqa: E402 + ImageParentMetadata, + VideoMetadata, + WorkflowImageData, +) +from inference_models.models.base.object_detection import Detections # noqa: E402 + +TimeInZoneBlockV3 = v3_tensor.TimeInZoneBlockV3 + + +def _native_detections(xyxy: np.ndarray, tracker_ids: np.ndarray) -> Detections: + n = len(xyxy) + return Detections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32), + class_id=torch.zeros(n, dtype=torch.int64), + confidence=torch.ones(n, dtype=torch.float32), + bboxes_metadata=[{"tracker_id": int(value)} for value in tracker_ids], + ) + + +def test_time_in_zone_v3_matches_supervision_and_old_bookkeeping() -> None: + zones = [ + [[10, 10], [10, 30], [30, 30], [30, 10]], + [[50, 50], [50, 70], [70, 70], [70, 50]], + ] + frames = [ + ( + np.array([[15, 15, 17, 17], [35, 35, 37, 37], [55, 55, 57, 57]]), + np.array([1, 2, 3]), + ), + ( + np.array([[35, 35, 37, 37], [15, 15, 17, 17], [55, 55, 57, 57]]), + np.array([1, 2, 3]), + ), + ( + np.array([[15, 15, 17, 17], [16, 16, 18, 18], [80, 80, 82, 82]]), + np.array([1, 2, 3]), + ), + ] + oracle_zones = [ + sv.PolygonZone(np.asarray(polygon), triggering_anchors=(sv.Position.CENTER,)) + for polygon in zones + ] + tracked_ids_in_zone = {} + block = TimeInZoneBlockV3() + + for frame_number, (xyxy, tracker_ids) in enumerate(frames, start=10): + metadata = VideoMetadata( + video_identifier="tensor-zone", + frame_number=frame_number, + fps=1, + frame_timestamp=datetime.datetime.now(datetime.timezone.utc), + comes_from_video_file=True, + ) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="root"), + numpy_image=np.zeros((100, 100, 3), dtype=np.uint8), + video_metadata=metadata, + ) + sv_input = sv.Detections(xyxy=xyxy.astype(float), tracker_id=tracker_ids) + is_in_any_zone = np.any( + [polygon_zone.trigger(sv_input) for polygon_zone in oracle_zones], axis=0 + ) + surviving_mask = np.zeros(len(xyxy), dtype=bool) + surviving_times = {} + ts_end = frame_number / metadata.fps + for index, is_in_zone, tracker_id in zip( + range(len(xyxy)), is_in_any_zone, tracker_ids + ): + time_in_zone = 0.0 + if is_in_zone: + ts_start = tracked_ids_in_zone.setdefault(int(tracker_id), ts_end) + time_in_zone = ts_end - ts_start + elif int(tracker_id) in tracked_ids_in_zone: + # Keep the old reset=False deletion quirk as the oracle. + del tracked_ids_in_zone[int(tracker_id)] + surviving_mask[index] = True + surviving_times[index] = time_in_zone + + result = block.run( + image=image, + detections=_native_detections(xyxy, tracker_ids), + zone=zones, + triggering_anchor="CENTER", + remove_out_of_zone_detections=False, + reset_out_of_zone_detections=False, + )["timed_detections"] + + np.testing.assert_array_equal(result.xyxy.numpy(), xyxy[surviving_mask]) + assert [item["tracker_id"] for item in result.bboxes_metadata] == ( + tracker_ids[surviving_mask].tolist() + ) + assert [ + item[TIME_IN_ZONE_KEY_IN_SV_DETECTIONS] for item in result.bboxes_metadata + ] == [surviving_times[index] for index in np.flatnonzero(surviving_mask)] diff --git a/tests/workflows/unit_tests/core_steps/analytics/test_velocity.py b/tests/workflows/unit_tests/core_steps/analytics/test_velocity.py index e3b6d9e38b..562af668f9 100644 --- a/tests/workflows/unit_tests/core_steps/analytics/test_velocity.py +++ b/tests/workflows/unit_tests/core_steps/analytics/test_velocity.py @@ -6,6 +6,7 @@ from inference.core.workflows.core_steps.analytics.velocity.v1 import VelocityBlockV1 from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, VideoMetadata, WorkflowImageData, ) @@ -881,3 +882,239 @@ def test_velocity_block_large_movement() -> None: frame2_result["velocity_detections"].data["smoothed_speed"], expected_data_frame2["smoothed_speed"], ) + + +# --- tensor-native sibling --------------------------------------------------- +# These tests pin numerical parity with the numpy block above plus the tensor +# block's device-resident state semantics. + + +def _tensor_velocity_imports(): + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.analytics.velocity.v1_tensor import ( + VelocityBlockV1 as TensorVelocityBlockV1, + ) + from inference_models.models.base.object_detection import ( + Detections as NativeDetections, + ) + + return torch, TensorVelocityBlockV1, NativeDetections + + +def _tensor_image(video_id: str, frame_number: int, fps: float = 10.0): + metadata = VideoMetadata( + video_identifier=video_id, + frame_number=frame_number, + frame_timestamp=datetime.datetime.fromtimestamp(1690000000 + frame_number), + fps=fps, + comes_from_video_file=True, + ) + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=np.zeros((100, 100, 3), dtype=np.uint8), + video_metadata=metadata, + ) + + +def _native_detections(torch, NativeDetections, xyxy_rows, tracker_ids): + return NativeDetections( + xyxy=torch.tensor(xyxy_rows, dtype=torch.float32), + class_id=torch.zeros(len(xyxy_rows), dtype=torch.long), + confidence=torch.full((len(xyxy_rows),), 0.9), + image_metadata={"class_names": {0: "object"}}, + bboxes_metadata=[ + {"detection_id": f"d{i}", "tracker_id": tid} + for i, tid in enumerate(tracker_ids) + ], + ) + + +def test_tensor_velocity_two_frame_calculation_matches_numpy_math() -> None: + # given - one object moving +10px in x over 1 frame @ 10fps (0.1 s) + torch, TensorVelocityBlockV1, NativeDetections = _tensor_velocity_imports() + block = TensorVelocityBlockV1() + + # when + block.run( + image=_tensor_image("vid", 1), + detections=_native_detections( + torch, NativeDetections, [[0.0, 0.0, 10.0, 10.0]], [1] + ), + smoothing_alpha=0.5, + pixels_per_meter=1.0, + ) + result = block.run( + image=_tensor_image("vid", 2), + detections=_native_detections( + torch, NativeDetections, [[10.0, 0.0, 20.0, 10.0]], [1] + ), + smoothing_alpha=0.5, + pixels_per_meter=1.0, + ) + + # then - dx=10px / 0.1s = 100 px/s; EMA(0.5) of [0,0] -> [100,0] = [50,0] + metadata = result["velocity_detections"].bboxes_metadata[0] + assert np.allclose(metadata["velocity"], [100.0, 0.0], atol=1e-3) + assert np.allclose(metadata["speed"], 100.0, atol=1e-3) + assert np.allclose(metadata["smoothed_velocity"], [50.0, 0.0], atol=1e-3) + assert np.allclose(metadata["smoothed_speed"], 50.0, atol=1e-3) + + +def test_tensor_velocity_new_tracker_and_units() -> None: + # given + torch, TensorVelocityBlockV1, NativeDetections = _tensor_velocity_imports() + block = TensorVelocityBlockV1() + + # when - first sighting, with a pixels_per_meter conversion in play + result = block.run( + image=_tensor_image("vid", 1), + detections=_native_detections( + torch, NativeDetections, [[0.0, 0.0, 10.0, 10.0]], [7] + ), + smoothing_alpha=0.5, + pixels_per_meter=100.0, + ) + + # then - no previous position: everything is zero + metadata = result["velocity_detections"].bboxes_metadata[0] + assert metadata["velocity"] == [0.0, 0.0] + assert metadata["speed"] == 0.0 + assert metadata["smoothed_velocity"] == [0.0, 0.0] + assert metadata["smoothed_speed"] == 0.0 + + +def test_tensor_velocity_absent_tracker_reappears_with_preserved_state() -> None: + # given - tracker 1 seen at frame 1, absent at frame 2, back at frame 3; + # velocity on reappearance must use the frame-1 position and timestamp + torch, TensorVelocityBlockV1, NativeDetections = _tensor_velocity_imports() + block = TensorVelocityBlockV1() + block.run( + image=_tensor_image("vid", 1), + detections=_native_detections( + torch, NativeDetections, [[0.0, 0.0, 10.0, 10.0]], [1] + ), + smoothing_alpha=1.0, + pixels_per_meter=1.0, + ) + block.run( + image=_tensor_image("vid", 2), + detections=_native_detections( + torch, NativeDetections, [[100.0, 100.0, 110.0, 110.0]], [2] + ), + smoothing_alpha=1.0, + pixels_per_meter=1.0, + ) + + # when - tracker 1 reappears at frame 3, +20px in x since frame 1 (0.2 s) + result = block.run( + image=_tensor_image("vid", 3), + detections=_native_detections( + torch, NativeDetections, [[20.0, 0.0, 30.0, 10.0]], [1] + ), + smoothing_alpha=1.0, + pixels_per_meter=1.0, + ) + + # then - 20px / 0.2s = 100 px/s + metadata = result["velocity_detections"].bboxes_metadata[0] + assert np.allclose(metadata["velocity"], [100.0, 0.0], atol=1e-3) + + +def test_tensor_velocity_zero_delta_time_yields_zero_velocity() -> None: + # given - the same frame timestamp twice + torch, TensorVelocityBlockV1, NativeDetections = _tensor_velocity_imports() + block = TensorVelocityBlockV1() + detections = _native_detections( + torch, NativeDetections, [[0.0, 0.0, 10.0, 10.0]], [1] + ) + block.run( + image=_tensor_image("vid", 5), + detections=detections, + smoothing_alpha=0.5, + pixels_per_meter=1.0, + ) + + # when + result = block.run( + image=_tensor_image("vid", 5), + detections=_native_detections( + torch, NativeDetections, [[50.0, 0.0, 60.0, 10.0]], [1] + ), + smoothing_alpha=0.5, + pixels_per_meter=1.0, + ) + + # then + metadata = result["velocity_detections"].bboxes_metadata[0] + assert metadata["velocity"] == [0.0, 0.0] + assert metadata["speed"] == 0.0 + + +def test_tensor_velocity_validations() -> None: + # given + torch, TensorVelocityBlockV1, NativeDetections = _tensor_velocity_imports() + block = TensorVelocityBlockV1() + untracked = _native_detections( + torch, NativeDetections, [[0.0, 0.0, 10.0, 10.0]], [None] + ) + + # when / then + with pytest.raises(ValueError, match="tracker_id"): + block.run( + image=_tensor_image("vid", 1), + detections=untracked, + smoothing_alpha=0.5, + pixels_per_meter=1.0, + ) + tracked = _native_detections(torch, NativeDetections, [[0.0, 0.0, 10.0, 10.0]], [1]) + with pytest.raises(ValueError, match="smoothing_alpha"): + block.run( + image=_tensor_image("vid", 1), + detections=tracked, + smoothing_alpha=0.0, + pixels_per_meter=1.0, + ) + with pytest.raises(ValueError, match="pixels_per_meter"): + block.run( + image=_tensor_image("vid", 1), + detections=tracked, + smoothing_alpha=0.5, + pixels_per_meter=0.0, + ) + + +def test_tensor_velocity_state_stays_on_device_when_mps_available() -> None: + # given + torch, TensorVelocityBlockV1, NativeDetections = _tensor_velocity_imports() + if not torch.backends.mps.is_available(): + pytest.skip("device-residency check needs MPS") + device = torch.device("mps") + block = TensorVelocityBlockV1() + + def on_device(xyxy_rows, tracker_ids): + detections = _native_detections(torch, NativeDetections, xyxy_rows, tracker_ids) + detections.xyxy = detections.xyxy.to(device) + detections.class_id = detections.class_id.to(device) + detections.confidence = detections.confidence.to(device) + return detections + + # when + block.run( + image=_tensor_image("vid", 1), + detections=on_device([[0.0, 0.0, 10.0, 10.0]], [1]), + smoothing_alpha=0.5, + pixels_per_meter=1.0, + ) + result = block.run( + image=_tensor_image("vid", 2), + detections=on_device([[10.0, 0.0, 20.0, 10.0]], [1]), + smoothing_alpha=0.5, + pixels_per_meter=1.0, + ) + + # then - the tracking state lives on the device + state = block._states["vid"] + assert state.positions.device.type == "mps" + assert state.smoothed_velocities.device.type == "mps" + metadata = result["velocity_detections"].bboxes_metadata[0] + assert np.allclose(metadata["velocity"], [100.0, 0.0], atol=1e-2) diff --git a/tests/workflows/unit_tests/core_steps/analytics/test_zone_geometry_parity.py b/tests/workflows/unit_tests/core_steps/analytics/test_zone_geometry_parity.py new file mode 100644 index 0000000000..b13a59f9ff --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/analytics/test_zone_geometry_parity.py @@ -0,0 +1,144 @@ +import importlib.util +from pathlib import Path + +import numpy as np +import supervision as sv + +_MODULE_PATH = ( + Path(__file__).parents[5] + / "inference/core/workflows/core_steps/analytics/_zone_geometry.py" +) +_MODULE_SPEC = importlib.util.spec_from_file_location("_zone_geometry", _MODULE_PATH) +_ZONE_GEOMETRY = importlib.util.module_from_spec(_MODULE_SPEC) +_MODULE_SPEC.loader.exec_module(_ZONE_GEOMETRY) +LeanLineZone = _ZONE_GEOMETRY.LeanLineZone +LeanPolygonZone = _ZONE_GEOMETRY.LeanPolygonZone +anchor_coordinates = _ZONE_GEOMETRY.anchor_coordinates + + +def test_anchor_coordinates_matches_supervision_for_every_position() -> None: + rng = np.random.default_rng(731) + top_left = rng.uniform(-100, 500, size=(50, 2)) + xyxy = np.concatenate( + [top_left, top_left + rng.uniform(0, 100, size=(50, 2))], axis=1 + ).astype(float) + detections = sv.Detections(xyxy=xyxy) + + for position in sv.Position: + if position == sv.Position.CENTER_OF_MASS: + try: + anchor_coordinates(xyxy, position) + except ValueError as error: + lean_error = error + else: + raise AssertionError("Lean anchor calculation did not raise") + try: + detections.get_anchors_coordinates(position) + except ValueError as error: + supervision_error = error + else: + raise AssertionError("Supervision anchor calculation did not raise") + assert str(lean_error) == str(supervision_error) + else: + np.testing.assert_array_equal( + anchor_coordinates(xyxy, position), + detections.get_anchors_coordinates(position), + ) + + +def test_lean_line_zone_matches_supervision() -> None: + rng = np.random.default_rng(904) + scenarios_per_anchor = 70 + + for triggering_anchors in ( + None, + [sv.Position.CENTER], + [sv.Position.BOTTOM_CENTER], + ): + for scenario in range(scenarios_per_anchor): + line_kind = scenario % 3 + if line_kind == 0: + start = (0.0, float(rng.uniform(20, 180))) + end = (200.0, start[1]) + elif line_kind == 1: + start = (float(rng.uniform(20, 180)), 0.0) + end = (start[0], 200.0) + else: + start = tuple(rng.uniform(-50, 50, size=2)) + delta = rng.uniform(-160, 160, size=2) + if np.linalg.norm(delta) < 1: + delta[0] += 40 + end = tuple(np.asarray(start) + delta) + + lean_zone = LeanLineZone(start, end, triggering_anchors) + supervision_kwargs = ( + {} + if triggering_anchors is None + else {"triggering_anchors": triggering_anchors} + ) + supervision_zone = sv.LineZone( + start=sv.Point(*start), end=sv.Point(*end), **supervision_kwargs + ) + + n = int(rng.integers(0, 31)) + tracker_ids = np.arange(n, dtype=int) + scenario * 100 + centers = rng.uniform(-100, 300, size=(n, 2)) + sizes = rng.uniform(0, 50, size=(n, 2)) + velocity = rng.normal(0, 8, size=(n, 2)) + if n: + centers[0] = np.asarray(start) + sizes[0] = 30 + + for frame in range(40): + centers += velocity + rng.normal(0, 2, size=(n, 2)) + if n and frame % 8 == 0: + velocity[: max(1, n // 5)] *= -1 + xyxy = np.concatenate( + [centers - sizes / 2, centers + sizes / 2], axis=1 + ).astype(float) + + lean_in, lean_out = lean_zone.trigger(xyxy, tracker_ids) + supervision_in, supervision_out = supervision_zone.trigger( + sv.Detections(xyxy=xyxy, tracker_id=tracker_ids) + ) + + np.testing.assert_array_equal(lean_in, supervision_in) + np.testing.assert_array_equal(lean_out, supervision_out) + assert lean_zone.in_count == supervision_zone.in_count + assert lean_zone.out_count == supervision_zone.out_count + + +def test_lean_polygon_zone_matches_supervision() -> None: + rng = np.random.default_rng(117) + + for triggering_anchors in ( + (sv.Position.CENTER,), + (sv.Position.BOTTOM_CENTER,), + ): + for scenario in range(220): + vertex_count = int(rng.integers(3, 9)) + center = rng.uniform(80, 220, size=2) + angles = np.sort(rng.uniform(0, 2 * np.pi, size=vertex_count)) + radii = rng.uniform(15, 75, size=vertex_count) + polygon = np.stack( + [ + center[0] + np.cos(angles) * radii, + center[1] + np.sin(angles) * radii, + ], + axis=1, + ) + + n = int(rng.integers(0, 31)) + top_left = rng.uniform(-80, 360, size=(n, 2)) + sizes = rng.uniform(0, 90, size=(n, 2)) + if n: + sizes[0] = 0 + xyxy = np.concatenate([top_left, top_left + sizes], axis=1).astype(float) + + lean_zone = LeanPolygonZone(polygon, triggering_anchors) + supervision_zone = sv.PolygonZone(polygon, triggering_anchors) + np.testing.assert_array_equal( + lean_zone.trigger(xyxy), + supervision_zone.trigger(sv.Detections(xyxy=xyxy)), + ) + assert lean_zone.current_count == supervision_zone.current_count diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_blur.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_blur.py index 4c39c97d79..4e2368e4f3 100644 --- a/tests/workflows/unit_tests/core_steps/classical_cv/test_blur.py +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_blur.py @@ -77,3 +77,254 @@ def test_image_blur_block() -> None: assert output.get("image").numpy_image.shape == (1000, 1000, 3) # check if the image is modified assert not np.array_equal(output.get("image").numpy_image, start_image) + + +# --- tensor-native sibling --------------------------------------------------- +# Parity contract: outputs match the numpy block bit-exactly on every path. + + +def _tensor_blur_imports(): + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.classical_cv.image_blur.v1_tensor import ( + ImageBlurBlockV1 as TensorImageBlurBlockV1, + ) + + return torch, TensorImageBlurBlockV1 + + +def _paired_images(torch, bgr: np.ndarray): + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + if bgr.ndim == 2: + chw = torch.from_numpy(bgr.copy()).unsqueeze(0) + else: + chw = torch.from_numpy(bgr[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=chw, + ) + return numpy_born, tensor_born + + +def _blur_case_image(case: str) -> np.ndarray: + rng = np.random.default_rng(42) + if case == "noise_color": + return rng.integers(0, 256, size=(37, 53, 3), dtype=np.uint8) + if case == "noise_gray": + return rng.integers(0, 256, size=(41, 29), dtype=np.uint8) + if case == "extremes": + return rng.choice([0, 255], size=(32, 24, 3)).astype(np.uint8) + if case == "gradient": + ramp = np.tile(np.arange(256, dtype=np.uint8), (8, 1)) + return np.stack([ramp, ramp[:, ::-1], 255 - ramp], axis=-1) + raise ValueError(case) + + +def _run_both_blur_blocks( + torch, tensor_block_class, bgr: np.ndarray, blur_type: str, kernel_size: int +): + numpy_born, tensor_born = _paired_images(torch, bgr) + numpy_result = ImageBlurBlockV1().run( + image=numpy_born, blur_type=blur_type, kernel_size=kernel_size + )["image"] + tensor_result = tensor_block_class().run( + image=tensor_born, blur_type=blur_type, kernel_size=kernel_size + )["image"] + return numpy_result, tensor_result + + +@pytest.mark.parametrize("kernel_size", [1, 3, 5, 9, 15]) +@pytest.mark.parametrize("case", ["noise_color", "noise_gray", "extremes"]) +def test_tensor_image_blur_average_bit_exact_parity(kernel_size, case) -> None: + # given - the same pixels as numpy-born BGR and tensor-born RGB CHW + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + bgr = _blur_case_image(case) + + # when + numpy_result, tensor_result = _run_both_blur_blocks( + torch, TensorImageBlurBlockV1, bgr, "average", kernel_size + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +@pytest.mark.parametrize("kernel_size", [1, 3, 4, 5, 6, 7]) +@pytest.mark.parametrize("case", ["noise_color", "noise_gray", "gradient"]) +def test_tensor_image_blur_gaussian_bit_exact_parity(kernel_size, case) -> None: + # given - even sizes exercise v1's _to_positive_odd coercion (4 -> 5, 6 -> 7) + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + bgr = _blur_case_image(case) + + # when + numpy_result, tensor_result = _run_both_blur_blocks( + torch, TensorImageBlurBlockV1, bgr, "gaussian", kernel_size + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +@pytest.mark.parametrize("kernel_size", [2, 3, 5, 9, 15]) +@pytest.mark.parametrize("case", ["noise_color", "noise_gray", "extremes"]) +def test_tensor_image_blur_median_bit_exact_parity(kernel_size, case) -> None: + # given - kernel_size=2 exercises v1's _to_positive_odd coercion (2 -> 3) + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + bgr = _blur_case_image(case) + + # when + numpy_result, tensor_result = _run_both_blur_blocks( + torch, TensorImageBlurBlockV1, bgr, "median", kernel_size + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +@pytest.mark.parametrize("kernel_size", [2, 4]) +def test_tensor_image_blur_average_even_kernel_ties_delegate(kernel_size) -> None: + # given - 0/1 checkerboard: every window mean is an exact .5 rounding tie; + # cv2's tie rounding is build-specific, so even kernels delegate + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + yy, xx = np.mgrid[0:16, 0:20] + checker = ((yy + xx) % 2).astype(np.uint8) + + # when + numpy_result, tensor_result = _run_both_blur_blocks( + torch, TensorImageBlurBlockV1, checker, "average", kernel_size + ) + + # then + assert not tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +def test_tensor_image_blur_gaussian_large_kernel_delegates() -> None: + # given - kernel 9 is beyond the small_gaussian_tab regime, so gaussian delegates + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + bgr = _blur_case_image("noise_color") + + # when + numpy_result, tensor_result = _run_both_blur_blocks( + torch, TensorImageBlurBlockV1, bgr, "gaussian", 9 + ) + + # then + assert not tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +def test_tensor_image_blur_median_large_kernel_delegates() -> None: + # given - median kernels above 15 delegate + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + bgr = _blur_case_image("noise_color") + + # when + numpy_result, tensor_result = _run_both_blur_blocks( + torch, TensorImageBlurBlockV1, bgr, "median", 17 + ) + + # then + assert not tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +def test_tensor_image_blur_bilateral_delegates() -> None: + # given - bilateral always delegates (data-dependent float weights) + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + bgr = _blur_case_image("noise_color") + + # when + numpy_result, tensor_result = _run_both_blur_blocks( + torch, TensorImageBlurBlockV1, bgr, "bilateral", 5 + ) + + # then + assert not tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +def test_tensor_image_blur_oversized_kernel_regime_delegates() -> None: + # given - kernel pad 7 reaches the 6-pixel image height, so this regime delegates + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + rng = np.random.default_rng(7) + bgr = rng.integers(0, 256, size=(6, 9, 3), dtype=np.uint8) + + # when + numpy_result, tensor_result = _run_both_blur_blocks( + torch, TensorImageBlurBlockV1, bgr, "average", 15 + ) + + # then + assert not tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +@pytest.mark.parametrize("blur_type", ["average", "gaussian", "median", "bilateral"]) +def test_tensor_image_blur_delegates_for_numpy_born_images(blur_type) -> None: + # given + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + bgr = _blur_case_image("noise_color") + numpy_born, _ = _paired_images(torch, bgr) + reference = ( + ImageBlurBlockV1() + .run(image=numpy_born, blur_type=blur_type, kernel_size=5)["image"] + .numpy_image + ) + + # when + result = TensorImageBlurBlockV1().run( + image=numpy_born, blur_type=blur_type, kernel_size=5 + )["image"] + + # then + assert np.array_equal(result.numpy_image, reference) + assert not numpy_born.is_tensor_materialised(), "delegate must not materialise" + + +def test_tensor_image_blur_unknown_blur_type_raises() -> None: + # given + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.zeros((3, 8, 8), dtype=torch.uint8), + ) + + # when / then + with pytest.raises(ValueError, match="Unknown blur type"): + TensorImageBlurBlockV1().run( + image=tensor_born, blur_type="motion", kernel_size=5 + ) + + +def test_tensor_image_blur_on_mps_device(monkeypatch) -> None: + # given - patch the global device so tensor-born images materialise on MPS + torch, TensorImageBlurBlockV1 = _tensor_blur_imports() + if not torch.backends.mps.is_available(): + pytest.skip("MPS device not available") + import inference.core.workflows.execution_engine.entities.base as base_module + + bgr = _blur_case_image("noise_color") + numpy_born, _ = _paired_images(torch, bgr) + reference = ( + ImageBlurBlockV1() + .run(image=numpy_born, blur_type="average", kernel_size=5)["image"] + .numpy_image + ) + monkeypatch.setattr(base_module, "WORKFLOWS_IMAGE_TENSOR_DEVICE", "mps") + _, tensor_born = _paired_images(torch, bgr) + assert tensor_born.tensor_image.device.type == "mps" + + # when + result = TensorImageBlurBlockV1().run( + image=tensor_born, blur_type="average", kernel_size=5 + )["image"] + + # then + assert result.tensor_image.device.type == "mps" + assert np.array_equal(result.numpy_image, reference) diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_contrast_enhancement.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_contrast_enhancement.py index b192a1d318..c42746a545 100644 --- a/tests/workflows/unit_tests/core_steps/classical_cv/test_contrast_enhancement.py +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_contrast_enhancement.py @@ -183,3 +183,169 @@ def test_contrast_enhancement_manifest_outputs(self): assert len(outputs) == 1 assert outputs[0].name == "image" + + +# --- tensor-native sibling --------------------------------------------------- +# Parity contract: a tensor-born image through the v1_tensor block must produce +# the same pixels as the numpy block on the equivalent BGR image; numpy-born +# images delegate to the numpy implementation without forcing materialization. + + +def _tensor_contrast_imports(): + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.classical_cv.contrast_enhancement.v1_tensor import ( + ContrastEnhancementBlock as TensorContrastEnhancementBlock, + ) + + return torch, TensorContrastEnhancementBlock + + +def _paired_images(torch, bgr: np.ndarray): + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + if bgr.ndim == 2: + chw = torch.from_numpy(bgr.copy()).unsqueeze(0) + else: + chw = torch.from_numpy(bgr[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=chw, + ) + return numpy_born, tensor_born + + +@pytest.mark.parametrize( + "clip_limit, contrast_multiplier, normalize_brightness, exact", + [ + (0, 1.0, False, True), + (0, 1.5, False, True), + (0, 1.0, True, True), + (3, 1.0, False, False), # np.percentile computes in float64 -> +-1 + (3, 1.8, True, False), + ], +) +def test_tensor_contrast_enhancement_parity_with_numpy_block( + clip_limit, contrast_multiplier, normalize_brightness, exact +) -> None: + # given - the same pixels as numpy-born BGR and tensor-born RGB CHW + torch, TensorContrastEnhancementBlock = _tensor_contrast_imports() + rng = np.random.default_rng(42) + bgr = rng.integers(20, 200, size=(24, 32, 3), dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = ( + ContrastEnhancementBlock() + .run( + image=numpy_born, + clip_limit=clip_limit, + contrast_multiplier=contrast_multiplier, + normalize_brightness=normalize_brightness, + )["image"] + .numpy_image + ) + tensor_result = ( + TensorContrastEnhancementBlock() + .run( + image=tensor_born, + clip_limit=clip_limit, + contrast_multiplier=contrast_multiplier, + normalize_brightness=normalize_brightness, + )["image"] + .numpy_image + ) + + # then + if exact: + assert np.array_equal(tensor_result, numpy_result) + else: + diff = np.abs(tensor_result.astype(np.int16) - numpy_result.astype(np.int16)) + assert diff.max() <= 1, f"max deviation {diff.max()} exceeds float tolerance" + + +def test_tensor_contrast_enhancement_grayscale_parity() -> None: + # given - a low-contrast grayscale ramp + torch, TensorContrastEnhancementBlock = _tensor_contrast_imports() + gray = np.tile(np.linspace(90, 160, 32, dtype=np.uint8), (24, 1)) + numpy_born, tensor_born = _paired_images(torch, gray) + + # when + numpy_result = ( + ContrastEnhancementBlock() + .run( + image=numpy_born, + clip_limit=0, + contrast_multiplier=1.0, + normalize_brightness=False, + )["image"] + .numpy_image + ) + tensor_result = ( + TensorContrastEnhancementBlock() + .run( + image=tensor_born, + clip_limit=0, + contrast_multiplier=1.0, + normalize_brightness=False, + )["image"] + .numpy_image + ) + + # then - tensor path emits (1, H, W) whose numpy view is the (H, W) shape + assert tensor_result.shape == numpy_result.shape == (24, 32) + assert np.array_equal(tensor_result, numpy_result) + + +def test_tensor_contrast_enhancement_flat_grayscale_returns_input() -> None: + # given - a flat histogram: nothing to stretch (numpy grayscale parity) + torch, TensorContrastEnhancementBlock = _tensor_contrast_imports() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.full((1, 8, 8), 127, dtype=torch.uint8), + ) + + # when + result = TensorContrastEnhancementBlock().run( + image=tensor_born, + clip_limit=0, + contrast_multiplier=1.0, + normalize_brightness=False, + )["image"] + + # then + assert result is tensor_born + + +def test_tensor_contrast_enhancement_delegates_for_numpy_born_images() -> None: + # given + torch, TensorContrastEnhancementBlock = _tensor_contrast_imports() + rng = np.random.default_rng(7) + bgr = rng.integers(40, 180, size=(16, 16, 3), dtype=np.uint8) + numpy_born, _ = _paired_images(torch, bgr) + reference = ( + ContrastEnhancementBlock() + .run( + image=WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ), + clip_limit=2, + contrast_multiplier=1.2, + normalize_brightness=True, + )["image"] + .numpy_image + ) + + # when + result = TensorContrastEnhancementBlock().run( + image=numpy_born, + clip_limit=2, + contrast_multiplier=1.2, + normalize_brightness=True, + )["image"] + + # then - identical output via the numpy delegate, and no forced H2D + assert np.array_equal(result.numpy_image, reference) + assert not numpy_born.is_tensor_materialised(), "delegate must not materialise" diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_contrast_equalization.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_contrast_equalization.py new file mode 100644 index 0000000000..b327d386c9 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_contrast_equalization.py @@ -0,0 +1,324 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from inference.core.workflows.core_steps.classical_cv.contrast_equalization.v1 import ( + ContrastEqualizationBlockV1, + ContrastEqualizationManifest, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) + +ALL_METHODS = [ + "Contrast Stretching", + "Histogram Equalization", + "Adaptive Equalization", +] + + +class TestContrastEqualizationManifest: + def test_contrast_equalization_validation_when_valid_manifest_is_given(self): + manifest = ContrastEqualizationManifest.model_validate( + { + "type": "roboflow_core/contrast_equalization@v1", + "name": "contrast_equalization", + "image": "$inputs.image", + "equalization_type": "Contrast Stretching", + } + ) + + assert manifest.type == "roboflow_core/contrast_equalization@v1" + assert manifest.name == "contrast_equalization" + assert manifest.equalization_type == "Contrast Stretching" + + def test_contrast_equalization_validation_when_image_is_missing(self): + with pytest.raises(ValidationError): + ContrastEqualizationManifest.model_validate( + { + "type": "roboflow_core/contrast_equalization@v1", + "name": "contrast_equalization", + } + ) + + def test_contrast_equalization_manifest_outputs(self): + outputs = ContrastEqualizationManifest.describe_outputs() + + assert len(outputs) == 1 + assert outputs[0].name == "image" + + +@pytest.mark.parametrize("equalization_type", ALL_METHODS) +def test_contrast_equalization_block_with_color_image(equalization_type) -> None: + # given + rng = np.random.default_rng(11) + image_data = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="test"), + numpy_image=rng.integers(60, 190, size=(48, 64, 3), dtype=np.uint8), + ) + + # when + result = ContrastEqualizationBlockV1().run( + image=image_data, equalization_type=equalization_type + ) + + # then + assert "image" in result + equalized = result["image"].numpy_image + assert equalized.shape == (48, 64, 3) + assert equalized.dtype == np.uint8 + + +@pytest.mark.parametrize("equalization_type", ALL_METHODS) +def test_contrast_equalization_block_with_grayscale_image(equalization_type) -> None: + # given + rng = np.random.default_rng(12) + image_data = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="test"), + numpy_image=rng.integers(90, 150, size=(48, 64), dtype=np.uint8), + ) + + # when + result = ContrastEqualizationBlockV1().run( + image=image_data, equalization_type=equalization_type + ) + + # then + assert "image" in result + equalized = result["image"].numpy_image + assert equalized.shape == (48, 64) + assert equalized.dtype == np.uint8 + + +def test_contrast_equalization_block_with_unknown_method() -> None: + # given + image_data = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="test"), + numpy_image=np.zeros((10, 10, 3), dtype=np.uint8), + ) + + # when / then + with pytest.raises(ValueError): + ContrastEqualizationBlockV1().run( + image=image_data, equalization_type="Unknown Method" + ) + + +# --- tensor-native sibling --------------------------------------------------- +# Parity contract: for the two LUT methods a tensor-born image must produce +# bit-identical pixels to the numpy block on the equivalent BGR image; CLAHE +# and numpy-born images delegate to the numpy implementation. + + +def _tensor_equalization_imports(): + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.classical_cv.contrast_equalization.v1_tensor import ( + ContrastEqualizationBlockV1 as TensorContrastEqualizationBlockV1, + ) + + return torch, TensorContrastEqualizationBlockV1 + + +def _paired_images(torch, bgr: np.ndarray): + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + if bgr.ndim == 2: + chw = torch.from_numpy(bgr.copy()).unsqueeze(0) + else: + chw = torch.from_numpy(bgr[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=chw, + ) + return numpy_born, tensor_born + + +def _case_image(case: str) -> np.ndarray: + rng = np.random.default_rng(42) + if case == "noise": + return rng.integers(0, 256, size=(24, 32, 3), dtype=np.uint8) + if case == "narrow_range": + return rng.integers(100, 141, size=(24, 32, 3), dtype=np.uint8) + if case == "gradient": + ramp = np.tile(np.arange(256, dtype=np.uint8), (8, 1)) + return np.stack([ramp, ramp[:, ::-1], ramp], axis=-1) + if case == "skewed": + image = rng.integers(0, 256, size=(32, 32, 3), dtype=np.uint8) + image[rng.random(image.shape) < 0.99] = 77 # p2 == p98 == 77 + return image + raise ValueError(case) + + +@pytest.mark.parametrize( + "equalization_type", ["Contrast Stretching", "Histogram Equalization"] +) +@pytest.mark.parametrize("case", ["noise", "narrow_range", "gradient", "skewed"]) +def test_tensor_contrast_equalization_bit_exact_parity(equalization_type, case) -> None: + # given - the same pixels as numpy-born BGR and tensor-born RGB CHW + torch, TensorContrastEqualizationBlockV1 = _tensor_equalization_imports() + bgr = _case_image(case) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = ( + ContrastEqualizationBlockV1() + .run(image=numpy_born, equalization_type=equalization_type)["image"] + .numpy_image + ) + tensor_result_image = TensorContrastEqualizationBlockV1().run( + image=tensor_born, equalization_type=equalization_type + )["image"] + + # then - bit-exact parity, and the output stays tensor-born + assert tensor_result_image.is_tensor_materialised() + assert np.array_equal(tensor_result_image.numpy_image, numpy_result) + + +@pytest.mark.parametrize( + "equalization_type", ["Contrast Stretching", "Histogram Equalization"] +) +@pytest.mark.parametrize("value", [0, 130, 255]) +def test_tensor_contrast_equalization_flat_grayscale_parity( + equalization_type, value +) -> None: + # given - degenerate histogram: single distinct value + torch, TensorContrastEqualizationBlockV1 = _tensor_equalization_imports() + gray = np.full((16, 20), value, dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, gray) + + # when + numpy_result = ( + ContrastEqualizationBlockV1() + .run(image=numpy_born, equalization_type=equalization_type)["image"] + .numpy_image + ) + tensor_result = ( + TensorContrastEqualizationBlockV1() + .run(image=tensor_born, equalization_type=equalization_type)["image"] + .numpy_image + ) + + # then - tensor path emits (1, H, W) whose numpy view is the (H, W) shape + assert tensor_result.shape == numpy_result.shape == (16, 20) + assert np.array_equal(tensor_result, numpy_result) + + +def test_tensor_contrast_equalization_grayscale_parity() -> None: + # given - a low-contrast grayscale ramp + torch, TensorContrastEqualizationBlockV1 = _tensor_equalization_imports() + gray = np.tile(np.linspace(90, 160, 32, dtype=np.uint8), (24, 1)) + numpy_born, tensor_born = _paired_images(torch, gray) + + for equalization_type in ["Contrast Stretching", "Histogram Equalization"]: + # when + numpy_result = ( + ContrastEqualizationBlockV1() + .run(image=numpy_born, equalization_type=equalization_type)["image"] + .numpy_image + ) + tensor_result = ( + TensorContrastEqualizationBlockV1() + .run(image=tensor_born, equalization_type=equalization_type)["image"] + .numpy_image + ) + + # then + assert tensor_result.shape == numpy_result.shape == (24, 32) + assert np.array_equal(tensor_result, numpy_result) + + +def test_tensor_contrast_equalization_adaptive_delegates_to_numpy_math() -> None: + # given - CLAHE's mapping is tile-local, so the tensor block delegates to + # the numpy implementation even for tensor-born images + torch, TensorContrastEqualizationBlockV1 = _tensor_equalization_imports() + bgr = _case_image("noise") + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = ( + ContrastEqualizationBlockV1() + .run(image=numpy_born, equalization_type="Adaptive Equalization")["image"] + .numpy_image + ) + tensor_result = ( + TensorContrastEqualizationBlockV1() + .run(image=tensor_born, equalization_type="Adaptive Equalization")["image"] + .numpy_image + ) + + # then + assert np.array_equal(tensor_result, numpy_result) + + +def test_tensor_contrast_equalization_delegates_for_numpy_born_images() -> None: + # given + torch, TensorContrastEqualizationBlockV1 = _tensor_equalization_imports() + bgr = _case_image("narrow_range") + numpy_born, _ = _paired_images(torch, bgr) + reference = ( + ContrastEqualizationBlockV1() + .run( + image=WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ), + equalization_type="Histogram Equalization", + )["image"] + .numpy_image + ) + + # when + result = TensorContrastEqualizationBlockV1().run( + image=numpy_born, equalization_type="Histogram Equalization" + )["image"] + + # then - identical output via the numpy delegate, and no forced H2D + assert np.array_equal(result.numpy_image, reference) + assert not numpy_born.is_tensor_materialised(), "delegate must not materialise" + + +def test_tensor_contrast_equalization_unknown_method_raises() -> None: + # given + torch, TensorContrastEqualizationBlockV1 = _tensor_equalization_imports() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.zeros((3, 8, 8), dtype=torch.uint8), + ) + + # when / then + with pytest.raises(ValueError): + TensorContrastEqualizationBlockV1().run( + image=tensor_born, equalization_type="Unknown Method" + ) + + +def test_tensor_contrast_equalization_on_mps_device(monkeypatch) -> None: + # given - tensor images land on the globally configured device; simulate an + # MPS deployment by patching the global + torch, TensorContrastEqualizationBlockV1 = _tensor_equalization_imports() + if not torch.backends.mps.is_available(): + pytest.skip("MPS device not available") + import inference.core.workflows.execution_engine.entities.base as base_module + + bgr = _case_image("noise") + numpy_born, _ = _paired_images(torch, bgr) + reference = ( + ContrastEqualizationBlockV1() + .run(image=numpy_born, equalization_type="Contrast Stretching")["image"] + .numpy_image + ) + monkeypatch.setattr(base_module, "WORKFLOWS_IMAGE_TENSOR_DEVICE", "mps") + _, tensor_born = _paired_images(torch, bgr) + assert tensor_born.tensor_image.device.type == "mps" + + # when + result = TensorContrastEqualizationBlockV1().run( + image=tensor_born, equalization_type="Contrast Stretching" + )["image"] + + # then - stays on device, and matches the numpy block bit-exactly + assert result.tensor_image.device.type == "mps" + assert np.array_equal(result.numpy_image, reference) diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_detections_nearest_neighbor.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_detections_nearest_neighbor.py index dcd77ca1c0..109a3deeb3 100644 --- a/tests/workflows/unit_tests/core_steps/classical_cv/test_detections_nearest_neighbor.py +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_detections_nearest_neighbor.py @@ -3,7 +3,9 @@ import numpy as np import pytest import supervision as sv +import torch +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.classical_cv.detections_nearest_neighbor.v1 import ( NEAREST_TARGET_DISTANCE_KEY, OUTPUT_KEY_MATCHED_QUERY_DETECTIONS, @@ -11,6 +13,28 @@ OUTPUT_KEY_QUERY_PREDICTIONS, DetectionsNearestNeighborBlockV1, ) +from inference.core.workflows.core_steps.classical_cv.detections_nearest_neighbor.v1_tensor import ( + DetectionsNearestNeighborBlockV1 as TensorNativeDetectionsNearestNeighborBlockV1, +) +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY +from inference_models.models.base.instance_segmentation import ( + InstanceDetections as NativeInstanceDetections, +) +from inference_models.models.base.keypoints_detection import ( + KeyPoints as NativeKeyPoints, +) +from inference_models.models.base.object_detection import Detections as NativeDetections + +# The sv-based tests below exercise the numpy v1 module directly, which is +# flag-agnostic - they pass in both flag directions, so no _NUMPY_ONLY marker is +# needed. Each has a `*_tensor_native` parity sibling (skipped when the flag is +# off) feeding the native inference_models representation into the v1_tensor +# block and asserting the same scenario, with per-box results read from +# `bboxes_metadata` (the native equivalent of the sv `.data` columns). +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) def make_detections(xyxy, detection_ids=None, class_name="object") -> sv.Detections: @@ -377,3 +401,487 @@ def test_keypoint_option_requires_keypoint_predictions() -> None: query_point="KEYPOINT", query_keypoint_name="left_shoulder", ) + + +def make_native_detections( + xyxy, detection_ids=None, class_name="object" +) -> NativeDetections: + n = len(xyxy) + bboxes_metadata = None + if detection_ids is not None: + bboxes_metadata = [ + {"detection_id": detection_id} for detection_id in detection_ids + ] + return NativeDetections( + xyxy=torch.tensor(xyxy, dtype=torch.float32).reshape(-1, 4), + class_id=torch.zeros(n, dtype=torch.long), + confidence=torch.full((n,), 0.9, dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: class_name}}, + bboxes_metadata=bboxes_metadata, + ) + + +def make_native_empty_detections() -> NativeDetections: + return NativeDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {}}, + bboxes_metadata=None, + ) + + +def make_native_keypoint_prediction( + xyxy, keypoints_xy, keypoints_class_name, detection_ids=None +): + """Build the native keypoint-detection shape: a ``(KeyPoints, Detections)`` + tuple whose bbox component carries the flattened per-box ``keypoints_xy`` / + ``keypoints_class_name`` payloads in ``bboxes_metadata`` (the native mirror + of the sv data columns).""" + detections = make_native_detections(xyxy=xyxy, detection_ids=detection_ids) + n = len(xyxy) + if detections.bboxes_metadata is None: + detections.bboxes_metadata = [{} for _ in range(n)] + max_keypoints = max(len(instance_xy) for instance_xy in keypoints_xy) + padded_xy = torch.zeros((n, max_keypoints, 2), dtype=torch.float32) + padded_confidence = torch.zeros((n, max_keypoints), dtype=torch.float32) + for index, (instance_xy, instance_names) in enumerate( + zip(keypoints_xy, keypoints_class_name) + ): + detections.bboxes_metadata[index]["keypoints_xy"] = [ + [float(x), float(y)] for x, y in instance_xy + ] + detections.bboxes_metadata[index]["keypoints_class_name"] = list(instance_names) + if len(instance_xy) > 0: + padded_xy[index, : len(instance_xy)] = torch.tensor( + instance_xy, dtype=torch.float32 + ) + padded_confidence[index, : len(instance_xy)] = 1.0 + key_points = NativeKeyPoints( + xy=padded_xy, + class_id=torch.zeros(n, dtype=torch.long), + confidence=padded_confidence, + image_metadata=detections.image_metadata, + ) + return key_points, detections + + +def bbox_component(prediction) -> NativeDetections: + if isinstance(prediction, tuple): + return prediction[1] + return prediction + + +def nearest_distances(prediction) -> list: + detections = bbox_component(prediction) + return [ + box_metadata[NEAREST_TARGET_DISTANCE_KEY] + for box_metadata in detections.bboxes_metadata + ] + + +def run_tensor_block( + query, + target, + query_point="CENTER", + target_point="CENTER", + query_keypoint_name=None, + target_keypoint_name=None, + max_distance=None, +): + block = TensorNativeDetectionsNearestNeighborBlockV1() + return block.run( + query_predictions=query, + target_predictions=target, + query_point=query_point, + target_point=target_point, + query_keypoint_name=query_keypoint_name, + target_keypoint_name=target_keypoint_name, + max_distance=max_distance, + ) + + +@_TENSOR_ONLY +def test_single_nearest_match_tensor_native() -> None: + # given + query = make_native_detections(xyxy=[[0, 0, 10, 10]], detection_ids=["q1"]) + target = make_native_detections( + xyxy=[ + [100, 100, 110, 110], # center (105, 105) -> far + [10, 10, 20, 20], # center (15, 15) -> near + ], + detection_ids=["t1", "t2"], + ) + + # when + result = run_tensor_block(query, target) + + # then + query_out = result[OUTPUT_KEY_QUERY_PREDICTIONS] + matched_query = result[OUTPUT_KEY_MATCHED_QUERY_DETECTIONS] + matched_target = result[OUTPUT_KEY_MATCHED_TARGET_DETECTIONS] + expected_distance = math.hypot(15 - 5, 15 - 5) + assert nearest_distances(query_out)[0] == pytest.approx( + expected_distance, rel=1e-6 + ) + assert len(matched_query) == 1 + assert len(matched_target) == 1 + assert matched_query.bboxes_metadata[0]["detection_id"] == "q1" + assert matched_target.bboxes_metadata[0]["detection_id"] == "t2" + # the sliced query rows carry the enriched per-box scalar + assert matched_query.bboxes_metadata[0][NEAREST_TARGET_DISTANCE_KEY] == ( + pytest.approx(expected_distance, rel=1e-6) + ) + + +@_TENSOR_ONLY +def test_exact_tie_within_epsilon_tensor_native() -> None: + # given: two targets equidistant (well within the 1px epsilon) from the query + query = make_native_detections(xyxy=[[0, 0, 10, 10]], detection_ids=["q1"]) + target = make_native_detections( + xyxy=[ + [10, 10, 20, 20], # center (15, 15) + [10, 10.2, 20, 20.2], # center (15, 15.2) -> distance differs by < 1px + [1000, 1000, 1010, 1010], # clearly far away + ], + detection_ids=["t1", "t2", "t3"], + ) + + # when + result = run_tensor_block(query, target) + + # then: the query row is duplicated once per tied match + matched_query = result[OUTPUT_KEY_MATCHED_QUERY_DETECTIONS] + matched_target = result[OUTPUT_KEY_MATCHED_TARGET_DETECTIONS] + assert len(matched_query) == 2 + assert len(matched_target) == 2 + assert [ + box_metadata["detection_id"] for box_metadata in matched_query.bboxes_metadata + ] == ["q1", "q1"] + assert { + box_metadata["detection_id"] for box_metadata in matched_target.bboxes_metadata + } == {"t1", "t2"} + + +@_TENSOR_ONLY +def test_no_tie_clear_separation_tensor_native() -> None: + # given + query = make_native_detections(xyxy=[[0, 0, 10, 10]], detection_ids=["q1"]) + target = make_native_detections( + xyxy=[ + [10, 10, 20, 20], # near + [1000, 1000, 1010, 1010], # far + ], + detection_ids=["t1", "t2"], + ) + + # when + result = run_tensor_block(query, target) + + # then + matched_query = result[OUTPUT_KEY_MATCHED_QUERY_DETECTIONS] + matched_target = result[OUTPUT_KEY_MATCHED_TARGET_DETECTIONS] + assert len(matched_query) == 1 + assert len(matched_target) == 1 + assert matched_target.bboxes_metadata[0]["detection_id"] == "t1" + + +@_TENSOR_ONLY +def test_max_distance_excludes_matches_beyond_the_limit_tensor_native() -> None: + # given: the only target is farther than max_distance + query = make_native_detections( + xyxy=[[0, 0, 10, 10]], detection_ids=["q1"] + ) # center (5, 5) + target = make_native_detections( + xyxy=[[100, 100, 110, 110]], detection_ids=["t1"] + ) # center (105, 105) -> distance ~141.4 + + # when + result = run_tensor_block(query, target, max_distance=50) + + # then + query_out = result[OUTPUT_KEY_QUERY_PREDICTIONS] + matched_query = result[OUTPUT_KEY_MATCHED_QUERY_DETECTIONS] + matched_target = result[OUTPUT_KEY_MATCHED_TARGET_DETECTIONS] + assert nearest_distances(query_out)[0] is None + assert len(matched_query) == 0 + assert len(matched_target) == 0 + + +@_TENSOR_ONLY +def test_max_distance_retains_in_range_target_while_excluding_out_of_range_target_tensor_native() -> ( + None +): + # given: t1 is beyond max_distance and t2 is within it + query = make_native_detections( + xyxy=[[0, 0, 10, 10]], detection_ids=["q1"] + ) # center (5, 5) + target = make_native_detections( + xyxy=[ + [100, 100, 110, 110], # center (105, 105) -> distance ~141.4, excluded + [40, 5, 50, 15], # center (45, 10) -> distance ~40.3, within limit + ], + detection_ids=["t1", "t2"], + ) + + # when + result = run_tensor_block(query, target, max_distance=45) + + # then + matched_target = result[OUTPUT_KEY_MATCHED_TARGET_DETECTIONS] + assert len(matched_target) == 1 + assert matched_target.bboxes_metadata[0]["detection_id"] == "t2" + + +@_TENSOR_ONLY +def test_max_distance_boundary_is_inclusive_tensor_native() -> None: + # given: the target sits exactly on the max_distance boundary + query = make_native_detections( + xyxy=[[0, 0, 0, 0]], detection_ids=["q1"] + ) # center (0, 0) + target = make_native_detections( + xyxy=[[50, 0, 50, 0]], detection_ids=["t1"] + ) # center (50, 0) -> distance exactly 50 + + # when + result = run_tensor_block(query, target, max_distance=50) + + # then: equal-to-limit is included, not excluded + matched_target = result[OUTPUT_KEY_MATCHED_TARGET_DETECTIONS] + assert len(matched_target) == 1 + assert matched_target.bboxes_metadata[0]["detection_id"] == "t1" + + +@_TENSOR_ONLY +def test_empty_target_set_tensor_native() -> None: + # given + query = make_native_detections( + xyxy=[[0, 0, 10, 10], [20, 20, 30, 30]], detection_ids=["q1", "q2"] + ) + target = make_native_empty_detections() + + # when + result = run_tensor_block(query, target) + + # then + query_out = result[OUTPUT_KEY_QUERY_PREDICTIONS] + matched_query = result[OUTPUT_KEY_MATCHED_QUERY_DETECTIONS] + matched_target = result[OUTPUT_KEY_MATCHED_TARGET_DETECTIONS] + assert nearest_distances(query_out) == [None, None] + assert len(matched_query) == 0 + assert len(matched_target) == 0 + + +@_TENSOR_ONLY +def test_self_match_exclusion_same_set_as_query_and_target_tensor_native() -> None: + # given: the same detections passed as both query and target + detections = make_native_detections( + xyxy=[ + [0, 0, 10, 10], # center (5, 5) + [10, 10, 20, 20], # center (15, 15) -> nearest to detection 0 + [1000, 1000, 1010, 1010], # center (1005, 1005) -> far + ], + detection_ids=["a", "b", "c"], + ) + + # when + result = run_tensor_block(detections, detections) + + # then + query_out = result[OUTPUT_KEY_QUERY_PREDICTIONS] + matched_query = result[OUTPUT_KEY_MATCHED_QUERY_DETECTIONS] + matched_target = result[OUTPUT_KEY_MATCHED_TARGET_DETECTIONS] + pairs = dict( + zip( + [entry["detection_id"] for entry in matched_query.bboxes_metadata], + [entry["detection_id"] for entry in matched_target.bboxes_metadata], + ) + ) + # detection "a" must not match itself, nearest is "b"; "b" must not match + # itself either, nearest is "a" + assert pairs["a"] == "b" + assert pairs["b"] == "a" + assert nearest_distances(query_out)[0] > 0 + + +@_TENSOR_ONLY +def test_self_match_not_excluded_when_detection_id_missing_tensor_native() -> None: + # given: detections lacking `detection_id` on both sides - self-exclusion + # must be skipped gracefully (not raise), so a detection matches itself + detections = NativeDetections( + xyxy=torch.tensor( + [[0, 0, 10, 10], [1000, 1000, 1010, 1010]], dtype=torch.float32 + ), + class_id=torch.zeros(2, dtype=torch.long), + confidence=torch.full((2,), 0.9, dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "object"}}, + bboxes_metadata=[{}, {}], + ) + + # when + result = run_tensor_block(detections, detections) + + # then + query_out = result[OUTPUT_KEY_QUERY_PREDICTIONS] + assert nearest_distances(query_out)[0] == 0.0 + + +@_TENSOR_ONLY +@pytest.mark.parametrize( + "point_option,expected_point", + [ + ("CENTER", (5, 10)), + ("CENTER_LEFT", (0, 10)), + ("CENTER_RIGHT", (10, 10)), + ("TOP_CENTER", (5, 0)), + ("TOP_LEFT", (0, 0)), + ("TOP_RIGHT", (10, 0)), + ("BOTTOM_LEFT", (0, 20)), + ("BOTTOM_CENTER", (5, 20)), + ("BOTTOM_RIGHT", (10, 20)), + ], +) +def test_bbox_anchor_point_options_tensor_native(point_option, expected_point) -> None: + # given + query = make_native_detections(xyxy=[[0, 0, 10, 20]], detection_ids=["q1"]) + target = make_native_detections(xyxy=[[100, 100, 100, 100]], detection_ids=["t1"]) + + # when + result = run_tensor_block( + query, target, query_point=point_option, target_point="CENTER" + ) + + # then + ex, ey = expected_point + expected_distance = math.hypot(100 - ex, 100 - ey) + query_out = result[OUTPUT_KEY_QUERY_PREDICTIONS] + assert nearest_distances(query_out)[0] == pytest.approx( + expected_distance, rel=1e-6 + ) + + +@_TENSOR_ONLY +def test_keypoint_anchor_point_option_tensor_native() -> None: + # given + query = make_native_keypoint_prediction( + xyxy=[[0, 0, 10, 10]], + keypoints_xy=[[[3, 4], [7, 8]]], + keypoints_class_name=[["left_shoulder", "right_shoulder"]], + detection_ids=["q1"], + ) + target = make_native_detections(xyxy=[[103, 104, 103, 104]], detection_ids=["t1"]) + + # when + result = run_tensor_block( + query, + target, + query_point="KEYPOINT", + target_point="CENTER", + query_keypoint_name="left_shoulder", + ) + + # then: the keypoint-detection tuple shape survives, with the per-box + # scalar written into the bbox component's metadata + query_out = result[OUTPUT_KEY_QUERY_PREDICTIONS] + assert isinstance(query_out, tuple) + assert nearest_distances(query_out)[0] == pytest.approx( + 100.0 * math.sqrt(2), rel=1e-6 + ) + + +@_TENSOR_ONLY +def test_keypoint_missing_on_detection_excludes_it_from_matching_tensor_native() -> ( + None +): + # given: target detection's keypoints do not include the requested name + query = make_native_detections(xyxy=[[0, 0, 10, 10]], detection_ids=["q1"]) + target = make_native_keypoint_prediction( + xyxy=[[10, 10, 20, 20], [200, 200, 210, 210]], + keypoints_xy=[[[15, 15]], [[0, 0]]], + keypoints_class_name=[["nose"], ["left_eye"]], + detection_ids=["t1", "t2"], + ) + + # when + result = run_tensor_block( + query, + target, + query_point="CENTER", + target_point="KEYPOINT", + target_keypoint_name="nose", + ) + + # then: only t1 has a "nose" keypoint, t2 must be excluded from matching; + # the matched target keeps the tuple shape with the KeyPoints component + # sliced in lockstep + matched_query = result[OUTPUT_KEY_MATCHED_QUERY_DETECTIONS] + matched_target = result[OUTPUT_KEY_MATCHED_TARGET_DETECTIONS] + assert len(matched_query) == 1 + assert isinstance(matched_target, tuple) + matched_target_key_points, matched_target_detections = matched_target + assert len(matched_target_detections) == 1 + assert int(matched_target_key_points.xy.shape[0]) == 1 + assert matched_target_detections.bboxes_metadata[0]["detection_id"] == "t1" + + +@_TENSOR_ONLY +def test_query_point_keypoint_requires_keypoint_name_tensor_native() -> None: + # given + query = make_native_detections(xyxy=[[0, 0, 10, 10]], detection_ids=["q1"]) + target = make_native_detections(xyxy=[[10, 10, 20, 20]], detection_ids=["t1"]) + + # when / then + with pytest.raises( + expected_exception=ValueError, + match="query_keypoint_name", + ): + run_tensor_block(query, target, query_point="KEYPOINT") + + +@_TENSOR_ONLY +def test_keypoint_option_requires_keypoint_predictions_tensor_native() -> None: + # given: query_point is KEYPOINT but query predictions have no keypoint data + query = make_native_detections(xyxy=[[0, 0, 10, 10]], detection_ids=["q1"]) + target = make_native_detections(xyxy=[[10, 10, 20, 20]], detection_ids=["t1"]) + + # when / then + with pytest.raises(expected_exception=ValueError, match="keypoint data"): + run_tensor_block( + query, + target, + query_point="KEYPOINT", + query_keypoint_name="left_shoulder", + ) + + +@_TENSOR_ONLY +def test_instance_segmentation_input_preserves_type_and_masks_tensor_native() -> None: + # given: instance-segmentation query - the matched output must stay an + # InstanceDetections with the surviving row's mask carried over + masks = torch.zeros((2, 40, 40), dtype=torch.bool) + masks[0, :10, :10] = True + masks[1, 20:30, 20:30] = True + query = NativeInstanceDetections( + xyxy=torch.tensor([[0, 0, 10, 10], [20, 20, 30, 30]], dtype=torch.float32), + class_id=torch.zeros(2, dtype=torch.long), + confidence=torch.full((2,), 0.9, dtype=torch.float32), + mask=masks, + image_metadata={CLASS_NAMES_KEY: {0: "object"}}, + bboxes_metadata=[{"detection_id": "q1"}, {"detection_id": "q2"}], + ) + target = make_native_detections(xyxy=[[10, 10, 20, 20]], detection_ids=["t1"]) + + # when + result = run_tensor_block(query, target) + + # then + query_out = result[OUTPUT_KEY_QUERY_PREDICTIONS] + matched_query = result[OUTPUT_KEY_MATCHED_QUERY_DETECTIONS] + assert isinstance(query_out, NativeInstanceDetections) + assert isinstance(matched_query, NativeInstanceDetections) + assert len(matched_query) == 2 + assert matched_query.mask.shape == (2, 40, 40) + assert all( + box_metadata[NEAREST_TARGET_DISTANCE_KEY] is not None + for box_metadata in matched_query.bboxes_metadata + ) diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_dominant_color.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_dominant_color.py index bde7ab648f..e667fbe1d2 100644 --- a/tests/workflows/unit_tests/core_steps/classical_cv/test_dominant_color.py +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_dominant_color.py @@ -78,3 +78,202 @@ def test_dominant_color_block() -> None: 0, 0, ), " Expected rgb_color to be [255, 0, 0], aka a red image" + + +# --- tensor-native sibling --------------------------------------------------- +# Parity contract: outputs equal the numpy block's exactly under a pinned RNG. + + +def _tensor_dominant_color_imports(): + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.classical_cv.dominant_color import ( + v1_tensor, + ) + + return torch, v1_tensor + + +def _paired_images(torch, bgr: np.ndarray): + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + chw = torch.from_numpy(bgr[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=chw, + ) + return numpy_born, tensor_born + + +def _dominant_case_image(case: str) -> np.ndarray: + rng = np.random.default_rng(42) + if case == "noise": + return rng.integers(0, 256, size=(240, 320, 3), dtype=np.uint8) + if case == "dominant_red": + image = np.zeros((250, 300, 3), dtype=np.uint8) + image[:, :, 2] = 255 # BGR red + image[rng.random((250, 300)) < 0.1] = (0, 255, 0) # ~10% green speckle + return image + if case == "non_divisible": + # 317x203: strides do not divide the dims evenly + return rng.integers(0, 256, size=(317, 203, 3), dtype=np.uint8) + raise ValueError(case) + + +@pytest.mark.parametrize("seed", [7, 1234]) +@pytest.mark.parametrize("case", ["noise", "dominant_red", "non_divisible"]) +def test_tensor_dominant_color_seeded_parity(seed: int, case: str) -> None: + # given - the same pixels as numpy-born BGR and tensor-born RGB CHW + torch, v1_tensor = _tensor_dominant_color_imports() + bgr = _dominant_case_image(case) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when - identical global-RNG state before each run forces identical draws + np.random.seed(seed) + numpy_result = DominantColorBlockV1().run( + image=numpy_born, color_clusters=4, max_iterations=100, target_size=100 + ) + np.random.seed(seed) + tensor_result = v1_tensor.DominantColorBlockV1().run( + image=tensor_born, color_clusters=4, max_iterations=100, target_size=100 + ) + + # then + assert isinstance(tensor_result["rgb_color"], tuple) + assert tensor_result["rgb_color"] == numpy_result["rgb_color"] + + +def test_tensor_dominant_color_seeded_parity_with_non_default_parameters() -> None: + # given + torch, v1_tensor = _tensor_dominant_color_imports() + bgr = _dominant_case_image("non_divisible") + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + np.random.seed(97) + numpy_result = DominantColorBlockV1().run( + image=numpy_born, color_clusters=6, max_iterations=30, target_size=64 + ) + np.random.seed(97) + tensor_result = v1_tensor.DominantColorBlockV1().run( + image=tensor_born, color_clusters=6, max_iterations=30, target_size=64 + ) + + # then + assert tensor_result["rgb_color"] == numpy_result["rgb_color"] + + +@pytest.mark.parametrize("target_size", [250, 100, 50]) +def test_tensor_dominant_color_downsample_matches_v1_slicing(target_size: int) -> None: + # given - 317x203 exercises strided slicing with a remainder + torch, v1_tensor = _tensor_dominant_color_imports() + bgr = _dominant_case_image("non_divisible") + numpy_born, tensor_born = _paired_images(torch, bgr) + height, width = bgr.shape[:2] + scale_factor = max(1, min(width, height) // target_size) + + # when + device_downsampled = v1_tensor._downsample_to_bgr_numpy( + chw=tensor_born.tensor_image, scale_factor=scale_factor + ) + + # then + assert scale_factor == {250: 1, 100: 2, 50: 4}[target_size] + assert device_downsampled.dtype == np.uint8 + assert np.array_equal( + device_downsampled, numpy_born.numpy_image[::scale_factor, ::scale_factor] + ) + + +def test_tensor_dominant_color_numpy_born_delegates_without_materialisation() -> None: + # given + torch, v1_tensor = _tensor_dominant_color_imports() + bgr = _dominant_case_image("noise") + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + np.random.seed(3) + reference = DominantColorBlockV1().run( + image=WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ), + color_clusters=4, + max_iterations=100, + target_size=100, + ) + + # when + np.random.seed(3) + result = v1_tensor.DominantColorBlockV1().run( + image=numpy_born, color_clusters=4, max_iterations=100, target_size=100 + ) + + # then + assert result["rgb_color"] == reference["rgb_color"] + assert not numpy_born.is_tensor_materialised(), "delegate must not materialise" + + +def test_tensor_dominant_color_obvious_dominant_color_on_tensor_path() -> None: + # given - ~90% red / ~10% green speckle, tensor-born + torch, v1_tensor = _tensor_dominant_color_imports() + bgr = _dominant_case_image("dominant_red") + _, tensor_born = _paired_images(torch, bgr) + + # when + np.random.seed(11) + result = v1_tensor.DominantColorBlockV1().run( + image=tensor_born, color_clusters=4, max_iterations=100, target_size=100 + ) + + # then - bounds, not equality: the green cluster split is seed-dependent + r, g, b = result["rgb_color"] + assert r >= 200 + assert g <= 60 + assert b <= 60 + + +def test_tensor_dominant_color_insufficient_pixels_raises_on_both_paths() -> None: + # given - 4 pixels < 10 clusters, so np.random.choice(replace=False) raises + torch, v1_tensor = _tensor_dominant_color_imports() + bgr = np.full((2, 2, 3), 128, dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when / then + with pytest.raises(ValueError): + DominantColorBlockV1().run( + image=numpy_born, color_clusters=10, max_iterations=100, target_size=100 + ) + with pytest.raises(ValueError): + v1_tensor.DominantColorBlockV1().run( + image=tensor_born, color_clusters=10, max_iterations=100, target_size=100 + ) + + +def test_tensor_dominant_color_on_mps_device(monkeypatch) -> None: + # given - patch the global device so tensor-born images materialise on MPS + torch, v1_tensor = _tensor_dominant_color_imports() + if not torch.backends.mps.is_available(): + pytest.skip("MPS device not available") + import inference.core.workflows.execution_engine.entities.base as base_module + + bgr = _dominant_case_image("non_divisible") + numpy_born, _ = _paired_images(torch, bgr) + np.random.seed(21) + reference = DominantColorBlockV1().run( + image=numpy_born, color_clusters=4, max_iterations=100, target_size=100 + ) + monkeypatch.setattr(base_module, "WORKFLOWS_IMAGE_TENSOR_DEVICE", "mps") + _, tensor_born = _paired_images(torch, bgr) + assert tensor_born.tensor_image.device.type == "mps" + + # when + np.random.seed(21) + result = v1_tensor.DominantColorBlockV1().run( + image=tensor_born, color_clusters=4, max_iterations=100, target_size=100 + ) + + # then + assert result["rgb_color"] == reference["rgb_color"] diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_grayscale.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_grayscale.py index 69df07085f..d3d8f5855c 100644 --- a/tests/workflows/unit_tests/core_steps/classical_cv/test_grayscale.py +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_grayscale.py @@ -68,3 +68,151 @@ def test_convert_grayscale_block() -> None: assert output.get("image").numpy_image.shape == (1000, 1000) # check if the image is modified assert not np.array_equal(output.get("image").numpy_image, start_image) + + +# --- tensor-native sibling --------------------------------------------------- +# Parity contract: outputs match the numpy block bit-exactly on every path. + + +def _tensor_grayscale_imports(): + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.classical_cv.convert_grayscale.v1_tensor import ( + ConvertGrayscaleBlockV1 as TensorConvertGrayscaleBlockV1, + ) + + return torch, TensorConvertGrayscaleBlockV1 + + +def _paired_images(torch, bgr: np.ndarray): + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + if bgr.ndim == 2: + chw = torch.from_numpy(bgr.copy()).unsqueeze(0) + else: + chw = torch.from_numpy(bgr[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=chw, + ) + return numpy_born, tensor_born + + +@pytest.mark.parametrize( + "height,width", [(24, 32), (17, 23), (1, 1), (3, 641), (127, 129)] +) +def test_tensor_convert_grayscale_bit_exact_parity(height, width) -> None: + # given - the same random pixels as numpy-born BGR and tensor-born RGB CHW + torch, TensorConvertGrayscaleBlockV1 = _tensor_grayscale_imports() + rng = np.random.default_rng(height * 1000 + width) + bgr = rng.integers(0, 256, size=(height, width, 3), dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = ConvertGrayscaleBlockV1().run(image=numpy_born)["image"].numpy_image + tensor_result_image = TensorConvertGrayscaleBlockV1().run(image=tensor_born)[ + "image" + ] + + # then + assert tensor_result_image.is_tensor_materialised() + assert tuple(tensor_result_image.tensor_image.shape) == (1, height, width) + assert tensor_result_image.numpy_image.shape == (height, width) + assert np.array_equal(tensor_result_image.numpy_image, numpy_result) + + +def test_tensor_convert_grayscale_rounding_boundary_pixels() -> None: + # given - triples whose luminance lands exactly on a .5 rounding tie, plus extremes + torch, TensorConvertGrayscaleBlockV1 = _tensor_grayscale_imports() + boundary_bgr_triples = [ + (4, 12, 0), + (36, 44, 32), + (60, 52, 64), + (147, 251, 95), + (92, 20, 128), + (116, 28, 160), + (203, 227, 191), + (227, 235, 223), + (255, 0, 0), + (0, 255, 0), + (0, 0, 255), + (255, 255, 255), + (0, 0, 0), + ] + expected_gray = [8, 40, 57, 193, 61, 78, 214, 231, 29, 150, 76, 255, 0] + bgr = np.array([boundary_bgr_triples], dtype=np.uint8) # (1, 13, 3) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = ConvertGrayscaleBlockV1().run(image=numpy_born)["image"].numpy_image + tensor_result = ( + TensorConvertGrayscaleBlockV1().run(image=tensor_born)["image"].numpy_image + ) + + # then - cv2 rounds half up on these ties + assert np.array_equal(tensor_result, numpy_result) + assert np.array_equal(tensor_result, np.array([expected_gray], dtype=np.uint8)) + + +def test_tensor_convert_grayscale_delegates_for_numpy_born_images() -> None: + # given + torch, TensorConvertGrayscaleBlockV1 = _tensor_grayscale_imports() + rng = np.random.default_rng(3) + bgr = rng.integers(0, 256, size=(19, 27, 3), dtype=np.uint8) + numpy_born, _ = _paired_images(torch, bgr) + reference = ( + ConvertGrayscaleBlockV1() + .run( + image=WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + )["image"] + .numpy_image + ) + + # when + result = TensorConvertGrayscaleBlockV1().run(image=numpy_born)["image"] + + # then + assert np.array_equal(result.numpy_image, reference) + assert not numpy_born.is_tensor_materialised(), "delegate must not materialise" + + +def test_tensor_convert_grayscale_raises_on_single_channel_input_like_v1() -> None: + # given - cv2.cvtColor rejects 2-D input + import cv2 + + torch, TensorConvertGrayscaleBlockV1 = _tensor_grayscale_imports() + gray = np.full((8, 8), 77, dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, gray) + + # when / then + with pytest.raises(cv2.error): + ConvertGrayscaleBlockV1().run(image=numpy_born) + with pytest.raises(cv2.error): + TensorConvertGrayscaleBlockV1().run(image=tensor_born) + + +def test_tensor_convert_grayscale_on_mps_device(monkeypatch) -> None: + # given - patch the global device so tensor-born images materialise on MPS + torch, TensorConvertGrayscaleBlockV1 = _tensor_grayscale_imports() + if not torch.backends.mps.is_available(): + pytest.skip("MPS device not available") + import inference.core.workflows.execution_engine.entities.base as base_module + + rng = np.random.default_rng(9) + bgr = rng.integers(0, 256, size=(33, 47, 3), dtype=np.uint8) + numpy_born, _ = _paired_images(torch, bgr) + reference = ConvertGrayscaleBlockV1().run(image=numpy_born)["image"].numpy_image + monkeypatch.setattr(base_module, "WORKFLOWS_IMAGE_TENSOR_DEVICE", "mps") + _, tensor_born = _paired_images(torch, bgr) + assert tensor_born.tensor_image.device.type == "mps" + + # when + result = TensorConvertGrayscaleBlockV1().run(image=tensor_born)["image"] + + # then + assert result.tensor_image.device.type == "mps" + assert np.array_equal(result.numpy_image, reference) diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_image_preprocessing.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_image_preprocessing.py new file mode 100644 index 0000000000..760d8437e4 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_image_preprocessing.py @@ -0,0 +1,433 @@ +import numpy as np +import pytest +from pydantic import ValidationError + +from inference.core.workflows.core_steps.classical_cv.image_preprocessing.v1 import ( + ImagePreprocessingBlockV1, + ImagePreprocessingManifest, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) + + +class TestImagePreprocessingManifest: + def test_image_preprocessing_validation_when_valid_manifest_is_given(self): + manifest = ImagePreprocessingManifest.model_validate( + { + "type": "roboflow_core/image_preprocessing@v1", + "name": "image_preprocessing", + "image": "$inputs.image", + "task_type": "resize", + "width": 320, + "height": 240, + } + ) + + assert manifest.type == "roboflow_core/image_preprocessing@v1" + assert manifest.name == "image_preprocessing" + assert manifest.task_type == "resize" + assert manifest.width == 320 + assert manifest.height == 240 + + def test_image_preprocessing_validation_when_image_is_missing(self): + with pytest.raises(ValidationError): + ImagePreprocessingManifest.model_validate( + { + "type": "roboflow_core/image_preprocessing@v1", + "name": "image_preprocessing", + "task_type": "flip", + } + ) + + def test_image_preprocessing_manifest_outputs(self): + outputs = ImagePreprocessingManifest.describe_outputs() + + assert len(outputs) == 1 + assert outputs[0].name == "image" + + +def _numpy_born_image(bgr: np.ndarray) -> WorkflowImageData: + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + + +def _run_v1(image: WorkflowImageData, **overrides) -> WorkflowImageData: + kwargs = { + "task_type": "flip", + "width": None, + "height": None, + "rotation_degrees": None, + "flip_type": None, + } + kwargs.update(overrides) + return ImagePreprocessingBlockV1().run(image=image, **kwargs)["image"] + + +def test_image_preprocessing_resize_to_exact_dimensions() -> None: + # given + rng = np.random.default_rng(3) + image = _numpy_born_image(rng.integers(0, 256, size=(48, 64, 3), dtype=np.uint8)) + + # when + result = _run_v1(image, task_type="resize", width=32, height=24) + + # then + assert result.numpy_image.shape == (24, 32, 3) + assert result.numpy_image.dtype == np.uint8 + + +def test_image_preprocessing_rotate_by_right_angle() -> None: + # given + rng = np.random.default_rng(4) + image = _numpy_born_image(rng.integers(0, 256, size=(48, 64, 3), dtype=np.uint8)) + + # when + result = _run_v1(image, task_type="rotate", rotation_degrees=90) + + # then - canvas swaps dimensions for a 90 degrees rotation + assert result.numpy_image.shape == (64, 48, 3) + + +def test_image_preprocessing_flip_vertical() -> None: + # given + rng = np.random.default_rng(5) + bgr = rng.integers(0, 256, size=(48, 64, 3), dtype=np.uint8) + image = _numpy_born_image(bgr) + + # when + result = _run_v1(image, task_type="flip", flip_type="vertical") + + # then + assert np.array_equal(result.numpy_image, bgr[::-1]) + + +def test_image_preprocessing_invalid_task_type_raises() -> None: + # given + image = _numpy_born_image(np.zeros((10, 10, 3), dtype=np.uint8)) + + # when / then + with pytest.raises(ValueError): + _run_v1(image, task_type="unknown") + + +# --- tensor-native sibling --------------------------------------------------- +# Parity contract: outputs match the numpy block bit-exactly on every path. + + +def _tensor_preprocessing_imports(): + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.classical_cv.image_preprocessing.v1_tensor import ( + ImagePreprocessingBlockV1 as TensorImagePreprocessingBlockV1, + ) + + return torch, TensorImagePreprocessingBlockV1 + + +def _paired_images(torch, bgr: np.ndarray): + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + if bgr.ndim == 2: + chw = torch.from_numpy(bgr.copy()).unsqueeze(0) + else: + chw = torch.from_numpy(bgr[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=chw, + ) + return numpy_born, tensor_born + + +def _run_tensor( + block_class, image: WorkflowImageData, **overrides +) -> WorkflowImageData: + kwargs = { + "task_type": "flip", + "width": None, + "height": None, + "rotation_degrees": None, + "flip_type": None, + } + kwargs.update(overrides) + return block_class().run(image=image, **kwargs)["image"] + + +def _case_image(case: str) -> np.ndarray: + rng = np.random.default_rng(42) + if case == "color_odd": + return rng.integers(0, 256, size=(7, 9, 3), dtype=np.uint8) + if case == "gray_odd": + return rng.integers(0, 256, size=(7, 9), dtype=np.uint8) + if case == "color_even": + return rng.integers(0, 256, size=(8, 8, 3), dtype=np.uint8) + if case == "color_mixed": + return rng.integers(0, 256, size=(8, 7, 3), dtype=np.uint8) + if case == "color_larger": + return rng.integers(0, 256, size=(48, 64, 3), dtype=np.uint8) + if case == "checker": + # 0/1 checkerboard: every 2-tap half-pixel sum is an exact rounding tie + checker = ((np.indices((7, 9)).sum(axis=0)) % 2).astype(np.uint8) + return np.stack([checker, 1 - checker, checker], axis=-1) + raise ValueError(case) + + +@pytest.mark.parametrize("flip_type", ["vertical", "horizontal", "both"]) +@pytest.mark.parametrize("case", ["color_odd", "gray_odd"]) +def test_tensor_flip_bit_exact_parity(flip_type, case) -> None: + # given - the same pixels as numpy-born BGR and tensor-born RGB CHW + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + bgr = _case_image(case) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = _run_v1(numpy_born, task_type="flip", flip_type=flip_type) + tensor_result = _run_tensor( + TensorImagePreprocessingBlockV1, + tensor_born, + task_type="flip", + flip_type=flip_type, + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +def test_tensor_flip_unrecognized_type_passes_image_through() -> None: + # given - flip_type=None is the only run()-reachable pass-through value in v1 + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + bgr = _case_image("color_odd") + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = _run_v1(numpy_born, task_type="flip", flip_type=None) + tensor_result = _run_tensor( + TensorImagePreprocessingBlockV1, tensor_born, task_type="flip", flip_type=None + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(numpy_result.numpy_image, bgr) + assert np.array_equal(tensor_result.numpy_image, bgr) + + +@pytest.mark.parametrize("rotation_degrees", [90, 180, 270, -90, -180, -270, 360, -360]) +@pytest.mark.parametrize("case", ["color_odd", "color_even", "color_mixed", "checker"]) +def test_tensor_rotate_right_angles_bit_exact_parity(rotation_degrees, case) -> None: + # given - the warpAffine mapping differs per axis parity, hence odd/even/mixed dims + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + bgr = _case_image(case) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = _run_v1( + numpy_born, task_type="rotate", rotation_degrees=rotation_degrees + ) + tensor_result = _run_tensor( + TensorImagePreprocessingBlockV1, + tensor_born, + task_type="rotate", + rotation_degrees=rotation_degrees, + ) + + # then + assert tensor_result.is_tensor_materialised() + assert tensor_result.numpy_image.shape == numpy_result.numpy_image.shape + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +@pytest.mark.parametrize("rotation_degrees", [90, -90, 180, 360]) +def test_tensor_rotate_right_angles_grayscale_parity(rotation_degrees) -> None: + # given + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + gray = _case_image("gray_odd") + numpy_born, tensor_born = _paired_images(torch, gray) + + # when + numpy_result = _run_v1( + numpy_born, task_type="rotate", rotation_degrees=rotation_degrees + ) + tensor_result = _run_tensor( + TensorImagePreprocessingBlockV1, + tensor_born, + task_type="rotate", + rotation_degrees=rotation_degrees, + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +def test_tensor_rotate_arbitrary_angle_delegates_to_numpy_math() -> None: + # given - arbitrary angles are a real bilinear warp, so the block delegates + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + bgr = _case_image("color_larger") + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = _run_v1(numpy_born, task_type="rotate", rotation_degrees=33) + tensor_result = _run_tensor( + TensorImagePreprocessingBlockV1, + tensor_born, + task_type="rotate", + rotation_degrees=33, + ) + + # then + assert not tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +def test_tensor_rotate_zero_degrees_returns_tensor_copy() -> None: + # given - v1 early-returns a copy for rotation_degrees=0 + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + bgr = _case_image("color_odd") + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = _run_v1(numpy_born, task_type="rotate", rotation_degrees=0) + tensor_result = _run_tensor( + TensorImagePreprocessingBlockV1, + tensor_born, + task_type="rotate", + rotation_degrees=0, + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + assert np.array_equal(tensor_result.numpy_image, bgr) + + +@pytest.mark.parametrize( + "width,height", + [ + (32, 24), # downscale to exact dims + (96, 128), # upscale to exact dims + (30, None), # height derived from aspect ratio: int(48 * 30 / 64) = 22 + (None, 25), # width derived from aspect ratio: int(64 * 25 / 48) = 33 + ], +) +def test_tensor_resize_delegates_with_identical_output(width, height) -> None: + # given - INTER_AREA is not bit-exactly replicable in torch, so resizes delegate + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + bgr = _case_image("color_larger") + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = _run_v1(numpy_born, task_type="resize", width=width, height=height) + tensor_result = _run_tensor( + TensorImagePreprocessingBlockV1, + tensor_born, + task_type="resize", + width=width, + height=height, + ) + + # then + assert tensor_result.numpy_image.shape == numpy_result.numpy_image.shape + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + + +def test_tensor_resize_without_target_dims_returns_tensor_copy() -> None: + # given - v1 early-returns a copy for width=height=None + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + bgr = _case_image("color_odd") + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_result = _run_v1(numpy_born, task_type="resize", width=None, height=None) + tensor_result = _run_tensor( + TensorImagePreprocessingBlockV1, + tensor_born, + task_type="resize", + width=None, + height=None, + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result.numpy_image) + assert np.array_equal(tensor_result.numpy_image, bgr) + + +@pytest.mark.parametrize( + "overrides", + [ + {"task_type": "resize", "width": 0, "height": 24}, + {"task_type": "resize", "width": -5, "height": 24}, + {"task_type": "resize", "width": 32, "height": 0}, + {"task_type": "rotate", "rotation_degrees": 361}, + {"task_type": "rotate", "rotation_degrees": -361}, + {"task_type": "flip", "flip_type": "diagonal"}, + {"task_type": "unknown"}, + ], +) +def test_tensor_validation_errors_match_v1(overrides) -> None: + # given + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + bgr = _case_image("color_odd") + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when / then + with pytest.raises(ValueError): + _run_v1(numpy_born, **overrides) + with pytest.raises(ValueError): + _run_tensor(TensorImagePreprocessingBlockV1, tensor_born, **overrides) + + +def test_tensor_block_delegates_for_numpy_born_images() -> None: + # given + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + bgr = _case_image("color_odd") + numpy_born, _ = _paired_images(torch, bgr) + reference = _run_v1( + _numpy_born_image(bgr), task_type="flip", flip_type="vertical" + ).numpy_image + + # when + result = _run_tensor( + TensorImagePreprocessingBlockV1, + numpy_born, + task_type="flip", + flip_type="vertical", + ) + + # then + assert np.array_equal(result.numpy_image, reference) + assert not numpy_born.is_tensor_materialised(), "delegate must not materialise" + + +def test_tensor_flip_on_mps_device(monkeypatch) -> None: + # given - patch the global device so tensor-born images materialise on MPS + torch, TensorImagePreprocessingBlockV1 = _tensor_preprocessing_imports() + if not torch.backends.mps.is_available(): + pytest.skip("MPS device not available") + import inference.core.workflows.execution_engine.entities.base as base_module + + bgr = _case_image("color_odd") + numpy_born, _ = _paired_images(torch, bgr) + reference = _run_v1( + numpy_born, task_type="flip", flip_type="horizontal" + ).numpy_image + monkeypatch.setattr(base_module, "WORKFLOWS_IMAGE_TENSOR_DEVICE", "mps") + _, tensor_born = _paired_images(torch, bgr) + assert tensor_born.tensor_image.device.type == "mps" + + # when + result = _run_tensor( + TensorImagePreprocessingBlockV1, + tensor_born, + task_type="flip", + flip_type="horizontal", + ) + + # then + assert result.tensor_image.device.type == "mps" + assert np.array_equal(result.numpy_image, reference) diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_pixel_color_count.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_pixel_color_count.py index 749c03bfea..5388e8c023 100644 --- a/tests/workflows/unit_tests/core_steps/classical_cv/test_pixel_color_count.py +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_pixel_color_count.py @@ -129,3 +129,257 @@ def test_convert_color_to_bgr_tuple_when_invalid_value() -> None: # when with pytest.raises(ValueError): _ = convert_color_to_bgr_tuple(color="invalid") + + +# --- tensor-native sibling --------------------------------------------------- +# Parity contract: exact count parity with v1; the frame never crosses device->host. + + +def _tensor_pixel_count_imports(): + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.classical_cv.pixel_color_count.v1_tensor import ( + PixelationCountBlockV1 as TensorPixelationCountBlockV1, + ) + + return torch, TensorPixelationCountBlockV1 + + +def _paired_images(torch, bgr: np.ndarray): + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + if bgr.ndim == 2: + chw = torch.from_numpy(bgr.copy()).unsqueeze(0) + else: + chw = torch.from_numpy(bgr[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=chw, + ) + return numpy_born, tensor_born + + +def _image_with_planted_target() -> np.ndarray: + # random background, a planted block of target RGB (68, 17, 34), +-10 offset pixels + rng = np.random.default_rng(42) + bgr = rng.integers(0, 256, size=(24, 32, 3), dtype=np.uint8) + bgr[0:4, 0:5] = (34, 17, 68) + bgr[10, 10] = (24, 7, 58) + bgr[10, 11] = (44, 27, 78) + return bgr + + +@pytest.mark.parametrize( + "target_color", + ["#441122", "#412", "(68, 17, 34)", (68, 17, 34)], + ids=["hex_6_digit", "hex_3_digit", "tuple_string", "rgb_tuple"], +) +@pytest.mark.parametrize("tolerance", [0, 10, 64]) +def test_tensor_pixel_color_count_exact_parity_across_formats( + target_color, tolerance +) -> None: + # given - every target format spells the same colour, RGB (68, 17, 34) + torch, TensorPixelationCountBlockV1 = _tensor_pixel_count_imports() + numpy_born, tensor_born = _paired_images(torch, _image_with_planted_target()) + + # when + numpy_count = PixelationCountBlockV1().run( + image=numpy_born, target_color=target_color, tolerance=tolerance + )["matching_pixels_count"] + tensor_count = TensorPixelationCountBlockV1().run( + image=tensor_born, target_color=target_color, tolerance=tolerance + )["matching_pixels_count"] + + # then + assert isinstance(tensor_count, int) + assert tensor_count == numpy_count + assert tensor_count >= 20, "planted 4x5 block must match at any tolerance" + assert tensor_born.is_tensor_materialised() + assert tensor_born._numpy_image is None, "tensor path must not materialise numpy" + + +def test_tensor_pixel_color_count_boundary_inclusivity() -> None: + # given - cv2.inRange bounds are inclusive: pixels exactly at a bound match + torch, TensorPixelationCountBlockV1 = _tensor_pixel_count_imports() + bgr = np.zeros((4, 4, 3), dtype=np.uint8) + bgr[0, 0] = (190, 140, 90) # exactly at lower bound -> match + bgr[0, 1] = (210, 160, 110) # exactly at upper bound -> match + bgr[0, 2] = (200, 150, 100) # the target itself -> match + bgr[0, 3] = (189, 140, 90) # one below lower on one channel -> no match + bgr[1, 0] = (211, 160, 110) # one above upper on one channel -> no match + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_count = PixelationCountBlockV1().run( + image=numpy_born, target_color=(100, 150, 200), tolerance=10 + )["matching_pixels_count"] + tensor_count = TensorPixelationCountBlockV1().run( + image=tensor_born, target_color=(100, 150, 200), tolerance=10 + )["matching_pixels_count"] + + # then + assert tensor_count == numpy_count == 3 + + +@pytest.mark.parametrize( + "target_rgb, planted_matching, planted_non_matching", + [ + # bounds are not clipped to uint8: lower bound -10 acts as "no lower limit" + ((0, 0, 0), [(0, 0, 0), (10, 10, 10)], [(11, 11, 11)]), + # upper bound 265 acts as "no upper limit" + ((255, 255, 255), [(255, 255, 255), (245, 245, 245)], [(244, 244, 244)]), + ], + ids=["lower_bound_below_zero", "upper_bound_above_255"], +) +def test_tensor_pixel_color_count_bounds_beyond_uint8_range( + target_rgb, planted_matching, planted_non_matching +) -> None: + # given - a mid-grey background that never matches, plus planted pixels + torch, TensorPixelationCountBlockV1 = _tensor_pixel_count_imports() + bgr = np.full((6, 6, 3), 128, dtype=np.uint8) + planted = planted_matching + planted_non_matching + for index, pixel in enumerate(planted): + bgr[0, index] = pixel + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_count = PixelationCountBlockV1().run( + image=numpy_born, target_color=target_rgb, tolerance=10 + )["matching_pixels_count"] + tensor_count = TensorPixelationCountBlockV1().run( + image=tensor_born, target_color=target_rgb, tolerance=10 + )["matching_pixels_count"] + + # then + assert tensor_count == numpy_count == len(planted_matching) + + +def test_tensor_pixel_color_count_tolerance_255_matches_everything() -> None: + # given - tolerance 255 makes every per-channel range cover all of [0, 255] + torch, TensorPixelationCountBlockV1 = _tensor_pixel_count_imports() + rng = np.random.default_rng(11) + bgr = rng.integers(0, 256, size=(17, 23, 3), dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_count = PixelationCountBlockV1().run( + image=numpy_born, target_color=(200, 3, 254), tolerance=255 + )["matching_pixels_count"] + tensor_count = TensorPixelationCountBlockV1().run( + image=tensor_born, target_color=(200, 3, 254), tolerance=255 + )["matching_pixels_count"] + + # then + assert tensor_count == numpy_count == 17 * 23 + + +def test_tensor_pixel_color_count_zero_matches() -> None: + # given - pixels capped at 200, so the target range [245, 265] never matches + torch, TensorPixelationCountBlockV1 = _tensor_pixel_count_imports() + rng = np.random.default_rng(3) + bgr = rng.integers(0, 200, size=(16, 16, 3), dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when + numpy_count = PixelationCountBlockV1().run( + image=numpy_born, target_color=(255, 0, 255), tolerance=10 + )["matching_pixels_count"] + tensor_count = TensorPixelationCountBlockV1().run( + image=tensor_born, target_color=(255, 0, 255), tolerance=10 + )["matching_pixels_count"] + + # then + assert tensor_count == numpy_count == 0 + + +def test_tensor_pixel_color_count_delegates_for_numpy_born_images() -> None: + # given + torch, TensorPixelationCountBlockV1 = _tensor_pixel_count_imports() + bgr = _image_with_planted_target() + numpy_born, _ = _paired_images(torch, bgr) + reference = PixelationCountBlockV1().run( + image=WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ), + target_color="#441122", + tolerance=10, + )["matching_pixels_count"] + + # when + result = TensorPixelationCountBlockV1().run( + image=numpy_born, target_color="#441122", tolerance=10 + )["matching_pixels_count"] + + # then + assert result == reference + assert not numpy_born.is_tensor_materialised(), "delegate must not materialise" + + +@pytest.mark.parametrize( + "invalid_color", + ["invalid", "#zzz", "(1, 2, a)", (255, 0, 0, 0)], + ids=["plain_string", "bad_hex", "bad_tuple_string", "four_element_tuple"], +) +def test_tensor_pixel_color_count_invalid_color_raises(invalid_color) -> None: + # given + torch, TensorPixelationCountBlockV1 = _tensor_pixel_count_imports() + bgr = np.zeros((4, 4, 3), dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when / then + with pytest.raises(ValueError): + PixelationCountBlockV1().run( + image=numpy_born, target_color=invalid_color, tolerance=10 + ) + with pytest.raises(ValueError): + TensorPixelationCountBlockV1().run( + image=tensor_born, target_color=invalid_color, tolerance=10 + ) + assert tensor_born._numpy_image is None, "failed run must not materialise numpy" + + +def test_tensor_pixel_color_count_grayscale_matches_v1_error() -> None: + # given - 3-element inRange bounds against a 1-channel image raise cv2.error + import cv2 + + torch, TensorPixelationCountBlockV1 = _tensor_pixel_count_imports() + gray = np.full((8, 9), 128, dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, gray) + + # when / then + with pytest.raises(cv2.error): + PixelationCountBlockV1().run( + image=numpy_born, target_color=(128, 128, 128), tolerance=10 + ) + with pytest.raises(cv2.error): + TensorPixelationCountBlockV1().run( + image=tensor_born, target_color=(128, 128, 128), tolerance=10 + ) + + +def test_tensor_pixel_color_count_on_mps_device(monkeypatch) -> None: + # given - patch the global device so tensor-born images materialise on MPS + torch, TensorPixelationCountBlockV1 = _tensor_pixel_count_imports() + if not torch.backends.mps.is_available(): + pytest.skip("MPS device not available") + import inference.core.workflows.execution_engine.entities.base as base_module + + bgr = _image_with_planted_target() + numpy_born, _ = _paired_images(torch, bgr) + reference = PixelationCountBlockV1().run( + image=numpy_born, target_color=(68, 17, 34), tolerance=10 + )["matching_pixels_count"] + monkeypatch.setattr(base_module, "WORKFLOWS_IMAGE_TENSOR_DEVICE", "mps") + _, tensor_born = _paired_images(torch, bgr) + assert tensor_born.tensor_image.device.type == "mps" + + # when + result = TensorPixelationCountBlockV1().run( + image=tensor_born, target_color=(68, 17, 34), tolerance=10 + )["matching_pixels_count"] + + # then + assert result == reference + assert tensor_born._numpy_image is None, "tensor path must not materialise numpy" diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_size_measurement.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_size_measurement.py index c29af73918..3341b885d7 100644 --- a/tests/workflows/unit_tests/core_steps/classical_cv/test_size_measurement.py +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_size_measurement.py @@ -94,3 +94,47 @@ def test_size_measurement_block_with_invalid_reference_dimensions(): match="reference_dimensions must be a string in the format 'width,height'", ): block.run(reference_predictions, object_predictions, reference_dimensions) + + +def test_size_measurement_tensor_block_skips_objects_with_empty_masks(): + # regression for the None-dimension guard ported from the numpy source + # (c03c2514b): an InstanceDetections row whose mask decodes to no contour + # yields (None, None) and must produce a None entry, not a TypeError. + pytest.importorskip("inference_models") + import torch + + from inference.core.workflows.core_steps.classical_cv.size_measurement.v1_tensor import ( + OUTPUT_KEY as TENSOR_OUTPUT_KEY, + ) + from inference.core.workflows.core_steps.classical_cv.size_measurement.v1_tensor import ( + SizeMeasurementBlockV1 as TensorSizeMeasurementBlockV1, + ) + from inference_models.models.base.instance_segmentation import InstanceDetections + from inference_models.models.base.object_detection import Detections + + # given + reference_predictions = Detections( + xyxy=torch.tensor([[10.0, 10.0, 50.0, 50.0]]), + class_id=torch.tensor([0]), + confidence=torch.tensor([0.9]), + ) + valid_mask = torch.zeros((100, 100), dtype=torch.bool) + valid_mask[10:60, 10:60] = True + empty_mask = torch.zeros((100, 100), dtype=torch.bool) + object_predictions = InstanceDetections( + xyxy=torch.tensor([[10.0, 10.0, 60.0, 60.0], [30.0, 30.0, 70.0, 70.0]]), + class_id=torch.tensor([0, 1]), + confidence=torch.tensor([0.8, 0.85]), + mask=torch.stack([valid_mask, empty_mask]), + ) + + # when + block = TensorSizeMeasurementBlockV1() + result = block.run(reference_predictions, object_predictions, "5.0,5.0") + + # then + dimensions = result[TENSOR_OUTPUT_KEY] + assert len(dimensions) == 2 + assert dimensions[0] is not None + assert set(dimensions[0].keys()) == {"width", "height", "longer", "shorter"} + assert dimensions[1] is None diff --git a/tests/workflows/unit_tests/core_steps/classical_cv/test_threshold.py b/tests/workflows/unit_tests/core_steps/classical_cv/test_threshold.py index bbd4adca89..088fd84cb1 100644 --- a/tests/workflows/unit_tests/core_steps/classical_cv/test_threshold.py +++ b/tests/workflows/unit_tests/core_steps/classical_cv/test_threshold.py @@ -78,3 +78,315 @@ def test_threshold_block(dogs_image: np.ndarray) -> None: assert output.get("image").numpy_image.shape == dogs_image.shape # check if the image is modified assert not np.array_equal(output.get("image").numpy_image, dogs_image) + + +# --- tensor-native sibling --------------------------------------------------- +# Parity contract: outputs match the numpy block bit-exactly on every path. + + +def _tensor_threshold_imports(): + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.classical_cv.threshold.v1_tensor import ( + ImageThresholdBlockV1 as TensorImageThresholdBlockV1, + ) + + return torch, TensorImageThresholdBlockV1 + + +def _paired_images(torch, bgr: np.ndarray): + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=bgr, + ) + if bgr.ndim == 2: + chw = torch.from_numpy(bgr.copy()).unsqueeze(0) + else: + chw = torch.from_numpy(bgr[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=chw, + ) + return numpy_born, tensor_born + + +def _run_threshold_pair( + block_cls, image: WorkflowImageData, threshold_type: str, thresh_value, max_value +): + return block_cls().run( + image=image, + threshold_type=threshold_type, + thresh_value=thresh_value, + max_value=max_value, + )["image"] + + +FIXED_THRESHOLD_TYPES = ["binary", "binary_inv", "trunc", "tozero", "tozero_inv"] + + +@pytest.mark.parametrize("threshold_type", FIXED_THRESHOLD_TYPES) +@pytest.mark.parametrize("thresh_value", [0, 127, 254, 255]) +@pytest.mark.parametrize("layout", ["grayscale", "color"]) +def test_tensor_threshold_fixed_types_bit_exact_parity( + threshold_type, thresh_value, layout +) -> None: + # given - same pixels numpy-born (BGR) and tensor-born (RGB CHW); half the + # pixels sit in {t-1, t, t+1} to exercise cv2's strict `>` comparison + torch, TensorImageThresholdBlockV1 = _tensor_threshold_imports() + rng = np.random.default_rng(31) + shape = (24, 32) if layout == "grayscale" else (24, 32, 3) + noise = rng.integers(0, 256, size=shape) + band = rng.integers( + max(0, thresh_value - 1), min(255, thresh_value + 1) + 1, size=shape + ) + pixels = np.where(rng.random(shape) < 0.5, band, noise).astype(np.uint8) + numpy_born, tensor_born = _paired_images(torch, pixels) + + # when + numpy_result = _run_threshold_pair( + ImageThresholdBlockV1, numpy_born, threshold_type, thresh_value, 255 + ).numpy_image + tensor_result = _run_threshold_pair( + TensorImageThresholdBlockV1, tensor_born, threshold_type, thresh_value, 255 + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result) + + +@pytest.mark.parametrize("threshold_type", FIXED_THRESHOLD_TYPES) +@pytest.mark.parametrize( + "thresh_value,max_value", + [ + (127, 300), # maxval saturates to 255 + (127, 254.5), # cvRound is half-to-even -> 254 + (254, 300), + (300, 255), # degenerate: threshold above the uint8 range + (-5, 255), # degenerate: negative threshold + ], +) +def test_tensor_threshold_fixed_types_non_standard_params_parity( + threshold_type, thresh_value, max_value +) -> None: + # given - selectors can feed values outside the documented 0-255 ints + torch, TensorImageThresholdBlockV1 = _tensor_threshold_imports() + rng = np.random.default_rng(7) + pixels = rng.integers(0, 256, size=(16, 20), dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, pixels) + + # when + numpy_result = _run_threshold_pair( + ImageThresholdBlockV1, numpy_born, threshold_type, thresh_value, max_value + ).numpy_image + tensor_result = _run_threshold_pair( + TensorImageThresholdBlockV1, + tensor_born, + threshold_type, + thresh_value, + max_value, + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result) + + +def _otsu_case_image(case: str) -> np.ndarray: + rng = np.random.default_rng(42) + if case == "bimodal": + dark = np.clip(rng.normal(80, 12, size=(16, 32)), 0, 255) + bright = np.clip(rng.normal(170, 12, size=(16, 32)), 0, 255) + return np.concatenate([dark, bright], axis=0).astype(np.uint8) + if case == "uniform_noise": + return rng.integers(0, 256, size=(32, 32), dtype=np.uint8) + if case == "flat_zero": + return np.zeros((16, 20), dtype=np.uint8) + if case == "flat_mid": + return np.full((16, 20), 130, dtype=np.uint8) + if case == "two_adjacent": + image = np.full((16, 16), 100, dtype=np.uint8) + image[::2, ::2] = 101 + return image + raise ValueError(case) + + +@pytest.mark.parametrize( + "case", ["bimodal", "uniform_noise", "flat_zero", "flat_mid", "two_adjacent"] +) +def test_tensor_threshold_otsu_bit_exact_parity(case) -> None: + # given + import cv2 + + torch, TensorImageThresholdBlockV1 = _tensor_threshold_imports() + from inference.core.workflows.core_steps.classical_cv.threshold.v1_tensor import ( + _otsu_threshold_from_counts, + ) + + gray = _otsu_case_image(case) + numpy_born, tensor_born = _paired_images(torch, gray) + + # when + numpy_result = _run_threshold_pair( + ImageThresholdBlockV1, numpy_born, "otsu", 0, 255 + ).numpy_image + tensor_result = _run_threshold_pair( + TensorImageThresholdBlockV1, tensor_born, "otsu", 0, 255 + ) + + # then + expected_threshold, _ = cv2.threshold( + gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU + ) + counts = np.bincount(gray.reshape(-1), minlength=256).tolist() + assert _otsu_threshold_from_counts(counts) == int(expected_threshold) + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result) + + +def test_tensor_threshold_otsu_three_channel_tensor_raises_like_v1() -> None: + # given - cv2 rejects multi-channel otsu; delegation surfaces the identical error + import cv2 + + torch, TensorImageThresholdBlockV1 = _tensor_threshold_imports() + rng = np.random.default_rng(3) + bgr = rng.integers(0, 256, size=(16, 16, 3), dtype=np.uint8) + numpy_born, tensor_born = _paired_images(torch, bgr) + + # when / then + with pytest.raises(cv2.error): + _run_threshold_pair(ImageThresholdBlockV1, numpy_born, "otsu", 0, 255) + with pytest.raises(cv2.error): + _run_threshold_pair(TensorImageThresholdBlockV1, tensor_born, "otsu", 0, 255) + + +def _adaptive_case_image(case: str) -> np.ndarray: + rng = np.random.default_rng(17) + if case == "noise": + return rng.integers(0, 256, size=(24, 32), dtype=np.uint8) + if case == "gradient": + return np.tile(np.arange(256, dtype=np.uint8), (8, 1)) + if case == "odd_dims": + return rng.integers(0, 256, size=(37, 53), dtype=np.uint8) + if case == "tiny": + # smaller than the 11x11 window - cv2's replicate border still works + return rng.integers(0, 256, size=(3, 5), dtype=np.uint8) + if case == "boundary_hugging": + # src within 2 of the local mean - off-by-one mean errors flip the decision + base = np.full((32, 32), 100, dtype=np.int32) + return np.clip(base + rng.integers(-2, 3, size=base.shape), 0, 255).astype( + np.uint8 + ) + raise ValueError(case) + + +@pytest.mark.parametrize( + "case", ["noise", "gradient", "odd_dims", "tiny", "boundary_hugging"] +) +@pytest.mark.parametrize("max_value", [255, 300]) +def test_tensor_threshold_adaptive_mean_bit_exact_parity(case, max_value) -> None: + # given + torch, TensorImageThresholdBlockV1 = _tensor_threshold_imports() + gray = _adaptive_case_image(case) + numpy_born, tensor_born = _paired_images(torch, gray) + + # when + numpy_result = _run_threshold_pair( + ImageThresholdBlockV1, numpy_born, "adaptive_mean", 127, max_value + ).numpy_image + tensor_result = _run_threshold_pair( + TensorImageThresholdBlockV1, tensor_born, "adaptive_mean", 127, max_value + ) + + # then + assert tensor_result.is_tensor_materialised() + assert np.array_equal(tensor_result.numpy_image, numpy_result) + + +def test_tensor_threshold_adaptive_gaussian_delegates_to_numpy_math() -> None: + # given - cv2's float32 Gaussian mean is SIMD-dispatch dependent, so it delegates + torch, TensorImageThresholdBlockV1 = _tensor_threshold_imports() + gray = _adaptive_case_image("noise") + numpy_born, tensor_born = _paired_images(torch, gray) + + # when + numpy_result = _run_threshold_pair( + ImageThresholdBlockV1, numpy_born, "adaptive_gaussian", 127, 255 + ).numpy_image + tensor_result = _run_threshold_pair( + TensorImageThresholdBlockV1, tensor_born, "adaptive_gaussian", 127, 255 + ) + + # then + assert np.array_equal(tensor_result.numpy_image, numpy_result) + + +def test_tensor_threshold_delegates_for_numpy_born_images() -> None: + # given + torch, TensorImageThresholdBlockV1 = _tensor_threshold_imports() + rng = np.random.default_rng(23) + gray = rng.integers(0, 256, size=(24, 32), dtype=np.uint8) + numpy_born, _ = _paired_images(torch, gray) + reference = _run_threshold_pair( + ImageThresholdBlockV1, + WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=gray, + ), + "binary", + 127, + 255, + ).numpy_image + + # when + result = _run_threshold_pair( + TensorImageThresholdBlockV1, numpy_born, "binary", 127, 255 + ) + + # then + assert np.array_equal(result.numpy_image, reference) + assert not numpy_born.is_tensor_materialised(), "delegate must not materialise" + + +def test_tensor_threshold_unknown_type_raises_on_tensor_path() -> None: + # given + torch, TensorImageThresholdBlockV1 = _tensor_threshold_imports() + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.zeros((1, 8, 8), dtype=torch.uint8), + ) + + # when / then + with pytest.raises(ValueError): + _run_threshold_pair( + TensorImageThresholdBlockV1, tensor_born, "unknown_type", 127, 255 + ) + + +def test_tensor_threshold_on_mps_device(monkeypatch) -> None: + # given - patch the global device so tensor-born images materialise on MPS + torch, TensorImageThresholdBlockV1 = _tensor_threshold_imports() + if not torch.backends.mps.is_available(): + pytest.skip("MPS device not available") + import inference.core.workflows.execution_engine.entities.base as base_module + + gray = _otsu_case_image("bimodal") + numpy_born, _ = _paired_images(torch, gray) + references = { + threshold_type: _run_threshold_pair( + ImageThresholdBlockV1, numpy_born, threshold_type, 127, 255 + ).numpy_image + for threshold_type in ["binary", "otsu"] + } + monkeypatch.setattr(base_module, "WORKFLOWS_IMAGE_TENSOR_DEVICE", "mps") + _, tensor_born = _paired_images(torch, gray) + assert tensor_born.tensor_image.device.type == "mps" + + for threshold_type, reference in references.items(): + # when + result = _run_threshold_pair( + TensorImageThresholdBlockV1, tensor_born, threshold_type, 127, 255 + ) + + # then + assert result.tensor_image.device.type == "mps" + assert np.array_equal(result.numpy_image, reference) diff --git a/tests/workflows/unit_tests/core_steps/common/query_language/operations/detection/test_base.py b/tests/workflows/unit_tests/core_steps/common/query_language/operations/detection/test_base.py index a299a39314..300a6b62e8 100644 --- a/tests/workflows/unit_tests/core_steps/common/query_language/operations/detection/test_base.py +++ b/tests/workflows/unit_tests/core_steps/common/query_language/operations/detection/test_base.py @@ -1,5 +1,7 @@ import numpy as np +import pytest +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.common.query_language.entities.enums import ( DetectionsProperty, ) @@ -9,9 +11,35 @@ from inference.core.workflows.execution_engine.constants import ( AREA_CONVERTED_KEY_IN_SV_DETECTIONS, AREA_KEY_IN_SV_DETECTIONS, + CLASS_NAMES_KEY, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + POLYGON_KEY_IN_SV_DETECTIONS, ) +# Under ENABLE_TENSOR_DATA_REPRESENTATION `extract_detection_property` routes to the +# tensor-native extractor, which expects a 7-element tuple +# (xyxy, mask, class_id, confidence, tracker_id, data, metadata) instead of the sv +# single-detection 6-tuple (note class_id/confidence are swapped vs sv). The 6-tuple +# tests below are skipped when the flag is on; each has a `*_tensor_native` parity test +# (skipped when the flag is off) exercising the same scenario with the 7-tuple form. +# `extract_detection_property` is an ELEMENT-level op (one detection tuple, not a +# collection), so it has no Detections/InstanceDetections/KeyPoints code path; the +# representation only surfaces as which optional `data` keys are populated โ€” hence the +# KEYPOINTS_XY (keypoint-origin) and POLYGON (instance-segmentation-origin) cases below. +_KEYPOINTS_XY = np.array([[10.0, 20.0], [30.0, 40.0]], dtype=np.float32) +_POLYGON = np.array([[0, 0], [10, 0], [10, 10], [0, 10]], dtype=np.float32) +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv 6-tuple input; extract_detection_property is native-only under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + +@_NUMPY_ONLY def test_extract_detection_property_with_area_px() -> None: # given detection = ( @@ -38,6 +66,7 @@ def test_extract_detection_property_with_area_px() -> None: assert result == 400.0 +@_NUMPY_ONLY def test_extract_detection_property_with_area_converted() -> None: # given detection = ( @@ -62,3 +91,152 @@ def test_extract_detection_property_with_area_converted() -> None: # then assert result == 4.0 + + +@_NUMPY_ONLY +def test_extract_detection_property_with_keypoints_xy() -> None: + # given - a detection from a keypoint model carries its keypoint coordinates in the + # per-detection data dict (index 5 of the sv 6-tuple). + detection = ( + np.array([0.0, 0.0, 50.0, 50.0], dtype=np.float32), + None, + np.float32(0.9), + np.int64(1), + None, + {KEYPOINTS_XY_KEY_IN_SV_DETECTIONS: _KEYPOINTS_XY}, + ) + + # when + result = extract_detection_property( + value=detection, + property_name=DetectionsProperty.KEYPOINTS_XY, + execution_context="", + ) + + # then + assert np.array_equal(result, _KEYPOINTS_XY) + + +@_NUMPY_ONLY +def test_extract_detection_property_with_polygon() -> None: + # given - a detection from an instance-segmentation model carries its mask polygon + # in the per-detection data dict (index 5 of the sv 6-tuple). + detection = ( + np.array([0.0, 0.0, 50.0, 50.0], dtype=np.float32), + None, + np.float32(0.9), + np.int64(1), + None, + {POLYGON_KEY_IN_SV_DETECTIONS: _POLYGON}, + ) + + # when + result = extract_detection_property( + value=detection, + property_name=DetectionsProperty.POLYGON, + execution_context="", + ) + + # then + assert np.array_equal(result, _POLYGON) + + +def _single_tensor_native_detection() -> tuple: + # 7-tuple (xyxy, mask, class_id, confidence, tracker_id, data, metadata) as yielded + # when iterating a native `inference_models.Detections`. AREA / AREA_CONVERTED are + # read from `data` (index 5); `metadata` (index 6) carries the class_id -> name map + # and is not consulted for the area properties. + return ( + np.array([0.0, 0.0, 50.0, 50.0], dtype=np.float32), + None, + np.int64(1), + np.float32(0.9), + None, + { + AREA_KEY_IN_SV_DETECTIONS: np.float32(400.0), + AREA_CONVERTED_KEY_IN_SV_DETECTIONS: np.float32(4.0), + }, + {CLASS_NAMES_KEY: {1: "leaf"}}, + ) + + +@_TENSOR_ONLY +def test_extract_detection_property_with_area_px_tensor_native() -> None: + # given + detection = _single_tensor_native_detection() + + # when + result = extract_detection_property( + value=detection, + property_name=DetectionsProperty.AREA, + execution_context="", + ) + + # then + assert result == 400.0 + + +@_TENSOR_ONLY +def test_extract_detection_property_with_area_converted_tensor_native() -> None: + # given + detection = _single_tensor_native_detection() + + # when + result = extract_detection_property( + value=detection, + property_name=DetectionsProperty.AREA_CONVERTED, + execution_context="", + ) + + # then + assert result == 4.0 + + +@_TENSOR_ONLY +def test_extract_detection_property_with_keypoints_xy_tensor_native() -> None: + # given - keypoint-origin detection; keypoints live in `data` (index 5 of the + # tensor-native 7-tuple). + detection = ( + np.array([0.0, 0.0, 50.0, 50.0], dtype=np.float32), + None, + np.int64(1), + np.float32(0.9), + None, + {KEYPOINTS_XY_KEY_IN_SV_DETECTIONS: _KEYPOINTS_XY}, + {CLASS_NAMES_KEY: {1: "person"}}, + ) + + # when + result = extract_detection_property( + value=detection, + property_name=DetectionsProperty.KEYPOINTS_XY, + execution_context="", + ) + + # then + assert np.array_equal(result, _KEYPOINTS_XY) + + +@_TENSOR_ONLY +def test_extract_detection_property_with_polygon_tensor_native() -> None: + # given - instance-segmentation-origin detection; mask polygon lives in `data` + # (index 5 of the tensor-native 7-tuple). + detection = ( + np.array([0.0, 0.0, 50.0, 50.0], dtype=np.float32), + None, + np.int64(1), + np.float32(0.9), + None, + {POLYGON_KEY_IN_SV_DETECTIONS: _POLYGON}, + {CLASS_NAMES_KEY: {1: "leaf"}}, + ) + + # when + result = extract_detection_property( + value=detection, + property_name=DetectionsProperty.POLYGON, + execution_context="", + ) + + # then + assert np.array_equal(result, _POLYGON) diff --git a/tests/workflows/unit_tests/core_steps/common/query_language/operations/detections/test_base.py b/tests/workflows/unit_tests/core_steps/common/query_language/operations/detections/test_base.py index e18238fd73..29299abce7 100644 --- a/tests/workflows/unit_tests/core_steps/common/query_language/operations/detections/test_base.py +++ b/tests/workflows/unit_tests/core_steps/common/query_language/operations/detections/test_base.py @@ -1,7 +1,9 @@ import numpy as np import pytest import supervision as sv +import torch +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.common.query_language.errors import ( InvalidInputTypeError, OperationError, @@ -9,6 +11,48 @@ from inference.core.workflows.core_steps.common.query_language.operations.detections.base import ( rename_detections, ) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, +) +from inference_models.models.base.instance_segmentation import ( + InstanceDetections as NativeInstanceDetections, +) +from inference_models.models.base.keypoints_detection import ( + KeyPoints as NativeKeyPoints, +) +from inference_models.models.base.object_detection import Detections as NativeDetections + +# Under ENABLE_TENSOR_DATA_REPRESENTATION `rename_detections` is native-only: it rejects +# sv.Detections. The sv-based tests below are skipped when the flag is on; each has a +# `*_tensor_native` parity test (skipped when the flag is off) exercising the same +# scenario with an `inference_models.Detections` input. Native renaming rewrites the +# `class_id` tensor and the `image_metadata[CLASS_NAMES_KEY]` class_id -> name map +# (there is no per-box `data["class_name"]`); the parity tests therefore resolve each +# box's name through that map, the same way consumers do. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections input; rename_detections is native-only under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + + +def _native_detections() -> NativeDetections: + return NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.float32), + class_id=torch.tensor([10, 11], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {10: "a", 11: "b"}}, + ) + + +def _resolved_class_names(detections: NativeDetections) -> list: + class_names = detections.image_metadata[CLASS_NAMES_KEY] + return [class_names[int(class_id)] for class_id in detections.class_id.tolist()] def test_rename_detections_when_not_sv_detections_provided() -> None: @@ -23,6 +67,7 @@ def test_rename_detections_when_not_sv_detections_provided() -> None: ) +@_NUMPY_ONLY def test_rename_detections_when_strict_mode_enabled_and_all_classes_present() -> None: # given detections = sv.Detections( @@ -55,6 +100,7 @@ def test_rename_detections_when_strict_mode_enabled_and_all_classes_present() -> ], "Expected to change with mapping" +@_NUMPY_ONLY def test_rename_detections_when_strict_mode_enabled_and_not_all_classes_present() -> ( None ): @@ -77,6 +123,7 @@ def test_rename_detections_when_strict_mode_enabled_and_not_all_classes_present( ) +@_NUMPY_ONLY def test_rename_detections_when_non_strict_mode_enabled_and_all_classes_present() -> ( None ): @@ -111,6 +158,7 @@ def test_rename_detections_when_non_strict_mode_enabled_and_all_classes_present( ], "Expected to change with mapping" +@_NUMPY_ONLY def test_rename_detections_when_non_strict_mode_enabled_and_not_all_classes_present() -> ( None ): @@ -145,6 +193,7 @@ def test_rename_detections_when_non_strict_mode_enabled_and_not_all_classes_pres ], "Expected to change with mapping" +@_NUMPY_ONLY def test_rename_detections_when_detections_have_no_class_name_data_and_not_strict() -> ( None ): @@ -179,26 +228,7 @@ def test_rename_detections_when_detections_have_no_class_name_data_and_not_stric ), "Expected no fabricated class_name field on the no-op path" -def test_rename_detections_when_detections_have_no_class_name_data_and_strict() -> None: - # given - detections = sv.Detections( - xyxy=np.array([[0, 1, 2, 3], [0, 1, 2, 3]]), - class_id=np.array([10, 11]), - confidence=np.array([0.3, 0.4]), - ) - - # when / then - strict mode cannot guarantee class_map coverage without class - # names, so it raises rather than emitting a length-mismatched result. - with pytest.raises(OperationError): - rename_detections( - detections=detections, - class_map={"a": "A"}, - strict=True, - new_classes_id_offset=1024, - global_parameters={}, - ) - - +@_NUMPY_ONLY def test_rename_detections_when_detections_are_empty_is_a_noop() -> None: # given detections = sv.Detections.empty() @@ -225,6 +255,7 @@ def test_rename_detections_when_detections_are_empty_is_a_noop() -> None: assert isinstance(strict, sv.Detections) and len(strict) == 0 +@_NUMPY_ONLY def test_rename_detections_when_mapping_is_parametrised() -> None: # given detections = sv.Detections( @@ -258,3 +289,479 @@ def test_rename_detections_when_mapping_is_parametrised() -> None: "A", "B", ], "Expected to change with mapping" + + +@_TENSOR_ONLY +def test_rename_detections_when_strict_mode_enabled_and_all_classes_present_tensor_native() -> ( + None +): + # given + detections = _native_detections() + + # when + result = rename_detections( + detections=detections, + class_map={"a": "A", "b": "B"}, + strict=True, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then + assert result.xyxy.tolist() == [[0, 1, 2, 3], [0, 1, 2, 3]], "Expected no to change" + assert torch.allclose( + result.confidence, torch.tensor([0.3, 0.4]) + ), "Expected no to change" + assert result.class_id.tolist() == [ + 0, + 1, + ], "Expected to change with mapping" + assert _resolved_class_names(result) == [ + "A", + "B", + ], "Expected to change with mapping" + + +@_TENSOR_ONLY +def test_rename_detections_when_strict_mode_enabled_and_not_all_classes_present_tensor_native() -> ( + None +): + # given + detections = _native_detections() + + # when + with pytest.raises(OperationError): + _ = rename_detections( + detections=detections, + class_map={"a": "A"}, + strict=True, + new_classes_id_offset=1024, + global_parameters={}, + ) + + +@_TENSOR_ONLY +def test_rename_detections_when_non_strict_mode_enabled_and_all_classes_present_tensor_native() -> ( + None +): + # given + detections = _native_detections() + + # when + result = rename_detections( + detections=detections, + class_map={"a": "A", "b": "B"}, + strict=False, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then + assert result.xyxy.tolist() == [[0, 1, 2, 3], [0, 1, 2, 3]], "Expected no to change" + assert torch.allclose( + result.confidence, torch.tensor([0.3, 0.4]) + ), "Expected no to change" + assert result.class_id.tolist() == [ + 1024, + 1025, + ], "Expected to change with mapping" + assert _resolved_class_names(result) == [ + "A", + "B", + ], "Expected to change with mapping" + + +@_TENSOR_ONLY +def test_rename_detections_when_non_strict_mode_enabled_and_not_all_classes_present_tensor_native() -> ( + None +): + # given + detections = _native_detections() + + # when + result = rename_detections( + detections=detections, + class_map={"a": "A"}, + strict=False, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then + assert result.xyxy.tolist() == [[0, 1, 2, 3], [0, 1, 2, 3]], "Expected no to change" + assert torch.allclose( + result.confidence, torch.tensor([0.3, 0.4]) + ), "Expected no to change" + assert result.class_id.tolist() == [ + 1024, + 11, + ], "Expected to change with mapping" + assert _resolved_class_names(result) == [ + "A", + "b", + ], "Expected to change with mapping" + + +@_TENSOR_ONLY +def test_rename_detections_when_detections_have_no_class_name_data_and_not_strict_tensor_native() -> ( + None +): + # given - the native analog of sv detections without `data["class_name"]` is a + # missing `image_metadata[CLASS_NAMES_KEY]` map + detections = NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.float32), + class_id=torch.tensor([10, 11], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={}, + ) + + # when + result = rename_detections( + detections=detections, + class_map={"a": "A"}, + strict=False, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then - with no class names to rename, non-strict mode is a no-op copy: the + # boxes pass through unchanged and no class_names map is fabricated + assert isinstance(result, NativeDetections) + assert result is not detections, "Expected a copy, not the input object" + assert result.xyxy.tolist() == [[0, 1, 2, 3], [0, 1, 2, 3]] + assert torch.allclose(result.confidence, torch.tensor([0.3, 0.4])) + assert result.class_id.tolist() == [ + 10, + 11, + ], "Expected class_id unchanged in the no-op path" + assert (result.image_metadata or {}).get( + CLASS_NAMES_KEY + ) is None, "Expected no fabricated class_names map on the no-op path" + + +@_TENSOR_ONLY +def test_rename_detections_when_detections_have_no_class_name_data_and_strict_tensor_native() -> ( + None +): + # given + detections = NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.float32), + class_id=torch.tensor([10, 11], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={}, + ) + + # when - strict mode cannot guarantee class_map coverage without class names, + # so the shared helper raises (same OperationError as the numpy arm) + with pytest.raises(OperationError): + _ = rename_detections( + detections=detections, + class_map={"a": "A"}, + strict=True, + new_classes_id_offset=1024, + global_parameters={}, + ) + + +@_TENSOR_ONLY +def test_rename_detections_when_detections_are_empty_is_a_noop_tensor_native() -> None: + # given - empty native detections without a class_names map (the mirror of + # sv.Detections.empty(), which carries no `class_name` data) + def _empty_native_detections() -> NativeDetections: + return NativeDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + image_metadata={}, + ) + + # when - empty detections have nothing to rename; both modes return a + # consistent (empty) result without raising + non_strict = rename_detections( + detections=_empty_native_detections(), + class_map={"a": "A"}, + strict=False, + new_classes_id_offset=1024, + global_parameters={}, + ) + strict = rename_detections( + detections=_empty_native_detections(), + class_map={"a": "A"}, + strict=True, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then + assert isinstance(non_strict, NativeDetections) + assert int(non_strict.xyxy.shape[0]) == 0 + assert isinstance(strict, NativeDetections) + assert int(strict.xyxy.shape[0]) == 0 + + +def _native_detections_with_overrides() -> NativeDetections: + # classes_replacement-style input: both rows share class_id 7 whose map name is + # "cat"; row 1 carries a per-box CLASS_NAME_KEY override ("dog") โ€” the label the + # tensor serializer would prefer (C1), hence the label rename must operate on. + return NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.float32), + class_id=torch.tensor([7, 7], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {7: "cat"}}, + bboxes_metadata=[ + {"detection_id": "d0"}, + {"detection_id": "d1", CLASS_NAME_KEY: "dog"}, + ], + ) + + +@_TENSOR_ONLY +def test_rename_detections_rewrites_per_box_class_overrides_tensor_native() -> None: + # given + detections = _native_detections_with_overrides() + + # when - only the override name is in the class_map; numpy on the equivalent + # per-row class_names ["cat", "dog"] yields names ["cat", "canine"], ids [7, 1024] + result = rename_detections( + detections=detections, + class_map={"dog": "canine"}, + strict=False, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then + assert result.class_id.tolist() == [7, 1024] + assert ( + result.bboxes_metadata[1][CLASS_NAME_KEY] == "canine" + ), "Expected the per-box override rewritten to the renamed label" + assert ( + CLASS_NAME_KEY not in result.bboxes_metadata[0] + ), "Expected no override fabricated on a row that had none" + assert result.image_metadata[CLASS_NAMES_KEY][7] == "cat" + assert ( + detections.bboxes_metadata[1][CLASS_NAME_KEY] == "dog" + ), "Expected the source object untouched" + + +@_TENSOR_ONLY +def test_rename_detections_strict_mode_covers_per_box_overrides_tensor_native() -> None: + # given + detections = _native_detections_with_overrides() + + # when - class_map covers the map name but not the override; numpy raises + # "Class 'dog' not found in class_map." for the equivalent per-row names + with pytest.raises(OperationError): + _ = rename_detections( + detections=detections, + class_map={"cat": "feline"}, + strict=True, + new_classes_id_offset=1024, + global_parameters={}, + ) + + +@_TENSOR_ONLY +def test_rename_detections_splits_shared_class_id_on_divergent_overrides_tensor_native() -> ( + None +): + # given + detections = _native_detections_with_overrides() + + # when - both effective names renamed to NEW targets; numpy assigns offset ids + # to the sorted new targets (canine=1024, feline=1025), splitting the shared id + result = rename_detections( + detections=detections, + class_map={"cat": "feline", "dog": "canine"}, + strict=False, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then + assert result.class_id.tolist() == [1025, 1024] + assert ( + result.image_metadata[CLASS_NAMES_KEY][1025] == "feline" + ), "Expected the override-less row to resolve via the class_names map" + assert result.bboxes_metadata[1][CLASS_NAME_KEY] == "canine" + + +@_TENSOR_ONLY +def test_rename_detections_keeps_uncovered_override_in_non_strict_mode_tensor_native() -> ( + None +): + # given + detections = _native_detections_with_overrides() + + # when - the override name is not in the class_map; numpy keeps the name and + # its original id on the equivalent per-row input + result = rename_detections( + detections=detections, + class_map={"cat": "feline"}, + strict=False, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then + assert result.class_id.tolist() == [1024, 7] + assert ( + result.bboxes_metadata[1][CLASS_NAME_KEY] == "dog" + ), "Expected the uncovered override kept intact" + assert result.image_metadata[CLASS_NAMES_KEY][1024] == "feline" + + +@_TENSOR_ONLY +def test_rename_detections_when_mapping_is_parametrised_tensor_native() -> None: + # given + detections = _native_detections() + + # when + result = rename_detections( + detections=detections, + class_map="class_map_param", + strict="strict_param", + new_classes_id_offset=1024, + global_parameters={ + "strict_param": True, + "class_map_param": {"a": "A", "b": "B"}, + }, + ) + + # then + assert result.xyxy.tolist() == [[0, 1, 2, 3], [0, 1, 2, 3]], "Expected no to change" + assert torch.allclose( + result.confidence, torch.tensor([0.3, 0.4]) + ), "Expected no to change" + assert result.class_id.tolist() == [ + 0, + 1, + ], "Expected to change with mapping" + assert _resolved_class_names(result) == [ + "A", + "B", + ], "Expected to change with mapping" + + +# --------------------------------------------------------------------------- +# Parity for the other native detection representations: InstanceDetections +# (rename carries the per-instance mask through unchanged) and the KeyPoints +# prediction (a (KeyPoints, Detections) tuple; rename touches only the bbox +# component's class ids / names, the KeyPoints component is preserved as-is). +# The class-id / class-name rewrite logic itself is identical across reps and +# is already covered by the plain-`Detections` cases above, so these focus on +# the representation-specific carry-through. +# --------------------------------------------------------------------------- + + +def _native_instance_detections(mask: torch.Tensor) -> NativeInstanceDetections: + return NativeInstanceDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.float32), + class_id=torch.tensor([10, 11], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + mask=mask, + image_metadata={CLASS_NAMES_KEY: {10: "a", 11: "b"}}, + ) + + +@_TENSOR_ONLY +def test_rename_detections_for_instance_segmentation_tensor_native() -> None: + # given - distinct per-instance pixel counts (3 / 5) so we can assert the masks + # ride through unchanged and in order. + mask = torch.zeros((2, 20, 20), dtype=torch.bool) + mask[0, 0, 0:3] = True # 3 px + mask[1, 0, 0:5] = True # 5 px + detections = _native_instance_detections(mask) + + # when + result = rename_detections( + detections=detections, + class_map={"a": "A", "b": "B"}, + strict=True, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then - same repr, class ids / names rewritten, mask carried through unchanged + assert isinstance(result, NativeInstanceDetections) + assert result.xyxy.tolist() == [[0, 1, 2, 3], [0, 1, 2, 3]], "Expected no to change" + assert result.class_id.tolist() == [0, 1], "Expected to change with mapping" + assert _resolved_class_names(result) == [ + "A", + "B", + ], "Expected to change with mapping" + assert result.mask.sum(dim=(1, 2)).cpu().numpy().tolist() == [ + 3, + 5, + ], "Expected mask to be carried through unchanged and in order" + assert torch.equal(result.mask, mask) + + +@_TENSOR_ONLY +def test_rename_detections_for_instance_segmentation_when_empty_tensor_native() -> None: + # given + detections = NativeInstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, 20, 20), dtype=torch.bool), + image_metadata={CLASS_NAMES_KEY: {}}, + ) + + # when + result = rename_detections( + detections=detections, + class_map={"a": "A", "b": "B"}, + strict=True, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then - empty object of the SAME repr + assert isinstance(result, NativeInstanceDetections) + assert result.xyxy.shape[0] == 0 + assert result.mask.shape[0] == 0 + + +@_TENSOR_ONLY +def test_rename_detections_for_keypoints_tensor_native() -> None: + # given - a keypoint prediction is a (KeyPoints, Detections) tuple. Rename rewrites + # the bbox component's class ids / names; the KeyPoints component is preserved. + key_points = NativeKeyPoints( + xy=torch.tensor( + [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]] + ), # (instances, kpts, 2) + class_id=torch.tensor([10, 11], dtype=torch.long), + confidence=torch.tensor([[0.9, 0.8], [0.7, 0.6]]), # (instances, kpts) + image_metadata={CLASS_NAMES_KEY: {10: "a", 11: "b"}}, + ) + bboxes = NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.float32), + class_id=torch.tensor([10, 11], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {10: "a", 11: "b"}}, + ) + prediction = (key_points, bboxes) + + # when + result = rename_detections( + detections=prediction, + class_map={"a": "A", "b": "B"}, + strict=True, + new_classes_id_offset=1024, + global_parameters={}, + ) + + # then - a (KeyPoints, Detections) tuple; bbox renamed, keypoints carried through + assert isinstance(result, tuple) and len(result) == 2 + result_key_points, result_bboxes = result + assert isinstance(result_key_points, NativeKeyPoints) + assert isinstance(result_bboxes, NativeDetections) + assert result_bboxes.class_id.tolist() == [0, 1], "Expected to change with mapping" + assert _resolved_class_names(result_bboxes) == [ + "A", + "B", + ], "Expected to change with mapping" + # the KeyPoints component is unchanged + assert torch.equal(result_key_points.xy, key_points.xy) + assert torch.equal(result_key_points.class_id, key_points.class_id) diff --git a/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_classification_results_operations.py b/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_classification_results_operations.py index c8a213765d..808f044842 100644 --- a/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_classification_results_operations.py +++ b/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_classification_results_operations.py @@ -1,4 +1,5 @@ import pytest +import torch from inference.core.entities.responses.inference import ( ClassificationInferenceResponse, @@ -7,12 +8,34 @@ MultiLabelClassificationInferenceResponse, MultiLabelClassificationPrediction, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.common.query_language.errors import ( InvalidInputTypeError, ) from inference.core.workflows.core_steps.common.query_language.operations.core import ( execute_operations, ) +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY +from inference_models import ClassificationPrediction as NativeClassificationPrediction +from inference_models import ( + MultiLabelClassificationPrediction as NativeMultiLabelClassificationPrediction, +) + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the UQL classification extractors are +# native-only: they reject the serialised dict form. The dict-based tests below are +# skipped when the flag is on; each has a `*_tensor_native` parity test (skipped when +# the flag is off) exercising the same scenario with a native `inference_models` +# prediction. Multi-class class names live in `images_metadata[i][CLASS_NAMES_KEY]`; +# multi-label in `image_metadata[CLASS_NAMES_KEY]`. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="dict input; UQL classification extractors are native-only under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) def test_classification_result_extraction_when_data_is_empty() -> None: @@ -29,6 +52,7 @@ def test_classification_result_extraction_when_data_is_empty() -> None: _ = execute_operations(value=None, operations=operations) +@_NUMPY_ONLY def test_classification_result_extraction_of_top_class_for_multi_class_classification_result() -> ( None ): @@ -61,6 +85,7 @@ def test_classification_result_extraction_of_top_class_for_multi_class_classific assert result == "cat" +@_NUMPY_ONLY def test_classification_result_extraction_of_top_class_for_multi_label_classification_result_when_no_class_detected() -> ( None ): @@ -87,6 +112,7 @@ def test_classification_result_extraction_of_top_class_for_multi_label_classific assert result == [] +@_NUMPY_ONLY def test_classification_result_extraction_of_top_class_for_multi_label_classification_result_when_classes_detected() -> ( None ): @@ -113,6 +139,7 @@ def test_classification_result_extraction_of_top_class_for_multi_label_classific assert result == ["cat", "dog"] +@_NUMPY_ONLY def test_classification_result_extraction_of_top_class_confidence_for_multi_class_classification_result() -> ( None ): @@ -145,6 +172,7 @@ def test_classification_result_extraction_of_top_class_confidence_for_multi_clas assert abs(result - 0.6) < 1e-5 +@_NUMPY_ONLY def test_classification_result_extraction_of_top_class_confidence_for_multi_label_classification_result_when_no_class_detected() -> ( None ): @@ -171,6 +199,7 @@ def test_classification_result_extraction_of_top_class_confidence_for_multi_labe assert result == [] +@_NUMPY_ONLY def test_classification_result_extraction_of_top_class_confidence_for_multi_label_classification_result_when_class_detected() -> ( None ): @@ -197,6 +226,7 @@ def test_classification_result_extraction_of_top_class_confidence_for_multi_labe assert result == [0.4] +@_NUMPY_ONLY def test_classification_result_extraction_of_top_class_confidence_single_for_multi_label_classification_result_when_class_detected() -> ( None ): @@ -223,6 +253,7 @@ def test_classification_result_extraction_of_top_class_confidence_single_for_mul assert result == 0.6 +@_NUMPY_ONLY def test_classification_result_extraction_of_top_class_confidence_single_for_multi_label_classification_result_when_no_classes_detected() -> ( None ): @@ -249,6 +280,7 @@ def test_classification_result_extraction_of_top_class_confidence_single_for_mul assert result == 0.0 +@_NUMPY_ONLY def test_classification_result_extraction_of_all_classes_for_multi_class_classification_result() -> ( None ): @@ -281,6 +313,7 @@ def test_classification_result_extraction_of_all_classes_for_multi_class_classif assert result == ["cat", "dog"] +@_NUMPY_ONLY def test_classification_result_extraction_of_all_classes_for_multi_label_classification_result() -> ( None ): @@ -308,6 +341,7 @@ def test_classification_result_extraction_of_all_classes_for_multi_label_classif assert result == ["cat", "dog", "animal"] +@_NUMPY_ONLY def test_classification_result_extraction_of_all_confidences_for_multi_class_classification_result() -> ( None ): @@ -340,6 +374,7 @@ def test_classification_result_extraction_of_all_confidences_for_multi_class_cla assert result == [0.6, 0.4] +@_NUMPY_ONLY def test_classification_result_extraction_of_all_confidences_for_multi_label_classification_result() -> ( None ): @@ -365,3 +400,278 @@ def test_classification_result_extraction_of_all_confidences_for_multi_label_cla # then assert result == [0.6, 0.4, 0.0] + + +# --------------------------------------------------------------------------- +# Tensor-native parity variants (run only under ENABLE_TENSOR_DATA_REPRESENTATION). +# Same scenarios as the dict-based tests above, but with native `inference_models` +# predictions. Multi-class: class_id (bs,), confidence (bs, num_classes), class names +# in images_metadata[i][CLASS_NAMES_KEY]. Multi-label: class_ids (selected,), +# confidence (num_classes,), class names in image_metadata[CLASS_NAMES_KEY]. +# --------------------------------------------------------------------------- + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_top_class_for_multi_class_classification_result_tensor_native() -> ( + None +): + # given + operations = [ + {"type": "ClassificationPropertyExtract", "property_name": "top_class"} + ] + prediction = NativeClassificationPrediction( + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.6, 0.4]], dtype=torch.float32), + images_metadata=[{CLASS_NAMES_KEY: {0: "cat", 1: "dog"}}], + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == "cat" + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_top_class_for_multi_label_classification_result_when_classes_detected_tensor_native() -> ( + None +): + # given + operations = [ + {"type": "ClassificationPropertyExtract", "property_name": "top_class"} + ] + prediction = NativeMultiLabelClassificationPrediction( + class_ids=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "cat", 1: "dog"}}, + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == ["cat", "dog"] + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_top_class_for_multi_label_classification_result_when_no_class_detected_tensor_native() -> ( + None +): + # given - no labels above threshold -> empty class_ids + operations = [ + {"type": "ClassificationPropertyExtract", "property_name": "top_class"} + ] + prediction = NativeMultiLabelClassificationPrediction( + class_ids=torch.tensor([], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "cat", 1: "dog"}}, + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == [] + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_top_class_confidence_for_multi_class_classification_result_tensor_native() -> ( + None +): + # given - confidence is the full (bs, num_classes) softmax; top class is id 0 + operations = [ + { + "type": "ClassificationPropertyExtract", + "property_name": "top_class_confidence", + } + ] + prediction = NativeClassificationPrediction( + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.6, 0.4]], dtype=torch.float32), + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == pytest.approx(0.6) + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_top_class_confidence_for_multi_label_classification_result_when_no_class_detected_tensor_native() -> ( + None +): + # given + operations = [ + { + "type": "ClassificationPropertyExtract", + "property_name": "top_class_confidence", + } + ] + prediction = NativeMultiLabelClassificationPrediction( + class_ids=torch.tensor([], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4], dtype=torch.float32), + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == [] + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_top_class_confidence_for_multi_label_classification_result_when_class_detected_tensor_native() -> ( + None +): + # given - only label id 1 is above threshold + operations = [ + { + "type": "ClassificationPropertyExtract", + "property_name": "top_class_confidence", + } + ] + prediction = NativeMultiLabelClassificationPrediction( + class_ids=torch.tensor([1], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4], dtype=torch.float32), + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == pytest.approx([0.4]) + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_top_class_confidence_single_for_multi_label_classification_result_when_class_detected_tensor_native() -> ( + None +): + # given - both labels above threshold; single -> max selected confidence + operations = [ + { + "type": "ClassificationPropertyExtract", + "property_name": "top_class_confidence_single", + } + ] + prediction = NativeMultiLabelClassificationPrediction( + class_ids=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4], dtype=torch.float32), + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == pytest.approx(0.6) + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_top_class_confidence_single_for_multi_label_classification_result_when_no_classes_detected_tensor_native() -> ( + None +): + # given - nothing above threshold; single -> 0.0 + operations = [ + { + "type": "ClassificationPropertyExtract", + "property_name": "top_class_confidence_single", + } + ] + prediction = NativeMultiLabelClassificationPrediction( + class_ids=torch.tensor([], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4], dtype=torch.float32), + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == 0.0 + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_all_classes_for_multi_class_classification_result_tensor_native() -> ( + None +): + # given - all_classes returns the full class map, ordered by class id + operations = [ + {"type": "ClassificationPropertyExtract", "property_name": "all_classes"} + ] + # class map deliberately NOT in id order, to exercise sorted-by-id output + prediction = NativeClassificationPrediction( + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.6, 0.4]], dtype=torch.float32), + images_metadata=[{CLASS_NAMES_KEY: {1: "dog", 0: "cat"}}], + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == ["cat", "dog"] + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_all_classes_for_multi_label_classification_result_tensor_native() -> ( + None +): + # given - all_classes returns the full class map ordered by id (0, 1, 3), + # regardless of which were detected; map insertion order is deliberately + # scrambled so the sorted-by-id output is actually exercised. Class id 3 is + # sparse, so the (num_classes,) confidence vector has 4 slots (id 2 unnamed). + operations = [ + {"type": "ClassificationPropertyExtract", "property_name": "all_classes"} + ] + prediction = NativeMultiLabelClassificationPrediction( + class_ids=torch.tensor([1], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4, 0.0, 0.0], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {3: "animal", 0: "cat", 1: "dog"}}, + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == ["cat", "dog", "animal"] + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_all_confidences_for_multi_class_classification_result_tensor_native() -> ( + None +): + # given - all_confidences returns the full softmax row + operations = [ + {"type": "ClassificationPropertyExtract", "property_name": "all_confidences"} + ] + prediction = NativeClassificationPrediction( + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.6, 0.4]], dtype=torch.float32), + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == pytest.approx([0.6, 0.4]) + + +@_TENSOR_ONLY +def test_classification_result_extraction_of_all_confidences_for_multi_label_classification_result_tensor_native() -> ( + None +): + # given - native returns the full sigmoid vector wholesale; the dict counterpart + # used sparse ids {0, 1, 3} for the same three confidences, mirrored here as a + # contiguous (num_classes,) tensor. + operations = [ + {"type": "ClassificationPropertyExtract", "property_name": "all_confidences"} + ] + prediction = NativeMultiLabelClassificationPrediction( + class_ids=torch.tensor([1], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4, 0.0], dtype=torch.float32), + ) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then + assert result == pytest.approx([0.6, 0.4, 0.0]) diff --git a/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_detections_operations.py b/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_detections_operations.py index 723ac1819c..ef686fd173 100644 --- a/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_detections_operations.py +++ b/tests/workflows/unit_tests/core_steps/common/query_language/operations/test_detections_operations.py @@ -1,7 +1,9 @@ import numpy as np import pytest import supervision as sv +import torch +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.common.query_language.errors import ( InvalidInputTypeError, OperationError, @@ -9,6 +11,33 @@ from inference.core.workflows.core_steps.common.query_language.operations.core import ( execute_operations, ) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + POLYGON_KEY, +) +from inference_models.models.base.instance_segmentation import ( + InstanceDetections as NativeInstanceDetections, +) +from inference_models.models.base.keypoints_detection import ( + KeyPoints as NativeKeyPoints, +) +from inference_models.models.base.object_detection import Detections as NativeDetections + +# Under ENABLE_TENSOR_DATA_REPRESENTATION the UQL detections ops are native-only: +# they reject sv.Detections / dict inputs. The sv-based tests below are skipped when +# the flag is on; each has a `*_tensor_native` parity test (skipped when the flag is +# off) that exercises the same scenario with an `inference_models.Detections` input. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="sv.Detections input; UQL detections ops are native-only under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) def test_detections_to_dictionary_when_invalid_input_is_provided() -> None: @@ -19,6 +48,7 @@ def test_detections_to_dictionary_when_invalid_input_is_provided() -> None: _ = execute_operations(value="invalid", operations=operations) +@_NUMPY_ONLY def test_detections_to_dictionary_when_valid_input_is_provided() -> None: # given operations = [{"type": "DetectionsToDictionary"}] @@ -69,6 +99,7 @@ def test_detections_to_dictionary_when_valid_input_is_provided() -> None: } +@_NUMPY_ONLY def test_detections_to_dictionary_when_malformed_input_is_provided() -> None: # given operations = [{"type": "DetectionsToDictionary"}] @@ -97,6 +128,7 @@ def test_picking_detections_by_parent_class_when_invalid_input_provided() -> Non _ = execute_operations(value="FOR SURE NOT DETECTIONS", operations=operations) +@_NUMPY_ONLY def test_picking_detections_by_parent_class_when_empty_detections_provided() -> None: # given operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] @@ -110,6 +142,7 @@ def test_picking_detections_by_parent_class_when_empty_detections_provided() -> assert len(result) == 0 +@_NUMPY_ONLY def test_picking_detections_by_parent_class_when_class_name_field_not_defined() -> None: # given operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] @@ -131,6 +164,7 @@ def test_picking_detections_by_parent_class_when_class_name_field_not_defined() assert len(result) == 0 +@_NUMPY_ONLY def test_picking_detections_by_parent_class_when_parent_class_not_fond() -> None: # given operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] @@ -153,6 +187,7 @@ def test_picking_detections_by_parent_class_when_parent_class_not_fond() -> None assert len(result) == 0 +@_NUMPY_ONLY def test_picking_detections_by_parent_class_when_no_child_detections_matching() -> None: # given operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] @@ -178,6 +213,7 @@ def test_picking_detections_by_parent_class_when_no_child_detections_matching() assert np.allclose(result.confidence, [0.3, 0.4]) +@_NUMPY_ONLY def test_picking_detections_by_parent_class_when_there_are_child_detections_matching() -> ( None ): @@ -207,6 +243,7 @@ def test_picking_detections_by_parent_class_when_there_are_child_detections_matc assert np.allclose(result.confidence, [0.3, 0.4, 0.5]) +@_NUMPY_ONLY def test_picking_detections_by_parent_class_when_there_are_child_detections_matching_different_parents() -> ( None ): @@ -246,3 +283,439 @@ def test_picking_detections_by_parent_class_when_there_are_child_detections_matc ), ) assert np.allclose(result.confidence, [0.3, 0.6, 0.4, 0.5, 0.7]) + + +# --------------------------------------------------------------------------- +# Tensor-native parity variants (run only under ENABLE_TENSOR_DATA_REPRESENTATION). +# Same scenarios as the sv.Detections tests above, but with native +# `inference_models.Detections` inputs. Class names live in +# image_metadata[CLASS_NAMES_KEY]; per-box detection_id in bboxes_metadata. +# --------------------------------------------------------------------------- + + +@_TENSOR_ONLY +def test_detections_to_dictionary_when_valid_input_is_provided_tensor_native() -> None: + # given + operations = [{"type": "DetectionsToDictionary"}] + detections = NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={ + CLASS_NAMES_KEY: {0: "cat", 1: "dog"}, + IMAGE_DIMENSIONS_KEY: [192, 168], + }, + bboxes_metadata=[{DETECTION_ID_KEY: "one"}, {DETECTION_ID_KEY: "two"}], + ) + + # when + result = execute_operations(value=detections, operations=operations) + + # then + assert result["image"] == {"width": 168, "height": 192} + first, second = result["predictions"] + assert (first["x"], first["y"], first["width"], first["height"]) == ( + 1.0, + 2.0, + 2.0, + 2.0, + ) + assert first["class_id"] == 0 and first["class"] == "cat" + assert first["detection_id"] == "one" + assert first["confidence"] == pytest.approx(0.3, abs=1e-6) + assert (second["x"], second["y"], second["width"], second["height"]) == ( + 5.0, + 6.0, + 2.0, + 2.0, + ) + assert second["class_id"] == 1 and second["class"] == "dog" + assert second["detection_id"] == "two" + assert second["confidence"] == pytest.approx(0.4, abs=1e-6) + + +@_TENSOR_ONLY +def test_detections_to_dictionary_when_malformed_input_is_provided_tensor_native() -> ( + None +): + # given - native detections without per-box detection_id -> serialiser raises + operations = [{"type": "DetectionsToDictionary"}] + detections = NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "cat", 1: "dog"}}, + bboxes_metadata=None, + ) + + # when + with pytest.raises(OperationError): + _ = execute_operations(value=detections, operations=operations) + + +@_TENSOR_ONLY +def test_picking_detections_by_parent_class_when_empty_detections_provided_tensor_native() -> ( + None +): + # given + operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] + detections = NativeDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {}}, + ) + + # when + result = execute_operations(value=detections, operations=operations) + + # then + assert result is not detections + assert result.xyxy.shape[0] == 0 + + +@_TENSOR_ONLY +def test_picking_detections_by_parent_class_when_class_name_field_not_defined_tensor_native() -> ( + None +): + # given - no CLASS_NAMES_KEY: the native op is stricter than the numpy sibling + # (which yields []) and raises, since class names cannot be resolved. + operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] + detections = NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={}, + ) + + # when + with pytest.raises(OperationError): + _ = execute_operations(value=detections, operations=operations) + + +@_TENSOR_ONLY +def test_picking_detections_by_parent_class_when_parent_class_not_fond_tensor_native() -> ( + None +): + # given + operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] + detections = NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "c", 1: "d"}}, + ) + + # when + result = execute_operations(value=detections, operations=operations) + + # then + assert result.xyxy.shape[0] == 0 + + +@_TENSOR_ONLY +def test_picking_detections_by_parent_class_when_no_child_detections_matching_tensor_native() -> ( + None +): + # given + operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] + detections = NativeDetections( + xyxy=torch.tensor( + [[0, 0, 10, 10], [20, 20, 30, 30], [40, 40, 50, 50]], dtype=torch.float32 + ), + class_id=torch.tensor([0, 0, 1], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4, 0.5], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "a", 1: "b"}}, + ) + + # when + result = execute_operations(value=detections, operations=operations) + + # then + assert result.xyxy.shape[0] == 2 + assert np.allclose( + result.xyxy.cpu().numpy(), np.array([[0, 0, 10, 10], [20, 20, 30, 30]]) + ) + assert np.allclose(result.confidence.cpu().numpy(), [0.3, 0.4]) + + +@_TENSOR_ONLY +def test_picking_detections_by_parent_class_when_there_are_child_detections_matching_tensor_native() -> ( + None +): + # given + operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] + detections = NativeDetections( + xyxy=torch.tensor( + [[0, 0, 50, 50], [20, 20, 30, 30], [40, 40, 50, 50]], dtype=torch.float32 + ), + class_id=torch.tensor([0, 1, 1], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4, 0.5], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "a", 1: "b"}}, + ) + + # when + result = execute_operations(value=detections, operations=operations) + + # then + assert result.xyxy.shape[0] == 3 + assert np.allclose( + result.xyxy.cpu().numpy(), + np.array([[0, 0, 50, 50], [20, 20, 30, 30], [40, 40, 50, 50]]), + ) + assert np.allclose(result.confidence.cpu().numpy(), [0.3, 0.4, 0.5]) + + +@_TENSOR_ONLY +def test_picking_detections_by_parent_class_when_there_are_child_detections_matching_different_parents_tensor_native() -> ( + None +): + # given + operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] + detections = NativeDetections( + xyxy=torch.tensor( + [ + [0, 0, 50, 50], + [20, 20, 30, 30], + [40, 40, 50, 50], + [100, 100, 200, 200], + [150, 100, 250, 200], + [400, 400, 600, 600], + ], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 1, 1, 0, 2, 3], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4, 0.5, 0.6, 0.7, 0.9], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "a", 1: "b", 2: "c", 3: "d"}}, + ) + + # when + result = execute_operations(value=detections, operations=operations) + + # then - parents first (original order), then contained dependents + assert result.xyxy.shape[0] == 5 + assert np.allclose( + result.xyxy.cpu().numpy(), + np.array( + [ + [0, 0, 50, 50], + [100, 100, 200, 200], + [20, 20, 30, 30], + [40, 40, 50, 50], + [150, 100, 250, 200], + ] + ), + ) + assert np.allclose(result.confidence.cpu().numpy(), [0.3, 0.6, 0.4, 0.5, 0.7]) + + +# --------------------------------------------------------------------------- +# Parity for the other native detection representations: InstanceDetections +# (masks -> polygons; pick preserves the mask) and the KeyPoints prediction +# (a (KeyPoints, Detections) tuple; pick slices both components together). +# --------------------------------------------------------------------------- + + +@_TENSOR_ONLY +def test_detections_to_dictionary_for_instance_segmentation_tensor_native() -> None: + # given - InstanceDetections serialises like Detections but adds a per-box + # polygon (POLYGON_KEY) derived from the mask. An instance whose mask has no + # contour is silently dropped, so the masks here are filled blobs. + operations = [{"type": "DetectionsToDictionary"}] + mask = torch.zeros((2, 20, 20), dtype=torch.bool) + mask[0, 2:12, 2:12] = True + mask[1, 5:15, 5:15] = True + detections = NativeInstanceDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + mask=mask, + image_metadata={ + CLASS_NAMES_KEY: {0: "cat", 1: "dog"}, + IMAGE_DIMENSIONS_KEY: [192, 168], + }, + bboxes_metadata=[{DETECTION_ID_KEY: "one"}, {DETECTION_ID_KEY: "two"}], + ) + + # when + result = execute_operations(value=detections, operations=operations) + + # then + assert result["image"] == {"width": 168, "height": 192} + first, second = result["predictions"] + assert first["class_id"] == 0 and first["class"] == "cat" + assert first["detection_id"] == "one" + assert first["confidence"] == pytest.approx(0.3, abs=1e-6) + assert second["class_id"] == 1 and second["class"] == "dog" + assert second["detection_id"] == "two" + # the InstanceDetections-specific bit: a polygon rides on every prediction + for prediction in (first, second): + assert POLYGON_KEY in prediction + assert len(prediction[POLYGON_KEY]) >= 3 + assert set(prediction[POLYGON_KEY][0].keys()) == {"x", "y"} + + +@_TENSOR_ONLY +def test_picking_detections_by_parent_class_for_instance_segmentation_tensor_native() -> ( + None +): + # given - parent is NOT first, so the result re-orders boxes; the mask rows must + # follow the same re-ordering. Each instance's mask carries a distinct pixel + # count (1 / 2 / 3) so we can assert the rows move with their boxes. + operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] + mask = torch.zeros((3, 20, 20), dtype=torch.bool) + mask[0, 0, 0:1] = True # idx0 child -> 1 px + mask[1, 0, 0:2] = True # idx1 parent -> 2 px + mask[2, 0, 0:3] = True # idx2 child -> 3 px + detections = NativeInstanceDetections( + xyxy=torch.tensor( + [[20, 20, 30, 30], [0, 0, 50, 50], [40, 40, 50, 50]], dtype=torch.float32 + ), + class_id=torch.tensor([1, 0, 1], dtype=torch.long), + confidence=torch.tensor([0.4, 0.3, 0.5], dtype=torch.float32), + mask=mask, + image_metadata={CLASS_NAMES_KEY: {0: "a", 1: "b"}}, + ) + + # when + result = execute_operations(value=detections, operations=operations) + + # then - same repr, parents-first ordering, mask sliced in lockstep + assert isinstance(result, NativeInstanceDetections) + assert result.xyxy.shape[0] == 3 + assert np.allclose( + result.xyxy.cpu().numpy(), + np.array([[0, 0, 50, 50], [20, 20, 30, 30], [40, 40, 50, 50]]), + ) + assert np.allclose(result.confidence.cpu().numpy(), [0.3, 0.4, 0.5]) + mask_pixel_counts = result.mask.sum(dim=(1, 2)).cpu().numpy().tolist() + assert mask_pixel_counts == [2, 1, 3] + + +@_TENSOR_ONLY +def test_picking_detections_by_parent_class_for_instance_segmentation_when_empty_tensor_native() -> ( + None +): + # given + operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] + detections = NativeInstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, 20, 20), dtype=torch.bool), + image_metadata={CLASS_NAMES_KEY: {}}, + ) + + # when + result = execute_operations(value=detections, operations=operations) + + # then - empty object of the SAME repr (not an sv / plain Detections) + assert isinstance(result, NativeInstanceDetections) + assert result.xyxy.shape[0] == 0 + assert result.mask.shape[0] == 0 + + +@_TENSOR_ONLY +def test_detections_to_dictionary_for_keypoints_tensor_native() -> None: + # given - a keypoint prediction is a (KeyPoints, Detections) tuple. The serialiser + # uses only the bbox Detections component (class names / detection ids live there); + # the raw keypoint payload is not part of the dictionary. + operations = [{"type": "DetectionsToDictionary"}] + key_points = NativeKeyPoints( + xy=torch.tensor( + [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]] + ), # (instances, kpts, 2) + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([[0.9, 0.8], [0.7, 0.6]]), # (instances, kpts) + image_metadata={CLASS_NAMES_KEY: {0: "cat", 1: "dog"}}, + ) + bboxes = NativeDetections( + xyxy=torch.tensor([[0, 1, 2, 3], [4, 5, 6, 7]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.3, 0.4], dtype=torch.float32), + image_metadata={ + CLASS_NAMES_KEY: {0: "cat", 1: "dog"}, + IMAGE_DIMENSIONS_KEY: [192, 168], + }, + bboxes_metadata=[{DETECTION_ID_KEY: "one"}, {DETECTION_ID_KEY: "two"}], + ) + prediction = (key_points, bboxes) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then - serialised from the bbox component + assert result["image"] == {"width": 168, "height": 192} + first, second = result["predictions"] + assert first["class_id"] == 0 and first["class"] == "cat" + assert first["detection_id"] == "one" + assert first["confidence"] == pytest.approx(0.3, abs=1e-6) + assert second["class_id"] == 1 and second["class"] == "dog" + assert second["detection_id"] == "two" + + +@_TENSOR_ONLY +def test_detections_to_dictionary_for_keypoints_when_bbox_missing_tensor_native() -> ( + None +): + # given - a keypoint tuple with no bbox component is not a valid input + operations = [{"type": "DetectionsToDictionary"}] + key_points = NativeKeyPoints( + xy=torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]), + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.9, 0.8]]), + image_metadata={CLASS_NAMES_KEY: {0: "cat"}}, + ) + prediction = (key_points, None) + + # when + with pytest.raises(InvalidInputTypeError): + _ = execute_operations(value=prediction, operations=operations) + + +@_TENSOR_ONLY +def test_picking_detections_by_parent_class_for_keypoints_tensor_native() -> None: + # given - pick on a keypoint tuple slices BOTH the KeyPoints and the bbox + # Detections with the same parents-first index order. Parent (class "a") is at + # index 1, so both components re-order to [parent, child, child]. Each instance's + # first keypoint x is a distinct marker so we can assert the keypoints follow. + operations = [{"type": "PickDetectionsByParentClass", "parent_class": "a"}] + key_points = NativeKeyPoints( + xy=torch.tensor( + [ + [[10.0, 0.0], [11.0, 0.0]], # idx0 child + [[20.0, 0.0], [21.0, 0.0]], # idx1 parent + [[30.0, 0.0], [31.0, 0.0]], # idx2 child + ] + ), + class_id=torch.tensor([1, 0, 1], dtype=torch.long), + confidence=torch.tensor([[0.9, 0.8], [0.7, 0.6], [0.5, 0.4]]), + image_metadata={CLASS_NAMES_KEY: {0: "a", 1: "b"}}, + ) + bboxes = NativeDetections( + xyxy=torch.tensor( + [[20, 20, 30, 30], [0, 0, 50, 50], [40, 40, 50, 50]], dtype=torch.float32 + ), + class_id=torch.tensor([1, 0, 1], dtype=torch.long), + confidence=torch.tensor([0.4, 0.3, 0.5], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "a", 1: "b"}}, + ) + prediction = (key_points, bboxes) + + # when + result = execute_operations(value=prediction, operations=operations) + + # then - a (KeyPoints, Detections) tuple, both sliced parents-first + assert isinstance(result, tuple) and len(result) == 2 + result_key_points, result_bboxes = result + assert isinstance(result_key_points, NativeKeyPoints) + assert isinstance(result_bboxes, NativeDetections) + assert np.allclose( + result_bboxes.xyxy.cpu().numpy(), + np.array([[0, 0, 50, 50], [20, 20, 30, 30], [40, 40, 50, 50]]), + ) + assert np.allclose(result_bboxes.confidence.cpu().numpy(), [0.3, 0.4, 0.5]) + # the KeyPoints component is reduced to the same instances, in the same order + # (len() here is the KeyPoints.__len__ added alongside these tests). + assert len(result_key_points) == 3 + assert np.allclose(result_key_points.xy[:, 0, 0].cpu().numpy(), [20.0, 10.0, 30.0]) diff --git a/tests/workflows/unit_tests/core_steps/common/test_deserializers.py b/tests/workflows/unit_tests/core_steps/common/test_deserializers.py index f8182992a7..05ba0613f2 100644 --- a/tests/workflows/unit_tests/core_steps/common/test_deserializers.py +++ b/tests/workflows/unit_tests/core_steps/common/test_deserializers.py @@ -1418,3 +1418,62 @@ def test_deserialize_labeled_points_kind_when_coordinates_are_not_numbers() -> N _ = deserialize_labeled_points_kind( parameter="some", value=[{"x": "a", "y": 100}] ) + + +def test_tensor_deserialize_detections_kind_round_trips_nearest_target_distance() -> ( + None +): + # given: one detection with a real match distance, one unmatched (None) + pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.common import ( + deserializers_tensor, + serializers_tensor, + ) + + detections = { + "image": { + "width": 168, + "height": 192, + }, + "predictions": [ + { + "width": 1.0, + "height": 1.0, + "x": 1.5, + "y": 1.5, + "confidence": 0.25, + "class_id": 1, + "class": "cat", + "detection_id": "first", + "nearest_target_distance": 12.5, + }, + { + "width": 1.0, + "height": 1.0, + "x": 3.5, + "y": 3.5, + "confidence": 0.5, + "class_id": 2, + "class": "dog", + "detection_id": "second", + "nearest_target_distance": None, + }, + ], + } + + # when + result = deserializers_tensor.deserialize_detections_kind( + parameter="my_param", + detections=detections, + ) + + # then + assert [ + entry["nearest_target_distance"] for entry in result.bboxes_metadata + ] == [12.5, None] + serialized = serializers_tensor.serialise_sv_detections(result) + assert [ + prediction["nearest_target_distance"] + for prediction in serialized["predictions"] + ] == [12.5, None] diff --git a/tests/workflows/unit_tests/core_steps/common/test_rle_compact.py b/tests/workflows/unit_tests/core_steps/common/test_rle_compact.py new file mode 100644 index 0000000000..b48120976b --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/common/test_rle_compact.py @@ -0,0 +1,140 @@ +"""Tests for the zero-densification InstancesRLEMasks -> CompactMask transcode. + +The parity contract: `compact_mask_from_coco_rle` must produce a CompactMask +that decodes identically to +`CompactMask.from_dense(coco_rle_masks_to_numpy_mask(rle), xyxy, image_shape)`, +without ever building the dense (N, H, W) stack. +""" + +import numpy as np +import pytest +import torch + +pytest.importorskip( + "supervision.detection.compact_mask", + reason="supervision build without CompactMask (needs the compact-masks release)", +) + +from supervision.detection.compact_mask import CompactMask # noqa: E402 + +from inference.core.workflows.core_steps.common.rle_compact import ( # noqa: E402 + _decode_coco_counts, + compact_mask_from_coco_rle, + instances_rle_to_compact_mask, +) +from inference_models.models.base.types import InstancesRLEMasks # noqa: E402 +from inference_models.models.common.rle_utils import ( # noqa: E402 + coco_rle_masks_to_numpy_mask, + torch_mask_to_coco_rle, +) + + +def _encode(masks: np.ndarray) -> InstancesRLEMasks: + """(N, H, W) bool -> InstancesRLEMasks using the same encoder inference uses.""" + h, w = masks.shape[1], masks.shape[2] + counts = [ + torch_mask_to_coco_rle(torch.from_numpy(masks[i]))["counts"] + for i in range(masks.shape[0]) + ] + return InstancesRLEMasks(image_size=(h, w), masks=counts) + + +def _tight_box(mask: np.ndarray) -> np.ndarray: + ys, xs = np.where(mask) + if len(xs) == 0: + return np.array([0, 0, 0, 0], dtype=np.float32) + return np.array([xs.min(), ys.min(), xs.max(), ys.max()], dtype=np.float32) + + +def test_coco_counts_decoder_inverts_pycocotools(): + rng = np.random.default_rng(0) + for _ in range(100): + h = int(rng.integers(1, 40)) + w = int(rng.integers(1, 40)) + mask = (rng.random((h, w)) < rng.random()).astype(np.uint8) + compressed = torch_mask_to_coco_rle(torch.from_numpy(mask))["counts"] + + decoded = _decode_coco_counts(compressed) + + flat = np.zeros(h * w, dtype=bool) + pos = 0 + for idx, run_len in enumerate(decoded): + if idx % 2 == 1: + flat[pos : pos + run_len] = True + pos += run_len + assert pos == h * w + # COCO is column-major (F-order): reshape (w, h) then transpose. + np.testing.assert_array_equal(flat.reshape(w, h).T, mask.astype(bool)) + + +@pytest.mark.parametrize("seed", range(25)) +def test_parity_with_from_dense(seed: int): + rng = np.random.default_rng(seed) + h = int(rng.integers(8, 64)) + w = int(rng.integers(8, 64)) + n = int(rng.integers(0, 6)) + + masks = np.zeros((n, h, w), dtype=bool) + xyxy = np.zeros((n, 4), dtype=np.float32) + for i in range(n): + kind = rng.random() + if kind < 0.15: + pass # all-False + elif kind < 0.30: + masks[i, :, :] = True # all-True + else: + y1, y2 = sorted(rng.integers(0, h, size=2)) + x1, x2 = sorted(rng.integers(0, w, size=2)) + masks[i, y1 : y2 + 1, x1 : x2 + 1] = ( + rng.random((y2 - y1 + 1, x2 - x1 + 1)) < 0.6 + ) + box = _tight_box(masks[i]) + # Exercise clipping: degenerate and out-of-bounds boxes. + jitter = rng.random() + if jitter < 0.2: + box = np.array([box[0], box[1], box[0] - 1, box[1] - 1], dtype=np.float32) + elif jitter < 0.35: + box = box + np.array([-3, -3, 5, 5], dtype=np.float32) + xyxy[i] = box + + rle = _encode(masks) if n else InstancesRLEMasks((h, w), []) + dense = coco_rle_masks_to_numpy_mask(rle) if n else np.zeros((0, h, w), dtype=bool) + + reference = CompactMask.from_dense(dense, xyxy, (h, w)) + candidate = compact_mask_from_coco_rle((h, w), rle.masks, xyxy) + + np.testing.assert_array_equal(candidate.to_dense(), reference.to_dense()) + for i in range(n): + assert candidate[i].shape == (h, w) + np.testing.assert_array_equal(candidate[i], reference[i]) + + +def test_adapter_matches_full_frame_decode(): + rng = np.random.default_rng(123) + h, w, n = 80, 100, 4 + masks = np.zeros((n, h, w), dtype=bool) + xyxy = np.zeros((n, 4), dtype=np.float32) + for i in range(n): + x1 = int(rng.integers(0, w - 30)) + y1 = int(rng.integers(0, h - 30)) + masks[i, y1 : y1 + 20, x1 : x1 + 25] = True + xyxy[i] = _tight_box(masks[i]) + + rle = _encode(masks) + compact = instances_rle_to_compact_mask(rle, xyxy) + + assert isinstance(compact, CompactMask) + # The compact form decodes to exactly the full-frame dense masks ... + np.testing.assert_array_equal(compact.to_dense(), masks) + np.testing.assert_array_equal(compact.to_dense(), coco_rle_masks_to_numpy_mask(rle)) + # ... while storing only crop-area pixels, never the dense stack. + crop_px = int(np.prod(compact._crop_shapes, axis=1).sum()) + assert crop_px < n * h * w + + +def test_empty_masks(): + compact = compact_mask_from_coco_rle( + (50, 50), [], np.empty((0, 4), dtype=np.float32) + ) + assert isinstance(compact, CompactMask) + assert compact.to_dense().shape == (0, 50, 50) diff --git a/tests/workflows/unit_tests/core_steps/common/test_rle_embed.py b/tests/workflows/unit_tests/core_steps/common/test_rle_embed.py new file mode 100644 index 0000000000..f4d40e8f23 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/common/test_rle_embed.py @@ -0,0 +1,174 @@ +from typing import List, Tuple + +import numpy as np +import pytest +import torch + +from inference.core.workflows.core_steps.common.tensor_native import ( + embed_rle_masks_in_larger_canvas, +) +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import ( + coco_rle_masks_to_numpy_mask, + torch_mask_to_coco_rle, +) + + +def _encode_dense_slices(dense_slices: List[np.ndarray]) -> InstancesRLEMasks: + """Encode a list of dense (h, w) bool/uint8 slices through the real encoder.""" + assert len(dense_slices) > 0 + h, w = dense_slices[0].shape + encoded = [ + torch_mask_to_coco_rle(torch.from_numpy(np.ascontiguousarray(s).astype(bool))) + for s in dense_slices + ] + return InstancesRLEMasks.from_coco_rle_masks(image_size=(h, w), masks=encoded) + + +def _numpy_reference_canvas( + dense_slice: np.ndarray, + offset_xy: Tuple[int, int], + target_size_hw: Tuple[int, int], +) -> np.ndarray: + """All-zeros (H, W) canvas with the slice written at [y0:y0+h, x0:x0+w].""" + h, w = dense_slice.shape + x0, y0 = offset_xy + target_h, target_w = target_size_hw + canvas = np.zeros((target_h, target_w), dtype=bool) + canvas[y0 : y0 + h, x0 : x0 + w] = dense_slice.astype(bool) + return canvas + + +def _assert_embed_matches_reference( + dense_slices: List[np.ndarray], + offset_xy: Tuple[int, int], + target_size_hw: Tuple[int, int], +) -> None: + wrapped = _encode_dense_slices(dense_slices) + embedded = embed_rle_masks_in_larger_canvas( + masks=wrapped, offset_xy=offset_xy, target_size_hw=target_size_hw + ) + assert embedded.image_size == target_size_hw + assert len(embedded.masks) == len(dense_slices) + decoded = coco_rle_masks_to_numpy_mask(embedded) + assert decoded.shape == (len(dense_slices), target_size_hw[0], target_size_hw[1]) + for i, dense_slice in enumerate(dense_slices): + expected = _numpy_reference_canvas(dense_slice, offset_xy, target_size_hw) + np.testing.assert_array_equal(decoded[i], expected) + + +def test_embed_empty_masks_returns_empty_on_canvas() -> None: + wrapped = InstancesRLEMasks(image_size=(10, 12), masks=[]) + embedded = embed_rle_masks_in_larger_canvas( + masks=wrapped, offset_xy=(2, 3), target_size_hw=(20, 30) + ) + assert embedded.image_size == (20, 30) + assert embedded.masks == [] + decoded = coco_rle_masks_to_numpy_mask(embedded) + assert decoded.shape == (0, 20, 30) + + +def test_embed_single_instance_interior_placement() -> None: + rng = np.random.default_rng(seed=1) + dense = rng.integers(0, 2, size=(6, 8)).astype(bool) + _assert_embed_matches_reference([dense], offset_xy=(5, 4), target_size_hw=(20, 25)) + + +def test_embed_single_instance_top_left_origin() -> None: + rng = np.random.default_rng(seed=2) + dense = rng.integers(0, 2, size=(7, 9)).astype(bool) + _assert_embed_matches_reference([dense], offset_xy=(0, 0), target_size_hw=(15, 18)) + + +def test_embed_multiple_instances_interior() -> None: + rng = np.random.default_rng(seed=3) + dense_slices = [rng.integers(0, 2, size=(5, 6)).astype(bool) for _ in range(4)] + _assert_embed_matches_reference( + dense_slices, offset_xy=(3, 7), target_size_hw=(20, 16) + ) + + +def test_embed_full_zero_slice() -> None: + dense = np.zeros((4, 5), dtype=bool) + _assert_embed_matches_reference([dense], offset_xy=(2, 1), target_size_hw=(12, 14)) + + +def test_embed_full_one_slice() -> None: + dense = np.ones((4, 5), dtype=bool) + _assert_embed_matches_reference([dense], offset_xy=(2, 1), target_size_hw=(12, 14)) + + +def test_embed_slice_with_full_zero_and_full_one_columns() -> None: + # column 0 all zeros, column 1 all ones, then mixed columns. + dense = np.array( + [ + [0, 1, 1, 0], + [0, 1, 0, 1], + [0, 1, 1, 1], + ], + dtype=bool, + ) + _assert_embed_matches_reference([dense], offset_xy=(4, 2), target_size_hw=(10, 12)) + + +def test_embed_slice_fills_entire_canvas() -> None: + rng = np.random.default_rng(seed=4) + dense = rng.integers(0, 2, size=(8, 11)).astype(bool) + _assert_embed_matches_reference([dense], offset_xy=(0, 0), target_size_hw=(8, 11)) + + +def test_embed_slice_flush_to_bottom_right() -> None: + rng = np.random.default_rng(seed=5) + h, w = 6, 7 + target_h, target_w = 20, 18 + dense = rng.integers(0, 2, size=(h, w)).astype(bool) + _assert_embed_matches_reference( + [dense], + offset_xy=(target_w - w, target_h - h), + target_size_hw=(target_h, target_w), + ) + + +def test_embed_single_pixel_slice() -> None: + dense = np.ones((1, 1), dtype=bool) + _assert_embed_matches_reference([dense], offset_xy=(9, 13), target_size_hw=(30, 25)) + + +@pytest.mark.parametrize("seed", list(range(12))) +def test_embed_random_small_masks_match_reference(seed: int) -> None: + rng = np.random.default_rng(seed=seed) + h = int(rng.integers(1, 9)) + w = int(rng.integers(1, 9)) + target_h = h + int(rng.integers(0, 12)) + target_w = w + int(rng.integers(0, 12)) + x0 = int(rng.integers(0, target_w - w + 1)) + y0 = int(rng.integers(0, target_h - h + 1)) + n = int(rng.integers(1, 4)) + dense_slices = [rng.integers(0, 2, size=(h, w)).astype(bool) for _ in range(n)] + _assert_embed_matches_reference( + dense_slices, offset_xy=(x0, y0), target_size_hw=(target_h, target_w) + ) + + +def test_embed_raises_when_slice_does_not_fit_in_width() -> None: + wrapped = _encode_dense_slices([np.ones((4, 6), dtype=bool)]) + with pytest.raises(ValueError): + embed_rle_masks_in_larger_canvas( + masks=wrapped, offset_xy=(5, 0), target_size_hw=(10, 10) + ) + + +def test_embed_raises_when_slice_does_not_fit_in_height() -> None: + wrapped = _encode_dense_slices([np.ones((6, 4), dtype=bool)]) + with pytest.raises(ValueError): + embed_rle_masks_in_larger_canvas( + masks=wrapped, offset_xy=(0, 5), target_size_hw=(10, 10) + ) + + +def test_embed_raises_on_negative_offset() -> None: + wrapped = _encode_dense_slices([np.ones((4, 4), dtype=bool)]) + with pytest.raises(ValueError): + embed_rle_masks_in_larger_canvas( + masks=wrapped, offset_xy=(-1, 0), target_size_hw=(10, 10) + ) diff --git a/tests/workflows/unit_tests/core_steps/common/test_serializers.py b/tests/workflows/unit_tests/core_steps/common/test_serializers.py index 05098dd25d..7d27c21187 100644 --- a/tests/workflows/unit_tests/core_steps/common/test_serializers.py +++ b/tests/workflows/unit_tests/core_steps/common/test_serializers.py @@ -3,6 +3,7 @@ import cv2 import numpy as np +import pytest import supervision as sv from inference.core.workflows.core_steps.common.serializers import ( @@ -1218,3 +1219,388 @@ def test_serialise_rle_sv_detections_with_parent_origin() -> None: }, ], } + + +def test_serialise_native_classification_key_ordering_matches_numpy_path() -> None: + # given + torch = pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_native_classification, + ) + from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, + ) + + base_metadata = { + "class_names": {0: "cat", 1: "dog"}, + "prediction_type": "classification", + "image_dimensions": [480, 640], + "inference_id": "iid", + "parent_id": "p1", + "root_parent_id": "r1", + } + + def single_label(metadata: dict) -> "ClassificationPrediction": + return ClassificationPrediction( + class_id=torch.tensor([1]), + confidence=torch.tensor([[0.25, 0.5]], dtype=torch.float32), + images_metadata=[metadata], + ) + + def multi_label(metadata: dict) -> "MultiLabelClassificationPrediction": + return MultiLabelClassificationPrediction( + class_ids=torch.tensor([0, 1]), + confidence=torch.tensor([0.5, 0.25], dtype=torch.float32), + image_metadata=metadata, + ) + + # when + single_label_result = serialise_native_classification( + single_label(dict(base_metadata)) + ) + multi_label_result = serialise_native_classification( + multi_label(dict(base_metadata)) + ) + single_label_timed_result = serialise_native_classification( + single_label({**base_metadata, "time": 0.0123}) + ) + multi_label_timed_result = serialise_native_classification( + multi_label({**base_metadata, "time": 0.0123}) + ) + + # then - exact key order matters: orjson byte-parity with the numpy path depends on it + assert list(single_label_result.keys()) == [ + "inference_id", + "image", + "predictions", + "top", + "confidence", + "prediction_type", + "parent_id", + "root_parent_id", + ] + assert [list(e.keys()) for e in single_label_result["predictions"]] == [ + ["class", "class_id", "confidence"], + ["class", "class_id", "confidence"], + ] + assert single_label_result["predictions"][0] == { + "class": "dog", + "class_id": 1, + "confidence": 0.5, + } + assert list(multi_label_result.keys()) == [ + "inference_id", + "image", + "predictions", + "predicted_classes", + "prediction_type", + "parent_id", + "root_parent_id", + ] + assert [list(e.keys()) for e in multi_label_result["predictions"].values()] == [ + ["confidence", "class_id"], + ["confidence", "class_id"], + ] + assert list(single_label_timed_result.keys()) == [ + "inference_id", + "time", + "image", + "predictions", + "top", + "confidence", + "prediction_type", + "parent_id", + "root_parent_id", + ] + assert single_label_timed_result["time"] == 0.0123 + assert list(multi_label_timed_result.keys()) == [ + "inference_id", + "time", + "image", + "predictions", + "predicted_classes", + "prediction_type", + "parent_id", + "root_parent_id", + ] + assert multi_label_timed_result["time"] == 0.0123 + + +def test_tensor_wildcard_serializer_dispatches_native_values_like_kind_serializers() -> ( + None +): + # given + torch = pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.common import serializers_tensor + from inference_models.models.base.classification import ClassificationPrediction + from inference_models.models.base.instance_segmentation import InstanceDetections + from inference_models.models.base.keypoints_detection import KeyPoints + from inference_models.models.base.object_detection import Detections + from inference_models.models.base.types import InstancesRLEMasks + from inference_models.models.common.rle_utils import torch_mask_to_coco_rle + + od = Detections( + xyxy=torch.tensor([[10.0, 20.0, 30.0, 40.0]]), + class_id=torch.tensor([1]), + confidence=torch.tensor([0.5]), + image_metadata={"class_names": {1: "dog"}, "image_dimensions": [100, 200]}, + bboxes_metadata=[{"detection_id": "det-1"}], + ) + dense_mask = torch.zeros((1, 15, 15), dtype=torch.bool) + dense_mask[0, 2:6, 3:9] = True + instance_dense = InstanceDetections( + xyxy=torch.tensor([[3.0, 2.0, 9.0, 6.0]]), + class_id=torch.tensor([0]), + confidence=torch.tensor([0.9]), + mask=dense_mask, + image_metadata={"class_names": {0: "cat"}, "image_dimensions": [15, 15]}, + bboxes_metadata=[{"detection_id": "det-2"}], + ) + rle = torch_mask_to_coco_rle(dense_mask[0]) + instance_rle = InstanceDetections( + xyxy=torch.tensor([[3.0, 2.0, 9.0, 6.0]]), + class_id=torch.tensor([0]), + confidence=torch.tensor([0.9]), + mask=InstancesRLEMasks(image_size=(15, 15), masks=[rle["counts"]]), + image_metadata={"class_names": {0: "cat"}, "image_dimensions": [15, 15]}, + bboxes_metadata=[{"detection_id": "det-3"}], + ) + key_points = KeyPoints( + xy=torch.tensor([[[11.0, 11.0], [12.0, 13.0]]]), + class_id=torch.tensor([0]), + confidence=torch.tensor([[0.9, 0.8]]), + ) + kp_tuple = (key_points, od) + classification = ClassificationPrediction( + class_id=torch.tensor([1]), + confidence=torch.tensor([[0.25, 0.5]], dtype=torch.float32), + images_metadata=[ + { + "class_names": {0: "cat", 1: "dog"}, + "prediction_type": "classification", + "image_dimensions": [480, 640], + "inference_id": "iid", + } + ], + ) + bare_tensor = torch.tensor([[0.25, 0.5], [0.75, 1.0]]) + + # when + result = serializers_tensor.serialize_wildcard_kind( + value={ + "od": od, + "nested": [instance_dense, {"deeper": instance_rle}], + "kp": kp_tuple, + "cls": classification, + "tensor": bare_tensor, + "untouched": "text", + "number": 42, + "none": None, + "plain_tuple": (1, 2), + } + ) + + # then + assert result["od"] == serializers_tensor.serialise_sv_detections(od) + assert result["nested"][0] == serializers_tensor.serialise_sv_detections( + instance_dense + ) + assert result["nested"][1]["deeper"] == serializers_tensor.serialise_sv_detections( + instance_rle + ) + assert result["kp"] == serializers_tensor.serialise_native_keypoint_detection( + prediction=kp_tuple + ) + assert result["cls"] == serializers_tensor.serialise_native_classification( + prediction=classification + ) + assert result["tensor"] == [[0.25, 0.5], [0.75, 1.0]] + assert result["untouched"] == "text" + assert result["number"] == 42 + assert result["none"] is None + assert result["plain_tuple"] == (1, 2), "non-KP tuples pass through like numpy" + + +def test_tensor_wildcard_serializer_matches_numpy_wildcard_for_equivalent_prediction() -> ( + None +): + # given - the same logical OD prediction as sv (numpy path) and native (tensor path) + torch = pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.common import serializers_tensor + from inference_models.models.base.object_detection import Detections + + sv_detections = sv.Detections( + xyxy=np.array([[10.0, 20.0, 30.0, 40.0]], dtype=np.float64), + class_id=np.array([1]), + confidence=np.array([0.5], dtype=np.float64), + data={ + "class_name": np.array(["dog"]), + "detection_id": np.array(["det-1"]), + "image_dimensions": np.array([[100, 200]]), + }, + ) + native_detections = Detections( + xyxy=torch.tensor([[10.0, 20.0, 30.0, 40.0]]), + class_id=torch.tensor([1]), + confidence=torch.tensor([0.5]), + image_metadata={"class_names": {1: "dog"}, "image_dimensions": [100, 200]}, + bboxes_metadata=[{"detection_id": "det-1"}], + ) + + # when + numpy_result = serialize_wildcard_kind(value={"predictions": sv_detections}) + tensor_result = serializers_tensor.serialize_wildcard_kind( + value={"predictions": native_detections} + ) + + # then + assert numpy_result == tensor_result + + +def test_tensor_wildcard_serializer_keeps_numpy_behavior_for_legacy_values() -> None: + # given - sv.Detections + datetime + image reaching the tensor wildcard + pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.common import serializers_tensor + from inference.core.workflows.execution_engine.entities.base import VideoMetadata + + sv_detections = sv.Detections( + xyxy=np.array([[10.0, 20.0, 30.0, 40.0]], dtype=np.float64), + class_id=np.array([1]), + confidence=np.array([0.5], dtype=np.float64), + data={ + "class_name": np.array(["dog"]), + "detection_id": np.array(["det-1"]), + }, + ) + timestamp = datetime.now() + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="origin"), + numpy_image=np.zeros((10, 10, 3), dtype=np.uint8), + # explicit metadata: without it, serialisation mints frame_timestamp + # per call, breaking the two-call comparison below + video_metadata=VideoMetadata( + video_identifier="vid", + frame_number=0, + frame_timestamp=timestamp, + ), + ) + + # when + result = serializers_tensor.serialize_wildcard_kind( + value={"sv": sv_detections, "ts": timestamp, "img": image} + ) + + # then + expected = serialize_wildcard_kind( + value={"sv": sv_detections, "ts": timestamp, "img": image} + ) + assert result == expected + + +def test_tensor_serialise_sv_detections_skips_padded_keypoint_slots() -> None: + # given the padded per-box keypoint rows the sv -> native conversion carries + # (detection 0 has 2 real keypoints, detection 1 has 1 real + 1 padding slot) + torch = pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.common import serializers_tensor + from inference_models.models.base.object_detection import Detections + + native_detections = Detections( + xyxy=torch.tensor([[0.0, 0.0, 10.0, 10.0], [20.0, 20.0, 30.0, 30.0]]), + class_id=torch.tensor([0, 0]), + confidence=torch.tensor([0.5, 0.25]), + image_metadata={"class_names": {0: "obj"}, "image_dimensions": [192, 168]}, + bboxes_metadata=[ + { + "detection_id": "first", + "keypoints_class_id": np.array([0, 1], dtype=int), + "keypoints_class_name": np.array(["nose", "eye"], dtype=object), + "keypoints_confidence": np.array([0.5, 0.25], dtype=np.float32), + "keypoints_xy": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=np.float32), + }, + { + "detection_id": "second", + "keypoints_class_id": np.array([0, 0], dtype=int), + "keypoints_class_name": np.array(["nose", ""], dtype=object), + "keypoints_confidence": np.array([0.75, 0.0], dtype=np.float32), + "keypoints_xy": np.array([[21.0, 22.0], [0.0, 0.0]], dtype=np.float32), + }, + ], + ) + sv_detections = sv.Detections( + xyxy=np.array([[0, 0, 10, 10], [20, 20, 30, 30]], dtype=np.float64), + class_id=np.array([0, 0]), + confidence=np.array([0.5, 0.25], dtype=np.float64), + data={ + "class_name": np.array(["obj", "obj"]), + "detection_id": np.array(["first", "second"]), + "image_dimensions": np.array([[192, 168], [192, 168]]), + "keypoints_xy": np.array( + [[[1.0, 2.0], [3.0, 4.0]], [[21.0, 22.0], [0.0, 0.0]]], + dtype=np.float32, + ), + "keypoints_confidence": np.array( + [[0.5, 0.25], [0.75, 0.0]], dtype=np.float32 + ), + "keypoints_class_id": np.array([[0, 1], [0, 0]], dtype=int), + "keypoints_class_name": np.array( + [["nose", "eye"], ["nose", ""]], dtype=object + ), + }, + ) + + # when + result = serializers_tensor.serialise_sv_detections(native_detections) + + # then the padding slot must not surface as a fabricated keypoint + assert len(result["predictions"][0]["keypoints"]) == 2 + assert len(result["predictions"][1]["keypoints"]) == 1 + assert result["predictions"][1]["keypoints"][0]["class"] == "nose" + for prediction in result["predictions"]: + for keypoint in prediction["keypoints"]: + assert keypoint["class"] != "", "No empty-named padding keypoint may leak" + assert result == serialise_sv_detections(detections=sv_detections) + + +def test_tensor_serialise_sv_detections_with_nearest_target_distance() -> None: + # given: one detection with a real match distance, one unmatched (None) + torch = pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.common import serializers_tensor + from inference_models.models.base.object_detection import Detections + + native_detections = Detections( + xyxy=torch.tensor([[1.0, 1.0, 2.0, 2.0], [3.0, 3.0, 4.0, 4.0]]), + class_id=torch.tensor([1, 2]), + confidence=torch.tensor([0.25, 0.5]), + image_metadata={"class_names": {1: "cat", 2: "dog"}}, + bboxes_metadata=[ + {"detection_id": "first", "nearest_target_distance": 12.5}, + {"detection_id": "second", "nearest_target_distance": None}, + ], + ) + sv_detections = sv.Detections( + xyxy=np.array([[1, 1, 2, 2], [3, 3, 4, 4]], dtype=np.float64), + class_id=np.array([1, 2]), + confidence=np.array([0.25, 0.5], dtype=np.float64), + data={ + "class_name": np.array(["cat", "dog"]), + "detection_id": np.array(["first", "second"]), + "nearest_target_distance": np.array([12.5, None], dtype=object), + }, + ) + + # when + result = serializers_tensor.serialise_sv_detections(native_detections) + + # then + predictions = result["predictions"] + assert predictions[0]["nearest_target_distance"] == 12.5 + assert isinstance(predictions[0]["nearest_target_distance"], float) + assert predictions[1]["nearest_target_distance"] is None + assert result == serialise_sv_detections(detections=sv_detections) diff --git a/tests/workflows/unit_tests/core_steps/common/test_tensor_native_host_mirror.py b/tests/workflows/unit_tests/core_steps/common/test_tensor_native_host_mirror.py new file mode 100644 index 0000000000..b654d2bf6d --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/common/test_tensor_native_host_mirror.py @@ -0,0 +1,448 @@ +"""Tests for the per-box host mirror of the detection tensors +(``HOST_MIRROR_KEYS`` in ``tensor_native.py``). + +``attach_native_detection_metadata`` stores a host-side copy of each box's +``xyxy`` / ``class_id`` / ``confidence`` in ``bboxes_metadata`` (one batched +device->host read per tensor); visualization-phase consumers prefer it over +device reads. These tests prove: + +* the mirror values match the tensors for every prediction shape attach handles, +* the implementation is batched (dispatch-op count independent of N, no + ``.item()`` loops), +* geometry-mutating helpers drop the mirror (staleness contract), +* the wire format is byte-identical with and without the mirror, and the + deserializers never resurrect it. +""" + +import json + +import numpy as np +import pytest +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from inference.core.workflows.core_steps.common.deserializers_tensor import ( + deserialize_detections_kind, +) +from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_sv_detections, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + HOST_CLASS_ID_KEY, + HOST_CONFIDENCE_KEY, + HOST_MIRROR_KEYS, + HOST_XYXY_KEY, + attach_native_detection_metadata, + native_detections_to_root_coordinates, + read_host_mirror, + strip_host_mirror_metadata, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +CLASS_NAMES = {0: "cat", 1: "dog", 2: "goggles"} + + +class _OpAudit(TorchDispatchMode): + """Record every dispatched aten op (with tensor-argument shapes/dtypes).""" + + def __init__(self): + super().__init__() + self.ops = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + shapes = [] + + def _collect(value): + if isinstance(value, torch.Tensor): + shapes.append((tuple(value.shape), str(value.dtype))) + elif isinstance(value, (list, tuple)): + for item in value: + _collect(item) + + _collect(args) + _collect(list((kwargs or {}).values())) + self.ops.append((str(func), shapes)) + return func(*args, **(kwargs or {})) + + def op_names(self): + return [name for name, _ in self.ops] + + +def _image(height: int = 48, width: int = 64) -> WorkflowImageData: + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=np.zeros((height, width, 3), dtype=np.uint8), + ) + + +def _od_detections(n: int = 3, existing_metadata=None) -> Detections: + xyxy = np.asarray( + [ + [10.5 + index, 20.25 + index, 110.75 + index, 220.125 + index] + for index in range(n) + ], + dtype=np.float32, + ).reshape(n, 4) + return Detections( + xyxy=torch.tensor(xyxy, dtype=torch.float32), + class_id=torch.tensor([index % 3 for index in range(n)], dtype=torch.long), + confidence=torch.tensor( + np.linspace(0.42, 0.99, max(n, 1))[:n], dtype=torch.float32 + ), + image_metadata=None, + bboxes_metadata=existing_metadata, + ) + + +def _is_detections(n: int = 2) -> InstanceDetections: + masks = torch.zeros((n, 48, 64), dtype=torch.bool) + for index in range(n): + masks[index, 5 + index : 15 + index, 10 : 30 + index] = True + base = _od_detections(n) + return InstanceDetections( + xyxy=base.xyxy, + class_id=base.class_id, + confidence=base.confidence, + mask=masks, + image_metadata=None, + bboxes_metadata=None, + ) + + +def _attach(detections, prediction_type: str = "object-detection"): + return attach_native_detection_metadata( + detections=detections, + image=_image(), + class_names=CLASS_NAMES, + prediction_type=prediction_type, + ) + + +def _assert_mirror_matches_tensors(detections) -> None: + xyxy = detections.xyxy.detach().cpu().numpy() + class_id = detections.class_id.detach().cpu().numpy() + confidence = detections.confidence.detach().cpu().numpy() + assert detections.bboxes_metadata is not None + for index, entry in enumerate(detections.bboxes_metadata): + assert entry[HOST_XYXY_KEY] == pytest.approx(xyxy[index].tolist(), abs=0) + assert entry[HOST_CLASS_ID_KEY] == int(class_id[index]) + assert entry[HOST_CONFIDENCE_KEY] == float(confidence[index]) + assert isinstance(entry[HOST_XYXY_KEY], list) + assert len(entry[HOST_XYXY_KEY]) == 4 + assert all(isinstance(value, float) for value in entry[HOST_XYXY_KEY]) + assert isinstance(entry[HOST_CLASS_ID_KEY], int) + assert isinstance(entry[HOST_CONFIDENCE_KEY], float) + + +# --------------------------------------------------------------------------- # +# attach: mirror presence and values per prediction shape +# --------------------------------------------------------------------------- # + + +def test_attach_writes_host_mirror_for_object_detection() -> None: + # when + detections = _attach(_od_detections(3)) + + # then + _assert_mirror_matches_tensors(detections) + for entry in detections.bboxes_metadata: + assert DETECTION_ID_KEY in entry + + +def test_attach_writes_host_mirror_for_instance_segmentation() -> None: + # when + detections = _attach(_is_detections(2), prediction_type="instance-segmentation") + + # then + _assert_mirror_matches_tensors(detections) + + +def test_attach_writes_host_mirror_for_keypoint_bbox_component() -> None: + # given - the keypoint-detection blocks attach on the bbox component of the + # (KeyPoints, Detections) tuple, then add per-box keypoint payload keys. + detections = _attach(_od_detections(2), prediction_type="keypoint-detection") + key_points = KeyPoints( + xy=torch.zeros((2, 4, 2), dtype=torch.float32), + class_id=detections.class_id.clone(), + confidence=torch.zeros((2, 4), dtype=torch.float32), + image_metadata=detections.image_metadata, + ) + for entry in detections.bboxes_metadata: + entry["keypoints_xy"] = [[1.0, 2.0]] + + # then - the mirror survives the payload write and the tuple keeps it + prediction = (key_points, detections) + _assert_mirror_matches_tensors(prediction[1]) + + +def test_attach_preserves_existing_keys_and_overwrites_stale_mirror() -> None: + # given - metadata carrying a model-set key and a STALE mirror + existing = [ + {"text": "reading", HOST_XYXY_KEY: [0.0, 0.0, 0.0, 0.0]}, + {"text": "glasses", HOST_CLASS_ID_KEY: 999}, + ] + detections = _od_detections(2, existing_metadata=existing) + + # when + detections = _attach(detections) + + # then - custom keys preserved, mirror refreshed from the tensors + assert [entry["text"] for entry in detections.bboxes_metadata] == [ + "reading", + "glasses", + ] + _assert_mirror_matches_tensors(detections) + + +def test_attach_empty_prediction_has_no_metadata() -> None: + # when + detections = _attach(_od_detections(0)) + + # then + assert detections.bboxes_metadata is None + + +# --------------------------------------------------------------------------- # +# attach: batched implementation (no per-box .item() loops) +# --------------------------------------------------------------------------- # + + +def _attach_op_names(n: int): + detections = _od_detections(n) + audit = _OpAudit() + with audit: + attach_native_detection_metadata( + detections=detections, + image=_image(), + class_names=CLASS_NAMES, + prediction_type="object-detection", + ) + return audit.op_names() + + +def test_attach_reads_are_batched_not_per_box() -> None: + # when + ops_small = _attach_op_names(2) + ops_large = _attach_op_names(16) + + # then - per-box `.item()` would dispatch `_local_scalar_dense` N times and + # scale the op count with N; the batched implementation is N-independent. + for ops in (ops_small, ops_large): + assert not any("_local_scalar_dense" in name for name in ops) + assert not any("nonzero" in name for name in ops) + assert sum("_to_copy" in name for name in ops) <= 3 + assert ops_small == ops_large + + +@pytest.mark.skipif(not torch.backends.mps.is_available(), reason="requires MPS") +def test_attach_makes_exactly_one_transfer_per_tensor_cross_device() -> None: + # given + detections = _od_detections(4) + detections.xyxy = detections.xyxy.to("mps") + detections.class_id = detections.class_id.to("mps") + detections.confidence = detections.confidence.to("mps") + + # when + audit = _OpAudit() + with audit: + attach_native_detection_metadata( + detections=detections, + image=_image(), + class_names=CLASS_NAMES, + prediction_type="object-detection", + ) + + # then - one device->host copy per tensor, nothing per box + assert sum("_to_copy" in name for name in audit.op_names()) == 3 + _assert_mirror_matches_tensors(detections) + + +# --------------------------------------------------------------------------- # +# read_host_mirror / strip_host_mirror_metadata +# --------------------------------------------------------------------------- # + + +def test_read_host_mirror_assembles_arrays_matching_tensor_reads() -> None: + # given + detections = _attach(_od_detections(3)) + + # when + mirror = read_host_mirror(detections.bboxes_metadata, 3) + + # then + assert mirror is not None + xyxy, class_id, confidence = mirror + expected_xyxy = detections.xyxy.detach().cpu().numpy().astype(np.float32) + expected_class_id = detections.class_id.detach().cpu().numpy().astype(int) + expected_confidence = ( + detections.confidence.detach().cpu().numpy().astype(np.float32) + ) + assert np.array_equal(xyxy, expected_xyxy) and xyxy.dtype == expected_xyxy.dtype + assert ( + np.array_equal(class_id, expected_class_id) + and class_id.dtype == expected_class_id.dtype + ) + assert ( + np.array_equal(confidence, expected_confidence) + and confidence.dtype == expected_confidence.dtype + ) + + +@pytest.mark.parametrize("missing_key", HOST_MIRROR_KEYS) +def test_read_host_mirror_requires_every_box_to_carry_all_keys(missing_key) -> None: + # given + detections = _attach(_od_detections(3)) + del detections.bboxes_metadata[1][missing_key] + + # when / then + assert read_host_mirror(detections.bboxes_metadata, 3) is None + + +def test_read_host_mirror_rejects_row_count_mismatch_and_none() -> None: + detections = _attach(_od_detections(3)) + assert read_host_mirror(detections.bboxes_metadata, 2) is None + assert read_host_mirror(None, 3) is None + assert read_host_mirror([], 0) is None + + +def test_strip_host_mirror_metadata_removes_keys_without_mutating_input() -> None: + # given + detections = _attach(_od_detections(2)) + original_entries = detections.bboxes_metadata + + # when + stripped = strip_host_mirror_metadata(original_entries) + + # then + assert stripped is not None and len(stripped) == 2 + for entry in stripped: + assert not any(key in entry for key in HOST_MIRROR_KEYS) + assert DETECTION_ID_KEY in entry + # caller-shared dicts are untouched + for entry in original_entries: + assert all(key in entry for key in HOST_MIRROR_KEYS) + assert strip_host_mirror_metadata(None) is None + # mirror-free entries are reused as-is + plain = [{"a": 1}] + assert strip_host_mirror_metadata(plain)[0] is plain[0] + + +# --------------------------------------------------------------------------- # +# staleness contract: root-coordinate shift drops the mirror +# --------------------------------------------------------------------------- # + + +def test_root_coordinate_shift_drops_host_mirror() -> None: + # given - a crop-anchored prediction with a fresh mirror + detections = _attach(_od_detections(2)) + detections.image_metadata = { + CLASS_NAMES_KEY: CLASS_NAMES, + PREDICTION_TYPE_KEY: "object-detection", + IMAGE_DIMENSIONS_KEY: [48, 64], + PARENT_ID_KEY: "crop-1", + PARENT_COORDINATES_KEY: [100, 50], + PARENT_DIMENSIONS_KEY: [480, 640], + ROOT_PARENT_ID_KEY: "root", + ROOT_PARENT_COORDINATES_KEY: [100, 50], + ROOT_PARENT_DIMENSIONS_KEY: [480, 640], + } + + # when + shifted = native_detections_to_root_coordinates(detections) + + # then - xyxy moved, mirror dropped (not stale-carried) + assert not torch.equal(shifted.xyxy, detections.xyxy) + for entry in shifted.bboxes_metadata: + assert not any(key in entry for key in HOST_MIRROR_KEYS) + assert DETECTION_ID_KEY in entry + + +# --------------------------------------------------------------------------- # +# wire format: the mirror never reaches serialized output +# --------------------------------------------------------------------------- # + + +def _assert_no_mirror_key_anywhere(payload) -> None: + if isinstance(payload, dict): + for key, value in payload.items(): + assert key not in HOST_MIRROR_KEYS + _assert_no_mirror_key_anywhere(value) + elif isinstance(payload, list): + for item in payload: + _assert_no_mirror_key_anywhere(item) + + +def test_wire_serialization_is_byte_identical_with_and_without_mirror() -> None: + # given - the same prediction with and without the mirror (identical + # detection ids, tensors shared) + mirrored = _attach(_od_detections(3)) + stripped = Detections( + xyxy=mirrored.xyxy, + class_id=mirrored.class_id, + confidence=mirrored.confidence, + image_metadata=mirrored.image_metadata, + bboxes_metadata=strip_host_mirror_metadata(mirrored.bboxes_metadata), + ) + + # when + serialized_mirrored = serialise_sv_detections(mirrored) + serialized_stripped = serialise_sv_detections(stripped) + + # then - byte-identical JSON, and no mirror key anywhere in the payload + assert json.dumps(serialized_mirrored) == json.dumps(serialized_stripped) + _assert_no_mirror_key_anywhere(serialized_mirrored) + + +def test_wire_serialization_is_byte_identical_for_instance_segmentation() -> None: + # given + mirrored = _attach(_is_detections(2), prediction_type="instance-segmentation") + stripped = InstanceDetections( + xyxy=mirrored.xyxy, + class_id=mirrored.class_id, + confidence=mirrored.confidence, + mask=mirrored.mask, + image_metadata=mirrored.image_metadata, + bboxes_metadata=strip_host_mirror_metadata(mirrored.bboxes_metadata), + ) + + # when + serialized_mirrored = serialise_sv_detections(mirrored) + serialized_stripped = serialise_sv_detections(stripped) + + # then + assert json.dumps(serialized_mirrored) == json.dumps(serialized_stripped) + _assert_no_mirror_key_anywhere(serialized_mirrored) + + +def test_deserializer_does_not_resurrect_the_mirror() -> None: + # given - wire payload produced from a mirrored prediction + serialized = serialise_sv_detections(_attach(_od_detections(3))) + + # when + deserialized = deserialize_detections_kind( + parameter="predictions", detections=serialized + ) + + # then - a deserialized prediction is mirror-less (consumers fall back to + # tensor reads on it) + assert deserialized.bboxes_metadata is not None + for entry in deserialized.bboxes_metadata: + assert not any(key in entry for key in HOST_MIRROR_KEYS) diff --git a/tests/workflows/unit_tests/core_steps/common/test_tensor_native_root_coordinates.py b/tests/workflows/unit_tests/core_steps/common/test_tensor_native_root_coordinates.py new file mode 100644 index 0000000000..7d0eb6ea51 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/common/test_tensor_native_root_coordinates.py @@ -0,0 +1,389 @@ +"""Tests for the tensor-native root-coordinate conversion +(``native_detections_to_root_coordinates``) - especially the mask re-anchoring +and per-box geometry-payload shifts that mirror numpy's +``sv_detections_to_root_coordinates``.""" + +import numpy as np +import pytest +import supervision as sv +import torch +from pycocotools import mask as mask_utils +from supervision.config import ORIENTED_BOX_COORDINATES + +from inference.core.workflows.core_steps.common.serializers import ( + serialise_sv_detections as numpy_serialise_sv_detections, +) +from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_sv_detections as tensor_serialise_sv_detections, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + native_detections_to_root_coordinates, +) +from inference.core.workflows.core_steps.common.utils import ( + sv_detections_to_root_coordinates, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + POLYGON_KEY_IN_SV_DETECTIONS, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks + +CROP_OFFSET_X, CROP_OFFSET_Y = 100, 50 +CROP_H, CROP_W = 40, 60 +ROOT_H, ROOT_W = 480, 640 + + +def _crop_local_dense_masks() -> np.ndarray: + masks = np.zeros((2, CROP_H, CROP_W), dtype=bool) + masks[0, 5:15, 10:30] = True + masks[1, 20:35, 40:55] = True + return masks + + +def _native_image_metadata() -> dict: + return { + CLASS_NAMES_KEY: {0: "a", 1: "b"}, + PREDICTION_TYPE_KEY: "instance-segmentation", + IMAGE_DIMENSIONS_KEY: [CROP_H, CROP_W], + PARENT_ID_KEY: "crop-1", + PARENT_COORDINATES_KEY: [CROP_OFFSET_X, CROP_OFFSET_Y], + PARENT_DIMENSIONS_KEY: [ROOT_H, ROOT_W], + ROOT_PARENT_ID_KEY: "root-image", + ROOT_PARENT_COORDINATES_KEY: [CROP_OFFSET_X, CROP_OFFSET_Y], + ROOT_PARENT_DIMENSIONS_KEY: [ROOT_H, ROOT_W], + } + + +def _polygons() -> np.ndarray: + # Uniform 4-point polygons so the numpy path's vectorized + # `data[POLYGON] + shift` broadcast applies cleanly. + return np.array( + [ + [[10, 5], [30, 5], [30, 15], [10, 15]], + [[40, 20], [55, 20], [55, 35], [40, 35]], + ], + dtype=np.int64, + ) + + +def _native_instance_detections( + mask: object, with_geometry_payloads: bool = True +) -> InstanceDetections: + polygons = _polygons() + bboxes_metadata = [] + for index in range(2): + entry = {DETECTION_ID_KEY: f"d{index}"} + if with_geometry_payloads: + entry[POLYGON_KEY_IN_SV_DETECTIONS] = polygons[index] + entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] = [ + [11.0, 6.0], + [29.0, 14.0], + ] + entry[ORIENTED_BOX_COORDINATES] = np.array( + [[10.0, 5.0], [30.0, 5.0], [30.0, 15.0], [10.0, 15.0]] + ) + bboxes_metadata.append(entry) + return InstanceDetections( + xyxy=torch.tensor( + [[10.0, 5.0, 30.0, 15.0], [40.0, 20.0, 55.0, 35.0]], dtype=torch.float32 + ), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.5, 0.25], dtype=torch.float32), + mask=mask, + image_metadata=_native_image_metadata(), + bboxes_metadata=bboxes_metadata, + ) + + +def _dense_masks_to_rle(masks: np.ndarray) -> InstancesRLEMasks: + encoded = [ + mask_utils.encode(np.asfortranarray(mask.astype(np.uint8)))["counts"] + for mask in masks + ] + return InstancesRLEMasks(image_size=(CROP_H, CROP_W), masks=encoded) + + +def _expected_root_masks() -> np.ndarray: + anchored = np.zeros((2, ROOT_H, ROOT_W), dtype=bool) + anchored[ + :, + CROP_OFFSET_Y : CROP_OFFSET_Y + CROP_H, + CROP_OFFSET_X : CROP_OFFSET_X + CROP_W, + ] = _crop_local_dense_masks() + return anchored + + +def _decode_native_masks(mask: object) -> np.ndarray: + if isinstance(mask, InstancesRLEMasks): + return np.stack( + [ + mask_utils.decode({"size": list(mask.image_size), "counts": counts}) + for counts in mask.masks + ] + ).astype(bool) + return mask.detach().cpu().numpy().astype(bool) + + +def test_root_conversion_re_anchors_dense_masks() -> None: + # given + detections = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + + # when + result = native_detections_to_root_coordinates(prediction=detections) + + # then + assert tuple(result.mask.shape) == (2, ROOT_H, ROOT_W) + assert result.mask.dtype == torch.bool + assert np.array_equal(_decode_native_masks(result.mask), _expected_root_masks()) + assert torch.allclose( + result.xyxy, + torch.tensor( + [[110.0, 55.0, 130.0, 65.0], [140.0, 70.0, 155.0, 85.0]], + dtype=torch.float32, + ), + ) + # source object is untouched (conversion returns a copy) + assert tuple(detections.mask.shape) == (2, CROP_H, CROP_W) + assert detections.image_metadata[ROOT_PARENT_COORDINATES_KEY] == [ + CROP_OFFSET_X, + CROP_OFFSET_Y, + ] + + +def test_root_conversion_re_anchors_rle_masks() -> None: + # given + detections = _native_instance_detections( + mask=_dense_masks_to_rle(_crop_local_dense_masks()) + ) + + # when + result = native_detections_to_root_coordinates(prediction=detections) + + # then + assert isinstance(result.mask, InstancesRLEMasks) + assert tuple(result.mask.image_size) == (ROOT_H, ROOT_W) + assert np.array_equal(_decode_native_masks(result.mask), _expected_root_masks()) + + +def test_root_conversion_shifts_geometry_payloads_including_obb() -> None: + # given + detections = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + + # when + result = native_detections_to_root_coordinates(prediction=detections) + + # then + entry = result.bboxes_metadata[0] + shifted_polygon = entry[POLYGON_KEY_IN_SV_DETECTIONS] + assert isinstance(shifted_polygon, np.ndarray) + assert np.issubdtype(shifted_polygon.dtype, np.integer) + assert np.array_equal( + shifted_polygon, + _polygons()[0] + np.array([CROP_OFFSET_X, CROP_OFFSET_Y]), + ) + shifted_keypoints = entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] + assert isinstance(shifted_keypoints, list) + assert shifted_keypoints == [ + [11.0 + CROP_OFFSET_X, 6.0 + CROP_OFFSET_Y], + [29.0 + CROP_OFFSET_X, 14.0 + CROP_OFFSET_Y], + ] + # OBB corners shifted like every other geometry payload (the crop side + # subtracts the crop origin from them, so root conversion adds it back) + shifted_obb = entry[ORIENTED_BOX_COORDINATES] + assert isinstance(shifted_obb, np.ndarray) + assert np.issubdtype(shifted_obb.dtype, np.floating) + assert np.array_equal( + shifted_obb, + np.array([[10.0, 5.0], [30.0, 5.0], [30.0, 15.0], [10.0, 15.0]]) + + np.array([CROP_OFFSET_X, CROP_OFFSET_Y]), + ) + # source metadata untouched + assert np.array_equal( + detections.bboxes_metadata[0][POLYGON_KEY_IN_SV_DETECTIONS], _polygons()[0] + ) + assert np.array_equal( + detections.bboxes_metadata[0][ORIENTED_BOX_COORDINATES], + np.array([[10.0, 5.0], [30.0, 5.0], [30.0, 15.0], [10.0, 15.0]]), + ) + + +def test_root_conversion_with_zero_offset_still_reanchors_crop_predictions() -> None: + # given - a crop taken at (0, 0): the offset is zero, but the dimensions are + # crop-sized and the parent is NOT the root. A zero offset alone must not be + # treated as proof of root anchoring (PR review): masks still need + # re-embedding onto the root canvas and the metadata still needs the root + # rewrite. + metadata = _native_image_metadata() + metadata[PARENT_COORDINATES_KEY] = [0, 0] + metadata[ROOT_PARENT_COORDINATES_KEY] = [0, 0] + detections = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + detections.image_metadata = metadata + + # when + result = native_detections_to_root_coordinates(prediction=detections) + + # then - boxes keep their values (the shift IS zero) but the masks are + # re-anchored onto the root canvas and the lineage collapses to the root + assert result is not detections + assert torch.equal(result.xyxy, detections.xyxy) + assert tuple(result.mask.shape[-2:]) == (ROOT_H, ROOT_W) + assert list(result.image_metadata[IMAGE_DIMENSIONS_KEY]) == [ROOT_H, ROOT_W] + assert result.image_metadata[PARENT_ID_KEY] == "root-image" + + +def test_root_conversion_is_noop_when_already_root_anchored() -> None: + # given - genuinely root-anchored: zero offset AND root-sized dimensions AND + # the parent IS the root + metadata = _native_image_metadata() + metadata[IMAGE_DIMENSIONS_KEY] = [ROOT_H, ROOT_W] + metadata[PARENT_ID_KEY] = "root-image" + metadata[PARENT_COORDINATES_KEY] = [0, 0] + metadata[ROOT_PARENT_COORDINATES_KEY] = [0, 0] + detections = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + detections.image_metadata = metadata + + # when + result = native_detections_to_root_coordinates(prediction=detections) + + # then - identity: the input object is returned untouched + assert result is detections + + +def test_root_conversion_raises_on_masks_without_root_dimensions() -> None: + # given + metadata = _native_image_metadata() + del metadata[ROOT_PARENT_DIMENSIONS_KEY] + detections = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + detections.image_metadata = metadata + + # when / then + with pytest.raises(ValueError, match="root"): + _ = native_detections_to_root_coordinates(prediction=detections) + + +def test_root_conversion_raises_when_masks_do_not_fit_root_canvas() -> None: + # given + metadata = _native_image_metadata() + metadata[ROOT_PARENT_DIMENSIONS_KEY] = [60, 120] # crop at (100, 50) cannot fit + detections = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + detections.image_metadata = metadata + + # when / then + with pytest.raises(ValueError, match="fit"): + _ = native_detections_to_root_coordinates(prediction=detections) + + +def test_plain_detections_root_conversion_shifts_polygon_payloads() -> None: + # given - OD (no masks) prediction carrying a declared polygon payload + detections = Detections( + xyxy=torch.tensor([[10.0, 5.0, 30.0, 15.0]], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([0.5], dtype=torch.float32), + image_metadata=_native_image_metadata(), + bboxes_metadata=[ + { + DETECTION_ID_KEY: "d0", + POLYGON_KEY_IN_SV_DETECTIONS: [[10, 5], [30, 5], [30, 15]], + } + ], + ) + + # when + result = native_detections_to_root_coordinates(prediction=detections) + + # then - list container and integer coordinates preserved + assert result.bboxes_metadata[0][POLYGON_KEY_IN_SV_DETECTIONS] == [ + [110, 55], + [130, 55], + [130, 65], + ] + + +def _sv_equivalent_detections() -> sv.Detections: + polygons = _polygons() + return sv.Detections( + xyxy=np.array( + [[10.0, 5.0, 30.0, 15.0], [40.0, 20.0, 55.0, 35.0]], dtype=np.float32 + ), + class_id=np.array([0, 1]), + confidence=np.array([0.5, 0.25], dtype=np.float32), + mask=_crop_local_dense_masks(), + data={ + "class_name": np.array(["a", "b"]), + DETECTION_ID_KEY: np.array(["d0", "d1"]), + PARENT_ID_KEY: np.array(["crop-1"] * 2), + PARENT_COORDINATES_KEY: np.array([[CROP_OFFSET_X, CROP_OFFSET_Y]] * 2), + PARENT_DIMENSIONS_KEY: np.array([[ROOT_H, ROOT_W]] * 2), + ROOT_PARENT_ID_KEY: np.array(["root-image"] * 2), + ROOT_PARENT_COORDINATES_KEY: np.array([[CROP_OFFSET_X, CROP_OFFSET_Y]] * 2), + ROOT_PARENT_DIMENSIONS_KEY: np.array([[ROOT_H, ROOT_W]] * 2), + IMAGE_DIMENSIONS_KEY: np.array([[CROP_H, CROP_W]] * 2), + POLYGON_KEY_IN_SV_DETECTIONS: polygons, + ORIENTED_BOX_COORDINATES: np.array( + [[[10.0, 5.0], [30.0, 5.0], [30.0, 15.0], [10.0, 15.0]]] * 2 + ), + }, + ) + + +def test_root_conversion_parity_with_numpy_path() -> None: + """End-to-end parity: the same crop-local prediction expressed as + sv.Detections vs native InstanceDetections must root-convert to the same + boxes, masks, polygons AND the same serialized response dict.""" + # given + sv_detections = _sv_equivalent_detections() + native_detections = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()), + with_geometry_payloads=True, + ) + # numpy path has no keypoints in this scenario; strip them from the native + # side too so both serialize the same payload set (OBB stays on BOTH sides - + # both root conversions shift it now) + for entry in native_detections.bboxes_metadata: + del entry[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS] + + # when + sv_root = sv_detections_to_root_coordinates(detections=sv_detections) + native_root = native_detections_to_root_coordinates(prediction=native_detections) + + # then - geometry parity + assert np.allclose(sv_root.xyxy, native_root.xyxy.detach().cpu().numpy()) + assert np.array_equal(sv_root.mask, _decode_native_masks(native_root.mask)) + for index in range(2): + assert np.array_equal( + sv_root.data[POLYGON_KEY_IN_SV_DETECTIONS][index], + native_root.bboxes_metadata[index][POLYGON_KEY_IN_SV_DETECTIONS], + ) + assert np.array_equal( + sv_root.data[ORIENTED_BOX_COORDINATES][index], + native_root.bboxes_metadata[index][ORIENTED_BOX_COORDINATES], + ) + + # then - serialized-output parity + numpy_serialized = numpy_serialise_sv_detections(sv_root) + tensor_serialized = tensor_serialise_sv_detections(native_root) + assert numpy_serialized == tensor_serialized diff --git a/tests/workflows/unit_tests/core_steps/common/test_tensor_native_selection.py b/tests/workflows/unit_tests/core_steps/common/test_tensor_native_selection.py new file mode 100644 index 0000000000..7c44618e10 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/common/test_tensor_native_selection.py @@ -0,0 +1,182 @@ +"""Behavioural contract of the tensor-native row-selection helpers +(`take_prediction_by_indices` / `take_prediction_by_mask`): torch masks select +tensor fields on-device with no host round trip, and every entry style +(index list, python mask, numpy mask, torch mask) yields identical results.""" + +import numpy as np +import pytest +import torch + +from inference.core.workflows.core_steps.common.tensor_native import ( + take_prediction_by_indices, + take_prediction_by_mask, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks + + +def _detections() -> Detections: + return Detections( + xyxy=torch.tensor( + [[0, 0, 10, 10], [10, 10, 20, 20], [20, 20, 30, 30]], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 1, 2]), + confidence=torch.tensor([0.9, 0.5, 0.7]), + image_metadata={"class_names": {0: "a", 1: "b", 2: "c"}}, + bboxes_metadata=[{"detection_id": f"d{i}"} for i in range(3)], + ) + + +def _instance_detections(rle: bool = False) -> InstanceDetections: + base = _detections() + if rle: + mask = InstancesRLEMasks( + image_size=(30, 30), + masks=[{"size": [30, 30], "counts": f"stub-{i}"} for i in range(3)], + ) + else: + dense = torch.zeros((3, 30, 30), dtype=torch.bool) + for i in range(3): + dense[i, i * 10 : (i + 1) * 10, i * 10 : (i + 1) * 10] = True + mask = dense + return InstanceDetections( + xyxy=base.xyxy, + class_id=base.class_id, + confidence=base.confidence, + mask=mask, + image_metadata=base.image_metadata, + bboxes_metadata=base.bboxes_metadata, + ) + + +def test_torch_mask_selection_matches_every_other_entry_style() -> None: + # given + detections = _instance_detections() + torch_mask = torch.tensor([True, False, True]) + + # when + by_torch_mask = take_prediction_by_mask(detections, torch_mask) + by_list_mask = take_prediction_by_mask(detections, [True, False, True]) + by_numpy_mask = take_prediction_by_mask(detections, np.array([True, False, True])) + by_indices = take_prediction_by_indices(detections, [0, 2]) + + # then - all four styles produce identical selections + for result in (by_torch_mask, by_list_mask, by_numpy_mask, by_indices): + assert torch.equal(result.xyxy, detections.xyxy[[0, 2]]) + assert torch.equal(result.class_id, torch.tensor([0, 2])) + assert torch.equal(result.mask, detections.mask[[0, 2]]) + assert [m["detection_id"] for m in result.bboxes_metadata] == ["d0", "d2"] + + +def test_torch_mask_selection_copies_metadata_dicts() -> None: + # given + detections = _detections() + + # when + result = take_prediction_by_mask(detections, torch.tensor([True, True, False])) + result.bboxes_metadata[0]["tracker_id"] = 7 + + # then - the source metadata must not observe downstream mutation + assert "tracker_id" not in detections.bboxes_metadata[0] + + +def test_torch_mask_identity_aliases_tensors() -> None: + # given + detections = _instance_detections() + + # when - all-True mask is the identity fast path + result = take_prediction_by_mask(detections, torch.tensor([True, True, True])) + + # then - tensor fields are aliased, not copied + assert result.xyxy is detections.xyxy + assert result.mask is detections.mask + + +def test_torch_mask_selection_handles_rle_masks() -> None: + # given + detections = _instance_detections(rle=True) + + # when + result = take_prediction_by_mask(detections, torch.tensor([False, True, True])) + + # then - the RLE list follows the surviving rows + assert [m["counts"] for m in result.mask.masks] == ["stub-1", "stub-2"] + + +def test_torch_mask_selection_on_keypoint_tuple_preserves_auxiliary_tensors() -> None: + # given - RF-DETR-shaped keypoints with covariance / detection_confidence + key_points = KeyPoints( + xy=torch.arange(3 * 2 * 2, dtype=torch.float32).reshape(3, 2, 2), + class_id=torch.tensor([0, 1, 2]), + confidence=torch.full((3, 2), 0.5), + image_metadata={"class_names": {0: "p"}}, + key_points_metadata=[{"detection_id": f"k{i}"} for i in range(3)], + covariance=torch.arange(3 * 2 * 2 * 2, dtype=torch.float32).reshape(3, 2, 2, 2), + detection_confidence=torch.tensor([0.1, 0.2, 0.3]), + ) + detections = _detections() + + # when + kp, det = take_prediction_by_mask( + (key_points, detections), torch.tensor([True, False, True]) + ) + + # then - both tuple components sliced consistently, aux tensors in lockstep + assert torch.equal(kp.xy, key_points.xy[[0, 2]]) + assert torch.equal(kp.covariance, key_points.covariance[[0, 2]]) + assert torch.equal(kp.detection_confidence, torch.tensor([0.1, 0.3])) + assert torch.equal(det.xyxy, detections.xyxy[[0, 2]]) + assert [m["detection_id"] for m in kp.key_points_metadata] == ["k0", "k2"] + + +def test_mismatched_length_torch_mask_keeps_legacy_index_semantics() -> None: + # given - a shorter mask: its nonzero positions are treated as indices + detections = _detections() + + # when + result = take_prediction_by_mask(detections, torch.tensor([False, True])) + + # then + assert torch.equal(result.xyxy, detections.xyxy[[1]]) + assert [m["detection_id"] for m in result.bboxes_metadata] == ["d1"] + + +def test_index_selection_can_reorder_rows() -> None: + # given + detections = _detections() + + # when - masks cannot reorder, index lists can + result = take_prediction_by_indices(detections, [2, 0]) + + # then + assert torch.equal(result.class_id, torch.tensor([2, 0])) + assert [m["detection_id"] for m in result.bboxes_metadata] == ["d2", "d0"] + + +@pytest.mark.skipif( + not torch.backends.mps.is_available(), reason="device-preservation check needs MPS" +) +def test_torch_mask_selection_stays_on_device() -> None: + # given - prediction and mask both device-resident + detections = _detections() + device = torch.device("mps") + on_device = Detections( + xyxy=detections.xyxy.to(device), + class_id=detections.class_id.to(device), + confidence=detections.confidence.to(device), + image_metadata=detections.image_metadata, + bboxes_metadata=detections.bboxes_metadata, + ) + + # when + result = take_prediction_by_mask( + on_device, torch.tensor([True, False, True], device=device) + ) + + # then - results stay on the device + assert result.xyxy.device.type == "mps" + assert result.class_id.device.type == "mps" + assert [m["detection_id"] for m in result.bboxes_metadata] == ["d0", "d2"] diff --git a/tests/workflows/unit_tests/core_steps/common/test_utils.py b/tests/workflows/unit_tests/core_steps/common/test_utils.py index 45e7d34a10..94bf38662e 100644 --- a/tests/workflows/unit_tests/core_steps/common/test_utils.py +++ b/tests/workflows/unit_tests/core_steps/common/test_utils.py @@ -3,6 +3,7 @@ import numpy as np import pytest import supervision as sv +from supervision.config import ORIENTED_BOX_COORDINATES from inference.core.workflows.core_steps.common.utils import ( add_inference_keypoints_to_sv_detections, @@ -476,6 +477,12 @@ def test_sv_detections_to_root_coordinates_when_shift_is_needed() -> None: [[50, 125], [100, 125], [100, 225], [50, 225]], ] ), + ORIENTED_BOX_COORDINATES: np.array( + [ + [[25.0, 50.0], [75.0, 50.0], [75.0, 150.0], [25.0, 150.0]], + [[50.0, 125.0], [100.0, 125.0], [100.0, 225.0], [50.0, 225.0]], + ] + ), }, ) @@ -572,6 +579,25 @@ def test_sv_detections_to_root_coordinates_when_shift_is_needed() -> None: ] ), ), "Expected polygon metadata to be shifted into root coordinates" + assert np.allclose( + result[ORIENTED_BOX_COORDINATES], + np.array( + [ + [ + [50 + 25.0, 100 + 50.0], + [50 + 75.0, 100 + 50.0], + [50 + 75.0, 100 + 150.0], + [50 + 25.0, 100 + 150.0], + ], + [ + [50 + 50.0, 100 + 125.0], + [50 + 100.0, 100 + 125.0], + [50 + 100.0, 100 + 225.0], + [50 + 50.0, 100 + 225.0], + ], + ] + ), + ), "Expected oriented-box corners to be shifted into root coordinates" def test_sv_detections_to_root_coordinates_when_scale_and_shift_is_needed() -> None: diff --git a/tests/workflows/unit_tests/core_steps/formatters/test_property_extraction.py b/tests/workflows/unit_tests/core_steps/formatters/test_property_extraction.py index 7ddac78eef..dd3052ac34 100644 --- a/tests/workflows/unit_tests/core_steps/formatters/test_property_extraction.py +++ b/tests/workflows/unit_tests/core_steps/formatters/test_property_extraction.py @@ -1,11 +1,14 @@ import numpy as np +import pytest import supervision as sv +import torch from inference.core.entities.responses.inference import ( ClassificationInferenceResponse, ClassificationPrediction, InferenceResponseImage, ) +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.common.query_language.entities.operations import ( OperationsChain, ) @@ -15,9 +18,26 @@ from inference.core.workflows.execution_engine.constants import ( AREA_CONVERTED_KEY_IN_SV_DETECTIONS, AREA_KEY_IN_SV_DETECTIONS, + CLASS_NAMES_KEY, +) +from inference_models import ClassificationPrediction as NativeClassificationPrediction +from inference_models.models.base.object_detection import Detections as NativeDetections + +# PropertyDefinitionBlockV1 runs UQL operation chains internally, which are +# native-only under ENABLE_TENSOR_DATA_REPRESENTATION - hence the flag-opposed +# _NUMPY_ONLY / _TENSOR_ONLY split below. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="dict / sv.Detections input; the UQL chain inside the block is native-only " + "under ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", ) +@_NUMPY_ONLY def test_property_extraction_block() -> None: # given data = ClassificationInferenceResponse( @@ -57,6 +77,7 @@ def test_property_extraction_block() -> None: assert result == {"output": "cat-mutated"} +@_NUMPY_ONLY def test_property_extraction_block_with_center() -> None: # given detections = sv.Detections( @@ -81,6 +102,7 @@ def test_property_extraction_block_with_center() -> None: assert result == {"output": [[20, 30], [40, 50]]} +@_NUMPY_ONLY def test_property_extraction_block_with_top_left() -> None: # given detections = sv.Detections( @@ -105,6 +127,7 @@ def test_property_extraction_block_with_top_left() -> None: assert result == {"output": [[10, 20], [30, 40]]} +@_NUMPY_ONLY def test_property_extraction_block_with_top_right() -> None: # given detections = sv.Detections( @@ -129,6 +152,7 @@ def test_property_extraction_block_with_top_right() -> None: assert result == {"output": [[30, 20], [50, 40]]} +@_NUMPY_ONLY def test_property_extraction_block_with_bottom_left() -> None: # given detections = sv.Detections( @@ -153,6 +177,7 @@ def test_property_extraction_block_with_bottom_left() -> None: assert result == {"output": [[10, 40], [30, 60]]} +@_NUMPY_ONLY def test_property_extraction_block_with_bottom_right() -> None: # given detections = sv.Detections( @@ -177,6 +202,7 @@ def test_property_extraction_block_with_bottom_right() -> None: assert result == {"output": [[30, 40], [50, 60]]} +@_NUMPY_ONLY def test_property_extraction_block_with_area_px() -> None: # given detections = sv.Detections( @@ -204,6 +230,7 @@ def test_property_extraction_block_with_area_px() -> None: assert result == {"output": [400.0, 1000.0]} +@_NUMPY_ONLY def test_property_extraction_block_with_area_converted() -> None: # given detections = sv.Detections( @@ -231,3 +258,209 @@ def test_property_extraction_block_with_area_converted() -> None: # then assert result == {"output": [4.0, 10.0]} + + +# --------------------------------------------------------------------------- +# Tensor-native variants (run only under ENABLE_TENSOR_DATA_REPRESENTATION) +# --------------------------------------------------------------------------- + + +@_TENSOR_ONLY +def test_property_extraction_block_tensor_native() -> None: + # given + data = NativeClassificationPrediction( + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.6, 0.4]], dtype=torch.float32), + images_metadata=[{CLASS_NAMES_KEY: {0: "cat", 1: "dog"}}], + ) + operations = OperationsChain.model_validate( + { + "operations": [ + { + "type": "ClassificationPropertyExtract", + "property_name": "top_class", + }, + { + "type": "LookupTable", + "lookup_table": {"cat": "cat-mutated"}, + }, + ] + } + ).operations + step = PropertyDefinitionBlockV1() + + # when + result = step.run(data=data, operations=operations) + + # then + assert result == {"output": "cat-mutated"} + + +def _native_detections_for_property_extraction() -> NativeDetections: + return NativeDetections( + xyxy=torch.tensor([[10, 20, 30, 40], [30, 40, 50, 60]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4], dtype=torch.float32), + ) + + +@_TENSOR_ONLY +def test_property_extraction_block_with_center_tensor_native() -> None: + # given + detections = _native_detections_for_property_extraction() + operations = OperationsChain.model_validate( + { + "operations": [ + {"type": "DetectionsPropertyExtract", "property_name": "center"} + ] + } + ).operations + step = PropertyDefinitionBlockV1() + + # when + result = step.run(data=detections, operations=operations) + + # then - anchor coords are rounded ints + assert result == {"output": [[20, 30], [40, 50]]} + + +@_TENSOR_ONLY +def test_property_extraction_block_with_top_left_tensor_native() -> None: + # given + detections = _native_detections_for_property_extraction() + operations = OperationsChain.model_validate( + { + "operations": [ + {"type": "DetectionsPropertyExtract", "property_name": "top_left"} + ] + } + ).operations + step = PropertyDefinitionBlockV1() + + # when + result = step.run(data=detections, operations=operations) + + # then + assert result == {"output": [[10, 20], [30, 40]]} + + +@_TENSOR_ONLY +def test_property_extraction_block_with_top_right_tensor_native() -> None: + # given + detections = _native_detections_for_property_extraction() + operations = OperationsChain.model_validate( + { + "operations": [ + {"type": "DetectionsPropertyExtract", "property_name": "top_right"} + ] + } + ).operations + step = PropertyDefinitionBlockV1() + + # when + result = step.run(data=detections, operations=operations) + + # then + assert result == {"output": [[30, 20], [50, 40]]} + + +@_TENSOR_ONLY +def test_property_extraction_block_with_bottom_left_tensor_native() -> None: + # given + detections = _native_detections_for_property_extraction() + operations = OperationsChain.model_validate( + { + "operations": [ + {"type": "DetectionsPropertyExtract", "property_name": "bottom_left"} + ] + } + ).operations + step = PropertyDefinitionBlockV1() + + # when + result = step.run(data=detections, operations=operations) + + # then + assert result == {"output": [[10, 40], [30, 60]]} + + +@_TENSOR_ONLY +def test_property_extraction_block_with_bottom_right_tensor_native() -> None: + # given + detections = _native_detections_for_property_extraction() + operations = OperationsChain.model_validate( + { + "operations": [ + {"type": "DetectionsPropertyExtract", "property_name": "bottom_right"} + ] + } + ).operations + step = PropertyDefinitionBlockV1() + + # when + result = step.run(data=detections, operations=operations) + + # then + assert result == {"output": [[30, 40], [50, 60]]} + + +@_TENSOR_ONLY +def test_property_extraction_block_with_area_px_tensor_native() -> None: + # given - non-geometric per-box scalars are read from bboxes_metadata + detections = NativeDetections( + xyxy=torch.tensor([[10, 20, 30, 40], [30, 40, 50, 60]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4], dtype=torch.float32), + bboxes_metadata=[ + {AREA_KEY_IN_SV_DETECTIONS: 400.0}, + {AREA_KEY_IN_SV_DETECTIONS: 1000.0}, + ], + ) + operations = OperationsChain.model_validate( + { + "operations": [ + { + "type": "DetectionsPropertyExtract", + "property_name": AREA_KEY_IN_SV_DETECTIONS, + } + ] + } + ).operations + step = PropertyDefinitionBlockV1() + + # when + result = step.run(data=detections, operations=operations) + + # then + assert result == {"output": [400.0, 1000.0]} + + +@_TENSOR_ONLY +def test_property_extraction_block_with_area_converted_tensor_native() -> None: + # given - non-geometric per-box scalars are read from bboxes_metadata + detections = NativeDetections( + xyxy=torch.tensor([[10, 20, 30, 40], [30, 40, 50, 60]], dtype=torch.float32), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.6, 0.4], dtype=torch.float32), + bboxes_metadata=[ + {AREA_CONVERTED_KEY_IN_SV_DETECTIONS: 4.0}, + {AREA_CONVERTED_KEY_IN_SV_DETECTIONS: 10.0}, + ], + ) + operations = OperationsChain.model_validate( + { + "operations": [ + { + "type": "DetectionsPropertyExtract", + "property_name": AREA_CONVERTED_KEY_IN_SV_DETECTIONS, + } + ] + } + ).operations + step = PropertyDefinitionBlockV1() + + # when + result = step.run(data=detections, operations=operations) + + # then + assert result == {"output": [4.0, 10.0]} diff --git a/tests/workflows/unit_tests/core_steps/formatters/test_tensor_formatters_no_materialization.py b/tests/workflows/unit_tests/core_steps/formatters/test_tensor_formatters_no_materialization.py new file mode 100644 index 0000000000..f6145af793 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/formatters/test_tensor_formatters_no_materialization.py @@ -0,0 +1,199 @@ +"""Guards that the tensor-native VLM formatter blocks read image H/W without +forcing a device->host materialisation of a tensor-source image. + +Previously these blocks called ``image.numpy_image.shape[:2]`` purely to read +(H, W). On a tensor-source ``WorkflowImageData`` that eagerly downloads the whole +frame to host (measured ~3-4 ms for 720p on a Jetson Orin Nano) and leaves a +numpy cache behind. They now use ``_read_shape_without_materialization()``, which +reads the shape straight off the already-present representation. + +Each test drives a tensor-source image (CPU torch tensor - no CUDA required) +through a formatter and asserts BOTH: + * the H/W-dependent output is correct (image_dimensions metadata, and for the + coordinate-scaling parsers, the scaled xyxy), and + * the image's internal numpy cache (``_numpy_image``) stays ``None`` - i.e. no + materialisation happened as a side effect of the shape read. + +Non-square dimensions (H=192, W=168) are used throughout so an accidental H/W +swap would flip both the dimensions metadata and the scaled coordinates. +""" + +import numpy as np +import pytest + +pytest.importorskip("torch") +pytest.importorskip("inference_models") + +import torch + +from inference.core.workflows.core_steps.formatters.vlm_as_classifier.v1_tensor import ( + VLMAsClassifierBlockV1, +) +from inference.core.workflows.core_steps.formatters.vlm_as_classifier.v2_tensor import ( + VLMAsClassifierBlockV2, +) +from inference.core.workflows.core_steps.formatters.vlm_as_detector.v1_tensor import ( + VLMAsDetectorBlockV1, +) +from inference.core.workflows.core_steps.formatters.vlm_as_detector.v2_tensor import ( + VLMAsDetectorBlockV2, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) + +IMAGE_HEIGHT = 192 +IMAGE_WIDTH = 168 + + +def _tensor_source_image() -> WorkflowImageData: + # CHW RGB uint8 tensor source - constructing with `tensor_image` leaves the + # numpy representation unmaterialised (`_numpy_image is None`). + tensor_image = torch.zeros((3, IMAGE_HEIGHT, IMAGE_WIDTH), dtype=torch.uint8) + image = WorkflowImageData( + tensor_image=tensor_image, + parent_metadata=ImageParentMetadata(parent_id="parent"), + ) + # Sanity: the source really is tensor-only before the block runs. + assert image._numpy_image is None + assert image.is_tensor_materialised() is True + return image + + +def _assert_not_materialised(image: WorkflowImageData) -> None: + assert image._numpy_image is None, ( + "Reading H/W must not materialise the numpy representation of a " + "tensor-source image (forces a full-frame device->host download)." + ) + + +# --------------------------------------------------------------------------- # +# vlm_as_classifier +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("block_cls", [VLMAsClassifierBlockV1, VLMAsClassifierBlockV2]) +def test_classifier_multi_class_reads_shape_without_materialization( + block_cls, +) -> None: + # given + image = _tensor_source_image() + vlm_output = '{"class_name": "cat", "confidence": 0.9}' + + # when + result = block_cls().run(image=image, vlm_output=vlm_output, classes=["cat", "dog"]) + + # then + assert result["error_status"] is False + dimensions = result["predictions"].images_metadata[0]["image_dimensions"] + assert list(dimensions) == [IMAGE_HEIGHT, IMAGE_WIDTH] + _assert_not_materialised(image) + + +@pytest.mark.parametrize("block_cls", [VLMAsClassifierBlockV1, VLMAsClassifierBlockV2]) +def test_classifier_multi_label_reads_shape_without_materialization( + block_cls, +) -> None: + # given + image = _tensor_source_image() + vlm_output = ( + '{"predicted_classes": [' + '{"class": "cat", "confidence": 0.8}, ' + '{"class": "dog", "confidence": 0.6}]}' + ) + + # when + result = block_cls().run(image=image, vlm_output=vlm_output, classes=["cat", "dog"]) + + # then + assert result["error_status"] is False + # MultiLabelClassificationPrediction exposes a single `image_metadata` dict + # (multi-class uses a per-image `images_metadata` list). + dimensions = result["predictions"].image_metadata["image_dimensions"] + assert list(dimensions) == [IMAGE_HEIGHT, IMAGE_WIDTH] + _assert_not_materialised(image) + + +# --------------------------------------------------------------------------- # +# vlm_as_detector +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("block_cls", [VLMAsDetectorBlockV1, VLMAsDetectorBlockV2]) +def test_detector_gemini_scales_coords_without_materialization(block_cls) -> None: + # given - gemini `box_2d` is [y_min, x_min, y_max, x_max] normalised to 1000, + # so the resulting pixel box is a direct function of (H, W). + image = _tensor_source_image() + vlm_output = '{"detections": [{"box_2d": [100, 200, 500, 800], "label": "cat"}]}' + + # when + result = block_cls().run( + image=image, + vlm_output=vlm_output, + classes=["cat", "dog"], + model_type="google-gemini", + task_type="object-detection", + ) + + # then + assert result["error_status"] is False + xyxy = result["predictions"].xyxy.cpu().numpy() + # x_min = 200/1000*168 = 33.6 -> 34 ; y_min = 100/1000*192 = 19.2 -> 19 + # x_max = 800/1000*168 = 134.4 -> 134 ; y_max = 500/1000*192 = 96.0 -> 96 + assert np.array_equal(xyxy, np.array([[34.0, 19.0, 134.0, 96.0]])) + dimensions = result["predictions"].image_metadata["image_dimensions"] + assert list(dimensions) == [IMAGE_HEIGHT, IMAGE_WIDTH] + _assert_not_materialised(image) + + +def test_detector_v2_llm_scales_coords_without_materialization() -> None: + # given - the openai/claude LLM parser (v2 only) multiplies normalised + # [0, 1] coords by (W, H); exercises the second changed site in v2_tensor. + image = _tensor_source_image() + vlm_output = ( + '{"detections": [{"x_min": 0.1, "y_min": 0.2, "x_max": 0.5, ' + '"y_max": 0.8, "class_name": "cat"}]}' + ) + + # when + result = VLMAsDetectorBlockV2().run( + image=image, + vlm_output=vlm_output, + classes=["cat", "dog"], + model_type="openai", + task_type="object-detection", + ) + + # then + assert result["error_status"] is False + xyxy = result["predictions"].xyxy.cpu().numpy() + # x_min = 0.1*168 = 16.8 -> 17 ; y_min = 0.2*192 = 38.4 -> 38 + # x_max = 0.5*168 = 84.0 -> 84 ; y_max = 0.8*192 = 153.6 -> 154 + assert np.array_equal(xyxy, np.array([[17.0, 38.0, 84.0, 154.0]])) + dimensions = result["predictions"].image_metadata["image_dimensions"] + assert list(dimensions) == [IMAGE_HEIGHT, IMAGE_WIDTH] + _assert_not_materialised(image) + + +@pytest.mark.parametrize("block_cls", [VLMAsDetectorBlockV1, VLMAsDetectorBlockV2]) +def test_detector_florence_reads_shape_without_materialization(block_cls) -> None: + # given - the florence-2 parser passes (W, H) as resolution_wh to + # sv.Detections.from_lmm; exercises the florence changed site in each block. + image = _tensor_source_image() + vlm_output = '{"bboxes": [[10.0, 20.0, 90.0, 120.0]], "bboxes_labels": ["cat"]}' + + # when + result = block_cls().run( + image=image, + vlm_output=vlm_output, + classes=["cat", "dog"], + model_type="florence-2", + task_type="open-vocabulary-object-detection", + ) + + # then + assert result["error_status"] is False + dimensions = result["predictions"].image_metadata["image_dimensions"] + assert list(dimensions) == [IMAGE_HEIGHT, IMAGE_WIDTH] + _assert_not_materialised(image) diff --git a/tests/workflows/unit_tests/core_steps/formatters/vlm_as_detector/test_v1_tensor.py b/tests/workflows/unit_tests/core_steps/formatters/vlm_as_detector/test_v1_tensor.py new file mode 100644 index 0000000000..cebd3681f7 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/formatters/vlm_as_detector/test_v1_tensor.py @@ -0,0 +1,58 @@ +import numpy as np +import pytest + +pytest.importorskip("torch") +pytest.importorskip("inference_models") + +from inference.core.workflows.core_steps.formatters.vlm_as_detector.v1_tensor import ( + VLMAsDetectorBlockV1, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) +from inference_models.models.base.object_detection import Detections + + +def test_formatter_for_florence2_open_vocabulary_object_detection() -> None: + # Locks the `florence_task_type == ""` exact-match + # semantics ported from the numpy source (previously an accidental substring + # test via `in`): the OVD branch must execute for the exact task type and map + # class ids from the caller-provided `classes` list by index. + # given + block = VLMAsDetectorBlockV1() + image = WorkflowImageData( + numpy_image=np.zeros((192, 168, 3), dtype=np.uint8), + parent_metadata=ImageParentMetadata(parent_id="parent"), + ) + vlm_output = """ +{"bboxes": [[434.0, 30.848499298095703, 760.4000244140625, 530.4144897460938], [0.4000000059604645, 96.13949584960938, 528.4000244140625, 564.5574951171875]], "bboxes_labels": ["cat", "dog"]} +""" + + # when + result = block.run( + image=image, + vlm_output=vlm_output, + classes=["cat", "dog"], + model_type="florence-2", + task_type="open-vocabulary-object-detection", + ) + + # then + assert result["error_status"] is False + assert isinstance(result["predictions"], Detections) + assert len(result["inference_id"]) > 0 + assert np.allclose( + result["predictions"].xyxy.cpu().numpy(), + np.array([[434, 30.848, 760.4, 530.41], [0.4, 96.139, 528.4, 564.56]]), + atol=1e-1, + ), "Expected coordinates to be the same as given in raw input" + assert result["predictions"].class_id.cpu().tolist() == [0, 1] + assert np.allclose( + result["predictions"].confidence.cpu().numpy(), np.array([1.0, 1.0]) + ) + assert [entry["class"] for entry in result["predictions"].bboxes_metadata] == [ + "cat", + "dog", + ] + assert result["predictions"].image_metadata["class_names"] == {0: "cat", 1: "dog"} diff --git a/tests/workflows/unit_tests/core_steps/fusion/test_detections_classes_replacement.py b/tests/workflows/unit_tests/core_steps/fusion/test_detections_classes_replacement.py index 6926b82599..41c72321d5 100644 --- a/tests/workflows/unit_tests/core_steps/fusion/test_detections_classes_replacement.py +++ b/tests/workflows/unit_tests/core_steps/fusion/test_detections_classes_replacement.py @@ -809,3 +809,121 @@ def test_classes_replacement_with_strings_and_none_with_fallback() -> None: ] assert result["predictions"].confidence[1] == 0.0 assert result["predictions"].class_id[1] == 99 + + +def test_classes_replacement_tensor_native_empty_prediction_with_fallback_and_none_class_id() -> ( + None +): + # fallback_class_id=None must not raise TypeError on int(None) + torch = pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.fusion.detections_classes_replacement.v1_tensor import ( + DetectionsClassesReplacementBlockV1 as TensorBlock, + ) + from inference_models.models.base.classification import ( + ClassificationPrediction as NativeClassificationPrediction, + ) + from inference_models.models.base.object_detection import ( + Detections as NativeDetections, + ) + + # given + step = TensorBlock() + detections = NativeDetections( + xyxy=torch.tensor( + [ + [10.0, 20.0, 30.0, 40.0], + [11.0, 21.0, 31.0, 41.0], + ] + ), + class_id=torch.tensor([7, 7]), + confidence=torch.tensor([0.36, 0.91]), + image_metadata={"class_names": {7: "animal"}}, + bboxes_metadata=[{"detection_id": "zero"}, {"detection_id": "one"}], + ) + first_cls_prediction = NativeClassificationPrediction( + class_id=torch.tensor([0]), + confidence=torch.tensor([[0.6, 0.4]]), + images_metadata=[{"parent_id": "zero", "class_names": {0: "cat", 1: "dog"}}], + ) + second_cls_prediction = NativeClassificationPrediction( + class_id=torch.tensor([0]), + confidence=torch.zeros((1, 0)), + images_metadata=[{"parent_id": "one"}], + ) + classification_predictions = Batch( + content=[first_cls_prediction, second_cls_prediction], + indices=[(0, 0), (0, 1)], + ) + + # when + result = step.run( + object_detection_predictions=detections, + classification_predictions=classification_predictions, + fallback_class_name="unknown", + fallback_class_id=None, + ) + + # then + predictions = result["predictions"] + assert ( + len(predictions) == 2 + ), "Expected both detections to be preserved via fallback class" + assert predictions.confidence[1] == 0, "Fallback confidence expected to be 0" + assert ( + int(predictions.class_id[1]) == sys.maxsize + ), "class id expected to fall back to sys.maxsize when fallback_class_id left as None" + assert ( + predictions.image_metadata["class_names"][sys.maxsize] == "unknown" + ), "class_names map expected to resolve the fallback id to fallback_class_name" + + +def test_classes_replacement_tensor_native_string_predictions_with_negative_fallback_class_id() -> ( + None +): + torch = pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.fusion.detections_classes_replacement.v1_tensor import ( + DetectionsClassesReplacementBlockV1 as TensorBlock, + ) + from inference_models.models.base.object_detection import ( + Detections as NativeDetections, + ) + + # given + step = TensorBlock() + detections = NativeDetections( + xyxy=torch.tensor( + [ + [10.0, 20.0, 30.0, 40.0], + [11.0, 21.0, 31.0, 41.0], + ] + ), + class_id=torch.tensor([7, 7]), + confidence=torch.tensor([0.36, 0.91]), + image_metadata={"class_names": {7: "animal"}}, + bboxes_metadata=[{"detection_id": "zero"}, {"detection_id": "one"}], + ) + classification_predictions = Batch( + content=["K619879", None], + indices=[(0, 0), (0, 1)], + ) + + # when + result = step.run( + object_detection_predictions=detections, + classification_predictions=classification_predictions, + fallback_class_name="unreadable", + fallback_class_id=-5, + ) + + # then + predictions = result["predictions"] + assert len(predictions) == 2 + assert ( + int(predictions.class_id[1]) == sys.maxsize + ), "negative fallback_class_id expected to normalise to sys.maxsize" + assert predictions.confidence[1] == 0.0 + assert ( + predictions.image_metadata["class_names"][sys.maxsize] == "unreadable" + ), "class_names map expected to resolve the fallback id to fallback_class_name" diff --git a/tests/workflows/unit_tests/core_steps/fusion/test_detections_consensus_v1_tensor.py b/tests/workflows/unit_tests/core_steps/fusion/test_detections_consensus_v1_tensor.py new file mode 100644 index 0000000000..bd92a8a98c --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/fusion/test_detections_consensus_v1_tensor.py @@ -0,0 +1,535 @@ +"""Fuzz parity between the vectorised tensor-native consensus matching and a +verbatim copy of the ORIGINAL per-pair logic (the oracle). + +The optimisation replaced the O(N^2) `calculate_iou` / `_resolve_class_name` +per-pair device syncs with a single vectorised `sv.box_iou_batch` matrix and a +host download of `class_id` per source (see `_precompute_pair_data`). Everything +downstream of the matching (mask padding, merge, presence checks) is unchanged +and shared between the oracle and the new code, so this test drives hundreds of +random multi-source scenarios and asserts every output field is identical. + +`uuid4` is patched to a deterministic counter (reset before each run) so the +merged-detection ``detection_id`` values are comparable: because the two paths +produce identical merges in identical order, the deterministic id stream lines +up exactly - any divergence in the matching would desync it. + +Runs on CPU torch; no CUDA required. +""" + +import random +from typing import Any, Dict, List, Optional, Set, Tuple, Union + +import numpy as np +import pytest +import torch + +from inference.core.workflows.core_steps.fusion.detections_consensus import ( + v1_tensor as tv, +) +from inference.core.workflows.core_steps.fusion.detections_consensus.v1_tensor import ( + AggregationMode, + MaskAggregationMode, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +TensorNativeDetections = Union[Detections, InstanceDetections] + + +# --------------------------------------------------------------------------- # +# Oracle: a verbatim copy of the ORIGINAL matching (per-pair calculate_iou / # +# _resolve_class_name over freshly-built single-row predictions). Everything # +# else reuses the unchanged shared helpers from the block module. # +# --------------------------------------------------------------------------- # +def _oracle_enumerate_detections( + detections_from_sources: List[TensorNativeDetections], + excluded_source_id: Optional[int] = None, +): + for source_id, detections in enumerate(detections_from_sources): + if excluded_source_id == source_id: + continue + for i in range(len(detections)): + yield source_id, tv.take_prediction_by_indices(detections, [i]) + + +def _oracle_max_overlap( + detection: TensorNativeDetections, + source: int, + detections_from_sources: List[TensorNativeDetections], + iou_threshold: float, + class_aware: bool, + detections_already_considered: Set[str], +) -> Dict[int, Tuple[TensorNativeDetections, float]]: + current_max_overlap: Dict[int, Tuple[TensorNativeDetections, float]] = {} + for other_source, other_detection in _oracle_enumerate_detections( + detections_from_sources=detections_from_sources, + excluded_source_id=source, + ): + if tv._resolve_detection_id(other_detection) in detections_already_considered: + continue + if class_aware and tv._resolve_class_name(detection) != tv._resolve_class_name( + other_detection + ): + continue + iou_value = tv.calculate_iou(detection_a=detection, detection_b=other_detection) + if iou_value <= iou_threshold: + continue + if current_max_overlap.get(other_source) is None: + current_max_overlap[other_source] = (other_detection, iou_value) + if current_max_overlap[other_source][1] < iou_value: + current_max_overlap[other_source] = (other_detection, iou_value) + return current_max_overlap + + +def _oracle_get_consensus_for_single_detection( + detection, + source_id, + detections_from_sources, + iou_threshold, + class_aware, + required_votes, + confidence, + detections_merge_confidence_aggregation, + detections_merge_coordinates_aggregation, + detections_merge_mask_aggregation, + detections_already_considered, +): + if ( + len(detection) + and tv._resolve_detection_id(detection) in detections_already_considered + ): + return [], detections_already_considered + consensus_detections = [] + detections_with_max_overlap = _oracle_max_overlap( + detection=detection, + source=source_id, + detections_from_sources=detections_from_sources, + iou_threshold=iou_threshold, + class_aware=class_aware, + detections_already_considered=detections_already_considered, + ) + if len(detections_with_max_overlap) < (required_votes - 1): + return consensus_detections, detections_already_considered + detection_mask = tv._single_mask(detection) + overlap_masks = { + other_source: tv._single_mask(matched_value[0]) + for other_source, matched_value in detections_with_max_overlap.items() + } + if detection_mask is not None: + for other_source, matched_mask in overlap_masks.items(): + if matched_mask is None: + overlap_masks[other_source] = np.zeros(detection_mask.shape, dtype=bool) + else: + shape = None + for matched_mask in overlap_masks.values(): + if matched_mask is not None: + shape = matched_mask.shape + break + if shape: + for other_source, matched_mask in overlap_masks.items(): + if matched_mask is None: + overlap_masks[other_source] = np.zeros(shape, dtype=bool) + detection_mask = np.zeros(shape, dtype=bool) + group_detections = [detection] + [ + matched_value[0] for matched_value in detections_with_max_overlap.values() + ] + group_masks = ( + [detection_mask] + + [overlap_masks[other_source] for other_source in detections_with_max_overlap] + if detection_mask is not None + else None + ) + merged_detection = tv.merge_detections( + detections=group_detections, + masks=group_masks, + confidence_aggregation_mode=detections_merge_confidence_aggregation, + boxes_aggregation_mode=detections_merge_coordinates_aggregation, + mask_aggregation_mode=detections_merge_mask_aggregation, + ) + if float(merged_detection.confidence[0]) < confidence: + return consensus_detections, detections_already_considered + consensus_detections.append(merged_detection) + detections_already_considered.add(tv._resolve_detection_id(detection)) + for matched_value in detections_with_max_overlap.values(): + detections_already_considered.add(tv._resolve_detection_id(matched_value[0])) + return consensus_detections, detections_already_considered + + +def oracle_agree_on_consensus( + detections_from_sources, + required_votes, + class_aware, + iou_threshold, + confidence, + classes_to_consider, + required_objects, + presence_confidence_aggregation, + detections_merge_confidence_aggregation, + detections_merge_coordinates_aggregation, + detections_merge_mask_aggregation, +): + if tv.does_not_detect_objects_in_any_source( + detections_from_sources=detections_from_sources + ): + return ( + "undefined", + False, + {}, + tv._empty_native_detections( + device=( + detections_from_sources[0].xyxy.device + if detections_from_sources + else None + ) + ), + ) + parent_id = tv.get_parent_id_of_detections_from_sources( + detections_from_sources=detections_from_sources, + ) + detections_from_sources = tv.filter_predictions( + predictions=detections_from_sources, + classes_to_consider=classes_to_consider, + ) + detections_already_considered = set() + consensus_detections = [] + for source_id, detection in _oracle_enumerate_detections( + detections_from_sources=detections_from_sources + ): + ( + consensus_detections_update, + detections_already_considered, + ) = _oracle_get_consensus_for_single_detection( + detection=detection, + source_id=source_id, + detections_from_sources=detections_from_sources, + iou_threshold=iou_threshold, + class_aware=class_aware, + required_votes=required_votes, + confidence=confidence, + detections_merge_confidence_aggregation=detections_merge_confidence_aggregation, + detections_merge_coordinates_aggregation=detections_merge_coordinates_aggregation, + detections_merge_mask_aggregation=detections_merge_mask_aggregation, + detections_already_considered=detections_already_considered, + ) + consensus_detections += consensus_detections_update + consensus_detections = tv._merge_native_detections(consensus_detections) + ( + object_present, + presence_confidence, + ) = tv.check_objects_presence_in_consensus_detections( + consensus_detections=consensus_detections, + aggregation_mode=presence_confidence_aggregation, + class_aware=class_aware, + required_objects=required_objects, + ) + return parent_id, object_present, presence_confidence, consensus_detections + + +# --------------------------------------------------------------------------- # +# Random native-detections scenario generation # +# --------------------------------------------------------------------------- # +_CLASS_POOL = ["a", "b", "c"] +_IMG_H, _IMG_W = 64, 64 + + +class _DetSpec: + """Plain-data description of one source, materialised fresh into a native + prediction for each run so the two paths never share (or mutate) tensors.""" + + def __init__( + self, + xyxy: np.ndarray, + class_ids: List[int], + confidences: List[float], + detection_ids: List[str], + class_names_map: Dict[int, str], + parent_id: str, + with_mask: bool, + ) -> None: + self.xyxy = xyxy + self.class_ids = class_ids + self.confidences = confidences + self.detection_ids = detection_ids + self.class_names_map = class_names_map + self.parent_id = parent_id + self.with_mask = with_mask + + def materialise(self) -> TensorNativeDetections: + n = self.xyxy.shape[0] + image_metadata = { + CLASS_NAMES_KEY: dict(self.class_names_map), + PARENT_ID_KEY: self.parent_id, + PARENT_COORDINATES_KEY: [0, 0], + PARENT_DIMENSIONS_KEY: [_IMG_H, _IMG_W], + ROOT_PARENT_ID_KEY: self.parent_id, + ROOT_PARENT_COORDINATES_KEY: [0, 0], + ROOT_PARENT_DIMENSIONS_KEY: [_IMG_H, _IMG_W], + IMAGE_DIMENSIONS_KEY: [_IMG_H, _IMG_W], + } + bboxes_metadata = [{DETECTION_ID_KEY: did} for did in self.detection_ids] + xyxy = torch.tensor(self.xyxy, dtype=torch.float32) + class_id = torch.tensor(self.class_ids, dtype=torch.long) + confidence = torch.tensor(self.confidences, dtype=torch.float32) + if not self.with_mask: + return Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + mask = torch.zeros((n, _IMG_H, _IMG_W), dtype=torch.bool) + for i in range(n): + x1, y1, x2, y2 = self.xyxy[i].astype(int) + x1 = max(0, min(_IMG_W, x1)) + x2 = max(0, min(_IMG_W, x2)) + y1 = max(0, min(_IMG_H, y1)) + y2 = max(0, min(_IMG_H, y2)) + if x2 > x1 and y2 > y1: + mask[i, y1:y2, x1:x2] = True + return InstanceDetections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + + +def _random_box(rng: random.Random) -> np.ndarray: + x1 = rng.uniform(0, _IMG_W - 8) + y1 = rng.uniform(0, _IMG_H - 8) + w = rng.uniform(4, _IMG_W - x1) + h = rng.uniform(4, _IMG_H - y1) + return np.array([x1, y1, x1 + w, y1 + h], dtype=np.float64) + + +def _jitter(box: np.ndarray, rng: random.Random) -> np.ndarray: + d = np.array([rng.uniform(-3, 3) for _ in range(4)]) + out = box + d + out[2] = max(out[2], out[0] + 2) + out[3] = max(out[3], out[1] + 2) + return np.clip(out, 0, max(_IMG_H, _IMG_W)) + + +def _build_sources(scenario_index: int) -> Tuple[List[_DetSpec], Dict[str, Any]]: + rng = random.Random(20260714 + scenario_index * 7919) + parent_id = "parent-image" + num_sources = rng.randint(2, 4) + pool_size = rng.randint(1, 3) + class_ids = list(range(pool_size)) + class_names_map = {cid: _CLASS_POOL[cid] for cid in class_ids} + # A handful of shared cluster anchors so cross-source boxes actually overlap + # (identical anchors exercise the equal-IoU tie rule). + anchors = [_random_box(rng) for _ in range(rng.randint(1, 4))] + homogeneous_mask = rng.random() < 0.5 + global_scenario_masks = rng.random() < 0.5 + + specs: List[_DetSpec] = [] + for source_id in range(num_sources): + count = rng.choice([0, 1, 1, 2, 3, 4, 5]) + boxes = [] + for _ in range(count): + mode = rng.random() + if anchors and mode < 0.4: + boxes.append(_jitter(rng.choice(anchors), rng)) + elif anchors and mode < 0.6: + boxes.append(np.array(rng.choice(anchors), dtype=np.float64)) + else: + boxes.append(_random_box(rng)) + xyxy = np.stack(boxes, axis=0) if boxes else np.zeros((0, 4), dtype=np.float64) + cids = [rng.choice(class_ids) for _ in range(count)] + confs = [round(rng.uniform(0.05, 0.99), 4) for _ in range(count)] + dids = [f"s{source_id}_d{i}" for i in range(count)] + if homogeneous_mask: + with_mask = global_scenario_masks + else: + with_mask = rng.random() < 0.5 + specs.append( + _DetSpec( + xyxy=xyxy, + class_ids=cids, + confidences=confs, + detection_ids=dids, + class_names_map=class_names_map, + parent_id=parent_id, + with_mask=with_mask, + ) + ) + + classes_to_consider = rng.choice( + [None, None, [_CLASS_POOL[0]], _CLASS_POOL[:pool_size], ["nonexistent"]] + ) + required_objects_choice = rng.choice([None, 1, 2, {"a": 1}, {"a": 1, "b": 1}]) + params = { + "required_votes": rng.randint(1, 3), + "class_aware": rng.random() < 0.5, + "iou_threshold": rng.choice([0.0, 0.2, 0.3, 0.5, 0.7]), + "confidence": rng.choice([0.0, 0.0, 0.3, 0.6]), + "classes_to_consider": classes_to_consider, + "required_objects": required_objects_choice, + "presence_confidence_aggregation": rng.choice(list(AggregationMode)), + "detections_merge_confidence_aggregation": rng.choice(list(AggregationMode)), + "detections_merge_coordinates_aggregation": rng.choice(list(AggregationMode)), + "detections_merge_mask_aggregation": rng.choice(list(MaskAggregationMode)), + } + return specs, params + + +# --------------------------------------------------------------------------- # +# Output comparison # +# --------------------------------------------------------------------------- # +def _assert_detections_equal( + got: TensorNativeDetections, exp: TensorNativeDetections +) -> None: + assert type(got) is type(exp), f"{type(got)} != {type(exp)}" + np.testing.assert_array_equal( + got.xyxy.detach().cpu().numpy(), exp.xyxy.detach().cpu().numpy() + ) + np.testing.assert_array_equal( + got.class_id.detach().cpu().numpy(), exp.class_id.detach().cpu().numpy() + ) + np.testing.assert_array_equal( + got.confidence.detach().cpu().numpy(), + exp.confidence.detach().cpu().numpy(), + ) + got_mask = getattr(got, "mask", None) + exp_mask = getattr(exp, "mask", None) + assert (got_mask is None) == (exp_mask is None) + if got_mask is not None: + np.testing.assert_array_equal( + got_mask.detach().cpu().numpy(), exp_mask.detach().cpu().numpy() + ) + got_meta = got.image_metadata or {} + exp_meta = exp.image_metadata or {} + assert got_meta.get(CLASS_NAMES_KEY) == exp_meta.get(CLASS_NAMES_KEY) + got_ids = [m.get(DETECTION_ID_KEY) for m in (got.bboxes_metadata or [])] + exp_ids = [m.get(DETECTION_ID_KEY) for m in (exp.bboxes_metadata or [])] + assert got_ids == exp_ids + + +def _run_with_deterministic_uuid(monkeypatch, fn, specs, params): + counter = {"n": 0} + + def _fake_uuid4(): + counter["n"] += 1 + return f"merged-{counter['n']:04d}" + + monkeypatch.setattr(tv, "uuid4", _fake_uuid4) + sources = [spec.materialise() for spec in specs] + return fn( + detections_from_sources=sources, + required_votes=params["required_votes"], + class_aware=params["class_aware"], + iou_threshold=params["iou_threshold"], + confidence=params["confidence"], + classes_to_consider=params["classes_to_consider"], + required_objects=params["required_objects"], + presence_confidence_aggregation=params["presence_confidence_aggregation"], + detections_merge_confidence_aggregation=params[ + "detections_merge_confidence_aggregation" + ], + detections_merge_coordinates_aggregation=params[ + "detections_merge_coordinates_aggregation" + ], + detections_merge_mask_aggregation=params["detections_merge_mask_aggregation"], + ) + + +@pytest.mark.parametrize("scenario_index", list(range(160))) +def test_vectorised_consensus_matches_original_per_pair_logic( + scenario_index: int, monkeypatch: pytest.MonkeyPatch +) -> None: + # given + specs, params = _build_sources(scenario_index) + + # when + exp_parent, exp_present, exp_conf, exp_dets = _run_with_deterministic_uuid( + monkeypatch, oracle_agree_on_consensus, specs, params + ) + got_parent, got_present, got_conf, got_dets = _run_with_deterministic_uuid( + monkeypatch, + tv.agree_on_consensus_for_all_detections_sources, + specs, + params, + ) + + # then + assert got_parent == exp_parent + assert got_present == exp_present + assert got_conf == pytest.approx(exp_conf) + assert len(got_dets) == len(exp_dets) + _assert_detections_equal(got_dets, exp_dets) + + +def test_scenarios_are_non_trivial() -> None: + """Guard the fuzz corpus: across the seeded scenarios there must be runs + that actually emit consensus detections (otherwise the parity assertions + would be vacuously satisfied by everything returning empty).""" + emitted = 0 + empty = 0 + for scenario_index in range(160): + specs, params = _build_sources(scenario_index) + sources = [spec.materialise() for spec in specs] + _, _, _, dets = tv.agree_on_consensus_for_all_detections_sources( + detections_from_sources=sources, + required_votes=params["required_votes"], + class_aware=params["class_aware"], + iou_threshold=params["iou_threshold"], + confidence=params["confidence"], + classes_to_consider=params["classes_to_consider"], + required_objects=params["required_objects"], + presence_confidence_aggregation=params["presence_confidence_aggregation"], + detections_merge_confidence_aggregation=params[ + "detections_merge_confidence_aggregation" + ], + detections_merge_coordinates_aggregation=params[ + "detections_merge_coordinates_aggregation" + ], + detections_merge_mask_aggregation=params[ + "detections_merge_mask_aggregation" + ], + ) + if len(dets) > 0: + emitted += 1 + else: + empty += 1 + assert emitted >= 20, f"too few non-empty consensus outputs: {emitted}" + assert empty >= 5, f"expected some empty consensus outputs too: {empty}" + + +def test_single_d2h_per_tensor_field_per_source() -> None: + """Exactly one host download per tensor field (xyxy, class_id) per source in + `_precompute_pair_data`, and none left in the O(N^2) matching loop.""" + specs, _ = _build_sources(3) + sources = [spec.materialise() for spec in specs if spec.xyxy.shape[0] > 0] + # Make sure the corpus for this check has several detections to match. + assert sum(len(s) for s in sources) >= 3 + + calls = {"count": 0} + original_to = torch.Tensor.to + + def _counting_to(self, *args, **kwargs): + target = args[0] if args else kwargs.get("device") + if isinstance(target, str) and target == "cpu": + calls["count"] += 1 + return original_to(self, *args, **kwargs) + + # xyxy + class_id => 2 downloads per non-empty source. + import unittest.mock as _mock + + with _mock.patch.object(torch.Tensor, "to", _counting_to): + tv._precompute_pair_data(sources) + assert calls["count"] == 2 * len(sources) diff --git a/tests/workflows/unit_tests/core_steps/fusion/test_detections_stitch_v1_tensor.py b/tests/workflows/unit_tests/core_steps/fusion/test_detections_stitch_v1_tensor.py new file mode 100644 index 0000000000..806fc6bb28 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/fusion/test_detections_stitch_v1_tensor.py @@ -0,0 +1,455 @@ +"""Parity tests for the torch NMM port of the tensor-native detections_stitch +block. + +`_oracle_with_nmm` below is a verbatim copy of the block's NMM branch BEFORE the +torch port (dense masks D2H -> sv.Detections.with_nmm on CPU -> re-upload). The +fuzz suite asserts the new `with_nmm` is value-identical to that oracle: same +surviving/merged boxes, class ids, confidences and exact bool-equal masks. +""" + +import random +from copy import deepcopy +from typing import List, Optional, Tuple, Union + +import numpy as np +import pytest +import supervision as sv +import torch + +import inference.core.workflows.core_steps.fusion.detections_stitch.v1_tensor as stitch_module +from inference.core.workflows.core_steps.fusion.detections_stitch.v1_tensor import ( + with_nmm, + with_nms, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import ( + coco_rle_masks_to_numpy_mask, + torch_mask_to_coco_rle, +) + +TensorNativeDetections = Union[Detections, InstanceDetections] + + +def _oracle_with_nmm( + detections: TensorNativeDetections, + threshold: float, +) -> TensorNativeDetections: + """Verbatim copy of the pre-port NMM branch of detections_stitch v1_tensor + (sv.Detections used as the NMM algorithm, full masks round-tripped through + host memory). Serves as the behavioral oracle for the torch port.""" + if len(detections) == 0: + return detections + is_instance_segmentation = isinstance(detections, InstanceDetections) + masks = None + if is_instance_segmentation and detections.mask is not None: + if isinstance(detections.mask, InstancesRLEMasks): + masks = coco_rle_masks_to_numpy_mask(detections.mask) + else: + masks = detections.mask.detach().to("cpu").numpy().astype(bool) + nmm_output = sv.Detections( + xyxy=detections.xyxy.detach().to("cpu").numpy().astype(float), + confidence=detections.confidence.detach().to("cpu").numpy().astype(float), + class_id=detections.class_id.detach().to("cpu").numpy().astype(int), + mask=masks, + ).with_nmm(threshold=threshold) + number_of_detections = len(nmm_output) + device = detections.xyxy.device + xyxy = torch.as_tensor( + np.asarray(nmm_output.xyxy), dtype=torch.float32, device=device + ).reshape(-1, 4) + class_id = torch.as_tensor( + np.asarray(nmm_output.class_id), dtype=torch.long, device=device + ) + confidence = torch.as_tensor( + np.asarray(nmm_output.confidence), dtype=torch.float32, device=device + ) + if is_instance_segmentation: + if nmm_output.mask is not None: + mask = torch.as_tensor( + np.asarray(nmm_output.mask), dtype=torch.bool, device=device + ) + else: + mask = torch.zeros((number_of_detections, 0, 0), dtype=torch.bool) + return InstanceDetections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + mask=mask, + image_metadata=None, + bboxes_metadata=None, + ) + return Detections( + xyxy=xyxy, + class_id=class_id, + confidence=confidence, + image_metadata=None, + bboxes_metadata=None, + ) + + +def _make_instance_detections( + masks: np.ndarray, + confidence: np.ndarray, + class_id: np.ndarray, + xyxy: Optional[np.ndarray] = None, +) -> InstanceDetections: + if xyxy is None: + xyxy = _boxes_from_masks(masks) + return InstanceDetections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32), + class_id=torch.as_tensor(class_id, dtype=torch.long), + confidence=torch.as_tensor(confidence, dtype=torch.float32), + mask=torch.as_tensor(masks, dtype=torch.bool), + image_metadata=None, + bboxes_metadata=None, + ) + + +def _boxes_from_masks(masks: np.ndarray) -> np.ndarray: + boxes = np.zeros((masks.shape[0], 4), dtype=np.float32) + for index, mask in enumerate(masks): + ys, xs = np.where(mask) + if len(ys) == 0: + continue + boxes[index] = [xs.min(), ys.min(), xs.max() + 1, ys.max() + 1] + return boxes + + +def _random_blob_mask( + rng: random.Random, height: int, width: int, kind: str +) -> np.ndarray: + mask = np.zeros((height, width), dtype=bool) + if kind == "empty": + return mask + if kind == "full": + mask[:] = True + return mask + if kind == "rect": + x0 = rng.randrange(0, max(1, width - 2)) + y0 = rng.randrange(0, max(1, height - 2)) + x1 = rng.randrange(x0 + 1, width + 1) + y1 = rng.randrange(y0 + 1, height + 1) + mask[y0:y1, x0:x1] = True + return mask + # "circle" + cy = rng.uniform(0, height) + cx = rng.uniform(0, width) + radius = rng.uniform(2, max(3, min(height, width) / 2)) + yy, xx = np.mgrid[0:height, 0:width] + mask[(yy - cy) ** 2 + (xx - cx) ** 2 <= radius**2] = True + return mask + + +def _random_case( + seed: int, +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + rng = random.Random(seed) + height, width = rng.choice([(64, 96), (128, 128), (150, 200), (97, 61), (700, 900)]) + n = rng.choice([1, 2, 3, 5, 8, 12, 20, 30]) + kinds = ["rect", "rect", "circle", "empty", "full"] + masks = np.zeros((n, height, width), dtype=bool) + for index in range(n): + kind = rng.choice(kinds) + masks[index] = _random_blob_mask(rng, height, width, kind) + # Sometimes nest: replace with a strict subset of a previous mask. + if index > 0 and rng.random() < 0.2 and masks[index - 1].any(): + ys, xs = np.where(masks[index - 1]) + y_mid = int(np.median(ys)) + x_mid = int(np.median(xs)) + nested = np.zeros_like(masks[index - 1]) + nested[ys.min() : y_mid + 1, xs.min() : x_mid + 1] = masks[index - 1][ + ys.min() : y_mid + 1, xs.min() : x_mid + 1 + ] + masks[index] = nested + # Sometimes duplicate a previous mask exactly (IoU == 1 pairs). + if index > 0 and rng.random() < 0.1: + masks[index] = masks[rng.randrange(0, index)] + confidence = np.array([rng.uniform(0.05, 0.999) for _ in range(n)]) + # Introduce exact-tie confidences occasionally. + if n > 1 and rng.random() < 0.3: + confidence = np.round(confidence, 1) + 0.05 + number_of_classes = rng.choice([1, 1, 2, 3]) + class_id = np.array([rng.randrange(0, number_of_classes) for _ in range(n)]) + return masks, confidence, class_id + + +def _assert_same_result( + result: InstanceDetections, expected: InstanceDetections +) -> None: + assert isinstance(result, InstanceDetections) + assert len(result) == len(expected) + assert torch.equal(result.xyxy, expected.xyxy) + assert torch.equal(result.class_id, expected.class_id) + assert torch.equal(result.confidence, expected.confidence) + assert torch.equal(result.mask, expected.mask) + + +def _forbid_sv_fallback(monkeypatch) -> None: + def _fail(*args, **kwargs): + raise AssertionError( + "torch NMM fast path unexpectedly fell back to the sv implementation" + ) + + monkeypatch.setattr(stitch_module, "_with_nmm_sv", _fail) + + +@pytest.mark.parametrize("seed", list(range(60))) +@pytest.mark.parametrize("threshold", [0.2, 0.5]) +def test_nmm_torch_port_fuzz_matches_sv_oracle( + seed: int, threshold: float, monkeypatch +) -> None: + # given + masks, confidence, class_id = _random_case(seed=seed) + detections = _make_instance_detections( + masks=masks, confidence=confidence, class_id=class_id + ) + expected = _oracle_with_nmm(detections=deepcopy(detections), threshold=threshold) + _forbid_sv_fallback(monkeypatch) + + # when + result = with_nmm(detections=deepcopy(detections), threshold=threshold) + + # then + _assert_same_result(result=result, expected=expected) + + +@pytest.mark.parametrize("seed", list(range(60, 90))) +@pytest.mark.parametrize("threshold", [0.0, 0.05, 0.3, 0.75, 0.9, 1.0]) +def test_nmm_torch_port_fuzz_varied_thresholds( + seed: int, threshold: float, monkeypatch +) -> None: + # given + masks, confidence, class_id = _random_case(seed=seed) + detections = _make_instance_detections( + masks=masks, confidence=confidence, class_id=class_id + ) + expected = _oracle_with_nmm(detections=deepcopy(detections), threshold=threshold) + _forbid_sv_fallback(monkeypatch) + + # when + result = with_nmm(detections=deepcopy(detections), threshold=threshold) + + # then + _assert_same_result(result=result, expected=expected) + + +def test_nmm_torch_port_transitive_union_chain(monkeypatch) -> None: + # given: A-B and B-C overlap heavily, A-C do not; with a low threshold the + # growing-union semantics of sv._group_overlapping_masks must chain them. + height, width = 100, 300 + masks = np.zeros((3, height, width), dtype=bool) + masks[0, 20:80, 0:120] = True + masks[1, 20:80, 60:180] = True + masks[2, 20:80, 120:300] = True + confidence = np.array([0.9, 0.8, 0.7]) + class_id = np.array([0, 0, 0]) + detections = _make_instance_detections( + masks=masks, confidence=confidence, class_id=class_id + ) + for threshold in [0.1, 0.2, 0.4, 0.6]: + expected = _oracle_with_nmm( + detections=deepcopy(detections), threshold=threshold + ) + _forbid_sv_fallback(monkeypatch) + + # when + result = with_nmm(detections=deepcopy(detections), threshold=threshold) + + # then + _assert_same_result(result=result, expected=expected) + + +def test_nmm_torch_port_single_detection_passthrough(monkeypatch) -> None: + # given + masks = np.zeros((1, 50, 70), dtype=bool) + masks[0, 10:30, 20:60] = True + detections = _make_instance_detections( + masks=masks, confidence=np.array([0.77]), class_id=np.array([1]) + ) + expected = _oracle_with_nmm(detections=deepcopy(detections), threshold=0.3) + _forbid_sv_fallback(monkeypatch) + + # when + result = with_nmm(detections=deepcopy(detections), threshold=0.3) + + # then + _assert_same_result(result=result, expected=expected) + assert len(result) == 1 + + +def test_nmm_torch_port_all_empty_masks(monkeypatch) -> None: + # given: unions are all zero -> IoU 0 -> nothing merges above a positive + # threshold, every detection survives untouched. + masks = np.zeros((4, 40, 40), dtype=bool) + xyxy = np.array( + [[0, 0, 10, 10], [5, 5, 15, 15], [20, 20, 30, 30], [1, 1, 2, 2]], + dtype=np.float32, + ) + detections = _make_instance_detections( + masks=masks, + confidence=np.array([0.9, 0.8, 0.7, 0.6]), + class_id=np.array([0, 0, 1, 1]), + xyxy=xyxy, + ) + expected = _oracle_with_nmm(detections=deepcopy(detections), threshold=0.3) + _forbid_sv_fallback(monkeypatch) + + # when + result = with_nmm(detections=deepcopy(detections), threshold=0.3) + + # then + _assert_same_result(result=result, expected=expected) + assert len(result) == 4 + + +def test_nmm_torch_port_disjoint_masks_survive(monkeypatch) -> None: + # given + masks = np.zeros((3, 60, 60), dtype=bool) + masks[0, 0:10, 0:10] = True + masks[1, 20:30, 20:30] = True + masks[2, 40:50, 40:50] = True + detections = _make_instance_detections( + masks=masks, + confidence=np.array([0.5, 0.6, 0.7]), + class_id=np.array([0, 0, 0]), + ) + expected = _oracle_with_nmm(detections=deepcopy(detections), threshold=0.3) + _forbid_sv_fallback(monkeypatch) + + # when + result = with_nmm(detections=deepcopy(detections), threshold=0.3) + + # then + _assert_same_result(result=result, expected=expected) + assert len(result) == 3 + + +def test_nmm_empty_detections_returned_as_is() -> None: + # given + detections = InstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + mask=torch.zeros((0, 0, 0), dtype=torch.bool), + image_metadata=None, + bboxes_metadata=None, + ) + + # when + result = with_nmm(detections=detections, threshold=0.3) + + # then + assert result is detections + + +def test_nmm_rle_masks_route_through_sv_path_and_match_dense_oracle() -> None: + # given: RLE-mask inputs keep the pre-port behavior (host-side decode + sv + # NMM); the result must match the oracle run on the equivalent dense masks. + masks, confidence, class_id = _random_case(seed=1234) + dense = _make_instance_detections( + masks=masks, confidence=confidence, class_id=class_id + ) + rle_counts = [ + torch_mask_to_coco_rle(torch.as_tensor(mask, dtype=torch.bool))["counts"] + for mask in masks + ] + rle = InstanceDetections( + xyxy=dense.xyxy.clone(), + class_id=dense.class_id.clone(), + confidence=dense.confidence.clone(), + mask=InstancesRLEMasks(image_size=masks.shape[1:], masks=rle_counts), + image_metadata=None, + bboxes_metadata=None, + ) + expected = _oracle_with_nmm(detections=deepcopy(dense), threshold=0.3) + + # when + result = with_nmm(detections=rle, threshold=0.3) + + # then + _assert_same_result(result=result, expected=expected) + + +def test_nmm_bbox_only_detections_match_oracle() -> None: + # given: no masks -> the sv box-NMM path is kept verbatim. + xyxy = np.array( + [ + [0, 0, 100, 100], + [10, 10, 110, 110], + [200, 200, 300, 300], + [205, 205, 295, 295], + [400, 0, 500, 80], + ], + dtype=np.float32, + ) + detections = Detections( + xyxy=torch.as_tensor(xyxy, dtype=torch.float32), + class_id=torch.as_tensor([0, 0, 1, 1, 0], dtype=torch.long), + confidence=torch.as_tensor([0.9, 0.85, 0.7, 0.95, 0.5], dtype=torch.float32), + image_metadata=None, + bboxes_metadata=None, + ) + expected = _oracle_with_nmm(detections=deepcopy(detections), threshold=0.3) + + # when + result = with_nmm(detections=deepcopy(detections), threshold=0.3) + + # then + assert isinstance(result, Detections) + assert torch.equal(result.xyxy, expected.xyxy) + assert torch.equal(result.class_id, expected.class_id) + assert torch.equal(result.confidence, expected.confidence) + + +def test_nms_branch_smoke() -> None: + # given: the NMS branch is untouched by the NMM port โ€” smoke-check that two + # heavily-overlapping same-class boxes collapse to the higher-confidence one + # and that surviving mask rows are the original rows. + masks = np.zeros((3, 50, 50), dtype=bool) + masks[0, 0:20, 0:20] = True + masks[1, 1:21, 1:21] = True + masks[2, 30:45, 30:45] = True + detections = _make_instance_detections( + masks=masks, + confidence=np.array([0.9, 0.6, 0.8]), + class_id=np.array([0, 0, 1]), + xyxy=np.array( + [[0, 0, 20, 20], [1, 1, 21, 21], [30, 30, 45, 45]], dtype=np.float32 + ), + ) + + # when + result = with_nms(detections=deepcopy(detections), threshold=0.5) + + # then + assert len(result) == 2 + assert torch.equal( + result.xyxy, + torch.as_tensor([[0, 0, 20, 20], [30, 30, 45, 45]], dtype=torch.float32), + ) + assert torch.equal(result.class_id, torch.as_tensor([0, 1], dtype=torch.long)) + assert torch.equal( + result.confidence, torch.as_tensor([0.9, 0.8], dtype=torch.float32) + ) + assert torch.equal(result.mask[0], torch.as_tensor(masks[0], dtype=torch.bool)) + assert torch.equal(result.mask[1], torch.as_tensor(masks[2], dtype=torch.bool)) + + +def test_nmm_chunked_pairwise_intersection_matches_unchunked(monkeypatch) -> None: + # given: force the tiled matmul path of the pairwise-intersection helper and + # confirm decisions stay identical (counts are exact either way). + masks, confidence, class_id = _random_case(seed=4321) + detections = _make_instance_detections( + masks=masks, confidence=confidence, class_id=class_id + ) + expected = _oracle_with_nmm(detections=deepcopy(detections), threshold=0.25) + monkeypatch.setattr(stitch_module, "_NMM_PAIRWISE_FLOAT_BUDGET_BYTES", 512 * 1024) + _forbid_sv_fallback(monkeypatch) + + # when + result = with_nmm(detections=deepcopy(detections), threshold=0.25) + + # then + _assert_same_result(result=result, expected=expected) diff --git a/tests/workflows/unit_tests/core_steps/fusion/test_frame_delay_v1_tensor.py b/tests/workflows/unit_tests/core_steps/fusion/test_frame_delay_v1_tensor.py new file mode 100644 index 0000000000..a8354936f2 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/fusion/test_frame_delay_v1_tensor.py @@ -0,0 +1,330 @@ +"""Tests for the tensor-data-representation sibling of the frame_delay block. + +The sibling reuses the numpy block's buffering machinery verbatim; the only +behavioral change is that image payloads are spilled to host memory at +buffer-insertion time so that CUDA tensors (e.g. Jetson bridge-pool buffers) +are released instead of being pinned for the delay window. The numpy test +module (`test_frame_delay.py`) is flag-agnostic and keeps covering the shared +delay semantics in both flag directions; the tests below cover the tensor +sibling only, hence the `_TENSOR_ONLY` marker. +""" + +import datetime +import gc +import weakref + +import numpy as np +import pytest +import torch +from pydantic import ValidationError + +from inference.core.env import ( + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_IMAGE_TENSOR_DEVICE, +) +from inference.core.workflows.core_steps.fusion.frame_delay.v1 import ( + MAX_OFFSET, + MAX_TRACKED_VIDEOS, +) +from inference.core.workflows.core_steps.fusion.frame_delay.v1 import ( + FrameDelayBlockV1 as NumpyFrameDelayBlockV1, +) +from inference.core.workflows.core_steps.fusion.frame_delay.v1_tensor import ( + BlockManifest, + FrameDelayBlockV1, + _spill_images_to_host, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + VideoMetadata, + WorkflowImageData, +) + +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + + +def _video_metadata(frame_number: int, video_id: str = "vid_1") -> VideoMetadata: + return VideoMetadata( + video_identifier=video_id, + frame_number=frame_number, + frame_timestamp=datetime.datetime.fromtimestamp(1726570800).astimezone( + tz=datetime.timezone.utc + ), + fps=30, + comes_from_video_file=True, + ) + + +def _metadata_image(frame_number: int, video_id: str = "vid_1") -> WorkflowImageData: + """The `image` input of the block: only provides video metadata.""" + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="metadata_provider"), + numpy_image=np.zeros((2, 2, 3), dtype=np.uint8), + video_metadata=_video_metadata(frame_number=frame_number, video_id=video_id), + ) + + +def _tensor_born_image(frame_number: int = 0) -> WorkflowImageData: + """A frame that exists only as a CHW RGB uint8 tensor (R=10, G=20, B=30), + like the frames handed over by the Jetson tensor bridge.""" + chw = torch.zeros((3, 4, 6), dtype=torch.uint8) + chw[0].fill_(10) + chw[1].fill_(20) + chw[2].fill_(30) + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="tensor_frame"), + tensor_image=chw, + video_metadata=_video_metadata(frame_number=frame_number), + ) + + +@_TENSOR_ONLY +def test_delay_semantics_match_numpy_sibling_for_non_image_payloads() -> None: + # given + numpy_block = NumpyFrameDelayBlockV1() + tensor_block = FrameDelayBlockV1() + + # when - identical monotonic frame sequence on both siblings + for n in range(6): + expected = numpy_block.run(image=_metadata_image(n), data=f"det-{n}", offset=-2) + actual = tensor_block.run(image=_metadata_image(n), data=f"det-{n}", offset=-2) + + # then - result dicts are identical frame by frame + assert actual == expected + + +@_TENSOR_ONLY +def test_non_image_payloads_are_buffered_untouched() -> None: + # given - a payload holding GPU-representable data that is NOT an image; + # per the block's documented policy it must be stored as-is (delaying + # mask-carrying predictions still retains their memory) + block = FrameDelayBlockV1() + payload = {"mask": torch.ones((2, 3), dtype=torch.bool), "value": 42} + + # when + result = block.run(image=_metadata_image(0), data=payload, offset=0) + + # then - the very same object is buffered and emitted + assert result["output"] is payload + + +@_TENSOR_ONLY +def test_image_payload_is_spilled_to_host_when_buffered() -> None: + # given + block = FrameDelayBlockV1() + image = _tensor_born_image(frame_number=0) + + # when + result = block.run(image=_metadata_image(0), data=image, offset=0) + emitted = result["output"] + + # then - the emitted object is a host-resident copy, not the input + assert isinstance(emitted, WorkflowImageData) + assert emitted is not image + assert emitted.is_tensor_materialised() is False + assert emitted._tensor_image is None + assert emitted._numpy_image is not None + # CHW RGB (10, 20, 30) -> HWC BGR (30, 20, 10) + assert emitted.numpy_image.shape == (4, 6, 3) + assert np.all(emitted.numpy_image == np.array([30, 20, 10], dtype=np.uint8)) + # lineage and video metadata are carried over exactly + assert emitted._parent_metadata is image._parent_metadata + assert emitted._workflow_root_ancestor_metadata is ( + image._workflow_root_ancestor_metadata + ) + assert emitted._video_metadata is image._video_metadata + # the buffer holds the spilled copy, and the input image is left untouched + assert block._buffers["vid_1"][0] is emitted + assert image.is_tensor_materialised() is True + + +@_TENSOR_ONLY +def test_spill_releases_the_buffered_tensor_storage() -> None: + # given - the block buffers the frame, then every external reference to + # the tensor is dropped (as happens when the bridge frame goes out of + # scope after the workflow step) + block = FrameDelayBlockV1() + tensor = torch.full((3, 4, 6), 7, dtype=torch.uint8) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="tensor_frame"), + tensor_image=tensor, + video_metadata=_video_metadata(frame_number=0), + ) + stored_tensor = image._tensor_image + tensor_ref = weakref.ref(stored_tensor) + + # when + result = block.run(image=_metadata_image(0), data=image, offset=0) + del image, tensor, stored_tensor + gc.collect() + + # then - nothing in the block retains the tensor (best-effort refcount + # check via weakref), and the emitted numpy buffer owns its memory rather + # than aliasing the (now released) tensor storage + assert tensor_ref() is None + assert result["output"].numpy_image.flags["OWNDATA"] is True + + +@_TENSOR_ONLY +def test_emitted_image_re_materialises_tensor_lazily() -> None: + # given + block = FrameDelayBlockV1() + image = _tensor_born_image(frame_number=0) + expected_pixels = image._tensor_image.detach().to("cpu").clone() + + # when + emitted = block.run(image=_metadata_image(0), data=image, offset=0)["output"] + + # then - the tensor is rebuilt lazily, on the configured device, with the + # exact original pixels (CHW RGB -> HWC BGR -> CHW RGB round trip) + assert emitted.is_tensor_materialised() is False + re_uploaded = emitted.tensor_image + assert emitted.is_tensor_materialised() is True + assert re_uploaded.device.type == WORKFLOWS_IMAGE_TENSOR_DEVICE.type + assert torch.equal(re_uploaded.detach().to("cpu"), expected_pixels) + + +@_TENSOR_ONLY +def test_images_inside_list_payloads_are_spilled() -> None: + # given - LIST_OF_VALUES payloads (e.g. a collapsed list of crops) may mix + # images with other values + block = FrameDelayBlockV1() + tensor_image = _tensor_born_image(frame_number=0) + host_image = _metadata_image(0) + payload = [tensor_image, "det-0", host_image] + + # when + emitted = block.run(image=_metadata_image(0), data=payload, offset=0)["output"] + + # then - only the tensor-materialised image is replaced by a spilled copy + assert emitted is not payload + assert emitted[0] is not tensor_image + assert emitted[0].is_tensor_materialised() is False + assert emitted[0]._numpy_image is not None + assert emitted[1] == "det-0" + assert emitted[2] is host_image + + +@_TENSOR_ONLY +def test_containers_without_tensor_images_keep_identity() -> None: + # given + host_image = _metadata_image(0) + payload = [host_image, "det-0", 42] + + # when + spilled = _spill_images_to_host(data=payload) + + # then - nothing to spill, the original container object is kept + assert spilled is payload + + +@_TENSOR_ONLY +def test_images_nested_in_dicts_are_not_spilled() -> None: + # given - documented policy: only top-level images and images inside + # list/tuple containers are spilled; dict payloads are opaque + tensor_image = _tensor_born_image(frame_number=0) + payload = {"frame": tensor_image} + + # when + spilled = _spill_images_to_host(data=payload) + + # then + assert spilled is payload + assert tensor_image.is_tensor_materialised() is True + + +@_TENSOR_ONLY +def test_buffer_is_bounded_on_tensor_sibling() -> None: + # given + block = FrameDelayBlockV1() + offset = -3 + + # when + for n in range(500): + block.run(image=_metadata_image(n), data=f"det-{n}", offset=offset) + + # then - eviction machinery inherited from the numpy sibling is intact + buffer = block._buffers["vid_1"] + assert len(buffer) == abs(offset) + 1 + assert sorted(buffer) == [496, 497, 498, 499] + + +@_TENSOR_ONLY +def test_buffer_is_cleared_when_frame_numbers_restart_on_tensor_sibling() -> None: + # given + block = FrameDelayBlockV1() + for n in range(1000, 1010): + block.run(image=_metadata_image(n), data=f"old-{n}", offset=-1) + + # when + result = block.run(image=_metadata_image(0), data="new-0", offset=-1) + + # then + assert result["is_available"] is False + assert list(block._buffers["vid_1"]) == [0] + + +@_TENSOR_ONLY +def test_inactive_video_buffers_are_evicted_on_tensor_sibling() -> None: + # given + block = FrameDelayBlockV1() + + # when + for stream in range(MAX_TRACKED_VIDEOS + 5): + block.run( + image=_metadata_image(0, video_id=f"vid_{stream}"), data="det-0", offset=-1 + ) + + # then + assert len(block._buffers) == MAX_TRACKED_VIDEOS + assert "vid_0" not in block._buffers + assert f"vid_{MAX_TRACKED_VIDEOS + 4}" in block._buffers + + +@_TENSOR_ONLY +def test_positive_offset_rejected_at_runtime_on_tensor_sibling() -> None: + # given + block = FrameDelayBlockV1() + + # when / then + with pytest.raises(ValueError): + block.run(image=_metadata_image(7), data="det-7", offset=10) + + +@_TENSOR_ONLY +def test_offset_validation_enforced_by_tensor_sibling_manifest() -> None: + # when / then - the inherited manifest validators are active + with pytest.raises(ValidationError): + BlockManifest( + type="roboflow_core/frame_delay@v1", + image="$inputs.image", + data="$steps.model.predictions", + offset=10, + ) + with pytest.raises(ValidationError): + BlockManifest( + type="roboflow_core/frame_delay@v1", + image="$inputs.image", + data="$steps.model.predictions", + offset=-(MAX_OFFSET + 1), + ) + manifest = BlockManifest( + type="roboflow_core/frame_delay@v1", + name="frame_delay", + image="$inputs.image", + data="$steps.model.predictions", + offset=-1, + ) + assert manifest.offset == -1 + + +@_TENSOR_ONLY +def test_tensor_sibling_manifest_masquerades_as_the_same_block() -> None: + # then - same block identity, extended documentation + schema_extra = BlockManifest.model_config["json_schema_extra"] + assert schema_extra["name"] == "Frame Delay" + assert schema_extra["version"] == "v1" + assert "Tensor Data Representation" in schema_extra["long_description"] diff --git a/tests/workflows/unit_tests/core_steps/fusion/test_overlap_analysis_v1_tensor.py b/tests/workflows/unit_tests/core_steps/fusion/test_overlap_analysis_v1_tensor.py new file mode 100644 index 0000000000..39d5166c55 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/fusion/test_overlap_analysis_v1_tensor.py @@ -0,0 +1,309 @@ +from typing import Any, Dict, List, Optional, Union + +import numpy as np +import pytest +import supervision as sv +import torch +from shapely.geometry import Polygon, box + +import inference.core.workflows.core_steps.fusion.overlap_analysis.v1_tensor as overlap_module +from inference.core.workflows.core_steps.common.tensor_native import ( + instance_mask_to_numpy, +) +from inference.core.workflows.core_steps.fusion.overlap_analysis.v1_tensor import ( + OverlapAnalysisBlockV1, + _class_names, + _detection_ids, + _safe_get, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks +from inference_models.models.common.rle_utils import torch_mask_to_coco_rle + +NativeDetections = Union[Detections, InstanceDetections] + + +def _oracle_detection_to_shapely( + detections: NativeDetections, xyxy: np.ndarray, idx: int +) -> Polygon: + x1, y1, x2, y2 = xyxy[idx] + bbox_poly = box(float(x1), float(y1), float(x2), float(y2)) + if isinstance(detections, InstanceDetections) and detections.mask is not None: + mask = instance_mask_to_numpy(detections, idx) + if np.any(mask): + polygons = sv.mask_to_polygons(mask=mask.astype(np.uint8)) + if polygons: + longest = max(polygons, key=len) + if len(longest) >= 3: + candidate = Polygon( + [(float(pt[0]), float(pt[1])) for pt in longest] + ) + if candidate.is_valid and not candidate.is_empty: + return candidate + return bbox_poly + + +def _oracle_run( + reference_predictions: NativeDetections, + candidate_predictions: NativeDetections, + min_overlap: float, +) -> Dict[str, List[Dict[str, Any]]]: + if len(reference_predictions) == 0 or len(candidate_predictions) == 0: + return {"overlaps": []} + + ref_xyxy = reference_predictions.xyxy.detach().to("cpu").numpy() + cand_xyxy = candidate_predictions.xyxy.detach().to("cpu").numpy() + + iou_matrix = sv.box_iou_batch(ref_xyxy, cand_xyxy) + + ref_ids = _detection_ids(reference_predictions) + cand_ids = _detection_ids(candidate_predictions) + ref_class_names = _class_names(reference_predictions) + cand_class_names = _class_names(candidate_predictions) + + results: List[Dict[str, Any]] = [] + ref_polys: Dict[int, Polygon] = {} + cand_polys: Dict[int, Polygon] = {} + + for i in range(len(reference_predictions)): + for j in range(len(candidate_predictions)): + if iou_matrix[i, j] <= 0.0: + continue + if i not in ref_polys: + ref_polys[i] = _oracle_detection_to_shapely( + reference_predictions, ref_xyxy, i + ) + if j not in cand_polys: + cand_polys[j] = _oracle_detection_to_shapely( + candidate_predictions, cand_xyxy, j + ) + ref_poly = ref_polys[i] + cand_poly = cand_polys[j] + if ref_poly.area <= 0: + continue + intersection_area = ref_poly.intersection(cand_poly).area + overlap_ratio = intersection_area / ref_poly.area + if overlap_ratio < min_overlap: + continue + record: Dict[str, Any] = { + "reference_class": _safe_get(ref_class_names, i), + "reference_confidence": ( + float(reference_predictions.confidence[i]) + if reference_predictions.confidence is not None + else None + ), + "candidate_class": _safe_get(cand_class_names, j), + "candidate_confidence": ( + float(candidate_predictions.confidence[j]) + if candidate_predictions.confidence is not None + else None + ), + "overlap_ratio": float(overlap_ratio), + } + if ref_ids is not None: + record["reference_detection_id"] = _safe_get(ref_ids, i) + if cand_ids is not None: + record["candidate_detection_id"] = _safe_get(cand_ids, j) + results.append(record) + return {"overlaps": results} + + +def _assert_results_equal( + actual: Dict[str, List[Dict[str, Any]]], + expected: Dict[str, List[Dict[str, Any]]], + *, + exact_ratio: bool, +) -> None: + assert len(actual["overlaps"]) == len(expected["overlaps"]) + for actual_record, expected_record in zip(actual["overlaps"], expected["overlaps"]): + assert list(actual_record) == list(expected_record) + for key in actual_record: + if key == "overlap_ratio" and not exact_ratio: + assert actual_record[key] == pytest.approx( + expected_record[key], rel=1e-9, abs=0.0 + ) + else: + assert actual_record[key] == expected_record[key] + + +def _random_boxes(rng: np.random.Generator, count: int) -> np.ndarray: + centers = rng.normal(30.0, 8.0, size=(count, 2)) + sizes = rng.uniform(0.25, 35.0, size=(count, 2)) + boxes = np.column_stack( + ( + centers[:, 0] - sizes[:, 0] / 2, + centers[:, 1] - sizes[:, 1] / 2, + centers[:, 0] + sizes[:, 0] / 2, + centers[:, 1] + sizes[:, 1] / 2, + ) + ).astype(np.float32) + if count: + reverse_x = rng.random(count) < 0.25 + reverse_y = rng.random(count) < 0.25 + boxes[reverse_x, 0], boxes[reverse_x, 2] = ( + boxes[reverse_x, 2].copy(), + boxes[reverse_x, 0].copy(), + ) + boxes[reverse_y, 1], boxes[reverse_y, 3] = ( + boxes[reverse_y, 3].copy(), + boxes[reverse_y, 1].copy(), + ) + degenerate = np.arange(count) % 7 == 0 + boxes[degenerate, 2] = boxes[degenerate, 0] + horizontal = np.arange(count) % 11 == 0 + boxes[horizontal, 3] = boxes[horizontal, 1] + return boxes + + +def _make_bbox_detections( + boxes: np.ndarray, + *, + seed: int, + instance_without_masks: bool, + include_class_names: bool, + include_ids: bool, + confidence_is_none: bool, +) -> NativeDetections: + count = len(boxes) + class_ids = torch.arange(count, dtype=torch.int64) % 5 + confidence: Optional[torch.Tensor] + confidence = ( + None + if confidence_is_none + else torch.linspace(0.05, 0.95, max(count, 1), dtype=torch.float32)[:count] + ) + image_metadata = ( + {"class_names": {idx: f"seed_{seed}_class_{idx}" for idx in range(5)}} + if include_class_names + else None + ) + bboxes_metadata = None + if include_ids: + bboxes_metadata = [ + {"detection_id": f"seed_{seed}_detection_{idx}"} if idx % 3 else {} + for idx in range(count) + ] + kwargs = dict( + xyxy=torch.from_numpy(boxes), + class_id=class_ids, + confidence=confidence, + image_metadata=image_metadata, + bboxes_metadata=bboxes_metadata, + ) + if instance_without_masks: + return InstanceDetections(mask=None, **kwargs) + return Detections(**kwargs) + + +@pytest.mark.parametrize("seed", range(120)) +def test_bbox_path_matches_shapely_oracle_across_seeded_scenarios(seed: int) -> None: + rng = np.random.default_rng(seed) + reference = _make_bbox_detections( + _random_boxes(rng, int(rng.integers(0, 61))), + seed=seed, + instance_without_masks=seed % 3 == 0, + include_class_names=seed % 4 != 0, + include_ids=seed % 5 in {0, 1}, + confidence_is_none=seed % 13 == 0, + ) + candidate = _make_bbox_detections( + _random_boxes(rng, int(rng.integers(0, 61))), + seed=seed + 1000, + instance_without_masks=seed % 3 == 1, + include_class_names=seed % 4 != 1, + include_ids=seed % 5 in {1, 2}, + confidence_is_none=seed % 17 == 0, + ) + min_overlap = (0.0, 0.1, 0.5, 1.0)[seed % 4] + + expected = _oracle_run(reference, candidate, min_overlap) + actual = OverlapAnalysisBlockV1().run(reference, candidate, min_overlap) + + _assert_results_equal(actual, expected, exact_ratio=False) + + +def _dense_mask_detections() -> InstanceDetections: + masks = torch.zeros((3, 32, 32), dtype=torch.bool) + masks[0, 2:7, 2:7] = True + masks[0, 12:29, 10:28] = True + masks[2, 7:25, 4:23] = True + return InstanceDetections( + xyxy=torch.tensor( + [[0, 0, 31, 31], [2, 2, 20, 20], [4, 5, 29, 29]], + dtype=torch.float32, + ), + class_id=torch.tensor([0, 1, 2]), + confidence=torch.tensor([0.25, 0.5, 0.75]), + mask=masks, + image_metadata={"class_names": {0: "zero", 1: "one", 2: "two"}}, + bboxes_metadata=[{"detection_id": "a"}, {}, {"detection_id": "c"}], + ) + + +def _as_rle(detections: InstanceDetections) -> InstanceDetections: + encoded = [torch_mask_to_coco_rle(mask) for mask in detections.mask] + return InstanceDetections( + xyxy=detections.xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + mask=InstancesRLEMasks.from_coco_rle_masks((32, 32), encoded), + image_metadata=detections.image_metadata, + bboxes_metadata=detections.bboxes_metadata, + ) + + +def _as_bboxes(detections: InstanceDetections) -> Detections: + return Detections( + xyxy=detections.xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + image_metadata=detections.image_metadata, + bboxes_metadata=detections.bboxes_metadata, + ) + + +@pytest.mark.parametrize( + ("reference_kind", "candidate_kind"), + [ + ("dense", "dense"), + ("rle", "rle"), + ("bbox", "dense"), + ("dense", "bbox"), + ], +) +def test_mask_paths_match_shapely_oracle( + reference_kind: str, candidate_kind: str +) -> None: + dense_reference = _dense_mask_detections() + dense_candidate = _dense_mask_detections() + dense_candidate.xyxy = dense_candidate.xyxy + torch.tensor([2.0, 0.0, 2.0, 0.0]) + variants = { + "dense": lambda detections: detections, + "rle": _as_rle, + "bbox": _as_bboxes, + } + reference = variants[reference_kind](dense_reference) + candidate = variants[candidate_kind](dense_candidate) + + expected = _oracle_run(reference, candidate, 0.1) + actual = OverlapAnalysisBlockV1().run(reference, candidate, 0.1) + + _assert_results_equal(actual, expected, exact_ratio=True) + + +def test_dense_mask_path_does_not_call_instance_mask_to_numpy(monkeypatch) -> None: + reference = _dense_mask_detections() + candidate = _dense_mask_detections() + expected = _oracle_run(reference, candidate, 0.1) + + def fail_if_called(*args, **kwargs): + raise AssertionError("instance_mask_to_numpy must not be called") + + monkeypatch.setattr( + overlap_module, "instance_mask_to_numpy", fail_if_called, raising=False + ) + + actual = OverlapAnalysisBlockV1().run(reference, candidate, 0.1) + + _assert_results_equal(actual, expected, exact_ratio=True) diff --git a/tests/workflows/unit_tests/core_steps/integrations/roboflow/visual_search_classifier/test_v1.py b/tests/workflows/unit_tests/core_steps/integrations/roboflow/visual_search_classifier/test_v1.py index 17b9879a5b..75671f334f 100644 --- a/tests/workflows/unit_tests/core_steps/integrations/roboflow/visual_search_classifier/test_v1.py +++ b/tests/workflows/unit_tests/core_steps/integrations/roboflow/visual_search_classifier/test_v1.py @@ -6,6 +6,7 @@ import numpy as np import pytest +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.utils.image_utils import load_image_base64 from inference.core.workflows.core_steps.common.query_language.operations.core import ( execute_operations, @@ -275,6 +276,13 @@ def test_run_returns_multi_label_predictions_when_candidate_has_multiple_classes } +@pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="numpy block's dict output through the UQL extractor is a numpy-only " + "combination; flag-on runs the v1_tensor sibling (native ClassificationPrediction) " + "and the UQL native extractor rejects dicts โ€” the known classification-deserializer " + "seam, tracked separately", +) def test_run_returns_multi_label_predictions_compatible_with_all_classes_extraction() -> ( None ): @@ -515,10 +523,19 @@ def test_run_returns_error_when_project_search_fails() -> None: def test_run_batch_returns_one_classification_per_image() -> None: block = RoboflowVisualSearchClassifierBlockV1(api_key="api-key") + # thread-pooled searches reach the mock out of submission order - key by image + responses_by_image = { + "image-1": { + "results": [make_candidate(image_id="img-1", class_name="a", class_id=0)] + }, + "image-2": { + "results": [make_candidate(image_id="img-2", class_name="b", class_id=1)] + }, + } + with mock.patch.object(v1, "search_project_images_at_roboflow") as search_mock: - search_mock.side_effect = [ - {"results": [make_candidate(image_id="img-1", class_name="a", class_id=0)]}, - {"results": [make_candidate(image_id="img-2", class_name="b", class_id=1)]}, + search_mock.side_effect = lambda **kwargs: responses_by_image[ + kwargs["image_base64"] ] result = block.run( @@ -529,8 +546,10 @@ def test_run_batch_returns_one_classification_per_image() -> None: assert [row["predictions"]["top"] for row in result] == ["a", "b"] assert [row["inference_id"] for row in result] - assert search_mock.call_args_list[0].kwargs["image_base64"] == "image-1" - assert search_mock.call_args_list[1].kwargs["image_base64"] == "image-2" + assert {call.kwargs["image_base64"] for call in search_mock.call_args_list} == { + "image-1", + "image-2", + } def test_run_batch_searches_images_in_parallel_and_preserves_output_order() -> None: diff --git a/tests/workflows/unit_tests/core_steps/models/foundation/test_depth_estimation.py b/tests/workflows/unit_tests/core_steps/models/foundation/test_depth_estimation.py index 74bad909e2..5199e45a69 100644 --- a/tests/workflows/unit_tests/core_steps/models/foundation/test_depth_estimation.py +++ b/tests/workflows/unit_tests/core_steps/models/foundation/test_depth_estimation.py @@ -4,18 +4,36 @@ import numpy as np import pytest +import torch from pydantic import ValidationError +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION +from inference.core.utils.depth_encoding import ( + decode_png_normalized_depth, + encode_normalized_depth_to_png16, +) from inference.core.workflows.core_steps.common.entities import StepExecutionMode from inference.core.workflows.core_steps.models.foundation.depth_estimation.v1 import ( BlockManifest, DepthEstimationBlockV1, ) +from inference.core.workflows.core_steps.models.foundation.depth_estimation.v1_tensor import ( + DepthEstimationBlockV1 as DepthEstimationBlockV1Tensor, +) from inference.core.workflows.execution_engine.entities.base import ( ImageParentMetadata, WorkflowImageData, ) +# The numpy tests below import `v1` explicitly and drive it through mocked +# managers/clients, so they pass in both flag directions and stay unmarked. +# The tensor-native sibling tests target `v1_tensor` and only need to hold +# under ENABLE_TENSOR_DATA_REPRESENTATION. +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + @pytest.fixture def mock_model_manager(): @@ -150,3 +168,207 @@ def test_run_locally_with_yolo26_depth_variant( model_id="yolo26n-depth-768", api_key="test_key" ) mock_model_manager.infer_from_request_sync.assert_called_once() + + +@pytest.fixture +def mock_tensor_native_model_manager(): + """ModelManager mock for the tensor-native path: the block calls + `run_tensor_native_inference` and receives raw per-image torch depth maps + (larger == closer - the DepthAnything convention all tensor-native depth + models / adapters share).""" + mock = MagicMock(spec=["add_model", "run_tensor_native_inference"]) + mock.run_tensor_native_inference.return_value = [ + torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + ] + return mock + + +@_TENSOR_ONLY +def test_run_locally_tensor_native_normalizes_per_image( + mock_tensor_native_model_manager, mock_workflow_image_data +): + block = DepthEstimationBlockV1Tensor( + model_manager=mock_tensor_native_model_manager, + api_key="test_api_key", + step_execution_mode=StepExecutionMode.LOCAL, + ) + + result = block.run( + images=[mock_workflow_image_data], + model_version="depth-anything-v3/small", + ) + + assert len(result) == 1 + mock_tensor_native_model_manager.add_model.assert_called_once_with( + model_id="depth-anything-v3/small", api_key="test_api_key" + ) + call = mock_tensor_native_model_manager.run_tensor_native_inference.call_args + assert call.args == ("depth-anything-v3/small",) + assert call.kwargs["input_color_format"] == "bgr" + assert len(call.kwargs["images"]) == 1 + assert isinstance(call.kwargs["images"][0], np.ndarray) + normalized_depth = result[0]["normalized_depth"] + assert isinstance(normalized_depth, torch.Tensor) + # min-max normalization of a larger-means-closer raw map: 1.0 == nearest + assert torch.allclose( + normalized_depth.cpu(), + torch.tensor([[0.0, 1 / 3], [2 / 3, 1.0]]), + atol=1e-6, + ) + assert isinstance(result[0]["image"], WorkflowImageData) + + +@_TENSOR_ONLY +def test_run_locally_tensor_native_uses_materialised_tensor_image( + mock_tensor_native_model_manager, +): + tensor_image = torch.zeros((3, 4, 4), dtype=torch.uint8) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="some"), + tensor_image=tensor_image, + ) + block = DepthEstimationBlockV1Tensor( + model_manager=mock_tensor_native_model_manager, + api_key="test_api_key", + step_execution_mode=StepExecutionMode.LOCAL, + ) + + result = block.run(images=[image], model_version="depth-anything-v3/small") + + assert len(result) == 1 + call = mock_tensor_native_model_manager.run_tensor_native_inference.call_args + assert call.kwargs["input_color_format"] == "rgb" + assert isinstance(call.kwargs["images"][0], torch.Tensor) + + +@_TENSOR_ONLY +def test_run_locally_tensor_native_raises_on_zero_variation_depth_map( + mock_tensor_native_model_manager, mock_workflow_image_data +): + # numpy parity: flag-off the same ValueError is raised inside the model's + # predict() and propagates out of run_locally() + mock_tensor_native_model_manager.run_tensor_native_inference.return_value = [ + torch.ones((4, 4)) + ] + block = DepthEstimationBlockV1Tensor( + model_manager=mock_tensor_native_model_manager, + api_key="test_api_key", + step_execution_mode=StepExecutionMode.LOCAL, + ) + + with pytest.raises(ValueError, match="min equals max"): + block.run( + images=[mock_workflow_image_data], + model_version="depth-anything-v3/small", + ) + + +@_TENSOR_ONLY +def test_run_locally_tensor_native_with_yolo26_depth_variant( + mock_tensor_native_model_manager, mock_workflow_image_data +): + """YOLO26 ids run natively through `run_tensor_native_inference`; the + metric-to-proximity flip lives in InferenceModelsDepthEstimationAdapter, + so the block-side handling is identical for every model version.""" + block = DepthEstimationBlockV1Tensor( + model_manager=mock_tensor_native_model_manager, + api_key="test_key", + step_execution_mode=StepExecutionMode.LOCAL, + ) + result = block.run( + images=[mock_workflow_image_data], + model_version="yolo26n-depth-768", + ) + assert len(result) == 1 + mock_tensor_native_model_manager.add_model.assert_called_once_with( + model_id="yolo26n-depth-768", api_key="test_key" + ) + call = mock_tensor_native_model_manager.run_tensor_native_inference.call_args + assert call.args == ("yolo26n-depth-768",) + + +@_TENSOR_ONLY +@patch( + "inference.core.workflows.core_steps.models.foundation.depth_estimation.v1_tensor.InferenceHTTPClient" +) +def test_run_remotely_tensor_decodes_png16_payload( + mock_client_cls, mock_workflow_image_data +): + """The remote path mirrors numpy v1: request png16, receive the ndarray the + SDK decoded from the base64 PNG16 payload, and (tensor-side) convert it to + a torch tensor.""" + original_map = np.linspace(0.0, 1.0, num=24, dtype=np.float32).reshape(4, 6) + decoded_by_sdk = decode_png_normalized_depth( + encode_normalized_depth_to_png16(original_map) + ) + mock_client = MagicMock() + mock_client.depth_estimation.return_value = { + "normalized_depth": decoded_by_sdk, + "image": "0000", # hex-encoded empty image + } + mock_client_cls.return_value = mock_client + + block = DepthEstimationBlockV1Tensor( + model_manager=MagicMock(), + api_key="test_api_key", + step_execution_mode=StepExecutionMode.REMOTE, + ) + + result = block.run( + images=[mock_workflow_image_data], + model_version="depth-anything-v3/small", + ) + + assert len(result) == 1 + mock_client.depth_estimation.assert_called_once_with( + inference_input=mock_workflow_image_data.base64_image, + model_id="depth-anything-v3/small", + model_id_in_path=True, + depth_map_format="png16", + ) + normalized_depth = result[0]["normalized_depth"] + assert isinstance(normalized_depth, torch.Tensor) + assert normalized_depth.dtype == torch.float32 + # png16 quantization step is 1/65535 + assert torch.allclose( + normalized_depth, + torch.as_tensor(original_map), + atol=1.0 / 65535, + ) + + +@_TENSOR_ONLY +@patch( + "inference.core.workflows.core_steps.models.foundation.depth_estimation.v1_tensor.InferenceHTTPClient" +) +def test_run_remotely_tensor_handles_legacy_float_list_payload( + mock_client_cls, mock_workflow_image_data +): + """Servers that predate `depth_map_format` return the nested float list; + `np.array` handles both shapes, exactly like numpy v1.""" + mock_client = MagicMock() + mock_client.depth_estimation.return_value = { + "normalized_depth": [[0.1, 0.2], [0.3, 0.4]], + "image": "0000", + } + mock_client_cls.return_value = mock_client + + block = DepthEstimationBlockV1Tensor( + model_manager=MagicMock(), + api_key="test_api_key", + step_execution_mode=StepExecutionMode.REMOTE, + ) + + result = block.run( + images=[mock_workflow_image_data], + model_version="depth-anything-v3/small", + ) + + assert len(result) == 1 + normalized_depth = result[0]["normalized_depth"] + assert isinstance(normalized_depth, torch.Tensor) + assert torch.allclose( + normalized_depth, + torch.tensor([[0.1, 0.2], [0.3, 0.4]]), + atol=1e-6, + ) diff --git a/tests/workflows/unit_tests/core_steps/models/foundation/test_pp_ocr.py b/tests/workflows/unit_tests/core_steps/models/foundation/test_pp_ocr.py index e81f47937a..07366acabe 100644 --- a/tests/workflows/unit_tests/core_steps/models/foundation/test_pp_ocr.py +++ b/tests/workflows/unit_tests/core_steps/models/foundation/test_pp_ocr.py @@ -171,9 +171,7 @@ def test_pp_ocr_run_locally_full_mode() -> None: images = _make_images() # when - result = block.run( - images=images, text_detection="small", text_recognition="small" - ) + result = block.run(images=images, text_detection="small", text_recognition="small") # then model_manager.infer_from_request_sync.assert_called_once() @@ -219,9 +217,7 @@ def test_pp_ocr_run_locally_detect_only() -> None: images = _make_images() # when - result = block.run( - images=images, text_detection="small", text_recognition="none" - ) + result = block.run(images=images, text_detection="small", text_recognition="none") # then assert len(result) == 1 @@ -247,9 +243,7 @@ def test_pp_ocr_run_locally_recognize_only() -> None: images = _make_images() # when - result = block.run( - images=images, text_detection="none", text_recognition="small" - ) + result = block.run(images=images, text_detection="none", text_recognition="small") # then assert len(result) == 1 diff --git a/tests/workflows/unit_tests/core_steps/sampling/test_identify_changes.py b/tests/workflows/unit_tests/core_steps/sampling/test_identify_changes.py index 33c436955b..63ccf820c4 100644 --- a/tests/workflows/unit_tests/core_steps/sampling/test_identify_changes.py +++ b/tests/workflows/unit_tests/core_steps/sampling/test_identify_changes.py @@ -1,4 +1,5 @@ import numpy as np +import pytest from inference.core.workflows.core_steps.sampling.identify_changes.v1 import ( IdentifyChangesBlockV1, @@ -63,3 +64,43 @@ def test_identify_changes() -> None: # average and std should not be zero anymore assert not np.allclose(result.get("average"), initial_value_normalized) assert not np.all(result.get("std") == [0, 0, 0, 0, 0]) + + +def test_identify_changes_zero_variance_is_not_outlier() -> None: + # given + identify_changes_block = IdentifyChangesBlockV1() + embedding = np.array([0.1, -0.4, 0.3, 0.9, -0.2]) + + # when - second run scores against fresh statistics (std == 0), hitting the + # zero-std guard; exactly two runs keeps std platform-exact (more steps flake) + inputs = {**default_inputs, "warmup": 1} + for _ in range(2): + result = identify_changes_block.run(**inputs, embedding=embedding) + + # then + assert not result.get("warming_up") + assert not result.get("is_outlier") + assert result.get("percentile") == 0.5 + assert result.get("z_score") == 0 + + +def test_identify_changes_tensor_sibling_zero_variance_is_not_outlier() -> None: + # given + torch = pytest.importorskip("torch") + from inference.core.workflows.core_steps.sampling.identify_changes.v1_tensor import ( + IdentifyChangesBlockV1 as IdentifyChangesTensorBlockV1, + ) + + identify_changes_block = IdentifyChangesTensorBlockV1() + embedding = torch.tensor([0.1, -0.4, 0.3, 0.9, -0.2]) + + # when - two runs against fresh statistics (std == 0) hit the zero-std guard + inputs = {**default_inputs, "warmup": 1} + for _ in range(2): + result = identify_changes_block.run(**inputs, embedding=embedding) + + # then + assert not result.get("warming_up") + assert not result.get("is_outlier") + assert result.get("percentile") == 0.5 + assert result.get("z_score") == 0 diff --git a/tests/workflows/unit_tests/core_steps/sinks/roboflow/roboflow_dataset_upload/test_v1.py b/tests/workflows/unit_tests/core_steps/sinks/roboflow/roboflow_dataset_upload/test_v1.py index 2d234cb5bd..2e931f5064 100644 --- a/tests/workflows/unit_tests/core_steps/sinks/roboflow/roboflow_dataset_upload/test_v1.py +++ b/tests/workflows/unit_tests/core_steps/sinks/roboflow/roboflow_dataset_upload/test_v1.py @@ -11,6 +11,7 @@ from fastapi import BackgroundTasks from inference.core.cache import MemoryCache +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload import v1 from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1 import ( BatchCreationFrequency, @@ -22,12 +23,20 @@ is_prediction_registration_forbidden, register_datapoint, ) +from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v1_tensor import ( + RoboflowDatasetUploadBlockV1 as TensorRoboflowDatasetUploadBlockV1, +) from inference.core.workflows.execution_engine.entities.base import ( Batch, ImageParentMetadata, WorkflowImageData, ) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + def test_encode_prediction_when_classification_prediction_provided() -> None: # given @@ -899,6 +908,42 @@ def test_execution_policy_noops_without_api_key() -> None: ] +@_TENSOR_ONLY +def test_execution_policy_noops_without_api_key_tensor_native() -> None: + data_collector_block = TensorRoboflowDatasetUploadBlockV1( + cache=MemoryCache(), + api_key=None, + background_tasks=None, + thread_pool_executor=None, + disable_sinks=True, + ) + + result = data_collector_block.run( + images=Batch(content=[MagicMock()], indices=[(0,)]), + predictions=Batch(content=[MagicMock()], indices=[(0,)]), + target_project="my_project", + usage_quota_name="my_quota", + persist_predictions=True, + minutely_usage_limit=10, + hourly_usage_limit=100, + daily_usage_limit=1000, + max_image_size=(128, 128), + compression_level=75, + registration_tags=["some"], + disable_sink=False, + fire_and_forget=True, + labeling_batch_prefix="my_batch", + labeling_batches_recreation_frequency="never", + ) + + assert result == [ + { + "error_status": False, + "message": "Sink was disabled by workflow execution policy", + } + ] + + def test_run_sink_when_sink_is_disabled_by_configuration() -> None: # given data_collector_block = RoboflowDatasetUploadBlockV1( diff --git a/tests/workflows/unit_tests/core_steps/sinks/roboflow/roboflow_dataset_upload/test_v2.py b/tests/workflows/unit_tests/core_steps/sinks/roboflow/roboflow_dataset_upload/test_v2.py index c2446dc437..72fa643d5d 100644 --- a/tests/workflows/unit_tests/core_steps/sinks/roboflow/roboflow_dataset_upload/test_v2.py +++ b/tests/workflows/unit_tests/core_steps/sinks/roboflow/roboflow_dataset_upload/test_v2.py @@ -8,18 +8,27 @@ from pydantic import ValidationError from inference.core.cache import MemoryCache +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload import v2 from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v2 import ( BlockManifest, RoboflowDatasetUploadBlockV2, maybe_register_datapoint_at_roboflow, ) +from inference.core.workflows.core_steps.sinks.roboflow.dataset_upload.v2_tensor import ( + RoboflowDatasetUploadBlockV2 as TensorRoboflowDatasetUploadBlockV2, +) from inference.core.workflows.execution_engine.entities.base import ( Batch, ImageParentMetadata, WorkflowImageData, ) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) + @mock.patch.object(v2, "register_datapoint_at_roboflow") @mock.patch.object(v2, "random") @@ -385,6 +394,43 @@ def test_execution_policy_noops_without_api_key() -> None: ] +@_TENSOR_ONLY +def test_execution_policy_noops_without_api_key_tensor_native() -> None: + data_collector_block = TensorRoboflowDatasetUploadBlockV2( + cache=MemoryCache(), + api_key=None, + background_tasks=None, + thread_pool_executor=None, + disable_sinks=True, + ) + + result = data_collector_block.run( + images=Batch(content=[MagicMock()], indices=[(0,)]), + predictions=None, + target_project="my_project", + usage_quota_name="my_quota", + data_percentage=100.0, + persist_predictions=True, + minutely_usage_limit=10, + hourly_usage_limit=100, + daily_usage_limit=1000, + max_image_size=(128, 128), + compression_level=75, + registration_tags=["some"], + disable_sink=False, + fire_and_forget=False, + labeling_batch_prefix="my_batch", + labeling_batches_recreation_frequency="never", + ) + + assert result == [ + { + "error_status": False, + "message": "Sink was disabled by workflow execution policy", + } + ] + + @mock.patch.object(v2, "register_datapoint_at_roboflow") def test_run_sink_when_data_sampled_off( register_datapoint_at_roboflow_mock: MagicMock, diff --git a/tests/workflows/unit_tests/core_steps/transformations/test_absolute_static_crop.py b/tests/workflows/unit_tests/core_steps/transformations/test_absolute_static_crop.py index e6d0ad29e2..2652615cf0 100644 --- a/tests/workflows/unit_tests/core_steps/transformations/test_absolute_static_crop.py +++ b/tests/workflows/unit_tests/core_steps/transformations/test_absolute_static_crop.py @@ -1,6 +1,7 @@ from datetime import datetime import numpy as np +import pytest from inference.core.workflows.core_steps.transformations.absolute_static_crop.v1 import ( take_static_crop, @@ -85,3 +86,123 @@ def test_take_absolute_static_crop_when_output_crop_is_empty() -> None: # then assert result is None, "Expected no crop as result" + + +def test_take_absolute_static_crop_clamps_out_of_bounds_crop() -> None: + # given + np_image = np.zeros((100, 100, 3), dtype=np.uint8) + np_image[0:20, 0:10] = 30 + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="origin_image"), + numpy_image=np_image, + ) + + # when - crop centred at the image corner extends past the top-left edge + result = take_static_crop( + image=image, + x_center=0, + y_center=0, + width=20, + height=40, + ) + + # then - clamped to the in-bounds region instead of wrapping via negative indices + assert result is not None, "Expected clamped crop as result" + assert result.numpy_image.shape == (20, 10, 3) + assert ( + result.numpy_image == (np.ones((20, 10, 3), dtype=np.uint8) * 30) + ).all(), "Expected the top-left in-bounds region to be returned" + offset = result.parent_metadata.origin_coordinates + assert (offset.left_top_x, offset.left_top_y) == (0, 0) + + +def test_take_absolute_static_crop_tensor_sibling_clamps_out_of_bounds_crop() -> None: + # given + pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.transformations.absolute_static_crop.v1_tensor import ( + take_static_crop as take_static_crop_tensor, + ) + + np_image = np.zeros((100, 100, 3), dtype=np.uint8) + np_image[0:20, 0:10] = 30 + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="origin_image"), + numpy_image=np_image, + ) + + # when + result = take_static_crop_tensor( + image=image, + x_center=0, + y_center=0, + width=20, + height=40, + ) + + # then + assert result is not None, "Expected clamped crop as result" + assert result.numpy_image.shape == (20, 10, 3) + assert ( + result.numpy_image == (np.ones((20, 10, 3), dtype=np.uint8) * 30) + ).all(), "Expected the top-left in-bounds region to be returned" + offset = result.parent_metadata.origin_coordinates + assert (offset.left_top_x, offset.left_top_y) == (0, 0) + + +def test_take_absolute_static_crop_tensor_sibling_clamps_on_tensor_path() -> None: + # given + torch = pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.transformations.absolute_static_crop.v1_tensor import ( + take_static_crop as take_static_crop_tensor, + ) + + tensor_image = torch.zeros((3, 100, 100), dtype=torch.uint8) + tensor_image[:, 0:20, 0:10] = 30 + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="origin_image"), + tensor_image=tensor_image, + ) + + # when + result = take_static_crop_tensor( + image=image, + x_center=0, + y_center=0, + width=20, + height=40, + ) + + # then + assert result is not None, "Expected clamped crop as result" + assert tuple(result.tensor_image.shape) == (3, 20, 10) + assert bool((result.tensor_image == 30).all()) + + +def test_take_absolute_static_crop_tensor_sibling_returns_none_when_crop_fully_outside() -> ( + None +): + # given + pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.transformations.absolute_static_crop.v1_tensor import ( + take_static_crop as take_static_crop_tensor, + ) + + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="origin_image"), + numpy_image=np.zeros((100, 100, 3), dtype=np.uint8), + ) + + # when - the whole crop lies left of / above the image + result = take_static_crop_tensor( + image=image, + x_center=-50, + y_center=-50, + width=20, + height=20, + ) + + # then + assert result is None, "Expected no crop as result" diff --git a/tests/workflows/unit_tests/core_steps/transformations/test_byte_tracker_tensor.py b/tests/workflows/unit_tests/core_steps/transformations/test_byte_tracker_tensor.py new file mode 100644 index 0000000000..8545142fbd --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/transformations/test_byte_tracker_tensor.py @@ -0,0 +1,367 @@ +"""Parity tests for the tensor-native ByteTracker blocks (v1/v2/v3). + +These blocks used to do a genuine GPU->CPU->GPU round trip: the surviving +detections were rebuilt by re-uploading ``sv_tracked.xyxy`` / ``class_id`` / +``confidence`` from numpy back to the device. ByteTrack only *filters and +reorders* detections (it never mutates coordinates), so the re-upload is pure +waste - the surviving rows can be sliced straight out of the original device +tensors with ``take_prediction_by_indices`` (masks included). + +The oracle embedded here (``_oracle_track``) is the ORIGINAL implementation: it +D2H-materialises the input, tags rows with a positional index, runs ByteTrack, +then re-uploads the numpy boxes. Each test drives the real (new) block and the +oracle over the *same* multi-frame sequence, each with its own fresh +deterministic tracker, and asserts: + +* identical surviving rows (by input ``detection_id``) in identical order, +* identical box / class / confidence tensors, +* identical ``tracker_id`` values written into ``bboxes_metadata``, +* for instance-segmentation inputs, masks that follow the SAME index selection + (the round-tripping oracle drops masks, so masks are checked against the input + masks sliced by the surviving rows), +* for v3, identical ``new_instances`` / ``already_seen_instances`` splits. + +Torch runs on CPU here, which is enough to exercise the selection logic; the +device-transfer win is orthogonal to correctness. +""" + +import datetime +from typing import List, Optional, Tuple, Union + +import numpy as np +import supervision as sv +import torch + +from inference.core.workflows.core_steps.transformations.byte_tracker.v1_tensor import ( + ByteTrackerBlockV1, +) +from inference.core.workflows.core_steps.transformations.byte_tracker.v2_tensor import ( + ByteTrackerBlockV2, +) +from inference.core.workflows.core_steps.transformations.byte_tracker.v3_tensor import ( + ByteTrackerBlockV3, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + VideoMetadata, + WorkflowImageData, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +_ORACLE_INDEX_KEY = "__oracle_input_index__" + + +# --------------------------------------------------------------------------- # +# Oracle (the original round-tripping implementation) +# --------------------------------------------------------------------------- # +def _oracle_track( + tracker: sv.ByteTrack, + detections: Union[Detections, InstanceDetections], +) -> Tuple[Detections, List[int], np.ndarray]: + """The pre-change implementation: D2H, run ByteTrack, re-upload numpy boxes. + + Returns the reconstructed plain ``Detections`` (masks dropped, as the old + block did) plus the surviving input indices and their tracker ids. + """ + n = int(detections.xyxy.shape[0]) + sv_input = sv.Detections( + xyxy=detections.xyxy.detach().to("cpu").numpy(), + class_id=detections.class_id.detach().to("cpu").numpy(), + confidence=detections.confidence.detach().to("cpu").numpy(), + data={_ORACLE_INDEX_KEY: np.arange(n, dtype=np.int64)}, + ) + sv_tracked = tracker.update_with_detections(sv_input) + kept_indices = ( + sv_tracked.data.get(_ORACLE_INDEX_KEY, np.empty((0,), dtype=np.int64)) + if sv_tracked.data + else np.empty((0,), dtype=np.int64) + ) + tracker_ids = ( + sv_tracked.tracker_id + if sv_tracked.tracker_id is not None + else np.full(len(sv_tracked), -1, dtype=np.int64) + ) + original_meta = detections.bboxes_metadata or [{} for _ in range(n)] + new_bboxes_meta: List[dict] = [] + for new_i, orig_i in enumerate(kept_indices): + base = dict(original_meta[int(orig_i)] or {}) + base["tracker_id"] = int(tracker_ids[new_i]) + new_bboxes_meta.append(base) + device = detections.xyxy.device + out = Detections( + xyxy=torch.from_numpy(sv_tracked.xyxy).to( + device=device, dtype=detections.xyxy.dtype + ), + class_id=torch.from_numpy(sv_tracked.class_id).to( + device=device, dtype=detections.class_id.dtype + ), + confidence=torch.from_numpy(sv_tracked.confidence).to( + device=device, dtype=detections.confidence.dtype + ), + image_metadata=detections.image_metadata, + bboxes_metadata=new_bboxes_meta if new_bboxes_meta else None, + ) + return out, [int(i) for i in kept_indices.tolist()], tracker_ids + + +class _OracleCache: + """Independent copy of the v3 InstanceCache new/seen bookkeeping.""" + + def __init__(self, size: int) -> None: + from collections import deque + + size = max(1, size) + self._order = deque(maxlen=size) + self._seen = set() + + def record(self, tracker_id: int) -> bool: + in_cache = tracker_id in self._seen + if not in_cache: + while len(self._seen) >= self._order.maxlen: + self._seen.remove(self._order.popleft()) + self._order.append(tracker_id) + self._seen.add(tracker_id) + return in_cache + + +# --------------------------------------------------------------------------- # +# Builders +# --------------------------------------------------------------------------- # +def _meta(video_id: str, frame_number: int, fps: float = 1.0) -> VideoMetadata: + return VideoMetadata( + video_identifier=video_id, + frame_number=frame_number, + fps=fps, + frame_timestamp=datetime.datetime.fromtimestamp(1726570875).astimezone( + tz=datetime.timezone.utc + ), + comes_from_video_file=True, + ) + + +def _wrap_image(metadata: VideoMetadata) -> WorkflowImageData: + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="root"), + numpy_image=np.zeros((128, 128, 3), dtype=np.uint8), + video_metadata=metadata, + ) + + +def _boxes_metadata(count: int, prefix: str) -> List[dict]: + return [{"detection_id": f"{prefix}-{i}"} for i in range(count)] + + +def _object_detections(boxes: List[List[float]], prefix: str) -> Detections: + n = len(boxes) + return Detections( + xyxy=torch.tensor(boxes, dtype=torch.float32), + class_id=torch.arange(n, dtype=torch.long) % 3, + confidence=torch.full((n,), 0.9, dtype=torch.float32), + image_metadata={"class_names": {0: "a", 1: "b", 2: "c"}}, + bboxes_metadata=_boxes_metadata(n, prefix), + ) + + +def _instance_detections(boxes: List[List[float]], prefix: str) -> InstanceDetections: + n = len(boxes) + height, width = 40, 40 + # Give every instance a distinct dense mask so the selection is verifiable. + mask = torch.zeros((n, height, width), dtype=torch.bool) + for i, box in enumerate(boxes): + x1, y1, x2, y2 = (int(round(v)) for v in box) + x1 = max(0, min(width - 1, x1)) + x2 = max(x1 + 1, min(width, x2)) + y1 = max(0, min(height - 1, y1)) + y2 = max(y1 + 1, min(height, y2)) + mask[i, y1:y2, x1:x2] = True + return InstanceDetections( + xyxy=torch.tensor(boxes, dtype=torch.float32), + class_id=torch.arange(n, dtype=torch.long) % 3, + confidence=torch.full((n,), 0.9, dtype=torch.float32), + mask=mask, + image_metadata={"class_names": {0: "a", 1: "b", 2: "c"}}, + bboxes_metadata=_boxes_metadata(n, prefix), + ) + + +def _keypoint_prediction( + boxes: List[List[float]], prefix: str +) -> Tuple[KeyPoints, Detections]: + det = _object_detections(boxes, prefix) + n = len(boxes) + key_points = KeyPoints( + xy=torch.zeros((n, 3, 2), dtype=torch.float32), + class_id=det.class_id.clone(), + confidence=torch.ones((n, 3), dtype=torch.float32), + image_metadata={"class_names": {0: "a", 1: "b", 2: "c"}}, + ) + return key_points, det + + +# A 3-frame sequence: 4 objects persist and drift, one (last) leaves after +# frame 1, and a brand-new object enters on frame 3 (entering / leaving / +# persisting all exercised). +_FRAME_BOXES: List[List[List[float]]] = [ + [[10, 10, 20, 20], [21, 10, 31, 20], [31, 10, 41, 20], [2, 2, 8, 8]], + [[12, 10, 22, 20], [23, 10, 33, 20], [33, 10, 43, 20]], + [[14, 10, 24, 20], [25, 10, 35, 20], [35, 10, 45, 20], [1, 30, 9, 38]], +] + + +def _det_ids(detections: Union[Detections, InstanceDetections]) -> List[str]: + if detections.bboxes_metadata is None: + return [] + return [m.get("detection_id") for m in detections.bboxes_metadata] + + +def _tracker_ids(detections: Union[Detections, InstanceDetections]) -> List[int]: + if detections.bboxes_metadata is None: + return [] + return [int(m["tracker_id"]) for m in detections.bboxes_metadata] + + +def _assert_box_parity( + produced: Union[Detections, InstanceDetections], + oracle: Detections, + input_detections: Union[Detections, InstanceDetections], + kept_indices: List[int], +) -> None: + assert torch.equal(produced.xyxy, oracle.xyxy) + assert torch.equal(produced.class_id, oracle.class_id) + assert torch.equal(produced.confidence, oracle.confidence) + assert _tracker_ids(produced) == _tracker_ids(oracle) + expected_ids = [_det_ids(input_detections)[i] for i in kept_indices] + assert _det_ids(produced) == expected_ids + if isinstance(input_detections, InstanceDetections): + # Old block dropped masks; the new block must carry them, sliced by the + # SAME surviving rows. + assert isinstance(produced, InstanceDetections) + selector = torch.tensor(kept_indices, dtype=torch.long) + assert torch.equal(produced.mask, input_detections.mask[selector]) + else: + assert isinstance(produced, Detections) + assert not isinstance(produced, InstanceDetections) + + +# --------------------------------------------------------------------------- # +# v1 +# --------------------------------------------------------------------------- # +def _run_v1(builder) -> None: + block = ByteTrackerBlockV1() + oracle_tracker = sv.ByteTrack(frame_rate=1) + for frame_number, boxes in enumerate(_FRAME_BOXES): + detections = builder(boxes, prefix=f"f{frame_number}") + produced = block.run( + metadata=_meta("vid", frame_number, fps=1), + detections=detections, + )["tracked_detections"] + oracle_out, kept, _ = _oracle_track(oracle_tracker, detections) + _assert_box_parity(produced, oracle_out, detections, kept) + + +def test_v1_object_detection_parity() -> None: + _run_v1(_object_detections) + + +def test_v1_instance_segmentation_mask_selection() -> None: + _run_v1(_instance_detections) + + +def test_v1_all_surviving_frame_reuses_input_tensor_no_reupload() -> None: + # Frame 1: every high-confidence detection is activated and matched, so all + # rows survive in input order -> identity selection returns the ORIGINAL + # device tensor object (proves no re-upload / no copy). + block = ByteTrackerBlockV1() + detections = _object_detections(_FRAME_BOXES[0], prefix="f0") + produced = block.run( + metadata=_meta("vid", 0, fps=1), + detections=detections, + )["tracked_detections"] + assert produced.xyxy is detections.xyxy + assert produced.class_id is detections.class_id + assert produced.confidence is detections.confidence + + +# --------------------------------------------------------------------------- # +# v2 +# --------------------------------------------------------------------------- # +def _run_v2(builder) -> None: + block = ByteTrackerBlockV2() + oracle_tracker = sv.ByteTrack(frame_rate=1) + for frame_number, boxes in enumerate(_FRAME_BOXES): + detections = builder(boxes, prefix=f"f{frame_number}") + produced = block.run( + image=_wrap_image(_meta("vid", frame_number, fps=1)), + detections=detections, + )["tracked_detections"] + oracle_out, kept, _ = _oracle_track(oracle_tracker, detections) + _assert_box_parity(produced, oracle_out, detections, kept) + + +def test_v2_object_detection_parity() -> None: + _run_v2(_object_detections) + + +def test_v2_instance_segmentation_mask_selection() -> None: + _run_v2(_instance_detections) + + +# --------------------------------------------------------------------------- # +# v3 +# --------------------------------------------------------------------------- # +def _run_v3(builder, keypoint_input: bool = False) -> None: + block = ByteTrackerBlockV3() + oracle_tracker = sv.ByteTrack(frame_rate=1) + oracle_cache = _OracleCache(size=16384) + for frame_number, boxes in enumerate(_FRAME_BOXES): + raw = builder(boxes, prefix=f"f{frame_number}") + # For keypoint input the bbox Detections drives tracking / parity. + bbox_input = raw[1] if keypoint_input else raw + result = block.run( + image=_wrap_image(_meta("vid", frame_number, fps=1)), + detections=raw, + ) + oracle_out, kept, tracker_ids = _oracle_track(oracle_tracker, bbox_input) + _assert_box_parity(result["tracked_detections"], oracle_out, bbox_input, kept) + # new / already-seen split parity. + not_seen, seen = [], [] + for position, tid in enumerate(tracker_ids.tolist()): + (seen if oracle_cache.record(int(tid)) else not_seen).append(position) + assert _tracker_ids(result["new_instances"]) == [ + _tracker_ids(oracle_out)[i] for i in not_seen + ] + assert _tracker_ids(result["already_seen_instances"]) == [ + _tracker_ids(oracle_out)[i] for i in seen + ] + assert _det_ids(result["new_instances"]) == [ + _det_ids(oracle_out)[i] for i in not_seen + ] + assert _det_ids(result["already_seen_instances"]) == [ + _det_ids(oracle_out)[i] for i in seen + ] + + +def test_v3_object_detection_parity_and_splits() -> None: + _run_v3(_object_detections) + + +def test_v3_instance_segmentation_mask_selection() -> None: + _run_v3(_instance_detections) + + +def test_v3_keypoint_input_tracks_on_bbox_and_drops_keypoints() -> None: + _run_v3(_keypoint_prediction, keypoint_input=True) + + +def test_v3_first_frame_all_new_none_seen() -> None: + block = ByteTrackerBlockV3() + detections = _object_detections(_FRAME_BOXES[0], prefix="f0") + result = block.run( + image=_wrap_image(_meta("vid", 0, fps=1)), + detections=detections, + ) + tracked_ids = _tracker_ids(result["tracked_detections"]) + assert _tracker_ids(result["new_instances"]) == tracked_ids + assert len(result["already_seen_instances"]) == 0 diff --git a/tests/workflows/unit_tests/core_steps/transformations/test_detection_offset.py b/tests/workflows/unit_tests/core_steps/transformations/test_detection_offset.py index 87f386c82c..678a8e5eec 100644 --- a/tests/workflows/unit_tests/core_steps/transformations/test_detection_offset.py +++ b/tests/workflows/unit_tests/core_steps/transformations/test_detection_offset.py @@ -226,3 +226,48 @@ def test_run_offsets_by_percentage_when_percent_units_selected() -> None: assert ( y2 == 620 ), "Bottom corner should be moved by 5% of detection height to the bottom" + + +def test_run_offsets_by_percentage_when_percent_units_selected_tensor_native() -> None: + # given - the tensor-native mirror of the numpy percent-mode test above + torch = pytest.importorskip("torch") + pytest.importorskip("inference_models") + from inference.core.workflows.core_steps.transformations.detection_offset.v1_tensor import ( + DetectionOffsetBlockV1 as TensorDetectionOffsetBlockV1, + ) + from inference_models.models.base.object_detection import ( + Detections as NativeDetections, + ) + + detections = NativeDetections( + xyxy=torch.tensor([[100.0, 200.0, 400.0, 600.0]], dtype=torch.float64), + class_id=torch.tensor([1]), + confidence=torch.tensor([0.5], dtype=torch.float64), + image_metadata={ + "class_names": {1: "truck"}, + "image_dimensions": [640, 640], + }, + bboxes_metadata=[ + {"detection_id": "three", "class": "truck", "parent_id": "p3"} + ], + ) + block = TensorDetectionOffsetBlockV1() + + # when + result = block.run( + predictions=Batch(content=[detections], indices=[(0,)]), + offset_width=10, + offset_height=10, + units="Percent (%)", + ) + + # then + x1, y1, x2, y2 = result[0]["predictions"].xyxy[0].tolist() + assert x1 == 85, "Left corner should be moved by 5% of detection width to the left" + assert y1 == 180, "Top corner should be moved by 5% of detection height to the top" + assert ( + x2 == 415 + ), "Right corner should be moved by 5% of detection width to the right" + assert ( + y2 == 620 + ), "Bottom corner should be moved by 5% of detection height to the bottom" diff --git a/tests/workflows/unit_tests/core_steps/transformations/test_dynamic_zones.py b/tests/workflows/unit_tests/core_steps/transformations/test_dynamic_zones.py index e3e2b29ba2..5ca3ae1447 100644 --- a/tests/workflows/unit_tests/core_steps/transformations/test_dynamic_zones.py +++ b/tests/workflows/unit_tests/core_steps/transformations/test_dynamic_zones.py @@ -1,4 +1,5 @@ import numpy as np +import pytest import supervision as sv from inference.core.workflows.core_steps.transformations.dynamic_zones.v1 import ( @@ -328,3 +329,139 @@ def test_calculate_least_squares_polygon_with_midpoint_fraction(): least_squares_polygon, np.array([[1639, 1550], [1567, 1726], [1668, 1837], [1761, 1556]]), ), "Correct least squares polygon should be calculated based on the contour and polygon." + + +def test_dynamic_zones_tensor_native_stores_scaled_polygon_in_bboxes_metadata(): + pytest.importorskip("torch") + pytest.importorskip("inference_models") + import torch + + from inference.core.workflows.core_steps.transformations.dynamic_zones.v1_tensor import ( + OUTPUT_KEY as TENSOR_OUTPUT_KEY, + ) + from inference.core.workflows.core_steps.transformations.dynamic_zones.v1_tensor import ( + OUTPUT_KEY_DETECTIONS as TENSOR_OUTPUT_KEY_DETECTIONS, + ) + from inference.core.workflows.core_steps.transformations.dynamic_zones.v1_tensor import ( + DynamicZonesBlockV1 as TensorDynamicZonesBlockV1, + ) + from inference.core.workflows.execution_engine.constants import ( + POLYGON_KEY_IN_SV_DETECTIONS, + ) + from inference.core.workflows.execution_engine.entities.base import Batch + from inference_models.models.base.instance_segmentation import InstanceDetections + + # given - a single square instance and a scale ratio that visibly moves vertices + mask = np.zeros((1, 100, 100), dtype=bool) + mask[0, 20:60, 30:70] = True + detections = InstanceDetections( + xyxy=torch.tensor([[30.0, 20.0, 70.0, 60.0]]), + class_id=torch.tensor([0]), + confidence=torch.tensor([0.9]), + mask=torch.from_numpy(mask), + image_metadata={"class_names": {0: "zone"}}, + bboxes_metadata=[{"detection_id": "det-1"}], + ) + block = TensorDynamicZonesBlockV1() + + # when + result = block.run( + predictions=Batch(content=[detections], indices=[(0,)]), + required_number_of_vertices=4, + scale_ratio=0.5, + ) + + # then + scaled_polygons = result[0][TENSOR_OUTPUT_KEY] + output_detections = result[0][TENSOR_OUTPUT_KEY_DETECTIONS] + assert len(scaled_polygons) == 1 + stored_polygon = output_detections.bboxes_metadata[0][POLYGON_KEY_IN_SV_DETECTIONS] + assert stored_polygon.shape == ( + len(scaled_polygons[0]), + 2, + ), "Per-box polygon payload must be the (V, 2) polygon itself - numpy's sv COLUMN wrapping does not apply to per-box bboxes_metadata" + assert ( + stored_polygon.tolist() == scaled_polygons[0] + ), "Per-box polygon payload must equal the scaled polygon emitted in the zones output" + unscaled_square = {(30, 20), (69, 20), (69, 59), (30, 59)} + assert ( + set(map(tuple, stored_polygon.tolist())) != unscaled_square + ), "Payload must not be the pre-scale polygon (scale_ratio=0.5 moves vertices)" + assert output_detections.bboxes_metadata[0]["detection_id"] == "det-1" + + +def test_dynamic_zones_tensor_native_serialized_polygon_payload_matches_numpy(): + # a (1, V, 2) per-box polygon payload would bypass the serializer's + # declared-polygon fast path and nest the response field one extra level + pytest.importorskip("torch") + pytest.importorskip("inference_models") + import torch + + from inference.core.workflows.core_steps.common.serializers import ( + serialise_sv_detections as numpy_serialise, + ) + from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_sv_detections as tensor_serialise, + ) + from inference.core.workflows.core_steps.transformations.dynamic_zones.v1 import ( + OUTPUT_KEY_DETECTIONS, + DynamicZonesBlockV1, + ) + from inference.core.workflows.core_steps.transformations.dynamic_zones.v1_tensor import ( + OUTPUT_KEY_DETECTIONS as TENSOR_OUTPUT_KEY_DETECTIONS, + ) + from inference.core.workflows.core_steps.transformations.dynamic_zones.v1_tensor import ( + DynamicZonesBlockV1 as TensorDynamicZonesBlockV1, + ) + from inference.core.workflows.execution_engine.constants import ( + POLYGON_KEY_IN_SV_DETECTIONS, + ) + from inference.core.workflows.execution_engine.entities.base import Batch + from inference_models.models.base.instance_segmentation import InstanceDetections + + # given - identical single-instance input in both representations + mask = np.zeros((1, 100, 100), dtype=bool) + mask[0, 20:60, 30:70] = True + sv_detections = sv.Detections( + xyxy=np.array([[30.0, 20.0, 70.0, 60.0]], dtype=np.float32), + class_id=np.array([0]), + confidence=np.array([0.9], dtype=np.float32), + mask=mask.copy(), + data={ + "class_name": np.array(["zone"]), + "detection_id": np.array(["det-1"]), + }, + ) + native_detections = InstanceDetections( + xyxy=torch.tensor([[30.0, 20.0, 70.0, 60.0]]), + class_id=torch.tensor([0]), + confidence=torch.tensor([0.9]), + mask=torch.from_numpy(mask.copy()), + image_metadata={"class_names": {0: "zone"}}, + bboxes_metadata=[{"detection_id": "det-1"}], + ) + + # when + numpy_result = DynamicZonesBlockV1().run( + predictions=Batch(content=[sv_detections], indices=[(0,)]), + required_number_of_vertices=4, + scale_ratio=0.5, + ) + tensor_result = TensorDynamicZonesBlockV1().run( + predictions=Batch(content=[native_detections], indices=[(0,)]), + required_number_of_vertices=4, + scale_ratio=0.5, + ) + + # then + numpy_serialised = numpy_serialise(numpy_result[0][OUTPUT_KEY_DETECTIONS]) + tensor_serialised = tensor_serialise(tensor_result[0][TENSOR_OUTPUT_KEY_DETECTIONS]) + numpy_prediction = numpy_serialised["predictions"][0] + tensor_prediction = tensor_serialised["predictions"][0] + assert ( + numpy_prediction[POLYGON_KEY_IN_SV_DETECTIONS] + == tensor_prediction[POLYGON_KEY_IN_SV_DETECTIONS] + ), "Declared-polygon response field must match numpy (no extra nesting level)" + assert ( + numpy_serialised == tensor_serialised + ), "Serialized dynamic_zones predictions must be identical across representations" diff --git a/tests/workflows/unit_tests/core_steps/transformations/test_image_slicer_v2.py b/tests/workflows/unit_tests/core_steps/transformations/test_image_slicer_v2.py index 933d063e34..18c6396810 100644 --- a/tests/workflows/unit_tests/core_steps/transformations/test_image_slicer_v2.py +++ b/tests/workflows/unit_tests/core_steps/transformations/test_image_slicer_v2.py @@ -1,5 +1,6 @@ import numpy as np import pytest +import torch from pydantic import ValidationError from inference.core.workflows.core_steps.transformations.image_slicer.v2 import ( @@ -196,6 +197,81 @@ def test_running_block() -> None: ), "Expected 9th crop to have the following coordinates regarding root" +def test_running_block_when_tensor_image_given_keeps_crops_on_device() -> None: + # given - a tensor-backed image (CHW RGB) must be sliced without forcing a + # full-frame device->host copy, and each crop must stay tensor-backed. + tensor_image = torch.arange(3 * 256 * 512, dtype=torch.uint8).reshape(3, 256, 512) + image = WorkflowImageData( + tensor_image=tensor_image, + parent_metadata=ImageParentMetadata(parent_id="parent"), + ) + assert image.is_tensor_materialised() + block = ImageSlicerBlockV2() + + # when + result = block.run( + image=image, + slice_width=200, + slice_height=100, + overlap_ratio_width=0.1, + overlap_ratio_height=0.2, + ) + + # then + assert len(result) == 9, "Expected exactly 9 crops" + # the parent frame must not have been materialised to host by the slicer + assert ( + image._numpy_image is None + ), "Expected parent frame to remain on device (no host round-trip)" + + expected_root_coordinates = [ + OriginCoordinatesSystem( + left_top_x=0, left_top_y=0, origin_width=512, origin_height=256 + ), + OriginCoordinatesSystem( + left_top_x=180, left_top_y=0, origin_width=512, origin_height=256 + ), + OriginCoordinatesSystem( + left_top_x=312, left_top_y=0, origin_width=512, origin_height=256 + ), + OriginCoordinatesSystem( + left_top_x=0, left_top_y=80, origin_width=512, origin_height=256 + ), + OriginCoordinatesSystem( + left_top_x=180, left_top_y=80, origin_width=512, origin_height=256 + ), + OriginCoordinatesSystem( + left_top_x=312, left_top_y=80, origin_width=512, origin_height=256 + ), + OriginCoordinatesSystem( + left_top_x=0, left_top_y=156, origin_width=512, origin_height=256 + ), + OriginCoordinatesSystem( + left_top_x=180, left_top_y=156, origin_width=512, origin_height=256 + ), + OriginCoordinatesSystem( + left_top_x=312, left_top_y=156, origin_width=512, origin_height=256 + ), + ] + for i, expected_coordinates in enumerate(expected_root_coordinates): + crop = result[i]["slices"] + assert crop.parent_metadata.parent_id.startswith( + "image_slicer." + ), f"Expected parent to be set properly for {i}th crop" + # crop must stay tensor-backed - no eager host materialisation + assert crop.is_tensor_materialised(), f"{i}th crop lost its tensor image" + assert crop._numpy_image is None, f"{i}th crop was materialised to host" + assert crop.tensor_image.shape == (3, 100, 200) + assert crop.numpy_image.shape == (100, 200, 3) + assert ( + crop.workflow_root_ancestor_metadata.origin_coordinates + == expected_coordinates + ), f"Expected {i}th crop to have the correct coordinates regarding root" + + # zero-copy content correctness: the first crop equals the source tensor slice + assert torch.equal(result[0]["slices"].tensor_image, tensor_image[:, 0:100, 0:200]) + + def test_running_block_when_slice_size_exceed_image_size() -> None: # given image = WorkflowImageData( diff --git a/tests/workflows/unit_tests/core_steps/transformations/test_stabilize_detections.py b/tests/workflows/unit_tests/core_steps/transformations/test_stabilize_detections.py index 86c01c6d25..98bc705823 100644 --- a/tests/workflows/unit_tests/core_steps/transformations/test_stabilize_detections.py +++ b/tests/workflows/unit_tests/core_steps/transformations/test_stabilize_detections.py @@ -1,4 +1,5 @@ import numpy as np +import pytest import supervision as sv from inference.core.workflows.core_steps.transformations.stabilize_detections.v1 import ( @@ -205,3 +206,48 @@ def test_stabilize_detections(): assert len(res_6["tracked_detections"]) == 2 assert len(res_7["tracked_detections"]) == 2 assert len(res_8["tracked_detections"]) == 0 + + +def test_stabilize_detections_tensor_native_accepts_empty_untracked_detections(): + # fresh block has no cached trackers to gap-fill - empty input yields empty output + pytest.importorskip("torch") + pytest.importorskip("inference_models") + import torch + + from inference.core.workflows.core_steps.transformations.stabilize_detections.v1_tensor import ( + OUTPUT_KEY as TENSOR_OUTPUT_KEY, + ) + from inference.core.workflows.core_steps.transformations.stabilize_detections.v1_tensor import ( + StabilizeTrackedDetectionsBlockV1 as TensorStabilizeBlockV1, + ) + from inference_models.models.base.object_detection import Detections + + # given + empty_detections = Detections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + image_metadata={"class_names": {}}, + bboxes_metadata=[], + ) + block = TensorStabilizeBlockV1() + img = WorkflowImageData( + parent_metadata=ImageParentMetadata(""), + video_metadata=VideoMetadata( + video_identifier="1", + frame_number=1, + frame_timestamp=0, + ), + numpy_image=np.zeros((10, 10, 3)), + ) + + # when - must not raise despite tracker ids being absent + result = block.run( + image=img, + detections=empty_detections, + smoothing_window_size=3, + bbox_smoothing_coefficient=0.5, + ) + + # then + assert int(result[TENSOR_OUTPUT_KEY].xyxy.shape[0]) == 0 diff --git a/tests/workflows/unit_tests/core_steps/visualizations/test_bounding_box_v1_tensor_gpu.py b/tests/workflows/unit_tests/core_steps/visualizations/test_bounding_box_v1_tensor_gpu.py new file mode 100644 index 0000000000..9fc6fb198b --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/visualizations/test_bounding_box_v1_tensor_gpu.py @@ -0,0 +1,361 @@ +import numpy as np +import pytest +import supervision as sv +import torch + +from inference.core.workflows.core_steps.visualizations.bounding_box.v1_tensor import ( + _gpu_box_draw_eligible, + gpu_draw_boxes, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) +from inference_models.models.base.object_detection import Detections + +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" +) + +PALETTE = sv.ColorPalette.DEFAULT +SCENE_H, SCENE_W = 240, 320 + + +def _paint( + xyxy: np.ndarray, + colors_rgb: np.ndarray, + thickness: int, + device: str = "cpu", + roundness: float = 0.0, +) -> tuple: + """Run the painter on a zero scene; return (chw tensor, painted-pixel mask).""" + scene = torch.zeros((3, SCENE_H, SCENE_W), dtype=torch.uint8, device=device) + annotated = gpu_draw_boxes( + scene, xyxy.astype(int), colors_rgb, thickness, roundness + ) + painted = (annotated != 0).any(dim=0).cpu().numpy() + return annotated, painted + + +def _sv_painted_mask(xyxy: np.ndarray, thickness: int) -> np.ndarray: + scene = np.zeros((SCENE_H, SCENE_W, 3), dtype=np.uint8) + annotator = sv.BoxAnnotator( + color=PALETTE, color_lookup=sv.ColorLookup.INDEX, thickness=thickness + ) + out = annotator.annotate(scene, sv.Detections(xyxy=xyxy.astype(np.float32))) + return (out != 0).any(axis=2) + + +_SCENARIOS = { + "plain": np.array([[30, 40, 120, 100], [150, 60, 290, 200]], dtype=float), + "overlap_chain": np.array( + [[30, 40, 160, 160], [80, 90, 200, 210], [100, 20, 140, 230]], dtype=float + ), + "off_image": np.array( + [[-20, -30, 60, 50], [250, 180, 400, 400], [-50, 100, 380, 140]], dtype=float + ), + "tiny_and_degenerate": np.array( + [[10, 10, 12, 12], [50, 50, 50, 60], [70, 70, 70, 70], [90, 5, 95, 9]], + dtype=float, + ), + "edge_touching": np.array( + [[0, 0, SCENE_W - 1, SCENE_H - 1], [5, 5, 15, SCENE_H - 10]], dtype=float + ), +} + + +def _index_colors(n: int) -> np.ndarray: + return np.asarray([PALETTE.by_idx(i).as_rgb() for i in range(n)], dtype=np.uint8) + + +@pytest.mark.parametrize("scenario", sorted(_SCENARIOS)) +@pytest.mark.parametrize("thickness", [1, 2, 3, 5, 8]) +def test_gpu_boxes_visually_match_sv(scenario: str, thickness: int) -> None: + # Approximate renderer: borders are square-cornered `thickness`-wide + # bands, cv2 draws round joins / slightly wider bands. Painted regions + # must still substantially agree (IoU), and t=1 is pixel-identical. + xyxy = _SCENARIOS[scenario] + _, painted = _paint(xyxy, _index_colors(xyxy.shape[0]), thickness) + reference = _sv_painted_mask(xyxy, thickness) + union = (painted | reference).sum() + if union == 0: + return + iou = (painted & reference).sum() / union + threshold = 1.0 if thickness == 1 else 0.5 + assert iou >= threshold, f"painted-region IoU {iou:.3f}" + + +def test_border_band_geometry_exact() -> None: + # One box, thickness 3: band spans [edge - 1, edge + 1] (outer = t // 2), + # interior and background stay untouched. + x1, y1, x2, y2 = 50, 60, 150, 140 + annotated, painted = _paint( + np.array([[x1, y1, x2, y2]], dtype=float), + np.array([[255, 0, 0]], np.uint8), + thickness=3, + ) + assert painted[y1 - 1 : y1 + 2, x1 - 1 : x2 + 2].all() # top band + assert painted[y2 - 1 : y2 + 2, x1 - 1 : x2 + 2].all() # bottom band + assert painted[y1 - 1 : y2 + 2, x1 - 1 : x1 + 2].all() # left band + assert painted[y1 - 1 : y2 + 2, x2 - 1 : x2 + 2].all() # right band + assert not painted[y1 + 2 : y2 - 1, x1 + 2 : x2 - 1].any() # interior clean + assert not painted[: y1 - 1].any() and not painted[y2 + 2 :].any() # outside + assert (annotated[0][painted] == 255).all() # right channel, right color + + +def _sv_round_painted_mask( + xyxy: np.ndarray, thickness: int, roundness: float +) -> np.ndarray: + scene = np.zeros((SCENE_H, SCENE_W, 3), dtype=np.uint8) + annotator = sv.RoundBoxAnnotator( + color=PALETTE, + color_lookup=sv.ColorLookup.INDEX, + thickness=thickness, + roundness=roundness, + ) + out = annotator.annotate(scene, sv.Detections(xyxy=xyxy.astype(np.float32))) + return (out != 0).any(axis=2) + + +@pytest.mark.parametrize("scenario", sorted(_SCENARIOS)) +@pytest.mark.parametrize("roundness", [0.2, 0.5, 1.0]) +def test_gpu_rounded_boxes_visually_match_sv(scenario: str, roundness: float) -> None: + xyxy = _SCENARIOS[scenario] + _, painted = _paint(xyxy, _index_colors(xyxy.shape[0]), 3, roundness=roundness) + reference = _sv_round_painted_mask(xyxy, 3, roundness) + union = (painted | reference).sum() + if union == 0: + return + iou = (painted & reference).sum() / union + assert iou >= 0.45, f"painted-region IoU {iou:.3f}" + + +def test_rounded_corners_are_actually_rounded() -> None: + # Large box, high roundness: the square corner pixel must stay unpainted, + # the edge midpoints and arc diagonals must be painted. + x1, y1, x2, y2 = 40, 40, 200, 200 + radius = (200 - 40) // 2 # roundness=1.0 + _, painted = _paint( + np.array([[x1, y1, x2, y2]], dtype=float), + np.array([[255, 0, 0]], np.uint8), + thickness=3, + roundness=1.0, + ) + assert not painted[y1, x1] # square corner cut away + assert painted[y1, (x1 + x2) // 2] # top edge midpoint + assert painted[(y1 + y2) // 2, x1] # left edge midpoint + diag = int(radius - radius / np.sqrt(2)) + assert painted[y1 + diag, x1 + diag] # 45-degree point on the TL arc + + +def test_overlapping_boxes_later_box_wins() -> None: + # sv paints sequentially: the higher detection index owns contested pixels. + xyxy = np.array([[20, 20, 100, 100], [60, 20, 140, 100]], dtype=float) + colors = np.array([[255, 0, 0], [0, 255, 0]], np.uint8) + annotated, _ = _paint(xyxy, colors, thickness=2) + # box 1's left band crosses box 0's interior row: contested column region + # around x=60 on box 0's top band must be box 1's color where box 1 paints. + top = annotated[:, 20, 60].cpu().numpy() # (3,) at shared corner pixel + assert tuple(top) == (0, 255, 0) + + +def test_fully_off_frame_box_is_noop() -> None: + annotated, painted = _paint( + np.array([[SCENE_W + 10, SCENE_H + 10, SCENE_W + 50, SCENE_H + 50]], float), + np.array([[255, 0, 0]], np.uint8), + thickness=2, + ) + assert not painted.any() + + +def test_gpu_box_paint_is_in_place() -> None: + scene = torch.zeros((3, 64, 64), dtype=torch.uint8) + annotated = gpu_draw_boxes( + scene, + np.array([[10, 10, 40, 40]], dtype=int), + np.array([[255, 0, 0]], dtype=np.uint8), + 2, + ) + assert annotated.data_ptr() == scene.data_ptr() + + +def test_non_contiguous_scene_raises() -> None: + # .view must fail (caught by the block's sv fallback) rather than write + # into a silent copy. Transposed spatial dims are not viewable as (3, -1). + scene = torch.zeros((64, 64, 3), dtype=torch.uint8).permute(2, 1, 0) + with pytest.raises(RuntimeError): + gpu_draw_boxes( + scene, + np.array([[10, 10, 40, 40]], dtype=int), + np.array([[255, 0, 0]], dtype=np.uint8), + 2, + ) + + +@requires_cuda +def test_gpu_boxes_on_cuda_match_cpu() -> None: + xyxy = _SCENARIOS["overlap_chain"] + colors = _index_colors(xyxy.shape[0]) + cpu_out, _ = _paint(xyxy, colors, thickness=2, device="cpu") + cuda_out, _ = _paint(xyxy, colors, thickness=2, device="cuda") + assert np.array_equal(cpu_out.numpy(), cuda_out.cpu().numpy()) + + +def _build_detections( + boxes: np.ndarray, class_id: np.ndarray, device: str +) -> Detections: + n = boxes.shape[0] + return Detections( + xyxy=torch.tensor(boxes, dtype=torch.float32, device=device), + class_id=torch.tensor(class_id, dtype=torch.int32, device=device), + confidence=torch.full((n,), 0.9, device=device), + image_metadata={"class_names": {i: f"c{i}" for i in range(10)}}, + ) + + +def _eligible_detections(device: str = "cpu") -> Detections: + boxes = np.array([[10, 10, 50, 50]], dtype=np.float32) + return _build_detections(boxes, np.array([1]), device=device) + + +def _tensor_backed_image() -> WorkflowImageData: + tensor = torch.zeros((3, 64, 64), dtype=torch.uint8) + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), tensor_image=tensor + ) + + +def test_gpu_box_draw_eligible_happy_path() -> None: + for axis in ("CLASS", "INDEX", "TRACK"): + assert ( + _gpu_box_draw_eligible(_eligible_detections(), axis, _tensor_backed_image()) + is True + ) + + +def test_gpu_box_draw_not_eligible_for_empty_detections() -> None: + empty = _build_detections( + np.zeros((0, 4), dtype=np.float32), np.zeros((0,), dtype=int), device="cpu" + ) + assert _gpu_box_draw_eligible(empty, "CLASS", _tensor_backed_image()) is False + + +def test_gpu_box_draw_not_eligible_for_numpy_sourced_image() -> None: + numpy_image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), + numpy_image=np.zeros((64, 64, 3), dtype=np.uint8), + ) + assert _gpu_box_draw_eligible(_eligible_detections(), "CLASS", numpy_image) is False + + +def test_track_color_axis_matches_sv_colors() -> None: + from inference.core.workflows.core_steps.visualizations.bounding_box.v1_tensor import ( + BoundingBoxVisualizationBlockV1, + ) + + boxes = np.array([[20, 20, 90, 90], [150, 30, 260, 140], [40, 150, 120, 220]]) + tracker_ids = [7, -1, 3] # -1 = sv pending track (gray) + detections = Detections( + xyxy=torch.tensor(boxes, dtype=torch.float32), + class_id=torch.tensor([0, 1, 2]), + confidence=torch.full((3,), 0.9), + bboxes_metadata=[{"tracker_id": tid} for tid in tracker_ids], + ) + image = _tensor_backed_image_sized(240, 320) + block = BoundingBoxVisualizationBlockV1() + out = block.run( + image=image, + predictions=detections, + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="TRACK", + thickness=2, + roundness=0.0, + )["image"] + assert out._tensor_image is not None and out._numpy_image is None # GPU path + annotated = out._tensor_image + for (x1, y1, _, _), tid in zip(boxes, tracker_ids): + expected = (128, 128, 128) if tid == -1 else PALETTE.by_idx(tid).as_rgb() + got = tuple(int(v) for v in annotated[:, y1, x1 + 10]) # top band pixel + assert got == expected, f"track {tid}: {got} != {expected}" + + +def _tensor_backed_image_sized(h: int, w: int) -> WorkflowImageData: + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), + tensor_image=torch.zeros((3, h, w), dtype=torch.uint8), + ) + + +def _empty_detections(device: str = "cpu") -> Detections: + return _build_detections( + np.zeros((0, 4), dtype=np.float32), np.zeros((0,), dtype=int), device=device + ) + + +def _run_block(image: WorkflowImageData, detections: Detections, copy_image: bool): + from inference.core.workflows.core_steps.visualizations.bounding_box.v1_tensor import ( + BoundingBoxVisualizationBlockV1, + ) + + return BoundingBoxVisualizationBlockV1().run( + image=image, + predictions=detections, + copy_image=copy_image, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="CLASS", + thickness=2, + roundness=0.0, + )["image"] + + +def test_empty_predictions_take_the_tensor_passthrough_with_copy() -> None: + image = _tensor_backed_image() + + out = _run_block(image, _empty_detections(), copy_image=True) + + # then - output stays on-device (an empty annotate must never pay the + # full-resolution numpy materialisation) with independent storage + assert out._tensor_image is not None and out._numpy_image is None + assert out._tensor_image.data_ptr() != image.tensor_image.data_ptr() + assert torch.equal(out._tensor_image, image.tensor_image) + + +def test_empty_predictions_passthrough_shares_backing_without_copy() -> None: + image = _tensor_backed_image() + + out = _run_block(image, _empty_detections(), copy_image=False) + + assert out._tensor_image is not None and out._numpy_image is None + assert out._tensor_image.data_ptr() == image.tensor_image.data_ptr() + + +def test_empty_predictions_on_numpy_sourced_image_stay_numpy() -> None: + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), + numpy_image=np.zeros((64, 64, 3), dtype=np.uint8), + ) + + out = _run_block(image, _empty_detections(), copy_image=True) + + assert out._numpy_image is not None and out._tensor_image is None + assert not np.shares_memory(out._numpy_image, image.numpy_image) + assert np.array_equal(out._numpy_image, image.numpy_image) + + +def test_empty_passthrough_helper_declines_non_empty_predictions() -> None: + from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + empty_predictions_passthrough, + ) + + result = empty_predictions_passthrough( + image=_tensor_backed_image(), + detections=_eligible_detections(), + copy_image=True, + ) + + assert result is None diff --git a/tests/workflows/unit_tests/core_steps/visualizations/test_host_mirror_viz.py b/tests/workflows/unit_tests/core_steps/visualizations/test_host_mirror_viz.py new file mode 100644 index 0000000000..3c84b6b6d9 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/visualizations/test_host_mirror_viz.py @@ -0,0 +1,424 @@ +"""Visualization-phase tests for the per-box host mirror +(``tensor_native.HOST_MIRROR_KEYS``). + +Proves that ``to_supervision_for_annotation`` builds the sv view with ZERO +device reads when every box carries the mirror (TorchDispatchMode audit), that +the mirror and tensor-read paths produce identical sv views, that mirror-less / +partially-mirrored predictions keep today's tensor-read behaviour, and that the +label and bounding-box blocks render identically (and without detection-tensor +reads) end to end. +""" + +import numpy as np +import pytest +import supervision as sv +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from inference.core.workflows.core_steps.common.tensor_native import ( + HOST_MIRROR_KEYS, + attach_native_detection_metadata, + strip_host_mirror_metadata, +) +from inference.core.workflows.core_steps.visualizations.bounding_box.v1_tensor import ( + BoundingBoxVisualizationBlockV1, +) +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + to_supervision_for_annotation, +) +from inference.core.workflows.core_steps.visualizations.label.v1_tensor import ( + LabelVisualizationBlockV1, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.v1.dynamic_blocks.representation_boundary import ( + native_detections_to_sv, + sv_detections_to_native, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections + +CLASS_NAMES = {0: "cat", 1: "dog", 2: "goggles"} +SCENE_H, SCENE_W = 480, 640 +N_BOXES = 3 + + +def _storage_ptr(tensor: torch.Tensor) -> int: + return tensor.untyped_storage().data_ptr() + + +class _OpAudit(TorchDispatchMode): + """Record every dispatched aten op with the storage pointers of its tensor + arguments โ€” lets a test assert that SPECIFIC tensors (the detection + xyxy/class_id/confidence) were never touched, while the annotator's own + tensors (scene, sprites, paste indices) stay unconstrained.""" + + def __init__(self): + super().__init__() + self.ops = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + pointers = [] + + def _collect(value): + if isinstance(value, torch.Tensor): + try: + pointers.append(_storage_ptr(value)) + except RuntimeError: + pass + elif isinstance(value, (list, tuple)): + for item in value: + _collect(item) + + _collect(args) + _collect(list((kwargs or {}).values())) + self.ops.append((str(func), pointers)) + return func(*args, **(kwargs or {})) + + def op_names(self): + return [name for name, _ in self.ops] + + def assert_never_touched(self, *tensors: torch.Tensor) -> None: + assert not any("_local_scalar_dense" in name for name in self.op_names()) + assert not any("nonzero" in name for name in self.op_names()) + forbidden = {_storage_ptr(tensor) for tensor in tensors} + touched = [ + name + for name, pointers in self.ops + if any(pointer in forbidden for pointer in pointers) + ] + assert touched == [], f"ops read the detection tensors: {touched}" + + +def _image_data() -> WorkflowImageData: + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=np.zeros((SCENE_H, SCENE_W, 3), dtype=np.uint8), + ) + + +def _mirrored_od_detections(n: int = N_BOXES) -> Detections: + boxes = np.array( + [ + [40.5, 60.25, 200.75, 200.125], + [250.0, 90.0, 420.0, 260.0], + [90.0, 260.0, 300.0, 430.0], + [10.0, 10.0, 30.0, 30.0], + ], + dtype=np.float32, + )[:n] + detections = Detections( + xyxy=torch.tensor(boxes, dtype=torch.float32), + class_id=torch.tensor([index % 3 for index in range(n)], dtype=torch.long), + confidence=torch.tensor( + np.linspace(0.42, 0.99, max(n, 1))[:n], dtype=torch.float32 + ), + image_metadata=None, + bboxes_metadata=None, + ) + return attach_native_detection_metadata( + detections=detections, + image=_image_data(), + class_names=CLASS_NAMES, + prediction_type="object-detection", + ) + + +def _stripped_twin(detections: Detections) -> Detections: + return Detections( + xyxy=detections.xyxy, + class_id=detections.class_id, + confidence=detections.confidence, + image_metadata=detections.image_metadata, + bboxes_metadata=strip_host_mirror_metadata(detections.bboxes_metadata), + ) + + +def _assert_same_sv_view(first: sv.Detections, second: sv.Detections) -> None: + assert np.array_equal(first.xyxy, second.xyxy) + assert first.xyxy.dtype == second.xyxy.dtype + assert np.array_equal(first.class_id, second.class_id) + assert first.class_id.dtype == second.class_id.dtype + assert np.array_equal(first.confidence, second.confidence) + assert first.confidence.dtype == second.confidence.dtype + assert (first.tracker_id is None) == (second.tracker_id is None) + if first.tracker_id is not None: + assert np.array_equal(first.tracker_id, second.tracker_id) + assert set(first.data.keys()) == set(second.data.keys()) + for key in first.data: + assert np.array_equal(first.data[key], second.data[key]), key + + +# --------------------------------------------------------------------------- # +# to_supervision_for_annotation: zero device reads on the mirror path +# --------------------------------------------------------------------------- # + + +def test_to_supervision_with_full_mirror_dispatches_zero_tensor_ops() -> None: + # given + detections = _mirrored_od_detections() + + # when + audit = _OpAudit() + with audit: + sv_view = to_supervision_for_annotation(detections) + + # then - the mirror path never touches a tensor: no aten op is dispatched + # at all (in particular no `_to_copy` D2H, no `_local_scalar_dense`, no + # `nonzero`). + assert audit.ops == [] + assert len(sv_view) == N_BOXES + + +def test_to_supervision_mirror_and_tensor_paths_are_bit_identical() -> None: + # given + mirrored = _mirrored_od_detections() + stripped = _stripped_twin(mirrored) + + # when + from_mirror = to_supervision_for_annotation(mirrored) + from_tensors = to_supervision_for_annotation(stripped) + + # then + _assert_same_sv_view(from_mirror, from_tensors) + + +def test_to_supervision_mirror_keys_never_surface_in_sv_data() -> None: + # when + sv_view = to_supervision_for_annotation(_mirrored_od_detections()) + + # then + assert not any(key in sv_view.data for key in HOST_MIRROR_KEYS) + + +@pytest.mark.parametrize("missing_key", HOST_MIRROR_KEYS) +def test_to_supervision_partial_mirror_falls_back_to_tensor_reads( + missing_key, +) -> None: + # given - one box lost one mirror key (e.g. a transformed row) + mirrored = _mirrored_od_detections() + reference = to_supervision_for_annotation(_stripped_twin(mirrored)) + del mirrored.bboxes_metadata[1][missing_key] + + # when + audit = _OpAudit() + with audit: + partial_view = to_supervision_for_annotation(mirrored) + + # then - identical output via the tensor-read path (detach ops dispatched) + assert any("detach" in name for name in audit.op_names()) + assert np.array_equal(partial_view.xyxy, reference.xyxy) + assert np.array_equal(partial_view.class_id, reference.class_id) + assert np.array_equal(partial_view.confidence, reference.confidence) + # the surviving mirror keys of other boxes still never surface in .data + assert not any(key in partial_view.data for key in HOST_MIRROR_KEYS) + + +def test_to_supervision_tracker_id_and_extra_keys_unaffected_by_mirror() -> None: + # given + mirrored = _mirrored_od_detections() + for index, entry in enumerate(mirrored.bboxes_metadata): + entry["tracker_id"] = 7 + index + entry["time_in_zone"] = float(index) + + # when + audit = _OpAudit() + with audit: + sv_view = to_supervision_for_annotation(mirrored) + + # then + assert audit.ops == [] + assert np.array_equal(sv_view.tracker_id, np.asarray([7, 8, 9])) + assert np.array_equal(sv_view.data["time_in_zone"], np.asarray([0.0, 1.0, 2.0])) + + +def test_to_supervision_keypoint_tuple_with_mirrored_bbox_is_read_free() -> None: + # given + detections = _mirrored_od_detections() + key_points = KeyPoints( + xy=torch.zeros((N_BOXES, 4, 2), dtype=torch.float32), + class_id=detections.class_id.clone(), + confidence=torch.zeros((N_BOXES, 4), dtype=torch.float32), + image_metadata=detections.image_metadata, + ) + + # when + audit = _OpAudit() + with audit: + sv_view = to_supervision_for_annotation((key_points, detections)) + + # then + assert audit.ops == [] + assert len(sv_view) == N_BOXES + + +def test_to_supervision_instance_segmentation_mirror_reads_only_the_mask() -> None: + # given - dense-mask instance segmentation with a full mirror + base = _mirrored_od_detections(2) + masks = torch.zeros((2, SCENE_H, SCENE_W), dtype=torch.bool) + masks[0, 5:15, 10:30] = True + masks[1, 20:35, 40:55] = True + detections = InstanceDetections( + xyxy=base.xyxy, + class_id=base.class_id, + confidence=base.confidence, + mask=masks, + image_metadata=base.image_metadata, + bboxes_metadata=base.bboxes_metadata, + ) + + # when + audit = _OpAudit() + with audit: + sv_view = to_supervision_for_annotation(detections, materialise_masks=True) + + # then - the mask materialisation is untouched (its bulk transfer remains), + # but xyxy/class_id/confidence never touch the device + assert sv_view.mask is not None + audit.assert_never_touched( + detections.xyxy, detections.class_id, detections.confidence + ) + + +# --------------------------------------------------------------------------- # +# custom-python boundary: the mirror is internal transport only +# --------------------------------------------------------------------------- # + + +def test_representation_boundary_excludes_mirror_from_sv_data_and_roundtrip() -> None: + # given + mirrored = _mirrored_od_detections() + + # when + sv_view = native_detections_to_sv(mirrored) + rebuilt = sv_detections_to_native(sv_view) + + # then - legacy user code never sees the mirror, and the round-trip does + # not re-attach one (user code may have edited the boxes) + assert not any(key in sv_view.data for key in HOST_MIRROR_KEYS) + assert rebuilt.bboxes_metadata is not None + for entry in rebuilt.bboxes_metadata: + assert not any(key in entry for key in HOST_MIRROR_KEYS) + + +# --------------------------------------------------------------------------- # +# label block end to end (GPU sprite path) +# --------------------------------------------------------------------------- # + +_LABEL_RUN_KWARGS = dict( + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="CLASS", + text="Class and Confidence", + text_position="TOP_LEFT", + text_color="WHITE", + text_scale=1.0, + text_thickness=1, + text_padding=10, + border_radius=0, +) + + +def _tensor_image(seed: int = 7) -> WorkflowImageData: + rng = np.random.default_rng(seed) + scene = rng.integers(0, 255, (3, SCENE_H, SCENE_W)).astype(np.uint8) + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.from_numpy(scene), + ) + + +def test_label_block_run_with_mirror_never_reads_detection_tensors() -> None: + # given + block = LabelVisualizationBlockV1() + detections = _mirrored_od_detections() + + # warm the sprite cache so the audited run measures the steady state + block.run(image=_tensor_image(), predictions=detections, **_LABEL_RUN_KWARGS) + + # when + audit = _OpAudit() + with audit: + result = block.run( + image=_tensor_image(), predictions=detections, **_LABEL_RUN_KWARGS + ) + + # then - the sprite compositor works on the scene / sprites only; the + # detection tensors are never read (no D2H outside the sprite uploads) + audit.assert_never_touched( + detections.xyxy, detections.class_id, detections.confidence + ) + assert result["image"]._tensor_image is not None + + +def test_label_block_output_identical_with_and_without_mirror() -> None: + # given + mirrored = _mirrored_od_detections() + stripped = _stripped_twin(mirrored) + + # when + out_mirrored = LabelVisualizationBlockV1().run( + image=_tensor_image(), predictions=mirrored, **_LABEL_RUN_KWARGS + )["image"] + out_stripped = LabelVisualizationBlockV1().run( + image=_tensor_image(), predictions=stripped, **_LABEL_RUN_KWARGS + )["image"] + + # then + assert torch.equal(out_mirrored._tensor_image, out_stripped._tensor_image) + + +# --------------------------------------------------------------------------- # +# bounding-box block (GPU painter path) +# --------------------------------------------------------------------------- # + +_BBOX_RUN_KWARGS = dict( + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="CLASS", + thickness=2, + roundness=0.0, +) + + +def test_bounding_box_block_run_with_mirror_never_reads_detection_tensors() -> None: + # given + block = BoundingBoxVisualizationBlockV1() + detections = _mirrored_od_detections() + + # when + audit = _OpAudit() + with audit: + result = block.run( + image=_tensor_image(), predictions=detections, **_BBOX_RUN_KWARGS + ) + + # then + audit.assert_never_touched( + detections.xyxy, detections.class_id, detections.confidence + ) + assert result["image"]._tensor_image is not None + + +def test_bounding_box_block_output_identical_with_and_without_mirror() -> None: + # given + mirrored = _mirrored_od_detections() + stripped = _stripped_twin(mirrored) + + # when + out_mirrored = BoundingBoxVisualizationBlockV1().run( + image=_tensor_image(), predictions=mirrored, **_BBOX_RUN_KWARGS + )["image"] + out_stripped = BoundingBoxVisualizationBlockV1().run( + image=_tensor_image(), predictions=stripped, **_BBOX_RUN_KWARGS + )["image"] + + # then + assert torch.equal(out_mirrored._tensor_image, out_stripped._tensor_image) diff --git a/tests/workflows/unit_tests/core_steps/visualizations/test_label_v1_tensor_gpu.py b/tests/workflows/unit_tests/core_steps/visualizations/test_label_v1_tensor_gpu.py new file mode 100644 index 0000000000..fe8b2a7d3f --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/visualizations/test_label_v1_tensor_gpu.py @@ -0,0 +1,731 @@ +import numpy as np +import pytest +import supervision as sv +import torch + +import inference.core.workflows.core_steps.visualizations.label.v1_tensor as label_v1_tensor +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + to_supervision_for_annotation, +) +from inference.core.workflows.core_steps.visualizations.common.label_text import ( + build_detection_labels, +) +from inference.core.workflows.core_steps.visualizations.common.utils import str_to_color +from inference.core.workflows.core_steps.visualizations.label.v1_tensor import ( + LabelVisualizationBlockV1, + _gpu_label_paste_eligible, + _measure_label, + _render_label_sprite, + _SceneDependentLabelError, + gpu_paste_label_sprites, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) +from inference_models.models.base.object_detection import Detections + +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" +) + +PALETTE = sv.ColorPalette.DEFAULT +SCENE_H, SCENE_W = 480, 640 +DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + +ALL_ANCHORS = [ + "TOP_LEFT", + "TOP_RIGHT", + "TOP_CENTER", + "CENTER", + "CENTER_LEFT", + "CENTER_RIGHT", + "BOTTOM_LEFT", + "BOTTOM_CENTER", + "BOTTOM_RIGHT", +] + +_DEFAULT_RUN_KWARGS = dict( + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="CLASS", + text="Class and Confidence", + text_position="TOP_LEFT", + text_color="WHITE", + text_scale=1.0, + text_thickness=1, + text_padding=10, + border_radius=0, +) + + +def _make_scene(seed: int, h: int = SCENE_H, w: int = SCENE_W) -> np.ndarray: + rng = np.random.default_rng(seed) + return rng.integers(0, 255, (h, w, 3)).astype(np.uint8) + + +def _build_detections( + boxes: np.ndarray, + class_id: np.ndarray, + device: str = "cpu", + confidence: np.ndarray = None, + bboxes_metadata=None, +) -> Detections: + n = boxes.shape[0] + if confidence is None: + confidence = np.linspace(0.42, 0.99, max(n, 1))[:n] + return Detections( + xyxy=torch.tensor(boxes, dtype=torch.float32, device=device), + class_id=torch.tensor(class_id, dtype=torch.int32, device=device), + confidence=torch.tensor(confidence, dtype=torch.float32, device=device), + image_metadata={"class_names": {0: "goggles", 1: "cat", 2: "dog"}}, + bboxes_metadata=bboxes_metadata, + ) + + +def _default_detections(device: str = "cpu") -> Detections: + boxes = np.array( + [ + [40, 60, 200, 200], + [250, 90, 420, 260], + [90, 260, 300, 430], + ], + dtype=np.float32, + ) + return _build_detections(boxes, np.array([0, 1, 2]), device=device) + + +def _tensor_image_from_bgr(scene_bgr: np.ndarray, device: str = "cpu"): + tensor = ( + torch.from_numpy(scene_bgr[:, :, ::-1].copy()) + .permute(2, 0, 1) + .contiguous() + .to(device) + ) + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), tensor_image=tensor + ) + + +def _run_block(image, detections, block=None, **overrides): + block = block or LabelVisualizationBlockV1() + kwargs = {**_DEFAULT_RUN_KWARGS, **overrides} + return block.run(image=image, predictions=detections, **kwargs)["image"] + + +def _sv_reference(scene_bgr: np.ndarray, detections: Detections, **overrides): + """What the sv path (and the flag-off block) draws for the same inputs.""" + kwargs = {**_DEFAULT_RUN_KWARGS, **overrides} + sv_view = to_supervision_for_annotation(detections) + labels = build_detection_labels(sv_view, kwargs["text"]) + annotator = sv.LabelAnnotator( + color=PALETTE, + color_lookup=getattr(sv.ColorLookup, kwargs["color_axis"]), + text_position=getattr(sv.Position, kwargs["text_position"]), + text_color=str_to_color(kwargs["text_color"]), + text_scale=kwargs["text_scale"], + text_thickness=kwargs["text_thickness"], + text_padding=kwargs["text_padding"], + border_radius=kwargs["border_radius"], + ) + return annotator.annotate(scene_bgr.copy(), sv_view, labels=labels) + + +def _to_bgr(tensor_chw: torch.Tensor) -> np.ndarray: + return tensor_chw.permute(1, 2, 0).cpu().numpy()[:, :, ::-1] + + +def _assert_gpu_bit_exact(out, expected_bgr: np.ndarray) -> None: + # the GPU path returns a tensor image and never materialises numpy + assert out._tensor_image is not None and out._numpy_image is None + actual = _to_bgr(out._tensor_image) + assert np.array_equal(actual, expected_bgr) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("text_position", ALL_ANCHORS) +def test_gpu_labels_match_sv_bit_exact_for_every_anchor( + device: str, text_position: str +) -> None: + # given: interior boxes plus boxes that push the label over each frame + # edge (the label patch of an edge box lands partially off-frame for most + # anchors, exercising the frame-clipped sprite variants) + scene = _make_scene(11) + boxes = np.array( + [ + [40, 60, 200, 200], + [250, 90, 420, 260], + [90, 260, 300, 430], + [5, 5, 60, 60], + [560, 400, 635, 475], + ], + dtype=np.float32, + ) + detections = _build_detections(boxes, np.array([0, 1, 2, 1, 0]), device=device) + expected = _sv_reference(scene, detections, text_position=text_position) + + # when + out = _run_block( + _tensor_image_from_bgr(scene, device=device), + detections, + text_position=text_position, + ) + + # then + assert out._tensor_image is not None and out._numpy_image is None + assert np.array_equal(_to_bgr(out._tensor_image), expected) + + +@pytest.mark.parametrize( + "text", ["Class", "Confidence", "Class and Confidence", "Index", "Dimensions"] +) +def test_gpu_labels_match_sv_for_text_options(text: str) -> None: + # given + scene = _make_scene(23) + detections = _default_detections() + expected = _sv_reference(scene, detections, text=text) + + # when + out = _run_block(_tensor_image_from_bgr(scene), detections, text=text) + + # then + _assert_gpu_bit_exact(out, expected) + + +def test_gpu_labels_match_sv_with_overlapping_labels() -> None: + # given: staggered boxes whose label patches overlap each other โ€” sv draws + # sequentially, so a later label must overwrite an earlier one; the last + # two detections share the same box, so their labels coincide entirely + scene = _make_scene(31) + boxes = np.array( + [ + [100, 100, 300, 300], + [130, 115, 330, 300], + [160, 130, 360, 300], + [160, 130, 360, 300], + ], + dtype=np.float32, + ) + detections = _build_detections(boxes, np.array([0, 1, 2, 1])) + expected = _sv_reference(scene, detections) + + # when + out = _run_block(_tensor_image_from_bgr(scene), detections) + + # then: paste order matches sv's draw order bit-for-bit + _assert_gpu_bit_exact(out, expected) + + +@pytest.mark.parametrize("text_position", ["TOP_LEFT", "CENTER", "BOTTOM_RIGHT"]) +def test_gpu_labels_match_sv_when_clipped_at_every_frame_edge( + text_position: str, +) -> None: + # given: boxes hugging each frame edge and corner so label patches are + # clipped left/top/right/bottom and across two edges at once โ€” cv2's AA + # rasterisation changes where strokes are cut by the frame border, which + # the frame-clipped sprite variants must reproduce exactly + scene = _make_scene(47) + boxes = np.array( + [ + [2, 2, 60, 40], # top-left corner + [580, 3, 637, 50], # top-right corner + [3, 430, 70, 477], # bottom-left corner + [590, 440, 636, 476], # bottom-right corner + [2, 200, 40, 260], # left edge + [600, 200, 638, 260], # right edge + ], + dtype=np.float32, + ) + detections = _build_detections(boxes, np.array([0, 1, 2, 0, 1, 2])) + expected = _sv_reference(scene, detections, text_position=text_position) + + # when + out = _run_block( + _tensor_image_from_bgr(scene), detections, text_position=text_position + ) + + # then + _assert_gpu_bit_exact(out, expected) + + +def test_gpu_labels_match_sv_when_label_is_wider_than_frame() -> None: + # given: a label so long it overflows both vertical frame edges at once + scene = _make_scene(53, h=240, w=320) + boxes = np.array([[100, 100, 220, 200]], dtype=np.float32) + detections = _build_detections(boxes, np.array([0])) + detections.bboxes_metadata = [ + {"note": "a very long annotation that cannot possibly fit this frame"} + ] + expected = _sv_reference(scene, detections, text="note", text_position="CENTER") + + # when + out = _run_block( + _tensor_image_from_bgr(scene), detections, text="note", text_position="CENTER" + ) + + # then + _assert_gpu_bit_exact(out, expected) + + +def test_gpu_labels_match_sv_for_fully_off_frame_labels() -> None: + # given: box far outside the frame โ€” sv draws nothing visible + scene = _make_scene(59, h=240, w=320) + boxes = np.array([[400, 400, 500, 500]], dtype=np.float32) + detections = _build_detections(boxes, np.array([0])) + + # when + out = _run_block(_tensor_image_from_bgr(scene), detections) + + # then: scene unchanged, still on the tensor path + _assert_gpu_bit_exact(out, scene) + + +@pytest.mark.parametrize("border_radius", [3, 8, 15]) +def test_gpu_labels_match_sv_with_border_radius(border_radius: int) -> None: + # given: rounded corners leave the scene visible through the corner cuts + scene = _make_scene(61) + detections = _default_detections() + expected = _sv_reference(scene, detections, border_radius=border_radius) + + # when + out = _run_block( + _tensor_image_from_bgr(scene), detections, border_radius=border_radius + ) + + # then + _assert_gpu_bit_exact(out, expected) + + +@pytest.mark.parametrize( + "text_scale,text_thickness,text_padding", + [(0.5, 1, 10), (0.5, 2, 8), (2.0, 2, 25), (1.0, 1, 15)], +) +def test_gpu_labels_match_sv_with_custom_typography( + text_scale: float, text_thickness: int, text_padding: int +) -> None: + # given + scene = _make_scene(67) + detections = _default_detections() + expected = _sv_reference( + scene, + detections, + text_scale=text_scale, + text_thickness=text_thickness, + text_padding=text_padding, + ) + + # when + out = _run_block( + _tensor_image_from_bgr(scene), + detections, + text_scale=text_scale, + text_thickness=text_thickness, + text_padding=text_padding, + ) + + # then + _assert_gpu_bit_exact(out, expected) + + +def test_gpu_labels_match_sv_for_multiline_custom_text() -> None: + # given: sv's wrap_text splits custom labels on newlines + scene = _make_scene(71) + boxes = np.array([[80, 120, 300, 300], [320, 120, 560, 300]], dtype=np.float32) + detections = _build_detections(boxes, np.array([0, 1])) + detections.bboxes_metadata = [ + {"note": "first line\nsecond line"}, + {"note": "single"}, + ] + expected = _sv_reference(scene, detections, text="note") + + # when + out = _run_block(_tensor_image_from_bgr(scene), detections, text="note") + + # then + _assert_gpu_bit_exact(out, expected) + + +@pytest.mark.parametrize("color_axis", ["INDEX", "TRACK"]) +def test_gpu_labels_match_sv_for_other_color_axes(color_axis: str) -> None: + # given: non-contiguous tracker ids incl. sv's pending track (-1), whose + # resolve_color turns BOTH the background and the text gray + scene = _make_scene(73) + boxes = np.array( + [[40, 60, 200, 200], [250, 90, 420, 260], [90, 260, 300, 430]], + dtype=np.float32, + ) + detections = _build_detections( + boxes, + np.array([0, 1, 2]), + bboxes_metadata=[{"tracker_id": 12}, {"tracker_id": -1}, {"tracker_id": 3}], + ) + expected = _sv_reference(scene, detections, color_axis=color_axis) + + # when + out = _run_block(_tensor_image_from_bgr(scene), detections, color_axis=color_axis) + + # then + _assert_gpu_bit_exact(out, expected) + + +def test_track_axis_without_tracker_ids_raises_sv_error() -> None: + # given: TRACK lookup but no tracker ids โ€” the sv annotator's exact + # ValueError must surface (via the sv fallback) + scene = _make_scene(79) + detections = _default_detections() + + # when / then + with pytest.raises(ValueError, match="resolve color by track"): + _run_block(_tensor_image_from_bgr(scene), detections, color_axis="TRACK") + + +def test_sprite_cache_reuses_sprites_across_runs(monkeypatch) -> None: + # given: a render-call counter around the real renderer + calls = {"n": 0} + real_render = label_v1_tensor._render_label_sprite + + def counting_render(*args, **kwargs): + calls["n"] += 1 + return real_render(*args, **kwargs) + + monkeypatch.setattr(label_v1_tensor, "_render_label_sprite", counting_render) + scene = _make_scene(83) + detections = _default_detections() # 3 distinct labels + block = LabelVisualizationBlockV1() + + # when: the same frame twice + _run_block(_tensor_image_from_bgr(scene), detections, block=block) + first_run_renders = calls["n"] + _run_block(_tensor_image_from_bgr(scene), detections, block=block) + + # then: one sprite per distinct label, and the second run renders nothing + assert first_run_renders == 3 + assert len(block._sprite_cache) == 3 + assert calls["n"] == first_run_renders + + +@pytest.mark.parametrize("device", DEVICES) +def test_cached_sprites_are_device_resident_tensors(device: str) -> None: + # given + scene = _make_scene(89) + detections = _default_detections(device=device) + block = LabelVisualizationBlockV1() + + # when + _run_block(_tensor_image_from_bgr(scene, device=device), detections, block=block) + + # then: sprite pixel payloads live on the scene device โ€” a cache-hit frame + # uploads no pixel data + assert len(block._sprite_cache) > 0 + for sprite in block._sprite_cache.values(): + assert isinstance(sprite.colors_dev, torch.Tensor) + assert sprite.colors_dev.device.type == device + for flat in sprite._flat_by_width.values(): + assert flat.device.type == device + + +def test_sprite_cache_is_bounded(monkeypatch) -> None: + # given + monkeypatch.setattr(label_v1_tensor, "_SPRITE_CACHE_MAX_ENTRIES", 2) + scene = _make_scene(97) + boxes = np.array( + [ + [40, 60, 200, 200], + [250, 90, 420, 260], + [90, 260, 300, 430], + [330, 280, 500, 430], + ], + dtype=np.float32, + ) + detections = _build_detections(boxes, np.array([0, 1, 2, 0])) + block = LabelVisualizationBlockV1() + + # when: 4 distinct labels through a cache capped at 2 + _run_block(_tensor_image_from_bgr(scene), detections, block=block) + + # then + assert len(block._sprite_cache) == 2 + + +def test_copy_image_true_leaves_input_untouched() -> None: + # given + scene = _make_scene(101) + image = _tensor_image_from_bgr(scene) + original = image.tensor_image.clone() + + # when + out = _run_block(image, _default_detections(), copy_image=True) + + # then: independent storage, input pixels intact, output annotated + assert out._tensor_image is not None + assert out._tensor_image.data_ptr() != image.tensor_image.data_ptr() + assert torch.equal(image.tensor_image, original) + assert not torch.equal(out._tensor_image, original) + + +def test_copy_image_false_mutates_input_in_place() -> None: + # given + scene = _make_scene(103) + image = _tensor_image_from_bgr(scene) + input_tensor = image.tensor_image + + # when + out = _run_block(image, _default_detections(), copy_image=False) + + # then: same storage annotated in place, and numpy is never materialised + assert out._tensor_image is not None + assert out._tensor_image.data_ptr() == input_tensor.data_ptr() + assert image._numpy_image is None + + +def _empty_detections(device: str = "cpu") -> Detections: + return _build_detections( + np.zeros((0, 4), dtype=np.float32), np.zeros((0,), dtype=int), device=device + ) + + +def test_empty_predictions_take_the_tensor_passthrough_with_copy() -> None: + # given + scene = _make_scene(107) + image = _tensor_image_from_bgr(scene) + + # when + out = _run_block(image, _empty_detections(), copy_image=True) + + # then: output stays on-device (an empty annotate must never pay the + # full-resolution numpy materialisation) with independent storage + assert out._tensor_image is not None and out._numpy_image is None + assert out._tensor_image.data_ptr() != image.tensor_image.data_ptr() + assert torch.equal(out._tensor_image, image.tensor_image) + assert image._numpy_image is None + + +def test_empty_predictions_passthrough_shares_backing_without_copy() -> None: + # given + image = _tensor_image_from_bgr(_make_scene(109)) + + # when + out = _run_block(image, _empty_detections(), copy_image=False) + + # then + assert out._tensor_image is not None and out._numpy_image is None + assert out._tensor_image.data_ptr() == image.tensor_image.data_ptr() + + +def test_empty_predictions_on_numpy_sourced_image_stay_numpy() -> None: + # given + scene = _make_scene(113) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), numpy_image=scene + ) + + # when + out = _run_block(image, _empty_detections(), copy_image=True) + + # then + assert out._numpy_image is not None and out._tensor_image is None + assert not np.shares_memory(out._numpy_image, scene) + assert np.array_equal(out._numpy_image, scene) + + +def test_numpy_sourced_image_takes_the_sv_path_unchanged() -> None: + # given: a numpy-backed image must behave exactly as before the GPU path + # existed โ€” same sv annotator, numpy output, no tensor materialisation + scene = _make_scene(127) + detections = _default_detections() + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), numpy_image=scene + ) + expected = _sv_reference(scene, detections) + + # when + out = _run_block(image, detections) + + # then + assert out._numpy_image is not None and out._tensor_image is None + assert image._tensor_image is None # never materialised as a side effect + assert np.array_equal(out._numpy_image, expected) + + +def test_mask_dependent_area_text_keeps_the_sv_path() -> None: + # given: `Area` reads masks in sv (falling back to box area for OD input) + # โ€” the sprite compositor must not take over that configuration + scene = _make_scene(131) + detections = _default_detections() + expected = _sv_reference(scene, detections, text="Area") + + # when + out = _run_block(_tensor_image_from_bgr(scene), detections, text="Area") + + # then: sv fallback produced a numpy image identical to flag-off behavior + assert out._numpy_image is not None + assert np.array_equal(out._numpy_image, expected) + + +def test_center_of_mass_on_od_input_raises_the_sv_error() -> None: + # given: CENTER_OF_MASS anchors on the mask centroid; without masks sv's + # get_anchors_coordinates raises โ€” the block must surface the same error + # it always has (via the sv path), not swallow it in the GPU branch + scene = _make_scene(131) + detections = _default_detections() + + # when / then + with pytest.raises(ValueError, match="CENTER_OF_MASS"): + _run_block( + _tensor_image_from_bgr(scene), + detections, + text_position="CENTER_OF_MASS", + ) + + +def test_scene_dependent_sprite_falls_back_to_sv_path() -> None: + # given: text_padding smaller than the font's descender extent โ€” the AA + # descender ink of "goggles" escapes the opaque background, so sv blends + # it with the scene; the sprite path must refuse and fall back, keeping + # the output bit-identical to sv + scene = _make_scene(137) + detections = _default_detections() + expected = _sv_reference(scene, detections, text="Class", text_padding=2) + + # when + out = _run_block( + _tensor_image_from_bgr(scene), detections, text="Class", text_padding=2 + ) + + # then + assert out._numpy_image is not None # sv fallback was taken + assert np.array_equal(out._numpy_image, expected) + + +def test_render_label_sprite_raises_when_text_ink_escapes_background() -> None: + # given + measurement = _measure_label("goggles", 1.0, 1, 2) + margin = measurement.margin + + # when / then + with pytest.raises(_SceneDependentLabelError): + _render_label_sprite( + measurement=measurement, + text_color_bgr=(255, 255, 255), + background_color_bgr=(0, 0, 255), + text_scale=1.0, + text_thickness=1, + text_padding=2, + border_radius=0, + device=torch.device("cpu"), + box_in_canvas=( + margin, + margin, + margin + measurement.width_padded, + margin + measurement.height_padded, + ), + canvas_hw=( + measurement.height_padded + 1 + 2 * margin, + measurement.width_padded + 1 + 2 * margin, + ), + frame_edge_sides=(False, False, False, False), + ) + + +def _interior_sprite(label: str = "ab", device: str = "cpu"): + measurement = _measure_label(label, 1.0, 1, 10) + margin = measurement.margin + return _render_label_sprite( + measurement=measurement, + text_color_bgr=(255, 255, 255), + background_color_bgr=(0, 0, 255), + text_scale=1.0, + text_thickness=1, + text_padding=10, + border_radius=0, + device=torch.device(device), + box_in_canvas=( + margin, + margin, + margin + measurement.width_padded, + margin + measurement.height_padded, + ), + canvas_hw=( + measurement.height_padded + 1 + 2 * margin, + measurement.width_padded + 1 + 2 * margin, + ), + frame_edge_sides=(False, False, False, False), + ) + + +def test_gpu_paste_is_in_place() -> None: + # given + sprite = _interior_sprite() + scene = torch.zeros((3, 200, 200), dtype=torch.uint8) + + # when + out = gpu_paste_label_sprites(scene, [sprite], [(5, 5)]) + + # then + assert out.data_ptr() == scene.data_ptr() + assert int((out != 0).any(dim=0).sum()) > 0 + + +def test_gpu_paste_rejects_non_contiguous_scene() -> None: + # given: .view must fail (caught by the block's sv fallback) rather than + # write into a silent copy + sprite = _interior_sprite() + scene = torch.zeros((200, 200, 3), dtype=torch.uint8).permute(2, 1, 0) + + # when / then + with pytest.raises(RuntimeError): + gpu_paste_label_sprites(scene, [sprite], [(5, 5)]) + + +def test_gpu_paste_rejects_out_of_frame_pixels() -> None: + # given: an origin that pushes sprite pixels off-frame โ€” the block should + # have picked a frame-clipped variant, so this is a hard error, raised + # before any scene write + sprite = _interior_sprite() + scene = torch.zeros((3, 200, 200), dtype=torch.uint8) + + # when / then + with pytest.raises(ValueError, match="outside the frame"): + gpu_paste_label_sprites(scene, [sprite], [(150, 150)]) + assert int(scene.sum()) == 0 # untouched + + +def test_gpu_label_paste_eligible_semantics() -> None: + # given + tensor_image = _tensor_image_from_bgr(_make_scene(139)) + numpy_image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), + numpy_image=_make_scene(139), + ) + detections = _default_detections() + + # when / then + for axis in ("CLASS", "INDEX", "TRACK"): + assert _gpu_label_paste_eligible(detections, axis, tensor_image) is True + assert _gpu_label_paste_eligible(detections, "SOMETHING", tensor_image) is False + assert _gpu_label_paste_eligible(detections, "CLASS", numpy_image) is False + assert ( + _gpu_label_paste_eligible(_empty_detections(), "CLASS", tensor_image) is False + ) + assert _gpu_label_paste_eligible(object(), "CLASS", tensor_image) is False + + +@requires_cuda +def test_gpu_labels_on_cuda_match_cpu() -> None: + # given + scene = _make_scene(149) + expected_out = _run_block(_tensor_image_from_bgr(scene), _default_detections()) + + # when + cuda_out = _run_block( + _tensor_image_from_bgr(scene, device="cuda"), + _default_detections(device="cuda"), + ) + + # then + assert cuda_out._tensor_image.is_cuda + assert np.array_equal( + _to_bgr(cuda_out._tensor_image), _to_bgr(expected_out._tensor_image) + ) diff --git a/tests/workflows/unit_tests/core_steps/visualizations/test_label_v2_and_font_schema.py b/tests/workflows/unit_tests/core_steps/visualizations/test_label_v2_and_font_schema.py index bf288b602f..3e1d719117 100644 --- a/tests/workflows/unit_tests/core_steps/visualizations/test_label_v2_and_font_schema.py +++ b/tests/workflows/unit_tests/core_steps/visualizations/test_label_v2_and_font_schema.py @@ -1,7 +1,9 @@ import numpy as np import pytest import supervision as sv +import torch +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.visualizations.common.fonts.registry import ( FONTS_REGISTRY, ) @@ -9,13 +11,31 @@ LabelManifestV2, LabelVisualizationBlockV2, ) +from inference.core.workflows.core_steps.visualizations.label.v2_tensor import ( + LabelVisualizationBlockV2 as LabelVisualizationBlockV2Tensor, +) from inference.core.workflows.core_steps.visualizations.rich_label.v1 import ( RichLabelManifest, ) +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.entities.base import ( ImageParentMetadata, WorkflowImageData, ) +from inference_models.models.base.object_detection import Detections as NativeDetections + +# The loader binds `_tensor` siblings under the same names when +# ENABLE_TENSOR_DATA_REPRESENTATION is set - hence the flag-opposed +# _NUMPY_ONLY / _TENSOR_ONLY split below. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="loader binds the tensor-native sibling under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) def test_rich_label_font_schema_is_inline_enum_of_display_names() -> None: @@ -51,12 +71,22 @@ def test_label_v2_manifest_defaults_to_manual_text_size_mode() -> None: assert schema["properties"]["text_scale"]["title"] == "Scale" +@_NUMPY_ONLY def test_label_v2_block_is_registered_in_loader() -> None: from inference.core.workflows.core_steps.loader import load_blocks assert LabelVisualizationBlockV2 in load_blocks() +@_TENSOR_ONLY +def test_label_v2_tensor_block_is_registered_in_loader() -> None: + from inference.core.workflows.core_steps.loader import load_blocks + + blocks = load_blocks() + assert LabelVisualizationBlockV2Tensor in blocks + assert LabelVisualizationBlockV2 not in blocks + + def test_label_v2_automatic_mode_scales_text_for_smaller_images(bundled_fonts) -> None: block = LabelVisualizationBlockV2() image = WorkflowImageData( @@ -108,3 +138,146 @@ def test_label_v2_automatic_mode_scales_text_for_smaller_images(bundled_fonts) - assert not np.array_equal( manual_output["image"].numpy_image, automatic_output["image"].numpy_image ) + + +def _native_predictions() -> NativeDetections: + return NativeDetections( + xyxy=torch.tensor([[10, 10, 100, 100]], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([0.9], dtype=torch.float32), + image_metadata={CLASS_NAMES_KEY: {0: "person"}}, + ) + + +def _run_label_v2_tensor_block(block, image, predictions, **overrides): + kwargs = { + "image": image, + "predictions": predictions, + "copy_image": True, + "color_palette": "DEFAULT", + "palette_size": 10, + "custom_colors": None, + "color_axis": "CLASS", + "text": "Class", + "text_position": "TOP_LEFT", + "text_color": "WHITE", + "text_size_mode": "Manual", + "text_scale": 1.0, + "text_thickness": 1, + "text_padding": 10, + "border_radius": 0, + } + kwargs.update(overrides) + return block.run(**kwargs) + + +@_TENSOR_ONLY +def test_label_v2_automatic_mode_scales_text_for_smaller_images_tensor_native( + bundled_fonts, +) -> None: + block = LabelVisualizationBlockV2Tensor() + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="some"), + numpy_image=np.zeros((540, 960, 3), dtype=np.uint8), + ) + + manual_output = _run_label_v2_tensor_block( + block, image, _native_predictions(), text_size_mode="Manual" + ) + automatic_output = _run_label_v2_tensor_block( + block, image, _native_predictions(), text_size_mode="Automatic" + ) + + assert not np.array_equal( + manual_output["image"].numpy_image, automatic_output["image"].numpy_image + ) + + +@_TENSOR_ONLY +def test_label_v2_tensor_native_output_matches_numpy_block(bundled_fonts) -> None: + # given - the same frame and semantically-identical predictions in both + # representations; the tensor sibling must reproduce the numpy block's + # rendering byte-for-byte (it reuses the numpy drawing internals). + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="some"), + numpy_image=np.zeros((540, 960, 3), dtype=np.uint8), + ) + sv_predictions = sv.Detections( + xyxy=np.array([[10, 10, 100, 100]], dtype=np.float64), + class_id=np.array([0]), + confidence=np.array([0.9]), + data={"class_name": np.array(["person"])}, + ) + + # when + numpy_output = LabelVisualizationBlockV2().run( + image=image, + predictions=sv_predictions, + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="CLASS", + text="Class", + text_position="TOP_LEFT", + text_color="WHITE", + text_size_mode="Automatic", + text_scale=1.0, + text_thickness=1, + text_padding=10, + border_radius=0, + ) + tensor_output = _run_label_v2_tensor_block( + LabelVisualizationBlockV2Tensor(), + image, + _native_predictions(), + text_size_mode="Automatic", + ) + + # then + assert np.array_equal( + numpy_output["image"].numpy_image, tensor_output["image"].numpy_image + ) + + +@_TENSOR_ONLY +def test_label_v2_tensor_native_empty_predictions_passthrough_is_device_resident() -> ( + None +): + # given - a tensor-source image and an empty prediction: the sibling must + # pass the tensor representation through without materialising numpy + block = LabelVisualizationBlockV2Tensor() + tensor = torch.zeros((3, 240, 320), dtype=torch.uint8) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="some"), + tensor_image=tensor, + ) + empty_predictions = NativeDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + ) + + # when + copied_output = _run_label_v2_tensor_block( + block, image, empty_predictions, copy_image=True + ) + shared_output = _run_label_v2_tensor_block( + block, image, empty_predictions, copy_image=False + ) + + # then + assert image._numpy_image is None, "passthrough must not materialise numpy" + for output in [copied_output, shared_output]: + assert output["image"].is_tensor_materialised() is True + assert output["image"]._numpy_image is None + assert torch.equal(output["image"].tensor_image, image.tensor_image) + # copy semantics: independent storage when copy_image=True, shared otherwise + assert ( + copied_output["image"].tensor_image.data_ptr() + != image.tensor_image.data_ptr() + ) + assert ( + shared_output["image"].tensor_image.data_ptr() + == image.tensor_image.data_ptr() + ) diff --git a/tests/workflows/unit_tests/core_steps/visualizations/test_mask_v1_tensor_gpu.py b/tests/workflows/unit_tests/core_steps/visualizations/test_mask_v1_tensor_gpu.py new file mode 100644 index 0000000000..7d3e844b63 --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/visualizations/test_mask_v1_tensor_gpu.py @@ -0,0 +1,976 @@ +import cv2 +import numpy as np +import pytest +import supervision as sv +import torch +from pycocotools import mask as mask_utils +from torch.utils._python_dispatch import TorchDispatchMode + +from inference.core.workflows.core_steps.visualizations.common.base_tensor import ( + to_supervision_for_annotation, +) +from inference.core.workflows.core_steps.visualizations.mask.v1_tensor import ( + MaskVisualizationBlockV1, + _coco_rle_counts_to_runs, + _resolve_color_ids, + _rle_to_dense_masks, + gpu_mask_composite, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.types import InstancesRLEMasks + +requires_cuda = pytest.mark.skipif( + not torch.cuda.is_available(), reason="requires CUDA" +) + +OPACITY = 0.5 +PALETTE = sv.ColorPalette.DEFAULT +SCENE_H, SCENE_W = 540, 960 + +DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +def _build_dense_detections( + masks: np.ndarray, + boxes: np.ndarray, + class_id: np.ndarray, + device: str, +) -> InstanceDetections: + n = masks.shape[0] + return InstanceDetections( + xyxy=torch.tensor(boxes, dtype=torch.int32, device=device), + class_id=torch.tensor(class_id, dtype=torch.int32, device=device), + confidence=torch.full((n,), 0.9, device=device), + mask=torch.from_numpy(masks).to(device), + image_metadata={"class_names": {i: f"c{i}" for i in range(10)}}, + ) + + +def _build_rle_detections( + masks: list, + boxes: np.ndarray, + class_id: np.ndarray, + device: str, +) -> InstanceDetections: + n = len(masks) + payloads = [ + mask_utils.encode(np.asfortranarray(mask.astype(np.uint8)))["counts"] + for mask in masks + ] + return InstanceDetections( + xyxy=torch.tensor(boxes, dtype=torch.int32, device=device), + class_id=torch.tensor(class_id, dtype=torch.int32, device=device), + confidence=torch.full((n,), 0.9, device=device), + mask=InstancesRLEMasks(image_size=tuple(masks[0].shape), masks=payloads), + image_metadata={"class_names": {i: f"c{i}" for i in range(10)}}, + ) + + +def _single_mask_inputs(h: int = 64, w: int = 64): + masks = np.zeros((1, h, w), dtype=bool) + masks[0, 10:30, 12:40] = True + boxes = np.array([[12, 10, 39, 29]], dtype=np.int32) + class_id = np.array([1], dtype=np.int32) + return masks, boxes, class_id + + +def test_resolve_color_ids_matches_sv_semantics() -> None: + # given + masks, boxes, class_id = _single_mask_inputs() + predictions = _build_dense_detections(masks, boxes, class_id, device="cpu") + predictions.bboxes_metadata = [{"tracker_id": 7}] + + # when / then: same palette indices sv's resolve_color_idx would use, + # returned as device tensors (never a device->host read) + for axis, expected in (("CLASS", class_id), ("INDEX", [0]), ("TRACK", [7])): + ids = _resolve_color_ids(predictions, axis, torch.device("cpu")) + assert isinstance(ids, torch.Tensor) and ids.dtype == torch.int64 + assert np.array_equal(ids.numpy(), expected) + + +def test_resolve_color_ids_raises_when_ids_are_missing() -> None: + # given: missing class_id / tracker_id crash with a clear ValueError, + # raised before any mask work + masks, boxes, class_id = _single_mask_inputs() + predictions = _build_dense_detections(masks, boxes, class_id, device="cpu") + predictions.class_id = None + + # when / then + with pytest.raises(ValueError, match="resolve color by class"): + _resolve_color_ids(predictions, "CLASS", torch.device("cpu")) + with pytest.raises(ValueError, match="resolve color by track"): + _resolve_color_ids(predictions, "TRACK", torch.device("cpu")) + + +def test_resolve_color_ids_passes_pending_track_sentinel_through() -> None: + # given: sv's pending-track id (-1) must reach the caller unmapped so it + # can be painted with sv's gray + masks, boxes, class_id = _single_mask_inputs() + predictions = _build_dense_detections(masks, boxes, class_id, device="cpu") + predictions.bboxes_metadata = [{"tracker_id": -1}] + + # when + ids = _resolve_color_ids(predictions, "TRACK", torch.device("cpu")) + + # then + assert ids.tolist() == [-1] + + +def _make_scene(seed: int, h: int = SCENE_H, w: int = SCENE_W) -> np.ndarray: + rng = np.random.default_rng(seed) + yy = np.linspace(0, 1, h, dtype=np.float32)[:, None] + xx = np.linspace(0, 1, w, dtype=np.float32)[None, :] + b = 180.0 * xx + 40.0 * yy + g = 150.0 * yy + 50.0 * (1.0 - xx) + r = 160.0 * (1.0 - xx) * (1.0 - yy) + 60.0 + grad = np.stack([b, g, r], axis=2) + noise = rng.integers(0, 255, (h, w, 3)).astype(np.float32) + noise = cv2.GaussianBlur(noise, (0, 0), sigmaX=9) + scene = grad + 0.45 * (noise - 127.0) + return np.clip(scene, 0, 255).astype(np.uint8) + + +def _ellipse_mask( + cx: float, cy: float, ax: float, ay: float, angle: float = 0.0 +) -> np.ndarray: + mask = np.zeros((SCENE_H, SCENE_W), dtype=np.uint8) + cv2.ellipse( + mask, (int(cx), int(cy)), (int(ax), int(ay)), float(angle), 0, 360, 1, -1 + ) + return mask.astype(bool) + + +def _rect_mask(x1: int, y1: int, x2: int, y2: int) -> np.ndarray: + mask = np.zeros((SCENE_H, SCENE_W), dtype=bool) + mask[y1 : y2 + 1, x1 : x2 + 1] = True + return mask + + +def _tight_xyxy(mask: np.ndarray) -> list: + ys, xs = np.where(mask) + return [int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())] + + +def _scenario_disjoint_masks() -> list: + return [ + _ellipse_mask(180, 140, 120, 90), + _rect_mask(430, 60, 640, 240), + _ellipse_mask(760, 400, 140, 100, angle=30), + ] + + +def _scenario_partial_overlap() -> list: + return [ + _ellipse_mask(360, 270, 240, 180), + _ellipse_mask(600, 270, 240, 180), + ] + + +def _scenario_three_level_nesting() -> list: + return [ + _ellipse_mask(480, 270, 320, 220), + _ellipse_mask(480, 270, 180, 130), + _ellipse_mask(480, 270, 80, 55), + ] + + +def _scenario_three_way_chain() -> list: + masks = [ + _ellipse_mask(350, 220, 200, 150, angle=15), + _ellipse_mask(510, 300, 200, 150, angle=-10), + _ellipse_mask(650, 210, 200, 150, angle=25), + ] + assert (masks[0] & masks[1] & masks[2]).sum() > 0 + return masks + + +def _scenario_twelve_mask_cluster() -> list: + masks = [] + center_x, center_y = 480, 270 + for i in range(12): + angle = 2 * np.pi * i / 12 + cx = center_x + 150 * np.cos(angle) + cy = center_y + 100 * np.sin(angle) + masks.append( + _ellipse_mask(cx, cy, 170 + 7 * i, 90 + 4 * i, angle=np.degrees(angle)) + ) + common = masks[0].copy() + for mask in masks[1:]: + common &= mask + assert common.sum() > 0 + return masks + + +def _scenario_edge_touching() -> list: + # masks clipped by every frame edge + return [ + _rect_mask(0, 0, 199, 159), + _rect_mask(SCENE_W - 240, SCENE_H - 180, SCENE_W - 1, SCENE_H - 1), + _ellipse_mask(0, SCENE_H // 2, 160, 120), + _ellipse_mask(SCENE_W - 1, 100, 180, 140), + ] + + +def _scenario_full_frame() -> list: + # one mask covering every pixel plus a nested one + return [ + _rect_mask(0, 0, SCENE_W - 1, SCENE_H - 1), + _ellipse_mask(480, 270, 180, 130), + ] + + +def _random_blob_masks(seed: int = 7, n: int = 15) -> list: + # 15 random box-bounded blobs, including a single-pixel near-empty mask + rng = np.random.default_rng(seed) + masks = [] + for i in range(n): + crop_w = int(rng.integers(90, 360)) + crop_h = int(rng.integers(90, 360)) + x1 = int(rng.integers(0, SCENE_W - crop_w - 1)) + y1 = int(rng.integers(0, SCENE_H - crop_h - 1)) + mask = np.zeros((SCENE_H, SCENE_W), dtype=bool) + if i != 7: + low = rng.random((max(2, crop_h // 24), max(2, crop_w // 24))) + up = torch.nn.functional.interpolate( + torch.from_numpy(low)[None, None].float(), + size=(crop_h, crop_w), + mode="bilinear", + align_corners=False, + )[0, 0].numpy() + blob = up > 0.45 + if not blob.any(): + blob[crop_h // 2, crop_w // 2] = True + mask[y1 : y1 + crop_h, x1 : x1 + crop_w] = blob + else: + mask[y1, x1] = True # single-pixel mask keeps a valid tight box + masks.append(mask) + return masks + + +OVERLAP_SCENARIOS = [ + _scenario_partial_overlap, + _scenario_three_level_nesting, + _scenario_three_way_chain, + _scenario_twelve_mask_cluster, + _scenario_edge_touching, + _scenario_full_frame, + _random_blob_masks, +] +OVERLAP_SCENARIO_IDS = [ + "partial_overlap", + "three_level_nesting", + "three_way_chain", + "twelve_mask_cluster", + "edge_touching", + "full_frame", + "random_blobs", +] + + +def _reference_blend_all( + scene: np.ndarray, masks: list, colors_bgr: np.ndarray, opacity: float +) -> np.ndarray: + """Order-independent blend-all reference: the overlay color of a pixel is + the mean of the covering masks' colors, alpha-composited once with the + scene (np.round is round-half-to-even, like the compositor).""" + stack = np.stack(masks).astype(np.float64) # (N, H, W) + count = stack.sum(axis=0) # (H, W) + premul = np.einsum("nhw,nc->hwc", stack, colors_bgr.astype(np.float64) * opacity) + hit = count > 0 + out = scene.astype(np.float64) + out[hit] = np.round(premul[hit] / count[hit][:, None] + (1.0 - opacity) * out[hit]) + return out.astype(np.uint8) + + +def _detections_and_colors(masks: list, device: str, rle: bool = False): + boxes = np.asarray([_tight_xyxy(mask) for mask in masks], dtype=np.int32) + class_id = (np.arange(len(masks)) % 10).astype(np.int32) + if rle: + detections = _build_rle_detections(masks, boxes, class_id, device=device) + else: + detections = _build_dense_detections( + np.stack(masks, axis=0), boxes, class_id, device=device + ) + colors_bgr = np.asarray( + [PALETTE.by_idx(int(c)).as_bgr() for c in class_id], dtype=np.uint8 + ) + return detections, colors_bgr + + +def _composite_bgr( + scene_bgr: np.ndarray, + detections: InstanceDetections, + colors_bgr: np.ndarray, + opacity: float, + device: str = "cpu", +) -> np.ndarray: + """Test adapter: the production compositor is tensor-only (CHW RGB uint8 + in, the same tensor out, mutated in place) โ€” wrap numpy HWC BGR scenes and + colors so parity checks against sv / the numpy reference stay convenient.""" + scene_t = ( + torch.from_numpy(scene_bgr[:, :, ::-1].copy()) + .permute(2, 0, 1) + .contiguous() + .to(device) + ) + colors_rgb_t = torch.from_numpy(np.ascontiguousarray(colors_bgr[:, ::-1])).to( + device + ) + out = gpu_mask_composite(scene_t, detections.mask, colors_rgb_t, opacity) + return out.permute(1, 2, 0).cpu().numpy()[:, :, ::-1] + + +def _runs_to_dense(runs: np.ndarray, h: int, w: int) -> np.ndarray: + flat = np.zeros(h * w, dtype=bool) + position, value = 0, False + for run in runs: + flat[position : position + run] = value + position += int(run) + value = not value + return flat.reshape((h, w), order="F") + + +@pytest.mark.parametrize("seed", [3, 11, 42]) +def test_coco_rle_counts_decoder_matches_pycocotools(seed: int) -> None: + # given: blobs with long background runs (multi-char varints) and jagged + # boundaries (negative deltas) + masks = _random_blob_masks(seed=seed, n=6) + + for mask in masks: + encoded = mask_utils.encode(np.asfortranarray(mask.astype(np.uint8))) + + # when + runs = _coco_rle_counts_to_runs(encoded["counts"]) + rebuilt = _runs_to_dense(runs, SCENE_H, SCENE_W) + + # then + assert np.array_equal(rebuilt, mask_utils.decode(encoded).astype(bool)) + + +def test_coco_rle_counts_decoder_accepts_uncompressed_lists() -> None: + assert np.array_equal( + _coco_rle_counts_to_runs([3, 2, 5]), np.array([3, 2, 5], dtype=np.int64) + ) + assert _coco_rle_counts_to_runs(b"").size == 0 + + +@pytest.mark.parametrize("device", DEVICES) +def test_rle_to_dense_masks_matches_pycocotools(device: str) -> None: + # given + masks = _random_blob_masks(seed=5, n=4) + payloads = [ + mask_utils.encode(np.asfortranarray(mask.astype(np.uint8)))["counts"] + for mask in masks + ] + rle = InstancesRLEMasks(image_size=(SCENE_H, SCENE_W), masks=payloads) + + # when + dense = _rle_to_dense_masks(rle, torch.device(device)) + + # then + assert dense.shape == (len(masks), SCENE_H, SCENE_W) + assert dense.dtype == torch.bool + assert np.array_equal(dense.cpu().numpy(), np.stack(masks)) + + +@pytest.mark.parametrize("device", DEVICES) +def test_gpu_mask_composite_matches_sv_annotator_on_disjoint_masks( + device: str, +) -> None: + # given: no overlaps โ€” single-covered pixels use the exact same + # premultiplied blend as sv (round-half-even like cvRound). Observed max + # abs diff on this scenario: 0 (bit-exact); the tolerance of 1 covers + # last-ulp division differences on other data. + scene = _make_scene(101) + masks = _scenario_disjoint_masks() + detections, colors_bgr = _detections_and_colors(masks, device=device) + annotator = sv.MaskAnnotator( + color=PALETTE, color_lookup=sv.ColorLookup.CLASS, opacity=OPACITY + ) + expected = annotator.annotate( + scene=scene.copy(), detections=to_supervision_for_annotation(detections) + ) + + # when + actual = _composite_bgr(scene, detections, colors_bgr, OPACITY, device=device) + + # then + max_diff = int(np.abs(expected.astype(np.int16) - actual.astype(np.int16)).max()) + assert max_diff <= 1 + mismatched_share = 1.0 - float((expected == actual).all(axis=2).mean()) + assert mismatched_share < 0.001 + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("rle", [False, True], ids=["dense", "rle"]) +@pytest.mark.parametrize("scenario", OVERLAP_SCENARIOS, ids=OVERLAP_SCENARIO_IDS) +def test_gpu_mask_composite_matches_blend_all_reference( + scenario, rle: bool, device: str +) -> None: + # given + scene = _make_scene(101) + masks = scenario() + detections, colors_bgr = _detections_and_colors(masks, device=device, rle=rle) + expected = _reference_blend_all(scene, masks, colors_bgr, OPACITY) + + # when + actual = _composite_bgr(scene, detections, colors_bgr, OPACITY, device=device) + + # then + assert float((expected == actual).all(axis=2).mean()) == 1.0 + assert int(np.abs(expected.astype(np.int16) - actual.astype(np.int16)).max()) == 0 + + +@pytest.mark.parametrize("device", DEVICES) +def test_gpu_mask_composite_leaves_unmasked_pixels_untouched(device: str) -> None: + # given + scene = _make_scene(404) + masks = _scenario_disjoint_masks() + detections, colors_bgr = _detections_and_colors(masks, device=device) + covered = np.stack(masks).any(axis=0) + + # when + actual = _composite_bgr(scene, detections, colors_bgr, OPACITY, device=device) + + # then: sv's addWeighted re-blends the whole frame; the compositor must + # leave uncovered pixels BITWISE untouched + assert np.array_equal(actual[~covered], scene[~covered]) + assert not np.array_equal(actual[covered], scene[covered]) + + +@pytest.mark.parametrize("device", DEVICES) +def test_gpu_mask_composite_rle_carrier_matches_dense_carrier(device: str) -> None: + # given: the same masks through both carriers + scene = _make_scene(505) + masks = _scenario_twelve_mask_cluster() + dense_detections, colors_bgr = _detections_and_colors(masks, device=device) + rle_detections, _ = _detections_and_colors(masks, device=device, rle=True) + + # when + from_dense = _composite_bgr( + scene, dense_detections, colors_bgr, OPACITY, device=device + ) + from_rle = _composite_bgr(scene, rle_detections, colors_bgr, OPACITY, device=device) + + # then + assert np.array_equal(from_dense, from_rle) + + +@pytest.mark.parametrize("device", DEVICES) +def test_gpu_mask_composite_is_order_independent(device: str) -> None: + # given: blend-all semantics must not depend on detection order + scene = _make_scene(606) + masks = _scenario_three_way_chain() + detections, colors_bgr = _detections_and_colors(masks, device=device) + permutation = [2, 0, 1] + permuted, permuted_colors = _detections_and_colors( + [masks[i] for i in permutation], device=device + ) + # keep each mask's color stable under the permutation + permuted.class_id = detections.class_id[permutation] + permuted_colors = colors_bgr[permutation] + + # when + original = _composite_bgr(scene, detections, colors_bgr, OPACITY, device=device) + shuffled = _composite_bgr(scene, permuted, permuted_colors, OPACITY, device=device) + + # then + assert np.array_equal(original, shuffled) + + +def test_gpu_mask_composite_accepts_numpy_mask_stack() -> None: + # given: a numpy-carried dense mask stack is uploaded and painted like the + # torch carrier (there is no sv fallback to route it to any more) + scene = _make_scene(111, h=128, w=128) + masks = np.zeros((1, 128, 128), dtype=bool) + masks[0, 20:60, 30:90] = True + colors_bgr = np.asarray([PALETTE.by_idx(1).as_bgr()], dtype=np.uint8) + expected = _reference_blend_all(scene, list(masks), colors_bgr, OPACITY) + + scene_t = torch.from_numpy(scene[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + out = gpu_mask_composite( + scene_t, + masks, + torch.from_numpy(np.ascontiguousarray(colors_bgr[:, ::-1])), + OPACITY, + ) + + # then + actual = out.permute(1, 2, 0).numpy()[:, :, ::-1] + assert np.array_equal(actual, expected) + + +def test_gpu_mask_composite_casts_non_bool_masks() -> None: + # given: float dense masks (previously served by the sv fallback's + # astype(bool)) are cast on device โ€” nonzero values paint + scene = _make_scene(222, h=96, w=96) + masks = np.zeros((1, 96, 96), dtype=bool) + masks[0, 10:40, 10:40] = True + colors_bgr = np.asarray([PALETTE.by_idx(2).as_bgr()], dtype=np.uint8) + expected = _reference_blend_all(scene, list(masks), colors_bgr, OPACITY) + + scene_t = torch.from_numpy(scene[:, :, ::-1].copy()).permute(2, 0, 1).contiguous() + out = gpu_mask_composite( + scene_t, + torch.from_numpy(masks).float(), + torch.from_numpy(np.ascontiguousarray(colors_bgr[:, ::-1])), + OPACITY, + ) + + # then + actual = out.permute(1, 2, 0).numpy()[:, :, ::-1] + assert np.array_equal(actual, expected) + + +def test_gpu_mask_composite_rejects_mismatched_mask_canvas() -> None: + # given: mask canvas smaller than the scene โ€” silent slicing would paint + # misaligned masks; with the sv fallback gone this must raise loudly + scene = _make_scene(707, h=256, w=256) + masks, boxes, class_id = _single_mask_inputs(h=64, w=64) + detections = _build_dense_detections(masks, boxes, class_id, device="cpu") + colors_bgr = np.asarray([[255, 0, 0]], dtype=np.uint8) + + # when / then + with pytest.raises(ValueError, match="does not match scene"): + _composite_bgr(scene, detections, colors_bgr, OPACITY) + + +def test_gpu_mask_composite_with_all_false_masks_leaves_scene_unchanged() -> None: + # given: no foreground pixels at all (boxes play no role in the composite) + scene = _make_scene(808, h=128, w=128) + masks = np.zeros((1, 128, 128), dtype=bool) + boxes = np.array([[-50, -50, -10, -10]], dtype=np.int32) + detections = _build_dense_detections( + masks, boxes, np.array([0], dtype=np.int32), device="cpu" + ) + colors_bgr = np.asarray([[255, 0, 0]], dtype=np.uint8) + + # when + actual = _composite_bgr(scene, detections, colors_bgr, OPACITY) + + # then + assert np.array_equal(actual, scene) + + +@requires_cuda +@pytest.mark.parametrize("rle", [False, True], ids=["dense", "rle"]) +def test_gpu_mask_composite_chw_rgb_tensor_scene_matches_reference( + rle: bool, +) -> None: + # given: the WorkflowImageData.tensor_image contract - CHW uint8 RGB on device + scene = _make_scene(303) + masks = _scenario_three_way_chain() + detections, colors_bgr = _detections_and_colors(masks, device="cuda", rle=rle) + scene_chw_rgb = ( + torch.from_numpy(scene[:, :, ::-1].copy()).permute(2, 0, 1).contiguous().cuda() + ) + expected = _reference_blend_all(scene, masks, colors_bgr, OPACITY) + + # when + annotated_tensor = gpu_mask_composite( + scene_chw_rgb, + detections.mask, + torch.from_numpy(np.ascontiguousarray(colors_bgr[:, ::-1])).cuda(), + OPACITY, + ) + + # then: result stays on device; converting back to HWC BGR matches the + # reference exactly + assert annotated_tensor.is_cuda + assert annotated_tensor.data_ptr() == scene_chw_rgb.data_ptr() # in-place + actual = annotated_tensor.permute(1, 2, 0).cpu().numpy()[:, :, ::-1] + assert float((expected == actual).all(axis=2).mean()) == 1.0 + assert int(np.abs(expected.astype(np.int16) - actual.astype(np.int16)).max()) == 0 + + +# -------------------------------------------------------------------------- +# Block-level tests (MaskVisualizationBlockV1.run) +# -------------------------------------------------------------------------- + + +def _tensor_backed_image( + scene_bgr: np.ndarray, device: str = "cpu" +) -> WorkflowImageData: + tensor = ( + torch.from_numpy(scene_bgr[:, :, ::-1].copy()) + .permute(2, 0, 1) + .contiguous() + .to(device) + ) + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), tensor_image=tensor + ) + + +def _numpy_backed_image(scene_bgr: np.ndarray) -> WorkflowImageData: + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="p"), + numpy_image=scene_bgr, + ) + + +def _run_block( + image: WorkflowImageData, + detections, + copy_image: bool = True, + color_axis: str = "CLASS", + opacity: float = OPACITY, +) -> WorkflowImageData: + return MaskVisualizationBlockV1().run( + image=image, + predictions=detections, + copy_image=copy_image, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis=color_axis, + opacity=opacity, + )["image"] + + +def _empty_detections(device: str = "cpu") -> InstanceDetections: + return InstanceDetections( + xyxy=torch.zeros((0, 4), dtype=torch.int32, device=device), + class_id=torch.zeros((0,), dtype=torch.int32, device=device), + confidence=torch.zeros((0,), device=device), + mask=torch.zeros((0, 8, 8), dtype=torch.bool, device=device), + ) + + +@pytest.mark.parametrize("device", DEVICES) +@pytest.mark.parametrize("rle", [False, True], ids=["dense", "rle"]) +def test_block_tensor_path_matches_reference(rle: bool, device: str) -> None: + # given + scene = _make_scene(120) + masks = _scenario_three_way_chain() + detections, colors_bgr = _detections_and_colors(masks, device=device, rle=rle) + image = _tensor_backed_image(scene, device=device) + expected = _reference_blend_all(scene, masks, colors_bgr, OPACITY) + + # when + out = _run_block(image, detections, copy_image=True) + + # then: tensor in -> tensor out, pixels match the reference + assert out._tensor_image is not None and out._numpy_image is None + actual = out._tensor_image.permute(1, 2, 0).cpu().numpy()[:, :, ::-1] + assert np.array_equal(actual, expected) + + +def test_block_copy_image_true_preserves_input_tensor() -> None: + # given + scene = _make_scene(121, h=128, w=128) + masks = np.zeros((1, 128, 128), dtype=bool) + masks[0, 20:60, 30:90] = True + detections = _build_dense_detections( + masks, + np.array([[30, 20, 89, 59]], dtype=np.int32), + np.array([1], dtype=np.int32), + device="cpu", + ) + image = _tensor_backed_image(scene) + before = image.tensor_image.clone() + + # when + out = _run_block(image, detections, copy_image=True) + + # then: independent storage, input untouched + assert out._tensor_image.data_ptr() != image.tensor_image.data_ptr() + assert torch.equal(image.tensor_image, before) + assert not torch.equal(out._tensor_image, before) + + +def test_block_copy_image_false_mutates_input_tensor_in_place() -> None: + # given + scene = _make_scene(122, h=128, w=128) + masks = np.zeros((1, 128, 128), dtype=bool) + masks[0, 20:60, 30:90] = True + detections = _build_dense_detections( + masks, + np.array([[30, 20, 89, 59]], dtype=np.int32), + np.array([1], dtype=np.int32), + device="cpu", + ) + image = _tensor_backed_image(scene) + input_tensor = image.tensor_image + before = input_tensor.clone() + + # when + out = _run_block(image, detections, copy_image=False) + + # then: same storage, annotated in place, numpy/base64 caches invalidated + assert out._tensor_image.data_ptr() == input_tensor.data_ptr() + assert not torch.equal(input_tensor, before) + assert image._numpy_image is None and image._base64_image is None + + +@pytest.mark.parametrize("copy_image", [True, False]) +def test_block_numpy_path_matches_sv_annotator_reference(copy_image: bool) -> None: + # given: numpy-sourced images take the pre-rewrite sv.MaskAnnotator path + # unchanged โ€” bit-exact vs a directly-constructed annotator (painter's + # algorithm on overlaps included, which the torch compositor diverges from) + scene = _make_scene(123) + masks = _scenario_three_way_chain() + detections, _ = _detections_and_colors(masks, device="cpu") + annotator = sv.MaskAnnotator( + color=PALETTE, color_lookup=sv.ColorLookup.CLASS, opacity=OPACITY + ) + expected = annotator.annotate( + scene=scene.copy(), detections=to_supervision_for_annotation(detections) + ) + image = _numpy_backed_image(scene.copy()) + + # when + out = _run_block(image, detections, copy_image=copy_image) + + # then: numpy in -> numpy out, sv's exact pixels + assert out._numpy_image is not None and out._tensor_image is None + assert np.array_equal(out._numpy_image, expected) + + +def test_block_numpy_path_copy_semantics() -> None: + # given + scene = _make_scene(124, h=128, w=128) + masks = np.zeros((1, 128, 128), dtype=bool) + masks[0, 20:60, 30:90] = True + detections = _build_dense_detections( + masks, + np.array([[30, 20, 89, 59]], dtype=np.int32), + np.array([1], dtype=np.int32), + device="cpu", + ) + + # when: copy_image=True leaves the caller's buffer untouched + image = _numpy_backed_image(scene.copy()) + out = _run_block(image, detections, copy_image=True) + assert not np.shares_memory(out._numpy_image, image.numpy_image) + assert np.array_equal(image.numpy_image, scene) + + # and: copy_image=False mutates the caller's buffer in place + image = _numpy_backed_image(scene.copy()) + buffer = image.numpy_image + out = _run_block(image, detections, copy_image=False) + assert np.shares_memory(out._numpy_image, buffer) + assert not np.array_equal(buffer, scene) + + +def test_block_empty_predictions_take_the_tensor_passthrough() -> None: + # given + scene = _make_scene(125, h=64, w=64) + image = _tensor_backed_image(scene) + + # when + out_copy = _run_block(image, _empty_detections(), copy_image=True) + out_share = _run_block(image, _empty_detections(), copy_image=False) + + # then: stays on-device; independent storage iff copy_image + assert out_copy._tensor_image is not None and out_copy._numpy_image is None + assert out_copy._tensor_image.data_ptr() != image.tensor_image.data_ptr() + assert torch.equal(out_copy._tensor_image, image.tensor_image) + assert out_share._tensor_image.data_ptr() == image.tensor_image.data_ptr() + + +def test_block_empty_predictions_on_numpy_sourced_image_stay_numpy() -> None: + # given + scene = _make_scene(126, h=64, w=64) + image = _numpy_backed_image(scene) + + # when + out = _run_block(image, _empty_detections(), copy_image=True) + + # then + assert out._numpy_image is not None and out._tensor_image is None + assert not np.shares_memory(out._numpy_image, image.numpy_image) + assert np.array_equal(out._numpy_image, image.numpy_image) + + +@pytest.mark.parametrize("device", DEVICES) +def test_block_track_colors_match_sv_annotator(device: str) -> None: + # given: disjoint masks colored by tracker_id (non-contiguous ids to + # catch any id-vs-index mixup) + scene = _make_scene(909) + masks = _scenario_disjoint_masks() + detections, _ = _detections_and_colors(masks, device=device) + tracker_ids = [12, 3, 27] + detections.bboxes_metadata = [{"tracker_id": tid} for tid in tracker_ids] + annotator = sv.MaskAnnotator( + color=PALETTE, color_lookup=sv.ColorLookup.TRACK, opacity=OPACITY + ) + expected = annotator.annotate( + scene=scene.copy(), detections=to_supervision_for_annotation(detections) + ) + image = _tensor_backed_image(scene, device=device) + + # when + out = _run_block(image, detections, color_axis="TRACK") + + # then + actual = out._tensor_image.permute(1, 2, 0).cpu().numpy()[:, :, ::-1] + max_diff = int(np.abs(expected.astype(np.int16) - actual.astype(np.int16)).max()) + assert max_diff <= 1 + mismatched_share = 1.0 - float((expected == actual).all(axis=2).mean()) + assert mismatched_share < 0.001 + + +def test_block_pending_track_id_paints_sv_gray() -> None: + # given: sv's pending-track sentinel (-1) maps to Color.GREY (128,128,128) + scene = np.full((128, 128, 3), 200, dtype=np.uint8) + masks = np.zeros((1, 128, 128), dtype=bool) + masks[0, 20:60, 30:90] = True + detections = _build_dense_detections( + masks, + np.array([[30, 20, 89, 59]], dtype=np.int32), + np.array([1], dtype=np.int32), + device="cpu", + ) + detections.bboxes_metadata = [{"tracker_id": -1}] + image = _tensor_backed_image(scene) + + # when + out = _run_block(image, detections, color_axis="TRACK") + + # then: masked pixels = round(128*0.5 + 200*0.5) = 164 on every channel + actual = out._tensor_image.permute(1, 2, 0).cpu().numpy()[:, :, ::-1] + assert np.array_equal( + actual[masks[0]], np.full((int(masks[0].sum()), 3), 164, dtype=np.uint8) + ) + assert np.array_equal(actual[~masks[0]], scene[~masks[0]]) + + +def test_block_missing_tracker_ids_raise_value_error() -> None: + # given + scene = _make_scene(127, h=64, w=64) + masks, boxes, class_id = _single_mask_inputs() + detections = _build_dense_detections(masks, boxes, class_id, device="cpu") + image = _tensor_backed_image(scene) + + # when / then: raised BEFORE any painting, input untouched + before = image.tensor_image.clone() + with pytest.raises(ValueError, match="resolve color by track"): + _run_block(image, detections, color_axis="TRACK", copy_image=False) + assert torch.equal(image.tensor_image, before) + + +def test_block_raises_for_unsupported_color_axis() -> None: + scene = _make_scene(128, h=64, w=64) + masks, boxes, class_id = _single_mask_inputs() + detections = _build_dense_detections(masks, boxes, class_id, device="cpu") + + with pytest.raises(ValueError, match="color_axis"): + _run_block(_tensor_backed_image(scene), detections, color_axis="SOMETHING") + + +def test_block_raises_for_non_instance_detections() -> None: + scene = _make_scene(129, h=64, w=64) + + class _NotDetections: + xyxy = torch.ones((1, 4)) + + with pytest.raises(ValueError, match="instance segmentation"): + _run_block(_tensor_backed_image(scene), _NotDetections()) + + +def test_block_raises_for_mask_count_mismatch() -> None: + # given: 1 box but 2 RLE payloads (previously silently sv-fallback-routed) + scene = _make_scene(130, h=64, w=64) + masks, boxes, class_id = _single_mask_inputs() + detections = InstanceDetections( + xyxy=torch.tensor(boxes, dtype=torch.int32), + class_id=torch.tensor(class_id, dtype=torch.int32), + confidence=torch.full((1,), 0.9), + mask=InstancesRLEMasks(image_size=(64, 64), masks=[b"", b""]), + ) + + with pytest.raises(ValueError, match="RLE masks"): + _run_block(_tensor_backed_image(scene), detections) + + +def test_block_raises_for_missing_mask_carrier() -> None: + scene = _make_scene(131, h=64, w=64) + masks, boxes, class_id = _single_mask_inputs() + detections = _build_dense_detections(masks, boxes, class_id, device="cpu") + detections.mask = None + + with pytest.raises(ValueError, match="no usable mask"): + _run_block(_tensor_backed_image(scene), detections) + + +def test_block_canvas_mismatch_raises_without_partial_mutation() -> None: + # given: the composite validates before its single staged write, so a + # raising run must leave a copy_image=False input bitwise untouched + scene = _make_scene(132, h=256, w=256) + masks, boxes, class_id = _single_mask_inputs(h=64, w=64) + detections = _build_dense_detections(masks, boxes, class_id, device="cpu") + image = _tensor_backed_image(scene) + before = image.tensor_image.clone() + + with pytest.raises(ValueError, match="does not match scene"): + _run_block(image, detections, copy_image=False) + assert torch.equal(image.tensor_image, before) + + +# -------------------------------------------------------------------------- +# Sync audit: the dense path must enqueue a fixed, N-independent op sequence +# with no device->host reads. +# -------------------------------------------------------------------------- + + +class _DispatchRecorder(TorchDispatchMode): + def __init__(self): + super().__init__() + self.calls = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + self.calls.append((str(func), kwargs)) + return func(*args, **kwargs) + + +def _dense_composite_trace(n: int) -> list: + scene = ( + torch.from_numpy(_make_scene(999, h=96, w=128)[:, :, ::-1].copy()) + .permute(2, 0, 1) + .contiguous() + ) + rng = np.random.default_rng(n) + masks = torch.from_numpy(rng.random((n, 96, 128)) > 0.6) + colors = torch.randint(0, 255, (n, 3), dtype=torch.uint8) + recorder = _DispatchRecorder() + with recorder: + gpu_mask_composite(scene, masks, colors, OPACITY) + return recorder.calls + + +def test_dense_composite_dispatches_no_sync_ops() -> None: + # The structural zero-sync invariant (asserted on op names so it holds on + # CPU-only runners too): no data-dependent indexing (`nonzero`, + # `masked_select`), no host readback (`_local_scalar_dense` is `.item()`, + # `aten.item` its alias), and no op asked to change a tensor's device + # (`_to_copy`/`to`/`copy_` are dtype-only on this path โ€” a `device=` kwarg + # would be the D2H/H2D tell). + calls = _dense_composite_trace(8) + op_names = [name for name, _ in calls] + assert calls, "dispatch trace is empty - the mode did not record" + forbidden_fragments = ("nonzero", "_local_scalar_dense", "masked_select") + for name in op_names: + for fragment in forbidden_fragments: + assert fragment not in name, f"sync-inducing op in dense path: {name}" + assert not name.startswith("aten.item"), f"host readback in dense path: {name}" + for name, kwargs in calls: + assert kwargs.get("device") is None, ( + f"{name} was asked to change device ({kwargs['device']}) inside " + "the dense composite - that is a cross-device copy" + ) + # the single staged write into the caller's storage is present + assert any(name.startswith("aten.copy_") for name in op_names) + + +def test_dense_composite_dispatch_count_is_n_independent() -> None: + # Fixed-shape contract: the op SEQUENCE (not just the count) must be + # identical for N=1/8/32 - nothing in the path branches on data or N. + traces = {n: [name for name, _ in _dense_composite_trace(n)] for n in (1, 8, 32)} + assert traces[1] == traces[8] == traces[32] + assert len(traces[1]) < 40 # bounded: a handful of fixed ops per frame diff --git a/tests/workflows/unit_tests/core_steps/visualizations/test_rich_label.py b/tests/workflows/unit_tests/core_steps/visualizations/test_rich_label.py index 359d4a6a28..9fbae4e432 100644 --- a/tests/workflows/unit_tests/core_steps/visualizations/test_rich_label.py +++ b/tests/workflows/unit_tests/core_steps/visualizations/test_rich_label.py @@ -1,8 +1,10 @@ import numpy as np import pytest import supervision as sv +import torch from pydantic import ValidationError +from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION from inference.core.workflows.core_steps.visualizations.common.fonts import ( FONTS_REGISTRY, ) @@ -10,10 +12,28 @@ RichLabelManifest, RichLabelVisualizationBlockV1, ) +from inference.core.workflows.core_steps.visualizations.rich_label.v1_tensor import ( + RichLabelVisualizationBlockV1 as RichLabelVisualizationBlockV1Tensor, +) +from inference.core.workflows.execution_engine.constants import CLASS_NAMES_KEY from inference.core.workflows.execution_engine.entities.base import ( ImageParentMetadata, WorkflowImageData, ) +from inference_models.models.base.object_detection import Detections as NativeDetections + +# The loader binds `_tensor` siblings under the same names when +# ENABLE_TENSOR_DATA_REPRESENTATION is set - hence the flag-opposed +# _NUMPY_ONLY / _TENSOR_ONLY split below. +_NUMPY_ONLY = pytest.mark.skipif( + ENABLE_TENSOR_DATA_REPRESENTATION, + reason="loader binds the tensor-native sibling under " + "ENABLE_TENSOR_DATA_REPRESENTATION โ€” see the *_tensor_native parity test", +) +_TENSOR_ONLY = pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-native variant; runs only with ENABLE_TENSOR_DATA_REPRESENTATION=True", +) def _build_image(height: int = 400, width: int = 400) -> WorkflowImageData: @@ -168,6 +188,7 @@ def test_manifest_font_family_enum_matches_fonts_registry() -> None: } +@_NUMPY_ONLY def test_rich_label_block_is_registered_in_loader() -> None: # given from inference.core.workflows.core_steps.loader import load_blocks @@ -176,6 +197,19 @@ def test_rich_label_block_is_registered_in_loader() -> None: assert RichLabelVisualizationBlockV1 in load_blocks() +@_TENSOR_ONLY +def test_rich_label_tensor_block_is_registered_in_loader() -> None: + # given + from inference.core.workflows.core_steps.loader import load_blocks + + # when + blocks = load_blocks() + + # then + assert RichLabelVisualizationBlockV1Tensor in blocks + assert RichLabelVisualizationBlockV1 not in blocks + + def test_rich_label_visualization_block(bundled_fonts) -> None: # given block = RichLabelVisualizationBlockV1() @@ -286,3 +320,152 @@ def test_rich_label_visualization_block_with_max_line_length(bundled_fonts) -> N assert not np.array_equal( output.get("image").numpy_image, np.zeros((400, 400, 3), dtype=np.uint8) ) + + +def _build_native_predictions(class_names=("person", "car")) -> NativeDetections: + boxes = [[10 + 120 * i, 10, 100 + 120 * i, 100] for i in range(len(class_names))] + + return NativeDetections( + xyxy=torch.tensor(boxes, dtype=torch.float32), + class_id=torch.arange(len(class_names), dtype=torch.long), + confidence=torch.full((len(class_names),), 0.9, dtype=torch.float32), + image_metadata={ + CLASS_NAMES_KEY: {i: name for i, name in enumerate(class_names)} + }, + ) + + +def _run_tensor_block(block: RichLabelVisualizationBlockV1Tensor, **overrides): + kwargs = { + "image": _build_image(), + "predictions": _build_native_predictions(), + "copy_image": True, + "color_palette": "DEFAULT", + "palette_size": 10, + "custom_colors": None, + "color_axis": "CLASS", + "text": "Class", + "text_position": "TOP_LEFT", + "text_color": "WHITE", + "font_family": "geist_mono", + "text_size_mode": "Manual", + "font_size": 14, + "text_padding": 10, + "border_radius": 0, + "max_line_length": None, + } + kwargs.update(overrides) + + return block.run(**kwargs) + + +@_TENSOR_ONLY +def test_rich_label_visualization_block_tensor_native(bundled_fonts) -> None: + # given + block = RichLabelVisualizationBlockV1Tensor() + + # when + output = _run_tensor_block(block) + + # then + assert output is not None + assert "image" in output + assert hasattr(output.get("image"), "numpy_image") + assert output.get("image").numpy_image.shape == (400, 400, 3) + assert output.get("image").numpy_image.dtype == np.uint8 + assert not np.array_equal( + output.get("image").numpy_image, np.zeros((400, 400, 3), dtype=np.uint8) + ), "Image should be modified by label rendering" + + +@_TENSOR_ONLY +def test_rich_label_tensor_native_output_matches_numpy_block(bundled_fonts) -> None: + # given - the same frame and semantically-identical predictions in both + # representations; the tensor sibling must reproduce the numpy block's + # rendering byte-for-byte (it reuses the numpy drawing internals). + numpy_block = RichLabelVisualizationBlockV1() + tensor_block = RichLabelVisualizationBlockV1Tensor() + + # when + numpy_output = _run_block(numpy_block, text_size_mode="Automatic") + tensor_output = _run_tensor_block(tensor_block, text_size_mode="Automatic") + + # then + assert np.array_equal( + numpy_output["image"].numpy_image, tensor_output["image"].numpy_image + ) + + +@_TENSOR_ONLY +def test_rich_label_visualization_block_tensor_native_raises_on_unknown_font() -> None: + # given - an unknown font id may reach run() through an input selector + block = RichLabelVisualizationBlockV1Tensor() + + # when + with pytest.raises(ValueError) as error: + _ = _run_tensor_block(block, font_family="comic_sans") + + # then + assert "comic_sans" in str(error.value) + + +@_TENSOR_ONLY +def test_rich_label_visualization_block_with_empty_detections_tensor_native() -> None: + # given + block = RichLabelVisualizationBlockV1Tensor() + empty_predictions = NativeDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + ) + + # when + output = _run_tensor_block(block, predictions=empty_predictions) + + # then + assert output.get("image").numpy_image.shape == (400, 400, 3) + assert np.array_equal( + output.get("image").numpy_image, np.zeros((400, 400, 3), dtype=np.uint8) + ), "Image should be unmodified when there are no detections" + + +@_TENSOR_ONLY +def test_rich_label_tensor_native_empty_predictions_passthrough_is_device_resident() -> ( + None +): + # given - a tensor-source image and an empty prediction: the sibling must + # pass the tensor representation through without materialising numpy + block = RichLabelVisualizationBlockV1Tensor() + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="some"), + tensor_image=torch.zeros((3, 240, 320), dtype=torch.uint8), + ) + empty_predictions = NativeDetections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + ) + + # when + copied_output = _run_tensor_block( + block, image=image, predictions=empty_predictions, copy_image=True + ) + shared_output = _run_tensor_block( + block, image=image, predictions=empty_predictions, copy_image=False + ) + + # then + assert image._numpy_image is None, "passthrough must not materialise numpy" + for output in [copied_output, shared_output]: + assert output["image"].is_tensor_materialised() is True + assert output["image"]._numpy_image is None + assert torch.equal(output["image"].tensor_image, image.tensor_image) + # copy semantics: independent storage when copy_image=True, shared otherwise + assert ( + copied_output["image"].tensor_image.data_ptr() + != image.tensor_image.data_ptr() + ) + assert ( + shared_output["image"].tensor_image.data_ptr() + == image.tensor_image.data_ptr() + ) diff --git a/tests/workflows/unit_tests/core_steps/visualizations/test_viz_phase_transfer_audit.py b/tests/workflows/unit_tests/core_steps/visualizations/test_viz_phase_transfer_audit.py new file mode 100644 index 0000000000..6d8d740ace --- /dev/null +++ b/tests/workflows/unit_tests/core_steps/visualizations/test_viz_phase_transfer_audit.py @@ -0,0 +1,329 @@ +"""Whole-viz-phase host<->device transfer audit. + +Runs the tensor-native mask compositor followed by the label sprite painter โ€” +the full viz phase of the tensor pipeline โ€” under ``TorchDispatchMode`` and +asserts the combined op trace is transfer-clean: + +* no ``aten.nonzero`` and no ``aten._local_scalar_dense`` (both force a + deviceโ†’host sync), +* no deviceโ†’host ``_to_copy`` / ``copy_`` at all, +* every hostโ†’device ``_to_copy`` / ``copy_`` carries ``non_blocking=True``. + +The dispatcher exposes the flag directly (``copy_`` passes it as the third +positional argument, ``_to_copy`` in kwargs). On the CPU test host every +tensor lives on the CPU, so cross-device classification is vacuous there โ€” +the flag is asserted anyway via the structural fact that on the label tensor +path EVERY ``aten.copy_`` is a staged hostโ†’device upload (the ring table +upload plus, cold, the sprite pixel/index payloads), while the mask phase +dispatches only same-device copies. The identical assertions become the real +H2D/D2H checks when this suite runs on a CUDA host. + +Also pins the warm-path dispatch budget: the label paste enqueues a fixed, +label-count-independent number of ops. +""" + +from collections import Counter + +import numpy as np +import torch +from torch.utils._python_dispatch import TorchDispatchMode + +from inference.core.workflows.core_steps.common.tensor_native import ( + attach_native_detection_metadata, +) +from inference.core.workflows.core_steps.visualizations.label.v1_tensor import ( + LabelVisualizationBlockV1, +) +from inference.core.workflows.core_steps.visualizations.mask.v1_tensor import ( + MaskVisualizationBlockV1, +) +from inference.core.workflows.execution_engine.entities.base import ( + ImageParentMetadata, + WorkflowImageData, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.object_detection import Detections + +CLASS_NAMES = {0: "cat", 1: "dog", 2: "goggles", 3: "bird"} +SCENE_H, SCENE_W = 480, 640 +N_BOXES = 3 + +_LABEL_RUN_KWARGS = dict( + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="CLASS", + text="Class and Confidence", + text_position="TOP_LEFT", + text_color="WHITE", + text_scale=1.0, + text_thickness=1, + text_padding=10, + border_radius=0, +) + +_MASK_RUN_KWARGS = dict( + copy_image=True, + color_palette="DEFAULT", + palette_size=10, + custom_colors=None, + color_axis="CLASS", + opacity=0.5, +) + + +class _TransferAudit(TorchDispatchMode): + """Record every dispatched aten op; for ``copy_`` / ``_to_copy`` also the + source/target devices and the ``non_blocking`` flag exactly as the + dispatcher exposes them.""" + + def __init__(self): + super().__init__() + self.ops = [] + self.copies = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + name = str(func) + self.ops.append(name) + if name == "aten.copy_.default": + self.copies.append( + { + "op": name, + "dst": args[0].device, + "src": args[1].device, + "non_blocking": ( + bool(args[2]) + if len(args) > 2 + else bool(kwargs.get("non_blocking", False)) + ), + } + ) + elif name == "aten._to_copy.default": + source = args[0] + target = kwargs.get("device") + self.copies.append( + { + "op": name, + "dst": ( + torch.device(target) if target is not None else source.device + ), + "src": source.device, + "non_blocking": bool(kwargs.get("non_blocking", False)), + } + ) + return func(*args, **kwargs) + + +def _assert_no_forbidden_ops(ops) -> None: + offenders = [ + name for name in ops if "nonzero" in name or "_local_scalar_dense" in name + ] + assert offenders == [], f"sync-forcing ops dispatched: {offenders}" + + +def _assert_no_device_to_host(copies) -> None: + downloads = [ + record + for record in copies + if record["src"].type == "cuda" and record["dst"].type == "cpu" + ] + assert downloads == [], f"device->host transfers dispatched: {downloads}" + + +def _assert_uploads_non_blocking(copies) -> None: + blocking = [ + record + for record in copies + if record["src"].type == "cpu" + and record["dst"].type == "cuda" + and not record["non_blocking"] + ] + assert blocking == [], f"blocking host->device transfers dispatched: {blocking}" + + +def _tensor_image(seed: int = 7) -> WorkflowImageData: + rng = np.random.default_rng(seed) + scene = rng.integers(0, 255, (3, SCENE_H, SCENE_W)).astype(np.uint8) + return WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.from_numpy(scene), + ) + + +def _mirrored_od_detections(n: int = N_BOXES) -> Detections: + boxes = np.array( + [ + [40.5, 60.25, 200.75, 200.125], + [250.0, 90.0, 420.0, 260.0], + [90.0, 260.0, 300.0, 430.0], + [10.0, 10.0, 30.0, 30.0], + ], + dtype=np.float32, + )[:n] + detections = Detections( + xyxy=torch.tensor(boxes, dtype=torch.float32), + class_id=torch.tensor([index % 3 for index in range(n)], dtype=torch.long), + confidence=torch.tensor( + np.linspace(0.42, 0.99, max(n, 1))[:n], dtype=torch.float32 + ), + image_metadata=None, + bboxes_metadata=None, + ) + return attach_native_detection_metadata( + detections=detections, + image=_tensor_image(), + class_names=CLASS_NAMES, + prediction_type="object-detection", + ) + + +def _disjoint_od_detections(n: int) -> Detections: + """Vertically separated boxes whose TOP_LEFT labels can never overlap, so + the paste's owner-resolution branch state is identical for every ``n``.""" + boxes = np.asarray( + [ + [60.0, 60.0 + 110.0 * index, 240.0, 130.0 + 110.0 * index] + for index in range(n) + ], + dtype=np.float32, + ) + detections = Detections( + xyxy=torch.tensor(boxes, dtype=torch.float32), + class_id=torch.tensor(list(range(n)), dtype=torch.long), + confidence=torch.tensor(np.linspace(0.5, 0.9, n), dtype=torch.float32), + image_metadata=None, + bboxes_metadata=None, + ) + return attach_native_detection_metadata( + detections=detections, + image=_tensor_image(), + class_names=CLASS_NAMES, + prediction_type="object-detection", + ) + + +def _mirrored_is_detections(n: int = 2) -> InstanceDetections: + base = _mirrored_od_detections(n) + masks = torch.zeros((n, SCENE_H, SCENE_W), dtype=torch.bool) + masks[0, 70:180, 50:190] = True + if n > 1: + masks[1, 100:250, 260:410] = True + return InstanceDetections( + xyxy=base.xyxy, + class_id=base.class_id, + confidence=base.confidence, + mask=masks, + image_metadata=base.image_metadata, + bboxes_metadata=base.bboxes_metadata, + ) + + +# --------------------------------------------------------------------------- # +# whole viz phase: mask composite -> label paste, warm caches +# --------------------------------------------------------------------------- # + + +def test_viz_phase_mask_then_label_combined_trace_is_transfer_clean() -> None: + # given + mask_block = MaskVisualizationBlockV1() + label_block = LabelVisualizationBlockV1() + segmentation = _mirrored_is_detections() + detections = _mirrored_od_detections() + # warm every cache: the mask palette LUT, the label sprite cache + flat + # templates + the pinned table ring slabs + mask_block.run(image=_tensor_image(), predictions=segmentation, **_MASK_RUN_KWARGS) + label_block.run(image=_tensor_image(), predictions=detections, **_LABEL_RUN_KWARGS) + + # when - the audited steady-state viz phase, label consuming mask's output + mask_audit = _TransferAudit() + with mask_audit: + masked = mask_block.run( + image=_tensor_image(), predictions=segmentation, **_MASK_RUN_KWARGS + ) + label_audit = _TransferAudit() + with label_audit: + result = label_block.run( + image=masked["image"], predictions=detections, **_LABEL_RUN_KWARGS + ) + + # then - the combined trace never syncs and never crosses the bus blocking + combined_ops = mask_audit.ops + label_audit.ops + combined_copies = mask_audit.copies + label_audit.copies + _assert_no_forbidden_ops(combined_ops) + _assert_no_device_to_host(combined_copies) + _assert_uploads_non_blocking(combined_copies) + # mask phase: everything is device-resident โ€” its only in-place copy is + # the same-device staged scene write, no host->device staging at all + for record in mask_audit.copies: + if record["op"] == "aten.copy_.default": + assert record["src"] == record["dst"] + # label phase: with a warm sprite cache EVERY `copy_` is a staged + # host->device upload โ€” exactly one (the packed table through the pinned + # ring) and it must carry non_blocking=True, also on the CPU host where + # the copy itself is trivial + label_copies = [ + record for record in label_audit.copies if record["op"] == "aten.copy_.default" + ] + assert len(label_copies) == 1, f"expected only the table upload: {label_copies}" + assert all(record["non_blocking"] for record in label_copies) + assert result["image"]._tensor_image is not None + + +def test_cold_cache_label_run_stages_every_upload_non_blocking() -> None: + # given - a fresh block: sprite cache, flat templates and ring all cold + block = LabelVisualizationBlockV1() + detections = _mirrored_od_detections() + image = _tensor_image() + + # when + audit = _TransferAudit() + with audit: + result = block.run(image=image, predictions=detections, **_LABEL_RUN_KWARGS) + + # then - even the cache-miss payloads (sprite colors + flat-index + # templates) and the first table upload are staged non-blocking + _assert_no_forbidden_ops(audit.ops) + _assert_no_device_to_host(audit.copies) + _assert_uploads_non_blocking(audit.copies) + cold_copies = [ + record for record in audit.copies if record["op"] == "aten.copy_.default" + ] + # one colors payload + one flat-index template per distinct sprite, plus + # the packed table + assert len(cold_copies) == 2 * N_BOXES + 1, f"unexpected copies: {cold_copies}" + assert all(record["non_blocking"] for record in cold_copies), cold_copies + assert result["image"]._tensor_image is not None + + +# --------------------------------------------------------------------------- # +# micro-check: warm paste dispatch budget is flat in the label count +# --------------------------------------------------------------------------- # + + +def test_warm_label_paste_dispatch_count_is_flat_in_label_count() -> None: + # given / when + counts = {} + traces = {} + for n in (2, 3, 4): + block = LabelVisualizationBlockV1() + detections = _disjoint_od_detections(n) + block.run( + image=_tensor_image(), predictions=detections, **_LABEL_RUN_KWARGS + ) # warm sprites, flat templates, ring slabs + audit = _TransferAudit() + audited_image = _tensor_image() + with audit: + block.run(image=audited_image, predictions=detections, **_LABEL_RUN_KWARGS) + counts[n] = len(audit.ops) + traces[n] = Counter(audit.ops) + + # then - flat in N, and no growth against the pre-ring baseline: the warm + # paste measured 15 dispatched ops before this change (its pageable packed + # upload was a from_numpy lift + `.to`); the ring swaps that for the one + # staged `copy_`, keeping the budget at 15 (clone, copy_, 3 slices, + # arange, repeat_interleave, index_select, 2 cats, index, add, t, view, + # index_put_). + assert counts[2] == counts[3] == counts[4], counts + assert counts[4] <= 15, traces[4] diff --git a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_block_assembler.py b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_block_assembler.py index 0232998298..34af218327 100644 --- a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_block_assembler.py +++ b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_block_assembler.py @@ -32,6 +32,7 @@ ManifestDescription, PythonCode, SelectorType, + TensorCompatibility, ValueType, ) @@ -502,3 +503,179 @@ def test_create_dynamic_block_specification() -> None: _ = result.manifest_class.model_validate( {"name": "some", "type": "MyBlock", "a": "$steps.some.a", "b": 1} ) # error expected - value "b" not a list + + +def _tensor_knob_block_definition(tensor_compatibility: str) -> DynamicBlockDefinition: + return DynamicBlockDefinition( + type="DynamicBlockDefinition", + manifest=ManifestDescription( + type="ManifestDescription", + block_type="MyTensorKnobBlock", + inputs={ + "a": DynamicInputDefinition( + type="DynamicInputDefinition", + selector_types=[SelectorType.STEP_OUTPUT], + ), + }, + outputs={"output": DynamicOutputDefinition(type="DynamicOutputDefinition")}, + tensor_compatibility=tensor_compatibility, + ), + code=PythonCode( + type="PythonCode", + run_function_code=PYTHON_CODE, + ), + ) + + +def test_manifest_description_tensor_compatibility_defaults_to_legacy() -> None: + # when - definition WITHOUT the field, as every pre-knob definition JSON + manifest = ManifestDescription.model_validate( + { + "type": "ManifestDescription", + "block_type": "MyBlock", + "inputs": {}, + } + ) + + # then + assert ( + manifest.tensor_compatibility is TensorCompatibility.LEGACY_COMPATIBILITY + ), "Expected additive field to default to legacy_compatibility for BC" + + +def test_manifest_description_tensor_compatibility_rejects_unknown_value() -> None: + # when + with pytest.raises(ValidationError): + _ = ManifestDescription.model_validate( + { + "type": "ManifestDescription", + "block_type": "MyBlock", + "inputs": {}, + "tensor_compatibility": "quantum_mode", + } + ) + + +@mock.patch.object(block_assembler, "ENABLE_TENSOR_DATA_REPRESENTATION", False) +def test_create_dynamic_block_specification_when_tensor_native_declared_and_flag_off() -> ( + None +): + # given + definition = _tensor_knob_block_definition(tensor_compatibility="tensor_native") + + # when + with pytest.raises(DynamicBlockError) as error: + _ = create_dynamic_block_specification( + dynamic_block_definition=definition, + kinds_lookup={"*": WILDCARD_KIND}, + ) + + # then + assert "numpy data representation" in str(error.value) + + +@mock.patch.object(block_assembler, "ENABLE_TENSOR_DATA_REPRESENTATION", True) +@mock.patch.object(block_assembler, "WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE", "modal") +def test_create_dynamic_block_specification_when_tensor_native_declared_and_modal_mode() -> ( + None +): + # given + definition = _tensor_knob_block_definition(tensor_compatibility="tensor_native") + + # when + with pytest.raises(DynamicBlockError) as error: + _ = create_dynamic_block_specification( + dynamic_block_definition=definition, + kinds_lookup={"*": WILDCARD_KIND}, + ) + + # then + assert "not yet supported for remote" in str(error.value) + + +@mock.patch.object(block_assembler, "ENABLE_TENSOR_DATA_REPRESENTATION", True) +@mock.patch.object(block_assembler, "WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE", "local") +def test_create_dynamic_block_specification_when_tensor_native_declared_flag_on_local() -> ( + None +): + # given + definition = _tensor_knob_block_definition(tensor_compatibility="tensor_native") + + # when + result = create_dynamic_block_specification( + dynamic_block_definition=definition, + kinds_lookup={"*": WILDCARD_KIND}, + ) + + # then - compiles, and the raw manifest description is threaded onto the block class + assert ( + result.block_class._manifest_description.tensor_compatibility + is TensorCompatibility.TENSOR_NATIVE + ) + + +@pytest.mark.parametrize("flag_value", [True, False]) +def test_create_dynamic_block_specification_default_knob_compiles_regardless_of_flag( + flag_value: bool, +) -> None: + # given - no tensor_compatibility declared anywhere (the pre-knob definition shape) + definition = DynamicBlockDefinition( + type="DynamicBlockDefinition", + manifest=ManifestDescription( + type="ManifestDescription", + block_type="MyLegacyBlock", + inputs={ + "a": DynamicInputDefinition( + type="DynamicInputDefinition", + selector_types=[SelectorType.STEP_OUTPUT], + ), + }, + ), + code=PythonCode(type="PythonCode", run_function_code=PYTHON_CODE), + ) + + # when + with mock.patch.object( + block_assembler, "ENABLE_TENSOR_DATA_REPRESENTATION", flag_value + ): + result = create_dynamic_block_specification( + dynamic_block_definition=definition, + kinds_lookup={"*": WILDCARD_KIND}, + ) + + # then + assert ( + result.block_class._manifest_description.tensor_compatibility + is TensorCompatibility.LEGACY_COMPATIBILITY + ) + + +@pytest.mark.parametrize("flag_value", [True, False]) +def test_create_dynamic_block_specification_default_knob_allowed_in_modal_mode( + flag_value: bool, +) -> None: + # given - D5 contract lock: legacy_compatibility (the default) over modal + # execution stays allowed regardless of the tensor flag; only tensor_native + # is blocked on the modal path. + definition = _tensor_knob_block_definition( + tensor_compatibility="legacy_compatibility" + ) + + # when - skip_class_eval short-circuits dynamic-module creation BEFORE the + # modal remote-validation branch, so no remote call is attempted here. + with mock.patch.object( + block_assembler, "ENABLE_TENSOR_DATA_REPRESENTATION", flag_value + ), mock.patch.object( + block_assembler, "WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE", "modal" + ): + result = create_dynamic_block_specification( + dynamic_block_definition=definition, + kinds_lookup={"*": WILDCARD_KIND}, + skip_class_eval=True, + ) + + # then + assert ( + result.block_class._manifest_description.tensor_compatibility + is TensorCompatibility.LEGACY_COMPATIBILITY + ) diff --git a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_block_scaffolding.py b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_block_scaffolding.py index 2d7992a555..c1532e679a 100644 --- a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_block_scaffolding.py +++ b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_block_scaffolding.py @@ -37,6 +37,18 @@ def _clear_modal_executor_cache() -> None: block_scaffolding._MODAL_EXECUTOR_CACHE.clear() +@pytest.fixture +def isolated_modal_executor_cache(): + # The modal arm reuses executors through the module-level + # _MODAL_EXECUTOR_CACHE, so a warm entry left by an earlier test starves + # the current test's patched ModalExecutor of the construction call its + # assertions depend on - and the current test's mock would leak to later + # ones. Clear on both sides, like the cache-lifecycle tests below do. + _clear_modal_executor_cache() + yield + _clear_modal_executor_cache() + + def test_create_dynamic_module_when_syntax_error_happens() -> None: # given init_function = """ @@ -465,6 +477,305 @@ def run_function(self, a, b) -> BlockResult: ] +def _boundary_wiring_manifest_description(): + from inference.core.workflows.execution_engine.v1.dynamic_blocks.entities import ( + DynamicInputDefinition, + DynamicOutputDefinition, + ManifestDescription, + SelectorType, + ) + + return ManifestDescription( + type="ManifestDescription", + block_type="BoundaryProbe", + inputs={ + "predictions": DynamicInputDefinition( + type="DynamicInputDefinition", + selector_types=[SelectorType.STEP_OUTPUT], + selector_data_kind={ + SelectorType.STEP_OUTPUT: ["object_detection_prediction"] + }, + ), + }, + outputs={ + "result": DynamicOutputDefinition( + type="DynamicOutputDefinition", kind=["object_detection_prediction"] + ), + "was_sv": DynamicOutputDefinition(type="DynamicOutputDefinition"), + }, + ) + + +def _native_od_fixture(): + import torch + + from inference_models.models.base.object_detection import Detections + + return Detections( + xyxy=torch.tensor([[10.0, 10.0, 20.0, 20.0]]), + class_id=torch.tensor([0]), + confidence=torch.tensor([0.9]), + image_metadata={"class_names": {0: "dog"}}, + bboxes_metadata=[{"detection_id": "det-1"}], + ) + + +def test_run_wrapper_applies_representation_boundary_both_ways() -> None: + # given + import supervision as sv + + from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( + representation_boundary, + ) + from inference_models.models.base.object_detection import Detections + + run_function = """ +def run_function(self, predictions) -> BlockResult: + return {"result": predictions, "was_sv": isinstance(predictions, sv.Detections)} +""" + python_code = PythonCode( + type="PythonCode", + run_function_code=run_function, + run_function_name="run_function", + ) + block_class = assembly_custom_python_block( + block_type_name="BoundaryProbe", + unique_identifier="boundary-probe-id", + manifest=BlockManifest, + python_code=python_code, + manifest_description=_boundary_wiring_manifest_description(), + ) + + # when + with mock.patch.object( + representation_boundary, "_TENSOR_REPRESENTATION_ACTIVE", True + ): + result = block_class().run(predictions=_native_od_fixture()) + + # then - IN boundary delivered sv to user code, OUT boundary restored native + assert result["was_sv"] is True, "User code must have received sv.Detections" + assert isinstance( + result["result"], Detections + ), "Declared-kind output must be converted back to native" + assert result["result"].bboxes_metadata[0]["detection_id"] == "det-1" + + +def test_run_wrapper_boundary_error_not_wrapped_as_user_code_error() -> None: + # given - user code succeeds, but returns a value that cannot satisfy the + # declared output kind; the boundary error must surface as itself, NOT as a + # DynamicBlockCodeError blaming the user's code. + from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( + representation_boundary, + ) + from inference.core.workflows.execution_engine.v1.dynamic_blocks.representation_boundary import ( + RepresentationBoundaryError, + ) + + run_function = """ +def run_function(self, predictions) -> BlockResult: + return {"result": object(), "was_sv": True} +""" + python_code = PythonCode( + type="PythonCode", + run_function_code=run_function, + run_function_name="run_function", + ) + block_class = assembly_custom_python_block( + block_type_name="BoundaryProbe", + unique_identifier="boundary-probe-err-id", + manifest=BlockManifest, + python_code=python_code, + manifest_description=_boundary_wiring_manifest_description(), + ) + + # when + with mock.patch.object( + representation_boundary, "_TENSOR_REPRESENTATION_ACTIVE", True + ), pytest.raises(RepresentationBoundaryError) as error: + _ = block_class().run(predictions=_native_od_fixture()) + + # then + assert not isinstance(error.value, DynamicBlockCodeError) + assert "BoundaryProbe" in str(error.value) + assert "result" in str(error.value) + + +def _legacy_sv_fixture(): + import numpy as np + import supervision as sv + + return sv.Detections( + xyxy=np.array([[10.0, 10.0, 20.0, 20.0]], dtype=np.float32), + class_id=np.array([0]), + confidence=np.array([0.9], dtype=np.float32), + data={ + "class_name": np.array(["dog"]), + "detection_id": np.array(["det-remote-1"]), + }, + ) + + +def test_run_wrapper_modal_arm_converts_kwargs_and_remote_result( + isolated_modal_executor_cache, +) -> None: + # given - D2-REVISED Option A: the Modal arm must ship CONVERTED (sv) inputs + # to the executor and convert the sv result coming back into native objects. + from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( + block_scaffolding, + modal_executor, + representation_boundary, + ) + from inference_models.models.base.object_detection import Detections + + python_code = PythonCode( + type="PythonCode", + run_function_code="def run_function(self, predictions) -> BlockResult:\n return None\n", + run_function_name="run_function", + ) + block_class = assembly_custom_python_block( + block_type_name="BoundaryProbe", + unique_identifier="boundary-probe-modal-id", + manifest=BlockManifest, + python_code=python_code, + manifest_description=_boundary_wiring_manifest_description(), + ) + executor_instance = mock.MagicMock() + executor_instance.execute_remote.return_value = { + "result": _legacy_sv_fixture(), + "was_sv": True, + } + + # when + with mock.patch.object( + representation_boundary, "_TENSOR_REPRESENTATION_ACTIVE", True + ), mock.patch.object( + block_scaffolding, "WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE", "modal" + ), mock.patch.object( + block_scaffolding, "get_roboflow_workspace", return_value="test-workspace" + ), mock.patch.object( + modal_executor, "ModalExecutor", return_value=executor_instance + ): + result = block_class().run(predictions=_native_od_fixture()) + + # then - inputs leg: executor received sv (the `_type='sv_detections'` arm + # territory), not native objects + import supervision as sv + + sent_inputs = executor_instance.execute_remote.call_args.kwargs["inputs"] + assert isinstance( + sent_inputs["predictions"], sv.Detections + ), "Modal arm must ship converted sv inputs" + # then - return leg: sv result converted to native before entering the engine + assert isinstance( + result["result"], Detections + ), "Modal arm must convert the sv result back to native" + assert result["result"].bboxes_metadata[0]["detection_id"] == "det-remote-1" + + +def test_run_wrapper_modal_arm_is_passthrough_when_flag_off( + isolated_modal_executor_cache, +) -> None: + # given + from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( + block_scaffolding, + modal_executor, + representation_boundary, + ) + + python_code = PythonCode( + type="PythonCode", + run_function_code="def run_function(self, predictions) -> BlockResult:\n return None\n", + run_function_name="run_function", + ) + block_class = assembly_custom_python_block( + block_type_name="BoundaryProbe", + unique_identifier="boundary-probe-modal-off-id", + manifest=BlockManifest, + python_code=python_code, + manifest_description=_boundary_wiring_manifest_description(), + ) + legacy_input = _legacy_sv_fixture() + remote_result = {"result": _legacy_sv_fixture(), "was_sv": True} + executor_instance = mock.MagicMock() + executor_instance.execute_remote.return_value = remote_result + + # when + with mock.patch.object( + representation_boundary, "_TENSOR_REPRESENTATION_ACTIVE", False + ), mock.patch.object( + block_scaffolding, "WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE", "modal" + ), mock.patch.object( + block_scaffolding, "get_roboflow_workspace", return_value="test-workspace" + ), mock.patch.object( + modal_executor, "ModalExecutor", return_value=executor_instance + ): + result = block_class().run(predictions=legacy_input) + + # then - flag-off byte-parity: both legs are is-identity + sent_inputs = executor_instance.execute_remote.call_args.kwargs["inputs"] + assert sent_inputs["predictions"] is legacy_input + assert result is remote_result + + +def test_run_wrapper_local_forbidden_gate_wins_over_conversion_error() -> None: + # given - local mode, custom python forbidden, and kwargs that WOULD raise a + # RepresentationBoundaryError at IN conversion (wrong type for the declared + # kind). The clearer misconfiguration error must fire first. + import torch + + from inference.core.workflows.errors import WorkflowEnvironmentConfigurationError + from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( + block_scaffolding, + representation_boundary, + ) + + python_code = PythonCode( + type="PythonCode", + run_function_code="def run_function(self, predictions) -> BlockResult:\n return None\n", + run_function_name="run_function", + ) + block_class = assembly_custom_python_block( + block_type_name="BoundaryProbe", + unique_identifier="boundary-probe-gate-id", + manifest=BlockManifest, + python_code=python_code, + manifest_description=_boundary_wiring_manifest_description(), + ) + + # when + with mock.patch.object( + representation_boundary, "_TENSOR_REPRESENTATION_ACTIVE", True + ), mock.patch.object( + block_scaffolding, "WORKFLOWS_CUSTOM_PYTHON_EXECUTION_MODE", "local" + ), mock.patch.object( + block_scaffolding, "ALLOW_CUSTOM_PYTHON_EXECUTION_IN_WORKFLOWS", False + ), pytest.raises( + WorkflowEnvironmentConfigurationError + ): + _ = block_class().run(predictions=torch.tensor([1.0])) + + +def test_imports_lines_tensor_native_extension_tracks_the_flag() -> None: + # given - the extension is resolved at import time (load-time swap philosophy): + # flag-on ships the tensor-native authoring imports, flag-off must keep the + # generated module source byte-identical to the legacy list. This test runs in + # both CI directions, locking each side of the contract. + from inference.core.env import ENABLE_TENSOR_DATA_REPRESENTATION + from inference.core.workflows.execution_engine.v1.dynamic_blocks.block_scaffolding import ( + IMPORTS_LINES, + TENSOR_NATIVE_IMPORTS_LINES, + ) + + # then + assert len(TENSOR_NATIVE_IMPORTS_LINES) > 0 + for line in TENSOR_NATIVE_IMPORTS_LINES: + assert (line in IMPORTS_LINES) == ENABLE_TENSOR_DATA_REPRESENTATION + # the legacy prefix is untouched in both directions + assert IMPORTS_LINES[0] == "from typing import Any, List, Dict, Set, Optional" + assert "import supervision as sv" in IMPORTS_LINES + assert "import numpy as np" in IMPORTS_LINES + + def test_modal_executor_cache_closes_idle_entries(monkeypatch) -> None: from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( modal_executor, diff --git a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_executor.py b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_executor.py new file mode 100644 index 0000000000..6a06d33e3c --- /dev/null +++ b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_modal_executor.py @@ -0,0 +1,148 @@ +"""Tests for the Modal wire serializer's tensor-pivot defense-in-depth +(Step 7 of the tensor_compatibility plan): native ``inference_models`` objects +must never be silently stringified on the wire โ€” flag-on they raise, flag-off +the pre-existing generic-object contract stays byte-identical.""" + +import json +from unittest import mock + +import numpy as np +import pytest +import supervision as sv +import torch + +from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( + representation_boundary, +) +from inference.core.workflows.execution_engine.v1.dynamic_blocks.modal_executor import ( + serialize_for_modal_remote_execution, +) +from inference.core.workflows.execution_engine.v1.dynamic_blocks.representation_boundary import ( + RepresentationBoundaryError, +) +from inference_models.models.base.object_detection import Detections + +_flag_on = mock.patch.object( + representation_boundary, "_TENSOR_REPRESENTATION_ACTIVE", True +) +_flag_off = mock.patch.object( + representation_boundary, "_TENSOR_REPRESENTATION_ACTIVE", False +) + + +def _native_detections() -> Detections: + return Detections( + xyxy=torch.tensor([[10.0, 10.0, 60.0, 60.0]]), + class_id=torch.tensor([0]), + confidence=torch.tensor([0.9]), + ) + + +class _ArbitraryObject: + def __init__(self): + self.field = "value" + + def __str__(self): + return "arbitrary-object" + + +def test_modal_serializer_raises_on_unconverted_native_detections_when_flag_on() -> ( + None +): + # when + with _flag_on, pytest.raises(RepresentationBoundaryError) as error: + _ = serialize_for_modal_remote_execution( + inputs={"predictions": _native_detections()} + ) + + # then + assert "Detections" in str(error.value) + assert "reached the Modal wire serializer unconverted" in str(error.value) + + +def test_modal_serializer_raises_on_bare_tensor_when_flag_on() -> None: + # when + with _flag_on, pytest.raises(RepresentationBoundaryError) as error: + _ = serialize_for_modal_remote_execution( + inputs={"value": torch.tensor([1.0, 2.0])} + ) + + # then + assert "Tensor" in str(error.value) + + +def test_modal_serializer_raises_on_native_nested_in_containers_when_flag_on() -> None: + # given - natives hide inside dicts/lists; the JSON encoder is the choke + # point, so nesting must not smuggle them past the guard + inputs = {"payload": {"items": [_native_detections()]}} + + # when + with _flag_on, pytest.raises(RepresentationBoundaryError): + _ = serialize_for_modal_remote_execution(inputs=inputs) + + +def test_modal_serializer_keeps_generic_fallback_for_arbitrary_objects_when_flag_on() -> ( + None +): + # when + with _flag_on: + payload = json.loads( + serialize_for_modal_remote_execution(inputs={"value": _ArbitraryObject()}) + ) + + # then - pre-existing generic contract untouched for non-native objects + assert payload["value"] == { + "_type": "object", + "class": "_ArbitraryObject", + "value": "arbitrary-object", + } + + +def test_modal_serializer_flag_off_behavior_is_byte_identical_to_legacy() -> None: + # given - flag-off natives cannot exist in real flows, but the guard must be + # provably inert: everything with __dict__ stringifies exactly as before + inputs = { + "native": _native_detections(), + "tensor": torch.tensor([1.0]), + "object": _ArbitraryObject(), + } + + # when + with _flag_off: + payload = json.loads(serialize_for_modal_remote_execution(inputs=inputs)) + + # then + assert payload["native"]["_type"] == "object" + assert payload["native"]["class"] == "Detections" + assert payload["tensor"]["_type"] == "object" + assert payload["tensor"]["class"] == "Tensor" + assert payload["object"] == { + "_type": "object", + "class": "_ArbitraryObject", + "value": "arbitrary-object", + } + + +def test_modal_serializer_sv_detections_ride_dedicated_arm_when_flag_on() -> None: + # given - the Option-A input leg: boundary-converted sv rides the dedicated + # wire arm, never the generic fallback + detections = sv.Detections( + xyxy=np.array([[10.0, 10.0, 60.0, 60.0]], dtype=np.float32), + class_id=np.array([0]), + confidence=np.array([0.9], dtype=np.float32), + data={ + "class_name": np.array(["widget"]), + "detection_id": np.array(["det-1"]), + "image_dimensions": np.array([[480, 640]]), + }, + ) + + # when + with _flag_on: + payload = json.loads( + serialize_for_modal_remote_execution(inputs={"predictions": detections}) + ) + + # then + assert payload["predictions"]["_type"] == "sv_detections" + assert payload["predictions"]["predictions"][0]["class"] == "widget" diff --git a/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_representation_boundary.py b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_representation_boundary.py new file mode 100644 index 0000000000..e0f0b702c5 --- /dev/null +++ b/tests/workflows/unit_tests/execution_engine/dynamic_blocs/test_representation_boundary.py @@ -0,0 +1,1307 @@ +"""Unit tests for the dynamic-block representation boundary (Step 2 - the IN +direction: native -> legacy). Pure converter/walker tests; engine wiring is +Step 4. + +The boundary resolves ``ENABLE_TENSOR_DATA_REPRESENTATION`` once at import time +into ``representation_boundary._TENSOR_REPRESENTATION_ACTIVE``; tests patch that +module constant (the same technique block_assembler tests use for the env flag) +so both CI flag directions exercise every case deterministically. +""" + +from typing import Optional +from unittest import mock +from uuid import UUID + +import numpy as np +import pytest +import supervision as sv + +torch = pytest.importorskip("torch") +pytest.importorskip("inference_models") + +from pycocotools import mask as mask_utils + +from inference.core.workflows.core_steps.common.utils import ( + sv_detections_to_root_coordinates, +) +from inference.core.workflows.execution_engine.constants import ( + CLASS_NAME_KEY, + CLASS_NAMES_KEY, + DETECTION_ID_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS, + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS, + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + PARENT_ID_KEY, + POLYGON_KEY_IN_SV_DETECTIONS, + PREDICTION_TYPE_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, + TRACKER_ID_KEY, +) +from inference.core.workflows.execution_engine.entities.base import ( + Batch, + ImageParentMetadata, + WorkflowImageData, +) +from inference.core.workflows.execution_engine.v1.dynamic_blocks import ( + representation_boundary, +) +from inference.core.workflows.execution_engine.v1.dynamic_blocks.entities import ( + DynamicInputDefinition, + ManifestDescription, + SelectorType, +) +from inference.core.workflows.execution_engine.v1.dynamic_blocks.representation_boundary import ( + RepresentationBoundaryError, + convert_kwargs_to_legacy, + native_detections_to_sv, +) +from inference_models.models.base.classification import ( + ClassificationPrediction, + MultiLabelClassificationPrediction, +) +from inference_models.models.base.instance_segmentation import InstanceDetections +from inference_models.models.base.keypoints_detection import KeyPoints +from inference_models.models.base.object_detection import Detections +from inference_models.models.base.types import InstancesRLEMasks + +CROP_OFFSET_X, CROP_OFFSET_Y = 100, 50 +CROP_H, CROP_W = 40, 60 +ROOT_H, ROOT_W = 480, 640 + +_boundary_on = mock.patch.object( + representation_boundary, "_TENSOR_REPRESENTATION_ACTIVE", True +) +_boundary_off = mock.patch.object( + representation_boundary, "_TENSOR_REPRESENTATION_ACTIVE", False +) + + +def _crop_local_dense_masks() -> np.ndarray: + masks = np.zeros((2, CROP_H, CROP_W), dtype=bool) + masks[0, 5:15, 10:30] = True + masks[1, 20:35, 40:55] = True + return masks + + +def _native_image_metadata() -> dict: + return { + CLASS_NAMES_KEY: {0: "a", 1: "b"}, + PREDICTION_TYPE_KEY: "instance-segmentation", + IMAGE_DIMENSIONS_KEY: [CROP_H, CROP_W], + INFERENCE_ID_KEY: "iid-1", + PARENT_ID_KEY: "crop-1", + PARENT_COORDINATES_KEY: [CROP_OFFSET_X, CROP_OFFSET_Y], + PARENT_DIMENSIONS_KEY: [ROOT_H, ROOT_W], + ROOT_PARENT_ID_KEY: "root-image", + ROOT_PARENT_COORDINATES_KEY: [CROP_OFFSET_X, CROP_OFFSET_Y], + ROOT_PARENT_DIMENSIONS_KEY: [ROOT_H, ROOT_W], + } + + +def _native_instance_detections( + mask: object, + with_polygons: bool = False, +) -> InstanceDetections: + bboxes_metadata = [ + {DETECTION_ID_KEY: "d0"}, + {DETECTION_ID_KEY: "d1"}, + ] + if with_polygons: + bboxes_metadata[0][POLYGON_KEY_IN_SV_DETECTIONS] = np.array( + [[10, 5], [30, 5], [30, 15], [10, 15]], dtype=np.int64 + ) + bboxes_metadata[1][POLYGON_KEY_IN_SV_DETECTIONS] = np.array( + [[40, 20], [55, 20], [55, 35], [40, 35]], dtype=np.int64 + ) + return InstanceDetections( + xyxy=torch.tensor( + [[10.0, 5.0, 30.0, 15.0], [40.0, 20.0, 55.0, 35.0]], dtype=torch.float32 + ), + class_id=torch.tensor([0, 1], dtype=torch.long), + confidence=torch.tensor([0.5, 0.25], dtype=torch.float32), + mask=mask, + image_metadata=_native_image_metadata(), + bboxes_metadata=bboxes_metadata, + ) + + +def _native_object_detections() -> Detections: + return Detections( + xyxy=torch.tensor([[10.0, 5.0, 30.0, 15.0]], dtype=torch.float32), + class_id=torch.tensor([1], dtype=torch.long), + confidence=torch.tensor([0.75], dtype=torch.float32), + image_metadata=_native_image_metadata(), + bboxes_metadata=[{DETECTION_ID_KEY: "d0"}], + ) + + +def _dense_masks_to_rle(masks: np.ndarray) -> InstancesRLEMasks: + encoded = [ + mask_utils.encode(np.asfortranarray(mask.astype(np.uint8)))["counts"] + for mask in masks + ] + return InstancesRLEMasks(image_size=(CROP_H, CROP_W), masks=encoded) + + +def _manifest(inputs: dict) -> ManifestDescription: + return ManifestDescription( + type="ManifestDescription", + block_type="MyBlock", + inputs=inputs, + ) + + +def _input_with_kinds(*kind_names: str) -> DynamicInputDefinition: + return DynamicInputDefinition( + type="DynamicInputDefinition", + selector_types=[SelectorType.STEP_OUTPUT], + selector_data_kind={SelectorType.STEP_OUTPUT: list(kind_names)}, + ) + + +def test_native_detections_to_sv_carries_full_lineage_key_set() -> None: + # given + native = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + + # when + converted = native_detections_to_sv(detections=native) + + # then - key-set superset of what attach_parents_coordinates_to_sv_detections + # attaches, plus prediction_type / image_dimensions / inference_id + expected_keys = { + "class_name", + DETECTION_ID_KEY, + PARENT_ID_KEY, + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + ROOT_PARENT_ID_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + PREDICTION_TYPE_KEY, + IMAGE_DIMENSIONS_KEY, + INFERENCE_ID_KEY, + } + assert expected_keys.issubset(set(converted.data.keys())) + # shapes/values match the numpy convention exactly: (n, 2) [x, y] coordinates, + # (n, 2) [h, w] dimensions, broadcast string ids + assert np.array_equal( + converted.data[PARENT_COORDINATES_KEY], + np.array([[CROP_OFFSET_X, CROP_OFFSET_Y]] * 2), + ) + assert np.array_equal( + converted.data[ROOT_PARENT_DIMENSIONS_KEY], np.array([[ROOT_H, ROOT_W]] * 2) + ) + assert converted.data[PARENT_ID_KEY].tolist() == ["crop-1", "crop-1"] + assert converted.data[ROOT_PARENT_ID_KEY].tolist() == ["root-image", "root-image"] + assert converted.data[PREDICTION_TYPE_KEY].tolist() == ["instance-segmentation"] * 2 + assert np.array_equal( + converted.data[IMAGE_DIMENSIONS_KEY], np.array([[CROP_H, CROP_W]] * 2) + ) + assert converted.data[INFERENCE_ID_KEY].tolist() == ["iid-1", "iid-1"] + assert np.array_equal( + converted.xyxy, + np.array([[10.0, 5.0, 30.0, 15.0], [40.0, 20.0, 55.0, 35.0]], np.float32), + ) + assert converted.mask.dtype == bool and converted.mask.shape == (2, CROP_H, CROP_W) + + +def test_native_detections_to_sv_output_survives_numpy_root_coordinates_shift() -> None: + # given - crop-local native prediction with dense masks and per-box polygons + dense = _crop_local_dense_masks() + native = _native_instance_detections( + mask=torch.from_numpy(dense), with_polygons=True + ) + converted = native_detections_to_sv(detections=native) + + # when - the NUMPY root-coordinates conversion runs on the converted object + shifted = sv_detections_to_root_coordinates(detections=converted) + + # then - boxes, masks and polygon payloads land in root coordinates + assert np.array_equal( + shifted.xyxy, + np.array( + [ + [110.0, 55.0, 130.0, 65.0], + [140.0, 70.0, 155.0, 85.0], + ], + dtype=np.float32, + ), + ) + expected_masks = np.zeros((2, ROOT_H, ROOT_W), dtype=bool) + expected_masks[ + :, + CROP_OFFSET_Y : CROP_OFFSET_Y + CROP_H, + CROP_OFFSET_X : CROP_OFFSET_X + CROP_W, + ] = dense + assert np.array_equal(shifted.mask, expected_masks) + assert np.array_equal( + np.asarray(shifted.data[POLYGON_KEY_IN_SV_DETECTIONS][0]), + np.array([[110, 55], [130, 55], [130, 65], [110, 65]]), + ) + + +def test_native_detections_to_sv_resolves_class_name_override_first() -> None: + # given - classes_replacement-style per-box override on row 0 only + native = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + native.bboxes_metadata[0][CLASS_NAME_KEY] = "replaced-label" + + # when + converted = native_detections_to_sv(detections=native) + + # then - override wins on row 0, map resolves row 1; no shadow 'class' column + assert converted.data["class_name"].tolist() == ["replaced-label", "b"] + assert CLASS_NAME_KEY not in converted.data + + +def test_native_detections_to_sv_mints_uuid_detection_ids_when_missing() -> None: + # given + native = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + del native.bboxes_metadata[1][DETECTION_ID_KEY] + + # when + converted = native_detections_to_sv(detections=native) + + # then + detection_ids = converted.data[DETECTION_ID_KEY].tolist() + assert detection_ids[0] == "d0" + assert detection_ids[1] != "" and UUID(detection_ids[1]) + + +def test_native_detections_to_sv_rle_and_dense_carriers_materialise_identically() -> ( + None +): + # given + dense = _crop_local_dense_masks() + via_dense = _native_instance_detections(mask=torch.from_numpy(dense)) + via_rle = _native_instance_detections(mask=_dense_masks_to_rle(dense)) + + # when + converted_dense = native_detections_to_sv(detections=via_dense) + converted_rle = native_detections_to_sv(detections=via_rle) + + # then - both carriers land as the SAME dense (N, H, W) bool ndarray + assert converted_dense.mask.dtype == bool and converted_rle.mask.dtype == bool + assert np.array_equal(converted_dense.mask, converted_rle.mask) + assert np.array_equal(converted_dense.mask, dense) + + +def test_native_detections_to_sv_pads_keypoint_payloads_like_numpy() -> None: + # given - 2 keypoints on row 0, none on row 1 (ragged input) + native = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + native.bboxes_metadata[0].update( + { + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS: [[11.0, 6.0], [29.0, 14.0]], + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS: [0.875, 0.75], + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS: [0, 1], + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS: ["nose", "tail"], + } + ) + native.bboxes_metadata[1].update( + { + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS: [], + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS: [], + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS: [], + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS: [], + } + ) + + # when + converted = native_detections_to_sv(detections=native) + + # then - proper padded N-d arrays, the add_inference_keypoints_to_sv_detections + # convention (ragged object arrays break supervision's is_data_equal) + assert converted.data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS].shape == (2, 2, 2) + assert converted.data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS].dtype == np.float32 + assert np.array_equal( + converted.data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS][0], + np.array([[11.0, 6.0], [29.0, 14.0]], dtype=np.float32), + ) + assert np.array_equal( + converted.data[KEYPOINTS_XY_KEY_IN_SV_DETECTIONS][1], + np.zeros((2, 2), dtype=np.float32), + ) + assert converted.data[KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS][1].tolist() == [ + "", + "", + ] + + +def test_native_detections_to_sv_tracker_ids() -> None: + # given + native = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + native.bboxes_metadata[0][TRACKER_ID_KEY] = 7 + native.bboxes_metadata[1][TRACKER_ID_KEY] = 11 + + # when + converted = native_detections_to_sv(detections=native) + + # then - tracker ids populate the sv field, not a data column + assert converted.tracker_id.tolist() == [7, 11] + assert TRACKER_ID_KEY not in converted.data + + +def test_native_detections_to_sv_handles_plain_and_empty_detections() -> None: + # given + plain = _native_object_detections() + empty = Detections( + xyxy=torch.zeros((0, 4), dtype=torch.float32), + class_id=torch.zeros((0,), dtype=torch.long), + confidence=torch.zeros((0,), dtype=torch.float32), + image_metadata=_native_image_metadata(), + bboxes_metadata=[], + ) + + # when + converted_plain = native_detections_to_sv(detections=plain) + converted_empty = native_detections_to_sv(detections=empty) + + # then + assert converted_plain.mask is None + assert len(converted_plain) == 1 + assert len(converted_empty) == 0 + + +def test_convert_kwargs_is_identity_when_tensor_representation_off() -> None: + # given + kwargs = {"predictions": _native_object_detections(), "threshold": 0.5} + manifest = _manifest(inputs={}) + + # when + with _boundary_off: + result = convert_kwargs_to_legacy( + kwargs=kwargs, manifest_description=manifest, block_name="block" + ) + + # then - the very same object, untouched (flag-off byte-parity guarantee) + assert result is kwargs + assert isinstance(result["predictions"], Detections) + + +def test_convert_kwargs_is_identity_for_tensor_native_mode() -> None: + # given + kwargs = {"predictions": _native_object_detections()} + manifest = ManifestDescription( + type="ManifestDescription", + block_type="MyBlock", + inputs={}, + tensor_compatibility="tensor_native", + ) + + # when + with _boundary_on: + result = convert_kwargs_to_legacy( + kwargs=kwargs, manifest_description=manifest, block_name="block" + ) + + # then + assert result is kwargs + assert isinstance(result["predictions"], Detections) + + +def test_convert_kwargs_declared_kind_converts_batch_preserving_indices() -> None: + # given + manifest = _manifest( + inputs={"predictions": _input_with_kinds("object_detection_prediction")} + ) + batch = Batch.init( + content=[_native_object_detections(), None, _native_object_detections()], + indices=[(0,), (1,), (2,)], + ) + + # when + with _boundary_on: + result = convert_kwargs_to_legacy( + kwargs={"predictions": batch}, + manifest_description=manifest, + block_name="block", + ) + + # then + converted = result["predictions"] + assert isinstance(converted, Batch) + assert converted.indices == [(0,), (1,), (2,)] + assert isinstance(converted[0], sv.Detections) + assert converted[1] is None + assert isinstance(converted[2], sv.Detections) + assert converted[0].data["class_name"].tolist() == ["b"] + + +def test_convert_kwargs_nested_batches_walked() -> None: + # given + manifest = _manifest( + inputs={"predictions": _input_with_kinds("object_detection_prediction")} + ) + inner_one = Batch.init(content=[_native_object_detections()], indices=[(0, 0)]) + inner_two = Batch.init(content=[None], indices=[(1, 0)]) + outer = Batch.init(content=[inner_one, inner_two], indices=[(0,), (1,)]) + + # when + with _boundary_on: + result = convert_kwargs_to_legacy( + kwargs={"predictions": outer}, + manifest_description=manifest, + block_name="block", + ) + + # then + outer_converted = result["predictions"] + assert outer_converted.indices == [(0,), (1,)] + assert outer_converted[0].indices == [(0, 0)] + assert isinstance(outer_converted[0][0], sv.Detections) + assert outer_converted[1][0] is None + + +def test_convert_kwargs_declared_kind_with_wrong_type_raises_loudly() -> None: + # given - declared classification receives a native detections object + manifest = _manifest( + inputs={"predictions": _input_with_kinds("classification_prediction")} + ) + + # when + with _boundary_on, pytest.raises(RepresentationBoundaryError) as error: + _ = convert_kwargs_to_legacy( + kwargs={"predictions": _native_object_detections()}, + manifest_description=manifest, + block_name="my_block", + ) + + # then + assert "my_block" in str(error.value) + assert "predictions" in str(error.value) + assert "Detections" in str(error.value) + + +def test_convert_kwargs_wildcard_sniffs_known_native_types() -> None: + # given + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="img"), + numpy_image=np.zeros((6, 8, 3), dtype=np.uint8), + ) + classification = ClassificationPrediction( + class_id=torch.tensor([1]), + confidence=torch.tensor([[0.25, 0.5]], dtype=torch.float32), + images_metadata=[ + { + CLASS_NAMES_KEY: {0: "cat", 1: "dog"}, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [480, 640], + } + ], + ) + key_points = KeyPoints( + xy=torch.tensor([[[11.0, 6.0]]], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.9]], dtype=torch.float32), + ) + kwargs = { + "detections": _native_object_detections(), + "keypoints": (key_points, _native_object_detections()), + "classes": classification, + "image": image, + "threshold": 0.4, + "label": "on", + "nested": {"inner": [_native_object_detections(), 3]}, + } + + # when + with _boundary_on: + result = convert_kwargs_to_legacy( + kwargs=kwargs, manifest_description=_manifest(inputs={}), block_name="b" + ) + + # then + assert isinstance(result["detections"], sv.Detections) + assert isinstance(result["keypoints"], sv.Detections) + assert result["classes"]["top"] == "dog" + assert result["image"] is image + assert result["threshold"] == 0.4 and result["label"] == "on" + assert isinstance(result["nested"]["inner"][0], sv.Detections) + assert result["nested"]["inner"][1] == 3 + + +def test_convert_kwargs_wildcard_bare_tensor_raises_actionable_error() -> None: + # when + with _boundary_on, pytest.raises(RepresentationBoundaryError) as error: + _ = convert_kwargs_to_legacy( + kwargs={"embedding": torch.zeros(4)}, + manifest_description=_manifest(inputs={}), + block_name="my_block", + ) + + # then - block, input, offending type and remediation all present + message = str(error.value) + assert "my_block" in message + assert "embedding" in message + assert "torch.Tensor" in message + assert "tensor_native" in message + + +def test_convert_kwargs_bare_keypoints_and_unknown_dataclass_raise() -> None: + # given + bare_key_points = KeyPoints( + xy=torch.tensor([[[1.0, 1.0]]], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.9]], dtype=torch.float32), + ) + unknown = InstancesRLEMasks(image_size=(4, 4), masks=[]) + + # when / then + with _boundary_on: + with pytest.raises(RepresentationBoundaryError): + _ = convert_kwargs_to_legacy( + kwargs={"kp": bare_key_points}, + manifest_description=_manifest(inputs={}), + block_name="b", + ) + with pytest.raises(RepresentationBoundaryError): + _ = convert_kwargs_to_legacy( + kwargs={"masks": unknown}, + manifest_description=_manifest(inputs={}), + block_name="b", + ) + + +def test_convert_kwargs_declared_embedding_and_tensor_kinds() -> None: + # given + manifest = _manifest( + inputs={ + "embedding": _input_with_kinds("embedding"), + "raw": _input_with_kinds("tensor"), + "static_embedding": _input_with_kinds("embedding"), + } + ) + kwargs = { + "embedding": torch.tensor([[1.0, 2.0], [3.0, 4.0]]), + "raw": torch.ones((2, 2)), + "static_embedding": [0.1, 0.2], + } + + # when + with _boundary_on: + result = convert_kwargs_to_legacy( + kwargs=kwargs, manifest_description=manifest, block_name="b" + ) + + # then - embedding flattens to List[float] (the numpy clip in-memory shape: + # `predictions.embeddings[0]`); `tensor` kind lands as ndarray; static values + # pass through untouched + assert result["embedding"] == [1.0, 2.0, 3.0, 4.0] + assert isinstance(result["raw"], np.ndarray) + assert np.array_equal(result["raw"], np.ones((2, 2))) + assert result["static_embedding"] == [0.1, 0.2] + + +def test_classification_conversion_matches_numpy_in_memory_dict() -> None: + # given - the exact in-memory dict the numpy classification block emits: + # response.model_dump(by_alias=True, exclude_none=True) + in-place + # prediction_type / parent_id / root_parent_id writes + from inference.core.entities.responses.inference import ( + ClassificationInferenceResponse, + ) + from inference.core.entities.responses.inference import ( + ClassificationPrediction as ResponseClassificationPrediction, + ) + from inference.core.entities.responses.inference import InferenceResponseImage + + response = ClassificationInferenceResponse( + image=InferenceResponseImage(width=640, height=480), + predictions=[ + ResponseClassificationPrediction( + **{"class": "dog", "class_id": 1, "confidence": 0.5} + ), + ResponseClassificationPrediction( + **{"class": "cat", "class_id": 0, "confidence": 0.25} + ), + ], + top="dog", + confidence=0.5, + time=0.0123, + inference_id="iid", + ) + numpy_in_memory = response.model_dump(by_alias=True, exclude_none=True) + numpy_in_memory[PREDICTION_TYPE_KEY] = "classification" + numpy_in_memory[PARENT_ID_KEY] = "p1" + numpy_in_memory[ROOT_PARENT_ID_KEY] = "r1" + + native = ClassificationPrediction( + class_id=torch.tensor([1]), + confidence=torch.tensor([[0.25, 0.5]], dtype=torch.float32), + images_metadata=[ + { + CLASS_NAMES_KEY: {0: "cat", 1: "dog"}, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [480, 640], + INFERENCE_ID_KEY: "iid", + PARENT_ID_KEY: "p1", + ROOT_PARENT_ID_KEY: "r1", + "time": 0.0123, + } + ], + ) + + # when + with _boundary_on: + result = convert_kwargs_to_legacy( + kwargs={"classes": native}, + manifest_description=_manifest( + inputs={"classes": _input_with_kinds("classification_prediction")} + ), + block_name="b", + ) + + # then + assert result["classes"] == numpy_in_memory + + +def test_convert_kwargs_multi_label_classification_sniffed() -> None: + # given + native = MultiLabelClassificationPrediction( + class_ids=torch.tensor([0, 1]), + confidence=torch.tensor([0.5, 0.25], dtype=torch.float32), + image_metadata={ + CLASS_NAMES_KEY: {0: "cat", 1: "dog"}, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [480, 640], + }, + ) + + # when + with _boundary_on: + result = convert_kwargs_to_legacy( + kwargs={"classes": native}, + manifest_description=_manifest(inputs={}), + block_name="b", + ) + + # then + assert result["classes"]["predicted_classes"] == ["cat", "dog"] + assert result["classes"]["predictions"]["cat"]["class_id"] == 0 + + +def test_convert_kwargs_wildcard_keypoint_tuple_without_bbox_component_raises() -> None: + # given - the native keypoint tuple is (KeyPoints, Optional[Detections]); a + # None bbox component has no legacy sv.Detections equivalent BY DESIGN. + key_points = KeyPoints( + xy=torch.tensor([[[1.0, 2.0]]]), + class_id=torch.tensor([0]), + confidence=torch.tensor([[0.9]]), + ) + + # when + with _boundary_on, pytest.raises(RepresentationBoundaryError) as error: + _ = convert_kwargs_to_legacy( + kwargs={"prediction": (key_points, None)}, + manifest_description=_manifest(inputs={}), + block_name="my_block", + ) + + # then + assert "missing the bounding-box component" in str(error.value) + assert "my_block" in str(error.value) + + +def test_convert_kwargs_empty_batch_passes_through_preserving_type() -> None: + # given + empty_batch = Batch.init(content=[], indices=[]) + + # when + with _boundary_on: + result = convert_kwargs_to_legacy( + kwargs={"predictions": empty_batch}, + manifest_description=_manifest( + inputs={"predictions": _input_with_kinds("object_detection_prediction")} + ), + block_name="my_block", + ) + + # then + converted = result["predictions"] + assert isinstance(converted, Batch) + assert len(converted) == 0 + assert list(converted.indices) == [] + + +# --------------------------------------------------------------------------- # +# Step 3: OUT direction (legacy -> native) + result walker + round trips # +# --------------------------------------------------------------------------- # + +from inference.core.workflows.core_steps.common.serializers_tensor import ( + serialise_native_classification, + serialise_sv_detections, +) +from inference.core.workflows.core_steps.common.tensor_native import ( + native_detections_to_root_coordinates, +) +from inference.core.workflows.execution_engine.v1.dynamic_blocks.entities import ( + DynamicOutputDefinition, +) +from inference.core.workflows.execution_engine.v1.dynamic_blocks.representation_boundary import ( + classification_dict_to_native, + convert_block_result_to_native, + sv_detections_to_native, + sv_detections_to_native_key_point_prediction, +) +from inference.core.workflows.execution_engine.v1.entities import FlowControl + + +def _manifest_with_outputs(outputs: dict) -> ManifestDescription: + return ManifestDescription( + type="ManifestDescription", + block_type="MyBlock", + inputs={}, + outputs=outputs, + ) + + +def _output_with_kinds(*kind_names: str) -> DynamicOutputDefinition: + return DynamicOutputDefinition( + type="DynamicOutputDefinition", kind=list(kind_names) + ) + + +def _assert_native_detections_close(left: "Detections", right: "Detections") -> None: + assert np.allclose(left.xyxy.cpu().numpy(), right.xyxy.cpu().numpy()) + assert left.class_id.cpu().tolist() == right.class_id.cpu().tolist() + assert np.allclose(left.confidence.cpu().numpy(), right.confidence.cpu().numpy()) + + +def test_sv_detections_to_native_round_trips_dense_instance_detections() -> None: + # given - the crop-lineage native fixture with dense masks + native = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + + # when - native -> legacy -> native + as_sv = native_detections_to_sv(detections=native) + round_tripped = sv_detections_to_native(sv_detections=as_sv) + + # then - values, ids, lineage and masks survive + assert isinstance(round_tripped, InstanceDetections) + _assert_native_detections_close(native, round_tripped) + assert [per_box[DETECTION_ID_KEY] for per_box in round_tripped.bboxes_metadata] == [ + "d0", + "d1", + ] + original_metadata = _native_image_metadata() + for key in ( + PARENT_ID_KEY, + ROOT_PARENT_ID_KEY, + PREDICTION_TYPE_KEY, + INFERENCE_ID_KEY, + ): + assert str(original_metadata[key]) == str(round_tripped.image_metadata[key]) + for key in ( + PARENT_COORDINATES_KEY, + PARENT_DIMENSIONS_KEY, + ROOT_PARENT_COORDINATES_KEY, + ROOT_PARENT_DIMENSIONS_KEY, + IMAGE_DIMENSIONS_KEY, + ): + assert list(original_metadata[key]) == list(round_tripped.image_metadata[key]) + assert isinstance(round_tripped.mask, torch.Tensor) + assert np.array_equal( + round_tripped.mask.cpu().numpy().astype(bool), _crop_local_dense_masks() + ) + # acceptance bar: the tensor serializer accepts it... + serialized = serialise_sv_detections(round_tripped) + assert [entry["class"] for entry in serialized["predictions"]] == ["a", "b"] + # ...and root-coordinates conversion works on it (crop lineage present) + at_root = native_detections_to_root_coordinates(prediction=round_tripped) + assert np.allclose( + at_root.xyxy.cpu().numpy()[:, 0], + native.xyxy.cpu().numpy()[:, 0] + CROP_OFFSET_X, + ) + + +def test_sv_detections_to_native_round_trips_rle_carrier_under_prefer_rle() -> None: + # given - native with an RLE carrier + dense = _crop_local_dense_masks() + native = _native_instance_detections(mask=_dense_masks_to_rle(dense)) + + # when - IN densifies; OUT re-encodes under prefer_rle (declared RLE kind) + as_sv = native_detections_to_sv(detections=native) + round_tripped = sv_detections_to_native(sv_detections=as_sv, prefer_rle=True) + + # then - carrier is RLE again and decodes to the same masks + assert isinstance(round_tripped, InstanceDetections) + assert isinstance(round_tripped.mask, InstancesRLEMasks) + from inference_models.models.common.rle_utils import coco_rle_masks_to_numpy_mask + + assert np.array_equal(coco_rle_masks_to_numpy_mask(round_tripped.mask), dense) + + +def test_sv_detections_to_native_round_trips_tracker_ids_and_class_overrides() -> None: + # given - two rows SHARING class_id 0, one carrying a per-box override, plus + # tracker ids (the classes_replacement -> rename regression shape) + native = InstanceDetections( + xyxy=torch.tensor( + [[10.0, 5.0, 30.0, 15.0], [40.0, 20.0, 55.0, 35.0]], dtype=torch.float32 + ), + class_id=torch.tensor([0, 0], dtype=torch.long), + confidence=torch.tensor([0.5, 0.25], dtype=torch.float32), + mask=torch.from_numpy(_crop_local_dense_masks()), + image_metadata=_native_image_metadata(), + bboxes_metadata=[ + {DETECTION_ID_KEY: "d0", TRACKER_ID_KEY: 11}, + {DETECTION_ID_KEY: "d1", TRACKER_ID_KEY: 14, CLASS_NAME_KEY: "override"}, + ], + ) + + # when + as_sv = native_detections_to_sv(detections=native) + round_tripped = sv_detections_to_native(sv_detections=as_sv) + + # then - serialized outputs are identical (the strongest effective-equality) + original_serialized = serialise_sv_detections(native) + round_tripped_serialized = serialise_sv_detections(round_tripped) + assert original_serialized == round_tripped_serialized + assert [per_box[TRACKER_ID_KEY] for per_box in round_tripped.bboxes_metadata] == [ + 11, + 14, + ] + + +def test_sv_detections_to_native_from_user_built_sv() -> None: + # given - a bare sv.Detections the way legacy user code builds one (no + # lineage, no detection ids) + user_built = sv.Detections( + xyxy=np.array([[1.0, 2.0, 11.0, 22.0]], dtype=np.float32), + class_id=np.array([3]), + confidence=np.array([0.9], dtype=np.float32), + data={"class_name": np.array(["widget"], dtype=object)}, + ) + + # when + as_native = sv_detections_to_native(sv_detections=user_built) + + # then - serializer hard requirements are met out of the box + assert isinstance(as_native, Detections) + assert as_native.image_metadata[CLASS_NAMES_KEY] == {3: "widget"} + minted = as_native.bboxes_metadata[0][DETECTION_ID_KEY] + UUID(minted) # parseable uuid + serialized = serialise_sv_detections(as_native) + assert serialized["predictions"][0]["class"] == "widget" + # and converting back preserves the user's view + back = native_detections_to_sv(detections=as_native) + assert np.allclose(back.xyxy, user_built.xyxy) + assert back.data["class_name"].tolist() == ["widget"] + + +def test_classification_single_label_round_trip() -> None: + # given - a native single-label prediction with threshold + time in metadata + metadata = { + CLASS_NAMES_KEY: {0: "cat", 1: "dog"}, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [480, 640], + INFERENCE_ID_KEY: "iid-9", + PARENT_ID_KEY: "img-1", + ROOT_PARENT_ID_KEY: "img-1", + "classification_confidence_threshold": 0.3, + "time": 0.0123, + } + native = ClassificationPrediction( + class_id=torch.tensor([1], dtype=torch.long), + confidence=torch.tensor([[0.3, 0.7]], dtype=torch.float32), + images_metadata=[metadata], + ) + + # when - native -> legacy dict -> native -> legacy dict + legacy_dict = serialise_native_classification(native) + rebuilt = classification_dict_to_native( + prediction=legacy_dict, block_name="b", value_name="o" + ) + re_serialized = serialise_native_classification(rebuilt) + + # then - byte-stable through the round trip, time included + assert isinstance(rebuilt, ClassificationPrediction) + assert legacy_dict == re_serialized + assert re_serialized["time"] == 0.0123 + + +def test_classification_multi_label_round_trip() -> None: + # given - a native multi-label prediction with a gap class id + metadata = { + CLASS_NAMES_KEY: {0: "cat", 1: "1", 2: "dog"}, + PREDICTION_TYPE_KEY: "classification", + IMAGE_DIMENSIONS_KEY: [480, 640], + INFERENCE_ID_KEY: "iid-9", + PARENT_ID_KEY: "img-1", + ROOT_PARENT_ID_KEY: "img-1", + } + native = MultiLabelClassificationPrediction( + class_ids=torch.tensor([0, 2], dtype=torch.long), + confidence=torch.tensor([0.9, 0.0, 0.8], dtype=torch.float32), + image_metadata=metadata, + ) + + # when + legacy_dict = serialise_native_classification(native) + rebuilt = classification_dict_to_native( + prediction=legacy_dict, block_name="b", value_name="o" + ) + re_serialized = serialise_native_classification(rebuilt) + + # then + assert isinstance(rebuilt, MultiLabelClassificationPrediction) + assert legacy_dict == re_serialized + assert rebuilt.class_ids.cpu().tolist() == [0, 2] + + +def test_keypoint_tuple_round_trip() -> None: + # given - a native keypoint prediction: bbox component carries per-box + # keypoint payloads (the serializer convention) + bboxes_metadata = [ + { + DETECTION_ID_KEY: "d0", + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS: [[11.0, 6.0], [21.0, 9.0]], + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS: [0.875, 0.75], + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS: [0, 1], + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS: ["nose", "eye"], + }, + ] + bbox_component = Detections( + xyxy=torch.tensor([[10.0, 5.0, 30.0, 15.0]], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([0.5], dtype=torch.float32), + image_metadata=_native_image_metadata(), + bboxes_metadata=bboxes_metadata, + ) + key_points = KeyPoints( + xy=torch.tensor([[[11.0, 6.0], [21.0, 9.0]]], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([[0.875, 0.75]], dtype=torch.float32), + image_metadata=_native_image_metadata(), + ) + + # when - IN via the declared keypoint kind, then OUT + with _boundary_on: + as_legacy = convert_kwargs_to_legacy( + kwargs={"prediction": (key_points, bbox_component)}, + manifest_description=_manifest( + inputs={ + "prediction": _input_with_kinds("keypoint_detection_prediction") + } + ), + block_name="b", + )["prediction"] + result = convert_block_result_to_native( + result={"prediction": as_legacy}, + manifest_description=_manifest_with_outputs( + outputs={ + "prediction": _output_with_kinds("keypoint_detection_prediction") + } + ), + block_name="b", + ) + + # then - the tuple shape is rebuilt with matching keypoints + rebuilt_key_points, rebuilt_bbox = result["prediction"] + assert isinstance(rebuilt_key_points, KeyPoints) + assert np.allclose(rebuilt_key_points.xy.cpu().numpy(), key_points.xy.cpu().numpy()) + assert np.allclose( + rebuilt_key_points.confidence.cpu().numpy(), + key_points.confidence.cpu().numpy(), + ) + assert ( + serialise_sv_detections(rebuilt_bbox)["predictions"][0]["keypoints"] + == serialise_sv_detections(bbox_component)["predictions"][0]["keypoints"] + ) + + +def test_convert_block_result_flow_control_passes_through() -> None: + # given + flow_control = FlowControl(context="$steps.a") + + # when + with _boundary_on: + scalar_result = convert_block_result_to_native( + result=flow_control, + manifest_description=_manifest_with_outputs(outputs={}), + block_name="b", + ) + listed_result = convert_block_result_to_native( + result=[flow_control, {"output": "text"}], + manifest_description=_manifest_with_outputs(outputs={}), + block_name="b", + ) + + # then + assert scalar_result is flow_control + assert listed_result[0] is flow_control + assert listed_result[1] == {"output": "text"} + + +def test_convert_block_result_walks_nested_lists_and_sniffs_wildcard() -> None: + # given - List[List[dict]] result (dimensionality +1) with an sv.Detections + # under a WILDCARD output plus a representation-invariant dict + native = _native_instance_detections( + mask=torch.from_numpy(_crop_local_dense_masks()) + ) + as_sv = native_detections_to_sv(detections=native) + result = [[{"detections": as_sv, "meta": {"key": "value"}}]] + + # when + with _boundary_on: + converted = convert_block_result_to_native( + result=result, + manifest_description=_manifest_with_outputs(outputs={}), + block_name="b", + ) + + # then - sv sniffed to native InstanceDetections; dict untouched + leaf = converted[0][0] + assert isinstance(leaf["detections"], InstanceDetections) + assert leaf["meta"] == {"key": "value"} + + +def test_convert_block_result_wildcard_sniff_rebuilds_keypoint_tuple() -> None: + # given - sv.Detections carrying all four keypoint payload columns + bboxes_metadata = [ + { + DETECTION_ID_KEY: "d0", + KEYPOINTS_XY_KEY_IN_SV_DETECTIONS: [[11.0, 6.0]], + KEYPOINTS_CONFIDENCE_KEY_IN_SV_DETECTIONS: [0.9], + KEYPOINTS_CLASS_ID_KEY_IN_SV_DETECTIONS: [0], + KEYPOINTS_CLASS_NAME_KEY_IN_SV_DETECTIONS: ["nose"], + }, + ] + bbox_component = Detections( + xyxy=torch.tensor([[10.0, 5.0, 30.0, 15.0]], dtype=torch.float32), + class_id=torch.tensor([0], dtype=torch.long), + confidence=torch.tensor([0.5], dtype=torch.float32), + image_metadata=_native_image_metadata(), + bboxes_metadata=bboxes_metadata, + ) + as_sv = native_detections_to_sv(detections=bbox_component) + + # when + with _boundary_on: + converted = convert_block_result_to_native( + result={"prediction": as_sv}, + manifest_description=_manifest_with_outputs(outputs={}), + block_name="b", + ) + + # then + rebuilt_key_points, rebuilt_bbox = converted["prediction"] + assert isinstance(rebuilt_key_points, KeyPoints) + assert isinstance(rebuilt_bbox, Detections) + + +def test_convert_block_result_declared_kind_wrong_type_raises_loudly() -> None: + # given - a declared classification output receiving sv.Detections + native = _native_object_detections() + as_sv = native_detections_to_sv(detections=native) + + # when + with _boundary_on, pytest.raises(RepresentationBoundaryError) as error: + _ = convert_block_result_to_native( + result={"output": as_sv}, + manifest_description=_manifest_with_outputs( + outputs={"output": _output_with_kinds("classification_prediction")} + ), + block_name="my_block", + ) + + # then + assert "my_block" in str(error.value) + assert "output" in str(error.value) + + +def test_convert_block_result_embedding_and_tensor_kinds() -> None: + # given + result = { + "embedding": [0.25, 0.5, 0.75], + "raw": np.array([[1, 2], [3, 4]], dtype=np.int64), + } + + # when + with _boundary_on: + converted = convert_block_result_to_native( + result=result, + manifest_description=_manifest_with_outputs( + outputs={ + "embedding": _output_with_kinds("embedding"), + "raw": _output_with_kinds("tensor"), + } + ), + block_name="b", + ) + + # then - embedding is float32; tensor preserves dtype + assert isinstance(converted["embedding"], torch.Tensor) + assert converted["embedding"].dtype == torch.float32 + assert np.allclose(converted["embedding"].cpu().numpy(), [0.25, 0.5, 0.75]) + assert isinstance(converted["raw"], torch.Tensor) + assert converted["raw"].dtype == torch.int64 + + +def test_convert_block_result_is_identity_when_boundary_inactive() -> None: + # given + result = {"output": object()} + manifest = _manifest_with_outputs(outputs={}) + + # when - flag off + with _boundary_off: + flag_off = convert_block_result_to_native( + result=result, manifest_description=manifest, block_name="b" + ) + # and - tensor_native mode (validated at construction into the enum) + manifest_native = ManifestDescription( + type="ManifestDescription", + block_type="MyBlock", + inputs={}, + outputs={}, + tensor_compatibility="tensor_native", + ) + with _boundary_on: + native_mode = convert_block_result_to_native( + result=result, manifest_description=manifest_native, block_name="b" + ) + + # then - the very same object, no copies + assert flag_off is result + assert native_mode is result + + +def test_sv_detections_to_native_takes_rle_mask_column_without_reencoding() -> None: + # given - an sv carrying the `rle_mask` column exactly as the numpy + # deserializer produces it (object array of COCO-RLE dicts) + from inference_models.models.common.rle_utils import ( + coco_rle_masks_to_numpy_mask, + torch_mask_to_coco_rle, + ) + + dense = _crop_local_dense_masks() + coco_dicts = [ + torch_mask_to_coco_rle(torch.as_tensor(instance_mask)) + for instance_mask in dense + ] + sv_detections = sv.Detections( + xyxy=np.asarray( + [[1.0, 2.0, 10.0, 12.0], [3.0, 4.0, 20.0, 22.0]], dtype=np.float32 + ), + class_id=np.asarray([0, 1]), + confidence=np.asarray([0.5, 0.75], dtype=np.float32), + mask=dense.astype(bool), + data={"rle_mask": np.array(coco_dicts, dtype=object)}, + ) + + # when + converted = sv_detections_to_native(sv_detections=sv_detections) + + # then - the carried RLE wins, verbatim (no re-encode), and decodes back + assert isinstance(converted, InstanceDetections) + assert isinstance(converted.mask, InstancesRLEMasks) + assert converted.mask.masks == [entry["counts"] for entry in coco_dicts] + assert np.array_equal( + coco_rle_masks_to_numpy_mask(converted.mask), dense.astype(bool) + ) + + +def test_sv_detections_to_native_malformed_rle_mask_column_raises_loudly() -> None: + # given - an entry missing the `counts` key + sv_detections = sv.Detections( + xyxy=np.asarray([[1.0, 2.0, 10.0, 12.0]], dtype=np.float32), + class_id=np.asarray([0]), + confidence=np.asarray([0.5], dtype=np.float32), + data={"rle_mask": np.array([{"size": [4, 4]}], dtype=object)}, + ) + + # when + with pytest.raises(ValueError) as error: + _ = sv_detections_to_native(sv_detections=sv_detections) + + # then + assert "rle_mask" in str(error.value) + assert "COCO-RLE" in str(error.value) + + +@pytest.mark.parametrize( + "declared_kind, expected_rle", + [ + ("instance_segmentation_prediction", False), + ("rle_instance_segmentation_prediction", True), + ], +) +def test_convert_block_result_empty_sv_declared_mask_kind_stays_instance_shaped( + declared_kind: str, + expected_rle: bool, +) -> None: + # given + manifest = _manifest_with_outputs( + outputs={"predictions": _output_with_kinds(declared_kind)} + ) + + # when + with _boundary_on: + converted = convert_block_result_to_native( + result={"predictions": sv.Detections.empty()}, + manifest_description=manifest, + block_name="my_block", + ) + + # then - empty output under a declared mask-carrying kind keeps the + # InstanceDetections shape (native empty convention), metadata is None, + # and the tensor serializer accepts it + prediction = converted["predictions"] + assert isinstance(prediction, InstanceDetections) + assert prediction.bboxes_metadata is None + assert isinstance(prediction.mask, InstancesRLEMasks) is expected_rle + serialized = serialise_sv_detections(prediction) + assert serialized["predictions"] == [] + + +def test_convert_block_result_empty_sv_wildcard_stays_plain_detections() -> None: + # when - wildcard output: intent unknowable, plain Detections is correct + with _boundary_on: + converted = convert_block_result_to_native( + result={"anything": sv.Detections.empty()}, + manifest_description=_manifest_with_outputs(outputs={}), + block_name="my_block", + ) + + # then + prediction = converted["anything"] + assert isinstance(prediction, Detections) + assert not isinstance(prediction, InstanceDetections) + assert prediction.bboxes_metadata is None + assert serialise_sv_detections(prediction)["predictions"] == [] + + +def test_convert_block_result_wildcard_bare_sv_keypoints_raises() -> None: + # given - a bare sv.KeyPoints has no native keypoint-prediction equivalent + bare_key_points = sv.KeyPoints( + xy=np.asarray([[[1.0, 2.0], [3.0, 4.0]]], dtype=np.float32) + ) + + # when + with _boundary_on, pytest.raises(RepresentationBoundaryError) as error: + _ = convert_block_result_to_native( + result={"key_points": bare_key_points}, + manifest_description=_manifest_with_outputs(outputs={}), + block_name="my_block", + ) + + # then - loud, actionable, mirrors the IN-side bare-KeyPoints rule + assert "my_block" in str(error.value) + assert "sv.KeyPoints" in str(error.value) or "detections component" in str( + error.value + ) diff --git a/tests/workflows/unit_tests/execution_engine/entities/test_base.py b/tests/workflows/unit_tests/execution_engine/entities/test_base.py index 3327dd4c0f..c49350e3ff 100644 --- a/tests/workflows/unit_tests/execution_engine/entities/test_base.py +++ b/tests/workflows/unit_tests/execution_engine/entities/test_base.py @@ -1122,3 +1122,532 @@ def test_parent_origin_validation_rejects_negative_height() -> None: width=100, height=-100, ) + + +# --------------------------------------------------------------------------- +# Tensor-native representation tests +# --------------------------------------------------------------------------- + + +import torch # noqa: E402 (kept low to avoid touching the existing import block) + +from inference.core.env import ( # noqa: E402 + ENABLE_TENSOR_DATA_REPRESENTATION, + WORKFLOWS_IMAGE_TENSOR_DEVICE, +) + + +@pytest.mark.skipif( + not ENABLE_TENSOR_DATA_REPRESENTATION, + reason="tensor-only: WORKFLOWS_IMAGE_TENSOR_DEVICE is None when the flag is off, " + "so the configured-device assertion does not apply", +) +def test_init_workflow_image_data_from_tensor_only() -> None: + # given + # Allocated on the configured device so the no-copy identity below holds even + # when WORKFLOWS_IMAGE_TENSOR_DEVICE is cuda. + tensor = torch.zeros( + (3, 10, 20), dtype=torch.uint8, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + + # when + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=tensor, + ) + + # then + assert image.tensor_image is tensor + assert image.tensor_image.device == WORKFLOWS_IMAGE_TENSOR_DEVICE + + +def test_workflow_image_data_numpy_fallback_does_rgb_to_bgr() -> None: + # given + # Tensor is CHW uint8 RGB by convention. Bake distinct R/G/B values so + # the channel order is asserted, not just the shape. + tensor = torch.zeros((3, 2, 2), dtype=torch.uint8) + tensor[0] = 10 # R + tensor[1] = 20 # G + tensor[2] = 30 # B + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=tensor, + ) + + # when + numpy_image = image.numpy_image + + # then + # numpy is HWC uint8 BGR, so [..., 0]=B=30, [..., 1]=G=20, [..., 2]=R=10. + assert numpy_image.shape == (2, 2, 3) + assert numpy_image.dtype == np.uint8 + assert np.all(numpy_image[..., 0] == 30) + assert np.all(numpy_image[..., 1] == 20) + assert np.all(numpy_image[..., 2] == 10) + + +def test_workflow_image_data_tensor_fallback_does_bgr_to_rgb() -> None: + # given + # numpy is HWC uint8 BGR. + numpy_image = np.zeros((2, 2, 3), dtype=np.uint8) + numpy_image[..., 0] = 30 # B + numpy_image[..., 1] = 20 # G + numpy_image[..., 2] = 10 # R + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=numpy_image, + ) + + # when + tensor = image.tensor_image + + # then + # tensor is CHW uint8 RGB, contiguous. + assert tuple(tensor.shape) == (3, 2, 2) + assert tensor.dtype == torch.uint8 + assert tensor.is_contiguous() + assert torch.all(tensor[0] == 10) + assert torch.all(tensor[1] == 20) + assert torch.all(tensor[2] == 30) + + +def test_workflow_image_data_chw_round_trip_is_pixel_exact_and_non_square() -> None: + # given a non-square BGR image so an H<->W transpose would be detectable + numpy_image = np.zeros((4, 6, 3), dtype=np.uint8) + numpy_image[..., 0] = 30 # B + numpy_image[..., 1] = 20 # G + numpy_image[..., 2] = 10 # R + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=numpy_image, + ) + + # when + tensor = image.tensor_image # CHW RGB + round_tripped = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=tensor, + ).numpy_image # back to HWC BGR + + # then + assert tuple(tensor.shape) == (3, 4, 6) # C, H, W + assert tensor.is_contiguous() + assert image._read_shape_without_materialization() == (4, 6) # H, W + assert round_tripped.shape == (4, 6, 3) + assert np.array_equal(round_tripped, numpy_image) + + +def test_workflow_image_data_numpy_fallback_caches() -> None: + # given + tensor = torch.zeros((3, 4, 3), dtype=torch.uint8) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=tensor, + ) + + # when + first = image.numpy_image + second = image.numpy_image + + # then + assert first is second, "lazy materialization should cache the numpy array" + + +def test_workflow_image_data_tensor_fallback_caches() -> None: + # given + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=np.zeros((3, 4, 3), dtype=np.uint8), + ) + + # when + first = image.tensor_image + second = image.tensor_image + + # then + assert first is second, "lazy materialization should cache the tensor" + + +def test_workflow_image_data_parent_metadata_without_forcing_materialization() -> None: + # given + # Tensor-only construction. Reading parent_metadata must populate + # origin coordinates without going through numpy materialization. + tensor = torch.zeros((3, 7, 11), dtype=torch.uint8) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=tensor, + ) + + # when + parent_metadata = image.parent_metadata + + # then + assert parent_metadata.origin_coordinates == OriginCoordinatesSystem( + left_top_x=0, + left_top_y=0, + origin_width=11, + origin_height=7, + ) + # Sanity: shape was read from the tensor; the numpy cache is still empty. + assert image._numpy_image is None + + +def test_workflow_image_create_crop_from_tensor() -> None: + # given + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + workflow_root_ancestor_metadata=ImageParentMetadata(parent_id="root"), + tensor_image=torch.zeros( + (3, 192, 168), dtype=torch.uint8, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + ) + cropped = torch.zeros( + (3, 64, 60), dtype=torch.uint8, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ) + + # when + result = WorkflowImageData.create_crop_from_tensor( + origin_image_data=image, + crop_identifier="my_crop", + cropped_tensor_image=cropped, + offset_x=100, + offset_y=20, + ) + + # then + assert result.tensor_image is cropped + assert result.parent_metadata.parent_id == "my_crop" + assert result.parent_metadata.origin_coordinates == OriginCoordinatesSystem( + left_top_x=100, + left_top_y=20, + origin_width=168, + origin_height=192, + ) + assert result.workflow_root_ancestor_metadata.parent_id == "root" + assert ( + result.workflow_root_ancestor_metadata.origin_coordinates + == OriginCoordinatesSystem( + left_top_x=100, + left_top_y=20, + origin_width=168, + origin_height=192, + ) + ) + + +def test_workflow_image_copy_and_replace_preserves_tensor() -> None: + # given + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.zeros( + (3, 4, 3), dtype=torch.uint8, device=WORKFLOWS_IMAGE_TENSOR_DEVICE + ), + ) + + # when + # copy_and_replace with no image-related kwargs must carry the tensor over. + result = WorkflowImageData.copy_and_replace( + origin_image_data=image, + parent_metadata=ImageParentMetadata(parent_id="new_parent"), + ) + + # then + assert result.tensor_image is image.tensor_image + assert result.parent_metadata.parent_id == "new_parent" + + +def test_workflow_image_copy_and_replace_swaps_tensor_for_numpy() -> None: + # given + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.zeros((3, 4, 3), dtype=torch.uint8), + ) + replacement = np.zeros((5, 6, 3), dtype=np.uint8) + + # when + result = WorkflowImageData.copy_and_replace( + origin_image_data=image, + numpy_image=replacement, + ) + + # then + # When any image-representation kwarg is provided, the others reset to + # whatever was passed (None by default). Tensor should be gone. + assert result._tensor_image is None + assert result.numpy_image is replacement + + +def test_workflow_image_data_single_channel_numpy_to_tensor() -> None: + # given - Convert Grayscale / Threshold blocks produce 2-D (H, W) arrays + gray = np.arange(24, dtype=np.uint8).reshape(4, 6) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=gray, + ) + + # when + tensor = image.tensor_image + + # then - (1, H, W), values untouched (no channel reversal for single-channel) + assert tuple(tensor.shape) == (1, 4, 6) + assert tensor.dtype == torch.uint8 + assert tensor.is_contiguous() + assert np.array_equal(tensor.squeeze(0).cpu().numpy(), gray) + + +def test_workflow_image_data_single_channel_tensor_to_numpy_round_trip() -> None: + # given + gray = np.arange(24, dtype=np.uint8).reshape(4, 6) + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.from_numpy(gray.copy()).unsqueeze(0), + ) + + # when + round_tripped = tensor_born.numpy_image + + # then - back to the 2-D (H, W) shape numpy-land blocks produce + assert round_tripped.shape == (4, 6) + assert np.array_equal(round_tripped, gray) + assert tensor_born._read_shape_without_materialization() == (4, 6) + + +def test_workflow_image_data_declared_numpy_mutation_refreshes_tensor() -> None: + # given - numpy-born image with the tensor sibling already derived; both + # representations stay cached for readers (fan-out reads are free) + numpy_image = np.zeros((4, 6, 3), dtype=np.uint8) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=numpy_image, + ) + first_tensor = image.tensor_image + assert image.tensor_image is first_tensor, "readers keep the cached tensor" + + # when - the copy_image=False visualization pattern: mutate numpy in place, + # then DECLARE the mutation per the class contract + image.numpy_image[0, 0] = (1, 2, 3) # BGR + undeclared_tensor = image.tensor_image + image.declare_numpy_image_mutated() + refreshed_tensor = image.tensor_image + + # then - before the declaration the sibling cache is stale (the documented + # limitation of undeclared in-place mutation); after it, the tensor is + # re-derived from the mutated pixels (RGB order at [:, 0, 0]) + assert undeclared_tensor is first_tensor and torch.all(undeclared_tensor == 0) + assert tuple(int(c) for c in refreshed_tensor[:, 0, 0]) == (3, 2, 1) + + +def test_workflow_image_data_declared_tensor_mutation_refreshes_numpy() -> None: + # given - tensor-born image (the tensor-mode video/crop case) with the numpy + # sibling already derived; tensor residency survives numpy reads + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.zeros((3, 4, 6), dtype=torch.uint8), + ) + _ = image.numpy_image + assert image.is_tensor_materialised(), "numpy reads must not evict the tensor" + + # when - mutate the tensor in place and declare it + image.tensor_image[:, 0, 0] = torch.tensor([3, 2, 1], dtype=torch.uint8) # RGB + image.declare_tensor_image_mutated() + + # then - numpy is re-derived from the mutated tensor (BGR order at [0, 0]) + assert np.array_equal(image.numpy_image[0, 0], np.array([1, 2, 3], dtype=np.uint8)) + + +def test_workflow_image_data_tensor_from_base64_does_not_cache_numpy() -> None: + # given - PNG so the decode is lossless and pixel assertions are exact + numpy_image = np.zeros((4, 6, 3), dtype=np.uint8) + numpy_image[..., 0] = 30 # B + numpy_image[..., 1] = 20 # G + numpy_image[..., 2] = 10 # R + encoded = base64.b64encode(cv2.imencode(".png", numpy_image)[1]).decode("ascii") + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + base64_image=encoded, + ) + + # when - the tensor is requested FIRST on a base64-born image + tensor = image.tensor_image + + # then - the source decodes straight into the tensor representation; the + # transient decode buffer must NOT be left behind as a cached numpy + assert image._numpy_image is None, "decode must not leave a cached numpy behind" + assert tuple(tensor.shape) == (3, 4, 6) + assert torch.all(tensor[0] == 10) # R + assert torch.all(tensor[1] == 20) # G + assert torch.all(tensor[2] == 30) # B + # numpy, when later needed, re-derives from the tensor - pixel-exact + assert np.array_equal(image.numpy_image, numpy_image) + + +def test_workflow_image_data_shape_read_fallback_materializes_per_flag() -> None: + # given - a base64-born image with no in-memory representation yet + encoded = base64.b64encode( + cv2.imencode(".png", np.zeros((7, 11, 3), dtype=np.uint8))[1] + ).decode("ascii") + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + base64_image=encoded, + ) + + # when + shape = image._read_shape_without_materialization() + + # then - the flag-appropriate representation (and only it) materialised + assert shape == (7, 11) + if ENABLE_TENSOR_DATA_REPRESENTATION: + assert image._tensor_image is not None, "tensor mode materialises the tensor" + assert image._numpy_image is None, "tensor mode must not cache numpy" + else: + assert image._numpy_image is not None, "numpy mode materialises numpy" + assert image._tensor_image is None, "numpy mode must not build the tensor" + + +def test_workflow_image_data_tensor_from_file_reference_does_not_cache_numpy( + tmp_path, +) -> None: + # given - a lossless PNG on disk referenced by path + numpy_image = np.zeros((4, 6, 3), dtype=np.uint8) + numpy_image[..., 0] = 30 # B + numpy_image[..., 1] = 20 # G + numpy_image[..., 2] = 10 # R + path = str(tmp_path / "reference.png") + assert cv2.imwrite(path, numpy_image) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + image_reference=path, + ) + + # when - the tensor is requested FIRST on a reference-born image + tensor = image.tensor_image + + # then - decoded straight to CHW RGB, no cached numpy left behind + assert image._numpy_image is None, "decode must not leave a cached numpy behind" + assert tuple(tensor.shape) == (3, 4, 6) + assert torch.all(tensor[0] == 10) # R + assert torch.all(tensor[1] == 20) # G + assert torch.all(tensor[2] == 30) # B + assert np.array_equal(image.numpy_image, numpy_image) + + +def test_workflow_image_data_declare_mutated_requires_materialised_representation() -> ( + None +): + # given - tensor-born image: numpy never materialised, and vice versa + tensor_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.zeros((3, 4, 6), dtype=torch.uint8), + ) + numpy_born = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + numpy_image=np.zeros((4, 6, 3), dtype=np.uint8), + ) + + # when / then - declaring a mutation of a representation that was never + # materialised is a caller bug and must raise + with pytest.raises(ValueError): + tensor_born.declare_numpy_image_mutated() + with pytest.raises(ValueError): + numpy_born.declare_tensor_image_mutated() + + +def test_workflow_image_data_declare_mutated_cuts_off_original_source( + tmp_path, +) -> None: + # given - a file-born image whose numpy got materialised and mutated + numpy_image = np.zeros((4, 6, 3), dtype=np.uint8) + path = str(tmp_path / "source.png") + assert cv2.imwrite(path, numpy_image) + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + image_reference=path, + ) + buffer = image.numpy_image + buffer[0, 0] = (1, 2, 3) + + # when + image.declare_numpy_image_mutated() + + # then - the original reference no longer describes the pixels and must be + # cut off: serialization falls back to the (mutated) numpy content + assert image._image_reference is None + serialized = image.to_inference_format() + assert serialized["type"] == "numpy_object" + assert np.array_equal(serialized["value"][0, 0], np.array([1, 2, 3], np.uint8)) + + +def test_workflow_image_data_fetch_then_declare_pattern() -> None: + # given - the caller pattern the visualization blocks use: fetch the buffer, + # copy or mutate-and-declare depending on copy_image + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + tensor_image=torch.zeros((3, 4, 6), dtype=torch.uint8), + ) + + # when - copy_image=True leg: mutate a private copy + scene = image.numpy_image.copy() + scene[0, 0] = (9, 9, 9) + + # then - the image is untouched + assert np.array_equal(image.numpy_image[0, 0], np.zeros(3, dtype=np.uint8)) + + # when - copy_image=False leg: mutate the live buffer, then declare + scene = image.numpy_image + scene[0, 0] = (1, 2, 3) # BGR + image.declare_numpy_image_mutated() + + # then - numpy is the sole source of truth; the tensor re-derives mutated + assert not image.is_tensor_materialised(), "declare must drop the tensor" + assert tuple(int(c) for c in image.tensor_image[:, 0, 0]) == (3, 2, 1) # RGB + + +def test_workflow_image_data_guards_source_fallback_after_declared_mutation() -> None: + # given - a base64-born image whose numpy got materialised, mutated and + # declared (original source invalidated), and then - simulating an + # internal-consistency bug - the sole-source-of-truth representation is + # dropped, so the property fallback chain would reach the source decoders + encoded = base64.b64encode( + cv2.imencode(".png", np.zeros((4, 6, 3), dtype=np.uint8))[1] + ).decode("ascii") + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + base64_image=encoded, + ) + image.numpy_image[0, 0] = (1, 2, 3) + image.declare_numpy_image_mutated() + image._numpy_image = None # the invariant-breaking drop + + # when / then - both properties must refuse to decode the stale source and + # name the mutation as the cause instead of resurrecting pre-mutation pixels + with pytest.raises(ValueError, match="mutation"): + _ = image.numpy_image + with pytest.raises(ValueError, match="mutation"): + _ = image.tensor_image + + +def test_workflow_image_data_base64_rederived_after_mutation_is_valid_source() -> None: + # given - a base64-born image, mutated brightly and declared; base64 + # accessed AFTERWARDS re-encodes from the mutated pixels (a valid value) + encoded = base64.b64encode( + cv2.imencode(".png", np.zeros((32, 32, 3), dtype=np.uint8))[1] + ).decode("ascii") + image = WorkflowImageData( + parent_metadata=ImageParentMetadata(parent_id="parent"), + base64_image=encoded, + ) + image.numpy_image[:, :] = (200, 200, 200) + image.declare_numpy_image_mutated() + rederived = image.base64_image + assert rederived != encoded, "base64 must be re-encoded from mutated pixels" + + # when - simulating the degenerate drop of the pixel caches: the re-derived + # base64 is now the only channel left + image._numpy_image = None + image._tensor_image = None + recovered = image.numpy_image + + # then - the re-derived base64 decodes (no staleness guard trip) and holds + # the POST-mutation content (JPEG re-encode is lossy, hence approximate) + assert recovered.mean() > 150, "decoded pixels must reflect the mutation" + assert tuple(recovered.shape) == (32, 32, 3)