File size: 3,594 Bytes
1772f26 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
# Copyright 2022 Google LLC
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# https://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
r"""A test script for mid frame interpolation from two input frames.
Usage example:
python3 -m frame_interpolation.eval.interpolator_test \
--frame1 <filepath of the first frame> \
--frame2 <filepath of the second frame> \
--model_path <The filepath of the TF2 saved model to use>
The output is saved to <the directory of the input frames>/output_frame.png. If
`--output_frame` filepath is provided, it will be used instead.
"""
import os
from typing import Sequence
from . import interpolator as interpolator_lib
from . import util
from absl import app
from absl import flags
import numpy as np
# Controls TF_CCP log level.
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '1'
_FRAME1 = flags.DEFINE_string(
name='frame1',
default=None,
help='The filepath of the first input frame.',
required=True)
_FRAME2 = flags.DEFINE_string(
name='frame2',
default=None,
help='The filepath of the second input frame.',
required=True)
_MODEL_PATH = flags.DEFINE_string(
name='model_path',
default=None,
help='The path of the TF2 saved model to use.')
_OUTPUT_FRAME = flags.DEFINE_string(
name='output_frame',
default=None,
help='The output filepath of the interpolated mid-frame.')
_ALIGN = flags.DEFINE_integer(
name='align',
default=64,
help='If >1, pad the input size so it is evenly divisible by this value.')
_BLOCK_HEIGHT = flags.DEFINE_integer(
name='block_height',
default=1,
help='An int >= 1, number of patches along height, '
'patch_height = height//block_height, should be evenly divisible.')
_BLOCK_WIDTH = flags.DEFINE_integer(
name='block_width',
default=1,
help='An int >= 1, number of patches along width, '
'patch_width = width//block_width, should be evenly divisible.')
def _run_interpolator() -> None:
"""Writes interpolated mid frame from a given two input frame filepaths."""
interpolator = interpolator_lib.Interpolator(
model_path=_MODEL_PATH.value,
align=_ALIGN.value,
block_shape=[_BLOCK_HEIGHT.value, _BLOCK_WIDTH.value])
# First batched image.
image_1 = util.read_image(_FRAME1.value)
image_batch_1 = np.expand_dims(image_1, axis=0)
# Second batched image.
image_2 = util.read_image(_FRAME2.value)
image_batch_2 = np.expand_dims(image_2, axis=0)
# Batched time.
batch_dt = np.full(shape=(1,), fill_value=0.5, dtype=np.float32)
# Invoke the model for one mid-frame interpolation.
mid_frame = interpolator(image_batch_1, image_batch_2, batch_dt)[0]
# Write interpolated mid-frame.
mid_frame_filepath = _OUTPUT_FRAME.value
if not mid_frame_filepath:
mid_frame_filepath = f'{os.path.dirname(_FRAME1.value)}/output_frame.png'
util.write_image(mid_frame_filepath, mid_frame)
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
_run_interpolator()
if __name__ == '__main__':
app.run(main)
|