Spaces:
Runtime error
Runtime error
| 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() | |