def Demo(opt):
	if opt.mode == 'video':
		if opt.modality == 'R':
			video_loader = VideoLoader(opt.video_path).start()
			(fourcc, fps, w, h) = video_loader.videoinfo()
			#load recognition loader
			print('Loading model...')
			sys.stdout.flush()
			Recognition= RecognitionLoader(video_loader, recognition_classes, opt.temporal_sample_length, step=opt.recognize_sample_step).start()
			#data writer
			if opt.save_video:
				if not os.path.exists(opt.out_path):
					os.makedirs(opt.out_path)
				save_path = os.path.join(opt.out_path, 'VideoDemo.avi')
				writer = DataWriter(Recognition, savepath=save_path, save_video=opt.save_video, vis=opt.show, fps=fps, frameSize=(w, h)).start()
				print('output video has been saved in %s'%save_path)
				writer.stop()
		elif opt.modality == 'DR':
			if opt.save_video:
				if not os.path.exists(opt.out_path):
					os.makedirs(opt.out_path)
				out_path = os.path.join(opt.out_path, 'result.avi')
				video_DR = Video_Detect_Recognition_Processor(opt.video_path, out_path, detect_classes, recognition_classes).start()		
				video_DR.read_detect()
	if opt.mode == 'camera':
		if opt.modality == 'R':
			CameraRecognitionLoader = CameraRecognition(classes).start()
		elif opt.modality== 'DR':
			#Camera = CameraLoader().start()
			#Detect = Detect_Recognition_Processor(classes).start()
			#Recognize = RecognizeProcessor(Detect, classes).start()
			DR = Read_Detect_Recognition_Processor(detect_classes, recognition_classes).start()
			DR.read_detect()
Esempio n. 2
0
def main():
    outSize = (64, 64)
    sp = False
    video_loader = VideoLoader("kunkun_nmsl_Trim.mp4",
                               batchSize=32,
                               wh_rate=(1, 1),
                               inDim=outSize,
                               sp=sp)

    (fourcc, fps, frameSize, shape_dst, ow, oh) = video_loader.videoinfo()
    # fps = 10
    # frameSize = outSize
    print(fourcc, fps, frameSize, shape_dst, ow, oh)
    num_frames = video_loader.videoLen()
    save_path = 'out_video/kunkun_nmsl_2.avi'
    fourcc = cv2.VideoWriter_fourcc(*'XVID')

    stream = cv2.VideoWriter(save_path, fourcc, fps, frameSize)
    assert stream.isOpened(), 'Cannot open video for writing'

    video_loader.start()
    save_loader = saveFrame(video_loader, sp=sp).start()
    ascend_loader = AscendHost(save_loader, sp=sp).start()
    collect_loader = DataCollector(ascend_loader, num_job=8, sp=sp).start()
    post_processor = PostProcessor(collect_loader,
                                   remap=(shape_dst, ow, oh),
                                   vis=True,
                                   nobg=True,
                                   sp=sp).start()
    # post_processor = PostProcessor(collect_loader, remap=(shape_dst, ow, oh), vis=True, sp=sp).start()

    cnt = 0
    while True:
        cnt += 1
        start = time.time()
        (img_render, img, orig_img, im_name) = post_processor.getitem()
        if orig_img is None:
            print('All done!')
            break

        # orig_img = img_render

        cv2.imshow("Demo", orig_img)
        cv2.waitKey(90)
        stream.write(orig_img)
        interval = time.time() - start
        print("Time: {:.4f}\t{:04d}/{:04d}".format(interval, cnt, num_frames))
        # sys.stdout.flush()

    stream.release()
