teklia-layout-reader 0.2.1__py3-none-any.whl

layout_reader/__init__.py

@@ -0,0 +1,9 @@
+import logging
+import random
+
+logging.basicConfig(
+    format="%(asctime)s %(levelname)s/%(name)s: %(message)s",
+    level=logging.INFO,
+)
+
+random.seed(42)

layout_reader/cli.py

@@ -0,0 +1,27 @@
+import argparse
+
+from layout_reader.datasets import add_dataset_command
+from layout_reader.inference import add_predict_command
+from layout_reader.train.sft import add_training_command
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="teklia-layout-reader",
+        description="Scripts for Layout Reader",
+    )
+
+    commands = parser.add_subparsers(metavar="subcommands")
+    add_dataset_command(commands)
+    add_training_command(commands)
+    add_predict_command(commands)
+
+    args = vars(parser.parse_args())
+    if "func" in args:
+        args.pop("func")(**args)
+    else:
+        parser.print_help()
+
+
+if __name__ == "__main__":
+    main()

layout_reader/datasets/__init__.py

@@ -0,0 +1,99 @@
+"""Manage datasets."""
+
+import argparse
+from pathlib import Path
+
+from layout_reader.datasets.analyze import run_dataset_analysis
+from layout_reader.datasets.extract import Mode, run_dataset_extraction
+
+
+def add_dataset_command(commands) -> None:
+    export = commands.add_parser(
+        "dataset",
+        description=__doc__,
+        help=__doc__,
+    )
+    subcommands = export.add_subparsers(metavar="subcommands")
+    add_extract_parser(subcommands)
+    add_analyze_parser(subcommands)
+
+
+def valid_rate(rate: float) -> float:
+    """Convert to a float between 0 and 1."""
+    try:
+        rate = float(rate)
+    except ValueError as e:
+        raise argparse.ArgumentTypeError(
+            f"`The shuffle rate ({rate}) should be a float."
+        ) from e
+
+    if rate > 1 or rate < 0:
+        raise argparse.ArgumentTypeError(
+            f"`The shuffle rate ({rate}) should be between 0 and 1."
+        )
+    return rate
+
+
+def add_extract_parser(subcommands) -> None:
+    parser = subcommands.add_parser(
+        "extract",
+        description="Convert a HF dataset into LayoutReader format.",
+        help="Convert a HF dataset into LayoutReader format.",
+    )
+
+    parser.add_argument(
+        "dataset",
+        type=str,
+        help="Name of the HuggingFace or YOLO dataset",
+    )
+    parser.add_argument(
+        "--output-dir",
+        type=Path,
+        help="Output directory.",
+        required=True,
+    )
+    parser.add_argument(
+        "--mode",
+        type=Mode,
+        help=f"Dataset type. Must be in {[e.value for e in Mode]}",
+        required=True,
+    )
+    # Optional
+    parser.add_argument(
+        "--shuffle-rate",
+        type=valid_rate,
+        help="Ratio of the data that will be shuffled (expected between 0 and 1).",
+        default=0.5,
+    )
+    parser.add_argument(
+        "--extract-classes",
+        action="store_true",
+        help="Whether to extract classes.",
+    )
+    parser.add_argument(
+        "--extract-separators",
+        action="store_true",
+        help="Whether to extract separators.",
+    )
+    parser.set_defaults(func=run_dataset_extraction)
+
+
+def add_analyze_parser(subcommands) -> None:
+    parser = subcommands.add_parser(
+        "analyze",
+        description="Analyze a LayoutReader dataset.",
+        help="Analyze a LayoutReader dataset.",
+    )
+
+    parser.add_argument(
+        "dataset_path",
+        type=Path,
+        help="Path of the LayoutReader dataset.",
+    )
+    parser.add_argument(
+        "--report-path",
+        type=Path,
+        help="Path to save the Markdown report.",
+        default="dataset_report.md",
+    )
+    parser.set_defaults(func=run_dataset_analysis)

layout_reader/datasets/analyze.py

