Views
No views yet
1def draw_bodypose(canvas: np.ndarray, keypoints: List[Keypoint]) -> np.ndarray:
2 """
3 Draw keypoints and limbs representing body pose on a given canvas.
4
5 Args:
6 canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the body pose.
7 keypoints (List[Keypoint]): A list of Keypoint objects representing the body keypoints to be drawn.
8
9 Returns:
10 np.ndarray: A 3D numpy array representing the modified canvas with the drawn body pose.
11
12 Note:
13 The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
14 """
15 H, W, C = canvas.shape
16
17
18 if max(W, H) < 500:
19 ratio = 1.0
20 elif max(W, H) >= 500 and max(W, H) < 1000:
21 ratio = 2.0
22 elif max(W, H) >= 1000 and max(W, H) < 2000:
23 ratio = 3.0
24 elif max(W, H) >= 2000 and max(W, H) < 3000:
25 ratio = 4.0
26 elif max(W, H) >= 3000 and max(W, H) < 4000:
27 ratio = 5.0
28 elif max(W, H) >= 4000 and max(W, H) < 5000:
29 ratio = 6.0
30 else:
31 ratio = 7.0
32
33 stickwidth = 4
34
35 limbSeq = [
36 [2, 3], [2, 6], [3, 4], [4, 5],
37 [6, 7], [7, 8], [2, 9], [9, 10],
38 [10, 11], [2, 12], [12, 13], [13, 14],
39 [2, 1], [1, 15], [15, 17], [1, 16],
40 [16, 18],
41 ]
42
43 colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
44 [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
45 [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85]]
46
47 for (k1_index, k2_index), color in zip(limbSeq, colors):
48 keypoint1 = keypoints[k1_index - 1]
49 keypoint2 = keypoints[k2_index - 1]
50
51 if keypoint1 is None or keypoint2 is None:
52 continue
53
54 Y = np.array([keypoint1.x, keypoint2.x]) * float(W)
55 X = np.array([keypoint1.y, keypoint2.y]) * float(H)
56 mX = np.mean(X)
57 mY = np.mean(Y)
58 length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
59 angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
60 polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), int(stickwidth * ratio)), int(angle), 0, 360, 1)
61 cv2.fillConvexPoly(canvas, polygon, [int(float(c) * 0.6) for c in color])
62
63 for keypoint, color in zip(keypoints, colors):
64 if keypoint is None:
65 continue
66
67 x, y = keypoint.x, keypoint.y
68 x = int(x * W)
69 y = int(y * H)
70 cv2.circle(canvas, (int(x), int(y)), int(4 * ratio), color, thickness=-1)
71
72 return canvas1from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
2from diffusers import DDIMScheduler, EulerAncestralDiscreteScheduler
3from controlnet_aux import OpenposeDetector
4from PIL import Image
5import torch
6import numpy as np
7import cv2
8
9
10
11controlnet_conditioning_scale = 1.0
12prompt = "your prompt, the longer the better, you can describe it as detail as possible"
13negative_prompt = 'longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality'
14
15
16
17eulera_scheduler = EulerAncestralDiscreteScheduler.from_pretrained("stabilityai/stable-diffusion-xl-base-1.0", subfolder="scheduler")
18
19
20controlnet = ControlNetModel.from_pretrained(
21 "xinsir/controlnet-openpose-sdxl-1.0",
22 torch_dtype=torch.float16
23)
24
25# when test with other base model, you need to change the vae also.
26vae = AutoencoderKL.from_pretrained("madebyollin/sdxl-vae-fp16-fix", torch_dtype=torch.float16)
27
28
29pipe = StableDiffusionXLControlNetPipeline.from_pretrained(
30 "stabilityai/stable-diffusion-xl-base-1.0",
31 controlnet=controlnet,
32 vae=vae,
33 safety_checker=None,
34 torch_dtype=torch.float16,
35 scheduler=eulera_scheduler,
36)
37
38processor = OpenposeDetector.from_pretrained('lllyasviel/ControlNet')
39
40
41controlnet_img = cv2.imread("your image path")
42controlnet_img = processor(controlnet_img, hand_and_face=False, output_type='cv2')
43
44
45# need to resize the image resolution to 1024 * 1024 or same bucket resolution to get the best performance
46height, width, _ = controlnet_img.shape
47ratio = np.sqrt(1024. * 1024. / (width * height))
48new_width, new_height = int(width * ratio), int(height * ratio)
49controlnet_img = cv2.resize(controlnet_img, (new_width, new_height))
50controlnet_img = Image.fromarray(controlnet_img)
51
52images = pipe(
53 prompt,
54 negative_prompt=negative_prompt,
55 image=controlnet_img,
56 controlnet_conditioning_scale=controlnet_conditioning_scale,
57 width=new_width,
58 height=new_height,
59 num_inference_steps=30,
60 ).images
61
62images[0].save(f"your image save path, png format is usually better than jpg or webp in terms of image quality but got much bigger")| metric | xinsir/controlnet-openpose-sdxl-1.0 | lllyasviel/control_v11p_sd15_openpose | thibaud/controlnet-openpose-sdxl-1.0 |
|---|---|---|---|
| mAP | 0.357 | 0.326 | 0.209 |