app.py

# Copyright 2025 lishiyan
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     https://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import shutil
import base64
import re
import json
import pytesseract
import openpyxl
import fitz  # PyMuPDF
from zai import ZhipuAiClient
from PIL import Image
from mimetypes import guess_type
from openpyxl.styles import numbers
from datetime import datetime
import gradio as gr
import tempfile
from pathlib import Path
import asyncio
from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
import functools
import traceback
# ====================== Hugging Face Tesseract 环境 ========================== 
import pytesseract
import os
import shutil

# 根据环境自动检测 Tesseract 路径
def setup_tesseract():
    # 尝试所有可能的 Tesseract 路径
    possible_paths = [
        '/usr/bin/tesseract',      # Hugging Face Spaces 或 Linux 标准路径
        '/usr/local/bin/tesseract', # Linux/macOS 本地安装
        '/opt/homebrew/bin/tesseract', # macOS Homebrew
        '/opt/tesseract/bin/tesseract', # 某些 Linux 发行版
        'tesseract',               # 系统PATH中
    ]
    
    # 首先尝试使用 which 命令查找 tesseract
    tesseract_path = shutil.which('tesseract')
    if tesseract_path:
        pytesseract.pytesseract.tesseract_cmd = tesseract_path
        print(f"Found Tesseract at: {tesseract_path}")
        return
    
    # 如果 which 找不到，尝试常见路径
    for path in possible_paths:
        if path != 'tesseract' and os.path.exists(path):
            pytesseract.pytesseract.tesseract_cmd = path
            print(f"Found Tesseract at: {path}")
            return
    
    # 如果是 'tesseract' 字符串，直接设置（让系统在PATH中查找）
    if 'tesseract' in possible_paths:
        pytesseract.pytesseract.tesseract_cmd = 'tesseract'
        print("Using 'tesseract' from system PATH")
        return
    
    # 如果都找不到，提供详细的错误信息
    print("Tesseract not found in any of the expected locations:")
    for path in possible_paths:
        print(f"  - {path}")
    print("\nTroubleshooting steps:")
    print("1. Ensure packages.txt contains 'tesseract-ocr'")
    print("2. Try rebuilding the Space")
    print("3. Check if the base image supports apt packages")
    
    # 不抛出错误，而是设置一个默认值，让 pytesseract 自己处理
    pytesseract.pytesseract.tesseract_cmd = 'tesseract'

# 设置 Tesseract
setup_tesseract()

