From 5fcd05babe2f687096d1e47ba932d73b59bf0afb Mon Sep 17 00:00:00 2001
From: "DESKTOP-UL3K6JV\\Daniq" <31.10.93@bk.ru>
Date: Sun, 16 Jun 2024 16:42:22 +0500
Subject: [PATCH] Changed to pil

---
 classifiers.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/classifiers.py b/classifiers.py
index 315c03c..62346d8 100644
--- a/classifiers.py
+++ b/classifiers.py
@@ -55,13 +55,12 @@ class YOLOClassify:
     CATEGORY = "SuperMasterBlasterLaser/ComfyUI_YOLO_Classifiers"
 
     def classify(self, yolo_classifier_model, image_to_classify):
-        transform = transforms.Compose(
-            [transforms.ToPILImage()]
-        )
+        transform = transforms.Compose([
+            transforms.ToPILImage(mode='LA')
+        ])
 
-        img_pil = transform(image_to_classify).unsqueeze(0)
-
-        result = yolo_classifier_model(img_pil)
+        img = transform(image_to_classify).unsqueeze(0)
+        result = yolo_classifier_model(img)
 
         class_name = None
         for r in result: