|
import argparse |
|
import os |
|
import time |
|
|
|
os.sys.path += ['expman'] |
|
import expman |
|
|
|
import numpy as np |
|
import pandas as pd |
|
import tensorflow as tf |
|
from tqdm import trange |
|
|
|
|
|
def main(args): |
|
is_run_dir = expman.is_exp_dir(args.model) |
|
if is_run_dir: |
|
exp = expman.from_dir(args.model) |
|
for model_file in ('best_savedmodel', 'best_model.h5', 'last_model.h5'): |
|
model_path = exp.path_to(model_file) |
|
if os.path.exists(model_path): |
|
break |
|
elif tf.saved_model.contains_saved_model(args.model): |
|
model_path = args.model |
|
else: |
|
print('Cannot find suitable model snapshot in {}'.format(args.model)) |
|
exit(1) |
|
|
|
model = tf.keras.models.load_model(model_path, compile=False, custom_objects={'tf': tf}) |
|
data = np.empty((1, args.rh, args.rw, 1), dtype=np.float32) |
|
|
|
|
|
model.predict(data) |
|
|
|
start = time.time() |
|
for _ in trange(args.n): |
|
model.predict(data) |
|
end = time.time() |
|
elapsed = end - start |
|
|
|
throughput = elapsed / args.n |
|
fps = args.n / elapsed |
|
print(f'Total: {elapsed:g}s ({throughput * 1000} ms/img, {fps} fps)') |
|
|
|
timings = pd.Series({'elapsed': elapsed, 'throughput': throughput, 'fps': fps}) |
|
|
|
if is_run_dir and not args.output: |
|
timings_path = exp.path_to('timings.csv') |
|
timings.to_csv(timings_path) |
|
|
|
if args.output: |
|
timings.to_csv(args.output) |
|
|
|
|
|
if __name__ == '__main__': |
|
parser = argparse.ArgumentParser(description='Predict on test video') |
|
parser.add_argument('model', help='path to model or run dir') |
|
parser.add_argument('-n', type=int, default=100, help='number of predictions') |
|
parser.add_argument('-rh', type=int, default=128, help='RoI height (-1 for full height)') |
|
parser.add_argument('-rw', type=int, default=128, help='RoI width (-1 for full width)') |
|
parser.add_argument('-o', '--output', help='CSV output file') |
|
|
|
args = parser.parse_args() |
|
main(args) |
|
|