@@ -0,0 +1,161 @@
+"""
+Analyze a local LayoutReader dataset.
+"""
+
+import logging
+from collections import Counter, defaultdict
+from functools import partial
+from pathlib import Path
+
+import numpy as np
+from mdutils.mdutils import MdUtils
+from prettytable import PrettyTable, TableStyle
+
+from layout_reader.datasets.utils import load_gzip_jsonl
+
+logger = logging.getLogger(__name__)
+
+METRIC_COLUMN = "Metric"
+
+
+def create_table(
+    data: dict,
+    count: bool = False,
+    total: bool = True,
+) -> PrettyTable:
+    """
+    Each keys will be made into a column.
+    We compute min, max, mean, median, total by default.
+    Total can be disabled. Count (length) computation can be enabled.
+    """
+
+    statistics = PrettyTable(field_names=[METRIC_COLUMN, *data.keys()])
+    statistics.align.update({METRIC_COLUMN: "l"})
+    statistics.set_style(TableStyle.MARKDOWN)
+
+    operations = []
+
+    if count:
+        operations.append(("Count", len, None))
+
+    operations.extend(
+        [
+            ("Min", np.min, None),
+            ("Max", np.max, None),
+            ("Mean", np.mean, 2),
+            ("Median", np.median, 2),
+        ]
+    )
+    if total:
+        operations.append(("Total", np.sum, None))
+
+    statistics.add_rows(
+        [
+            [
+                col_name,
+                *list(
+                    map(
+                        # Round values if needed
+                        partial(round, ndigits=digits),
+                        map(operator, data.values()),
+                    )
+                ),
+            ]
+            for col_name, operator, digits in operations
+        ]
+    )
+
+    return statistics
+
+
+class Statistics:
+    HEADERS = {
+        "Coordinates": "Coordinates statistics",
+        "Classes": "Classes statistics",
+        "Separators": "Separators statistics",
+    }
+
+    def __init__(self, output: Path) -> None:
+        self.document = MdUtils(file_name=str(output), title="Statistics")
+
+    def _write_section(self, table: PrettyTable, title: str, level: int = 2) -> None:
+        """
+        Write the new section in the file.
+
+        <title with appropriate level>
+
+        <table>
+
+        """
+        self.document.new_header(level=level, title=title, add_table_of_contents="n")
+        self.document.write("\n")
+
+        logger.info(f"{title}\n\n{table}\n")
+
+        self.document.write(table.get_string())
+        self.document.write("\n")
+
+    def create_classes_statistics(self, labels: list[str], title: str) -> None:
+        """
+        Compute statistics on class labels and write them to file.
+        """
+        class_counter = Counter()
+
+        for classes in labels:
+            class_counter.update(classes)
+
+        statistics = PrettyTable(
+            field_names=[METRIC_COLUMN] + [f"Class ID {k}" for k in class_counter]
+        )
+        statistics.set_style(TableStyle.MARKDOWN)
+
+        statistics.add_row(["Count"] + [v for v in class_counter.values()])
+
+        self._write_section(
+            table=statistics,
+            title=title,
+        )
+
+    def create_boxes_statistics(self, labels: list[list[float]], title: str) -> None:
+        """
+        Compute statistics on bounding boxes and write them to file.
+        """
+        data = defaultdict(list)
+
+        for page in labels:
+            for box in page:
+                x1, y1, x2, y2 = box
+                data["Box width"].append(x2 - x1)
+                data["Box height"].append(y2 - y1)
+                data["Box surface (%)"].append((x2 - x1) * (y2 - y1) / 10000)
+
+        self._write_section(
+            table=create_table(data=data, total=False),
+            title=title,
+        )
+
+    def run(self, filenames: list[Path]) -> None:
+        """Iterate over each split and create report"""
+        for filename in filenames:
+            self.document.new_header(level=1, title=str(filename.stem).capitalize())
+            labels = load_gzip_jsonl(filename)
+            self.create_classes_statistics(
+                labels=[page["target_classes"] for page in labels],
+                title=Statistics.HEADERS["Classes"],
+            )
+            self.create_boxes_statistics(
+                labels=[page["target_boxes"] for page in labels],
+                title=Statistics.HEADERS["Coordinates"],
+            )
+            self.create_boxes_statistics(
+                labels=[page["separators"] for page in labels],
+                title=Statistics.HEADERS["Separators"],
+            )
+        self.document.create_md_file()
+
+
+def run_dataset_analysis(dataset_path: Path, report_path: Path) -> None:
+    """
+    Compute statistics on a dataset in LayoutReader format.
+    """
+    Statistics(output=report_path).run(filenames=dataset_path.glob("*.jsonl.gz"))

layout_reader/datasets/extract.py

