Spaces:
Configuration error
Configuration error
File size: 3,344 Bytes
9b35681 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 |
import glob
import os
from PIL import Image
import numpy as np
from matplotlib import pyplot as plt
from io import BytesIO
def visualize_hand(all_joints, img, side=["right", "left"], n_avail_joints=21):
# Define the connections between joints for drawing lines and their corresponding colors
connections = [
((0, 1), "red"),
((1, 2), "green"),
((2, 3), "blue"),
((3, 4), "purple"),
((0, 5), "orange"),
((5, 6), "pink"),
((6, 7), "brown"),
((7, 8), "cyan"),
((0, 9), "yellow"),
((9, 10), "magenta"),
((10, 11), "lime"),
((11, 12), "indigo"),
((0, 13), "olive"),
((13, 14), "teal"),
((14, 15), "navy"),
((15, 16), "gray"),
((0, 17), "lavender"),
((17, 18), "silver"),
((18, 19), "maroon"),
((19, 20), "fuchsia"),
]
H, W, C = img.shape
# Create a figure and axis
plt.figure()
ax = plt.gca()
# Plot joints as points
ax.imshow(img)
start_is = []
if "right" in side:
start_is.append(0)
if "left" in side:
start_is.append(21)
for start_i in start_is:
joints = all_joints[start_i : start_i + n_avail_joints]
if len(joints) == 1:
ax.scatter(joints[0][0], joints[0][1], color="red", s=10)
else:
for connection, color in connections[: len(joints) - 1]:
joint1 = joints[connection[0]]
joint2 = joints[connection[1]]
ax.plot([joint1[0], joint2[0]], [joint1[1], joint2[1]], color=color, linewidth=4)
ax.set_xlim([0, W])
ax.set_ylim([0, H])
ax.grid(False)
ax.set_axis_off()
ax.invert_yaxis()
# plt.subplots_adjust(wspace=0.01)
# plt.show()
buf = BytesIO()
plt.savefig(buf, format="png", bbox_inches="tight", pad_inches=0)
plt.close()
# Convert BytesIO object to numpy array
buf.seek(0)
img_pil = Image.open(buf)
img_pil = img_pil.resize((W, H))
numpy_img = np.array(img_pil)
return numpy_img
'''put brush example at alpha channel'''
# img_dir = "bad_hands"
# masked_paths = sorted(glob.glob(os.path.join(img_dir, "*_mask.jpg")))
# for masked_pth in masked_paths:
# img_path = masked_pth.replace("_mask.jpg", ".jpg")
# assert os.path.exists(img_path), f"Image path {img_path} does not exist."
# masked = np.array(Image.open(masked_pth))
# mask = (np.all(masked > 245, axis=-1)).astype(np.uint8)*128 + 64
# img = np.array(Image.open(img_path))
# composite = np.concatenate((img, mask[..., None]), axis=-1)
# composite = Image.fromarray(composite)
# composite.save(masked_pth.replace("_mask.jpg", "_composite.png"))
# print(f"Saved composite image {masked_pth.replace('_mask.jpg', '_composite.png')}")
'''visualize keypoint example'''
data_dir = "bad_hands"
kpts_paths = sorted(glob.glob(os.path.join(data_dir, "*.npy")))
for kpts_pth in kpts_paths:
img_pth = kpts_pth.replace(".npy", ".jpg")
kpts = np.load(kpts_pth)
img = np.array(Image.open(img_pth))
h, w = img.shape[:2]
kpts = kpts / np.array([256,256]) * np.array([w, h])
kpts_vis = visualize_hand(kpts, img)
save_path = kpts_pth.replace(".npy", "_kpts.png")
kpts_vis = Image.fromarray(kpts_vis).save(save_path)
print(f"Saved {save_path}") |