File size: 2,176 Bytes
0e359f1 ac957a1 0e359f1 ac957a1 0e359f1 89e197c 0e359f1 89e197c 5b13b85 0e359f1 89e197c 5b13b85 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 51 52 53 54 55 56 57 58 59 60 61 |
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
import os
class EndpointHandler():
def __init__(self, path=""):
# 换个顺序就行了？？？？？离谱
os.system("ln -s /repository /app/repository")
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')
texts=info.pop('texts',data)
target=texts[0]
replacement=texts[1]
if replacement=="":
img=self.model.inference_remove(raw_image,target)
img.save("./1.png")
with open('./1.png','rb') as img_file:
encoded_string = base64.b64encode(img_file.read()).decode('utf-8')
return {'image':encoded_string}
else:
img=self.model.inference_replace_sam(raw_image,target,replacement)
img.save("./1.png")
with open('./1.png','rb') as img_file:
encoded_string = base64.b64encode(img_file.read()).decode('utf-8')
return {'image':encoded_string}
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")
|