AIStoryBoard/nodejs/sort-dialog.js

/**
 * 步骤：
 * 1. 创建startSortDialog()函数
 * 2. 创建ocrTxtResultArr用来接收外部传入的识别结果txt文件路径数组
 * 3. 根据check文件夹里面的txt文件先排序，文件名的index越小越靠数组前位，为每一个check文件夹创建一个数组，把文件夹内文件名排序结果保存到数组里。
 * 4. 再新建一个外围数组，根据check名字的index顺序，把每一个check文件夹内的数组保存到外围数组里。index越小越靠前。
 * 5. 返回最终的二维数组
 */

import fs from 'fs';
import path from 'path';

function startSortDialog(ocrTxtResultArr) {
  try {
    console.log('🚀 开始整理对话排序流程...');
    
    // 步骤2: 创建变量ocrTxtResultArr 用来接收外部传入的识别结果txt文件路径数组
    console.log('\n📋 步骤2: 验证txt文件路径数组参数');
    if (!ocrTxtResultArr || !Array.isArray(ocrTxtResultArr)) {
      throw new Error('步骤2失败: ocrTxtResultArr 必须是一个数组');
    }
    if (ocrTxtResultArr.length === 0) {
      throw new Error('步骤2失败: ocrTxtResultArr 不能为空数组');
    }
    console.log(`✅ txt文件路径数组长度: ${ocrTxtResultArr.length}`);
    
    // 步骤3: 根据check文件夹里面的txt文件先排序
    console.log('\n📁 步骤3: 按check文件夹分组并排序txt文件...');
    
    // 按check文件夹分组
    const checkGroups = {}; // { checkIndex: [txtPaths] }
    
    for (const txtPath of ocrTxtResultArr) {
      // 获取txt文件所在的文件夹路径
      const txtDir = path.dirname(txtPath);
      const txtFileName = path.basename(txtPath);
      
      // 提取check文件夹名称（如 check1, check2 等）
      const checkFolderName = path.basename(txtDir);
      const checkMatch = checkFolderName.match(/^check(\d+)$/);
      
      if (checkMatch) {
        const checkIndex = parseInt(checkMatch[1]);
        
        if (!checkGroups[checkIndex]) {
          checkGroups[checkIndex] = [];
        }
        
        checkGroups[checkIndex].push(txtPath);
      } else {
        console.log(`   ⚠️ 跳过非check文件夹: ${checkFolderName}`);
      }
    }
    
    console.log(`✅ 找到 ${Object.keys(checkGroups).length} 个check文件夹`);
    
    // 为每一个check文件夹创建一个数组，把文件夹内文件名排序结果保存到数组里
    const sortedCheckGroups = {};
    
    for (const [checkIndexStr, txtPaths] of Object.entries(checkGroups)) {
      const checkIndex = parseInt(checkIndexStr);
      
      // 按文件名中的index排序（文件名的index越小越靠数组前位）
      const sortedPaths = txtPaths.sort((a, b) => {
        const fileNameA = path.basename(a, path.extname(a));
        const fileNameB = path.basename(b, path.extname(b));
        
        // 提取文件名中的数字（如 dialog_1.txt -> 1）
        const matchA = fileNameA.match(/(\d+)$/);
        const matchB = fileNameB.match(/(\d+)$/);
        
        const indexA = matchA ? parseInt(matchA[1]) : 0;
        const indexB = matchB ? parseInt(matchB[1]) : 0;
        
        return indexA - indexB; // index越小越靠前
      });
      
      sortedCheckGroups[checkIndex] = sortedPaths;
      
      console.log(`   📁 check${checkIndex}: ${sortedPaths.length} 个txt文件`);
      sortedPaths.forEach((p, i) => {
        console.log(`      ${i + 1}. ${path.basename(p)}`);
      });
    }
    
    // 步骤4: 再新建一个外围数组，根据check名字的index顺序，把每一个check文件夹内的数组保存到外围数组里
    console.log('\n📦 步骤4: 按check文件夹index顺序组织二维数组...');
    
    // 获取所有check索引并排序
    const checkIndices = Object.keys(sortedCheckGroups).map(Number).sort((a, b) => a - b);
    
    // 创建外围数组（二维数组）
    const resultArray = [];
    
    for (const checkIndex of checkIndices) {
      resultArray.push(sortedCheckGroups[checkIndex]);
      console.log(`   ✅ check${checkIndex}: 已添加到外围数组 (${sortedCheckGroups[checkIndex].length} 个文件)`);
    }
    
    // 步骤5: 返回最终的二维数组
    console.log('\n📋 步骤5: 准备返回最终的二维数组...');
    console.log(`📊 二维数组结构:`);
    console.log(`   外层数组长度: ${resultArray.length} (${checkIndices.length} 个check文件夹)`);
    resultArray.forEach((checkArray, index) => {
      console.log(`   [${index}] check${checkIndices[index]}: ${checkArray.length} 个txt文件`);
    });
    console.log(`✅ 步骤5完成: 二维数组已准备就绪`);
    
    console.log('\n🎉 所有步骤完成！');
    console.log(`📊 共处理 ${ocrTxtResultArr.length} 个txt文件`);
    console.log(`📊 共组织成 ${resultArray.length} 个check文件夹数组`);
    
    return resultArray;  // 步骤5: 返回最终的二维数组
    
  } catch (error) {
    console.error(`\n❌ 整理对话排序失败: ${error.message}`);
    throw error;
  }
}

