initial commit

synthdog/template.py

+"""
+Donut
+Copyright (c) 2022-present NAVER Corp.
+MIT License
+"""
+import json
+import os
+import re
+from typing import Any, List
+
+import numpy as np
+from elements import Background, Document
+from PIL import Image
+from synthtiger import components, layers, templates
+
+
+class SynthDoG(templates.Template):
+    def __init__(self, config=None, split_ratio: List[float] = [0.8, 0.1, 0.1]):
+        super().__init__(config)
+        if config is None:
+            config = {}
+
+        self.quality = config.get("quality", [50, 95])
+        self.landscape = config.get("landscape", 0.5)
+        self.short_size = config.get("short_size", [720, 1024])
+        self.aspect_ratio = config.get("aspect_ratio", [1, 2])
+        self.background = Background(config.get("background", {}))
+        self.document = Document(config.get("document", {}))
+        self.effect = components.Iterator(
+            [
+                components.Switch(components.RGB()),
+                components.Switch(components.Shadow()),
+                components.Switch(components.Contrast()),
+                components.Switch(components.Brightness()),
+                components.Switch(components.MotionBlur()),
+                components.Switch(components.GaussianBlur()),
+            ],
+            **config.get("effect", {}),
+        )
+
+        # config for splits (output_filename, split_ratio etc)
+        self.splits = ["train", "validation", "test"]
+        self.split_indexes = [0, 0, 0]
+        self.split_ratio = [sum(split_ratio[: i + 1]) for i in range(0, len(split_ratio))]
+
+    def generate(self):
+        landscape = np.random.rand() < self.landscape
+        short_size = np.random.randint(self.short_size[0], self.short_size[1] + 1)
+        aspect_ratio = np.random.uniform(self.aspect_ratio[0], self.aspect_ratio[1])
+        long_size = int(short_size * aspect_ratio)
+        size = (long_size, short_size) if landscape else (short_size, long_size)
+
+        bg_layer = self.background.generate(size)
+        paper_layer, text_layers, texts = self.document.generate(size)
+
+        document_group = layers.Group([*text_layers, paper_layer])
+        document_space = np.clip(size - document_group.size, 0, None)
+        document_group.left = np.random.randint(document_space[0] + 1)
+        document_group.top = np.random.randint(document_space[1] + 1)
+        roi = np.array(paper_layer.quad, dtype=int)
+
+        layer = layers.Group([*document_group.layers, bg_layer]).merge()
+        self.effect.apply([layer])
+
+        image = layer.output(bbox=[0, 0, *size])
+        label = " ".join(texts)
+        label = label.strip()
+        label = re.sub(r"\s+", " ", label)
+        quality = np.random.randint(self.quality[0], self.quality[1] + 1)
+
+        data = {
+            "image": image,
+            "label": label,
+            "quality": quality,
+            "roi": roi,
+        }
+
+        return data
+
+    def init_save(self, root):
+        if not os.path.exists(root):
+            os.makedirs(root, exist_ok=True)
+
+    def save(self, root, data, idx):
+        image = data["image"]
+        label = data["label"]
+        quality = data["quality"]
+        roi = data["roi"]
+
+        # split
+        output_dirpath = os.path.join(root, "train")
+        file_idx = idx
+
+        split_prob = np.random.rand()
+        for _idx, (split, ratio) in enumerate(zip(self.splits, self.split_ratio)):
+            if split_prob < ratio:
+                output_dirpath = os.path.join(root, split)
+                file_idx = self.split_indexes[_idx]
+                self.split_indexes[_idx] += 1
+                break
+
+        # save image
+        image_filename = f"image_{file_idx}.jpg"
+        image_filepath = os.path.join(output_dirpath, image_filename)
+        os.makedirs(os.path.dirname(image_filepath), exist_ok=True)
+        image = Image.fromarray(image[..., :3].astype(np.uint8))
+        image.save(image_filepath, quality=quality)
+
+        # save metadata (gt_json)
+        metadata_filename = "metadata.jsonl"
+        metadata_filepath = os.path.join(output_dirpath, metadata_filename)
+        os.makedirs(os.path.dirname(metadata_filepath), exist_ok=True)
+
+        metadata = self.format_metadata(image_filename=image_filename, keys=["text_sequence"], values=[label])
+        with open(metadata_filepath, "a") as fp:
+            json.dump(metadata, fp, ensure_ascii=False)
+            fp.write("\n")
+
+    def end_save(self, root):
+        pass
+
+    def format_metadata(self, image_filename: str, keys: List[str], values: List[Any]):
+        """
+        Fit gt_parse contents to huggingface dataset's format
+        keys and values, whose lengths are equal, are used to constrcut 'gt_parse' field in 'ground_truth' field
+        Args:
+            keys: List of task_name
+            values: List of actual gt data corresponding to each task_name
+        """
+        assert len(keys) == len(values), "Length does not match: keys({}), values({})".format(len(keys), len(values))
+
+        _gt_parse_v = dict()
+        for k, v in zip(keys, values):
+            _gt_parse_v[k] = v
+        gt_parse = {"gt_parse": _gt_parse_v}
+        gt_parse_str = json.dumps(gt_parse, ensure_ascii=False)
+        metadata = {"file_name": image_filename, "ground_truth": gt_parse_str}
+        return metadata

