55 lines
2.3 KiB
Python
55 lines
2.3 KiB
Python
#! /usr/bin/env python
|
|
# -*- coding: utf-8 -*-
|
|
|
|
# Copyright 2021 Imperial College London (Pingchuan Ma)
|
|
# Apache 2.0 (http://www.apache.org/licenses/LICENSE-2.0)
|
|
import cv2
|
|
|
|
|
|
def extract_opencv_generator(filename):
|
|
"""extract_opencv_generator.
|
|
|
|
:param filename: str, the filename for video.
|
|
"""
|
|
cap = cv2.VideoCapture(filename)
|
|
|
|
while(cap.isOpened()):
|
|
ret, frame = cap.read() # BGR
|
|
if ret:
|
|
yield frame
|
|
else:
|
|
break
|
|
cap.release()
|
|
|
|
|
|
def get_landmarks(multi_sub_landmarks):
|
|
"""get_landmarks.
|
|
|
|
:param multi_sub_landmarks: dict, a dictionary contains landmarks, bbox, and landmarks_scores.
|
|
"""
|
|
|
|
landmarks = [None] * len( multi_sub_landmarks["landmarks"])
|
|
for frame_idx in range(len(landmarks)):
|
|
if len(multi_sub_landmarks["landmarks"][frame_idx]) == 0:
|
|
continue
|
|
else:
|
|
# -- decide person id using maximal bounding box 0: Left, 1: top, 2: right, 3: bottom, probability
|
|
max_bbox_person_id = 0
|
|
max_bbox_len = multi_sub_landmarks["bbox"][frame_idx][max_bbox_person_id][2] + \
|
|
multi_sub_landmarks["bbox"][frame_idx][max_bbox_person_id][3] - \
|
|
multi_sub_landmarks["bbox"][frame_idx][max_bbox_person_id][0] - \
|
|
multi_sub_landmarks["bbox"][frame_idx][max_bbox_person_id][1]
|
|
landmark_scores = multi_sub_landmarks["landmarks_scores"][frame_idx][max_bbox_person_id]
|
|
for temp_person_id in range(1, len(multi_sub_landmarks["bbox"][frame_idx])):
|
|
temp_bbox_len = multi_sub_landmarks["bbox"][frame_idx][temp_person_id][2] + \
|
|
multi_sub_landmarks["bbox"][frame_idx][temp_person_id][3] - \
|
|
multi_sub_landmarks["bbox"][frame_idx][temp_person_id][0] - \
|
|
multi_sub_landmarks["bbox"][frame_idx][temp_person_id][1]
|
|
if temp_bbox_len > max_bbox_len:
|
|
max_bbox_person_id = temp_person_id
|
|
max_bbox_len = temp_bbox_len
|
|
landmark_scores = multi_sub_landmarks['landmarks_scores'][frame_idx][temp_person_id]
|
|
if landmark_scores[17:].min() >= 0.2:
|
|
landmarks[frame_idx] = multi_sub_landmarks["landmarks"][frame_idx][max_bbox_person_id]
|
|
return landmarks
|