Spaces:
Runtime error
Runtime error
File size: 1,243 Bytes
0c7479d |
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 |
from PIL import Image
import json
from transformers import AutoProcessor, Blip2ForConditionalGeneration
import torch
import os
processor = AutoProcessor.from_pretrained("Salesforce/blip2-opt-2.7b")
model = Blip2ForConditionalGeneration.from_pretrained("Salesforce/blip2-opt-2.7b", torch_dtype=torch.float16)
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
def get_blip2_text(image):
inputs = processor(image, return_tensors="pt").to(device, torch.float16)
generated_ids = model.generate(**inputs, max_new_tokens=50)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0].strip()
return generated_text
data_path = "files"
save_path = ""
image_names = os.listdir(data_path)
image_names = sorted(image_names)
text_data = {}
f = open("data.txt","w")
for each in image_names:
if '.jpg' in each:
this_data = {}
this_data['target'] = each
this_data['source'] = each[:-4]+'.json'
this_image = Image.open(os.path.join(data_path, each))
print(each)
generated_text = get_blip2_text(this_image)
this_data['prompt'] = generated_text
print(this_data)
f.write(str(this_data)+"\n")
f.close()
|