【Lab】提取result的bits和Y-PSNR数据并整理到Excel

【Lab】提取result的bits和Y-PSNR数据并整理到Excel

更新：使用openpyxl库直接将数据写入Excel中

注意：openpyxl是第三方库，如果没有安装。请命令行里键入pip install openpyxl

正则表达式提取数据

比如文件夹路径为C:\Users\Administrator\Desktop\mainbuilding33，其下有若干个文件夹，文件结构大概这样。

└─new
    ├─QP22
    │  └─result
    │          result12_11.txt
    │          result1_2.txt
    │          result3.txt
    │          result3_0.txt
    │          result3_12.txt
    │          result3_4.txt
    │          result4_22.txt
    │          result4_5.txt
    │          result5_1.txt
    │          result5_8.txt
    │
    ├─QP27
    │  └─result
    │          result12_11.txt
    │          result1_2.txt
    │          result3.txt
    │          result3_0.txt
    │          result3_12.txt
    │          result3_4.txt
    │          result4_22.txt
    │          result4_5.txt
    │          result5_1.txt
    │          result5_8.txt
    │
    └─QP32
        └─result
                result12_11.txt
                result1_2.txt
                result3.txt
                result3_0.txt
                result3_12.txt
                result3_4.txt
                result4_22.txt
                result4_5.txt
                result5_1.txt
                result5_8.txt

其中new为采用的方法，这样的文件夹在C:\Users\Administrator\Desktop\mainbuilding33下有多个，我只是给出了new。每个方法文件夹下面都有QP22、QP27、QP32三个文件夹，其下都有一个result文件夹，再往下就是结果集了，以txt格式保存的。

首先要将路径切换到要处理的*.txt目录下。这里用os.walk(path)函数可以递归遍历指定路径下的所有文件，再使用new_path = os.path.join(path, folder)拼合成绝对路径就行了。

先看下最终录入表格中的数据。这个文件就是使用该程序自动生成的。表格标题new和reference1就是方法文件夹，QP32、QP27、QP22是方法文件夹下的文件夹。通常来说，三个QP文件夹下的txt文件个数和文件名都相同（看上面的文件结构）。只是里面的的数据有所不同。比特率和psnr是我们需要的数据。注意到第一列的数字，它是txt文件名下划线右边的数字部分，按照从小到大的顺序录入数据。好了知道了要求，可以开始干活了。

• 按照Excel要求需要把下划线后的数字按从小到大的顺序排列。文件命名格式是这样的: \d+_\d+.txt,用正则表达式的方式\d+代表至少一个数字。按排序后的顺序依次打开文件。
• 按上述顺序打开文件后，提取某行的两个信息。其中有一个文件是\d+.txt的命名格式，为根节点文件，需要单独处理。每打开一个文件获取到数据后立刻存入到列表中，这样保证了提取到的数据也保持顺序。

这是\d+_\d+.txt类型的文件所需提取的部分。

下面是\d+.txt类型（根节点）文件里面要提取的部分。

这两个数据对应的模式为

found_rate = re.findall(r'POC\s+4.*?(\d+).bits', text)
found_psnr = re.findall(r'POC\s+4.*?Y\s+(.*?).dB', text)

re.findall返回列表，而提取到的数据只有一个，所以若是要取得列表里的数据，加上索引0

rate = found_rate[0]
psnr = found_psnr[0]

采用列表存这些数据是肯定的，因为这能保证存入的数据和文件打开的顺序一致。但是文件夹层次有点复杂，for循环中很容易导致列表被重新初始化。所以决定采用dict即字典的数据结构存储数据。用文件夹名给dict指定key，value里存数据。这个dict应该是全局的。大概就是这样的，order是表格的第一列数字，表示文件名下划线的右边的数字。psnr和rate也是列表。因为按顺序打开文件，故这些数据已经有序。

{'new_QP22': {'order': ['0', '2', '5', '7', '8', '9'],
              'psnr': ['40.7843',
                       '42.0684',
                       '40.6155',
                       '41.2164',
                       '41.5846',
                       '43.8196'],
              'rate': ['754592',
                       '1028656',
                       '817128',
                       '801176',
                       '852344',
                       '1019080']},
 'new_QP27': {'order': ['0', '2', '5', '7', '8', '9'],
              'psnr': ['36.2940',
                       '37.0931',
                       '36.0402',
                       '36.5335',
                       '36.8345',
                       '38.9266'],
              'rate': ['385512',
                       '575288',
                       '413672',
                       '448480',
                       '424440',
                       '610656']},
 'new_QP32': {'order': ['0', '2', '5', '7', '8', '9'],
              'psnr': ['34.8091',
                       '33.9552',
                       '34.5914',
                       '34.9151',
                       '34.7430',
                       '34.9727'],
              'rate': ['276680',
                       '319760',
                       '301488',
                       '330904',
                       '258832',
                       '328872']}
}

