pose-normalization-v3 / json_to_openpose.py
Harsha909's picture
Upload folder using huggingface_hub
ac5f9d6 verified
Raw
History Blame Contribute Delete
1.55 kB
import os, json, cv2, numpy as np, glob
INPUT_DIR = r"C:\Users\harsh\DWPose\output_normalized"
OUTPUT_DIR = r"C:\Users\harsh\DWPose\pose_frames"
W, H = 512, 512
CONF = 0.3
# body connections
EDGES = [
(0,1),(1,2),(2,3),(3,4),
(1,5),(5,6),(6,7),
(1,8),(8,9),(9,10),
(1,11),(11,12),(12,13)
]
def draw_pose(coords, scores):
canvas = np.zeros((H, W, 3), dtype=np.uint8)
for i, (x,y) in enumerate(coords):
if scores[i] > CONF:
cv2.circle(canvas, (int(x), int(y)), 4, (0,255,0), -1)
for i,j in EDGES:
if scores[i] > CONF and scores[j] > CONF:
cv2.line(canvas,
(int(coords[i][0]), int(coords[i][1])),
(int(coords[j][0]), int(coords[j][1])),
(255,255,255), 2)
return canvas
def process(json_file):
name = os.path.basename(json_file).replace("_norm_kps.json", "")
out_folder = os.path.join(OUTPUT_DIR, name)
os.makedirs(out_folder, exist_ok=True)
with open(json_file) as f:
data = json.load(f)
for i, frame in enumerate(data["frames"]):
coords = np.array(frame["body"]["coords"])
scores = np.array(frame["body"]["scores"])
img = draw_pose(coords, scores)
cv2.imwrite(os.path.join(out_folder, f"{i:05d}.png"), img)
print("Done:", name)
def main():
files = glob.glob(os.path.join(INPUT_DIR, "*_norm_kps.json"))
for f in files:
process(f)
if __name__ == "__main__":
main()