todd-tasks 0.7.1__py3-none-any.whl

todd_tasks/__init__.py

@@ -0,0 +1 @@
+from .registries import *

todd_tasks/image_classification/__init__.py

@@ -0,0 +1 @@
+from . import datasets, models

todd_tasks/image_classification/datasets/__init__.py

@@ -0,0 +1,2 @@
+from .imagenet import *
+from .satin import *

todd_tasks/image_classification/datasets/imagenet.py

@@ -0,0 +1,125 @@
+__all__ = [
+    'ImageNetDataset',
+]
+
+import json
+import os
+import pathlib
+from abc import ABC
+from typing import Literal, TypedDict
+
+import torch
+
+from todd_torch.datasets import KeysProtocol, PILDataset
+from todd_torch.datasets.access_layers import PILAccessLayer
+from todd_torch.registries import DatasetRegistry
+
+Split = Literal['train', 'val']
+
+
+class Synset(TypedDict):
+    WNID: str
+    words: str
+    gloss: str
+    num_children: int
+    children: list[int]
+    wordnet_height: int
+    num_train_images: int
+
+
+Synsets = dict[int, Synset]
+
+
+class Annotation(TypedDict):
+    name: str
+    synset_id: int
+
+
+Annotations = list[Annotation]
+
+
+class Keys(KeysProtocol[str]):  # pylint: disable=unsubscriptable-object
+
+    def __init__(
+        self,
+        annotations: Annotations,
+        synsets: Synsets,
+        suffix: str,
+    ) -> None:
+        self._annotations = annotations
+        self._synsets = synsets
+        self._suffix = suffix
+
+    def __len__(self) -> int:
+        return len(self._annotations)
+
+    def __getitem__(self, index: int) -> str:
+        annotation = self._annotations[index]
+        return os.path.join(
+            self._synsets[annotation['synset_id']]['WNID'],
+            annotation['name'].removesuffix(f'.{self._suffix}'),
+        )
+
+
+class T(TypedDict):
+    id_: str
+    image: torch.Tensor
+    category: int
+
+
+@DatasetRegistry.register_()
+class ImageNetDataset(PILDataset[T], ABC):
+    DATA_ROOT = pathlib.Path('data/imagenet')
+    ANNOTATIONS_ROOT = DATA_ROOT / 'annotations'
+    SYNSETS_FILE = DATA_ROOT / 'synsets.json'
+    SUFFIX = 'JPEG'
+
+    def __init__(
+        self,
+        *args,
+        split: Split,
+        access_layer: PILAccessLayer | None = None,
+        annotations_file: pathlib.Path | str | None = None,
+        synsets_file: pathlib.Path | str | None = None,
+        **kwargs,
+    ) -> None:
+        if access_layer is None:
+            access_layer = PILAccessLayer(
+                data_root=str(self.DATA_ROOT),
+                task_name=split,
+                subfolder_action='walk',
+                suffix=self.SUFFIX,
+            )
+        if annotations_file is None:
+            annotations_file = self.ANNOTATIONS_ROOT / f'{split}.json'
+        elif isinstance(annotations_file, str):
+            annotations_file = pathlib.Path(annotations_file)
+        if synsets_file is None:
+            synsets_file = self.SYNSETS_FILE
+        elif isinstance(synsets_file, str):
+            synsets_file = pathlib.Path(synsets_file)
+
+        with annotations_file.open() as f:
+            self._annotations: Annotations = json.load(f)
+
+        with self.SYNSETS_FILE.open() as f:
+            synsets: dict[str, Synset] = json.load(f)
+        synsets_: Synsets = {int(k): v for k, v in synsets.items()}
+        self._synsets = synsets_
+
+        self._categories = {
+            synset_id: i
+            for i, synset_id in enumerate(sorted(synsets_))
+        }
+
+        super().__init__(*args, access_layer=access_layer, **kwargs)
+
+    def build_keys(self) -> Keys:
+        return Keys(self._annotations, self._synsets, self.SUFFIX)
+
+    def __getitem__(self, index: int) -> T:
+        key, image = self._access(index)
+        tensor = self._transform(image)
+        synset_id = self._annotations[index]['synset_id']
+        category = self._categories[synset_id]
+        return T(id_=key, image=tensor, category=category)

todd_tasks/image_classification/datasets/satin.py

@@ -0,0 +1,81 @@
+__all__ = [
+    'SATINDataset',
+]
+
+import io
+import pathlib
+from typing import Any, Literal, TypedDict
+
+import datasets
+import torch
+import torchvision.transforms.functional as F
+from PIL import Image
+
+from todd import Config
+from todd_torch.datasets import AccessLayerRegistry, BaseDataset, IndexKeys
+from todd_torch.datasets.access_layers import HFAccessLayer
+from todd_torch.patches.pil import convert_rgb
+from todd_torch.registries import DatasetRegistry
+
+
+class T(TypedDict):
+    id_: int
+    image: torch.Tensor
+    data: dict[str, Any]
+
+
+Split = Literal['SAT-4', 'SAT-6', 'NASC-TG2', 'WHU-RS19', 'RSSCN7', 'RS_C11',
+                'SIRI-WHU', 'EuroSAT', 'NWPU-RESISC45', 'PatternNet',
+                'RSD46-WHU', 'GID', 'CLRS', 'Optimal-31',
+                'Airbus-Wind-Turbines-Patches', 'USTC_SmokeRS',
+                'Canadian_Cropland', 'Ships-In-Satellite-Imagery',
+                'Satellite-Images-of-Hurricane-Damage',
+                'Brazilian_Coffee_Scenes', 'Brazilian_Cerrado-Savanna_Scenes',
+                'Million-AID', 'UC_Merced_LandUse_MultiLabel', 'MLRSNet',
+                'MultiScene', 'RSI-CB256', 'AID_MultiLabel']
+
+
+@DatasetRegistry.register_()
+class SATINDataset(BaseDataset[T, int, dict[str, Any]]):
+    DATA_ROOT = pathlib.Path('data/satin')
+
+    def __init__(
+        self,
+        *args,
+        split: Split,
+        access_layer: HFAccessLayer | None = None,
+        **kwargs,
+    ) -> None:
+        if access_layer is None:
+            access_layer = AccessLayerRegistry.build(
+                Config(
+                    type=HFAccessLayer.__name__,
+                    data_root=str(self.DATA_ROOT),
+                    task_name=str(datasets.Split.TRAIN),
+                    datasets=dict(
+                        path='jonathan-roberts1/satin',
+                        name=split,
+                        trust_remote_code=True,
+                    ),
+                ),
+            )
+
+        super().__init__(*args, access_layer=access_layer, **kwargs)
+        self._split = split
+
+    def build_keys(self) -> IndexKeys:
+        return IndexKeys(len(self._access_layer))
+
+    def _transform(self, image: Image.Image) -> torch.Tensor:
+        if self._transforms is None:
+            return F.pil_to_tensor(image)
+        return self._transforms(image)
+
+    def __getitem__(self, index: int) -> T:
+        key, data = self._access(index)
+        image = data.pop('image')
+        if not isinstance(image, Image.Image):
+            image = Image.open(io.BytesIO(image['bytes']))
+        image = convert_rgb(image)
+        tensor = self._transform(image)
+        return T(id_=key, image=tensor, data=data)

todd_tasks/image_classification/models/__init__.py

@@ -0,0 +1 @@
+from .ram import *