# ====================== Module 1: PDF/Image to Markdown ======================
class PDFImageToMarkdown:
    def __init__(self, zhipu_client):
        self.client = zhipu_client

    def detect_text_direction(self, image):
        """使用 OCR 检测图片中文本的方向。"""
        ocr_result = pytesseract.image_to_osd(image)
        rotation_match = re.search(r'Rotate:\s+(\d+)', ocr_result)
        if rotation_match:
            return int(rotation_match.group(1))
        return 0

    def correct_image_orientation(self, image_path):
        """读取图片，根据 OCR 检测的方向信息摆正图片，并保存到指定路径。"""
        image = Image.open(image_path)
        rotation = self.detect_text_direction(image)
        if rotation != 0:
            while rotation != 0:
                image = image.rotate(rotation, expand=True)
                rotation = self.detect_text_direction(image)
            image.save(image_path)
            print(f"{image_path}: image corrected to the right direction")

    def extract_imginfo_from_zhipu(self, file_path):
        """通过模型获取图片信息"""
        with open(file_path, 'rb') as img_file:
            file_base = base64.b64encode(img_file.read()).decode('utf-8')
        
        prompt_txt = "Precisely identify the content within an image, specifically extracting all information and table content from an invoice. Convert the extracted information into Markdown format, and ensure the output is in Chinese."

        response = self.client.chat.completions.create(
            model="glm-4.5v",
            temperature=0,
            messages=[
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/png;base64,{file_base}"
                            }
                        },
                        {
                            "type": "text",
                            "text": prompt_txt
                        }
                    ]
                }
            ]
        )
        return response.choices[0].message.content

    def remove_empty_lines(self, input_file):
        """删除 markdown 中的空行和 markdown 标记"""
        temp_file = input_file + '.tmp'
        with open(input_file, 'r', encoding='utf-8') as infile, open(temp_file, 'w', encoding='utf-8') as outfile:
            for line in infile:
                clean_line = line.replace('```markdown', '').replace('```', '')
                if clean_line.strip():
                    outfile.write(clean_line)
        os.replace(temp_file, input_file)

    def detect_file_type(self, file_path):
        """检测文件类型，返回主要类型和子类型"""
        mime_type, _ = guess_type(file_path)
        if mime_type:
            return mime_type.split('/')
        ext = os.path.splitext(file_path)[1].lower()
        if ext in ['.pdf']:
            return ('application', 'pdf')
        elif ext in ['.jpg', '.jpeg']:
            return ('image', 'jpeg')
        elif ext in ['.png']:
            return ('image', 'png')
        else:
            raise ValueError(f"Unsupported file format: {file_path}")

    def process_image_file(self, image_path):
        """处理单个图片文件，校正方向并提取信息"""
        self.correct_image_orientation(image_path)
        return self.extract_imginfo_from_zhipu(image_path)

    def pdf_to_md(self, pdf_path, image_folder, finish_folder, dpi=300):
        """处理PDF文件，转换为Markdown文本"""
        pdf_document = fitz.open(pdf_path)
        pdf_name = os.path.splitext(os.path.basename(pdf_path))[0]
        all_text = ""
        
        for page_num in range(len(pdf_document)):
            page = pdf_document.load_page(page_num)
            zoom = dpi / 72.0
            mat = fitz.Matrix(zoom, zoom)
            pix = page.get_pixmap(matrix=mat)
            image_filename = f"{pdf_name}_page_{page_num + 1}.png"
            image_path = os.path.join(image_folder, image_filename)
            pix.save(image_path)
            image_text = self.process_image_file(image_path)
            all_text += image_text + '\n'
            shutil.move(image_path, os.path.join(finish_folder, image_filename))
        
        return all_text

    def file_to_md(self, input_path, image_folder, md_folder, finish_folder, dpi=300):
        """主处理函数，支持多种输入文件格式"""
        file_type, file_subtype = self.detect_file_type(input_path)
        base_name = os.path.splitext(os.path.basename(input_path))[0]
        markdown_file_path = os.path.join(md_folder, f"{base_name}.md")
        
        if file_type == 'application' and file_subtype == 'pdf':
            all_text = self.pdf_to_md(input_path, image_folder, finish_folder, dpi)
        elif file_type == 'image':
            all_text = self.process_image_file(input_path)
        else:
            raise ValueError(f"Unsupported file type: {file_type}/{file_subtype}")
        
        with open(markdown_file_path, 'w', encoding='utf-8') as f:
            f.write(all_text)
        self.remove_empty_lines(markdown_file_path)
        print(f"File {input_path} converted to markdown successfully.")
        return markdown_file_path

