File size: 1,156 Bytes
0d92e3f ee61ccb 0d92e3f ee61ccb 0d92e3f ee61ccb 0d92e3f |
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 |
from typing import Dict, List, Any
from transformers import Pipeline
from transformers import BlipProcessor, BlipForConditionalGeneration
from PIL import Image
from io import BytesIO
import base64
import json
class EndpointHandler():
def __init__(self, path=""):
self.processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
self.model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to("cuda")
def __call__(self, data):
info=data['inputs']
img=info.pop('image',data)
image_bytes=base64.b64decode(img)
raw_images = Image.open(BytesIO(image_bytes))
inputs = self.processor(raw_images, return_tensors="pt").to("cuda")
out = self.model.generate(**inputs)
# print(self.processor.decode(out[0], skip_special_tokens=True))
return {'text':self.processor.decode(out[0], skip_special_tokens=True)}
if __name__=="__main__":
my_handler=EndpointHandler(path='.')
test_payload={"inputs": "/home/ubuntu/guoling/1.png"}
test_result=my_handler(test_payload)
print(test_result)
|