added movedim

2024-06-16 16:56:18 +05:00
parent 5fcd05babe
commit b0eab338d1
1 changed files with 3 additions and 1 deletions
@@ -55,8 +55,10 @@ class YOLOClassify:
    CATEGORY = "SuperMasterBlasterLaser/ComfyUI_YOLO_Classifiers"

    def classify(self, yolo_classifier_model, image_to_classify):
+        image_to_classify = image_to_classify.movedim(3, 1)
+
        transform = transforms.Compose([
-            transforms.ToPILImage(mode='LA')
+            transforms.Resize(size=(640, 640))
        ])

        img = transform(image_to_classify).unsqueeze(0)