Spaces:
Running
Running
| 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) | |
| 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 | |