# ====================== Module 2: Markdown to JSON ======================
class MarkdownToJSON:
    DEFAULT_PROMPT = """
    请将以下 Markdown 格式的发票内容转换为 JSON 格式。要求如下：
    1. **提取以下字段**：
       - 发票代码（invoice_code）
       - 开票日期（invoice_date）
       - 购买方信息（buyer）：包含名称（name）和统一社会信用代码/纳税人识别号（tax_id）
       - 销售方信息（seller）：包含名称（name）和统一社会信用代码/纳税人识别号（tax_id）
       - 发票项目（items）：每个项目包含以下字段：
         - 项目名称（item_name）
         - 规格型号（specification）
         - 单位（unit）
         - 数量（quantity）
         - 单价（unit_price）
         - 金额（amount）
         - 税率/征收率（tax_rate）
         - 税额（tax_amount）
       - 合计金额（total_amount）
       - 合计税额（total_tax）
       - 价税合计（total_including_tax）：包含大写金额（capitalized）和小写金额（numeric）
       - 备注（remarks）：以列表形式存储
       - 开票人（issuer）
    2. **JSON 格式要求**：
       - 字段名称必须与上述要求一致。
       - 金额和税额字段的值应为字符串类型。
       - 发票项目（items）应为数组，每个项目为一个对象。
       - 备注（remarks）应为数组，每个备注为字符串。
    """

    def __init__(self, zhipu_client):
        self.client = zhipu_client

    def extract_mdinfo_from_zhipu(self, content_md, prompt_txt=None):
        """调用智谱AI API，将 Markdown 内容转换为 JSON 格式"""
        if prompt_txt is None:
            prompt_txt = self.DEFAULT_PROMPT

        messages_to_zhipu = [
            {"role": "system", "content": prompt_txt},
            {"role": "user", "content": content_md}
        ]

        try:
            response = self.client.chat.completions.create(
                model="glm-4.6",
                temperature=0,
                top_p=0.1,
                max_tokens=4095,
                messages=messages_to_zhipu,
            )

            content_from_glm = response.choices[0].message.content
            print(f"API返回原始内容:\n{content_from_glm}")  # 调试输出
            
            # 更健壮的JSON提取方式
            json_str = content_from_glm.strip()
            if '```json' in json_str:
                json_str = json_str.split('```json')[1].split('```')[0].strip()
            elif '```' in json_str:
                json_str = json_str.split('```')[1].strip()
            
            # 尝试解析JSON
            try:
                json_data = json.loads(json_str)
                return json_data
            except json.JSONDecodeError as e:
                print(f"JSON解析错误，尝试修复格式: {e}")
                # 尝试修复常见的JSON格式问题
                json_str = re.sub(r',\s*}', '}', json_str)  # 修复多余的逗号
                json_str = re.sub(r',\s*]', ']', json_str)
                json_str = re.sub(r'([{,]\s*)(\w+)(\s*:)', r'\1"\2"\3', json_str)  # 添加缺失的引号
                return json.loads(json_str)

        except Exception as e:
            print(f"Error during API call or JSON parsing: {e}")
            print(f"Problematic content:\n{content_from_glm}")
            raise ValueError(f"无法解析API返回的JSON数据: {str(e)}")

    def clean_numeric_string(self, value: str) -> str:
        """清理含货币符号、千位分隔符和空格的字符串"""
        return value.replace('¥', '').replace(',', '').strip()

    def remove_symbols(self, data):
        """将空字符串、空格、无效值强制转换为 0.0"""
        if 'items' in data:
            for item in data['items']:
                for key in ['unit_price', 'amount', 'tax_amount']:
                    if key in item and isinstance(item[key], str):
                        cleaned_value = self.clean_numeric_string(item[key])
                        if cleaned_value and cleaned_value.replace('.', '', 1).isdigit():
                            item[key] = float(cleaned_value)
                        else:
                            item[key] = 0.0
        
        for key in ['total_amount', 'total_tax']:
            if key in data and isinstance(data[key], str):
                cleaned_value = self.clean_numeric_string(data[key])
                if cleaned_value and cleaned_value.replace('.', '', 1).isdigit():
                    data[key] = float(cleaned_value)
                else:
                    data[key] = 0.0
        
        if 'total_including_tax' in data and isinstance(data['total_including_tax'], dict):
            numeric_value = data['total_including_tax'].get('numeric')
            if isinstance(numeric_value, str):
                cleaned_value = self.clean_numeric_string(numeric_value)
                if cleaned_value and cleaned_value.replace('.', '', 1).isdigit():
                    data['total_including_tax']['numeric'] = float(cleaned_value)
                else:
                    data['total_including_tax']['numeric'] = 0.0
        
        return data

    def md_to_json(self, md_dir, json_dir, finish_folder, prompt_text=None):
        """将指定目录下的 Markdown 文件转换为 JSON 文件"""
        for filename in os.listdir(md_dir):
            if filename.endswith(".md"):
                filepath = os.path.join(md_dir, filename)
                try:
                    with open(filepath, 'r', encoding='utf-8') as file:
                        markdown_content = file.read()

                    # 增加调试输出
                    print(f"处理文件: {filename}")
                    print(f"Markdown内容预览:\n{markdown_content[:500]}...")

                    extracted_info = self.extract_mdinfo_from_zhipu(markdown_content, prompt_text)
                    extracted_info = self.remove_symbols(extracted_info)

                    json_filename = os.path.splitext(filename)[0] + ".json"
                    json_filepath = os.path.join(json_dir, json_filename)

                    with open(json_filepath, 'w', encoding='utf-8') as json_file:
                        json.dump(extracted_info, json_file, indent=4, ensure_ascii=False)

                    print(f"JSON文件保存成功: {json_filepath}")
                    shutil.move(filepath, os.path.join(finish_folder, filename))

                except Exception as e:
                    print(f"处理文件 {filename} 失败: {str(e)}")
                    # 将失败文件移动到特定目录以便检查
                    error_folder = os.path.join(os.path.dirname(finish_folder), "errors")
                    os.makedirs(error_folder, exist_ok=True)
                    shutil.move(filepath, os.path.join(error_folder, filename))
                    continue

