app.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, Body, Query
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
from gradio_client import Client, handle_file
import requests
from PIL import Image
from utils.model import load_detection_model, load_structure_model
from utils.preprocessing import prepare_image, prepare_cropped_image
from utils.detection import (
    detect_tables,
    detect_cells,
    outputs_to_objects,
    objects_to_crops,
    get_cell_coordinates_by_row,
)
from utils.visualization import visualize_detected_tables, plot_results
from utils.ocr import apply_ocr, save_csv
from utils.title_scoring import score_article_titles
from utils.summarization import summarize_text
from utils.git import get_git_commit_message


# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- 配置部分 ---
MODEL_PATH = "best_model.pth"
MAP_PATH = "char_map.json"
hf_token = os.getenv("read_token")
FEISHU_WEBHOOK_URL = os.getenv("FEISHU_WEBHOOK_URL")
WEB_HOOK = os.getenv("WEB_HOOK")


def check_hf_token(token):
    headers = {"Authorization": f"Bearer {token}"}
    # 访问 HF 的身份验证接口
    response = requests.get("https://huggingface.co/api/whoami-v2", headers=headers)

    if response.status_code == 200:
        user_info = response.json()
        logger.info(f"✅ Token 有效！")
        logger.info(f"👤 用户名: {user_info.get('name')}")
        logger.info(f"🔑 权限类型: {user_info.get('auth', {}).get('type', 'N/A')}")

        # 向飞书 webhook 发送消息
        if WEB_HOOK:
            try:
                # 获取Git commit message
                commit_message = get_git_commit_message()

                # 构建消息内容
                message = f"你好，{user_info.get('name')}，我已启动"
                if commit_message:
                    message += f"\n\{commit_message}"

                webhook_headers = {"Content-Type": "application/json"}
                webhook_data = {
                    "msg_type": "text",
                    "content": {"text": message},
                }
                webhook_response = requests.post(
                    WEB_HOOK, json=webhook_data, headers=webhook_headers
                )
                logger.info(f"✅ 飞书消息发送成功: {webhook_response.status_code}")
            except Exception as e:
                logger.error(f"❌ 飞书消息发送失败: {e}")
    else:
        logger.error(f"❌ Token 无效或已过期。")
        logger.error(f"🚫 错误代码: {response.status_code}")
        logger.error(f"📄 返回信息: {response.text}")


# --- 全局初始化模型（启动时加载一次） ---
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()
    check_hf_token(hf_token)
    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))


@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}


# 全局初始化表格检测模型
detection_model, device = load_detection_model()
structure_model = load_structure_model(device)

# 初始化翻译器
from transformers import pipeline

# 中文到英文翻译器
translator_zh_en = pipeline("translation", model="Helsinki-NLP/opus-mt-zh-en")
# 英文到中文翻译器
translator_en_zh = pipeline("translation", model="Helsinki-NLP/opus-mt-en-zh")

# 初始化文本转语音模型
from fastapi.responses import FileResponse
import os
import wave
import numpy as np

tts = pipeline("text-to-speech", model="facebook/mms-tts-eng")


def add_silence_between_sentences(audio_data, sampling_rate, silence_duration_ms=500):
    """
    在句子之间添加静音，增加停顿时间。

    - audio_data: 原始音频数据 (numpy array)
    - sampling_rate: 采样率
    - silence_duration_ms: 静音持续时间（毫秒），默认500ms
    """
    import re

    # 这里假设句子之间已经有某种分隔，我们通过检测音量低谷来分段
    # 或者更简单：返回一个可以拼接静音的函数

    # 计算静音样本数
    silence_samples = int(sampling_rate * silence_duration_ms / 1000)
    silence = np.zeros(silence_samples, dtype=np.float32)

    return silence


def split_text_into_sentences(text):
    """将文本按句子分割，保留标点符号"""
    import re
    # 按句号、问号、感叹号分割，但保留分隔符
    sentences = re.findall(r'[^.!?。！？]+[.!?。！？]*', text.strip())
    return [s.strip() for s in sentences if s.strip()]


def generate_tts_with_pauses(text, silence_duration_ms=500):
    """
    生成带停顿的TTS音频。

    - text: 输入文本
    - silence_duration_ms: 句子间停顿时间（毫秒）
    """
    sentences = split_text_into_sentences(text)

    if len(sentences) <= 1:
        # 只有一句话，直接生成
        return tts(text)

    # 逐句生成音频
    audio_segments = []
    sampling_rate = None

    for sentence in sentences:
        if sentence:
            result = tts(sentence)
            audio_segments.append(result["audio"])
            if sampling_rate is None:
                sampling_rate = result["sampling_rate"]

    # 拼接音频，在之间添加静音
    silence_samples = int(sampling_rate * silence_duration_ms / 1000)
    silence = np.zeros(silence_samples, dtype=np.float32)

    combined_audio = []
    for i, segment in enumerate(audio_segments):
        # 统一转换为 1D numpy 数组，避免不同环境返回的格式不一致（如 torch Tensor / numpy 多维数组）
        segment = np.asarray(segment).flatten()
        combined_audio.append(segment)
        if i < len(audio_segments) - 1:  # 最后一段后面不加静音
            combined_audio.append(silence)

    final_audio = np.concatenate(combined_audio)

    return {"audio": final_audio, "sampling_rate": sampling_rate}


