chugr 0.2.0.dev0__py3-none-any.whl

chugr/image/transforms_alb.py

@@ -0,0 +1,112 @@
+
+try:
+    import albumentations as alb
+    from albumentations.pytorch import ToTensorV2
+    has_albumentations = True
+except ImportError:
+    has_albumentations = False
+
+try:
+    import cv2
+    has_cv2 = True
+except ImportError:
+    has_cv2 = False
+
+import random
+
+import numpy as np
+
+
+class AlbWrapper:
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, im):
+        return self.transforms(image=np.asarray(im))["image"]
+
+    def __repr__(self) -> str:
+        format_string = self.__class__.__name__ + "("
+        for t in self.transforms:
+            format_string += "\n"
+            format_string += f"        {t}"
+        format_string += "\n)"
+        return format_string
+
+
+if has_albumentations and has_cv2:
+
+    class ErosionAlb(alb.ImageOnlyTransform):
+        def __init__(self, scale, p=0.5):
+            super().__init__(p=p)
+            if type(scale) is tuple or type(scale) is list:
+                assert len(scale) == 2
+                self.scale = scale
+            else:
+                self.scale = (scale, scale)
+
+        def get_transform_init_args_names(self):
+            return ('scale',)
+
+        def apply(self, img, **params):
+            kernel = cv2.getStructuringElement(
+                cv2.MORPH_ELLIPSE, tuple(int(k) for k in random.choices(range(self.scale[0], self.scale[1] + 1), k=2))
+            )
+            img = cv2.erode(img, kernel, iterations=1)
+            return img
+
+
+    class DilationAlb(alb.ImageOnlyTransform):
+        def __init__(self, scale, p=0.5):
+            super().__init__(p=p)
+            if type(scale) is tuple or type(scale) is list:
+                assert len(scale) == 2
+                self.scale = scale
+            else:
+                self.scale = (scale, scale)
+
+        def get_transform_init_args_names(self):
+            return ('scale',)
+
+        def apply(self, img, **params):
+            kernel = cv2.getStructuringElement(
+                cv2.MORPH_ELLIPSE,
+                tuple(int(k) for k in random.choices(range(self.scale[0], self.scale[1] + 1), k=2))
+            )
+            img = cv2.dilate(img, kernel, iterations=1)
+            return img
+
+
+    class BitmapAlb(alb.ImageOnlyTransform):
+        def __init__(self, value=0, lower=200, p=0.5):
+            super().__init__(p=p)
+            self.lower = lower
+            self.value = value
+
+        def get_transform_init_args_names(self):
+            return ('value', 'lower')
+
+        def apply(self, img, **params):
+            img = img.copy()
+            img[img < self.lower] = self.value
+            return img
+
+
+    class CropMarginCv2:
+
+        def __init__(self):
+            pass
+
+        def __call__(self, img):
+            data = np.array(img.convert("L"))
+            data = data.astype(np.uint8)
+            max_val = data.max()
+            min_val = data.min()
+            if max_val == min_val:
+                return img
+            data = (data - min_val) / (max_val - min_val) * 255
+            gray = 255 * (data < 200).astype(np.uint8)
+
+            coords = cv2.findNonZero(gray)  # Find all non-zero points (text)
+            a, b, w, h = cv2.boundingRect(coords)  # Find minimum spanning bounding box
+            return img.crop((a, b, w + a, h + b))
+

chugr/image/transforms_factory.py

