kcai-data-sampling-lama 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (24) hide show
  1. kcai_data_sampling_lama-0.1.0/PKG-INFO +126 -0
  2. kcai_data_sampling_lama-0.1.0/README.md +107 -0
  3. kcai_data_sampling_lama-0.1.0/pyproject.toml +44 -0
  4. kcai_data_sampling_lama-0.1.0/requirements.txt +2 -0
  5. kcai_data_sampling_lama-0.1.0/setup.cfg +4 -0
  6. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/__init__.py +12 -0
  7. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/_version_.py +24 -0
  8. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/api/__init__.py +9 -0
  9. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/api/models/__init__.py +10 -0
  10. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/api/models/lama.py +108 -0
  11. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/api/models/weights.py +56 -0
  12. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/api/transformations/__init__.py +12 -0
  13. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/api/transformations/inpaint.py +65 -0
  14. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/configs.py +72 -0
  15. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/py.typed +0 -0
  16. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/transformations/__init__.py +11 -0
  17. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama/transformations/inpaint.py +50 -0
  18. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama.egg-info/PKG-INFO +126 -0
  19. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama.egg-info/SOURCES.txt +22 -0
  20. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama.egg-info/dependency_links.txt +1 -0
  21. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama.egg-info/entry_points.txt +5 -0
  22. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama.egg-info/requires.txt +2 -0
  23. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama.egg-info/scm_file_list.json +18 -0
  24. kcai_data_sampling_lama-0.1.0/src/kcai_data_sampling_lama.egg-info/top_level.txt +1 -0
