From 4f0778038b2f0380188715fb37509cc9b1194ab9 Mon Sep 17 00:00:00 2001 From: Tactical-Camell Date: Sat, 12 Sep 2026 12:28:03 +0400 Subject: [PATCH 1/2] feat: add batch processing for image directories --- README.md | 15 ++- batch_process.py | 219 ++++++++++++++++++++++++++++++++++++ main.py | 62 ++++------ pytest.ini | 2 + tests/test_batch_process.py | 150 ++++++++++++++++++++++++ 5 files changed, 405 insertions(+), 43 deletions(-) create mode 100644 batch_process.py create mode 100644 pytest.ini create mode 100644 tests/test_batch_process.py diff --git a/README.md b/README.md index 9b1baa8..957b28e 100644 --- a/README.md +++ b/README.md @@ -34,10 +34,17 @@ Powered by a fine-tuned Vision-Encoder-Decoder model hosted on Hugging Face: ```source .venv/bin/activate``` 3. Install dependencies: ```pip install -r requirements.txt``` -4. Run Transcription: transcribe an image containing handwriting. You can use the default one, or add a new one into the ```images/``` folder and change the directory on ```IMAGE_FILE```. Then run: -``` -python main.py -``` +4. **Run Transcription:** transcribe an image or an entire folder of handwriting images: + - **Single image:** + ```bash + python main.py --image images/test3.png + ``` + - **Batch process an entire directory:** + ```bash + python main.py --folder images/ + # or using batch_process.py: + python batch_process.py images/ --output output/ + ``` --- diff --git a/batch_process.py b/batch_process.py new file mode 100644 index 0000000..48ac9f1 --- /dev/null +++ b/batch_process.py @@ -0,0 +1,219 @@ +import os +import sys +from pathlib import Path +from count_lines import measure_text_lines, crop_and_save_lines, TextLine + +SUPPORTED_EXTENSIONS = {'.jpg', '.jpeg', '.png'} + + +def get_image_files(directory_path: str | Path) -> list[Path]: + """ + Returns a sorted list of all supported image files (.jpg, .jpeg, .png) + found inside directory_path (non-recursive). + """ + path = Path(directory_path) + if not path.exists(): + raise FileNotFoundError(f"Directory not found: {directory_path}") + if not path.is_dir(): + raise NotADirectoryError(f"Path is not a directory: {directory_path}") + + images = [ + f for f in path.iterdir() + if f.is_file() and f.suffix.lower() in SUPPORTED_EXTENSIONS + ] + return sorted(images, key=lambda p: p.name) + + +def cleanup_line_crops(output_dir: str | Path = "output", prefix: str | None = None) -> None: + """Deletes temporary cropped line images while preserving final transcriptions.""" + out = Path(output_dir) + if not out.exists(): + return + for item in out.glob("*.png"): + if "_line_" in item.name: + if prefix is None or item.name.startswith(f"{prefix}_line_"): + try: + item.unlink() + except Exception as e: + print(f"Failed to delete {item}: {e}") + + +def export_to_txt(lines: list[str], output_path: str | Path) -> None: + """Export lines to a UTF-8 text file.""" + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w", encoding="utf-8") as f: + f.write("\n".join(lines)) + print(f"Saved transcript to: {output_path}") + + +def export_to_pdf(lines: list[str], output_path: str | Path) -> None: + """Export lines to a PDF file handling Spanish characters and accents.""" + from fpdf import FPDF + + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + pdf = FPDF(format='letter') + pdf.add_page() + pdf.set_font("Helvetica", size=16) + + for line in lines: + pdf.cell(200, 10, text=line, align='C', new_x="LMARGIN", new_y="NEXT") + pdf.ln() + + pdf.output(str(output_path)) + print(f"Saved transcript to: {output_path}") + + +def transcribe_lines( + line_image_paths: list[str | Path], + processor=None, + model=None, + device: str = "cpu", + transcribe_fn=None, +) -> tuple[list[str], list[str]]: + """ + Transcribes a list of line image crops into raw and NLP-corrected texts. + Accepts either a custom transcribe_fn(img_path, i) or (processor, model, device). + """ + raw_lines: list[str] = [] + nlp_lines: list[str] = [] + + for i, line_img in enumerate(line_image_paths, start=1): + if transcribe_fn is not None: + raw, nlp = transcribe_fn(str(line_img), i) + else: + from main import test_single_phrase + raw, nlp = test_single_phrase(str(line_img), processor, model, i, device=device) + raw_lines.append(raw) + nlp_lines.append(nlp) + + return raw_lines, nlp_lines + + +def process_image( + image_path: str | Path, + processor=None, + model=None, + output_dir: str | Path = "output", + device: str = "cpu", + transcribe_fn=None, +) -> dict[str, str]: + """ + Processes a single image: + 1. Measures text lines using measure_text_lines() + 2. Crops and saves line segments into output_dir + 3. Transcribes each line segment + 4. Cleans up intermediate line crops + 5. Exports raw and NLP transcriptions as .txt and .pdf named after the input image + """ + img_path = Path(image_path) + if not img_path.exists(): + fallback = Path("images") / image_path + if fallback.exists(): + img_path = fallback + else: + raise FileNotFoundError(f"Could not find image '{image_path}'.") + + out_dir = Path(output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + print(f"\n--- Processing Image: {img_path.name} ---") + detected_lines = measure_text_lines(str(img_path)) + print(f"Detected {len(detected_lines)} line(s).") + + if not detected_lines: + print(f"No lines detected in {img_path.name}.") + return {} + + crop_and_save_lines(str(img_path), detected_lines, output_dir=str(out_dir)) + + base_name = img_path.stem + line_image_paths = [ + out_dir / f"{base_name}_line_{i}.png" + for i in range(1, len(detected_lines) + 1) + ] + + raw_lines, nlp_lines = transcribe_lines( + line_image_paths, + processor=processor, + model=model, + device=device, + transcribe_fn=transcribe_fn, + ) + + # Clean up intermediate line crops + cleanup_line_crops(output_dir=out_dir, prefix=base_name) + + txt_raw = out_dir / f"{base_name}_transcription_raw.txt" + pdf_raw = out_dir / f"{base_name}_transcription_raw.pdf" + txt_nlp = out_dir / f"{base_name}_transcription_nlp.txt" + pdf_nlp = out_dir / f"{base_name}_transcription_nlp.pdf" + + export_to_txt(raw_lines, txt_raw) + export_to_pdf(raw_lines, pdf_raw) + export_to_txt(nlp_lines, txt_nlp) + export_to_pdf(nlp_lines, pdf_nlp) + + return { + "txt_raw": str(txt_raw), + "pdf_raw": str(pdf_raw), + "txt_nlp": str(txt_nlp), + "pdf_nlp": str(pdf_nlp), + } + + +def batch_process( + directory_path: str | Path, + processor=None, + model=None, + output_dir: str | Path = "output", + device: str = "cpu", + transcribe_fn=None, +) -> list[dict[str, str]]: + """ + Processes all supported images in a directory. + """ + image_files = get_image_files(directory_path) + print(f"Found {len(image_files)} image(s) to process in '{directory_path}'.") + + results = [] + for img_file in image_files: + res = process_image( + img_file, + processor=processor, + model=model, + output_dir=output_dir, + device=device, + transcribe_fn=transcribe_fn, + ) + if res: + results.append(res) + + return results + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser(description="Batch process directory of images for OCR line extraction and transcription.") + parser.add_argument("directory", nargs="?", default="images", help="Path to folder containing images (default: images)") + parser.add_argument("--output", "-o", default="output", help="Output directory (default: output)") + parser.add_argument("--model-dir", "-m", default="./trocr_spanish_final", help="Path to local TrOCR model directory") + args = parser.parse_args() + + if not os.path.exists(args.model_dir): + raise FileNotFoundError(f"Local model directory '{args.model_dir}' missing.") + + import torch + from transformers import RobertaTokenizer, TrOCRProcessor, VisionEncoderDecoderModel, ViTImageProcessor + + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Loading model from '{args.model_dir}' on {device}...") + image_processor = ViTImageProcessor.from_pretrained(args.model_dir) + tokenizer = RobertaTokenizer.from_pretrained(args.model_dir) + processor = TrOCRProcessor(image_processor=image_processor, tokenizer=tokenizer) + model = VisionEncoderDecoderModel.from_pretrained(args.model_dir).to(device) + model.eval() + + batch_process(args.directory, processor=processor, model=model, output_dir=args.output, device=device) diff --git a/main.py b/main.py index 87a5ad7..f13df52 100644 --- a/main.py +++ b/main.py @@ -41,7 +41,7 @@ def correct_spanish_text(text: str) -> str: return " ".join(corrected_words) -def test_single_phrase(image_path: str, processor: TrOCRProcessor, model: VisionEncoderDecoderModel, i: int) -> tuple[str, str]: +def test_single_phrase(image_path: str, processor: TrOCRProcessor, model: VisionEncoderDecoderModel, i: int, device: str = "cpu") -> tuple[str, str]: # Locate image target_path = image_path if not Path(target_path).exists(): @@ -108,50 +108,34 @@ def export_to_pdf(lines: list[str], output_path: str): print(f"Saved transcript to: {output_path}") if __name__ == "__main__": - detected_lines = run_as_main("images/"+IMAGE_FILE+".png") + import argparse + from batch_process import batch_process, process_image + + parser = argparse.ArgumentParser(description="OCR handwritten line extraction and transcription.") + parser.add_argument("--folder", "-f", default=None, help="Directory containing images to batch process") + parser.add_argument("--image", "-i", default=None, help="Path to single image file to process") + parser.add_argument("--output", "-o", default="output", help="Output directory (default: output)") + parser.add_argument("--model-dir", "-m", default=LOCAL_MODEL_DIR, help="Local model directory") + args = parser.parse_args() device = "cuda" if torch.cuda.is_available() else "cpu" - if not os.path.exists(LOCAL_MODEL_DIR): - raise FileNotFoundError(f"Local model directory '{LOCAL_MODEL_DIR}' missing.") + if not os.path.exists(args.model_dir): + raise FileNotFoundError(f"Local model directory '{args.model_dir}' missing.") - print(f"\nLoading local model from '{LOCAL_MODEL_DIR}' on {device}...") + print(f"\nLoading local model from '{args.model_dir}' on {device}...") # Load processor, tokenizer, and model directly - image_processor = ViTImageProcessor.from_pretrained(LOCAL_MODEL_DIR) - tokenizer = RobertaTokenizer.from_pretrained(LOCAL_MODEL_DIR) + image_processor = ViTImageProcessor.from_pretrained(args.model_dir) + tokenizer = RobertaTokenizer.from_pretrained(args.model_dir) processor = TrOCRProcessor(image_processor=image_processor, tokenizer=tokenizer) - model = VisionEncoderDecoderModel.from_pretrained(LOCAL_MODEL_DIR).to(device) + model = VisionEncoderDecoderModel.from_pretrained(args.model_dir).to(device) model.eval() - RAW_LINES = [] - NLP_LINES = [] - - for i, line in enumerate(detected_lines, start=1): - txt_tuple = test_single_phrase(f"output/{IMAGE_FILE}_line_{i}.png", processor, model, i) - - raw_text = txt_tuple[0] - corrected_text = txt_tuple[1] - - RAW_LINES.append(raw_text) - NLP_LINES.append(corrected_text) - - print(f"\nRaw Lines:\n{RAW_LINES}") - print(f"\nNLP Lines:\n{NLP_LINES}") - - delete_output_content() - print("\nOutput content deleted.") - - base_name = Path(IMAGE_FILE).stem - - txt_out_raw = f"output/{base_name}_transcription_raw.txt" - pdf_out_raw = f"output/{base_name}_transcription_raw.pdf" - txt_out_nlp = f"output/{base_name}_transcription_nlp.txt" - pdf_out_nlp = f"output/{base_name}_transcription_nlp.pdf" - - export_to_txt(RAW_LINES, txt_out_raw) - export_to_pdf(RAW_LINES, pdf_out_raw) - export_to_txt(NLP_LINES, txt_out_nlp) - export_to_pdf(NLP_LINES, pdf_out_nlp) - - print("Transcription files saved.") \ No newline at end of file + if args.folder: + print(f"Starting batch processing on folder: {args.folder}") + batch_process(args.folder, processor=processor, model=model, output_dir=args.output, device=device) + else: + target_image = args.image or f"images/{IMAGE_FILE}.png" + print(f"Processing single image: {target_image}") + process_image(target_image, processor=processor, model=model, output_dir=args.output, device=device) \ No newline at end of file diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..a635c5c --- /dev/null +++ b/pytest.ini @@ -0,0 +1,2 @@ +[pytest] +pythonpath = . diff --git a/tests/test_batch_process.py b/tests/test_batch_process.py new file mode 100644 index 0000000..e02f2cc --- /dev/null +++ b/tests/test_batch_process.py @@ -0,0 +1,150 @@ +import os +from pathlib import Path +import pytest +from batch_process import ( + get_image_files, + cleanup_line_crops, + export_to_txt, + export_to_pdf, + transcribe_lines, + process_image, + batch_process, + SUPPORTED_EXTENSIONS, +) + + +def test_get_image_files(tmp_path): + # Create supported and unsupported test files + (tmp_path / "img1.png").touch() + (tmp_path / "img2.JPG").touch() + (tmp_path / "img3.jpeg").touch() + (tmp_path / "doc.txt").touch() + (tmp_path / "notes.pdf").touch() + (tmp_path / "subfolder").mkdir() + + files = get_image_files(tmp_path) + filenames = [f.name for f in files] + + assert len(files) == 3 + assert filenames == ["img1.png", "img2.JPG", "img3.jpeg"] + + +def test_get_image_files_nonexistent(): + with pytest.raises(FileNotFoundError): + get_image_files("/nonexistent/directory/path/12345") + + +def test_get_image_files_not_a_directory(tmp_path): + file_path = tmp_path / "single_file.png" + file_path.touch() + with pytest.raises(NotADirectoryError): + get_image_files(file_path) + + +def test_cleanup_line_crops(tmp_path): + (tmp_path / "test_line_1.png").touch() + (tmp_path / "test_line_2.png").touch() + (tmp_path / "other_line_1.png").touch() + (tmp_path / "test_transcription_raw.txt").touch() + (tmp_path / "test_transcription_raw.pdf").touch() + + # Clean only 'test' prefix + cleanup_line_crops(tmp_path, prefix="test") + remaining = [f.name for f in tmp_path.iterdir()] + + assert "test_line_1.png" not in remaining + assert "test_line_2.png" not in remaining + assert "other_line_1.png" in remaining + assert "test_transcription_raw.txt" in remaining + assert "test_transcription_raw.pdf" in remaining + + +def test_export_to_txt_and_pdf(tmp_path): + lines = ["Primera línea", "Segunda línea con acento: áéíóú", "Tercera línea: ñ y ¿preguntas?"] + txt_path = tmp_path / "test.txt" + pdf_path = tmp_path / "test.pdf" + + export_to_txt(lines, txt_path) + assert txt_path.exists() + content = txt_path.read_text(encoding="utf-8") + assert "Segunda línea con acento: áéíóú" in content + + export_to_pdf(lines, pdf_path) + assert pdf_path.exists() + assert pdf_path.stat().st_size > 0 + + +def test_transcribe_lines_with_custom_fn(): + mock_images = ["output/sample_line_1.png", "output/sample_line_2.png"] + + def mock_transcribe(img_path, idx): + return f"Raw text {idx}", f"NLP text {idx}" + + raw, nlp = transcribe_lines(mock_images, transcribe_fn=mock_transcribe) + assert raw == ["Raw text 1", "Raw text 2"] + assert nlp == ["NLP text 1", "NLP text 2"] + + +def test_process_image_end_to_end(tmp_path): + output_dir = tmp_path / "output" + sample_image = "images/test3.png" + + def mock_transcribe(img_path, idx): + return f"Raw line {idx}", f"Corrected line {idx}" + + result = process_image( + sample_image, + output_dir=output_dir, + transcribe_fn=mock_transcribe, + ) + + assert result != {} + assert "txt_raw" in result + assert "pdf_raw" in result + assert "txt_nlp" in result + assert "pdf_nlp" in result + + # Check generated files + assert Path(result["txt_raw"]).exists() + assert Path(result["pdf_raw"]).exists() + assert Path(result["txt_nlp"]).exists() + assert Path(result["pdf_nlp"]).exists() + + # Ensure output files are named after the input image stem + assert "test3_transcription_raw.txt" in result["txt_raw"] + assert "test3_transcription_raw.pdf" in result["pdf_raw"] + + # Verify transcript contents + raw_content = Path(result["txt_raw"]).read_text(encoding="utf-8") + assert "Raw line 1" in raw_content + + # Ensure temporary crop images are cleaned up + remaining_crops = list(output_dir.glob("*_line_*.png")) + assert len(remaining_crops) == 0 + + +def test_batch_process_directory(tmp_path): + # Setup test input folder with copies of small test images + in_dir = tmp_path / "input_images" + in_dir.mkdir() + out_dir = tmp_path / "output_dir" + + import shutil + shutil.copy("images/test1.png", in_dir / "sample1.png") + shutil.copy("images/test3.png", in_dir / "sample2.png") + + def mock_transcribe(img_path, idx): + return f"Text {idx}", f"NLP {idx}" + + results = batch_process( + in_dir, + output_dir=out_dir, + transcribe_fn=mock_transcribe, + ) + + assert len(results) == 2 + # Verify outputs for both images exist + assert (out_dir / "sample1_transcription_raw.txt").exists() + assert (out_dir / "sample1_transcription_raw.pdf").exists() + assert (out_dir / "sample2_transcription_raw.txt").exists() + assert (out_dir / "sample2_transcription_raw.pdf").exists() From 97c801c330df0d4a2aa3cb543b1de960235d7602 Mon Sep 17 00:00:00 2001 From: Tactical-Camell Date: Sat, 12 Sep 2026 12:39:18 +0400 Subject: [PATCH 2/2] fix(batch): output exact-named files, sanitize PDF unicode, and support positional CLI paths --- batch_process.py | 33 +++++++++++++++++++++++++----- main.py | 40 ++++++++++++++++++++++++++++--------- tests/test_batch_process.py | 28 ++++++++++++++++++++++---- 3 files changed, 83 insertions(+), 18 deletions(-) diff --git a/batch_process.py b/batch_process.py index 48ac9f1..a4c88f8 100644 --- a/batch_process.py +++ b/batch_process.py @@ -1,7 +1,6 @@ import os -import sys from pathlib import Path -from count_lines import measure_text_lines, crop_and_save_lines, TextLine +from count_lines import measure_text_lines, crop_and_save_lines SUPPORTED_EXTENSIONS = {'.jpg', '.jpeg', '.png'} @@ -47,8 +46,21 @@ def export_to_txt(lines: list[str], output_path: str | Path) -> None: print(f"Saved transcript to: {output_path}") +def _sanitize_pdf_text(text: str) -> str: + """Replaces Unicode punctuation not present in Latin-1 with safe substitutes.""" + replacements = { + '—': '-', '–': '-', '―': '-', + '“': '"', '”': '"', '«': '"', '»': '"', + '‘': "'", '’': "'", '`': "'", + '…': '...', '•': '*', '\xa0': ' ' + } + for orig, repl in replacements.items(): + text = text.replace(orig, repl) + return text.encode('latin-1', errors='replace').decode('latin-1') + + def export_to_pdf(lines: list[str], output_path: str | Path) -> None: - """Export lines to a PDF file handling Spanish characters and accents.""" + """Export lines to a PDF file handling Spanish characters, accents, and punctuation.""" from fpdf import FPDF output_path = Path(output_path) @@ -58,7 +70,8 @@ def export_to_pdf(lines: list[str], output_path: str | Path) -> None: pdf.set_font("Helvetica", size=16) for line in lines: - pdf.cell(200, 10, text=line, align='C', new_x="LMARGIN", new_y="NEXT") + safe_line = _sanitize_pdf_text(line) + pdf.cell(200, 10, text=safe_line, align='C', new_x="LMARGIN", new_y="NEXT") pdf.ln() pdf.output(str(output_path)) @@ -145,17 +158,25 @@ def process_image( # Clean up intermediate line crops cleanup_line_crops(output_dir=out_dir, prefix=base_name) + # Primary output files named exactly after input image + txt_main = out_dir / f"{base_name}.txt" + pdf_main = out_dir / f"{base_name}.pdf" txt_raw = out_dir / f"{base_name}_transcription_raw.txt" pdf_raw = out_dir / f"{base_name}_transcription_raw.pdf" txt_nlp = out_dir / f"{base_name}_transcription_nlp.txt" pdf_nlp = out_dir / f"{base_name}_transcription_nlp.pdf" + final_lines = nlp_lines if any(nlp_lines) else raw_lines + export_to_txt(final_lines, txt_main) + export_to_pdf(final_lines, pdf_main) export_to_txt(raw_lines, txt_raw) export_to_pdf(raw_lines, pdf_raw) export_to_txt(nlp_lines, txt_nlp) export_to_pdf(nlp_lines, pdf_nlp) return { + "txt": str(txt_main), + "pdf": str(pdf_main), "txt_raw": str(txt_raw), "pdf_raw": str(pdf_raw), "txt_nlp": str(txt_nlp), @@ -196,7 +217,9 @@ def batch_process( if __name__ == "__main__": import argparse - parser = argparse.ArgumentParser(description="Batch process directory of images for OCR line extraction and transcription.") + parser = argparse.ArgumentParser( + description="Batch process directory of images for OCR line extraction and transcription." + ) parser.add_argument("directory", nargs="?", default="images", help="Path to folder containing images (default: images)") parser.add_argument("--output", "-o", default="output", help="Output directory (default: output)") parser.add_argument("--model-dir", "-m", default="./trocr_spanish_final", help="Path to local TrOCR model directory") diff --git a/main.py b/main.py index f13df52..1c6a3ba 100644 --- a/main.py +++ b/main.py @@ -1,4 +1,5 @@ -import os, shutil +import os +import shutil from pathlib import Path import re import torch @@ -11,12 +12,12 @@ ViTImageProcessor, ) from spellchecker import SpellChecker -from count_lines import run_as_main # Target a single-phrase test image directly IMAGE_FILE = "test3" LOCAL_MODEL_DIR = "./trocr_spanish_final" + def correct_spanish_text(text: str) -> str: """Certifies words against the Spanish vocabulary and corrects typos/confusions.""" spell = SpellChecker(language='es') @@ -41,7 +42,14 @@ def correct_spanish_text(text: str) -> str: return " ".join(corrected_words) -def test_single_phrase(image_path: str, processor: TrOCRProcessor, model: VisionEncoderDecoderModel, i: int, device: str = "cpu") -> tuple[str, str]: + +def test_single_phrase( + image_path: str, + processor: TrOCRProcessor, + model: VisionEncoderDecoderModel, + i: int, + device: str = "cpu" +) -> tuple[str, str]: # Locate image target_path = image_path if not Path(target_path).exists(): @@ -73,6 +81,7 @@ def test_single_phrase(image_path: str, processor: TrOCRProcessor, model: Vision return text, corrected_text + def delete_output_content(): folder = "output" for filename in os.listdir(folder): @@ -85,6 +94,7 @@ def delete_output_content(): except Exception as e: print('Failed to delete %s. Reason: %s' % (file_path, e)) + def export_to_txt(lines: list[str], output_path: str): """Export lines to a UTF-8 text file.""" with open(output_path, "w", encoding="utf-8") as f: @@ -95,7 +105,7 @@ def export_to_txt(lines: list[str], output_path: str): def export_to_pdf(lines: list[str], output_path: str): """Export lines to a PDF file handling Spanish characters and accents.""" pdf = FPDF(format='letter') - #pdf.set_auto_page_break(auto=True, margin=15) + # pdf.set_auto_page_break(auto=True, margin=15) pdf.add_page() pdf.set_font("Helvetica", size=16) @@ -107,11 +117,13 @@ def export_to_pdf(lines: list[str], output_path: str): pdf.output(output_path) print(f"Saved transcript to: {output_path}") + if __name__ == "__main__": import argparse from batch_process import batch_process, process_image parser = argparse.ArgumentParser(description="OCR handwritten line extraction and transcription.") + parser.add_argument("path", nargs="?", default=None, help="Path to directory or image file to process") parser.add_argument("--folder", "-f", default=None, help="Directory containing images to batch process") parser.add_argument("--image", "-i", default=None, help="Path to single image file to process") parser.add_argument("--output", "-o", default="output", help="Output directory (default: output)") @@ -132,10 +144,20 @@ def export_to_pdf(lines: list[str], output_path: str): model = VisionEncoderDecoderModel.from_pretrained(args.model_dir).to(device) model.eval() - if args.folder: - print(f"Starting batch processing on folder: {args.folder}") - batch_process(args.folder, processor=processor, model=model, output_dir=args.output, device=device) + folder_target = args.folder + image_target = args.image + + if args.path: + p = Path(args.path) + if p.is_dir() or (not p.exists() and not p.suffix): + folder_target = args.path + else: + image_target = args.path + + if folder_target: + print(f"Starting batch processing on folder: {folder_target}") + batch_process(folder_target, processor=processor, model=model, output_dir=args.output, device=device) else: - target_image = args.image or f"images/{IMAGE_FILE}.png" + target_image = image_target or f"images/{IMAGE_FILE}.png" print(f"Processing single image: {target_image}") - process_image(target_image, processor=processor, model=model, output_dir=args.output, device=device) \ No newline at end of file + process_image(target_image, processor=processor, model=model, output_dir=args.output, device=device) diff --git a/tests/test_batch_process.py b/tests/test_batch_process.py index e02f2cc..9b007dc 100644 --- a/tests/test_batch_process.py +++ b/tests/test_batch_process.py @@ -1,4 +1,3 @@ -import os from pathlib import Path import pytest from batch_process import ( @@ -9,7 +8,6 @@ transcribe_lines, process_image, batch_process, - SUPPORTED_EXTENSIONS, ) @@ -74,6 +72,16 @@ def test_export_to_txt_and_pdf(tmp_path): assert pdf_path.stat().st_size > 0 +def test_export_to_pdf_unicode_punctuation(tmp_path): + # Tests non-Latin-1 characters like em-dashes and curly quotes + lines = ["Nota con guión — y comillas “españolas”", "Puntos suspensivos… y viñeta •"] + pdf_path = tmp_path / "unicode_test.pdf" + + export_to_pdf(lines, pdf_path) + assert pdf_path.exists() + assert pdf_path.stat().st_size > 0 + + def test_transcribe_lines_with_custom_fn(): mock_images = ["output/sample_line_1.png", "output/sample_line_2.png"] @@ -99,22 +107,30 @@ def mock_transcribe(img_path, idx): ) assert result != {} + assert "txt" in result + assert "pdf" in result assert "txt_raw" in result assert "pdf_raw" in result assert "txt_nlp" in result assert "pdf_nlp" in result # Check generated files + assert Path(result["txt"]).exists() + assert Path(result["pdf"]).exists() assert Path(result["txt_raw"]).exists() assert Path(result["pdf_raw"]).exists() assert Path(result["txt_nlp"]).exists() assert Path(result["pdf_nlp"]).exists() - # Ensure output files are named after the input image stem + # Ensure output files are named exactly after the input image stem + assert result["txt"].endswith("test3.txt") + assert result["pdf"].endswith("test3.pdf") assert "test3_transcription_raw.txt" in result["txt_raw"] assert "test3_transcription_raw.pdf" in result["pdf_raw"] # Verify transcript contents + main_content = Path(result["txt"]).read_text(encoding="utf-8") + assert "Corrected line 1" in main_content raw_content = Path(result["txt_raw"]).read_text(encoding="utf-8") assert "Raw line 1" in raw_content @@ -143,8 +159,12 @@ def mock_transcribe(img_path, idx): ) assert len(results) == 2 - # Verify outputs for both images exist + # Verify outputs for both images exist, including exact-named files + assert (out_dir / "sample1.txt").exists() + assert (out_dir / "sample1.pdf").exists() assert (out_dir / "sample1_transcription_raw.txt").exists() assert (out_dir / "sample1_transcription_raw.pdf").exists() + assert (out_dir / "sample2.txt").exists() + assert (out_dir / "sample2.pdf").exists() assert (out_dir / "sample2_transcription_raw.txt").exists() assert (out_dir / "sample2_transcription_raw.pdf").exists()