@@ -0,0 +1,58 @@
+from typing import Any, Dict, Optional, Union
+
+from chugr.common import ImageInputCfg, ImageAugCfg
+from .build_transforms_doc import build_transforms_doc_better, build_transforms_doc_nougat, build_transforms_doc_basic, \
+    build_transforms_doc_naflex
+from .build_transforms_image import build_transforms_image_timm, build_transforms_image_basic
+
+_transform_factories = {
+    "image_basic": build_transforms_image_basic,
+    "image_timm": build_transforms_image_timm,
+    "doc_basic": build_transforms_doc_basic,
+    "doc_nougat": build_transforms_doc_nougat,
+    "doc_better": build_transforms_doc_better,
+    "doc_naflex": build_transforms_doc_naflex,
+}
+
+def build_image_transforms(
+        input_cfg: ImageInputCfg,
+        is_training=True,
+        do_normalize=True,
+        do_convert=False,
+        composed=True,
+        aug_cfg: Optional[Union[Dict[str, Any], ImageAugCfg]] = None,
+):
+    common_args = dict(
+        input_cfg=input_cfg,
+        is_training=is_training,
+        do_normalize=do_normalize,
+        aug_cfg=aug_cfg,
+        composed=composed,
+    )
+
+    tt = input_cfg.transform_type
+    assert tt in _transform_factories, \
+        f"Unrecognized transform type: {tt}. Must be one of {list(_transform_factories.keys())}."
+    transforms = _transform_factories[tt](**common_args)
+
+    return transforms
+
+
+def create_image_preprocessor(
+        input_cfg: ImageInputCfg,
+        is_training=True,
+        do_normalize=True,
+        do_convert=False,
+        aug_cfg: Optional[Union[Dict[str, Any], ImageAugCfg]] = None,
+):
+    transforms = build_image_transforms(
+        input_cfg=input_cfg,
+        is_training=is_training,
+        do_normalize=do_normalize,
+        do_convert=do_convert,
+        aug_cfg=aug_cfg,
+        composed=True,
+    )
+    # NOTE for now, a stack of composed transforms are the image pre-processor
+    return transforms
+

chugr/image/transforms_torch.py

