143 lines
4.4 KiB
Python
143 lines
4.4 KiB
Python
import os
|
|
import io
|
|
import json
|
|
import requests
|
|
from io import BytesIO
|
|
from PIL import Image
|
|
import sys
|
|
import torch
|
|
import numpy as np
|
|
import base64
|
|
|
|
p = os.path.dirname(os.path.realpath(__file__))
|
|
|
|
globalAccessToken = ''
|
|
|
|
def tensor2pil(image):
|
|
return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
|
|
|
|
def pil2tensor(image):
|
|
return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)
|
|
|
|
def getBabelAccessToken():
|
|
try:
|
|
global globalAccessToken
|
|
if globalAccessToken != '':
|
|
return globalAccessToken
|
|
config_path = os.path.join(p, 'config.json')
|
|
with open(config_path, 'r') as f:
|
|
config = json.load(f)
|
|
apikey = config["BABEL_KEY"]
|
|
apisecret = config["BABEL_SECRET"]
|
|
|
|
if apikey == '' or apisecret == '':
|
|
raise Exception(f"apikey and apisecret is neccessary")
|
|
|
|
response = requests.post(
|
|
f"https://removebg-api.babel-hz.com/oauth/token",
|
|
headers={
|
|
"accept": "application/json",
|
|
},
|
|
json={
|
|
"grantType":"client_credentials",
|
|
"apikey":apikey,
|
|
"apisecret":apisecret
|
|
},
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
jsonData = response.json()
|
|
print(jsonData)
|
|
if jsonData['code'] == 0:
|
|
globalAccessToken = jsonData['data']['accessToken']
|
|
return globalAccessToken
|
|
else:
|
|
print("Response Msg:", jsonData['msg'])
|
|
raise Exception(f"Failed to fetch accessToken msg: {jsonData['msg']}")
|
|
else:
|
|
print("Response Status Code:", response.status_code)
|
|
raise Exception(f"Failed to fetch accessToken status code: {response.status_code}")
|
|
|
|
except Exception as e:
|
|
print(e)
|
|
return ""
|
|
|
|
|
|
class BabelRemovebg:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "generate_removebg"
|
|
CATEGORY = "Babel/removebg/body"
|
|
|
|
def generate_removebg(self, image=None):
|
|
|
|
accessToken = getBabelAccessToken()
|
|
|
|
image = tensor2pil(image)
|
|
# 将图像转换为字节流
|
|
image_byte_arr = io.BytesIO()
|
|
image.save(image_byte_arr, format='JPEG', quality = 90)
|
|
image_byte_arr = image_byte_arr.getvalue()
|
|
# 将字节流编码为 base64
|
|
imageBase64 = base64.b64encode(image_byte_arr).decode('utf-8')
|
|
print(f"image size is {len(imageBase64)}")
|
|
|
|
response = requests.post(
|
|
f"https://removebg-api.babel-hz.com/task/cutoff/body",
|
|
headers={
|
|
"Authorization": f"accessToken {accessToken}",
|
|
"accept": "application/json"
|
|
},
|
|
json={
|
|
"imageData":imageBase64,
|
|
"type":"foreground",
|
|
"returnType":"url"
|
|
},
|
|
)
|
|
|
|
if response.status_code == 200:
|
|
jsonData = response.json()
|
|
if jsonData['code'] == 0:
|
|
taskId = jsonData['data']['taskId']
|
|
foreground = jsonData['data']['foreground']
|
|
response = requests.get(foreground)
|
|
|
|
if response.status_code == 200:
|
|
imagefile = f"{taskId}.png"
|
|
with open(imagefile, "wb") as f:
|
|
f.write(response.content)
|
|
image_data = Image.open(imagefile)
|
|
output_t = pil2tensor(image_data)
|
|
os.remove(imagefile)
|
|
print(output_t.shape)
|
|
return (output_t,)
|
|
else:
|
|
print(f"Failed to download image. Status code: {response.status_code}")
|
|
|
|
else:
|
|
print("Response Msg:", jsonData['msg'])
|
|
raise Exception(f"Failed to removebg msg: {jsonData['msg']}")
|
|
else:
|
|
print("Response Status Code:", response.status_code)
|
|
raise Exception(f"Failed to removebg status code: {response.status_code}")
|
|
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"BabelRemovebg": BabelRemovebg,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"BabelRemovebg": "Babel-Removebg",
|
|
}
|