数据获取到了，只差写入Excel中了。

使用openpyxl将数据写到Excel

这个库的使用就不再这里介绍了，可以看openpyxl文档，也可以看我之前的一篇博客。由于全部数据我都获取到了，使用openpyxl的一个append方法可以将数据成行成行写入表格中。但是根据表格要求，每行的数据来自于不同的上面dict中的不同key，这意味着不做任何处理写入数据，得到的表格和要求的正好是行列对调了的（想象一下矩阵的转置）。

所以需要使用到Python内置的zip函数，将这个“矩阵”先转置再写入，就能得到上面Excel的样子。

全部写入后还做了一些额外的工作。

比如先将每个单元格的字体都变成了等线，为了使数据显示不那么紧凑，增加了每个单元格的列宽。同时还合并了一些单元格，加粗或者倾斜了某些单元格的字体。那么问题来了，当全部数据写入完毕的情况下，会有好几个Table。怎么知道那些应该合并的数据位于哪一行呢？

以new这个表格标题的单元格合并为例（毕竟找到这行，往下数1到2行就能定位QP和比特率/psnr那些行了）。首先第一行必然就是第一个表格标题的行号。怎么确定下面几个表格标题的行号呢？

这是有规律可循的，一个表格是由标题、下面的一行比特率/psnr、再下一行QP、程序里加上的两行空行共5行。再加上方法文件加下任意一个QP文件夹（上面有说到里面的文件个数相同）下的txt文件的个数。假如new方法下任一QP文件夹里有6个txt文件，因为第1行是第一个表格的标题行。这样可以推算处下一个表格的标题行号为

1 + 5 + 6 = 12

没错真的是第12行。假如所有方法文件夹下的文件个数用列表表示为[6， 8， 8]分别加上5得到[11, 13, 13]，再在索引为0的地方插入1，表示第一个表格的标题行号。得到[1, 11, 13, 13]，最后将这里面的数两两相加。即a[0] + a[1] 放入a[1]中（不要动a[0]=1 ,第一行是第一个表格的标题这是肯定的），然后将新的a[1]与a[2]相加放入a[2]中...用代码表示如下。为了防止脚标越界，故下面的代码中有减去1。

for i in range(len(step_list) - 1):
    step_list[i + 1] = step_list[i] + step_list[i + 1]

这样之后列表中的数据为[1, 12, 25, 38]，很好，就是这些位置。

代码中的step_list就是存了这些信息。

最后，打开那么多文件，有可能某个文件是空的，没有数据，这就导致re.findall()返回[]空列表，不加判断取脚标[0]会发生脚标越界错误。所以程序中假如了非空判断，当某个文件数据缺失时，提醒用户去异常信息中给出的路径查看该问题文件是否正常。

found_rate = re.findall(r'POC\s+0.*?(\d+).bits', text)
found_psnr = re.findall(r'POC\s+0.*?Y\s+(.*?).dB', text)
# 如果没有数据，抛出异常
if not found_rate or not found_psnr:
    raise Exception(f'数据缺失！请前往{os.path.join(current_path, filename)}查看')

还有一点，代码中某些列表中会看到很多None，表示写入excel中时，这个单元格不写数据。但还是必须的，为了整行写入，作为占位存在。

好了，上代码

import os
import re
import pprint
from openpyxl import Workbook
from openpyxl.styles import Font, Alignment
from openpyxl.utils import get_column_letter

# TODO 1.改成自己的路径，r表示原始字符串，保留
original_path = r'C:\Users\Administrator\Desktop\mainbuilding33'

