# Init model and load weights.
nets = prediction_helpers.LoadPredictionModels(
args.in_model_weights.split(','), net_settings, args.cuda_device_id)
# Initialize the base frame generator, either capturing frames live from a
# USB camera, or reading from a video file.
base_frame_generator = None
video_capture = None
if args.in_video_device_id is not None:
assert base_frame_generator is None
video_capture = cv2.VideoCapture(args.in_video_device_id)
base_frame_generator = image_helpers.VideoCaptureFrameGenerator(
video_capture)
elif args.in_video_file:
assert base_frame_generator is None
base_frame_generator = image_helpers.VideoFrameGenerator(
args.in_video_file)
# Make sure the frames arrive at roughly the frame rate we need.
# First add delays between captured frames (needed when reading from a video
# file to simulate the original camera frame rate).
# Then add frame skipping. This can be effective both for video files (to e.g.
# get a 10 FPS subsequence from a 30 FPS video) and live USB camera capture.
assert base_frame_generator is not None
frame_generator_delay = image_helpers.VideoFrameGeneratorLimitedFpsDelay(
base_frame_generator, args.delay_max_fps)
final_frame_generator = image_helpers.VideoFrameGeneratorLimitedFpsSkip(
frame_generator_delay, args.skip_max_fps)
trajectory_prediction = None
print('Live prediction started.')
for raw_frame, _ in final_frame_generator:
parser.add_argument('--target_width', type=int, default=-1)
args = parser.parse_args()
# Init model and load weights.
net = models.MakeNetwork(
args.net_name,
in_shape=[3, args.target_height, args.target_width],
out_dims=args.net_out_total_dimensions,
dropout_prob=0.0)
net.load_state_dict(torch.load(args.in_model_weights))
net.eval()
net.cuda()
result_data = []
frames_generator = image_helpers.VideoFrameGenerator(args.in_video)
for raw_frame, frame_index in frames_generator:
frame_cropped = image_helpers.CropHWC(raw_frame, args.crop_top,
args.crop_bottom, args.crop_left,
args.crop_right)
frame_resized = image_helpers.MaybeResizeHWC(frame_cropped,
args.target_height,
args.target_width)
frame_chw = np.transpose(frame_resized, (2, 0, 1))
frame_float = frame_chw.astype(np.float32) / 255.0
# Add a dummy dimension to make it a "batch" of size 1.
frame_tensor = Variable(torch.from_numpy(frame_float[np.newaxis,
...])).cuda()
result_tensor = net(frame_tensor).cpu()
result_value = result_tensor.data.numpy()[
0, args.net_out_dimension_to_use].item()