Esempio n. 3
0
def get_det_processor(file_name):
    videofile = file_name
    mode = args.mode
    if not os.path.exists(args.outputpath):
        os.mkdir(args.outputpath)

    if not len(videofile):
        raise IOError('Error: must contain --video')

    # Load input video
    data_loader = VideoLoader(videofile, batchSize=args.detbatch).start()
    (fourcc, fps, frameSize) = data_loader.videoinfo()

    # Load detection loader
    print('Loading YOLO model..')
    sys.stdout.flush()
    det_loader = DetectionLoader(data_loader, batchSize=args.detbatch).start()
    det_processor = DetectionProcessor(det_loader).start()
    return det_processor
 def __init__(self, videofile, mode='normal'):
     self.videofile = videofile
     self.mode = mode
     self.data_loader = VideoLoader(self.videofile,
                                    batchSize=args.detbatch).start()
     (fourcc, fps, frameSize) = self.data_loader.videoinfo()
     self.fourcc = fourcc
     self.fps = fps
     self.frameSize = frameSize
     self.det_loader = DetectionLoader(self.data_loader,
                                       batchSize=args.detbatch).start()
     self.det_processor = DetectionProcessor(self.det_loader).start()
     self.pose_dataset = Mscoco()
     save_path = os.path.join(
         args.outputpath, 'AlphaPose_' +
         ntpath.basename(self.videofile).split('.')[0] + '.mp4')
     self.writer = DataWriter(args.save_video, save_path,
                              cv2.VideoWriter_fourcc(*'DIVX'), self.fps,
                              self.frameSize).start()
     self.results = list()
Esempio n. 5
0
def handle_video(videofile):
    args.video = videofile
    videofile = args.video
    mode = args.mode

    if not len(videofile):
        raise IOError('Error: must contain --video')

    # Load input video
    data_loader = VideoLoader(videofile, batchSize=args.detbatch).start()
    (fourcc, fps, frameSize) = data_loader.videoinfo()

    print('the video is {} f/s'.format(fps))

    # Load detection loader
    print('Loading YOLO model..')
    sys.stdout.flush()
    det_loader = DetectionLoader(data_loader, batchSize=args.detbatch).start()
    #  start a thread to read frames from the file video stream
    det_processor = DetectionProcessor(det_loader).start()

    # Load pose model
    pose_dataset = Mscoco()
    if args.fast_inference:
        pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
    else:
        pose_model = InferenNet(4 * 1 + 1, pose_dataset)
    pose_model.cuda()
    pose_model.eval()

    runtime_profile = {'dt': [], 'pt': [], 'pn': []}

    # Data writer
    save_path = os.path.join(
        args.outputpath,
        'AlphaPose_' + ntpath.basename(videofile).split('.')[0] + '.avi')
    writer = DataWriter(args.save_video, save_path,
                        cv2.VideoWriter_fourcc(*'XVID'), fps,
                        frameSize).start()

    im_names_desc = tqdm(range(data_loader.length()))
    batchSize = args.posebatch
    for i in im_names_desc:
        start_time = getTime()
        with torch.no_grad():
            (inps, orig_img, im_name, boxes, scores, pt1,
             pt2) = det_processor.read()
            if orig_img is None:
                break
            if boxes is None or boxes.nelement() == 0:
                writer.save(None, None, None, None, None, orig_img,
                            im_name.split('/')[-1])
                continue

            ckpt_time, det_time = getTime(start_time)
            runtime_profile['dt'].append(det_time)
            # Pose Estimation

            datalen = inps.size(0)
            leftover = 0
            if (datalen) % batchSize:
                leftover = 1
            num_batches = datalen // batchSize + leftover
            hm = []
            for j in range(num_batches):
                inps_j = inps[j * batchSize:min((j + 1) *
                                                batchSize, datalen)].cuda()
                hm_j = pose_model(inps_j)
                hm.append(hm_j)
            hm = torch.cat(hm)
            ckpt_time, pose_time = getTime(ckpt_time)
            runtime_profile['pt'].append(pose_time)

            hm = hm.cpu().data
            writer.save(boxes, scores, hm, pt1, pt2, orig_img,
                        im_name.split('/')[-1])

            ckpt_time, post_time = getTime(ckpt_time)
            runtime_profile['pn'].append(post_time)

        if args.profile:
            # TQDM
            im_names_desc.set_description(
                'det time: {dt:.4f} | pose time: {pt:.4f} | post processing: {pn:.4f}'
                .format(dt=np.mean(runtime_profile['dt']),
                        pt=np.mean(runtime_profile['pt']),
                        pn=np.mean(runtime_profile['pn'])))

    if (args.save_img or args.save_video) and not args.vis_fast:
        print(
            '===========================> Rendering remaining images in the queue...'
        )
        print(
            '===========================> If this step takes too long, you can enable the --vis_fast flag to use fast rendering (real-time).'
        )
    while (writer.running()):
        pass
    writer.stop()
    final_result = writer.results()

    # 获取第 0 个框的人
    kpts = []
    for i in range(len(final_result)):
        try:
            preds = final_result[i]['result']
            # preds[i]['keypoints'] (17,2)
            # preds[i]['kp_score'] (17,1)
            # preds[i]['proposal_score'] (1)
            # 选择 y 坐标最大的人 —— 用于打羽毛球视频
            max_index = 0
            min_index = 0
            # max_y = np.mean(preds[0]['keypoints'].data.numpy()[:, 1])
            min_x = np.mean(preds[0]['keypoints'].data.numpy()[:, 0])
            max_x = np.mean(preds[0]['keypoints'].data.numpy()[:, 0])

            for k in range(len(preds)):
                # tmp_y = np.mean(preds[k]['keypoints'].data.numpy()[:, 1])
                tmp_x = np.mean(preds[k]['keypoints'].data.numpy()[:, 0])
                # if tmp_y > max_y:
                if tmp_x < min_x:
                    min_index = k
                    # max_y = tmp_y
                    min_x = tmp_x
            for k in range(len(preds)):
                # tmp_y = np.mean(preds[k]['keypoints'].data.numpy()[:, 1])
                tmp_x = np.mean(preds[k]['keypoints'].data.numpy()[:, 0])
                # if tmp_y > max_y:
                if tmp_x > max_x:
                    max_index = k
                    max_x = tmp_x
            mid_index = 0
            for k in range(len(preds)):
                if k == max_index or k == min_index:
                    continue
                mid_index = k
            kpt = preds[mid_index]['keypoints']
            # kpt = final_result[i]['result'][0]['keypoints']
            kpts.append(kpt.data.numpy())

        except:
            # print(sys.exc_info())
            print('error...')

    filename = os.path.basename(args.video).split('.')[0]
    name = filename + '.npz'
    kpts = np.array(kpts).astype(np.float32)
    # print('kpts npz save in ', name)
    # np.savez_compressed(name, kpts=kpts)
    return kpts