def find_data(path):
    """
    找到每个result文本里面的rate和psnr_Y数据，存到一个字典中
    :param path: 待处理数据的文件夹路径
    :return: data_dict 数据存放在该字典
    """
    data_dict = {}
    for current_path, subfolders, filesname in os.walk(path):
        # 不是空列表才执行以下
        if filesname:
            # 找出所有根节点文件
            root_list = re.findall(r'result\d+.txt', str(filesname))
            # 排序，key参数指定对'[\d].txt'里靠近txt的数字排序.这里file_name为列表的每个子项
            # 由于字符串的排序方法会先比较第一个数字，再比较下一个，所以会导致'10'比'2'小，需转型为int来比较
            sorted_file_list = sorted(filesname, key=lambda file_name: int(re.findall(r'(\d+).txt', file_name)[0]))
            # excel中A列的数据
            order_data = [i for i in re.findall(r'(\d+).txt', str(sorted_file_list))]
            # 方法名
            method = current_path.split('\\')[-3]
            # QP名
            qp = current_path.split('\\')[-2]
            # 用来存放每个文件中提取的psnr和rate
            rate_list, psnr_list = [], []
            for filename in sorted_file_list:
                with open(os.path.join(current_path, filename)) as f:
                    # 读取全部内容
                    text = f.read()
                    # 特殊处理根节点文件
                    if filename in root_list:
                        # (\d+)为捕获组，只会返回括号里面的内容
                        found_rate = re.findall(r'POC\s+0.*?(\d+).bits', text)
                        found_psnr = re.findall(r'POC\s+0.*?Y\s+(.*?).dB', text)
                        # 如果没有数据，抛出异常
                        if not found_rate or not found_psnr:
                            raise Exception(f'数据缺失！请前往{os.path.join(current_path, filename)}查看')

                        # 如果抛出异常直接终止程序，if执行了这下面的肯定执行不了，故无需else语句
                        rate = found_rate[0]
                        psnr = found_psnr[0]
                        rate_list.append(rate)
                        psnr_list.append(psnr)
                    # 其他文件
                    else:
                        found_rate = re.findall(r'POC\s+4.*?(\d+).bits', text)
                        found_psnr = re.findall(r'POC\s+4.*?Y\s+(.*?).dB', text)
                        if not found_rate or not found_psnr:
                            raise Exception(f'数据缺失！请前往{os.path.join(current_path, filename)}查看')

                        rate = found_rate[0]
                        psnr = found_psnr[0]
                        rate_list.append(rate)
                        psnr_list.append(psnr)

            # 所有数据收集完毕，存入字典并返回
            data_dict[method + '_' + qp] = {'rate': rate_list, 'psnr': psnr_list, 'order': order_data}

    return data_dict

def save_to_excel(filename=None):
    """
    将提取到的rate和psnr数据写入excel中
    可以自定义保存的路径以及文件名,如C:\\Users\\Administrator\\Desktop\\abc.xlsx'
    :param filename:目录及文件名
    """
    basename = os.path.basename(original_path)
    # 新建一个工作表
    wb = Workbook()
    # 获取当前活动的工作表
    sheet = wb.active
    # 修改工作表名称
    sheet.title = basename
    data = find_data(original_path)

    # 用来定位表格标题的行号
    step_list = []
    method = None
    for key in data.keys():
        # 遇到方法相同的情况跳过，保证只写入一次, 此为表格标题的那行
        if key.split('_')[0] != method:
            # 写入一次后，值更新
            step_list.append(len(data[key]['order']) + 5)
            method = key.split('_')[0]
            # 换成32在前的顺序
            qp_list = sorted(set([qp.split('_')[1] for qp in data.keys()]), reverse=True)
            # 一行共两个QP32、QP27、QP22
            qp_list.extend(qp_list)
            qp_list.insert(0, None)
            data_list = [[None, method],
                         [None, '比特率', None, None, 'psnr'],
                         qp_list]
            k32 = method + '_' + qp_list[1]
            k27 = method + '_' + qp_list[2]
            k22 = method + '_' + qp_list[3]
            rows = zip(data[k32]['order'], data[k32]['rate'], data[k27]['rate'], data[k22]['rate'],
                       data[k32]['psnr'], data[k27]['psnr'], data[k22]['psnr'])

            for row in rows:
                # 数据是文本格式，转为数字写入
                data_list.append(list(map(float, row)))
            # 增加行两空行
            data_list.append([])
            data_list.append([])
            # 将每一行写入excel
            for row in data_list:
                sheet.append(row)

    # 全局单元格设置
    # 每个单元格的字体设置
    for row in sheet.rows:
        for cell in row:
            cell.font = Font(name='等线')

    # 为了使数据不那么紧凑，增加每列的列宽。列要用到字母，get_column_letter(i)转换
    # 都是从1开始数的，注意range
    for i in range(1, sheet.max_column + 1):
        sheet.column_dimensions[get_column_letter(i)].width = 14

    # 针对需要合并的单元格
    step_list.insert(0, 1)
    for i in range(len(step_list) - 1):
        step_list[i + 1] = step_list[i] + step_list[i + 1]

    # step_list最后一个数值是将要写入的最新一行的行号。暂时用不到，所以减1
    for row_num in step_list[:-1]:
        # 1. 设置表格大标题
        start_index = 'B' + str(row_num)
        end_index = get_column_letter(sheet.max_column) + str(row_num)
        sheet[start_index].font = Font(name='等线', size=24, italic=True, bold=True)
        sheet.merge_cells(start_index + ':' + end_index)
        sheet.row_dimensions[row_num].height = 36
        sheet[start_index].alignment = Alignment(horizontal='center', vertical='center')
        # 合并rate，合并psnr
        rate_start_index = 'B' + str(row_num + 1)
        rate_end_index = get_column_letter((sheet.max_column - 1) / 2 + 1) + str(row_num + 1)
        psnr_start_index = get_column_letter((sheet.max_column - 1) / 2 + 2) + str(row_num + 1)
        psnr_end_index = get_column_letter(sheet.max_column) + str(row_num + 1)
        sheet.row_dimensions[row_num + 1].height = 20
        sheet.merge_cells(rate_start_index + ':' + rate_end_index)
        sheet.merge_cells(psnr_start_index + ':' + psnr_end_index)
        sheet[rate_start_index].font = Font(name='等线', bold=True)
        sheet[psnr_start_index].font = Font(name='等线', bold=True)
        sheet[rate_start_index].alignment = Alignment(horizontal='center', vertical='center')
        sheet[psnr_start_index].alignment = Alignment(horizontal='center', vertical='center')
        # 3. QP那行.
        # sheet.rows和sheet.columns是生成器，没有下标，先转为list
        for qp_cell in list(sheet.rows)[row_num + 1]:
            qp_cell.alignment = Alignment(horizontal='right', vertical='center')

    # 用户可以自定义文件名和存放的目录，若是没有填入参数，采用如下默认方案
    if filename is None:
        # 默认当前工作目录和默认的文件名
        filename = basename + '.xlsx'

    wb.save(filename)