synthdog/utils/__init__.py

+"""
+Donut
+Copyright (c) 2022-present NAVER Corp.
+MIT License
+"""
+from utils.text_reader import TextReader
+
+__all__ = ["TextReader"]

synthdog/utils/text_reader.py

+"""
+Donut
+Copyright (c) 2022-present NAVER Corp.
+MIT License
+"""
+from collections import OrderedDict
+
+
+class TextReader:
+    def __init__(self, path, cache_size=2 ** 28, block_size=2 ** 20):
+        self.fp = open(path, "r", encoding="utf-8")
+        self.length = 0
+        self.offsets = [0]
+        self.cache = OrderedDict()
+        self.cache_size = cache_size
+        self.block_size = block_size
+        self.bucket_size = cache_size // block_size
+        self.idx = 0
+
+        while True:
+            text = self.fp.read(self.block_size)
+            if not text:
+                break
+            self.length += len(text)
+            self.offsets.append(self.fp.tell())
+
+    def __len__(self):
+        return self.length
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        char = self.get()
+        self.next()
+        return char
+
+    def move(self, idx):
+        self.idx = idx
+
+    def next(self):
+        self.idx = (self.idx + 1) % self.length
+
+    def prev(self):
+        self.idx = (self.idx - 1) % self.length
+
+    def get(self):
+        key = self.idx // self.block_size
+
+        if key in self.cache:
+            text = self.cache[key]
+        else:
+            if len(self.cache) >= self.bucket_size:
+                self.cache.popitem(last=False)
+
+            offset = self.offsets[key]
+            self.fp.seek(offset, 0)
+            text = self.fp.read(self.block_size)
+            self.cache[key] = text
+
+        self.cache.move_to_end(key)
+        char = text[self.idx % self.block_size]
+        return char

test.py

+"""
+Donut
+Copyright (c) 2022-present NAVER Corp.
+MIT License
+"""
+import argparse
+import json
+import os
+import re
+from pathlib import Path
+
+import numpy as np
+import torch
+from datasets import load_dataset
+from PIL import Image
+from tqdm import tqdm
+
+from donut import DonutModel, JSONParseEvaluator, load_json, save_json
+
+
+def test(args):
+    pretrained_model = DonutModel.from_pretrained(args.pretrained_path)
+
+    if torch.cuda.is_available():
+        pretrained_model.half()
+        pretrained_model.to("cuda")
+    else:
+        pretrained_model.encoder.to(torch.bfloat16)
+
+    pretrained_model.eval()
+
+    if args.save_path:
+        os.makedirs(os.path.dirname(args.save_path), exist_ok=True)
+
+    output_list = []
+    accs = []
+
+    dataset = load_dataset(args.dataset_name_or_path, split=args.split)
+
+    for idx, sample in tqdm(enumerate(dataset), total=len(dataset)):
+        ground_truth = json.loads(sample["ground_truth"])
+
+        if args.task_name == "docvqa":
+            output = pretrained_model.inference(
+                image=sample["image"],
+                prompt=f"<s_{args.task_name}><s_question>{ground_truth["gt_parses"][0]['question'].lower()}</s_question><s_answer>",
+            )["predictions"][0]
+        else:
+            output = pretrained_model.inference(image=sample["image"], prompt=f"<s_{args.task_name}>")["predictions"][0]
+
+        if args.task_name == "rvlcdip":
+            gt = ground_truth["gt_parse"]
+            score = float(output["class"] == gt["class"])
+        elif args.task_name == "docvqa":
+            score = 0.0  # note: docvqa is evaluated on the official website
+        else:
+            gt = ground_truth["gt_parse"]
+            evaluator = JSONParseEvaluator()
+            score = evaluator.cal_acc(output, gt)
+
+        accs.append(score)
+
+        output_list.append(output)
+
+    scores = {"accuracies": accs, "mean_accuracy": np.mean(accs)}
+    print(scores, f"length : {len(accs)}")
+
+    if args.save_path:
+        scores["predictions"] = output_list
+        save_json(args.save_path, scores)
+
+    return output_list
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--pretrained_path", type=str)
+    parser.add_argument("--dataset_name_or_path", type=str)
+    parser.add_argument("--split", type=str, default="test")
+    parser.add_argument("--task_name", type=str, default=None)
+    parser.add_argument("--save_path", type=str, default=None)
+    args, left_argv = parser.parse_known_args()
+
+    if args.task_name is None:
+        args.task_name = os.path.basename(args.dataset_name_or_path)
+
+    predicts = test(args)

