# main.py import os import json import torch import torch.nn as nn from torchvision.models import resnet50 from torchvision.transforms import transforms from PIL import Image, ImageOps, ImageEnhance from typing import List, Dict from fastapi import FastAPI, File, UploadFile, HTTPException, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse import io import cv2 import numpy as np import hashlib import logging from PIL.ExifTags import TAGS # 设置日志 logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # --- 配置部分 --- MODEL_PATH = "best_model.pth" MAP_PATH = "char_map.json" # --- 全局初始化模型（启动时加载一次） --- class Recognizer: def __init__(self, model_path: str = MODEL_PATH, map_path: str = MAP_PATH): if not os.path.exists(model_path) or not os.path.exists(map_path): raise FileNotFoundError( f"模型文件 '{model_path}' 或字符映射表 '{map_path}' 不存在，" "请先运行 train_model.py 进行模型训练。" ) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 加载字符映射 with open(map_path, "r", encoding="utf-8") as f: self.char_to_idx = json.load(f) self.idx_to_char = {v: k for k, v in self.char_to_idx.items()} num_classes = len(self.char_to_idx) # 构建并加载模型 self.model = self._get_model(num_classes) try: ckpt = torch.load(model_path, map_location=self.device, weights_only=True) except Exception: ckpt = torch.load(model_path, map_location=self.device, weights_only=False) if isinstance(ckpt, nn.Module): state_dict = ckpt.state_dict() elif isinstance(ckpt, dict): state_dict = ckpt.get("state_dict", ckpt.get("model_state", ckpt)) else: raise ValueError("不支持的模型文件格式") self.model.load_state_dict(state_dict, strict=False) self.model.to(self.device) self.model.eval() # 定义图像预处理 self.transform = transforms.Compose( [ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ] ) def _get_model(self, num_classes: int) -> nn.Module: model = resnet50(weights=None) num_ftrs = model.fc.in_features model.fc = nn.Sequential(nn.Dropout(0.3), nn.Linear(num_ftrs, num_classes)) return model def preprocess_image(self, image_bytes: bytes, user_agent: str = "") -> Image.Image: """统一的图像预处理，特别处理移动设备上传的图片""" try: image = Image.open(io.BytesIO(image_bytes)) # 记录原始图像信息用于调试 logger.info( f"原始图像 - 格式: {image.format}, 模式: {image.mode}, 尺寸: {image.size}" ) # 处理EXIF方向信息（手机照片常有旋转问题） try: exif = image._getexif() if exif: for tag, value in exif.items(): decoded = TAGS.get(tag, tag) if decoded == "Orientation": if value == 3: image = image.rotate(180, expand=True) elif value == 6: image = image.rotate(270, expand=True) elif value == 8: image = image.rotate(90, expand=True) break except Exception as e: logger.warning(f"EXIF处理失败: {e}") # 转换为RGB if image.mode != "RGB": image = image.convert("RGB") # 检测是否为移动设备并应用增强处理 is_mobile = any( mobile_indicator in user_agent.lower() for mobile_indicator in ["mobile", "iphone", "android", "ipad"] ) if is_mobile: logger.info("检测到移动设备，应用增强预处理") # 增强对比度 enhancer = ImageEnhance.Contrast(image) image = enhancer.enhance(1.2) # 轻微锐化 enhancer = ImageEnhance.Sharpness(image) image = enhancer.enhance(1.1) return image except Exception as e: raise ValueError(f"无法处理图片: {e}") def assess_image_quality(self, image: Image.Image) -> Dict: """评估图像质量""" # 转换为numpy数组进行处理 img_array = np.array(image) if len(img_array.shape) == 3: gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY) else: gray = img_array # 计算清晰度（拉普拉斯方差） clarity = cv2.Laplacian(gray, cv2.CV_64F).var() # 计算亮度和对比度 brightness = np.mean(gray) contrast = np.std(gray) quality_info = { "clarity": float(clarity), "brightness": float(brightness), "contrast": float(contrast), "is_acceptable": clarity > 50 and 30 < brightness < 220, } logger.info(f"图像质量评估: {quality_info}") return quality_info def recognize( self, image_bytes: bytes, top_k: int = 5, user_agent: str = "" ) -> List[Dict[str, str]]: """ 识别上传的图像。 Args: image_bytes: 图片的二进制数据。 top_k: 返回前k个结果。 user_agent: 用户代理字符串，用于设备检测 Returns: 一个字典列表，每个字典包含 `char` 和 `prob` 键。 """ # 记录图像哈希用于调试 image_hash = hashlib.md5(image_bytes).hexdigest() logger.info(f"处理图像 - 哈希: {image_hash}, 设备: {user_agent}") # 预处理图像 image = self.preprocess_image(image_bytes, user_agent) # 评估图像质量 quality_info = self.assess_image_quality(image) if not quality_info["is_acceptable"]: logger.warning(f"图像质量可能影响识别: {quality_info}") # 应用模型预处理 image_tensor = self.transform(image).unsqueeze(0).to(self.device) with torch.no_grad(): outputs = self.model(image_tensor) probabilities = torch.nn.functional.softmax(outputs, dim=1) top_probs, top_indices = torch.topk(probabilities, top_k) results = [] top_probs_np = top_probs.cpu().numpy().flatten() top_indices_np = top_indices.cpu().numpy().flatten() for i in range(top_k): char_idx = top_indices_np[i] char_name = self.idx_to_char.get(char_idx, "?") probability = top_probs_np[i] results.append({"char": char_name, "prob": f"{probability:.2%}"}) logger.info(f"识别结果: {results}") return results # --- FastAPI 应用初始化 --- app = FastAPI( title="汉字书法字体识别", description="上传一张汉字图片，返回识别出的汉字及其置信度。", version="1.0.0", ) # 添加 CORS 中间件 app.add_middleware( CORSMiddleware, allow_origins=["*"], # 生产环境请替换为具体域名 allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # 初始化识别器（全局单例） try: recognizer = Recognizer() logger.info("模型加载成功，服务已启动") except Exception as e: logger.error(f"启动失败: {e}") recognizer = None # --- API 路由 --- @app.post("/upload", response_model=List[Dict[str, str]]) async def upload_image(request: Request, file: UploadFile = File(...)): """ 上传图片进行汉字识别。 - **file**: 需要识别的图片文件 (jpg, png, etc.) """ if not recognizer: raise HTTPException(status_code=503, detail="服务暂不可用，模型文件未找到。") # 检查文件类型 if not file.content_type.startswith("image/"): raise HTTPException(status_code=400, detail="上传的文件必须是图片格式。") try: image_bytes = await file.read() user_agent = request.headers.get("User-Agent", "") results = recognizer.recognize(image_bytes, top_k=5, user_agent=user_agent) return results except ValueError as ve: logger.error(f"图片处理错误: {ve}") raise HTTPException(status_code=400, detail=str(ve)) except Exception as e: logger.error(f"识别错误: {e}") raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}") # 调试接口 @app.post("/debug_upload") async def debug_upload(request: Request, file: UploadFile = File(...)): """调试接口，返回上传图片的详细信息""" try: image_bytes = await file.read() user_agent = request.headers.get("User-Agent", "") # 使用recognizer的预处理方法来分析图片 image = recognizer.preprocess_image(image_bytes, user_agent) quality_info = recognizer.assess_image_quality(image) debug_info = { "file_size": len(image_bytes), "file_hash": hashlib.md5(image_bytes).hexdigest(), "image_size": image.size, "image_mode": image.mode, "quality_assessment": quality_info, "user_agent": user_agent, } return JSONResponse(content=debug_info) except Exception as e: logger.error(f"调试接口错误: {e}") raise HTTPException(status_code=500, detail=str(e)) from fastapi.staticfiles import StaticFiles app.mount("/", StaticFiles(directory="static", html=True), name="web") @app.get("/") async def root(): return {"message": "欢迎使用汉字书法识别模型，请使用 POST /upload 接口上传图片。"} @app.get("/health") async def health_check(): """健康检查接口""" status = "healthy" if recognizer else "unhealthy" return {"status": status, "model_loaded": recognizer is not None} if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")