From a3750ffce4cc63a34cb8367da5b00907bf4387bc Mon Sep 17 00:00:00 2001
From: Admin <admin@viaeon.com>
Date: Thu, 16 Jul 2026 19:51:13 +0000
Subject: [PATCH] =?UTF-8?q?feat:=20=E5=90=AF=E5=8A=A8=E6=97=B6=E8=87=AA?=
 =?UTF-8?q?=E5=8A=A8=E4=B8=8B=E8=BD=BD=E7=BC=BA=E5=A4=B1=E7=9A=84=E6=A8=A1?=
 =?UTF-8?q?=E5=9E=8B=E6=96=87=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 Dockerfile                  |  7 ++--
 internal/captcha/handler.go | 67 +++++++++++++++++++++++++++++++++++--
 2 files changed, 68 insertions(+), 6 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 0351504..cb02bce 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -33,13 +33,12 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
 
 WORKDIR /app
 
-# 复制二进制文件、前端和模型
+# 复制二进制文件和前端
 COPY --from=builder /app/anticaptcha .
 COPY --from=builder /app/web/dist ./web/dist
-COPY --from=builder /app/models ./models
 
-# 创建数据目录
-RUN mkdir -p /app/data
+# 创建数据和模型目录（模型启动时自动下载）
+RUN mkdir -p /app/data /app/models
 
 ENV GIN_MODE=release
 ENV TZ=Asia/Shanghai
diff --git a/internal/captcha/handler.go b/internal/captcha/handler.go
index 0ca4ed3..620bb82 100644
--- a/internal/captcha/handler.go
+++ b/internal/captcha/handler.go
@@ -1,6 +1,10 @@
 package captcha
 
 import (
+	"fmt"
+	"net/http"
+	"os"
+	"path/filepath"
 	"sync"
 
 	"anticaptcha/pkg/opencv"
@@ -11,10 +15,69 @@ type Handler struct {
 	mu        sync.RWMutex
 }
 
+// 模型配置
+var modelConfigs = map[string]string{
+	"math.onnx":    "[AntiCAP]-CRNN_Math.onnx",
+	"ocr.onnx":     "[Dddd]-OCR.onnx",
+	"rotate.onnx":  "[AntiCAP]-Rotation-RotNetR.onnx",
+	"siamese.onnx": "[AntiCAP]-Siamese-ResNet18.onnx",
+	"charsets.txt": "[Dddd]-CharSets.txt",
+}
+
+const modelBaseURL = "https://raw.githubusercontent.com/81NewArk/AntiCAP/main/AntiCAP/AntiCAP-Models"
+
 func NewHandler(modelPath string) *Handler {
-	return &Handler{
+	h := &Handler{
 		modelPath: modelPath,
 	}
+	// 确保模型目录存在并下载缺失的模型
+	h.ensureModels()
+	return h
+}
+
+// ensureModels 检查并下载缺失的模型
+func (h *Handler) ensureModels() {
+	// 确保目录存在
+	if err := os.MkdirAll(h.modelPath, 0755); err != nil {
+		fmt.Printf("警告: 创建模型目录失败: %v\n", err)
+		return
+	}
+
+	for localName, remoteName := range modelConfigs {
+		localPath := filepath.Join(h.modelPath, localName)
+		if _, err := os.Stat(localPath); os.IsNotExist(err) {
+			fmt.Printf("下载模型: %s -> %s\n", remoteName, localName)
+			if err := h.downloadModel(remoteName, localPath); err != nil {
+				fmt.Printf("警告: 下载模型 %s 失败: %v\n", remoteName, err)
+			} else {
+				fmt.Printf("模型下载完成: %s\n", localName)
+			}
+		}
+	}
+}
+
+// downloadModel 下载模型文件
+func (h *Handler) downloadModel(remoteName, localPath string) error {
+	url := fmt.Sprintf("%s/%s", modelBaseURL, remoteName)
+
+	resp, err := http.Get(url)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("HTTP %d", resp.StatusCode)
+	}
+
+	out, err := os.Create(localPath)
+	if err != nil {
+		return err
+	}
+	defer out.Close()
+
+	_, err = out.ReadFrom(resp.Body)
+	return err
 }
 
 // OCR 文字识别（需要 ONNX 模型）
@@ -143,4 +206,4 @@ func (h *Handler) DetectionIconOrder(orderImgBase64, targetImgBase64 string) ([]
 func (h *Handler) DetectionTextOrder(orderImgBase64, targetImgBase64 string) ([]map[string]int, error) {
 	// 暂时返回空结果
 	return []map[string]int{}, nil
-}
\ No newline at end of file
+}