args.dataset = 'coco'
if not args.sp:
    torch.multiprocessing.set_start_method('forkserver', force=True)
    torch.multiprocessing.set_sharing_strategy('file_system')

if __name__ == "__main__":
    videofile = args.video
    mode = args.mode
    if not os.path.exists(args.outputpath):
        os.mkdir(args.outputpath)

    if not len(videofile):
        raise IOError('Error: must contain --video')

    # Load input video
    data_loader = VideoLoader(videofile, batchSize=args.detbatch).start()
    (fourcc, fps, frameSize) = data_loader.videoinfo()

    # Load detection loader
    print('Loading YOLO model..')
    sys.stdout.flush()
    det_loader = DetectionLoader(data_loader, batchSize=args.detbatch).start()
    det_processor = DetectionProcessor(det_loader).start()

    # Load pose model
    pose_dataset = Mscoco()
    if args.fast_inference:
        pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
    else:
        pose_model = InferenNet(4 * 1 + 1, pose_dataset)
    pose_model.cuda()
Esempio n. 7
0
def handle_video(video_file):
    # =========== common ===============
    args.video = video_file
    base_name = os.path.basename(args.video)
    video_name = base_name[:base_name.rfind('.')]
    # =========== end common ===============
    # =========== image ===============
    # img_path = f'outputs/alpha_pose_{video_name}/split_image/'
    # args.inputpath = img_path
    # args.outputpath = f'outputs/alpha_pose_{video_name}'
    # if os.path.exists(args.outputpath):
    #     shutil.rmtree(f'{args.outputpath}/vis', ignore_errors=True)
    # else:
    #     os.mkdir(args.outputpath)

    # # if not len(video_file):
    # #     raise IOError('Error: must contain --video')

    # if len(img_path) and img_path != '/':
    #     for root, dirs, files in os.walk(img_path):
    #         im_names = sorted([f for f in files if 'png' in f or 'jpg' in f])
    # else:
    #     raise IOError('Error: must contain either --indir/--list')

    # # Load input images
    # data_loader = ImageLoader(im_names, batchSize=args.detbatch, format='yolo').start()
    # print(f'Totally {data_loader.datalen} images')
    # =========== end image ===============
    # =========== video ===============
    args.outputpath = f'outputs/alpha_pose_{video_name}'
    if os.path.exists(args.outputpath):
        shutil.rmtree(f'{args.outputpath}/vis', ignore_errors=True)
    else:
        os.mkdir(args.outputpath)
    videofile = args.video
    mode = args.mode
    if not len(videofile):
        raise IOError('Error: must contain --video')
    # Load input video
    data_loader = VideoLoader(videofile, batchSize=args.detbatch).start()
    (fourcc, fps, frameSize) = data_loader.videoinfo()
    print('the video is {} f/s'.format(fps))
    # =========== end video ===============
    # Load detection loader
    print('Loading YOLO model..')
    sys.stdout.flush()
    det_loader = DetectionLoader(data_loader, batchSize=args.detbatch).start()
    #  start a thread to read frames from the file video stream
    det_processor = DetectionProcessor(det_loader).start()
    # Load pose model
    pose_dataset = Mscoco()
    if args.fast_inference:
        pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
    else:
        pose_model = InferenNet(4 * 1 + 1, pose_dataset)
    pose_model  #.cuda()
    pose_model.eval()
    runtime_profile = {'dt': [], 'pt': [], 'pn': []}
    # Data writer
    save_path = os.path.join(
        args.outputpath,
        'AlphaPose_' + ntpath.basename(video_file).split('.')[0] + '.avi')
    # writer = DataWriter(args.save_video, save_path, cv2.VideoWriter_fourcc(*'XVID'), fps, frameSize).start()
    writer = DataWriter(args.save_video).start()
    print('Start pose estimation...')
    im_names_desc = tqdm(range(data_loader.length()))
    batchSize = args.posebatch
    for i in im_names_desc:

        start_time = getTime()
        with torch.no_grad():
            (inps, orig_img, im_name, boxes, scores, pt1,
             pt2) = det_processor.read()
            if orig_img is None:
                print(f'{i}-th image read None: handle_video')
                break
            if boxes is None or boxes.nelement() == 0:
                writer.save(None, None, None, None, None, orig_img,
                            im_name.split('/')[-1])
                continue

            ckpt_time, det_time = getTime(start_time)
            runtime_profile['dt'].append(det_time)
            # Pose Estimation

            datalen = inps.size(0)
            leftover = 0
            if datalen % batchSize:
                leftover = 1
            num_batches = datalen // batchSize + leftover
            hm = []
            for j in range(num_batches):
                inps_j = inps[j * batchSize:min((j + 1) *
                                                batchSize, datalen)]  #.cuda()
                hm_j = pose_model(inps_j)
                hm.append(hm_j)
            hm = torch.cat(hm)
            ckpt_time, pose_time = getTime(ckpt_time)
            runtime_profile['pt'].append(pose_time)

            hm = hm.cpu().data
            writer.save(boxes, scores, hm, pt1, pt2, orig_img,
                        im_name.split('/')[-1])

            ckpt_time, post_time = getTime(ckpt_time)
            runtime_profile['pn'].append(post_time)

        if args.profile:
            # TQDM
            im_names_desc.set_description(
                'det time: {dt:.4f} | pose time: {pt:.4f} | post processing: {pn:.4f}'
                .format(dt=np.mean(runtime_profile['dt']),
                        pt=np.mean(runtime_profile['pt']),
                        pn=np.mean(runtime_profile['pn'])))
    if (args.save_img or args.save_video) and not args.vis_fast:
        print(
            '===========================> Rendering remaining images in the queue...'
        )
        print(
            '===========================> If this step takes too long, you can enable the --vis_fast flag to use fast rendering (real-time).'
        )
    while writer.running():
        pass
    writer.stop()
    final_result = writer.results()
    write_json(final_result, args.outputpath)

    return final_result, video_name