train.py

+"""
+Donut
+Copyright (c) 2022-present NAVER Corp.
+MIT License
+"""
+import argparse
+import datetime
+import json
+import os
+import random
+from io import BytesIO
+from os.path import basename
+from pathlib import Path
+
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from pytorch_lightning.callbacks import LearningRateMonitor, ModelCheckpoint
+from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
+from pytorch_lightning.plugins import CheckpointIO
+from pytorch_lightning.utilities import rank_zero_only
+from sconf import Config
+
+from donut import DonutDataset
+from lightning_module import DonutDataPLModule, DonutModelPLModule
+
+
+class CustomCheckpointIO(CheckpointIO):
+    def save_checkpoint(self, checkpoint, path, storage_options=None):
+        del checkpoint["state_dict"]
+        torch.save(checkpoint, path)
+
+    def load_checkpoint(self, path, storage_options=None):
+        checkpoint = torch.load(path + "artifacts.ckpt")
+        state_dict = torch.load(path + "pytorch_model.bin")
+        checkpoint["state_dict"] = {"model." + key: value for key, value in state_dict.items()}
+        return checkpoint
+
+    def remove_checkpoint(self, path) -> None:
+        return super().remove_checkpoint(path)
+
+
+@rank_zero_only
+def save_config_file(config, path):
+    if not Path(path).exists():
+        os.makedirs(path)
+    save_path = Path(path) / "config.yaml"
+    print(config.dumps())
+    with open(save_path, "w") as f:
+        f.write(config.dumps(modified_color=None, quote_str=True))
+        print(f"Config is saved at {save_path}")
+
+
+def train(config):
+    pl.utilities.seed.seed_everything(config.get("seed", 42), workers=True)
+
+    model_module = DonutModelPLModule(config)
+    data_module = DonutDataPLModule(config)
+
+    # add datasets to data_module
+    datasets = {"train": [], "validation": []}
+    for i, dataset_name_or_path in enumerate(config.dataset_name_or_paths):
+        task_name = os.path.basename(dataset_name_or_path)  # e.g., cord-v2, docvqa, rvlcdip, ...
+        for split in ["train", "validation"]:
+            datasets[split].append(
+                DonutDataset(
+                    dataset_name_or_path=dataset_name_or_path,
+                    donut_model=model_module.model,
+                    max_length=config.max_length,
+                    split=split,
+                    task_start_token=config.task_start_tokens[i]
+                    if config.get("task_start_tokens", None)
+                    else f"<s_{task_name}>",
+                    prompt_end_token="<s_answer>" if "docvqa" in dataset_name_or_path else f"<s_{task_name}>",
+                    sort_json_key=config.sort_json_key,
+                )
+            )
+            # prompt_end_token is used for ignoring a given prompt in a loss function
+            # for docvqa task, i.e., {"question": {used as a prompt}, "answer": {prediction target}},
+            # set prompt_end_token to "<s_answer>"
+    data_module.train_datasets = datasets["train"]
+    data_module.val_datasets = datasets["validation"]
+
+    logger = TensorBoardLogger(
+        save_dir=config.result_path,
+        name=config.exp_name,
+        version=config.exp_version,
+        default_hp_metric=False,
+    )
+
+    lr_callback = LearningRateMonitor(logging_interval="step")
+
+    checkpoint_callback = ModelCheckpoint(
+        monitor="val_metric",
+        dirpath=Path(config.result_path) / config.exp_name / config.exp_version,
+        filename="artifacts",
+        save_top_k=1,
+        save_last=False,
+        mode="min",
+    )
+
+    custom_ckpt = CustomCheckpointIO()
+    trainer = pl.Trainer(
+        resume_from_checkpoint=config.get("resume_from_checkpoint_path", None),
+        num_nodes=config.get("num_nodes", 1),
+        gpus=torch.cuda.device_count(),
+        strategy="ddp",
+        accelerator="gpu",
+        plugins=custom_ckpt,
+        max_epochs=config.max_epochs,
+        max_steps=config.max_steps,
+        val_check_interval=config.val_check_interval,
+        check_val_every_n_epoch=config.check_val_every_n_epoch,
+        gradient_clip_val=config.gradient_clip_val,
+        precision=16,
+        num_sanity_val_steps=0,
+        logger=logger,
+        callbacks=[lr_callback, checkpoint_callback],
+    )
+
+    trainer.fit(model_module, data_module)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, required=True)
+    parser.add_argument("--exp_version", type=str, required=False)
+    args, left_argv = parser.parse_known_args()
+
+    config = Config(args.config)
+    config.argv_update(left_argv)
+
+    config.exp_name = basename(args.config).split(".")[0]
+    config.exp_version = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") if not args.exp_version else args.exp_version
+
+    save_config_file(config, Path(config.result_path) / config.exp_name / config.exp_version)
+    train(config)