cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Internal change

Browse Source 
PiperOrigin-RevId: 482259130
This commit is contained in:
MediaPipe Team 2022-10-19 11:32:42 -07:00 committed by Copybara-Service
parent c260074abb
commit 3d588bae8b
3 changed files with 30 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mediapipe/model_maker/python/vision/image_classifier/image_classifier_test.py
View File

@ -66,8 +66,13 @@ class ImageClassifierTest(tf.test.TestCase, parameterized.TestCase):
hparams=image_classifier.HParams( hparams=image_classifier.HParams(
train_epochs=1, batch_size=1, shuffle=True)), train_epochs=1, batch_size=1, shuffle=True)),
dict( dict(
testcase_name='efficientnet_lite1', testcase_name='efficientnet_lite2',
model_spec=image_classifier.SupportedModels.EFFICIENTNET_LITE1, model_spec=image_classifier.SupportedModels.EFFICIENTNET_LITE2,
hparams=image_classifier.HParams(
train_epochs=1, batch_size=1, shuffle=True)),
dict(
testcase_name='efficientnet_lite4',
model_spec=image_classifier.SupportedModels.EFFICIENTNET_LITE4,
hparams=image_classifier.HParams( hparams=image_classifier.HParams(
train_epochs=1, batch_size=1, shuffle=True)), train_epochs=1, batch_size=1, shuffle=True)),
) )

17
mediapipe/model_maker/python/vision/image_classifier/model_spec.py
View File

@ -53,11 +53,17 @@ efficientnet_lite0_spec = functools.partial(
uri='https://tfhub.dev/tensorflow/efficientnet/lite0/feature-vector/2', uri='https://tfhub.dev/tensorflow/efficientnet/lite0/feature-vector/2',
name='efficientnet_lite0') name='efficientnet_lite0')
efficientnet_lite1_spec = functools.partial( efficientnet_lite2_spec = functools.partial(
ModelSpec, ModelSpec,
uri='https://tfhub.dev/tensorflow/efficientnet/lite1/feature-vector/2', uri='https://tfhub.dev/tensorflow/efficientnet/lite2/feature-vector/2',
input_image_shape=[240, 240], input_image_shape=[260, 260],
name='efficientnet_lite1') name='efficientnet_lite2')
efficientnet_lite4_spec = functools.partial(
ModelSpec,
uri='https://tfhub.dev/tensorflow/efficientnet/lite4/feature-vector/2',
input_image_shape=[300, 300],
name='efficientnet_lite4')
# TODO: Document the exposed models. # TODO: Document the exposed models.
@ -66,7 +72,8 @@ class SupportedModels(enum.Enum):
"""Image classifier model supported by model maker.""" """Image classifier model supported by model maker."""
MOBILENET_V2 = mobilenet_v2_spec MOBILENET_V2 = mobilenet_v2_spec
EFFICIENTNET_LITE0 = efficientnet_lite0_spec EFFICIENTNET_LITE0 = efficientnet_lite0_spec
EFFICIENTNET_LITE1 = efficientnet_lite1_spec EFFICIENTNET_LITE2 = efficientnet_lite2_spec
EFFICIENTNET_LITE4 = efficientnet_lite4_spec
@classmethod @classmethod
def get(cls, spec: 'SupportedModels') -> 'ModelSpec': def get(cls, spec: 'SupportedModels') -> 'ModelSpec':

16
mediapipe/model_maker/python/vision/image_classifier/model_spec_test.py
View File

@ -37,11 +37,17 @@ class ModelSpecTest(tf.test.TestCase, parameterized.TestCase):
expected_name='efficientnet_lite0', expected_name='efficientnet_lite0',
expected_input_image_shape=[224, 224]), expected_input_image_shape=[224, 224]),
dict( dict(
testcase_name='efficientnet_lite1_spec_test', testcase_name='efficientnet_lite2_spec_test',
model_spec=ms.efficientnet_lite1_spec, model_spec=ms.efficientnet_lite2_spec,
expected_uri='https://tfhub.dev/tensorflow/efficientnet/lite1/feature-vector/2', expected_uri='https://tfhub.dev/tensorflow/efficientnet/lite2/feature-vector/2',
expected_name='efficientnet_lite1', expected_name='efficientnet_lite2',
expected_input_image_shape=[240, 240]), expected_input_image_shape=[260, 260]),
dict(
testcase_name='efficientnet_lite4_spec_test',
model_spec=ms.efficientnet_lite4_spec,
expected_uri='https://tfhub.dev/tensorflow/efficientnet/lite4/feature-vector/2',
expected_name='efficientnet_lite4',
expected_input_image_shape=[300, 300]),
) )
def test_predefiend_spec(self, model_spec: Callable[..., ms.ModelSpec], def test_predefiend_spec(self, model_spec: Callable[..., ms.ModelSpec],
expected_uri: str, expected_name: str, expected_uri: str, expected_name: str,