Esempio n. 8
0
def alphapose_process_video(videofile, pose_result_handler, inference_steps=1):
    if not len(videofile):
        raise IOError("Error: must contain --video")

    # Load input video
    print(f"Opening video {videofile}")
    data_loader = VideoLoader(videofile, batchSize=args.detbatch).start()
    (fourcc, fps, frameSize) = data_loader.videoinfo()

    # Load detection loader
    print("Loading YOLO model..")
    sys.stdout.flush()
    det_loader = DetectionLoader(
        data_loader,
        batchSize=args.detbatch,
        path=ALPHAPOSE_DIR,
        inference_steps=inference_steps,
    ).start()
    det_processor = DetectionProcessor(det_loader).start()

    # Load pose model
    pose_dataset = Mscoco()
    if args.fast_inference:
        pose_model = InferenNet_fast(4 * 1 + 1,
                                     pose_dataset,
                                     path=ALPHAPOSE_DIR)
    else:
        pose_model = InferenNet(4 * 1 + 1, pose_dataset, path=ALPHAPOSE_DIR)
    pose_model.cuda()
    pose_model.eval()

    runtime_profile = {"dt": [], "pt": [], "pn": []}

    # Data writer
    args.save_video = args.video_savefile is not None
    writer = DataWriter(
        save_video=args.
        save_video,  # Note: DataWriter uses args.save_video internally as well
        savepath=args.video_savefile,
        fourcc=cv2.VideoWriter_fourcc(*"XVID"),
        fps=fps,
        frameSize=frameSize,
        result_handler=pose_result_handler,
    ).start()

    im_names_desc = tqdm(range(data_loader.length()))
    batchSize = args.posebatch
    for i in im_names_desc:
        start_time = getTime()
        with torch.no_grad():
            (inps, orig_img, im_name, boxes, scores, pt1,
             pt2) = det_processor.read()
            if orig_img is None:
                break
            if boxes is None or boxes.nelement() == 0:
                writer.save(None, None, None, None, None, orig_img,
                            im_name.split("/")[-1])
                continue

            ckpt_time, det_time = getTime(start_time)
            runtime_profile["dt"].append(det_time)
            # Pose Estimation

            datalen = inps.size(0)
            leftover = 0
            if (datalen) % batchSize:
                leftover = 1
            num_batches = datalen // batchSize + leftover
            hm = []
            for j in range(num_batches):
                inps_j = inps[j * batchSize:min((j + 1) *
                                                batchSize, datalen)].cuda()
                hm_j = pose_model(inps_j)
                hm.append(hm_j)
            hm = torch.cat(hm)
            ckpt_time, pose_time = getTime(ckpt_time)
            runtime_profile["pt"].append(pose_time)

            hm = hm.cpu().data
            writer.save(boxes, scores, hm, pt1, pt2, orig_img,
                        im_name.split("/")[-1])

            ckpt_time, post_time = getTime(ckpt_time)
            runtime_profile["pn"].append(post_time)

        if args.profile:
            # TQDM
            im_names_desc.set_description(
                "det time: {dt:.3f} | pose time: {pt:.2f} | post processing: {pn:.4f}"
                .format(
                    dt=np.mean(runtime_profile["dt"]),
                    pt=np.mean(runtime_profile["pt"]),
                    pn=np.mean(runtime_profile["pn"]),
                ))

    print("===========================> Finish Model Running.")
    if (args.save_img or args.video_savefile) and not args.vis_fast:
        print(
            "===========================> Rendering remaining images in the queue..."
        )
        print(
            "===========================> If this step takes too long, you can enable "
            "the --vis_fast flag to use fast rendering (real-time).")
    while writer.running():
        pass
    writer.stop()
    return writer.results()
