106 lines
4.4 KiB
Python
106 lines
4.4 KiB
Python
import cv2
|
|
import mediapipe as mp
|
|
|
|
async def listen_node_output(prompt_id, target_node_id):
|
|
uri = "ws://localhost:8188/ws"
|
|
async with websockets.connect(uri) as websocket:
|
|
await websocket.send(json.dumps({"prompt_id": prompt_id}))
|
|
|
|
while True:
|
|
message = await websocket.recv()
|
|
|
|
# 1. 处理二进制数据
|
|
if isinstance(message, bytes):
|
|
try:
|
|
message = message.decode("utf-8")
|
|
except UnicodeDecodeError:
|
|
print(f"[{datetime.now().strftime('%H:%M:%S')}] 收到二进制数据 (长度: {len(message)} bytes)")
|
|
continue
|
|
|
|
# 2. 解析JSON
|
|
try:
|
|
data = json.loads(message)
|
|
print(f"cur data {data}")
|
|
except json.JSONDecodeError:
|
|
print(f"[{datetime.now().strftime('%H:%M:%S')}] 非JSON数据: {message[:100]}...")
|
|
continue
|
|
|
|
if data.get("type") == "progress":
|
|
pdata = data.get('data', {})
|
|
print(f"cur progress {pdata}")
|
|
# if value >= max:
|
|
# return f"finished value:{value} max:{max}"
|
|
|
|
if data.get("type") == "progress_state":
|
|
nodes = data.get('data', {}).get('nodes', {})
|
|
for node in nodes:
|
|
if node.get("state", None) != 'finished':
|
|
print(f"progress_state node:{node}")
|
|
|
|
# 3. 输出运行状态
|
|
if data.get("type") == "status":
|
|
print(f"[{datetime.now().strftime('%H:%M:%S')}] 系统状态: {data.get('data', {}).get('status', {})}")
|
|
status = data.get('data', {}).get('status', None)
|
|
print(f"statue:{status}")
|
|
if status:
|
|
exec_info = status.get('exec_info', None)
|
|
print(f"exec_info:{exec_info}")
|
|
if exec_info:
|
|
queue_remaining = exec_info.get('queue_remaining', None)
|
|
print(f"queue_remaining:{queue_remaining}")
|
|
if queue_remaining == 0:
|
|
return "process end"
|
|
continue
|
|
|
|
if data.get("type") == "executing":
|
|
node_id = data.get("data", {}).get("node")
|
|
progress = data.get("data", {}).get("progress", 0)
|
|
print(f"[{datetime.now().strftime('%H:%M:%S')}] 正在执行节点 {node_id} (进度: {progress:.0%})")
|
|
|
|
# 4. 目标节点完成
|
|
if data.get("type") == "executed" and data.get("data", {}).get("node") == target_node_id:
|
|
output = data.get("data", {}).get("output")
|
|
print(f"[{datetime.now().strftime('%H:%M:%S')}] 节点 {target_node_id} 完成!")
|
|
return output
|
|
|
|
|
|
|
|
# output = asyncio.run(listen_node_output(prompt_id, "104"))
|
|
# print("最终输出:", output)
|
|
|
|
# def is_thigh_visible(image_path):
|
|
# # 初始化MediaPipe Pose模型
|
|
# mp_pose = mp.solutions.pose
|
|
# pose = mp_pose.Pose(static_image_mode=True, min_detection_confidence=0.5)
|
|
|
|
# # 读取图像
|
|
# image = cv2.imread(image_path)
|
|
# image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
|
|
|
# # 进行姿态估计
|
|
# results = pose.process(image_rgb)
|
|
|
|
# if not results.pose_landmarks:
|
|
# return False # 未检测到人体
|
|
|
|
# # 获取关键点
|
|
# landmarks = results.pose_landmarks.landmark
|
|
|
|
# # 检查髋部和膝盖之间的关键点（大腿部分）
|
|
# # MediaPipe关键点索引：
|
|
# # 23: 左髋, 24: 右髋
|
|
# # 25: 左膝, 26: 右膝
|
|
|
|
# # 检查左大腿是否可见（髋部和膝盖之间）
|
|
# left_hip_visible = landmarks[mp_pose.PoseLandmark.LEFT_HIP.value].visibility > 0.6
|
|
# left_knee_visible = landmarks[mp_pose.PoseLandmark.LEFT_KNEE.value].visibility > 0.6
|
|
# left_thigh_visible = left_hip_visible and left_knee_visible
|
|
|
|
# # 检查右大腿是否可见
|
|
# right_hip_visible = landmarks[mp_pose.PoseLandmark.RIGHT_HIP.value].visibility > 0.6
|
|
# right_knee_visible = landmarks[mp_pose.PoseLandmark.RIGHT_KNEE.value].visibility > 0.6
|
|
# right_thigh_visible = right_hip_visible and right_knee_visible
|
|
|
|
# return left_thigh_visible and right_thigh_visible
|
|
|
|
# print(is_thigh_visible("/home/xsl/code/ComfyUI/change_cloth/static/imgs/fc74eded_00001_.jpg")) |