dtourolle
06422a07db
AI Denoise's Apply switch becomes Method: Bilinear, Fast, Medium, Best, default Best, so an untouched raw writes nothing and develops through the mixture. `apply` is still read and never written: 0 is Bilinear, 1 keeps a network already chosen. - Best is the mixture of a flat and an edge expert with a learned gate; Medium and Fast are students distilled from it. 2.48 s, 0.79 s and 0.57 s for a 20 MP frame on TensorRT fp16. - Each network carries its own tile border (256 for the mixture, 192 for the students) through `dr_denoise::Shipped` and `TileNet::halo`. - The file is hashed once at open and each network keys its own cached result; Bilinear keeps the result in memory for the way back. - Each has an .a16w16 sibling for the Hexagon: 0.00 dB on the 6D gate, at most 0.11 dB with the noise scaled x0.5 to x4. - APK BUNDLED 19 -> 23; the PKGBUILD installs all three.
266 lines
12 KiB
Python
266 lines
12 KiB
Python
"""The body of quantise-models.sh; see there. Run through it, not directly."""
|
||
|
||
import glob
|
||
import os
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
import numpy as np
|
||
import onnx
|
||
import onnxruntime as ort
|
||
from onnx import version_converter
|
||
from onnxruntime.quantization import CalibrationDataReader, CalibrationMethod, QuantType, quantize_static
|
||
from onnxruntime.quantization.calibrate import create_calibrator, save_tensors_data
|
||
from onnxruntime.quantization.execution_providers.qnn import get_qnn_qdq_config
|
||
from PIL import Image
|
||
|
||
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
||
import htp_graph # noqa: E402
|
||
|
||
MODELS = Path(__file__).resolve().parents[1] / "models"
|
||
PHOTOS = 300 # calibration photographs; face crops and keypoints come from fewer
|
||
CHUNK = 4 # inputs whose activations are held at once (scrfd_10g: ~1 GB each)
|
||
|
||
Q = QuantType
|
||
FORMS = {"int8": (Q.QUInt8, Q.QInt8), "a16w8": (Q.QUInt16, Q.QInt8), "a16w16": (Q.QUInt16, Q.QInt16)}
|
||
|
||
# Per model: where it lives, the form `Rung::form` gives its role, the exact
|
||
# rewrites its graph needs, and nodes that stay float on the CPU because one
|
||
# scale cannot serve the tensor (the segmenter's rows: boxes in pixels beside
|
||
# scores in 0..1) or because the HTP's 16-bit arithmetic drifts there (the
|
||
# scene model's attention). Measured, inference.md §1.5.
|
||
TABLE = {
|
||
"scrfd_500m_640": dict(dir="face", form="a16w8", feed="scrfd"),
|
||
"scrfd_2.5g_640": dict(dir="face", form="a16w8", feed="scrfd"),
|
||
"scrfd_10g_640": dict(dir="face", form="a16w8", feed="scrfd"),
|
||
"2d106det_b1": dict(dir="face", form="a16w8", feed="landmarks"),
|
||
"yolo26n-seg": dict(dir="segment", form="a16w16", feed="yolo", float_from="/model.23/Concat_4"),
|
||
"yolo26s-sem-ade20k": dict(dir="scene", form="a16w16", feed="yolo",
|
||
float_nodes=["/model.10/m/m.0/attn/MatMul", "/model.10/m/m.0/attn/Softmax",
|
||
"/model.10/m/m.0/attn/MatMul_1"]),
|
||
"migan-512": dict(dir="inpaint", form="a16w16", feed="migan"),
|
||
"xfeat-1024": dict(dir="keypoints", form="int8", feed="xfeat", rewrites=["unfold", "resize"]),
|
||
"xfeat-768": dict(dir="keypoints", form="int8", feed="xfeat", rewrites=["unfold", "resize"]),
|
||
"mosaic-fast-1408": dict(dir="denoise", form="a16w16", feed=None, rewrites=["bayer"]),
|
||
"mosaic-medium-1408": dict(dir="denoise", form="a16w16", feed=None, rewrites=["bayer"]),
|
||
# Exported with its packs as SpaceToDepth already; the rewrite finds
|
||
# nothing to do.
|
||
"mosaic-best-1408": dict(dir="denoise", form="a16w16", feed=None, rewrites=["bayer"]),
|
||
}
|
||
|
||
|
||
# ---- the app's samplers (dr-face Letterbox, dr-segment Letterbox::sample, align.rs) ----
|
||
|
||
def load(p):
|
||
return np.asarray(Image.open(p).convert("RGB"), np.float32) / 255.0
|
||
|
||
|
||
def bilinear(img, sx, sy):
|
||
h, w = img.shape[:2]
|
||
x0, y0 = np.floor(sx).astype(np.int64), np.floor(sy).astype(np.int64)
|
||
fx, fy = (sx - x0)[..., None], (sy - y0)[..., None]
|
||
c = lambda a, n: np.clip(a, 0, n - 1) # noqa: E731 - neighbours clamp at the edge
|
||
top = img[c(y0, h), c(x0, w)] * (1 - fx) + img[c(y0, h), c(x0 + 1, w)] * fx
|
||
bot = img[c(y0 + 1, h), c(x0, w)] * (1 - fx) + img[c(y0 + 1, h), c(x0 + 1, w)] * fx
|
||
return top * (1 - fy) + bot * fy
|
||
|
||
|
||
def chw(x):
|
||
return np.ascontiguousarray(x.transpose(2, 0, 1))[None].astype(np.float32)
|
||
|
||
|
||
def letterbox(img, edge, yolo):
|
||
h, w = img.shape[:2]
|
||
s = min(edge / w, edge / h)
|
||
px, py = (edge - w * s) / 2, (edge - h * s) / 2
|
||
ix, iy = np.meshgrid(np.arange(edge) + 0.5, np.arange(edge) + 0.5)
|
||
if yolo: # semantic.rs: no -0.5, pad 0.5, 0..1
|
||
sx, sy = (ix - px) / s, (iy - py) / s
|
||
out = bilinear(img, sx, sy)
|
||
out[(sx < 0) | (sx >= w) | (sy < 0) | (sy >= h)] = 0.5
|
||
else: # detect.rs: -0.5, pad 114, (v·255 − 127.5)/128
|
||
sx, sy = (ix - px) / s - 0.5, (iy - py) / s - 0.5
|
||
out = (bilinear(img, sx, sy) * 255 - 127.5) / 128
|
||
out[(sx < -0.5) | (sx > w - 0.5) | (sy < -0.5) | (sy > h - 0.5)] = (114 - 127.5) / 128
|
||
return chw(out), (s, px, py)
|
||
|
||
|
||
def crop_box(img, x0, y0, bw, bh, ow, oh):
|
||
"""align.rs crop_box: output (u+.5) → source, −0.5, bilinear, outside black."""
|
||
u, v = np.meshgrid(np.arange(ow) + 0.5, np.arange(oh) + 0.5)
|
||
sx, sy = x0 + u * bw / ow - 0.5, y0 + v * bh / oh - 0.5
|
||
out = bilinear(img, sx, sy)
|
||
h, w = img.shape[:2]
|
||
out[(sx < -1) | (sx > w) | (sy < -1) | (sy > h)] = 0
|
||
return out
|
||
|
||
|
||
def scrfd_boxes(outs, s, px, py):
|
||
"""detect.rs decode: score ≥ 0.5, greedy NMS at 0.4, min side 24 px."""
|
||
fmc = len(outs) // 3
|
||
boxes, scores = [], []
|
||
for i, st in enumerate([8, 16, 32, 64][:fmc]):
|
||
sc, bx = outs[i].reshape(-1), outs[fmc + i].reshape(-1, 4)
|
||
idx = np.nonzero(sc >= 0.5)[0]
|
||
cell = idx // 2
|
||
cx, cy = (cell % (640 // st)) * st, (cell // (640 // st)) * st
|
||
boxes.append(np.stack([cx - bx[idx, 0] * st, cy - bx[idx, 1] * st, cx + bx[idx, 2] * st, cy + bx[idx, 3] * st], 1))
|
||
scores.append(sc[idx])
|
||
b, sc = np.concatenate(boxes), np.concatenate(scores)
|
||
keep = []
|
||
for i in np.argsort(-sc):
|
||
x0 = np.maximum(b[i, 0], b[keep, 0]); y0 = np.maximum(b[i, 1], b[keep, 1])
|
||
x1 = np.minimum(b[i, 2], b[keep, 2]); y1 = np.minimum(b[i, 3], b[keep, 3])
|
||
inter = np.clip(x1 - x0, 0, None) * np.clip(y1 - y0, 0, None)
|
||
area = lambda r: (r[..., 2] - r[..., 0]) * (r[..., 3] - r[..., 1]) # noqa: E731
|
||
if not keep or (inter / (area(b[i]) + area(b[keep]) - inter)).max() <= 0.4:
|
||
keep.append(i)
|
||
b = (b[keep] - [px, py, px, py]) / s
|
||
return b[np.minimum(b[:, 2] - b[:, 0], b[:, 3] - b[:, 1]) >= 32]
|
||
|
||
|
||
# ---- one calibration input per photograph (or per face), as the app makes it ----
|
||
|
||
def feeds(kind, photos, model):
|
||
name = model.get_inputs()[0].name
|
||
if kind == "scrfd":
|
||
for p in photos:
|
||
yield {name: letterbox(load(p), 640, False)[0]}
|
||
elif kind == "landmarks": # landmarks.rs: 1.5× the box, square, 0..255
|
||
det = ort.InferenceSession(str(MODELS / "face/scrfd_10g_640.onnx"), providers=["CPUExecutionProvider"])
|
||
for p in photos:
|
||
img = load(p)
|
||
x, ctx = letterbox(img, 640, False)
|
||
for b in scrfd_boxes(det.run(None, {"input.1": x}), *ctx):
|
||
cx, cy, side = (b[0] + b[2]) / 2, (b[1] + b[3]) / 2, 1.5 * max(b[2] - b[0], b[3] - b[1])
|
||
yield {name: chw(crop_box(img, cx - side / 2, cy - side / 2, side, side, 192, 192) * 255)}
|
||
elif kind == "yolo":
|
||
for p in photos:
|
||
yield {name: letterbox(load(p), 640, True)[0]}
|
||
elif kind == "migan": # migan.rs: ch0 = known − 0.5, ch1–3 = (rgb·2 − 1)·known
|
||
rng = np.random.default_rng(7)
|
||
for p in photos:
|
||
img = load(p)
|
||
h, w = img.shape[:2]
|
||
e = min(h, w)
|
||
sq = Image.fromarray((img[(h - e) // 2:(h + e) // 2, (w - e) // 2:(w + e) // 2] * 255).astype(np.uint8))
|
||
img = np.asarray(sq.resize((512, 512), Image.BILINEAR), np.float32) / 255
|
||
known = np.ones((512, 512), np.float32)
|
||
for _ in range(rng.integers(1, 3)): # a panorama's unknown border: a wedge along one edge
|
||
side = rng.integers(4)
|
||
depth = np.linspace(rng.integers(20, 110), rng.integers(20, 110), 512).astype(int)
|
||
edge = np.arange(512)[:, None] < depth[None, :] # [depth, along]: inside the wedge
|
||
wedge = edge if side % 2 == 0 else edge[::-1] # top / bottom of a column
|
||
known[wedge if side < 2 else wedge.T] = 0 # or left / right of a row
|
||
x = np.concatenate([(known - 0.5)[None], ((img * 2 - 1) * known[..., None]).transpose(2, 0, 1)])[None]
|
||
yield {name: x.astype(np.float32)}
|
||
elif kind == "xfeat": # xfeat.rs: grey 0..1, shrink to fit, top-left, zero pad
|
||
_, _, H, W = [d if isinstance(d, int) else 1 for d in model.get_inputs()[0].shape]
|
||
for p in photos:
|
||
g = load(p).mean(2) # a display-rendered photograph is already the app's (R+G+B)/3 ^ 1/2.2
|
||
h, w = g.shape
|
||
s = min(W / w, H / h, 1.0)
|
||
if s < 1:
|
||
g = np.asarray(Image.fromarray(g).resize((round(w * s), round(h * s)), Image.BOX))
|
||
pad = np.zeros((H, W), np.float32)
|
||
pad[:g.shape[0], :g.shape[1]] = g
|
||
yield {name: pad[None, None]}
|
||
|
||
|
||
class Items(CalibrationDataReader):
|
||
def __init__(self, items):
|
||
self.it = iter(items)
|
||
|
||
def get_next(self):
|
||
return next(self.it, None)
|
||
|
||
|
||
def calibrate(path, items, cache):
|
||
"""Min/max ranges in chunks — every ORT calibrator holds all activations
|
||
until it folds them, and the others measurably degrade the result."""
|
||
cal = create_calibrator(Path(path), None, augmented_model_path=f"{path}.aug.onnx",
|
||
calibrate_method=CalibrationMethod.MinMax)
|
||
batch, n = [], 0
|
||
for item in items:
|
||
batch.append(item)
|
||
n += 1
|
||
if len(batch) == CHUNK:
|
||
cal.collect_data(Items(batch))
|
||
batch = []
|
||
if batch:
|
||
cal.collect_data(Items(batch))
|
||
save_tensors_data(cal.compute_data(), cache)
|
||
os.remove(f"{path}.aug.onnx")
|
||
return n
|
||
|
||
|
||
def downstream(m, start):
|
||
names, live = set(), set()
|
||
for n in m.graph.node:
|
||
if n.name == start or any(i in live for i in n.input):
|
||
names.add(n.name)
|
||
live.update(n.output)
|
||
return sorted(names)
|
||
|
||
|
||
def main():
|
||
args = sys.argv[1:]
|
||
ranges = None
|
||
if args[:1] == ["--ranges"]:
|
||
ranges, args = args[1], args[2:]
|
||
photo_dir = None
|
||
else:
|
||
photo_dir, args = args[0], args[1:]
|
||
names = args or [n for n in TABLE if TABLE[n]["feed"]]
|
||
photos = []
|
||
if photo_dir:
|
||
photos = sorted(p for e in ("jpg", "jpeg", "JPG", "JPEG", "png")
|
||
for p in glob.glob(os.path.join(photo_dir, "**", f"*.{e}"), recursive=True))[:PHOTOS]
|
||
if len(photos) < 50:
|
||
sys.exit(f"only {len(photos)} photographs under {photo_dir}; calibration wants hundreds")
|
||
print(f"==> calibrating on {len(photos)} photographs")
|
||
|
||
for stem in names:
|
||
spec = TABLE[stem]
|
||
src = MODELS / spec["dir"] / f"{stem}.onnx"
|
||
out = src.with_name(f"{stem}.{spec['form']}.onnx")
|
||
work = src.with_name(f"{stem}.quant-work.onnx")
|
||
print(f" {stem} -> {out.name}")
|
||
m = onnx.load(src)
|
||
if next(o.version for o in m.opset_import if o.domain in ("", "ai.onnx")) < 13:
|
||
m = version_converter.convert_version(m, 17) # per-channel QDQ needs 13
|
||
m.ir_version = 8
|
||
onnx.save(m, work)
|
||
if spec.get("rewrites"):
|
||
onnx.save(htp_graph.rewrite(str(work), spec["rewrites"]), work)
|
||
model = ort.InferenceSession(str(work), providers=["CPUExecutionProvider"])
|
||
cache = str(work) + ".ranges"
|
||
if spec["feed"] is None:
|
||
if not ranges:
|
||
sys.exit(f"{stem} is calibrated on mosaics, not photographs: pass --ranges")
|
||
cache = ranges
|
||
else:
|
||
n = calibrate(str(work), feeds(spec["feed"], photos, model), cache)
|
||
print(f" {n} calibration inputs")
|
||
act, wt = FORMS[spec["form"]]
|
||
# The config needs a reader only to exist; the ranges come from `cache`.
|
||
zeros = {i.name: np.zeros([d if isinstance(d, int) else 1 for d in i.shape], np.float32)
|
||
for i in model.get_inputs()}
|
||
cfg = get_qnn_qdq_config(str(work), Items([zeros]), activation_type=act, weight_type=wt, per_channel=True)
|
||
exclude = list(cfg.nodes_to_exclude or []) + spec.get("float_nodes", [])
|
||
if spec.get("float_from"):
|
||
exclude += downstream(onnx.load(work), spec["float_from"])
|
||
quantize_static(str(work), str(out), None, quant_format=cfg.quant_format,
|
||
op_types_to_quantize=cfg.op_types_to_quantize, per_channel=True,
|
||
activation_type=act, weight_type=wt, nodes_to_exclude=exclude,
|
||
calibrate_method=CalibrationMethod.MinMax, extra_options=cfg.extra_options,
|
||
calibration_cache_path=cache)
|
||
os.remove(work)
|
||
if cache != ranges:
|
||
os.remove(cache)
|
||
print(f" {out.name}: {out.stat().st_size // 1024} KB")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|