Esempio n. 9
0
        args.outputpath, ('%s' % datetime.today().strftime('%y%m%d_%H%M%S')))

    #init_detector
    # detector = GroundTruthDetections()

    #init_tracker
    # tracker = Sort(use_dlib=use_dlibTracker)  # create instance of the SORT tracker

    if not os.path.exists(args.outputpath):
        os.mkdir(args.outputpath)

    if not len(videofile):
        raise IOError('Error: must contain --video')

    # Load input video
    fvs = VideoLoader(videofile).start()
    (fourcc, fps, frameSize) = fvs.videoinfo()

    # Data writer
    save_path = os.path.join(
        args.outputpath,
        'AlphaPose_' + videofile.split('/')[-1].split('.')[0] + '.avi')
    writer = DataWriter(args.save_video, save_path,
                        cv2.VideoWriter_fourcc(*'XVID'), fps,
                        frameSize).start()

    # Load YOLO model
    print('Loading YOLO model..')
    det_model = Darknet("yolo/cfg/yolov3.cfg")
    det_model.load_weights('models/yolo/yolov3.weights')
    det_model.net_info['height'] = args.inp_dim
Esempio n. 10
0
def main(file_name):
    #  videofile = args.video
    videofile = file_name
    mode = args.mode
    if not os.path.exists(args.outputpath):
        os.mkdir(args.outputpath)

    if not len(videofile):
        raise IOError('Error: must contain --video')

    # Load input video
    data_loader = VideoLoader(videofile, batchSize=args.detbatch).start()
    (fourcc, fps, frameSize) = data_loader.videoinfo()

    # Load detection loader
    print('Loading YOLO model..')
    sys.stdout.flush()
    det_loader = DetectionLoader(data_loader, batchSize=args.detbatch).start()
    det_processor = DetectionProcessor(det_loader).start()

    # Load pose model
    pose_dataset = Mscoco()
    if args.fast_inference:
        pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
    else:
        pose_model = InferenNet(4 * 1 + 1, pose_dataset)
    pose_model.cuda()
    pose_model.eval()

    runtime_profile = {'dt': [], 'pt': [], 'pn': []}

    # Data writer
    save_path = os.path.join(
        args.outputpath,
        'AlphaPose_' + ntpath.basename(videofile).split('.')[0] + '.avi')
    writer = DataWriter(args.save_video, save_path,
                        cv2.VideoWriter_fourcc(*'XVID'), fps,
                        frameSize).start()

    im_names_desc = tqdm(range(data_loader.length()))
    batchSize = args.posebatch
    for i in im_names_desc:
        start_time = getTime()
        with torch.no_grad():
            (inps, orig_img, im_name, boxes, scores, pt1,
             pt2) = det_processor.read()
            if orig_img is None:
                break
            if boxes is None or boxes.nelement() == 0:
                writer.save(None, None, None, None, None, orig_img,
                            im_name.split('/')[-1])
                continue

            ckpt_time, det_time = getTime(start_time)
            runtime_profile['dt'].append(det_time)
            # Pose Estimation

            datalen = inps.size(0)
            leftover = 0
            if (datalen) % batchSize:
                leftover = 1
            num_batches = datalen // batchSize + leftover
            hm = []
            for j in range(num_batches):
                inps_j = inps[j * batchSize:min((j + 1) *
                                                batchSize, datalen)].cuda()
                hm_j = pose_model(inps_j)
                hm.append(hm_j)
            hm = torch.cat(hm)
            ckpt_time, pose_time = getTime(ckpt_time)
            runtime_profile['pt'].append(pose_time)

            hm = hm.cpu().data
            import ipdb
            ipdb.set_trace()
            writer.save(boxes, scores, hm, pt1, pt2, orig_img,
                        im_name.split('/')[-1])

            ckpt_time, post_time = getTime(ckpt_time)
            runtime_profile['pn'].append(post_time)

        if args.profile:
            # TQDM
            im_names_desc.set_description(
                'det time: {dt:.4f} | pose time: {pt:.4f} | post processing: {pn:.4f}'
                .format(dt=np.mean(runtime_profile['dt']),
                        pt=np.mean(runtime_profile['pt']),
                        pn=np.mean(runtime_profile['pn'])))

    print('===========================> Finish Model Running.')
    if (args.save_img or args.save_video) and not args.vis_fast:
        print(
            '===========================> Rendering remaining images in the queue...'
        )
        print(
            '===========================> If this step takes too long, you can enable the --vis_fast flag to use fast rendering (real-time).'
        )
    while (writer.running()):
        pass
    writer.stop()
    final_result = writer.results()
    write_json(final_result, args.outputpath)
