describe.py

describe.py

from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
from PIL import Image
import requests


import warnings
warnings.filterwarnings("ignore")


def load_model():
model = Pix2StructForConditionalGeneration.from_pretrained('google/pix2struct-textcaps-base').to("mps")
processor = Pix2StructProcessor.from_pretrained('google/pix2struct-textcaps-base')
return model, processor


def process_data(image, model, processor):
inputs = processor(images=image, return_tensors='pt').to("mps")
generated_text = model.generate(**inputs, max_length= 60)
response = processor.decode(generated_text[0], skip_special_tokens=True)
return response




image_path = './logo.png'
image = Image.open(image_path)


# Load the model and processor
model, processor = load_model()


# Process the image and generate text
response = process_data(image, model, processor)


print(response)