diff --git a/node/DLNode/node_object_detection.py b/node/DLNode/node_object_detection.py index 1eb1a4a8..c3b29ba5 100644 --- a/node/DLNode/node_object_detection.py +++ b/node/DLNode/node_object_detection.py @@ -44,6 +44,18 @@ os.path.dirname(os.path.abspath(__file__)), 'object_detection' ) + +def _resolve_builtin_model_path(*relative_candidates): + """Return the first existing built-in model path, else the first candidate.""" + if not relative_candidates: + return "" + for rel in relative_candidates: + candidate = os.path.join(_OBJECT_DETECTION_BASE, rel) + if os.path.isfile(candidate): + return candidate + return os.path.join(_OBJECT_DETECTION_BASE, relative_candidates[0]) + + # Directory where user-uploaded ONNX models are stored permanently if getattr(sys, 'frozen', False): from src.utils.paths import get_models_dir @@ -119,6 +131,19 @@ def build_class_names_from_source(source_label, num_classes): 'num_classes': 80, 'class_names': _COCO_CLASSES, }, + { + 'name': 'yolo8n_B', + 'path': _resolve_builtin_model_path( + os.path.join('YOLO', 'model', 'yolov8n.onnx'), + os.path.join('YOLO', 'model', 'yolo11_n.onnx'), + ), + 'output_format': 'yolo11', + 'input_width': 640, + 'input_height': 640, + 'num_classes': 80, + 'class_names': _COCO_CLASSES, + 'supports_batched_detection': True, + }, { 'name': 'YOLOv8m(640x640)', 'path': os.path.join(_OBJECT_DETECTION_BASE, 'YOLO', 'model', 'yolov8m.onnx'), @@ -128,6 +153,19 @@ def build_class_names_from_source(source_label, num_classes): 'num_classes': 80, 'class_names': _COCO_CLASSES, }, + { + 'name': 'yolo8s_B', + 'path': _resolve_builtin_model_path( + os.path.join('YOLO', 'model', 'yolov8s.onnx'), + os.path.join('YOLO', 'model', 'yolo11_n.onnx'), + ), + 'output_format': 'yolo11', + 'input_width': 640, + 'input_height': 640, + 'num_classes': 80, + 'class_names': _COCO_CLASSES, + 'supports_batched_detection': True, + }, { 'name': 'FreeYOLO-Nano(640x640)', 'path': os.path.join(_OBJECT_DETECTION_BASE, 'FreeYOLO', 'model', 'yolo_free_nano_640x640.onnx'), @@ -205,6 +243,11 @@ def get_class_rejection_dropdown_items(class_name_dict): for class_id in sorted(class_name_dict.keys())] +def get_batch_badge_label(is_batch_capable): + """Return the label shown beside the model combo for true batch-capable models.""" + return "B" if is_batch_capable else "" + + class FactoryNode: node_label = 'ObjectDetection' @@ -232,6 +275,7 @@ def add_node( node.tag_node_input_text_name = node.tag_node_name + ':' + node.TYPE_TEXT + ':Input02' node.tag_node_input_text_value_name = node.tag_node_name + ':' + node.TYPE_TEXT + ':Input02Value' + node.tag_node_batch_badge_value_name = node.tag_node_name + ':BatchBadgeValue' node.tag_node_input_float_name = node.tag_node_name + ':' + node.TYPE_FLOAT + ':Input03' @@ -270,6 +314,7 @@ def on_model_change(sender, app_data, user_data): dpg.configure_item(node.tag_delete_btn, enabled=not is_builtin) except Exception: pass + node._update_batch_badge(selected_model) node.tag_node_output_image_name = node.tag_node_name + ':' + node.TYPE_IMAGE + ':Output01' node.tag_node_output_image = node.tag_node_name + ':' + node.TYPE_IMAGE + ':Output01Value' @@ -440,13 +485,20 @@ def _on_class_source_change(sender, app_data, user_data): tag=node.tag_node_input_text_name, attribute_type=dpg.mvNode_Attr_Static, ): - dpg.add_combo( - list(node._model_class.keys()), - default_value=list(node._model_class.keys())[0], - width=small_window_w, - tag=node.tag_node_input_text_value_name, - callback=on_model_change, - ) + default_model = list(node._model_class.keys())[0] + with dpg.group(horizontal=True): + dpg.add_combo( + list(node._model_class.keys()), + default_value=default_model, + width=small_window_w - 32, + tag=node.tag_node_input_text_value_name, + callback=on_model_change, + ) + dpg.add_text( + get_batch_badge_label(node._model_batch_capable.get(default_model, False)), + tag=node.tag_node_batch_badge_value_name, + color=(120, 255, 120, 255), + ) # ---- Collapse / expand toggle for the settings section ---------- # The configuration widgets below (provider, score, reject, draw, @@ -686,6 +738,7 @@ class Node(Node): _model_class: dict = {} # name → CustomONNX factory callable _model_path_setting: dict = {} # name → onnx file path _model_class_name_list: dict = {} # name → {int_id: str_name} + _model_batch_capable: dict = {} # name → supports true batched OD inference _model_instance: dict = {} @@ -721,6 +774,16 @@ def _ensure_builtin_models(cls): if not os.path.isfile(path): logger.debug(f"[Builtin] Skipping '{name}' — ONNX file not found: {path}") continue + supports_batch = meta.get('supports_batched_detection', None) + if supports_batch is None: + supports_batch = False + try: + inspected = onnx_inspector.inspect_onnx_model(path) + supports_batch = bool(inspected.get('supports_batched_detection', False)) + except Exception as exc: + logger.warning(f"[Builtin] Could not inspect batch support for '{name}': {exc}") + else: + supports_batch = bool(supports_batch) # Registry always stores class_names with string keys entry = { 'name': name, @@ -730,6 +793,7 @@ def _ensure_builtin_models(cls): 'input_height': meta['input_height'], 'num_classes': meta['num_classes'], 'class_names': {str(k): v for k, v in meta['class_names'].items()}, + 'supports_batched_detection': supports_batch, } if meta.get('disable_optimizations'): entry['disable_optimizations'] = True @@ -777,14 +841,30 @@ def _load_custom_models_from_registry(cls): nanodet_reg_first = entry.get('nanodet_reg_first', None) if nanodet_reg_first is not None: nanodet_reg_first = bool(nanodet_reg_first) + supports_batch = entry.get('supports_batched_detection', None) + if supports_batch is None: + try: + meta = onnx_inspector.inspect_onnx_model(path) + supports_batch = bool(meta.get('supports_batched_detection', False)) + entry['supports_batched_detection'] = supports_batch + except Exception as exc: + logger.warning(f"Could not inspect batch support for '{name}': {exc}") + supports_batch = False + else: + try: + custom_models_registry.save_entry(entry) + except Exception as exc: + logger.warning(f"Could not persist batch support for '{name}': {exc}") cls._register_custom_model(name, path, class_names, output_fmt, in_w, in_h, disable_optimizations=disable_opt, - nanodet_reg_first=nanodet_reg_first) + nanodet_reg_first=nanodet_reg_first, + supports_batch=bool(supports_batch)) logger.info(f"Loaded model from registry: {name}") @classmethod def _register_custom_model(cls, name, path, class_names, output_fmt, in_w, in_h, - disable_optimizations=False, nanodet_reg_first=None): + disable_optimizations=False, nanodet_reg_first=None, + supports_batch=False): """Add a model to the class-level runtime dictionaries.""" def _make_factory(p, fmt, w, h, disable_opt, nd_reg_first): def factory(model_path, providers=None, provider_options=None): @@ -810,6 +890,7 @@ def factory(model_path, providers=None, provider_options=None): cls._model_class[name] = _make_factory(path, output_fmt, in_w, in_h, disable_optimizations, nanodet_reg_first) cls._model_path_setting[name] = path cls._model_class_name_list[name] = class_names + cls._model_batch_capable[name] = bool(supports_batch) # ------------------------------------------------------------------ # Upload callback @@ -839,7 +920,8 @@ def _callback_onnx_select(self, sender, data, user_data=None): f"[Upload] Inspection result: format='{meta.get('output_format')}', " f"input={meta.get('input_width')}x{meta.get('input_height')}, " f"num_classes={meta.get('num_classes')}, " - f"class_names_count={len(meta.get('class_names', {}))}" + f"class_names_count={len(meta.get('class_names', {}))}, " + f"supports_batch={meta.get('supports_batched_detection', False)}" ) except Exception as exc: logger.error(f"[Upload] ONNX inspection failed: {exc}", exc_info=True) @@ -926,6 +1008,16 @@ def _render_preview_details(self, meta, class_names): f"Number of classes: {num_cls}", parent=self.tag_preview_details, ) + dpg.add_text( + "Dynamic batch : " + + ("yes" if meta.get("has_dynamic_batch", False) else "no"), + parent=self.tag_preview_details, + ) + dpg.add_text( + "ObjectDetection B : " + + ("yes" if meta.get("supports_batched_detection", False) else "no"), + parent=self.tag_preview_details, + ) if class_names: dpg.add_text("Class list:", parent=self.tag_preview_details) @@ -1058,7 +1150,16 @@ def _finalise_upload(node, onnx_path: str, meta: dict, class_names: dict, custom f"input={in_w}x{in_h}, classes={num_classes}" ) - Node._register_custom_model(name, onnx_path, class_names, output_fmt, in_w, in_h) + supports_batch = bool(meta.get("supports_batched_detection", False)) + Node._register_custom_model( + name, + onnx_path, + class_names, + output_fmt, + in_w, + in_h, + supports_batch=supports_batch, + ) registry_entry = { "name": name, @@ -1068,6 +1169,7 @@ def _finalise_upload(node, onnx_path: str, meta: dict, class_names: dict, custom "input_width": in_w, "input_height": in_h, "num_classes": num_classes, + "supports_batched_detection": supports_batch, } try: custom_models_registry.save_entry(registry_entry) @@ -1082,6 +1184,7 @@ def _finalise_upload(node, onnx_path: str, meta: dict, class_names: dict, custom if name not in current_items: current_items = list(current_items) + [name] dpg.configure_item(model_combo_tag, items=current_items, default_value=name) + node._update_batch_badge(name) logger.info(f"[Upload] Model dropdown updated — '{name}' selected.") except Exception as exc: logger.warning(f"[Upload] Could not update model dropdown: {exc}") @@ -1131,6 +1234,7 @@ def _delete_custom_model(cls, name: str) -> bool: cls._model_class.pop(name, None) cls._model_path_setting.pop(name, None) cls._model_class_name_list.pop(name, None) + cls._model_batch_capable.pop(name, None) # Drop any cached inference instances for this model (name_provider keys). for key in [k for k in cls._model_instance if k == name or k.startswith(name + '_')]: cls._model_instance.pop(key, None) @@ -1154,6 +1258,7 @@ def _delete_selected_model(self, name: str): new_default = remaining[0] if remaining else "" try: dpg.configure_item(model_combo_tag, items=remaining, default_value=new_default) + self._update_batch_badge(new_default) logger.info(f"[Delete] Model '{name}' deleted — '{new_default}' now selected.") except Exception as exc: logger.warning(f"[Delete] Could not update model dropdown: {exc}") @@ -1252,6 +1357,16 @@ def _build_trt_provider_options(batch_size=1): "trt_max_workspace_size": str(workspace_bytes), } + def _update_batch_badge(self, model_name): + """Refresh the small 'B' marker beside the model combo.""" + try: + dpg_set_value( + self.tag_node_batch_badge_value_name, + get_batch_badge_label(self._model_batch_capable.get(model_name, False)), + ) + except Exception as exc: + logger.debug(f"Could not update batch badge for '{model_name}': {exc}") + @staticmethod def _run_batch_inference(model_instance, frames): """Run inference on a list of frames, returning one result per frame. @@ -1681,6 +1796,7 @@ def set_setting_dict(self, node_id, setting_dict): entry.get('output_format', 'yolo11'), int(entry.get('input_width', 640)), int(entry.get('input_height', 640)), + supports_batch=bool(entry.get('supports_batched_detection', False)), ) logger.info(f"Restored custom model from registry on set_setting_dict: {model_name}") else: @@ -1696,6 +1812,7 @@ def set_setting_dict(self, node_id, setting_dict): pass dpg_set_value(self.tag_node_input_text_value_name, model_name) + self._update_batch_badge(model_name) dpg_set_value(self.tag_node_input_float_value_name, score_th) # Update the dropdown items to match the loaded model's classes diff --git a/node/DLNode/object_detection/custom_models_registry.py b/node/DLNode/object_detection/custom_models_registry.py index d2e17e2a..6cec29ee 100644 --- a/node/DLNode/object_detection/custom_models_registry.py +++ b/node/DLNode/object_detection/custom_models_registry.py @@ -17,6 +17,7 @@ "input_width": int "input_height": int "num_classes": int + "supports_batched_detection": bool - true when ObjectDetection can batch it } """ diff --git a/node/DLNode/object_detection/onnx_inspector.py b/node/DLNode/object_detection/onnx_inspector.py index af66c8a3..45f20b4b 100644 --- a/node/DLNode/object_detection/onnx_inspector.py +++ b/node/DLNode/object_detection/onnx_inspector.py @@ -20,6 +20,25 @@ logger = logging.getLogger(__name__) +def has_dynamic_batch_dim(input_shape: list) -> bool: + """True when the first input dimension is dynamic/symbolic.""" + if not input_shape: + return False + batch_dim = input_shape[0] + return ( + batch_dim is None + or isinstance(batch_dim, str) + or (isinstance(batch_dim, int) and batch_dim < 0) + ) + + +def supports_batched_detection(input_shape: list, output_format: str) -> bool: + """True when ObjectDetection can run a real batched ONNX call for this model.""" + if not has_dynamic_batch_dim(input_shape): + return False + return output_format in {"yolo11", "yolo11_obb", "yolox", "nanodet"} + + def _dim_value(dim) -> object: """Extract an int value from an ONNX TensorShapeProto.Dimension, or a string param name.""" if dim.HasField("dim_value"): @@ -201,6 +220,8 @@ def _inspect_onnx_static(model_path: str) -> dict: "class_names": class_names, "input_width": input_width, "input_height": input_height, + "has_dynamic_batch": has_dynamic_batch_dim(input_shape), + "supports_batched_detection": supports_batched_detection(input_shape, output_format), } logger.info( f"[ONNX Inspector/static] Inspection complete — " @@ -500,6 +521,8 @@ def inspect_onnx_model(model_path: str) -> dict: "class_names": class_names, "input_width": input_width, "input_height": input_height, + "has_dynamic_batch": has_dynamic_batch_dim(input_shape), + "supports_batched_detection": supports_batched_detection(input_shape, output_format), } logger.info( f"[ONNX Inspector] Inspection complete — " diff --git a/tests/test_model_class_dropdown_link.py b/tests/test_model_class_dropdown_link.py index c03a2d63..466c1361 100644 --- a/tests/test_model_class_dropdown_link.py +++ b/tests/test_model_class_dropdown_link.py @@ -5,11 +5,24 @@ import pytest import sys import os +import importlib.util # Add parent directory to path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from node.DLNode.object_detection.coco_class_names import coco_class_names + +_COCO_PATH = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + 'node', 'DLNode', 'object_detection', 'coco_class_names.py' +) +_COCO_SPEC = importlib.util.spec_from_file_location('_test_coco_class_names', _COCO_PATH) +_COCO_MODULE = importlib.util.module_from_spec(_COCO_SPEC) +_COCO_SPEC.loader.exec_module(_COCO_MODULE) +coco_class_names = _COCO_MODULE.coco_class_names +_OD_PATH = os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + 'node', 'DLNode', 'node_object_detection.py' +) def test_coco_class_names(): @@ -21,16 +34,8 @@ def test_coco_class_names(): def test_get_class_rejection_dropdown_items(): """Test the function that generates dropdown items from class names""" - import unittest.mock as mock - for mod in ('dearpygui', 'dearpygui.dearpygui', 'node_editor', 'node_editor.util', - 'src', 'src.utils', 'src.utils.logging', 'src.utils.gpu_utils', - 'node.DLNode.object_detection.CustomONNX', - 'node.DLNode.object_detection.CustomONNX.custom_onnx', - 'node.DLNode.object_detection.custom_models_registry', - 'node.DLNode.object_detection.onnx_inspector'): - sys.modules.setdefault(mod, mock.MagicMock()) - - from node.DLNode.node_object_detection import get_class_rejection_dropdown_items + module = _load_od_module() + get_class_rejection_dropdown_items = module.get_class_rejection_dropdown_items # Test with COCO classes coco_items = get_class_rejection_dropdown_items(coco_class_names) @@ -53,26 +58,96 @@ def test_get_class_rejection_dropdown_items(): assert person_items[0] == "0: person" -def _mock_dpg_modules(): - """Helper: install mocks for dpg and related modules.""" +def test_get_batch_badge_label(): + """Only true batch-capable models should show the B marker.""" + module = _load_od_module() + get_batch_badge_label = module.get_batch_badge_label + + assert get_batch_badge_label(True) == "B" + assert get_batch_badge_label(False) == "" + + +def _load_od_module(): + """Load node_object_detection with lightweight stubs.""" import unittest.mock as mock - for mod in ('dearpygui', 'dearpygui.dearpygui', 'node_editor', 'node_editor.util', - 'src', 'src.utils', 'src.utils.logging', 'src.utils.gpu_utils', - 'node.DLNode.object_detection.CustomONNX', - 'node.DLNode.object_detection.CustomONNX.custom_onnx', - 'node.DLNode.object_detection.custom_models_registry', - 'node.DLNode.object_detection.onnx_inspector'): - sys.modules.setdefault(mod, mock.MagicMock()) + import types + + dpg_mock = mock.MagicMock() + mocked = { + 'cv2': mock.MagicMock(), + 'numpy': mock.MagicMock(), + 'onnxruntime': mock.MagicMock(), + 'dearpygui': types.ModuleType('dearpygui'), + 'dearpygui.dearpygui': dpg_mock, + 'node_editor': types.ModuleType('node_editor'), + 'node_editor.util': types.SimpleNamespace( + dpg_get_value=mock.MagicMock(), + dpg_set_value=mock.MagicMock(), + ), + 'node.basenode': types.SimpleNamespace(Node=type('BaseNode', (), {})), + 'src': types.ModuleType('src'), + 'src.utils': types.ModuleType('src.utils'), + 'src.utils.logging': types.SimpleNamespace(get_logger=lambda name: mock.MagicMock()), + 'src.utils.gpu_utils': types.SimpleNamespace(get_execution_providers=lambda: ['CPUExecutionProvider']), + 'node.DLNode.object_detection': types.ModuleType('node.DLNode.object_detection'), + 'node.DLNode.object_detection.coco_class_names': _COCO_MODULE, + 'node.DLNode.object_detection.BlazeFace': types.ModuleType('node.DLNode.object_detection.BlazeFace'), + 'node.DLNode.object_detection.BlazeFace.blazeface': types.SimpleNamespace(BlazeFace=mock.MagicMock()), + 'node.DLNode.object_detection.CustomONNX': mock.MagicMock(), + 'node.DLNode.object_detection.CustomONNX.custom_onnx': types.SimpleNamespace(CustomONNX=mock.MagicMock()), + 'node.DLNode.object_detection.custom_models_registry': types.SimpleNamespace( + load_registry=lambda: [], + save_entry=mock.MagicMock(), + remove_entry=mock.MagicMock(), + get_entry=lambda name: None, + ), + 'node.DLNode.object_detection.onnx_inspector': types.SimpleNamespace( + inspect_onnx_model=lambda path: { + 'output_format': 'yolo11', + 'input_width': 640, + 'input_height': 640, + 'num_classes': 80, + 'class_names': {}, + 'supports_batched_detection': False, + 'has_dynamic_batch': False, + } + ), + } + mocked['dearpygui'].dearpygui = dpg_mock + mocked['node_editor'].util = mocked['node_editor.util'] + mocked['src'].utils = mocked['src.utils'] + mocked['node.DLNode.object_detection'].coco_class_names = _COCO_MODULE + mocked['node.DLNode.object_detection'].onnx_inspector = mocked['node.DLNode.object_detection.onnx_inspector'] + mocked['node.DLNode.object_detection'].custom_models_registry = mocked['node.DLNode.object_detection.custom_models_registry'] + + saved = {} + for name, mod in mocked.items(): + saved[name] = sys.modules.get(name) + sys.modules[name] = mod + + spec = importlib.util.spec_from_file_location('_od_dropdown_test', _OD_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + for name, orig in saved.items(): + if orig is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = orig + + return module def test_builtin_models_defined(): """Test that _BUILTIN_MODELS contains expected entries.""" - _mock_dpg_modules() - from node.DLNode.node_object_detection import _BUILTIN_MODELS + module = _load_od_module() + _BUILTIN_MODELS = module._BUILTIN_MODELS names = {m['name'] for m in _BUILTIN_MODELS} assert 'YOLOX-Nano(416x416)' in names assert 'YOLO11Nano' in names + assert 'yolo8n_B' in names + assert 'yolo8s_B' in names assert 'YOLOTENNIS' in names assert 'Light-Weight Person Detector' in names @@ -90,14 +165,15 @@ def test_builtin_models_defined(): assert tennis['class_names'].get(1) == 'player2' assert tennis['class_names'].get(2) == 'ball' + for name in ('yolo8n_B', 'yolo8s_B'): + model = next((m for m in _BUILTIN_MODELS if m['name'] == name), None) + assert model is not None, f"{name} should be in _BUILTIN_MODELS" + assert model.get('supports_batched_detection') is True + def test_callback_function_exists(): """Test that on_model_change callback and related helpers are present.""" - file_path = os.path.join( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - 'node', 'DLNode', 'node_object_detection.py' - ) - with open(file_path, 'r') as f: + with open(_OD_PATH, 'r') as f: content = f.read() assert 'def on_model_change' in content, "Should have on_model_change callback" @@ -108,16 +184,35 @@ def test_callback_function_exists(): def test_rejected_classes_cleared_on_model_change(): """Test that rejected classes are cleared when model changes.""" - file_path = os.path.join( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - 'node', 'DLNode', 'node_object_detection.py' - ) - with open(file_path, 'r') as f: + with open(_OD_PATH, 'r') as f: content = f.read() assert 'dpg_set_value(node.tag_node_rejected_classes_value_name, "")' in content, \ "Should clear rejected classes when model changes" +def test_update_batch_badge_uses_batch_capability_map(): + """The runtime badge should reflect the selected model batch capability.""" + module = _load_od_module() + node = module.Node() + node.tag_node_batch_badge_value_name = 'badge' + node._model_batch_capable = {'dynamic': True, 'static': False} + + calls = [] + + def fake_set_value(tag, value): + calls.append((tag, value)) + + original = module.dpg_set_value + module.dpg_set_value = fake_set_value + try: + node._update_batch_badge('dynamic') + node._update_batch_badge('static') + finally: + module.dpg_set_value = original + + assert calls == [('badge', 'B'), ('badge', '')] + + if __name__ == '__main__': pytest.main([__file__, '-v']) diff --git a/tests/test_onnx_batch_metadata.py b/tests/test_onnx_batch_metadata.py new file mode 100644 index 00000000..d090718a --- /dev/null +++ b/tests/test_onnx_batch_metadata.py @@ -0,0 +1,75 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- + +import importlib.util +import os +import sys +import types +from unittest import mock + +import pytest + + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +INSPECTOR_PATH = os.path.join( + REPO_ROOT, 'node', 'DLNode', 'object_detection', 'onnx_inspector.py' +) + + +def _load_inspector_module(): + mocked = { + 'onnxruntime': mock.MagicMock(), + 'node.DLNode.object_detection.onnx_session_utils': types.SimpleNamespace( + make_session=mock.MagicMock() + ), + } + + saved = {} + for name, mod in mocked.items(): + saved[name] = sys.modules.get(name) + sys.modules[name] = mod + + spec = importlib.util.spec_from_file_location('_onnx_batch_meta_test', INSPECTOR_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + + for name, orig in saved.items(): + if orig is None: + sys.modules.pop(name, None) + else: + sys.modules[name] = orig + + return module + + +INSPECTOR = _load_inspector_module() + + +@pytest.mark.parametrize( + ("input_shape", "expected"), + [ + ([1, 3, 640, 640], False), + (["batch", 3, 640, 640], True), + ([-1, 3, 640, 640], True), + ([0, 3, 640, 640], False), + ([None, 3, 640, 640], True), + ([], False), + ], +) +def test_has_dynamic_batch_dim(input_shape, expected): + assert INSPECTOR.has_dynamic_batch_dim(input_shape) is expected + + +@pytest.mark.parametrize( + ("input_shape", "output_format", "expected"), + [ + (["batch", 3, 640, 640], "yolo11", True), + (["batch", 3, 640, 640], "yolox", True), + (["batch", 3, 640, 640], "nanodet", True), + (["batch", 3, 640, 640], "ssd", False), + ([1, 3, 640, 640], "yolo11", False), + (["batch", 3, 640, 640], "unknown", False), + ], +) +def test_supports_batched_detection(input_shape, output_format, expected): + assert INSPECTOR.supports_batched_detection(input_shape, output_format) is expected