def fix_proto_input_format(fromfile, tofile):
    """
    Utility to fix .prototxt files generated by pytorch-caffe-darknet-converter
    fromfile: source file to fix
    tofile  : destination fixed file (will be created if does not exist)

    Fix:
                          input_shape {
    input_dim: x              dim: x
    input_dim: y    --->      dim: y
    input_dim: z              dim: z
    input_dim: w              dim: w
                          }
    """
    with open(fromfile) as f:
        lines = f.readlines()
    dim = [e for e in lines if 'input_dim' in e]
    dim = [''.join([u for u in w if u.isdigit()]) for w in dim]
    dim = ['input_shape {\n']+['    dim: '+e+'\n' for e in dim]+['}\n']
    ind = list(range(len(lines)))
    cut = [i for i in ind if 'input_dim' in lines[i]]
    newlines = lines[:cut[0]]+dim+lines[cut[-1]+1:]
    with open(tofile, 'w') as f:
        f.writelines(newlines)

if __name__ == '__main__':
    import sys
    if len(sys.argv) != 3:
        print('USAGE: python fix_proto_input_format.py src.prototxt dst.prototxt')
        exit()
    fromfile = sys.argv[1]
    tofile = sys.argv[2]
    fix_proto_input_format(fromfile, tofile)