diff --git a/comfy_api/latest/_sdk.py b/comfy_api/latest/_sdk.py index 90d533852..2503d5144 100644 --- a/comfy_api/latest/_sdk.py +++ b/comfy_api/latest/_sdk.py @@ -1423,55 +1423,6 @@ class ClipSegRef(_TypedRef): return result[0], result[1] -class ImageClassifierRef(_TypedRef): - KIND = "IMAGE_CLASSIFIER" - - async def classify( - self, images: ImageRef, use_accelerator: bool = True, - top_k: int = 5, - ) -> list[list[dict[str, Any]]]: - """Classify a host-side image batch and return bounded label scores.""" - return await current_runtime().ops.apply( - "image_classifier.classify", self, { - "images": images, - "use_accelerator": bool(use_accelerator), - "top_k": int(top_k), - }) - - async def predict_scores( - self, images: ImageRef, - ) -> "ClassifierScoresRef": - """Run a multi-label classifier and retain its score matrix host-side.""" - return await current_runtime().ops.apply( - "image_classifier.predict_scores", self, {"images": images}) - - -class ClassifierScoresRef(_TypedRef): - """Opaque bounded batch-by-class scores from an image classifier.""" - - KIND = "CLASSIFIER_SCORES" - - async def shape(self) -> tuple[int, int]: - result = await current_runtime().ops.apply( - "classifier_scores.shape", self, {}) - return int(result[0]), int(result[1]) - - async def select_above( - self, batch_index: int, start: int, end: int, threshold: float, - offset: int = 0, limit: int = 512, - ) -> dict[str, Any]: - """Page score/index pairs above a threshold in one class range.""" - return await current_runtime().ops.apply( - "classifier_scores.select_above", self, { - "batch_index": int(batch_index), - "start": int(start), - "end": int(end), - "threshold": float(threshold), - "offset": int(offset), - "limit": int(limit), - }) - - class SemanticSegmentationRef(_TypedRef): """Opaque fixed-architecture semantic segmentation model.""" @@ -2314,18 +2265,6 @@ class ModelsDomain(Protocol): dtype: str = "float16", ) -> PowerPaintRef: ... async def load_clipseg(self, model: str) -> ClipSegRef: ... - async def load_image_classifier( - self, model: str, architecture: str, labels: list[str], - ) -> ImageClassifierRef: ... - async def load_onnx_image_classifier( - self, model: str, input_layout: str = "NHWC", - channel_order: str = "BGR", resize_mode: str = "fit_pad", - input_scale: float = 255.0, - pad_color: tuple[float, float, float] = (1.0, 1.0, 1.0), - mean: tuple[float, float, float] = (0.0, 0.0, 0.0), - std: tuple[float, float, float] = (1.0, 1.0, 1.0), - activation: str = "identity", resize_filter: str = "lanczos", - ) -> ImageClassifierRef: ... async def load_segformer( self, model: str, variant: str, num_labels: int, ) -> SemanticSegmentationRef: ... @@ -3268,170 +3207,6 @@ _CLIPSEG_CACHE = WeightCache( ) -@dataclass -class _ImageClassifierEntry: - model: Any - processor: Any - architecture: str - num_labels: int - lock: threading.Lock = field(default_factory=threading.Lock) - - -def _load_image_classifier_weight( - path: str, architecture: str, -) -> _ImageClassifierEntry: - """Build one closed image-classifier architecture from SafeTensors.""" - import torch - from safetensors.torch import load_file - from transformers import ( - BeitConfig, - BeitForImageClassification, - BeitImageProcessor, - ConvNextImageProcessor, - ResNetConfig, - ResNetForImageClassification, - ViTConfig, - ViTForImageClassification, - ViTImageProcessor, - ) - - state = load_file(path, device="cpu") - if not state: - raise ValueError("classifier SafeTensors file contains no weights") - heads = { - "vit-base-patch16-224": "classifier.weight", - "beit-base-patch16-224": "classifier.weight", - "resnet-50-224": "classifier.1.weight", - } - if architecture not in heads: - raise ValueError("image classifier architecture is not supported") - head = state.get(heads[architecture]) - if not isinstance(head, torch.Tensor) or head.ndim != 2: - raise ValueError("classifier weights have no compatible output head") - num_labels = int(head.shape[0]) - if not 1 <= num_labels <= 10_000: - raise ValueError("classifier output count is outside the safe range") - floating_dtypes = { - value.dtype for value in state.values() - if isinstance(value, torch.Tensor) and value.is_floating_point() - } - if len(floating_dtypes) != 1: - raise ValueError("classifier weights must use one floating-point dtype") - dtype = next(iter(floating_dtypes)) - if dtype not in (torch.float16, torch.bfloat16, torch.float32): - raise ValueError("classifier weights use an unsupported dtype") - - if architecture == "vit-base-patch16-224": - config = ViTConfig( - num_labels=num_labels, - attention_probs_dropout_prob=0.0, - encoder_stride=16, - hidden_act="gelu", - hidden_dropout_prob=0.0, - hidden_size=768, - image_size=224, - initializer_range=0.02, - intermediate_size=3072, - layer_norm_eps=1e-12, - num_attention_heads=12, - num_channels=3, - num_hidden_layers=12, - patch_size=16, - qkv_bias=True, - ) - model = ViTForImageClassification(config) - processor = ViTImageProcessor( - do_resize=True, - size={"height": 224, "width": 224}, - resample=2, - do_rescale=True, - rescale_factor=1.0 / 255.0, - do_normalize=True, - image_mean=(0.5, 0.5, 0.5), - image_std=(0.5, 0.5, 0.5), - ) - elif architecture == "beit-base-patch16-224": - config = BeitConfig( - num_labels=num_labels, - attention_probs_dropout_prob=0.0, - drop_path_rate=0.1, - hidden_act="gelu", - hidden_dropout_prob=0.0, - hidden_size=768, - image_size=224, - initializer_range=0.02, - intermediate_size=3072, - layer_norm_eps=1e-12, - layer_scale_init_value=0.1, - num_attention_heads=12, - num_channels=3, - num_hidden_layers=12, - patch_size=16, - use_absolute_position_embeddings=False, - use_mask_token=False, - use_mean_pooling=True, - use_relative_position_bias=True, - use_shared_relative_position_bias=False, - ) - model = BeitForImageClassification(config) - processor = BeitImageProcessor( - do_resize=True, - size={"height": 224, "width": 224}, - resample=2, - do_rescale=True, - rescale_factor=1.0 / 255.0, - do_normalize=True, - do_center_crop=False, - crop_size={"height": 224, "width": 224}, - do_reduce_labels=False, - image_mean=(0.5, 0.5, 0.5), - image_std=(0.5, 0.5, 0.5), - ) - else: - config = ResNetConfig( - num_labels=num_labels, - depths=[3, 4, 6, 3], - downsample_in_first_stage=False, - embedding_size=64, - hidden_act="relu", - hidden_sizes=[256, 512, 1024, 2048], - layer_type="bottleneck", - num_channels=3, - out_features=["stage4"], - out_indices=[4], - ) - model = ResNetForImageClassification(config) - processor = ConvNextImageProcessor( - do_resize=True, - size={"shortest_edge": 224}, - resample=3, - do_rescale=True, - rescale_factor=1.0 / 255.0, - do_normalize=True, - image_mean=(0.485, 0.456, 0.406), - image_std=(0.229, 0.224, 0.225), - ) - - model = model.to(dtype=dtype) - model.load_state_dict(state, strict=True) - model.eval() - return _ImageClassifierEntry( - model=model, - processor=processor, - architecture=architecture, - num_labels=num_labels, - ) - - - - -_IMAGE_CLASSIFIER_CACHE = WeightCache( - load=_loader("_load_image_classifier_weight"), - max_entries=3, - release=_release_model_to_cpu, -) - - @dataclass class _TextEncoderEntry: clip: Any @@ -3856,99 +3631,6 @@ def _validate_onnx_weight_file(path: str) -> None: raise ValueError("ONNX model failed structural validation") from exc -@dataclass -class _OnnxImageClassifierEntry: - session: Any - input_name: str - output_name: str - input_height: int - input_width: int - class_count: int - input_layouts: frozenset[str] - lock: threading.Lock = field(default_factory=threading.Lock) - - -def _load_onnx_image_classifier(path: str) -> _OnnxImageClassifierEntry: - _validate_onnx_weight_file(path) - try: - import onnxruntime as ort - except ImportError as exc: - raise RuntimeError( - "ONNX image classification requires onnxruntime") from exc - - options = ort.SessionOptions() - options.log_severity_level = 3 - available = set(ort.get_available_providers()) - providers = [ - provider for provider in ( - "CUDAExecutionProvider", "CPUExecutionProvider") - if provider in available - ] - if not providers: - raise RuntimeError("ONNX Runtime has no supported execution provider") - try: - session = ort.InferenceSession( - path, sess_options=options, providers=providers) - except Exception as exc: - raise ValueError("ONNX image classifier could not be loaded") from exc - inputs = session.get_inputs() - outputs = session.get_outputs() - if len(inputs) != 1 or len(outputs) != 1: - raise ValueError("ONNX image classifier must have one input and output") - model_input = inputs[0] - model_output = outputs[0] - if model_input.type != "tensor(float)" or model_output.type not in { - "tensor(float)", "tensor(float16)", "tensor(double)", - }: - raise ValueError("ONNX image classifier must use floating-point tensors") - input_shape = model_input.shape - output_shape = model_output.shape - if len(input_shape) != 4 or len(output_shape) != 2: - raise ValueError("ONNX image classifier has an invalid tensor rank") - - # WD-style NHWC and common NCHW models are both admitted. The selected - # layout is checked again when the loader binds preprocessing options. - nhwc = input_shape[3] == 3 - nchw = input_shape[1] == 3 - if not nhwc and not nchw: - raise ValueError("ONNX image classifier must consume three channels") - if nhwc and nchw: - raise ValueError("ONNX classifier channel layout is ambiguous") - height = input_shape[1] if nhwc else input_shape[2] - width = input_shape[2] if nhwc else input_shape[3] - class_count = output_shape[1] - if (type(height) is not int or type(width) is not int - or not 1 <= height <= 4096 or not 1 <= width <= 4096): - raise ValueError("ONNX classifier spatial dimensions must be fixed") - if type(class_count) is not int or not 1 <= class_count <= 16_384: - raise ValueError("ONNX classifier output count is outside the safe range") - return _OnnxImageClassifierEntry( - session=session, - input_name=model_input.name, - output_name=model_output.name, - input_height=height, - input_width=width, - class_count=class_count, - input_layouts=frozenset( - layout for layout, valid in (("NHWC", nhwc), ("NCHW", nchw)) - if valid), - ) - - - - -_ONNX_IMAGE_CLASSIFIER_CACHE = WeightCache( - load=_loader("_load_onnx_image_classifier"), max_entries=3) - - - - - - - - - - @dataclass class _SamEntry: model: Any @@ -5176,8 +4858,6 @@ class _InProcessModels: _TEXT_GENERATOR_CACHE.clear() _INPAINT_MODEL_CACHE.clear() _CLIPSEG_CACHE.clear() - _IMAGE_CLASSIFIER_CACHE.clear() - _ONNX_IMAGE_CLASSIFIER_CACHE.clear() _TEXT_ENCODER_CACHE.clear() _LANGUAGE_MODEL_CACHE.clear() InProcessLlamaCpp().clear() @@ -5199,128 +4879,6 @@ class _InProcessModels: return ClipSegRef._wrap(await current_runtime().refs.create( "CLIPSEGMODEL", value)) # type: ignore[return-value] - async def load_image_classifier( - self, model: str, architecture: str, labels: list[str], - ) -> ImageClassifierRef: - import folder_paths - - model = self._model_name(model, "image classifier weight") - if not model.lower().endswith(".safetensors"): - raise ValueError("image classifier weights must use SafeTensors") - architecture = str(architecture) - if architecture not in { - "vit-base-patch16-224", - "beit-base-patch16-224", - "resnet-50-224", - }: - raise ValueError("image classifier architecture is not supported") - if not isinstance(labels, (list, tuple)): - raise TypeError("image classifier labels must be a list") - labels = tuple(str(label) for label in labels) - if (not labels or len(labels) > 10_000 - or any(not label or len(label) > 256 for label in labels)): - raise ValueError("image classifier labels are invalid") - path = folder_paths.get_full_path_or_raise("detection", model) - entry = await asyncio.to_thread( - _IMAGE_CLASSIFIER_CACHE.get, path, architecture) - if len(labels) != entry.num_labels: - raise ValueError( - "image classifier labels do not match the weight output count") - value = { - "model": entry.model, - "processor": entry.processor, - "architecture": entry.architecture, - "labels": labels, - "lock": entry.lock, - } - return ImageClassifierRef._wrap(await current_runtime().refs.create( - "IMAGE_CLASSIFIER", value)) # type: ignore[return-value] - - async def load_onnx_image_classifier( - self, model: str, input_layout: str = "NHWC", - channel_order: str = "BGR", resize_mode: str = "fit_pad", - input_scale: float = 255.0, - pad_color: tuple[float, float, float] = (1.0, 1.0, 1.0), - mean: tuple[float, float, float] = (0.0, 0.0, 0.0), - std: tuple[float, float, float] = (1.0, 1.0, 1.0), - activation: str = "identity", resize_filter: str = "lanczos", - ) -> ImageClassifierRef: - """Bind a self-contained standard ONNX image classifier. - - Preprocessing is a closed, reusable transform. Labels, category - ranges, thresholds, exclusions, and output formatting remain node - code; the host only retains and pages the numeric score matrix. - """ - import math - import folder_paths - - model = self._model_name(model, "ONNX image classifier") - if not model.lower().endswith(".onnx"): - raise ValueError("ONNX image classifiers must use .onnx files") - input_layout = str(input_layout).upper() - channel_order = str(channel_order).upper() - resize_mode = str(resize_mode).lower() - activation = str(activation).lower() - resize_filter = str(resize_filter).lower() - if input_layout not in {"NHWC", "NCHW"}: - raise ValueError("ONNX classifier layout must be NHWC or NCHW") - if channel_order not in {"RGB", "BGR"}: - raise ValueError("ONNX classifier channel order must be RGB or BGR") - if resize_mode not in {"fit_pad", "stretch"}: - raise ValueError("ONNX classifier resize mode is not supported") - if activation not in {"identity", "sigmoid", "softmax"}: - raise ValueError("ONNX classifier activation is not supported") - if resize_filter not in {"nearest", "bilinear", "bicubic", "lanczos"}: - raise ValueError("ONNX classifier resize filter is not supported") - input_scale = float(input_scale) - if not math.isfinite(input_scale) or not 0 < input_scale <= 65_535: - raise ValueError("ONNX classifier input scale is invalid") - - def triple( - value: Any, field_name: str, *, nonzero: bool = False, - unit: bool = False, - ) -> tuple[float, float, float]: - if not isinstance(value, (list, tuple)) or len(value) != 3: - raise ValueError( - f"ONNX classifier {field_name} must have three values") - result = tuple(float(item) for item in value) - if (any(not math.isfinite(item) or abs(item) > 1_000_000 - for item in result) - or (nonzero and any(item == 0 for item in result)) - or (unit and any(not 0 <= item <= 1 for item in result))): - raise ValueError(f"ONNX classifier {field_name} is invalid") - return result # type: ignore[return-value] - - pad_color = triple(pad_color, "pad color", unit=True) - mean = triple(mean, "mean") - std = triple(std, "standard deviation", nonzero=True) - path = folder_paths.get_full_path_or_raise("onnx", model) - entry = await asyncio.to_thread(_ONNX_IMAGE_CLASSIFIER_CACHE.get, path) - if input_layout not in entry.input_layouts: - raise ValueError( - f"ONNX classifier tensor is not laid out as {input_layout}") - value = { - "secure_kind": "image_classifier.onnx", - "session": entry.session, - "input_name": entry.input_name, - "output_name": entry.output_name, - "input_height": entry.input_height, - "input_width": entry.input_width, - "class_count": entry.class_count, - "input_layout": input_layout, - "channel_order": channel_order, - "resize_mode": resize_mode, - "input_scale": input_scale, - "pad_color": pad_color, - "mean": mean, - "std": std, - "activation": activation, - "resize_filter": resize_filter, - "lock": entry.lock, - } - return ImageClassifierRef._wrap(await current_runtime().refs.create( - "IMAGE_CLASSIFIER", value)) # type: ignore[return-value] - async def load_segformer( self, model: str, variant: str, num_labels: int, ) -> SemanticSegmentationRef: @@ -9696,12 +9254,6 @@ class InProcessOps: "style_model.apply": self._style_model_apply, "clipseg.predict_mask": _vendor_ops.clipseg_predict_mask, "clipseg.segment": _vendor_ops.clipseg_segment, - "image_classifier.classify": _vendor_ops.image_classifier_classify, - "image_classifier.predict_scores": - _vendor_ops.image_classifier_predict_scores, - "classifier_scores.shape": _vendor_ops.classifier_scores_shape, - "classifier_scores.select_above": - _vendor_ops.classifier_scores_select_above, "semantic_segmentation.mask": _vendor_ops.semantic_segmentation_mask, "object_detector.detect": self._object_detector_detect, "inpaint_model.inpaint": _vendor_ops.inpaint_model_inpaint, @@ -13751,18 +13303,10 @@ def _ref_type_for(v: Any) -> tuple[type, str]: return SamModelRef, "SAM_MODEL" if v.get("secure_kind") == "object_detector.rt_detr" and "model" in v: return ObjectDetectorRef, "OBJECT_DETECTOR" - if v.get("secure_kind") == "classifier_scores.v1" and "scores" in v: - return ClassifierScoresRef, "CLASSIFIER_SCORES" - if v.get("secure_kind") == "image_classifier.onnx" and set(v) >= { - "session", "input_name", "output_name", "class_count", "lock", - }: - return ImageClassifierRef, "IMAGE_CLASSIFIER" if v.get("secure_kind") == "powerpaint.pipeline" and set(v) >= { "powerpaint", "clip", }: return PowerPaintRef, "POWERPAINT_MODEL" - if set(v) >= {"model", "processor", "architecture", "labels"}: - return ImageClassifierRef, "IMAGE_CLASSIFIER" if set(v) >= {"model", "processor"}: return ClipSegRef, "CLIPSEGMODEL" if "samples" in v: diff --git a/comfy_api/latest/_sdk_public.py b/comfy_api/latest/_sdk_public.py index c712faf2b..eef32c22d 100644 --- a/comfy_api/latest/_sdk_public.py +++ b/comfy_api/latest/_sdk_public.py @@ -17,7 +17,6 @@ from ._sdk import ( # noqa: F401 ControlNetWeightsRef, Context, CivitaiDomain, - ClassifierScoresRef, ClosureRef, ClosuresDomain, ExecutionDomain, @@ -28,7 +27,6 @@ from ._sdk import ( # noqa: F401 GuiderRef, HuggingFaceWeight, ImgBBDomain, - ImageClassifierRef, InpaintModelRef, ImagePreprocessorRef, ImageRef, @@ -81,8 +79,6 @@ __all__ = [ "UpscaleModelRef", "ValueRef", "ImageRef", - "ImageClassifierRef", - "ClassifierScoresRef", "InpaintModelRef", "ImagePreprocessorRef", "ImgBBDomain", diff --git a/comfy_api/latest/_sdk_public.pyi b/comfy_api/latest/_sdk_public.pyi index 29a7dc5d9..d5c7a3260 100644 --- a/comfy_api/latest/_sdk_public.pyi +++ b/comfy_api/latest/_sdk_public.pyi @@ -390,30 +390,6 @@ class AssetRef(Ref): class ClipSegRef(Ref): KIND: str -class ImageClassifierRef(Ref): - KIND: str - async def classify( - self, - images: ImageRef, - use_accelerator: bool = ..., - top_k: int = ..., - ) -> list[list[dict[str, Any]]]: ... - async def predict_scores( - self, images: ImageRef, - ) -> ClassifierScoresRef: ... - -class ClassifierScoresRef(Ref): - async def shape(self) -> tuple[int, int]: ... - async def select_above( - self, - batch_index: int, - start: int, - end: int, - threshold: float, - offset: int = ..., - limit: int = ..., - ) -> dict[str, Any]: ... - class SemanticSegmentationRef(Ref): KIND: str async def mask( @@ -705,19 +681,6 @@ class ModelsDomain(Protocol): revision: str = ..., sha256: Optional[str] = ..., ) -> str: ... - async def load_onnx_image_classifier( - self, - model: str, - input_layout: str = ..., - channel_order: str = ..., - resize_mode: str = ..., - input_scale: float = ..., - pad_color: tuple[float, float, float] = ..., - mean: tuple[float, float, float] = ..., - std: tuple[float, float, float] = ..., - activation: str = ..., - resize_filter: str = ..., - ) -> ImageClassifierRef: ... async def list_diffusion_models( self, include_connectors: bool = ... ) -> list[str]: ... @@ -797,12 +760,6 @@ class ModelsDomain(Protocol): dtype: str = ..., ) -> PowerPaintRef: ... async def load_clipseg(self, model: str) -> ClipSegRef: ... - async def load_image_classifier( - self, - model: str, - architecture: str, - labels: list[str], - ) -> ImageClassifierRef: ... async def load_segformer( self, model: str, diff --git a/comfy_api/latest/_vendor_ops.py b/comfy_api/latest/_vendor_ops.py index 59d2dad9b..912bc2d38 100644 --- a/comfy_api/latest/_vendor_ops.py +++ b/comfy_api/latest/_vendor_ops.py @@ -15,10 +15,8 @@ from . import _sdk if TYPE_CHECKING: from ._sdk import ( - ClassifierScoresRef, ClipSegRef, ControlNetWeightsRef, - ImageClassifierRef, ImagePreprocessorRef, ImageRef, InpaintModelRef, @@ -416,230 +414,6 @@ async def clipseg_predict_mask(clipseg: "ClipSegRef", images: "ImageRef", text: await rt.refs.create("MASK", result) ) # type: ignore[return-value] -async def image_classifier_classify(classifier: "ImageClassifierRef", images: "ImageRef", - use_accelerator: bool = True, top_k: int = 5, -) -> list[list[dict[str, Any]]]: - from contextlib import nullcontext - import numpy as np - import torch - from PIL import Image - import comfy.model_management - - top_k = int(top_k) - if not 1 <= top_k <= 1000: - raise ValueError("image classifier top_k must be in [1, 1000]") - rt = _sdk.current_runtime() - bundle = await rt.refs.resolve(classifier) - pixels = await rt.refs.resolve(images) - if pixels.ndim != 4 or pixels.shape[-1] < 3: - raise ValueError("classifier images must be a non-empty BHWC batch") - if not 1 <= len(pixels) <= 4096: - raise ValueError("classifier batch size must be in [1, 4096]") - labels = tuple(bundle["labels"]) - if not labels: - raise ValueError("image classifier has no labels") - - model = bundle["model"] - processor = bundle["processor"] - model_lock = bundle.get("lock") - offload_device = comfy.model_management.unet_offload_device() - if use_accelerator: - device = comfy.model_management.get_torch_device() - dtype = comfy.model_management.unet_dtype() - else: - device = torch.device("cpu") - dtype = torch.float32 - - with model_lock if model_lock is not None else nullcontext(): - model.to(dtype).to(device) - try: - source = [Image.fromarray(np.clip( - image.detach().cpu().numpy()[..., :3] * 255.0, - 0, 255, - ).astype(np.uint8), mode="RGB") for image in pixels] - inputs = processor(images=source, return_tensors="pt") - inputs = { - key: value.to(device) for key, value in inputs.items() - } - autocast = ( - dtype != torch.float32 - and not comfy.model_management.is_device_mps(device) - ) - scope = ( - torch.autocast( - comfy.model_management.get_autocast_device(device), - dtype=dtype, - ) if autocast else nullcontext() - ) - with scope, torch.inference_mode(): - logits = model(**inputs).logits - scores = torch.softmax(logits.float(), dim=-1) - finally: - model.to(offload_device) - - if scores.ndim != 2 or scores.shape[1] != len(labels): - raise RuntimeError("image classifier returned an invalid score shape") - count = min(top_k, len(labels)) - values, indices = torch.topk(scores.cpu(), count, dim=-1) - return [[ - {"label": labels[int(index)], "score": float(score)} - for score, index in zip(row_scores, row_indices) - ] for row_scores, row_indices in zip(values, indices)] - -async def image_classifier_predict_scores(classifier: "ImageClassifierRef", images: "ImageRef", -) -> "ClassifierScoresRef": - import numpy as np - import torch - from PIL import Image - - rt = _sdk.current_runtime() - bundle = await rt.refs.resolve(classifier) - if (not isinstance(bundle, dict) - or bundle.get("secure_kind") != "image_classifier.onnx"): - raise TypeError( - "predict_scores requires a validated ONNX image classifier") - pixels = await rt.refs.resolve(images) - if (not isinstance(pixels, torch.Tensor) or pixels.ndim != 4 - or pixels.shape[-1] < 3 or not 1 <= len(pixels) <= 64): - raise ValueError( - "ONNX classifier images must be a 1-64 item BHWC RGB batch") - height, width = map(int, pixels.shape[1:3]) - if (height <= 0 or width <= 0 - or height * width * len(pixels) > 268_435_456 - or not bool(torch.isfinite(pixels[..., :3]).all())): - raise ValueError("ONNX classifier image values are invalid") - - target_height = int(bundle["input_height"]) - target_width = int(bundle["input_width"]) - resampling = { - "nearest": Image.Resampling.NEAREST, - "bilinear": Image.Resampling.BILINEAR, - "bicubic": Image.Resampling.BICUBIC, - "lanczos": Image.Resampling.LANCZOS, - }[bundle["resize_filter"]] - pad = tuple( - int(round(float(value) * 255.0)) for value in bundle["pad_color"]) - mean = np.asarray(bundle["mean"], dtype=np.float32) - std = np.asarray(bundle["std"], dtype=np.float32) - input_scale = float(bundle["input_scale"]) - - def infer() -> np.ndarray: - rows = [] - with bundle["lock"]: - for frame in pixels: - source_array = np.clip( - frame.detach().cpu().numpy()[..., :3] * 255.0, - 0, 255, - ).astype(np.uint8) - source = Image.fromarray(source_array) - if bundle["resize_mode"] == "fit_pad": - ratio = min( - target_width / source.width, - target_height / source.height, - ) - resized_size = ( - max(1, int(source.width * ratio)), - max(1, int(source.height * ratio)), - ) - resized = source.resize(resized_size, resampling) - prepared = Image.new( - "RGB", (target_width, target_height), pad) - prepared.paste(resized, ( - (target_width - resized_size[0]) // 2, - (target_height - resized_size[1]) // 2, - )) - else: - prepared = source.resize( - (target_width, target_height), resampling) - array = np.asarray(prepared, dtype=np.float32) - array = array * (input_scale / 255.0) - if bundle["channel_order"] == "BGR": - array = array[..., ::-1] - array = (array - mean) / std - if bundle["input_layout"] == "NCHW": - array = np.transpose(array, (2, 0, 1)) - model_input = np.ascontiguousarray( - array[None, ...], dtype=np.float32) - output = bundle["session"].run( - [bundle["output_name"]], - {bundle["input_name"]: model_input}, - )[0] - output = np.asarray(output) - if output.shape != (1, int(bundle["class_count"])): - raise RuntimeError( - "ONNX classifier returned an invalid score shape") - row = output[0].astype(np.float32, copy=False) - if bundle["activation"] == "sigmoid": - row = 1.0 / (1.0 + np.exp(-np.clip(row, -80, 80))) - elif bundle["activation"] == "softmax": - shifted = row - np.max(row) - exponent = np.exp(shifted) - row = exponent / np.sum(exponent) - if not np.isfinite(row).all(): - raise RuntimeError( - "ONNX classifier returned non-finite scores") - rows.append(row.astype(np.float32, copy=True)) - return np.stack(rows, axis=0) - - scores = await asyncio.to_thread(infer) - value = {"secure_kind": "classifier_scores.v1", "scores": scores} - return _sdk.ClassifierScoresRef._wrap(await rt.refs.create( - "CLASSIFIER_SCORES", value)) # type: ignore[return-value] - -async def classifier_scores_shape(scores: "ClassifierScoresRef", -) -> tuple[int, int]: - import numpy as np - - bundle = await _sdk.current_runtime().refs.resolve(scores) - value = bundle.get("scores") if isinstance(bundle, dict) else None - if (not isinstance(value, np.ndarray) or value.ndim != 2 - or not 1 <= value.shape[0] <= 64 - or not 1 <= value.shape[1] <= 16_384): - raise TypeError("CLASSIFIER_SCORES handle is invalid") - return int(value.shape[0]), int(value.shape[1]) - -async def classifier_scores_select_above(scores: "ClassifierScoresRef", batch_index: int, - start: int, end: int, threshold: float, - offset: int = 0, limit: int = 512, -) -> dict[str, Any]: - import math - import numpy as np - - bundle = await _sdk.current_runtime().refs.resolve(scores) - value = bundle.get("scores") if isinstance(bundle, dict) else None - if (not isinstance(value, np.ndarray) or value.ndim != 2 - or not 1 <= value.shape[0] <= 64 - or not 1 <= value.shape[1] <= 16_384): - raise TypeError("CLASSIFIER_SCORES handle is invalid") - batch_index = int(batch_index) - start = int(start) - end = int(end) - offset = int(offset) - limit = int(limit) - threshold = float(threshold) - if not 0 <= batch_index < value.shape[0]: - raise ValueError("classifier score batch index is invalid") - if not 0 <= start <= end <= value.shape[1]: - raise ValueError("classifier score class range is invalid") - if not math.isfinite(threshold) or abs(threshold) > 1_000_000: - raise ValueError("classifier score threshold is invalid") - if not 0 <= offset <= value.shape[1]: - raise ValueError("classifier score page offset is invalid") - if not 1 <= limit <= 512: - raise ValueError("classifier score page limit must be in [1, 512]") - matches = np.flatnonzero(value[batch_index, start:end] > threshold) - matches = matches.astype(np.int64, copy=False) + start - selected = matches[offset:offset + limit] - next_offset = offset + len(selected) - return { - "items": [{ - "index": int(index), - "score": float(value[batch_index, index]), - } for index in selected], - "next_offset": ( - next_offset if next_offset < len(matches) else None), - } - async def semantic_segmentation_mask(segmentation: "SemanticSegmentationRef", image: "ImageRef", classes: list[int], ) -> "MaskRef":