commit f76dd0c32555690391b04de63496245a81a114ce Author: gear_machine Date: Sat Aug 16 12:37:26 2025 +0900 Initial commit: ProportionChanger ComfyUI custom node - Rename WanVideoUniAnimateDWPoseDetector to ProportionChangerUniAnimateDWPoseDetector - Update category from "WanVideoWrapper" to "ProportionChanger" - Update display name to "ProportionChanger UniAnimate DWPose Detector" - Remove unused WanVideoUniAnimatePoseInput node - Add development gitignore and example_workflows structure Based on kijai's WanVideoWrapper - adapted for body proportion changing. diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3854ff2 --- /dev/null +++ b/.gitignore @@ -0,0 +1,18 @@ +output/ +*__pycache__/ +samples*/ +runs/ +checkpoints/ +master_ip +logs/ +*.DS_Store +.idea +tools/ +.vscode/ +convert_* +*.pt + +# Development files +issues/ +.claude/ +CLAUDE.md \ No newline at end of file diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..261eeb9 --- /dev/null +++ b/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..22cfb73 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .unianimate.nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file diff --git a/example_workflows/.gitkeep b/example_workflows/.gitkeep new file mode 100644 index 0000000..5aad412 --- /dev/null +++ b/example_workflows/.gitkeep @@ -0,0 +1,3 @@ +# Example workflows directory + +This directory will contain example workflows for the ProportionChanger nodes. \ No newline at end of file diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..48f64e5 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,9 @@ +ftfy +accelerate>=1.2.1 +einops +diffusers>=0.33.0 +peft>=0.15.0 +sentencepiece>=0.2.0 +protobuf +pyloudnorm +gguf>=0.14.0 \ No newline at end of file diff --git a/unianimate/dwpose/__init__.py b/unianimate/dwpose/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/unianimate/dwpose/jit_det.py b/unianimate/dwpose/jit_det.py new file mode 100644 index 0000000..2ae8635 --- /dev/null +++ b/unianimate/dwpose/jit_det.py @@ -0,0 +1,125 @@ +import cv2 +import numpy as np +import torch + +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(model, oriImg, detect_classes=[0]): + input_shape = (640,640) + img, ratio = preprocess(oriImg, input_shape) + + device, dtype = next(model.parameters()).device, next(model.parameters()).dtype + input = img[None, :, :, :] + input = torch.from_numpy(input).to(device, dtype) + + output = model(input).float().cpu().detach().numpy() + predictions = demo_postprocess(output[0], input_shape) + + 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 None: + return None + final_boxes, final_scores, final_cls_inds = dets[:, :4], dets[:, 4], dets[:, 5] + isscore = final_scores>0.3 + iscat = np.isin(final_cls_inds, detect_classes) + isbbox = [ i and j for (i, j) in zip(isscore, iscat)] + final_boxes = final_boxes[isbbox] + return final_boxes \ No newline at end of file diff --git a/unianimate/dwpose/jit_pose.py b/unianimate/dwpose/jit_pose.py new file mode 100644 index 0000000..7a7b6ac --- /dev/null +++ b/unianimate/dwpose/jit_pose.py @@ -0,0 +1,363 @@ +from typing import List, Tuple + +import cv2 +import numpy as np +import torch + +def preprocess( + img: np.ndarray, out_bbox, input_size: Tuple[int, int] = (192, 256) +) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: + """Do preprocessing for DWPose 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(model, img, bs=5): + """Inference DWPose model implemented in TorchScript. + + Args: + model : TorchScript Model. + img : Input image in shape. + + Returns: + outputs : Output of DWPose model. + """ + all_out = [] + # build input + orig_img_count = len(img) + #Pad zeros to fit batch size + for _ in range(bs - (orig_img_count % bs)): + img.append(np.zeros_like(img[0])) + input = np.stack(img, axis=0).transpose(0, 3, 1, 2) + device, dtype = next(model.parameters()).device, next(model.parameters()).dtype + input = torch.from_numpy(input).to(device, dtype) + + out1, out2 = [], [] + for i in range(input.shape[0] // bs): + curr_batch_output = model(input[i*bs:(i+1)*bs]) + out1.append(curr_batch_output[0].float()) + out2.append(curr_batch_output[1].float()) + out1, out2 = torch.cat(out1, dim=0)[:orig_img_count], torch.cat(out2, dim=0)[:orig_img_count] + out1, out2 = out1.float().cpu().detach().numpy(), out2.float().cpu().detach().numpy() + all_outputs = out1, out2 + + for batch_idx in range(len(all_outputs[0])): + outputs = [all_outputs[i][batch_idx:batch_idx+1,...] for i in range(len(all_outputs))] + 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 DWPose 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(model, out_bbox, oriImg, model_input_size=(288, 384)): + resized_img, center, scale = preprocess(oriImg, out_bbox, model_input_size) + #outputs = inference(session, resized_img, dtype) + outputs = inference(model, 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/onnxdet.py b/unianimate/dwpose/onnxdet.py new file mode 100644 index 0000000..15ae797 --- /dev/null +++ b/unianimate/dwpose/onnxdet.py @@ -0,0 +1,127 @@ +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 new file mode 100644 index 0000000..79cd4a0 --- /dev/null +++ b/unianimate/dwpose/onnxpose.py @@ -0,0 +1,360 @@ +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 new file mode 100644 index 0000000..ded1ecc --- /dev/null +++ b/unianimate/dwpose/util.py @@ -0,0 +1,401 @@ +import math +import numpy as np +import colorsys +import cv2 + + +eps = 0.01 + + +def smart_resize(x, s): + Ht, Wt = s + if x.ndim == 2: + Ho, Wo = x.shape + Co = 1 + else: + Ho, Wo, Co = x.shape + if Co == 3 or Co == 1: + k = float(Ht + Wt) / float(Ho + Wo) + return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4) + else: + return np.stack([smart_resize(x[:, :, i], s) for i in range(Co)], axis=2) + + +def smart_resize_k(x, fx, fy): + if x.ndim == 2: + Ho, Wo = x.shape + Co = 1 + else: + Ho, Wo, Co = x.shape + Ht, Wt = Ho * fy, Wo * fx + if Co == 3 or Co == 1: + k = float(Ht + Wt) / float(Ho + Wo) + return cv2.resize(x, (int(Wt), int(Ht)), interpolation=cv2.INTER_AREA if k < 1 else cv2.INTER_LANCZOS4) + else: + return np.stack([smart_resize_k(x[:, :, i], fx, fy) for i in range(Co)], axis=2) + + +def padRightDownCorner(img, stride, padValue): + h = img.shape[0] + w = img.shape[1] + + pad = 4 * [None] + pad[0] = 0 # up + pad[1] = 0 # left + pad[2] = 0 if (h % stride == 0) else stride - (h % stride) # down + pad[3] = 0 if (w % stride == 0) else stride - (w % stride) # right + + img_padded = img + pad_up = np.tile(img_padded[0:1, :, :]*0 + padValue, (pad[0], 1, 1)) + img_padded = np.concatenate((pad_up, img_padded), axis=0) + pad_left = np.tile(img_padded[:, 0:1, :]*0 + padValue, (1, pad[1], 1)) + img_padded = np.concatenate((pad_left, img_padded), axis=1) + pad_down = np.tile(img_padded[-2:-1, :, :]*0 + padValue, (pad[2], 1, 1)) + img_padded = np.concatenate((img_padded, pad_down), axis=0) + pad_right = np.tile(img_padded[:, -2:-1, :]*0 + padValue, (1, pad[3], 1)) + img_padded = np.concatenate((img_padded, pad_right), axis=1) + + return img_padded, pad + + +def transfer(model, model_weights): + transfered_model_weights = {} + for weights_name in model.state_dict().keys(): + transfered_model_weights[weights_name] = model_weights['.'.join(weights_name.split('.')[1:])] + return transfered_model_weights + + +def draw_bodypose(canvas, candidate, subset): + H, W, C = canvas.shape + candidate = np.array(candidate) + subset = np.array(subset) + + stickwidth = 4 + + limbSeq = [[2, 3], [2, 6], [3, 4], [4, 5], [6, 7], [7, 8], [2, 9], [9, 10], \ + [10, 11], [2, 12], [12, 13], [13, 14], [2, 1], [1, 15], [15, 17], \ + [1, 16], [16, 18], [3, 17], [6, 18]] + + colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \ + [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \ + [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]] + + for i in range(17): + for n in range(len(subset)): + index = subset[n][np.array(limbSeq[i]) - 1] + if -1 in index: + continue + Y = candidate[index.astype(int), 0] * float(W) + X = candidate[index.astype(int), 1] * float(H) + mX = np.mean(X) + mY = np.mean(Y) + length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5 + angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1])) + polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1) + cv2.fillConvexPoly(canvas, polygon, colors[i]) + + canvas = (canvas * 0.6).astype(np.uint8) + + for i in range(18): + for n in range(len(subset)): + index = int(subset[n][i]) + if index == -1: + continue + x, y = candidate[index][0:2] + x = int(x * W) + y = int(y * H) + cv2.circle(canvas, (int(x), int(y)), 4, colors[i], thickness=-1) + + return canvas + +def alpha_blend_color(color, alpha): + """blend color according to point conf + """ + return [int(c * alpha) for c in color] + +def draw_body_and_foot(canvas, candidate, subset, stick_width=4, draw_body=True, draw_feet=True, body_keypoint_size=4, draw_head=True): + H, W, C = canvas.shape + candidate = np.array(candidate) + subset = np.array(subset) + limbSeq_and_colors = [] + + if draw_body: + limbSeq_and_colors = [ + ([2, 3], [255, 0, 0]), # Neck to Right Shoulder + ([2, 6], [255, 85, 0]), # Neck to Left Shoulder + ([3, 4], [255, 170, 0]), # Right Shoulder to Right Elbow + ([4, 5], [255, 255, 0]), # Right Elbow to Right Wrist + ([6, 7], [170, 255, 0]), # Left Shoulder to Left Elbow + ([7, 8], [85, 255, 0]), # Left Elbow to Left Wrist + ([2, 9], [0, 255, 0]), # Neck to Right Hip + ([9, 10], [0, 255, 85]), # Right Hip to Right Knee + ([10, 11], [0, 255, 170]), # Right Knee to Right Ankle + ([2, 12], [0, 255, 255]), # Neck to Left Hip + ([12, 13], [0, 170, 255]), # Left Hip to Left Knee + ([13, 14], [0, 85, 255]), # Left Knee to Left Ankle + ] + else: + limbSeq_and_colors = [ + ([2, 3], [0, 0, 0]), # Neck to Right Shoulder + ([2, 6], [0, 0, 0]), # Neck to Left Shoulder + ([3, 4], [0, 0, 0]), # Right Shoulder to Right Elbow + ([4, 5], [0, 0, 0]), # Right Elbow to Right Wrist + ([6, 7], [0, 0, 0]), # Left Shoulder to Left Elbow + ([7, 8], [0, 0, 0]), # Left Elbow to Left Wrist + ([2, 9], [0, 0, 0]), # Neck to Right Hip + ([9, 10], [0, 0, 0]), # Right Hip to Right Knee + ([10, 11], [0, 0, 0]), # Right Knee to Right Ankle + ([2, 12], [0, 0, 0]), # Neck to Left Hip + ([12, 13], [0, 0, 0]), # Left Hip to Left Knee + ([13, 14], [0, 0, 0]), # Left Knee to Left Ankle + ] + + # Conditionally add head-related elements + if draw_head: + head_elements = [ + ([2, 1], [0, 0, 255]), # Neck to Nose + ([1, 15], [0, 0, 255]), # Nose to Right Eye + ([15, 17], [85, 0, 255]), # Right Eye to Right Ear + ([1, 16], [170, 0, 255]), # Nose to Left Eye + ([16, 18], [255, 0, 255]), # Left Eye to Left Ear + ] + else: + head_elements = [ + ([2, 1], [0, 0, 0]), # Neck to Nose + ([1, 15], [0, 0, 0]), # Nose to Right Eye + ([15, 17], [0, 0, 0]), # Right Eye to Right Ear + ([1, 16], [0, 0, 0]), # Nose to Left Eye + ([16, 18], [0, 0, 0]), # Left Eye to Left Ear + ] + if draw_feet: + limbSeq_and_colors += [ + ([14, 19], [170, 255, 255]), # Left Ankle to Right Foot + ([11, 20], [255, 255, 0]), # Right Ankle to Left Foot + ] + + # Append head elements based on the condition + limbSeq_and_colors += head_elements + + for limb_info in limbSeq_and_colors: + limbSeq, color = limb_info + for n in range(len(subset)): + index = subset[n][np.array(limbSeq) - 1] + if -1 in index: + continue + Y = candidate[index.astype(int), 0] * float(W) + X = candidate[index.astype(int), 1] * float(H) + mX = np.mean(X) + mY = np.mean(Y) + 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, color) + + canvas = (canvas * 0.6).astype(np.uint8) + + for limb_info in limbSeq_and_colors: + limbSeq, color = limb_info + for i in limbSeq: + for n in range(len(subset)): + index = int(subset[n][i - 1]) + if index == -1: + continue + x, y = candidate[index][0:2] + if not np.isfinite(x) or not np.isfinite(y): + continue + 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, color, thickness=-1) + + return canvas + + +def draw_handpose(canvas, all_hand_peaks, draw_hands=True, hand_keypoint_size=4): + H, W, C = canvas.shape + + edges = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [0, 9], [9, 10], \ + [10, 11], [11, 12], [0, 13], [13, 14], [14, 15], [15, 16], [0, 17], [17, 18], [18, 19], [19, 20]] + + for peaks in all_hand_peaks: + peaks = np.array(peaks) + + if draw_hands: + for ie, e in enumerate(edges): + x1, y1 = peaks[e[0]] + x2, y2 = peaks[e[1]] + x1 = int(x1 * W) + y1 = int(y1 * H) + x2 = int(x2 * W) + y2 = int(y2 * H) + if x1 > eps and y1 > eps and x2 > eps and y2 > eps: + h = (ie / float(len(edges))) % 1.0 + s, v = 1.0, 1.0 + r, g, b = colorsys.hsv_to_rgb(h, s, v) + color = (int(255 * r), int(255 * g), int(255 * b)) + cv2.line(canvas, (x1, y1), (x2, y2), color, thickness=2) + if hand_keypoint_size > 0: + for i, keypoint in enumerate(peaks): + x, y = keypoint + x = int(x * W) + y = int(y * H) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), hand_keypoint_size, (0, 0, 255), thickness=-1) + return canvas + + +def draw_facepose(canvas, all_lmks): + H, W, C = canvas.shape + for lmks in all_lmks: + lmks = np.array(lmks) + for lmk in lmks: + x, y = lmk + x = int(x * W) + y = int(y * H) + if x > eps and y > eps: + cv2.circle(canvas, (x, y), 3, (255, 255, 255), thickness=-1) + return canvas + + +# detect hand according to body pose keypoints +# please refer to https://github.com/CMU-Perceptual-Computing-Lab/openpose/blob/master/src/openpose/hand/handDetector.cpp +def handDetect(candidate, subset, oriImg): + # right hand: wrist 4, elbow 3, shoulder 2 + # left hand: wrist 7, elbow 6, shoulder 5 + ratioWristElbow = 0.33 + detect_result = [] + image_height, image_width = oriImg.shape[0:2] + for person in subset.astype(int): + # if any of three not detected + has_left = np.sum(person[[5, 6, 7]] == -1) == 0 + has_right = np.sum(person[[2, 3, 4]] == -1) == 0 + if not (has_left or has_right): + continue + hands = [] + #left hand + if has_left: + left_shoulder_index, left_elbow_index, left_wrist_index = person[[5, 6, 7]] + x1, y1 = candidate[left_shoulder_index][:2] + x2, y2 = candidate[left_elbow_index][:2] + x3, y3 = candidate[left_wrist_index][:2] + hands.append([x1, y1, x2, y2, x3, y3, True]) + # right hand + if has_right: + right_shoulder_index, right_elbow_index, right_wrist_index = person[[2, 3, 4]] + x1, y1 = candidate[right_shoulder_index][:2] + x2, y2 = candidate[right_elbow_index][:2] + x3, y3 = candidate[right_wrist_index][:2] + hands.append([x1, y1, x2, y2, x3, y3, False]) + + for x1, y1, x2, y2, x3, y3, is_left in hands: + + x = x3 + ratioWristElbow * (x3 - x2) + y = y3 + ratioWristElbow * (y3 - y2) + distanceWristElbow = math.sqrt((x3 - x2) ** 2 + (y3 - y2) ** 2) + distanceElbowShoulder = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2) + width = 1.5 * max(distanceWristElbow, 0.9 * distanceElbowShoulder) + # x-y refers to the center --> offset to topLeft point + # handRectangle.x -= handRectangle.width / 2.f; + # handRectangle.y -= handRectangle.height / 2.f; + x -= width / 2 + y -= width / 2 # width = height + # overflow the image + if x < 0: x = 0 + if y < 0: y = 0 + width1 = width + width2 = width + if x + width > image_width: width1 = image_width - x + if y + width > image_height: width2 = image_height - y + width = min(width1, width2) + # the max hand box value is 20 pixels + if width >= 20: + detect_result.append([int(x), int(y), int(width), is_left]) + + ''' + return value: [[x, y, w, True if left hand else False]]. + width=height since the network require squared input. + x, y is the coordinate of top left + ''' + return detect_result + + +# Written by Lvmin +def faceDetect(candidate, subset, oriImg): + # left right eye ear 14 15 16 17 + detect_result = [] + image_height, image_width = oriImg.shape[0:2] + for person in subset.astype(int): + has_head = person[0] > -1 + if not has_head: + continue + + has_left_eye = person[14] > -1 + has_right_eye = person[15] > -1 + has_left_ear = person[16] > -1 + has_right_ear = person[17] > -1 + + if not (has_left_eye or has_right_eye or has_left_ear or has_right_ear): + continue + + head, left_eye, right_eye, left_ear, right_ear = person[[0, 14, 15, 16, 17]] + + width = 0.0 + x0, y0 = candidate[head][:2] + + if has_left_eye: + x1, y1 = candidate[left_eye][:2] + d = max(abs(x0 - x1), abs(y0 - y1)) + width = max(width, d * 3.0) + + if has_right_eye: + x1, y1 = candidate[right_eye][:2] + d = max(abs(x0 - x1), abs(y0 - y1)) + width = max(width, d * 3.0) + + if has_left_ear: + x1, y1 = candidate[left_ear][:2] + d = max(abs(x0 - x1), abs(y0 - y1)) + width = max(width, d * 1.5) + + if has_right_ear: + x1, y1 = candidate[right_ear][:2] + d = max(abs(x0 - x1), abs(y0 - y1)) + width = max(width, d * 1.5) + + x, y = x0, y0 + + x -= width + y -= width + + if x < 0: + x = 0 + + if y < 0: + y = 0 + + width1 = width * 2 + width2 = width * 2 + + if x + width > image_width: + width1 = image_width - x + + if y + width > image_height: + width2 = image_height - y + + width = min(width1, width2) + + if width >= 20: + detect_result.append([int(x), int(y), int(width)]) + + return detect_result + + +# get max index of 2d array +def npmax(array): + arrayindex = array.argmax(1) + arrayvalue = array.max(1) + i = arrayvalue.argmax() + j = arrayindex[i] + return i, j diff --git a/unianimate/dwpose/wholebody.py b/unianimate/dwpose/wholebody.py new file mode 100644 index 0000000..b838661 --- /dev/null +++ b/unianimate/dwpose/wholebody.py @@ -0,0 +1,42 @@ +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: + def __init__(self, model_det, model_pose): + self.model_det = model_det + self.model_pose = model_pose + + + def __call__(self, oriImg): + det_result = inference_jit_yolox(self.model_det, oriImg, detect_classes=[0]) + keypoints, scores = inference_jit_pose(self.model_pose, det_result, oriImg) + + keypoints_info = np.concatenate( + (keypoints, scores[..., None]), axis=-1) + # compute neck joint + neck = np.mean(keypoints_info[:, [5, 6]], axis=1) + # neck score when visualizing pred + neck[:, 2:4] = np.logical_and( + keypoints_info[:, 5, 2:4] > 0.3, + keypoints_info[:, 6, 2:4] > 0.3).astype(int) + new_keypoints_info = np.insert( + keypoints_info, 17, neck, axis=1) + mmpose_idx = [ + 17, 6, 8, 10, 7, 9, 12, 14, 16, 13, 15, 2, 1, 4, 3 + ] + openpose_idx = [ + 1, 2, 3, 4, 6, 7, 8, 9, 10, 12, 13, 14, 15, 16, 17 + ] + new_keypoints_info[:, openpose_idx] = \ + new_keypoints_info[:, mmpose_idx] + keypoints_info = new_keypoints_info + + keypoints, scores = keypoints_info[ + ..., :2], keypoints_info[..., 2] + + return keypoints, scores + + diff --git a/unianimate/nodes.py b/unianimate/nodes.py new file mode 100644 index 0000000..3c2e1fb --- /dev/null +++ b/unianimate/nodes.py @@ -0,0 +1,795 @@ + +import torch.nn as nn +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): + + concat_dim = 4 + transformer.dwpose_embedding = nn.Sequential( + nn.Conv3d(3, concat_dim * 4, (3,3,3), stride=(1,1,1), padding=(1,1,1)), + nn.SiLU(), + nn.Conv3d(concat_dim * 4, concat_dim * 4, (3,3,3), stride=(1,1,1), padding=(1,1,1)), + nn.SiLU(), + nn.Conv3d(concat_dim * 4, concat_dim * 4, (3,3,3), stride=(1,1,1), padding=(1,1,1)), + nn.SiLU(), + nn.Conv3d(concat_dim * 4, concat_dim * 4, (3,3,3), stride=(1,2,2), padding=(1,1,1)), + nn.SiLU(), + nn.Conv3d(concat_dim * 4, concat_dim * 4, 3, stride=(2,2,2), padding=1), + nn.SiLU(), + nn.Conv3d(concat_dim * 4, concat_dim * 4, 3, stride=(2,2,2), padding=1), + nn.SiLU(), + nn.Conv3d(concat_dim * 4, 5120, (1,2,2), stride=(1,2,2), padding=0)) + + randomref_dim = 20 + transformer.randomref_embedding_pose = nn.Sequential( + nn.Conv2d(3, concat_dim * 4, 3, stride=1, padding=1), + nn.SiLU(), + nn.Conv2d(concat_dim * 4, concat_dim * 4, 3, stride=1, padding=1), + nn.SiLU(), + nn.Conv2d(concat_dim * 4, concat_dim * 4, 3, stride=1, padding=1), + nn.SiLU(), + nn.Conv2d(concat_dim * 4, concat_dim * 4, 3, stride=2, padding=1), + nn.SiLU(), + nn.Conv2d(concat_dim * 4, concat_dim * 4, 3, stride=2, padding=1), + nn.SiLU(), + nn.Conv2d(concat_dim * 4, randomref_dim, 3, stride=2, padding=1), + ) + state_dict_new = {} + for key in list(state_dict.keys()): + if "dwpose_embedding" in key: + state_dict_new[key.split("dwpose_embedding.")[1]] = state_dict.pop(key) + transformer.dwpose_embedding.load_state_dict(state_dict_new, strict=True) + state_dict_new = {} + for key in list(state_dict.keys()): + if "randomref_embedding_pose" in key: + state_dict_new[key.split("randomref_embedding_pose.")[1]] = state_dict.pop(key) + transformer.randomref_embedding_pose.load_state_dict(state_dict_new,strict=True) + return transformer + +# Openpose +# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose +# 2nd Edited by https://github.com/Hzzone/pytorch-openpose +# 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): + r = 2 * math.pi * cutoff * t_e + return r / (r + 1) + + +def exponential_smoothing(a, x, x_prev): + return a * x + (1 - a) * x_prev + + +class OneEuroFilter: + def __init__(self, t0, x0, dx0=0.0, min_cutoff=1.0, beta=0.0, + d_cutoff=1.0): + """Initialize the one euro filter.""" + # The parameters. + self.min_cutoff = float(min_cutoff) + self.beta = float(beta) + self.d_cutoff = float(d_cutoff) + # Previous values. + self.x_prev = x0 + self.dx_prev = float(dx0) + self.t_prev = float(t0) + + def __call__(self, t, x): + """Compute the filtered signal.""" + t_e = t - self.t_prev + + # The filtered derivative of the signal. + a_d = smoothing_factor(t_e, self.d_cutoff) + dx = (x - self.x_prev) / t_e + dx_hat = exponential_smoothing(a_d, dx, self.dx_prev) + + # The filtered signal. + cutoff = self.min_cutoff + self.beta * abs(dx_hat) + a = smoothing_factor(t_e, cutoff) + x_hat = exponential_smoothing(a, x, self.x_prev) + + # Memorize the previous values. + self.x_prev = x_hat + self.dx_prev = dx_hat + self.t_prev = t + + return x_hat + +class DWposeDetector: + def __init__(self, model_det, model_pose): + self.pose_estimation = Wholebody(model_det, model_pose) + + def __call__(self, oriImg, score_threshold=0.3): + oriImg = oriImg.copy() + H, W, C = oriImg.shape + with torch.no_grad(): + candidate, subset = self.pose_estimation(oriImg) + candidate = candidate[0][np.newaxis, :, :] + subset = subset[0][np.newaxis, :] + nums, keys, locs = candidate.shape + candidate[..., 0] /= float(W) + candidate[..., 1] /= float(H) + body = candidate[:,:18].copy() + body = body.reshape(nums*18, locs) + score = subset[:,:18].copy() + + for i in range(len(score)): + for j in range(len(score[i])): + if score[i][j] > score_threshold: + score[i][j] = int(18*i+j) + else: + score[i][j] = -1 + + un_visible = subset score_threshold: + bodyfoot_score[i][j] = int(18*i+j) + else: + bodyfoot_score[i][j] = -1 + if -1 not in bodyfoot_score[:,18] and -1 not in bodyfoot_score[:,19]: + bodyfoot_score[:,18] = np.array([18.]) + else: + bodyfoot_score[:,18] = np.array([-1.]) + if -1 not in bodyfoot_score[:,21] and -1 not in bodyfoot_score[:,22]: + bodyfoot_score[:,19] = np.array([19.]) + else: + bodyfoot_score[:,19] = np.array([-1.]) + bodyfoot_score = bodyfoot_score[:, :20] + + bodyfoot = candidate[:,:24].copy() + + for i in range(nums): + if -1 not in bodyfoot[i][18] and -1 not in bodyfoot[i][19]: + bodyfoot[i][18] = (bodyfoot[i][18]+bodyfoot[i][19])/2 + else: + bodyfoot[i][18] = np.array([-1., -1.]) + if -1 not in bodyfoot[i][21] and -1 not in bodyfoot[i][22]: + bodyfoot[i][19] = (bodyfoot[i][21]+bodyfoot[i][22])/2 + else: + bodyfoot[i][19] = np.array([-1., -1.]) + + bodyfoot = bodyfoot[:,:20,:] + bodyfoot = bodyfoot.reshape(nums*20, locs) + + foot = candidate[:,18:24] + + faces = candidate[:,24:92] + + hands = candidate[:,92:113] + hands = np.vstack([hands, candidate[:,113:]]) + + # bodies = dict(candidate=body, subset=score) + bodies = dict(candidate=bodyfoot, subset=bodyfoot_score) + pose = dict(bodies=bodies, hands=hands, faces=faces) + + # return draw_pose(pose, H, W) + return pose + +def draw_pose(pose, H, W, stick_width=4,draw_body=True, draw_hands=True, draw_feet=True, + body_keypoint_size=4, hand_keypoint_size=4, draw_head=True): + from .dwpose.util import draw_body_and_foot, draw_handpose, draw_facepose + bodies = pose['bodies'] + faces = pose['faces'] + hands = pose['hands'] + candidate = bodies['candidate'] + subset = bodies['subset'] + + canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8) + canvas = draw_body_and_foot(canvas, candidate, subset, draw_body=draw_body, stick_width=stick_width, draw_feet=draw_feet, draw_head=draw_head, body_keypoint_size=body_keypoint_size) + canvas = draw_handpose(canvas, hands, draw_hands=draw_hands, hand_keypoint_size=hand_keypoint_size) + canvas_without_face = copy.deepcopy(canvas) + canvas = draw_facepose(canvas, faces) + + return canvas_without_face, canvas + + +def pose_extract(pose_images, ref_image, dwpose_model, height, width, score_threshold, stick_width, + draw_body=True, draw_hands=True, hand_keypoint_size=4, draw_feet=True, + body_keypoint_size=4, handle_not_detected="repeat", draw_head=True): + + results_vis = [] + comfy_pbar = ProgressBar(len(pose_images)) + + if ref_image is not None: + try: + pose_ref = dwpose_model(ref_image.squeeze(0), score_threshold=score_threshold) + except: + raise ValueError("No pose detected in reference image") + prev_pose = None + for img in tqdm(pose_images, desc="Pose Extraction", unit="image", total=len(pose_images)): + try: + pose = dwpose_model(img, score_threshold=score_threshold) + if handle_not_detected == "repeat": + prev_pose = pose + except: + if prev_pose is not None: + pose = prev_pose + else: + pose = np.zeros_like(img) + results_vis.append(pose) + comfy_pbar.update(1) + + bodies = results_vis[0]['bodies'] + faces = results_vis[0]['faces'] + hands = results_vis[0]['hands'] + candidate = bodies['candidate'] + + if ref_image is not None: + ref_bodies = pose_ref['bodies'] + ref_faces = pose_ref['faces'] + ref_hands = pose_ref['hands'] + ref_candidate = ref_bodies['candidate'] + + + ref_2_x = ref_candidate[2][0] + ref_2_y = ref_candidate[2][1] + ref_5_x = ref_candidate[5][0] + ref_5_y = ref_candidate[5][1] + ref_8_x = ref_candidate[8][0] + ref_8_y = ref_candidate[8][1] + ref_11_x = ref_candidate[11][0] + ref_11_y = ref_candidate[11][1] + ref_center1 = 0.5*(ref_candidate[2]+ref_candidate[5]) + ref_center2 = 0.5*(ref_candidate[8]+ref_candidate[11]) + + zero_2_x = candidate[2][0] + zero_2_y = candidate[2][1] + zero_5_x = candidate[5][0] + zero_5_y = candidate[5][1] + zero_8_x = candidate[8][0] + zero_8_y = candidate[8][1] + zero_11_x = candidate[11][0] + zero_11_y = candidate[11][1] + zero_center1 = 0.5*(candidate[2]+candidate[5]) + zero_center2 = 0.5*(candidate[8]+candidate[11]) + + x_ratio = (ref_5_x-ref_2_x)/(zero_5_x-zero_2_x) + y_ratio = (ref_center2[1]-ref_center1[1])/(zero_center2[1]-zero_center1[1]) + + results_vis[0]['bodies']['candidate'][:,0] *= x_ratio + results_vis[0]['bodies']['candidate'][:,1] *= y_ratio + results_vis[0]['faces'][:,:,0] *= x_ratio + results_vis[0]['faces'][:,:,1] *= y_ratio + results_vis[0]['hands'][:,:,0] *= x_ratio + results_vis[0]['hands'][:,:,1] *= y_ratio + + ########neck######## + l_neck_ref = ((ref_candidate[0][0] - ref_candidate[1][0]) ** 2 + (ref_candidate[0][1] - ref_candidate[1][1]) ** 2) ** 0.5 + l_neck_0 = ((candidate[0][0] - candidate[1][0]) ** 2 + (candidate[0][1] - candidate[1][1]) ** 2) ** 0.5 + neck_ratio = l_neck_ref / l_neck_0 + + x_offset_neck = (candidate[1][0]-candidate[0][0])*(1.-neck_ratio) + y_offset_neck = (candidate[1][1]-candidate[0][1])*(1.-neck_ratio) + + results_vis[0]['bodies']['candidate'][0,0] += x_offset_neck + results_vis[0]['bodies']['candidate'][0,1] += y_offset_neck + results_vis[0]['bodies']['candidate'][14,0] += x_offset_neck + results_vis[0]['bodies']['candidate'][14,1] += y_offset_neck + results_vis[0]['bodies']['candidate'][15,0] += x_offset_neck + results_vis[0]['bodies']['candidate'][15,1] += y_offset_neck + results_vis[0]['bodies']['candidate'][16,0] += x_offset_neck + results_vis[0]['bodies']['candidate'][16,1] += y_offset_neck + results_vis[0]['bodies']['candidate'][17,0] += x_offset_neck + results_vis[0]['bodies']['candidate'][17,1] += y_offset_neck + + ########shoulder2######## + l_shoulder2_ref = ((ref_candidate[2][0] - ref_candidate[1][0]) ** 2 + (ref_candidate[2][1] - ref_candidate[1][1]) ** 2) ** 0.5 + l_shoulder2_0 = ((candidate[2][0] - candidate[1][0]) ** 2 + (candidate[2][1] - candidate[1][1]) ** 2) ** 0.5 + + shoulder2_ratio = l_shoulder2_ref / l_shoulder2_0 + + x_offset_shoulder2 = (candidate[1][0]-candidate[2][0])*(1.-shoulder2_ratio) + y_offset_shoulder2 = (candidate[1][1]-candidate[2][1])*(1.-shoulder2_ratio) + + results_vis[0]['bodies']['candidate'][2,0] += x_offset_shoulder2 + results_vis[0]['bodies']['candidate'][2,1] += y_offset_shoulder2 + results_vis[0]['bodies']['candidate'][3,0] += x_offset_shoulder2 + results_vis[0]['bodies']['candidate'][3,1] += y_offset_shoulder2 + results_vis[0]['bodies']['candidate'][4,0] += x_offset_shoulder2 + results_vis[0]['bodies']['candidate'][4,1] += y_offset_shoulder2 + results_vis[0]['hands'][1,:,0] += x_offset_shoulder2 + results_vis[0]['hands'][1,:,1] += y_offset_shoulder2 + + ########shoulder5######## + l_shoulder5_ref = ((ref_candidate[5][0] - ref_candidate[1][0]) ** 2 + (ref_candidate[5][1] - ref_candidate[1][1]) ** 2) ** 0.5 + l_shoulder5_0 = ((candidate[5][0] - candidate[1][0]) ** 2 + (candidate[5][1] - candidate[1][1]) ** 2) ** 0.5 + + shoulder5_ratio = l_shoulder5_ref / l_shoulder5_0 + + x_offset_shoulder5 = (candidate[1][0]-candidate[5][0])*(1.-shoulder5_ratio) + y_offset_shoulder5 = (candidate[1][1]-candidate[5][1])*(1.-shoulder5_ratio) + + results_vis[0]['bodies']['candidate'][5,0] += x_offset_shoulder5 + results_vis[0]['bodies']['candidate'][5,1] += y_offset_shoulder5 + results_vis[0]['bodies']['candidate'][6,0] += x_offset_shoulder5 + results_vis[0]['bodies']['candidate'][6,1] += y_offset_shoulder5 + results_vis[0]['bodies']['candidate'][7,0] += x_offset_shoulder5 + results_vis[0]['bodies']['candidate'][7,1] += y_offset_shoulder5 + results_vis[0]['hands'][0,:,0] += x_offset_shoulder5 + results_vis[0]['hands'][0,:,1] += y_offset_shoulder5 + + ########arm3######## + l_arm3_ref = ((ref_candidate[3][0] - ref_candidate[2][0]) ** 2 + (ref_candidate[3][1] - ref_candidate[2][1]) ** 2) ** 0.5 + l_arm3_0 = ((candidate[3][0] - candidate[2][0]) ** 2 + (candidate[3][1] - candidate[2][1]) ** 2) ** 0.5 + + arm3_ratio = l_arm3_ref / l_arm3_0 + + x_offset_arm3 = (candidate[2][0]-candidate[3][0])*(1.-arm3_ratio) + y_offset_arm3 = (candidate[2][1]-candidate[3][1])*(1.-arm3_ratio) + + results_vis[0]['bodies']['candidate'][3,0] += x_offset_arm3 + results_vis[0]['bodies']['candidate'][3,1] += y_offset_arm3 + results_vis[0]['bodies']['candidate'][4,0] += x_offset_arm3 + results_vis[0]['bodies']['candidate'][4,1] += y_offset_arm3 + results_vis[0]['hands'][1,:,0] += x_offset_arm3 + results_vis[0]['hands'][1,:,1] += y_offset_arm3 + + ########arm4######## + l_arm4_ref = ((ref_candidate[4][0] - ref_candidate[3][0]) ** 2 + (ref_candidate[4][1] - ref_candidate[3][1]) ** 2) ** 0.5 + l_arm4_0 = ((candidate[4][0] - candidate[3][0]) ** 2 + (candidate[4][1] - candidate[3][1]) ** 2) ** 0.5 + + arm4_ratio = l_arm4_ref / l_arm4_0 + + x_offset_arm4 = (candidate[3][0]-candidate[4][0])*(1.-arm4_ratio) + y_offset_arm4 = (candidate[3][1]-candidate[4][1])*(1.-arm4_ratio) + + results_vis[0]['bodies']['candidate'][4,0] += x_offset_arm4 + results_vis[0]['bodies']['candidate'][4,1] += y_offset_arm4 + results_vis[0]['hands'][1,:,0] += x_offset_arm4 + results_vis[0]['hands'][1,:,1] += y_offset_arm4 + + ########arm6######## + l_arm6_ref = ((ref_candidate[6][0] - ref_candidate[5][0]) ** 2 + (ref_candidate[6][1] - ref_candidate[5][1]) ** 2) ** 0.5 + l_arm6_0 = ((candidate[6][0] - candidate[5][0]) ** 2 + (candidate[6][1] - candidate[5][1]) ** 2) ** 0.5 + + arm6_ratio = l_arm6_ref / l_arm6_0 + + x_offset_arm6 = (candidate[5][0]-candidate[6][0])*(1.-arm6_ratio) + y_offset_arm6 = (candidate[5][1]-candidate[6][1])*(1.-arm6_ratio) + + results_vis[0]['bodies']['candidate'][6,0] += x_offset_arm6 + results_vis[0]['bodies']['candidate'][6,1] += y_offset_arm6 + results_vis[0]['bodies']['candidate'][7,0] += x_offset_arm6 + results_vis[0]['bodies']['candidate'][7,1] += y_offset_arm6 + results_vis[0]['hands'][0,:,0] += x_offset_arm6 + results_vis[0]['hands'][0,:,1] += y_offset_arm6 + + ########arm7######## + l_arm7_ref = ((ref_candidate[7][0] - ref_candidate[6][0]) ** 2 + (ref_candidate[7][1] - ref_candidate[6][1]) ** 2) ** 0.5 + l_arm7_0 = ((candidate[7][0] - candidate[6][0]) ** 2 + (candidate[7][1] - candidate[6][1]) ** 2) ** 0.5 + + arm7_ratio = l_arm7_ref / l_arm7_0 + + x_offset_arm7 = (candidate[6][0]-candidate[7][0])*(1.-arm7_ratio) + y_offset_arm7 = (candidate[6][1]-candidate[7][1])*(1.-arm7_ratio) + + results_vis[0]['bodies']['candidate'][7,0] += x_offset_arm7 + results_vis[0]['bodies']['candidate'][7,1] += y_offset_arm7 + results_vis[0]['hands'][0,:,0] += x_offset_arm7 + results_vis[0]['hands'][0,:,1] += y_offset_arm7 + + ########head14######## + l_head14_ref = ((ref_candidate[14][0] - ref_candidate[0][0]) ** 2 + (ref_candidate[14][1] - ref_candidate[0][1]) ** 2) ** 0.5 + l_head14_0 = ((candidate[14][0] - candidate[0][0]) ** 2 + (candidate[14][1] - candidate[0][1]) ** 2) ** 0.5 + + head14_ratio = l_head14_ref / l_head14_0 + + x_offset_head14 = (candidate[0][0]-candidate[14][0])*(1.-head14_ratio) + y_offset_head14 = (candidate[0][1]-candidate[14][1])*(1.-head14_ratio) + + results_vis[0]['bodies']['candidate'][14,0] += x_offset_head14 + results_vis[0]['bodies']['candidate'][14,1] += y_offset_head14 + results_vis[0]['bodies']['candidate'][16,0] += x_offset_head14 + results_vis[0]['bodies']['candidate'][16,1] += y_offset_head14 + + ########head15######## + l_head15_ref = ((ref_candidate[15][0] - ref_candidate[0][0]) ** 2 + (ref_candidate[15][1] - ref_candidate[0][1]) ** 2) ** 0.5 + l_head15_0 = ((candidate[15][0] - candidate[0][0]) ** 2 + (candidate[15][1] - candidate[0][1]) ** 2) ** 0.5 + + head15_ratio = l_head15_ref / l_head15_0 + + x_offset_head15 = (candidate[0][0]-candidate[15][0])*(1.-head15_ratio) + y_offset_head15 = (candidate[0][1]-candidate[15][1])*(1.-head15_ratio) + + results_vis[0]['bodies']['candidate'][15,0] += x_offset_head15 + results_vis[0]['bodies']['candidate'][15,1] += y_offset_head15 + results_vis[0]['bodies']['candidate'][17,0] += x_offset_head15 + results_vis[0]['bodies']['candidate'][17,1] += y_offset_head15 + + ########head16######## + l_head16_ref = ((ref_candidate[16][0] - ref_candidate[14][0]) ** 2 + (ref_candidate[16][1] - ref_candidate[14][1]) ** 2) ** 0.5 + l_head16_0 = ((candidate[16][0] - candidate[14][0]) ** 2 + (candidate[16][1] - candidate[14][1]) ** 2) ** 0.5 + + head16_ratio = l_head16_ref / l_head16_0 + + x_offset_head16 = (candidate[14][0]-candidate[16][0])*(1.-head16_ratio) + y_offset_head16 = (candidate[14][1]-candidate[16][1])*(1.-head16_ratio) + + results_vis[0]['bodies']['candidate'][16,0] += x_offset_head16 + results_vis[0]['bodies']['candidate'][16,1] += y_offset_head16 + + ########head17######## + l_head17_ref = ((ref_candidate[17][0] - ref_candidate[15][0]) ** 2 + (ref_candidate[17][1] - ref_candidate[15][1]) ** 2) ** 0.5 + l_head17_0 = ((candidate[17][0] - candidate[15][0]) ** 2 + (candidate[17][1] - candidate[15][1]) ** 2) ** 0.5 + + head17_ratio = l_head17_ref / l_head17_0 + + x_offset_head17 = (candidate[15][0]-candidate[17][0])*(1.-head17_ratio) + y_offset_head17 = (candidate[15][1]-candidate[17][1])*(1.-head17_ratio) + + results_vis[0]['bodies']['candidate'][17,0] += x_offset_head17 + results_vis[0]['bodies']['candidate'][17,1] += y_offset_head17 + + ########MovingAverage######## + + ########left leg######## + l_ll1_ref = ((ref_candidate[8][0] - ref_candidate[9][0]) ** 2 + (ref_candidate[8][1] - ref_candidate[9][1]) ** 2) ** 0.5 + l_ll1_0 = ((candidate[8][0] - candidate[9][0]) ** 2 + (candidate[8][1] - candidate[9][1]) ** 2) ** 0.5 + ll1_ratio = l_ll1_ref / l_ll1_0 + + x_offset_ll1 = (candidate[9][0]-candidate[8][0])*(ll1_ratio-1.) + y_offset_ll1 = (candidate[9][1]-candidate[8][1])*(ll1_ratio-1.) + + results_vis[0]['bodies']['candidate'][9,0] += x_offset_ll1 + results_vis[0]['bodies']['candidate'][9,1] += y_offset_ll1 + results_vis[0]['bodies']['candidate'][10,0] += x_offset_ll1 + results_vis[0]['bodies']['candidate'][10,1] += y_offset_ll1 + results_vis[0]['bodies']['candidate'][19,0] += x_offset_ll1 + results_vis[0]['bodies']['candidate'][19,1] += y_offset_ll1 + + l_ll2_ref = ((ref_candidate[9][0] - ref_candidate[10][0]) ** 2 + (ref_candidate[9][1] - ref_candidate[10][1]) ** 2) ** 0.5 + l_ll2_0 = ((candidate[9][0] - candidate[10][0]) ** 2 + (candidate[9][1] - candidate[10][1]) ** 2) ** 0.5 + ll2_ratio = l_ll2_ref / l_ll2_0 + + x_offset_ll2 = (candidate[10][0]-candidate[9][0])*(ll2_ratio-1.) + y_offset_ll2 = (candidate[10][1]-candidate[9][1])*(ll2_ratio-1.) + + results_vis[0]['bodies']['candidate'][10,0] += x_offset_ll2 + results_vis[0]['bodies']['candidate'][10,1] += y_offset_ll2 + results_vis[0]['bodies']['candidate'][19,0] += x_offset_ll2 + results_vis[0]['bodies']['candidate'][19,1] += y_offset_ll2 + + ########right leg######## + l_rl1_ref = ((ref_candidate[11][0] - ref_candidate[12][0]) ** 2 + (ref_candidate[11][1] - ref_candidate[12][1]) ** 2) ** 0.5 + l_rl1_0 = ((candidate[11][0] - candidate[12][0]) ** 2 + (candidate[11][1] - candidate[12][1]) ** 2) ** 0.5 + rl1_ratio = l_rl1_ref / l_rl1_0 + + x_offset_rl1 = (candidate[12][0]-candidate[11][0])*(rl1_ratio-1.) + y_offset_rl1 = (candidate[12][1]-candidate[11][1])*(rl1_ratio-1.) + + results_vis[0]['bodies']['candidate'][12,0] += x_offset_rl1 + results_vis[0]['bodies']['candidate'][12,1] += y_offset_rl1 + results_vis[0]['bodies']['candidate'][13,0] += x_offset_rl1 + results_vis[0]['bodies']['candidate'][13,1] += y_offset_rl1 + results_vis[0]['bodies']['candidate'][18,0] += x_offset_rl1 + results_vis[0]['bodies']['candidate'][18,1] += y_offset_rl1 + + l_rl2_ref = ((ref_candidate[12][0] - ref_candidate[13][0]) ** 2 + (ref_candidate[12][1] - ref_candidate[13][1]) ** 2) ** 0.5 + l_rl2_0 = ((candidate[12][0] - candidate[13][0]) ** 2 + (candidate[12][1] - candidate[13][1]) ** 2) ** 0.5 + rl2_ratio = l_rl2_ref / l_rl2_0 + + x_offset_rl2 = (candidate[13][0]-candidate[12][0])*(rl2_ratio-1.) + y_offset_rl2 = (candidate[13][1]-candidate[12][1])*(rl2_ratio-1.) + + results_vis[0]['bodies']['candidate'][13,0] += x_offset_rl2 + results_vis[0]['bodies']['candidate'][13,1] += y_offset_rl2 + results_vis[0]['bodies']['candidate'][18,0] += x_offset_rl2 + results_vis[0]['bodies']['candidate'][18,1] += y_offset_rl2 + + offset = ref_candidate[1] - results_vis[0]['bodies']['candidate'][1] + + results_vis[0]['bodies']['candidate'] += offset[np.newaxis, :] + results_vis[0]['faces'] += offset[np.newaxis, np.newaxis, :] + results_vis[0]['hands'] += offset[np.newaxis, np.newaxis, :] + + for i in range(1, len(results_vis)): + results_vis[i]['bodies']['candidate'][:,0] *= x_ratio + results_vis[i]['bodies']['candidate'][:,1] *= y_ratio + results_vis[i]['faces'][:,:,0] *= x_ratio + results_vis[i]['faces'][:,:,1] *= y_ratio + results_vis[i]['hands'][:,:,0] *= x_ratio + results_vis[i]['hands'][:,:,1] *= y_ratio + + ########neck######## + x_offset_neck = (results_vis[i]['bodies']['candidate'][1][0]-results_vis[i]['bodies']['candidate'][0][0])*(1.-neck_ratio) + y_offset_neck = (results_vis[i]['bodies']['candidate'][1][1]-results_vis[i]['bodies']['candidate'][0][1])*(1.-neck_ratio) + + results_vis[i]['bodies']['candidate'][0,0] += x_offset_neck + results_vis[i]['bodies']['candidate'][0,1] += y_offset_neck + results_vis[i]['bodies']['candidate'][14,0] += x_offset_neck + results_vis[i]['bodies']['candidate'][14,1] += y_offset_neck + results_vis[i]['bodies']['candidate'][15,0] += x_offset_neck + results_vis[i]['bodies']['candidate'][15,1] += y_offset_neck + results_vis[i]['bodies']['candidate'][16,0] += x_offset_neck + results_vis[i]['bodies']['candidate'][16,1] += y_offset_neck + results_vis[i]['bodies']['candidate'][17,0] += x_offset_neck + results_vis[i]['bodies']['candidate'][17,1] += y_offset_neck + + ########shoulder2######## + + + x_offset_shoulder2 = (results_vis[i]['bodies']['candidate'][1][0]-results_vis[i]['bodies']['candidate'][2][0])*(1.-shoulder2_ratio) + y_offset_shoulder2 = (results_vis[i]['bodies']['candidate'][1][1]-results_vis[i]['bodies']['candidate'][2][1])*(1.-shoulder2_ratio) + + results_vis[i]['bodies']['candidate'][2,0] += x_offset_shoulder2 + results_vis[i]['bodies']['candidate'][2,1] += y_offset_shoulder2 + results_vis[i]['bodies']['candidate'][3,0] += x_offset_shoulder2 + results_vis[i]['bodies']['candidate'][3,1] += y_offset_shoulder2 + results_vis[i]['bodies']['candidate'][4,0] += x_offset_shoulder2 + results_vis[i]['bodies']['candidate'][4,1] += y_offset_shoulder2 + results_vis[i]['hands'][1,:,0] += x_offset_shoulder2 + results_vis[i]['hands'][1,:,1] += y_offset_shoulder2 + + ########shoulder5######## + + x_offset_shoulder5 = (results_vis[i]['bodies']['candidate'][1][0]-results_vis[i]['bodies']['candidate'][5][0])*(1.-shoulder5_ratio) + y_offset_shoulder5 = (results_vis[i]['bodies']['candidate'][1][1]-results_vis[i]['bodies']['candidate'][5][1])*(1.-shoulder5_ratio) + + results_vis[i]['bodies']['candidate'][5,0] += x_offset_shoulder5 + results_vis[i]['bodies']['candidate'][5,1] += y_offset_shoulder5 + results_vis[i]['bodies']['candidate'][6,0] += x_offset_shoulder5 + results_vis[i]['bodies']['candidate'][6,1] += y_offset_shoulder5 + results_vis[i]['bodies']['candidate'][7,0] += x_offset_shoulder5 + results_vis[i]['bodies']['candidate'][7,1] += y_offset_shoulder5 + results_vis[i]['hands'][0,:,0] += x_offset_shoulder5 + results_vis[i]['hands'][0,:,1] += y_offset_shoulder5 + + ########arm3######## + + x_offset_arm3 = (results_vis[i]['bodies']['candidate'][2][0]-results_vis[i]['bodies']['candidate'][3][0])*(1.-arm3_ratio) + y_offset_arm3 = (results_vis[i]['bodies']['candidate'][2][1]-results_vis[i]['bodies']['candidate'][3][1])*(1.-arm3_ratio) + + results_vis[i]['bodies']['candidate'][3,0] += x_offset_arm3 + results_vis[i]['bodies']['candidate'][3,1] += y_offset_arm3 + results_vis[i]['bodies']['candidate'][4,0] += x_offset_arm3 + results_vis[i]['bodies']['candidate'][4,1] += y_offset_arm3 + results_vis[i]['hands'][1,:,0] += x_offset_arm3 + results_vis[i]['hands'][1,:,1] += y_offset_arm3 + + ########arm4######## + + x_offset_arm4 = (results_vis[i]['bodies']['candidate'][3][0]-results_vis[i]['bodies']['candidate'][4][0])*(1.-arm4_ratio) + y_offset_arm4 = (results_vis[i]['bodies']['candidate'][3][1]-results_vis[i]['bodies']['candidate'][4][1])*(1.-arm4_ratio) + + results_vis[i]['bodies']['candidate'][4,0] += x_offset_arm4 + results_vis[i]['bodies']['candidate'][4,1] += y_offset_arm4 + results_vis[i]['hands'][1,:,0] += x_offset_arm4 + results_vis[i]['hands'][1,:,1] += y_offset_arm4 + + ########arm6######## + + x_offset_arm6 = (results_vis[i]['bodies']['candidate'][5][0]-results_vis[i]['bodies']['candidate'][6][0])*(1.-arm6_ratio) + y_offset_arm6 = (results_vis[i]['bodies']['candidate'][5][1]-results_vis[i]['bodies']['candidate'][6][1])*(1.-arm6_ratio) + + results_vis[i]['bodies']['candidate'][6,0] += x_offset_arm6 + results_vis[i]['bodies']['candidate'][6,1] += y_offset_arm6 + results_vis[i]['bodies']['candidate'][7,0] += x_offset_arm6 + results_vis[i]['bodies']['candidate'][7,1] += y_offset_arm6 + results_vis[i]['hands'][0,:,0] += x_offset_arm6 + results_vis[i]['hands'][0,:,1] += y_offset_arm6 + + ########arm7######## + + x_offset_arm7 = (results_vis[i]['bodies']['candidate'][6][0]-results_vis[i]['bodies']['candidate'][7][0])*(1.-arm7_ratio) + y_offset_arm7 = (results_vis[i]['bodies']['candidate'][6][1]-results_vis[i]['bodies']['candidate'][7][1])*(1.-arm7_ratio) + + results_vis[i]['bodies']['candidate'][7,0] += x_offset_arm7 + results_vis[i]['bodies']['candidate'][7,1] += y_offset_arm7 + results_vis[i]['hands'][0,:,0] += x_offset_arm7 + results_vis[i]['hands'][0,:,1] += y_offset_arm7 + + ########head14######## + + x_offset_head14 = (results_vis[i]['bodies']['candidate'][0][0]-results_vis[i]['bodies']['candidate'][14][0])*(1.-head14_ratio) + y_offset_head14 = (results_vis[i]['bodies']['candidate'][0][1]-results_vis[i]['bodies']['candidate'][14][1])*(1.-head14_ratio) + + results_vis[i]['bodies']['candidate'][14,0] += x_offset_head14 + results_vis[i]['bodies']['candidate'][14,1] += y_offset_head14 + results_vis[i]['bodies']['candidate'][16,0] += x_offset_head14 + results_vis[i]['bodies']['candidate'][16,1] += y_offset_head14 + + ########head15######## + + x_offset_head15 = (results_vis[i]['bodies']['candidate'][0][0]-results_vis[i]['bodies']['candidate'][15][0])*(1.-head15_ratio) + y_offset_head15 = (results_vis[i]['bodies']['candidate'][0][1]-results_vis[i]['bodies']['candidate'][15][1])*(1.-head15_ratio) + + results_vis[i]['bodies']['candidate'][15,0] += x_offset_head15 + results_vis[i]['bodies']['candidate'][15,1] += y_offset_head15 + results_vis[i]['bodies']['candidate'][17,0] += x_offset_head15 + results_vis[i]['bodies']['candidate'][17,1] += y_offset_head15 + + ########head16######## + + x_offset_head16 = (results_vis[i]['bodies']['candidate'][14][0]-results_vis[i]['bodies']['candidate'][16][0])*(1.-head16_ratio) + y_offset_head16 = (results_vis[i]['bodies']['candidate'][14][1]-results_vis[i]['bodies']['candidate'][16][1])*(1.-head16_ratio) + + results_vis[i]['bodies']['candidate'][16,0] += x_offset_head16 + results_vis[i]['bodies']['candidate'][16,1] += y_offset_head16 + + ########head17######## + x_offset_head17 = (results_vis[i]['bodies']['candidate'][15][0]-results_vis[i]['bodies']['candidate'][17][0])*(1.-head17_ratio) + y_offset_head17 = (results_vis[i]['bodies']['candidate'][15][1]-results_vis[i]['bodies']['candidate'][17][1])*(1.-head17_ratio) + + results_vis[i]['bodies']['candidate'][17,0] += x_offset_head17 + results_vis[i]['bodies']['candidate'][17,1] += y_offset_head17 + + # ########MovingAverage######## + + ########left leg######## + x_offset_ll1 = (results_vis[i]['bodies']['candidate'][9][0]-results_vis[i]['bodies']['candidate'][8][0])*(ll1_ratio-1.) + y_offset_ll1 = (results_vis[i]['bodies']['candidate'][9][1]-results_vis[i]['bodies']['candidate'][8][1])*(ll1_ratio-1.) + + results_vis[i]['bodies']['candidate'][9,0] += x_offset_ll1 + results_vis[i]['bodies']['candidate'][9,1] += y_offset_ll1 + results_vis[i]['bodies']['candidate'][10,0] += x_offset_ll1 + results_vis[i]['bodies']['candidate'][10,1] += y_offset_ll1 + results_vis[i]['bodies']['candidate'][19,0] += x_offset_ll1 + results_vis[i]['bodies']['candidate'][19,1] += y_offset_ll1 + + + + x_offset_ll2 = (results_vis[i]['bodies']['candidate'][10][0]-results_vis[i]['bodies']['candidate'][9][0])*(ll2_ratio-1.) + y_offset_ll2 = (results_vis[i]['bodies']['candidate'][10][1]-results_vis[i]['bodies']['candidate'][9][1])*(ll2_ratio-1.) + + results_vis[i]['bodies']['candidate'][10,0] += x_offset_ll2 + results_vis[i]['bodies']['candidate'][10,1] += y_offset_ll2 + results_vis[i]['bodies']['candidate'][19,0] += x_offset_ll2 + results_vis[i]['bodies']['candidate'][19,1] += y_offset_ll2 + + ########right leg######## + + x_offset_rl1 = (results_vis[i]['bodies']['candidate'][12][0]-results_vis[i]['bodies']['candidate'][11][0])*(rl1_ratio-1.) + y_offset_rl1 = (results_vis[i]['bodies']['candidate'][12][1]-results_vis[i]['bodies']['candidate'][11][1])*(rl1_ratio-1.) + + results_vis[i]['bodies']['candidate'][12,0] += x_offset_rl1 + results_vis[i]['bodies']['candidate'][12,1] += y_offset_rl1 + results_vis[i]['bodies']['candidate'][13,0] += x_offset_rl1 + results_vis[i]['bodies']['candidate'][13,1] += y_offset_rl1 + results_vis[i]['bodies']['candidate'][18,0] += x_offset_rl1 + results_vis[i]['bodies']['candidate'][18,1] += y_offset_rl1 + + + x_offset_rl2 = (results_vis[i]['bodies']['candidate'][13][0]-results_vis[i]['bodies']['candidate'][12][0])*(rl2_ratio-1.) + y_offset_rl2 = (results_vis[i]['bodies']['candidate'][13][1]-results_vis[i]['bodies']['candidate'][12][1])*(rl2_ratio-1.) + + results_vis[i]['bodies']['candidate'][13,0] += x_offset_rl2 + results_vis[i]['bodies']['candidate'][13,1] += y_offset_rl2 + results_vis[i]['bodies']['candidate'][18,0] += x_offset_rl2 + results_vis[i]['bodies']['candidate'][18,1] += y_offset_rl2 + + results_vis[i]['bodies']['candidate'] += offset[np.newaxis, :] + results_vis[i]['faces'] += offset[np.newaxis, np.newaxis, :] + results_vis[i]['hands'] += offset[np.newaxis, np.newaxis, :] + + dwpose_woface_list = [] + for i in range(len(results_vis)): + #try: + dwpose_woface, dwpose_wface = draw_pose(results_vis[i], H=height, W=width, stick_width=stick_width, + draw_body=draw_body, draw_hands=draw_hands, hand_keypoint_size=hand_keypoint_size, + draw_feet=draw_feet, body_keypoint_size=body_keypoint_size, draw_head=draw_head) + result = torch.from_numpy(dwpose_woface) + #except: + # result = torch.zeros((height, width, 3), dtype=torch.uint8) + dwpose_woface_list.append(result) + dwpose_woface_tensor = torch.stack(dwpose_woface_list, dim=0) + + dwpose_woface_ref_tensor = None + if ref_image is not None: + dwpose_woface_ref, dwpose_wface_ref = draw_pose(pose_ref, H=height, W=width, stick_width=stick_width, + draw_body=draw_body, draw_hands=draw_hands, hand_keypoint_size=hand_keypoint_size, + draw_feet=draw_feet, body_keypoint_size=body_keypoint_size, draw_head=draw_head) + dwpose_woface_ref_tensor = torch.from_numpy(dwpose_woface_ref) + + return dwpose_woface_tensor, dwpose_woface_ref_tensor + +class ProportionChangerUniAnimateDWPoseDetector: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "pose_images": ("IMAGE", {"tooltip": "Pose images"}), + "score_threshold": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Score threshold for pose detection"}), + "stick_width": ("INT", {"default": 4, "min": 1, "max": 100, "step": 1, "tooltip": "Stick width for drawing keypoints"}), + "draw_body": ("BOOLEAN", {"default": True, "tooltip": "Draw body keypoints"}), + "body_keypoint_size": ("INT", {"default": 4, "min": 0, "max": 100, "step": 1, "tooltip": "Body keypoint size"}), + "draw_feet": ("BOOLEAN", {"default": True, "tooltip": "Draw feet keypoints"}), + "draw_hands": ("BOOLEAN", {"default": True, "tooltip": "Draw hand keypoints"}), + "hand_keypoint_size": ("INT", {"default": 4, "min": 0, "max": 100, "step": 1, "tooltip": "Hand keypoint size"}), + "colorspace": (["RGB", "BGR"], {"tooltip": "Color space for the output image"}), + "handle_not_detected": (["empty", "repeat"], {"default": "empty", "tooltip": "How to handle undetected poses, empty inserts black and repeat inserts previous detection"}), + "draw_head": ("BOOLEAN", {"default": True, "tooltip": "Draw head keypoints"}), + }, + "optional": { + "reference_pose_image": ("IMAGE", {"tooltip": "Reference pose image"}), + }, + } + + RETURN_TYPES = ("IMAGE", "IMAGE", ) + RETURN_NAMES = ("poses", "reference_pose",) + FUNCTION = "process" + CATEGORY = "ProportionChanger" + + def process(self, pose_images, score_threshold, stick_width, reference_pose_image=None, draw_body=True, body_keypoint_size=4, + draw_feet=True, draw_hands=True, hand_keypoint_size=4, colorspace="RGB", handle_not_detected="empty", draw_head=True): + + device = mm.get_torch_device() + + #model loading + dw_pose_model = "dw-ll_ucoco_384_bs5.torchscript.pt" + yolo_model = "yolox_l.torchscript.pt" + + script_directory = os.path.dirname(os.path.abspath(__file__)) + model_base_path = os.path.join(script_directory, "models", "DWPose") + + model_det=os.path.join(model_base_path, yolo_model) + model_pose=os.path.join(model_base_path, dw_pose_model) + + if not os.path.exists(model_det): + log.info(f"Downloading yolo model to: {model_base_path}") + from huggingface_hub import snapshot_download + snapshot_download(repo_id="hr16/yolox-onnx", + allow_patterns=[f"*{yolo_model}*"], + local_dir=model_base_path, + local_dir_use_symlinks=False) + + if not os.path.exists(model_pose): + log.info(f"Downloading dwpose model to: {model_base_path}") + from huggingface_hub import snapshot_download + snapshot_download(repo_id="hr16/DWPose-TorchScript-BatchSize5", + allow_patterns=[f"*{dw_pose_model}*"], + local_dir=model_base_path, + local_dir_use_symlinks=False) + + if not hasattr(self, "det") or not hasattr(self, "pose"): + self.det = torch.jit.load(model_det, map_location=device) + self.pose = torch.jit.load(model_pose, map_location=device) + self.dwpose_detector = DWposeDetector(self.det, self.pose) + + #model inference + height, width = pose_images.shape[1:3] + + pose_np = pose_images.cpu().numpy() * 255 + ref_np = None + if reference_pose_image is not None: + ref = reference_pose_image + ref_np = ref.cpu().numpy() * 255 + + 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 + if reference_pose_image is not None: + reference_pose = reference_pose.unsqueeze(0) / 255.0 + else: + reference_pose = torch.zeros(1, 64, 64, 3, device=torch.device("cpu")) + + if colorspace == "BGR": + poses=torch.flip(poses, dims=[-1]) + + return (poses, reference_pose, ) + + +NODE_CLASS_MAPPINGS = { + "ProportionChangerUniAnimateDWPoseDetector": ProportionChangerUniAnimateDWPoseDetector, + } +NODE_DISPLAY_NAME_MAPPINGS = { + "ProportionChangerUniAnimateDWPoseDetector": "ProportionChanger UniAnimate DWPose Detector", + } + + \ No newline at end of file diff --git a/utils.py b/utils.py new file mode 100644 index 0000000..6fbc5c0 --- /dev/null +++ b/utils.py @@ -0,0 +1,311 @@ +import importlib.metadata +import torch +import logging +from tqdm import tqdm +import types, collections +from comfy.utils import ProgressBar, copy_to_param, set_attr_param +from comfy.model_patcher import get_key_weight, string_to_seed +from comfy.lora import calculate_weight +from comfy.model_management import cast_to_device +from comfy.float import stochastic_rounding +logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') +log = logging.getLogger(__name__) + +from accelerate.utils import set_module_tensor_to_device +def check_diffusers_version(): + try: + version = importlib.metadata.version('diffusers') + required_version = '0.31.0' + if version < required_version: + raise AssertionError(f"diffusers version {version} is installed, but version {required_version} or higher is required.") + except importlib.metadata.PackageNotFoundError: + raise AssertionError("diffusers is not installed.") + +def print_memory(device): + memory = torch.cuda.memory_allocated(device) / 1024**3 + max_memory = torch.cuda.max_memory_allocated(device) / 1024**3 + max_reserved = torch.cuda.max_memory_reserved(device) / 1024**3 + log.info(f"Allocated memory: {memory=:.3f} GB") + log.info(f"Max allocated memory: {max_memory=:.3f} GB") + log.info(f"Max reserved memory: {max_reserved=:.3f} GB") + #memory_summary = torch.cuda.memory_summary(device=device, abbreviated=False) + #log.info(f"Memory Summary:\n{memory_summary}") + +def get_module_memory_mb(module): + memory = 0 + for param in module.parameters(): + if param.data is not None: + memory += param.nelement() * param.element_size() + return memory / (1024 * 1024) # Convert to MB + +def get_tensor_memory(tensor): + memory_bytes = tensor.element_size() * tensor.nelement() + return f"{memory_bytes / (1024 * 1024):.2f} MB" + +def patch_weight_to_device(self, key, device_to=None, inplace_update=False, backup_keys=False, scale_weight=None): + if key not in self.patches: + return + + weight, set_func, convert_func = get_key_weight(self.model, key) + inplace_update = self.weight_inplace_update or inplace_update + + if backup_keys and key not in self.backup: + self.backup[key] = collections.namedtuple('Dimension', ['weight', 'inplace_update'])(weight.to(device=self.offload_device, copy=inplace_update), inplace_update) + + if device_to is not None: + temp_weight = cast_to_device(weight, device_to, torch.float32, copy=True) + else: + temp_weight = weight.to(torch.float32, copy=True) + if convert_func is not None: + temp_weight = convert_func(temp_weight, inplace=True) + + if scale_weight is not None: + temp_weight = temp_weight * scale_weight.to(temp_weight.device, temp_weight.dtype) + + out_weight = calculate_weight(self.patches[key], temp_weight, key) + + if set_func is None: + out_weight = stochastic_rounding(out_weight, weight.dtype, seed=string_to_seed(key)) + if inplace_update: + copy_to_param(self.model, key, out_weight) + else: + set_attr_param(self.model, key, out_weight) + else: + set_func(out_weight, inplace_update=inplace_update, seed=string_to_seed(key)) + +def apply_lora(model, device_to, transformer_load_device, params_to_keep=None, dtype=None, base_dtype=None, state_dict=None, low_mem_load=False, control_lora=False, scale_weights={}): + model.patch_weight_to_device = types.MethodType(patch_weight_to_device, model) + to_load = [] + for n, m in model.model.named_modules(): + params = [] + skip = False + for name, param in m.named_parameters(recurse=False): + params.append(name) + for name, param in m.named_parameters(recurse=True): + if name not in params: + skip = True # skip random weights in non leaf modules + break + if not skip and (hasattr(m, "comfy_cast_weights") or len(params) > 0): + to_load.append((n, m, params)) + + to_load.sort(reverse=True) + pbar = ProgressBar(len(to_load)) + for x in tqdm(to_load, desc="Loading model and applying LoRA weights:", leave=True): + name = x[0] + m = x[1] + params = x[2] + if hasattr(m, "comfy_patched_weights"): + if m.comfy_patched_weights == True: + continue + for param in params: + name = name.replace("._orig_mod.", ".") # torch compiled modules have this prefix + if low_mem_load: + dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype + if "patch_embedding" in name: + dtype_to_use = torch.float32 + key = f"{name.replace('diffusion_model.', '')}.{param}" + try: + set_module_tensor_to_device(model.model.diffusion_model, key, device=transformer_load_device, dtype=dtype_to_use, value=state_dict[key]) + except: + continue + key = f"{name}.{param}" + if scale_weights is not None: + scale_key = key.replace("weight", "scale_weight").replace("diffusion_model.", "") if "weight" in key else None + if low_mem_load: + model.patch_weight_to_device(f"{name}.{param}", device_to=device_to, inplace_update=True, backup_keys=control_lora, scale_weight=scale_weights.get(scale_key, None)) + else: + model.patch_weight_to_device(f"{name}.{param}", device_to=device_to, backup_keys=control_lora, scale_weight=scale_weights.get(scale_key, None)) + if device_to != transformer_load_device: + set_module_tensor_to_device(m, param, device=transformer_load_device) + if low_mem_load: + try: + set_module_tensor_to_device(model.model.diffusion_model, key, device=transformer_load_device, dtype=dtype_to_use, value=model.model.diffusion_model.state_dict()[key]) + except: + continue + m.comfy_patched_weights = True + pbar.update(1) + + model.current_weight_patches_uuid = model.patches_uuid + if low_mem_load: + for name, param in model.model.diffusion_model.named_parameters(): + if param.device != transformer_load_device: + dtype_to_use = base_dtype if any(keyword in name for keyword in params_to_keep) else dtype + if "patch_embedding" in name: + dtype_to_use = torch.float32 + try: + set_module_tensor_to_device(model.model.diffusion_model, name, device=transformer_load_device, dtype=dtype_to_use, value=state_dict[name]) + except: + continue + return model + + +# from https://github.com/cubiq/ComfyUI_IPAdapter_plus/blob/9d076a3df0d2763cef5510ec5ab807f6632c39f5/utils.py#L181 +def split_tiles(embeds, num_split): + _, H, W, _ = embeds.shape + out = [] + for x in embeds: + x = x.unsqueeze(0) + h, w = H // num_split, W // num_split + x_split = torch.cat([x[:, i*h:(i+1)*h, j*w:(j+1)*w, :] for i in range(num_split) for j in range(num_split)], dim=0) + out.append(x_split) + + x_split = torch.stack(out, dim=0) + + return x_split + +def merge_hiddenstates(x, tiles): + chunk_size = tiles*tiles + x = x.split(chunk_size) + + out = [] + for embeds in x: + num_tiles = embeds.shape[0] + tile_size = int((embeds.shape[1]-1) ** 0.5) + grid_size = int(num_tiles ** 0.5) + + # Extract class tokens + class_tokens = embeds[:, 0, :] # Save class tokens: [num_tiles, embeds[-1]] + avg_class_token = class_tokens.mean(dim=0, keepdim=True).unsqueeze(0) # Average token, shape: [1, 1, embeds[-1]] + + patch_embeds = embeds[:, 1:, :] # Shape: [num_tiles, tile_size^2, embeds[-1]] + reshaped = patch_embeds.reshape(grid_size, grid_size, tile_size, tile_size, embeds.shape[-1]) + + merged = torch.cat([torch.cat([reshaped[i, j] for j in range(grid_size)], dim=1) + for i in range(grid_size)], dim=0) + + merged = merged.unsqueeze(0) # Shape: [1, grid_size*tile_size, grid_size*tile_size, embeds[-1]] + + # Pool to original size + pooled = torch.nn.functional.adaptive_avg_pool2d(merged.permute(0, 3, 1, 2), (tile_size, tile_size)).permute(0, 2, 3, 1) + flattened = pooled.reshape(1, tile_size*tile_size, embeds.shape[-1]) + + # Add back the class token + with_class = torch.cat([avg_class_token, flattened], dim=1) # Shape: original shape + out.append(with_class) + + out = torch.cat(out, dim=0) + + return out + +from comfy.clip_vision import clip_preprocess, ClipVisionModel + +def clip_encode_image_tiled(clip_vision, image, tiles=1, ratio=1.0): + embeds = encode_image_(clip_vision, image) + tiles = min(tiles, 16) + + if tiles > 1: + # split in tiles + image_split = split_tiles(image, tiles) + + # get the embeds for each tile + embeds_split = {} + for i in image_split: + encoded = encode_image_(clip_vision, i) + if not hasattr(embeds_split, "last_hidden_state"): + embeds_split["last_hidden_state"] = encoded + else: + embeds_split["last_hidden_state"] = torch.cat(embeds_split["last_hidden_state"], encoded, dim=0) + + embeds_split['last_hidden_state'] = merge_hiddenstates(embeds_split['last_hidden_state'], tiles) + + if embeds.shape[0] > 1: # if we have more than one image we need to average the embeddings for consistency + embeds = embeds * ratio + embeds_split['last_hidden_state']*(1-ratio) + else: # otherwise we can concatenate them, they can be averaged later + embeds = torch.cat([embeds * ratio, embeds_split['last_hidden_state']]) + + return embeds + +def encode_image_(clip_vision, image): + if isinstance(clip_vision, ClipVisionModel): + out = clip_vision.encode_image(image).last_hidden_state + else: + pixel_values = clip_preprocess(image, size=224, crop=True).float() + out = clip_vision.visual(pixel_values) + + return out + +# Code based on https://github.com/WikiChao/FreSca (MIT License) +import torch +import torch.fft as fft + +def fourier_filter(x, scale_low=1.0, scale_high=1.5, freq_cutoff=20): + """ + Apply frequency-dependent scaling to an image tensor using Fourier transforms. + + Parameters: + x: Input tensor of shape (B, C, H, W) + scale_low: Scaling factor for low-frequency components (default: 1.0) + scale_high: Scaling factor for high-frequency components (default: 1.5) + freq_cutoff: Number of frequency indices around center to consider as low-frequency (default: 20) + + Returns: + x_filtered: Filtered version of x in spatial domain with frequency-specific scaling applied. + """ + # Preserve input dtype and device + dtype, device = x.dtype, x.device + + # Convert to float32 for FFT computations + x = x.to(torch.float32) + + # 1) Apply FFT and shift low frequencies to center + x_freq = fft.fftn(x, dim=(-2, -1)) + x_freq = fft.fftshift(x_freq, dim=(-2, -1)) + + # 2) Create a mask to scale frequencies differently + C, B, H, W = x_freq.shape + crow, ccol = H // 2, W // 2 + + # Initialize mask with high-frequency scaling factor + mask = torch.ones((C, B, H, W), device=device) * scale_high + + # Apply low-frequency scaling factor to center region + mask[ + ..., + crow - freq_cutoff : crow + freq_cutoff, + ccol - freq_cutoff : ccol + freq_cutoff, + ] = scale_low + + # 3) Apply frequency-specific scaling + x_freq = x_freq * mask + + # 4) Convert back to spatial domain + x_freq = fft.ifftshift(x_freq, dim=(-2, -1)) + x_filtered = fft.ifftn(x_freq, dim=(-2, -1)).real + + # 5) Restore original dtype + x_filtered = x_filtered.to(dtype) + + return x_filtered + +def is_image_black(image, threshold=1e-3): + if image.min() < 0: + image = (image + 1) / 2 + return torch.all(image < threshold).item() + +def add_noise_to_reference_video(image, ratio=None): + sigma = torch.ones((image.shape[0],)).to(image.device, image.dtype) * ratio + image_noise = torch.randn_like(image) * sigma[:, None, None, None] + image_noise = torch.where(image==-1, torch.zeros_like(image), image_noise) + image = image + image_noise + return image + +def optimized_scale(positive_flat, negative_flat): + + # Calculate dot production + dot_product = torch.sum(positive_flat * negative_flat, dim=1, keepdim=True) + + # Squared norm of uncondition + squared_norm = torch.sum(negative_flat ** 2, dim=1, keepdim=True) + 1e-8 + + # st_star = v_cond^T * v_uncond / ||v_uncond||^2 + st_star = dot_product / squared_norm + + return st_star + +def find_closest_valid_dim(fixed_dim, var_dim, block_size): + for delta in range(1, 17): + for sign in [-1, 1]: + candidate = var_dim + sign * delta + if candidate > 0 and ((fixed_dim * candidate) // 4) % block_size == 0: + return candidate + return var_dim \ No newline at end of file