From 680870bed3a3ab0a16903a05de4c6d4600893cbe Mon Sep 17 00:00:00 2001 From: aaron <> Date: Sat, 8 Mar 2025 12:12:29 +0800 Subject: [PATCH] =?UTF-8?q?=E5=A2=9E=E5=8A=A0=20url=20ai=20=E6=8E=A5?= =?UTF-8?q?=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- app/api/endpoints/ai.py | 30 ++++++++++++++++++++++++++++-- app/api/endpoints/upload.py | 2 -- app/core/ai_client.py | 8 ++++---- app/core/qwen_client.py | 11 ++++------- 4 files changed, 36 insertions(+), 15 deletions(-) diff --git a/app/api/endpoints/ai.py b/app/api/endpoints/ai.py index 3c0415d..8f5709b 100644 --- a/app/api/endpoints/ai.py +++ b/app/api/endpoints/ai.py @@ -5,14 +5,40 @@ from app.api.deps import get_current_user from app.models.user import UserDB import logging from app.core.qcloud import qcloud_manager - +from pydantic import BaseModel router = APIRouter() +class ExtractPickupCodeRequest(BaseModel): + url: str + +@router.post("/extract_pickup_code_with_url", response_model=ResponseModel) +async def extract_pickup_code_with_url( + data: ExtractPickupCodeRequest +): + """从图片 URL 中提取取件码""" + + if not data.url: + return error_response(code=400, message="URL 不能为空") + + try: + # 调用 AI 客户端提取取件码 + result = await ai_client.extract_pickup_code(url=data.url) + if "error" in result: + return error_response(code=500, message=result.get("message", "提取取件码失败")) + + # 返回原始数据和格式化文本 + return success_response(data=result) + except Exception as e: + logging.exception(f"提取取件码失败: {str(e)}") + return error_response(code=500, message=f"提取取件码失败: {str(e)}") + + @router.post("/extract_pickup_code", response_model=ResponseModel) async def extract_pickup_code( file: UploadFile = File(...) ): """从图片中提取取件码""" + try: # 检查文件类型 if not file.content_type.startswith('image/'): @@ -22,7 +48,7 @@ async def extract_pickup_code( image_data = await file.read() # 调用 AI 客户端提取取件码 - result = await ai_client.extract_pickup_code(image_data) + result = await ai_client.extract_pickup_code(image_data=image_data) if "error" in result: return error_response(code=500, message=result.get("message", "提取取件码失败")) diff --git a/app/api/endpoints/upload.py b/app/api/endpoints/upload.py index 7a8741a..c5804cf 100644 --- a/app/api/endpoints/upload.py +++ b/app/api/endpoints/upload.py @@ -12,7 +12,6 @@ router = APIRouter() @router.post("/image", response_model=ResponseModel) async def upload_image( file: UploadFile = File(...), - current_user: UserDB = Depends(get_current_user) ): """上传单张图片""" if not file.content_type.startswith('image/'): @@ -27,7 +26,6 @@ async def upload_image( @router.post("/images", response_model=ResponseModel) async def upload_images( files: List[UploadFile] = File(...), - current_user: UserDB = Depends(get_current_user) ): """上传多张图片""" if len(files) > 5: diff --git a/app/core/ai_client.py b/app/core/ai_client.py index 130242b..da75bb9 100644 --- a/app/core/ai_client.py +++ b/app/core/ai_client.py @@ -12,7 +12,7 @@ class AIClient: def __init__(self): self.timeout = 15 # 请求超时时间（秒） - async def extract_pickup_code(self, image_data: bytes) -> Dict[str, Any]: + async def extract_pickup_code(self, image_data: bytes = None, url: str = None) -> Dict[str, Any]: """ 从图片中提取取件码 @@ -23,7 +23,7 @@ class AIClient: Dict: 提取结果，包含取件码信息 """ try: - primary_result = await self._extract_with_qwen(image_data) + primary_result = await self._extract_with_qwen(image_data, url) # 检查结果是否有效 if self._is_valid_result(primary_result): @@ -35,12 +35,12 @@ class AIClient: logging.exception(f"提取取件码异常: {str(e)}") return {"error": "处理失败", "message": str(e)} - async def _extract_with_qwen(self, image_data: bytes) -> Dict[str, Any]: + async def _extract_with_qwen(self, image_data: bytes = None, url: str = None) -> Dict[str, Any]: """使用千问提取取件码""" try: # 添加超时控制 return await asyncio.wait_for( - qwen_client.extract_pickup_code(image_data), + qwen_client.extract_pickup_code(image_data, url), timeout=self.timeout ) except asyncio.TimeoutError: diff --git a/app/core/qwen_client.py b/app/core/qwen_client.py index c239bd3..a1aa2ec 100644 --- a/app/core/qwen_client.py +++ b/app/core/qwen_client.py @@ -20,7 +20,7 @@ class QwenClient: self.api_key = settings.QWEN_API_KEY self.model = "qwen-vl-max" # 使用千问视觉语言大模型 - async def extract_pickup_code(self, imageData: bytes) -> Dict[str, Any]: + async def extract_pickup_code(self, imageData: bytes = None, url: str = None) -> Dict[str, Any]: """ 从图片中提取取件码 @@ -30,10 +30,7 @@ class QwenClient: Returns: Dict: 提取结果，包含取件码信息 """ - try: - # 将图片转换为 base64 - image_base64 = base64.b64encode(imageData).decode('utf-8') - + try: # 构建消息 messages = [ { @@ -48,8 +45,8 @@ class QwenClient: "text": "请识别图中驿站的所有取件码，以[{\"station\":\"驿站名字\",\"pickup_codes\":[\"3232\",\"2323\"]}]的格式返回。只返回JSON格式数据，不要其他解释。" }, { - "type": "image", - "image": f"data:image/jpeg;base64,{image_base64}" + "type": "image", + "image": f"data:image/jpeg;base64,{base64.b64encode(imageData).decode('utf-8')}" if imageData else url, } ] }