如何让SynthText数据集生成的掩码与原图像同名？

问题背景 #

现有基于MATLAB标注文件gt.mat生成文本掩码的脚本可正常运行，但需要让生成的掩码文件名与原RGB图像完全一致——原图像文件名存储在gt.mat的imnames字段中。原图像分布在/home/SynthText/下的200个子文件夹（每个约4000张图），当前脚本将所有掩码输出至单个文件夹，不清楚如何在cv2.imwrite()中调用gt.mat['imnames']来指定掩码文件名。

核心解决方案 #

直接从gt.mat的imnames数组中提取对应样本的原始文件名，以此作为掩码的输出文件名，无需关心原图像所在的子文件夹（只需保证文件名一致，输出到统一掩码目录即可）。

关键步骤 #

1. 正确读取imnames字段：gt.mat中的imnames通常是MATLAB格式的字符串数组，读取后需转换为Python可直接使用的字符串类型（使用scipy.io.loadmat时添加squeeze_me=True可自动压缩维度，避免嵌套数组）。
2. 提取纯文件名：若imnames中存储的是完整路径（如/home/SynthText/sub_dir/abc.jpg），需用os.path.basename()提取仅含文件名的部分，避免把原路径结构带入掩码输出目录。
3. 循环匹配生成并保存：在掩码生成循环中，为每个样本匹配对应的文件名，拼接输出路径后调用cv2.imwrite()保存。

修改后的代码示例 #

假设原有SynthTextDataset类结构如下，重点修改掩码保存的循环逻辑：

import scipy.io as sio
import cv2
import numpy as np
import os

class SynthTextDataset:
    def __init__(self, mat_path):
        # 加载mat文件，squeeze_me=True处理MATLAB字符串数组的维度问题
        self.mat_data = sio.loadmat(mat_path, squeeze_me=True)
        self.imnames = self.mat_data['imnames']  # 获取所有图像文件名（可能带路径）
        self.word_bboxes = self.mat_data['wordBB']  # 文本框标注
        # 其他初始化逻辑（如获取图像尺寸等）...

    def generate_text_mask(self, idx):
        # 原有掩码生成逻辑：根据文本框绘制掩码
        # 示例：从原图像路径读取尺寸，也可从mat的其他字段获取
        raw_img_path = os.path.join("/home/SynthText", self.imnames[idx])
        img = cv2.imread(raw_img_path)
        mask = np.zeros(img.shape[:2], dtype=np.uint8)
        
        # 绘制文本区域为白色（255）
        bboxes = self.word_bboxes[idx]
        # 处理SynthText的wordBB格式：可能是3x2xN或2xN，需根据实际情况调整
        if len(bboxes.shape) == 3:
            bboxes = bboxes.transpose(2, 0, 1)
        for bbox in bboxes:
            pts = bbox.astype(np.int32).reshape((-1, 1, 2))
            cv2.fillPoly(mask, [pts], 255)
        
        return mask

# 主执行逻辑
mask_output_dir = "./synthtext_masks"
os.makedirs(mask_output_dir, exist_ok=True)

# 加载数据集
dataset = SynthTextDataset("/path/to/gt.mat")

# 循环生成并保存掩码
for idx in range(len(dataset.imnames)):
    # 提取纯文件名（去掉路径部分）
    raw_filename = os.path.basename(dataset.imnames[idx])
    # 生成掩码
    text_mask = dataset.generate_text_mask(idx)
    # 拼接保存路径
    save_path = os.path.join(mask_output_dir, raw_filename)
    # 若需要PNG无损格式，可替换后缀：raw_filename = os.path.splitext(raw_filename)[0] + ".png"
    cv2.imwrite(save_path, text_mask)

注意事项 #

• 若imnames中的条目不带路径仅为纯文件名，可直接使用，无需os.path.basename()处理。
• 确保imnames的索引与掩码生成的索引严格对应，避免出现文件名与掩码不匹配的情况。

内容的提问来源于stack exchange，提问作者Leon