|
|
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
|
|
#
|
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
|
# you may not use this file except in compliance with the License.
|
|
|
# You may obtain a copy of the License at
|
|
|
#
|
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
|
#
|
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
# See the License for the specific language governing permissions and
|
|
|
# limitations under the License.
|
|
|
|
|
|
import argparse
|
|
|
import paddle
|
|
|
import os, sys
|
|
|
import copy as cp
|
|
|
import cv2
|
|
|
import math
|
|
|
try:
|
|
|
import ppdet
|
|
|
except ImportError as e:
|
|
|
print(
|
|
|
f"Warning! {e}, [paddledet] package and it's dependencies is required for AVA."
|
|
|
)
|
|
|
|
|
|
__dir__ = os.path.dirname(os.path.abspath(__file__))
|
|
|
sys.path.append(os.path.abspath(os.path.join(__dir__, '../')))
|
|
|
|
|
|
from paddlevideo.modeling.builder import build_model
|
|
|
from paddlevideo.utils import get_config
|
|
|
from paddlevideo.loader.builder import build_dataloader, build_dataset, build_pipeline
|
|
|
from paddlevideo.metrics.ava_utils import read_labelmap
|
|
|
|
|
|
import time
|
|
|
from os import path as osp
|
|
|
import numpy as np
|
|
|
from paddlevideo.utils import get_config
|
|
|
import pickle
|
|
|
|
|
|
from paddlevideo.utils import (get_logger, load, mkdir, save)
|
|
|
import shutil
|
|
|
|
|
|
FONTFACE = cv2.FONT_HERSHEY_DUPLEX
|
|
|
FONTSCALE = 0.5
|
|
|
FONTCOLOR = (255, 255, 255) # BGR, white
|
|
|
MSGCOLOR = (128, 128, 128) # BGR, gray
|
|
|
THICKNESS = 1
|
|
|
LINETYPE = 1
|
|
|
|
|
|
|
|
|
def hex2color(h):
|
|
|
"""Convert the 6-digit hex string to tuple of 3 int value (RGB)"""
|
|
|
return (int(h[:2], 16), int(h[2:4], 16), int(h[4:], 16))
|
|
|
|
|
|
|
|
|
plate_blue = '03045e-023e8a-0077b6-0096c7-00b4d8-48cae4'
|
|
|
plate_blue = plate_blue.split('-')
|
|
|
plate_blue = [hex2color(h) for h in plate_blue]
|
|
|
plate_green = '004b23-006400-007200-008000-38b000-70e000'
|
|
|
plate_green = plate_green.split('-')
|
|
|
plate_green = [hex2color(h) for h in plate_green]
|
|
|
|
|
|
|
|
|
def abbrev(name):
|
|
|
"""Get the abbreviation of label name:
|
|
|
'take (an object) from (a person)' -> 'take ... from ...'
|
|
|
"""
|
|
|
while name.find('(') != -1:
|
|
|
st, ed = name.find('('), name.find(')')
|
|
|
name = name[:st] + '...' + name[ed + 1:]
|
|
|
return name
|
|
|
|
|
|
|
|
|
# annotations is pred results
|
|
|
def visualize(frames, annotations, plate=plate_blue, max_num=5):
|
|
|
"""Visualize frames with predicted annotations.
|
|
|
Args:
|
|
|
frames (list[np.ndarray]): Frames for visualization, note that
|
|
|
len(frames) % len(annotations) should be 0.
|
|
|
annotations (list[list[tuple]]): The predicted results.
|
|
|
plate (str): The plate used for visualization. Default: plate_blue.
|
|
|
max_num (int): Max number of labels to visualize for a person box.
|
|
|
Default: 5，目前不能大于5.
|
|
|
Returns:
|
|
|
list[np.ndarray]: Visualized frames.
|
|
|
"""
|
|
|
|
|
|
assert max_num + 1 <= len(plate)
|
|
|
plate = [x[::-1] for x in plate]
|
|
|
frames_ = cp.deepcopy(frames)
|
|
|
nf, na = len(frames), len(annotations)
|
|
|
assert nf % na == 0
|
|
|
nfpa = len(frames) // len(annotations)
|
|
|
anno = None
|
|
|
h, w, _ = frames[0].shape
|
|
|
# proposals被归一化需要还原真实坐标值
|
|
|
scale_ratio = np.array([w, h, w, h])
|
|
|
|
|
|
for i in range(na):
|
|
|
anno = annotations[i]
|
|
|
if anno is None:
|
|
|
continue
|
|
|
for j in range(nfpa):
|
|
|
ind = i * nfpa + j
|
|
|
frame = frames_[ind]
|
|
|
for ann in anno:
|
|
|
box = ann[0]
|
|
|
label = ann[1]
|
|
|
if not len(label):
|
|
|
continue
|
|
|
score = ann[2]
|
|
|
box = (box * scale_ratio).astype(np.int64)
|
|
|
st, ed = tuple(box[:2]), tuple(box[2:])
|
|
|
cv2.rectangle(frame, st, ed, plate[0], 2)
|
|
|
for k, lb in enumerate(label):
|
|
|
if k >= max_num:
|
|
|
break
|
|
|
text = abbrev(lb)
|
|
|
text = ': '.join([text, str(score[k])])
|
|
|
location = (0 + st[0], 18 + k * 18 + st[1])
|
|
|
textsize = cv2.getTextSize(text, FONTFACE, FONTSCALE,
|
|
|
THICKNESS)[0]
|
|
|
textwidth = textsize[0]
|
|
|
diag0 = (location[0] + textwidth, location[1] - 14)
|
|
|
diag1 = (location[0], location[1] + 2)
|
|
|
cv2.rectangle(frame, diag0, diag1, plate[k + 1], -1)
|
|
|
cv2.putText(frame, text, location, FONTFACE, FONTSCALE,
|
|
|
FONTCOLOR, THICKNESS, LINETYPE)
|
|
|
|
|
|
return frames_
|
|
|
|
|
|
|
|
|
def frame_extraction(video_path, target_dir):
|
|
|
"""Extract frames given video_path.
|
|
|
Args:
|
|
|
video_path (str): The video_path.
|
|
|
"""
|
|
|
|
|
|
if not os.path.exists(target_dir):
|
|
|
os.makedirs(target_dir, exist_ok=True)
|
|
|
|
|
|
# Should be able to handle videos up to several hours
|
|
|
frame_tmpl = osp.join(target_dir, '{:05d}.jpg')
|
|
|
vid = cv2.VideoCapture(video_path)
|
|
|
|
|
|
FPS = int(vid.get(5))
|
|
|
|
|
|
frames = []
|
|
|
frame_paths = []
|
|
|
|
|
|
flag, frame = vid.read()
|
|
|
index = 1
|
|
|
while flag:
|
|
|
frames.append(frame)
|
|
|
frame_path = frame_tmpl.format(index)
|
|
|
frame_paths.append(frame_path)
|
|
|
cv2.imwrite(frame_path, frame)
|
|
|
index += 1
|
|
|
flag, frame = vid.read()
|
|
|
return frame_paths, frames, FPS
|
|
|
|
|
|
|
|
|
def parse_args():
|
|
|
def str2bool(v):
|
|
|
return v.lower() in ("true", "t", "1")
|
|
|
|
|
|
# general params
|
|
|
parser = argparse.ArgumentParser("PaddleVideo Inference model script")
|
|
|
parser.add_argument('-c',
|
|
|
'--config',
|
|
|
type=str,
|
|
|
default='configs/example.yaml',
|
|
|
help='config file path')
|
|
|
|
|
|
parser.add_argument('--video_path', help='video file/url')
|
|
|
|
|
|
parser.add_argument('-o',
|
|
|
'--override',
|
|
|
action='append',
|
|
|
default=[],
|
|
|
help='config options to be overridden')
|
|
|
parser.add_argument('-w',
|
|
|
'--weights',
|
|
|
type=str,
|
|
|
help='weights for finetuning or testing')
|
|
|
|
|
|
#detection_model_name
|
|
|
parser.add_argument('--detection_model_name',
|
|
|
help='the name of detection model ')
|
|
|
# detection_model_weights
|
|
|
parser.add_argument('--detection_model_weights',
|
|
|
help='the weights path of detection model ')
|
|
|
|
|
|
# params for predict
|
|
|
parser.add_argument('--out-filename',
|
|
|
default='ava_det_demo.mp4',
|
|
|
help='output filename')
|
|
|
parser.add_argument('--predict-stepsize',
|
|
|
default=8,
|
|
|
type=int,
|
|
|
help='give out a prediction per n frames')
|
|
|
parser.add_argument(
|
|
|
'--output-stepsize',
|
|
|
default=4,
|
|
|
type=int,
|
|
|
help=('show one frame per n frames in the demo, we should have: '
|
|
|
'predict_stepsize % output_stepsize == 0'))
|
|
|
parser.add_argument('--output-fps',
|
|
|
default=6,
|
|
|
type=int,
|
|
|
help='the fps of demo video output')
|
|
|
|
|
|
return parser.parse_args()
|
|
|
|
|
|
|
|
|
# 一帧的结果。根据概率大小进行排序
|
|
|
def pack_result(human_detection, result):
|
|
|
"""Short summary.
|
|
|
Args:
|
|
|
human_detection (np.ndarray): Human detection result.
|
|
|
result (type): The predicted label of each human proposal.
|
|
|
Returns:
|
|
|
tuple: Tuple of human proposal, label name and label score.
|
|
|
"""
|
|
|
results = []
|
|
|
if result is None:
|
|
|
return None
|
|
|
|
|
|
for prop, res in zip(human_detection, result):
|
|
|
res.sort(key=lambda x: -x[1])
|
|
|
|
|
|
results.append((prop, [x[0] for x in res], [x[1] for x in res]))
|
|
|
|
|
|
return results
|
|
|
|
|
|
|
|
|
# 构造数据处理需要的results
|
|
|
def get_timestep_result(frame_dir, timestamp, clip_len, frame_interval, FPS):
|
|
|
result = {}
|
|
|
|
|
|
result["frame_dir"] = frame_dir
|
|
|
|
|
|
frame_num = len(os.listdir(frame_dir))
|
|
|
|
|
|
dir_name = frame_dir.split("/")[-1]
|
|
|
result["video_id"] = dir_name
|
|
|
|
|
|
result['timestamp'] = timestamp
|
|
|
|
|
|
timestamp_str = '{:04d}'.format(timestamp)
|
|
|
img_key = dir_name + "," + timestamp_str
|
|
|
result['img_key'] = img_key
|
|
|
|
|
|
result['shot_info'] = (1, frame_num)
|
|
|
result['fps'] = FPS
|
|
|
|
|
|
result['suffix'] = '{:05}.jpg'
|
|
|
|
|
|
result['timestamp_start'] = 1
|
|
|
result['timestamp_end'] = int(frame_num / result['fps'])
|
|
|
|
|
|
return result
|
|
|
|
|
|
|
|
|
def detection_inference(frame_paths, output_dir, model_name, weights_path):
|
|
|
"""Detect human boxes given frame paths.
|
|
|
Args:
|
|
|
frame_paths (list[str]): The paths of frames to do detection inference.
|
|
|
Returns:
|
|
|
list[np.ndarray]: The human detection results.
|
|
|
"""
|
|
|
|
|
|
detection_cfg = ppdet.model_zoo.get_config_file(model_name)
|
|
|
detection_cfg = ppdet.core.workspace.load_config(detection_cfg)
|
|
|
detection_trainer = ppdet.engine.Trainer(detection_cfg, mode='test')
|
|
|
detection_trainer.load_weights(weights_path)
|
|
|
|
|
|
print('Performing Human Detection for each frame')
|
|
|
|
|
|
detection_trainer.predict(frame_paths, output_dir=output_dir, save_txt=True)
|
|
|
|
|
|
print("finish object detection")
|
|
|
|
|
|
results = []
|
|
|
|
|
|
for frame_path in frame_paths:
|
|
|
(file_dir, file_name) = os.path.split(frame_path)
|
|
|
(file_path, ext) = os.path.splitext(frame_path)
|
|
|
|
|
|
txt_file_name = file_name.replace(ext, ".txt")
|
|
|
txt_path = os.path.join(output_dir, txt_file_name)
|
|
|
results.append(txt_path)
|
|
|
|
|
|
return results
|
|
|
|
|
|
|
|
|
def get_detection_result(txt_file_path, img_h, img_w, person_det_score_thr):
|
|
|
"""
|
|
|
根据检测结果文件得到图像中人的检测框(proposals)和置信度（scores）
|
|
|
txt_file_path:检测结果存放路径
|
|
|
img_h:图像高度
|
|
|
img_w:图像宽度
|
|
|
"""
|
|
|
|
|
|
proposals = []
|
|
|
scores = []
|
|
|
|
|
|
with open(txt_file_path, 'r') as detection_file:
|
|
|
lines = detection_file.readlines()
|
|
|
for line in lines: # person 0.9842637181282043 0.0 469.1407470703125 944.7770385742188 831.806396484375
|
|
|
items = line.split(" ")
|
|
|
if items[0] != 'person': #只要人
|
|
|
continue
|
|
|
|
|
|
score = items[1]
|
|
|
|
|
|
if (float)(score) < person_det_score_thr:
|
|
|
continue
|
|
|
|
|
|
x1 = (float(items[2])) / img_w
|
|
|
y1 = ((float)(items[3])) / img_h
|
|
|
box_w = ((float)(items[4]))
|
|
|
box_h = ((float)(items[5]))
|
|
|
|
|
|
x2 = (float(items[2]) + box_w) / img_w
|
|
|
y2 = (float(items[3]) + box_h) / img_h
|
|
|
|
|
|
scores.append(score)
|
|
|
|
|
|
proposals.append([x1, y1, x2, y2])
|
|
|
|
|
|
return np.array(proposals), np.array(scores)
|
|
|
|
|
|
|
|
|
@paddle.no_grad()
|
|
|
def main(args):
|
|
|
config = get_config(args.config, show=False) #parse config file
|
|
|
|
|
|
# extract frames from video
|
|
|
video_path = args.video_path
|
|
|
frame_dir = 'tmp_frames'
|
|
|
frame_paths, frames, FPS = frame_extraction(video_path, frame_dir)
|
|
|
|
|
|
num_frame = len(frame_paths) #视频秒数*FPS
|
|
|
assert num_frame != 0
|
|
|
print("Frame Number：", num_frame)
|
|
|
|
|
|
# 帧图像高度和宽度
|
|
|
h, w, _ = frames[0].shape
|
|
|
|
|
|
# Get clip_len, frame_interval and calculate center index of each clip
|
|
|
data_process_pipeline = build_pipeline(config.PIPELINE.test) #测试时输出处理流水配置
|
|
|
|
|
|
clip_len = config.PIPELINE.test.sample['clip_len']
|
|
|
assert clip_len % 2 == 0, 'We would like to have an even clip_len'
|
|
|
frame_interval = config.PIPELINE.test.sample['frame_interval']
|
|
|
|
|
|
# 此处关键帧每秒取一个
|
|
|
clip_len = config.PIPELINE.test.sample['clip_len']
|
|
|
assert clip_len % 2 == 0, 'We would like to have an even clip_len'
|
|
|
frame_interval = config.PIPELINE.test.sample['frame_interval']
|
|
|
window_size = clip_len * frame_interval
|
|
|
timestamps = np.arange(window_size // 2, (num_frame + 1 - window_size // 2),
|
|
|
args.predict_stepsize)
|
|
|
print("timetamps number:", len(timestamps))
|
|
|
|
|
|
# get selected frame list according to timestamps
|
|
|
selected_frame_list = []
|
|
|
for timestamp in timestamps:
|
|
|
selected_frame_list.append(frame_paths[timestamp - 1])
|
|
|
|
|
|
# Load label_map
|
|
|
label_map_path = config.DATASET.test['label_file']
|
|
|
categories, class_whitelist = read_labelmap(open(label_map_path))
|
|
|
label_map = {}
|
|
|
for item in categories:
|
|
|
id = item['id']
|
|
|
name = item['name']
|
|
|
label_map[id] = name
|
|
|
|
|
|
# Construct model.
|
|
|
if config.MODEL.backbone.get('pretrained'):
|
|
|
config.MODEL.backbone.pretrained = '' # disable pretrain model init
|
|
|
model = build_model(config.MODEL)
|
|
|
|
|
|
model.eval()
|
|
|
state_dicts = load(args.weights)
|
|
|
model.set_state_dict(state_dicts)
|
|
|
|
|
|
detection_result_dir = 'tmp_detection'
|
|
|
detection_model_name = args.detection_model_name
|
|
|
detection_model_weights = args.detection_model_weights
|
|
|
detection_txt_list = detection_inference(selected_frame_list,
|
|
|
detection_result_dir,
|
|
|
detection_model_name,
|
|
|
detection_model_weights)
|
|
|
assert len(detection_txt_list) == len(timestamps)
|
|
|
|
|
|
print('Performing SpatioTemporal Action Detection for each clip')
|
|
|
human_detections = []
|
|
|
predictions = []
|
|
|
|
|
|
index = 0
|
|
|
for timestamp, detection_txt_path in zip(timestamps, detection_txt_list):
|
|
|
proposals, scores = get_detection_result(
|
|
|
detection_txt_path, h, w,
|
|
|
(float)(config.DATASET.test['person_det_score_thr']))
|
|
|
if proposals.shape[0] == 0:
|
|
|
predictions.append(None)
|
|
|
human_detections.append(None)
|
|
|
continue
|
|
|
|
|
|
human_detections.append(proposals)
|
|
|
|
|
|
result = get_timestep_result(frame_dir,
|
|
|
timestamp,
|
|
|
clip_len,
|
|
|
frame_interval,
|
|
|
FPS=FPS)
|
|
|
result["proposals"] = proposals
|
|
|
result["scores"] = scores
|
|
|
|
|
|
new_result = data_process_pipeline(result)
|
|
|
proposals = new_result['proposals']
|
|
|
|
|
|
img_slow = new_result['imgs'][0]
|
|
|
img_slow = img_slow[np.newaxis, :]
|
|
|
img_fast = new_result['imgs'][1]
|
|
|
img_fast = img_fast[np.newaxis, :]
|
|
|
|
|
|
proposals = proposals[np.newaxis, :]
|
|
|
|
|
|
scores = scores[np.newaxis, :]
|
|
|
|
|
|
img_shape = np.asarray(new_result['img_shape'])
|
|
|
img_shape = img_shape[np.newaxis, :]
|
|
|
|
|
|
data = [
|
|
|
paddle.to_tensor(img_slow, dtype='float32'),
|
|
|
paddle.to_tensor(img_fast, dtype='float32'),
|
|
|
paddle.to_tensor(proposals, dtype='float32'), scores,
|
|
|
paddle.to_tensor(img_shape, dtype='int32')
|
|
|
]
|
|
|
|
|
|
with paddle.no_grad():
|
|
|
result = model(data, mode='infer')
|
|
|
|
|
|
result = result[0]
|
|
|
prediction = []
|
|
|
|
|
|
person_num = proposals.shape[1]
|
|
|
# N proposals
|
|
|
for i in range(person_num):
|
|
|
prediction.append([])
|
|
|
|
|
|
# Perform action score thr
|
|
|
for i in range(len(result)):
|
|
|
if i + 1 not in class_whitelist:
|
|
|
continue
|
|
|
for j in range(person_num):
|
|
|
if result[i][j, 4] > config.MODEL.head['action_thr']:
|
|
|
prediction[j].append((label_map[i + 1], result[i][j,
|
|
|
4]))
|
|
|
predictions.append(prediction)
|
|
|
|
|
|
index = index + 1
|
|
|
if index % 10 == 0:
|
|
|
print(index, "/", len(timestamps))
|
|
|
|
|
|
results = []
|
|
|
for human_detection, prediction in zip(human_detections, predictions):
|
|
|
results.append(pack_result(human_detection, prediction))
|
|
|
|
|
|
def dense_timestamps(timestamps, n):
|
|
|
"""Make it nx frames."""
|
|
|
old_frame_interval = (timestamps[1] - timestamps[0])
|
|
|
start = timestamps[0] - old_frame_interval / n * (n - 1) / 2
|
|
|
new_frame_inds = np.arange(
|
|
|
len(timestamps) * n) * old_frame_interval / n + start
|
|
|
return new_frame_inds.astype(np.int)
|
|
|
|
|
|
dense_n = int(args.predict_stepsize / args.output_stepsize) #30
|
|
|
frames = [
|
|
|
cv2.imread(frame_paths[i - 1])
|
|
|
for i in dense_timestamps(timestamps, dense_n)
|
|
|
]
|
|
|
|
|
|
vis_frames = visualize(frames, results)
|
|
|
|
|
|
try:
|
|
|
import moviepy.editor as mpy
|
|
|
except ImportError:
|
|
|
raise ImportError('Please install moviepy to enable output file')
|
|
|
|
|
|
vid = mpy.ImageSequenceClip([x[:, :, ::-1] for x in vis_frames],
|
|
|
fps=args.output_fps)
|
|
|
vid.write_videofile(args.out_filename)
|
|
|
print("finish write !")
|
|
|
|
|
|
# delete tmp files and dirs
|
|
|
shutil.rmtree(frame_dir)
|
|
|
shutil.rmtree(detection_result_dir)
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
args = parse_args() #解析参数
|
|
|
main(args)
|