/**
 * 旧的函数：根据panels排序dialogues（保留用于向后兼容）
 * @param {Array} dialogues - 对话列表（外部传入）
 * @param {Array} panels - 格子列表（外部传入）
 * @param {number} imageWidth - 图片宽度（外部传入）
 * @param {number} imageHeight - 图片高度（外部传入）
 * @param {string} imageName - 图片名称（用于生成txt文件名）
 * @param {string} outputDir - 输出目录（外部传入）
 * @returns {Object} 包含排序后的对话列表和生成的文件信息
 */
async function startSortDialogByPanels(dialogues, panels, imageWidth, imageHeight, imageName, outputDir) {
  try {
    // 步骤2: 创建变量dialogues 用来接收外部传入的对话列表
    if (!dialogues || !Array.isArray(dialogues)) {
      throw new Error('dialogues 参数必须是一个数组');
    }
    
    // 步骤3: 创建变量panels 用来接收外部传入的格子列表
    if (!panels || !Array.isArray(panels)) {
      throw new Error('panels 参数必须是一个数组');
    }
    
    // 步骤4: 创建变量imageWidth 用来接收外部传入的图片宽度
    if (!imageWidth || typeof imageWidth !== 'number') {
      throw new Error('imageWidth 参数必须是一个数字');
    }
    
    // 步骤5: 创建变量imageHeight 用来接收外部传入的图片高度
    if (!imageHeight || typeof imageHeight !== 'number') {
      throw new Error('imageHeight 参数必须是一个数字');
    }
    
    if (!imageName) {
      throw new Error('imageName 参数不能为空');
    }
    
    if (!outputDir) {
      throw new Error('outputDir 参数不能为空');
    }
    
    // 确保输出目录存在
    if (!fs.existsSync(outputDir)) {
      fs.mkdirSync(outputDir, { recursive: true });
    }
    
    console.log('='.repeat(60));
    console.log('📝 排序对话并生成txt文件');
    console.log('='.repeat(60));
    console.log(`📋 对话数量: ${dialogues.length}`);
    console.log(`📦 格子数量: ${panels.length}`);
    console.log(`📐 图片尺寸: ${imageWidth}x${imageHeight}`);
    console.log(`📂 输出目录: ${outputDir}`);
    
    // 步骤6: 根据dialogues和panels调用sortDialoguesByPanels()函数排序对话
    console.log('\n📋 步骤6: 调用sortDialoguesByPanels()函数排序对话...');
    const sortedDialogues = sortDialoguesByPanels(dialogues, panels, imageWidth, imageHeight);
    
    console.log(`✅ 已排序 ${sortedDialogues.length} 段对话`);
    
    // 步骤7: 根据outputDir路径保存结果txt文件
    console.log('\n💾 步骤7: 根据outputDir路径保存结果txt文件...');
    const result = generateCheckTxtFiles(sortedDialogues, imageName, outputDir, 5);
    
    // 统计生成的文件
    const txtFiles = fs.readdirSync(outputDir)
      .filter(f => f.startsWith(imageName) && f.endsWith('.txt'))
      .sort();
    
    console.log(`✅ 已生成 ${txtFiles.length} 个txt文件`);
    
    console.log('\n' + '='.repeat(60));
    console.log('✅ 处理完成！');
    console.log('='.repeat(60));
    
    return {
      sortedDialogues: sortedDialogues,
      txtFiles: txtFiles,
      outputDir: outputDir
    };
    
  } catch (error) {
    console.error(`\n❌ 处理失败: ${error.message}`);
    if (error.stack) {
      console.error(error.stack);
    }
    throw error;
  }
}

