diff --git a/src/utils/__init__.py b/src/utils/__init__.py index fabdc03..5d1677d 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -5,7 +5,7 @@ from PIL import Image -def calculate_iou(box1, box2): +def calculate_iou(box1: tuple[float], box2: tuple[float]) -> float: """ Calculate Intersection over Union (IoU) between two bounding boxes. @@ -34,8 +34,8 @@ def calculate_iou(box1, box2): return iou -def pdf_to_img(pdf_path, folder, resolution=300): - file_path = os.path.join(folder, f"{pdf_path}") +def pdf_to_img(file_name: str, folder: str, resolution: int = 300): + file_path = os.path.join(folder, file_name) file_handle = fitz.open(file_path) # Get the first page @@ -55,8 +55,8 @@ def pdf_to_img(pdf_path, folder, resolution=300): def generate_id(): - id = str(uuid.uuid4()) - return id + id = str(uuid.uuid4()) + return id def contains_number(s): """ diff --git a/test/unit/src/test_utils.py b/test/unit/src/test_utils.py index fd25d2d..5285a4c 100644 --- a/test/unit/src/test_utils.py +++ b/test/unit/src/test_utils.py @@ -1,6 +1,12 @@ -from src.utils import calculate_iou - +from src.utils import ( + calculate_iou, + pdf_to_img, +) +from PIL import Image +import fitz +import os import pytest +import shutil @pytest.mark.parametrize( ("box1", "box2", "expected"), @@ -10,7 +16,47 @@ ((0.2, 0.2, 0.0, 0.0), (0.2, 0.2, 0.1, 0.0), 1/3,), ], ) -def test_calculate_iou(box1: tuple, box2: tuple, expected: float) -> None: +def test_calculate_iou(box1: tuple[float], box2: tuple[float], expected: float) -> None: iou = calculate_iou(box1, box2) assert iou >= 0 and iou <= 1 - assert pytest.approx(iou) == expected \ No newline at end of file + assert pytest.approx(iou) == expected + + +def _create_sample_pdf(pdf_path: str) -> None: + doc = fitz.open() + doc.insert_page(0) + doc.save(str(pdf_path)) + + +@pytest.mark.parametrize( + ("resolution", "expected_width", "expected_height", "expected_format"), + [ + (300, 2480, 3508, "RGB"), + (150, 1240, 1754, "RGB"), + ], +) +def test_pdf_to_img(resolution: int, expected_width: int, expected_height: int, expected_format: str) -> None: + folder = r"./test/unit/temp/" + + # Create the output folder if it doesn't exist + os.makedirs(folder, exist_ok=True) + + file_name = "sample.pdf" + + pdf_path = os.path.join(folder, file_name) + + _create_sample_pdf(pdf_path) + + img = pdf_to_img(file_name, folder, resolution) + + # Check if the output image is an instance of PIL Image + assert isinstance(img, Image.Image) + + # Check image dimensions + assert pytest.approx(img.size, rel=1.0) == (expected_width, expected_height) + + # Check image format + assert img.mode == expected_format + + # Clean up: Remove the temporary folder and its contents + shutil.rmtree(folder)