if __name__ == '__main__':
    # 2.可以一个填入参数，自定义保存的路径以及文件名
    # 如'C:\\Users\\Administrator\\Desktop\\abc.xlsx'
    save_to_excel()

运行结果就是直接生成Excel，一键运行啥都不用做了。不过最好还是打开稍微检查下，有错误了交到Boss那儿后果不太好。

旧的处理方法以及代码

这时以前写的，只是提取到了txt中，还是得手动赋值整列数据到excel中。上面的可以说最终版，十分方便。这个还是麻烦了。但是本着记录自己的学习过程，还是保留在博客中。

方法很简单，每打开一个文件提获取到数据后，立刻写入输出文件中，这样也保证了提取到的数据也保持顺序。

正则表达式提取即可。

import os
import re

# 1. 改成自己的路径，r表示原始字符串，保留
original_path = r'C:\Users\Administrator\Desktop\mainbuilding33'

def find_save(file_path, filename, root_list, out):
    with open(file_path) as f:
        # 读取全部内容
        text = f.read()
        # 特殊处理根节点文件
        if filename in root_list:
            # (\d+)为捕获组，只会返回括号里面的内容
            rate = re.findall(r'POC\s+0.*?(\d+).bits', text)[0]
            psnrY = re.findall(r'POC\s+0.*?Y\s+(.*?).dB', text)[0]
            out.write(rate + '\t\t' + psnrY + '\n')
        # 其他文件
        else:
            rate = re.findall(r'POC\s+4.*?(\d+).bits', text)[0]
            psnrY = re.findall(r'POC\s+4.*?Y\s+(.*?).dB', text)[0]
            out.write(rate + '\t\t' + psnrY + '\n')

def run(path):
    filename = path.split('\\')[-1] + '.txt'
    # 输出文件
    with open(filename, 'w', encoding='utf-8') as f:
        for current_path, subfolders, filesname in os.walk(path):
            # 不是空列表才执行以下
            if filesname:
                # 排序，key参数指定对'[\d].txt'里靠近txt的数字排序.这里file_name为列表的每个子项
                # 由于字符串的排序方法会先比较第一个数字，再比较下一个，所以会导致'10'比'2'小，需转型为int来比较
                sorted_file_list = sorted(filesname, key=lambda file_name: int(re.findall(r'(\d+).txt', file_name)[0]))
                special_list = re.findall(r'result\d+.txt', str(filesname))

                method_name = current_path.split('\\')[-3]
                qp_name = current_path.split('\\')[-2]
                f.write('\n')
                f.write('*' * 6 + method_name + ' ' + qp_name + '*' * 6 + '\n')
                f.write('Rate' + '\t\t' + 'PSNR' + '\n\n')
                for file in sorted_file_list:
                    find_save(os.path.join(current_path, file), file, special_list, f)

if __name__ == '__main__':
    run(original_path)

结果如下：

很好，两列数据直接复制到Excel，OK！

---

by @sunhaiyu

2017.6.20