validate-models.js 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115
  1. // scripts/validate-models.js
  2. import fs from 'fs-extra';
  3. import path from 'path';
  4. import { fileURLToPath } from 'url';
  5. const __dirname = path.dirname(fileURLToPath(import.meta.url));
  6. class ModelValidator {
  7. constructor() {
  8. this.modelDir = path.join(process.cwd(), 'models', 'ocr');
  9. this.requiredFiles = {
  10. detection: path.join(this.modelDir, 'Det', '中文_OCRv3.onnx'),
  11. recognition: path.join(this.modelDir, 'Rec', '中文简体_OCRv3.onnx'),
  12. classification: path.join(this.modelDir, 'Cls', '原始分类器模型.onnx'),
  13. keys: path.join(this.modelDir, 'Keys', '中文简体_OCRv3.txt')
  14. };
  15. }
  16. async validateModels() {
  17. console.log('🔍 验证模型文件...\n');
  18. let allValid = true;
  19. const results = [];
  20. for (const [name, filePath] of Object.entries(this.requiredFiles)) {
  21. const exists = await fs.pathExists(filePath);
  22. const result = {
  23. name: this.getDisplayName(name),
  24. path: filePath,
  25. exists,
  26. size: 0,
  27. valid: false
  28. };
  29. if (exists) {
  30. try {
  31. const stats = await fs.stat(filePath);
  32. result.size = stats.size;
  33. result.valid = this.validateFileSize(name, stats.size);
  34. if (name === 'keys') {
  35. result.valid = await this.validateKeysFile(filePath);
  36. }
  37. } catch (error) {
  38. result.valid = false;
  39. }
  40. }
  41. results.push(result);
  42. allValid = allValid && result.valid;
  43. }
  44. // 显示验证结果
  45. this.displayResults(results);
  46. if (!allValid) {
  47. console.log('\n❌ 模型文件不完整或损坏');
  48. console.log('💡 运行以下命令重新下载:');
  49. console.log(' yarn download-models');
  50. process.exit(1);
  51. } else {
  52. console.log('\n✅ 所有模型文件验证通过!');
  53. }
  54. return allValid;
  55. }
  56. getDisplayName(name) {
  57. const names = {
  58. detection: '检测模型',
  59. recognition: '识别模型',
  60. classification: '分类模型',
  61. keys: '字符集文件'
  62. };
  63. return names[name] || name;
  64. }
  65. validateFileSize(name, size) {
  66. const minSizes = {
  67. detection: 1000000, // 1MB
  68. recognition: 5000000, // 5MB
  69. classification: 1000000, // 1MB
  70. keys: 1000 // 1KB
  71. };
  72. return size > (minSizes[name] || 0);
  73. }
  74. async validateKeysFile(filePath) {
  75. try {
  76. const content = await fs.readFile(filePath, 'utf8');
  77. const lines = content.split('\n').filter(line => line.trim());
  78. return lines.length > 1000; // 至少1000个字符
  79. } catch {
  80. return false;
  81. }
  82. }
  83. displayResults(results) {
  84. results.forEach(result => {
  85. const status = result.exists && result.valid ? '✅' : '❌';
  86. const sizeInfo = result.exists ? ` (${(result.size / 1024 / 1024).toFixed(2)} MB)` : '';
  87. console.log(`${status} ${result.name}${sizeInfo}`);
  88. if (!result.exists) {
  89. console.log(` 文件不存在: ${result.path}`);
  90. } else if (!result.valid) {
  91. console.log(` 文件可能损坏: ${result.path}`);
  92. }
  93. });
  94. }
  95. }
  96. // 执行验证
  97. const validator = new ModelValidator();
  98. validator.validateModels().catch(console.error);