@@ -0,0 +1,137 @@
+import random
+
+import numpy as np
+import torch
+import torchvision.transforms.functional as F
+from torchvision import transforms
+from PIL import Image, ImageFilter
+
+
+class AlignLongAxis:
+    def __init__(
+            self,
+            input_size,
+            interpolation=transforms.InterpolationMode.BICUBIC
+    ):
+        self.input_size = input_size
+        self.interpolation = interpolation
+
+    def __call__(self, img):
+        img_width, img_height = F.get_image_size(img)
+        if (
+            (self.input_size[0] > self.input_size[1] and img_width > img_height) or
+            (self.input_size[0] < self.input_size[1] and img_width < img_height)
+        ):
+            img = F.rotate(img, angle=-90, expand=True, interpolation=self.interpolation)
+        return img
+
+
+class Bitmap:
+    def __init__(self, threshold=200):
+        self.lut = [0 if i < threshold else i for i in range(256)]
+
+    def __call__(self, img):
+        if img.mode == "RGB" and len(self.lut) == 256:
+            lut = self.lut + self.lut + self.lut
+        else:
+            lut = self.lut
+        return img.point(lut)
+
+
+class Erosion:
+    def __init__(self, scale=3):
+        super().__init__()
+        if type(scale) is tuple or type(scale) is list:
+            assert len(scale) == 2
+            self.scale = scale
+        else:
+            self.scale = (scale, scale)
+
+    @staticmethod
+    def get_params(scale):
+        if type(scale) is tuple or type(scale) is list:
+            assert len(scale) == 2
+            scale = random.choice(scale)
+        return scale
+
+    def __call__(self, img):
+        kernel_size = self.get_params(self.scale)
+        if isinstance(img, torch.Tensor):
+            padding = kernel_size // 2
+            img = -torch.nn.functional.max_pool2d(-img, kernel_size=kernel_size, stride=1, padding=padding)  # minpool
+        elif isinstance(img, Image.Image):
+            img = img.filter(ImageFilter.MinFilter(kernel_size))
+        return img
+
+
+class Dilation:
+    def __init__(self, scale=3):
+        super().__init__()
+        self.scale = scale
+
+    @staticmethod
+    def get_params(scale):
+        if type(scale) is tuple or type(scale) is list:
+            assert len(scale) == 2
+            scale = random.choice(scale)
+        return scale
+
+    def __call__(self, img):
+        kernel_size = self.get_params(self.scale)
+        if isinstance(img, torch.Tensor):
+            padding = kernel_size // 2
+            img = torch.nn.functional.max_pool2d(img, kernel_size=kernel_size, stride=1, padding=padding)
+        elif isinstance(img, Image.Image):
+            img = img.filter(ImageFilter.MaxFilter(kernel_size))
+        return img
+
+
+def python_find_non_zero(image: np.array):
+    """This is a reimplementation of a findNonZero function equivalent to cv2."""
+    non_zero_indices = np.column_stack(np.nonzero(image))
+    idxvec = non_zero_indices[:, [1, 0]]
+    idxvec = idxvec.reshape(-1, 1, 2)
+    return idxvec
+
+
+def python_bounding_rect(coordinates):
+    """This is a reimplementation of a BoundingRect function equivalent to cv2."""
+    min_values = np.min(coordinates, axis=(0, 1)).astype(int)
+    max_values = np.max(coordinates, axis=(0, 1)).astype(int)
+    x_min, y_min = min_values[0], min_values[1]
+    width = max_values[0] - x_min + 1
+    height = max_values[1] - y_min + 1
+    return x_min, y_min, width, height
+
+
+class CropMargin:
+    def __init__(self) -> None:
+        pass
+
+    def __call__(
+        self,
+        image,
+        gray_threshold: int = 200,
+    ) -> np.array:
+        # FIXME check tensor vs PIL and convert as needed, this is assuming PIL right now
+        assert not isinstance(image, torch.Tensor)
+        data = np.array(image.convert("L")).astype(np.uint8)
+        max_val = data.max()
+        min_val = data.min()
+        if max_val == min_val:
+            return image
+        data = (data - min_val) / (max_val - min_val) * 255
+        gray = data < gray_threshold
+        coords = python_find_non_zero(gray)
+        x_min, y_min, width, height = python_bounding_rect(coords)
+        image = F.crop(image, y_min, x_min, height, width)
+        return image
+
+
+class ConvertColor:
+    def __init__(self, mode='RGB'):
+        self.mode = mode
+
+    def __call__(self, image):
+        assert isinstance(image, Image.Image)
+        return image.convert(self.mode)

chugr/loader.py

