Update app.py

app.py

@@ -4,7 +4,7 @@ import base64
 import json
 from io import BytesIO
 from PIL import Image, ImageDraw
-from mistralai.client import Mistral
+from mistralai import Mistral
 
 def process_document(file, api_key, progress=gr.Progress()):
     if not api_key:
@@ -24,7 +24,6 @@ def process_document(file, api_key, progress=gr.Progress()):
             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)
@@ -37,56 +36,93 @@ def process_document(file, api_key, progress=gr.Progress()):
     drawn_images = []
     all_json_results = []
 
-    # 2. 升級版：極度嚴格的 Prompt (防幻覺、防抓印刷體)
-    prompt = """
-    你是一個專精於「中文手寫字跡辨識」的 AI 專家。
-    這張圖片是一份學生的考卷，上面同時包含「電腦打字的印刷體題目」以及「學生手寫的作答字跡」。
-    學生的字跡混合了繁體中文與簡體中文，且字體較為潦草，通常寫在橫線上或空白處。
+    # ==========================================
+    # Prompt 1: 給 Pixtral (專注看圖、免 JSON 壓力)
+    # ==========================================
+    prompt_step1_pixtral = """
+    你是一個專精於「中文手寫字跡辨識」的 AI 視覺專家。
+    這張圖片是一份學生的考卷。你的唯一任務是找出「所有學生手寫的文字」，並給出位置。
 
-    你的任務是：
-    1. **絕對忽略印刷體**：只尋找並提取「手寫字」。絕對不要轉錄任何電腦印刷體的題目內容！
-    2. **精準辨識（拒絕幻覺）**：請盡最大努力辨識學生的手寫字（包含簡體與繁體）。如果你遇到真的看不懂的草字，請嚴格使用「[無法辨識]」四個字代替。絕對不要自己發明、猜測或聯想詞彙。
-    3. **輸出邊界框**：給出該段手寫文字在圖片上的相對邊界框 [ymin, xmin, ymax, xmax]（範圍 0.0 到 1.0）。
+    請嚴格遵守以下規則：
+    1. 絕對忽略所有電腦打字的印刷體！只提取手寫字。
+    2. 盡力辨識潦草字跡（包含簡繁體），若完全無法辨識請寫 [無法辨識]。
+    3. 給出每段手寫字在圖片上的相對邊界框 [ymin, xmin, ymax, xmax]（範圍 0.0 到 1.0）。
+    
+    【重要】請使用「純文字」輸出，絕對不要輸出 JSON 格式！請用以下格式條列：
+    內容：[手寫字]
+    座標：[ymin, xmin, ymax, xmax]
+    ---
+    內容：[手寫字]
+    座標：[ymin, xmin, ymax, xmax]
+    """
+
+    # ==========================================
+    # Prompt 2: 給 Mistral Large (專注格式轉換)
+    # ==========================================
+    prompt_step2_large = """
+    你是一個資料格式化專家。請將以下由 OCR 模型提取出的純文字紀錄，轉換為嚴格的 JSON 格式。
+    請確保座標被轉換為包含 4 個浮點數的陣列。
 
-    請嚴格以 JSON 格式輸出：
+    必須輸出的 JSON 格式如下：
     {
       "handwriting": [
-        {"text": "第一段手寫內容", "box": [0.12, 0.34, 0.15, 0.55]},
-        {"text": "第二段手寫內容", "box": [0.60, 0.10, 0.65, 0.80]}
+        {"text": "提取的內容", "box": [0.12, 0.34, 0.15, 0.55]}
       ]
     }
+
+    這是需要轉換的原始資料：
+    \n\n
     """
 
     for page_num, img in enumerate(images):
-        progress(0.3, desc=f"處理第 {page_num + 1} 頁：準備圖片資料...")
+        progress(0.2, 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(
+            # ---------------------------------------------------------
+            # 步驟一：呼叫 Pixtral 進行視覺提取 (純文字輸出)
+            # ---------------------------------------------------------
+            progress(0.3, desc=f"第 {page_num + 1} 頁 [Step 1/2]：Pixtral 視覺提取中 (約 30 秒)...")
+            response_step1 = client.chat.complete(
                 model="pixtral-large-latest",
                 messages=[
                     {
                         "role": "user",
                         "content": [
-                            {"type": "text", "text": prompt},
+                            {"type": "text", "text": prompt_step1_pixtral},
                             {"type": "image_url", "image_url": base64_image}
                         ]
                     }
+                ]
+            )
+            pixtral_raw_text = response_step1.choices[0].message.content
+            print(f"--- Pixtral 原始輸出 (Page {page_num+1}) ---\n{pixtral_raw_text}\n---------------------------")
+            
+            # ---------------------------------------------------------
+            # 步驟二：呼叫 Mistral Large 進行 JSON 結構化
+            # ---------------------------------------------------------
+            progress(0.7, desc=f"第 {page_num + 1} 頁 [Step 2/2]：Mistral Large 格式轉換中 (約 5 秒)...")
+            response_step2 = client.chat.complete(
+                model="mistral-large-latest",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": prompt_step2_large + pixtral_raw_text
+                    }
                 ],
                 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)
+            final_json_content = response_step2.choices[0].message.content
             
+            # ---------------------------------------------------------
+            # 解析並畫框
+            # ---------------------------------------------------------
+            progress(0.9, desc=f"第 {page_num + 1} 頁：正在圖片上繪製邊界框...")
+            data = json.loads(final_json_content)
             all_json_results.append({f"Page {page_num + 1}": data})
             
             draw = ImageDraw.Draw(img)
@@ -106,7 +142,7 @@ def process_document(file, api_key, progress=gr.Progress()):
             drawn_images.append(img)
             
         except Exception as e:
-            error_msg = f"API 或解析發生錯誤: {str(e)}"
+            error_msg = f"發生錯誤: {str(e)}"
             all_json_results.append({f"Page {page_num + 1} Error": error_msg})
             drawn_images.append(img)
 
@@ -114,19 +150,19 @@ def process_document(file, api_key, progress=gr.Progress()):
     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.Blocks(title="雙模型協作：手寫字偵測 Demo") as demo:
+    gr.Markdown("## 📝 雙模型協作：手寫字偵測與畫框 Demo")
+    gr.Markdown("**架構：** `Pixtral-large` (視覺提取純文字) $\\rightarrow$ `Mistral-large` (轉換標準 JSON) $\\rightarrow$ 畫框")
     
     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")
+            submit_btn = gr.Button("開始雙階段偵測", variant="primary")
             
         with gr.Column():
             output_gallery = gr.Gallery(label="畫框結果預覽", columns=1, height="auto")
-            output_json = gr.JSON(label="模型輸出的 JSON (文字與座標)")
+            output_json = gr.JSON(label="Step 2 轉換後的標準 JSON")
 
     submit_btn.click(
         fn=process_document,