arya / image_processing.py
ArdaKaratas's picture
Update image_processing.py
391aa6c verified
"""
Image Processing Tool
Provides basic image analysis capabilities for GAIA benchmark questions.
"""
from smolagents import Tool
from PIL import Image
import requests
from io import BytesIO
class ImageProcessingTool(Tool):
name = "image_processing"
description = "Analyzes images from URLs. Can extract basic information about images including dimensions, format, and basic properties."
inputs = {
"image_url": {
"type": "string",
"description": "URL of the image to analyze"
},
"task": {
"type": "string",
"description": "What to do with the image: 'info' (get basic info), 'describe' (describe the image)",
"default": "info",
"nullable": True
}
}
output_type = "string"
def forward(self, image_url: str, task: str = "info") -> str:
"""
Process an image from a URL.
Args:
image_url: URL of the image
task: What task to perform ('info' or 'describe')
Returns:
Information about the image
"""
try:
# Download the image
response = requests.get(image_url, timeout=10)
response.raise_for_status()
# Open image
image = Image.open(BytesIO(response.content))
if task == "info":
# Return basic image information
info = {
"format": image.format,
"mode": image.mode,
"size": image.size,
"width": image.width,
"height": image.height,
}
return (
f"Image Information:\n"
f"Format: {info['format']}\n"
f"Mode: {info['mode']}\n"
f"Dimensions: {info['width']}x{info['height']} pixels\n"
f"Size: {info['size']}"
)
elif task == "describe":
# Basic description based on image properties
return (
f"This is a {image.format} image with dimensions {image.width}x{image.height} pixels. "
f"Color mode: {image.mode}. "
f"For detailed visual description, use a vision model."
)
else:
return f"Unknown task: {task}. Use 'info' or 'describe'."
except requests.exceptions.RequestException as e:
return f"Error downloading image: {str(e)}"
except Exception as e:
return f"Error processing image: {str(e)}"