From 428cc2a01a10d536a60c92a880672660a2f18952 Mon Sep 17 00:00:00 2001 From: Kenneth Verbeure Date: Thu, 26 Feb 2026 22:07:22 +0100 Subject: [PATCH 1/4] feat: add return_page_numbers parameter to segment_chemical_structures_from_file --- decimer_segmentation/decimer_segmentation.py | 69 +++++++++++--- tests/conftest.py | 40 +++++++-- tests/helpers.py | 36 ++++++++ tests/test_segmentation.py | 95 +++++++++++++++----- 4 files changed, 195 insertions(+), 45 deletions(-) create mode 100644 tests/helpers.py diff --git a/decimer_segmentation/decimer_segmentation.py b/decimer_segmentation/decimer_segmentation.py index fb71bc6..fd55288 100644 --- a/decimer_segmentation/decimer_segmentation.py +++ b/decimer_segmentation/decimer_segmentation.py @@ -127,7 +127,7 @@ def _load_model_internal() -> modellib.MaskRCNN: logger.debug(f"Some optimizer options not available: {e}") model_path = pystow.join("DECIMER-Segmentation_model") - print(model_path) + logger.debug("Model path: %s", model_path) if not os.path.exists(str(model_path) + "/" + MODEL_FILENAME): logger.info("Downloading model weights...") download_trained_weights(MODEL_DOWNLOAD_URL, str(model_path)) @@ -184,8 +184,14 @@ def segment_chemical_structures_from_file( file_path: str, expand: bool = True, return_bboxes: bool = False, + return_page_numbers: bool = False, **kwargs, -) -> Union[List[np.ndarray], Tuple[List[np.ndarray], List[Tuple[int, int, int, int]]]]: +) -> Union[ + List[np.ndarray], + Tuple[List[np.ndarray], List[Tuple[int, int, int, int]]], + Tuple[List[np.ndarray], List[int]], + Tuple[List[np.ndarray], List[Tuple[int, int, int, int]], List[int]], +]: """ Segment chemical structures from a PDF or image file. @@ -193,10 +199,13 @@ def segment_chemical_structures_from_file( file_path: Path to input file (PDF or image) expand: Whether to expand masks to capture complete structures return_bboxes: Whether to return bounding boxes along with segments + return_page_numbers: Whether to return 1-indexed page numbers per segment Returns: List of segmented chemical structure images as numpy arrays. If return_bboxes is True, returns a tuple of (segments, bboxes). + If return_page_numbers is True, returns (segments, page_numbers). + If both are True, returns (segments, bboxes, page_numbers). """ if not os.path.exists(file_path): raise FileNotFoundError(f"Input file not found: {file_path}") @@ -204,35 +213,61 @@ def segment_chemical_structures_from_file( if "poppler_path" in kwargs: logger.warning("poppler_path parameter is deprecated and ignored") - images = _load_images_from_file(file_path) + images_with_pages = _load_images_from_file(file_path) - if not images: + if not images_with_pages: logger.warning(f"No images could be extracted from {file_path}") - return ([], []) if return_bboxes else [] + if return_bboxes and return_page_numbers: + return ([], [], []) + elif return_bboxes: + return ([], []) + elif return_page_numbers: + return ([], []) + else: + return [] # Process all images sequentially (model can't parallelize) all_segments = [] all_bboxes = [] - for image in images: + all_page_numbers = [] + for image, page_number in images_with_pages: segments, bboxes = segment_chemical_structures( image, expand, return_bboxes=True ) all_segments.extend(segments) all_bboxes.extend(bboxes) + all_page_numbers.extend([page_number] * len(segments)) + + if return_bboxes and return_page_numbers: + return (all_segments, all_bboxes, all_page_numbers) + elif return_bboxes: + return (all_segments, all_bboxes) + elif return_page_numbers: + return (all_segments, all_page_numbers) + else: + return all_segments - return (all_segments, all_bboxes) if return_bboxes else all_segments +def _load_images_from_file( + file_path: str, +) -> List[Tuple[np.ndarray, int]]: + """Load images from PDF or image file. -def _load_images_from_file(file_path: str) -> List[np.ndarray]: - """Load images from PDF or image file.""" + Returns: + List of (image, page_number) tuples. Page numbers are 1-indexed. + """ if file_path.lower().endswith(".pdf"): return _load_pdf_pages(file_path) else: - return _load_single_image(file_path) + return [(img, 1) for img in _load_single_image(file_path)] + +def _load_pdf_pages(pdf_path: str) -> List[Tuple[np.ndarray, int]]: + """Load all pages from a PDF as images using PyMuPDF. -def _load_pdf_pages(pdf_path: str) -> List[np.ndarray]: - """Load all pages from a PDF as images using PyMuPDF.""" + Returns: + List of (image, page_number) tuples. Page numbers are 1-indexed. + """ pdf_document = pymupdf.open(pdf_path) page_count = pdf_document.page_count @@ -247,7 +282,7 @@ def _load_pdf_pages(pdf_path: str) -> List[np.ndarray]: if pix.n == 3: img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR) pdf_document.close() - return [img_array.copy()] + return [(img_array.copy(), 1)] # Multiple pages - use threading for I/O images = [None] * page_count @@ -270,7 +305,13 @@ def render_page(page_num: int) -> Tuple[int, np.ndarray]: images[page_num] = img_array pdf_document.close() - return [img for img in images if img is not None] + + # Convert to (image, page_number) tuples with 1-indexed page numbers + result = [] + for page_num, img in enumerate(images): + if img is not None: + result.append((img, page_num + 1)) + return result def _load_single_image(image_path: str) -> List[np.ndarray]: diff --git a/tests/conftest.py b/tests/conftest.py index 7c24f69..75baaba 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -9,6 +9,7 @@ import warnings import pytest import numpy as np +import pystow import tempfile # Suppress SWIG-related deprecation warnings from third-party libraries (e.g., OpenCV) @@ -91,14 +92,40 @@ def temp_image_file(sample_image): os.unlink(f.name) +@pytest.fixture +def mock_model_detection(monkeypatch): + """ + Mock the model detection to return fake results without running the actual model. + Returns 2 fake detections per image for testing. + """ + + def mock_get_mrcnn_results(image): + """Return fake detection results.""" + h, w = image.shape[:2] + # Create 2 fake masks + masks = np.zeros((h, w, 2), dtype=bool) + masks[50:150, 50:150, 0] = True + masks[200:300, 200:300, 1] = True + + # Fake bounding boxes (y0, x0, y1, x1) + bboxes = [(50, 50, 150, 150), (200, 200, 300, 300)] + + # Fake confidence scores + scores = [0.95, 0.92] + + return masks, bboxes, scores + + # Patch the get_mrcnn_results function + import decimer_segmentation.decimer_segmentation as ds + + monkeypatch.setattr(ds, "get_mrcnn_results", mock_get_mrcnn_results) + + @pytest.fixture(scope="session") def model_available(): """Check if the model weights are available.""" - import os - - model_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) model_path = os.path.join( - model_dir, "decimer_segmentation", "mask_rcnn_molecule.h5" + str(pystow.join("DECIMER-Segmentation_model")), "mask_rcnn_molecule.h5" ) return os.path.exists(model_path) @@ -118,11 +145,8 @@ def pytest_configure(config): # Skip model-dependent tests if model not available def pytest_collection_modifyitems(config, items): """Modify test collection based on available resources.""" - import os - - model_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) model_path = os.path.join( - model_dir, "decimer_segmentation", "mask_rcnn_molecule.h5" + str(pystow.join("DECIMER-Segmentation_model")), "mask_rcnn_molecule.h5" ) model_available = os.path.exists(model_path) diff --git a/tests/helpers.py b/tests/helpers.py new file mode 100644 index 0000000..fae2d67 --- /dev/null +++ b/tests/helpers.py @@ -0,0 +1,36 @@ +""" +Test helper functions for DECIMER Segmentation tests. +""" + +import tempfile +import pymupdf + + +def create_test_pdf(num_pages: int = 1) -> str: + """ + Create a temporary PDF with black rectangles (simulating chemical structures). + + Args: + num_pages: Number of pages to create + + Returns: + Path to temporary PDF file (caller must delete) + """ + f = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) + pdf = pymupdf.open() + + for i in range(num_pages): + page = pdf.new_page(width=500, height=500) + # Place structure at different positions on each page + x_offset = 50 + (i * 100) % 300 + y_offset = 50 + (i * 100) % 300 + page.draw_rect( + pymupdf.Rect(x_offset, y_offset, x_offset + 100, y_offset + 100), + color=(0, 0, 0), + fill=(0, 0, 0), + ) + + pdf.save(f.name) + pdf.close() + f.close() + return f.name diff --git a/tests/test_segmentation.py b/tests/test_segmentation.py index 01a25a0..c0aab50 100644 --- a/tests/test_segmentation.py +++ b/tests/test_segmentation.py @@ -16,6 +16,8 @@ # Add parent directory to path for imports sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from tests.helpers import create_test_pdf + class TestImageProcessing: """Tests for image processing functions.""" @@ -541,32 +543,79 @@ def test_full_pipeline_empty_image(self): # Expected if model not available pass - def test_segment_from_file_with_return_bboxes(self): - """Test segment_chemical_structures_from_file returns bboxes when return_bboxes=True.""" + def test_segment_from_file_with_return_bboxes(self, mock_model_detection): + """ + Test segment_chemical_structures_from_file returns bboxes + when return_bboxes=True. + """ from decimer_segmentation import segment_chemical_structures_from_file - image = np.ones((500, 500, 3), dtype=np.uint8) * 255 - cv2.rectangle(image, (50, 50), (150, 150), (0, 0, 0), -1) + pdf_path = create_test_pdf(num_pages=1) + try: + result = segment_chemical_structures_from_file( + pdf_path, expand=False, return_bboxes=True + ) + assert isinstance(result, tuple) + assert len(result) == 2 + segments, bboxes = result + assert isinstance(segments, list) + assert isinstance(bboxes, list) + assert len(segments) == len(bboxes) + assert len(segments) == 2, "Mocked model should return 2 detections" + for bbox in bboxes: + assert len(bbox) == 4 + finally: + os.unlink(pdf_path) + + def test_segment_multipage_pdf_with_page_numbers(self, mock_model_detection): + """ + Test segment_chemical_structures_from_file returns page numbers + for multi-page PDFs. + """ + from decimer_segmentation import segment_chemical_structures_from_file - with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f: - cv2.imwrite(f.name, image) - try: - result = segment_chemical_structures_from_file( - f.name, expand=False, return_bboxes=True - ) - assert isinstance(result, tuple) - assert len(result) == 2 - segments, bboxes = result - assert isinstance(segments, list) - assert isinstance(bboxes, list) - assert len(segments) == len(bboxes) - for bbox in bboxes: - assert len(bbox) == 4 - except Exception: - # Expected if model not available - pass - finally: - os.unlink(f.name) + pdf_path = create_test_pdf(num_pages=3) + try: + result = segment_chemical_structures_from_file( + pdf_path, expand=False, return_page_numbers=True + ) + segments, page_numbers = result + + # Mock returns 2 detections per page: pages should be [1, 1, 2, 2, 3, 3] + assert page_numbers == [1, 1, 2, 2, 3, 3], ( + f"Expected page numbers [1, 1, 2, 2, 3, 3], got {page_numbers}" + ) + + # Verify we got 6 valid segments + assert len(segments) == 6 + for i, seg in enumerate(segments): + assert isinstance(seg, np.ndarray), f"Segment {i} should be numpy array" + assert seg.size > 0, f"Segment {i} should not be empty" + finally: + os.unlink(pdf_path) + + def test_segment_pdf_with_bboxes_and_page_numbers(self, mock_model_detection): + """ + Test segment_chemical_structures_from_file returns both bboxes + and page numbers. + """ + from decimer_segmentation import segment_chemical_structures_from_file + + pdf_path = create_test_pdf(num_pages=1) + try: + result = segment_chemical_structures_from_file( + pdf_path, expand=False, return_bboxes=True, return_page_numbers=True + ) + segments, bboxes, page_numbers = result + + # Mock returns exact known values + assert page_numbers == [1, 1] + assert bboxes == [(50, 50, 150, 150), (200, 200, 300, 300)] + assert len(segments) == 2 + for seg in segments: + assert isinstance(seg, np.ndarray) and seg.size > 0 + finally: + os.unlink(pdf_path) # Benchmark tests (optional, for performance monitoring) From 44c7fe5a6a81cdc002c2ff9eee54519febbfba72 Mon Sep 17 00:00:00 2001 From: Kenneth Verbeure Date: Fri, 27 Feb 2026 09:12:04 +0100 Subject: [PATCH 2/4] fix: improve thread safety in PDF loading by using separate document handles per thread --- decimer_segmentation/decimer_segmentation.py | 41 ++++++++++++-------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/decimer_segmentation/decimer_segmentation.py b/decimer_segmentation/decimer_segmentation.py index fd55288..150ed2c 100644 --- a/decimer_segmentation/decimer_segmentation.py +++ b/decimer_segmentation/decimer_segmentation.py @@ -268,11 +268,19 @@ def _load_pdf_pages(pdf_path: str) -> List[Tuple[np.ndarray, int]]: Returns: List of (image, page_number) tuples. Page numbers are 1-indexed. """ - pdf_document = pymupdf.open(pdf_path) - page_count = pdf_document.page_count + # Quick check for page count to determine strategy + with pymupdf.open(pdf_path) as doc: + page_count = doc.page_count if page_count == 1: - # Single page - no threading overhead + return _load_pdf_single_page(pdf_path) + else: + return _load_pdf_multipage(pdf_path, page_count) + + +def _load_pdf_single_page(pdf_path: str) -> List[Tuple[np.ndarray, int]]: + """Load a single-page PDF without threading overhead.""" + with pymupdf.open(pdf_path) as pdf_document: page = pdf_document[0] matrix = pymupdf.Matrix(300 / 72, 300 / 72) pix = page.get_pixmap(matrix=matrix, alpha=False) @@ -281,22 +289,25 @@ def _load_pdf_pages(pdf_path: str) -> List[Tuple[np.ndarray, int]]: ) if pix.n == 3: img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR) - pdf_document.close() return [(img_array.copy(), 1)] - # Multiple pages - use threading for I/O + +def _load_pdf_multipage(pdf_path: str, page_count: int) -> List[Tuple[np.ndarray, int]]: + """Load multi-page PDF using threading with separate document handles per thread.""" images = [None] * page_count def render_page(page_num: int) -> Tuple[int, np.ndarray]: - page = pdf_document[page_num] - matrix = pymupdf.Matrix(300 / 72, 300 / 72) - pix = page.get_pixmap(matrix=matrix, alpha=False) - img_array = np.frombuffer(pix.samples, dtype=np.uint8).reshape( - pix.h, pix.w, pix.n - ) - if pix.n == 3: - img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR) - return page_num, img_array.copy() + # Each thread opens its own document handle for thread safety + with pymupdf.open(pdf_path) as doc: + page = doc[page_num] + matrix = pymupdf.Matrix(300 / 72, 300 / 72) + pix = page.get_pixmap(matrix=matrix, alpha=False) + img_array = np.frombuffer(pix.samples, dtype=np.uint8).reshape( + pix.h, pix.w, pix.n + ) + if pix.n == 3: + img_array = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR) + return page_num, img_array.copy() with ThreadPoolExecutor(max_workers=min(4, page_count)) as executor: futures = [executor.submit(render_page, i) for i in range(page_count)] @@ -304,8 +315,6 @@ def render_page(page_num: int) -> Tuple[int, np.ndarray]: page_num, img_array = future.result() images[page_num] = img_array - pdf_document.close() - # Convert to (image, page_number) tuples with 1-indexed page numbers result = [] for page_num, img in enumerate(images): From 72737f49c9e85e487e1baab3235d0ffb80d0fbad Mon Sep 17 00:00:00 2001 From: Kohulan Rajan Date: Wed, 11 Mar 2026 22:07:08 +0100 Subject: [PATCH 3/4] Update tests/helpers.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/helpers.py | 33 +++++++++++++++++---------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/tests/helpers.py b/tests/helpers.py index fae2d67..654e657 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -17,20 +17,21 @@ def create_test_pdf(num_pages: int = 1) -> str: Path to temporary PDF file (caller must delete) """ f = tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) - pdf = pymupdf.open() - - for i in range(num_pages): - page = pdf.new_page(width=500, height=500) - # Place structure at different positions on each page - x_offset = 50 + (i * 100) % 300 - y_offset = 50 + (i * 100) % 300 - page.draw_rect( - pymupdf.Rect(x_offset, y_offset, x_offset + 100, y_offset + 100), - color=(0, 0, 0), - fill=(0, 0, 0), - ) - - pdf.save(f.name) - pdf.close() + pdf_path = f.name f.close() - return f.name + + with pymupdf.open() as pdf: + for i in range(num_pages): + page = pdf.new_page(width=500, height=500) + # Place structure at different positions on each page + x_offset = 50 + (i * 100) % 300 + y_offset = 50 + (i * 100) % 300 + page.draw_rect( + pymupdf.Rect(x_offset, y_offset, x_offset + 100, y_offset + 100), + color=(0, 0, 0), + fill=(0, 0, 0), + ) + + pdf.save(pdf_path) + + return pdf_path From 901117b94634ea5e2d863477281176bdea799bec Mon Sep 17 00:00:00 2001 From: Kohulan Rajan Date: Wed, 11 Mar 2026 22:07:21 +0100 Subject: [PATCH 4/4] Update decimer_segmentation/decimer_segmentation.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- decimer_segmentation/decimer_segmentation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/decimer_segmentation/decimer_segmentation.py b/decimer_segmentation/decimer_segmentation.py index 150ed2c..40c2054 100644 --- a/decimer_segmentation/decimer_segmentation.py +++ b/decimer_segmentation/decimer_segmentation.py @@ -128,7 +128,7 @@ def _load_model_internal() -> modellib.MaskRCNN: model_path = pystow.join("DECIMER-Segmentation_model") logger.debug("Model path: %s", model_path) - if not os.path.exists(str(model_path) + "/" + MODEL_FILENAME): + if not os.path.exists(os.path.join(model_path, MODEL_FILENAME)): logger.info("Downloading model weights...") download_trained_weights(MODEL_DOWNLOAD_URL, str(model_path)) logger.info("Successfully downloaded the segmentation model weights!")