232 lines
8.3 KiB
Python
232 lines
8.3 KiB
Python
'''
|
||
@ Date: 2021-01-13 20:38:33
|
||
@ Author: Qing Shuai
|
||
@ LastEditors: Qing Shuai
|
||
@ LastEditTime: 2021-01-25 14:41:56
|
||
@ FilePath: /EasyMocap/scripts/preprocess/extract_video.py
|
||
'''
|
||
import os, sys
|
||
import cv2
|
||
from os.path import join
|
||
from tqdm import tqdm
|
||
from glob import glob
|
||
import numpy as np
|
||
code_path = join(os.path.dirname(__file__), '..', '..', 'code')
|
||
sys.path.append(code_path)
|
||
|
||
mkdir = lambda x: os.makedirs(x, exist_ok=True)
|
||
|
||
def extract_video(videoname, path, start, end, step):
|
||
base = os.path.basename(videoname).replace('.mp4', '')
|
||
if not os.path.exists(videoname):
|
||
return base
|
||
outpath = join(path, 'images', base)
|
||
if os.path.exists(outpath) and len(os.listdir(outpath)) > 0:
|
||
return base
|
||
else:
|
||
os.makedirs(outpath)
|
||
video = cv2.VideoCapture(videoname)
|
||
totalFrames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
|
||
for cnt in tqdm(range(totalFrames)):
|
||
ret, frame = video.read()
|
||
if cnt < start:continue
|
||
if cnt > end:break
|
||
if not ret:break
|
||
cv2.imwrite(join(outpath, '{:06d}.jpg'.format(cnt)), frame)
|
||
video.release()
|
||
return base
|
||
|
||
def extract_2d(openpose, image, keypoints, render):
|
||
if not os.path.exists(keypoints):
|
||
os.makedirs(keypoints, exist_ok=True)
|
||
cmd = './build/examples/openpose/openpose.bin --image_dir {} --write_json {} --display 0'.format(image, keypoints)
|
||
if args.handface:
|
||
cmd = cmd + ' --hand --face'
|
||
if args.render:
|
||
cmd = cmd + ' --write_images {}'.format(render)
|
||
else:
|
||
cmd = cmd + ' --render_pose 0'
|
||
os.chdir(openpose)
|
||
os.system(cmd)
|
||
|
||
import json
|
||
def read_json(path):
|
||
with open(path) as f:
|
||
data = json.load(f)
|
||
return data
|
||
|
||
def save_json(file, data):
|
||
if not os.path.exists(os.path.dirname(file)):
|
||
os.makedirs(os.path.dirname(file))
|
||
with open(file, 'w') as f:
|
||
json.dump(data, f, indent=4)
|
||
|
||
def create_annot_file(annotname, imgname):
|
||
assert os.path.exists(imgname), imgname
|
||
img = cv2.imread(imgname)
|
||
height, width = img.shape[0], img.shape[1]
|
||
imgnamesep = imgname.split(os.sep)
|
||
filename = os.sep.join(imgnamesep[imgnamesep.index('images'):])
|
||
annot = {
|
||
'filename':filename,
|
||
'height':height,
|
||
'width':width,
|
||
'annots': [],
|
||
'isKeyframe': False
|
||
}
|
||
save_json(annotname, annot)
|
||
return annot
|
||
|
||
def bbox_from_openpose(keypoints, rescale=1.2, detection_thresh=0.01):
|
||
"""Get center and scale for bounding box from openpose detections."""
|
||
valid = keypoints[:,-1] > detection_thresh
|
||
valid_keypoints = keypoints[valid][:,:-1]
|
||
center = valid_keypoints.mean(axis=0)
|
||
bbox_size = valid_keypoints.max(axis=0) - valid_keypoints.min(axis=0)
|
||
# adjust bounding box tightness
|
||
bbox_size = bbox_size * rescale
|
||
bbox = [
|
||
center[0] - bbox_size[0]/2,
|
||
center[1] - bbox_size[1]/2,
|
||
center[0] + bbox_size[0]/2,
|
||
center[1] + bbox_size[1]/2,
|
||
keypoints[valid, 2].mean()
|
||
]
|
||
return bbox
|
||
|
||
def load_openpose(opname):
|
||
mapname = {'face_keypoints_2d':'face2d', 'hand_left_keypoints_2d':'handl2d', 'hand_right_keypoints_2d':'handr2d'}
|
||
assert os.path.exists(opname), opname
|
||
data = read_json(opname)
|
||
out = []
|
||
pid = 0
|
||
for i, d in enumerate(data['people']):
|
||
keypoints = d['pose_keypoints_2d']
|
||
keypoints = np.array(keypoints).reshape(-1, 3)
|
||
annot = {
|
||
'bbox': bbox_from_openpose(keypoints),
|
||
'personID': pid + i,
|
||
'keypoints': keypoints.tolist(),
|
||
'isKeyframe': False
|
||
}
|
||
for key in ['face_keypoints_2d', 'hand_left_keypoints_2d', 'hand_right_keypoints_2d']:
|
||
if len(d[key]) == 0:
|
||
continue
|
||
kpts = np.array(d[key]).reshape(-1, 3)
|
||
annot[mapname[key]] = kpts.tolist()
|
||
out.append(annot)
|
||
return out
|
||
|
||
def convert_from_openpose(src, dst):
|
||
# convert the 2d pose from openpose
|
||
inputlist = sorted(os.listdir(src))
|
||
for inp in tqdm(inputlist):
|
||
annots = load_openpose(join(src, inp))
|
||
base = inp.replace('_keypoints.json', '')
|
||
annotname = join(dst, base+'.json')
|
||
imgname = annotname.replace('annots', 'images').replace('.json', '.jpg')
|
||
if not os.path.exists(imgname):
|
||
os.remove(join(src, inp))
|
||
continue
|
||
annot = create_annot_file(annotname, imgname)
|
||
annot['annots'] = annots
|
||
save_json(annotname, annot)
|
||
|
||
def detect_frame(detector, img, pid=0):
|
||
lDetections = detector.detect([img])[0]
|
||
annots = []
|
||
for i in range(len(lDetections)):
|
||
annot = {
|
||
'bbox': [float(d) for d in lDetections[i]['bbox']],
|
||
'personID': pid + i,
|
||
'keypoints': lDetections[i]['keypoints'].tolist(),
|
||
'isKeyframe': True
|
||
}
|
||
annots.append(annot)
|
||
return annots
|
||
|
||
def extract_yolo_hrnet(image_root, annot_root):
|
||
imgnames = sorted(glob(join(image_root, '*.jpg')))
|
||
import torch
|
||
device = torch.device('cuda')
|
||
from estimator.detector import Detector
|
||
config = {
|
||
'yolov4': {
|
||
'ckpt_path': 'data/models/yolov4.weights',
|
||
'conf_thres': 0.3,
|
||
'box_nms_thres': 0.5 # 阈值=0.9，表示IOU 0.9的不会被筛掉
|
||
},
|
||
'hrnet':{
|
||
'nof_joints': 17,
|
||
'c': 48,
|
||
'checkpoint_path': 'data/models/pose_hrnet_w48_384x288.pth'
|
||
},
|
||
'detect':{
|
||
'MIN_PERSON_JOINTS': 10,
|
||
'MIN_BBOX_AREA': 5000,
|
||
'MIN_JOINTS_CONF': 0.3,
|
||
'MIN_BBOX_LEN': 150
|
||
}
|
||
}
|
||
detector = Detector('yolo', 'hrnet', device, config)
|
||
for nf, imgname in enumerate(tqdm(imgnames)):
|
||
annotname = join(annot_root, os.path.basename(imgname).replace('.jpg', '.json'))
|
||
annot = create_annot_file(annotname, imgname)
|
||
img0 = cv2.imread(imgname)
|
||
annot['annots'] = detect_frame(detector, img0, 0)
|
||
for i in range(len(annot['annots'])):
|
||
x = annot['annots'][i]
|
||
x['area'] = max(x['bbox'][2] - x['bbox'][0], x['bbox'][3] - x['bbox'][1])**2
|
||
annot['annots'].sort(key=lambda x:-x['area'])
|
||
# 重新赋值人的ID
|
||
for i in range(len(annot['annots'])):
|
||
annot['annots'][i]['personID'] = i
|
||
save_json(annotname, annot)
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument('path', type=str, default=None)
|
||
parser.add_argument('--mode', type=str, default='openpose', choices=['openpose', 'yolo-hrnet'])
|
||
parser.add_argument('--handface', action='store_true')
|
||
parser.add_argument('--openpose', type=str,
|
||
default='/media/qing/Project/openpose')
|
||
parser.add_argument('--render', action='store_true', help='use to render the openpose 2d')
|
||
parser.add_argument('--no2d', action='store_true')
|
||
parser.add_argument('--start', type=int, default=0,
|
||
help='frame start')
|
||
parser.add_argument('--end', type=int, default=10000,
|
||
help='frame end')
|
||
parser.add_argument('--step', type=int, default=1,
|
||
help='frame step')
|
||
parser.add_argument('--debug', action='store_true')
|
||
args = parser.parse_args()
|
||
mode = args.mode
|
||
|
||
if os.path.isdir(args.path):
|
||
videos = sorted(glob(join(args.path, 'videos', '*.mp4')))
|
||
subs = []
|
||
for video in videos:
|
||
basename = extract_video(video, args.path, start=args.start, end=args.end, step=args.step)
|
||
subs.append(basename)
|
||
print('cameras: ', ' '.join(subs))
|
||
if not args.no2d:
|
||
for sub in subs:
|
||
image_root = join(args.path, 'images', sub)
|
||
annot_root = join(args.path, 'annots', sub)
|
||
if os.path.exists(annot_root):
|
||
print('skip ', annot_root)
|
||
continue
|
||
if mode == 'openpose':
|
||
extract_2d(args.openpose, image_root,
|
||
join(args.path, 'openpose', sub),
|
||
join(args.path, 'openpose_render', sub))
|
||
convert_from_openpose(
|
||
src=join(args.path, 'openpose', sub),
|
||
dst=annot_root
|
||
)
|
||
elif mode == 'yolo-hrnet':
|
||
extract_yolo_hrnet(image_root, annot_root)
|
||
else:
|
||
print(args.path, ' not exists')
|