Вы можете использовать регулярное выражение для анализа сигнатур функций в строках, преобразования параметров в требуемый синтаксис, а затем и затем используйте re.sub
для добавления новых сигнатур:
import re
import itertools
def signature(header):
s = re.findall('(?<=\()[\w\W]+(?=\)$)', header)
return re.split(',\s(?=\w\w)', s[0]) if s else ''
def combine_args(d):
if '=' in d:
return '{}x{}'.format(*re.findall('\d+', d)) if 'kernel_size' in d and len(re.findall('\d+', d)) == 2 else '{},{}'.format(*re.findall('\d+', d)) if len(re.findall('\d+', d)) == 2 else d
return d
def combine_header(d):
vals = [[a, list(b)] for a, b in itertools.groupby(d, key=lambda x:x.isdigit())]
return list(itertools.chain(*[[' -> '.join(b)] if a else [combine_args(i) for i in b] for a, b in vals]))
lines = ['TVLoss()', 'Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())', 'Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)', 'Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())', 'Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)', 'Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())', 'Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)', 'Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())', 'Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'ContentLoss((crit): MSELoss())', 'Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)', 'Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())']
final_lines = [re.sub('(?<=\()[\w\W]+(?=\)$)', ', '.join(combine_header(c)), a) for a, c in zip(lines, map(signature, lines))]
Вывод:
['TVLoss()', 'Conv2d(3 -> 64, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())', 'Conv2d(64 -> 64, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)', 'Conv2d(64 -> 128, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())', 'Conv2d(128 -> 128, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)', 'Conv2d(128 -> 256, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())', 'Conv2d(256 -> 256, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'Conv2d(256 -> 256, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'Conv2d(256 -> 256, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)', 'Conv2d(256 -> 512, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())', 'Conv2d(512 -> 512, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'ContentLoss((crit): MSELoss())', 'Conv2d(512 -> 512, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'Conv2d(512 -> 512, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)', 'Conv2d(512 -> 512, 3x3, 1,1, 1,1)', 'ReLU(inplace)', 'StyleLoss((gram): GramMatrix()(crit): MSELoss())']