/**
 * 步骤6: 根据dialogues和panels调用sortDialoguesByPanels()函数排序对话
 * 按格子位置排序对话
 * 排序规则：
 * 1. 先按格子的Y坐标（从上到下）- 越往上越靠前
 * 2. 再按格子的X坐标（从右到左）- 越往右越靠前
 * 3. 同一格子内，按文字块的X坐标从右到左
 * 
 * @param {Array} dialogues - 对话列表，每个包含bbox和text信息
 * @param {Array} panels - 格子列表
 * @param {number} imageWidth - 图片宽度
 * @param {number} imageHeight - 图片高度
 * @returns {Array} 排序后的对话列表
 */
function sortDialoguesByPanels(dialogues, panels, imageWidth, imageHeight) {
  if (!dialogues || dialogues.length === 0) {
    return [];
  }
  
  // 计算每个对话的中心点和所属格子
  const dialoguesWithInfo = dialogues.map((dialogue, index) => {
    const bbox = dialogue.bbox;
    const centerX = (bbox.x1 + bbox.x2) / 2;
    const centerY = (bbox.y1 + bbox.y2) / 2;
    const panelIndex = getTextBlockPanel(dialogue, panels);
    
    return {
      ...dialogue,
      centerX,
      centerY,
      panelIndex,
      originalIndex: index
    };
  });
  
  // 排序规则：从右到左、从上到下（日式漫画阅读顺序）
  // 1. 先按格子的Y坐标（从上到下）- Y坐标越小（越往上）越靠前
  // 2. 再按格子的X坐标（从右到左）- X坐标越大（越靠右）越靠前
  // 3. 同一格子内，先按文字块的Y坐标从上到下，再按X坐标从右到左
  dialoguesWithInfo.sort((a, b) => {
    const panelA = a.panelIndex >= 0 && a.panelIndex < panels.length 
      ? panels[a.panelIndex] 
      : null;
    const panelB = b.panelIndex >= 0 && b.panelIndex < panels.length 
      ? panels[b.panelIndex] 
      : null;
    
    if (panelA && panelB) {
      // 两个都在格子里
      // 先按格子的Y坐标（从上到下）- Y坐标越小（越往上）越靠前
      const yDiff = panelA.y - panelB.y;
      if (Math.abs(yDiff) > imageHeight * 0.05) {
        return yDiff;
      }
      
      // 再按格子的X坐标（从右到左）- X坐标越大（越靠右）越靠前
      const xDiff = panelB.center_x - panelA.center_x;
      if (Math.abs(xDiff) > imageWidth * 0.1) {
        return xDiff;
      }
      
      // 同一格子内，先按文字块的Y坐标从上到下，再按X坐标从右到左
      const textYDiff = a.centerY - b.centerY;
      if (Math.abs(textYDiff) > (panelA.height || panelB.height) * 0.1) {
        return textYDiff;
      }
      // 同一行内，按X坐标从右到左
      return b.centerX - a.centerX;
    } else if (panelA) {
      // A在格子里，B不在
      return -1;
    } else if (panelB) {
      // B在格子里，A不在
      return 1;
    } else {
      // 两个都不在格子里，按Y坐标从上到下，X坐标从右到左
      const yDiff = a.centerY - b.centerY;
      if (Math.abs(yDiff) > imageHeight * 0.05) {
        return yDiff;
      }
      return b.centerX - a.centerX;
    }
  });
  
  // 移除临时添加的属性
  return dialoguesWithInfo.map(({ centerX, centerY, panelIndex, originalIndex, ...rest }) => rest);
}