Esempio n. 11
0
def handle_video(videofile, no_nan=True):
    args.video = videofile
    videofile = args.video
    mode = args.mode

    if not len(videofile):
        raise IOError('Error: must contain --video')

    # Load input video
    data_loader = VideoLoader(videofile, batchSize=args.detbatch).start()
    (fourcc, fps, frameSize) = data_loader.videoinfo()
    cam_w = frameSize[0]
    cam_h = frameSize[1]

    print('the video is {} f/s'.format(fps))

    # Load detection loader
    print('Loading YOLO model..')
    sys.stdout.flush()
    det_loader = DetectionLoader(data_loader, batchSize=args.detbatch).start()
    #  start a thread to read frames from the file video stream
    det_processor = DetectionProcessor(det_loader).start()

    # Load pose model
    pose_dataset = Mscoco()
    if args.fast_inference:
        pose_model = InferenNet_fast(4 * 1 + 1, pose_dataset)
    else:
        pose_model = InferenNet(4 * 1 + 1, pose_dataset)
    pose_model.cuda()
    pose_model.eval()

    runtime_profile = {'dt': [], 'pt': [], 'pn': []}

    # Data writer
    save_path = os.path.join(
        args.outputpath,
        'AlphaPose_' + ntpath.basename(videofile).split('.')[0] + '.avi')
    writer = DataWriter(args.save_video, save_path,
                        cv2.VideoWriter_fourcc(*'XVID'), fps,
                        frameSize).start()

    im_names_desc = tqdm(range(data_loader.length()))
    batchSize = args.posebatch
    frames_w_pose = []
    frame_cnt = 0
    for i in im_names_desc:
        start_time = getTime()
        with torch.no_grad():
            (inps, orig_img, im_name, boxes, scores, pt1,
             pt2) = det_processor.read()
            if orig_img is None:
                break

            frame_cnt += 1
            if boxes is None or boxes.nelement() == 0:
                writer.save(None, None, None, None, None, orig_img,
                            im_name.split('/')[-1])
                continue

            frames_w_pose.append(frame_cnt - 1)
            ckpt_time, det_time = getTime(start_time)
            runtime_profile['dt'].append(det_time)

            # Pose Estimation
            datalen = inps.size(0)
            leftover = 0
            if (datalen) % batchSize:
                leftover = 1
            num_batches = datalen // batchSize + leftover
            hm = []
            for j in range(num_batches):
                inps_j = inps[j * batchSize:min((j + 1) *
                                                batchSize, datalen)].cuda()
                hm_j = pose_model(inps_j)
                hm.append(hm_j)
            hm = torch.cat(hm)
            ckpt_time, pose_time = getTime(ckpt_time)
            runtime_profile['pt'].append(pose_time)

            hm = hm.cpu().data
            writer.save(boxes, scores, hm, pt1, pt2, orig_img,
                        im_name.split('/')[-1])

            ckpt_time, post_time = getTime(ckpt_time)
            runtime_profile['pn'].append(post_time)

        if args.profile:
            # TQDM
            im_names_desc.set_description(
                'det time: {dt:.4f} | pose time: {pt:.4f} | post processing: {pn:.4f}'
                .format(dt=np.mean(runtime_profile['dt']),
                        pt=np.mean(runtime_profile['pt']),
                        pn=np.mean(runtime_profile['pn'])))

    if (args.save_img or args.save_video) and not args.vis_fast:
        print(
            '===========================> Rendering remaining images in the queue...'
        )
        print(
            '===========================> If this step takes too long, you can enable the --vis_fast flag to use fast rendering (real-time).'
        )
    while (writer.running()):
        pass
    writer.stop()
    final_result = writer.results()

    kpts = []
    if not no_nan:
        for i in range(frame_cnt):
            # initialize to NaN so we can interpolate later
            kpts.append(np.full((17, 2), np.nan, dtype=np.float32))

    for i in range(len(final_result)):
        try:
            kpt = final_result[i]['result'][0]['keypoints']
            if not no_nan:
                kpts[frames_w_pose[i]] = kpt.data.numpy()
            else:
                kpts.append(kpt.data.numpy())
        except:
            print('error...')

    kpts = np.array(kpts).astype(np.float32)

    #filename = os.path.basename(args.video).split('.')[0]
    #name = filename + '.npz'
    #print('kpts npz save in ', name)
    #np.savez_compressed(name, kpts=kpts, fps=fps, cam_w=cam_w, cam_h=cam_h)

    return kpts, fps, cam_w, cam_h
