pdf_translator_web.py

import os
import json
import hashlib
from pathlib import Path
import streamlit as st
import pymupdf
from deep_translator import (
    GoogleTranslator,
    ChatGptTranslator,
)

# Constants
DEFAULT_PAGES_PER_LOAD = 2
DEFAULT_MODEL = "default_model"
DEFAULT_API_BASE = "http://localhost:8080/v1"

# Supported translators
TRANSLATORS = {
    'google': GoogleTranslator,
    'chatgpt': ChatGptTranslator,
}

# Color options
COLOR_MAP = {
    "darkred": (0.8, 0, 0),
    "black": (0, 0, 0),
    "blue": (0, 0, 0.8),
    "darkgreen": (0, 0.5, 0),
    "purple": (0.5, 0, 0.5),
}

# Target language options for ChatGPT
LANGUAGE_OPTIONS = {
    "简体中文": "zh-CN",
    "繁體中文": "zh-TW",
    "English": "en",
    "日本語": "ja",
    "한국어": "ko",
    "Español": "es",
    "Français": "fr",
    "Deutsch": "de",
}

# Add source language options
SOURCE_LANGUAGE_OPTIONS = {
    "English": "en",
    "简体中文": "zh-CN",
    "繁體中文": "zh-TW",
    "日本語": "ja",
    "한국어": "ko",
    "Español": "es",
    "Français": "fr",
    "Deutsch": "de",
    "Auto": "auto",
}

def get_cache_dir():
    """Get or create cache directory"""
    cache_dir = Path('.cached')
    cache_dir.mkdir(exist_ok=True)
    return cache_dir

def get_cache_key(file_content: bytes, page_num: int, translator_name: str, target_lang: str):
    """Generate cache key for a specific page translation"""
    # 使用文件内容的hash作为缓存key的一部分
    file_hash = hashlib.md5(file_content).hexdigest()
    return f"{file_hash}_page{page_num}_{translator_name}_{target_lang}.pdf"

def get_cached_translation(cache_key: str) -> pymupdf.Document:
    """Get cached translation if exists"""
    cache_path = get_cache_dir() / cache_key
    if cache_path.exists():
        return pymupdf.open(str(cache_path))
    return None

def save_translation_cache(doc: pymupdf.Document, cache_key: str):
    """Save translation to cache"""
    cache_path = get_cache_dir() / cache_key
    doc.save(str(cache_path))

def translate_pdf_pages(doc, doc_bytes, start_page, num_pages, translator, text_color, translator_name, target_lang):
    """Translate specific pages of a PDF document with progress and caching"""
    WHITE = pymupdf.pdfcolor["white"]
    rgb_color = COLOR_MAP.get(text_color.lower(), COLOR_MAP["darkred"])
    
    translated_pages = []
    
    # Create a progress bar
    progress_bar = st.progress(0)
    status_text = st.empty()
    
    for i, page_num in enumerate(range(start_page, min(start_page + num_pages, doc.page_count))):
        status_text.text(f"Translating page {page_num + 1}...")
        
        # Check cache first
        cache_key = get_cache_key(doc_bytes, page_num, translator_name, target_lang)
        cached_doc = get_cached_translation(cache_key)
        
        if cached_doc is not None:
            translated_pages.append(cached_doc)
        else:
            # Create a new PDF document for this page
            new_doc = pymupdf.open()
            new_doc.insert_pdf(doc, from_page=page_num, to_page=page_num)
            page = new_doc[0]
            
            # Extract and translate text blocks
            blocks = page.get_text("blocks", flags=pymupdf.TEXT_DEHYPHENATE)
            
            for block in blocks:
                bbox = block[:4]
                text = block[4]
                
                # Translate the text
                translated = translator.translate(text)
                
                # Cover original text with white and add translation in color
                page.draw_rect(bbox, color=None, fill=WHITE)
                page.insert_htmlbox(
                    bbox,
                    translated,
                    css=f"* {{font-family: sans-serif; color: rgb({int(rgb_color[0]*255)}, {int(rgb_color[1]*255)}, {int(rgb_color[2]*255)});}}"
                )
            
            # Save to cache
            save_translation_cache(new_doc, cache_key)
            translated_pages.append(new_doc)
        
        # Update progress
        progress = (i + 1) / min(num_pages, doc.page_count - start_page)
        progress_bar.progress(progress)
    
    # Clear progress indicators
    progress_bar.empty()
    status_text.empty()
    
    return translated_pages

def get_page_image(page, scale=2.0):
    """Get high quality image from PDF page"""
    # 计算缩放后的尺寸
    zoom = scale
    mat = pymupdf.Matrix(zoom, zoom)
    
    # 使用高分辨率渲染页面
    pix = page.get_pixmap(matrix=mat, alpha=False)
    
    return pix