@@ -0,0 +1,249 @@
+import warnings
+from functools import partial
+from typing import Callable, List, Optional
+
+from chugr.common import DataCfg, DataTaskCfg, DistributedCfg, LoaderBundle, source_to_shard_spec, SourceSpec
+from chugr.hfds import create_loader_hf
+from chugr.task_pipeline import create_task_pipeline
+from chugr.wds import create_loader_wds, get_error_handler
+
+
+def create_loader(
+        data_cfg: DataCfg,
+        task_cfg: DataTaskCfg,
+        task_pipeline: Optional[List[Callable]] = None,
+        is_training: bool = False,
+        start_interval: int = 0,
+        seed: int = 0,
+        distributed: DistributedCfg = DistributedCfg(),
+) -> LoaderBundle:
+    """
+    Creates a dataloader for training or validation based on configuration settings.
+
+    Parameters:
+        data_cfg: Configuration object for the dataset.
+        task_cfg: Configuration object for the task specific processing.
+        task_pipeline: Task specific processing pipeline (takes priority over task_cfg).
+        is_training : Indicates if the loader is for training data (True) or validation data (False).
+        start_interval: The starting interval (epoch for full passes) for setting seed, etc. appropriately.
+        seed: Seed for random operations to ensure reproducibility.
+        distributed: Distributed device information.
+
+    Returns:
+        DataLoader: A PyTorch DataLoader instance configured according to the provided settings.
+
+    Note:
+        Currently supports "wds" and "hf_dataset" as dataset formats.
+    """
+    if data_cfg.format == "wds":
+        loader = create_loader_from_config_wds(
+            data_cfg=data_cfg,
+            task_cfg=task_cfg,
+            task_pipeline=task_pipeline,
+            is_training=is_training,
+            start_interval=start_interval,
+            seed=seed,
+            distributed=distributed,
+        )
+
+    elif data_cfg.format.startswith("hf"):
+        loader = create_loader_from_config_hf(
+            data_cfg=data_cfg,
+            task_cfg=task_cfg,
+            task_pipeline=task_pipeline,
+            is_training=is_training,
+            start_interval=start_interval,
+            seed=seed,
+            distributed=distributed,
+        )
+
+    else:
+        assert False, f"Unsupported dataset format ({data_cfg.format})."
+
+    return loader
+
+
+def _validate_cfgs(
+    data_cfg: DataCfg,
+    task_cfg: Optional[DataTaskCfg],
+    is_training: bool = False,
+):
+    batch_size = data_cfg.batch_size
+    if batch_size is not None:
+        if task_cfg.decode_and_process_fn is None:
+            # images must be processed into tensors to batch, text may be left as strings (collated as lists)
+            assert task_cfg.image_process_fn is not None, \
+                'task_cfg.image_process_fn must be set if batching enabled'
+        else:
+            assert task_cfg.decode_fn is None, \
+                'task_cfg.decode_fn should not be set at the same time as task_cfg.decode_and_process_fn'
+
+
+def create_loader_from_config_wds(
+        data_cfg: DataCfg,
+        task_cfg: Optional[DataTaskCfg],
+        task_pipeline: Optional[List[Callable]] = None,
+        is_training: bool = False,
+        start_interval: int = 0,
+        seed: int = 0,
+        collate_fn: Optional[Callable] = None,
+        distributed: DistributedCfg = DistributedCfg(),
+):
+    """
+
+    Args:
+        data_cfg:
+        task_cfg:
+        task_pipeline:
+        is_training:
+        start_interval:
+        seed:
+        collate_fn:
+        distributed:
+
+    Returns:
+
+    """
+    _validate_cfgs(data_cfg, task_cfg, is_training=is_training)
+
+    if task_pipeline is None:
+        assert task_cfg is not None
+        task_pipeline = create_task_pipeline(
+            task_cfg,
+        )
+
+    handler = get_error_handler(task_cfg.error_handler)
+
+    naflex_scheduler = None
+    naflex_stage_factory = None
+    if is_training and data_cfg.naflex_seq_lens:
+        from chugr.naflex import LengthBucketer, NaFlexBatcher
+
+        assert data_cfg.num_samples, 'NaFlex batching requires num_samples (per interval).'
+        naflex_scheduler, processors, image_key = _create_naflex_batching(
+            data_cfg, task_cfg, num_samples=data_cfg.num_samples, seed=seed, world_size=distributed.world_size)
+
+        def naflex_stage_factory(interval):
+            stages = []
+            if data_cfg.length_bucketing:
+                stages.append(LengthBucketer(seed=seed, interval=interval))
+            stages.append(NaFlexBatcher(
+                naflex_scheduler, processors, image_key=image_key, interval=interval, handler=handler))
+            return stages
+
+    return create_loader_wds(
+        shards=data_cfg.shard_spec,
+        task_pipeline=task_pipeline,
+        num_samples=data_cfg.num_samples,
+        is_training=is_training,
+        resampled=data_cfg.resampled,
+        multi_interval=True, # FIXME via config?
+        num_workers=data_cfg.num_workers,
+        batch_size=data_cfg.batch_size,
+        persistent_workers=data_cfg.persistent_workers,
+        collate_fn=collate_fn,
+        start_interval=start_interval,
+        seed=seed,
+        handler=handler,
+        distributed=distributed,
+        naflex_scheduler=naflex_scheduler,
+        naflex_stage_factory=naflex_stage_factory,
+    )
+
+
+def _create_naflex_batching(data_cfg: DataCfg, task_cfg: DataTaskCfg, num_samples: int, seed: int, world_size: int):
+    """ NaFlex batch schedule (for `num_samples` per interval), per seq_len task processors, and image key. """
+    from chugr.naflex import NaFlexBatchScheduler
+    from chugr.task_pipeline import create_task_processor
+
+    assert task_cfg.image_process_factory is not None, \
+        'NaFlex batching requires task_cfg.image_process_factory (seq_len -> image process fn).'
+    max_tokens = data_cfg.naflex_max_tokens or (data_cfg.batch_size or 1) * max(data_cfg.naflex_seq_lens)
+    scheduler = NaFlexBatchScheduler(
+        seq_lens=data_cfg.naflex_seq_lens,
+        max_tokens_per_batch=max_tokens,
+        num_samples=num_samples,
+        seq_len_probs=data_cfg.naflex_seq_len_probs,
+        per_row_text_tokens=data_cfg.naflex_text_tokens,
+        batch_divisor=data_cfg.naflex_batch_divisor,
+        seed=seed,
+        world_size=world_size,
+    )
+    processors = {
+        s: create_task_processor(task_cfg, image_process_fn=task_cfg.image_process_factory(s))
+        for s in scheduler.seq_lens
+    }
+    image_key = getattr(task_cfg, 'image_input_feat').output_name
+    return scheduler, processors, image_key
+
+
+def create_loader_from_config_hf(
+        data_cfg: DataCfg,
+        task_cfg: DataTaskCfg,
+        task_pipeline: Optional[List[Callable]] = None,
+        is_training: bool = False,
+        start_interval: int = 0,
+        seed: int = 0,
+        distributed: DistributedCfg = DistributedCfg(),
+):
+    """
+
+    Args:
+        data_cfg:
+        task_cfg:
+        task_pipeline:
+        is_training:
+        start_interval:
+        seed:
+        distributed:
+
+    Returns:
+
+    """
+    assert not isinstance(data_cfg.source, (list, tuple)), "Multiple sources not supported for HF datasets."
+    assert isinstance(data_cfg.source, (str, SourceSpec)), \
+        "The specified source for HF dataset must be a string or SourceSpec."
+
+    _validate_cfgs(data_cfg, task_cfg, is_training=is_training)
+
+    if isinstance(data_cfg.source, SourceSpec):
+        source = data_cfg.source.url
+        split = data_cfg.source.split
+        assert split, "Split must be set in SourceSpec with HF datasets."
+    else:
+        source = data_cfg.source
+        split = data_cfg.split
+        assert split, "Split must be set in DataCfg when string source is used with HF datasets."
+
+    if task_pipeline is None:
+        assert task_cfg is not None
+        task_pipeline = create_task_pipeline(
+            task_cfg,
+        )
+
+    streaming = 'hfids' in data_cfg.format
+
+    naflex_factory = None
+    if is_training and data_cfg.naflex_seq_lens:
+        # schedule driven NaFlex batching, the schedule is built once the dataset length is known (map-style)
+        if data_cfg.length_bucketing:
+            warnings.warn('length_bucketing is not supported for HF datasets, ignored.')
+        naflex_factory = partial(
+            _create_naflex_batching, data_cfg, task_cfg, seed=seed, world_size=distributed.world_size)
+
+    return create_loader_hf(
+        source=source,
+        split=split,
+        task_pipeline=task_pipeline,
+        streaming=streaming,
+        is_training=is_training,
+        batch_size=data_cfg.batch_size,
+        data_dir=data_cfg.data_dir,
+        num_samples=data_cfg.num_samples,
+        num_workers=data_cfg.num_workers,
+        persistent_workers=data_cfg.persistent_workers,
+        seed=seed,
+        distributed=distributed,
+        naflex_factory=naflex_factory,
+        handler=get_error_handler(task_cfg.error_handler),
+    )