Add Interleaved filters

nemo_curator/models/clip.py

@@ -106,6 +106,24 @@ def __call__(self, images: torch.Tensor | npt.NDArray[np.uint8] | list[np.ndarra
         # Normalize embeddings
         return embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True)  # type: ignore[no-any-return]
 
+    @torch.no_grad()
+    def encode_text(self, texts: list[str]) -> torch.Tensor:
+        """Encode text(s) to normalized CLIP text embeddings.
+
+        Args:
+            texts: List of strings to encode.
+
+        Returns:
+            Normalized text embeddings, shape (len(texts), dim).
+        """
+        if not texts:
+            msg = "encode_text requires at least one text"
+            raise ValueError(msg)
+        inputs = self.processor(text=texts, return_tensors="pt", padding=True, truncation=True)
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        embed = self.clip.get_text_features(**inputs)
+        return embed / torch.linalg.vector_norm(embed, dim=-1, keepdim=True)  # type: ignore[no-any-return]
+
     @classmethod
     def download_weights_on_node(cls, model_dir: str) -> None:
         """Download the weights for the CLIPImageEmbeddings model on the node."""

nemo_curator/stages/image/filters/__init__.py

@@ -15,4 +15,7 @@
 from .aesthetic_filter import ImageAestheticFilterStage
 from .nsfw_filter import ImageNSFWFilterStage
 
-__all__ = ["ImageAestheticFilterStage", "ImageNSFWFilterStage"]
+__all__ = [
+    "ImageAestheticFilterStage",
+    "ImageNSFWFilterStage",
+]

nemo_curator/stages/interleaved/filter/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from nemo_curator.stages.interleaved.filter.blur_filter import InterleavedBlurFilterStage
+from nemo_curator.stages.interleaved.filter.clip_score_filter import InterleavedCLIPScoreFilterStage
+from nemo_curator.stages.interleaved.filter.image_to_text_ratio_filter import (
+    InterleavedImageToTextRatioFilterStage,
+)
+from nemo_curator.stages.interleaved.filter.qrcode_filter import InterleavedQRCodeFilterStage
+
+__all__ = [
+    "InterleavedBlurFilterStage",
+    "InterleavedCLIPScoreFilterStage",
+    "InterleavedImageToTextRatioFilterStage",
+    "InterleavedQRCodeFilterStage",
+]

nemo_curator/stages/interleaved/filter/blur_filter.py

@@ -0,0 +1,119 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from multiprocessing import Process, Queue
+
+import cv2
+import numpy as np
+import pandas as pd
+
+from nemo_curator.stages.interleaved.stages import BaseInterleavedFilterStage
+from nemo_curator.tasks import InterleavedBatch
+
+
+def _sharpness_score(image: np.ndarray) -> float:
+    """Compute Laplacian variance as sharpness score; higher is sharper."""
+    return float(cv2.Laplacian(image, cv2.CV_64F).var())
+
+
+def _image_bytes_to_array(image_bytes: bytes) -> np.ndarray | None:
+    """Decode image bytes to RGB numpy array for OpenCV, or None on failure."""
+    try:
+        arr = np.frombuffer(image_bytes, dtype=np.uint8)
+        image = cv2.imdecode(arr, cv2.IMREAD_COLOR)
+    except Exception:  # noqa: BLE001
+        return None
+    if image is None:
+        return None
+    return cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+
+def _process_one_blur(item: tuple, score_threshold: float) -> tuple:
+    """Process one (idx, image_bytes); return (idx, keep)."""
+    idx, image_bytes = item
+    if image_bytes is None:
+        return (idx, False)
+    image = _image_bytes_to_array(image_bytes)
+    if image is None:
+        return (idx, False)
+    return (idx, _sharpness_score(image) >= score_threshold)
+
+
+def _worker_blur(
+    work_queue: Queue,
+    result_queue: Queue,
+    score_threshold: float,
+) -> None:
+    """Module-level worker for multiprocessing; runs until sentinel."""
+    while True:
+        item = work_queue.get()
+        idx, image_bytes = item
+        _, keep = _process_one_blur((idx, image_bytes), score_threshold)
+        result_queue.put((idx, keep))
+
+
+@dataclass
+class InterleavedBlurFilterStage(BaseInterleavedFilterStage):
+    """Filter interleaved image rows by sharpness (Laplacian variance); drop blurry images."""
+
+    score_threshold: float = 100.0
+    image_content_types: tuple[str, ...] = ("image/jpeg", "image/jpg", "image/png")
+    max_workers: int | None = None
+    name: str = "interleaved_blur_filter"
+
+    def _run_workers(
+        self,
+        task: InterleavedBatch,
+        df: pd.DataFrame,
+        image_mask: pd.Series,
+        keep_mask: pd.Series,
+    ) -> None:
+        """Run blur filter in worker processes; updates keep_mask in place."""
+        num_items = int(image_mask.sum())
+        work_queue = Queue()
+        result_queue = Queue()
+
+        processes = [
+            Process(
+                target=_worker_blur,
+                args=(work_queue, result_queue, self.score_threshold),
+            )
+            for _ in range(self.max_workers)
+        ]
+        for p in processes:
+            p.start()
+        for idx, image_bytes in self.iter_materialized_bytes(task=task, df=df, row_mask=image_mask):
+            work_queue.put((idx, image_bytes))
+        for _ in range(num_items):
+            i, keep = result_queue.get()
+            keep_mask.loc[i] = keep
+        for p in processes:
+            p.terminate()
+            p.join()
+
+    def content_keep_mask(self, task: InterleavedBatch, df: pd.DataFrame) -> pd.Series:
+        keep_mask = pd.Series(True, index=df.index, dtype=bool)
+        image_mask = (df["modality"] == "image") & (df["content_type"].isin(self.image_content_types))
+        if not image_mask.any():
+            return keep_mask
+        if self.max_workers is not None and self.max_workers > 1:
+            self._run_workers(task, df, image_mask, keep_mask)
+        else:
+            for idx, image_bytes in self.iter_materialized_bytes(task=task, df=df, row_mask=image_mask):
+                _, keep = _process_one_blur((idx, image_bytes), self.score_threshold)
+                keep_mask.loc[idx] = keep
+        return keep_mask

nemo_curator/stages/interleaved/filter/clip_score_filter.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import pandas as pd
+
+from nemo_curator.models.clip import CLIPImageEmbeddings
+from nemo_curator.stages.interleaved.filter.blur_filter import _image_bytes_to_array
+from nemo_curator.stages.interleaved.stages import BaseInterleavedFilterStage
+from nemo_curator.stages.resources import Resources
+from nemo_curator.tasks import InterleavedBatch
+
+
+def _sample_texts_list_from_df(df: pd.DataFrame, sample_id: str) -> list[str]:
+    """Return list of text_content from all text rows for the given sample_id (non-empty)."""
+    if "text_content" not in df.columns or "modality" not in df.columns:
+        return []
+    subset = df[(df["sample_id"] == sample_id) & (df["modality"] == "text")]
+    if subset.empty:
+        return []
+    return [s.strip() for s in subset["text_content"].dropna().astype(str).tolist() if s.strip()]
+
+
+@dataclass
+class InterleavedCLIPScoreFilterStage(BaseInterleavedFilterStage):
+    """Filter interleaved image rows by CLIP image-text relevance score.
+
+    For each image row, all text rows with the same sample_id form (image, text)
+    pairs. CLIP similarity is computed for each pair. An image is kept only if at
+    least one pair has score >= min_score; otherwise it is dropped.
+    """
+
+    model_dir: str | None = None
+    min_score: float = 0.15
+    image_content_types: tuple[str, ...] = ("image/jpeg", "image/jpg", "image/png")
+    name: str = "interleaved_clip_score_filter"
+    resources: Resources = field(default_factory=lambda: Resources(gpus=0.25))
+
+    def __post_init__(self) -> None:
+        self._model: CLIPImageEmbeddings | None = None
+
+    def _ensure_model(self) -> None:
+        if self._model is not None:
+            return
+        if self.model_dir is None:
+            msg = "InterleavedCLIPScoreFilterStage requires model_dir to be set"
+            raise RuntimeError(msg)
+        CLIPImageEmbeddings.download_weights_on_node(self.model_dir)
+        self._model = CLIPImageEmbeddings(self.model_dir)
+        self._model.setup()
+
+    def content_keep_mask(self, task: InterleavedBatch, df: pd.DataFrame) -> pd.Series:
+        keep_mask = pd.Series(True, index=df.index, dtype=bool)
+        image_mask = (df["modality"] == "image") & (df["content_type"].isin(self.image_content_types))
+        if not image_mask.any():
+            return keep_mask
+
+        self._ensure_model()
+        assert self._model is not None
+
+        for idx, image_bytes in self.iter_materialized_bytes(task=task, df=df, row_mask=image_mask):
+            if image_bytes is None:
+                keep_mask.loc[idx] = False
+                continue
+            image = _image_bytes_to_array(image_bytes)
+            if image is None:
+                keep_mask.loc[idx] = False
+                continue
+            sample_id = df.loc[idx, "sample_id"]
+            texts = _sample_texts_list_from_df(df, sample_id)
+            if not texts:
+                keep_mask.loc[idx] = False
+                continue
+            img_emb = self._model([image])
+            text_emb = self._model.encode_text(texts)
+            scores = img_emb @ text_emb.T
+            max_score = scores.max()
+            keep_mask.loc[idx] = (max_score >= self.min_score).item()
+        return keep_mask

nemo_curator/stages/interleaved/filter/image_to_text_ratio_filter.py

@@ -0,0 +1,65 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pandas as pd
+
+from nemo_curator.stages.interleaved.stages import BaseInterleavedFilterStage
+from nemo_curator.tasks import InterleavedBatch
+
+
+def _text_word_count(text: str | None) -> int:
+    """Count words in text by splitting on whitespace."""
+    if text is None or (isinstance(text, float) and pd.isna(text)):
+        return 0
+    return len(str(text).split())
+
+
+@dataclass
+class InterleavedImageToTextRatioFilterStage(BaseInterleavedFilterStage):
+    """Filter interleaved samples by image-to-text ratio (images per word).
+
+    Groups rows by sample_id. For each sample:
+    - image_count = number of rows with modality == 'image'
+    - text_word_count = sum of len(text_content.split()) over text rows
+    - ratio = image_count / max(text_word_count, 1)
+
+    Samples with ratio outside [min_ratio, max_ratio] are dropped (all their rows).
+    """
+
+    min_ratio: float = 0.0
+    max_ratio: float = float("inf")
+    name: str = "interleaved_image_to_text_ratio_filter"
+
+    def content_keep_mask(self, task: InterleavedBatch, df: pd.DataFrame) -> pd.Series:
+        keep_mask = pd.Series(True, index=df.index, dtype=bool)
+        if "sample_id" not in df.columns:
+            return keep_mask
+
+        sample_keep: dict[str, bool] = {}
+        for sample_id, group in df.groupby("sample_id"):
+            image_count = int((group["modality"] == "image").sum())
+            text_mask = group["modality"] == "text"
+            text_word_count = 0
+            if text_mask.any() and "text_content" in group.columns:
+                text_word_count = sum(_text_word_count(t) for t in group.loc[text_mask, "text_content"].tolist())
+            ratio = image_count / max(text_word_count, 1)
+            sample_keep[sample_id] = self.min_ratio <= ratio <= self.max_ratio
+
+        keep_mask = df["sample_id"].map(sample_keep)
+        keep_mask = keep_mask.fillna(True)
+        return keep_mask.astype(bool)