def main():
    st.set_page_config(layout="wide", page_title="PDF Translator for Human: with Local-LLM/GPT")
    st.title("PDF Translator for Human: with Local-LLM/GPT")

    # Sidebar configuration
    with st.sidebar:
        st.header("Settings")
        
        uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
        
        # Add source language selection
        source_lang_name = st.selectbox(
            "Source Language",
            options=list(SOURCE_LANGUAGE_OPTIONS.keys()),
            index=0  # Default to English
        )
        source_lang = SOURCE_LANGUAGE_OPTIONS[source_lang_name]
        
        translator_name = st.selectbox(
            "Select Translator",
            options=list(TRANSLATORS.keys()),
            index=0
        )
        
        pages_per_load = st.number_input(
            "Pages per load",
            min_value=1,
            max_value=5,
            value=DEFAULT_PAGES_PER_LOAD
        )
        
        text_color = st.selectbox(
            "Translation Color",
            options=list(COLOR_MAP.keys()),
            index=0
        )
        
        # ChatGPT specific settings
        if translator_name == 'chatgpt':
            st.subheader("ChatGPT Settings")
            target_lang = st.selectbox(
                "Target Language",
                options=list(LANGUAGE_OPTIONS.keys()),
                index=0
            )
            api_key = st.text_input(
                "OpenAI API Key",
                value=os.getenv("OPENAI_API_KEY", ""),
                type="password"
            )
            api_base = st.text_input(
                "API Base URL",
                value=os.getenv("OPENAI_API_BASE", DEFAULT_API_BASE)
            )
            model = st.text_input(
                "Model Name",
                value=os.getenv("OPENAI_MODEL", DEFAULT_MODEL)
            )
            
            # Update environment variables
            os.environ["OPENAI_API_KEY"] = api_key
            os.environ["OPENAI_API_BASE"] = api_base
            os.environ["OPENAI_MODEL"] = model
            target_lang = LANGUAGE_OPTIONS[target_lang]
        else:
            # For Google Translator, also show target language selection
            target_lang_name = st.selectbox(
                "Target Language",
                options=list(SOURCE_LANGUAGE_OPTIONS.keys())[:-1],  # Remove "Auto" option
                index=0  # Default to first language
            )
            target_lang = SOURCE_LANGUAGE_OPTIONS[target_lang_name]

    # Main content area
    if uploaded_file is not None:
        doc_bytes = uploaded_file.read()
        doc = pymupdf.open(stream=doc_bytes)
        
        # Create two columns for side-by-side display
        col1, col2 = st.columns(2)
        
        # Initialize session state
        if 'current_page' not in st.session_state:
            st.session_state.current_page = 0
            st.session_state.translation_started = True  # 自动开始翻译
        
        # Display original pages immediately
        with col1:
            st.header("Original")
            for page_num in range(st.session_state.current_page,
                                min(st.session_state.current_page + pages_per_load, doc.page_count)):
                page = doc[page_num]
                pix = get_page_image(page)
                st.image(pix.tobytes(), caption=f"Page {page_num + 1}", use_container_width=True)
        
        # Translation column
        with col2:
            st.header("Translated")
            
            # Configure translator with selected source language
            TranslatorClass = TRANSLATORS[translator_name]
            translator = TranslatorClass(source=source_lang, target=target_lang)
            
            # Translate current batch of pages
            translated_pages = translate_pdf_pages(
                doc,
                doc_bytes,
                st.session_state.current_page,
                pages_per_load,
                translator,
                text_color,
                translator_name,
                target_lang
            )
            
            # Display translated pages
            for i, trans_doc in enumerate(translated_pages):
                page = trans_doc[0]
                pix = get_page_image(page)
                st.image(pix.tobytes(), caption=f"Page {st.session_state.current_page + i + 1}", use_container_width=True)
        
        # Navigation buttons
        nav_col1, nav_col2 = st.columns(2)
        with nav_col1:
            if st.session_state.current_page > 0:
                if st.button("Previous Pages"):
                    st.session_state.current_page = max(0, st.session_state.current_page - pages_per_load)
                    st.rerun()
        
        with nav_col2:
            if st.session_state.current_page + pages_per_load < doc.page_count:
                if st.button("Next Pages"):
                    st.session_state.current_page = min(
                        doc.page_count - 1,
                        st.session_state.current_page + pages_per_load
                    )
                    st.rerun()
    else:
        st.info("Please upload a PDF file to begin translation")

if __name__ == "__main__":
    main()