(file, prefix=colorstr('TensorFlow.js:'))
| 427 | |
| 428 | @try_export |
| 429 | def export_tfjs(file, prefix=colorstr('TensorFlow.js:')): |
| 430 | # YOLOv5 TensorFlow.js export |
| 431 | check_requirements('tensorflowjs') |
| 432 | import tensorflowjs as tfjs |
| 433 | |
| 434 | LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...') |
| 435 | f = str(file).replace('.pt', '_web_model') # js dir |
| 436 | f_pb = file.with_suffix('.pb') # *.pb path |
| 437 | f_json = f'{f}/model.json' # *.json path |
| 438 | |
| 439 | cmd = f'tensorflowjs_converter --input_format=tf_frozen_model ' \ |
| 440 | f'--output_node_names=Identity,Identity_1,Identity_2,Identity_3 {f_pb} {f}' |
| 441 | subprocess.run(cmd.split()) |
| 442 | |
| 443 | json = Path(f_json).read_text() |
| 444 | with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order |
| 445 | subst = re.sub( |
| 446 | r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, ' |
| 447 | r'"Identity.?.?": {"name": "Identity.?.?"}, ' |
| 448 | r'"Identity.?.?": {"name": "Identity.?.?"}, ' |
| 449 | r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, ' |
| 450 | r'"Identity_1": {"name": "Identity_1"}, ' |
| 451 | r'"Identity_2": {"name": "Identity_2"}, ' |
| 452 | r'"Identity_3": {"name": "Identity_3"}}}', json) |
| 453 | j.write(subst) |
| 454 | return f, None |
| 455 | |
| 456 | |
| 457 | def add_tflite_metadata(file, metadata, num_outputs): |
no test coverage detected