diff --git a/README.md b/README.md index db148d1..d561d8a 100644 --- a/README.md +++ b/README.md @@ -177,6 +177,17 @@ pytest # Run with coverage pytest --cov=hff_remover ``` +# Surya HFF detector + +- **What:** A detector that uses the Surya layout model to find headers, footers, footnotes, and text in document images. Output uses our 4 classes only: **0=text, 1=footer, 2=header, 3=footnote** (no raw Surya labels). + +- **Where:** `SuryaLayoutDetector` in `src/hff_remover/detector.py`. Same interface as other detectors: `detect()`, `detect_batch()`, `get_all_detections()`. + +- **Saving:** YOLO output is written via the existing **processor** class (`YOLOInferenceDatasetWriter`). The Surya class has `save_to_yolo()` which calls that writer (optional same-class merge before saving). + +- **Run:** `python test.py [inference_dir]` — detects and saves to `inference_data/` (images, labels/*.txt, data.yaml). Use `--merge` to merge same-class boxes into one per class. + +- **Dependencies:** `surya-ocr`, `transformers>=4.30.2,<5` (in pyproject.toml). ## License diff --git a/pyproject.toml b/pyproject.toml index 86d03b0..73b8253 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +26,8 @@ dependencies = [ "huggingface-hub>=0.20.0", "paddleocr==2.9.0", "paddlepaddle==3.0.0", + "surya-ocr>=0.17.0", + "transformers>=4.30.2,<5", ] [project.optional-dependencies] diff --git a/src/hff_remover/detector.py b/src/hff_remover/detector.py index d2ad3c7..87a59d2 100644 --- a/src/hff_remover/detector.py +++ b/src/hff_remover/detector.py @@ -606,6 +606,231 @@ def get_all_detections( return detections +# Surya layout detector (our labels: 0=text, 1=footer, 2=header, 3=footnote) +SURYA_HFF_CLASS_IDS = { + 0: "text", + 1: "footer", + 2: "header", + 3: "footnote", +} +SURYA_LABEL_TO_OUR_CLASS = { + "text": (0, "text"), + "caption": (0, "text"), + "list-item": (0, "text"), + "formula": (0, "text"), + "text-inline-math": (0, "text"), + "handwriting": (0, "text"), + "table-of-contents": (0, "text"), + "form": (0, "text"), + "table": (0, "text"), + "picture": (0, "text"), + "figure": (0, "text"), + "section-header": (2, "header"), + "page-header": (2, "header"), + "page-footer": (1, "footer"), + "footnote": (3, "footnote"), +} + + +def _surya_label_to_our_class(surya_label: str, default_to_text: bool = False) -> Optional[tuple]: + key = surya_label.strip().lower().replace(" ", "-") if surya_label else "" + out = SURYA_LABEL_TO_OUR_CLASS.get(key) + if out is not None: + return out + return (0, "text") if default_to_text else None + + +class SuryaLayoutDetector(BaseHFFDetector): + def __init__( + self, + confidence_threshold: float = 0.5, + layout_checkpoint: Optional[str] = None, + ): + from surya.layout import LayoutPredictor + from surya.foundation import FoundationPredictor + from surya.settings import settings + + self.confidence_threshold = confidence_threshold + checkpoint = layout_checkpoint or getattr(settings, "LAYOUT_MODEL_CHECKPOINT", None) + foundation = FoundationPredictor() if checkpoint is None else FoundationPredictor(checkpoint=checkpoint) + self.layout_predictor = LayoutPredictor(foundation) + + def _load_image(self, image: Union[str, Path, np.ndarray]) -> Any: + from PIL import Image + + if isinstance(image, (str, Path)): + return Image.open(str(image)).convert("RGB") + if isinstance(image, np.ndarray): + if image.ndim == 2: + return Image.fromarray(image).convert("RGB") + if image.shape[2] == 3: + return Image.fromarray(image[:, :, ::-1]) + return Image.fromarray(image).convert("RGB") + return image + + def _layout_result_to_detections( + self, + layout_pred: Any, + apply_confidence_filter: bool = True, + filter_to_hff_only: bool = True, + ) -> List[Dict[str, Any]]: + detections: List[Dict[str, Any]] = [] + if hasattr(layout_pred, "bboxes"): + bboxes_data = layout_pred.bboxes or [] + elif isinstance(layout_pred, dict): + bboxes_data = layout_pred.get("bboxes") or layout_pred.get("boxes") or [] + else: + bboxes_data = [] + + for item in bboxes_data: + raw_label = (item.get("label") or item.get("type") or "") if isinstance(item, dict) else getattr(item, "label", "") + bbox = item.get("bbox") if isinstance(item, dict) else getattr(item, "bbox", None) + top_k = item.get("top_k") if isinstance(item, dict) else getattr(item, "top_k", None) or {} + conf = 1.0 + if isinstance(item, dict) and "confidence" in item: + conf = float(item["confidence"]) + elif hasattr(item, "confidence") and item.confidence is not None: + conf = float(item.confidence) + elif isinstance(top_k, dict) and top_k: + key = str(raw_label).strip().lower().replace(" ", "-") if raw_label else "" + conf = float(top_k.get(raw_label) or top_k.get(key) or list(top_k.values())[0]) + + if apply_confidence_filter and conf < self.confidence_threshold: + continue + + raw_label = str(raw_label or "") + mapped = _surya_label_to_our_class(raw_label, default_to_text=not filter_to_hff_only) + if mapped is None: + continue + our_class_id, our_class_name = mapped + + if not bbox or len(bbox) < 4: + continue + bbox = list(bbox) + if len(bbox) == 4: + bbox = list(map(float, bbox)) + elif len(bbox) == 8: + xs, ys = bbox[0::2], bbox[1::2] + bbox = [min(xs), min(ys), max(xs), max(ys)] + else: + continue + + detections.append({ + "bbox": bbox, + "class_id": our_class_id, + "class_name": our_class_name, + "confidence": conf, + }) + + return detections + + def detect( + self, + image: Union[str, Path, np.ndarray], + image_size: int = 1024, + ) -> List[Dict[str, Any]]: + pil_image = self._load_image(image) + layout_predictions = self.layout_predictor([pil_image]) + if not layout_predictions: + return [] + return self._layout_result_to_detections( + layout_predictions[0], + apply_confidence_filter=True, + filter_to_hff_only=True, + ) + + def detect_batch( + self, + images: List[Union[str, Path, np.ndarray]], + image_size: int = 1024, + batch_size: int = 8, + ) -> List[List[Dict[str, Any]]]: + pil_list = [self._load_image(im) for im in images] + layout_predictions = self.layout_predictor(pil_list) + + all_detections: List[List[Dict[str, Any]]] = [] + for pred in layout_predictions: + all_detections.append( + self._layout_result_to_detections( + pred, + apply_confidence_filter=True, + filter_to_hff_only=True, + ) + ) + return all_detections + + def get_all_detections( + self, + image: Union[str, Path, np.ndarray], + image_size: int = 1024, + ) -> List[Dict[str, Any]]: + _ = image_size + pil_image = self._load_image(image) + layout_predictions = self.layout_predictor([pil_image]) + if not layout_predictions: + return [] + return self._layout_result_to_detections( + layout_predictions[0], + apply_confidence_filter=True, + filter_to_hff_only=False, + ) + + def save_to_yolo( + self, + image: Union[str, Path, np.ndarray], + image_rel_path: Union[str, Path], + inference_dir: Union[str, Path] = "inference_data", + detections: Optional[List[Dict[str, Any]]] = None, + merge_same_class: bool = True, + ) -> tuple: + """ + Save detections in YOLO format using processor.YOLOInferenceDatasetWriter. + Same output as processor: inference_dir/images/, labels/*.txt, data.yaml. + Returns (image_out_path, label_out_path). + """ + from hff_remover.processor import YOLOInferenceDatasetWriter + from hff_remover.utils import load_image + + if detections is None: + detections = self.detect(image) + if isinstance(image, (str, Path)): + image = load_image(image) + if merge_same_class: + by_class: Dict[int, List[Dict[str, Any]]] = {} + for d in detections: + cid = d.get("class_id") + if cid is None: + continue + by_class.setdefault(cid, []).append(d) + merged = [] + for cid, group in by_class.items(): + boxes = [b.get("bbox") for b in group if b.get("bbox") and len(b.get("bbox", [])) == 4] + if not boxes: + continue + x1 = min(b[0] for b in boxes) + y1 = min(b[1] for b in boxes) + x2 = max(b[2] for b in boxes) + y2 = max(b[3] for b in boxes) + best = max(group, key=lambda b: b.get("confidence", 0)) + merged.append({ + "bbox": [x1, y1, x2, y2], + "class_id": cid, + "class_name": best.get("class_name", ""), + "confidence": best.get("confidence", 0), + }) + detections = merged + class_name_to_id = {name: idx for idx, name in SURYA_HFF_CLASS_IDS.items()} + writer = YOLOInferenceDatasetWriter( + base_dir=inference_dir, + class_name_to_id=class_name_to_id, + ) + return writer.write_sample( + image=image, + detections=detections, + image_rel_path=image_rel_path, + ) + + # ============================================================================= # Ensemble Detector # ============================================================================= diff --git a/tests/test_detector.py b/tests/test_detector.py index 0871c73..feee596 100644 --- a/tests/test_detector.py +++ b/tests/test_detector.py @@ -5,7 +5,14 @@ import numpy as np import pytest -from hff_remover.detector import HFFDetector, HFF_CLASSES, CLASS_NAMES +from hff_remover.detector import ( + HFFDetector, + HFF_CLASSES, + CLASS_NAMES, + SuryaLayoutDetector, + SURYA_HFF_CLASS_IDS, + _surya_label_to_our_class, +) class TestHFFClasses: @@ -203,3 +210,179 @@ def test_confidence_threshold(self, mock_yolo, mock_download): # Check that predict was called with the confidence threshold call_kwargs = mock_model.predict.call_args[1] assert call_kwargs["conf"] == 0.7 + + +# ============================================================================= +# Surya Layout Detector tests (our labels only; no YOLO/save_results in class) +# ============================================================================= + +class TestSuryaHFFClassMapping: + """Tests for Surya HFF class constants and mapping (no surya import needed).""" + + def test_surya_hff_class_ids(self): + """Our class_id mapping: 0=text, 1=footer, 2=header, 3=footnote.""" + assert SURYA_HFF_CLASS_IDS[0] == "text" + assert SURYA_HFF_CLASS_IDS[1] == "footer" + assert SURYA_HFF_CLASS_IDS[2] == "header" + assert SURYA_HFF_CLASS_IDS[3] == "footnote" + assert len(SURYA_HFF_CLASS_IDS) == 4 + + def test_surya_label_to_our_class_mapping(self): + """Normalized Surya labels map to our (class_id, class_name).""" + assert _surya_label_to_our_class("Text") == (0, "text") + assert _surya_label_to_our_class("Page-footer") == (1, "footer") + assert _surya_label_to_our_class("Page-header") == (2, "header") + assert _surya_label_to_our_class("Footnote") == (3, "footnote") + assert _surya_label_to_our_class("Section-header") == (2, "header") + + def test_surya_label_unknown_returns_none_when_not_default(self): + """Unknown Surya labels return None when default_to_text=False.""" + assert _surya_label_to_our_class("Table") is None + assert _surya_label_to_our_class("Picture") is None + + def test_surya_label_unknown_maps_to_text_when_default(self): + """Unknown Surya labels map to (0, 'text') when default_to_text=True.""" + assert _surya_label_to_our_class("Table", default_to_text=True) == (0, "text") + + +class TestSuryaLayoutDetector: + """Tests for SuryaLayoutDetector: our class_id/class_name only, no raw Surya labels.""" + + @pytest.fixture(autouse=True) + def skip_if_no_surya(self): + pytest.importorskip("surya") + + @patch("surya.layout.LayoutPredictor") + @patch("surya.foundation.FoundationPredictor") + def test_detect_returns_only_our_classes(self, mock_foundation, mock_layout_cls): + """detect() returns only our HFF/text classes with our class_id and class_name.""" + mock_layout = MagicMock() + mock_layout.return_value = [ + { + "bboxes": [ + {"bbox": [10, 10, 100, 30], "label": "Page-header", "top_k": {"Page-header": 0.9}}, + {"bbox": [10, 500, 100, 530], "label": "Page-footer", "top_k": {"Page-footer": 0.85}}, + {"bbox": [10, 100, 400, 400], "label": "Text", "top_k": {"Text": 0.95}}, + ] + } + ] + mock_layout_cls.return_value = mock_layout + mock_foundation.return_value = MagicMock() + + detector = SuryaLayoutDetector(confidence_threshold=0.3) + detections = detector.detect(np.zeros((600, 500, 3), dtype=np.uint8)) + + assert len(detections) == 3 + class_ids = {d["class_id"] for d in detections} + class_names = {d["class_name"] for d in detections} + assert class_ids <= {0, 1, 2, 3} + assert class_names <= {"text", "footer", "header", "footnote"} + for d in detections: + assert "bbox" in d and len(d["bbox"]) == 4 + assert "confidence" in d + assert isinstance(d["class_id"], int) + assert d["class_name"] in ("text", "footer", "header", "footnote") + + @patch("surya.layout.LayoutPredictor") + @patch("surya.foundation.FoundationPredictor") + def test_detect_filters_by_confidence(self, mock_foundation, mock_layout_cls): + """detect() filters out detections below confidence threshold.""" + mock_layout = MagicMock() + mock_layout.return_value = [ + { + "bboxes": [ + {"bbox": [10, 10, 100, 30], "label": "Page-header", "top_k": {"Page-header": 0.9}}, + {"bbox": [10, 500, 100, 530], "label": "Footnote", "top_k": {"Footnote": 0.2}}, + ] + } + ] + mock_layout_cls.return_value = mock_layout + mock_foundation.return_value = MagicMock() + + detector = SuryaLayoutDetector(confidence_threshold=0.5) + detections = detector.detect(np.zeros((600, 500, 3), dtype=np.uint8)) + + assert len(detections) == 1 + assert detections[0]["class_name"] == "header" + + @patch("surya.layout.LayoutPredictor") + @patch("surya.foundation.FoundationPredictor") + def test_detect_no_raw_surya_labels(self, mock_foundation, mock_layout_cls): + """Output never contains raw Surya label names (e.g. Page-header, Text).""" + mock_layout = MagicMock() + mock_layout.return_value = [ + { + "bboxes": [ + {"bbox": [10, 10, 100, 30], "label": "Page-header", "top_k": {"Page-header": 0.9}}, + ] + } + ] + mock_layout_cls.return_value = mock_layout + mock_foundation.return_value = MagicMock() + + detector = SuryaLayoutDetector(confidence_threshold=0.3) + detections = detector.detect(np.zeros((100, 100, 3), dtype=np.uint8)) + + raw_names = {"Page-header", "Page-footer", "Footnote", "Text", "Section-header", "Caption"} + for d in detections: + assert d["class_name"] not in raw_names + assert d["class_name"] in ("text", "footer", "header", "footnote") + + @patch("surya.layout.LayoutPredictor") + @patch("surya.foundation.FoundationPredictor") + def test_detect_batch_same_format(self, mock_foundation, mock_layout_cls): + """detect_batch() returns list of lists with same dict format (our class_id/class_name).""" + mock_layout = MagicMock() + mock_layout.return_value = [ + {"bboxes": [{"bbox": [0, 0, 50, 20], "label": "Text", "top_k": {"Text": 0.8}}]}, + {"bboxes": [{"bbox": [0, 0, 50, 20], "label": "Footnote", "top_k": {"Footnote": 0.7}}]}, + ] + mock_layout_cls.return_value = mock_layout + mock_foundation.return_value = MagicMock() + + detector = SuryaLayoutDetector(confidence_threshold=0.3) + results = detector.detect_batch( + [np.zeros((100, 100, 3), dtype=np.uint8), np.zeros((100, 100, 3), dtype=np.uint8)] + ) + + assert len(results) == 2 + assert len(results[0]) == 1 and results[0][0]["class_id"] == 0 and results[0][0]["class_name"] == "text" + assert len(results[1]) == 1 and results[1][0]["class_id"] == 3 and results[1][0]["class_name"] == "footnote" + + @patch("surya.layout.LayoutPredictor") + @patch("surya.foundation.FoundationPredictor") + def test_get_all_detections_same_format(self, mock_foundation, mock_layout_cls): + """get_all_detections() returns our class_id and class_name only.""" + mock_layout = MagicMock() + mock_layout.return_value = [ + { + "bboxes": [ + {"bbox": [0, 0, 50, 20], "label": "Text", "top_k": {"Text": 0.8}}, + {"bbox": [0, 80, 50, 100], "label": "Table", "top_k": {"Table": 0.9}}, + ] + } + ] + mock_layout_cls.return_value = mock_layout + mock_foundation.return_value = MagicMock() + + detector = SuryaLayoutDetector(confidence_threshold=0.3) + detections = detector.get_all_detections(np.zeros((100, 100, 3), dtype=np.uint8)) + + # Table maps to (0, "text") when filter_to_hff_only=False (default_to_text=True) + assert len(detections) == 2 + for d in detections: + assert d["class_id"] in (0, 1, 2, 3) + assert d["class_name"] in ("text", "footer", "header", "footnote") + + @patch("surya.layout.LayoutPredictor") + @patch("surya.foundation.FoundationPredictor") + def test_detect_empty_bboxes(self, mock_foundation, mock_layout_cls): + """detect() returns empty list when layout returns no bboxes.""" + mock_layout = MagicMock() + mock_layout.return_value = [{"bboxes": []}] + mock_layout_cls.return_value = mock_layout + mock_foundation.return_value = MagicMock() + + detector = SuryaLayoutDetector() + detections = detector.detect(np.zeros((100, 100, 3), dtype=np.uint8)) + assert detections == []