# ====================== Module 3: JSON to Excel ======================
class JSONToExcel:
    def json_to_exl(self, json_folder, output_folder):
        """从指定文件夹中读取 JSON 文件，并将所有相关信息填入新建的 Excel 表格中"""
        workbook = openpyxl.Workbook()
        sheet = workbook.active
        sheet.title = "发票信息"
        
        headers = [
            "发票代码", "开票日期", "项目名称", "规格型号", "单位", "数量", "单价", 
            "金额", "税率", "税额", "金额+税额", "卖方名称", "卖方税号", "买方名称", 
            "买方税号", "发票总额", "总税额", "价税合计(大写)", "价税合计(数值)", 
            "备注", "开票人", "来源文件"
        ]
        
        for col, header in enumerate(headers, start=1):
            sheet.cell(row=1, column=col, value=header)
        
        for filename in os.listdir(json_folder):
            if filename.endswith(".json"):
                filepath = os.path.join(json_folder, filename)
                with open(filepath, 'r', encoding='utf-8') as json_file:
                    data = json.load(json_file)
                    
                    invoice_code = data.get("invoice_code", "")
                    invoice_date = data.get("invoice_date", "")
                    buyer = data.get("buyer", {})
                    seller = data.get("seller", {})
                    remarks = "\n".join(data.get("remarks", [])) if data.get("remarks") else ""
                    issuer = data.get("issuer", "")
                    
                    items = data.get("items", [])
                    for item in items:
                        row = sheet.max_row + 1
                        
                        sheet.cell(row=row, column=1, value=invoice_code)
                        sheet.cell(row=row, column=2, value=invoice_date)
                        sheet.cell(row=row, column=3, value=item.get("item_name", ""))
                        sheet.cell(row=row, column=4, value=item.get("specification", ""))
                        sheet.cell(row=row, column=5, value=item.get("unit", ""))
                        sheet.cell(row=row, column=6, value=item.get("quantity", ""))
                        
                        sheet.cell(row=row, column=7, value=item.get("unit_price", 0))
                        sheet.cell(row=row, column=7).number_format = numbers.FORMAT_NUMBER_00
                        
                        sheet.cell(row=row, column=8, value=item.get("amount", 0))
                        sheet.cell(row=row, column=8).number_format = numbers.FORMAT_NUMBER_00
                        
                        sheet.cell(row=row, column=9, value=item.get("tax_rate", ""))
                        
                        sheet.cell(row=row, column=10, value=item.get("tax_amount", 0))
                        sheet.cell(row=row, column=10).number_format = numbers.FORMAT_NUMBER_00
                        
                        total_amount = item.get("amount", 0) + item.get("tax_amount", 0)
                        sheet.cell(row=row, column=11, value=total_amount)
                        sheet.cell(row=row, column=11).number_format = numbers.FORMAT_NUMBER_00
                        
                        sheet.cell(row=row, column=12, value=seller.get("name", ""))
                        sheet.cell(row=row, column=13, value=seller.get("tax_id", ""))
                        
                        sheet.cell(row=row, column=14, value=buyer.get("name", ""))
                        sheet.cell(row=row, column=15, value=buyer.get("tax_id", ""))
                        
                        sheet.cell(row=row, column=16, value=data.get("total_amount", 0))
                        sheet.cell(row=row, column=16).number_format = numbers.FORMAT_NUMBER_00
                        
                        sheet.cell(row=row, column=17, value=data.get("total_tax", 0))
                        sheet.cell(row=row, column=17).number_format = numbers.FORMAT_NUMBER_00
                        
                        total_including_tax = data.get("total_including_tax", {})
                        sheet.cell(row=row, column=18, value=total_including_tax.get("capitalized", ""))
                        sheet.cell(row=row, column=19, value=total_including_tax.get("numeric", 0))
                        sheet.cell(row=row, column=19).number_format = numbers.FORMAT_NUMBER_00
                        
                        sheet.cell(row=row, column=20, value=remarks)
                        sheet.cell(row=row, column=21, value=issuer)
                        sheet.cell(row=row, column=22, value=os.path.splitext(filename)[0])

        for column in sheet.columns:
            max_length = 0
            column_letter = column[0].column_letter
            for cell in column:
                try:
                    if len(str(cell.value)) > max_length:
                        max_length = len(str(cell.value))
                except:
                    pass
            adjusted_width = (max_length + 2) * 1.2
            sheet.column_dimensions[column_letter].width = adjusted_width

        os.makedirs(output_folder, exist_ok=True)
        # 修改时间戳格式为更易读的形式
        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
        excel_filename = f"发票信息汇总_{timestamp}.xlsx"
        excel_path = os.path.join(output_folder, excel_filename)
        
        workbook.save(excel_path)
        print(f"数据已成功写入新建的 Excel 文件: {excel_path}")
        return excel_path


