node_classification.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
import copy
import shutil
import time
import os
import sys

import numpy as np
import cv2
import dearpygui.dearpygui as dpg

from node_editor.util import dpg_get_value, dpg_set_value

from node.node_abc import DpgNodeABC

from node.DLNode.classification.MobileNetV3.mobilenet_v3 import MobileNetV3
from node.DLNode.classification.EfficientNetB0.efficientnet import EfficientNetB0
from node.DLNode.classification.ResNet50.resnet50 import ResNet50

# Import YoloCls using importlib.util due to hyphenated directory name
import importlib.util
_yolo_cls_init_path = os.path.join(
    os.path.dirname(os.path.abspath(__file__)),
    'classification', 'Yolo-cls', '__init__.py'
)
_yolo_cls_spec = importlib.util.spec_from_file_location(
    'yolo_cls_init_module', _yolo_cls_init_path
)
_yolo_cls_module = importlib.util.module_from_spec(_yolo_cls_spec)
_yolo_cls_spec.loader.exec_module(_yolo_cls_module)
YoloCls = _yolo_cls_module.YoloCls

from node.DLNode.classification.imagenet_class_names import imagenet_class_names
from node.DLNode.classification.esc50_class_names import esc50_class_names
from node.DLNode.classification.CustomONNX.custom_onnx import CustomONNX as CustomONNXClassification
from node.DLNode.classification import custom_models_registry as _cls_registry
from node.DLNode.object_detection import onnx_inspector

from node.basenode import Node

from src.utils.logging import get_logger
logger = get_logger(__name__)

if getattr(sys, 'frozen', False):
    from src.utils.paths import get_models_dir
    _CLS_UPLOADS_DIR = get_models_dir('classification')
else:
    _CLS_UPLOADS_DIR = os.path.join(
        os.path.dirname(os.path.abspath(__file__)), 'classification', 'CustomONNX', 'models'
    )


