Batchwise prediction of pose.
(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes, batchsize)
| 718 | |
| 719 | |
| 720 | def GetPoseF(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes, batchsize): |
| 721 | """Batchwise prediction of pose.""" |
| 722 | PredictedData = np.zeros((nframes, dlc_cfg["num_outputs"] * 3 * len(dlc_cfg["all_joints_names"]))) |
| 723 | batch_ind = 0 # keeps track of which image within a batch should be written to |
| 724 | batch_num = 0 # keeps track of which batch you are at |
| 725 | ny, nx = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)), int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) |
| 726 | if cfg["cropping"]: |
| 727 | ny, nx = checkcropping(cfg, cap) |
| 728 | |
| 729 | frames = np.empty((batchsize, ny, nx, 3), dtype="ubyte") # this keeps all frames in a batch |
| 730 | pbar = tqdm(total=nframes) |
| 731 | counter = 0 |
| 732 | step = max(10, int(nframes / 100)) |
| 733 | inds = [] |
| 734 | while cap.isOpened(): |
| 735 | if counter != 0 and counter % step == 0: |
| 736 | pbar.update(step) |
| 737 | ret, frame = cap.read() |
| 738 | if ret: |
| 739 | frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) |
| 740 | if cfg["cropping"]: |
| 741 | frames[batch_ind] = img_as_ubyte(frame[cfg["y1"] : cfg["y2"], cfg["x1"] : cfg["x2"]]) |
| 742 | else: |
| 743 | frames[batch_ind] = img_as_ubyte(frame) |
| 744 | inds.append(counter) |
| 745 | if batch_ind == batchsize - 1: |
| 746 | pose = predict.getposeNP(frames, dlc_cfg, sess, inputs, outputs) |
| 747 | PredictedData[inds] = pose |
| 748 | batch_ind = 0 |
| 749 | inds.clear() |
| 750 | batch_num += 1 |
| 751 | else: |
| 752 | batch_ind += 1 |
| 753 | elif counter >= nframes: |
| 754 | if batch_ind > 0: |
| 755 | pose = predict.getposeNP( |
| 756 | frames, dlc_cfg, sess, inputs, outputs |
| 757 | ) # process the whole batch (some frames might be from previous batch!) |
| 758 | PredictedData[inds[:batch_ind]] = pose[:batch_ind] |
| 759 | break |
| 760 | counter += 1 |
| 761 | |
| 762 | pbar.close() |
| 763 | return PredictedData, nframes |
| 764 | |
| 765 | |
| 766 | def GetPoseS(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes): |
no test coverage detected