feat: Expose classification filters for picture description (#105)

Preserve legacy picture description filters

Signed-off-by: drk <drukpa1455@gmail.com>

Author: drukpa1455

docling_jobkit/datamodel/convert.py

@@ -39,7 +39,7 @@
     DEFAULT_PAGE_RANGE,
     PageRange,
 )
-from docling_core.types.doc import ImageRefMode
+from docling_core.types.doc import ImageRefMode, PictureClassificationLabel
 
 
 class PictureDescriptionLocal(BaseModel):
@@ -70,6 +70,26 @@ class PictureDescriptionLocal(BaseModel):
             examples=[{"max_new_tokens": 200, "do_sample": False}],
         ),
     ] = {"max_new_tokens": 200, "do_sample": False}
+    classification_allow: Annotated[
+        Optional[list[PictureClassificationLabel]],
+        Field(
+            description="Only describe pictures whose predicted class is in this allow-list."
+        ),
+    ] = None
+    classification_deny: Annotated[
+        Optional[list[PictureClassificationLabel]],
+        Field(
+            description="Do not describe pictures whose predicted class is in this deny-list."
+        ),
+    ] = None
+    classification_min_confidence: Annotated[
+        float,
+        Field(
+            description="Minimum classification confidence required before a picture can be described.",
+            ge=0.0,
+            le=1.0,
+        ),
+    ] = 0.0
 
 
 class PictureDescriptionApi(BaseModel):
@@ -133,6 +153,26 @@ class PictureDescriptionApi(BaseModel):
             ],
         ),
     ] = "Describe this image in a few sentences."
+    classification_allow: Annotated[
+        Optional[list[PictureClassificationLabel]],
+        Field(
+            description="Only describe pictures whose predicted class is in this allow-list."
+        ),
+    ] = None
+    classification_deny: Annotated[
+        Optional[list[PictureClassificationLabel]],
+        Field(
+            description="Do not describe pictures whose predicted class is in this deny-list."
+        ),
+    ] = None
+    classification_min_confidence: Annotated[
+        float,
+        Field(
+            description="Minimum classification confidence required before a picture can be described.",
+            ge=0.0,
+            le=1.0,
+        ),
+    ] = 0.0
 
 
 class VlmModelLocal(BaseModel):

tests/test_options.py

@@ -13,7 +13,7 @@
     VlmPipelineOptions,
 )
 from docling.pipeline.vlm_pipeline import VlmPipeline
-from docling_core.types.doc import ImageRefMode
+from docling_core.types.doc import ImageRefMode, PictureClassificationLabel
 
 from docling_jobkit.convert.manager import (
     DoclingConverterManager,
@@ -24,6 +24,7 @@
 from docling_jobkit.datamodel.convert import (
     ConvertDocumentsOptions,
     PictureDescriptionApi,
+    PictureDescriptionLocal,
 )
 
 
@@ -281,6 +282,48 @@ def test_options_cache_key_with_presets():
     hashes.add(hash)
 
 
+def test_legacy_picture_description_filters_reach_pipeline_options():
+    m = DoclingConverterManager(config=DoclingConverterManagerConfig())
+
+    with pytest.warns(DeprecationWarning):
+        opts = ConvertDocumentsOptions(
+            do_picture_description=True,
+            picture_description_local=PictureDescriptionLocal(
+                repo_id="HuggingFaceTB/SmolVLM-256M-Instruct",
+                classification_allow=[PictureClassificationLabel.BAR_CHART],
+                classification_deny=[PictureClassificationLabel.LOGO],
+                classification_min_confidence=0.55,
+            ),
+        )
+
+    pipeline_opts = m.get_pdf_pipeline_opts(opts)
+    picture_opts = pipeline_opts.pipeline_options.picture_description_options
+
+    assert picture_opts.classification_allow == [PictureClassificationLabel.BAR_CHART]
+    assert picture_opts.classification_deny == [PictureClassificationLabel.LOGO]
+    assert picture_opts.classification_min_confidence == 0.55
+
+
+def test_legacy_picture_description_api_filters_reach_pipeline_options():
+    m = DoclingConverterManager(config=DoclingConverterManagerConfig())
+
+    with pytest.warns(DeprecationWarning):
+        opts = ConvertDocumentsOptions(
+            do_picture_description=True,
+            picture_description_api={
+                "url": "http://localhost",
+                "classification_allow": ["map"],
+                "classification_min_confidence": 0.8,
+            },
+        )
+
+    pipeline_opts = m.get_pdf_pipeline_opts(opts)
+    picture_opts = pipeline_opts.pipeline_options.picture_description_options
+
+    assert [label.value for label in picture_opts.classification_allow] == ["map"]
+    assert picture_opts.classification_min_confidence == 0.8
+
+
 def test_image_pipeline_uses_vlm_pipeline_when_requested():
     m = DoclingConverterManager(config=DoclingConverterManagerConfig())
     opts = ConvertDocumentsOptions(pipeline=ProcessingPipeline.VLM)