chugr 0.2.0.dev0__py3-none-any.whl

chugr/common/types.py

@@ -0,0 +1,122 @@
+from dataclasses import dataclass
+from multiprocessing import Value
+from typing import Any, Dict, List, Optional, Tuple, Union
+from numbers import Number
+
+from torch.utils.data import DataLoader, DistributedSampler
+
+
+class SharedCount:
+    def __init__(self, count: int = 0):
+        self.count = Value('i', count)
+
+    def set_value(self, epoch):
+        self.count.value = epoch
+
+    def get_value(self):
+        return self.count.value
+
+
+@dataclass
+class LoaderBundle:
+    """
+    Bundle a DataLoader with num_batch / num_sample limits, sampler or shared_interval counter exposed
+    to allow easy seed control per-interval.
+    """
+    loader: DataLoader
+    num_batches: int = 0
+    num_samples: int = 0
+    sampler: DistributedSampler = None
+    shared_interval: SharedCount = None
+    mean_batch_size: Optional[float] = None  # set for variable batch size (NaFlex schedule) loaders
+
+    def set_interval(self, interval):
+        if self.shared_interval is not None:
+            self.shared_interval.set_value(interval)
+        dataset = getattr(self.loader, 'dataset', None)
+        if hasattr(dataset, 'set_interval_count'):
+            # main process side interval state of the dataset (e.g. the HF streaming dataset epoch)
+            dataset.set_interval_count(interval)
+        if self.sampler is not None and hasattr(self.sampler, 'set_epoch'):
+            self.sampler.set_epoch(interval)
+
+    def __iter__(self):
+        return self.loader.__iter__()
+
+
+@dataclass
+class SplitInfo:
+    filenames: Tuple[str]
+    num_samples: int
+    shard_lengths: Tuple[int] = ()
+    name: str = ''
+
+
+# @dataclass
+# class ShardInfo:
+#     url: str
+#     weight: float = 1.0
+#     num_samples: Optional[int] = None
+#
+#
+
+
+@dataclass
+class SourceSpec:
+    url: str
+    split: Optional[str] = None  # dataset split
+    template: Optional[str] = None  # template to transform url -> usage
+    sampling_weight: Optional[float] = None
+    num_samples: Optional[int] = None
+
+    # TODO resolve dataset info and track base url, shard info (sizes, etc)
+    # base_url: str = None
+    # info_url: str = None
+
+
+@dataclass
+class SourceInfo(SourceSpec):
+    split_info: Dict[str, SplitInfo] = None
+    shard_info: Dict[str, Dict[str, Any]] = None
+
+
+@dataclass
+class ShardSpec:
+    urls: List[str]
+    weights: Optional[Union[float, List[float]]] = None
+    sizes: Optional[List[int]] = None
+
+    def __post_init__(self):
+        num_shards = len(self.urls)
+        if self.weights is not None:
+            if isinstance(self.weights, Number):
+                self.weights = [self.weights] * num_shards
+            assert len(self.weights) == num_shards
+        if self.sizes is not None:
+            assert len(self.sizes) == num_shards
+
+
+@dataclass(frozen=True)
+class FeatureInfo:
+    """ Feature Information
+
+    Attributes:
+        output_name: output feature name, None if an intermediary feature
+        input_key: input dataset key(s), ';' delimited for multiple options
+    """
+    output_name: Optional[str] = 'image'
+    input_key: Optional[str] = 'jpg;png'
+    #parent: Optional[str] = None
+
+
+@dataclass(frozen=True)
+class ImageFeatureInfo(FeatureInfo):
+    """ Image Feature Information
+
+    Attributes:
+        image_mode: Image colour mode (e.g. 'RGB', 'RGBA', 'L')
+        output_name: output feature name, None if an intermediary feature
+        input_key: input dataset key(s), ';' delimited for multiple options
+        parent: parent key to search for input_key, e.g. 'json'
+    """
+    image_mode: str = 'RGB'

chugr/common/urls.py