@@ -0,0 +1,126 @@
1
+ Metadata-Version: 2.4
2
+ Name: kcai-data-sampling-lama
3
+ Version: 0.1.0
4
+ Summary: Generative inpaint transformation (LaMa) and the associated tool-model plugin for kcai data-sampling
5
+ Author-email: Safenai <support@safenai.io>
6
+ License-Expression: Apache-2.0
7
+ Keywords: ml,data,augmentation,robustness,images,inpainting
8
+ Classifier: Development Status :: 3 - Alpha
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
11
+ Classifier: Programming Language :: Python :: 3.11
12
+ Classifier: Programming Language :: Python :: 3.12
13
+ Classifier: Programming Language :: Python :: 3.13
14
+ Classifier: Programming Language :: Python :: 3.14
15
+ Requires-Python: >=3.11
16
+ Description-Content-Type: text/markdown
17
+ Requires-Dist: torch>=2.0.0
18
+ Requires-Dist: kcai-data-sampling-core
19
+
20
+ # kcai-data-sampling-lama
21
+
22
+ The **generative** slot of kcai data-sampling: transformations where a model
23
+ *produces* the output. Here the model is a **tool** — it fills in content that
24
+ was not in the source sample, so the family's defining property is that the
25
+ output is genuinely new rather than a rearrangement of the input.
26
+
27
+ This package implements inpainting with
28
+ [LaMa](https://github.com/advimman/lama) (Suvorov et al., WACV 2022) as a tool
29
+ model, and exposes it through two entry-point groups: `inpaint` under
30
+ `kcai_data_sampling.transformations`, and `lama_inpaint` under
31
+ `kcai_data_sampling.models`.
32
+
33
+ A tool model declares `model_role = "tool"`, which is what makes the framework
34
+ treat it as generative. The map is deterministic — no randomness is drawn, so
35
+ the recorded seed is `None` — and not reversible, because what was in the filled
36
+ region is gone.
37
+
38
+ ## What it ships today
39
+
40
+ - `inpaint` — one method: erase a rectangular region and let the model refill it.
41
+ - `lama_inpaint` — the LaMa tool model behind it.
42
+
43
+ One method, as an example of the slot rather than a catalogue.
44
+
45
+ ## Install
46
+
47
+ ```bash
48
+ pip install "kcai-data-sampling[lama]"
49
+ ```
50
+
51
+ **This is the package that brings torch.** It is the only one of the
52
+ transformation packages that depends on a deep-learning framework, so an install
53
+ that does not ask for it stays torch-free. `torch` is imported lazily — importing
54
+ this package does not itself require torch — and the checkpoint is a weights
55
+ file fetched at run time, not a dependency.
56
+
57
+ The released *big-lama* TorchScript checkpoint (~200 MB) is downloaded into the
58
+ weight cache on first use, from the upstream release. The cache is
59
+ `$KCAI_WEIGHTS_DIR` when that is set; otherwise `.cache/` beside the working
60
+ directory — inside a checkout that resolves to the repository's own `.cache/`.
61
+ Checkpoints are never committed.
62
+
63
+ ## Example
64
+
65
+ The `models` and `operations` sections of a job configuration:
66
+
67
+ ```yaml
68
+ models:
69
+ lama:
70
+ type: lama_inpaint
71
+ weights: big-lama.pt
72
+ params: { margin: 256 }
73
+ operations:
74
+ transformations:
75
+ - name: inpaint
76
+ type: inpaint
77
+ tool_model: lama
78
+ top: 350
79
+ left: 700
80
+ height: 300
81
+ width: 500
82
+ ```
83
+
84
+ The transformation names its model with `tool_model: lama`, referring to the
85
+ key under `models:`. `top`/`left`/`height`/`width` describe the rectangle to
86
+ erase; the window must fit inside the frame, which is only knowable at run time,
87
+ so a window that would leave it is refused rather than silently wrapped.
88
+ `margin` is how many pixels of surrounding context are cropped around the mask
89
+ for the network.
90
+
91
+ Only the channels the network consumes are sent to it — three, RGB — so the
92
+ alpha plane of an RGBA batch is passed through untouched.
93
+
94
+ ## Bring your own model
95
+
96
+ Any object reachable through a `kcai_data_sampling.models` entry point is a
97
+ model, exactly like a transformation. To use a custom inference routine, ship a
98
+ small importable package that registers your adapter under that entry-point group
99
+ and lets it declare its own dependencies (`torch`, your inference library) in its
100
+ own metadata. They never belong to the core contracts or to the job package.
101
+ A classic-sketch adapter to paste into your package:
102
+
103
+ ```python
104
+ class MyModel:
105
+ channels = 3 # RGB only: the adapter receives exactly those planes
106
+
107
+ def __init__(self, weights: str, **params):
108
+ self.name = weights # the ledger records this string
109
+ self._load(weights, **params)
110
+
111
+ def inpaint(self, xs, masks):
112
+ """(B, H, W, C) uint8 image batch in, (B, H, W, C) uint8 out."""
113
+ ... # fill every masked pixel; leave the rest unchanged
114
+
115
+
116
+ def plugin() -> MyModel:
117
+ return MyModel
118
+ ```
119
+
120
+ and point `[project.entry-points."kcai_data_sampling.models"]` at it (e.g.
121
+ `mymodel = "my_package.models:plugin"`). A tool model is recognised by a
122
+ non-empty `name` plus an `inpaint` method, and conformance is checked against
123
+ the `tool` role the transformation declares.
124
+
125
+ Known-architecture custom weights need no code at all: keep `type: lama_inpaint`
126
+ and set `weights` to your own checkpoint file.
@@ -0,0 +1,107 @@
1
+ # kcai-data-sampling-lama
2
+
3
+ The **generative** slot of kcai data-sampling: transformations where a model
4
+ *produces* the output. Here the model is a **tool** — it fills in content that
5
+ was not in the source sample, so the family's defining property is that the
6
+ output is genuinely new rather than a rearrangement of the input.
7
+
8
+ This package implements inpainting with
9
+ [LaMa](https://github.com/advimman/lama) (Suvorov et al., WACV 2022) as a tool
10
+ model, and exposes it through two entry-point groups: `inpaint` under
11
+ `kcai_data_sampling.transformations`, and `lama_inpaint` under
12
+ `kcai_data_sampling.models`.
13
+
14
+ A tool model declares `model_role = "tool"`, which is what makes the framework
15
+ treat it as generative. The map is deterministic — no randomness is drawn, so
16
+ the recorded seed is `None` — and not reversible, because what was in the filled
17
+ region is gone.
18
+
19
+ ## What it ships today
20
+
21
+ - `inpaint` — one method: erase a rectangular region and let the model refill it.
22
+ - `lama_inpaint` — the LaMa tool model behind it.
23
+
24
+ One method, as an example of the slot rather than a catalogue.
25
+
26
+ ## Install
27
+
28
+ ```bash
29
+ pip install "kcai-data-sampling[lama]"
30
+ ```
31
+
32
+ **This is the package that brings torch.** It is the only one of the
33
+ transformation packages that depends on a deep-learning framework, so an install
34
+ that does not ask for it stays torch-free. `torch` is imported lazily — importing
35
+ this package does not itself require torch — and the checkpoint is a weights
36
+ file fetched at run time, not a dependency.
37
+
38
+ The released *big-lama* TorchScript checkpoint (~200 MB) is downloaded into the
39
+ weight cache on first use, from the upstream release. The cache is
40
+ `$KCAI_WEIGHTS_DIR` when that is set; otherwise `.cache/` beside the working
41
+ directory — inside a checkout that resolves to the repository's own `.cache/`.
42
+ Checkpoints are never committed.
43
+
44
+ ## Example
45
+
46
+ The `models` and `operations` sections of a job configuration:
47
+
48
+ ```yaml
49
+ models:
50
+ lama:
51
+ type: lama_inpaint
52
+ weights: big-lama.pt
53
+ params: { margin: 256 }
54
+ operations:
55
+ transformations:
56
+ - name: inpaint
57
+ type: inpaint
58
+ tool_model: lama
59
+ top: 350
60
+ left: 700
61
+ height: 300
62
+ width: 500
63
+ ```
64
+
65
+ The transformation names its model with `tool_model: lama`, referring to the
66
+ key under `models:`. `top`/`left`/`height`/`width` describe the rectangle to
67
+ erase; the window must fit inside the frame, which is only knowable at run time,
68
+ so a window that would leave it is refused rather than silently wrapped.
69
+ `margin` is how many pixels of surrounding context are cropped around the mask
70
+ for the network.
71
+
72
+ Only the channels the network consumes are sent to it — three, RGB — so the
73
+ alpha plane of an RGBA batch is passed through untouched.
74
+
75
+ ## Bring your own model
76
+
77
+ Any object reachable through a `kcai_data_sampling.models` entry point is a
78
+ model, exactly like a transformation. To use a custom inference routine, ship a
79
+ small importable package that registers your adapter under that entry-point group
80
+ and lets it declare its own dependencies (`torch`, your inference library) in its
81
+ own metadata. They never belong to the core contracts or to the job package.
82
+ A classic-sketch adapter to paste into your package:
83
+
84
+ ```python
85
+ class MyModel:
86
+ channels = 3 # RGB only: the adapter receives exactly those planes
87
+
88
+ def __init__(self, weights: str, **params):
89
+ self.name = weights # the ledger records this string
90
+ self._load(weights, **params)
91
+
92
+ def inpaint(self, xs, masks):
93
+ """(B, H, W, C) uint8 image batch in, (B, H, W, C) uint8 out."""
94
+ ... # fill every masked pixel; leave the rest unchanged
95
+
96
+
97
+ def plugin() -> MyModel:
98
+ return MyModel
99
+ ```
100
+
101
+ and point `[project.entry-points."kcai_data_sampling.models"]` at it (e.g.
102
+ `mymodel = "my_package.models:plugin"`). A tool model is recognised by a
103
+ non-empty `name` plus an `inpaint` method, and conformance is checked against
104
+ the `tool` role the transformation declares.
105
+
106
+ Known-architecture custom weights need no code at all: keep `type: lama_inpaint`
107
+ and set `weights` to your own checkpoint file.
@@ -0,0 +1,44 @@
1
+ [project]
2
+ name = "kcai-data-sampling-lama"
3
+ dynamic = ["dependencies", "version"]
4
+ description = "Generative inpaint transformation (LaMa) and the associated tool-model plugin for kcai data-sampling"
5
+ authors = [
6
+ {name = "Safenai", email = "support@safenai.io"},
7
+ ]
8
+ license = "Apache-2.0"
9
+ readme = "README.md"
10
+ requires-python = ">=3.11"
11
+ keywords = ["ml", "data", "augmentation", "robustness", "images", "inpainting"]
12
+ classifiers = [
13
+ "Development Status :: 3 - Alpha",
14
+ "Intended Audience :: Developers",
15
+ "Topic :: Software Development :: Libraries :: Python Modules",
16
+ "Programming Language :: Python :: 3.11",
17
+ "Programming Language :: Python :: 3.12",
18
+ "Programming Language :: Python :: 3.13",
19
+ "Programming Language :: Python :: 3.14",
20
+ ]
21
+
22
+ [project.entry-points."kcai_data_sampling.transformations"]
23
+ inpaint = "kcai_data_sampling_lama.api.transformations.inpaint:Inpaint"
24
+
25
+ [project.entry-points."kcai_data_sampling.models"]
26
+ lama_inpaint = "kcai_data_sampling_lama.api.models.lama:LamaTool"
27
+
28
+ [build-system]
29
+ requires = ["setuptools>=80", "setuptools-scm[simple]>=8"]
30
+ build-backend = "setuptools.build_meta"
31
+
32
+ [tool.setuptools_scm]
33
+ version_file = "src/kcai_data_sampling_lama/_version_.py"
34
+ root = "../.."
35
+ fallback_version = "0.0.1"
36
+
37
+ [tool.setuptools.dynamic]
38
+ dependencies = {file = ["requirements.txt"]}
39
+
40
+ [tool.setuptools.packages.find]
41
+ where = ["src"]
42
+
43
+ [tool.setuptools.package-data]
44
+ kcai_data_sampling_lama = ["py.typed"]
@@ -0,0 +1,2 @@
1
+ torch>=2.0.0
2
+ kcai-data-sampling-core
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,12 @@
1
+ """Generative inpainting with LaMa.
2
+
3
+ Ships the ``inpaint`` transformation and its ``lama_inpaint`` tool model: the
4
+ transformation erases a rectangle of each image and a LaMa artifact fills the
5
+ hole with content that was not in the image. The transformation class lives
6
+ under ``api/transformations/``, the model under ``api/models/``, and the pure
7
+ mask math under ``transformations/``. ``torch`` is needed only here; the
8
+ checkpoint downloads into the weight cache (``$KCAI_WEIGHTS_DIR`` or the
9
+ repository's ``.cache/``) on first use.
10
+ """
11
+
12
+ from kcai_data_sampling_lama._version_ import __version__ as __version__
@@ -0,0 +1,24 @@
1
+ # file generated by vcs-versioning
2
+ # don't change, don't track in version control
3
+ from __future__ import annotations
4
+
5
+ __all__ = [
6
+ "__version__",
7
+ "__version_tuple__",
8
+ "version",
9
+ "version_tuple",
10
+ "__commit_id__",
11
+ "commit_id",
12
+ ]
13
+
14
+ version: str
15
+ __version__: str
16
+ __version_tuple__: tuple[int | str, ...]
17
+ version_tuple: tuple[int | str, ...]
18
+ commit_id: str | None
19
+ __commit_id__: str | None
20
+
21
+ __version__ = version = '0.1.0'
22
+ __version_tuple__ = version_tuple = (0, 1, 0)
23
+
24
+ __commit_id__ = commit_id = None
@@ -0,0 +1,9 @@
1
+ """Model and transformation subclasses owned by this package.
2
+
3
+ The generic contracts live in ``kcai_data_sampling_core.api``; this package
4
+ holds its own model adapter (``api/models/``) and its own transformation class
5
+ (``api/transformations/``), each implementing the generic core contract for
6
+ the LaMa inpainting tool. Nothing here reads or writes images: batches flow in
7
+ as ``(B, H, W, 4)`` uint8 RGBA arrays and the only disk I/O in the package is
8
+ the weight cache.
9
+ """
@@ -0,0 +1,10 @@
1
+ """Tool-model adapters for this package.
2
+
3
+ Adapters implement the generic ``ToolModel`` protocol from
4
+ ``kcai_data_sampling_core.api.roles``; the LaMa adapter is the package's first
5
+ one. The weight-cache helper lives here too.
6
+ """
7
+
8
+ from kcai_data_sampling_lama.api.models.lama import LamaTool
9
+
10
+ __all__ = ["LamaTool"]
@@ -0,0 +1,108 @@
1
+ """LaMa (Suvorov et al., WACV 2022) as a tool model.
2
+
3
+ Wraps the official *big-lama* TorchScript checkpoint (~200 MB, downloaded into
4
+ the weight cache on first use) behind the ``inpaint(xs, masks)`` tool-model
5
+ contract. ``torch`` is imported lazily so that importing this package never
6
+ requires it; ``big-lama`` itself is a weights file fetched at run time, not a
7
+ dependency.
8
+ """
9
+
10
+ import warnings
11
+
12
+ import numpy as np
13
+
14
+ from kcai_data_sampling_lama.api.models.weights import fetch
15
+
16
+ URL = "https://github.com/enesmsahin/simple-lama-inpainting/releases/download/v0.1.0/big-lama.pt"
17
+
18
+
19
+ class LamaTool:
20
+ """The ``inpaint`` tool model backed by LaMa.
21
+
22
+ Attributes:
23
+ name: The checkpoint's file stem (``"big-lama"``), recorded on every
24
+ output row.
25
+ weights: The local path of the checkpoint in the weight cache.
26
+ margin: Pixels of context cropped around the mask union for the network.
27
+ channels: Image planes the network consumes (``3`` = RGB). The alpha
28
+ plane of an RGBA batch is then never sent to the network and is
29
+ returned untouched.
30
+ """
31
+
32
+ channels = 3
33
+
34
+ def __init__(self, weights: str = "big-lama.pt", margin: int = 256):
35
+ """Load the LaMa checkpoint.
36
+
37
+ Args:
38
+ weights: The checkpoint's file name in the weight cache; downloaded
39
+ from the upstream release on first use.
40
+ margin: Pixels of context cropped around the mask union for the
41
+ network (sides are padded to multiples of 8).
42
+
43
+ Raises:
44
+ ImportError: If ``torch`` is not installed.
45
+ """
46
+ import torch
47
+
48
+ self._torch = torch
49
+ self.weights = fetch(weights, URL)
50
+ self.name = self.weights.stem
51
+ self.margin = margin
52
+ with warnings.catch_warnings(): # TorchScript checkpoint; torch >= 2.9 deprecates jit.load
53
+ warnings.simplefilter("ignore", FutureWarning)
54
+ self.net = torch.jit.load(str(self.weights), map_location="cpu").eval()
55
+
56
+ def inpaint(self, xs: np.ndarray, masks: np.ndarray) -> np.ndarray:
57
+ """Fill every masked pixel with network output, leave the rest untouched.
58
+
59
+ The ``(B, H, W, 4)`` uint8 RGBA batch is converted to the network's
60
+ NCHW float input (the ``channels`` first planes, in ``[0, 1]``) and back.
61
+ The network runs once, on a crop around the union of the masks plus
62
+ ``margin`` pixels of context (sides padded to multiples of 8); only the
63
+ masked pixels of the selected planes are rewritten.
64
+
65
+ Args:
66
+ xs: Batch of image arrays shaped ``(B, H, W, 4)``, uint8 RGBA.
67
+ masks: One boolean mask per sample ``(B, H, W)``; True = erase and fill in.
68
+
69
+ Returns:
70
+ A batch of the same shape, uint8 in ``[0, 255]``. Unmasked pixels
71
+ and (when ``channels`` is 3) the whole alpha plane are exactly as
72
+ in ``xs``.
73
+
74
+ Raises:
75
+ ValueError: If ``xs`` does not carry at least ``channels`` planes.
76
+ """
77
+ if xs.ndim != 4 or xs.shape[-1] < self.channels:
78
+ raise ValueError(f"LamaTool.inpaint expects a (B, H, W, 4) uint8 RGBA batch, got shape {xs.shape}")
79
+ rows, cols = np.where(masks.any(axis=0))
80
+ if rows.size == 0:
81
+ return xs.copy()
82
+ h, w = masks.shape[-2:]
83
+ top, bottom = max(rows.min() - self.margin, 0), min(rows.max() + 1 + self.margin, h)
84
+ left, right = max(cols.min() - self.margin, 0), min(cols.max() + 1 + self.margin, w)
85
+ pad_h, pad_w = (-(bottom - top)) % 8, (-(right - left)) % 8
86
+
87
+ image = np.pad(
88
+ xs[:, top:bottom, left:right, : self.channels],
89
+ ((0, 0), (0, pad_h), (0, pad_w), (0, 0)),
90
+ mode="symmetric",
91
+ )
92
+ image = image.astype("float32").transpose(0, 3, 1, 2) / 255.0
93
+ hole = np.pad(
94
+ masks[:, top:bottom, left:right],
95
+ ((0, 0), (0, pad_h), (0, pad_w)),
96
+ mode="symmetric",
97
+ ).astype("float32")
98
+ torch = self._torch
99
+ with torch.inference_mode():
100
+ out = self.net(torch.from_numpy(image), torch.from_numpy(hole)[:, None])
101
+ filled = out[:, :, : bottom - top, : right - left].numpy()
102
+
103
+ result = xs.copy()
104
+ region = result[:, top:bottom, left:right, : self.channels]
105
+ hole_pixels = np.broadcast_to(masks[:, top:bottom, left:right, None], region.shape)
106
+ filled_pixels = np.transpose(filled, (0, 2, 3, 1))[hole_pixels]
107
+ region[hole_pixels] = np.clip(filled_pixels * 255.0, 0.0, 255.0).astype(np.uint8)
108
+ return result
@@ -0,0 +1,56 @@
1
+ """Where model checkpoints live: ``$KCAI_WEIGHTS_DIR`` or the repository's ``.cache/``.
2
+
3
+ The default cache directory is the repo-relative ``.cache/`` (discovered by
4
+ walking up from the working directory to the nearest ancestor that contains a
5
+ ``.git`` directory, falling back to the working directory), so a checkout's
6
+ checkpoints stay together with its code. ``$KCAI_WEIGHTS_DIR`` overrides it for
7
+ shared or pre-fetched caches.
8
+ """
9
+
10
+ import os
11
+ from pathlib import Path
12
+ import urllib.request
13
+
14
+
15
+ def weights_path(name: str) -> Path:
16
+ """The cache path of a checkpoint, creating the directory if needed.
17
+
18
+ Args:
19
+ name: The checkpoint's file name inside the cache directory.
20
+
21
+ Returns:
22
+ The absolute path where the checkpoint is (or will be) cached.
23
+ """
24
+ directory = Path(os.environ.get("KCAI_WEIGHTS_DIR", _default_cache_dir()))
25
+ directory.mkdir(parents=True, exist_ok=True)
26
+ return directory / name
27
+
28
+
29
+ def fetch(name: str, url: str) -> Path:
30
+ """The checkpoint's path, downloaded from ``url`` on first use.
31
+
32
+ Args:
33
+ name: The checkpoint's file name inside the cache directory.
34
+ url: The remote location to download from when the file is missing.
35
+
36
+ Returns:
37
+ The path the checkpoint was cached at.
38
+ """
39
+ path = weights_path(name)
40
+ if not path.exists():
41
+ urllib.request.urlretrieve(url, path)
42
+ return path
43
+
44
+
45
+ def _default_cache_dir() -> Path:
46
+ """The repository root's ``.cache/`` (nearest ancestor with a ``.git`` dir).
47
+
48
+ Returns:
49
+ ``<repo root>/.cache``, or ``<cwd>/.cache`` when no ancestor has a
50
+ ``.git`` directory.
51
+ """
52
+ cwd = Path.cwd()
53
+ for candidate in (cwd, *cwd.parents):
54
+ if (candidate / ".git").exists():
55
+ return candidate / ".cache"
56
+ return cwd / ".cache"
@@ -0,0 +1,12 @@
1
+ """Unary transformation classes in this package.
2
+
3
+ The generic core contract lives in ``kcai_data_sampling_core.api``; the class
4
+ here (``api.transformations.inpaint.Inpaint``) is the package's own subclass,
5
+ wrapping a pure algorithm function from
6
+ ``kcai_data_sampling_lama.transformations`` into the unary ``apply``
7
+ contract.
8
+ """
9
+
10
+ from kcai_data_sampling_lama.api.transformations.inpaint import Inpaint
11
+
12
+ __all__ = ["Inpaint"]
@@ -0,0 +1,65 @@
1
+ """The ``inpaint`` transformation class.
2
+
3
+ Wraps the region-mask math
4
+ (:func:`kcai_data_sampling_lama.transformations.inpaint.build_region_mask`)
5
+ into the unary transformation contract: each sample's window is erased and the
6
+ tool model fills it with content that was not in the image.
7
+ """
8
+
9
+ from kcai_data_sampling_core.api.roles import check_output
10
+ from kcai_data_sampling_core.api.unary import UnaryTransformation
11
+ import numpy as np
12
+ from typing_extensions import override
13
+
14
+ from kcai_data_sampling_lama.configs import InpaintTransformationConfig
15
+ from kcai_data_sampling_lama.transformations.inpaint import build_region_mask
16
+
17
+
18
+ class Inpaint(UnaryTransformation):
19
+ """Erase a rectangle of each image and let a tool model fill it in.
20
+
21
+ ``top``/``left`` default to 0; ``height``/``width`` are required (a
22
+ degenerate empty region is refused). The replacement is invented by the
23
+ ``tool_model``, so the map is not reversible: what was in the region is
24
+ gone. ``apply`` is deterministic (``stochastic`` not set), so the seed
25
+ records ``None``.
26
+ """
27
+
28
+ algorithm = "inpaint"
29
+
30
+ #: The registered config schema this algorithm validates against.
31
+ Config = InpaintTransformationConfig
32
+
33
+ #: The model slot this algorithm fills, and the methods it must expose.
34
+ model_role = "tool"
35
+ model_methods = ("inpaint",)
36
+
37
+ @override
38
+ def apply(
39
+ self,
40
+ xs: np.ndarray, # (b, h, w, c) uint8
41
+ rngs: list[np.random.Generator] | None = None,
42
+ ) -> np.ndarray:
43
+ """Erase the configured rectangle and fill it via the tool model.
44
+
45
+ Args:
46
+ xs: Batch of image arrays shaped ``(B, H, W, C)``.
47
+ rngs: Unused for this deterministic operation.
48
+
49
+ Returns:
50
+ A batch of the same shape; the masked rectangle is rewritten from
51
+ the tool model's output, every other pixel returns untouched.
52
+
53
+ Raises:
54
+ ValueError: If the region leaves the frame, or the tool model's
55
+ output breaks the numeric contract.
56
+ """
57
+ del rngs
58
+ masks = build_region_mask(xs.shape, **self.params)
59
+ return check_output(
60
+ self.algorithm,
61
+ self.tool_model,
62
+ "inpaint",
63
+ xs,
64
+ self.tool_model.inpaint(xs, masks),
65
+ )
@@ -0,0 +1,72 @@
1
+ """Per-algorithm config schemas for the inpaint transformation.
2
+
3
+ The schema subclasses the core base ``TransformationConfig`` and pins ``type``
4
+ to the algorithm's literal, so the registry-resolved validator in ``JobConfig``
5
+ picks it by `type`. A parameter that can sweep accepts ``SweepConfig``; the
6
+ ranges the algorithm allows are enforced here, on both the plain value and the
7
+ sweep interval (``top``/``left`` within ``[0, +oo)``, ``height``/``width``
8
+ within ``(0, +oo)``). ``tool_model`` names a model from the ``models:`` section
9
+ of the job config; the CLI resolves the name to an instance.
10
+ """
11
+
12
+ from typing import Literal, Self
13
+
14
+ from kcai_data_sampling_core.models.config import TransformationConfig
15
+ from kcai_data_sampling_core.models.sweep import SweepConfig
16
+ from pydantic import Field, model_validator
17
+
18
+
19
+ class InpaintTransformationConfig(TransformationConfig):
20
+ """Configuration of the ``inpaint`` transformation.
21
+
22
+ Attributes:
23
+ tool_model: Name of a tool model from the job's ``models:`` section;
24
+ the CLI resolves it to an instance. Optional here because the
25
+ transformation base consumes it before re-validating the remaining
26
+ parameters against this schema; a model-role transformation without
27
+ one is refused loudly at construction (the slot check).
28
+ top: Region top offset in pixels (default 0), or a sweep (its interval
29
+ must be ``>= 0``).
30
+ left: Region left offset in pixels (default 0), or a sweep (its
31
+ interval must be ``>= 0``).
32
+ height: Region height in pixels; required (a degenerate empty region is
33
+ refused), or a sweep (its interval must be ``> 0``).
34
+ width: Region width in pixels; required (a degenerate empty region is
35
+ refused), or a sweep (its interval must be ``> 0``).
36
+ """
37
+
38
+ type: Literal["inpaint"] = "inpaint"
39
+ tool_model: str | None = Field(
40
+ default=None,
41
+ exclude=True,
42
+ description="Name of a tool model from the job's models: section; excluded"
43
+ " from the resolved parameters (the base consumes it before validation).",
44
+ )
45
+ top: int | SweepConfig = Field(
46
+ default=0,
47
+ description="Region top offset in pixels; a SweepConfig expands it.",
48
+ )
49
+ left: int | SweepConfig = Field(
50
+ default=0,
51
+ description="Region left offset in pixels; a SweepConfig expands it.",
52
+ )
53
+ height: int | SweepConfig = Field(description="Region height in pixels; a SweepConfig expands it.")
54
+ width: int | SweepConfig = Field(description="Region width in pixels; a SweepConfig expands it.")
55
+
56
+ @model_validator(mode="after")
57
+ def _parameter_bounds(self) -> Self:
58
+ """Enforce each parameter's algorithm range on the value or the sweep.
59
+
60
+ Returns:
61
+ The validated config.
62
+
63
+ Raises:
64
+ ValueError: If ``top``/``left`` go below ``0`` or
65
+ ``height``/``width`` do not stay positive, whether given
66
+ directly or as a sweep interval.
67
+ """
68
+ for name in ("top", "left"):
69
+ self._check_parameter_bounds(name, minimum=0)
70
+ for name in ("height", "width"):
71
+ self._check_parameter_bounds(name, minimum=0, exclusive_min=True)
72
+ return self
@@ -0,0 +1,11 @@
1
+ """Transformation math functions for this package.
2
+
3
+ The module-level function here (``transformations.inpaint.build_region_mask``)
4
+ is the pure mask math; the transformation class that wraps it into the unary
5
+ contract lives in ``kcai_data_sampling_lama.api.transformations``. All
6
+ batches share the RGBA-U8 array convention ``(B, H, W, 4) uint8``.
7
+ """
8
+
9
+ from kcai_data_sampling_lama.transformations.inpaint import build_region_mask
10
+
11
+ __all__ = ["build_region_mask"]
@@ -0,0 +1,50 @@
1
+ """Inpainting: generate a region fill — the mask math.
2
+
3
+ ``build_region_mask`` is the pure function that marks the rectangle a tool
4
+ model should erase and refill; the transformation class that wraps it into the
5
+ unary contract lives in ``kcai_data_sampling_lama.api.transformations.inpaint``.
6
+ """
7
+
8
+ import numpy as np
9
+
10
+
11
+ def build_region_mask(
12
+ xs_shape: tuple[int, ...],
13
+ top: int,
14
+ left: int,
15
+ height: int,
16
+ width: int,
17
+ ) -> np.ndarray:
18
+ """The boolean ``(B, H, W)`` mask of the rectangle to erase.
19
+
20
+ ``True`` marks every pixel of the ``top:top+height`` by
21
+ ``left:left+width`` window of each sample. The window must fit inside the
22
+ frame: the frame size is only known at run time (from the ``(H, W)`` axes
23
+ of ``xs_shape``), so this check is apply-time, and it refuses out-of-frame
24
+ windows loudly instead of silently wrapping.
25
+
26
+ Args:
27
+ xs_shape: The batch shape ``(B, H, W, C)``; the frame is read from the
28
+ ``(H, W)`` axes.
29
+ top: Window top offset in pixels, ``>= 0``.
30
+ left: Window left offset in pixels, ``>= 0``.
31
+ height: Window height in pixels, ``> 0``.
32
+ width: Window width in pixels, ``> 0``.
33
+
34
+ Returns:
35
+ A boolean ``(B, H, W)`` array; ``True`` inside the window.
36
+
37
+ Raises:
38
+ ValueError: If ``top``/``left`` are negative or ``height``/``width``
39
+ are not positive, or if the window leaves the frame.
40
+ """
41
+ if top < 0 or left < 0 or height <= 0 or width <= 0:
42
+ raise ValueError(
43
+ f"inpaint region needs top/left >= 0 and height/width > 0, got {top=}, {left=}, {height=}, {width=}"
44
+ )
45
+ frame_h, frame_w = int(xs_shape[1]), int(xs_shape[2])
46
+ if top + height > frame_h or left + width > frame_w:
47
+ raise ValueError(f"inpaint region {height}x{width} at ({top},{left}) leaves the {frame_h}x{frame_w} frame")
48
+ masks = np.zeros((int(xs_shape[0]), frame_h, frame_w), dtype=bool)
49
+ masks[:, top : top + height, left : left + width] = True
50
+ return masks
@@ -0,0 +1,126 @@
1
+ Metadata-Version: 2.4
2
+ Name: kcai-data-sampling-lama
3
+ Version: 0.1.0
4
+ Summary: Generative inpaint transformation (LaMa) and the associated tool-model plugin for kcai data-sampling
5
+ Author-email: Safenai <support@safenai.io>
6
+ License-Expression: Apache-2.0
7
+ Keywords: ml,data,augmentation,robustness,images,inpainting
8
+ Classifier: Development Status :: 3 - Alpha
9
+ Classifier: Intended Audience :: Developers
10
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
11
+ Classifier: Programming Language :: Python :: 3.11
12
+ Classifier: Programming Language :: Python :: 3.12
13
+ Classifier: Programming Language :: Python :: 3.13
14
+ Classifier: Programming Language :: Python :: 3.14
15
+ Requires-Python: >=3.11
16
+ Description-Content-Type: text/markdown
17
+ Requires-Dist: torch>=2.0.0
18
+ Requires-Dist: kcai-data-sampling-core
19
+
20
+ # kcai-data-sampling-lama
21
+
22
+ The **generative** slot of kcai data-sampling: transformations where a model
23
+ *produces* the output. Here the model is a **tool** — it fills in content that
24
+ was not in the source sample, so the family's defining property is that the
25
+ output is genuinely new rather than a rearrangement of the input.
26
+
27
+ This package implements inpainting with
28
+ [LaMa](https://github.com/advimman/lama) (Suvorov et al., WACV 2022) as a tool
29
+ model, and exposes it through two entry-point groups: `inpaint` under
30
+ `kcai_data_sampling.transformations`, and `lama_inpaint` under
31
+ `kcai_data_sampling.models`.
32
+
33
+ A tool model declares `model_role = "tool"`, which is what makes the framework
34
+ treat it as generative. The map is deterministic — no randomness is drawn, so
35
+ the recorded seed is `None` — and not reversible, because what was in the filled
36
+ region is gone.
37
+
38
+ ## What it ships today
39
+
40
+ - `inpaint` — one method: erase a rectangular region and let the model refill it.
41
+ - `lama_inpaint` — the LaMa tool model behind it.
42
+
43
+ One method, as an example of the slot rather than a catalogue.
44
+
45
+ ## Install
46
+
47
+ ```bash
48
+ pip install "kcai-data-sampling[lama]"
49
+ ```
50
+
51
+ **This is the package that brings torch.** It is the only one of the
52
+ transformation packages that depends on a deep-learning framework, so an install
53
+ that does not ask for it stays torch-free. `torch` is imported lazily — importing
54
+ this package does not itself require torch — and the checkpoint is a weights
55
+ file fetched at run time, not a dependency.
56
+
57
+ The released *big-lama* TorchScript checkpoint (~200 MB) is downloaded into the
58
+ weight cache on first use, from the upstream release. The cache is
59
+ `$KCAI_WEIGHTS_DIR` when that is set; otherwise `.cache/` beside the working
60
+ directory — inside a checkout that resolves to the repository's own `.cache/`.
61
+ Checkpoints are never committed.
62
+
63
+ ## Example
64
+
65
+ The `models` and `operations` sections of a job configuration:
66
+
67
+ ```yaml
68
+ models:
69
+ lama:
70
+ type: lama_inpaint
71
+ weights: big-lama.pt
72
+ params: { margin: 256 }
73
+ operations:
74
+ transformations:
75
+ - name: inpaint
76
+ type: inpaint
77
+ tool_model: lama
78
+ top: 350
79
+ left: 700
80
+ height: 300
81
+ width: 500
82
+ ```
83
+
84
+ The transformation names its model with `tool_model: lama`, referring to the
85
+ key under `models:`. `top`/`left`/`height`/`width` describe the rectangle to
86
+ erase; the window must fit inside the frame, which is only knowable at run time,
87
+ so a window that would leave it is refused rather than silently wrapped.
88
+ `margin` is how many pixels of surrounding context are cropped around the mask
89
+ for the network.
90
+
91
+ Only the channels the network consumes are sent to it — three, RGB — so the
92
+ alpha plane of an RGBA batch is passed through untouched.
93
+
94
+ ## Bring your own model
95
+
96
+ Any object reachable through a `kcai_data_sampling.models` entry point is a
97
+ model, exactly like a transformation. To use a custom inference routine, ship a
98
+ small importable package that registers your adapter under that entry-point group
99
+ and lets it declare its own dependencies (`torch`, your inference library) in its
100
+ own metadata. They never belong to the core contracts or to the job package.
101
+ A classic-sketch adapter to paste into your package:
102
+
103
+ ```python
104
+ class MyModel:
105
+ channels = 3 # RGB only: the adapter receives exactly those planes
106
+
107
+ def __init__(self, weights: str, **params):
108
+ self.name = weights # the ledger records this string
109
+ self._load(weights, **params)
110
+
111
+ def inpaint(self, xs, masks):
112
+ """(B, H, W, C) uint8 image batch in, (B, H, W, C) uint8 out."""
113
+ ... # fill every masked pixel; leave the rest unchanged
114
+
115
+
116
+ def plugin() -> MyModel:
117
+ return MyModel
118
+ ```
119
+
120
+ and point `[project.entry-points."kcai_data_sampling.models"]` at it (e.g.
121
+ `mymodel = "my_package.models:plugin"`). A tool model is recognised by a
122
+ non-empty `name` plus an `inpaint` method, and conformance is checked against
123
+ the `tool` role the transformation declares.
124
+
125
+ Known-architecture custom weights need no code at all: keep `type: lama_inpaint`
126
+ and set `weights` to your own checkpoint file.
@@ -0,0 +1,22 @@
1
+ README.md
2
+ pyproject.toml
3
+ requirements.txt
4
+ src/kcai_data_sampling_lama/__init__.py
5
+ src/kcai_data_sampling_lama/_version_.py
6
+ src/kcai_data_sampling_lama/configs.py
7
+ src/kcai_data_sampling_lama/py.typed
8
+ src/kcai_data_sampling_lama.egg-info/PKG-INFO
9
+ src/kcai_data_sampling_lama.egg-info/SOURCES.txt
10
+ src/kcai_data_sampling_lama.egg-info/dependency_links.txt
11
+ src/kcai_data_sampling_lama.egg-info/entry_points.txt
12
+ src/kcai_data_sampling_lama.egg-info/requires.txt
13
+ src/kcai_data_sampling_lama.egg-info/scm_file_list.json
14
+ src/kcai_data_sampling_lama.egg-info/top_level.txt
15
+ src/kcai_data_sampling_lama/api/__init__.py
16
+ src/kcai_data_sampling_lama/api/models/__init__.py
17
+ src/kcai_data_sampling_lama/api/models/lama.py
18
+ src/kcai_data_sampling_lama/api/models/weights.py
19
+ src/kcai_data_sampling_lama/api/transformations/__init__.py
20
+ src/kcai_data_sampling_lama/api/transformations/inpaint.py
21
+ src/kcai_data_sampling_lama/transformations/__init__.py
22
+ src/kcai_data_sampling_lama/transformations/inpaint.py
@@ -0,0 +1,5 @@
1
+ [kcai_data_sampling.models]
2
+ lama_inpaint = kcai_data_sampling_lama.api.models.lama:LamaTool
3
+
4
+ [kcai_data_sampling.transformations]
5
+ inpaint = kcai_data_sampling_lama.api.transformations.inpaint:Inpaint
@@ -0,0 +1,2 @@
1
+ torch>=2.0.0
2
+ kcai-data-sampling-core
@@ -0,0 +1,18 @@
1
+ {
2
+ "files": [
3
+ "README.md",
4
+ "pyproject.toml",
5
+ "requirements.txt",
6
+ "src/kcai_data_sampling_lama/__init__.py",
7
+ "src/kcai_data_sampling_lama/api/__init__.py",
8
+ "src/kcai_data_sampling_lama/api/models/__init__.py",
9
+ "src/kcai_data_sampling_lama/api/models/lama.py",
10
+ "src/kcai_data_sampling_lama/api/models/weights.py",
11
+ "src/kcai_data_sampling_lama/api/transformations/__init__.py",
12
+ "src/kcai_data_sampling_lama/api/transformations/inpaint.py",
13
+ "src/kcai_data_sampling_lama/configs.py",
14
+ "src/kcai_data_sampling_lama/py.typed",
15
+ "src/kcai_data_sampling_lama/transformations/__init__.py",
16
+ "src/kcai_data_sampling_lama/transformations/inpaint.py"
17
+ ]
18
+ }