From f3d84c4545283d6f28e904a99a0fa342744c84a5 Mon Sep 17 00:00:00 2001
From: "DESKTOP-UL3K6JV\\Daniq" <31.10.93@bk.ru>
Date: Sun, 16 Jun 2024 16:15:26 +0500
Subject: [PATCH] Added tensor conversion

---
 classifiers.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/classifiers.py b/classifiers.py
index adb3bfe..2877dd9 100644
--- a/classifiers.py
+++ b/classifiers.py
@@ -1,6 +1,7 @@
 import os
 
 import folder_paths
+from torchvision import transforms
 from ultralytics import YOLO
 
 model_path = folder_paths.models_dir
@@ -54,7 +55,13 @@ class YOLOClassify:
     CATEGORY = "SuperMasterBlasterLaser/ComfyUI_YOLO_Classifiers"
 
     def classify(self, yolo_classifier_model, image_to_classify):
-        result = yolo_classifier_model.predict(image_to_classify)
+        transform = transforms.Compose(
+            [transforms.ToTensor(), transforms.Resize(size=(640, 640))]
+        )
+
+        img_tensor = transform(image_to_classify).unsqueeze(0)
+
+        result = yolo_classifier_model(img_tensor)
 
         class_name = None
         for r in result: