Update extract_video.py

Added fix when choosing to render the images from openpose (creating the folder)
This commit is contained in:
Carlos Barreto 2021-04-22 18:51:12 -03:00 committed by GitHub
parent 8092486693
commit e6b29a583d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -53,12 +53,17 @@ def extract_2d(openpose, image, keypoints, render, args):
if args.handface:
cmd = cmd + ' --hand --face'
if args.render:
cmd = cmd + ' --write_images {}'.format(render)
if os.path.exists(join(os.getcwd(),render)):
cmd = cmd + ' --write_images {}'.format(join(os.getcwd(),render))
else:
os.makedirs(join(os.getcwd(),render), exist_ok=True)
cmd = cmd + ' --write_images {}'.format(join(os.getcwd(),render))
else:
cmd = cmd + ' --render_pose 0'
os.chdir(openpose)
os.system(cmd)
import json
def read_json(path):
with open(path) as f:
@ -216,7 +221,7 @@ def extract_yolo_hrnet(image_root, annot_root, ext='jpg', use_low=False):
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('path', type=str, default=None, help="the path of data")
parser.add_argument('path', type=str, help="the path of data")
parser.add_argument('--mode', type=str, default='openpose', choices=['openpose', 'yolo-hrnet'], help="model to extract joints from image")
parser.add_argument('--ext', type=str, default='jpg', choices=['jpg', 'png'], help="image file extension")
parser.add_argument('--annot', type=str, default='annots', help="sub directory name to store the generated annotation files, default to be annots")