File size: 970 Bytes
6220346
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
from smolagents import Tool
from transformers import pipeline

class ConvertImageToTextTool(Tool):
    name = "convert_image_to_text"
    description = "Transcribe an image file to text using a free Hugging Face template."
    inputs = {
        "image_path": {
            "type": "string",
            "description": "The path of the image file to elaborate"
        }
    }
    output_type = "string"

    def __init__(self):
        super().__init__()
        self.model = "nlpconnect/vit-gpt2-image-captioning"
        self.transcriber = pipeline(
            "image-to-text",
            model=self.model,
            use_fast=True
        )

    def forward(self, image_path: str) -> str:
        try:
            result = self.transcriber(image_path)
            return f"Image description: {result[0]['generated_text']}"
        except Exception as e:
            return f"Error convert_image_to_text is not working properly, error: {e}, please skip this tool"