init

.gitignore

@@ -0,0 +1,6 @@
+__pycache__
+*.py[cod]
+*.ini
+.vscode/
+.idea/
+*.yaml

__init__.py

@@ -0,0 +1,6 @@
+from .classifiers import *
+
+NODE_CLASS_MAPPINGS = {
+    'YOLO Classifier Model Loader': YOLOClassifierModelLoader,
+    'YOLO Classify': YOLOClassify
+}

classifiers.py

@@ -0,0 +1,65 @@
+import os
+
+import folder_paths
+from ultralytics import YOLO
+
+model_path = folder_paths.models_dir
+folder_paths.add_model_folder_path("ultralytics_classifiers", os.path.join(model_path, "ultralytics", "classifiers"))
+
+
+class YOLOClassifierModelLoader:
+    # Download these models from https://docs.ultralytics.com/tasks/classify/
+    # Place them on models/ultralytics/classification folder
+    model_names = (
+        'yolov8n-cls.pt',
+        'yolov8s-cls.pt',
+        'yolov8m-cls.pt',
+        'yolov8l-cls.pt',
+        'yolov8x-cls.pt'
+    )
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_name": cls.model_names
+            }
+        }
+
+    RETURN_TYPES = ("YOLO_CLASSIFIER_MODEL",)
+    RETURN_NAMES = ("YOLO classifier model",)
+    FUNCTION = "load_model"
+    CATEGORY = "SuperMasterBlasterLaser/ComfyUI_YOLO_Classifiers"
+
+    def load_model(self, model_name):
+        model_path = folder_paths.get_full_path("ultralytics_classifiers", model_name)
+
+        model = YOLO(model_path)
+        return model
+
+
+class YOLOClassify:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "yolo_classifier_model": "YOLO_CLASSIFIER_MODEL",
+                "image_to_classify": "IMAGE"
+            }
+        }
+
+    RETURN_TYPES = ("STRING",)
+    RETURN_NAMES = ("Class name with highest score",)
+    FUNCTION = "classify"
+    CATEGORY = "SuperMasterBlasterLaser/ComfyUI_YOLO_Classifiers"
+
+    def classify(self, yolo_classifier_model, image_to_classify):
+        result = yolo_classifier_model(image_to_classify)
+
+        class_name = None
+        for r in result:
+            top1 = r.probs.top1
+            class_name = r.names[top1]
+            break
+
+        return class_name

requirements.txt

@@ -0,0 +1 @@
+ultralytics==8.2.32