app.py

import gradio as gr
import fitz  # PyMuPDF
import base64
import json
from io import BytesIO
from PIL import Image, ImageDraw
from mistralai import Mistral

def process_document(file, api_key, progress=gr.Progress()):
    if not api_key:
        return None, "錯誤：請提供 Mistral API 金鑰。"
    if not file:
        return None, "錯誤：請上傳檔案。"

    progress(0.05, desc="正在初始化 Mistral 客戶端...")
    client = Mistral(api_key=api_key)
    images = []
    
    # 1. 處理檔案上傳
    file_name = file.name.lower()
    try:
        progress(0.1, desc="正在讀取與轉換圖片...")
        if file_name.endswith('.pdf'):
            doc = fitz.open(file.name)
            for i in range(len(doc)):
                page = doc.load_page(i)
                # 設定 DPI，保持清晰度
                pix = page.get_pixmap(dpi=150)
                img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
                images.append(img)
        else:
            img = Image.open(file.name).convert("RGB")
            images.append(img)
    except Exception as e:
        return None, f"檔案讀取錯誤：{str(e)}"

    drawn_images = []
    all_json_results = []

    # 2. 升級版：極度嚴格的 Prompt (防幻覺、防抓印刷體)
    prompt = """
    你是一個專精於「中文手寫字跡辨識」的 AI 專家。
    這張圖片是一份學生的考卷，上面同時包含「電腦打字的印刷體題目」以及「學生手寫的作答字跡」。
    學生的字跡混合了繁體中文與簡體中文，且字體較為潦草，通常寫在橫線上或空白處。

    你的任務是：
    1. **絕對忽略印刷體**：只尋找並提取「手寫字」。絕對不要轉錄任何電腦印刷體的題目內容！
    2. **精準辨識（拒絕幻覺）**：請盡最大努力辨識學生的手寫字（包含簡體與繁體）。如果你遇到真的看不懂的草字，請嚴格使用「[無法辨識]」四個字代替。絕對不要自己發明、猜測或聯想詞彙。
    3. **輸出邊界框**：給出該段手寫文字在圖片上的相對邊界框 [ymin, xmin, ymax, xmax]（範圍 0.0 到 1.0）。

    請嚴格以 JSON 格式輸出：
    {
      "handwriting": [
        {"text": "第一段手寫內容", "box": [0.12, 0.34, 0.15, 0.55]},
        {"text": "第二段手寫內容", "box": [0.60, 0.10, 0.65, 0.80]}
      ]
    }
    """

    for page_num, img in enumerate(images):
        progress(0.3, desc=f"處理第 {page_num + 1} 頁：準備圖片資料...")
        buffered = BytesIO()
        img.save(buffered, format="JPEG")
        img_b64 = base64.b64encode(buffered.getvalue()).decode("utf-8")
        base64_image = f"data:image/jpeg;base64,{img_b64}"
        
        try:
            progress(0.4, desc=f"第 {page_num + 1} 頁：Pixtral-large-latest 運算中 (約 30~60 秒)...")
            response = client.chat.complete(
                model="pixtral-large-latest",
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                            {"type": "image_url", "image_url": base64_image}
                        ]
                    }
                ],
                response_format={"type": "json_object"}
            )
            
            progress(0.8, desc=f"第 {page_num + 1} 頁：正在解析 JSON 並畫框...")
            content = response.choices[0].message.content
            
            # 清理模型可能多加的 Markdown JSON 標記 (防錯機制)
            clean_content = content.replace("```json", "").replace("```", "").strip()
            data = json.loads(clean_content)
            
            all_json_results.append({f"Page {page_num + 1}": data})
            
            draw = ImageDraw.Draw(img)
            width, height = img.size
            
            for item in data.get("handwriting", []):
                box = item.get("box")
                if box and len(box) == 4:
                    ymin, xmin, ymax, xmax = box
                    abs_xmin = xmin * width
                    abs_ymin = ymin * height
                    abs_xmax = xmax * width
                    abs_ymax = ymax * height
                    # 畫紅框
                    draw.rectangle([abs_xmin, abs_ymin, abs_xmax, abs_ymax], outline="red", width=3)
                    
            drawn_images.append(img)
            
        except Exception as e:
            error_msg = f"API 或解析發生錯誤: {str(e)}"
            all_json_results.append({f"Page {page_num + 1} Error": error_msg})
            drawn_images.append(img)

    progress(1.0, desc="處理完成！")
    return drawn_images, json.dumps(all_json_results, ensure_ascii=False, indent=2)

# --- Gradio 介面設計 ---
with gr.Blocks(title="手寫字偵測與畫框 Demo") as demo:
    gr.Markdown("## 📝 手寫字偵測與畫框 Demo (Pixtral-large-latest)")
    gr.Markdown("上傳考卷 PDF 或圖片，系統會自動排除印刷體，找出所有的「手寫文字」並用紅框標示出來。")
    
    with gr.Row():
        with gr.Column():
            api_key_input = gr.Textbox(label="Mistral API Key", type="password")
            file_input = gr.File(label="上傳 PDF 或圖片", file_types=[".pdf", ".jpg", ".png", ".jpeg"])
            submit_btn = gr.Button("開始偵測畫框", variant="primary")
            
        with gr.Column():
            output_gallery = gr.Gallery(label="畫框結果預覽", columns=1, height="auto")
            output_json = gr.JSON(label="模型輸出的 JSON (文字與座標)")

    submit_btn.click(
        fn=process_document,
        inputs=[file_input, api_key_input],
        outputs=[output_gallery, output_json]
    )

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