class FactoryNode:
    node_label = 'Classification'
    node_tag = 'Classification'
    

    def __init__(self):
        pass

    
    def add_node(self, parent, node_id, pos=[0, 0], callback=None, opencv_setting_dict=None):
        """Adds a node to the processing graph."""
        node = Node()
        
        # タグ名
        tag_node_name = str(node_id) + ':' + node.node_tag
        tag_node_input01_name = tag_node_name + ':' + node.TYPE_IMAGE + ':Input01'
        tag_node_input01_value_name = tag_node_name + ':' + node.TYPE_IMAGE + ':Input01Value'
        tag_node_input02_name = tag_node_name + ':' + node.TYPE_TEXT + ':Input02'
        tag_node_input02_value_name = tag_node_name + ':' + node.TYPE_TEXT + ':Input02Value'
        tag_node_output01_name = tag_node_name + ':' + node.TYPE_IMAGE + ':Output01'
        tag_node_output01_value_name = tag_node_name + ':' + node.TYPE_IMAGE + ':Output01Value'
        tag_node_output02_name = tag_node_name + ':' + node.TYPE_TIME_MS + ':Output02'
        tag_node_output02_value_name = tag_node_name + ':' + node.TYPE_TIME_MS + ':Output02Value'
        tag_node_output_json_name = tag_node_name + ':' + node.TYPE_JSON + ':OutputJson'
        tag_node_output_json_value_name = tag_node_name + ':' + node.TYPE_JSON + ':OutputJsonValue'

        tag_provider_select_name = tag_node_name + ':' + node.TYPE_TEXT + ':Provider'
        tag_provider_select_value_name = tag_node_name + ':' + node.TYPE_IMAGE + ':ProviderValue'

        # OpenCV向け設定
        node._opencv_setting_dict = opencv_setting_dict
        small_window_w = node._opencv_setting_dict['process_width']
        small_window_h = node._opencv_setting_dict['process_height']
        use_pref_counter = node._opencv_setting_dict['use_pref_counter']
        use_gpu = node._opencv_setting_dict['use_gpu']

        # 初期化用黒画像
        black_image = np.zeros((small_window_w, small_window_h, 3))
        black_texture = node.convert_cv_to_dpg(
            black_image,
            small_window_w,
            small_window_h,
        )

        # テクスチャ登録
        with dpg.texture_registry(show=False):
            dpg.add_raw_texture(
                small_window_w,
                small_window_h,
                black_texture,
                tag=tag_node_output01_value_name,
                format=dpg.mvFormat_Float_rgb,
            )

        # Create yellow theme for JSON button
        with dpg.theme() as yellow_button_theme:
            with dpg.theme_component(dpg.mvButton):
                dpg.add_theme_color(dpg.mvThemeCol_Button, (255, 255, 153, 255))
                dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, (255, 255, 153, 255))
                dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, (255, 255, 153, 255))

        # ---- ONNX upload file dialog ----------------------------------------
        onnx_file_dialog_tag = "cls_onnx_select:" + str(node_id)
        with dpg.file_dialog(
            directory_selector=False,
            show=False,
            modal=True,
            height=400,
            callback=node._callback_onnx_select,
            tag=onnx_file_dialog_tag,
        ):
            dpg.add_file_extension("ONNX (*.onnx){.onnx}")
            dpg.add_file_extension("", color=(150, 255, 150, 255))
        node.tag_upload_file_dialog = onnx_file_dialog_tag

        # ---- ONNX preview / confirmation dialog -----------------------------
        preview_window_tag  = "cls_onnx_preview_window:"  + str(node_id)
        preview_name_tag    = "cls_onnx_preview_name:"    + str(node_id)
        preview_details_tag = "cls_onnx_preview_details:" + str(node_id)
        preview_status_tag  = "cls_onnx_preview_status:"  + str(node_id)
        preview_confirm_tag = "cls_onnx_preview_confirm:" + str(node_id)
        preview_cancel_tag  = "cls_onnx_preview_cancel:"  + str(node_id)
        preview_quit_tag    = "cls_onnx_preview_quit:"    + str(node_id)

        node.tag_preview_window  = preview_window_tag
        node.tag_preview_name    = preview_name_tag
        node.tag_preview_details = preview_details_tag
        node.tag_preview_status  = preview_status_tag
        node.tag_preview_confirm = preview_confirm_tag
        node.tag_preview_cancel  = preview_cancel_tag
        node.tag_preview_quit    = preview_quit_tag

        def _on_upload_confirm(sender, app_data, user_data):
            node._do_confirm_upload()

        def _on_close_preview(sender, app_data, user_data):
            node._close_upload_preview()

        with dpg.window(
            label="ONNX Model Preview",
            tag=preview_window_tag,
            modal=True,
            show=False,
            width=430,
            no_close=True,
        ):
            dpg.add_text("Model name (editable):")
            dpg.add_input_text(tag=preview_name_tag, width=410)
            dpg.add_separator()
            dpg.add_group(tag=preview_details_tag)
            dpg.add_separator()
            dpg.add_text("", tag=preview_status_tag)
            dpg.add_spacer(height=4)
            with dpg.group(horizontal=True):
                dpg.add_button(label="  Confirm Upload  ", tag=preview_confirm_tag, callback=_on_upload_confirm)
                dpg.add_spacer(width=10)
                dpg.add_button(label="  Cancel  ", tag=preview_cancel_tag, callback=_on_close_preview)
                dpg.add_spacer(width=10)
                dpg.add_button(label="  Quit  ", tag=preview_quit_tag, callback=_on_close_preview, show=False)

        # ノード
        with dpg.node(
                tag=tag_node_name,
                parent=parent,
                label=node.node_label,
                pos=pos,
        ):
            # 入力端子
            with dpg.node_attribute(
                    tag=tag_node_input01_name,
                    attribute_type=dpg.mvNode_Attr_Input,
            ):
                dpg.add_text(
                    tag=tag_node_input01_value_name,
                    default_value='Input BGR image',
                )
            # 画像
            with dpg.node_attribute(
                    tag=tag_node_output01_name,
                    attribute_type=dpg.mvNode_Attr_Output,
            ):
                dpg.add_image(tag_node_output01_value_name)
            # 使用アルゴリズム
            with dpg.node_attribute(
                    tag=tag_node_input02_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=tag_node_input02_value_name,
                )
            if use_gpu:
                # CPU/GPU切り替え
                with dpg.node_attribute(
                        tag=tag_provider_select_name,
                        attribute_type=dpg.mvNode_Attr_Static,
                ):
                    dpg.add_radio_button(
                        ("CPU", "GPU"),
                        tag=tag_provider_select_value_name,
                        default_value='CPU',
                        horizontal=True,
                    )
            # 処理時間
            if use_pref_counter:
                with dpg.node_attribute(
                        tag=tag_node_output02_name,
                        attribute_type=dpg.mvNode_Attr_Output,
                ):
                    dpg.add_text(
                        tag=tag_node_output02_value_name,
                        default_value='elapsed time(ms)',
                    )

            # JSON output button
            with dpg.node_attribute(
                    tag=tag_node_output_json_name,
                    attribute_type=dpg.mvNode_Attr_Output,
            ):
                btn = dpg.add_button(
                    label="JSON",
                    tag=tag_node_output_json_value_name,
                    width=small_window_w,
                    enabled=False,
                )
                dpg.bind_item_theme(btn, yellow_button_theme)

            # ---- Add Model button (yellow, opens ONNX upload dialog) --------
            node.tag_upload_btn = tag_node_name + ':UploadONNX'

            def _on_upload_clicked(sender, app_data, user_data):
                dpg.show_item(onnx_file_dialog_tag)

            with dpg.node_attribute(
                    tag=tag_node_name + ':UploadAttr',
                    attribute_type=dpg.mvNode_Attr_Static,
            ):
                with dpg.theme() as add_model_btn_theme:
                    with dpg.theme_component(dpg.mvButton):
                        dpg.add_theme_color(dpg.mvThemeCol_Button, (255, 220, 0, 255))
                        dpg.add_theme_color(dpg.mvThemeCol_ButtonHovered, (255, 235, 50, 255))
                        dpg.add_theme_color(dpg.mvThemeCol_ButtonActive, (220, 190, 0, 255))
                        dpg.add_theme_color(dpg.mvThemeCol_Text, (0, 0, 0, 255))

                add_model_btn = dpg.add_button(
                    label=u"Add Model",
                    tag=node.tag_upload_btn,
                    width=small_window_w,
                    callback=_on_upload_clicked,
                )
                dpg.bind_item_theme(add_model_btn, add_model_btn_theme)

        node.tag_node_name = tag_node_name
        return node



