(exported_graph_path, model_structure, output_path)
| 32 | |
| 33 | |
| 34 | def convert_ssd(exported_graph_path, model_structure, output_path): |
| 35 | num_anchors = 1917 |
| 36 | |
| 37 | saved_model_path = os.path.join(exported_graph_path, 'saved_model') |
| 38 | coreml_model_path = os.path.join(output_path, 'Model.mlmodel') |
| 39 | |
| 40 | json_labels = os.path.join(exported_graph_path, 'labels.json') |
| 41 | with open(json_labels) as f: |
| 42 | labels = json.load(f) |
| 43 | |
| 44 | # Strip the model down to something usable by Core ML. |
| 45 | # Instead of `concat_1`, use `Postprocessor/convert_scores`, because it |
| 46 | # applies the sigmoid to the class scores. |
| 47 | frozen_model_path = '.tmp/tmp_frozen_graph.pb' |
| 48 | input_node = 'Preprocessor/sub' |
| 49 | bbox_output_node = 'concat' |
| 50 | class_output_node = 'Postprocessor/convert_scores' |
| 51 | graph = optimize_graph(saved_model_path, frozen_model_path, [input_node], [bbox_output_node, class_output_node]) |
| 52 | |
| 53 | # conversion tensors have a `:0` at the end of the name |
| 54 | input_tensor = input_node + ':0' |
| 55 | bbox_output_tensor = bbox_output_node + ':0' |
| 56 | class_output_tensor = class_output_node + ':0' |
| 57 | |
| 58 | # Convert to Core ML model. |
| 59 | ssd_model = tfcoreml.convert( |
| 60 | tf_model_path=frozen_model_path, |
| 61 | mlmodel_path=coreml_model_path, |
| 62 | input_name_shape_dict={ input_tensor: [1, 300, 300, 3] }, |
| 63 | image_input_names=input_tensor, |
| 64 | output_feature_names=[bbox_output_tensor, class_output_tensor], |
| 65 | is_bgr=False, |
| 66 | red_bias=-1.0, |
| 67 | green_bias=-1.0, |
| 68 | blue_bias=-1.0, |
| 69 | image_scale=2./255) |
| 70 | |
| 71 | spec = ssd_model.get_spec() |
| 72 | |
| 73 | # Rename the inputs and outputs to something more readable. |
| 74 | spec.description.input[0].name = 'image' |
| 75 | spec.description.input[0].shortDescription = 'Input image' |
| 76 | spec.description.output[0].name = 'scores' |
| 77 | spec.description.output[0].shortDescription = 'Predicted class scores for each bounding box' |
| 78 | spec.description.output[1].name = 'boxes' |
| 79 | spec.description.output[1].shortDescription = 'Predicted coordinates for each bounding box' |
| 80 | |
| 81 | input_mlmodel = input_tensor.replace(':', '__').replace('/', '__') |
| 82 | class_output_mlmodel = class_output_tensor.replace(':', '__').replace('/', '__') |
| 83 | bbox_output_mlmodel = bbox_output_tensor.replace(':', '__').replace('/', '__') |
| 84 | |
| 85 | for i in range(len(spec.neuralNetwork.layers)): |
| 86 | if spec.neuralNetwork.layers[i].input[0] == input_mlmodel: |
| 87 | spec.neuralNetwork.layers[i].input[0] = 'image' |
| 88 | if spec.neuralNetwork.layers[i].output[0] == class_output_mlmodel: |
| 89 | spec.neuralNetwork.layers[i].output[0] = 'scores' |
| 90 | if spec.neuralNetwork.layers[i].output[0] == bbox_output_mlmodel: |
| 91 | spec.neuralNetwork.layers[i].output[0] = 'boxes' |
no test coverage detected