@@ -0,0 +1,88 @@
+import os
+from numbers import Number
+from typing import Sequence
+
+
+import braceexpand
+import re
+
+
+def envlookup(m):
+    """Look up match in the environment with prefix WDS_.
+
+    Args:
+        m: a match object
+
+    Returns:
+        str: the value of the environment variable WDS_<m.group(1)>
+    """
+    key = m.group(1)
+    for prefix in ('WDS_', 'CHUG_'):
+        key = prefix + key
+        if key in os.environ:
+            return os.environ[key]
+    assert key in os.environ, f"missing WDS/CHUG environment variable for {key}"
+
+
+def envsubst(s):
+    """Substitute ${var} with the value of the environment variable WDS_var.
+
+    Args:
+        s (str): string to be substituted
+
+    Returns:
+        str: the substituted string
+    """
+    return re.sub(r"\$\{(\w+)\}", envlookup, s)
+
+
+def _subst_and_expand(url: str):
+    for i in range(10):
+        last = url
+        url = envsubst(url)
+        if url == last:
+            break
+    return braceexpand.braceexpand(url)
+
+
+def expand_urls(urls, weights=None):
+    """ Expand urls (and optionally weights) if they are strings, otherwise return as lists.
+    """
+    if weights is None:
+        if isinstance(urls, str):
+            url_list = urls.split("::")
+            result = []
+            for url in url_list:
+                result.extend(_subst_and_expand(url))
+            return result, None
+        else:
+            return list(urls), None
+
+    if isinstance(urls, str):
+        url_list = urls.split('::')
+
+        if isinstance(weights, str):
+            weights = weights.split('::')
+        elif isinstance(weights, Number):
+            weights = [weights] * len(url_list)
+        assert len(weights) == len(url_list), \
+            f"Expected the number of data components ({len(url_list)}) and weights({len(weights)}) to match."
+        weights = [float(weight) for weight in weights]
+        all_urls, all_weights = [], []
+        for url, weight in zip(url_list, weights):
+            expanded_url = list(_subst_and_expand(url))
+            expanded_weights = [weight] * len(expanded_url)
+            all_urls.extend(expanded_url)
+            all_weights.extend(expanded_weights)
+    else:
+        all_urls = list(urls)
+        if isinstance(weights, Number):
+            # if weights is a scalar, expand to url list
+            all_weights = [float(weights)] * len(all_urls)
+        else:
+            assert len(weights) == len(all_urls), \
+                f"Expected the number of data components ({len(all_urls)}) and weights({len(weights)}) to match."
+            all_weights = list(weights)
+
+    return all_urls, all_weights
+

chugr/doc/__init__.py

@@ -0,0 +1,17 @@
+from .constants import (
+    DEFAULT_DOC_KEY,
+    DEFAULT_QUESTION_KEY,
+    DEFAULT_QUESTION_ID_KEY,
+    DEFAULT_ANSWER_KEY,
+    DEFAULT_DOC_KEY_TUPLE,
+    DEFAULT_QUESTION_KEY_TUPLE,
+    DEFAULT_ANSWER_KEY_TUPLE,
+    DEFAULT_DOC_FEAT,
+    DEFAULT_QUESTION_FEAT,
+    DEFAULT_QUESTION_ID_FEAT,
+    DEFAULT_ANSWER_FEAT
+)
+from .doc_processor import DocProcessor
+from .doc_read_processor import DocReadProcessor
+from .doc_vqa_processor import DocVqaProcessor
+from .synth_text import SynthTextCfg, SynthTextRenderer

chugr/doc/constants.py

@@ -0,0 +1,15 @@
+from chugr import FeatureInfo, ImageFeatureInfo
+
+DEFAULT_DOC_KEY = "pdf;tif;png;jpeg;jpg;webp;image"
+DEFAULT_QUESTION_KEY = "question;query"
+DEFAULT_QUESTION_ID_KEY = "question_id;query_id"
+DEFAULT_ANSWER_KEY = "answer;answers"
+
+DEFAULT_DOC_KEY_TUPLE = tuple(DEFAULT_DOC_KEY.split(';'))
+DEFAULT_QUESTION_KEY_TUPLE = tuple(DEFAULT_QUESTION_KEY.split(';'))
+DEFAULT_ANSWER_KEY_TUPLE = tuple(DEFAULT_ANSWER_KEY.split(';'))
+
+DEFAULT_DOC_FEAT = ImageFeatureInfo('image_input', input_key=DEFAULT_DOC_KEY, image_mode='L')
+DEFAULT_QUESTION_FEAT = FeatureInfo(None, input_key=DEFAULT_QUESTION_KEY)
+DEFAULT_QUESTION_ID_FEAT = FeatureInfo(None, input_key=DEFAULT_QUESTION_ID_KEY)
+DEFAULT_ANSWER_FEAT = FeatureInfo(None, input_key=DEFAULT_ANSWER_KEY)