@@ -0,0 +1,289 @@
+"""
+Convert a HuggingFace dataset to LayoutReader format.
+Example: https://huggingface.co/datasets/Teklia/Newspapers-finlam.
+"""
+
+import logging
+import random
+import tempfile
+from enum import Enum
+from pathlib import Path
+
+import numpy as np
+import tqdm
+from datasets import Dataset, load_dataset, load_from_disk
+
+from layout_reader.datasets.lsd import LineDetector
+from layout_reader.datasets.utils import (
+    BBX_FACTOR,
+    CLS_SHIFT,
+    UNK_TOKEN_ID,
+    check_is_valid_bbx,
+    check_too_many_zones,
+    convert_to_bbx,
+    load_yolo_line,
+    make_bbx_valid,
+    save_gzip_jsonl,
+)
+
+random.seed(42)
+np.random.seed(42)
+
+logger = logging.getLogger(__name__)
+
+
+class Mode(Enum):
+    HF = "HF"
+    YOLO = "YOLO"
+
+
+SPLIT_MAPPING = {
+    "val": "dev",
+}
+
+
+def sort_and_filter_zones(
+    orders: list,
+    boxes: list,
+    classes: list,
+) -> tuple:
+    """
+    Order all boxes and classes. Extract separators if needed.
+    """
+    filtered_boxes = []
+    filtered_classes = []
+    filtered_orders = []
+
+    # Iterate on zones
+    for order, box, classif in zip(
+        orders,
+        boxes,
+        classes,
+        strict=True,
+    ):
+        # Normalize box
+        box = make_bbx_valid(box)
+        if not check_is_valid_bbx(box):
+            logger.warning(
+                f"Ignoring box {box} ({order}) for element as bounding box is invalid"
+            )
+            continue
+        filtered_boxes.append(box)
+        filtered_classes.append(classif)
+        filtered_orders.append(order)
+
+    # Sort boxes and classes based on order
+    target_boxes = np.array(filtered_boxes)[np.argsort(orders)].tolist()
+    target_classes = np.array(filtered_classes)[np.argsort(orders)].tolist()
+    return target_boxes, target_classes
+
+
+def prepare_training_data(
+    target_boxes, target_classes, shuffle_rate: float = 0.5
+) -> tuple[list[list[int]], list[int], list[int]]:
+    """
+    Create dataset of zones from the target order.
+    """
+    # Shuffle zones or sort zones by XY
+    zones = [
+        (order, features[0], features[1])
+        for order, features in enumerate(zip(target_boxes, target_classes, strict=True))
+    ]
+    if random.random() < shuffle_rate:
+        random.shuffle(zones)
+    else:
+        zones = sorted(zones, key=lambda z: (z[1][0], z[1][1]))
+
+    # Build unsorted boxes & classes from target zones
+    source_orders, source_boxes, source_classes = map(list, zip(*zones, strict=True))
+
+    # Get sorted index and start at 1
+    target_index = (np.argsort(source_orders) + 1).tolist()
+    return source_boxes, source_classes, target_index
+
+
+def extract_split_hf(
+    subset: Dataset,
+    shuffle_rate: float,
+    extract_classes: bool = True,
+    extract_separators: bool = True,
+) -> list[dict]:
+    """
+    Extract data from a Dataset split.
+    """
+    samples = []
+    # Iterate on pages
+    for page in tqdm.tqdm(subset):
+        # Convert polygons to boxes
+        boxes = [
+            convert_to_bbx(polygon, BBX_FACTOR) for polygon in page["zone_polygons"]
+        ]
+        classes = [
+            classif + CLS_SHIFT if extract_classes else UNK_TOKEN_ID
+            for classif in page["zone_classes"]
+        ]
+        orders = page["zone_orders"]
+
+        # Filter invalid zones and sort them by order
+        target_boxes, target_classes = sort_and_filter_zones(orders, boxes, classes)
+
+        # Get separators
+        separators = []
+        if extract_separators:
+            with tempfile.NamedTemporaryFile(suffix=".jpg") as tmp_file:
+                image_path = tmp_file.name
+                page["page_image"].save(image_path)
+                detector = LineDetector()
+                separators = detector.process(image_path)
+
+        # Check that the number of zones is < MAX_LEN
+        if check_too_many_zones(boxes=target_boxes, separators=separators):
+            logger.warning("Too many zones on the current page - skipping.")
+            continue
+
+        # Create input data for Layout Reader - sort or shuffle
+        source_boxes, source_classes, target_index = prepare_training_data(
+            target_boxes, target_classes, shuffle_rate=shuffle_rate
+        )
+        samples.append(
+            {
+                "sample_id": page["page_arkindex_id"],
+                "source_boxes": source_boxes,
+                "source_classes": source_classes,
+                "target_boxes": target_boxes,
+                "target_classes": target_classes,
+                "target_index": target_index,
+                "separators": separators,
+            }
+        )
+    return samples
+
+
+def run_dataset_extraction_hf(
+    dataset: Dataset,
+    shuffle_rate: float,
+    extract_classes: bool = True,
+    extract_separators: bool = True,
+) -> Path:
+    """
+    Extract a dataset from HuggingFace or YOLO format and format it in LayoutReader format.
+    """
+    data = {}
+    for split in dataset:
+        # Extract data
+        logger.info(f"Extracting split {split}.")
+        data[split] = extract_split_hf(
+            dataset[split],
+            shuffle_rate,
+            extract_classes,
+            extract_separators,
+        )
+    return data
+
+
+def convert_yolo_labels(lines: list, extract_classes: bool):
+    """
+    Convert YOLO labels to LayoutReader features.
+    Assumes the list of zones is ordered.
+    """
+    target_boxes = []
+    target_classes = []
+    for line in lines:
+        classif, box = load_yolo_line(line, BBX_FACTOR)
+        box = make_bbx_valid(box)
+        if not check_is_valid_bbx(box):
+            print(f"Ignoring box {box}) as bounding box is invalid")
+            continue
+        target_boxes.append(box)
+        target_classes.append(classif + CLS_SHIFT if extract_classes else UNK_TOKEN_ID)
+    return target_boxes, target_classes
+
+
+def run_dataset_extraction_yolo(
+    dataset: str,
+    shuffle_rate: float,
+    extract_classes: bool = True,
+    extract_separators: bool = True,
+) -> Path:
+    data = {}
+    for subset in Path(dataset).iterdir():
+        samples = []
+        labels_path = sorted((Path(dataset) / subset.name / "labels").rglob("*.txt"))
+        images_base_path = Path(dataset) / subset.name / "images"
+
+        for label_path in tqdm.tqdm(labels_path):
+            image_path = images_base_path / label_path.with_suffix(".jpg").name
+            if not image_path.exists():
+                continue
+
+            # Remove duplicate lines (YOLO finds many of them)
+            yolo_lines = label_path.read_text().splitlines()
+            target_boxes, target_classes = convert_yolo_labels(
+                yolo_lines, extract_classes
+            )
+
+            # Extract separators
+            separators = []
+            if extract_separators:
+                detector = LineDetector()
+                separators = detector.process(image_path)
+
+            # Check that the number of zones is < MAX_LEN
+            if check_too_many_zones(boxes=target_boxes, separators=separators):
+                logger.warning("Too many zones on the current page - skipping.")
+                continue
+
+            # Create input data for Layout Reader - sort or shuffle
+            source_boxes, source_classes, target_index = prepare_training_data(
+                target_boxes, target_classes, shuffle_rate=shuffle_rate
+            )
+            samples.append(
+                {
+                    "sample_id": label_path.stem,
+                    "source_boxes": source_boxes,
+                    "source_classes": source_classes,
+                    "target_boxes": target_boxes,
+                    "target_classes": target_classes,
+                    "target_index": target_index,
+                    "separators": separators,
+                }
+            )
+        data[subset.name] = samples
+    return data
+
+
+def run_dataset_extraction(
+    dataset: str,
+    mode: Mode,
+    shuffle_rate: float,
+    output_dir: Path,
+    extract_classes: bool = True,
+    extract_separators: bool = True,
+) -> Path:
+    output_dir.mkdir(exist_ok=True)
+
+    if mode == Mode.HF:
+        if Path(dataset).exists():
+            dataset = load_from_disk(dataset)
+        else:
+            dataset = load_dataset(dataset)
+
+        data = run_dataset_extraction_hf(
+            dataset,
+            shuffle_rate,
+            extract_classes,
+            extract_separators,
+        )
+    else:
+        data = run_dataset_extraction_yolo(
+            dataset,
+            shuffle_rate,
+            extract_classes,
+            extract_separators,
+        )
+
+    # Write to JSONL in gzip format
+    for subset in data:
+        out_name = output_dir / f"{SPLIT_MAPPING.get(subset, subset)}.jsonl.gz"
+        logger.info(f"Saving {subset} data to file: {out_name}.")
+        save_gzip_jsonl(filename=out_name, content=data[subset])