class Node(Node):
    _ver = '0.0.1'

    node_label = 'Classification'
    node_tag = 'Classification'

    _opencv_setting_dict = None

    # モデル設定
    _model_class = {
        'MobileNetV3 Small': MobileNetV3,
        'MobileNetV3 Large': MobileNetV3,
        'EfficientNet B0': EfficientNetB0,
        'ResNet50': ResNet50,
        'Yolo-cls': YoloCls,
    }
    _model_base_path = os.path.dirname(os.path.abspath(__file__)) + '/classification/'
    _model_path_setting = {
        'MobileNetV3 Small':
        _model_base_path + 'MobileNetV3/model/MobileNetV3Small.onnx',
        'MobileNetV3 Large':
        _model_base_path + 'MobileNetV3/model/MobileNetV3Large.onnx',
        'EfficientNet B0':
        _model_base_path + 'EfficientNetB0/model/EfficientNetB0.onnx',
        'ResNet50':
        _model_base_path + 'ResNet50/model/ResNet50.onnx',
        'Yolo-cls':
        _model_base_path + 'Yolo-cls/model/son.onnx',
    }
    _model_class_name_dict = {
        'MobileNetV3 Small': imagenet_class_names,
        'MobileNetV3 Large': imagenet_class_names,
        'EfficientNet B0': imagenet_class_names,
        'ResNet50': imagenet_class_names,
        'Yolo-cls': esc50_class_names,
    }

    _model_instance = {}
    _class_name_dict = None

    def __init__(self):
        pass

    @classmethod
    def _load_custom_models_from_registry(cls):
        """Load user-uploaded models from the registry into the class dicts."""
        try:
            entries = _cls_registry.load_registry()
        except Exception as exc:
            logger.warning(f"[Classification] Could not load custom models registry: {exc}")
            return
        for entry in entries:
            name = entry.get('name', '')
            path = entry.get('path', '')
            if not name or not path:
                continue
            if name in cls._model_class:
                continue
            in_w = int(entry.get('input_width', 224))
            in_h = int(entry.get('input_height', 224))
            num_classes = int(entry.get('num_classes', 0))
            raw_class_names = entry.get('class_names', {})
            class_names = {int(k): str(v) for k, v in raw_class_names.items()} if raw_class_names else {}
            if not class_names and num_classes > 0:
                class_names = {i: f"class_{i}" for i in range(num_classes)}
            cls._register_custom_model(name, path, in_w, in_h, class_names)
            logger.info(f"[Classification] Loaded custom model from registry: {name}")

    @classmethod
    def _register_custom_model(cls, name, path, in_w, in_h, class_names=None):
        """Add a custom ONNX model to the class-level runtime dictionaries."""
        if class_names is None:
            class_names = {}

        def _make_factory(p, w, h):
            def factory(model_path, providers=None):
                if providers is None:
                    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
                return CustomONNXClassification(
                    model_path=p, input_width=w, input_height=h, providers=providers
                )
            return factory

        cls._model_class[name] = _make_factory(path, in_w, in_h)
        cls._model_path_setting[name] = path
        cls._model_class_name_dict[name] = class_names if class_names else {0: 'class_0'}

    # ------------------------------------------------------------------
    # Upload callbacks
    # ------------------------------------------------------------------

    def _callback_onnx_select(self, sender, data, user_data=None):
        """Handle ONNX file selection from the file dialog."""
        if data.get("file_name") == ".":
            return
        onnx_path = data.get("file_path_name", "")
        if not onnx_path or not os.path.isfile(onnx_path):
            return

        try:
            meta = onnx_inspector.inspect_onnx_model(onnx_path)
        except Exception as exc:
            logger.error(f"[Classification Upload] ONNX inspection failed: {exc}")
            self._pending_onnx_path = None
            self._pending_meta = None
            try:
                dpg.delete_item(self.tag_preview_details, children_only=True)
                dpg.add_text(
                    f"Inspection error: {exc}",
                    parent=self.tag_preview_details,
                    color=(255, 100, 100, 255),
                )
                self._set_upload_preview_actions(upload_succeeded=False)
                dpg.set_value(self.tag_preview_status, "")
                dpg.show_item(self.tag_preview_window)
            except Exception:
                pass
            return

        self._pending_onnx_path = onnx_path
        self._pending_meta = meta

        base_name = os.path.splitext(os.path.basename(onnx_path))[0]
        dpg.set_value(self.tag_preview_name, base_name)

        dpg.delete_item(self.tag_preview_details, children_only=True)
        in_w = meta.get("input_width", 224)
        in_h = meta.get("input_height", 224)
        num_cls = meta.get("num_classes", 0)
        class_names = meta.get("class_names", {})
        dpg.add_text(f"Input dimensions : {in_w} x {in_h} px (W x H)", parent=self.tag_preview_details)
        dpg.add_text(f"Number of classes: {num_cls}", parent=self.tag_preview_details)
        if class_names:
            max_show = 10
            dpg.add_text("Class list (first 10):", parent=self.tag_preview_details)
            for cid, cname in sorted(class_names.items(), key=lambda x: x[0])[:max_show]:
                dpg.add_text(f"  {cid}: {cname}", parent=self.tag_preview_details)

        self._set_upload_preview_actions(upload_succeeded=False)
        dpg.set_value(self.tag_preview_status, "")
        dpg.show_item(self.tag_preview_window)

    def _set_upload_preview_actions(self, upload_succeeded: bool):
        dpg.configure_item(self.tag_preview_confirm, show=not upload_succeeded)
        dpg.configure_item(self.tag_preview_cancel, show=not upload_succeeded)
        dpg.configure_item(self.tag_preview_quit, show=upload_succeeded)

    def _close_upload_preview(self):
        dpg.hide_item(self.tag_preview_window)

    def _do_confirm_upload(self):
        onnx_path = getattr(self, '_pending_onnx_path', None)
        meta = getattr(self, '_pending_meta', None)
        if not onnx_path or meta is None:
            dpg.set_value(self.tag_preview_status, "No pending upload — please select a file first.")
            return

        custom_name = dpg.get_value(self.tag_preview_name).strip()
        if not custom_name:
            custom_name = os.path.splitext(os.path.basename(onnx_path))[0]

        os.makedirs(_CLS_UPLOADS_DIR, exist_ok=True)
        dest_path = onnx_path
        try:
            basename = os.path.basename(onnx_path)
            candidate = os.path.join(_CLS_UPLOADS_DIR, basename)
            if os.path.abspath(onnx_path) != os.path.abspath(candidate):
                shutil.copy2(onnx_path, candidate)
                dest_path = candidate
        except Exception as exc:
            logger.warning(f"[Classification Upload] Could not copy ONNX: {exc}")

        try:
            Node._finalise_upload(self, dest_path, meta, custom_name=custom_name)
            dpg.set_value(self.tag_preview_status, f"\u2713 Model '{custom_name}' uploaded successfully!")
            self._set_upload_preview_actions(upload_succeeded=True)
        except Exception as exc:
            logger.error(f"[Classification Upload] Finalise failed: {exc}", exc_info=True)
            dpg.set_value(self.tag_preview_status, f"\u2717 Upload failed: {exc}")
            self._set_upload_preview_actions(upload_succeeded=False)

        self._pending_onnx_path = None
        self._pending_meta = None

    @staticmethod
    def _finalise_upload(node, onnx_path: str, meta: dict, custom_name: str = None):
        base = custom_name if custom_name else os.path.splitext(os.path.basename(onnx_path))[0]
        name = base
        counter = 1
        while name in Node._model_class:
            name = f"{base}_{counter}"
            counter += 1

        in_w = meta.get("input_width", 224)
        in_h = meta.get("input_height", 224)
        num_classes = meta.get("num_classes", 0)
        class_names = meta.get("class_names", {})
        if not class_names and num_classes > 0:
            class_names = {i: f"class_{i}" for i in range(num_classes)}

        Node._register_custom_model(name, onnx_path, in_w, in_h, class_names)

        registry_entry = {
            "name": name,
            "path": onnx_path,
            "input_width": in_w,
            "input_height": in_h,
            "output_format": meta.get("output_format", "unknown"),
            "num_classes": num_classes,
            "class_names": {str(k): v for k, v in class_names.items()},
        }
        try:
            _cls_registry.save_entry(registry_entry)
        except Exception as exc:
            logger.warning(f"[Classification Upload] Could not save registry entry: {exc}")

        model_combo_tag = node.tag_node_name + ':' + node.TYPE_TEXT + ':Input02Value'
        try:
            current_items = dpg.get_item_configuration(model_combo_tag).get("items", [])
            if name not in current_items:
                current_items = list(current_items) + [name]
            dpg.configure_item(model_combo_tag, items=current_items, default_value=name)
        except Exception as exc:
            logger.warning(f"[Classification Upload] Could not update model dropdown: {exc}")

    def update(
        self,
        node_id,
        connection_list,
        node_image_dict,
        node_result_dict,
        node_audio_dict,
    ):
        tag_node_name = str(node_id) + ':' + self.node_tag
        input_value02_tag = tag_node_name + ':' + self.TYPE_TEXT + ':Input02Value'
        output_value01_tag = tag_node_name + ':' + self.TYPE_IMAGE + ':Output01Value'
        output_value02_tag = tag_node_name + ':' + self.TYPE_TIME_MS + ':Output02Value'

        tag_provider_select_value_name = tag_node_name + ':' + self.TYPE_IMAGE + ':ProviderValue'

        small_window_w = self._opencv_setting_dict['process_width']
        small_window_h = self._opencv_setting_dict['process_height']
        use_pref_counter = self._opencv_setting_dict['use_pref_counter']
        use_gpu = self._opencv_setting_dict['use_gpu']

        # 接続情報確認
        src_node_name = ''
        connection_info_src = ''
        for connection_info in connection_list:
            connection_type = connection_info[0].split(':')[2]
            if connection_type == self.TYPE_INT:
                # 接続タグ取得
                source_tag = connection_info[0] + 'Value'
                destination_tag = connection_info[1] + 'Value'
                # 値更新
                input_value = int(dpg_get_value(source_tag))
                input_value = max([self._min_val, input_value])
                input_value = min([self._max_val, input_value])
                dpg_set_value(destination_tag, input_value)
            if connection_type == self.TYPE_IMAGE or connection_type == self.TYPE_AUDIO:
                # 画像取得元のノード名(ID付き)を取得
                connection_info_src = connection_info[0]
                connection_info_src = connection_info_src.split(':')[:2]
                src_node_name = connection_info_src[1]
                connection_info_src = ':'.join(connection_info_src)

        # 画像取得
        frame = self.get_input_frame(connection_list, node_image_dict, node_audio_dict)

        # CPU/GPU選択状態取得
        provider = 'CPU'
        if use_gpu:
            provider = dpg_get_value(tag_provider_select_value_name)

        # モデル情報取得
        model_name = dpg_get_value(input_value02_tag)
        model_path = self._model_path_setting[model_name]
        model_class = self._model_class[model_name]

        class_name_dict = self._model_class_name_dict[model_name]

        model_name_with_provider = model_name + '_' + provider

        # モデル取得
        if frame is not None:
            if model_name_with_provider not in self._model_instance:
                if provider == 'CPU':
                    providers = ['CPUExecutionProvider']
                    self._model_instance[
                        model_name_with_provider] = model_class(
                            model_path,
                            providers=providers,
                        )
                else:
                    self._model_instance[
                        model_name_with_provider] = model_class(model_path)

        # 計測開始
        if frame is not None and use_pref_counter:
            start_time = time.monotonic()

        # 接続元がObjectDetectionノードの場合、各バウンディングボックスに対して推論
        result = {}
        frame_list, class_id_list, score_list = [], [], []
        od_target_bboxes = []
        od_target_scores = []
        od_target_class_ids = []
        if frame is not None:
            if src_node_name == 'ObjectDetection':
                # 物体検出情報取得
                node_result = node_result_dict.get(connection_info_src, [])
                od_bboxes = node_result.get('bboxes', [])
                od_scores = node_result.get('scores', [])
                od_class_ids = node_result.get('class_ids', [])
                od_class_names = node_result.get('class_names', [])
                od_score_th = node_result.get('score_th', [])

                # バウンディングボックスで切り抜き
                for od_bbox, od_score, od_class_id in zip(
                        od_bboxes, od_scores, od_class_ids):
                    x1, y1 = int(od_bbox[0]), int(od_bbox[1])
                    x2, y2 = int(od_bbox[2]), int(od_bbox[3])

                    if od_score_th > od_score:
                        continue

                    frame_list.append(copy.deepcopy(frame[y1:y2, x1:x2]))
                    od_target_bboxes.append([x1, y1, x2, y2])
                    od_target_scores.append(od_score)
                    od_target_class_ids.append(od_class_id)

                # 各バウンディングボックスに対しClassification推論
                for temp_frame in frame_list:
                    class_scores, class_ids = self._model_instance[
                        model_name_with_provider](temp_frame)
                    score_list.append(class_scores[0])
                    class_id_list.append(class_ids[0])
                result['use_object_detection'] = True
                result['class_ids'] = class_id_list
                result['class_scores'] = score_list
                result['class_names'] = class_name_dict
                result['od_bboxes'] = od_target_bboxes
                result['od_scores'] = od_target_scores
                result['od_class_ids'] = od_target_class_ids
                result['od_class_names'] = od_class_names
                result['od_score_th'] = od_score_th
            else:
                class_scores, class_ids = self._model_instance[
                    model_name_with_provider](frame)
                result['use_object_detection'] = False
                result['class_ids'] = class_ids.tolist()
                result['class_scores'] = class_scores.tolist()
                result['class_names'] = class_name_dict

        # 計測終了
        if frame is not None and use_pref_counter:
            elapsed_time = time.monotonic() - start_time
            elapsed_time = int(elapsed_time * 1000)
            dpg_set_value(output_value02_tag,
                          str(elapsed_time).zfill(4) + 'ms')

        # 描画
        output_frame = frame
        if frame is not None:
            # Create debug_frame with original dimensions for output
            debug_frame = copy.deepcopy(frame)
            
            # Draw labels on the original frame
            if result['use_object_detection']:
                debug_frame = self.draw_classification_with_od_info(
                    debug_frame,
                    result['class_ids'],
                    result['class_scores'],
                    result['class_names'],
                    result['od_bboxes'],
                    result['od_scores'],
                    result['od_class_ids'],
                    result['od_class_names'],
                    result['od_score_th'],
                )
            else:
                debug_frame = self.draw_classification_info(
                    debug_frame,
                    result['class_ids'],
                    result['class_scores'],
                    result['class_names'],
                )
            
            output_frame = debug_frame

            texture = self.convert_cv_to_dpg(
                debug_frame,
                small_window_w,
                small_window_h,
            )
            dpg_set_value(output_value01_tag, texture)

        return {"image": output_frame, "json": result, "audio": None}

    def close(self, node_id):
        pass
    
    def draw_classification_info(
        self,
        image,
        class_ids,
        class_scores,
        class_names,
    ):
        """
        Override base class method to add color differentiation based on ranking.
        Position 1 (index 0, highest score): Red
        Position 2 (index 1): Yellow
        Position 3 (index 2): Blue
        Position 4 (index 3): Violet
        Position 5 (index 4): Magenta
        """
        debug_image = copy.deepcopy(image)
        
        # Define colors for top 5 positions (BGR format)
        rank_colors = [
            (0, 0, 255),      # Position 1 (index 0): Red (highest score)
            (0, 255, 255),    # Position 2 (index 1): Yellow
            (255, 0, 0),      # Position 3 (index 2): Blue
            (255, 0, 128),    # Position 4 (index 3): Violet (purple-ish)
            (255, 0, 255),    # Position 5 (index 4): Magenta
        ]
        
        for index, (class_score, class_id) in enumerate(zip(class_scores, class_ids)):
            score = "%.2f" % class_score
            text = "%s:%s(%s)" % (str(class_id), str(class_names[int(class_id)]), score)
            
            # Select color based on position (1, 2, 3, 4, 5)
            # Use default green for positions beyond 5
            if index < len(rank_colors):
                color = rank_colors[index]
            else:
                color = (0, 255, 0)  # Default green for lower rankings
            
            debug_image = cv2.putText(
                debug_image,
                text,
                (15, 25 + (index * 20)),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.6,
                color,
                thickness=2,
            )

        return debug_image

    def get_setting_dict(self, node_id):
        tag_node_name = str(node_id) + ':' + self.node_tag
        input_value02_tag = tag_node_name + ':' + self.TYPE_TEXT + ':Input02Value'

        # 選択モデル
        model_name = dpg_get_value(input_value02_tag)

        pos = dpg.get_item_pos(tag_node_name)

        setting_dict = {}
        setting_dict['ver'] = self._ver
        setting_dict['pos'] = pos
        setting_dict[input_value02_tag] = model_name

        return setting_dict

    def set_setting_dict(self, node_id, setting_dict):
        tag_node_name = str(node_id) + ':' + self.node_tag
        input_value02_tag = tag_node_name + ':' + self.TYPE_TEXT + ':Input02Value'

        model_name = setting_dict[input_value02_tag]

        dpg_set_value(input_value02_tag, model_name)


# Load user-uploaded custom models from registry at import time
Node._load_custom_models_from_registry()