chugr/doc/doc_processor.py

@@ -0,0 +1,248 @@
+import abc
+import json
+import logging
+import random
+from typing import Callable, Dict, List, Optional, Tuple
+
+from chugr.common import FeatureInfo, ImageFeatureInfo
+from chugr.wds import decode_image_pages, decode_pdf_pages
+
+
+from .constants import (
+    DEFAULT_DOC_FEAT,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+def get_next_valid_page_index(
+        current_index: int,
+        num_pages: int,
+        page_annos: list,
+        retries: int = 10,
+        wanted: str = 'lines',
+):
+    """
+    Get the index of the next valid page which contains text. If it doesn't find any non-empty page
+    after 'retries' attempts, it raises a RuntimeError.
+
+    Parameters:
+    current_index (int): Current page index.
+    num_pages (int): Total number of pages.
+    page_annos (list): List of page annotations.
+    retries (int): Number of maximum retries for a given document.
+
+    Returns:
+    int: The index of the next non-empty page.
+    """
+    for _ in range(retries):
+        # Get the next index, wrap around to 0 if it exceeds num_pages (in case of random init)
+        current_index = (current_index + 1) % num_pages
+        anno_page = page_annos[current_index]
+        anno_page = anno_page.get(wanted, anno_page)  # use 'lines' / 'words' level if exists
+        if anno_page["text"]:
+            return current_index
+    raise RuntimeError(f"No non-empty page found after {retries} attempts")
+
+
+def _get_value(keys, sample):
+    if isinstance(keys, (list, tuple)):
+        value = None
+        for k in keys:
+            if (value := sample.get(k, None)) is not None:
+                break
+        return value
+    else:
+        return sample.get(keys, None)
+
+
+class DocProcessor:
+    """ Process documents w/ OCR annotation for reading tasks.
+    """
+
+    def __init__(
+            self,
+            image_process_fn: Optional[Callable] = None,
+            text_process_fn: Optional[Callable] = None,
+            image_input_feat: ImageFeatureInfo = DEFAULT_DOC_FEAT,
+            text_input_feat: FeatureInfo = FeatureInfo('text_input', input_key='pages'),
+            text_target_feat: FeatureInfo = FeatureInfo('text_target', input_key=None),
+            page_sampling: str = 'random',
+            render_dpi: int = 150,
+            squeeze_pages: bool = True,
+            expand_pages: bool = False,
+            flatten_json: bool = True,
+            seed: int = 0,
+    ):
+        """
+
+        Args:
+            image_process_fn:
+            text_process_fn:
+            page_sampling:
+            render_dpi:
+            seed:
+        """
+        self.image_process_fn = image_process_fn
+        self.text_process_fn = text_process_fn
+
+        self.image_input_feat = image_input_feat
+        self.image_input_name = image_input_feat.output_name
+        self.image_input_key = image_input_feat.input_key.split(';')
+        self.text_input_feat = text_input_feat
+        self.text_input_name = text_input_feat.output_name
+        self.text_input_key = text_input_feat.input_key.split(';') if text_input_feat.input_key else None
+        self.text_target_feat = text_target_feat
+        self.text_target_name = text_target_feat.output_name if text_target_feat is not None else None
+
+        self.page_sampling = page_sampling
+        self.render_dpi = render_dpi
+        self.squeeze_pages = squeeze_pages
+        self.expand_pages = expand_pages
+        self.flatten_json = flatten_json
+        self.seed = seed
+        self._generator = random.Random(seed)
+        self._generator_worker = None
+        # FIXME note, should move to torchvision v2 annotations at some point
+        #  * they should all eventually have a generator arg for better handling random state
+        #  * they have forms that accept bbox/points args to transform annotations in sync with image
+
+    @property
+    def generator(self) -> random.Random:
+        # (re)seed per dataloader worker so workers don't produce identical random sequences
+        from torch.utils.data import get_worker_info
+        worker_info = get_worker_info()
+        worker_id = worker_info.id if worker_info is not None else -1
+        if worker_id != self._generator_worker:
+            if worker_info is not None:
+                self._generator = random.Random(worker_info.seed + self.seed)
+            self._generator_worker = worker_id
+        return self._generator
+
+    def _preprocess_image_pages(self, decoded_pages, page_image_info=None):
+        if self.image_process_fn is None:
+            return decoded_pages
+
+        if page_image_info is not None:
+            # FIXME, WIP. If train objective involves masking or otherwise processing image
+            #  with knowledge of annotations / text content, anno info should contain
+            #  mask locations, etc. For such a task, we need to pass it to image preprocess
+            decoded_pages = [self.image_process_fn(dp, page_info=pi) for dp, pi in zip(decoded_pages, page_image_info)]
+        else:
+            decoded_pages = [self.image_process_fn(dp) for dp in decoded_pages]
+
+        return decoded_pages
+
+    def _decode_image_pages(
+            self,
+            sample,
+            ext,
+            page_indices,
+            num_anno_pages,
+    ):
+        image_mode = self.image_input_feat.image_mode
+
+        decoded_pages, num_image_pages = decode_image_pages(
+            sample[ext],
+            image_mode=image_mode,
+            page_indices=page_indices,
+        )
+        if num_image_pages != num_anno_pages:
+            _logger.warning(
+                f'Mismatch between num image and num annotation pages {num_image_pages} != {num_anno_pages}'
+                f' for sample {sample["__url__"]}, {sample["__key__"]}.')
+
+        decoded_pages = self._preprocess_image_pages(decoded_pages)
+
+        return decoded_pages, num_image_pages
+
+    def _decode_pdf_pages(
+            self,
+            sample,
+            ext,
+            page_indices,
+            num_anno_pages,
+    ):
+        image_mode = self.image_input_feat.image_mode
+        decoded_pages, num_image_pages = decode_pdf_pages(
+            sample[ext],
+            image_mode=image_mode,
+            page_indices=page_indices,
+            render_dpi=self.render_dpi,
+        )
+        if num_anno_pages is not None and num_image_pages != num_anno_pages:
+            _logger.warning(
+                f'Mismatch between num image and num annotation pages {num_image_pages} != {num_anno_pages}'
+                f' for sample {sample["__url__"]}, {sample["__key__"]}.')
+
+        decoded_pages = self._preprocess_image_pages(decoded_pages)
+
+        return decoded_pages, num_image_pages
+
+    @abc.abstractmethod
+    def _decode_anno(self, sample) -> Tuple[Dict, List[int], int]:
+        pass
+
+    def _expand_anno(self, anno, count: int):
+        expanded_annos = [
+            {k: v[i] if isinstance(v, (list, tuple)) else v for k, v in anno.items()}
+            for i in range(count)
+        ]
+        return expanded_annos
+
+    def _squeeze_anno(self, anno):
+        anno = {k: v[0] if isinstance(v, (list, tuple)) else v for k, v in anno.items()}
+        return anno
+
+    def __call__(self, sample):
+        if 'json' in sample and isinstance(sample['json'], bytes):
+            # decode json if present and in undecoded state
+            sample['json'] = json.loads(sample['json'])
+
+        if self.flatten_json and 'json' in sample:
+            # flatten json into sample
+            sample.update(sample.pop('json'))
+
+        # FIXME separate decode & preprocess interfaces
+
+        # decode page annotations / text
+        page_anno, page_indices, num_anno_pages = self._decode_anno(sample)
+
+        # decode page images
+        page_images = []
+        for ext in self.image_input_key:
+            if ext in sample:
+                if ext == 'pdf':
+                    images, num_image_pages = self._decode_pdf_pages(
+                        sample,
+                        ext,
+                        page_indices,
+                        num_anno_pages,
+                    )
+                else:
+                    images, num_image_pages = self._decode_image_pages(
+                        sample,
+                        ext,
+                        page_indices,
+                        num_anno_pages,
+                    )
+                page_images.extend(images)
+                # process one document type per doc, should be ordered by priority
+                break
+
+        assert len(page_images), 'No page images present'
+
+        if self.expand_pages and len(page_images) > 1:
+            # expand pages and page annotations into multiple samples (return list of sample dicts)
+            page_anno = self._expand_anno(page_anno, len(page_images))
+            decoded = [{self.image_input_name: pi, **pa} for pi, pa in zip(page_images, page_anno)]
+        else:
+            if self.squeeze_pages and len(page_images) == 1:
+                # squeeze page & annotation lists into singular items
+                page_images = page_images[0]
+                page_anno = self._squeeze_anno(page_anno)
+            decoded = {self.image_input_name: page_images, **page_anno}
+
+        return decoded
+
+