Esempio n. 12
0
def main():
    save_path = './video_cap'
    video_name = 'webcam.avi'
    fourcc= cv2.VideoWriter_fourcc(*'XVID')
    fps = 24
    input("Press Any Key to Open WebCam (press 'q' to stop): ")
    cap = cv2.VideoCapture(0)
    frameSize = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)),
        int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    writer = cv2.VideoWriter(save_path+'/'+video_name, fourcc, fps, frameSize)
    for i in range(42):
        cap.read()
    start_write = False
    frame_cnt = 0
    while True:
        ret, frame = cap.read()
        cv2.imshow("capture", frame)
        if start_write:
            writer.write(frame)
            frame_cnt += 1

        if cv2.waitKey(1) & 0xFF == ord('s'):
            print('Start Rec!')
            start_write = True
        
        if cv2.waitKey(1) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()
    print("Start Processing!")
    outSize = (128,128)
    sp = False

    video_loader = VideoLoader(save_path+'/'+video_name, batchSize=32, wh_rate=(1,1), inDim=outSize, sp=sp, dataLen=frame_cnt)

    (fourcc, fps, frameSize, shape_dst, ow, oh) = video_loader.videoinfo()
    # fps = 10
    print(fourcc, fps, frameSize, shape_dst, ow, oh, video_loader.videoLen())
    num_frames = video_loader.videoLen()
    out_path = 'out_video/'+ video_name    
    fourcc=cv2.VideoWriter_fourcc(*'XVID')
    
    stream = cv2.VideoWriter(out_path, fourcc, fps, frameSize)
    assert stream.isOpened(), 'Cannot open video for writing'

    video_loader.start()
    save_loader = saveFrame(video_loader, sp=sp).start()
    ascend_loader = AscendHost(save_loader, sp=sp).start()
    collect_loader = DataCollector(ascend_loader, num_job=8, sp=sp).start()
    post_processor = PostProcessor(collect_loader, remap=(shape_dst, ow, oh), vis=True, nobg=True, sp=sp).start()
    # post_processor = PostProcessor(collect_loader, remap=(shape_dst, ow, oh), vis=True, sp=sp).start()

    cnt = 0
    while True:
        cnt += 1
        start = time.time()
        (img_render, img, orig_img, im_name) = post_processor.getitem()
        if orig_img is None:
            print('All done!')
            break
        
        orig_img = img_render

        cv2.imshow("Demo", orig_img)
        cv2.waitKey(160)
        stream.write(orig_img)
        interval = time.time() - start
        print("Time: {:.4f}\t{:04d}/{:04d}".format(interval, cnt, num_frames))
        # sys.stdout.flush()

    stream.release()