|
|
from smolagents import Tool |
|
|
from PIL import Image |
|
|
import requests |
|
|
from io import BytesIO |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class ImageLoadTool(Tool): |
|
|
name = "_my_tool_image_load" |
|
|
description = """ |
|
|
Load image for the provided task id |
|
|
To invoke the tool use code as below |
|
|
<code> |
|
|
loaded_image = _my_tool_image_load(task_id="dummy") |
|
|
</code> |
|
|
""" |
|
|
|
|
|
inputs = { |
|
|
"task_id": { |
|
|
"type": "string", |
|
|
"description": "task id to load image", |
|
|
} |
|
|
} |
|
|
|
|
|
output_type = "image" |
|
|
api_url = "https://agents-course-unit4-scoring.hf.space" |
|
|
|
|
|
def forward(self, task_id: str) -> Image: |
|
|
headers = { |
|
|
"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36" |
|
|
} |
|
|
url = f"{self.api_url}/files/{task_id}" |
|
|
response = requests.get(url, headers=headers) |
|
|
image = Image.open(BytesIO(response.content)).convert("RGB") |
|
|
print(f"***KS*** Loaded image for \n\ttask id: {task_id} \n\timage: {image}") |
|
|
return image |
|
|
|
