From a590595ef1970efa5c49ab92fd7af0bea4af4164 Mon Sep 17 00:00:00 2001
From: kadirnar <kadir.nar@hotmail.com>
Date: Thu, 25 Jan 2024 16:17:03 +0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=8C=A0=20Add=20image-classification=20pip?=
 =?UTF-8?q?eline=20and=20update=20req.txt=20file?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../transformers/image_classification/node.py | 53 +++++++++++++++++++
 requirements.txt                              |  7 +--
 2 files changed, 57 insertions(+), 3 deletions(-)
 create mode 100644 computer_vision/transformers/image_classification/node.py

diff --git a/computer_vision/transformers/image_classification/node.py b/computer_vision/transformers/image_classification/node.py
new file mode 100644
index 0000000..fbb2584
--- /dev/null
+++ b/computer_vision/transformers/image_classification/node.py
@@ -0,0 +1,53 @@
+from transformers import AutoImageProcessor, ResNetForImageClassification
+import torch
+from PIL import Image
+from typing import Tuple
+
+image_classification_model_name_list = [
+    "microsoft/resnet-50"
+]
+
+class ImageClassificationPipeline:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "model_name": (image_classification_model_name_list, {"default": image_classification_model_name_list[0]}),
+            },
+        }
+
+    RETURN_TYPES = ("STRING","FLOAT",)
+    FUNCTION = "classify_pipeline"
+
+    CATEGORY = "ComputerVision/Transformers"
+
+    def classify_image(image: Image.Image, model_name: str) -> Tuple[str, float]:
+        """
+        Classifies a given image using the specified model and returns the name and score of the predicted label.
+
+        Args:
+        image (PIL.Image.Image): The image to be classified.
+        model_name (str): The name of the model to be used for classification.
+
+        Returns:
+        Tuple[str, float]: The name of the predicted label and its corresponding score.
+        """
+        
+        # Load the specified model and corresponding image processor
+        image_processor = AutoImageProcessor.from_pretrained(model_name)
+        model = ResNetForImageClassification.from_pretrained(model_name)
+
+        # Process the image and prepare it for the model
+        inputs = image_processor(image, return_tensors="pt")
+
+        with torch.no_grad():
+            logits = model(**inputs).logits
+
+        # Predict the class with the highest score
+        predicted_label_id = logits.argmax(-1).item()
+        predicted_label_name = model.config.id2label[predicted_label_id]
+        predicted_score = torch.softmax(logits, -1)[0, predicted_label_id].item()
+
+        return predicted_label_name, predicted_score
+
diff --git a/requirements.txt b/requirements.txt
index dbba6ba..7e85007 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
-numpy
-Pillow
-torch>=2.0.0
+numpy=>1.26.3
+Pillow=>10.2.0
+torch=>2.1.2
+transformers=>4.37.1