descriptron-gui 2.8.0__tar.gz → 2.8.2__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.
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/PKG-INFO +3 -3
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/pyproject.toml +3 -3
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/cli.py +19 -1
- descriptron_gui-2.8.2/src/descriptron_gui/tools/biorag_annotate_tokens_v2.py +585 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/biorag_annotate_v1.py +229 -17
- descriptron_gui-2.8.2/src/descriptron_gui/tools/biorag_annotate_v2.py +3677 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron-v2-v91.py +19930 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron-v2-v92.py +20097 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron-v2-v93.py +20265 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron-v2-v94.py +20291 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron-v2-v95.py +20293 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron-v2-v96.py +20659 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron-v2-v97.py +20696 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron_category_nav_v1.py +109 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron_dinomask_dialog_v1.py +358 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron_ontology.py +768 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron_ontology_templates.json +18306 -0
- descriptron_gui-2.8.2/src/descriptron_gui/tools/descriptron_undo_v1.py +313 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/.gitignore +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/LICENSE +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/NOTICE +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/README.md +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/__init__.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/data/marmot.jpg +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/biorag_annotate_tokens_v1.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v73.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v74.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v75.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v76.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v77.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v78.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v79.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v80.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v81.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v82.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v83.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v84.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v85.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v86.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v87.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v88.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v89.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron-v2-v90.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron_category_edit.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron_descriptive_characters.json +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron_descriptive_characters.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron_folder_save_v1.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/descriptron_sam3_instances.py +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/sam3_assets/README.txt +0 -0
- {descriptron_gui-2.8.0 → descriptron_gui-2.8.2}/src/descriptron_gui/tools/sam3_assets/bpe_simple_vocab_16e6.txt.gz +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: descriptron-gui
|
|
3
|
-
Version: 2.8.
|
|
3
|
+
Version: 2.8.2
|
|
4
4
|
Summary: The Descriptron annotation GUI: SAM2-assisted segmentation, keypoints and COCO export
|
|
5
5
|
Project-URL: Homepage, https://github.com/alexrvandam/Descriptron
|
|
6
6
|
Author: Alex R. Van Dam
|
|
@@ -14,8 +14,8 @@ Classifier: Operating System :: OS Independent
|
|
|
14
14
|
Classifier: Programming Language :: Python :: 3
|
|
15
15
|
Classifier: Topic :: Scientific/Engineering :: Image Processing
|
|
16
16
|
Requires-Python: >=3.9
|
|
17
|
-
Requires-Dist: descriptron-core==2.8.
|
|
18
|
-
Requires-Dist: descriptron-vision==2.8.
|
|
17
|
+
Requires-Dist: descriptron-core==2.8.2
|
|
18
|
+
Requires-Dist: descriptron-vision==2.8.2
|
|
19
19
|
Requires-Dist: tkinterdnd2>=0.4
|
|
20
20
|
Requires-Dist: tkmacosx>=1.0.5; sys_platform == 'darwin'
|
|
21
21
|
Provides-Extra: dev
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "descriptron-gui"
|
|
7
|
-
version = "2.8.
|
|
7
|
+
version = "2.8.2"
|
|
8
8
|
description = "The Descriptron annotation GUI: SAM2-assisted segmentation, keypoints and COCO export"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.9"
|
|
@@ -24,8 +24,8 @@ classifiers = [
|
|
|
24
24
|
# fails to start. The GUI drives the analysis scripts directly, so it needs core
|
|
25
25
|
# as well as vision — it is not a thin front end.
|
|
26
26
|
dependencies = [
|
|
27
|
-
"descriptron-core==2.8.
|
|
28
|
-
"descriptron-vision==2.8.
|
|
27
|
+
"descriptron-core==2.8.2",
|
|
28
|
+
"descriptron-vision==2.8.2",
|
|
29
29
|
# coloured buttons on macOS (Aqua ignores tk.Button bg); pure Python, Mac only
|
|
30
30
|
"tkmacosx>=1.0.5; sys_platform == 'darwin'",
|
|
31
31
|
# drag and drop of images/folders onto the window (v78); the GUI starts without it
|
|
@@ -18,6 +18,12 @@ from importlib.resources import files
|
|
|
18
18
|
from pathlib import Path
|
|
19
19
|
|
|
20
20
|
|
|
21
|
+
def _gui_number(p: Path) -> int:
|
|
22
|
+
"""'descriptron-v2-v96.py' -> 96; a name without a number sorts first (-1)."""
|
|
23
|
+
tail = p.stem.rsplit("-v", 1)[-1]
|
|
24
|
+
return int(tail) if tail.isdigit() else -1
|
|
25
|
+
|
|
26
|
+
|
|
21
27
|
def main() -> None:
|
|
22
28
|
try:
|
|
23
29
|
import tkinter # noqa: F401
|
|
@@ -29,11 +35,23 @@ def main() -> None:
|
|
|
29
35
|
f"(original error: {exc})") from exc
|
|
30
36
|
|
|
31
37
|
tools = Path(str(files("descriptron_gui") / "tools"))
|
|
32
|
-
|
|
38
|
+
# 2.8.1: newest by the NUMBER after '-v' (a string sort put v100 before v99); DESCRIPTRON_GUI_VERSION=v94 picks an older one
|
|
39
|
+
candidates = sorted(tools.glob("descriptron-v2-*.py"), key=_gui_number, reverse=True)
|
|
33
40
|
if not candidates:
|
|
34
41
|
raise SystemExit(f"no GUI program found in {tools} — "
|
|
35
42
|
"was sync_sources.py run before building?")
|
|
36
43
|
path = candidates[0]
|
|
44
|
+
want = os.environ.get("DESCRIPTRON_GUI_VERSION", "").strip()
|
|
45
|
+
if want:
|
|
46
|
+
chosen = [p for p in candidates if _gui_number(p) == _gui_number(Path(f"descriptron-v2-v{want.lstrip('vV')}.py"))]
|
|
47
|
+
if not chosen:
|
|
48
|
+
raise SystemExit(f"DESCRIPTRON_GUI_VERSION={want}: no such GUI. Installed: "
|
|
49
|
+
+ ", ".join(p.stem.split("-")[-1] for p in candidates))
|
|
50
|
+
path = chosen[0]
|
|
51
|
+
if sys.argv[1:2] == ["--which"]:
|
|
52
|
+
from importlib.metadata import version
|
|
53
|
+
print(f"{path}\ndescriptron-gui {version('descriptron-gui')}")
|
|
54
|
+
return
|
|
37
55
|
sys.path.insert(0, str(tools))
|
|
38
56
|
os.environ.setdefault("DESCRIPTRON_GUI_DATA",
|
|
39
57
|
str(Path(str(files("descriptron_gui") / "data"))))
|
|
@@ -0,0 +1,585 @@
|
|
|
1
|
+
#!/usr/bin/env python
|
|
2
|
+
"""BioRAG Annotate - DINOv3 patch-token extraction (v2; 22 Oct 2026: --size native = the image's own long side,
|
|
3
|
+
snapped to a multiple of 16, tiled above --tile like any large image; the UI module is biorag_annotate_v2).
|
|
4
|
+
|
|
5
|
+
Run as a subprocess by the BioRAG Annotate window (descriptron-v2-v89.py), in an env whose
|
|
6
|
+
transformers knows DINOv3 (the biorag env). One .npz per image in --cache_dir:
|
|
7
|
+
|
|
8
|
+
tokens (gh*gw, D) float16 patch tokens of layer --layer (register / CLS tokens dropped)
|
|
9
|
+
grid (gh, gw) patch grid
|
|
10
|
+
image_hw (H, W) original image size
|
|
11
|
+
work_hw (nh, nw) size the image was resized to (multiples of 16)
|
|
12
|
+
scale (sy, sx) nh / H, nw / W
|
|
13
|
+
layer, size, model, source, source_size, source_mtime
|
|
14
|
+
|
|
15
|
+
Resolution rule (AVD): the image long side is resized to --size (1024 / 1536 / 2048) with cubic
|
|
16
|
+
interpolation - small images are UPSAMPLED, never kept small - and both sides are snapped to
|
|
17
|
+
multiples of 16. Images whose long side exceeds --tile are processed in overlapping tiles
|
|
18
|
+
(128 px overlap, tile centres kept), as in measure/biorag_probe_v1.py.
|
|
19
|
+
|
|
20
|
+
Progress goes to stdout as one line per image:
|
|
21
|
+
BRA_PROGRESS<TAB>i<TAB>n<TAB>ok|cached|failed<TAB>seconds<TAB>file name<TAB>message
|
|
22
|
+
and a final line BRA_DONE <n_ok> <n_failed>.
|
|
23
|
+
|
|
24
|
+
The pure helpers (load_rgb8, working_size, cache_path) import without torch, so the GUI side can
|
|
25
|
+
use them too.
|
|
26
|
+
"""
|
|
27
|
+
import argparse
|
|
28
|
+
import json
|
|
29
|
+
import os
|
|
30
|
+
import sys
|
|
31
|
+
import time
|
|
32
|
+
|
|
33
|
+
import numpy as np
|
|
34
|
+
|
|
35
|
+
MEAN = np.array([0.485, 0.456, 0.406], np.float32)
|
|
36
|
+
STD = np.array([0.229, 0.224, 0.225], np.float32)
|
|
37
|
+
DEFAULT_MODEL = "facebook/dinov3-vitl16-pretrain-lvd1689m"
|
|
38
|
+
CACHE_VERSION = 1
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
# ---------------------------------------------------------------- pure helpers (no torch)
|
|
42
|
+
def _to_uint8(arr):
|
|
43
|
+
"""Any numeric array -> uint8. 8-bit data is kept as is; deeper data is scaled by its maximum."""
|
|
44
|
+
a = np.asarray(arr)
|
|
45
|
+
if a.dtype == np.uint8:
|
|
46
|
+
return a
|
|
47
|
+
a = a.astype(np.float32)
|
|
48
|
+
mx = float(np.nanmax(a)) if a.size else 0.0
|
|
49
|
+
if mx <= 0:
|
|
50
|
+
return np.zeros(a.shape, np.uint8)
|
|
51
|
+
if mx <= 1.0 and a.dtype.kind == "f":
|
|
52
|
+
mx = 1.0
|
|
53
|
+
elif mx <= 255:
|
|
54
|
+
mx = 255.0
|
|
55
|
+
return np.clip(a / mx * 255.0 + 0.5, 0, 255).astype(np.uint8)
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def load_rgb8(path):
|
|
59
|
+
"""Image file -> HxWx3 uint8 RGB. Same pixel grid as the GUI (PIL, no EXIF rotation).
|
|
60
|
+
8-bit images: PIL convert('RGB') exactly as the GUI does (RGBA drops alpha).
|
|
61
|
+
16-bit / 32-bit / float images (where convert('RGB') would clip): scaled to 8 bit by their max.
|
|
62
|
+
Falls back to OpenCV for files PIL cannot open (e.g. 48-bit RGB tif)."""
|
|
63
|
+
from PIL import Image
|
|
64
|
+
try:
|
|
65
|
+
im = Image.open(path)
|
|
66
|
+
if im.mode in ("I;16", "I;16B", "I;16L", "I;16N", "I", "F"):
|
|
67
|
+
a = _to_uint8(np.array(im))
|
|
68
|
+
return np.ascontiguousarray(np.stack([a] * 3, -1))
|
|
69
|
+
return np.asarray(im.convert("RGB"))
|
|
70
|
+
except Exception:
|
|
71
|
+
import cv2
|
|
72
|
+
a = cv2.imread(path, cv2.IMREAD_UNCHANGED)
|
|
73
|
+
if a is None:
|
|
74
|
+
raise
|
|
75
|
+
a = _to_uint8(a)
|
|
76
|
+
if a.ndim == 2:
|
|
77
|
+
a = np.stack([a] * 3, -1)
|
|
78
|
+
elif a.shape[2] == 4:
|
|
79
|
+
a = cv2.cvtColor(a, cv2.COLOR_BGRA2RGB)
|
|
80
|
+
else:
|
|
81
|
+
a = cv2.cvtColor(a, cv2.COLOR_BGR2RGB)
|
|
82
|
+
return np.ascontiguousarray(a)
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
def working_size(h, w, long_side):
|
|
86
|
+
"""(H, W) -> (nh, nw): long side resized to long_side (up OR down), both sides multiples of 16.
|
|
87
|
+
v2: long_side <= 0 ('native') keeps the image's own long side (only snapped to multiples of 16)."""
|
|
88
|
+
if long_side is None or int(long_side) <= 0:
|
|
89
|
+
long_side = max(h, w)
|
|
90
|
+
s = float(long_side) / max(h, w)
|
|
91
|
+
nh = max(16, int(round(h * s / 16.0)) * 16)
|
|
92
|
+
nw = max(16, int(round(w * s / 16.0)) * 16)
|
|
93
|
+
return nh, nw
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def cache_path(cache_dir, image_path, layer, size, box=None):
|
|
97
|
+
"""Cache file of one image (or of one crop box x0,y0,x1,y1 of it) for one (layer, size). The full
|
|
98
|
+
file name (with extension) is kept, so a.tif and a.jpg never share a cache file."""
|
|
99
|
+
crop = "" if box is None else "__crop_%d_%d_%d_%d" % tuple(int(v) for v in box)
|
|
100
|
+
return os.path.join(cache_dir, f"{os.path.basename(image_path)}{crop}__L{int(layer)}_s{int(size)}.npz")
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
MAX_CROP_SCALE = 8.0
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
def crop_working_size(h, w, long_side, max_scale=MAX_CROP_SCALE):
|
|
107
|
+
"""Crop (h, w) -> working size: long side resized to long_side (cubic UPsampling for small crops),
|
|
108
|
+
but at most max_scale times the crop's own pixels; sides multiples of 16."""
|
|
109
|
+
s = min(float(long_side) / max(h, w), float(max_scale))
|
|
110
|
+
return max(16, int(round(h * s / 16.0)) * 16), max(16, int(round(w * s / 16.0)) * 16)
|
|
111
|
+
|
|
112
|
+
|
|
113
|
+
def parse_size(v):
|
|
114
|
+
"""--size value: an int (long side, px) or 'native' (-> 0 = the image's own long side). v2, 22 Oct."""
|
|
115
|
+
s = str(v).strip().lower()
|
|
116
|
+
return 0 if s in ("native", "0") else int(s)
|
|
117
|
+
|
|
118
|
+
|
|
119
|
+
def parse_list_size(line):
|
|
120
|
+
"""Optional third field of a list line: this crop's own size (long side, px), e.g. to extract a panned
|
|
121
|
+
view at the same scale as the region it came from. None if absent."""
|
|
122
|
+
parts = line.rstrip("\n").split("\t")
|
|
123
|
+
return int(parts[2]) if len(parts) > 2 and parts[2].strip() else None
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
def parse_list_line(line):
|
|
127
|
+
"""'path' or 'path<TAB>x0,y0,x1,y1' -> (path, box or None)."""
|
|
128
|
+
parts = line.rstrip("\n").split("\t")
|
|
129
|
+
box = None
|
|
130
|
+
if len(parts) > 1 and parts[1].strip():
|
|
131
|
+
box = tuple(int(round(float(v))) for v in parts[1].split(","))
|
|
132
|
+
return parts[0], box
|
|
133
|
+
|
|
134
|
+
|
|
135
|
+
def cache_is_valid(npz_path, image_path, layer, size, model):
|
|
136
|
+
"""True if the cache file exists and was made from this image file with these settings."""
|
|
137
|
+
if not os.path.isfile(npz_path):
|
|
138
|
+
return False
|
|
139
|
+
try:
|
|
140
|
+
with np.load(npz_path, allow_pickle=False) as z:
|
|
141
|
+
st = os.stat(image_path)
|
|
142
|
+
return (int(z["layer"]) == int(layer) and int(z["size"]) == int(size)
|
|
143
|
+
and str(z["model"]) == str(model)
|
|
144
|
+
and int(z["source_size"]) == int(st.st_size)
|
|
145
|
+
and abs(float(z["source_mtime"]) - float(st.st_mtime)) < 1.0
|
|
146
|
+
and int(z.get("version", 0)) == CACHE_VERSION)
|
|
147
|
+
except Exception:
|
|
148
|
+
return False
|
|
149
|
+
|
|
150
|
+
|
|
151
|
+
def save_cache(npz_path, tokens, grid, image_hw, work_hw, layer, size, model, image_path, box=None):
|
|
152
|
+
st = os.stat(image_path)
|
|
153
|
+
if box is None:
|
|
154
|
+
box = (0, 0, int(image_hw[1]), int(image_hw[0]))
|
|
155
|
+
tmp = npz_path + ".tmp.npz"
|
|
156
|
+
np.savez(tmp, tokens=np.asarray(tokens, np.float16), grid=np.asarray(grid, np.int32),
|
|
157
|
+
image_hw=np.asarray(image_hw, np.int32), work_hw=np.asarray(work_hw, np.int32),
|
|
158
|
+
scale=np.array([work_hw[0] / image_hw[0], work_hw[1] / image_hw[1]], np.float64),
|
|
159
|
+
layer=np.int32(layer), size=np.int32(size), model=np.str_(model),
|
|
160
|
+
source=np.str_(os.path.abspath(image_path)), source_size=np.int64(st.st_size),
|
|
161
|
+
source_mtime=np.float64(st.st_mtime), version=np.int32(CACHE_VERSION),
|
|
162
|
+
box=np.asarray(box, np.int32))
|
|
163
|
+
os.replace(tmp, npz_path)
|
|
164
|
+
|
|
165
|
+
|
|
166
|
+
def load_cache(npz_path):
|
|
167
|
+
with np.load(npz_path, allow_pickle=False) as z:
|
|
168
|
+
return {k: z[k] for k in z.files}
|
|
169
|
+
|
|
170
|
+
|
|
171
|
+
# ---------------------------------------------------------------- extraction (torch)
|
|
172
|
+
class Extractor:
|
|
173
|
+
def __init__(self, model_id, device, layer, tile=2048, overlap=128):
|
|
174
|
+
import torch
|
|
175
|
+
from transformers import AutoModel
|
|
176
|
+
self.torch = torch
|
|
177
|
+
self.dev = device
|
|
178
|
+
self.net = AutoModel.from_pretrained(model_id).to(device).eval()
|
|
179
|
+
self.dim = self.net.config.hidden_size
|
|
180
|
+
self.nreg = getattr(self.net.config, "num_register_tokens", 4)
|
|
181
|
+
self.layer, self.tile, self.ov = int(layer), int(tile), int(overlap)
|
|
182
|
+
|
|
183
|
+
def _one(self, x):
|
|
184
|
+
torch = self.torch
|
|
185
|
+
th, tw = x.shape[0] // 16, x.shape[1] // 16
|
|
186
|
+
px = torch.from_numpy(np.ascontiguousarray(x.transpose(2, 0, 1)))[None].to(self.dev)
|
|
187
|
+
with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16, enabled=self.dev == "cuda"):
|
|
188
|
+
hs = self.net(pixel_values=px, output_hidden_states=True).hidden_states
|
|
189
|
+
t = hs[self.layer + 1][0, 1 + self.nreg:, :].float().reshape(th, tw, self.dim)
|
|
190
|
+
out = t.cpu().numpy()
|
|
191
|
+
del hs, t, px
|
|
192
|
+
return out
|
|
193
|
+
|
|
194
|
+
def grid(self, x):
|
|
195
|
+
"""x: normalised HxWx3, multiples of 16 -> [gh, gw, D] float32."""
|
|
196
|
+
H, W = x.shape[:2]
|
|
197
|
+
gh, gw = H // 16, W // 16
|
|
198
|
+
if H <= self.tile and W <= self.tile:
|
|
199
|
+
return self._one(x)
|
|
200
|
+
OV, STEP = self.ov, self.tile - 2 * self.ov
|
|
201
|
+
out = np.zeros((gh, gw, self.dim), np.float32)
|
|
202
|
+
for y0 in range(0, H, STEP):
|
|
203
|
+
for x0 in range(0, W, STEP):
|
|
204
|
+
ya, xa = max(0, y0 - OV), max(0, x0 - OV)
|
|
205
|
+
t = x[ya:min(H, y0 + STEP + OV), xa:min(W, x0 + STEP + OV)]
|
|
206
|
+
t = t[:t.shape[0] // 16 * 16, :t.shape[1] // 16 * 16]
|
|
207
|
+
T = self._one(t)
|
|
208
|
+
ky, kx = (y0 - ya) // 16, (x0 - xa) // 16
|
|
209
|
+
ny, nx = min(STEP, H - y0) // 16, min(STEP, W - x0) // 16
|
|
210
|
+
out[y0 // 16:y0 // 16 + ny, x0 // 16:x0 // 16 + nx] = T[ky:ky + ny, kx:kx + nx]
|
|
211
|
+
return out
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
def tile_edges(n, size, anchor):
|
|
215
|
+
"""Tile boundaries along one axis of length n (multiples of 16): tiles of `size` laid so that one tile
|
|
216
|
+
starts at `anchor`; partial tiles at both ends."""
|
|
217
|
+
a = int(anchor) % int(size)
|
|
218
|
+
edges = sorted({0, int(n)} | set(range(a, int(n), int(size))))
|
|
219
|
+
return [(e0, e1) for e0, e1 in zip(edges[:-1], edges[1:]) if e1 > e0]
|
|
220
|
+
|
|
221
|
+
|
|
222
|
+
def grid_tiled_like(ex, x, th, tw, ay, ax):
|
|
223
|
+
"""Tokens of x (multiples of 16) in NON-overlapping tiles of the zoomed region's own size (th x tw), laid
|
|
224
|
+
so that one tile coincides with the region (anchor ay, ax; multiples of 16). The region's tokens are then
|
|
225
|
+
reproduced (same pixels, same patch grid phase, same context) and the other tiles see the same kind of
|
|
226
|
+
context as the region did (18 Oct: DINOv3 tokens changed by cosine ~0.13 when only the context window
|
|
227
|
+
changed, and ~0.11 with a half-patch grid shift)."""
|
|
228
|
+
H, W = x.shape[:2]
|
|
229
|
+
out = np.zeros((H // 16, W // 16, ex.dim), np.float32)
|
|
230
|
+
for y0, y1 in tile_edges(H, th, ay):
|
|
231
|
+
for x0, x1 in tile_edges(W, tw, ax):
|
|
232
|
+
out[y0 // 16:y1 // 16, x0 // 16:x1 // 16] = ex._one(np.ascontiguousarray(x[y0:y1, x0:x1]))
|
|
233
|
+
return out
|
|
234
|
+
|
|
235
|
+
|
|
236
|
+
def extract_one(ex, image_path, size, box=None):
|
|
237
|
+
"""Whole image (box None) or the crop box x0,y0,x1,y1 (upsampled, at most MAX_CROP_SCALE x).
|
|
238
|
+
The returned (H, W) is always the whole image's size."""
|
|
239
|
+
import cv2
|
|
240
|
+
rgb = load_rgb8(image_path)
|
|
241
|
+
H, W = rgb.shape[:2]
|
|
242
|
+
if box is not None:
|
|
243
|
+
x0, y0, x1, y1 = box
|
|
244
|
+
x0, y0, x1, y1 = max(0, x0), max(0, y0), min(W, x1), min(H, y1)
|
|
245
|
+
if x1 - x0 < 4 or y1 - y0 < 4:
|
|
246
|
+
raise ValueError(f"crop box {box} is empty")
|
|
247
|
+
rgb = rgb[y0:y1, x0:x1]
|
|
248
|
+
nh, nw = crop_working_size(y1 - y0, x1 - x0, size)
|
|
249
|
+
else:
|
|
250
|
+
nh, nw = working_size(H, W, size)
|
|
251
|
+
rgb_ = cv2.resize(rgb, (nw, nh), interpolation=cv2.INTER_CUBIC)
|
|
252
|
+
x = (rgb_.astype(np.float32) / 255.0 - MEAN) / STD
|
|
253
|
+
G = ex.grid(x)
|
|
254
|
+
gh, gw = G.shape[:2]
|
|
255
|
+
return G.reshape(gh * gw, -1).astype(np.float16), (gh, gw), (H, W), (nh, nw)
|
|
256
|
+
|
|
257
|
+
|
|
258
|
+
def _ui_module():
|
|
259
|
+
"""biorag_annotate_v2 (or biorag_annotate_v2_dev when this file is the _dev copy) for its pure helpers."""
|
|
260
|
+
import importlib
|
|
261
|
+
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
262
|
+
sfx = "_dev" if os.path.basename(__file__).endswith("_dev.py") else ""
|
|
263
|
+
return importlib.import_module("biorag_annotate_v2" + sfx)
|
|
264
|
+
|
|
265
|
+
|
|
266
|
+
# ---------------------------------------------------------------- trained DINOMask head (optional view)
|
|
267
|
+
# A DINOMask run folder holds head.pt + results.json (classes). The head was trained on the LAST-layer
|
|
268
|
+
# tokens (last_hidden_state) of DINOv3 ViT-L/16 at the run's size (2048 by default) with 4-shift
|
|
269
|
+
# averaging, so its features are computed here exactly like that (dinomask_predict_rot4_v1.py probs(),
|
|
270
|
+
# without the test-time rotation) - NOT from the layer-18 token cache.
|
|
271
|
+
# classes : argmax of the guided-filtered 4-shift-averaged probabilities, uint8, at the working size
|
|
272
|
+
# hidden : the head's 512-d activations after GELU, (gh*gw, 512) float16, from the UNSHIFTED pass
|
|
273
|
+
def dinomask_working_size(h, w, long_side):
|
|
274
|
+
"""As DINOMask probs(): s = size / max(H, W); sides floored to multiples of 16."""
|
|
275
|
+
s = float(long_side) / max(h, w)
|
|
276
|
+
return max(16, int(h * s) // 16 * 16), max(16, int(w * s) // 16 * 16)
|
|
277
|
+
|
|
278
|
+
|
|
279
|
+
def dinomask_run_id(run_dir):
|
|
280
|
+
st = os.stat(os.path.join(run_dir, "head.pt"))
|
|
281
|
+
import hashlib
|
|
282
|
+
return hashlib.md5(f"{os.path.abspath(run_dir)}|{st.st_size}|{int(st.st_mtime)}".encode()).hexdigest()[:10]
|
|
283
|
+
|
|
284
|
+
|
|
285
|
+
def dinomask_cache_path(cache_dir, image_path, run_dir, size):
|
|
286
|
+
return os.path.join(cache_dir, f"{os.path.basename(image_path)}__dinomask_{dinomask_run_id(run_dir)}_s{int(size)}.npz")
|
|
287
|
+
|
|
288
|
+
|
|
289
|
+
def dinomask_classes(run_dir):
|
|
290
|
+
import ast
|
|
291
|
+
c = json.load(open(os.path.join(run_dir, "results.json")))["classes"]
|
|
292
|
+
return ast.literal_eval(c) if isinstance(c, str) else list(c)
|
|
293
|
+
|
|
294
|
+
|
|
295
|
+
def dinomask_cache_is_valid(npz_path, image_path):
|
|
296
|
+
if not os.path.isfile(npz_path):
|
|
297
|
+
return False
|
|
298
|
+
try:
|
|
299
|
+
with np.load(npz_path, allow_pickle=False) as z:
|
|
300
|
+
st = os.stat(image_path)
|
|
301
|
+
return int(z["source_size"]) == int(st.st_size) and abs(float(z["source_mtime"]) - float(st.st_mtime)) < 1.0
|
|
302
|
+
except Exception:
|
|
303
|
+
return False
|
|
304
|
+
|
|
305
|
+
|
|
306
|
+
class DinoMaskRunner:
|
|
307
|
+
def __init__(self, run_dir, model_id, device):
|
|
308
|
+
import torch
|
|
309
|
+
import torch.nn as nn
|
|
310
|
+
from transformers import AutoModel
|
|
311
|
+
self.torch, self.dev = torch, device
|
|
312
|
+
self.classes = dinomask_classes(run_dir)
|
|
313
|
+
K = len(self.classes)
|
|
314
|
+
self.net = AutoModel.from_pretrained(model_id).to(device).eval()
|
|
315
|
+
self.nreg = self.net.config.num_register_tokens
|
|
316
|
+
D = self.net.config.hidden_size
|
|
317
|
+
self.head = nn.Sequential(nn.LayerNorm(D), nn.Linear(D, 512), nn.GELU(), nn.Dropout(0.2), nn.Linear(512, K)).to(device)
|
|
318
|
+
sd = torch.load(os.path.join(run_dir, "head.pt"), map_location=device)
|
|
319
|
+
self.head.load_state_dict({k.replace("net.", ""): v for k, v in sd.items()})
|
|
320
|
+
self.head.eval()
|
|
321
|
+
self.K = K
|
|
322
|
+
|
|
323
|
+
def run(self, rgb, size):
|
|
324
|
+
import cv2
|
|
325
|
+
gfilter = _ui_module().gfilter
|
|
326
|
+
torch = self.torch
|
|
327
|
+
H, W = rgb.shape[:2]
|
|
328
|
+
nh, nw = dinomask_working_size(H, W, size)
|
|
329
|
+
rgb_ = cv2.resize(rgb, (nw, nh), interpolation=cv2.INTER_CUBIC)
|
|
330
|
+
base = (rgb_.astype(np.float32) / 255 - MEAN) / STD
|
|
331
|
+
acc = np.zeros((nh, nw, self.K), np.float32)
|
|
332
|
+
hidden = None
|
|
333
|
+
for dx, dy in ((0, 0), (8, 0), (0, 8), (8, 8)):
|
|
334
|
+
x = np.pad(base, ((dy, 16 - dy if dy else 0), (dx, 16 - dx if dx else 0), (0, 0)), mode="edge")
|
|
335
|
+
gh, gw = x.shape[0] // 16, x.shape[1] // 16
|
|
336
|
+
px = torch.from_numpy(x[:gh * 16, :gw * 16].transpose(2, 0, 1).copy())[None].to(self.dev)
|
|
337
|
+
with torch.no_grad(), torch.autocast("cuda", dtype=torch.float16, enabled=self.dev == "cuda"):
|
|
338
|
+
T = self.net(pixel_values=px).last_hidden_state[0, 1 + self.nreg:].float()
|
|
339
|
+
with torch.no_grad():
|
|
340
|
+
h = self.head[2](self.head[1](self.head[0](T))) # LayerNorm -> Linear -> GELU
|
|
341
|
+
pr = torch.softmax(self.head[4](h), -1).cpu().numpy().reshape(gh, gw, self.K)
|
|
342
|
+
if dx == 0 and dy == 0:
|
|
343
|
+
hidden = (h.cpu().numpy().astype(np.float16), (gh, gw))
|
|
344
|
+
acc += cv2.resize(pr, (gw * 16, gh * 16), interpolation=cv2.INTER_LINEAR)[dy:dy + nh, dx:dx + nw]
|
|
345
|
+
P = gfilter(cv2.cvtColor(rgb_, cv2.COLOR_RGB2GRAY).astype(np.float32) / 255, acc / 4)
|
|
346
|
+
return P.argmax(-1).astype(np.uint8), hidden[0], hidden[1], (H, W), (nh, nw)
|
|
347
|
+
|
|
348
|
+
|
|
349
|
+
def main_dinomask(a, paths):
|
|
350
|
+
import json as _json # noqa: F401
|
|
351
|
+
os.makedirs(a.cache_dir, exist_ok=True)
|
|
352
|
+
n = len(paths)
|
|
353
|
+
todo, n_ok, n_fail = [], 0, 0
|
|
354
|
+
for i, p in enumerate(paths, 1):
|
|
355
|
+
cp = dinomask_cache_path(a.cache_dir, p, a.dinomask_run, a.size)
|
|
356
|
+
if not a.force and dinomask_cache_is_valid(cp, p):
|
|
357
|
+
n_ok += 1
|
|
358
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tcached\t0.0\t{os.path.basename(p)}\t", flush=True)
|
|
359
|
+
else:
|
|
360
|
+
todo.append((i, p, cp))
|
|
361
|
+
if todo:
|
|
362
|
+
dev = a.device
|
|
363
|
+
if dev == "auto":
|
|
364
|
+
import torch
|
|
365
|
+
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
|
366
|
+
print(f"BRA_INFO loading DINOMask head {a.dinomask_run} + {a.model} on {dev} (last layer, size {a.size})", flush=True)
|
|
367
|
+
try:
|
|
368
|
+
R = DinoMaskRunner(a.dinomask_run, a.model, dev)
|
|
369
|
+
except Exception as e:
|
|
370
|
+
for i, p, _ in todo:
|
|
371
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tfailed\t0.0\t{os.path.basename(p)}\tDINOMask: {e}", flush=True)
|
|
372
|
+
print(f"BRA_DONE {n_ok} {len(todo)}", flush=True)
|
|
373
|
+
return 1
|
|
374
|
+
for i, p, cp in todo:
|
|
375
|
+
t0 = time.time()
|
|
376
|
+
try:
|
|
377
|
+
lab, hid, grid, hw, whw = R.run(load_rgb8(p), a.size)
|
|
378
|
+
st = os.stat(p)
|
|
379
|
+
tmp = cp + ".tmp.npz"
|
|
380
|
+
np.savez(tmp, classes_map=lab, hidden=hid, grid=np.asarray(grid, np.int32),
|
|
381
|
+
image_hw=np.asarray(hw, np.int32), work_hw=np.asarray(whw, np.int32),
|
|
382
|
+
classes=np.array(R.classes), run=np.str_(os.path.abspath(a.dinomask_run)),
|
|
383
|
+
size=np.int32(a.size), model=np.str_(a.model), source_size=np.int64(st.st_size),
|
|
384
|
+
source_mtime=np.float64(st.st_mtime))
|
|
385
|
+
os.replace(tmp, cp)
|
|
386
|
+
n_ok += 1
|
|
387
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tok\t{time.time() - t0:.1f}\t{os.path.basename(p)}\t", flush=True)
|
|
388
|
+
except Exception as e:
|
|
389
|
+
n_fail += 1
|
|
390
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tfailed\t{time.time() - t0:.1f}\t{os.path.basename(p)}\t{str(e)[:200]}", flush=True)
|
|
391
|
+
if dev == "cuda":
|
|
392
|
+
print(f"BRA_INFO peak GPU memory {R.torch.cuda.max_memory_allocated() / 2**30:.2f} GB", flush=True)
|
|
393
|
+
print(f"BRA_DONE {n_ok} {n_fail}", flush=True)
|
|
394
|
+
return 0 if n_fail == 0 else 2
|
|
395
|
+
|
|
396
|
+
|
|
397
|
+
# ---------------------------------------------------------------- apply a selection rule at fine scale (18 Oct)
|
|
398
|
+
# The rule (biorag_annotate_v1.save_rule) was built on zoomed-region tokens. Here every target image is
|
|
399
|
+
# tokenised WHOLE at the SAME scale as that region (tiled, --tile px with 128 px overlap), the rule is applied,
|
|
400
|
+
# and only the per-patch result (selected bool + score float16) is written; the fine tokens are not cached
|
|
401
|
+
# (a 5400 x 3600 image at 1.5x is ~160k tokens, ~330 MB as float16).
|
|
402
|
+
def applied_cache_path(cache_dir, image_path, rule_id):
|
|
403
|
+
return os.path.join(cache_dir, f"{os.path.basename(image_path)}__applied_{rule_id}.npz")
|
|
404
|
+
|
|
405
|
+
|
|
406
|
+
def fine_working_size(h, w, sy, sx):
|
|
407
|
+
"""Whole image at the region's scale (sy, sx = working px per image px), sides multiples of 16."""
|
|
408
|
+
return max(16, int(round(h * float(sy) / 16.0)) * 16), max(16, int(round(w * float(sx) / 16.0)) * 16)
|
|
409
|
+
|
|
410
|
+
|
|
411
|
+
def _rss_gb():
|
|
412
|
+
try:
|
|
413
|
+
import resource
|
|
414
|
+
return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 2**20
|
|
415
|
+
except Exception:
|
|
416
|
+
return float("nan")
|
|
417
|
+
|
|
418
|
+
|
|
419
|
+
def main_apply(a, paths):
|
|
420
|
+
import cv2
|
|
421
|
+
load_rule, rule_scores = _ui_module().load_rule, _ui_module().rule_scores
|
|
422
|
+
rule = load_rule(a.apply_rule)
|
|
423
|
+
rid = rule["rule_id"]
|
|
424
|
+
sy, sx = (float(v) for v in rule["scale"])
|
|
425
|
+
th, tw = (int(v) for v in rule.get("region_work_hw", (a.tile, a.tile)))
|
|
426
|
+
src = rule.get("from_region") or [None, None]
|
|
427
|
+
src_name, src_box = src[0], (tuple(int(v) for v in src[1]) if src[1] is not None else None)
|
|
428
|
+
os.makedirs(a.cache_dir, exist_ok=True)
|
|
429
|
+
n = len(paths)
|
|
430
|
+
todo, n_ok, n_fail = [], 0, 0
|
|
431
|
+
for i, p in enumerate(paths, 1):
|
|
432
|
+
cp = applied_cache_path(a.cache_dir, p, rid)
|
|
433
|
+
if not a.force and dinomask_cache_is_valid(cp, p): # same source-file check
|
|
434
|
+
n_ok += 1
|
|
435
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tcached\t0.0\t{os.path.basename(p)}\t", flush=True)
|
|
436
|
+
else:
|
|
437
|
+
todo.append((i, p, cp))
|
|
438
|
+
if todo:
|
|
439
|
+
dev = a.device
|
|
440
|
+
if dev == "auto":
|
|
441
|
+
import torch
|
|
442
|
+
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
|
443
|
+
layer = int(rule["layer"])
|
|
444
|
+
print(f"BRA_INFO rule {rule['kind']} ({rid}): loading {a.model} on {dev} (layer {layer}, scale {sx:.3f})", flush=True)
|
|
445
|
+
try:
|
|
446
|
+
ex = Extractor(a.model, dev, layer, tile=a.tile)
|
|
447
|
+
except Exception as e:
|
|
448
|
+
for i, p, _ in todo:
|
|
449
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tfailed\t0.0\t{os.path.basename(p)}\tmodel: {e}", flush=True)
|
|
450
|
+
print(f"BRA_DONE {n_ok} {len(todo)}", flush=True)
|
|
451
|
+
return 1
|
|
452
|
+
for i, p, cp in todo:
|
|
453
|
+
t0 = time.time()
|
|
454
|
+
try:
|
|
455
|
+
rgb = load_rgb8(p)
|
|
456
|
+
H, W = rgb.shape[:2]
|
|
457
|
+
nh, nw = fine_working_size(H, W, sy, sx)
|
|
458
|
+
x = (cv2.resize(rgb, (nw, nh), interpolation=cv2.INTER_CUBIC).astype(np.float32) / 255.0 - MEAN) / STD
|
|
459
|
+
del rgb
|
|
460
|
+
# tile like the region: on the region's own image the patch grid is put in phase with the region
|
|
461
|
+
# (the <= 15 fine px before the first full patch are left out) and one tile IS the region
|
|
462
|
+
ox = oy = 0
|
|
463
|
+
if os.path.basename(p) == src_name and src_box is not None:
|
|
464
|
+
ox, oy = int(round(src_box[0] * nw / W)), int(round(src_box[1] * nh / H))
|
|
465
|
+
px, py = ox % 16, oy % 16
|
|
466
|
+
Hx, Wx = (nh - py) // 16 * 16, (nw - px) // 16 * 16
|
|
467
|
+
G = grid_tiled_like(ex, x[py:py + Hx, px:px + Wx], th, tw, oy - py, ox - px)
|
|
468
|
+
del x
|
|
469
|
+
gh, gw = G.shape[:2]
|
|
470
|
+
box = (int(round(px * W / nw)), int(round(py * H / nh)),
|
|
471
|
+
min(W, int(round((px + Wx) * W / nw))), min(H, int(round((py + Hx) * H / nh))))
|
|
472
|
+
t1 = time.time()
|
|
473
|
+
sel, score = rule_scores(G.reshape(gh * gw, -1), (gh, gw), rule)
|
|
474
|
+
del G
|
|
475
|
+
st = os.stat(p)
|
|
476
|
+
tmp = cp + ".tmp.npz"
|
|
477
|
+
np.savez(tmp, sel=sel.reshape(gh, gw), score=score.reshape(gh, gw).astype(np.float16),
|
|
478
|
+
grid=np.asarray((gh, gw), np.int32), image_hw=np.asarray((H, W), np.int32),
|
|
479
|
+
work_hw=np.asarray((Hx, Wx), np.int32), rule_id=np.str_(rid), box=np.asarray(box, np.int32),
|
|
480
|
+
source_size=np.int64(st.st_size), source_mtime=np.float64(st.st_mtime))
|
|
481
|
+
os.replace(tmp, cp)
|
|
482
|
+
n_ok += 1
|
|
483
|
+
print(f"BRA_INFO {os.path.basename(p)}: {gh}x{gw} fine patches ({Wx}x{Hx} px, tiles {tw}x{th}"
|
|
484
|
+
f"{', anchored on the region' if ox or oy else ''}), tokens "
|
|
485
|
+
f"{t1 - t0:.1f} s, rule {time.time() - t1:.1f} s, {int(sel.sum())} selected, peak RSS "
|
|
486
|
+
f"{_rss_gb():.2f} GB", flush=True)
|
|
487
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tok\t{time.time() - t0:.1f}\t{os.path.basename(p)}\t", flush=True)
|
|
488
|
+
except Exception as e:
|
|
489
|
+
n_fail += 1
|
|
490
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tfailed\t{time.time() - t0:.1f}\t{os.path.basename(p)}\t{str(e)[:200]}", flush=True)
|
|
491
|
+
if dev == "cuda":
|
|
492
|
+
try:
|
|
493
|
+
ex.torch.cuda.empty_cache()
|
|
494
|
+
except Exception:
|
|
495
|
+
pass
|
|
496
|
+
if dev == "cuda":
|
|
497
|
+
print(f"BRA_INFO peak GPU memory {ex.torch.cuda.max_memory_allocated() / 2**30:.2f} GB", flush=True)
|
|
498
|
+
print(f"BRA_DONE {n_ok} {n_fail}", flush=True)
|
|
499
|
+
return 0 if n_fail == 0 else 2
|
|
500
|
+
|
|
501
|
+
|
|
502
|
+
def main(argv=None):
|
|
503
|
+
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
504
|
+
ap.add_argument("--images", nargs="*", default=[], help="image files")
|
|
505
|
+
ap.add_argument("--image_list", default="", help="text file, one image path per line (alternative to --images)")
|
|
506
|
+
ap.add_argument("--cache_dir", required=True)
|
|
507
|
+
ap.add_argument("--size", type=parse_size, default=1536, help="long side in px (1024 / 1536 / 2048) or 'native' "
|
|
508
|
+
"(the image's own long side; tiled above --tile)")
|
|
509
|
+
ap.add_argument("--layer", type=int, default=18, help="0-based transformer block, 0-23 for ViT-L (e.g. 12 / 18 / 23)")
|
|
510
|
+
ap.add_argument("--model", default=DEFAULT_MODEL)
|
|
511
|
+
ap.add_argument("--device", default="auto", help="auto / cuda / cpu")
|
|
512
|
+
ap.add_argument("--tile", type=int, default=2048, help="tile larger working images (px)")
|
|
513
|
+
ap.add_argument("--force", action="store_true", help="re-extract even if a valid cache exists")
|
|
514
|
+
ap.add_argument("--apply_rule", default="", help="selection rule npz (BioRAG Annotate): tokenise each image "
|
|
515
|
+
"whole at the rule's scale and write only the per-patch result")
|
|
516
|
+
ap.add_argument("--dinomask_run", default="", help="DINOMask run folder (head.pt + results.json): compute its "
|
|
517
|
+
"class map + hidden activations instead of tokens (last layer, --size, 4-shift averaging)")
|
|
518
|
+
a = ap.parse_args(argv)
|
|
519
|
+
|
|
520
|
+
items = [(p, None) for p in a.images]
|
|
521
|
+
sizes = [None] * len(items)
|
|
522
|
+
if a.image_list:
|
|
523
|
+
with open(a.image_list) as fh:
|
|
524
|
+
lines = [l for l in fh if l.strip()]
|
|
525
|
+
items += [parse_list_line(l) for l in lines]
|
|
526
|
+
sizes += [parse_list_size(l) for l in lines]
|
|
527
|
+
paths = [p for p, _ in items]
|
|
528
|
+
if not paths:
|
|
529
|
+
print("BRA_DONE 0 0", flush=True)
|
|
530
|
+
return 0
|
|
531
|
+
if a.dinomask_run:
|
|
532
|
+
return main_dinomask(a, paths)
|
|
533
|
+
if a.apply_rule:
|
|
534
|
+
return main_apply(a, paths)
|
|
535
|
+
os.makedirs(a.cache_dir, exist_ok=True)
|
|
536
|
+
todo = []
|
|
537
|
+
n = len(paths)
|
|
538
|
+
for i, ((p, box), sz) in enumerate(zip(items, sizes), 1):
|
|
539
|
+
sz = sz or a.size
|
|
540
|
+
cp = cache_path(a.cache_dir, p, a.layer, sz, box)
|
|
541
|
+
if not a.force and cache_is_valid(cp, p, a.layer, sz, a.model):
|
|
542
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tcached\t0.0\t{os.path.basename(p)}\t", flush=True)
|
|
543
|
+
else:
|
|
544
|
+
todo.append((i, p, cp, box, sz))
|
|
545
|
+
n_ok, n_fail = n - len(todo), 0
|
|
546
|
+
if todo:
|
|
547
|
+
dev = a.device
|
|
548
|
+
if dev == "auto":
|
|
549
|
+
import torch
|
|
550
|
+
dev = "cuda" if torch.cuda.is_available() else "cpu"
|
|
551
|
+
print(f"BRA_INFO loading {a.model} on {dev} (layer {a.layer}, size {a.size})", flush=True)
|
|
552
|
+
try:
|
|
553
|
+
ex = Extractor(a.model, dev, a.layer, tile=a.tile)
|
|
554
|
+
except Exception as e:
|
|
555
|
+
for i, p, _, _, _ in todo:
|
|
556
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tfailed\t0.0\t{os.path.basename(p)}\tmodel: {e}", flush=True)
|
|
557
|
+
print(f"BRA_DONE {n_ok} {len(todo)}", flush=True)
|
|
558
|
+
return 1
|
|
559
|
+
for i, p, cp, box, sz in todo:
|
|
560
|
+
t0 = time.time()
|
|
561
|
+
try:
|
|
562
|
+
T, grid, hw, whw = extract_one(ex, p, sz, box)
|
|
563
|
+
save_cache(cp, T, grid, hw, whw, a.layer, sz, a.model, p, box)
|
|
564
|
+
n_ok += 1
|
|
565
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tok\t{time.time() - t0:.1f}\t{os.path.basename(p)}\t", flush=True)
|
|
566
|
+
except Exception as e:
|
|
567
|
+
n_fail += 1
|
|
568
|
+
msg = str(e).replace("\n", " ")[:200]
|
|
569
|
+
print(f"BRA_PROGRESS\t{i}\t{n}\tfailed\t{time.time() - t0:.1f}\t{os.path.basename(p)}\t{msg}", flush=True)
|
|
570
|
+
if dev == "cuda":
|
|
571
|
+
try:
|
|
572
|
+
ex.torch.cuda.empty_cache()
|
|
573
|
+
except Exception:
|
|
574
|
+
pass
|
|
575
|
+
if dev == "cuda":
|
|
576
|
+
try:
|
|
577
|
+
print(f"BRA_INFO peak GPU memory {ex.torch.cuda.max_memory_allocated() / 2**30:.2f} GB", flush=True)
|
|
578
|
+
except Exception:
|
|
579
|
+
pass
|
|
580
|
+
print(f"BRA_DONE {n_ok} {n_fail}", flush=True)
|
|
581
|
+
return 0 if n_fail == 0 else 2
|
|
582
|
+
|
|
583
|
+
|
|
584
|
+
if __name__ == "__main__":
|
|
585
|
+
sys.exit(main())
|