Determine SSD input shape

5 years ago · c5a2d28367
parent b2464e3379
commit c5a2d28367
1 changed files with 6 additions and 0 deletions
--- a/samples/dnn/tf_text_graph_ssd.py
+++ b/samples/dnn/tf_text_graph_ssd.py
@ -234,6 +234,12 @@ def createSSDGraph(modelPath, configPath, outputPath):

    # Connect input node to the first layer
    assert(graph_def.node[0].op == 'Placeholder')
+    try:
+        input_shape = graph_def.node[0].attr['shape']['shape'][0]['dim']
+        input_shape[1]['size'] = image_height
+        input_shape[2]['size'] = image_width
+    except:
+        print("Input shapes are undefined")
    # assert(graph_def.node[1].op == 'Conv2D')
    weights = graph_def.node[1].input[-1]
    for i in range(len(graph_def.node[1].input)):