<reponame>hukkai/mmaction2
import os, pickle, argparse, numpy
from mmdet.apis import init_detector, inference_detector
import mmcv
from PIL import Image
import multiprocessing as mp



def get_detect_list():
    all_videos = open('../../../data/ava_kinetics/annotations/kinetics_timestamps.txt').readlines()
    start_lookup = {}
    for line in all_videos:
        video_id, start_time, _ = line.split(',')
        start_lookup[video_id] = int(start_time)

    detect_list = []
    for video_id in os.listdir(rawframe_path):
        path = '%s/%s/'%(rawframe_path, video_id)
        frames = os.listdir(path)
        num_frames = len(frames) // 30 + 2
        for k in range(1, num_frames):
            frame = path + 'img_%05d.jpg'%(k*30+1)
            key = '%s,%04d'%(video_id, k + start_lookup[video_id])
            if os.path.isfile(frame):
                detect_list.append((key, frame))
    return tuple(detect_list)



def gpu_worker(device, detect_list, num_gpus):
    detect_sublist = [detect_list[i] for i in range(len(detect_list)) if i % num_gpus  == device]
    num_frames_per_gpu = len(detect_sublist)
    model = init_detector('cascade_x101.py',
                          'https://download.openmmlab.com/mmdetection/v2.0/cascade_rcnn/cascade_rcnn_x101_64x4d_fpn_20e_coco/'
                          'cascade_rcnn_x101_64x4d_fpn_20e_coco_20200509_224357-051557b1.pth',
                          'cuda:%d'%device)
    lookup = {}
    for idx, (key, frame_path) in enumerate(detect_sublist):
        result = inference_detector(model, frame_path)
        result = result[0][result[0][:, 4] >= 0.7]
        if len(result) > 0:
            img = Image.open(frame_path)
            result_ = []
            for record in result:
                h1, w1, h2, w2, score = record
                h1 /= img.size[0]
                h2 /= img.size[0]
                w1 /= img.size[1]
                w2 /= img.size[1]
                result_.append((h1, w1, h2, w2, score))
            lookup[key] = numpy.array(result_)

    with open('tmp_person_%d.pkl'%device, 'wb') as f:
        pickle.dump(lookup, f)
    return


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--rawframe_path', type=str, default='../../../data/ava_kinetics/rawframes/',
                             help='the path for kinetics videos')
    p.add_argument('--num_gpus', type=int, default=4)

    args = p.parse_args()
    rawframe_path, num_gpus = args.rawframe_path, args.num_gpus

    detect_list = get_detect_list()

    processes = []
    for rank in range(num_gpus):
        ctx = mp.get_context('spawn')
        p = ctx.Process(target=gpu_worker, args=(rank,detect_list, num_gpus))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

    lookup = {}
    for k in range(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('../../../data/ava_kinetics/annotations/kinetics_proposal.pkl', 'wb') as f:
        pickle.dump(lookup, f)
