Update app.py

app.py

@@ -17,8 +17,6 @@ import numpy as np
 ###############################
 os.system('pip uninstall -y magic-pdf')
 os.system('pip install git+https://github.com/opendatalab/MinerU.git@dev')
-
-# ★ OpenCV 설치 (headless 버전)
 os.system('pip install opencv-python-headless')
 
 os.system('wget https://github.com/opendatalab/MinerU/raw/dev/scripts/download_models_hf.py -O download_models_hf.py')
@@ -102,6 +100,74 @@ def read_fn(path):
     disk_rw = FileBasedDataReader(os.path.dirname(path))
     return disk_rw.read(os.path.basename(path))
 
+###############################
+# 이미지 전처리 함수 (이진화 + Deskew)
+###############################
+def preprocess_image(image_path):
+    """
+    1) Grayscale + Otsu Binarization
+    2) Deskew(기울임 보정)
+    """
+    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
+    if img is None:
+        # 이미지가 아니거나 로딩 실패 시 그대로 반환
+        return image_path
+
+    # (a) 이진화(Otsu)
+    _, img_bin = cv2.threshold(img, 0, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY)
+
+    # (b) 기울임 보정(deskew)
+    coords = np.column_stack(np.where(img_bin > 0))
+    angle = cv2.minAreaRect(coords)[-1]
+    if angle < -45:
+        angle = -(90 + angle)
+    else:
+        angle = -angle
+
+    (h, w) = img_bin.shape[:2]
+    center = (w // 2, h // 2)
+    M = cv2.getRotationMatrix2D(center, angle, 1.0)
+    img_rotated = cv2.warpAffine(
+        img_bin, M, (w, h),
+        flags=cv2.INTER_CUBIC,
+        borderMode=cv2.BORDER_CONSTANT,
+        borderValue=255
+    )
+
+    # 임시 파일로 저장
+    preprocessed_path = image_path + "_preprocessed.png"
+    cv2.imwrite(preprocessed_path, img_rotated)
+    return preprocessed_path
+
+###############################
+# PDF or Image -> PDF 변환
+# (불필요한 f.close() 제거)
+###############################
+def to_pdf(file_path):
+    """
+    이미지(JPG/PNG 등)라면 전처리 후 PDF로 변환.
+    이미 PDF라면 그대로 반환.
+    """
+    with pymupdf.open(file_path) as f:
+        # PDF인 경우
+        if f.is_pdf:
+            return file_path
+        # 이미지 파일인 경우
+        # (중첩된 with문에 들어가기 전, 그냥 블록이 끝나면 자동 close)
+    # 블록 종료 시 f는 이미 close됨.
+
+    # 이미지 전처리 -> 새 이미지 -> PDF 변환
+    preprocessed_path = preprocess_image(file_path)
+    with pymupdf.open(preprocessed_path) as img_doc:
+        pdf_bytes = img_doc.convert_to_pdf()
+
+    unique_filename = f"{uuid.uuid4()}.pdf"
+    tmp_file_path = os.path.join(os.path.dirname(file_path), unique_filename)
+    with open(tmp_file_path, 'wb') as tmp_pdf_file:
+        tmp_pdf_file.write(pdf_bytes)
+
+    return tmp_file_path
+
 def parse_pdf(doc_path, output_dir, end_page_id, is_ocr, layout_mode, formula_enable, table_enable, language):
     os.makedirs(output_dir, exist_ok=True)
     try:
@@ -146,84 +212,19 @@ def image_to_base64(image_path):
 
 def replace_image_with_base64(markdown_text, image_dir_path):
     pattern = r'\!\[(?:[^\]]*)\]\(([^)]+)\)'
+
     def replace(match):
         relative_path = match.group(1)
         full_path = os.path.join(image_dir_path, relative_path)
         base64_image = image_to_base64(full_path)
         return f"![{relative_path}](data:image/jpeg;base64,{base64_image})"
-    return re.sub(pattern, replace, markdown_text)
 
-###############################
-# 이미지 전처리 함수 (Grayscale/Binarization + Deskew)
-###############################
-def preprocess_image(image_path):
-    """
-    1) Grayscale + Binarization(OTSU)
-    2) Deskew(기울임 보정)
-    전처리된 이미지를 임시 경로에 저장 후 해당 경로를 반환
-    """
-    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
-    if img is None:
-        # 이미지 파일이 아닌 경우 혹은 로딩 실패 시 원본 경로 그대로 반환
-        return image_path
-
-    # (a) 이진화(binarization)
-    _, img_bin = cv2.threshold(img, 0, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY)
-
-    # (b) 기울임 보정(deskew)
-    coords = np.column_stack(np.where(img_bin > 0))
-    angle = cv2.minAreaRect(coords)[-1]
-    # OpenCV는 회전 각도를 [-90, 0)로 반환할 때가 많으므로 보정
-    if angle < -45:
-        angle = -(90 + angle)
-    else:
-        angle = -angle
-
-    (h, w) = img_bin.shape[:2]
-    center = (w // 2, h // 2)
-    M = cv2.getRotationMatrix2D(center, angle, 1.0)
-    img_rotated = cv2.warpAffine(
-        img_bin,
-        M,
-        (w, h),
-        flags=cv2.INTER_CUBIC,
-        borderMode=cv2.BORDER_CONSTANT,
-        borderValue=255
-    )
-
-    # 임시 파일로 저장
-    preprocessed_path = image_path + "_preprocessed.png"
-    cv2.imwrite(preprocessed_path, img_rotated)
-
-    return preprocessed_path
-
-def to_pdf(file_path):
-    """
-    이미지(JPG/PNG 등)를 PDF로 컨버팅하되,
-    이미지일 경우 전처리(Grayscale/Binarization + Deskew)를 먼저 적용
-    """
-    with pymupdf.open(file_path) as f:
-        if f.is_pdf:
-            return file_path
-        else:
-            # 이미지 파일인 경우, 전처리 수행 후 PDF 생성
-            f.close()
-            preprocessed_path = preprocess_image(file_path)
-
-            # 전처리된 이미지를 다시 PyMuPDF로 열어서 PDF 변환
-            with pymupdf.open(preprocessed_path) as img_doc:
-                pdf_bytes = img_doc.convert_to_pdf()
-
-            unique_filename = f"{uuid.uuid4()}.pdf"
-            tmp_file_path = os.path.join(os.path.dirname(file_path), unique_filename)
-            with open(tmp_file_path, 'wb') as tmp_pdf_file:
-                tmp_pdf_file.write(pdf_bytes)
-            return tmp_file_path
+    return re.sub(pattern, replace, markdown_text)
 
 def to_markdown(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language, progress=gr.Progress(track_tqdm=False)):
     """
     업로드된 PDF/이미지 -> PDF 변환 -> 마크다운 변환
-    (프로그레스 바 표시용)
+    (프로그레스 표시)
     """
     progress(0, "PDF로 변환 중...")
     file_path = to_pdf(file_path)
@@ -347,7 +348,6 @@ def convert_chat_messages_to_gradio_format(messages):
 def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
     """
     Gemini 응답 스트리밍
-    (user_message가 공백이면 기본 문구로 대체)
     """
     if not user_message.strip():
         user_message = "...(No content from user)..."
@@ -376,7 +376,7 @@ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
             parts = chunk.candidates[0].content.parts
             current_chunk = parts[0].text
 
-            # 만약 parts 가 2개라면, parts[0]는 thinking, parts[1]은 최종답변
+            # 만약 parts가 2개이면 [thinking, 최종답변]
             if len(parts) == 2 and not thinking_complete:
                 thought_buffer += current_chunk
                 messages[-1] = ChatMessage(
@@ -413,7 +413,7 @@ def stream_gemini_response(user_message: str, messages: list) -> Iterator[list]:
 
 def user_message(msg: str, history: list, doc_text: str) -> tuple[str, list]:
     """
-    doc_text(마크다운) 사용해 질문 자동 변형
+    doc_text(마크다운) 사용해 질문을 살짝 변형
     """
     if doc_text.strip():
         user_query = f"다음 문서를 참고하여 답변:\n\n{doc_text}\n\n질문: {msg}"
@@ -435,7 +435,6 @@ def reset_states(_):
 ###############################
 # (2) OCR FLEX 전용 (스니펫)
 ###############################
-# 별도의 LaTeX 설정
 latex_delimiters = [
     {"left": "$$", "right": "$$", "display": True},
     {"left": '$', "right": '$', "display": False}
@@ -444,8 +443,7 @@ latex_delimiters = [
 def to_markdown_ocr_flex(file_path, end_pages, is_ocr, layout_mode, formula_enable, table_enable, language):
     """
     스니펫에서 사용: 
-    업로드된 PDF/이미지 -> PDF 변환 -> 마크다운 변환
-    (마크다운 렌더링 / 마크다운 텍스트 / 압축파일 / PDF미리보기) 반환
+    업로드된 PDF/이미지 -> 전처리 -> PDF 변환 -> 마크다운 변환
     """
     file_path = to_pdf(file_path)
     if end_pages > 20:
@@ -475,11 +473,10 @@ def to_markdown_ocr_flex(file_path, end_pages, is_ocr, layout_mode, formula_enab
 ###############################
 if __name__ == "__main__":
     with gr.Blocks(title="VisionOCR", css=create_css()) as demo:
-        # 탭 영역
         with gr.Tabs():
-            #########################################################
+            ###############################
             # Tab (1) : PDF -> Markdown 변환 + Chat
-            #########################################################
+            ###############################
             with gr.Tab("PDF Chat with LLM"):
                 gr.HTML("""
                 <div class="title-area">
@@ -491,21 +488,19 @@ if __name__ == "__main__":
                 md_state = gr.State("")      # 변환된 마크다운 텍스트
                 chat_history = gr.State([])  # ChatMessage 리스트
 
-                # 업로드 & 변환
                 with gr.Row():
                     file = gr.File(label="PDF/이미지 업로드", file_types=[".pdf", ".png", ".jpeg", ".jpg"], interactive=True)
                     convert_btn = gr.Button("변환하기")
 
                 chatbot = gr.Chatbot(height=600)
 
-                # 새 파일 업로드 시: 이전 대화/마크다운/챗봇 초기화
                 file.change(
                     fn=reset_states,
                     inputs=file,
                     outputs=[chat_history, md_state, chatbot]
                 )
 
-                # 숨김 요소들
+                # 숨김 옵션들
                 max_pages = gr.Slider(1, 20, 10, visible=False, elem_classes="invisible")
                 layout_mode = gr.Dropdown(["layoutlmv3","doclayout_yolo"], value="doclayout_yolo", visible=False, elem_classes="invisible")
                 language = gr.Dropdown(all_lang, value='auto', visible=False, elem_classes="invisible")
@@ -520,7 +515,6 @@ if __name__ == "__main__":
                     show_progress=True
                 )
 
-                # Gemini Chat
                 gr.Markdown("## 추론 LLM과 대화")
 
                 with gr.Row():
@@ -546,9 +540,9 @@ if __name__ == "__main__":
                     outputs=[chat_history, md_state, chatbot]
                 )
 
-            #########################################################
-            # Tab (2) : OCR FLEX (스니펫 코드)
-            #########################################################
+            ###############################
+            # Tab (2) : OCR FLEX
+            ###############################
             with gr.Tab("OCR FLEX"):
                 gr.HTML("""
                 <div class="title-area">
@@ -558,87 +552,38 @@ if __name__ == "__main__":
                 """)
 
                 with gr.Row():
-                    # 왼쪽 패널
                     with gr.Column(variant='panel', scale=5):
-                        file_ocr = gr.File(
-                            label="PDF 또는 이미지 파일을 업로드하세요",
-                            file_types=[".pdf", ".png", ".jpeg", ".jpg"]
-                        )
-
-                        max_pages_ocr = gr.Slider(
-                            1, 20, 10,
-                            step=1,
-                            label='최대 변환 페이지 수'
-                        )
+                        file_ocr = gr.File(label="PDF 또는 이미지 파일", file_types=[".pdf", ".png", ".jpeg", ".jpg"])
+                        max_pages_ocr = gr.Slider(1, 20, 10, step=1, label='최대 변환 페이지 수')
 
                         with gr.Row():
-                            layout_mode_ocr = gr.Dropdown(
-                                ["layoutlmv3", "doclayout_yolo"],
-                                label="레이아웃 모델",
-                                value="doclayout_yolo"
-                            )
-                            language_ocr = gr.Dropdown(
-                                all_lang,
-                                label="언어",
-                                value='auto'
-                            )
+                            layout_mode_ocr = gr.Dropdown(["layoutlmv3", "doclayout_yolo"], label="레이아웃 모델", value="doclayout_yolo")
+                            language_ocr = gr.Dropdown(all_lang, label="언어", value='auto')
 
                         with gr.Row():
-                            formula_enable_ocr = gr.Checkbox(
-                                label="수식 인식 활성화",
-                                value=True
-                            )
-                            is_ocr_ocr = gr.Checkbox(
-                                label="OCR 강제 활성화",
-                                value=False
-                            )
-                            table_enable_ocr = gr.Checkbox(
-                                label="표 인식 활성화(테스트)",
-                                value=True
-                            )
+                            formula_enable_ocr = gr.Checkbox(label="수식 인식 활성화", value=True)
+                            is_ocr_ocr = gr.Checkbox(label="OCR 강제 활성화", value=False)
+                            table_enable_ocr = gr.Checkbox(label="표 인식 활성화(테스트)", value=True)
 
                         with gr.Row():
                             change_bu_ocr = gr.Button("변환")
-
-                            # ★ ClearButton 수정 ★
-                            # 첫 번째 인자 -> clear할 대상(컴포넌트),
-                            # 버튼에 표시될 텍스트는 value="초기화"
                             clear_bu_ocr = gr.ClearButton(
-                                components=[file_ocr, max_pages_ocr, layout_mode_ocr, language_ocr,
-                                            formula_enable_ocr, is_ocr_ocr, table_enable_ocr],
+                                components=[
+                                    file_ocr,
+                                    max_pages_ocr,
+                                    layout_mode_ocr,
+                                    language_ocr,
+                                    formula_enable_ocr,
+                                    is_ocr_ocr,
+                                    table_enable_ocr
+                                ],
                                 value="초기화"
                             )
 
-                        pdf_show_ocr = PDF(
-                            label='PDF 미리보기',
-                            interactive=False,
-                            visible=True,
-                            height=800
-                        )
-
-                        # 예제 폴더가 있다면 사용 (실제 실행환경에 따라 주의)
-                        with gr.Accordion("예제:", open=False):
-                            example_root = (
-                                os.path.join(os.path.dirname(__file__), "examples")
-                                if "__file__" in globals() else "./examples"
-                            )
-                            if os.path.exists(example_root):
-                                gr.Examples(
-                                    examples=[
-                                        os.path.join(example_root, _) for _ in os.listdir(example_root)
-                                        if _.endswith("pdf")
-                                    ],
-                                    inputs=file_ocr
-                                )
-                            else:
-                                gr.Markdown("예제 폴더가 존재하지 않습니다.")
+                        pdf_show_ocr = PDF(label='PDF 미리보기', interactive=False, visible=True, height=800)
 
-                    # 오른쪽 패널
                     with gr.Column(variant='panel', scale=5):
-                        output_file_ocr = gr.File(
-                            label="변환 결과",
-                            interactive=False
-                        )
+                        output_file_ocr = gr.File(label="변환 결과", interactive=False)
 
                         with gr.Tabs():
                             with gr.Tab("마크다운 렌더링"):
@@ -646,17 +591,13 @@ if __name__ == "__main__":
                                     label="마크다운 렌더링",
                                     height=1100,
                                     show_copy_button=True,
-                                    latex_delimiters=latex_delimiters,
-                                    line_breaks=True
+                                    line_breaks=True,
+                                    latex_delimiters=latex_delimiters
                                 )
 
                             with gr.Tab("마크다운 텍스트"):
-                                md_text_ocr = gr.TextArea(
-                                    lines=45,
-                                    show_copy_button=True
-                                )
+                                md_text_ocr = gr.TextArea(lines=45, show_copy_button=True)
 
-                # 이벤트 핸들러 (OCR FLEX)
                 file_ocr.change(
                     fn=to_pdf,
                     inputs=file_ocr,
@@ -685,5 +626,4 @@ if __name__ == "__main__":
                     ]
                 )
 
-    # 전체 앱 실행
     demo.launch(server_name="0.0.0.0", server_port=7860, debug=True, ssr_mode=True)