# ====================== 多线程处理函数 ======================
async def process_single_file(file_path, client, image_folder, md_folder, finish_folder):
    """异步处理单个文件"""
    pdf_to_md = PDFImageToMarkdown(client)
    try:
        # 步骤1: 文件转为Markdown
        md_path = pdf_to_md.file_to_md(file_path, image_folder, md_folder, finish_folder)
        return md_path
    except Exception as e:
        print(f"Error processing file {file_path}: {e}")
        return None

async def process_md_to_json(md_path, client, json_folder, finish_folder):
    """异步处理Markdown转JSON"""
    md_to_json = MarkdownToJSON(client)
    try:
        md_to_json.md_to_json(os.path.dirname(md_path), json_folder, finish_folder)
        return True
    except Exception as e:
        print(f"Error converting {md_path} to JSON: {e}")
        return False

async def process_invoices(api_key, files, progress=gr.Progress()):
    """修改后的批量处理函数，支持并发"""
    # 进度条
    progress(0, desc="开始处理")

    if not api_key:
        raise gr.Error("请输入智谱 API Key")
    
    try:
        client = ZhipuAiClient(api_key=api_key)
        # 测试API连接
        test_response = client.chat.completions.create(
            model="glm-4.6",
            messages=[{"role": "user", "content": "test"}],
            temperature=0
        )
        print(f"API测试响应: {test_response}")
    except Exception as e:
        raise gr.Error(f"API Key 无效或验证失败: {str(e)}")

    # 创建临时目录并检查权限
    try:
        temp_dir = tempfile.mkdtemp(prefix="invoice_processing_")
        print(f"临时目录创建成功: {temp_dir}")
        
        image_folder = os.path.join(temp_dir, "images")
        md_folder = os.path.join(temp_dir, "markdown")
        json_folder = os.path.join(temp_dir, "json")
        finish_folder = os.path.join(temp_dir, "finished")
        
        for folder in [image_folder, md_folder, json_folder, finish_folder]:
            os.makedirs(folder, exist_ok=True)
            print(f"创建文件夹: {folder}")
            
            # 检查文件夹是否可写
            test_file = os.path.join(folder, "test.txt")
            with open(test_file, "w") as f:
                f.write("test")
            os.remove(test_file)
    except Exception as e:
        raise gr.Error(f"无法创建临时目录: {str(e)}")

    try:
        progress(0.05, desc="准备并发处理...")
        
        # 初始化处理器
        pdf_to_md = PDFImageToMarkdown(client)
        md_to_json = MarkdownToJSON(client)
        json_to_excel = JSONToExcel()
        
        # 先同步处理一个文件测试流程
        test_success = False
        if files:
            first_file = files[0].name
            print(f"测试处理第一个文件: {first_file}")
            
            try:
                # 测试PDF/Image转Markdown
                md_path = pdf_to_md.file_to_md(first_file, image_folder, md_folder, finish_folder)
                print(f"成功生成Markdown: {md_path}")
                
                # 测试Markdown转JSON
                md_to_json.md_to_json(md_folder, json_folder, finish_folder)
                print(f"成功生成JSON文件在: {json_folder}")
                
                test_success = True
            except Exception as e:
                print(f"测试文件处理失败: {str(e)}")
                traceback.print_exc()
        
        if not test_success:
            raise gr.Error("测试文件处理失败，请检查文件格式是否正确")

        # 如果有多个文件，继续处理剩余文件
        if len(files) > 1:
            with ThreadPoolExecutor(max_workers=4) as executor:
                loop = asyncio.get_event_loop()
                
                # 第一阶段: 并行处理文件到Markdown
                md_tasks = []
                for i, file in enumerate(files[1:]):  # 跳过已测试的第一个文件
                    file_path = file.name
                    task = loop.run_in_executor(
                        executor,
                        functools.partial(
                            pdf_to_md.file_to_md,
                            input_path=file_path,
                            image_folder=image_folder,
                            md_folder=md_folder,
                            finish_folder=finish_folder
                        )
                    )
                    md_tasks.append(task)
                    progress(0.1 + 0.4*(i/len(files)), desc=f"处理文件中 ({i+2}/{len(files)})")
                
                await asyncio.gather(*md_tasks)

            # 处理Markdown到JSON
            progress(0.6, desc="转换JSON中...")
            md_to_json.md_to_json(md_folder, json_folder, finish_folder)

        # 生成Excel
        progress(0.9, desc="生成Excel...")
        if len(os.listdir(json_folder)) > 0:
            excel_path = json_to_excel.json_to_exl(json_folder, temp_dir)
            print(f"处理完成，结果保存在: {excel_path}")
            progress(1.0, desc="处理完成")
            return excel_path, 100
        else:
            raise gr.Error("没有生成有效的JSON文件")

    except Exception as e:
        print(f"处理过程中发生错误: {str(e)}")
        traceback.print_exc()
        raise gr.Error(f"处理失败: {str(e)}")
    finally:
        # 调试时可以先不删除临时文件
        pass

