1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950 |
- import paddle
- def get_input_spec(model_name, shape, trimap):
- """
- Get the input spec accoring the model_name.
- Args:
- model_name (str): The model name
- shape (str): The shape of input image
- trimap (str): Whether a trimap is required
- """
- input_spec = [{"img": paddle.static.InputSpec(shape=shape, name='img')}]
- if trimap:
- shape[1] = 1
- input_spec[0]['trimap'] = paddle.static.InputSpec(
- shape=shape, name='trimap')
- if model_name == 'RVM':
- input_spec.append(
- paddle.static.InputSpec(
- shape=[None, 16, None, None], name='r1'))
- input_spec.append(
- paddle.static.InputSpec(
- shape=[None, 20, None, None], name='r2'))
- input_spec.append(
- paddle.static.InputSpec(
- shape=[None, 40, None, None], name='r3'))
- input_spec.append(
- paddle.static.InputSpec(
- shape=[None, 64, None, None], name='r4'))
- input_spec.append(
- paddle.static.InputSpec(
- shape=[1], name='downsample_ratio'))
- return input_spec
|