413 lines
15 KiB
Python
413 lines
15 KiB
Python
from PIL import Image,ImageOps
|
|
import os
|
|
import folder_paths
|
|
import hashlib
|
|
import torch
|
|
import numpy as np
|
|
import json
|
|
from .src.openPose.pose_body import Body
|
|
from .src.openPose.util import draw_bodypose
|
|
import server
|
|
from aiohttp import web
|
|
import os
|
|
import json
|
|
import cv2
|
|
from tools import annotator_ckpts_path, load_file_from_url
|
|
|
|
|
|
body_estimation = None
|
|
body_model_path = "https://hf-mirror.com/lllyasviel/Annotators/resolve/main/body_pose_model.pth"
|
|
|
|
dir = os.path.abspath(os.path.join(__file__, "../../docs"))
|
|
if not os.path.exists(dir):
|
|
os.mkdir(dir)
|
|
|
|
# Directory node save settings
|
|
CHUNK_SIZE = 1024
|
|
basePath = folder_paths.folder_names_and_paths['custom_nodes'][0][0]
|
|
extension_path = os.path.join(basePath, 'ComfyUI_Lam', 'config')
|
|
file_settings_path = os.path.join(extension_path,"settings_nodes.json")
|
|
|
|
# Function create file json file
|
|
def create_settings_json(filename="settings_nodes.json"):
|
|
json_file = os.path.join(extension_path, filename)
|
|
if not os.path.exists(json_file):
|
|
print("找不到文件settings_nodes.json！创建文件！")
|
|
with open(json_file, "w") as f:
|
|
json.dump({}, f)
|
|
|
|
def get_settings_json(filename="settings_nodes.json", notExistCreate=True):
|
|
json_file = os.path.join(extension_path, filename)
|
|
if os.path.isfile(json_file):
|
|
f = open(json_file, "rb")
|
|
try:
|
|
load_data = json.load(f)
|
|
return load_data
|
|
except Exception as e:
|
|
print("Error load json file: ",e)
|
|
if notExistCreate:
|
|
f.close()
|
|
os.remove(json_file)
|
|
create_settings_json()
|
|
finally:
|
|
f.close()
|
|
|
|
return {}
|
|
|
|
# Load json file
|
|
@server.PromptServer.instance.routes.get("/lam/loading_node_settings")
|
|
async def loadingSettings(request):
|
|
load_data = get_settings_json()
|
|
return web.json_response({"settings_nodes": load_data})
|
|
|
|
# Save data to json file
|
|
@server.PromptServer.instance.routes.post("/lam/save_node_settings")
|
|
async def saveSettings(request):
|
|
try:
|
|
with open(file_settings_path, "wb") as f:
|
|
while True:
|
|
chunk = await request.content.read(CHUNK_SIZE)
|
|
if not chunk:
|
|
break
|
|
f.write(chunk)
|
|
|
|
return web.json_response({"message": "Painter data saved successfully"}, status=200)
|
|
|
|
except Exception as e:
|
|
print("Error save json file: ", e)
|
|
|
|
# create file json
|
|
create_settings_json()
|
|
|
|
@server.PromptServer.instance.routes.get("/lam/getImage")
|
|
async def get_groupNode(request):
|
|
if "name" in request.rel_url.query:
|
|
type = request.rel_url.query["type"]
|
|
name = request.rel_url.query["name"]
|
|
file = os.path.join(dir, type, name)
|
|
if os.path.isfile(file):
|
|
return web.FileResponse(file)
|
|
return web.Response(status=404)
|
|
|
|
|
|
@server.PromptServer.instance.routes.get("/lam/getHeads")
|
|
def getStyles(request):
|
|
filePath = os.path.join(dir, 'hands')
|
|
names = []
|
|
for root, dirs, files in os.walk(filePath):
|
|
for file in files:
|
|
if file.endswith(".png") or file.endswith(".jpg"):
|
|
names.append(file)
|
|
return web.json_response(names)
|
|
|
|
|
|
@server.PromptServer.instance.routes.post("/lam/getImagePose")
|
|
async def getImagePose(request):
|
|
global body_estimation # 声明 body_estimation 是全局变量
|
|
post = await request.post()
|
|
image = post.get("image")
|
|
if image and image.file:
|
|
# 将文件转换为 OpenCV 可以处理的格式
|
|
nparr = np.frombuffer(image.file.read(), np.uint8)
|
|
oriImg = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
|
if body_estimation == None:
|
|
body_modelpath = os.path.join(
|
|
annotator_ckpts_path, "body_pose_model.pth")
|
|
if not os.path.exists(body_modelpath):
|
|
load_file_from_url(
|
|
body_model_path, model_dir=annotator_ckpts_path)
|
|
body_estimation = Body(body_modelpath)
|
|
H, W, C = oriImg.shape
|
|
YW = 512*(1 if W > H else W/H)
|
|
YH = 512*(1 if W < H else H/W)
|
|
candidate, subset = body_estimation(oriImg)
|
|
if candidate.ndim == 2 and candidate.shape[1] == 4:
|
|
candidate = candidate[:, :2]
|
|
candidate[:, 0] /= float(W)
|
|
candidate[:, 1] /= float(H)
|
|
groups = []
|
|
subsets = []
|
|
for n in range(len(subset)):
|
|
groups.append([[-1, -1] if i < 0 else [int(candidate[int(i)][0]*YW),
|
|
int(candidate[int(i)][1]*YH)] for i in subset[n][:18]])
|
|
sub = []
|
|
index = 0
|
|
for s in subset[n][:18]:
|
|
if s >= 0:
|
|
sub.append(index)
|
|
index += 1
|
|
else:
|
|
sub.append(s)
|
|
subsets.append(sub)
|
|
|
|
return web.json_response({'groups': groups, 'subsets': subsets})
|
|
|
|
return web.Response(status=404)
|
|
|
|
|
|
def get_bounding_box(coords, yWidth, yHeight, xReund=(0, 0), YReund=(0, 0)):
|
|
try:
|
|
coords = [x for x in coords if x[0] >= 0 and x[1] >= 0]
|
|
if len(coords) == 0:
|
|
return None
|
|
coords = np.array(coords).flatten().tolist()
|
|
min_x = min(coords[0::2]) # 获取x坐标的最小值
|
|
min_y = min(coords[1::2]) # 获取y坐标的最小值
|
|
max_x = max(coords[0::2]) # 获取x坐标的最大值
|
|
max_y = max(coords[1::2]) # 获取y坐标的最大值
|
|
|
|
width = max_x - min_x # 宽度
|
|
height = max_y - min_y # 高度
|
|
wd = np.array(xReund) * width
|
|
hd = np.array(YReund) * height
|
|
min_x = int(min_x-wd[0])
|
|
min_y = int(min_y-hd[0])
|
|
width = int(width+wd[0]+wd[1])
|
|
height = int(height+hd[0]+hd[1])
|
|
min_x = min_x if min_x > 0 else 0
|
|
min_y = min_y if min_y > 0 else 0
|
|
width = width if width < yWidth else yWidth
|
|
height = height if height < yHeight else yHeight
|
|
return (min_x, min_y, width, height)
|
|
except Exception as e:
|
|
return None
|
|
|
|
|
|
def solid_mask(value, width, height):
|
|
out = torch.full((1, height, width), value,
|
|
dtype=torch.float32, device="cpu")
|
|
return out
|
|
|
|
|
|
def mask_combine(destination, source, x, y, operation="add"):
|
|
output = destination.reshape(
|
|
(-1, destination.shape[-2], destination.shape[-1])).clone()
|
|
source = source.reshape((-1, source.shape[-2], source.shape[-1]))
|
|
|
|
left, top = (x, y,)
|
|
right, bottom = (min(left + source.shape[-1], destination.shape[-1]), min(
|
|
top + source.shape[-2], destination.shape[-2]))
|
|
visible_width, visible_height = (right - left, bottom - top,)
|
|
|
|
source_portion = source[:, :visible_height, :visible_width]
|
|
destination_portion = destination[:, top:bottom, left:right]
|
|
|
|
if operation == "multiply":
|
|
output[:, top:bottom, left:right] = destination_portion * source_portion
|
|
elif operation == "add":
|
|
output[:, top:bottom, left:right] = destination_portion + source_portion
|
|
elif operation == "subtract":
|
|
output[:, top:bottom, left:right] = destination_portion - source_portion
|
|
elif operation == "and":
|
|
output[:, top:bottom, left:right] = torch.bitwise_and(
|
|
destination_portion.round().bool(), source_portion.round().bool()).float()
|
|
elif operation == "or":
|
|
output[:, top:bottom, left:right] = torch.bitwise_or(
|
|
destination_portion.round().bool(), source_portion.round().bool()).float()
|
|
elif operation == "xor":
|
|
output[:, top:bottom, left:right] = torch.bitwise_xor(
|
|
destination_portion.round().bool(), source_portion.round().bool()).float()
|
|
|
|
output = torch.clamp(output, 0.0, 1.0)
|
|
return output
|
|
|
|
|
|
def create_mask(data, width, height):
|
|
destination = solid_mask(0.0, width, height)
|
|
source = solid_mask(1.0, data[-2], data[-1])
|
|
return mask_combine(destination, source, data[0], data[1], operation="add")
|
|
|
|
# 在大图片上添加小图片的函数
|
|
|
|
|
|
def add_image_to_background(background_image, small_image, location, size, angle, flipX=False, flipY=False):
|
|
# 将小图片重新缩放
|
|
resized_image = cv2.resize(small_image, size)
|
|
|
|
# 如果需要翻转小图片，则进行翻转
|
|
if flipX:
|
|
resized_image = cv2.flip(resized_image, 1) # 1 表示水平翻转
|
|
if flipY:
|
|
resized_image = cv2.flip(resized_image, 0) # 0 表示垂直翻转
|
|
|
|
# 将小图片旋转指定角度
|
|
(h, w) = resized_image.shape[:2]
|
|
center = (w // 2, h // 2)
|
|
M = cv2.getRotationMatrix2D(center, angle*-1, 1.0)
|
|
rotated_image = cv2.warpAffine(resized_image, M, (w, h))
|
|
|
|
# 计算小图片的放置位置
|
|
(bx, by, bn) = background_image.shape
|
|
(sy, sx, sn) = rotated_image.shape
|
|
start_x = location[0]
|
|
start_y = location[1]
|
|
end_x = start_x + sx
|
|
end_y = start_y + sy
|
|
|
|
# 确保小图片不会超出大图片的边界
|
|
if start_x < 0:
|
|
dstart_x = -start_x
|
|
start_x = 0
|
|
end_x -= dstart_x
|
|
if start_y < 0:
|
|
dstart_y = -start_y
|
|
start_y = 0
|
|
end_y -= dstart_y
|
|
if end_x > bx:
|
|
dend_x = end_x - bx
|
|
end_x = bx
|
|
start_x -= dend_x
|
|
if end_y > by:
|
|
dend_y = end_y - by
|
|
end_y = by
|
|
start_y -= dend_y
|
|
|
|
background_image[start_y:end_y,
|
|
start_x:end_x] = rotated_image[0:end_y-start_y, 0:end_x-start_x]
|
|
|
|
return background_image
|
|
|
|
|
|
class openPoseEditorPlus:
|
|
@classmethod
|
|
def INPUT_TYPES(self):
|
|
work_dir = folder_paths.get_input_directory()
|
|
images = [img for img in os.listdir(work_dir) if os.path.isfile(os.path.join(work_dir, img))]
|
|
|
|
return {"required":
|
|
{
|
|
"image": (sorted(images), ),
|
|
"index": ("INT", {"default": 0, "min": 0, "max": 99999}),
|
|
"width": ("INT", {"default": 512, "min": 0, "max": 99999}),
|
|
"height": ("INT", {"default": 512, "min": 0, "max": 99999}),
|
|
},
|
|
"hidden": {"unique_id": "UNIQUE_ID", "wprompt": "PROMPT"},
|
|
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE","IMAGE", "INT", "INT", "MASKS", "MASKS","BOXS",)
|
|
RETURN_NAMES = ("image", "handImg","painterImg","width",
|
|
"height", "head_masks", "body_masks","body_boxs")
|
|
FUNCTION = "output_pose"
|
|
|
|
CATEGORY = "lam"
|
|
|
|
def output_pose(self,image, width, height, unique_id, wprompt, **kwargs):
|
|
groups = []
|
|
hands = []
|
|
if unique_id in wprompt:
|
|
if wprompt[unique_id]["inputs"]:
|
|
for arg in wprompt[unique_id]["inputs"]:
|
|
if arg.startswith("wPose_"):
|
|
data = json.loads(wprompt[unique_id]["inputs"][arg])
|
|
groups = data["groups"]
|
|
hands = data["hands"]
|
|
break
|
|
head_masks = []
|
|
body_masks = []
|
|
body_boxs = []
|
|
candidate = []
|
|
subsets = []
|
|
index = 0
|
|
for g in groups:
|
|
if isinstance(g, dict):
|
|
min_x, min_y, mwidth, mheight = (int(g['left']), int(
|
|
g['top']), int(g['width']), int(g['height']))
|
|
if min_x < 0:
|
|
mwidth = mwidth-min_x
|
|
if min_y < 0:
|
|
mheight = mheight-min_y
|
|
if min_x+mwidth > width:
|
|
mwidth = width-min_x
|
|
if min_y+mheight > height:
|
|
mheight = height-min_y
|
|
min_x = min_x if min_x > 0 else 0
|
|
min_y = min_y if min_y > 0 else 0
|
|
if mwidth <= 0 or mheight <= 0:
|
|
continue
|
|
body_masks.append(create_mask((min_x, min_y, mwidth, mheight), width, height))
|
|
body_boxs.append([mwidth, mheight,min_x, min_y])
|
|
continue
|
|
|
|
head = get_bounding_box(
|
|
g[-4:]+g[:2], width, height, (0.5, 0.5), (1, 0))
|
|
if head:
|
|
head_masks.append(create_mask(head, width, height))
|
|
whole = get_bounding_box(g, width, height, (0.2, 0.2), (0.2, 0.2))
|
|
if whole:
|
|
body_masks.append(create_mask(whole, width, height))
|
|
min_x, min_y, mwidth, mheight=whole
|
|
body_boxs.append([mwidth, mheight,min_x, min_y])
|
|
subset = []
|
|
for i in range(len(g)):
|
|
if g[i][0] < 0 or g[i][1] < 0:
|
|
subset.append(-1)
|
|
else:
|
|
candidate.append(g[i])
|
|
subset.append(index)
|
|
index += 1
|
|
|
|
subsets.append(subset)
|
|
|
|
candidate = [[i[0]/float(width), i[1]/float(height)]
|
|
for i in candidate]
|
|
|
|
canvas = np.zeros(shape=(height, width, 3), dtype=np.uint8)
|
|
canvasHands = np.zeros(shape=(height, width, 3), dtype=np.uint8)
|
|
for i in range(len(hands)):
|
|
zhand = hands[i]
|
|
for hand in zhand:
|
|
imgPath = os.path.join(dir, 'hands', hand['name'])
|
|
small = cv2.imread(imgPath) # 小图片路径
|
|
(h, w) = small.shape[:2]
|
|
# 小图片的新尺寸
|
|
size = (int(w*abs(float(hand['scaleX']))),
|
|
int(h*abs(float(hand['scaleY']))))
|
|
location = (int(hand['x']-size[0]/2),
|
|
int(hand['y']-size[1]/2)) # 小图片的左上角位置
|
|
flipX = hand['flipX']
|
|
flipY = hand['flipY']
|
|
angle = float(hand['angle'])
|
|
add_image_to_background(
|
|
canvasHands, small, location, size, angle, flipX, flipY)
|
|
poseImage = draw_bodypose(canvas, candidate, subsets)
|
|
poseImage = np.array(poseImage).astype(np.float32) / 255.0
|
|
poseImage = torch.from_numpy(poseImage)[None,]
|
|
|
|
handImage = np.array(canvasHands).astype(np.float32) / 255.0
|
|
handImage = torch.from_numpy(handImage)[None,]
|
|
|
|
image_path = folder_paths.get_annotated_filepath(image)
|
|
|
|
if os.path.exists(image_path):
|
|
i = Image.open(image_path)
|
|
i = ImageOps.exif_transpose(i)
|
|
painterImage = i.convert("RGB")
|
|
else:
|
|
painterImage = np.zeros(shape=(height, width, 3), dtype=np.uint8)
|
|
|
|
painterImage = np.array(painterImage).astype(np.float32) / 255.0
|
|
painterImage = torch.from_numpy(painterImage)[None,]
|
|
|
|
return (poseImage, handImage,painterImage, width, height, head_masks, body_masks,body_boxs,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(self, image):
|
|
image_path = os.path.join(
|
|
folder_paths.get_temp_directory(), image)
|
|
# print(f'Change: {image_path}')
|
|
|
|
m = hashlib.sha256()
|
|
with open(image_path, 'rb') as f:
|
|
m.update(f.read())
|
|
return m.digest().hex()
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LAM.OpenPoseEditorPlus": openPoseEditorPlus
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"LAM.OpenPoseEditorPlus": "姿态编辑器",
|
|
}
|