Spaces:
Sleeping
Sleeping
File size: 2,040 Bytes
860424e |
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 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
from PIL import Image
from transformers import BlipProcessor, BlipForQuestionAnswering
from smolagents.tools import tool
import torch
import requests
import os
DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
def _download_file(file_name: str) -> None:
"""Download file if it doesn't exist locally"""
try:
# Try to open the file to check if it exists
with open(file_name, 'rb') as f:
pass # File exists, do nothing
except FileNotFoundError:
# File doesn't exist, download it
url = f"{DEFAULT_API_URL}/files/{file_name.split('.')[-2]}"
r = requests.get(url)
with open(file_name, "wb") as f:
f.write(r.content)
@tool
def ask_question_about_image(question: str, path_to_image: str) -> str:
"""
Ask a question about an image and return the answer.
Args:
question: the question to ask about the image.
path_to_image: The path to the image to ask the question about.
Returns:
A string with the answer to the question.
"""
# Download the file if it doesn't exist
_download_file(path_to_image)
# Check if CUDA is available and use GPU if possible, otherwise use CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# Load the processor and model (using BLIP for more stable VQA)
processor = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
model = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base")
model = model.to(device)
# Load and process the image
image = Image.open(path_to_image).convert('RGB')
# Process the inputs
inputs = processor(image, question, return_tensors="pt")
inputs = {k: v.to(device) for k, v in inputs.items()}
# Generate the answer
with torch.no_grad():
outputs = model.generate(**inputs, max_length=50, num_beams=5)
# Decode and return the answer
answer = processor.decode(outputs[0], skip_special_tokens=True)
return answer
|