/**
 * 步骤7: 根据outputDir路径保存结果txt文件
 * 生成check.txt文件
 * 对对话进行排序和分组，然后生成txt文件
 * 
 * @param {Array} dialogues - 对话列表，每个包含bbox和text信息
 * @param {string} imageName - 图片名称（不含扩展名）
 * @param {string} outputDir - 输出目录
 * @param {number} yThreshold - y1分组阈值，默认5（正负5以内的对话归为同一气泡框）
 * @returns {Array} 排序后的对话列表
 */
function generateCheckTxtFiles(dialogues, imageName, outputDir, yThreshold = 5) {
  if (!dialogues || dialogues.length === 0) {
    return [];
  }
  
  // 确保输出目录存在
  if (!fs.existsSync(outputDir)) {
    fs.mkdirSync(outputDir, { recursive: true });
  }
  
  // 检查是否有panel_index和block_index，如果有则按panel_index和block_index分组
  const hasPanelIndex = dialogues.some(d => d.panel_index !== undefined);
  const hasBlockIndex = dialogues.some(d => d.block_index !== undefined);
  
  if (hasPanelIndex && hasBlockIndex) {
    // 先清理所有旧的txt文件（按panel分组清理）
    const panelsToClean = new Set();
    for (const dialogue of dialogues) {
      const panelIndex = dialogue.panel_index || 1;
      panelsToClean.add(panelIndex);
    }
    
    for (const panelIndex of panelsToClean) {
      const checkNFolder = path.join(outputDir, `check${panelIndex}`);
      if (fs.existsSync(checkNFolder)) {
        try {
          const files = fs.readdirSync(checkNFolder);
          for (const file of files) {
            // 删除所有checkN*.txt文件（包括checkN.txt和checkN_M.txt）
            if (file.match(/^check\d+(_\d+)?\.txt$/)) {
              const filePath = path.join(checkNFolder, file);
              fs.unlinkSync(filePath);
            }
          }
        } catch (error) {
          console.log(`    ⚠️  清理check${panelIndex}文件夹的旧txt文件失败: ${error.message}`);
        }
      }
    }
    
    // 按panel_index和block_index分组
    const dialoguesByPanelAndBlock = {};
    for (const dialogue of dialogues) {
      const panelIndex = dialogue.panel_index || 1;
      const blockIndex = dialogue.block_index || 1;
      const key = `${panelIndex}_${blockIndex}`;
      if (!dialoguesByPanelAndBlock[key]) {
        dialoguesByPanelAndBlock[key] = {
          panelIndex,
          blockIndex,
          dialogues: []
        };
      }
      dialoguesByPanelAndBlock[key].dialogues.push(dialogue);
    }
    
    // 为每个panel和block组合生成txt文件
    const generatedFiles = [];
    const sortedKeys = Object.keys(dialoguesByPanelAndBlock).sort((a, b) => {
      const [panelA, blockA] = a.split('_').map(Number);
      const [panelB, blockB] = b.split('_').map(Number);
      if (panelA !== panelB) {
        return panelA - panelB;
      }
      return blockA - blockB;
    });
    
    for (const key of sortedKeys) {
      const { panelIndex, blockIndex, dialogues: blockDialogues } = dialoguesByPanelAndBlock[key];
      
      // 对同一个block内的对话进行排序：先按y坐标从上到下，再按x坐标从右到左
      const sortedDialogues = [...blockDialogues].sort((a, b) => {
        const y1A = a.bbox?.y1 || 0;
        const y1B = b.bbox?.y1 || 0;
        if (Math.abs(y1A - y1B) > 5) {
          return y1A - y1B;  // 从上到下
        }
        const x1A = a.bbox?.x1 || 0;
        const x1B = b.bbox?.x1 || 0;
        return x1B - x1A;  // 从右到左
      });
      
      // 合并所有对话为一个txt文件
      const txtContent = sortedDialogues.map(d => d.text).join('');
      
      // 将txt文件保存到对应的checkN文件夹下
      const checkNFolder = path.join(outputDir, `check${panelIndex}`);
      if (!fs.existsSync(checkNFolder)) {
        fs.mkdirSync(checkNFolder, { recursive: true });
      }
      
      // 文件名格式：checkN_M.txt，其中N是panel_index，M是block_index
      const txtFileName = `check${panelIndex}_${blockIndex}.txt`;
      const txtPath = path.join(checkNFolder, txtFileName);
      
      fs.writeFileSync(txtPath, txtContent, 'utf-8');
      generatedFiles.push(path.join(`check${panelIndex}`, txtFileName));
      console.log(`    📄 已生成txt文件: check${panelIndex}/${txtFileName} (${sortedDialogues.length}段对话)`);
    }
    
    // 返回所有对话（已排序）
    return dialogues;
  } else if (hasPanelIndex) {
    // 只有panel_index，没有block_index，使用原来的逻辑（按y坐标分组）
    // 按panel_index分组
    const dialoguesByPanel = {};
    for (const dialogue of dialogues) {
      const panelIndex = dialogue.panel_index || 1;
      if (!dialoguesByPanel[panelIndex]) {
        dialoguesByPanel[panelIndex] = [];
      }
      dialoguesByPanel[panelIndex].push(dialogue);
    }
    
    // 为每个panel生成txt文件
    const generatedFiles = [];
    for (const panelIndex of Object.keys(dialoguesByPanel).sort((a, b) => parseInt(a) - parseInt(b))) {
      const panelDialogues = dialoguesByPanel[panelIndex];
      
      // 根据 bbox.x1 排序（x1 越大越靠前，从右到左的阅读顺序）
      const sortedDialogues = [...panelDialogues].sort((a, b) => {
        const x1A = a.bbox?.x1 || 0;
        const x1B = b.bbox?.x1 || 0;
        return x1B - x1A;  // 从右到左：x1 越大越靠前
      });
      
      // 按 y1 分组（y1 相同或相近的归为一组）
      const groups = [];
      let currentGroup = [];
      let currentY1 = null;
      
      for (const dialogue of sortedDialogues) {
        const y1 = dialogue.bbox?.y1 || 0;
        
        if (currentY1 === null || Math.abs(y1 - currentY1) <= yThreshold) {
          currentGroup.push(dialogue);
          if (currentY1 === null) {
            currentY1 = y1;
          }
        } else {
          if (currentGroup.length > 0) {
            groups.push(currentGroup);
          }
          currentGroup = [dialogue];
          currentY1 = y1;
        }
      }
      
      if (currentGroup.length > 0) {
        groups.push(currentGroup);
      }
      
      // 为每组生成一个 txt 文件
      // 将txt文件保存到对应的checkN文件夹下
      const checkNFolder = path.join(outputDir, `check${panelIndex}`);
      if (!fs.existsSync(checkNFolder)) {
        fs.mkdirSync(checkNFolder, { recursive: true });
      }
      
      for (let i = 0; i < groups.length; i++) {
        const group = groups[i];
        const txtContent = group.map(d => d.text).join('');
        
        // 如果只有一组，使用checkN.txt；否则使用checkN_M.txt
        const txtFileName = groups.length === 1 
          ? `check${panelIndex}.txt` 
          : `check${panelIndex}_${i + 1}.txt`;
        const txtPath = path.join(checkNFolder, txtFileName);
        
        fs.writeFileSync(txtPath, txtContent, 'utf-8');
        generatedFiles.push(path.join(`check${panelIndex}`, txtFileName));
        console.log(`    📄 已生成txt文件: check${panelIndex}/${txtFileName} (${group.length}段对话, y1≈${group[0].bbox?.y1 || 0})`);
      }
    }
    
    // 返回所有对话（已排序）
    return dialogues;
  } else {
    // 没有panel_index，使用原来的逻辑
    // 根据 bbox.x1 排序（x1 越大越靠前，从右到左的阅读顺序）
    const sortedDialogues = [...dialogues].sort((a, b) => {
      const x1A = a.bbox?.x1 || 0;
      const x1B = b.bbox?.x1 || 0;
      return x1B - x1A;  // 从右到左：x1 越大越靠前
    });
    
    // 按 y1 分组（y1 相同或相近的归为一组）
    const groups = [];
    let currentGroup = [];
    let currentY1 = null;
    
    for (const dialogue of sortedDialogues) {
      const y1 = dialogue.bbox?.y1 || 0;
      
      if (currentY1 === null || Math.abs(y1 - currentY1) <= yThreshold) {
        currentGroup.push(dialogue);
        if (currentY1 === null) {
          currentY1 = y1;
        }
      } else {
        if (currentGroup.length > 0) {
          groups.push(currentGroup);
        }
        currentGroup = [dialogue];
        currentY1 = y1;
      }
    }
    
    if (currentGroup.length > 0) {
      groups.push(currentGroup);
    }
    
    // 为每组生成一个 txt 文件
    const generatedFiles = [];
    for (let i = 0; i < groups.length; i++) {
      const group = groups[i];
      const txtContent = group.map(d => d.text).join('');
      
      // 如果只有一组，使用原文件名；否则使用带序号的文件名
      const txtFileName = groups.length === 1 
        ? `${imageName}.txt` 
        : `${imageName}_${i + 1}.txt`;
      const txtPath = path.join(outputDir, txtFileName);
      
      fs.writeFileSync(txtPath, txtContent, 'utf-8');
      generatedFiles.push(txtFileName);
      console.log(`    📄 已生成txt文件: ${txtFileName} (${group.length}段对话, y1≈${group[0].bbox?.y1 || 0})`);
    }
    
    return sortedDialogues;
  }
}

/**
 * 判断文字块属于哪个格子（辅助函数，被sortDialoguesByPanels调用）
 * @param {Object} textBlock - 文字块，包含bbox信息
 * @param {Array} panels - 格子列表
 * @returns {number} 格子索引，如果不在任何格子内返回-1
 */
function getTextBlockPanel(textBlock, panels) {
  const bbox = textBlock.bbox;
  const centerX = (bbox.x1 + bbox.x2) / 2;
  const centerY = (bbox.y1 + bbox.y2) / 2;
  
  for (let i = 0; i < panels.length; i++) {
    const panel = panels[i];
    if (
      panel.x <= centerX && centerX <= panel.x + panel.width &&
      panel.y <= centerY && centerY <= panel.y + panel.height
    ) {
      return i;
    }
  }
  
  return -1;
}

// 导出函数
export { startSortDialog, sortDialoguesByPanels };