Update app.py

app.py

@@ -1,11 +1,602 @@
+# 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 zhipuai import ZhipuAI
+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
+
+# ====================== 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-4v-plus-0111",
+            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-plus",
+                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 = ZhipuAI(api_key=api_key)
+        # 测试API连接
+        test_response = client.chat.completions.create(
+            model="glm-4",
+            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)
 
-def greet(name):
-    return "Hello " + name + "!"
+        # 生成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
 
-demo = gr.Interface(fn=greet, inputs="textbox", outputs="textbox")
+# ====================== 界面配置 ======================
+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()
+    demo.launch()