From b7d7f9afe5b0cab01b7214e6191aee2a1039f186 Mon Sep 17 00:00:00 2001 From: kijai <40791699+kijai@users.noreply.github.com> Date: Sat, 23 Aug 2025 18:59:21 +0300 Subject: [PATCH] cleanup unianimate stuff --- unianimate/dwpose/onnxdet.py | 127 ------------ unianimate/dwpose/onnxpose.py | 360 --------------------------------- unianimate/dwpose/util.py | 16 +- unianimate/dwpose/wholebody.py | 1 - unianimate/nodes.py | 21 +- 5 files changed, 17 insertions(+), 508 deletions(-) delete mode 100644 unianimate/dwpose/onnxdet.py delete mode 100644 unianimate/dwpose/onnxpose.py diff --git a/unianimate/dwpose/onnxdet.py b/unianimate/dwpose/onnxdet.py deleted file mode 100644 index 15ae797..0000000 --- a/unianimate/dwpose/onnxdet.py +++ /dev/null @@ -1,127 +0,0 @@ -import cv2 -import numpy as np - -import onnxruntime - -def nms(boxes, scores, nms_thr): - """Single class NMS implemented in Numpy.""" - x1 = boxes[:, 0] - y1 = boxes[:, 1] - x2 = boxes[:, 2] - y2 = boxes[:, 3] - - areas = (x2 - x1 + 1) * (y2 - y1 + 1) - order = scores.argsort()[::-1] - - keep = [] - while order.size > 0: - i = order[0] - keep.append(i) - xx1 = np.maximum(x1[i], x1[order[1:]]) - yy1 = np.maximum(y1[i], y1[order[1:]]) - xx2 = np.minimum(x2[i], x2[order[1:]]) - yy2 = np.minimum(y2[i], y2[order[1:]]) - - w = np.maximum(0.0, xx2 - xx1 + 1) - h = np.maximum(0.0, yy2 - yy1 + 1) - inter = w * h - ovr = inter / (areas[i] + areas[order[1:]] - inter) - - inds = np.where(ovr <= nms_thr)[0] - order = order[inds + 1] - - return keep - -def multiclass_nms(boxes, scores, nms_thr, score_thr): - """Multiclass NMS implemented in Numpy. Class-aware version.""" - final_dets = [] - num_classes = scores.shape[1] - for cls_ind in range(num_classes): - cls_scores = scores[:, cls_ind] - valid_score_mask = cls_scores > score_thr - if valid_score_mask.sum() == 0: - continue - else: - valid_scores = cls_scores[valid_score_mask] - valid_boxes = boxes[valid_score_mask] - keep = nms(valid_boxes, valid_scores, nms_thr) - if len(keep) > 0: - cls_inds = np.ones((len(keep), 1)) * cls_ind - dets = np.concatenate( - [valid_boxes[keep], valid_scores[keep, None], cls_inds], 1 - ) - final_dets.append(dets) - if len(final_dets) == 0: - return None - return np.concatenate(final_dets, 0) - -def demo_postprocess(outputs, img_size, p6=False): - grids = [] - expanded_strides = [] - strides = [8, 16, 32] if not p6 else [8, 16, 32, 64] - - hsizes = [img_size[0] // stride for stride in strides] - wsizes = [img_size[1] // stride for stride in strides] - - for hsize, wsize, stride in zip(hsizes, wsizes, strides): - xv, yv = np.meshgrid(np.arange(wsize), np.arange(hsize)) - grid = np.stack((xv, yv), 2).reshape(1, -1, 2) - grids.append(grid) - shape = grid.shape[:2] - expanded_strides.append(np.full((*shape, 1), stride)) - - grids = np.concatenate(grids, 1) - expanded_strides = np.concatenate(expanded_strides, 1) - outputs[..., :2] = (outputs[..., :2] + grids) * expanded_strides - outputs[..., 2:4] = np.exp(outputs[..., 2:4]) * expanded_strides - - return outputs - -def preprocess(img, input_size, swap=(2, 0, 1)): - if len(img.shape) == 3: - padded_img = np.ones((input_size[0], input_size[1], 3), dtype=np.uint8) * 114 - else: - padded_img = np.ones(input_size, dtype=np.uint8) * 114 - - r = min(input_size[0] / img.shape[0], input_size[1] / img.shape[1]) - resized_img = cv2.resize( - img, - (int(img.shape[1] * r), int(img.shape[0] * r)), - interpolation=cv2.INTER_LINEAR, - ).astype(np.uint8) - padded_img[: int(img.shape[0] * r), : int(img.shape[1] * r)] = resized_img - - padded_img = padded_img.transpose(swap) - padded_img = np.ascontiguousarray(padded_img, dtype=np.float32) - return padded_img, r - -def inference_detector(session, oriImg): - input_shape = (640,640) - img, ratio = preprocess(oriImg, input_shape) - - ort_inputs = {session.get_inputs()[0].name: img[None, :, :, :]} - - output = session.run(None, ort_inputs) - - predictions = demo_postprocess(output[0], input_shape)[0] - - boxes = predictions[:, :4] - scores = predictions[:, 4:5] * predictions[:, 5:] - - boxes_xyxy = np.ones_like(boxes) - boxes_xyxy[:, 0] = boxes[:, 0] - boxes[:, 2]/2. - boxes_xyxy[:, 1] = boxes[:, 1] - boxes[:, 3]/2. - boxes_xyxy[:, 2] = boxes[:, 0] + boxes[:, 2]/2. - boxes_xyxy[:, 3] = boxes[:, 1] + boxes[:, 3]/2. - boxes_xyxy /= ratio - dets = multiclass_nms(boxes_xyxy, scores, nms_thr=0.45, score_thr=0.1) - if dets is not None: - final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5] - isscore = final_scores>0.3 - iscat = final_cls_inds == 0 - isbbox = [ i and j for (i, j) in zip(isscore, iscat)] - final_boxes = final_boxes[isbbox] - else: - final_boxes = np.array([]) - - return final_boxes diff --git a/unianimate/dwpose/onnxpose.py b/unianimate/dwpose/onnxpose.py deleted file mode 100644 index 79cd4a0..0000000 --- a/unianimate/dwpose/onnxpose.py +++ /dev/null @@ -1,360 +0,0 @@ -from typing import List, Tuple - -import cv2 -import numpy as np -import onnxruntime as ort - -def preprocess( - img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256) -) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: - """Do preprocessing for RTMPose model inference. - - Args: - img (np.ndarray): Input image in shape. - input_size (tuple): Input image size in shape (w, h). - - Returns: - tuple: - - resized_img (np.ndarray): Preprocessed image. - - center (np.ndarray): Center of image. - - scale (np.ndarray): Scale of image. - """ - # get shape of image - img_shape = img.shape[:2] - out_img, out_center, out_scale = [], [], [] - if len(out_bbox) == 0: - out_bbox = [[0, 0, img_shape[1], img_shape[0]]] - for i in range(len(out_bbox)): - x0 = out_bbox[i][0] - y0 = out_bbox[i][1] - x1 = out_bbox[i][2] - y1 = out_bbox[i][3] - bbox = np.array([x0, y0, x1, y1]) - - # get center and scale - center, scale = bbox_xyxy2cs(bbox, padding=1.25) - - # do affine transformation - resized_img, scale = top_down_affine(input_size, scale, center, img) - - # normalize image - mean = np.array([123.675, 116.28, 103.53]) - std = np.array([58.395, 57.12, 57.375]) - resized_img = (resized_img - mean) / std - - out_img.append(resized_img) - out_center.append(center) - out_scale.append(scale) - - return out_img, out_center, out_scale - - -def inference(sess: ort.InferenceSession, img: np.ndarray) -> np.ndarray: - """Inference RTMPose model. - - Args: - sess (ort.InferenceSession): ONNXRuntime session. - img (np.ndarray): Input image in shape. - - Returns: - outputs (np.ndarray): Output of RTMPose model. - """ - all_out = [] - # build input - for i in range(len(img)): - input = [img[i].transpose(2, 0, 1)] - - # build output - sess_input = {sess.get_inputs()[0].name: input} - sess_output = [] - for out in sess.get_outputs(): - sess_output.append(out.name) - - # run model - outputs = sess.run(sess_output, sess_input) - all_out.append(outputs) - - return all_out - - -def postprocess(outputs: List[np.ndarray], - model_input_size: Tuple[int, int], - center: Tuple[int, int], - scale: Tuple[int, int], - simcc_split_ratio: float = 2.0 - ) -> Tuple[np.ndarray, np.ndarray]: - """Postprocess for RTMPose model output. - - Args: - outputs (np.ndarray): Output of RTMPose model. - model_input_size (tuple): RTMPose model Input image size. - center (tuple): Center of bbox in shape (x, y). - scale (tuple): Scale of bbox in shape (w, h). - simcc_split_ratio (float): Split ratio of simcc. - - Returns: - tuple: - - keypoints (np.ndarray): Rescaled keypoints. - - scores (np.ndarray): Model predict scores. - """ - all_key = [] - all_score = [] - for i in range(len(outputs)): - # use simcc to decode - simcc_x, simcc_y = outputs[i] - keypoints, scores = decode(simcc_x, simcc_y, simcc_split_ratio) - - # rescale keypoints - keypoints = keypoints / model_input_size * scale[i] + center[i] - scale[i] / 2 - all_key.append(keypoints[0]) - all_score.append(scores[0]) - - return np.array(all_key), np.array(all_score) - - -def bbox_xyxy2cs(bbox: np.ndarray, - padding: float = 1.) -> Tuple[np.ndarray, np.ndarray]: - """Transform the bbox format from (x,y,w,h) into (center, scale) - - Args: - bbox (ndarray): Bounding box(es) in shape (4,) or (n, 4), formatted - as (left, top, right, bottom) - padding (float): BBox padding factor that will be multilied to scale. - Default: 1.0 - - Returns: - tuple: A tuple containing center and scale. - - np.ndarray[float32]: Center (x, y) of the bbox in shape (2,) or - (n, 2) - - np.ndarray[float32]: Scale (w, h) of the bbox in shape (2,) or - (n, 2) - """ - # convert single bbox from (4, ) to (1, 4) - dim = bbox.ndim - if dim == 1: - bbox = bbox[None, :] - - # get bbox center and scale - x1, y1, x2, y2 = np.hsplit(bbox, [1, 2, 3]) - center = np.hstack([x1 + x2, y1 + y2]) * 0.5 - scale = np.hstack([x2 - x1, y2 - y1]) * padding - - if dim == 1: - center = center[0] - scale = scale[0] - - return center, scale - - -def _fix_aspect_ratio(bbox_scale: np.ndarray, - aspect_ratio: float) -> np.ndarray: - """Extend the scale to match the given aspect ratio. - - Args: - scale (np.ndarray): The image scale (w, h) in shape (2, ) - aspect_ratio (float): The ratio of ``w/h`` - - Returns: - np.ndarray: The reshaped image scale in (2, ) - """ - w, h = np.hsplit(bbox_scale, [1]) - bbox_scale = np.where(w > h * aspect_ratio, - np.hstack([w, w / aspect_ratio]), - np.hstack([h * aspect_ratio, h])) - return bbox_scale - - -def _rotate_point(pt: np.ndarray, angle_rad: float) -> np.ndarray: - """Rotate a point by an angle. - - Args: - pt (np.ndarray): 2D point coordinates (x, y) in shape (2, ) - angle_rad (float): rotation angle in radian - - Returns: - np.ndarray: Rotated point in shape (2, ) - """ - sn, cs = np.sin(angle_rad), np.cos(angle_rad) - rot_mat = np.array([[cs, -sn], [sn, cs]]) - return rot_mat @ pt - - -def _get_3rd_point(a: np.ndarray, b: np.ndarray) -> np.ndarray: - """To calculate the affine matrix, three pairs of points are required. This - function is used to get the 3rd point, given 2D points a & b. - - The 3rd point is defined by rotating vector `a - b` by 90 degrees - anticlockwise, using b as the rotation center. - - Args: - a (np.ndarray): The 1st point (x,y) in shape (2, ) - b (np.ndarray): The 2nd point (x,y) in shape (2, ) - - Returns: - np.ndarray: The 3rd point. - """ - direction = a - b - c = b + np.r_[-direction[1], direction[0]] - return c - - -def get_warp_matrix(center: np.ndarray, - scale: np.ndarray, - rot: float, - output_size: Tuple[int, int], - shift: Tuple[float, float] = (0., 0.), - inv: bool = False) -> np.ndarray: - """Calculate the affine transformation matrix that can warp the bbox area - in the input image to the output size. - - Args: - center (np.ndarray[2, ]): Center of the bounding box (x, y). - scale (np.ndarray[2, ]): Scale of the bounding box - wrt [width, height]. - rot (float): Rotation angle (degree). - output_size (np.ndarray[2, ] | list(2,)): Size of the - destination heatmaps. - shift (0-100%): Shift translation ratio wrt the width/height. - Default (0., 0.). - inv (bool): Option to inverse the affine transform direction. - (inv=False: src->dst or inv=True: dst->src) - - Returns: - np.ndarray: A 2x3 transformation matrix - """ - shift = np.array(shift) - src_w = scale[0] - dst_w = output_size[0] - dst_h = output_size[1] - - # compute transformation matrix - rot_rad = np.deg2rad(rot) - src_dir = _rotate_point(np.array([0., src_w * -0.5]), rot_rad) - dst_dir = np.array([0., dst_w * -0.5]) - - # get four corners of the src rectangle in the original image - src = np.zeros((3, 2), dtype=np.float32) - src[0, :] = center + scale * shift - src[1, :] = center + src_dir + scale * shift - src[2, :] = _get_3rd_point(src[0, :], src[1, :]) - - # get four corners of the dst rectangle in the input image - dst = np.zeros((3, 2), dtype=np.float32) - dst[0, :] = [dst_w * 0.5, dst_h * 0.5] - dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir - dst[2, :] = _get_3rd_point(dst[0, :], dst[1, :]) - - if inv: - warp_mat = cv2.getAffineTransform(np.float32(dst), np.float32(src)) - else: - warp_mat = cv2.getAffineTransform(np.float32(src), np.float32(dst)) - - return warp_mat - - -def top_down_affine(input_size: dict, bbox_scale: dict, bbox_center: dict, - img: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: - """Get the bbox image as the model input by affine transform. - - Args: - input_size (dict): The input size of the model. - bbox_scale (dict): The bbox scale of the img. - bbox_center (dict): The bbox center of the img. - img (np.ndarray): The original image. - - Returns: - tuple: A tuple containing center and scale. - - np.ndarray[float32]: img after affine transform. - - np.ndarray[float32]: bbox scale after affine transform. - """ - w, h = input_size - warp_size = (int(w), int(h)) - - # reshape bbox to fixed aspect ratio - bbox_scale = _fix_aspect_ratio(bbox_scale, aspect_ratio=w / h) - - # get the affine matrix - center = bbox_center - scale = bbox_scale - rot = 0 - warp_mat = get_warp_matrix(center, scale, rot, output_size=(w, h)) - - # do affine transform - img = cv2.warpAffine(img, warp_mat, warp_size, flags=cv2.INTER_LINEAR) - - return img, bbox_scale - - -def get_simcc_maximum(simcc_x: np.ndarray, - simcc_y: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: - """Get maximum response location and value from simcc representations. - - Note: - instance number: N - num_keypoints: K - heatmap height: H - heatmap width: W - - Args: - simcc_x (np.ndarray): x-axis SimCC in shape (K, Wx) or (N, K, Wx) - simcc_y (np.ndarray): y-axis SimCC in shape (K, Wy) or (N, K, Wy) - - Returns: - tuple: - - locs (np.ndarray): locations of maximum heatmap responses in shape - (K, 2) or (N, K, 2) - - vals (np.ndarray): values of maximum heatmap responses in shape - (K,) or (N, K) - """ - N, K, Wx = simcc_x.shape - simcc_x = simcc_x.reshape(N * K, -1) - simcc_y = simcc_y.reshape(N * K, -1) - - # get maximum value locations - x_locs = np.argmax(simcc_x, axis=1) - y_locs = np.argmax(simcc_y, axis=1) - locs = np.stack((x_locs, y_locs), axis=-1).astype(np.float32) - max_val_x = np.amax(simcc_x, axis=1) - max_val_y = np.amax(simcc_y, axis=1) - - # get maximum value across x and y axis - mask = max_val_x > max_val_y - max_val_x[mask] = max_val_y[mask] - vals = max_val_x - locs[vals <= 0.] = -1 - - # reshape - locs = locs.reshape(N, K, 2) - vals = vals.reshape(N, K) - - return locs, vals - - -def decode(simcc_x: np.ndarray, simcc_y: np.ndarray, - simcc_split_ratio) -> Tuple[np.ndarray, np.ndarray]: - """Modulate simcc distribution with Gaussian. - - Args: - simcc_x (np.ndarray[K, Wx]): model predicted simcc in x. - simcc_y (np.ndarray[K, Wy]): model predicted simcc in y. - simcc_split_ratio (int): The split ratio of simcc. - - Returns: - tuple: A tuple containing center and scale. - - np.ndarray[float32]: keypoints in shape (K, 2) or (n, K, 2) - - np.ndarray[float32]: scores in shape (K,) or (n, K) - """ - keypoints, scores = get_simcc_maximum(simcc_x, simcc_y) - keypoints /= simcc_split_ratio - - return keypoints, scores - - -def inference_pose(session, out_bbox, oriImg): - h, w = session.get_inputs()[0].shape[2:] - model_input_size = (w, h) - resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size) - outputs = inference(session, resized_img) - keypoints, scores = postprocess(outputs, model_input_size, center, scale) - - return keypoints, scores \ No newline at end of file diff --git a/unianimate/dwpose/util.py b/unianimate/dwpose/util.py index f721442..c8aa250 100644 --- a/unianimate/dwpose/util.py +++ b/unianimate/dwpose/util.py @@ -180,12 +180,11 @@ def draw_body_and_foot(canvas, candidate, subset, score, stick_width=4, draw_bod # Append head elements based on the condition limbSeq_and_colors += head_elements - for limb_info in limbSeq_and_colors[:17]: + for limb_info in limbSeq_and_colors[:19]: limbSeq, color = limb_info for n in range(len(subset)): index = subset[n][np.array(limbSeq) - 1] - conf = score[n][np.array(limbSeq) - 1] - if conf[0] < 0.3 or conf[1] < 0.3: + if index[0] < 0.3 or index[1] < 0.3: continue Y = candidate[index.astype(int), 0] * float(W) X = candidate[index.astype(int), 1] * float(H) @@ -194,11 +193,11 @@ def draw_body_and_foot(canvas, candidate, subset, score, stick_width=4, draw_bod length = np.sqrt((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stick_width), int(angle), 0, 360, 1) - cv2.fillConvexPoly(canvas, polygon, alpha_blend_color(color, conf[0] * conf[1])) + cv2.fillConvexPoly(canvas, polygon, alpha_blend_color(color, index[0] * index[1])) canvas = (canvas * 0.6).astype(np.uint8) - for limb_info in limbSeq_and_colors[:18]: + for limb_info in limbSeq_and_colors[:19]: limbSeq, color = limb_info for i in limbSeq: for n in range(len(subset)): @@ -209,12 +208,9 @@ def draw_body_and_foot(canvas, candidate, subset, score, stick_width=4, draw_bod conf = score[n][i - 1] if not np.isfinite(x) or not np.isfinite(y): continue - x = int(np.clip(x * W, 0, W - 1 - )) + x = int(np.clip(x * W, 0, W - 1)) y = int(np.clip(y * H, 0, H - 1)) - # x = int(x * W) - # y = int(y * H) - cv2.circle(canvas, (x, y), 4, alpha_blend_color(color, conf), thickness=-1) + cv2.circle(canvas, (x, y), body_keypoint_size, alpha_blend_color(color, conf), thickness=-1) return canvas diff --git a/unianimate/dwpose/wholebody.py b/unianimate/dwpose/wholebody.py index b838661..ab2022c 100644 --- a/unianimate/dwpose/wholebody.py +++ b/unianimate/dwpose/wholebody.py @@ -1,7 +1,6 @@ import numpy as np from .jit_det import inference_detector as inference_jit_yolox from .jit_pose import inference_pose as inference_jit_pose -import os class Wholebody: diff --git a/unianimate/nodes.py b/unianimate/nodes.py index eb54de4..9fb97d0 100644 --- a/unianimate/nodes.py +++ b/unianimate/nodes.py @@ -1,9 +1,14 @@ - +import torch import torch.nn as nn +import os, copy, math +import numpy as np +from tqdm import tqdm + from ..utils import log + import comfy.model_management as mm from comfy.utils import ProgressBar -from tqdm import tqdm + def update_transformer(transformer, state_dict): @@ -56,14 +61,6 @@ def update_transformer(transformer, state_dict): # 3rd Edited by ControlNet # 4th Edited by ControlNet (added face and correct hands) -import os -import torch -import numpy as np -import copy -import torch -import numpy as np -import math - from .dwpose.wholebody import Wholebody def smoothing_factor(t_e, cutoff): @@ -772,10 +769,14 @@ class WanVideoUniAnimateDWPoseDetector: ref = reference_pose_image ref_np = ref.cpu().numpy() * 255 + prev_fuser_state = torch._C._jit_texpr_fuser_enabled() + torch._C._jit_set_texpr_fuser_enabled(False) # removes warmup delay, may want to enable later poses, reference_pose = pose_extract(pose_np, ref_np, self.dwpose_detector, height, width, score_threshold, stick_width=stick_width, draw_body=draw_body, body_keypoint_size=body_keypoint_size, draw_feet=draw_feet, draw_hands=draw_hands, hand_keypoint_size=hand_keypoint_size, handle_not_detected=handle_not_detected, draw_head=draw_head) poses = poses / 255.0 + torch._C._jit_set_texpr_fuser_enabled(prev_fuser_state) + if reference_pose_image is not None: reference_pose = reference_pose.unsqueeze(0) / 255.0 else: