acmyu/KeyframesAI

1from src.controlnet_aux import DWposeDetector
2from PIL import Image
3import torchvision.transforms as transforms
4import torch
5
6def init_dwpose_detector(device):
7    # specify configs, ckpts and device, or it will be downloaded automatically and use cpu by default
8    det_config = './src/configs/yolox_l_8xb8-300e_coco.py'
9    det_ckpt = './ckpts/yolox_l_8x8_300e_coco_20211126_140236-d3bd2b23.pth'
10    pose_config = './src/configs/dwpose-l_384x288.py'
11    pose_ckpt = './ckpts/dw-ll_ucoco_384.pth'
12
13    dwpose_model = DWposeDetector(
14        det_config=det_config,
15        det_ckpt=det_ckpt,
16        pose_config=pose_config,
17        pose_ckpt=pose_ckpt,
18        device=device
19    )
20    return dwpose_model.to(device)
21
22
23def inference_pose(img_path, image_size=(1024, 1024)):
24    device = torch.device(f"cuda:{0}")
25    model = init_dwpose_detector(device=device)
26    pil_image = Image.open(img_path).convert("RGB").resize(image_size, Image.BICUBIC)
27    dwpose_image = model(pil_image, output_type='np', image_resolution=image_size[1])
28    save_dwpose_image = Image.fromarray(dwpose_image)
29    return save_dwpose_image
30
31
32
33inference_pose('imgs/test.png').save("pose.png")
34
35
36