| import numpy as np | |
| import os | |
| import torch.nn as nn | |
| os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" | |
| import cv2 | |
| import torch | |
| from . import util | |
| from .body import Body | |
| remote_model_path = "https://huggingface.co./TencentARC/T2I-Adapter/blob/main/third-party-models/body_pose_model.pth" | |
| class OpenposeInference(nn.Module): | |
| def __init__(self): | |
| super().__init__() | |
| body_modelpath = os.path.join('models', "body_pose_model.pth") | |
| if not os.path.exists(body_modelpath): | |
| from basicsr.utils.download_util import load_file_from_url | |
| load_file_from_url(remote_model_path, model_dir='models') | |
| self.body_estimation = Body(body_modelpath) | |
| def forward(self, x): | |
| x = x[:, :, ::-1].copy() | |
| with torch.no_grad(): | |
| candidate, subset = self.body_estimation(x) | |
| canvas = np.zeros_like(x) | |
| canvas = util.draw_bodypose(canvas, candidate, subset) | |
| canvas = cv2.cvtColor(canvas, cv2.COLOR_RGB2BGR) | |
| return canvas | |