(visualise, cache_scoremaps)
| 33 | |
| 34 | |
| 35 | def test_net(visualise, cache_scoremaps): |
| 36 | logging.basicConfig(level=logging.INFO) |
| 37 | |
| 38 | cfg = load_config() |
| 39 | dataset = PoseDatasetFactory.create(cfg) |
| 40 | dataset.set_shuffle(False) |
| 41 | dataset.set_test_mode(True) |
| 42 | |
| 43 | sess, inputs, outputs = setup_pose_prediction(cfg) |
| 44 | |
| 45 | if cache_scoremaps: |
| 46 | out_dir = cfg["scoremap_dir"] |
| 47 | if not os.path.exists(out_dir): |
| 48 | os.makedirs(out_dir) |
| 49 | |
| 50 | num_images = dataset.num_images |
| 51 | predictions = np.zeros((num_images,), dtype=np.object) |
| 52 | |
| 53 | for k in range(num_images): |
| 54 | print(f"processing image {k}/{num_images - 1}") |
| 55 | |
| 56 | batch = dataset.next_batch() |
| 57 | |
| 58 | outputs_np = sess.run(outputs, feed_dict={inputs: batch[Batch.inputs]}) |
| 59 | |
| 60 | scmap, locref = extract_cnn_output(outputs_np, cfg) |
| 61 | |
| 62 | pose = argmax_pose_predict(scmap, locref, cfg["stride"]) |
| 63 | |
| 64 | pose_refscale = np.copy(pose) |
| 65 | pose_refscale[:, 0:2] /= cfg["global_scale"] |
| 66 | predictions[k] = pose_refscale |
| 67 | |
| 68 | if visualise: |
| 69 | img = np.squeeze(batch[Batch.inputs]).astype("uint8") |
| 70 | visualize.show_heatmaps(cfg, img, scmap, pose) |
| 71 | visualize.waitforbuttonpress() |
| 72 | |
| 73 | if cache_scoremaps: |
| 74 | base = os.path.basename(batch[Batch.data_item].im_path) |
| 75 | raw_name = os.path.splitext(base)[0] |
| 76 | out_fn = os.path.join(out_dir, raw_name + ".mat") |
| 77 | scipy.io.savemat(out_fn, mdict={"scoremaps": scmap.astype("float32")}) |
| 78 | |
| 79 | out_fn = os.path.join(out_dir, raw_name + "_locreg" + ".mat") |
| 80 | if cfg["location_refinement"]: |
| 81 | scipy.io.savemat(out_fn, mdict={"locreg_pred": locref.astype("float32")}) |
| 82 | |
| 83 | scipy.io.savemat("predictions.mat", mdict={"joints": predictions}) |
| 84 | |
| 85 | sess.close() |
| 86 | |
| 87 | |
| 88 | if __name__ == "__main__": |
no test coverage detected