From 363dbd9a220d7c59786cea964fa0c2be13f492a4 Mon Sep 17 00:00:00 2001
From: "DESKTOP-UL3K6JV\\Daniq" <31.10.93@bk.ru>
Date: Sun, 16 Jun 2024 14:36:06 +0500
Subject: [PATCH] init

---
 .gitignore       |  6 +++++
 __init__.py      |  6 +++++
 classifiers.py   | 65 ++++++++++++++++++++++++++++++++++++++++++++++++
 requirements.txt |  1 +
 4 files changed, 78 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 __init__.py
 create mode 100644 classifiers.py
 create mode 100644 requirements.txt

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..1373c89
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,6 @@
+__pycache__
+*.py[cod]
+*.ini
+.vscode/
+.idea/
+*.yaml
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..2a7cfda
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,6 @@
+from .classifiers import *
+
+NODE_CLASS_MAPPINGS = {
+    'YOLO Classifier Model Loader': YOLOClassifierModelLoader,
+    'YOLO Classify': YOLOClassify
+}
diff --git a/classifiers.py b/classifiers.py
new file mode 100644
index 0000000..e8c519d
--- /dev/null
+++ b/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
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..9729bc9
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1 @@
+ultralytics==8.2.32
\ No newline at end of file