关注博客和手机以及互联网，记录生活点滴，分享所见所闻。
// 初始化模型 await PPOCRv5.init(); print("✅ 模型初始化成功"); final startTime = DateTime.now(); print("⏱️ 开始识别：${startTime.toLocal()}"); // 识别 assets 图片 final result = await PPOCRv5.fromAsset('assets/images/test4.png'); print("识别结果：$result"); // 结束计时+打印 final endTime = DateTime.now(); final cost = endTime.difference(startTime); print("🏁 结束识别：${endTime.toLocal()}"); print("⏱️ 耗时：${cost.inMilliseconds}ms / ${cost.inSeconds}s"); print("结果：$result"); import 'dart:math' as math; import 'dart:typed_data'; import 'package:flutter/services.dart' show rootBundle; import 'package:flutter_onnxruntime/flutter_onnxruntime.dart'; import 'package:image/image.dart' as img; import 'package:opencv_dart/opencv.dart' as cv; class PPOCRArgs { String detModel = 'assets/models/det_infer.onnx'; String recModel = 'assets/models/rec_infer.onnx'; String dictPath = 'assets/models/ppocrv5_dict.txt'; String? clsModel; int detLimitSideLen = 960; String detLimitType = 'max'; double detDbThresh = 0.3; double detDbBoxThresh = 0.6; double detDbUnclipRatio = 1.5; bool useDilation = false; String detDbScoreMode = 'fast'; String detBoxType = 'quad'; List<int> recImageShape = [3, 48, 320]; int recBatchNum = 6; bool useSpaceChar = true; double dropScore = 0.5; List<int> clsImageShape = [3, 48, 192]; List<String> labelList = ['0', '180']; int clsBatchNum = 6; double clsThresh = 0.9; bool get useAngleCls => clsModel != null && clsModel!.isNotEmpty; } class _DetPreprocessResult { final Float32List input; final List<int> inputShape; final List<double> shape; _DetPreprocessResult({ required this.input, required this.inputShape, required this.shape, }); } class _RecDecodeResult { final String text; final double score; const _RecDecodeResult(this.text, this.score); } class _ClsDecodeResult { final String label; final double score; const _ClsDecodeResult(this.label, this.score); } class _ImageConverters { static cv.Mat imageToMat(img.Image image) { final rgb = image.hasPalette ? image.convert(numChannels: 3) : image; final bytes = rgb.getBytes(order: img.ChannelOrder.rgb); return cv.Mat.fromList( rgb.height, rgb.width, cv.MatType.CV_8UC3, bytes, ); } static img.Image matToImage(cv.Mat mat) { final data = mat.toList3D(); final image = img.Image(width: mat.cols, height: mat.rows); for (int y = 0; y < mat.rows; y++) { for (int x = 0; x < mat.cols; x++) { final pixel = data[y][x]; image.setPixelRgb( x, y, pixel[0].toInt(), pixel[1].toInt(), pixel[2].toInt(), ); } } return image; } } class _DetResizeForTest { final int limitSideLen; final String limitType; const _DetResizeForTest({required this.limitSideLen, required this.limitType}); ({img.Image image, List<double> shape}) call(img.Image image) { final srcH = image.height; final srcW = image.width; final limitSide = limitType == 'max' ? math.max(srcH, srcW) : math.min(srcH, srcW); final ratio = limitSide > limitSideLen ? limitSideLen / limitSide : 1.0; int resizeH = (srcH * ratio).toInt(); int resizeW = (srcW * ratio).toInt(); resizeH = math.max(((resizeH / 32).round()) * 32, 32); resizeW = math.max(((resizeW / 32).round()) * 32, 32); final resized = img.copyResize(image, width: resizeW, height: resizeH); return ( image: resized, shape: [ srcH.toDouble(), srcW.toDouble(), resizeH / srcH, resizeW / srcW, ], ); } } class _DBPostProcess { final double thresh; final double boxThresh; final int maxCandidates; final double unclipRatio; final bool useDilation; final String scoreMode; final String boxType; final int minSize = 3; const _DBPostProcess({ required this.thresh, required this.boxThresh, required this.maxCandidates, required this.unclipRatio, required this.useDilation, required this.scoreMode, required this.boxType, }); List<List<List<double>>> call({ required List<double> pred, required int predH, required int predW, required List<double> shape, }) { final srcH = shape[0].toInt(); final srcW = shape[1].toInt(); final bitmapData = Uint8List(predH * predW); for (int i = 0; i < pred.length; i++) { bitmapData[i] = pred[i] > thresh ? 255 : 0; } var bitmap = cv.Mat.fromList(predH, predW, cv.MatType.CV_8UC1, bitmapData); if (useDilation) { final kernel = cv.Mat.fromList(2, 2, cv.MatType.CV_8UC1, [1, 1, 1, 1]); bitmap = cv.dilate(bitmap, kernel); } if (boxType == 'poly') { return _polygonsFromBitmap(pred, bitmap, predH, predW, srcW, srcH); } return _boxesFromBitmap(pred, bitmap, predH, predW, srcW, srcH); } List<List<List<double>>> _boxesFromBitmap( List<double> pred, cv.Mat bitmap, int predH, int predW, int destW, int destH, ) { final (contours, _) = cv.findContours(bitmap, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE); final numContours = math.min(contours.length, maxCandidates); final boxes = <List<List<double>>>[]; for (int i = 0; i < numContours; i++) { final contour = contours[i]; final mini = _getMiniBoxes(contour); if (mini == null || mini.$2 < minSize) { continue; } final points = mini.$1; final score = scoreMode == 'fast' ? _boxScoreFast(pred, predH, predW, points) : _boxScoreSlow(pred, predH, predW, _vecPointToDoubleList(contour)); if (score < boxThresh) { continue; } final expandedPaths = _unclip(points, unclipRatio); if (expandedPaths.length != 1) { continue; } final expanded = expandedPaths.first; final expandedVec = cv.VecPoint.fromList( expanded.map((e) => cv.Point(e[0].round(), e[1].round())).toList(), ); final expandedMini = _getMiniBoxes(expandedVec); if (expandedMini == null || expandedMini.$2 < minSize + 2) { continue; } final box = expandedMini.$1 .map( (p) => [ _clipDouble((p[0] / predW * destW).roundToDouble(), 0, destW.toDouble()), _clipDouble((p[1] / predH * destH).roundToDouble(), 0, destH.toDouble()), ], ) .toList(); boxes.add(box); } return boxes; } List<List<List<double>>> _polygonsFromBitmap( List<double> pred, cv.Mat bitmap, int predH, int predW, int destW, int destH, ) { final (contours, _) = cv.findContours(bitmap, cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE); final boxes = <List<List<double>>>[]; final limit = math.min(contours.length, maxCandidates); for (int i = 0; i < limit; i++) { final contour = contours[i]; final epsilon = 0.002 * cv.arcLength(contour, true); final approx = cv.approxPolyDP(contour, epsilon, true); if (approx.length < 4) { continue; } final points = approx.map((e) => [e.x.toDouble(), e.y.toDouble()]).toList(); final score = _boxScoreFast(pred, predH, predW, points); if (score < boxThresh) { continue; } final expandedPaths = _unclip(points, unclipRatio); if (expandedPaths.length != 1) { continue; } final expanded = expandedPaths.first; final expandedVec = cv.VecPoint.fromList( expanded.map((e) => cv.Point(e[0].round(), e[1].round())).toList(), ); final mini = _getMiniBoxes(expandedVec); if (mini == null || mini.$2 < minSize + 2) { continue; } final scaled = expanded .map( (p) => [ _clipDouble((p[0] / predW * destW).roundToDouble(), 0, destW.toDouble()), _clipDouble((p[1] / predH * destH).roundToDouble(), 0, destH.toDouble()), ], ) .toList(); boxes.add(scaled); } return boxes; } (List<List<double>>, double)? _getMiniBoxes(cv.VecPoint contour) { if (contour.length < 4) { return null; } final rotatedRect = cv.minAreaRect(contour); final points = cv.boxPoints(rotatedRect) .map((e) => [e.x, e.y]) .toList() ..sort((a, b) => a[0].compareTo(b[0])); int index1; int index2; int index3; int index4; if (points[1][1] > points[0][1]) { index1 = 0; index4 = 1; } else { index1 = 1; index4 = 0; } if (points[3][1] > points[2][1]) { index2 = 2; index3 = 3; } else { index2 = 3; index3 = 2; } final box = [ points[index1], points[index2], points[index3], points[index4], ]; final shortSide = math.min(rotatedRect.size.width, rotatedRect.size.height); return (box, shortSide); } double _boxScoreFast( List<double> bitmap, int h, int w, List<List<double>> box, ) { int xmin = box.map((p) => p[0].floor()).reduce(math.min).clamp(0, w - 1); int xmax = box.map((p) => p[0].ceil()).reduce(math.max).clamp(0, w - 1); int ymin = box.map((p) => p[1].floor()).reduce(math.min).clamp(0, h - 1); int ymax = box.map((p) => p[1].ceil()).reduce(math.max).clamp(0, h - 1); final width = xmax - xmin + 1; final height = ymax - ymin + 1; if (width <= 0 || height <= 0) { return 0; } final mask = cv.Mat.fromScalar(height, width, cv.MatType.CV_8UC1, cv.Scalar()); final shifted = box .map((p) => cv.Point((p[0] - xmin).round(), (p[1] - ymin).round())) .toList(); cv.fillPoly(mask, cv.VecVecPoint.fromList([shifted]), cv.Scalar(1)); final cropData = <double>[]; for (int y = ymin; y <= ymax; y++) { for (int x = xmin; x <= xmax; x++) { cropData.add(bitmap[y * w + x]); } } final crop = cv.Mat.fromList(height, width, cv.MatType.CV_32FC1, cropData); return crop.mean(mask: mask).val1; } double _boxScoreSlow( List<double> bitmap, int h, int w, List<List<double>> contour, ) { int xmin = contour.map((p) => p[0]).reduce(math.min).floor().clamp(0, w - 1); int xmax = contour.map((p) => p[0]).reduce(math.max).ceil().clamp(0, w - 1); int ymin = contour.map((p) => p[1]).reduce(math.min).floor().clamp(0, h - 1); int ymax = contour.map((p) => p[1]).reduce(math.max).ceil().clamp(0, h - 1); final width = xmax - xmin + 1; final height = ymax - ymin + 1; if (width <= 0 || height <= 0) { return 0; } final mask = cv.Mat.fromScalar(height, width, cv.MatType.CV_8UC1, cv.Scalar()); final shifted = contour .map((p) => cv.Point((p[0] - xmin).round(), (p[1] - ymin).round())) .toList(); cv.fillPoly(mask, cv.VecVecPoint.fromList([shifted]), cv.Scalar(1)); final cropData = <double>[]; for (int y = ymin; y <= ymax; y++) { for (int x = xmin; x <= xmax; x++) { cropData.add(bitmap[y * w + x]); } } final crop = cv.Mat.fromList(height, width, cv.MatType.CV_32FC1, cropData); return crop.mean(mask: mask).val1; } List<List<double>> _vecPointToDoubleList(cv.VecPoint contour) { return contour.map((p) => [p.x.toDouble(), p.y.toDouble()]).toList(); } List<List<List<double>>> _unclip(List<List<double>> box, double ratio) { final area = _polygonArea(box).abs(); final perimeter = _polygonPerimeter(box); if (box.length < 3 || perimeter <= 0 || area <= 0) { return []; } final distance = area * ratio / perimeter; if (distance <= 0) { return [box]; } final expanded = _offsetPolygon(box, distance); if (expanded.length < 3 || !_isFinitePolygon(expanded)) { return []; } return [expanded]; } List<List<double>> _offsetPolygon(List<List<double>> polygon, double distance) { final area = _polygonArea(polygon); final sign = area >= 0 ? 1.0 : -1.0; final count = polygon.length; final result = <List<double>>[]; for (int i = 0; i < count; i++) { final prev = polygon[(i - 1 + count) % count]; final curr = polygon[i]; final next = polygon[(i + 1) % count]; final prevLine = _offsetLine(prev, curr, distance, sign); final nextLine = _offsetLine(curr, next, distance, sign); final intersection = _lineIntersection( prevLine.$1, prevLine.$2, nextLine.$1, nextLine.$2, ); if (intersection != null && intersection[0].isFinite && intersection[1].isFinite) { result.add(intersection); } else { final n1 = _edgeNormal(prev, curr, sign); final n2 = _edgeNormal(curr, next, sign); final nx = n1[0] + n2[0]; final ny = n1[1] + n2[1]; final nl = math.sqrt(nx * nx + ny * ny); if (nl > 1e-6) { result.add([ curr[0] + nx / nl * distance, curr[1] + ny / nl * distance, ]); } else { result.add([ curr[0] + n1[0] * distance, curr[1] + n1[1] * distance, ]); } } } return result; } (List<double>, List<double>) _offsetLine( List<double> a, List<double> b, double distance, double sign, ) { final normal = _edgeNormal(a, b, sign); return ( [a[0] + normal[0] * distance, a[1] + normal[1] * distance], [b[0] + normal[0] * distance, b[1] + normal[1] * distance], ); } List<double> _edgeNormal(List<double> a, List<double> b, double sign) { final dx = b[0] - a[0]; final dy = b[1] - a[1]; final len = math.sqrt(dx * dx + dy * dy); if (len <= 1e-6) { return [0, 0]; } if (sign >= 0) { return [dy / len, -dx / len]; } return [-dy / len, dx / len]; } List<double>? _lineIntersection( List<double> a1, List<double> a2, List<double> b1, List<double> b2, ) { final x1 = a1[0]; final y1 = a1[1]; final x2 = a2[0]; final y2 = a2[1]; final x3 = b1[0]; final y3 = b1[1]; final x4 = b2[0]; final y4 = b2[1]; final denom = (x1 - x2) * (y3 - y4) - (y1 - y2) * (x3 - x4); if (denom.abs() <= 1e-6) { return null; } final px = ((x1 * y2 - y1 * x2) * (x3 - x4) - (x1 - x2) * (x3 * y4 - y3 * x4)) / denom; final py = ((x1 * y2 - y1 * x2) * (y3 - y4) - (y1 - y2) * (x3 * y4 - y3 * x4)) / denom; return [px, py]; } bool _isFinitePolygon(List<List<double>> polygon) { for (final point in polygon) { if (point.length != 2 || !point[0].isFinite || !point[1].isFinite) { return false; } } return true; } double _polygonArea(List<List<double>> points) { double area = 0; for (int i = 0; i < points.length; i++) { final j = (i + 1) % points.length; area += points[i][0] * points[j][1] - points[j][0] * points[i][1]; } return area / 2.0; } double _polygonPerimeter(List<List<double>> points) { double perimeter = 0; for (int i = 0; i < points.length; i++) { final j = (i + 1) % points.length; final dx = points[i][0] - points[j][0]; final dy = points[i][1] - points[j][1]; perimeter += math.sqrt(dx * dx + dy * dy); } return perimeter; } double _clipDouble(double value, double min, double max) { if (value < min) return min; if (value > max) return max; return value; } } class _TextDetector { final PPOCRArgs args; late final OrtSession session; late final String inputName; late final _DetResizeForTest _resize; late final _DBPostProcess _postProcess; _TextDetector(this.args); Future<void> init(OnnxRuntime runtime) async { session = await runtime.createSessionFromAsset(args.detModel); final inputs = await session.getInputInfo(); if (inputs.isEmpty) { throw Exception('检测模型没有输入信息'); } inputName = inputs.first['name']!; _resize = _DetResizeForTest( limitSideLen: args.detLimitSideLen, limitType: args.detLimitType, ); _postProcess = _DBPostProcess( thresh: args.detDbThresh, boxThresh: args.detDbBoxThresh, maxCandidates: 1000, unclipRatio: args.detDbUnclipRatio, useDilation: args.useDilation, scoreMode: args.detDbScoreMode, boxType: args.detBoxType, ); } Future<List<List<List<double>>>> detect(img.Image image) async { final pre = _preprocess(image); final inputTensor = await OrtValue.fromList(pre.input, pre.inputShape); final outputs = await session.run({inputName: inputTensor}); await inputTensor.dispose(); final output = outputs.values.first; final flattened = await output.asFlattenedList(); final pred = flattened.map((e) => (e as num).toDouble()).toList(); final shape = output.shape; if (shape.length < 4) { throw Exception('检测模型输出 shape 不符合预期: $shape'); } final predH = shape[2]; final predW = shape[3]; final boxes = _postProcess( pred: pred, predH: predH, predW: predW, shape: pre.shape, ); return args.detBoxType == 'poly' ? _filterTagDetResOnlyClip(boxes, image.height, image.width) : _filterTagDetRes(boxes, image.height, image.width); } _DetPreprocessResult _preprocess(img.Image image) { final resized = _resize(image); final normalized = _normalizeDetImage(resized.image); return _DetPreprocessResult( input: normalized, inputShape: [1, 3, resized.image.height, resized.image.width], shape: resized.shape, ); } Float32List _normalizeDetImage(img.Image image) { const mean = [0.485, 0.456, 0.406]; const std = [0.229, 0.224, 0.225]; final out = Float32List(3 * image.height * image.width); int offset = 0; for (int c = 0; c < 3; c++) { for (int y = 0; y < image.height; y++) { for (int x = 0; x < image.width; x++) { final pixel = image.getPixel(x, y); final value = switch (c) { 0 => pixel.r, 1 => pixel.g, _ => pixel.b, }; out[offset++] = ((value / 255.0) - mean[c]) / std[c]; } } } return out; } List<List<List<double>>> _filterTagDetRes( List<List<List<double>>> dtBoxes, int imgHeight, int imgWidth, ) { final result = <List<List<double>>>[]; for (final rawBox in dtBoxes) { final box = _clipDetRes(_orderPointsClockwise(rawBox), imgHeight, imgWidth); final rectWidth = _distance(box[0], box[1]).round(); final rectHeight = _distance(box[0], box[3]).round(); if (rectWidth <= 3 || rectHeight <= 3) { continue; } result.add(box); } return result; } List<List<List<double>>> _filterTagDetResOnlyClip( List<List<List<double>>> dtBoxes, int imgHeight, int imgWidth, ) { return dtBoxes.map((box) => _clipDetRes(box, imgHeight, imgWidth)).toList(); } List<List<double>> _orderPointsClockwise(List<List<double>> pts) { final sortedBySum = [...pts]..sort((a, b) => (a[0] + a[1]).compareTo(b[0] + b[1])); final topLeft = sortedBySum.first; final bottomRight = sortedBySum.last; final remain = [...pts] ..remove(topLeft) ..remove(bottomRight); remain.sort((a, b) => ((a[1] - a[0])).compareTo(b[1] - b[0])); final topRight = remain.first; final bottomLeft = remain.last; return [topLeft, topRight, bottomRight, bottomLeft]; } List<List<double>> _clipDetRes(List<List<double>> points, int imgHeight, int imgWidth) { return points .map( (p) => [ p[0].clamp(0, imgWidth - 1).toDouble(), p[1].clamp(0, imgHeight - 1).toDouble(), ], ) .toList(); } double _distance(List<double> a, List<double> b) { final dx = a[0] - b[0]; final dy = a[1] - b[1]; return math.sqrt(dx * dx + dy * dy); } } class _CTCLabelDecoder { late final List<String> character; Future<void> init(String dictPath, {bool useSpaceChar = false}) async { final dictContent = await rootBundle.loadString(dictPath); final chars = dictContent .split('\n') .map((e) => e.endsWith('\r') ? e.substring(0, e.length - 1) : e) .where((e) => e.isNotEmpty) .toList(); if (useSpaceChar) { chars.add(' '); } character = ['blank', ...chars]; } List<_RecDecodeResult> decode(List<double> preds, List<int> shape) { if (shape.length != 3) { throw Exception('识别模型输出 shape 不符合预期: $shape'); } final batch = shape[0]; final timeSteps = shape[1]; final classes = shape[2]; final results = <_RecDecodeResult>[]; for (int b = 0; b < batch; b++) { final indices = <int>[]; final scores = <double>[]; for (int t = 0; t < timeSteps; t++) { double maxScore = -double.infinity; int maxIdx = 0; for (int c = 0; c < classes; c++) { final index = b * timeSteps * classes + t * classes + c; final score = preds[index]; if (score > maxScore) { maxScore = score; maxIdx = c; } } indices.add(maxIdx); scores.add(maxScore); } final chars = <String>[]; final confs = <double>[]; for (int i = 0; i < indices.length; i++) { final idx = indices[i]; if (idx == 0) continue; if (i > 0 && indices[i] == indices[i - 1]) continue; if (idx >= 0 && idx < character.length) { chars.add(character[idx]); confs.add(scores[i]); } } final text = chars.join(); final score = confs.isEmpty ? 0.0 : confs.reduce((a, b) => a + b) / confs.length; results.add(_RecDecodeResult(text, score)); } return results; } } class _ClsPostProcess { final List<String> labelList; const _ClsPostProcess(this.labelList); List<_ClsDecodeResult> decode(List<double> preds, List<int> shape) { if (shape.length != 2) { throw Exception('方向分类模型输出 shape 不符合预期: $shape'); } final batch = shape[0]; final classes = shape[1]; final results = <_ClsDecodeResult>[]; for (int b = 0; b < batch; b++) { double maxScore = -double.infinity; int maxIdx = 0; for (int c = 0; c < classes; c++) { final score = preds[b * classes + c]; if (score > maxScore) { maxScore = score; maxIdx = c; } } results.add(_ClsDecodeResult(labelList[maxIdx], maxScore)); } return results; } } class _TextRecognizer { final PPOCRArgs args; final _CTCLabelDecoder _decoder = _CTCLabelDecoder(); late final OrtSession session; late final String inputName; _TextRecognizer(this.args); Future<void> init(OnnxRuntime runtime) async { session = await runtime.createSessionFromAsset(args.recModel); final inputs = await session.getInputInfo(); if (inputs.isEmpty) { throw Exception('识别模型没有输入信息'); } inputName = inputs.first['name']!; await _decoder.init(args.dictPath, useSpaceChar: args.useSpaceChar); } Future<List<_RecDecodeResult>> recognize(List<img.Image> imageList) async { if (imageList.isEmpty) { return const []; } final indexed = imageList.asMap().entries.toList() ..sort((a, b) { final wa = a.value.width / a.value.height; final wb = b.value.width / b.value.height; return wa.compareTo(wb); }); final results = List<_RecDecodeResult>.filled( imageList.length, const _RecDecodeResult('', 0.0), ); for (int beg = 0; beg < indexed.length; beg += args.recBatchNum) { final end = math.min(indexed.length, beg + args.recBatchNum); double maxWhRatio = 0; for (int i = beg; i < end; i++) { final image = indexed[i].value; maxWhRatio = math.max(maxWhRatio, image.width / image.height); } final batchData = Float32List((end - beg) * _singleImageElementCount(maxWhRatio)); int offset = 0; for (int i = beg; i < end; i++) { final normalized = _resizeNormImg(indexed[i].value, maxWhRatio); batchData.setRange(offset, offset + normalized.length, normalized); offset += normalized.length; } final inputShape = [ end - beg, args.recImageShape[0], args.recImageShape[1], math.max(args.recImageShape[2], (args.recImageShape[1] * maxWhRatio).ceil()), ]; final inputTensor = await OrtValue.fromList(batchData, inputShape); final outputs = await session.run({inputName: inputTensor}); await inputTensor.dispose(); final output = outputs.values.first; final preds = (await output.asFlattenedList()).map((e) => (e as num).toDouble()).toList(); final decoded = _decoder.decode(preds, output.shape); for (int i = 0; i < decoded.length; i++) { results[indexed[beg + i].key] = decoded[i]; } } return results; } int _singleImageElementCount(double maxWhRatio) { final imgC = args.recImageShape[0]; final imgH = args.recImageShape[1]; final imgW = math.max(args.recImageShape[2], (imgH * maxWhRatio).ceil()); return imgC * imgH * imgW; } Float32List _resizeNormImg(img.Image image, double maxWhRatio) { final imgC = args.recImageShape[0]; final imgH = args.recImageShape[1]; final imgW = math.max(args.recImageShape[2], (imgH * maxWhRatio).ceil()); final ratio = image.width / image.height; final resizedW = math.min(imgW, (imgH * ratio).ceil()); final resized = img.copyResize(image, width: resizedW, height: imgH); final output = Float32List(imgC * imgH * imgW); for (int c = 0; c < imgC; c++) { for (int y = 0; y < imgH; y++) { for (int x = 0; x < imgW; x++) { final outIndex = c * imgH * imgW + y * imgW + x; if (x >= resizedW) { output[outIndex] = 0; continue; } final pixel = resized.getPixel(x, y); final value = switch (c) { 0 => pixel.r, 1 => pixel.g, _ => pixel.b, }; output[outIndex] = ((value / 255.0) - 0.5) / 0.5; } } } return output; } } class _TextClassifier { final PPOCRArgs args; final _ClsPostProcess _postProcess; late final OrtSession session; late final String inputName; _TextClassifier(this.args) : _postProcess = _ClsPostProcess(args.labelList); Future<void> init(OnnxRuntime runtime) async { final model = args.clsModel; if (model == null || model.isEmpty) { throw Exception('未配置方向分类模型'); } session = await runtime.createSessionFromAsset(model); final inputs = await session.getInputInfo(); if (inputs.isEmpty) { throw Exception('方向分类模型没有输入信息'); } inputName = inputs.first['name']!; } Future<({List<img.Image> images, List<_ClsDecodeResult> results})> classify(List<img.Image> images) async { if (images.isEmpty) { return (images: <img.Image>[], results: <_ClsDecodeResult>[]); } final working = [...images]; final indexed = working.asMap().entries.toList() ..sort((a, b) { final wa = a.value.width / a.value.height; final wb = b.value.width / b.value.height; return wa.compareTo(wb); }); final clsRes = List<_ClsDecodeResult>.filled( working.length, const _ClsDecodeResult('', 0.0), ); for (int beg = 0; beg < indexed.length; beg += args.clsBatchNum) { final end = math.min(indexed.length, beg + args.clsBatchNum); final imgC = args.clsImageShape[0]; final imgH = args.clsImageShape[1]; final imgW = args.clsImageShape[2]; final batchData = Float32List((end - beg) * imgC * imgH * imgW); int offset = 0; for (int i = beg; i < end; i++) { final normalized = _resizeNormImg(indexed[i].value); batchData.setRange(offset, offset + normalized.length, normalized); offset += normalized.length; } final inputTensor = await OrtValue.fromList(batchData, [end - beg, imgC, imgH, imgW]); final outputs = await session.run({inputName: inputTensor}); await inputTensor.dispose(); final output = outputs.values.first; final preds = (await output.asFlattenedList()).map((e) => (e as num).toDouble()).toList(); final decoded = _postProcess.decode(preds, output.shape); for (int i = 0; i < decoded.length; i++) { final originalIndex = indexed[beg + i].key; clsRes[originalIndex] = decoded[i]; if (decoded[i].label.contains('180') && decoded[i].score > args.clsThresh) { working[originalIndex] = img.copyRotate(working[originalIndex], angle: 180); } } } return (images: working, results: clsRes); } Float32List _resizeNormImg(img.Image image) { final imgC = args.clsImageShape[0]; final imgH = args.clsImageShape[1]; final imgW = args.clsImageShape[2]; final ratio = image.width / image.height; final resizedW = math.min(imgW, (imgH * ratio).ceil()); final resized = img.copyResize(image, width: resizedW, height: imgH); final output = Float32List(imgC * imgH * imgW); for (int c = 0; c < imgC; c++) { for (int y = 0; y < imgH; y++) { for (int x = 0; x < imgW; x++) { final outIndex = c * imgH * imgW + y * imgW + x; if (x >= resizedW) { output[outIndex] = 0; continue; } final pixel = resized.getPixel(x, y); final value = switch (c) { 0 => pixel.r, 1 => pixel.g, _ => pixel.b, }; output[outIndex] = ((value / 255.0) - 0.5) / 0.5; } } } return output; } } class PPOCRv5 { static final PPOCRArgs args = PPOCRArgs(); static late final OnnxRuntime runtime; static late final _TextDetector _detector; static late final _TextRecognizer _recognizer; static _TextClassifier? _classifier; static bool _initialized = false; static Future<void> init() async { if (_initialized) return; runtime = OnnxRuntime(); _detector = _TextDetector(args); _recognizer = _TextRecognizer(args); await _detector.init(runtime); await _recognizer.init(runtime); if (args.useAngleCls) { _classifier = _TextClassifier(args); await _classifier!.init(runtime); } _initialized = true; } static Future<List<Map<String, dynamic>>> recognize( img.Image image, { bool det = true, bool rec = true, bool cls = true, }) async { await init(); if (det && rec) { final dtBoxes = await _detector.detect(image); if (dtBoxes.isEmpty) { return []; } final sorted = _sortedBoxes(dtBoxes); var crops = sorted.map((box) => _getRotateCropImage(image, box)).toList(); if (args.useAngleCls && cls && _classifier != null) { final clsOut = await _classifier!.classify(crops); crops = clsOut.images; } final recRes = await _recognizer.recognize(crops); final ocrRes = <Map<String, dynamic>>[]; for (int i = 0; i < sorted.length; i++) { final res = recRes[i]; if (res.score >= args.dropScore) { ocrRes.add({ 'box': sorted[i] .map((p) => [p[0].round(), p[1].round()]) .toList(), 'text': res.text, 'score': res.score, }); } } return ocrRes; } if (det && !rec) { final dtBoxes = await _detector.detect(image); return dtBoxes .map( (box) => { 'box': box.map((p) => [p[0].round(), p[1].round()]).toList(), }, ) .toList(); } var images = [image]; if (args.useAngleCls && cls && _classifier != null) { final clsOut = await _classifier!.classify(images); images = clsOut.images; } final recRes = await _recognizer.recognize(images); return recRes .map( (res) => { 'text': res.text, 'score': res.score, }, ) .toList(); } static Future<List<Map<String, dynamic>>> fromAsset(String assetPath) async { final data = await rootBundle.load(assetPath); final image = img.decodeImage(data.buffer.asUint8List()); if (image == null) { throw Exception('图像解码失败: $assetPath'); } return recognize(image); } static List<List<List<double>>> _sortedBoxes(List<List<List<double>>> dtBoxes) { final boxes = [...dtBoxes] ..sort((a, b) { final byY = a[0][1].compareTo(b[0][1]); if (byY != 0) return byY; return a[0][0].compareTo(b[0][0]); }); for (int i = 0; i < boxes.length - 1; i++) { for (int j = i; j >= 0; j--) { if ((boxes[j + 1][0][1] - boxes[j][0][1]).abs() < 10 && boxes[j + 1][0][0] < boxes[j][0][0]) { final tmp = boxes[j]; boxes[j] = boxes[j + 1]; boxes[j + 1] = tmp; } else { break; } } } return boxes; } static img.Image _getRotateCropImage(img.Image image, List<List<double>> points) { if (points.length != 4) { throw ArgumentError('shape of points must be 4*2'); } final cropWidth = math.max( _distance(points[0], points[1]), _distance(points[2], points[3]), ).toInt(); final cropHeight = math.max( _distance(points[0], points[3]), _distance(points[1], points[2]), ).toInt(); final safeWidth = math.max(cropWidth, 1); final safeHeight = math.max(cropHeight, 1); final src = cv.VecPoint2f.fromList( points.map((p) => cv.Point2f(p[0], p[1])).toList(), ); final dst = cv.VecPoint2f.fromList([ cv.Point2f(0, 0), cv.Point2f(safeWidth.toDouble(), 0), cv.Point2f(safeWidth.toDouble(), safeHeight.toDouble()), cv.Point2f(0, safeHeight.toDouble()), ]); final matrix = cv.getPerspectiveTransform2f(src, dst); var warped = cv.warpPerspective( _ImageConverters.imageToMat(image), matrix, (safeWidth, safeHeight), flags: cv.INTER_CUBIC, borderMode: cv.BORDER_REPLICATE, ); if (warped.rows / warped.cols >= 1.5) { warped = warped.rotate(cv.ROTATE_90_COUNTERCLOCKWISE); } return _ImageConverters.matToImage(warped); } static double _distance(List<double> a, List<double> b) { final dx = a[0] - b[0]; final dy = a[1] - b[1]; return math.sqrt(dx * dx + dy * dy); } }
https://github.com/microsoft/onnxruntime/archive/refs/tags/v1.24.1.zipmacos_x64平台的程序export HTTP_PROXY=http://127.0.0.1:6789export HTTPS_PROXY=http://127.0.0.1:6789rm -rf build/macos_x64./build.sh \--osx_arch x86_64 \--config Release \--skip_tests \--skip_submodule_sync \--build_shared_lib \--build_dir build/macos_x64 \--compile_no_warning_as_error \--parallel 8 \--cmake_extra_defines FETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVERiosrm -rf build/iosexport HTTP_PROXY=http://127.0.0.1:6789export HTTPS_PROXY=http://127.0.0.1:6789ONNX Runtime 1.24.1 iOS 终极完整命令./build.sh \--ios \--cmake_generator Xcode \--apple_sysroot $(xcrun --sdk iphoneos --show-sdk-path) \--apple_deploy_target 13.0 \--config Release \--build_shared_lib \--skip_tests \--skip_submodule_sync \--build_dir build/ios \--cmake_extra_defines FETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVER \--parallel 8 \--disable_contrib_ops \--disable_ml_ops \--compile_no_warning_as_errorandroid./build.sh \--android \--android_abi arm64-v8a \--android_api 21 \--android_sdk_path /Users/chenhuan/Library/Android/sdk \--android_ndk_path /Users/chenhuan/Library/Android/sdk/ndk/28.2.13676358 \--config Release \--skip_tests \--skip_submodule_sync \--build_shared_lib \--build_dir build/android \--compile_no_warning_as_error \--parallel 8 \--disable_contrib_ops \--disable_ml_ops \--cmake_extra_defines FETCHCONTENT_TRY_FIND_PACKAGE_MODE=NEVER
控制器var captchaRequest request.VerifyCodeRequest if err := ctx.ShouldBindWith(&captchaRequest, binding.JSON); err != nil { helper.Info(err.Error()) ReturnError(ctx, defined.ERROR_500) return } var jsonStr, err = service.CaptchaAESDecrypt(captchaRequest.Ciphertext) if err != nil { helper.Error(err) ReturnSuccess(ctx, nil) return } // 1. 初始化 RawData 变量 var rawData service.CaptchaRawData // 2. 解析 JSON 字符串到结构体 // Unmarshal 要求 JSON 字段类型与结构体字段类型匹配（如数字 vs int64，字符串 vs string） err = json.Unmarshal([]byte(jsonStr), &rawData) if err != nil { ReturnError(ctx, defined.ERROR_500) return } service.CaptchaConsole(rawData, rawData.Env.SessionId, ctx) ReturnSuccess(ctx, rawData) 服务package service import ( "activity/cache" "activity/dao" "activity/etc" "activity/helper" "activity/model" "activity/third/idc" "activity/third/ip2region" "activity/utils" "crypto/aes" "crypto/cipher" "encoding/hex" "errors" "fmt" "github.com/gin-gonic/gin" "github.com/google/uuid" log "github.com/sirupsen/logrus" "gonum.org/v1/gonum/stat" "math" ) // 定义常量：统一阈值，便于维护 const ( // 通用规则阈值（无类型差异的基础规则） DeviceIPBindCountMaxNormal = 2 // 正常IP绑定数上限 DeviceIPBindCountMinRisk = 3 // 风险IP绑定数下限 DeviceIPBindCountMaxRisk = 4 // 风险IP绑定数上限 DeviceIPBindCountMinAbnormal = 5 // 异常IP绑定数下限 IPVerifyFreqMinRisk = 10 // 风险验证频率下限 IPVerifyFreqMaxRisk = 19 // 风险验证频率上限 IPVerifyFreqMinAbnormal = 20 // 异常验证频率下限 // ------------ 按验证类型划分的 load_to_first_delay 阈值（ms）------------ // 滑动（swipe）：操作简单，延迟最短 SwipeDelayMinRisk = 80 // 滑动-风险延迟下限（略短） SwipeDelayMinNormal = 200 // 滑动-正常延迟下限 SwipeDelayMaxNormal = 2000 // 滑动-正常延迟上限 SwipeDelayMaxRisk = 3000 // 滑动-风险延迟上限（略长） SwipeDelayMinAbnormal = 50 // 滑动-异常延迟下限（过短） SwipeDelayMaxAbnormal = 4000 // 滑动-异常延迟上限（过长） // 文字点选（click）：需识别目标，延迟中等 ClickDelayMinRisk = 100 // 点选-风险延迟下限 ClickDelayMinNormal = 300 // 点选-正常延迟下限 ClickDelayMaxNormal = 7000 // 点选-正常延迟上限 ClickDelayMaxRisk = 8000 // 点选-风险延迟上限 ClickDelayMinAbnormal = 70 // 点选-异常延迟下限 ClickDelayMaxAbnormal = 9000 // 点选-异常延迟上限 // 文字输入（input）：需阅读+输入，延迟最长 InputDelayMinRisk = 200 // 输入-风险延迟下限 InputDelayMinNormal = 500 // 输入-正常延迟下限 InputDelayMaxNormal = 5000 // 输入-正常延迟上限 InputDelayMaxRisk = 7000 // 输入-风险延迟上限 InputDelayMinAbnormal = 100 // 输入-异常延迟下限 InputDelayMaxAbnormal = 10000 // 输入-异常延迟上限（允许更长犹豫时间） // 其他类型（drag等）：默认中等延迟 DefaultDelayMinRisk = 150 // 默认-风险延迟下限 DefaultDelayMinNormal = 350 // 默认-正常延迟下限 DefaultDelayMaxNormal = 3000 // 默认-正常延迟上限 DefaultDelayMaxRisk = 4000 // 默认-风险延迟上限 DefaultDelayMinAbnormal = 80 // 默认-异常延迟下限 DefaultDelayMaxAbnormal = 5000 // 默认-异常延迟上限 // Click类型阈值（其他原有阈值保持不变） ClickCVLowAbnormal = 0.08 // 点击CV过低（异常） ClickCVHighAbnormal = 0.85 // 点击CV过高（异常） ClickCVNormalMin = 0.09 // 点击CV正常下限 // Swipe类型阈值（单位：ms，统一数值与注释） SwipeIntervalMinAbnormal = 2 // 滑动间隔过低（异常） SwipeIntervalMaxAbnormal = 12 // 滑动间隔过高（异常） SwipeIntervalMinRiskFast = 2 // 滑动略快下限（风险） SwipeIntervalMaxRiskFast = 3 // 滑动略快上限（风险） SwipeIntervalMinRiskSlow = 10 // 滑动略慢下限（风险） SwipeIntervalMaxRiskSlow = 12 // 滑动略慢上限（风险） SwipeIntervalMinNormal = 3 // 滑动正常下限 SwipeIntervalMaxNormal = 10 // 滑动正常上限 // Input类型阈值（单位：ms） InputIntervalMinAbnormal = 150 // 输入间隔过低（异常） InputIntervalMaxAbnormal = 2500 // 输入间隔过高（异常） InputIntervalMinRiskFast = 150 // 输入略快下限（风险） InputIntervalMaxRiskFast = 400 // 输入略快上限（风险） InputIntervalMinRiskSlow = 2000 // 输入略慢下限（风险） InputIntervalMaxRiskSlow = 2500 // 输入略慢上限（风险） InputIntervalMinNormal = 400 // 输入正常下限 InputIntervalMaxNormal = 2000 // 输入正常上限 ) func CaptchaConsole(raw CaptchaRawData, sessionId string, req *gin.Context) { if !validateTimestamps(raw) { helper.Warning("验证码校验数据异常") return } var feature, err = ExtractFeatures(&raw, sessionId, req) if err != nil { helper.Error("验证码收集错误", err.Error()) return } swipeData := map[string]interface{}{ "verify_type": feature.VerifyType, "os": feature.OS, "device_type": feature.DeviceType, "ip_type": feature.IPType, "load_to_first_delay": feature.LoadToFirstDelay, "total_duration": feature.TotalDuration, "avg_swipe_interval": feature.AvgSwipeInterval, "ip_verify_freq": feature.IPVerifyFreq, "device_ip_bind_count": feature.DeviceIPBindCount, } swipeResult, err := utils.CaptchaPredict.PredictCaptcha(swipeData) if err != nil { log.Fatalf("❌ 验证预测失败: %v", err) } helper.Info("预测模型", swipeResult.AbnormalProbability, swipeResult.LabelDescription) // 3. 分数用途：日志打印/人工复核参考/缓存等 helper.Printf("样本[%s] 反作弊得分：%d，标签：%s", feature.SampleID, feature.TempScore, feature.Label) dao.GetPlatformDB().Model(&model.CaptchaVerifyFeature{}).Save(&feature) } // validateTimestamps 校验时间戳字段是否符合验证类型要求 func validateTimestamps(raw CaptchaRawData) bool { switch raw.Env.VerifyType { case "swipe", "click": // 滑动/点选验证 if len(raw.Timing.OperateTimestamps) == 0 { log.Printf("样本[%s]：滑动/点选验证缺少operateTimestamps", raw.Env.SessionId) return false } case "input", "sms", "text": // 输入类验证 if len(raw.Timing.InputTimestamps) == 0 { log.Printf("样本[%s]：输入类验证缺少inputTimestamps", raw.Env.SessionId) return false } case "none": // 无感验证 return true // 无要求 default: // 未知验证类型 log.Printf("样本[%s]：未知验证类型[%s]", raw.Env.SessionId, raw.Env.VerifyType) return false } return true } // ========== 结构体定义（对应原始数据结构） ========== // Env 原始数据中的环境信息 type Env struct { VerifyType string `json:"verifyType"` // 验证码类型 DeviceFingerprint string `json:"deviceFingerprint"` // 设备指纹 OS string `json:"os"` // 操作系统 DeviceType string `json:"deviceType"` // 设备类型 // 新增字段 Browser string `json:"browser"` // 浏览器（含版本） CookieEnabled bool `json:"cookieEnabled"` // 是否启用Cookie PluginCount int `json:"pluginCount"` // 插件数量 ScreenRes string `json:"screenRes"` // 屏幕分辨率 SessionId string `json:"sessionId"` // 会话ID UserAgent string `json:"userAgent"` // 用户代理（加密/哈希后） } // Timing 原始数据中的时间统计信息 type Timing struct { FirstOperateTime int64 `json:"firstOperateTime"` // 首次操作时间戳 LoadTime int64 `json:"loadTime"` // 加载时间戳 SubmitTime int64 `json:"submitTime"` // 提交时间戳 OperateTimestamps []int64 `json:"operateTimestamps"` // 操作时间戳（滑动/点击/拖拽） InputTimestamps []int64 `json:"inputTimestamps"` // 输入时间戳（输入型验证码） CorrectCount int `json:"correctCount"` // 输入正确次数 } // RawData 原始数据根结构 type CaptchaRawData struct { Env Env `json:"env"` Timing Timing `json:"timing"` } // Request 简化的请求结构体（按需扩展） type Request struct { IP string `json:"ip"` // IP地址（用于IP类型/频率查询） } // ========== 外部依赖函数（占位，需按业务实现） ========== // getIPType 第三方接口获取IP类型（IPv4/IPv6/代理等） func getIPType(req *gin.Context, realIP string) string { // 实际实现：调用IP解析接口（如ip2region/阿里云IP库） var tag, err = idc.Search(realIP) if err != nil { return "异常" } if tag != "" { return "IDC" } if etc.CommonConfigVar.Common.Env != "test" { if helper.IsViaUntrustedProxy(req) { return "代理" } } _, ipIfo, _ := ip2region.Search(realIP) if ipIfo.ISP == "" { return "未知" } return ipIfo.ISP } // getDeviceIPBindCount redis查询设备绑定7天的IP数量 func getDeviceIPBindCount(deviceFP string, realIP string) int { // 实际实现：查询设备指纹-IP绑定表 var key = fmt.Sprintf(cache.CAPTCHA_VERIFY_DEVICE_IP, deviceFP) cache.Cache.SAddWithExpire(key, realIP, 3600*24*7) var count, _ = cache.Cache.SCard(key) return count } // getIPVerifyFreq redis查询IP近期4小时验证频率 func getIPVerifyFreq(realIP string) int { var key = fmt.Sprintf(cache.CAPTCHA_VERIFY_IP_FREQ, realIP) count, _ := cache.Cache.IncrWithExpire(key, 3600) return int(count) } // DelayThreshold 按验证类型获取 load_to_first_delay 的阈值结构体 type DelayThreshold struct { MinAbnormal int // 异常延迟下限（低于此值为异常） MaxAbnormal int // 异常延迟上限（高于此值为异常） MinRisk int // 风险延迟下限（高于异常下限、低于正常下限） MaxRisk int // 风险延迟上限（高于正常上限、低于异常上限） MinNormal int // 正常延迟下限 MaxNormal int // 正常延迟上限 } // getDelayThresholdByType 根据验证类型获取对应的延迟阈值 func getDelayThresholdByType(verifyType string) DelayThreshold { switch verifyType { case "swipe": return DelayThreshold{ MinAbnormal: SwipeDelayMinAbnormal, MaxAbnormal: SwipeDelayMaxAbnormal, MinRisk: SwipeDelayMinRisk, MaxRisk: SwipeDelayMaxRisk, MinNormal: SwipeDelayMinNormal, MaxNormal: SwipeDelayMaxNormal, } case "click": return DelayThreshold{ MinAbnormal: ClickDelayMinAbnormal, MaxAbnormal: ClickDelayMaxAbnormal, MinRisk: ClickDelayMinRisk, MaxRisk: ClickDelayMaxRisk, MinNormal: ClickDelayMinNormal, MaxNormal: ClickDelayMaxNormal, } case "input": return DelayThreshold{ MinAbnormal: InputDelayMinAbnormal, MaxAbnormal: InputDelayMaxAbnormal, MinRisk: InputDelayMinRisk, MaxRisk: InputDelayMaxRisk, MinNormal: InputDelayMinNormal, MaxNormal: InputDelayMaxNormal, } default: // drag 等其他类型 return DelayThreshold{ MinAbnormal: DefaultDelayMinAbnormal, MaxAbnormal: DefaultDelayMaxAbnormal, MinRisk: DefaultDelayMinRisk, MaxRisk: DefaultDelayMaxRisk, MinNormal: DefaultDelayMinNormal, MaxNormal: DefaultDelayMaxNormal, } } } func autoLabel(features *model.CaptchaVerifyFeature, timing *Timing) (string, int) { // 入参校验：特征为空时，返回risk+0分 if features == nil { return "risk", 0 } // 消除未使用参数警告 _ = timing // 初始化反作弊基础得分（100分） score := 100 verifyType := features.VerifyType // 获取当前类型的专属延迟阈值 delayThreshold := getDelayThresholdByType(verifyType) // 当前样本的加载到首次操作延迟（转换为int类型，解决类型不匹配） currentDelay := int(features.LoadToFirstDelay) // 关键修改：int64 -> int // ========== 第一步：通用异常规则（触发直接返回abnormal） ========== // 规则1：加载到首次操作延迟 低于类型专属异常下限（过短，非人类操作） if currentDelay < delayThreshold.MinAbnormal { // 现在都是int类型，可直接比较 score -= 30 return "abnormal", score } // 规则2：加载到首次操作延迟 高于类型专属异常上限（过长，疑似机器卡顿时效） if currentDelay > delayThreshold.MaxAbnormal { // 类型一致 score -= 25 return "abnormal", score } // 规则3：设备绑定IP数 ≥5（多IP作弊） if features.DeviceIPBindCount >= DeviceIPBindCountMinAbnormal { score -= 20 return "abnormal", score } // 规则4：IP验证频率 ≥20（高频验证作弊） if features.IPVerifyFreq >= IPVerifyFreqMinAbnormal { score -= 15 return "abnormal", score } // 规则5：IP类型为高危类型（代理/异常等） if helper.Contains([]string{"异常", "代理", "未知", "IDC"}, features.IPType) { score -= 50 return "abnormal", score } // ========== 第二步：按验证码类型判定专属异常规则 ========== // 标记是否触发专属异常规则 hasExclusiveAbnormal := false switch verifyType { case "click": // 文字点选 // 负类规则：点击间隔变异系数超出异常区间 if features.ClickIntervalCV != -1 { if features.ClickIntervalCV < ClickCVLowAbnormal || features.ClickIntervalCV > ClickCVHighAbnormal { score -= 25 hasExclusiveAbnormal = true } else if features.ClickIntervalCV >= ClickCVLowAbnormal && features.ClickIntervalCV < ClickCVNormalMin { // 模糊规则：接近异常区间，先扣分（不直接返回） score -= 10 } } case "swipe": // 滑动验证 // 负类规则：平均滑动间隔 <2ms 或 >12ms（机械操作/异常慢滑） if features.AvgSwipeInterval != -1 && (features.AvgSwipeInterval < SwipeIntervalMinAbnormal || features.AvgSwipeInterval > SwipeIntervalMaxAbnormal) { score -= 25 hasExclusiveAbnormal = true } else if features.AvgSwipeInterval != -1 { // 模糊规则：2~3ms（略快）或 10~12ms（略慢） if (features.AvgSwipeInterval >= SwipeIntervalMinRiskFast && features.AvgSwipeInterval < SwipeIntervalMaxRiskFast) || (features.AvgSwipeInterval >= SwipeIntervalMinRiskSlow && features.AvgSwipeInterval <= SwipeIntervalMaxRiskSlow) { score -= 8 } } case "input": // 文字输入 // 负类规则：平均输入间隔 <150ms（自动填充）或 >2500ms（异常慢） if features.AvgInputInterval != -1 && (features.AvgInputInterval < InputIntervalMinAbnormal || features.AvgInputInterval > InputIntervalMaxAbnormal) { score -= 25 hasExclusiveAbnormal = true } else if features.AvgInputInterval != -1 { // 模糊规则：150~400ms（略快）或 2000~2500ms（略慢） if (features.AvgInputInterval >= InputIntervalMinRiskFast && features.AvgInputInterval < InputIntervalMaxRiskFast) || (features.AvgInputInterval >= InputIntervalMinRiskSlow && features.AvgInputInterval <= InputIntervalMaxRiskSlow) { score -= 8 } } } // 触发专属异常规则，直接返回abnormal if hasExclusiveAbnormal { return "abnormal", score } // ========== 第三步：正类规则（真实用户特征，触发返回normal） ========== // 正类基础条件：加载延迟正常 + 设备绑定IP数正常（类型一致） baseNormalCondition := currentDelay >= delayThreshold.MinNormal && currentDelay <= delayThreshold.MaxNormal && features.DeviceIPBindCount <= DeviceIPBindCountMaxNormal // 按验证码类型判定正类 switch verifyType { case "click": if baseNormalCondition && features.ClickIntervalCV != -1 && features.ClickIntervalCV > ClickCVNormalMin { score += 10 return "normal", score } case "swipe": if baseNormalCondition && features.AvgSwipeInterval != -1 && features.AvgSwipeInterval >= SwipeIntervalMinNormal && features.AvgSwipeInterval <= SwipeIntervalMaxNormal { score += 10 return "normal", score } case "input": if baseNormalCondition && features.AvgInputInterval != -1 && features.AvgInputInterval >= InputIntervalMinNormal && features.AvgInputInterval <= InputIntervalMaxNormal { score += 10 return "normal", score } default: // 其他类型（如drag）：仅判断基础条件 if baseNormalCondition { score += 10 return "normal", score } } // ========== 第四步：核心风险判定（基于类型专属延迟区间） ========== // 风险判定核心条件：满足任一即归为risk（类型一致） var isRisk bool // 条件1：延迟在「异常下限~正常下限」之间（略短） delayRiskLow := currentDelay >= delayThreshold.MinAbnormal && currentDelay < delayThreshold.MinNormal // 条件2：延迟在「正常上限~异常上限」之间（略长） delayRiskHigh := currentDelay > delayThreshold.MaxNormal && currentDelay <= delayThreshold.MaxAbnormal // 条件3：设备绑定IP数超出正常区间（3-4个） ipBindRisk := features.DeviceIPBindCount >= DeviceIPBindCountMinRisk && features.DeviceIPBindCount <= DeviceIPBindCountMaxRisk // 条件4：IP验证频率超出正常区间（10-19次） ipFreqRisk := features.IPVerifyFreq >= IPVerifyFreqMinRisk && features.IPVerifyFreq <= IPVerifyFreqMaxRisk // 满足任一风险条件，标记为risk isRisk = delayRiskLow || delayRiskHigh || ipBindRisk || ipFreqRisk // ========== 第五步：风险样本微调扣分（按延迟类型差异化扣分） ========== if isRisk { // 延迟略短：根据类型调整扣分（滑动操作对延迟敏感，扣分更多） if delayRiskLow { switch verifyType { case "swipe": score -= 8 // 滑动延迟略短，疑似脚本 case "click": score -= 6 // 点选延迟略短，中等风险 case "input": score -= 4 // 输入延迟略短，风险较低 default: score -= 5 } } // 延迟略长：输入类型允许更长犹豫，扣分更少 if delayRiskHigh { switch verifyType { case "swipe": score -= 6 // 滑动延迟略长，疑似卡顿/机器 case "click": score -= 5 // 点选延迟略长，中等风险 case "input": score -= 3 // 输入延迟略长，正常犹豫 default: score -= 4 } } // 设备绑定IP数略高：中度扣分 if ipBindRisk { score -= 8 } // IP验证频率略高：轻微扣分 if ipFreqRisk { score -= 5 } // 风险样本直接返回（兜底前的明确判定） return "risk", score } // ========== 兜底规则：未触发任何规则的边缘样本，仍归为risk ========== return "risk", score } // ========== 核心特征提取函数 ========== // ExtractFeatures 提取并加工特征 func ExtractFeatures(rawData *CaptchaRawData, sessionID string, req *gin.Context) (*model.CaptchaVerifyFeature, error) { // 1. 提取原始基础数据 env := rawData.Env timing := rawData.Timing verifyType := env.VerifyType realIP := helper.GetIp(req) helper.Info("ip:", realIP) // 2. 构建基础特征 features := &model.CaptchaVerifyFeature{ SampleID: uuid.NewString(), // 生成UUID（对应Python uuid.uuid4()） SessionID: sessionID, VerifyType: verifyType, LoadToFirstDelay: timing.FirstOperateTime - timing.LoadTime, TotalDuration: timing.SubmitTime - timing.LoadTime, DeviceFingerprint: env.DeviceFingerprint, OS: env.OS, DeviceType: env.DeviceType, IPType: getIPType(req, realIP), DeviceIPBindCount: getDeviceIPBindCount(env.DeviceFingerprint, realIP), IPVerifyFreq: getIPVerifyFreq(realIP), Ip: realIP, } // 3. 按验证码类型添加专属特征 switch verifyType { case "swipe", "click", "drag": operateTS := timing.OperateTimestamps features.OperateStepCount = len(operateTS) if len(operateTS) > 1 { // 计算操作时间间隔（对应Python np.diff） intervals := make([]float64, len(operateTS)-1) for i := 0; i < len(operateTS)-1; i++ { intervals[i] = float64(operateTS[i+1] - operateTS[i]) } switch verifyType { case "click": // 计算点击间隔变异系数（标准差/均值，除零保护） mean := stat.Mean(intervals, nil) if mean == 0 { features.ClickIntervalCV = 0 } else { std := math.Sqrt(stat.Variance(intervals, nil)) features.ClickIntervalCV = std / mean } case "swipe": // 计算平均滑动间隔（转int） mean := stat.Mean(intervals, nil) features.AvgSwipeInterval = int(mean) } } case "input": inputTS := timing.InputTimestamps features.InputCount = len(inputTS) // 计算正确率（除零保护） if len(inputTS) > 0 { features.CorrectRate = (float64(timing.CorrectCount) / float64(len(inputTS))) * 100 } else { features.CorrectRate = 0 } // 计算平均输入间隔 if len(inputTS) > 1 { intervals := make([]float64, len(inputTS)-1) for i := 0; i < len(inputTS)-1; i++ { intervals[i] = float64(inputTS[i+1] - inputTS[i]) } mean := stat.Mean(intervals, nil) features.AvgInputInterval = int(mean) } } // 4. 自动标注标签 features.Label, features.TempScore = autoLabel(features, &timing) return features, nil } // 配置参数（与前端保持完全一致） const ( appKey = "53D8DBC5DIK3436A" // 加密密钥 appIv = "KI5JL2SKE9883365" // 加密IV（不再从key截取，直接使用独立的appIv） ) // PKCS7Unpad 去除PKCS7填充（CryptoJS的Pkcs7对应标准PKCS7填充） // AES块大小固定为16字节，填充规则：填充n个字节则每个字节值为n func pKCS7Unpad(data []byte) ([]byte, error) { length := len(data) if length == 0 { return nil, errors.New("填充数据为空") } // 获取最后一个字节的填充长度 unpadLen := int(data[length-1]) // 验证填充长度合法性（必须在1~16之间） if unpadLen < 1 || unpadLen > aes.BlockSize { return nil, errors.New("无效的填充长度") } // 验证所有填充字节是否符合PKCS7规则 for i := 0; i < unpadLen; i++ { if data[length-1-i] != byte(unpadLen) { return nil, errors.New("填充数据不合法") } } // 去除填充部分，返回原始数据 return data[:length-unpadLen], nil } // AESDecrypt 解密前端加密的数据（适配最新的JS逻辑：IV使用appIv） // hexCiphertext: 前端返回的ciphertext（Hex编码字符串） // 返回值：解密后的原始数据字节、错误信息 func CaptchaAESDecrypt(hexCiphertext string) ([]byte, error) { // 1. 解码Hex字符串为密文字节（对应前端的toString(CryptoJS.enc.Hex)） cipherData, err := hex.DecodeString(hexCiphertext) if err != nil { return nil, fmt.Errorf("Hex解码失败: %w", err) } // 2. 处理密钥和IV（与最新JS逻辑严格对齐） // 密钥：appKey的UTF8字节（对应CryptoJS.enc.Utf8.parse(appKey)） key := []byte(appKey) // IV：appIv的UTF8字节（对应CryptoJS.enc.Utf8.parse(appIv)） iv := []byte(appIv) // 验证IV长度（AES-CBC要求IV必须等于块大小16字节） if len(iv) != aes.BlockSize { return nil, fmt.Errorf("IV长度非法，要求16字节，实际%d字节", len(iv)) } // 3. 创建AES密码块 block, err := aes.NewCipher(key) if err != nil { return nil, fmt.Errorf("创建AES块失败: %w", err) } // 4. 验证密文长度（必须是AES块大小的整数倍） if len(cipherData)%aes.BlockSize != 0 { return nil, errors.New("密文长度不是AES块大小的整数倍") } // 5. 创建CBC模式解密器（对应前端的CryptoJS.mode.CBC） mode := cipher.NewCBCDecrypter(block, iv) // 6. 执行解密（CBC模式解密会直接修改目标切片） plainData := make([]byte, len(cipherData)) mode.CryptBlocks(plainData, cipherData) // 7. 去除PKCS7填充（对应前端的CryptoJS.pad.Pkcs7） plainData, err = pKCS7Unpad(plainData) if err != nil { return nil, fmt.Errorf("去除填充失败: %w", err) } return plainData, nil } 训练import pandas as pd import numpy as np import warnings import matplotlib.pyplot as plt import optuna import os import pickle import json # 新增JSON库 import platform from typing import Dict, List, Optional, Tuple import time import re # 替换为XGBoost相关库 import xgboost as xgb from xgboost.callback import TrainingCallback # 导入回调基类 from sklearn.model_selection import train_test_split from sklearn.preprocessing import LabelEncoder from sklearn.metrics import classification_report, confusion_matrix, accuracy_score, roc_auc_score from imblearn.over_sampling import SMOTE # 忽略警告 warnings.filterwarnings('ignore') # ====================== 修复字体异常（核心） ====================== # 适配不同系统的中文字体，解决SimHei找不到问题 def set_matplotlib_chinese_font(): system = platform.system() try: if system == 'Darwin': # macOS plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'PingFang SC'] elif system == 'Windows': # Windows plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei'] elif system == 'Linux': # Linux plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'WenQuanYi Micro Hei'] plt.rcParams['axes.unicode_minus'] = False # 解决负号显示 print(f"✅ 字体设置完成（系统：{system}），当前字体：{plt.rcParams['font.sans-serif']}") except Exception as e: print(f"⚠️ 字体设置失败：{e}，使用默认字体") # 执行字体设置 set_matplotlib_chinese_font() # 开启Matplotlib交互式模式（实时绘图） plt.ion() # ====================== 新增：差异化配置 ====================== # 不同verify_type的差异化参数配置 VERIFY_TYPE_CONFIGS = { 'swipe': { # Swipe（滑动验证）参数：更关注时间序列特征 'model_params': { 'max_depth': 5, 'learning_rate': 0.015, 'subsample': 0.75, 'colsample_bytree': 0.75, 'reg_alpha': 0.2, 'reg_lambda': 0.2, 'min_child_weight': 4, 'gamma': 0.01, 'n_estimators': 250 }, 'core_features': ['load_to_first_delay', 'total_duration', 'avg_swipe_interval','ip_verify_freq','device_ip_bind_count'], 'smote_strategy': 0.7, # 过采样比例 'extreme_quantile': (0.02, 0.98) # 极端值过滤分位数 }, 'click': { # Click（点击验证）参数：更关注点击行为特征 'model_params': { 'max_depth': 4, 'learning_rate': 0.02, 'subsample': 0.7, 'colsample_bytree': 0.7, 'reg_alpha': 0.3, 'reg_lambda': 0.3, 'min_child_weight': 6, 'gamma': 0.03, 'n_estimators': 200 }, 'core_features': ['load_to_first_delay', 'total_duration', 'click_interval_cv','ip_verify_freq','device_ip_bind_count'], 'smote_strategy': 0.8, # 过采样比例 'extreme_quantile': (0.01, 0.99) # 极端值过滤分位数 }, 'default': { # 默认配置 'model_params': { 'max_depth': 4, 'learning_rate': 0.02, 'subsample': 0.7, 'colsample_bytree': 0.7, 'reg_alpha': 0.3, 'reg_lambda': 0.3, 'min_child_weight': 5, 'gamma': 0.02, 'n_estimators': 200 }, 'core_features': ['load_to_first_delay', 'total_duration', 'ip_verify_freq','device_ip_bind_count'], 'smote_strategy': 0.8, 'extreme_quantile': (0.01, 0.99) } } # 需要过滤的无用列（包含字符串ID、IP、时间等） USELESS_COLS = [ 'id', 'sample_id', 'session_id', 'ip', 'device_fingerprint', 'temp_score', 'created_at', 'updated_at', 'user_id', 'device_id', 'phone_number', 'email', 'cookie_id', 'session_token' ] # 全局变量：存储不同verify_type的训练过程数据 train_metrics = { 'swipe': { 'train_loss': [], 'val_loss': [], 'train_auc': [], 'val_auc': [], 'fig': None, 'ax1': None, 'ax2': None, 'ax3': None, 'ax4': None, 'lines': {}, 'core_feature_importance': {}, 'importance_iterations': [] }, 'click': { 'train_loss': [], 'val_loss': [], 'train_auc': [], 'val_auc': [], 'fig': None, 'ax1': None, 'ax2': None, 'ax3': None, 'ax4': None, 'lines': {}, 'core_feature_importance': {}, 'importance_iterations': [] } } # 全局变量：存储不同verify_type的预处理参数 preprocessing_params = { 'swipe': {}, 'click': {}, 'default': {} } # 初始化各verify_type的核心特征重要性存储 for vt in ['swipe', 'click']: config = VERIFY_TYPE_CONFIGS.get(vt, VERIFY_TYPE_CONFIGS['default']) core_features = config['core_features'] for feat in core_features: train_metrics[vt]['core_feature_importance'][feat] = [] # ====================== 新增：XGBoost模型转JSON工具函数 ====================== def convert_xgb_model_to_json(model: xgb.Booster, output_path: str): """ 将XGBoost模型转换为标准JSON格式（兼容Go解析） Args: model: 训练好的XGBoost Booster模型 output_path: JSON文件输出路径 """ try: # 1. 导出模型为JSON字符串（XGBoost原生JSON格式） model_json_str = model.save_raw("json") # 2. 解析为Python字典（确保格式规范） model_json = json.loads(model_json_str) # 3. 保存为格式化的JSON文件 with open(output_path, 'w', encoding='utf-8') as f: json.dump(model_json, f, indent=4, ensure_ascii=False) print(f"✅ XGBoost模型已转换为JSON格式：{output_path}") return True except Exception as e: print(f"❌ XGBoost模型转JSON失败：{e}") return False # ====================== 工具函数 ====================== def is_ip_address(s): """判断字符串是否为IP地址""" if not isinstance(s, str): return False ip_pattern = r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$' return re.match(ip_pattern, s) is not None def is_numeric_column(series): """判断列是否为数值类型（可以安全转换为数字）""" # 先移除空值 non_null_series = series.dropna() if len(non_null_series) == 0: return False # 检查是否包含IP地址 if non_null_series.apply(lambda x: is_ip_address(x) if isinstance(x, str) else False).any(): return False # 尝试转换为数值 try: pd.to_numeric(non_null_series, errors='raise') return True except (ValueError, TypeError): return False def get_verify_type_config(verify_type: str) -> dict: """获取指定verify_type的配置""" return VERIFY_TYPE_CONFIGS.get(verify_type, VERIFY_TYPE_CONFIGS['default']) def init_realtime_plot(verify_type: str): """初始化指定verify_type的实时绘图画布""" if train_metrics[verify_type]['fig'] is None: # 获取该类型的核心特征 config = get_verify_type_config(verify_type) core_features = config['core_features'] # 调整布局 fig = plt.figure(figsize=(15, 10)) fig.suptitle(f'模型训练监控 - Verify Type: {verify_type.upper()}', fontsize=16) # 子图1：LogLoss ax1 = plt.subplot(2, 2, 1) ax1.set_xlabel('迭代次数') ax1.set_ylabel('LogLoss') ax1.set_title('训练/验证集LogLoss（越小越好）') ax1.grid(True, alpha=0.3) line1, = ax1.plot([], [], label='训练集', color='blue', linewidth=1.5) line2, = ax1.plot([], [], label='验证集', color='red', linewidth=1.5) ax1.legend() # 子图2：AUC ax2 = plt.subplot(2, 2, 2) ax2.set_xlabel('迭代次数') ax2.set_ylabel('AUC') ax2.set_title('训练/验证集AUC（越大越好）') ax2.grid(True, alpha=0.3) line3, = ax2.plot([], [], label='训练集', color='blue', linewidth=1.5) line4, = ax2.plot([], [], label='验证集', color='red', linewidth=1.5) ax2.legend() # 子图3：核心特征重要性趋势 ax3 = plt.subplot(2, 2, 3) ax3.set_xlabel('迭代次数') ax3.set_ylabel('特征重要性（Gain）') ax3.set_title(f'核心特征重要性趋势（{"/".join(core_features)}）') ax3.grid(True, alpha=0.3) # 为每个核心特征创建线条 lines = {} colors = ['green', 'orange', 'purple', 'brown'] for i, feat in enumerate(core_features): line, = ax3.plot([], [], label=feat, color=colors[i % len(colors)], linewidth=1.5) lines[feat] = line ax3.legend() # 子图4：核心特征分布对比 ax4 = plt.subplot(2, 2, 4) ax4.set_title('核心特征分布（正常vs异常）') ax4.set_visible(False) # 保存到全局变量 train_metrics[verify_type]['fig'] = fig train_metrics[verify_type]['ax1'] = ax1 train_metrics[verify_type]['ax2'] = ax2 train_metrics[verify_type]['ax3'] = ax3 train_metrics[verify_type]['ax4'] = ax4 train_metrics[verify_type]['lines']['train_loss'] = line1 train_metrics[verify_type]['lines']['val_loss'] = line2 train_metrics[verify_type]['lines']['train_auc'] = line3 train_metrics[verify_type]['lines']['val_auc'] = line4 for feat, line in lines.items(): train_metrics[verify_type]['lines'][feat] = line def update_realtime_plot(verify_type: str, iteration, train_loss, val_loss, train_auc, val_auc, model=None): """更新指定verify_type的实时绘图数据""" metrics = train_metrics[verify_type] config = get_verify_type_config(verify_type) core_features = config['core_features'] # 添加基础指标 metrics['train_loss'].append(train_loss) metrics['val_loss'].append(val_loss) metrics['train_auc'].append(train_auc) metrics['val_auc'].append(val_auc) # 更新基础线条 x_data = list(range(1, len(metrics['train_loss']) + 1)) if len(x_data) == len(metrics['train_loss']): metrics['lines']['train_loss'].set_data(x_data, metrics['train_loss']) if len(x_data) == len(metrics['val_loss']): metrics['lines']['val_loss'].set_data(x_data, metrics['val_loss']) if len(x_data) == len(metrics['train_auc']): metrics['lines']['train_auc'].set_data(x_data, metrics['train_auc']) if len(x_data) == len(metrics['val_auc']): metrics['lines']['val_auc'].set_data(x_data, metrics['val_auc']) # 更新核心特征重要性 if model and iteration % 10 == 0: try: importance_dict = model.get_score(importance_type='gain') metrics['importance_iterations'].append(iteration) for feat in core_features: imp_value = importance_dict.get(feat, 0) metrics['core_feature_importance'][feat].append(imp_value) # 绘制核心特征重要性 imp_iter = metrics['importance_iterations'] for feat in core_features: imp_vals = metrics['core_feature_importance'][feat] if len(imp_iter) == len(imp_vals) and len(imp_iter) > 0: metrics['lines'][feat].set_data(imp_iter, imp_vals) metrics['ax3'].relim() metrics['ax3'].autoscale_view() except Exception as e: print(f"\n⚠️ {verify_type} 特征重要性绘图更新失败：{e}") # 自动调整坐标轴 metrics['ax1'].relim() metrics['ax1'].autoscale_view() metrics['ax2'].relim() metrics['ax2'].autoscale_view() # 刷新画布 try: metrics['fig'].canvas.draw() metrics['fig'].canvas.flush_events() except Exception as e: print(f"\n⚠️ {verify_type} 绘图刷新失败：{e}") time.sleep(0.01) # ====================== 差异化回调类 ====================== class XGBoostRealtimePlotCallback(TrainingCallback): """差异化的XGBoost实时绘图回调类""" def __init__(self, verify_type: str): super().__init__() self.iteration = 0 self.verify_type = verify_type def after_iteration(self, model, epoch, evals_log): """每轮迭代后执行的回调方法""" self.iteration = epoch try: # 提取训练和验证指标 train_loss = val_loss = train_auc = val_auc = None if 'train' in evals_log and 'logloss' in evals_log['train']: train_loss = evals_log['train']['logloss'][-1] if 'eval' in evals_log and 'logloss' in evals_log['eval']: val_loss = evals_log['eval']['logloss'][-1] if 'train' in evals_log and 'auc' in evals_log['train']: train_auc = evals_log['train']['auc'][-1] if 'eval' in evals_log and 'auc' in evals_log['eval']: val_auc = evals_log['eval']['auc'][-1] # 更新绘图 if all([train_loss, val_loss, train_auc, val_auc]): update_realtime_plot(self.verify_type, self.iteration, train_loss, val_loss, train_auc, val_auc, model) except Exception as e: print(f"\n⚠️ {self.verify_type} 回调函数执行失败（迭代{epoch}）：{e}") return False # ====================== 数据处理函数（修复版本） ====================== def load_and_validate_data(csv_file: str) -> pd.DataFrame: """加载并校验验证码行为数据""" try: df = pd.read_csv(csv_file, encoding='utf-8') print(f"✅ 数据加载成功，形状: {df.shape}") print(f"✅ verify_type分布：\n{df['verify_type'].value_counts()}") # 检查关键列 required_cols = [ 'label', 'verify_type', 'load_to_first_delay', 'total_duration' ] missing_cols = [col for col in required_cols if col not in df.columns] if missing_cols: raise ValueError(f"❌ 缺少关键列：{missing_cols}") # 过滤无效的verify_type valid_verify_types = ['swipe', 'click'] df = df[df['verify_type'].isin(valid_verify_types)] print(f"✅ 过滤后有效数据量：{len(df)}（仅保留swipe/click）") # 提前过滤无用列（修复核心问题） df = filter_useless_columns(df) return df except Exception as e: raise Exception(f"❌ 加载数据失败：{str(e)}") def filter_useless_columns(df: pd.DataFrame) -> pd.DataFrame: """过滤无用列（包含字符串ID、IP、时间等）""" df_filtered = df.copy() # 找出所有需要删除的列 cols_to_drop = [col for col in USELESS_COLS if col in df_filtered.columns] if cols_to_drop: df_filtered = df_filtered.drop(columns=cols_to_drop) print(f"\n🗑️ 已删除无用列：{cols_to_drop}") # 额外检查：识别并移除包含IP地址的列 ip_columns = [] for col in df_filtered.columns: # 检查列名是否包含ip相关关键词 if 'ip' in col.lower() and col != 'ip_type' and col != 'ip_verify_freq' and col != 'device_ip_bind_count': # 抽样检查列内容是否包含IP地址 sample_data = df_filtered[col].dropna().head(10) if sample_data.apply(lambda x: is_ip_address(x) if isinstance(x, str) else False).any(): ip_columns.append(col) if ip_columns: df_filtered = df_filtered.drop(columns=ip_columns) print(f"🗑️ 已删除包含IP地址的列：{ip_columns}") return df_filtered def filter_extreme_samples(df: pd.DataFrame, verify_type: str, cols: list) -> pd.DataFrame: """差异化过滤极端值样本""" config = get_verify_type_config(verify_type) lower_q, upper_q = config['extreme_quantile'] df_filtered = df.copy() total_samples = len(df_filtered) # 初始化该类型的预处理参数 preprocessing_params[verify_type] = {} for col in cols: if col not in df_filtered.columns: print(f"⚠️ {verify_type} 列{col}不存在，跳过过滤") continue # 确保列是数值类型 if not is_numeric_column(df_filtered[col]): print(f"⚠️ {verify_type} 列{col}不是数值类型，跳过过滤") continue q_low = df_filtered[col].quantile(lower_q) q_high = df_filtered[col].quantile(upper_q) # 保存分位数参数 preprocessing_params[verify_type][col] = { 'q_low': q_low, 'q_high': q_high, 'mean': df_filtered[col].mean(), 'std': df_filtered[col].std(), 'median': df_filtered[col].median() # 新增中位数保存 } # 过滤条件 filter_condition = (df_filtered[col] >= q_low) & (df_filtered[col] <= q_high) df_filtered = df_filtered[filter_condition] filtered_count = total_samples - len(df_filtered) total_samples = len(df_filtered) print(f"\n🔧 [{verify_type}] {col} 极端值过滤：") print(f" 分位数范围：{q_low:.2f} ~ {q_high:.2f}") print(f" 过滤样本数：{filtered_count}，剩余样本数：{total_samples}") return df_filtered def process_labels(df: pd.DataFrame) -> pd.DataFrame: """处理验证码标签""" print(f"\n📊 原始标签分布：") label_counts = df['label'].value_counts() print(label_counts) label_mapping = {'normal': 0, 'risk': 1, 'abnormal': 1} df['binary_label'] = df['label'].map(lambda x: label_mapping.get(x, 1)) print(f"\n📊 转换后的二进制标签分布：") print(df['binary_label'].value_counts()) return df def select_features(df: pd.DataFrame, verify_type: str) -> Tuple[List[str], List[str], List[str]]: """差异化选择特征（修复版本）""" config = get_verify_type_config(verify_type) core_features = config['core_features'] # 排除标签列和verify_type feature_cols = [col for col in df.columns if col not in ['label', 'binary_label', 'verify_type']] # 自动识别分类特征和数值特征 categorical_cols = [] numeric_cols = [] for col in feature_cols: # 检查是否为数值列 if is_numeric_column(df[col]): numeric_cols.append(col) else: categorical_cols.append(col) print(f"\n🔧 [{verify_type}] 自动识别特征类型：") print(f" 分类特征: {categorical_cols}") # 标记核心特征 marked_numeric = [f"⭐{feat}" if feat in core_features else feat for feat in numeric_cols] print(f" 数值特征: {marked_numeric}") # 检查核心特征是否存在 missing_core = [feat for feat in core_features if feat not in numeric_cols] if missing_core: print(f"⚠️ [{verify_type}] 核心特征缺失：{missing_core}") return feature_cols, categorical_cols, numeric_cols def fix_numeric_outliers(df: pd.DataFrame, verify_type: str, numeric_cols: list) -> pd.DataFrame: """差异化修复数值特征异常值""" config = get_verify_type_config(verify_type) core_features = config['core_features'] df_fixed = df.copy() # 处理核心特征 for col in core_features: if col in numeric_cols and col in df_fixed.columns: # 确保列是数值类型 df_fixed[col] = pd.to_numeric(df_fixed[col], errors='coerce') # 负数替换为0 neg_count = (df_fixed[col] < 0).sum() if neg_count > 0: df_fixed.loc[df_fixed[col] < 0, col] = 0 print(f"\n🔧 [{verify_type}] 修复{col}负数：{neg_count} 个值") # 处理其他数值特征 other_numeric = [col for col in numeric_cols if col not in core_features and col in df_fixed.columns] for col in other_numeric: # 确保列是数值类型 df_fixed[col] = pd.to_numeric(df_fixed[col], errors='coerce') Q1 = df_fixed[col].quantile(0.25) Q3 = df_fixed[col].quantile(0.75) IQR = Q3 - Q1 lower_bound = Q1 - 1.5 * IQR upper_bound = Q3 + 1.5 * IQR outlier_mask = (df_fixed[col] < lower_bound) | (df_fixed[col] > upper_bound) outlier_count = outlier_mask.sum() if outlier_count > 0: df_fixed.loc[outlier_mask, col] = df_fixed[col].clip(lower_bound, upper_bound) print(f"🔧 [{verify_type}] 修复{col}异常值：{outlier_count} 个值") return df_fixed def handle_missing_values(df: pd.DataFrame, verify_type: str, numeric_cols: list, categorical_cols: list) -> pd.DataFrame: """差异化填充缺失值（修复版本）""" config = get_verify_type_config(verify_type) core_features = config['core_features'] df_filled = df.copy() # 记录数值特征的中位数（用于Go端填充） numeric_medians = {} # 记录分类特征的众数（用于Go端填充） categorical_modes = {} # 填充核心特征 for col in core_features: if col in numeric_cols and col in df_filled.columns: # 确保列是数值类型 df_filled[col] = pd.to_numeric(df_filled[col], errors='coerce') numeric_medians[col] = df_filled[col].median() missing_count = df_filled[col].isnull().sum() if missing_count > 0: if 'delay' in col or 'duration' in col: fill_val = df_filled[col].median() df_filled[col] = df_filled[col].fillna(fill_val) print(f"\n🔧 [{verify_type}] 填充{col}缺失值：{missing_count} 个值，中位数={fill_val:.2f}") else: fill_val = 0 if 'freq' in col else df_filled[col].median() df_filled[col] = df_filled[col].fillna(fill_val) print(f"\n🔧 [{verify_type}] 填充{col}缺失值：{missing_count} 个值，默认值={fill_val:.2f}") # 填充其他数值特征 for col in [c for c in numeric_cols if c not in core_features and c in df_filled.columns]: df_filled[col] = pd.to_numeric(df_filled[col], errors='coerce') numeric_medians[col] = df_filled[col].median() missing_count = df_filled[col].isnull().sum() if missing_count > 0: fill_val = df_filled[col].median() df_filled[col] = df_filled[col].fillna(fill_val) # 填充分类特征 for col in categorical_cols: if col in df_filled.columns: if df_filled[col].isnull().sum() > 0: mode_val = df_filled[col].mode()[0] categorical_modes[col] = mode_val df_filled[col] = df_filled[col].fillna(mode_val) else: categorical_modes[col] = 'Unknown' # 保存中位数和众数到预处理参数 preprocessing_params[verify_type]['numeric_medians'] = numeric_medians preprocessing_params[verify_type]['categorical_modes'] = categorical_modes return df_filled def encode_categorical_features(df: pd.DataFrame, categorical_cols: list) -> Tuple[ pd.DataFrame, Dict[str, LabelEncoder]]: """编码分类特征""" label_encoders = {} df_encoded = df.copy() for col in categorical_cols: if col in df_encoded.columns: le = LabelEncoder() # 填充未知值为'Unknown' df_encoded[col] = df_encoded[col].fillna('Unknown').astype(str) # 确保'Unknown'在类别中 if 'Unknown' not in df_encoded[col].unique(): df_encoded[col] = df_encoded[col].cat.add_categories('Unknown') if df_encoded[ col].dtype.name == 'category' else \ df_encoded[col] df_encoded[col] = le.fit_transform(df_encoded[col]) label_encoders[col] = le print(f"\n🔤 编码 {col}：原始值={le.classes_[:10]}...（显示前10个）") return df_encoded, label_encoders def balance_samples(X: pd.DataFrame, y: pd.Series, verify_type: str) -> Tuple[pd.DataFrame, pd.Series]: """差异化样本平衡""" config = get_verify_type_config(verify_type) smote_strategy = config['smote_strategy'] if y.nunique() < 2: print(f"\n⚠️ [{verify_type}] 标签只有1类，跳过SMOTE采样") return X, y pos_ratio = y.sum() / len(y) print(f"\n📊 [{verify_type}] 原始样本分布：异常样本占比 {pos_ratio:.2%}") if pos_ratio < smote_strategy: try: smote = SMOTE(random_state=42, sampling_strategy=smote_strategy) X_resampled, y_resampled = smote.fit_resample(X, y) new_pos_ratio = y_resampled.sum() / len(y_resampled) print(f"✅ [{verify_type}] 采样后异常样本占比 {new_pos_ratio:.2%}，样本量 {len(X_resampled)}") return X_resampled, y_resampled except Exception as e: print(f"❌ [{verify_type}] SMOTE采样失败：{e}") return X, y def create_core_feature_interactions(df: pd.DataFrame, verify_type: str) -> pd.DataFrame: """差异化创建核心特征交互项""" config = get_verify_type_config(verify_type) core_features = config['core_features'] df_interact = df.copy() # 根据不同verify_type创建差异化的交互特征 if verify_type == 'swipe': # 滑动验证：关注滑动相关交互 if all(feat in df_interact.columns for feat in ['load_to_first_delay', 'avg_swipe_interval']): df_interact['load_swipe_ratio'] = df_interact['load_to_first_delay'] / ( df_interact['avg_swipe_interval'] + 1e-6) elif verify_type == 'click': # 点击验证：关注点击相关交互 if all(feat in df_interact.columns for feat in ['total_duration', 'click_interval_cv']): df_interact['duration_click_cv'] = df_interact['total_duration'] * df_interact['click_interval_cv'] # 通用交互特征 if all(feat in df_interact.columns for feat in ['load_to_first_delay', 'total_duration']): df_interact['load_to_total_ratio'] = df_interact['load_to_first_delay'] / (df_interact['total_duration'] + 1e-6) print(f"\n🔧 [{verify_type}] 创建核心特征交互项完成") return df_interact # ====================== 模型训练函数（差异化版本） ====================== def tune_hyperparameters(X_train, y_train, X_val, y_val, verify_type: str): """差异化调参""" config = get_verify_type_config(verify_type) base_params = config['model_params'] def objective(trial): params = { 'booster': 'gbtree', 'objective': 'binary:logistic', 'eval_metric': ['logloss', 'auc'], 'max_depth': trial.suggest_int('max_depth', base_params['max_depth'] - 1, base_params['max_depth'] + 1), 'learning_rate': trial.suggest_float('learning_rate', base_params['learning_rate'] * 0.8, base_params['learning_rate'] * 1.2, log=True), 'subsample': trial.suggest_float('subsample', base_params['subsample'] - 0.05, base_params['subsample'] + 0.05), 'colsample_bytree': trial.suggest_float('colsample_bytree', base_params['colsample_bytree'] - 0.05, base_params['colsample_bytree'] + 0.05), 'reg_alpha': trial.suggest_float('reg_alpha', base_params['reg_alpha'] * 0.8, base_params['reg_alpha'] * 1.2), 'reg_lambda': trial.suggest_float('reg_lambda', base_params['reg_lambda'] * 0.8, base_params['reg_lambda'] * 1.2), 'random_state': 42, 'verbosity': 0 } dtrain = xgb.DMatrix(X_train, label=y_train) dval = xgb.DMatrix(X_val, label=y_val) model = xgb.train( params, dtrain, num_boost_round=100, evals=[(dval, 'eval')], early_stopping_rounds=15, verbose_eval=False ) return model.best_score study = optuna.create_study(direction='maximize', study_name=f'{verify_type}_xgb_tune') study.optimize(objective, n_trials=5) print(f"\n🏆 [{verify_type}] 最优参数：{study.best_params}") print(f"🏆 [{verify_type}] 最优验证AUC：{study.best_value:.4f}") return study.best_params def train_xgboost_model(X_train, y_train, X_test, y_test, verify_type: str, use_optuna: bool = True): """差异化训练XGBoost模型""" # 初始化绘图 init_realtime_plot(verify_type) # 获取差异化配置 config = get_verify_type_config(verify_type) model_params = config['model_params'] # 基础参数 base_params = { 'booster': 'gbtree', 'objective': 'binary:logistic', 'eval_metric': ['logloss', 'auc'], 'verbosity': 1, 'random_state': 42, 'scale_pos_weight': 1, 'importance_type': 'gain', **model_params } # 调参 if use_optuna and y_train.nunique() >= 2: X_train_tune, X_val_tune, y_train_tune, y_val_tune = train_test_split( X_train, y_train, test_size=0.2, random_state=42, stratify=y_train ) best_params = tune_hyperparameters(X_train_tune, y_train_tune, X_val_tune, y_val_tune, verify_type) base_params = {**base_params, **best_params} # 准备数据 dtrain = xgb.DMatrix(X_train, label=y_train) dtest = xgb.DMatrix(X_test, label=y_test) # 保存特征列 preprocessing_params[verify_type]['feature_cols'] = X_train.columns.tolist() # 创建回调 plot_callback = XGBoostRealtimePlotCallback(verify_type) # 训练模型 print(f"\n🚀 开始训练 {verify_type} 类型验证码模型...") evals_result = {} try: model = xgb.train( base_params, dtrain, num_boost_round=base_params['n_estimators'], evals=[(dtrain, 'train'), (dtest, 'eval')], evals_result=evals_result, early_stopping_rounds=20, verbose_eval=10, callbacks=[plot_callback] ) except Exception as e: print(f"\n⚠️ [{verify_type}] 训练异常：{e}") # 重试（无回调） model = xgb.train( base_params, dtrain, num_boost_round=base_params['n_estimators'], evals=[(dtrain, 'train'), (dtest, 'eval')], evals_result=evals_result, early_stopping_rounds=20, verbose_eval=10 ) # 绘制特征分布 try: if y_test.nunique() >= 2: ax4 = train_metrics[verify_type]['ax4'] ax4.clear() ax4.set_title(f'{verify_type} 核心特征分布（正常vs异常）') X_test_df = pd.DataFrame(X_test, columns=X_train.columns) normal_mask = y_test == 0 abnormal_mask = y_test == 1 core_features = config['core_features'] positions = list(range(1, len(core_features) + 1)) has_data = False for i, feat in enumerate(core_features): if feat in X_test_df.columns: normal_data = X_test_df.loc[normal_mask, feat].values abnormal_data = X_test_df.loc[abnormal_mask, feat].values if len(normal_data) > 0 and len(abnormal_data) > 0: has_data = True ax4.boxplot(normal_data, positions=[positions[i] - 0.15], widths=0.2, patch_artist=True, boxprops=dict(facecolor='green', alpha=0.5), label='正常' if i == 0 else "") ax4.boxplot(abnormal_data, positions=[positions[i] + 0.15], widths=0.2, patch_artist=True, boxprops=dict(facecolor='red', alpha=0.5), label='异常' if i == 0 else "") if has_data: ax4.set_xticks(positions) ax4.set_xticklabels(core_features) ax4.legend() ax4.grid(True, alpha=0.3) ax4.set_visible(True) except Exception as e: print(f"\n⚠️ [{verify_type}] 特征分布绘图失败：{e}") return model, evals_result, base_params def evaluate_model(model, X_test, y_test, feature_cols, verify_type: str): """差异化评估模型""" config = get_verify_type_config(verify_type) core_features = config['core_features'] dtest = xgb.DMatrix(X_test) y_pred_prob = model.predict(dtest) y_pred = (y_pred_prob > 0.5).astype(int) # 计算指标 accuracy = accuracy_score(y_test, y_pred) auc = roc_auc_score(y_test, y_pred_prob) if y_test.nunique() >= 2 else -1 print(f"\n📈 [{verify_type}] 模型评估结果：") print(f"🏆 最佳迭代次数：{model.best_iteration}") print(f"🎯 准确率：{accuracy:.4f}") print(f"📊 AUC：{auc:.4f}" if auc != -1 else "📊 AUC：无意义") # 特征重要性 importance_dict = model.get_score(importance_type='gain') all_importance = {feat: importance_dict.get(feat, 0) for feat in feature_cols} importance_df = pd.DataFrame({ 'feature': list(all_importance.keys()), 'importance': list(all_importance.values()) }).sort_values('importance', ascending=False) # 标记核心特征 importance_df['is_core'] = importance_df['feature'].apply(lambda x: '⭐' if x in core_features else '') # 核心特征贡献 total_importance = importance_df['importance'].sum() core_total = importance_df[importance_df['feature'].isin(core_features)]['importance'].sum() core_contribution = (core_total / total_importance * 100) if total_importance > 0 else 0 print(f"\n🎯 [{verify_type}] 核心特征总贡献：{core_contribution:.1f}%") print(f"\n✨ [{verify_type}] 特征重要性（前10）：") print(importance_df.head(10)'is_core', 'feature', 'importance') return { 'accuracy': accuracy, 'auc': auc, 'feature_importance': importance_df, 'core_feature_contribution': core_contribution, 'best_iteration': model.best_iteration, 'feature_importance_dict': all_importance # 新增：便于JSON序列化 } def save_results(models: dict, label_encoders: dict, feature_cols: dict, eval_results: dict, params: dict, save_dir='./captcha_results'): """保存所有模型结果（包含JSON格式）""" # 创建目录 os.makedirs(save_dir, exist_ok=True) os.makedirs('./captcha_models', exist_ok=True) # 保存每个类型的模型 for verify_type in ['swipe', 'click']: if verify_type in models: # 保存模型二进制文件 model_bin_path = f'./captcha_models/captcha_xgb_model_{verify_type}.bin' models[verify_type].save_model(model_bin_path) print(f"\n💾 {verify_type} 模型已保存到：{model_bin_path}") # ========== 新增：转换模型为JSON格式 ========== model_json_path = f'./captcha_models/captcha_xgb_model_{verify_type}.json' convert_xgb_model_to_json(models[verify_type], model_json_path) # 保存预处理参数（pickle格式） preprocess_path = f'./captcha_models/preprocessing_params_{verify_type}.pkl' with open(preprocess_path, 'wb') as f: pickle.dump(preprocessing_params[verify_type], f) # ========== 完善JSON配置文件 ========== # 1. 准备JSON序列化的数据 json_data = { 'verify_type': verify_type, 'model_config': params.get(verify_type, {}), 'preprocessing_params': preprocessing_params.get(verify_type, {}), 'evaluation_metrics': { 'accuracy': float(eval_results[verify_type]['accuracy']), 'auc': float(eval_results[verify_type]['auc']), 'core_feature_contribution': float(eval_results[verify_type]['core_feature_contribution']), 'best_iteration': int(eval_results[verify_type]['best_iteration']) }, 'feature_importance': eval_results[verify_type]['feature_importance_dict'], 'core_features': get_verify_type_config(verify_type)['core_features'], 'feature_columns': feature_cols.get(verify_type, []), 'training_timestamp': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()), 'numeric_medians': preprocessing_params[verify_type].get('numeric_medians', {}), 'categorical_modes': preprocessing_params[verify_type].get('categorical_modes', {}) } # 2. 处理LabelEncoder（只保存类别映射） le_json = {} if verify_type in label_encoders: for col, le in label_encoders[verify_type].items(): le_json[col] = { 'classes': le.classes_.tolist(), 'mapping': {str(cls): int(idx) for idx, cls in enumerate(le.classes_)}, 'num_classes': len(le.classes_) } json_data['label_encoders'] = le_json # 3. 保存JSON配置文件 json_config_path = f'./captcha_models/captcha_model_{verify_type}.json' with open(json_config_path, 'w', encoding='utf-8') as f: json.dump(json_data, f, ensure_ascii=False, indent=4) print(f"💾 {verify_type} 模型JSON配置文件已保存到：{json_config_path}") # 保存评估结果（CSV格式） eval_df = pd.DataFrame({ 'metric': ['accuracy', 'auc', 'core_feature_contribution', 'best_iteration'], 'value': [ eval_results[verify_type]['accuracy'], eval_results[verify_type]['auc'], eval_results[verify_type]['core_feature_contribution'], eval_results[verify_type]['best_iteration'] ] }) eval_df.to_csv(f'{save_dir}/evaluation_results_{verify_type}.csv', index=False) # 保存全局配置JSON文件 global_config = { 'verify_type_configs': VERIFY_TYPE_CONFIGS, 'useless_columns': USELESS_COLS, 'training_timestamp': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()) } global_config_path = './captcha_models/global_config.json' with open(global_config_path, 'w', encoding='utf-8') as f: json.dump(global_config, f, ensure_ascii=False, indent=4) print(f"💾 全局配置JSON文件已保存到：{global_config_path}") # 保存标签编码器（pickle格式） le_path = './captcha_models/label_encoders.pkl' with open(le_path, 'wb') as f: pickle.dump(label_encoders, f) print(f"\n✅ 所有模型和配置已保存完成！") # ====================== 统一预测接口（兼容JSON配置） ====================== def predict_captcha(input_data: dict): """ 统一的预测接口，自动识别verify_type并使用对应模型 """ # 获取验证类型 verify_type = input_data.get('verify_type', 'default').lower() if verify_type not in ['swipe', 'click']: verify_type = 'default' # 加载对应模型 model_path = f'./captcha_models/captcha_xgb_model_{verify_type}.bin' preprocess_path = f'./captcha_models/preprocessing_params_{verify_type}.pkl' json_config_path = f'./captcha_models/captcha_model_{verify_type}.json' if not os.path.exists(model_path): raise Exception(f"❌ 未找到 {verify_type} 类型的模型文件") # 加载模型 model = xgb.Booster() model.load_model(model_path) # 优先从JSON加载配置 if os.path.exists(json_config_path): with open(json_config_path, 'r', encoding='utf-8') as f: json_config = json.load(f) preprocess_params = json_config.get('preprocessing_params', {}) feature_cols = json_config.get('feature_columns', []) label_encoders_json = json_config.get('label_encoders', {}) else: # 兼容旧的pickle格式 with open(preprocess_path, 'rb') as f: preprocess_params = pickle.load(f) feature_cols = preprocess_params.get('feature_cols', []) le_path = './captcha_models/label_encoders.pkl' with open(le_path, 'rb') as f: label_encoders = pickle.load(f) label_encoders_json = {} # 数据预处理（与训练时一致） df = pd.DataFrame([input_data]) # 过滤无用列 df = filter_useless_columns(df) # 处理数值特征 for col in df.columns: if df[col].dtype in ['int64', 'float64'] and df[col].iloc[0] < 0: df[col] = 0 # 编码分类特征（兼容JSON格式的编码器配置） for col in ['os', 'device_type', 'ip_type']: if col in df.columns and col in label_encoders_json: le_mapping = label_encoders_json[col]['mapping'] df[col] = df[col].fillna('Unknown').astype(str) # 处理未知类别 df[col] = df[col].apply(lambda x: le_mapping.get(x, le_mapping.get('Unknown', 0))) # 只保留训练时的特征列 for col in feature_cols: if col not in df.columns: df[col] = 0 df = df[feature_cols] # 预测 dmatrix = xgb.DMatrix(df) pred_prob = model.predict(dmatrix)[0] pred_label = 1 if pred_prob > 0.5 else 0 return { 'verify_type': verify_type, 'abnormal_probability': float(pred_prob), 'predicted_label': int(pred_label), 'label_description': '异常/疑似' if pred_label == 1 else '正常', 'prediction_timestamp': time.strftime('%Y-%m-%d %H:%M:%S', time.localtime()) } # ====================== 主函数 ====================== def main(): """主函数：差异化训练不同verify_type的模型""" # 配置 CSV_FILE_PATH = './captcha_verify_feature.csv' USE_OPTUNA = False try: # 1. 加载数据 df = load_and_validate_data(CSV_FILE_PATH) # 2. 按verify_type拆分数据 swipe_df = df[df['verify_type'] == 'swipe'].copy() click_df = df[df['verify_type'] == 'click'].copy() # 存储结果的字典 models = {} eval_results = {} all_params = {} label_encoders = {} feature_cols_dict = {} # 3. 训练swipe模型 if len(swipe_df) > 0: print("\n" + "=" * 60) print("开始训练 SWIPE (滑动) 验证模型") print("=" * 60) # 数据预处理 swipe_df = process_labels(swipe_df) swipe_df = filter_extreme_samples(swipe_df, 'swipe', ['load_to_first_delay', 'total_duration']) # 特征处理 feature_cols, categorical_cols, numeric_cols = select_features(swipe_df, 'swipe') swipe_df = fix_numeric_outliers(swipe_df, 'swipe', numeric_cols) swipe_df = handle_missing_values(swipe_df, 'swipe', numeric_cols, categorical_cols) swipe_df = create_core_feature_interactions(swipe_df, 'swipe') # 更新特征列表 feature_cols = [col for col in swipe_df.columns if col not in ['label', 'binary_label', 'verify_type']] swipe_df_encoded, le = encode_categorical_features(swipe_df, categorical_cols) label_encoders['swipe'] = le # 分离特征和标签 X = swipe_df_encoded[feature_cols] y = swipe_df_encoded['binary_label'] # 平衡样本 X_balanced, y_balanced = balance_samples(X, y, 'swipe') # 划分数据集 X_train, X_test, y_train, y_test = train_test_split( X_balanced, y_balanced, test_size=0.2, random_state=42, stratify=y_balanced if y_balanced.nunique() >= 2 else None ) # 训练模型 model, evals_result, params = train_xgboost_model( X_train, y_train, X_test, y_test, 'swipe', USE_OPTUNA ) # 评估模型 eval_res = evaluate_model(model, X_test, y_test, feature_cols, 'swipe') # 保存结果 models['swipe'] = model eval_results['swipe'] = eval_res all_params['swipe'] = params feature_cols_dict['swipe'] = feature_cols # 4. 训练click模型 if len(click_df) > 0: print("\n" + "=" * 60) print("开始训练 CLICK (点击) 验证模型") print("=" * 60) # 数据预处理 click_df = process_labels(click_df) click_df = filter_extreme_samples(click_df, 'click', ['load_to_first_delay', 'total_duration']) # 特征处理 feature_cols, categorical_cols, numeric_cols = select_features(click_df, 'click') click_df = fix_numeric_outliers(click_df, 'click', numeric_cols) click_df = handle_missing_values(click_df, 'click', numeric_cols, categorical_cols) click_df = create_core_feature_interactions(click_df, 'click') # 更新特征列表 feature_cols = [col for col in click_df.columns if col not in ['label', 'binary_label', 'verify_type']] click_df_encoded, le = encode_categorical_features(click_df, categorical_cols) label_encoders['click'] = le # 分离特征和标签 X = click_df_encoded[feature_cols] y = click_df_encoded['binary_label'] # 平衡样本 X_balanced, y_balanced = balance_samples(X, y, 'click') # 划分数据集 X_train, X_test, y_train, y_test = train_test_split( X_balanced, y_balanced, test_size=0.2, random_state=42, stratify=y_balanced if y_balanced.nunique() >= 2 else None ) # 训练模型 model, evals_result, params = train_xgboost_model( X_train, y_train, X_test, y_test, 'click', USE_OPTUNA ) # 评估模型 eval_res = evaluate_model(model, X_test, y_test, feature_cols, 'click') # 保存结果 models['click'] = model eval_results['click'] = eval_res all_params['click'] = params feature_cols_dict['click'] = feature_cols # 5. 保存所有结果（包含JSON格式） save_results(models, label_encoders, feature_cols_dict, eval_results, all_params) # 6. 测试预测 print("\n" + "=" * 60) print("测试差异化预测功能") print("=" * 60) # 测试滑动验证 test_swipe_data = { 'verify_type': 'swipe', 'os': 'Android', 'device_type': 'Mobile', 'ip_type': '电信', 'load_to_first_delay': 20000, 'total_duration': 30000, 'avg_swipe_interval': 100, 'ip_verify_freq': 10 } # 测试点击验证 test_click_data = { 'verify_type': 'click', 'os': 'iOS', 'device_type': 'Mobile', 'ip_type': '联通', 'load_to_first_delay': 15000, 'total_duration': 25000, 'click_interval_cv': 0.5, 'ip_verify_freq': 8 } # 预测（仅当模型存在时） if 'swipe' in models: swipe_result = predict_captcha(test_swipe_data) print(f"\n📱 滑动验证预测结果：") print(f" 验证类型：{swipe_result['verify_type']}") print(f" 异常概率：{swipe_result['abnormal_probability']:.4f}") print(f" 预测结果：{swipe_result['label_description']}") print(f" 预测时间：{swipe_result['prediction_timestamp']}") if 'click' in models: click_result = predict_captcha(test_click_data) print(f"\n🖱️ 点击验证预测结果：") print(f" 验证类型：{click_result['verify_type']}") print(f" 异常概率：{click_result['abnormal_probability']:.4f}") print(f" 预测结果：{click_result['label_description']}") print(f" 预测时间：{click_result['prediction_timestamp']}") print("\n🎉 差异化模型训练和预测完成！") # 保持绘图窗口 plt.ioff() plt.show() except Exception as e: print(f"\n❌ 训练失败：{str(e)}") import traceback traceback.print_exc() plt.ioff() raise if __name__ == "__main__": main() 推理/**预测模型**/ var CaptchaPredict *CaptchaPredictor // ========== 1. 适配Python生成的JSON配置结构体（补充中位数/众数字段） ========== type ModelJSONConfig struct { VerifyType string `json:"verify_type"` ModelConfig map[string]interface{} `json:"model_config"` PreprocessingParams map[string]interface{} `json:"preprocessing_params"` EvaluationMetrics struct { Accuracy float64 `json:"accuracy"` Auc float64 `json:"auc"` CoreFeatureContribution float64 `json:"core_feature_contribution"` BestIteration int `json:"best_iteration"` } `json:"evaluation_metrics"` FeatureImportance map[string]float64 `json:"feature_importance"` CoreFeatures []string `json:"core_features"` FeatureColumns []string `json:"feature_columns"` TrainingTimestamp string `json:"training_timestamp"` LabelEncoders map[string]struct { Classes []string `json:"classes"` Mapping map[string]int `json:"mapping"` NumClasses int `json:"num_classes"` } `json:"label_encoders"` NumericMedians map[string]float64 `json:"numeric_medians"` // 数值特征中位数（填充缺失值） CategoricalModes map[string]string `json:"categorical_modes"` // 分类特征众数（填充缺失值） } // ========== 2. XGBoost模型结构体（适配2.x原生JSON格式） ========== type XGBoostModel struct { Version []int `json:"version"` Learner struct { Attributes struct { BestIteration string `json:"best_iteration"` BestScore string `json:"best_score"` } `json:"attributes"` FeatureNames []string `json:"feature_names"` FeatureTypes []string `json:"feature_types"` GradientBooster struct { Model struct { GBTreeModelParam struct { NumParallelTree string `json:"num_parallel_tree"` NumTrees string `json:"num_trees"` } `json:"gbtree_model_param"` IterationIndptr []int `json:"iteration_indptr"` TreeInfo []int `json:"tree_info"` Trees []Tree `json:"trees"` // 树数组 } `json:"model"` TreeParam struct { NumFeature string `json:"num_feature"` } `json:"tree_param"` } `json:"gradient_booster"` LearnerModelParam struct { BaseScore string `json:"base_score"` NumClass string `json:"num_class"` NumFeature string `json:"num_feature"` NumTarget string `json:"num_target"` } `json:"learner_model_param"` Objective struct { Name string `json:"name"` } `json:"objective"` } `json:"learner"` } // Tree 适配XGBoost 2.x原生JSON格式的树结构 type Tree struct { BaseWeights []float64 `json:"base_weights"` Categories []string `json:"categories"` CategoriesNodes []int `json:"categories_nodes"` CategoriesSegments []int `json:"categories_segments"` CategoriesSizes []int `json:"categories_sizes"` DefaultLeft []int `json:"default_left"` ID int `json:"id"` LeftChildren []int `json:"left_children"` LossChanges []float64 `json:"loss_changes"` Parents []int `json:"parents"` RightChildren []int `json:"right_children"` SplitConditions []float64 `json:"split_conditions"` SplitIndices []int `json:"split_indices"` SplitType []int `json:"split_type"` SumHessian []float64 `json:"sum_hessian"` TreeParam struct { NumDeleted string `json:"num_deleted"` NumFeature string `json:"num_feature"` NumNodes string `json:"num_nodes"` SizeLeafVector string `json:"size_leaf_vector"` } `json:"tree_param"` } // ========== 3. 差异化预测器结构体 ========== type CaptchaPredictor struct { // 模型根路径 modelRootPath string // 加载的模型和配置 models map[string]*XGBoostModel configs map[string]*ModelJSONConfig // 初始化状态 initialized bool } // ========== 4. 预测结果结构体 ========== type PredictResult struct { VerifyType string `json:"verify_type"` AbnormalProbability float64 `json:"abnormal_probability"` PredictedLabel int `json:"predicted_label"` LabelDescription string `json:"label_description"` PredictionTimestamp string `json:"prediction_timestamp"` } // ========== 5. 构造函数：创建差异化预测器 ========== func NewCaptchaPredictor(modelRootPath string) (*CaptchaPredictor, error) { predictor := &CaptchaPredictor{ modelRootPath: modelRootPath, models: make(map[string]*XGBoostModel), configs: make(map[string]*ModelJSONConfig), initialized: false, } // 加载swipe和click两种模型 verifyTypes := []string{"swipe", "click"} for _, vt := range verifyTypes { // 加载JSON配置（Python生成的captcha_model_{vt}.json） configPath := fmt.Sprintf("%s/captcha_model_%s.json", modelRootPath, vt) config, err := loadModelConfig(configPath) if err != nil { log.Printf("⚠️ 加载%s配置失败: %v，跳过该模型", vt, err) continue } predictor.configs[vt] = config // 加载XGBoost JSON模型（Python转换的captcha_xgb_model_{vt}.json） modelPath := fmt.Sprintf("%s/captcha_xgb_model_%s.json", modelRootPath, vt) model, err := loadXGBoostModel(modelPath) if err != nil { log.Printf("⚠️ 加载%s模型失败: %v，跳过该模型", vt, err) delete(predictor.configs, vt) continue } predictor.models[vt] = model log.Printf("✅ 成功加载%s模型：特征数=%d，决策树数=%d", vt, len(config.FeatureColumns), len(model.Learner.GradientBooster.Model.Trees)) } if len(predictor.models) == 0 { return nil, fmt.Errorf("未加载到任何验证码模型") } predictor.initialized = true log.Println("✅ 验证码预测器初始化完成") return predictor, nil } // ========== 6. 内部加载方法 ========== func loadModelConfig(configPath string) (*ModelJSONConfig, error) { data, err := os.ReadFile(configPath) if err != nil { return nil, err } var config ModelJSONConfig if err := json.Unmarshal(data, &config); err != nil { return nil, fmt.Errorf("解析配置失败: %v", err) } return &config, nil } func loadXGBoostModel(modelPath string) (*XGBoostModel, error) { data, err := os.ReadFile(modelPath) if err != nil { return nil, err } var model XGBoostModel if err := json.Unmarshal(data, &model); err != nil { return nil, fmt.Errorf("解析模型失败: %v", err) } return &model, nil } // ========== 7. 预处理方法（完全匹配Python逻辑） ========== func (p *CaptchaPredictor) preprocess(rawData map[string]interface{}, verifyType string) ([]float64, error) { if !p.initialized { return nil, fmt.Errorf("预测器未初始化") } // 获取对应配置 config, ok := p.configs[verifyType] if !ok { return nil, fmt.Errorf("未找到%s类型的配置", verifyType) } // 初始化特征数组 featureCols := config.FeatureColumns features := make([]float64, len(featureCols)) // 构建特征索引映射 featureIndex := make(map[string]int, len(featureCols)) for idx, col := range featureCols { featureIndex[col] = idx } // 过滤无用列（完全匹配Python的USELESS_COLS） uselessCols := map[string]bool{ "id": true, "sample_id": true, "session_id": true, "ip": true, "device_fingerprint": true, "temp_score": true, "created_at": true, "updated_at": true, "user_id": true, "device_id": true, "phone_number": true, "email": true, "cookie_id": true, "session_token": true, } for col, idx := range featureIndex { // 跳过无用列 if uselessCols[col] { features[idx] = 0.0 continue } val, exists := rawData[col] // ========== 关键修改1：缺失值填充（匹配Python逻辑） ========== if !exists || val == nil { // 分类特征：使用众数填充 if le, ok := config.LabelEncoders[col]; ok { modeVal := config.CategoricalModes[col] if modeVal == "" { modeVal = "Unknown" } code, ok := le.Mapping[modeVal] if !ok { code = le.Mapping["Unknown"] } features[idx] = float64(code) log.Printf("⚠️ 特征[%s]缺失，填充众数[%s]编码: %d", col, modeVal, code) continue } // 数值特征：使用中位数填充 medianVal, ok := config.NumericMedians[col] if !ok { // 核心特征特殊处理 if strings.Contains(col, "delay") || strings.Contains(col, "duration") { medianVal = 0.0 } else if strings.Contains(col, "freq") { medianVal = 0.0 } else { medianVal = 0.0 } } features[idx] = medianVal //log.Printf("⚠️ 特征[%s]缺失，填充中位数: %.2f", col, medianVal) continue } // ========== 关键修改2：分类特征处理（完全匹配Python） ========== if le, ok := config.LabelEncoders[col]; ok { var strVal string switch v := val.(type) { case string: strVal = v case nil: strVal = "Unknown" default: strVal = fmt.Sprintf("%v", v) } // 使用编码映射 code, ok := le.Mapping[strVal] if !ok { // 未知值使用Unknown的编码 code = le.Mapping["Unknown"] log.Printf("⚠️ 特征[%s]未知值: %s，使用Unknown编码: %d", col, strVal, code) } features[idx] = float64(code) continue } // ========== 关键修改3：数值特征处理（完全匹配Python） ========== var numVal float64 switch v := val.(type) { case float64: numVal = v case int: numVal = float64(v) case int64: numVal = float64(v) case string: n, err := strconv.ParseFloat(v, 64) if err != nil { // 解析失败：使用中位数填充 medianVal, ok := config.NumericMedians[col] if !ok { medianVal = 0.0 } numVal = medianVal log.Printf("⚠️ 特征[%s]值错误: %s，填充中位数: %.2f", col, v, medianVal) } else { numVal = n } default: // 类型错误：使用中位数填充 medianVal, ok := config.NumericMedians[col] if !ok { medianVal = 0.0 } numVal = medianVal log.Printf("⚠️ 特征[%s]类型错误: %T，填充中位数: %.2f", col, v, medianVal) } // 负数修复（匹配Python逻辑：所有负数替换为0） if numVal < 0 { numVal = 0.0 log.Printf("⚠️ 特征[%s]负数修复为0（原值: %.2f）", col, numVal) } // ========== 关键修改4：极端值过滤（使用Python保存的分位数） ========== if colParams, ok := config.PreprocessingParams[col].(map[string]interface{}); ok { // 下界过滤 if qLow, ok := colParams["q_low"].(float64); ok && numVal < qLow { log.Printf("⚠️ 特征[%s]低于下界%.2f，截断为%.2f（原值: %.2f）", col, qLow, qLow, numVal) numVal = qLow } // 上界过滤 if qHigh, ok := colParams["q_high"].(float64); ok && numVal > qHigh { log.Printf("⚠️ 特征[%s]高于上界%.2f，截断为%.2f（原值: %.2f）", col, qHigh, qHigh, numVal) numVal = qHigh } } // ========== 关键修改5：特征交互项计算（匹配Python） ========== // 先保存原始值，后续计算交互项 features[idx] = numVal } // 计算特征交互项（完全匹配Python逻辑） p.calculateFeatureInteractions(features, featureIndex, rawData, verifyType) return features, nil } // calculateFeatureInteractions 计算特征交互项（匹配Python的create_core_feature_interactions） func (p *CaptchaPredictor) calculateFeatureInteractions(features []float64, featureIndex map[string]int, rawData map[string]interface{}, verifyType string) { // 通用交互项：load_to_total_ratio = load_to_first_delay / (total_duration + 1e-6) if loadIdx, ok := featureIndex["load_to_first_delay"]; ok { if totalIdx, ok := featureIndex["total_duration"]; ok { if ratioIdx, ok := featureIndex["load_to_total_ratio"]; ok { loadVal := features[loadIdx] totalVal := features[totalIdx] ratio := loadVal / (totalVal + 1e-6) features[ratioIdx] = ratio log.Printf("✅ 计算通用交互项 load_to_total_ratio = %.4f", ratio) } } } // 差异化交互项 switch verifyType { case "swipe": // 滑动验证：load_swipe_ratio = load_to_first_delay / (avg_swipe_interval + 1e-6) if loadIdx, ok := featureIndex["load_to_first_delay"]; ok { if swipeIdx, ok := featureIndex["avg_swipe_interval"]; ok { if ratioIdx, ok := featureIndex["load_swipe_ratio"]; ok { loadVal := features[loadIdx] swipeVal := features[swipeIdx] ratio := loadVal / (swipeVal + 1e-6) features[ratioIdx] = ratio log.Printf("✅ 计算滑动交互项 load_swipe_ratio = %.4f", ratio) } } } case "click": // 点击验证：duration_click_cv = total_duration * click_interval_cv if totalIdx, ok := featureIndex["total_duration"]; ok { if clickIdx, ok := featureIndex["click_interval_cv"]; ok { if cvIdx, ok := featureIndex["duration_click_cv"]; ok { totalVal := features[totalIdx] clickVal := features[clickIdx] product := totalVal * clickVal features[cvIdx] = product log.Printf("✅ 计算点击交互项 duration_click_cv = %.4f", product) } } } } } // ========== 8. 预测方法（优化XGBoost推理逻辑） ========== func (p *CaptchaPredictor) predict(features []float64, verifyType string) (float64, int, error) { if !p.initialized { return 0, 0, fmt.Errorf("预测器未初始化") } // 获取对应模型和配置 model, ok := p.models[verifyType] if !ok { return 0, 0, fmt.Errorf("未找到%s类型的模型", verifyType) } config, ok := p.configs[verifyType] if !ok { return 0, 0, fmt.Errorf("未找到%s类型的配置", verifyType) } // ========== 关键修改6：正确解析BaseScore ========== actualBaseScore := 0.5 if model.Learner.LearnerModelParam.BaseScore != "" { bs, err := strconv.ParseFloat(model.Learner.LearnerModelParam.BaseScore, 64) if err == nil { actualBaseScore = bs } } // 检查特征长度 featureCols := config.FeatureColumns if len(features) != len(featureCols) { return 0, 0, fmt.Errorf("特征长度不匹配：期望%d，实际%d", len(featureCols), len(features)) } // 获取所有决策树 trees := model.Learner.GradientBooster.Model.Trees if len(trees) == 0 { log.Println("⚠️ 模型中未找到决策树，返回默认概率0.5") return 0.5, 0, nil } // ========== 关键修改7：优化树遍历逻辑 ========== totalScore := 0.0 validTreeCount := 0 for treeIdx, tree := range trees { currentNodeID := 0 // 从根节点开始 nodeCount, err := strconv.Atoi(tree.TreeParam.NumNodes) if err != nil || nodeCount == 0 { log.Printf("⚠️ 树%d节点数无效: %s，跳过", treeIdx, tree.TreeParam.NumNodes) continue } // 遍历节点直到叶节点 for { // 边界检查 if currentNodeID < 0 || currentNodeID >= len(tree.LeftChildren) { break } // 检查是否是叶节点（左子节点为-1） leftChild := tree.LeftChildren[currentNodeID] if leftChild == -1 { // 累加叶节点权重 if currentNodeID < len(tree.BaseWeights) { totalScore += tree.BaseWeights[currentNodeID] validTreeCount++ } break } // 获取分裂特征和阈值 splitFeatIdx := tree.SplitIndices[currentNodeID] splitThreshold := tree.SplitConditions[currentNodeID] // 特征索引检查 if splitFeatIdx < 0 || splitFeatIdx >= len(features) { // 特征不存在，走右分支 currentNodeID = tree.RightChildren[currentNodeID] continue } // 判断分支（匹配XGBoost原生逻辑） featVal := features[splitFeatIdx] if math.IsNaN(featVal) || math.IsInf(featVal, 0) || featVal <= splitThreshold { currentNodeID = leftChild } else { currentNodeID = tree.RightChildren[currentNodeID] } } } // ========== 关键修改8：正确计算逻辑回归概率 ========== // XGBoost二分类公式：prob = 1 / (1 + exp(-(base_score_log_odds + sum(tree_scores)))) baseLogOdds := math.Log(actualBaseScore / (1 - actualBaseScore)) finalLogOdds := baseLogOdds + totalScore predProb := 1.0 / (1.0 + math.Exp(-finalLogOdds)) // 生成标签（阈值0.5，匹配Python） predLabel := 0 if predProb > 0.5 { predLabel = 1 } log.Printf("✅ %s模型推理完成：有效树数=%d，BaseScore=%.4f，总得分=%.4f，概率=%.4f，标签=%d", verifyType, validTreeCount, actualBaseScore, totalScore, predProb, predLabel) return predProb, predLabel, nil } // ========== 9. 统一预测接口（完全匹配Python的predict_captcha） ========== func (p *CaptchaPredictor) PredictCaptcha(rawData map[string]interface{}) (*PredictResult, error) { if !p.initialized { return nil, fmt.Errorf("预测器未初始化") } // 获取验证类型 verifyType, ok := rawData["verify_type"].(string) if !ok { verifyType = "default" } verifyType = strings.ToLower(verifyType) // 处理默认类型（匹配Python逻辑） if verifyType != "swipe" && verifyType != "click" { // 优先使用swipe模型，没有则使用click if _, ok := p.models["swipe"]; ok { verifyType = "swipe" } else if _, ok := p.models["click"]; ok { verifyType = "click" } else { return nil, fmt.Errorf("无可用的验证码模型") } log.Printf("⚠️ 验证类型%s不支持，自动切换为%s", rawData["verify_type"], verifyType) } // 预处理（完全匹配Python） features, err := p.preprocess(rawData, verifyType) if err != nil { return nil, fmt.Errorf("预处理失败: %v", err) } // 预测 predProb, predLabel, err := p.predict(features, verifyType) if err != nil { return nil, fmt.Errorf("预测失败: %v", err) } // 构建结果（完全匹配Python返回格式） result := &PredictResult{ VerifyType: verifyType, AbnormalProbability: predProb, PredictedLabel: predLabel, PredictionTimestamp: time.Now().Format("2006-01-02 15:04:05"), } if predLabel == 1 { result.LabelDescription = "异常/疑似" } else { result.LabelDescription = "正常" } return result, nil }
from fastapi import FastAPI, HTTPException from pydantic import BaseModel from mlx_lm import load, generate import uvicorn import os # 禁用 Flask 相关环境变量（避免环境污染） os.environ.pop("FLASK_APP", None) os.environ.pop("FLASK_ENV", None) # 初始化纯 FastAPI 应用，确保无 Flask 混入 app = FastAPI(title="MLX 翻译服务", version="1.0") # 预加载模型（只加载一次） print("正在加载翻译模型...") try: model, tokenizer = load("alexgusevski/HY-MT1.5-1.8B-q4-mlx") print("模型加载完成！") except Exception as e: raise RuntimeError(f"模型加载失败：{str(e)}") # 定义请求体模型 class TranslationRequest(BaseModel): source_term: str = "英文" target_term: str = "中文" source_text: str # 翻译接口（异步封装同步函数，避免阻塞事件循环） @app.post("/translate") async def translate(request: TranslationRequest): try: # 构建提示词 prompt = f"""{request.source_term} 翻译成 {request.target_term} 将以下文本翻译为{request.target_term}，注意只需要输出翻译结果，不要额外解释： {request.source_text}""" # 处理聊天模板 if hasattr(tokenizer, "apply_chat_template") and tokenizer.chat_template is not None: messages = [{"role": "user", "content": prompt}] prompt = tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) # mlx_lm.generate 是同步函数，用 run_in_threadpool 避免阻塞 import asyncio response = await asyncio.get_event_loop().run_in_executor( None, lambda: generate( model, tokenizer, prompt=prompt, verbose=False, max_tokens=1024, ) ) return { "success": True, "source_text": request.source_text, "target_text": response.strip(), "source_language": request.source_term, "target_language": request.target_term } except Exception as e: raise HTTPException(status_code=500, detail=f"翻译失败：{str(e)}") # 健康检查接口 @app.get("/health") async def health_check(): return {"status": "healthy", "message": "翻译服务正常运行"} # 启动服务 if __name__ == "__main__": uvicorn.run( app="main:app", # 若文件名为其他（如 translate_api.py），改为 "translate_api:app" host="0.0.0.0", port=8000, reload=False, # 关闭重载，避免异步冲突 workers=1 # 单进程运行，适配 mlx 模型 )
import os import subprocess import cv2 import numpy as np import re import time import random import sys import io from datetime import datetime from pathlib import Path from PIL import Image from difflib import SequenceMatcher from collections import defaultdict from functools import lru_cache from concurrent.futures import ThreadPoolExecutor # 导入ImgOcr类（请确保imgocr.py在当前目录） from imgocr import ImgOcr # ===================== 全局变量（优化：复用OCR实例） ===================== ocr_det = None # 检测用OCR实例（效率模式） ocr_rec = None # 识别用OCR实例（精准模式） group_text_cache = {} # 缓存分组最优文本结果 # ===================== 新增：有效单个英文字母定义 ===================== # 仅保留有语义的单个英文字母，其余单个字母全部过滤 VALID_SINGLE_ENGLISH_CHARS = {'a', 'i', 'o', 'A', 'I', 'O'} # ===================== 日志重定向类（保留原有功能） ===================== class Tee: """重定向stdout/stderr，同时输出到控制台和日志文件""" def __init__(self, filename, mode='a', encoding='utf-8'): self.file = open(filename, mode, encoding=encoding) self.stdout = sys.stdout self.stderr = sys.stderr sys.stdout = self sys.stderr = self def write(self, message): self.file.write(message) self.stdout.write(message) self.file.flush() self.stdout.flush() def flush(self): self.file.flush() self.stdout.flush() def close(self): sys.stdout = self.stdout sys.stderr = self.stderr self.file.close() # ===================== 配置参数（优化：调整关键参数） ===================== VIDEO_PATH = "/Users/chenhuan/Downloads/中英双字幕版第1季52集/CE003 Best Friend.mp4" FRAME_OUTPUT_DIR = "extracted_frames" CROPPED_20_DIR = "cropped_20_percent" # 兼容旧逻辑，实际内存处理为主 FINAL_SUBTITLE_DIR = "final_subtitle_boxes" SRT_OUTPUT_PATH = "subtitles.srt" LOG_FILE_PATH = f"subtitle_extract_{datetime.now().strftime('%Y%m%d_%H%M%S')}.log" FPS = 1 # 提取帧率（每秒1帧） AREA_THRESHOLD = 4000 MIN_SIZE_THRESHOLD = 100 SIMILARITY_THRESHOLD = 0.8 CONFIDENCE_THRESHOLD = 0.65 INVALID_CHARS = r'[\[\]\(\)\\\/\^\$\#\@\!\~\_\*]' LINE_Y_THRESHOLD = 10 EXTRA_WIDTH = 50 EXTRA_HEIGHT = 20 SAMPLE_COUNT = 20 # 优化：从35改为20（足够判断字幕类型） BILINGUAL_THRESHOLD = 0.6 SRT_END_OFFSET = 0.1 # 英文高频错误补全规则库 ENGLISH_CORRECTION_RULES = [ (r"^'m\s", "I'm "), (r"^Im\s", "I'm "), (r"^m\s", "I'm "), ] ENGLISH_FIX_RULES = [] # ===================== 核心优化函数：初始化OCR（仅加载一次） ===================== def init_ocr_instance(): """初始化OCR实例（全局复用，避免重复加载模型）""" global ocr_det, ocr_rec print("🔄 初始化OCR模型（仅加载一次）...") start_time = time.time() # 检测用效率模式（速度优先），识别用精准模式（精度优先） ocr_det = ImgOcr(model_version="v5", is_efficiency_mode=True) ocr_rec = ImgOcr(model_version="v5", is_efficiency_mode=True, use_angle_cls=True) print(f"✅ OCR模型初始化完成，耗时：{time.time() - start_time:.2f}s") # ===================== 最终修复：FFmpeg提取帧（兼容所有版本） ===================== def extract_frames_with_ffmpeg_optimized(video_path, output_dir, fps=1): """兼容所有FFmpeg版本的帧提取：移除冲突参数，保证基础功能可用""" os.makedirs(output_dir, exist_ok=True) # 最终简化版FFmpeg命令： # 1. 移除-skip_frame nokey（解码参数，编码侧报错） # 2. 移除-fps_mode cfr（新版兼容问题） # 3. 用-vsync 0（兼容模式）+ filter:v fps={fps} 控制帧率 # 4. 保证参数顺序：输入 → 过滤 → 编码 → 输出 ffmpeg_cmd = [ "ffmpeg", "-i", video_path, # 输入视频 "-vsync", "0", # 兼容模式，避免帧重复/丢失 "-filter:v", f"fps={fps}", # 控制输出帧率（每秒1帧） "-c:v", "mjpeg", # 编码格式（MJPEG，速度快） "-q:v", "3", # 画质（3为平衡，数值越小画质越好） "-y", # 覆盖已有文件（无需确认） f"{output_dir}/frame_%04d.jpg" # 输出路径（必须最后） ] try: # 执行命令并捕获输出 result = subprocess.run( ffmpeg_cmd, check=True, capture_output=True, text=True, encoding='utf-8' ) # 统计提取的帧数 frame_files = [f for f in os.listdir(output_dir) if f.startswith("frame_") and f.endswith(".jpg")] print(f"✅ 成功提取 {len(frame_files)} 帧到 {output_dir}") except subprocess.CalledProcessError as e: print(f"❌ 提取帧失败：{e.stderr}") raise # ===================== 优化：内存中裁剪（减少IO） ===================== def crop_bottom_20_percent_in_memory(img): """内存中裁剪底部20%区域，不保存文件""" width, height = img.size crop_y_start = int(height * 0.8) crop_region = (0, crop_y_start, width, height) return img.crop(crop_region), crop_region def crop_bottom_20_percent(image_path, output_dir): """兼容旧逻辑的裁剪函数（保留）""" os.makedirs(output_dir, exist_ok=True) img = Image.open(image_path) width, height = img.size crop_y_start = int(height * 0.8) crop_region = (0, crop_y_start, width, height) cropped_img = img.crop(crop_region) img_name = os.path.basename(image_path) cropped_path = os.path.join(output_dir, img_name) cropped_img.save(cropped_path) return cropped_path, crop_region # ===================== 优化：简化预处理（提升速度） ===================== def preprocess_image_optimized(image_path): """简化预处理：仅灰度化，移除锐化/阈值/滤波（速度提升50%+）""" img = cv2.imread(image_path) gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) return Image.fromarray(gray) # ===================== 原有工具函数（保留核心逻辑） ===================== def convert_box_to_bounding_rect(box): box_np = np.array(box) xmin = int(np.min(box_np[:, 0])) ymin = int(np.min(box_np[:, 1])) xmax = int(np.max(box_np[:, 0])) ymax = int(np.max(box_np[:, 1])) return xmin, ymin, xmax, ymax def filter_small_boxes(boxes, area_threshold): filtered_boxes = [] for box in boxes: xmin, ymin, xmax, ymax = convert_box_to_bounding_rect(box) width = xmax - xmin height = ymax - ymin area = width * height if area >= area_threshold and not (width < MIN_SIZE_THRESHOLD and height < MIN_SIZE_THRESHOLD): filtered_boxes.append(box) return filtered_boxes def merge_boxes(boxes): if not boxes: return None all_x = [] all_y = [] for box in boxes: box_np = np.array(box) all_x.extend(box_np[:, 0]) all_y.extend(box_np[:, 1]) xmin = int(np.min(all_x)) ymin = int(np.min(all_y)) xmax = int(np.max(all_x)) ymax = int(np.max(all_y)) return xmin, ymin, xmax, ymax def crop_and_save_subtitle_box(image_path, crop_region, merged_box, output_dir): os.makedirs(output_dir, exist_ok=True) img = Image.open(image_path) img_width, img_height = img.size crop_y_start = crop_region[1] xmin, ymin, xmax, ymax = merged_box original_xmin = xmin original_ymin = ymin + crop_y_start original_xmax = xmax original_ymax = ymax + crop_y_start new_xmin = max(0, original_xmin - EXTRA_WIDTH) new_xmax = min(img_width, original_xmax + EXTRA_WIDTH) new_ymin = max(0, original_ymin - EXTRA_HEIGHT) new_ymax = min(img_height, original_ymax + EXTRA_HEIGHT) subtitle_box = (new_xmin, new_ymin, new_xmax, new_ymax) subtitle_img = img.crop(subtitle_box) img_name = os.path.basename(image_path) subtitle_path = os.path.join(output_dir, img_name) subtitle_img.save(subtitle_path) print( f"📸 已保存字幕区域：{subtitle_path} | 调整后范围：宽[{new_xmin}→{new_xmax}] 高[{new_ymin}→{new_ymax}]") def has_chinese_char(text): chinese_pattern = re.compile(r'[\u4e00-\u9fff]') return bool(chinese_pattern.search(text)) def is_pure_english(text): english_pattern = re.compile(r'^[a-zA-Z0-9\s,.;!?\':"]+$') return bool(english_pattern.match(text.strip())) and not has_chinese_char(text) # ===================== 修复核心：merge_ocr_fragments_to_lines 函数 ===================== def merge_ocr_fragments_to_lines(ocr_result): """将OCR识别的零散片段按坐标合并为完整的字幕行（修复lines未初始化问题）""" if not ocr_result: return [] # 1. 计算每个片段的中心y坐标和置信度 fragments = [] for item in ocr_result: text = item['text'].strip() if not text: continue box = item['box'] conf = item['score'] # 计算片段的中心y坐标 y_coords = [point[1] for point in box] center_y = np.mean(y_coords) # 计算片段的x起始坐标（用于排序同一行的片段） x_coords = [point[0] for point in box] start_x = np.min(x_coords) fragments.append({ 'text': text, 'conf': conf, 'center_y': center_y, 'start_x': start_x }) # 2. 按y坐标聚类合并为行（y差≤LINE_Y_THRESHOLD视为同一行） lines = [] # 关键修复：提前初始化lines为空列表，解决Unresolved reference错误 if fragments: # 先按y坐标排序 fragments_sorted_by_y = sorted(fragments, key=lambda x: x['center_y']) current_line = [fragments_sorted_by_y[0]] current_line_y = fragments_sorted_by_y[0]['center_y'] for frag in fragments_sorted_by_y[1:]: if abs(frag['center_y'] - current_line_y) <= LINE_Y_THRESHOLD: current_line.append(frag) else: lines.append(current_line) current_line = [frag] current_line_y = frag['center_y'] # 加入最后一行 lines.append(current_line) else: lines = [] # 3. 每行内按x坐标排序，合并文本并计算平均置信度 merged_lines = [] for line in lines: # 按x起始坐标排序（保证从左到右的顺序） line_sorted_by_x = sorted(line, key=lambda x: x['start_x']) # 合并文本 line_text = ' '.join([frag['text'] for frag in line_sorted_by_x]) # 用空格连接片段，避免拼接错误 # 计算平均置信度 line_conf = np.mean([frag['conf'] for frag in line_sorted_by_x]) # 记录行的y坐标（用于后续排序） line_y = np.mean([frag['center_y'] for frag in line_sorted_by_x]) merged_lines.append({ 'text': line_text, 'conf': line_conf, 'y': line_y }) # 4. 按y坐标排序所有行（保证从上到下的顺序） merged_lines_sorted = sorted(merged_lines, key=lambda x: x['y']) return merged_lines_sorted # ===================== 核心修改：增强文本清洗函数（过滤无含义单个字母） ===================== def clean_text(text, is_chinese=False): # 第一步：移除无效字符 text = re.sub(INVALID_CHARS, '', text) # 第二步：针对英文文本，过滤无含义的单个字母 if not is_chinese and len(text.strip()) == 1: # 如果是单个字母，检查是否在有效列表中 single_char = text.strip() if single_char not in VALID_SINGLE_ENGLISH_CHARS: return "" # 过滤无含义单个字母 # 第三步：常规文本格式化 if is_chinese: text = re.sub(r'\s+', '', text) text = re.sub(r'([，。；！？：])\s+', r'\1', text) text = re.sub(r'\s+([，。；！？：])', r'\1', text) else: text = re.sub(r'\s+([,.;!?\':"])', r'\1', text) text = re.sub(r'([,.;!?\':"])\s+', r'\1 ', text) text = re.sub(r'\s+', ' ', text).strip() text = re.sub(r',\s+', ',', text) return text.strip() def fix_english_spelling(text): fixed_text = text for pattern, replacement in ENGLISH_FIX_RULES: fixed_text = re.sub(pattern, replacement, fixed_text) fixed_text = re.sub(r'\s+', ' ', fixed_text).strip() return fixed_text def correct_english_text(text): corrected_text = fix_english_spelling(text) for pattern, replacement in ENGLISH_CORRECTION_RULES: corrected_text = re.sub(pattern, replacement, corrected_text) return corrected_text def standardize_text(text, is_chinese=False): if not text: return "" chinese_punc = r'[，。；！？：“”‘’（）【】《》、…—]' english_punc = r'[,.!?;:"\'()\[\]<>-\u2026]' punctuation_pattern = f"{chinese_punc}{english_punc}" text_no_punc = re.sub(punctuation_pattern, '', text) text_clean_space = re.sub(r'\s+', ' ', text_no_punc).strip() if not is_chinese: text_standardized = text_clean_space.lower() else: text_standardized = text_clean_space return text_standardized def calculate_similarity(text1, text2): return SequenceMatcher(None, text1, text2).ratio() def group_similar_frames(ocr_results_list): if not ocr_results_list: return [] groups = [] current_group = [ocr_results_list[0]] for result in ocr_results_list[1:]: prev_result = current_group[-1] chinese_sim = calculate_similarity(prev_result[1], result[1]) english_sim = calculate_similarity(prev_result[3], result[3]) if chinese_sim >= SIMILARITY_THRESHOLD or english_sim >= SIMILARITY_THRESHOLD: current_group.append(result) else: groups.append(current_group) current_group = [result] groups.append(current_group) return groups # ===================== 优化：缓存分组最优文本结果 ===================== def get_group_key(group, text_type): """生成分组的唯一缓存key""" frame_names = tuple([item[0] for item in group]) return (frame_names, text_type) def select_best_text_in_group(group, text_type): """原始最优文本选择逻辑（保留）""" if not group: return "无" all_candidates = [] for item in group: frame_name, chinese_text, chinese_conf, english_text, english_conf = item if text_type == 'chinese': if chinese_text: standardized = standardize_text(chinese_text, is_chinese=True) all_candidates.append({ 'original': chinese_text, 'standardized': standardized, 'conf': chinese_conf, 'frame_name': frame_name }) else: if english_text: corrected_eng = correct_english_text(english_text) standardized = standardize_text(corrected_eng, is_chinese=False) all_candidates.append({ 'original': corrected_eng, 'standardized': standardized, 'conf': english_conf, 'frame_name': frame_name }) if not all_candidates: return "无" count_dict = defaultdict(int) for res in all_candidates: if res['standardized']: count_dict[res['standardized']] += 1 for res in all_candidates: res['count'] = count_dict.get(res['standardized'], 0) all_candidates.sort( key=lambda x: (x['count'], x['conf'], len(x['original'])), reverse=True ) max_count = all_candidates[0]['count'] top_candidates = [res for res in all_candidates if res['count'] == max_count] top_candidates.sort( key=lambda x: (x['conf'], len(x['original'])), reverse=True ) return top_candidates[0]['original'] def select_best_text_in_group_cached(group, text_type): """缓存版最优文本选择（避免重复计算）""" key = get_group_key(group, text_type) if key in group_text_cache: return group_text_cache[key] result = select_best_text_in_group(group, text_type) group_text_cache[key] = result return result def sample_groups(groups, sample_count=SAMPLE_COUNT): if len(groups) <= sample_count: sampled_groups = groups print(f"ℹ️ 总组数{len(groups)}≤采样数{sample_count}，使用全部组进行字幕类型判断") else: random.seed(42) sampled_groups = random.sample(groups, sample_count) print(f"ℹ️ 从{len(groups)}组中随机采样{sample_count}组进行字幕类型判断") return sampled_groups # ===================== 优化：缓存版字幕类型判断 ===================== def judge_subtitle_type_optimized(sampled_groups, bilingual_threshold=BILINGUAL_THRESHOLD): bilingual_count = 0 chinese_only_count = 0 english_only_count = 0 empty_count = 0 for group in sampled_groups: best_chinese = select_best_text_in_group_cached(group, 'chinese') best_english = select_best_text_in_group_cached(group, 'english') has_valid_chinese = best_chinese and best_chinese != "无" has_valid_english = best_english and best_english != "无" if has_valid_chinese and has_valid_english: bilingual_count += 1 elif has_valid_chinese and not has_valid_english: chinese_only_count += 1 elif not has_valid_chinese and has_valid_english: english_only_count += 1 else: empty_count += 1 total_valid = len(sampled_groups) - empty_count if total_valid == 0: print("❌ 采样组中无有效字幕，默认判定为单语（英文）") return "english_only" bilingual_ratio = bilingual_count / total_valid print(f"\n📊 采样组字幕类型统计：") print(f" - 双语组：{bilingual_count}个") print(f" - 仅中文组：{chinese_only_count}个") print(f" - 仅英文组：{english_only_count}个") print(f" - 空组：{empty_count}个") print(f" - 双语占比：{bilingual_ratio:.2%}（阈值：{bilingual_threshold:.2%}）") if bilingual_ratio >= bilingual_threshold: subtitle_type = "bilingual" print(f"✅ 判定字幕类型为：双语") elif chinese_only_count > english_only_count: subtitle_type = "chinese_only" print(f"✅ 判定字幕类型为：仅中文") else: subtitle_type = "english_only" print(f"✅ 判定字幕类型为：仅英文") return subtitle_type # ===================== 优化：缓存版分组过滤 ===================== def filter_groups_by_type_optimized(groups, subtitle_type): filtered_groups = [] filtered_count = 0 for group in groups: best_chinese = select_best_text_in_group_cached(group, 'chinese') best_english = select_best_text_in_group_cached(group, 'english') has_valid_chinese = best_chinese and best_chinese != "无" has_valid_english = best_english and best_english != "无" if subtitle_type == "bilingual": if has_valid_chinese and has_valid_english: filtered_groups.append(group) else: filtered_count += 1 elif subtitle_type == "chinese_only": if has_valid_chinese: filtered_groups.append(group) else: filtered_count += 1 else: if has_valid_english: filtered_groups.append(group) else: filtered_count += 1 print(f"\n🗂️ 分组过滤结果：") print(f" - 原始组数：{len(groups)}") print(f" - 过滤后组数：{len(filtered_groups)}") print(f" - 过滤掉的组数：{filtered_count}") return filtered_groups def seconds_to_srt_time(seconds): hours = int(seconds // 3600) minutes = int((seconds % 3600) // 60) secs = int(seconds % 60) ms = int((seconds % 1) * 1000) return f"{hours:02d}:{minutes:02d}:{secs:02d},{ms:03d}" def generate_srt(groups, subtitle_type, fps=FPS): srt_content = [] for idx, group in enumerate(groups, 1): frame_numbers = [] for item in group: frame_name = item[0] match = re.search(r'frame_(\d+)\.jpg', frame_name) if match: frame_numbers.append(int(match.group(1))) if not frame_numbers: continue first_frame = min(frame_numbers) last_frame = max(frame_numbers) start_time = (first_frame - 1) / fps end_time = (last_frame / fps) + SRT_END_OFFSET print(f"📌 第{idx}组时间：首帧{first_frame}({start_time:.2f}s) → 尾帧{last_frame}({end_time:.2f}s)") best_chinese = select_best_text_in_group_cached(group, 'chinese') best_english = select_best_text_in_group_cached(group, 'english') srt_text_lines = [] if subtitle_type == "bilingual": srt_text_lines.append(best_english) srt_text_lines.append(best_chinese) elif subtitle_type == "chinese_only": srt_text_lines.append(best_chinese) else: srt_text_lines.append(best_english) # 过滤空行 srt_text_lines = [line for line in srt_text_lines if line and line != "无"] if not srt_text_lines: continue srt_entry = [ str(idx), f"{seconds_to_srt_time(start_time)} --> {seconds_to_srt_time(end_time)}", "\n".join(srt_text_lines) ] srt_content.append("\n".join(srt_entry) + "\n") if srt_content: with open(SRT_OUTPUT_PATH, 'w', encoding='utf-8') as f: f.write("\n".join(srt_content).strip()) print(f"\n✅ SRT字幕文件已生成：{SRT_OUTPUT_PATH}") print(f"📊 最终生成字幕条目数：{len(srt_content)}") else: print("\n❌ 无有效字幕条目，无法生成SRT文件") # ===================== 优化：批量OCR识别（简化预处理） ===================== def ocr_subtitle_boxes_optimized(subtitle_dir): if not os.path.exists(subtitle_dir): print(f"❌ 目录 {subtitle_dir} 不存在，跳过OCR识别") return [] subtitle_files = sorted([f for f in os.listdir(subtitle_dir) if f.endswith((".jpg", ".png"))]) total_files = len(subtitle_files) ocr_results_list = [] if total_files == 0: print(f"⚠️ 目录 {subtitle_dir} 下无图片文件，跳过OCR识别") return [] print(f"\n📊 开始识别 {total_files} 张字幕图片（优化版批量处理）...") start_time = time.time() # 批量处理（减少模型调用开销） batch_size = 10 for batch_idx in range(0, total_files, batch_size): batch_files = subtitle_files[batch_idx:batch_idx + batch_size] batch_results = [] for sub_file in batch_files: sub_path = os.path.join(subtitle_dir, sub_file) # 简化预处理：仅灰度化（速度提升50%+） preprocessed_img = preprocess_image_optimized(sub_path) ocr_result = ocr_rec.ocr(preprocessed_img, det=True, rec=True, cls=False) merged_lines = merge_ocr_fragments_to_lines(ocr_result) chinese_lines = [] chinese_confs = [] english_lines = [] english_confs = [] for line in merged_lines: line_text = line['text'].strip() line_conf = line['conf'] if not line_text: continue if has_chinese_char(line_text): chinese_lines.append(line_text) chinese_confs.append(line_conf) elif is_pure_english(line_text): english_lines.append(line_text) english_confs.append(line_conf) # 清洗文本（包含无含义单个字母过滤） chinese_text_raw = clean_text(''.join(chinese_lines), is_chinese=True) chinese_conf = np.mean(chinese_confs) if chinese_confs else 0.0 english_text_raw = clean_text(' '.join(english_lines), is_chinese=False) english_conf = np.mean(english_confs) if english_confs else 0.0 batch_results.append((sub_file, chinese_text_raw, chinese_conf, english_text_raw, english_conf)) ocr_results_list.extend(batch_results) # 进度提示 processed = min(batch_idx + batch_size, total_files) progress = (processed / total_files) * 100 elapsed_time = time.time() - start_time print(f"\r🔄 识别进度：{processed}/{total_files} ({progress:.1f}%) | 已用时：{elapsed_time:.1f}s", end="") print("\n") groups = group_similar_frames(ocr_results_list) print(f"✅ 共分为 {len(groups)} 组重复字幕") # 打印分组详情（保留核心调试信息） print("\n" + "=" * 150) print("最终字幕识别结果（最优结果 + 组内统计）") print("=" * 150) for i, group in enumerate(groups, 1): frame_names = [item[0] for item in group] frame_range = f"{frame_names[0]} ~ {frame_names[-1]}" if len(frame_names) > 1 else frame_names[0] print(f"\n📦 第{i}组（帧范围：{frame_range} | 共{len(group)}帧）") best_chinese = select_best_text_in_group_cached(group, 'chinese') best_english = select_best_text_in_group_cached(group, 'english') print(f" 🎯 最优中文文本：{best_chinese if best_chinese else '无'}") print(f" 🎯 最优英文文本：{best_english if best_english else '无'}") print("\n" + "=" * 150) total_time = time.time() - start_time print(f"OCR识别优化完成！总耗时：{total_time:.1f}s") print("=" * 150) return groups # ===================== 优化：并行处理帧（多线程） ===================== def process_single_frame(args): """单帧处理函数（用于并行）""" frame_file, frame_dir, cropped_dir, subtitle_dir = args frame_path = os.path.join(frame_dir, frame_file) try: # 内存中裁剪（减少IO） img = Image.open(frame_path) cropped_img, crop_region = crop_bottom_20_percent_in_memory(img) # OCR检测（复用全局实例） cropped_img_bytes = io.BytesIO() cropped_img.save(cropped_img_bytes, format='JPEG') cropped_img_bytes.seek(0) ocr_result = ocr_det.ocr(cropped_img, det=True, rec=False, cls=False) if not ocr_result: return None detected_boxes = [item['box'] for item in ocr_result] filtered_boxes = filter_small_boxes(detected_boxes, AREA_THRESHOLD) if not filtered_boxes: return None merged_box = merge_boxes(filtered_boxes) if merged_box: # 保存最终字幕区域图片（必要IO） crop_and_save_subtitle_box(frame_path, crop_region, merged_box, subtitle_dir) return frame_file except Exception as e: print(f"⚠️ 处理帧 {frame_file} 出错：{e}") return None def process_frames_parallel(frame_dir, cropped_dir, subtitle_dir, max_workers=4): """并行处理帧（多线程，IO密集型任务最优）""" frame_files = sorted([f for f in os.listdir(frame_dir) if f.endswith((".jpg", ".png"))]) if not frame_files: print("⚠️ 无帧文件需要处理") return args_list = [(f, frame_dir, cropped_dir, subtitle_dir) for f in frame_files] print(f"🔄 开始并行处理 {len(frame_files)} 帧（线程数：{max_workers}）...") start_time = time.time() with ThreadPoolExecutor(max_workers=max_workers) as executor: results = list(executor.map(process_single_frame, args_list)) processed_frames = [r for r in results if r is not None] elapsed_time = time.time() - start_time print(f"✅ 并行处理完成！耗时：{elapsed_time:.2f}s | 成功处理 {len(processed_frames)} 帧") # ===================== 主函数（整合所有优化） ===================== def main(): tee = None try: # 初始化日志重定向 tee = Tee(LOG_FILE_PATH, 'w', 'utf-8') print(f"📝 日志文件已创建：{LOG_FILE_PATH}") print(f"⏰ 程序启动时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") print("=" * 80 + "\n") # 1. 初始化OCR实例（仅加载一次） init_ocr_instance() # 2. 提取帧（最终兼容版FFmpeg命令） extract_frames_with_ffmpeg_optimized(VIDEO_PATH, FRAME_OUTPUT_DIR, FPS) # 3. 并行处理帧（多线程+内存裁剪） process_frames_parallel(FRAME_OUTPUT_DIR, CROPPED_20_DIR, FINAL_SUBTITLE_DIR, max_workers=4) # 4. 优化版OCR识别（批量+缓存） groups = ocr_subtitle_boxes_optimized(FINAL_SUBTITLE_DIR) if not groups: print("❌ 无有效字幕分组，无法生成SRT文件") return # 5. 缓存版字幕类型判断+过滤 sampled_groups = sample_groups(groups) subtitle_type = judge_subtitle_type_optimized(sampled_groups) filtered_groups = filter_groups_by_type_optimized(groups, subtitle_type) if not filtered_groups: print("❌ 过滤后无有效分组，无法生成SRT文件") return # 6. 生成SRT文件 generate_srt(filtered_groups, subtitle_type, FPS) except Exception as e: print(f"❌ 程序执行出错：{str(e)}", file=sys.stderr) raise finally: # 恢复输出并关闭日志 if tee: print("\n" + "=" * 80) print(f"⏰ 程序结束时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}") print(f"📝 日志已保存至：{LOG_FILE_PATH}") tee.close() if __name__ == "__main__": main()
import torch from transformers import AutoTokenizer, AutoModelForSeq2SeqLM import re import time class GECTORErrorCorrectorM1: def __init__(self, model_name="512duncanl/gec-flan-t5-large"): """初始化GECToR模型，适配M1 Mac GPU加速""" # 设备优先级：MPS > CUDA > CPU if torch.backends.mps.is_available() and torch.backends.mps.is_built(): self.device = torch.device("mps") elif torch.cuda.is_available(): self.device = torch.device("cuda") else: self.device = torch.device("cpu") print(f"使用设备: {self.device}") # 加载模型和tokenizer（无量化，避免MPS报错） self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(self.device) self.model.eval() # 推理模式 def preprocess_text(self, text): """预处理：仅过滤句尾无意义数字，保留语义数字""" text = re.sub(r"\s+\d+$", "", text) # 只删句尾冗余数字 text = re.sub(r",(\S)", r", \1", text) # 逗号后加空格 text = re.sub(r"\s+", " ", text).strip() # 合并多余空格 return text def _get_diff(self, original, corrected): """对比文本差异，展示纠错点""" import difflib diff = difflib.ndiff(original.split(), corrected.split()) return [d for d in diff if d.startswith(('+', '-'))] def batch_correct(self, text_list, batch_size=4): """M1 优化的批量推理函数（含单条耗时统计）""" total_start_time = time.time() # 总耗时起始时间 results = [] processed_texts = [self.preprocess_text(t) for t in text_list] # 分批次处理，统计每个批次的耗时 for i in range(0, len(processed_texts), batch_size): batch_start = time.time() # 单个批次起始时间 batch_texts = processed_texts[i:i + batch_size] batch_origs = text_list[i:i + batch_size] # 批量编码 inputs = self.tokenizer( batch_texts, return_tensors="pt", padding=True, truncation=True, max_length=128 ).to(self.device) # 批量推理 with torch.no_grad(): outputs = self.model.generate( **inputs, max_length=128, num_beams=5, early_stopping=True, use_cache=False # MPS内存优化 ) # 批量解码 corrected_batch = self.tokenizer.batch_decode(outputs, skip_special_tokens=True) corrected_batch = [t.strip() for t in corrected_batch] # 计算当前批次总耗时，均分至每条文本 batch_total_time = (time.time() - batch_start) * 1000 # 毫秒 single_text_time = round(batch_total_time / len(batch_texts), 2) # 单条耗时 # 生成结果（含单条耗时） for j, (orig, processed, corrected) in enumerate(zip(batch_origs, batch_texts, corrected_batch)): idx = i + j + 1 errors = [] if processed != corrected: errors.append({ "original": processed, "corrected": corrected, "diff": self._get_diff(processed, corrected) }) results.append({ "序号": idx, "原始文本": orig, "预处理后": processed, "修复后文本": corrected, "错误详情": errors if errors else "无错误", "单条推理耗时(ms)": single_text_time # 新增单条耗时 }) # 统计总耗时 total_time = round((time.time() - total_start_time) * 1000, 2) results.append({ "序号": "总计", "原始文本": f"共{len(text_list)}条文本 | 批量大小{batch_size}", "预处理后": "-", "修复后文本": "-", "错误详情": "-", "单条推理耗时(ms)": f"总耗时 {total_time} ms" }) return results # ---------------------- 测试 M1 加速+单条耗时 ---------------------- if __name__ == "__main__": # 初始化 M1 优化版纠错器 corrector = GECTORErrorCorrectorM1() # 测试文本 test_texts = [ "I have 5 apples", "And turn you into a frog. 4", "Peppa and Suzy love playing in Peppa'a bedroom.", "No,George. This game is just for big girls", "I'm a tiny,little fair princess.", "George,I need some help:", "I'm making chocolate chip:cookies.", "Hmmm,I think your heart's a bit loose.", "Peppa and Suzy loves playing doctors and nurses.", "He r might get tired.", "Then I think I need Iots of cookies to make me better.", "Im a tiny,little fair princess." ] # 批量推理（batch_size=4） correction_results = corrector.batch_correct(test_texts, batch_size=4) # 输出结果（含单条耗时） print("=" * 100) print("GECToR M1 优化版 英文文本纠错结果（MPS加速+单条耗时统计）") print("=" * 100) for res in correction_results: print(f"\n【序号】{res['序号']}") print(f"原始文本：{res['原始文本']}") if res['序号'] != "总计": print(f"预处理后：{res['预处理后']}") print(f"修复后：{res['修复后文本']}") print(f"错误详情：{res['错误详情']}") print(f"单条推理耗时：{res['单条推理耗时(ms)']} 毫秒") else: print(f"{res['单条推理耗时(ms)']}")
1、获取最新的qqwry.dat 可以从https://github.com/metowolf/qqwry.dat/releases 这里下载最新版 2、使用c++把 qqwry.dat转成文本 原文章有问题不能把内容全部提取完，调整为 #include <stdio.h> #define DAT_FILE "路径/qqwry.dat" #define OUT_FILE "./out.txt" FILE *out = NULL; void showIp(unsigned int ip) { char buffer[20] = {0}; sprintf(buffer, "%u.%u.%u.%u ", (ip >> 24) & 0xff, (ip >> 16) & 0xff, (ip >> 8) & 0xff, ip & 0xff); fputs(buffer, out); } void getString(FILE *fp) { char temp[128] = {0}; int index = 0; while ((temp[index] = fgetc(fp)) != 0) { index++; } fputs(temp, out); } // 重构 getInfo() 函数（非递归版本） void getInfo(FILE *fp) { unsigned char mod; long offset; long preOffset; int count = 0; // 记录已读取的字符串数量 while (count < 2) { mod = fgetc(fp); if (mod == 0x01 || mod == 0x02) { offset = 0; fread(&offset, 3, 1, fp); preOffset = ftell(fp); fseek(fp, offset, SEEK_SET); } else { fseek(fp, -1, SEEK_CUR); // 回退非重定向字节 getString(fp); // 读取字符串 count++; if (count == 1 && mod != 0x01 && mod != 0x02) { // 若第一个字符串后无重定向，继续读第二个 getString(fp); count++; } } if (mod == 0x02 && count == 1) { // 模式2且只读了一个字符串时，返回原位置继续 fseek(fp, preOffset, SEEK_SET); } } } void parseRecord(FILE *fp, long int recordOffset) { unsigned int endIp = 0; fseek(fp, recordOffset, SEEK_SET); //seek到数据区 fread(&endIp, 4, 1, fp); //读取结束IP showIp(endIp); //显示结束IP getInfo(fp); //获取信息 fputs("\n", out); } void parseDat(FILE *fp) { long int startIndexOffset = 0; long int endIndexOffset = 0; long int indexOffset = 0; long int recordOffset = 0; unsigned int startIp = 0; fseek(fp, 0L, SEEK_SET); //seek文件头 fread(&startIndexOffset, 4, 1, fp); //读索引开始偏移量 fread(&endIndexOffset, 4, 1, fp); //读索引结束偏移量 indexOffset = startIndexOffset; // 修改索引区遍历逻辑 while (indexOffset < endIndexOffset) // 使用 < 而不是 <= { fseek(fp, indexOffset, SEEK_SET); fread(&startIp, 4, 1, fp); recordOffset = 0; // 关键：清空高位字节 fread(&recordOffset, 3, 1, fp); indexOffset = ftell(fp); showIp(startIp); parseRecord(fp, recordOffset); } } int main(void) { FILE *fp = NULL; out = fopen(OUT_FILE, "wb"); fp = fopen(DAT_FILE, "rb"); if (fp == NULL ) { printf("open file error!\n"); return 0; } parseDat(fp); fclose(fp); return 0; } 3、去高德地图获取省市区json 行政区域查询-基础 API 文档-开发指南-Web服务 API 接口地址 https://restapi.amap.com/v3/config/district 只保留国家下面的省市区 保存好region.json 4、使用python把输出的out.txt 和区域json转成ip2region.txt 原文章去掉了国外，我这里做了保留，并做了调整，可根据自己的需求做调整，调整了out.txt读取编码 不需要区域代码也可以不用region.json，可根据自己的需要调整代码 ip2region.txt的文本格式需要是：ip|ip|国家|区域|省份|城市|ISP import json,re # 加载省市县数据 with open("region.json", 'r', encoding='utf-8') as file: data = json.load(file) print("---省市县加载完成---") with open('ip2region.txt', 'w', encoding='utf-8') as outfile: with open('out.txt', 'r', encoding='gbk') as file: line = file.readline() while line: line2 = line.split(" ") area = line2[2].replace("\n","").replace("CZ88.NET","").split("–") containProvince = None containCity = None containCounty = None # 匹配省 for province in data: if(len(area)>1): if re.search(province["name"], area[1]) is not None: # 找到了省 containProvince = province # 有第二级地区 if(len(area)>2): for city in province["districts"]: if re.search(city["name"], area[2]) is not None: # 找到了省 containCity = city if(len(area)>3): for county in city["districts"]: if re.search(county["name"], area[3]) is not None: # 找到了省 containCounty = county if containProvince is None: #没有省的数据 跳过 #continue l = len(area) if l==1: result=line2[0]+"|"+line2[1]+"|"+area[0]+"|0|0|0|0" elif l==2: result=line2[0]+"|"+line2[1]+"|"+area[0]+"|0|"+area[1]+"|0|0" elif l==3: result=line2[0]+"|"+line2[1]+"|"+area[0]+"|0|"+area[1]+"|"+area[2]+"|0" elif l>=4: result=line2[0]+"|"+line2[1]+"|"+area[0]+"|0|"+area[1]+"|"+area[2]+"|"+area[l-1] else: result=line2[0]+"|"+line2[1]+"|"+area[0]+"|"+str(containProvince["adcode"])[:2]+"|" # if containCity is not None: # result = result +"|"+containCity["name"]+"|"+containCity["adcode"] # if containCounty is not None: # result = result +"|"+containCounty["name"]+"|"+containCounty["adcode"] l = len(area) if l==2: result=result+area[1]+"|0|0" elif l==3: result=result+area[1]+"|"+area[2]+"|0" elif l>=4: result=result+area[1]+"|"+area[2]+"|"+area[l-1] # result.replace("\n","") result=result+"\n" outfile.write(result) line = file.readline() 执行命令 python ipdo.py ip2region.txt生成有124M 5、把ip2region.txt转成ip2region.xdb 使用https://gitee.com/lionsoul/ip2region/tree/master/maker/python 这个实现，也按照原文章把make.py155-161行，上一行的结束IP和下一行的起始IP不连续，直接注释掉这段代码，main.py改名为genxdb.py 执行命令python genxdb.py gen --src=./ip2region.txt --dst=./ip2region.xdb 生成ip2region.xdb 大小有22.8M py项目结构如图 py项目结构 6、项目里直接使用 java为例： ip2region.external=false ip2region.location=classpath:ip2region/ip2region.xdb logging.level.com.github.hiwepy.ip2region.spring.boot.ext=warn <dependency> <groupId>com.github.hiwepy</groupId> <artifactId>ip2region-spring-boot-starter</artifactId> <version>2.7.x.20240823.RELEASE</version> </dependency> @Autowired private IP2regionTemplate ip2regionTemplate; String region=ip2regionTemplate.getRegion(ip); 最后结果类似 中国|33|浙江|宁波Apple|0，经测试可用
导入脚本../mongodb/bin/mongorestore --uri "mongodb://admin:hudongtang2020@localhost:27011/" --drop --dir tmp/backup/data
#include <WiFiManager.h> #include <WiFi.h> #include <TFT_eSPI.h> #include <PubSubClient.h> #include "esp_system.h" #include <Arduino.h> #include <TJpg_Decoder.h> #include <HTTPClient.h> #include <src/CommonUtils.h> TFT_eSPI tft = TFT_eSPI(); // 创建TFT对象 WiFiManager wifiManager; unsigned long buttonPressTime = 0; bool isConnected = false; // 正在连接标记 bool resetTriggered = false; #define BOOT_BUTTON_PIN 0 // 显示状态常量 #define STATUS_WIFI_CONNECTING 0 #define STATUS_WIFI_CONNECTED 1 #define STATUS_WIFI_DISCONNECTED 2 #define STATUS_MQTT_CONNECTING 3 #define STATUS_MQTT_CONNECTED 4 #define STATUS_DOWNLOADING 5 #define STATUS_DISPLAYING_IMAGE 6 #define STATUS_PLAYING_MJPEG 7 int currentStatus = STATUS_WIFI_CONNECTING; // MQTT 配置 const char* mqttServer = "xxxx"; const int mqttPort = 1883; const char* mqttUser = "xxxx"; // 无需认证留空 const char* mqttPassword = "xxxx"; const char* clientId = "arduino_client_1232"; // 客户端ID需唯一 WiFiClient espClient; PubSubClient client(espClient); bool isVideo = false; // 内存缓冲区 uint8_t* jpegBuffer = NULL; // JPEG原始数据缓冲区 (PSRAM) size_t jpegSize = 0; // JPEG图片大小 uint8_t* mjpegBuffer = NULL; // MJPEG文件缓冲区 (PSRAM) size_t mjpegSize = 0; // MJPEG文件大小 size_t mjpegPosition = 0; // 当前MJPEG文件位置 // 帧缓冲区 uint8_t* frameBuffer = NULL; // 单帧数据缓冲区 size_t frameSize = 0; // 当前帧大小 // 流缓冲区 uint8_t* streamBuffer = NULL; // 流数据缓冲区 size_t streamBufferSize = 0; // 流缓冲区大小 size_t streamBufferPos = 0; // 流缓冲区位置 // 定义常量以提高可读性和可维护性 #define BUFFER_SIZE 4096 // 增大缓冲区大小 #define DOWNLOAD_TIMEOUT 15000 // 下载超时时间(毫秒) #define DATA_TIMEOUT 5000 // 数据超时时间(毫秒) #define PROGRESS_INTERVAL 500 // 进度显示间隔(毫秒) // 帧处理相关 #define MAX_FRAME_SIZE (1024 * 1024) // 最大帧大小1MB #define STREAM_BUFFER_SIZE (16 * 1024) // 流缓冲区大小16KB #define FRAME_DELAY 40 // 帧延迟(毫秒)，影响播放速度 // 默认MJPEG URL const char* defaultMjpegUrl = "https://chenhuan-esp.oss-cn-beijing.aliyuncs.com/output.avi"; // 当前播放状态 bool isPlayingMjpeg = false; unsigned long lastFrameTime = 0; unsigned long frameCount = 0; unsigned long mjpegStartTime = 0; // TJpgDec解码回调函数 - 用于将解码后的块数据绘制到屏幕 bool tft_output(int16_t x, int16_t y, uint16_t w, uint16_t h, uint16_t* bitmap) { // 检查块是否在屏幕范围内 if (y >= TFT_HEIGHT) return 0; // 绘制到TFT屏幕 tft.pushImage(x, y, w, h, bitmap); return 1; // 返回1继续解码，返回0终止解码 } // 向屏幕写入状态信息 void writeToScreen(const char* message, bool clearScreen = true) { if (clearScreen) { tft.fillScreen(TFT_BLACK); tft.setCursor(0, 0); } tft.printf(message); } // 更新屏幕状态显示 void updateStatusDisplay() { tft.fillScreen(TFT_BLACK); tft.setCursor(10, 10); tft.setTextColor(TFT_WHITE); tft.setTextSize(1); // 显示基本信息 tft.print("ESP32 MJPEG Player - "); switch (currentStatus) { case STATUS_WIFI_CONNECTING: tft.println("WiFi Connecting"); break; case STATUS_WIFI_CONNECTED: tft.println("WiFi Connected"); break; case STATUS_WIFI_DISCONNECTED: tft.println("WiFi Disconnected"); break; case STATUS_MQTT_CONNECTING: tft.println("MQTT Connecting"); break; case STATUS_MQTT_CONNECTED: tft.println("MQTT Connected"); break; case STATUS_DOWNLOADING: tft.println("Downloading..."); break; case STATUS_DISPLAYING_IMAGE: tft.println("Displaying Image"); break; case STATUS_PLAYING_MJPEG: tft.println("Playing MJPEG"); break; } // 显示WiFi信息 tft.setCursor(10, 30); if (WiFi.status() == WL_CONNECTED) { tft.print("WiFi: "); tft.print(WiFi.SSID()); tft.print(" ("); tft.print(WiFi.RSSI()); tft.println(" dBm)"); } else { tft.println("WiFi: Not Connected"); } // 显示MQTT信息 tft.setCursor(10, 50); if (client.connected()) { tft.println("MQTT: Connected"); } else { tft.println("MQTT: Not Connected"); } // 显示MJPEG播放信息 if (isPlayingMjpeg) { tft.setCursor(10, 70); float fps = frameCount / ((millis() - mjpegStartTime) / 1000.0); tft.print("FPS: "); tft.print(fps, 1); tft.print(" Frame: "); tft.print(frameCount); } } // 解码并显示JPEG图片 bool decodeAndDisplayJpeg(uint8_t* buffer, size_t size) { if (!buffer || size == 0) { Serial.println("错误：JPEG缓冲区为空"); return false; } Serial.println("开始解码JPEG图片..."); // 注册TJpgDec输出回调函数 TJpgDec.setCallback(tft_output); // 调用TJpgDec解码函数 bool result = TJpgDec.drawJpg(0, 0, buffer, size); if (result) { Serial.println("JPEG图片显示成功"); } return result; } // 检查WiFi状态并更新 void checkWifiStatus() { // 按钮检测（保持原功能） if (digitalRead(BOOT_BUTTON_PIN) == LOW && !resetTriggered) { delay(50); if (digitalRead(BOOT_BUTTON_PIN) == LOW) { buttonPressTime = millis(); resetTriggered = true; while (digitalRead(BOOT_BUTTON_PIN) == LOW) { if (millis() - buttonPressTime > 3000) { Serial.println("\n重置WiFi配置..."); writeToScreen("Resetting WiFi configuration..."); wifiManager.resetSettings(); delay(2000); ESP.restart(); } } resetTriggered = false; } } // WiFi状态监测 if (WiFi.status() != WL_CONNECTED && isConnected) { isConnected = false; currentStatus = STATUS_WIFI_DISCONNECTED; Serial.println("WiFi连接丢失，尝试重连..."); writeToScreen("WiFi connection lost, trying to reconnect..."); } else if (WiFi.status() == WL_CONNECTED && !isConnected) { isConnected = true; currentStatus = STATUS_WIFI_CONNECTED; Serial.println("WiFi已连接"); writeToScreen("WiFi is connected"); } } // 重连MQTT服务器 void reconnectMQTT() { if (WiFi.status() != WL_CONNECTED) return; currentStatus = STATUS_MQTT_CONNECTING; updateStatusDisplay(); while (!client.connected()) { Serial.print("尝试连接MQTT服务器..."); if (client.connect(clientId, mqttUser, mqttPassword)) { Serial.println("连接成功"); currentStatus = STATUS_MQTT_CONNECTED; updateStatusDisplay(); // 订阅主题 char topic[100]; sprintf(topic, "%s/%s", clientId, "downloadImages"); client.subscribe(topic); sprintf(topic, "%s/%s", clientId, "downloadVideos"); client.subscribe(topic); sprintf(topic, "%s/%s", clientId, "playMjpeg"); client.subscribe(topic); } else { Serial.print("连接失败，状态码="); Serial.print(client.state()); Serial.println(" 重试中..."); delay(2000); } } } // 核心下载函数（下载到PSRAM） bool downloadToPSRAM(const String& url, uint8_t*& buffer, size_t& size, bool showProgress = true) { // 初始化HTTP客户端 HTTPClient http; http.setReuse(true); // 启用连接复用 http.begin(url); // 设置超时和重试 http.setTimeout(DOWNLOAD_TIMEOUT); http.setFollowRedirects(HTTPC_STRICT_FOLLOW_REDIRECTS); // 执行请求 int httpCode = http.GET(); if (httpCode != HTTP_CODE_OK) { Serial.printf("HTTP错误: %d\n", httpCode); http.end(); return false; } // 获取数据流和内容长度 WiFiClient* stream = http.getStreamPtr(); size_t contentLength = http.getSize(); size_t totalSize = 0; // 确保PSRAM缓冲区已分配（如果未分配则分配） if (!buffer) { buffer = (uint8_t*)ps_malloc(contentLength); if (!buffer) { Serial.println("PSRAM分配失败"); http.end(); return false; } size = contentLength; } else if (contentLength > size) { // 如果已有缓冲区不够大，则重新分配 free(buffer); buffer = (uint8_t*)ps_malloc(contentLength); if (!buffer) { Serial.println("PSRAM重新分配失败"); http.end(); return false; } size = contentLength; } unsigned long timeoutMark = millis(); unsigned long lastProgressTime = 0; // 更新状态显示 currentStatus = STATUS_DOWNLOADING; updateStatusDisplay(); // 核心下载循环 while (http.connected() && (totalSize < contentLength || contentLength == 0)) { size_t available = stream->available(); if (available > 0) { // 计算本次读取大小 size_t readSize = min(available, contentLength - totalSize); // 读取数据到PSRAM缓冲区 size_t readBytes = stream->readBytes(buffer + totalSize, readSize); if (readBytes > 0) { totalSize += readBytes; // 重置超时计数器 timeoutMark = millis(); // 定期显示进度 if (showProgress && millis() - lastProgressTime > PROGRESS_INTERVAL) { if (contentLength > 0) { Serial.printf("进度: %.1f%% (%zu/%zu bytes)\r", (totalSize * 100.0) / contentLength, totalSize, contentLength); } else { Serial.printf("已下载: %zu bytes\r", totalSize); } // 更新进度显示 tft.setCursor(10, 90); tft.fillRect(10, 90, 220, 20, TFT_BLACK); if (contentLength > 0) { tft.printf("Downloading: %.1f%%", (totalSize * 100.0) / contentLength); } else { tft.printf("Downloading: %zu bytes", totalSize); } lastProgressTime = millis(); } } } else if (millis() - timeoutMark > DATA_TIMEOUT) { // 数据超时 Serial.println("\n下载超时"); break; } } // 关闭资源 http.end(); // 最终校验 if (contentLength > 0 && totalSize != contentLength) { Serial.printf("\n下载不完整 (%zu/%zu bytes)\n", totalSize, contentLength); // 释放不完整数据的缓冲区 if (buffer) { free(buffer); buffer = NULL; size = 0; } return false; } if (showProgress) { Serial.println("\n下载成功"); } // 更新实际大小 size = totalSize; return true; } // 查找MJPEG中的下一帧 bool findNextFrame() { if (!mjpegBuffer || mjpegPosition >= mjpegSize) { return false; } // JPEG帧以0xFF 0xD8开始，以0xFF 0xD9结束 const uint8_t JPEG_START[] = {0xFF, 0xD8}; const uint8_t JPEG_END[] = {0xFF, 0xD9}; // 寻找下一帧的开始 size_t frameStart = mjpegPosition; while (frameStart < mjpegSize - 1) { if (mjpegBuffer[frameStart] == JPEG_START[0] && mjpegBuffer[frameStart + 1] == JPEG_START[1]) { break; } frameStart++; } if (frameStart >= mjpegSize - 1) { // 没有找到更多帧 return false; } // 寻找帧的结束 size_t frameEnd = frameStart + 2; while (frameEnd < mjpegSize - 1) { if (mjpegBuffer[frameEnd] == JPEG_END[0] && mjpegBuffer[frameEnd + 1] == JPEG_END[1]) { frameEnd += 2; // 包括结束标记 break; } frameEnd++; } if (frameEnd >= mjpegSize) { // 不完整的帧 return false; } // 设置帧数据和大小 if (frameBuffer) { free(frameBuffer); } frameSize = frameEnd - frameStart; frameBuffer = (uint8_t*)ps_malloc(frameSize); if (!frameBuffer) { Serial.println("无法分配帧缓冲区"); return false; } memcpy(frameBuffer, mjpegBuffer + frameStart, frameSize); // 更新当前位置为下一帧开始 mjpegPosition = frameEnd; return true; } // 播放MJPEG文件 void playMjpeg() { if (!mjpegBuffer || mjpegSize == 0) { Serial.println("没有MJPEG数据可播放"); return; } Serial.println("开始播放MJPEG..."); isPlayingMjpeg = true; currentStatus = STATUS_PLAYING_MJPEG; frameCount = 0; mjpegStartTime = millis(); // 重置位置以开始播放 mjpegPosition = 0; // 主播放循环 while (isPlayingMjpeg && findNextFrame()) { // 检查WiFi连接 if (WiFi.status() != WL_CONNECTED) { Serial.println("WiFi断开，暂停MJPEG播放"); isPlayingMjpeg = false; break; } // 处理MQTT消息 client.loop(); // 计算帧延迟 unsigned long currentTime = millis(); if (currentTime - lastFrameTime < FRAME_DELAY) { delay(FRAME_DELAY - (currentTime - lastFrameTime)); } lastFrameTime = millis(); // 更新状态显示 if (frameCount % 10 == 0) { //updateStatusDisplay(); } // 显示当前帧 decodeAndDisplayJpeg(frameBuffer, frameSize); frameCount++; // 释放帧缓冲区 if (frameBuffer) { free(frameBuffer); frameBuffer = NULL; } } // 释放资源 if (frameBuffer) { free(frameBuffer); frameBuffer = NULL; } isPlayingMjpeg = false; Serial.printf("MJPEG播放完成，共播放 %lu 帧\n", frameCount); } // 消息回调函数（收到订阅消息时触发） void callback(char* topic, byte* payload, unsigned int length) { Serial.print("收到消息 ["); Serial.print(topic); Serial.print("]: "); String message = ""; for (unsigned int i = 0; i < length; i++) { message += (char)payload[i]; } Serial.println(message); // 处理接收到的消息 if (CommonUtils::endsWith(topic, "downloadImages")) { // 下载并显示图片 currentStatus = STATUS_DOWNLOADING; updateStatusDisplay(); bool downloadSuccess = downloadToPSRAM(message, jpegBuffer, jpegSize, true); if (downloadSuccess) { currentStatus = STATUS_DISPLAYING_IMAGE; updateStatusDisplay(); // 解码并显示JPEG图片 decodeAndDisplayJpeg(jpegBuffer, jpegSize); } else { tft.fillScreen(TFT_RED); tft.setCursor(10, 10); tft.setTextColor(TFT_WHITE); tft.println("Download Failed"); } isVideo = false; } else if (CommonUtils::endsWith(topic, "downloadVideos")) { // 下载视频（可能是MJPEG格式） currentStatus = STATUS_DOWNLOADING; updateStatusDisplay(); bool downloadSuccess = downloadToPSRAM(message, mjpegBuffer, mjpegSize, true); if (downloadSuccess) { Serial.println("视频下载成功，使用playMjpeg命令播放"); tft.fillScreen(TFT_GREEN); tft.setCursor(10, 10); tft.setTextColor(TFT_WHITE); tft.println("Video Downloaded"); tft.setCursor(10, 30); tft.println("Use playMjpeg to play"); } else { tft.fillScreen(TFT_RED); tft.setCursor(10, 10); tft.setTextColor(TFT_WHITE); tft.println("Download Failed"); } isVideo = true; } else if (CommonUtils::endsWith(topic, "playMjpeg")) { // 播放MJPEG if (message == "default" || message == "") { // 播放默认MJPEG if (!mjpegBuffer || mjpegSize == 0) { // 如果没有下载过，先下载 currentStatus = STATUS_DOWNLOADING; updateStatusDisplay(); bool downloadSuccess = downloadToPSRAM(defaultMjpegUrl, mjpegBuffer, mjpegSize, true); if (!downloadSuccess) { tft.fillScreen(TFT_RED); tft.setCursor(10, 10); tft.setTextColor(TFT_WHITE); tft.println("Download Failed"); return; } } // 播放MJPEG playMjpeg(); } else { // 从指定URL下载并播放MJPEG currentStatus = STATUS_DOWNLOADING; updateStatusDisplay(); bool downloadSuccess = downloadToPSRAM(message, mjpegBuffer, mjpegSize, true); if (downloadSuccess) { // 播放MJPEG playMjpeg(); } else { tft.fillScreen(TFT_RED); tft.setCursor(10, 10); tft.setTextColor(TFT_WHITE); tft.println("Download Failed"); } } isVideo = true; } } void setup() { Serial.begin(115200); delay(1000); // 等待串口连接稳定 pinMode(BOOT_BUTTON_PIN, INPUT_PULLUP); // 初始化显示屏 tft.init(); tft.setRotation(0); tft.setSwapBytes(true); tft.fillScreen(TFT_BLUE); // 控制背光 (如果需要) pinMode(TFT_BL, OUTPUT); digitalWrite(TFT_BL, HIGH); // 点亮背光 // 显示初始信息 tft.setCursor(10, 10); tft.setTextColor(TFT_WHITE); tft.setTextSize(2); tft.println("ESP32 MJPEG Player"); // 显示PSRAM信息 uint32_t psramSize = ESP.getPsramSize(); uint32_t freePsram = ESP.getFreePsram(); Serial.printf("PSRAM 可用: %u 字节\n", psramSize); Serial.printf("PSRAM 空闲: %u 字节\n", freePsram); tft.setTextSize(1); tft.setCursor(10, 40); tft.printf("PSRAM: %u KB 空闲\n", freePsram / 1024); // 连接WiFi Serial.println("开始连接WiFi..."); tft.setCursor(10, 60); tft.print("Connecting to WiFi..."); if (!wifiManager.autoConnect("ESP32-Config")) { Serial.println("进入配置模式"); tft.fillScreen(TFT_RED); tft.setCursor(10, 10); tft.setTextColor(TFT_WHITE); tft.println("WiFi Config Mode"); tft.setCursor(10, 30); tft.println("Connect to ESP32-Config"); tft.setCursor(10, 50); tft.println("and open 192.168.4.1"); while (1); // 停止执行 } // WiFi连接成功 isConnected = true; currentStatus = STATUS_WIFI_CONNECTED; Serial.println("WiFi已连接"); tft.fillScreen(TFT_GREEN); tft.setCursor(10, 10); tft.setTextColor(TFT_WHITE); tft.println("WiFi Connected"); tft.setCursor(10, 30); tft.print("IP: "); tft.println(WiFi.localIP()); delay(2000); // 初始化MQTT client.setServer(mqttServer, mqttPort); client.setCallback(callback); // 设置消息回调函数 // 分配内存缓冲区 jpegBuffer = (uint8_t*)ps_malloc(1024 * 1024); // 1MB JPEG缓冲区 if (!jpegBuffer) { Serial.println("JPEG缓冲区分配失败!"); } else { Serial.println("JPEG缓冲区分配成功"); } mjpegBuffer = (uint8_t*)ps_malloc(1024 * 1024 * 4); // 4MB MJPEG缓冲区 if (!mjpegBuffer) { Serial.println("MJPEG缓冲区分配失败!"); } else { Serial.println("MJPEG缓冲区分配成功"); } // 初始化TJpgDec TJpgDec.setJpgScale(1); // 缩放比例 TJpgDec.setCallback(tft_output); // 设置输出回调函数 // 显示初始状态 updateStatusDisplay(); Serial.println("程序初始化完成"); } void loop() { // 检查WiFi状态 checkWifiStatus(); // 保持MQTT连接 if (isConnected && !client.connected()) { reconnectMQTT(); } // 处理MQTT消息 if (client.connected()) { client.loop(); } // 显示状态更新 static unsigned long lastStatusUpdate = 0; if (millis() - lastStatusUpdate > 5000) { //updateStatusDisplay(); lastStatusUpdate = millis(); } // 如果正在播放MJPEG，则继续播放 if (!isPlayingMjpeg&&frameCount!=0&&isVideo==true) { playMjpeg(); } delay(10); // 短暂延迟，避免CPU占用过高 }
git clone https://github.com/emqx/emqx-usage-example cd emqx-usage-example/mqtt-lb-nginx/ docker compose up -d 部署验证部署成功我们可以访问服务器的 18083 端口，访问 MQTT 集群的 dashboard, 默认账号是 admin，密码是 public，登录以后会提示你修改密码连接验证我们可以使用 MQTT cli 进行连接验证，下面是验证方法安装curl -LO https://www.emqx.com/en/downloads/MQTTX/v1.9.10/mqttx-cli-linux-x64 install ./mqttx-cli-linux-x64 /usr/local/bin/mqttx 发起 10 个连接mqtt bench conn -c 10 -h loclhost -p 1883