File size: 1,695 Bytes
0e359f1 |
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 47 48 49 50 |
from typing import Dict, List, Any
from transformers import Pipeline
from PIL import Image
from io import BytesIO
import base64
import json
from visual_chatgpt import ImageEditing, Text2Box, Segmenting, Inpainting
class EndpointHandler():
def __init__(self, path=""):
# 换个顺序就行了？？？？？离谱
self.sam = Segmenting('cuda')
self.inpaint = Inpainting('cuda')
self.grounding = Text2Box('cuda')
self.model = ImageEditing(self.grounding,self.sam,self.inpaint)
def __call__(self, data):
# data=json.loads(data)
# inputs=data.pop("inputs",data)
# inputs=base64.b64decode(inputs)
# raw_images = Image.open(BytesIO(inputs))
info=data['inputs']
image=info.pop('image',data)
image=base64.b64decode(image)
raw_image=Image.open(BytesIO(image)).convert('RGB')
target=info.pop('target',data)
replacement=info.pop('replacement',data)
if replacement=="":
return self.model.inference_remove(raw_image,target)
else:
return self.model.inference_replace_sam(raw_image,target,replacement)
if __name__=="__main__":
my_handler=EndpointHandler(path='.')
# test_payload={"inputs": "/home/ubuntu/guoling/1.png"}
with open("/home/ubuntu/guoling/1.png",'rb') as img:
image_bytes=img.read()
image_base64=base64.b64encode(image_bytes).decode('utf-8')
target="the pig"
replacement=""
data={
'inputs':{
"image":image_base64,
"target":target,
"replacement":replacement
}
}
result=my_handler(data)
result.save("new1.png")
|