121 lines
3.9 KiB
Python
121 lines
3.9 KiB
Python
import os
|
|
import shutil
|
|
import subprocess
|
|
import uuid
|
|
from flask import Flask, request, render_template, send_file, jsonify
|
|
from werkzeug.utils import secure_filename
|
|
|
|
app = Flask(__name__)
|
|
|
|
BASE_UPLOAD_FOLDER = 'uploads'
|
|
BASE_OUTPUT_FOLDER = 'output'
|
|
os.makedirs(BASE_UPLOAD_FOLDER, exist_ok=True)
|
|
os.makedirs(BASE_OUTPUT_FOLDER, exist_ok=True)
|
|
|
|
|
|
def cleanup_session_folders(session_id):
|
|
"""Remove any existing session folders to ensure a clean slate."""
|
|
upload_folder = os.path.join(BASE_UPLOAD_FOLDER, session_id)
|
|
output_folder = os.path.join(BASE_OUTPUT_FOLDER, session_id)
|
|
if os.path.exists(upload_folder):
|
|
shutil.rmtree(upload_folder)
|
|
if os.path.exists(output_folder):
|
|
shutil.rmtree(output_folder)
|
|
|
|
|
|
def create_session_folders(session_id):
|
|
"""
|
|
Clears previous session folders if they exist and creates new ones.
|
|
This ensures that files from previous uploads won't mix with new ones.
|
|
"""
|
|
cleanup_session_folders(session_id)
|
|
upload_folder = os.path.join(BASE_UPLOAD_FOLDER, session_id)
|
|
output_folder = os.path.join(BASE_OUTPUT_FOLDER, session_id)
|
|
os.makedirs(upload_folder, exist_ok=True)
|
|
os.makedirs(output_folder, exist_ok=True)
|
|
return upload_folder, output_folder
|
|
|
|
|
|
def run_ocrmypdf(input_pdf, output_pdf):
|
|
cmd = [
|
|
'ocrmypdf',
|
|
'--rotate-pages',
|
|
'--deskew',
|
|
'--output-type', 'pdfa',
|
|
'--jobs', '4',
|
|
input_pdf,
|
|
output_pdf
|
|
]
|
|
result = subprocess.run(cmd, capture_output=True, text=True)
|
|
if result.returncode != 0:
|
|
if "PriorOcrFoundError" in result.stderr:
|
|
raise ValueError("This file already contains selectable text and was skipped.")
|
|
else:
|
|
raise RuntimeError(f"OCRmyPDF failed: {result.stderr}")
|
|
|
|
|
|
@app.route('/')
|
|
def index():
|
|
# Generate a new session id for each visitor.
|
|
session_id = str(uuid.uuid4())
|
|
return render_template('index.html', session_id=session_id)
|
|
|
|
|
|
@app.route('/upload/<session_id>', methods=['POST'])
|
|
def upload_files(session_id):
|
|
# Create fresh session folders, clearing any previous data.
|
|
upload_folder, output_folder = create_session_folders(session_id)
|
|
|
|
if 'files' not in request.files:
|
|
return jsonify({'error': 'No files were uploaded.'}), 400
|
|
|
|
files = request.files.getlist('files')
|
|
processed_files = []
|
|
skipped_files = []
|
|
|
|
for file in files:
|
|
if file.filename == '':
|
|
continue
|
|
|
|
filename = secure_filename(file.filename)
|
|
input_path = os.path.join(upload_folder, filename)
|
|
output_path = os.path.join(output_folder, filename)
|
|
|
|
file.save(input_path)
|
|
|
|
try:
|
|
run_ocrmypdf(input_path, output_path)
|
|
processed_files.append(filename)
|
|
except ValueError:
|
|
skipped_files.append(filename)
|
|
except Exception as e:
|
|
# If there's an error, cleanup the session folders
|
|
cleanup_session_folders(session_id)
|
|
return jsonify({'error': f"Failed to process {filename}: {str(e)}"}), 500
|
|
|
|
if not processed_files:
|
|
return jsonify({
|
|
'error': 'All files were skipped because they already contain selectable text.',
|
|
'skipped_files': skipped_files
|
|
})
|
|
|
|
# Create a zip file containing only the current upload's processed files.
|
|
zip_filename = os.path.join(output_folder, 'processed_files.zip')
|
|
subprocess.run(['zip', '-j', zip_filename] + [os.path.join(output_folder, f) for f in processed_files])
|
|
|
|
return jsonify({
|
|
'download_url': f'/download/{session_id}/processed_files.zip',
|
|
'processed_files': processed_files,
|
|
'skipped_files': skipped_files
|
|
})
|
|
|
|
|
|
@app.route('/download/<session_id>/<filename>')
|
|
def download_file(session_id, filename):
|
|
file_path = os.path.join(BASE_OUTPUT_FOLDER, session_id, filename)
|
|
return send_file(file_path, as_attachment=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
app.run(host='0.0.0.0', port=5000, debug=False)
|