patchcraft 0.2.0__py3-none-any.whl
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.
- patchcraft/__init__.py +39 -0
- patchcraft/cache.py +247 -0
- patchcraft/extract.py +116 -0
- patchcraft/geometry.py +290 -0
- patchcraft/metrics.py +151 -0
- patchcraft/pair.py +161 -0
- patchcraft/py.typed +0 -0
- patchcraft/reconstruct.py +118 -0
- patchcraft/resize.py +189 -0
- patchcraft/stitch.py +215 -0
- patchcraft-0.2.0.dist-info/METADATA +266 -0
- patchcraft-0.2.0.dist-info/RECORD +14 -0
- patchcraft-0.2.0.dist-info/WHEEL +4 -0
- patchcraft-0.2.0.dist-info/licenses/LICENSE +21 -0
patchcraft/__init__.py
ADDED
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
"""PatchCraft — image patch extraction, pairing and reconstruction utilities."""
|
|
2
|
+
|
|
3
|
+
from patchcraft.cache import Cache
|
|
4
|
+
from patchcraft.extract import Patchify, extract
|
|
5
|
+
from patchcraft.geometry import (
|
|
6
|
+
PairedTilingSpec,
|
|
7
|
+
TilingSpec,
|
|
8
|
+
num_patches,
|
|
9
|
+
paired_tilings,
|
|
10
|
+
scale_factor,
|
|
11
|
+
tilings,
|
|
12
|
+
)
|
|
13
|
+
from patchcraft.metrics import patch_metrics, per_patch_mse, per_patch_psnr
|
|
14
|
+
from patchcraft.pair import PatchMeta, PatchPair, pair
|
|
15
|
+
from patchcraft.reconstruct import reconstruct
|
|
16
|
+
from patchcraft.resize import resize
|
|
17
|
+
from patchcraft.stitch import stitch
|
|
18
|
+
|
|
19
|
+
__version__ = "0.2.0"
|
|
20
|
+
__all__ = [
|
|
21
|
+
"Cache",
|
|
22
|
+
"PairedTilingSpec",
|
|
23
|
+
"PatchMeta",
|
|
24
|
+
"PatchPair",
|
|
25
|
+
"Patchify",
|
|
26
|
+
"TilingSpec",
|
|
27
|
+
"extract",
|
|
28
|
+
"num_patches",
|
|
29
|
+
"pair",
|
|
30
|
+
"paired_tilings",
|
|
31
|
+
"patch_metrics",
|
|
32
|
+
"per_patch_mse",
|
|
33
|
+
"per_patch_psnr",
|
|
34
|
+
"reconstruct",
|
|
35
|
+
"resize",
|
|
36
|
+
"scale_factor",
|
|
37
|
+
"stitch",
|
|
38
|
+
"tilings",
|
|
39
|
+
]
|
patchcraft/cache.py
ADDED
|
@@ -0,0 +1,247 @@
|
|
|
1
|
+
"""Content-addressed disk cache.
|
|
2
|
+
|
|
3
|
+
Bytes in, bytes out. Optional ``zstandard`` compression (transparent fallback
|
|
4
|
+
when not installed). Retry on transient ``PermissionError`` from OneDrive,
|
|
5
|
+
antivirus, or the Windows Search indexer. Atomic write via ``*.tmp`` plus
|
|
6
|
+
``os.replace``. Sidecar JSON carries the full key, version, content checksum.
|
|
7
|
+
|
|
8
|
+
Contract: docs/THEORY.md §4 and §9.5.
|
|
9
|
+
"""
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
import hashlib
|
|
13
|
+
import json
|
|
14
|
+
import os
|
|
15
|
+
import time
|
|
16
|
+
from collections.abc import Callable
|
|
17
|
+
from pathlib import Path
|
|
18
|
+
from typing import Any
|
|
19
|
+
|
|
20
|
+
__all__ = ["Cache"]
|
|
21
|
+
|
|
22
|
+
_WRITE_BACKOFF: tuple[float, ...] = (0.25, 0.5, 1.0, 2.0, 4.0)
|
|
23
|
+
_READ_BACKOFF: tuple[float, ...] = (0.25,)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def _retry[T](
|
|
27
|
+
op: Callable[[], T],
|
|
28
|
+
backoff: tuple[float, ...],
|
|
29
|
+
) -> T:
|
|
30
|
+
"""Retry ``op`` on ``PermissionError`` with the given backoff schedule.
|
|
31
|
+
|
|
32
|
+
First attempt is immediate; subsequent attempts sleep for
|
|
33
|
+
``backoff[i-1]`` seconds. After exhausting the schedule the last
|
|
34
|
+
``PermissionError`` is re-raised. All other exceptions propagate
|
|
35
|
+
immediately.
|
|
36
|
+
"""
|
|
37
|
+
last: PermissionError | None = None
|
|
38
|
+
for attempt in range(len(backoff) + 1):
|
|
39
|
+
try:
|
|
40
|
+
return op()
|
|
41
|
+
except PermissionError as exc:
|
|
42
|
+
last = exc
|
|
43
|
+
if attempt < len(backoff):
|
|
44
|
+
time.sleep(backoff[attempt])
|
|
45
|
+
assert last is not None # invariant from loop above
|
|
46
|
+
raise last
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def _try_zstandard() -> Any | None:
|
|
50
|
+
try:
|
|
51
|
+
import zstandard
|
|
52
|
+
except ImportError:
|
|
53
|
+
return None
|
|
54
|
+
return zstandard
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def _normalize_part(part: Any) -> Any:
|
|
58
|
+
"""Coerce a key-part into something JSON-stable.
|
|
59
|
+
|
|
60
|
+
Lists and tuples both serialize as JSON arrays — order matters. Dicts
|
|
61
|
+
are sorted by key. Bytes are hashed (avoids embedding binary blobs in
|
|
62
|
+
the key). Anything else falls back to ``repr`` — explicit but uglier."""
|
|
63
|
+
if isinstance(part, (str, int, float, bool)) or part is None:
|
|
64
|
+
return part
|
|
65
|
+
if isinstance(part, (list, tuple)):
|
|
66
|
+
return [_normalize_part(p) for p in part]
|
|
67
|
+
if isinstance(part, dict):
|
|
68
|
+
return {str(k): _normalize_part(v) for k, v in sorted(part.items())}
|
|
69
|
+
if isinstance(part, bytes):
|
|
70
|
+
return {"__bytes_sha256__": hashlib.sha256(part).hexdigest()}
|
|
71
|
+
return {"__repr__": repr(part)}
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
class Cache:
|
|
75
|
+
"""Single-namespace content-addressed cache on disk.
|
|
76
|
+
|
|
77
|
+
``root`` is created on construction if missing. ``namespace`` is used
|
|
78
|
+
as a subdirectory and as part of the key (so two namespaces never
|
|
79
|
+
collide even if a caller produces identical key parts). ``version``
|
|
80
|
+
is the invalidation lever: bump it, and old entries become
|
|
81
|
+
unreadable by construction without any delete.
|
|
82
|
+
"""
|
|
83
|
+
|
|
84
|
+
def __init__(
|
|
85
|
+
self,
|
|
86
|
+
root: str | os.PathLike[str],
|
|
87
|
+
namespace: str,
|
|
88
|
+
version: int = 1,
|
|
89
|
+
) -> None:
|
|
90
|
+
if not isinstance(namespace, str) or not namespace:
|
|
91
|
+
raise ValueError(
|
|
92
|
+
f"namespace must be a non-empty str, got {namespace!r}"
|
|
93
|
+
)
|
|
94
|
+
if not isinstance(version, int) or isinstance(version, bool) or version <= 0:
|
|
95
|
+
raise ValueError(f"version must be a positive int, got {version!r}")
|
|
96
|
+
|
|
97
|
+
self._namespace = namespace
|
|
98
|
+
self._version = version
|
|
99
|
+
self._root = Path(root) / namespace
|
|
100
|
+
self._root.mkdir(parents=True, exist_ok=True)
|
|
101
|
+
self._zstd = _try_zstandard()
|
|
102
|
+
|
|
103
|
+
@property
|
|
104
|
+
def root(self) -> Path:
|
|
105
|
+
return self._root
|
|
106
|
+
|
|
107
|
+
@property
|
|
108
|
+
def namespace(self) -> str:
|
|
109
|
+
return self._namespace
|
|
110
|
+
|
|
111
|
+
@property
|
|
112
|
+
def version(self) -> int:
|
|
113
|
+
return self._version
|
|
114
|
+
|
|
115
|
+
def key_for(self, *parts: Any) -> str:
|
|
116
|
+
"""SHA-256 over a canonical JSON of ``parts``, namespace, version."""
|
|
117
|
+
canonical = {
|
|
118
|
+
"namespace": self._namespace,
|
|
119
|
+
"version": self._version,
|
|
120
|
+
"parts": [_normalize_part(p) for p in parts],
|
|
121
|
+
}
|
|
122
|
+
blob = json.dumps(canonical, sort_keys=True, separators=(",", ":")).encode()
|
|
123
|
+
return hashlib.sha256(blob).hexdigest()
|
|
124
|
+
|
|
125
|
+
def _paths(self, key: str) -> tuple[Path, Path]:
|
|
126
|
+
short = key[:16]
|
|
127
|
+
return self._root / f"{short}.bin", self._root / f"{short}.json"
|
|
128
|
+
|
|
129
|
+
def put(self, key: str, data: bytes) -> None:
|
|
130
|
+
"""Store ``data`` under ``key``. Atomic; retries on transient races."""
|
|
131
|
+
if not isinstance(key, str):
|
|
132
|
+
raise TypeError(f"key must be str, got {type(key).__name__}")
|
|
133
|
+
if not isinstance(data, (bytes, bytearray, memoryview)):
|
|
134
|
+
raise TypeError(f"data must be bytes-like, got {type(data).__name__}")
|
|
135
|
+
data_bytes = bytes(data)
|
|
136
|
+
|
|
137
|
+
bin_path, sidecar_path = self._paths(key)
|
|
138
|
+
|
|
139
|
+
if self._zstd is not None:
|
|
140
|
+
cctx = self._zstd.ZstdCompressor(level=3)
|
|
141
|
+
payload = cctx.compress(data_bytes)
|
|
142
|
+
compressed = True
|
|
143
|
+
else:
|
|
144
|
+
payload = data_bytes
|
|
145
|
+
compressed = False
|
|
146
|
+
|
|
147
|
+
checksum = hashlib.sha256(data_bytes).hexdigest()
|
|
148
|
+
sidecar: dict[str, Any] = {
|
|
149
|
+
"key": key,
|
|
150
|
+
"namespace": self._namespace,
|
|
151
|
+
"version": self._version,
|
|
152
|
+
"checksum": checksum,
|
|
153
|
+
"compressed": compressed,
|
|
154
|
+
"size": len(data_bytes),
|
|
155
|
+
}
|
|
156
|
+
|
|
157
|
+
tmp_bin = bin_path.with_suffix(bin_path.suffix + ".tmp")
|
|
158
|
+
tmp_side = sidecar_path.with_suffix(sidecar_path.suffix + ".tmp")
|
|
159
|
+
|
|
160
|
+
def _write_payload() -> None:
|
|
161
|
+
with open(tmp_bin, "wb") as f:
|
|
162
|
+
f.write(payload)
|
|
163
|
+
os.replace(tmp_bin, bin_path)
|
|
164
|
+
|
|
165
|
+
def _write_sidecar() -> None:
|
|
166
|
+
with open(tmp_side, "w", encoding="utf-8") as f:
|
|
167
|
+
json.dump(sidecar, f)
|
|
168
|
+
os.replace(tmp_side, sidecar_path)
|
|
169
|
+
|
|
170
|
+
_retry(_write_payload, _WRITE_BACKOFF)
|
|
171
|
+
_retry(_write_sidecar, _WRITE_BACKOFF)
|
|
172
|
+
|
|
173
|
+
def get(self, key: str) -> bytes | None:
|
|
174
|
+
"""Return cached bytes, or ``None`` if absent / version-mismatched.
|
|
175
|
+
|
|
176
|
+
Raises ``IOError`` on sidecar/payload mismatch (corrupt entry).
|
|
177
|
+
"""
|
|
178
|
+
if not isinstance(key, str):
|
|
179
|
+
raise TypeError(f"key must be str, got {type(key).__name__}")
|
|
180
|
+
|
|
181
|
+
bin_path, sidecar_path = self._paths(key)
|
|
182
|
+
if not (bin_path.exists() and sidecar_path.exists()):
|
|
183
|
+
return None
|
|
184
|
+
|
|
185
|
+
def _read_sidecar() -> dict[str, Any]:
|
|
186
|
+
with open(sidecar_path, "rb") as f:
|
|
187
|
+
parsed = json.loads(f.read())
|
|
188
|
+
if not isinstance(parsed, dict):
|
|
189
|
+
raise OSError(
|
|
190
|
+
f"sidecar at {sidecar_path} is not a JSON object"
|
|
191
|
+
)
|
|
192
|
+
return parsed
|
|
193
|
+
|
|
194
|
+
try:
|
|
195
|
+
sidecar = _retry(_read_sidecar, _READ_BACKOFF)
|
|
196
|
+
except FileNotFoundError:
|
|
197
|
+
return None
|
|
198
|
+
|
|
199
|
+
# Different key collided into the same 16-hex prefix? Treat as miss.
|
|
200
|
+
if sidecar.get("key") != key:
|
|
201
|
+
return None
|
|
202
|
+
# Stale version (cache invalidation lever). Transparent miss.
|
|
203
|
+
if sidecar.get("version") != self._version:
|
|
204
|
+
return None
|
|
205
|
+
if sidecar.get("namespace") != self._namespace:
|
|
206
|
+
return None
|
|
207
|
+
|
|
208
|
+
def _read_payload() -> bytes:
|
|
209
|
+
with open(bin_path, "rb") as f:
|
|
210
|
+
return f.read()
|
|
211
|
+
|
|
212
|
+
try:
|
|
213
|
+
payload = _retry(_read_payload, _READ_BACKOFF)
|
|
214
|
+
except FileNotFoundError:
|
|
215
|
+
return None
|
|
216
|
+
|
|
217
|
+
if sidecar.get("compressed"):
|
|
218
|
+
if self._zstd is None:
|
|
219
|
+
raise OSError(
|
|
220
|
+
f"cache entry {key[:16]!r} is zstd-compressed but "
|
|
221
|
+
"zstandard is not installed in this environment"
|
|
222
|
+
)
|
|
223
|
+
dctx = self._zstd.ZstdDecompressor()
|
|
224
|
+
try:
|
|
225
|
+
data: bytes = dctx.decompress(payload)
|
|
226
|
+
except self._zstd.ZstdError as exc:
|
|
227
|
+
# Compressed-payload corruption manifests here before checksum.
|
|
228
|
+
raise OSError(
|
|
229
|
+
f"cache entry {key[:16]!r} payload corrupt (zstd decode "
|
|
230
|
+
f"failed); remove it from {self._root} to invalidate"
|
|
231
|
+
) from exc
|
|
232
|
+
else:
|
|
233
|
+
data = payload
|
|
234
|
+
|
|
235
|
+
if hashlib.sha256(data).hexdigest() != sidecar.get("checksum"):
|
|
236
|
+
raise OSError(
|
|
237
|
+
f"cache entry {key[:16]!r} checksum mismatch — corrupt; "
|
|
238
|
+
f"remove {bin_path.name} and {sidecar_path.name} from "
|
|
239
|
+
f"{self._root} to invalidate"
|
|
240
|
+
)
|
|
241
|
+
return data
|
|
242
|
+
|
|
243
|
+
def __repr__(self) -> str:
|
|
244
|
+
return (
|
|
245
|
+
f"Cache(root={str(self._root)!r}, namespace={self._namespace!r}, "
|
|
246
|
+
f"version={self._version})"
|
|
247
|
+
)
|
patchcraft/extract.py
ADDED
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
"""Patch extraction via torch.nn.functional.unfold.
|
|
2
|
+
|
|
3
|
+
Contract: docs/THEORY.md §1 and §9.1, docs/ADR/0001-patch-extraction-api.md,
|
|
4
|
+
docs/ADR/0002-patchify-transform.md.
|
|
5
|
+
"""
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
import torch
|
|
9
|
+
import torch.nn.functional as F # noqa: N812 (torch convention)
|
|
10
|
+
|
|
11
|
+
__all__ = ["Patchify", "extract"]
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def _as_pair(value: int | tuple[int, int], name: str) -> tuple[int, int]:
|
|
15
|
+
if isinstance(value, int) and not isinstance(value, bool):
|
|
16
|
+
if value <= 0:
|
|
17
|
+
raise ValueError(f"{name} must be positive, got {value}")
|
|
18
|
+
return (value, value)
|
|
19
|
+
if isinstance(value, tuple) and len(value) == 2:
|
|
20
|
+
h, w = value
|
|
21
|
+
h_ok = isinstance(h, int) and not isinstance(h, bool)
|
|
22
|
+
w_ok = isinstance(w, int) and not isinstance(w, bool)
|
|
23
|
+
if not (h_ok and w_ok):
|
|
24
|
+
raise ValueError(f"{name} must contain ints, got {value!r}")
|
|
25
|
+
if h <= 0 or w <= 0:
|
|
26
|
+
raise ValueError(f"{name} must be positive, got {value!r}")
|
|
27
|
+
return (h, w)
|
|
28
|
+
raise ValueError(f"{name} must be int or (int, int), got {value!r}")
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def extract(
|
|
32
|
+
image: torch.Tensor,
|
|
33
|
+
patch_size: int | tuple[int, int],
|
|
34
|
+
stride: int | tuple[int, int],
|
|
35
|
+
dilation: int | tuple[int, int] = 1,
|
|
36
|
+
) -> torch.Tensor:
|
|
37
|
+
"""Extract rectangular patches from a `(C, H, W)` image.
|
|
38
|
+
|
|
39
|
+
Returns `Tensor[L, C, ph, pw]` in row-major order. Patch `k` has its
|
|
40
|
+
top-left at `(k // num_w * sh, k % num_w * sw)`. Truncation is the only
|
|
41
|
+
boundary policy: if the geometry fits no patch, returns `Tensor[0, C, ph, pw]`.
|
|
42
|
+
Dtype and device of `image` are preserved.
|
|
43
|
+
"""
|
|
44
|
+
if not isinstance(image, torch.Tensor):
|
|
45
|
+
raise TypeError(f"image must be torch.Tensor, got {type(image).__name__}")
|
|
46
|
+
if image.ndim != 3:
|
|
47
|
+
raise ValueError(f"image must have ndim==3 (C, H, W), got ndim={image.ndim}")
|
|
48
|
+
|
|
49
|
+
ph, pw = _as_pair(patch_size, "patch_size")
|
|
50
|
+
sh, sw = _as_pair(stride, "stride")
|
|
51
|
+
dh, dw = _as_pair(dilation, "dilation")
|
|
52
|
+
|
|
53
|
+
c, h, w = image.shape
|
|
54
|
+
eff_h = dh * (ph - 1) + 1
|
|
55
|
+
eff_w = dw * (pw - 1) + 1
|
|
56
|
+
|
|
57
|
+
if h < eff_h or w < eff_w:
|
|
58
|
+
return torch.empty(0, c, ph, pw, dtype=image.dtype, device=image.device)
|
|
59
|
+
|
|
60
|
+
unfolded = F.unfold(
|
|
61
|
+
image.unsqueeze(0),
|
|
62
|
+
kernel_size=(ph, pw),
|
|
63
|
+
dilation=(dh, dw),
|
|
64
|
+
stride=(sh, sw),
|
|
65
|
+
)
|
|
66
|
+
return unfolded[0].view(c, ph, pw, -1).permute(3, 0, 1, 2).contiguous()
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
class Patchify:
|
|
70
|
+
"""Callable that extracts patches with a frozen geometry.
|
|
71
|
+
|
|
72
|
+
Drop-in for `torchvision.transforms.Compose([...])`:
|
|
73
|
+
|
|
74
|
+
transform = Compose([
|
|
75
|
+
ToTensor(),
|
|
76
|
+
GaussianBlur(kernel_size=3),
|
|
77
|
+
Patchify(patch_size=4, stride=2),
|
|
78
|
+
])
|
|
79
|
+
patches = transform(pil_image) # Tensor[L, C, 4, 4]
|
|
80
|
+
|
|
81
|
+
Equivalent to `lambda img: extract(img, patch_size, stride, dilation)`,
|
|
82
|
+
but composable, repr-friendly, and validates the geometry at construction
|
|
83
|
+
instead of at first call. Holds no state beyond the geometry — no cache,
|
|
84
|
+
no fixed image size, no device. See ADR 0002.
|
|
85
|
+
|
|
86
|
+
Output shape is `(L, C, ph, pw)`, the same as `extract`. Subsequent
|
|
87
|
+
transforms in the Compose chain receive the patch stack, not a single
|
|
88
|
+
patch; they must accept `(N, C, H, W)` or be wrapped.
|
|
89
|
+
"""
|
|
90
|
+
|
|
91
|
+
__slots__ = ("_dh", "_dw", "_ph", "_pw", "_sh", "_sw")
|
|
92
|
+
|
|
93
|
+
def __init__(
|
|
94
|
+
self,
|
|
95
|
+
patch_size: int | tuple[int, int],
|
|
96
|
+
stride: int | tuple[int, int],
|
|
97
|
+
dilation: int | tuple[int, int] = 1,
|
|
98
|
+
) -> None:
|
|
99
|
+
self._ph, self._pw = _as_pair(patch_size, "patch_size")
|
|
100
|
+
self._sh, self._sw = _as_pair(stride, "stride")
|
|
101
|
+
self._dh, self._dw = _as_pair(dilation, "dilation")
|
|
102
|
+
|
|
103
|
+
def __call__(self, image: torch.Tensor) -> torch.Tensor:
|
|
104
|
+
return extract(
|
|
105
|
+
image,
|
|
106
|
+
patch_size=(self._ph, self._pw),
|
|
107
|
+
stride=(self._sh, self._sw),
|
|
108
|
+
dilation=(self._dh, self._dw),
|
|
109
|
+
)
|
|
110
|
+
|
|
111
|
+
def __repr__(self) -> str:
|
|
112
|
+
return (
|
|
113
|
+
f"Patchify(patch_size=({self._ph}, {self._pw}), "
|
|
114
|
+
f"stride=({self._sh}, {self._sw}), "
|
|
115
|
+
f"dilation=({self._dh}, {self._dw}))"
|
|
116
|
+
)
|
patchcraft/geometry.py
ADDED
|
@@ -0,0 +1,290 @@
|
|
|
1
|
+
"""Pre-flight geometry helpers: enumerate valid patch tilings, count patches.
|
|
2
|
+
|
|
3
|
+
Pure number-only API — does not touch tensors or images. Useful for:
|
|
4
|
+
- asking the lib "what patch sizes fit my image cleanly?" before extracting;
|
|
5
|
+
- precomputing patch counts for memory planning;
|
|
6
|
+
- driving parametrized tests over the full space of valid geometries.
|
|
7
|
+
|
|
8
|
+
Contract: docs/THEORY.md §1.5 and §9.6.
|
|
9
|
+
"""
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
from typing import NamedTuple
|
|
13
|
+
|
|
14
|
+
from patchcraft.extract import _as_pair
|
|
15
|
+
|
|
16
|
+
__all__ = [
|
|
17
|
+
"PairedTilingSpec",
|
|
18
|
+
"TilingSpec",
|
|
19
|
+
"num_patches",
|
|
20
|
+
"paired_tilings",
|
|
21
|
+
"scale_factor",
|
|
22
|
+
"tilings",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
class TilingSpec(NamedTuple):
|
|
27
|
+
"""One valid patch geometry for an image.
|
|
28
|
+
|
|
29
|
+
``overlap=False`` means an *exact tile*: ``patch_size == stride`` and the
|
|
30
|
+
image is divided into a clean grid with no overlap and no waste.
|
|
31
|
+
``overlap=True`` means ``stride < patch_size`` and full coverage is still
|
|
32
|
+
achieved — adjacent patches share pixels.
|
|
33
|
+
"""
|
|
34
|
+
|
|
35
|
+
patch_size: tuple[int, int]
|
|
36
|
+
stride: tuple[int, int]
|
|
37
|
+
dilation: tuple[int, int]
|
|
38
|
+
num_patches: tuple[int, int]
|
|
39
|
+
total_patches: int
|
|
40
|
+
overlap: bool
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
class PairedTilingSpec(NamedTuple):
|
|
44
|
+
"""A pair of tilings that align across two resolutions of the same image.
|
|
45
|
+
|
|
46
|
+
``hr.patch_size == scale_factor * lr.patch_size`` and same for stride;
|
|
47
|
+
``lr.total_patches == hr.total_patches`` by construction. Patch ``k`` on
|
|
48
|
+
the LR side and patch ``k`` on the HR side cover the same image region
|
|
49
|
+
at different resolutions.
|
|
50
|
+
"""
|
|
51
|
+
|
|
52
|
+
lr: TilingSpec
|
|
53
|
+
hr: TilingSpec
|
|
54
|
+
scale_factor: int
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def num_patches(
|
|
58
|
+
image_shape: tuple[int, ...],
|
|
59
|
+
patch_size: int | tuple[int, int],
|
|
60
|
+
stride: int | tuple[int, int],
|
|
61
|
+
dilation: int | tuple[int, int] = 1,
|
|
62
|
+
) -> tuple[int, int]:
|
|
63
|
+
"""Return ``(num_h, num_w)`` — how many patches `extract` would produce.
|
|
64
|
+
|
|
65
|
+
Accepts ``(H, W)`` or ``(C, H, W)`` for ``image_shape``. Returns ``(0, 0)``
|
|
66
|
+
in either axis when the effective patch does not fit (mirroring `extract`'s
|
|
67
|
+
empty-tensor behavior). Does not allocate or touch any tensor.
|
|
68
|
+
"""
|
|
69
|
+
if not (isinstance(image_shape, tuple) and len(image_shape) in (2, 3)):
|
|
70
|
+
raise ValueError(
|
|
71
|
+
f"image_shape must be (H, W) or (C, H, W), got {image_shape!r}"
|
|
72
|
+
)
|
|
73
|
+
h, w = (image_shape[-2], image_shape[-1])
|
|
74
|
+
for axis_name, val in zip(("H", "W"), (h, w), strict=True):
|
|
75
|
+
if not isinstance(val, int) or isinstance(val, bool) or val <= 0:
|
|
76
|
+
raise ValueError(
|
|
77
|
+
f"image_shape[{axis_name}] must be a positive int, got {val!r}"
|
|
78
|
+
)
|
|
79
|
+
|
|
80
|
+
ph, pw = _as_pair(patch_size, "patch_size")
|
|
81
|
+
sh, sw = _as_pair(stride, "stride")
|
|
82
|
+
dh, dw = _as_pair(dilation, "dilation")
|
|
83
|
+
|
|
84
|
+
eff_h = dh * (ph - 1) + 1
|
|
85
|
+
eff_w = dw * (pw - 1) + 1
|
|
86
|
+
nh = (h - eff_h) // sh + 1 if h >= eff_h else 0
|
|
87
|
+
nw = (w - eff_w) // sw + 1 if w >= eff_w else 0
|
|
88
|
+
return (nh, nw)
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def tilings(
|
|
92
|
+
image_shape: tuple[int, int] | tuple[int, int, int],
|
|
93
|
+
*,
|
|
94
|
+
allow_overlap: bool = False,
|
|
95
|
+
min_patch_size: int = 2,
|
|
96
|
+
max_patch_size: int | None = None,
|
|
97
|
+
) -> list[TilingSpec]:
|
|
98
|
+
"""Enumerate square geometries that fully cover an image (``dilation==1``).
|
|
99
|
+
|
|
100
|
+
Always emits geometries with **full coverage** (no truncated rows/cols).
|
|
101
|
+
Patches are square: ``ph == pw``, ``sh == sw``.
|
|
102
|
+
|
|
103
|
+
Parameters
|
|
104
|
+
----------
|
|
105
|
+
image_shape
|
|
106
|
+
``(H, W)`` or ``(C, H, W)``. Only H, W matter.
|
|
107
|
+
allow_overlap
|
|
108
|
+
If False (default), emit only exact tilings (``stride == patch_size``
|
|
109
|
+
and ``H % p == 0`` and ``W % p == 0``). If True, also emit
|
|
110
|
+
``stride < patch_size`` geometries with ``(H - p) % s == 0`` and
|
|
111
|
+
``(W - p) % s == 0`` — overlap-with-clean-edges.
|
|
112
|
+
min_patch_size
|
|
113
|
+
Smallest patch to consider. Defaults to 2 (skips the trivial pixel-wise
|
|
114
|
+
tiling at ``p == 1``).
|
|
115
|
+
max_patch_size
|
|
116
|
+
Largest patch to consider. Defaults to ``min(H, W)``.
|
|
117
|
+
|
|
118
|
+
Returns
|
|
119
|
+
-------
|
|
120
|
+
list[TilingSpec]
|
|
121
|
+
Sorted by ``(patch_size[0], stride[0])`` ascending. Always at least
|
|
122
|
+
one entry: ``(p=min(H, W), s=p)`` when image is square.
|
|
123
|
+
|
|
124
|
+
Raises
|
|
125
|
+
------
|
|
126
|
+
ValueError
|
|
127
|
+
On malformed ``image_shape``, non-positive bounds, or
|
|
128
|
+
``min_patch_size > max_patch_size``.
|
|
129
|
+
"""
|
|
130
|
+
if not (isinstance(image_shape, tuple) and len(image_shape) in (2, 3)):
|
|
131
|
+
raise ValueError(
|
|
132
|
+
f"image_shape must be (H, W) or (C, H, W), got {image_shape!r}"
|
|
133
|
+
)
|
|
134
|
+
h, w = (image_shape[-2], image_shape[-1])
|
|
135
|
+
for axis_name, val in zip(("H", "W"), (h, w), strict=True):
|
|
136
|
+
if not isinstance(val, int) or isinstance(val, bool) or val <= 0:
|
|
137
|
+
raise ValueError(
|
|
138
|
+
f"image_shape[{axis_name}] must be a positive int, got {val!r}"
|
|
139
|
+
)
|
|
140
|
+
|
|
141
|
+
if (
|
|
142
|
+
not isinstance(min_patch_size, int)
|
|
143
|
+
or isinstance(min_patch_size, bool)
|
|
144
|
+
or min_patch_size <= 0
|
|
145
|
+
):
|
|
146
|
+
raise ValueError(
|
|
147
|
+
f"min_patch_size must be a positive int, got {min_patch_size!r}"
|
|
148
|
+
)
|
|
149
|
+
if max_patch_size is None:
|
|
150
|
+
max_patch_size = min(h, w)
|
|
151
|
+
if (
|
|
152
|
+
not isinstance(max_patch_size, int)
|
|
153
|
+
or isinstance(max_patch_size, bool)
|
|
154
|
+
or max_patch_size <= 0
|
|
155
|
+
):
|
|
156
|
+
raise ValueError(
|
|
157
|
+
f"max_patch_size must be a positive int or None, got {max_patch_size!r}"
|
|
158
|
+
)
|
|
159
|
+
if min_patch_size > max_patch_size:
|
|
160
|
+
raise ValueError(
|
|
161
|
+
f"min_patch_size ({min_patch_size}) > max_patch_size ({max_patch_size})"
|
|
162
|
+
)
|
|
163
|
+
|
|
164
|
+
upper = min(max_patch_size, h, w)
|
|
165
|
+
results: list[TilingSpec] = []
|
|
166
|
+
for p in range(min_patch_size, upper + 1):
|
|
167
|
+
# Exact tile: stride == patch_size; requires divisibility on both axes.
|
|
168
|
+
if h % p == 0 and w % p == 0:
|
|
169
|
+
nh, nw = h // p, w // p
|
|
170
|
+
results.append(TilingSpec(
|
|
171
|
+
patch_size=(p, p),
|
|
172
|
+
stride=(p, p),
|
|
173
|
+
dilation=(1, 1),
|
|
174
|
+
num_patches=(nh, nw),
|
|
175
|
+
total_patches=nh * nw,
|
|
176
|
+
overlap=False,
|
|
177
|
+
))
|
|
178
|
+
if allow_overlap:
|
|
179
|
+
for s in range(1, p):
|
|
180
|
+
if (h - p) % s == 0 and (w - p) % s == 0:
|
|
181
|
+
nh = (h - p) // s + 1
|
|
182
|
+
nw = (w - p) // s + 1
|
|
183
|
+
results.append(TilingSpec(
|
|
184
|
+
patch_size=(p, p),
|
|
185
|
+
stride=(s, s),
|
|
186
|
+
dilation=(1, 1),
|
|
187
|
+
num_patches=(nh, nw),
|
|
188
|
+
total_patches=nh * nw,
|
|
189
|
+
overlap=True,
|
|
190
|
+
))
|
|
191
|
+
return results
|
|
192
|
+
|
|
193
|
+
|
|
194
|
+
def scale_factor(
|
|
195
|
+
lr_shape: tuple[int, int] | tuple[int, int, int],
|
|
196
|
+
hr_shape: tuple[int, int] | tuple[int, int, int],
|
|
197
|
+
) -> int | None:
|
|
198
|
+
"""Return the integer scale factor between two image shapes, or ``None``.
|
|
199
|
+
|
|
200
|
+
Accepts ``(H, W)`` or ``(C, H, W)`` for either argument (channels are
|
|
201
|
+
ignored). Returns ``k`` such that
|
|
202
|
+
``hr_shape[-2:] == (k * lr_shape[-2], k * lr_shape[-1])``, or ``None``
|
|
203
|
+
when no such integer ``k >= 1`` exists (non-divisible, anisotropic, or
|
|
204
|
+
LR larger than HR).
|
|
205
|
+
|
|
206
|
+
Pure shape math; no tensor, no allocation. Use it before calling
|
|
207
|
+
:func:`patchcraft.pair` to discover the scale factor from data instead
|
|
208
|
+
of hard-coding it.
|
|
209
|
+
"""
|
|
210
|
+
for name, shape in (("lr_shape", lr_shape), ("hr_shape", hr_shape)):
|
|
211
|
+
if not (isinstance(shape, tuple) and len(shape) in (2, 3)):
|
|
212
|
+
raise ValueError(
|
|
213
|
+
f"{name} must be (H, W) or (C, H, W), got {shape!r}"
|
|
214
|
+
)
|
|
215
|
+
h, w = shape[-2], shape[-1]
|
|
216
|
+
for axis, val in (("H", h), ("W", w)):
|
|
217
|
+
if not isinstance(val, int) or isinstance(val, bool) or val <= 0:
|
|
218
|
+
raise ValueError(
|
|
219
|
+
f"{name}[{axis}] must be a positive int, got {val!r}"
|
|
220
|
+
)
|
|
221
|
+
|
|
222
|
+
h_lr, w_lr = lr_shape[-2], lr_shape[-1]
|
|
223
|
+
h_hr, w_hr = hr_shape[-2], hr_shape[-1]
|
|
224
|
+
if h_hr % h_lr != 0 or w_hr % w_lr != 0:
|
|
225
|
+
return None
|
|
226
|
+
sf_h = h_hr // h_lr
|
|
227
|
+
sf_w = w_hr // w_lr
|
|
228
|
+
if sf_h != sf_w or sf_h < 1:
|
|
229
|
+
return None
|
|
230
|
+
return sf_h
|
|
231
|
+
|
|
232
|
+
|
|
233
|
+
def paired_tilings(
|
|
234
|
+
lr_shape: tuple[int, int] | tuple[int, int, int],
|
|
235
|
+
hr_shape: tuple[int, int] | tuple[int, int, int],
|
|
236
|
+
*,
|
|
237
|
+
allow_overlap: bool = False,
|
|
238
|
+
min_patch_size: int = 2,
|
|
239
|
+
max_patch_size: int | None = None,
|
|
240
|
+
) -> list[PairedTilingSpec]:
|
|
241
|
+
"""Enumerate aligned tiling pairs between two resolutions of the same image.
|
|
242
|
+
|
|
243
|
+
Requires ``hr_shape`` to be an integer multiple of ``lr_shape`` (see
|
|
244
|
+
:func:`scale_factor`). For each LR tiling emitted by :func:`tilings`,
|
|
245
|
+
derives the matching HR tiling by multiplying patch size and stride by
|
|
246
|
+
the scale factor. Both sides have identical ``total_patches`` and patch
|
|
247
|
+
``k`` covers the same image region on both sides.
|
|
248
|
+
|
|
249
|
+
Use the result to drive :func:`patchcraft.pair` with confidence that the
|
|
250
|
+
parameters produce sound, aligned LR/HR patch sets.
|
|
251
|
+
|
|
252
|
+
Raises
|
|
253
|
+
------
|
|
254
|
+
ValueError
|
|
255
|
+
If ``lr_shape`` and ``hr_shape`` are not related by an integer
|
|
256
|
+
scale factor, or on the same input validation cases as
|
|
257
|
+
:func:`tilings`.
|
|
258
|
+
"""
|
|
259
|
+
sf = scale_factor(lr_shape, hr_shape)
|
|
260
|
+
if sf is None:
|
|
261
|
+
raise ValueError(
|
|
262
|
+
f"lr_shape={lr_shape} and hr_shape={hr_shape} are not related "
|
|
263
|
+
"by a positive integer scale factor on both spatial axes"
|
|
264
|
+
)
|
|
265
|
+
h_hr, w_hr = hr_shape[-2], hr_shape[-1]
|
|
266
|
+
lr_specs = tilings(
|
|
267
|
+
lr_shape,
|
|
268
|
+
allow_overlap=allow_overlap,
|
|
269
|
+
min_patch_size=min_patch_size,
|
|
270
|
+
max_patch_size=max_patch_size,
|
|
271
|
+
)
|
|
272
|
+
|
|
273
|
+
pairs: list[PairedTilingSpec] = []
|
|
274
|
+
for lr in lr_specs:
|
|
275
|
+
ph_hr = lr.patch_size[0] * sf
|
|
276
|
+
pw_hr = lr.patch_size[1] * sf
|
|
277
|
+
sh_hr = lr.stride[0] * sf
|
|
278
|
+
sw_hr = lr.stride[1] * sf
|
|
279
|
+
nh_hr = (h_hr - ph_hr) // sh_hr + 1
|
|
280
|
+
nw_hr = (w_hr - pw_hr) // sw_hr + 1
|
|
281
|
+
hr = TilingSpec(
|
|
282
|
+
patch_size=(ph_hr, pw_hr),
|
|
283
|
+
stride=(sh_hr, sw_hr),
|
|
284
|
+
dilation=(1, 1),
|
|
285
|
+
num_patches=(nh_hr, nw_hr),
|
|
286
|
+
total_patches=nh_hr * nw_hr,
|
|
287
|
+
overlap=lr.overlap,
|
|
288
|
+
)
|
|
289
|
+
pairs.append(PairedTilingSpec(lr=lr, hr=hr, scale_factor=sf))
|
|
290
|
+
return pairs
|