|
|
|
|
|
import argparse
|
|
|
import multiprocessing as mp
|
|
|
import os
|
|
|
import pickle
|
|
|
|
|
|
import numpy as np
|
|
|
from mmdet.apis import inference_detector, init_detector
|
|
|
from mmdet.utils import register_all_modules
|
|
|
from PIL import Image
|
|
|
|
|
|
|
|
|
def get_vid_from_path(path):
|
|
|
video_id = path.split('/')[-1]
|
|
|
video_id = video_id.split('_')[:-2]
|
|
|
return '_'.join(video_id)
|
|
|
|
|
|
|
|
|
def prepare_det_lookup(datalist, frame_root):
|
|
|
with open(datalist) as f:
|
|
|
records = f.readlines()
|
|
|
det_lookup = {}
|
|
|
for record in records:
|
|
|
record = record.split(',')
|
|
|
folder_path = record[0]
|
|
|
video_id = get_vid_from_path(folder_path)
|
|
|
frame_id = int(record[1])
|
|
|
for idx in range(frame_id - 1, frame_id + 2):
|
|
|
proposal_id = '%s,%04d' % (video_id, idx)
|
|
|
det_lookup[proposal_id] = '%s/%s' % (frame_root, folder_path)
|
|
|
return det_lookup
|
|
|
|
|
|
|
|
|
def single_worker(rank, det_lookup, args):
|
|
|
detect_list = list(det_lookup)
|
|
|
detect_sublist = [
|
|
|
detect_list[i] for i in range(len(detect_list))
|
|
|
if i % args.num_gpus == rank
|
|
|
]
|
|
|
|
|
|
|
|
|
register_all_modules()
|
|
|
model = init_detector(
|
|
|
args.config, args.checkpoint, device='cuda:%d' % rank)
|
|
|
|
|
|
lookup = {}
|
|
|
for count, key in enumerate(detect_sublist):
|
|
|
try:
|
|
|
folder_path = det_lookup[key]
|
|
|
start = int(folder_path.split('/')[-1].split('_')[-2])
|
|
|
time = int(key.split(',')[1])
|
|
|
frame_id = (time - start) * 30 + 1
|
|
|
frame_path = '%s/img_%05d.jpg' % (folder_path, frame_id)
|
|
|
img = Image.open(frame_path)
|
|
|
result = inference_detector(model, frame_path)
|
|
|
bboxes = result._pred_instances.bboxes.cpu()
|
|
|
scores = result._pred_instances.scores.cpu()
|
|
|
labels = result._pred_instances.labels.cpu()
|
|
|
|
|
|
bboxes = bboxes[labels == 0]
|
|
|
scores = scores[labels == 0]
|
|
|
|
|
|
bboxes = bboxes[scores > 0.7].numpy()
|
|
|
scores = scores[scores > 0.7]
|
|
|
if scores.numel() > 0:
|
|
|
result_ = []
|
|
|
for idx, (h1, w1, h2, w2) in enumerate(bboxes):
|
|
|
h1 /= img.size[0]
|
|
|
h2 /= img.size[0]
|
|
|
w1 /= img.size[1]
|
|
|
w2 /= img.size[1]
|
|
|
score = scores[idx].item()
|
|
|
result_.append((h1, w1, h2, w2, score))
|
|
|
lookup[key] = np.array(result_)
|
|
|
except:
|
|
|
pass
|
|
|
|
|
|
with open('tmp_person_%d.pkl' % rank, 'wb') as f:
|
|
|
pickle.dump(lookup, f)
|
|
|
return
|
|
|
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
p = argparse.ArgumentParser()
|
|
|
p.add_argument(
|
|
|
'--avakinetics_root',
|
|
|
type=str,
|
|
|
default='../../../data/ava_kinetics',
|
|
|
help='the path to save ava-kinetics dataset')
|
|
|
p.add_argument(
|
|
|
'--datalist',
|
|
|
type=str,
|
|
|
default='../../../data/ava_kinetics/kinetics_train.csv',
|
|
|
help='the list for kinetics videos')
|
|
|
p.add_argument(
|
|
|
'--config',
|
|
|
type=str,
|
|
|
default='X-101-64x4d-FPN.py',
|
|
|
help='the human detector')
|
|
|
p.add_argument(
|
|
|
'--checkpoint',
|
|
|
type=str,
|
|
|
default='https://download.openmmlab.com/mmdetection/v2.0/'
|
|
|
'cascade_rcnn/cascade_rcnn_x101_64x4d_fpn_1x_coco/'
|
|
|
'cascade_rcnn_x101_64x4d_fpn_1x_coco_20200515_'
|
|
|
'075702-43ce6a30.pth',
|
|
|
help='the human detector checkpoint')
|
|
|
p.add_argument(
|
|
|
'--picklepath',
|
|
|
type=str,
|
|
|
default='../../../data/ava_kinetics/kinetics_proposal.pkl')
|
|
|
p.add_argument('--num_gpus', type=int, default=8)
|
|
|
|
|
|
args = p.parse_args()
|
|
|
|
|
|
frame_root = args.avakinetics_root + '/rawframes/'
|
|
|
det_lookup = prepare_det_lookup(args.datalist, frame_root)
|
|
|
|
|
|
processes = []
|
|
|
for rank in range(args.num_gpus):
|
|
|
ctx = mp.get_context('spawn')
|
|
|
p = ctx.Process(target=single_worker, args=(rank, det_lookup, args))
|
|
|
p.start()
|
|
|
processes.append(p)
|
|
|
|
|
|
for p in processes:
|
|
|
p.join()
|
|
|
|
|
|
lookup = {}
|
|
|
for k in range(args.num_gpus):
|
|
|
one_lookup = pickle.load(open('tmp_person_%d.pkl' % k, 'rb'))
|
|
|
os.remove('tmp_person_%d.pkl' % k)
|
|
|
for key in one_lookup:
|
|
|
lookup[key] = one_lookup[key]
|
|
|
|
|
|
with open(args.picklepath, 'wb') as f:
|
|
|
pickle.dump(lookup, f)
|
|
|
|