@app.post("/upload/text")
async def upload_text(file: UploadFile = File(...)):
    """
    上传图片进行表格检测和OCR文本识别。

    - **file**: 需要识别的图片文件 (jpg, png, etc.)
    """
    try:
        # 读取上传的图片
        image_bytes = await file.read()
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

        # Add "no object" to id2label if not present
        if len(detection_model.config.id2label) not in detection_model.config.id2label:
            detection_model.config.id2label[len(detection_model.config.id2label)] = (
                "no object"
            )

        # Detect tables
        pixel_values = prepare_image(image, device)
        outputs = detect_tables(detection_model, pixel_values)
        objects = outputs_to_objects(
            outputs, image.size, detection_model.config.id2label
        )

        # Crop the table
        tokens = []
        detection_class_thresholds = {
            "table": 0.5,
            "table rotated": 0.5,
            "no object": 10,
        }
        tables_crops = objects_to_crops(
            image, tokens, objects, detection_class_thresholds, padding=0
        )

        if not tables_crops:
            return JSONResponse(
                content={"error": "No tables detected in the image"}, status_code=400
            )

        cropped_table = tables_crops[0]["image"].convert("RGB")

        # Detect cells in the cropped table
        pixel_values = prepare_cropped_image(cropped_table, device)
        outputs = detect_cells(structure_model, pixel_values)
        cells = outputs_to_objects(
            outputs, cropped_table.size, structure_model.config.id2label
        )

        # Apply OCR
        cell_coordinates = get_cell_coordinates_by_row(cells)
        data = apply_ocr(cell_coordinates, cropped_table)

        # Save results as CSV
        save_csv(data)

        return JSONResponse(
            content={
                "tables": objects,
                "cells": cells,
                "extracted_data": data,
                "message": "CSV file saved as output.csv",
            }
        )
    except Exception as e:
        logger.error(f"Table OCR error: {e}")
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")


@app.get("/translate")
async def translate(
    text: str = Query(..., description="Text to translate"),
    toLang: str = Query("English", description="Target language: English or Chinese"),
):
    """
    翻译文本到指定语言。

    - **text**: 需要翻译的文本
    - **toLang**: 目标语言 (English 或 Chinese)
    """
    try:
        if toLang.lower() == "chinese":
            # 英文到中文翻译
            translation = translator_en_zh(text)
        else:
            # 中文到英文翻译
            translation = translator_zh_en(text)
        return {"translation": translation}
    except Exception as e:
        logger.error(f"Translation error: {e}")
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")


@app.get("/text-to-speech")
async def text_to_speech(
    text: str = Query(..., description="Text to convert to speech"),
    pause_ms: int = Query(500, description="Pause duration between sentences in milliseconds (default: 500ms)")
):
    """
    将文本转换为语音并返回音频文件。

    - **text**: 需要转换为语音的文本
    - **pause_ms**: 句子之间的停顿时间（毫秒），默认500毫秒
    """
    try:
        # 生成音频，带句子间停顿
        output = generate_tts_with_pauses(text, silence_duration_ms=pause_ms)

        # 使用内置的 wave 模块处理音频数据
        output_path = "output.wav"
        audio_data = output["audio"]
        sampling_rate = output["sampling_rate"]

        # 确保音频数据是正确的格式
        if isinstance(audio_data, np.ndarray):
            # 转换为 int16 格式
            audio_data = (audio_data * 32767).astype(np.int16)

            # 使用 wave 模块创建 WAV 文件
            with wave.open(output_path, "wb") as wf:
                wf.setnchannels(1)  # 单声道
                wf.setsampwidth(2)  # 16位
                wf.setframerate(sampling_rate)
                wf.writeframes(audio_data.tobytes())
        else:
            # 处理其他格式的音频数据
            raise HTTPException(status_code=500, detail="Unsupported audio data format")

        # 检查文件是否生成
        if not os.path.exists(output_path):
            raise HTTPException(status_code=500, detail="Failed to generate audio file")

        # 返回 WAV 文件
        return FileResponse(
            path=output_path, media_type="audio/wav", filename="output.wav"
        )
    except Exception as e:
        logger.error(f"Text-to-speech error: {e}")
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")


@app.post("/text/score")
async def score_titles(
    request_data: Dict[str, List[str]] = Body(
        ..., description="Object containing titles list"
    )
):
    """
    为文章标题列表打分，返回标题和对应的分数。

    - **titles**: 文章标题列表
    """
    try:
        # 获取标题列表
        titles = request_data.get("titles", [])

        # 调用打分函数
        scores = score_article_titles(titles)

        # 组合标题和分数
        results = [
            {"title": title, "score": score} for title, score in zip(titles, scores)
        ]

        # 按分数倒序排序
        results.sort(key=lambda x: x["score"], reverse=True)

        return {"results": results}
    except Exception as e:
        logger.error(f"Title scoring error: {e}")
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")


@app.post("/text/summary")
async def generate_summary(
    text: str = Body(..., description="Text to summarize"),
    max_length: int = Body(300, description="Maximum length of the summary"),
    min_length: int = Body(50, description="Minimum length of the summary"),
):
    """
    生成文本摘要。

    - **text**: 需要摘要的文本（建议为英文）
    - **max_length**: 摘要最大长度（token数）
    - **min_length**: 摘要最小长度（token数）
    """
    try:
        # 调用摘要函数
        summary = summarize_text(text, max_length=max_length, min_length=min_length)

        return {"summary": summary}
    except Exception as e:
        logger.error(f"Summarization error: {e}")
        raise HTTPException(status_code=500, detail=f"服务器内部错误: {e}")


from fastapi.staticfiles import StaticFiles

app.mount("/", StaticFiles(directory="static", html=True), name="web")


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="info")