# ====================== 界面配置 ======================
with gr.Blocks(title="发票识别系统") as demo:
    gr.Markdown("# 发票识别系统")
    gr.Markdown("上传PDF或图片格式的发票，系统将自动提取信息并生成Excel文件")
    
    with gr.Row():
        with gr.Column():
            api_key = gr.Textbox(
                label="输入智谱 API Key",
                placeholder="请输入您的智谱 API Key", 
                type="password"
            )
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(
                label="上传发票文件（可多选）",
                file_types=[".pdf", ".jpg", ".jpeg", ".png"],
                file_count="multiple"
            )
        with gr.Column():
            # 将进度条和下载结果组合在一起
            with gr.Group():
                file_output = gr.File(label="下载Excel结果")
                progress_bar = gr.Slider(
                    minimum=0, maximum=100, value=0,
                    label="处理进度",
                    interactive=False,
                    visible=False,
                    elem_classes=["compact-progress"]  # 添加自定义样式类
                )
    
    submit_btn = gr.Button("开始处理", size="lg")
    
    # CSS样式调整
    demo.css = """
    .compact-progress {
        margin-top: 8px;
        margin-bottom: 8px;
    }
    .compact-progress .wrap {
        padding: 0 !important;
    }
    """
    
    submit_btn.click(
        fn=lambda api_key, files: asyncio.run(process_invoices(api_key, files, progress=gr.Progress())),
        inputs=[api_key, file_input],
        outputs=[file_output, progress_bar]
    )

if __name__ == "__main__":
    demo.launch()