# 图片转为文字描述（自然语言）

2024-04-06 00:24:56 +08:00
commit f7061533a1
6 changed files with 227 additions and 0 deletions
@@ -0,0 +1,51 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from PIL import Image
+
+import numpy as np
+
+model_id = "vikhyatk/moondream2"
+revision = "2024-04-02"
+
+class ComfyUI_ImageToText:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "images": ("IMAGE",),
+                "log_prompt": (["No", "Yes"], {"default":"Yes"}),
+            },
+        }
+
+    RETURN_TYPES = ('STRING',)
+    RETURN_NAMES = ('text_positive',)
+    FUNCTION = "image2text"
+    OUTPUT_NODE = True
+    CATEGORY = "ComfyUI_Mexx"
+
+    def image2text(self, images, log_prompt):
+        pil_images = []
+        for image in images:
+            i = 255. * image.cpu().numpy()
+            img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
+            pil_images.append(img)
+        image = pil_images[0]
+        model = AutoModelForCausalLM.from_pretrained(
+            model_id, trust_remote_code=True, revision=revision
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_id, revision=revision)
+        enc_image = model.encode_image(image)
+        en = model.answer_question(enc_image, "Describe this image.", tokenizer)
+        if log_prompt == "Yes":
+            print(f"ImageToText: {en}")
+        return [en]
+
+NODE_CLASS_MAPPINGS = {
+    "ComfyUI_ImageToText": ComfyUI_ImageToText
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "ComfyUI_ImageToText": "ComfyUI_ImageToText"
+}