feat: organize base class (#36)

* feat: add LayoutTransform class and update visualizer batch method

* feat: refactor transforms to inherit from LayoutTransform

* refactor: clean up unused imports and improve test structure

Author: shunk031

src/layout_prompter/transforms/__init__.py

@@ -1,10 +1,12 @@
+from .base import LayoutTransform
 from .discretize_bboxes import DiscretizeBboxes
 from .label_dict_sort import LabelDictSort
 from .lexicographic_sort import LexicographicSort
 from .saliency_map_to_bboxes import SaliencyMapToBboxes
 from .shuffle_elements import ShuffleElements
 
 __all__ = [
+    "LayoutTransform",
     "DiscretizeBboxes",
     "LabelDictSort",
     "LexicographicSort",

src/layout_prompter/transforms/base.py

@@ -0,0 +1,30 @@
+import abc
+from typing import Any, List, Optional, Union
+
+from langchain_core.runnables import RunnableSerializable
+from langchain_core.runnables.config import RunnableConfig
+
+from layout_prompter.models import LayoutData, ProcessedLayoutData
+
+
+class LayoutTransform(RunnableSerializable):
+    @abc.abstractmethod
+    def invoke(
+        self,
+        input: Union[LayoutData, ProcessedLayoutData],
+        config: Optional[RunnableConfig] = None,
+        **kwargs: Any,
+    ) -> Any:
+        raise NotImplementedError
+
+    def batch(
+        self,
+        inputs: Union[List[LayoutData], List[ProcessedLayoutData]],
+        config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
+        *,
+        return_exceptions: bool = False,
+        **kwargs: Any | None,
+    ) -> List:
+        return super().batch(
+            inputs, config, return_exceptions=return_exceptions, **kwargs
+        )

src/layout_prompter/transforms/discretize_bboxes.py

@@ -1,21 +1,22 @@
 import copy
 from typing import Any, Optional, Union
 
-from langchain_core.runnables import RunnableSerializable
 from langchain_core.runnables.config import RunnableConfig
 from loguru import logger
 
 from layout_prompter.models import CanvasSize, LayoutData, ProcessedLayoutData
 from layout_prompter.utils import Configuration
 
+from .base import LayoutTransform
+
 
 class DiscretizeBboxesConfig(Configuration):
     """Configuration for Transform classes."""
 
     target_canvas_size: CanvasSize
 
 
-class DiscretizeBboxes(RunnableSerializable):
+class DiscretizeBboxes(LayoutTransform):
     name: str = "discretize-bboxes"
 
     def invoke(

src/layout_prompter/transforms/label_dict_sort.py

@@ -1,14 +1,15 @@
 import copy
 from typing import Any, List, Tuple, Union, cast
 
-from langchain_core.runnables import Runnable
 from langchain_core.runnables.config import RunnableConfig
 from loguru import logger
 
 from layout_prompter.models import LayoutData, NormalizedBbox, ProcessedLayoutData
 
+from .base import LayoutTransform
 
-class LabelDictSort(Runnable):
+
+class LabelDictSort(LayoutTransform):
     name: str = "label-dict-sort"
 
     def invoke(

src/layout_prompter/transforms/lexicographic_sort.py

@@ -1,14 +1,15 @@
 import copy
 from typing import Any, List, Tuple, Union, cast
 
-from langchain_core.runnables import Runnable
 from langchain_core.runnables.config import RunnableConfig
 from loguru import logger
 
 from layout_prompter.models import LayoutData, NormalizedBbox, ProcessedLayoutData
 
+from .base import LayoutTransform
 
-class LexicographicSort(Runnable):
+
+class LexicographicSort(LayoutTransform):
     name: str = "lexicographic-sort"
 
     def invoke(

src/layout_prompter/transforms/saliency_map_to_bboxes.py

@@ -2,12 +2,13 @@
 
 import cv2
 import numpy as np
-from langchain_core.runnables import RunnableSerializable
 
 from layout_prompter.typehints import PilImage
 
+from .base import LayoutTransform
 
-class SaliencyMapToBboxes(RunnableSerializable):
+
+class SaliencyMapToBboxes(LayoutTransform):
     name: str = "saliency-map-to-bboxes"
 
     threshold: int = 100

src/layout_prompter/transforms/shuffle_elements.py

@@ -1,13 +1,12 @@
 # import copy
 # from typing import Any, Union
 
-from langchain_core.runnables import RunnableSerializable
 
 # from langchain_core.runnables.config import RunnableConfig
 # from loguru import logger
-
 # from layout_prompter.models import CanvasSize, LayoutData, ProcessedLayoutData
+from .base import LayoutTransform
 
 
-class ShuffleElements(RunnableSerializable):
+class ShuffleElements(LayoutTransform):
     name: str = "shuffle-elements"

src/layout_prompter/visualizers/base.py

@@ -107,5 +107,17 @@ def invoke(
         input: Union[ProcessedLayoutData, LayoutSerializedOutputData],
         config: Optional[RunnableConfig] = None,
         **kwargs: Any,
-    ) -> Any:
+    ) -> PilImage:
         raise NotImplementedError
+
+    def batch(
+        self,
+        inputs: Union[List[ProcessedLayoutData], List[LayoutSerializedOutputData]],
+        config: Optional[Union[RunnableConfig, List[RunnableConfig]]] = None,
+        *,
+        return_exceptions: bool = False,
+        **kwargs: Any | None,
+    ) -> List[PilImage]:
+        return super().batch(
+            inputs, config, return_exceptions=return_exceptions, **kwargs
+        )

tests/layout_prompter_test.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Type, cast
+from typing import Dict, List, Type
 
 import pytest
 from langchain.chat_models import init_chat_model
@@ -12,7 +12,6 @@
     LayoutSerializedOutputData,
     PosterLayoutSerializedData,
     PosterLayoutSerializedOutputData,
-    ProcessedLayoutData,
     Rico25SerializedData,
     Rico25SerializedOutputData,
 )
@@ -24,7 +23,6 @@
 from layout_prompter.preprocessors import ContentAwareProcessor
 from layout_prompter.settings import PosterLayoutSettings, Rico25Settings, TaskSettings
 from layout_prompter.transforms import DiscretizeBboxes
-from layout_prompter.typehints import PilImage
 from layout_prompter.utils import get_num_workers
 from layout_prompter.utils.testing import LayoutPrompterTestCase
 from layout_prompter.visualizers import ContentAwareVisualizer
@@ -69,24 +67,7 @@ def test_gen_type_task(
         input_schema: Type[LayoutSerializedData],
         output_schema: Type[LayoutSerializedOutputData],
     ):
-        # tng_dataset = layout_dataset["train"]
-        # val_dataset = layout_dataset["validation"]
-        # tst_dataset = layout_dataset["test"]
-
-        # processor = GenTypeProcessor()
-
-        # examples = cast(
-        #     List[ProcessedLayoutData],
-        #     processor.batch(
-        #         inputs=tng_dataset,
-        #         # config={
-        #         #     "max_concurrency": 4,
-        #         # },
-        #     ),
-        # )
-
-        # breakpoint()
-        pass
+        raise NotImplementedError
 
     @pytest.mark.parametrize(
         argnames=("layout_dataset", "settings", "input_schema", "output_schema"),
@@ -119,15 +100,12 @@ def test_content_aware_generation(
         processor = ContentAwareProcessor()
 
         # Process the training dataset to get candidate examples
-        candidate_examples = cast(
-            List[ProcessedLayoutData],
-            processor.batch(
-                inputs=tng_dataset,
-                config={
-                    "max_concurrency": get_num_workers(max_concurrency=4),
-                    "callbacks": [ProgressBarCallback(total=len(tng_dataset))],
-                },
-            ),
+        candidate_examples = processor.batch(
+            inputs=tng_dataset,
+            config={
+                "max_concurrency": get_num_workers(max_concurrency=4),
+                "callbacks": [ProgressBarCallback(total=len(tng_dataset))],
+            },
         )
 
         # Select a random test example
@@ -142,13 +120,11 @@ def test_content_aware_generation(
         bbox_discretizer = DiscretizeBboxes()
 
         # Apply the bbox discretizer to candidate examples and test data
-        candidate_examples = cast(
-            List[ProcessedLayoutData],
-            bbox_discretizer.batch(
-                candidate_examples,
-                config={"configurable": {"target_canvas_size": target_canvas_size}},
-            ),
+        candidate_examples = bbox_discretizer.batch(
+            candidate_examples,
+            config={"configurable": {"target_canvas_size": target_canvas_size}},
         )
+
         processed_test_data = bbox_discretizer.invoke(
             processed_test_data,
             config={"configurable": {"target_canvas_size": target_canvas_size}},
@@ -187,18 +163,16 @@ def test_content_aware_generation(
             canvas_size=settings.canvas_size,
             labels=settings.labels,
         )
-        visualizations = cast(
-            List[PilImage],
-            visualizer.batch(
-                inputs=output.ranked_outputs,
-                config={
-                    "configurable": {
-                        "resize_ratio": 2.0,
-                        "bg_image": test_data.content_image.copy(),
-                        "content_bboxes": processed_test_data.discrete_content_bboxes,
-                    }
-                },
-            ),
+        # Perform the visualization
+        visualizations = visualizer.batch(
+            inputs=output.ranked_outputs,
+            config={
+                "configurable": {
+                    "resize_ratio": 2.0,
+                    "bg_image": test_data.content_image.copy(),
+                    "content_bboxes": processed_test_data.discrete_content_bboxes,
+                }
+            },
         )
 
         # Create the save directory