camstack-models / source /build-camstack-models.py
apocaliss92's picture
docs: notices, licences and source/ after the plate detector swap (D718)
90c9878 verified
Raw History Blame
25 kB
#!/usr/bin/env python3
"""
Rebuild every Ultralytics-derived (AGPL-3.0) model the CamStack mirror hosts.
This is the Corresponding Source for the AGPL-3.0 weights of
https://huggingface.co/camstack/camstack-models in
objectDetection/yolo26 yolo26{n,s,m,l,x} @640, yolo26{n,s,m,l} @320 and @256
objectDetection/yolov9 yolov9{t,s,m,c} @320 and @640
segmentation/yolo26-seg yolo26{n,s,m}-seg @640
Every hosted file is one output of ONE recipe in `RECIPES` below, and the recipe
holds exactly the settings that file records about itself: the Ultralytics
`export()` arguments (stored in each ONNX / CoreML file's metadata), the OpenVINO
weight compression and the NNCF calibration subset size (stored in each IR's
`rt_info`). `--check-hosted` proves the mapping is complete against the live
mirror: every hosted weight file must be an output of some recipe.
Usage:
python build-camstack-models.py --list # every recipe
python build-camstack-models.py --check-hosted # recipes == mirror?
python build-camstack-models.py --variant camstack-yolo26n-320
python build-camstack-models.py --family yolov9 # one mirror folder
python build-camstack-models.py --all # the whole mirror
python build-camstack-models.py --all --formats onnx,openvino --dry-run
# Something not in the mirror (same pipeline, your own settings):
python build-camstack-models.py --model yolo26n --imgsz 416 \
--openvino fp16,int8 --calib-samples 128
Outputs land in `--out` (default ./exported-models) under the MIRROR's paths,
e.g. exported-models/objectDetection/yolo26/openvino/camstack-yolo26n-320-int8.xml.
Formats:
onnx Ultralytics export(format="onnx") from the upstream .pt
coreml Ultralytics export(format="coreml") from the upstream .pt (macOS only)
openvino convert_lib.export_openvino: ov.convert_model(ONNX) + ov.save_model
(compress_to_fp16 per precision) — never the Ultralytics OV exporter
int8 convert_lib.quantize_openvino_int8: NNCF post-training quantisation
of the ONNX, calibrated on the first N images (sorted by file name)
of --calib-dir (default $CAMSTACK_CALIB_DIR or ./calibration-images)
INT8 calibration: the hosted INT8 IRs record subset_size 128 (YOLO26) and 300
(YOLOv9), drawn from COCO; an IR does not record WHICH images. The documented
set is COCO val2017: unpack http://images.cocodataset.org/zips/val2017.zip and
point --calib-dir at it — the recipe's N is taken in file-name order. Images are letterboxed to
the recipe's input size with grey (114) padding and scaled to [0, 1], exactly
as the CamStack runtime feeds the model. A different calibration set yields an
equivalent but not bit-identical IR (NNCF picks ranges from the data).
Toolchain: each recipe names the Ultralytics / torch / coremltools / OpenVINO
versions its hosted files record. Other versions produce the same model with a
possibly different graph layout; the script warns when they differ.
Environment: `pip install -r requirements.txt ultralytics` (requirements.txt is
convert_lib's). convert_lib.py is found next to this script (the mirror's
`source/` folder) or in packages/addon-model-studio/python/ (the repository).
"""
from __future__ import annotations
import argparse
import json
import os
import shutil
import subprocess
import sys
import urllib.request
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional
_HERE = Path(__file__).resolve().parent
_CONVERT_LIB_DIRS = (_HERE, _HERE.parent / 'packages' / 'addon-model-studio' / 'python')
_CONVERT_LIB_DIR = next((d for d in _CONVERT_LIB_DIRS if (d / 'convert_lib.py').is_file()), None)
if _CONVERT_LIB_DIR is None:
sys.exit('convert_lib.py not found next to this script or in packages/addon-model-studio/python/')
sys.path.insert(0, str(_CONVERT_LIB_DIR))
HF_REPO = 'camstack/camstack-models'
DEFAULT_CALIB_DIR = Path(os.environ.get('CAMSTACK_CALIB_DIR', 'calibration-images'))
# --- Recipe model ------------------------------------------------------------
@dataclass(frozen=True)
class OnnxExport:
"""Ultralytics export(format="onnx") arguments, as the hosted ONNX records them."""
dynamic: bool
opset: int
@dataclass(frozen=True)
class CoremlExport:
"""Ultralytics export(format="coreml") arguments, as the hosted package records them."""
dynamic: bool
batch: int
half: bool
@dataclass(frozen=True)
class OpenvinoOutput:
"""One hosted IR: file suffix after the model name, and how it is produced.
`precision` is `fp32` (compress_to_fp16=False), `fp16` (compress_to_fp16=True)
or `int8` (NNCF on the ONNX, `calib_samples` images).
"""
suffix: str
precision: str
@dataclass(frozen=True)
class Toolchain:
ultralytics: str
torch: str
coremltools: str
openvino: str
@dataclass(frozen=True)
class Recipe:
name: str # hosted file stem, e.g. camstack-yolo26n-320
folder: str # mirror folder, e.g. objectDetection/yolo26
imgsz: int
toolchain: Toolchain
# The PyTorch checkpoint: an Ultralytics asset id (`yolo26n` -> yolo26n.pt,
# auto-downloaded by Ultralytics).
checkpoint: str
onnx: Optional[OnnxExport]
coreml: Optional[CoremlExport]
openvino: tuple[OpenvinoOutput, ...] = ()
calib_samples: int = 0
note: str = ''
def outputs(self) -> list[str]:
"""Mirror paths this recipe produces (an .mlpackage is a directory)."""
out: list[str] = []
if self.onnx is not None:
out.append(f'{self.folder}/onnx/{self.name}.onnx')
if self.coreml is not None:
out.append(f'{self.folder}/coreml/{self.name}.mlpackage/')
for ov in self.openvino:
stem = f'{self.folder}/openvino/{self.name}{ov.suffix}'
out.extend([f'{stem}.xml', f'{stem}.bin'])
return out
# --- The hosted set ------------------------------------------------------------
# Toolchains as recorded by the hosted files themselves (ONNX/CoreML metadata
# `version`, CoreML `com.github.apple.coremltools.*` and `torch==`, IR rt_info
# `Runtime_version`).
TC_2604 = Toolchain('8.4.24', '2.10.0', '9.0', '2026.2.1')
TC_2607_REDUCED = Toolchain('8.4.92', 'onnx 2.13.0, coreml 2.5.1', '8.3.0', '2026.2.1')
TC_2607_YOLOV9 = Toolchain('8.4.98', '2.13.0', '9.0', '2026.2.1')
FP16_IR = OpenvinoOutput('', 'fp16')
def _yolo26_full(tier: str) -> Recipe:
"""yolo26<tier> @640: static ONNX, static fp32-compute CoreML, IR fp16 twice + INT8.
The unsuffixed IR and `-fp16` IR are the same compressed-to-FP16 weights;
the catalog addresses both names.
"""
return Recipe(
name=f'camstack-yolo26{tier}', folder='objectDetection/yolo26', imgsz=640,
toolchain=TC_2604, checkpoint=f'yolo26{tier}',
onnx=OnnxExport(dynamic=False, opset=17),
coreml=CoremlExport(dynamic=False, batch=1, half=False),
openvino=(FP16_IR, OpenvinoOutput('-fp16', 'fp16'), OpenvinoOutput('-int8', 'int8')),
calib_samples=128,
)
def _yolo26_reduced(tier: str, res: int) -> Recipe:
"""yolo26<tier>-<res>: dynamic-axes ONNX, batch-flexible (1..16) FP16 CoreML, IR fp32 + INT8."""
return Recipe(
name=f'camstack-yolo26{tier}-{res}', folder='objectDetection/yolo26', imgsz=res,
toolchain=TC_2607_REDUCED, checkpoint=f'yolo26{tier}',
onnx=OnnxExport(dynamic=True, opset=17),
coreml=CoremlExport(dynamic=True, batch=16, half=True),
openvino=(OpenvinoOutput('', 'fp32'), OpenvinoOutput('-int8', 'int8')),
calib_samples=128,
)
def _yolov9(tier: str, res: int) -> Recipe:
"""yolov9<tier>-<res>: static opset-13 ONNX, static CoreML, IR fp16 + INT8."""
return Recipe(
name=f'camstack-yolov9{tier}-{res}', folder='objectDetection/yolov9', imgsz=res,
toolchain=TC_2607_YOLOV9, checkpoint=f'yolov9{tier}',
onnx=OnnxExport(dynamic=False, opset=13),
coreml=CoremlExport(dynamic=False, batch=1, half=False),
openvino=(OpenvinoOutput('-fp16', 'fp16'), OpenvinoOutput('-int8', 'int8')),
calib_samples=300,
)
def _yolo26_seg(tier: str) -> Recipe:
return Recipe(
name=f'camstack-yolo26{tier}-seg', folder='segmentation/yolo26-seg', imgsz=640,
toolchain=TC_2604, checkpoint=f'yolo26{tier}-seg',
onnx=OnnxExport(dynamic=False, opset=17),
coreml=CoremlExport(dynamic=False, batch=1, half=False),
openvino=(FP16_IR,),
)
RECIPES: tuple[Recipe, ...] = (
*(_yolo26_full(t) for t in 'nsmlx'),
*(_yolo26_reduced(t, r) for t in 'nsml' for r in (320, 256)),
*(_yolov9(t, r) for t in ('t', 's', 'm', 'c') for r in (320, 640)),
*(_yolo26_seg(t) for t in 'nsm'),
)
FAMILIES = {
'yolo26': 'objectDetection/yolo26',
'yolov9': 'objectDetection/yolov9',
'yolo26-seg': 'segmentation/yolo26-seg',
}
ALL_FORMATS = ('onnx', 'coreml', 'openvino')
# --- Helpers ----------------------------------------------------------------
def size_mb(path: Path) -> float:
if path.is_file():
return path.stat().st_size / 1_048_576
return sum(f.stat().st_size for f in path.rglob('*') if f.is_file()) / 1_048_576
def warn_toolchain(recipe: Recipe) -> None:
try:
import ultralytics
except ImportError:
return
want = recipe.toolchain.ultralytics.split(' ')[0]
if ultralytics.__version__ != want:
print(f' ! the hosted {recipe.name} files were made with ultralytics {want}; '
f'this is {ultralytics.__version__} (same weights, the graph may differ)')
def place(src: Path, dest: Path) -> Path:
dest.parent.mkdir(parents=True, exist_ok=True)
if dest.exists():
shutil.rmtree(dest) if dest.is_dir() else dest.unlink()
(shutil.copytree if src.is_dir() else shutil.copy2)(src, dest)
print(f' ✓ {dest} ({size_mb(dest):.1f} MB)')
return dest
def checkpoint_path(recipe: Recipe) -> str:
return f'{recipe.checkpoint}.pt' # Ultralytics downloads its own assets
# --- Child processes ----------------------------------------------------------
# Every backend runs in its OWN interpreter. On macOS two of them in one process
# abort it: an ONNX export followed by a CoreML export ("libc++abi: mutex lock
# failed"), and OpenVINO beside torch / onnx / tensorflow (each initialises its
# own protobuf). A child per step keeps a crash to that step, reported as a
# failure, and the parent imports none of them.
_CHILD_MARKER = '@@CAMSTACK_CHILD@@'
_ULTRALYTICS_CHILD = f"""
import json, sys
from ultralytics import YOLO
args = json.loads(sys.argv[1])
path = YOLO(args.pop('checkpoint')).export(**args)
print('\\n{_CHILD_MARKER}' + str(path), flush=True)
"""
# Only openvino (and, for INT8, nncf + cv2) is imported here: no torch, no onnx,
# no tensorflow. The conversion is convert_lib's, the same code Model Studio runs.
_OPENVINO_CHILD = f"""
import json, sys
args = json.loads(sys.argv[1])
sys.path.insert(0, args['convert_lib_dir'])
from convert_lib import export_openvino, quantize_openvino_int8, validate_openvino_ir
if args['precision'] == 'int8':
quantize_openvino_int8(args['onnx'], args['xml'], args['calib_dir'], args['imgsz'],
max_samples=args['calib_samples'])
else:
export_openvino(args['onnx'], args['xml'], fp16=args['precision'] == 'fp16')
validate_openvino_ir(args['xml'])
print('\\n{_CHILD_MARKER}' + args['xml'], flush=True)
"""
_COREML_CHECK_CHILD = f"""
import json, sys
import coremltools as ct
import numpy as np
args = json.loads(sys.argv[1])
model = ct.models.MLModel(args['package'])
name = model.get_spec().description.input[0].name
for batch in (1, 4):
out = model.predict({{name: np.random.rand(batch, 3, args['imgsz'], args['imgsz']).astype(np.float32)}})
print(f' batch={{batch}} -> {{getattr(next(iter(out.values())), "shape", "?")}}')
print('\\n{_CHILD_MARKER}ok', flush=True)
"""
def run_child(code: str, payload: dict[str, object], what: str) -> str:
"""Run one backend step in a fresh interpreter; return what it reported.
The step counts as done only when the child printed its marker line; a
child that crashes (a signal, an abort) before that is a failed step and is
never retried.
"""
proc = subprocess.run(
[sys.executable, '-c', code, json.dumps(payload)],
stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
)
marked = [ln[len(_CHILD_MARKER):] for ln in proc.stdout.splitlines() if ln.startswith(_CHILD_MARKER)]
if not marked:
sys.stdout.write('\n'.join(proc.stdout.splitlines()[-25:]) + '\n')
raise RuntimeError(f'{what} failed (child exit {proc.returncode})')
if proc.returncode != 0:
print(f' ! {what}: the child exited {proc.returncode} after finishing')
return marked[-1]
def ultralytics_export(checkpoint: str, **kwargs: object) -> Path:
path = Path(run_child(_ULTRALYTICS_CHILD, {'checkpoint': checkpoint, **kwargs},
f'ultralytics export {kwargs} of {checkpoint}'))
if not path.exists():
raise RuntimeError(f'ultralytics export reported {path}, which does not exist')
return path
# --- Format steps -------------------------------------------------------------
def build_onnx(recipe: Recipe, out: Path) -> Path:
dest = out / recipe.folder / 'onnx' / f'{recipe.name}.onnx'
assert recipe.onnx is not None
print(f' → ONNX (imgsz={recipe.imgsz}, dynamic={recipe.onnx.dynamic}, opset={recipe.onnx.opset})')
src = ultralytics_export(
checkpoint_path(recipe), format='onnx', imgsz=recipe.imgsz,
dynamic=recipe.onnx.dynamic, simplify=True, opset=recipe.onnx.opset,
nms=False, batch=1, half=False,
)
return place(src, dest)
def build_coreml(recipe: Recipe, out: Path, verify: bool) -> Path:
assert recipe.coreml is not None
c = recipe.coreml
print(f' → CoreML (imgsz={recipe.imgsz}, dynamic={c.dynamic}, batch={c.batch}, half={c.half})')
src = ultralytics_export(
checkpoint_path(recipe), format='coreml', imgsz=recipe.imgsz,
dynamic=c.dynamic, batch=c.batch, half=c.half, nms=False,
)
dest = place(src, out / recipe.folder / 'coreml' / f'{recipe.name}.mlpackage')
if verify and c.dynamic:
verify_coreml_batch(dest, recipe.imgsz)
return dest
def build_openvino(recipe: Recipe, onnx_path: Path, out: Path, calib_dir: Path,
calib_samples: int) -> list[str]:
"""Every IR of the recipe, each in its own child. Returns the outputs it skipped."""
skipped: list[str] = []
for ov in recipe.openvino:
xml = out / recipe.folder / 'openvino' / f'{recipe.name}{ov.suffix}.xml'
if ov.precision == 'int8' and not (calib_dir.is_dir() and any(calib_dir.iterdir())):
print(f' ⊘ {xml.name}: no calibration images in {calib_dir} (--calib-dir)')
skipped.append(str(xml))
continue
detail = f'NNCF, {calib_samples} images from {calib_dir}' if ov.precision == 'int8' else 'ov.convert_model'
print(f' → OpenVINO {ov.precision.upper()} {xml.name} ({detail})')
xml.parent.mkdir(parents=True, exist_ok=True)
run_child(_OPENVINO_CHILD, {
'convert_lib_dir': str(_CONVERT_LIB_DIR), 'onnx': str(onnx_path.resolve()),
'xml': str(xml.resolve()), 'precision': ov.precision, 'imgsz': recipe.imgsz,
'calib_dir': str(calib_dir.resolve()), 'calib_samples': calib_samples,
}, f'OpenVINO {ov.precision} {xml.name}')
print(f' ✓ {xml} ({size_mb(xml.with_suffix(".bin")):.1f} MB)')
return skipped
def verify_coreml_batch(mlpackage: Path, imgsz: int) -> None:
"""Predict with N=1 then N=4 on a flexible-batch package, in a throwaway child.
Opt-in (--verify-coreml): CoreML's runtime has aborted the whole process on
one of these packages (Espresso gather_nd shape computation) on a Mac whose
toolchain differs from the recorded one. A crash here is a failed CHECK, not
a failed build — the package on disk is the export's output either way.
"""
try:
run_child(_COREML_CHECK_CHILD, {'package': str(mlpackage.resolve()), 'imgsz': imgsz},
f'CoreML batch check of {mlpackage.name}')
print(' ✓ flexible batch predicts at N=1 and N=4')
except RuntimeError as exc:
print(f' ! {exc}; the package is kept, check it on the target device')
# --- Orchestration ------------------------------------------------------------
@dataclass
class Report:
failed: list[str] = field(default_factory=list)
skipped: list[str] = field(default_factory=list)
def run_step(report: Report, label: str, step: Callable[[], object]) -> object:
try:
return step()
except Exception as exc: # one failed format must not stop the others
print(f' ✗ {label} failed: {exc}')
report.failed.append(label)
return None
def build(recipe: Recipe, formats: set[str], out: Path, calib_dir: Path,
calib_samples: Optional[int], verify_coreml: bool, report: Report) -> None:
print('\n' + '=' * 70 + f'\n{recipe.name} ({recipe.folder}, {recipe.imgsz}×{recipe.imgsz})\n' + '=' * 70)
if recipe.note:
print(f' {recipe.note}')
warn_toolchain(recipe)
onnx_path: Optional[Path] = None
if formats & {'onnx', 'openvino'} and recipe.onnx:
result = run_step(report, f'{recipe.name} onnx', lambda: build_onnx(recipe, out))
onnx_path = result if isinstance(result, Path) else None
if 'coreml' in formats and recipe.coreml is not None:
if sys.platform != 'darwin':
print(' ⊘ CoreML needs macOS')
report.skipped.append(f'{recipe.name} coreml')
else:
run_step(report, f'{recipe.name} coreml', lambda: build_coreml(recipe, out, verify_coreml))
if 'openvino' in formats and recipe.openvino and onnx_path is not None:
samples = calib_samples or recipe.calib_samples or 300
skipped = run_step(report, f'{recipe.name} openvino',
lambda: build_openvino(recipe, onnx_path, out, calib_dir, samples))
report.skipped.extend(skipped if isinstance(skipped, list) else [])
def adhoc_recipe(args: argparse.Namespace) -> Recipe:
"""A recipe from the command line, for a size/resolution the mirror does not host."""
precisions = [p for p in (args.openvino or 'fp16,int8').split(',') if p]
bad = [p for p in precisions if p not in ('fp32', 'fp16', 'int8')]
if bad:
sys.exit(f'--openvino accepts fp32, fp16, int8; got {bad}')
name = args.name or f'camstack-{args.model}-{args.imgsz}'
return Recipe(
name=name, folder=args.folder, imgsz=args.imgsz,
toolchain=Toolchain('any', 'any', 'any', 'any'), checkpoint=args.model,
onnx=OnnxExport(dynamic=args.dynamic, opset=args.opset),
coreml=CoremlExport(dynamic=args.dynamic, batch=16 if args.dynamic else 1, half=args.half),
openvino=tuple(OpenvinoOutput(f'-{p}', p) for p in precisions),
calib_samples=args.calib_samples or 300,
)
def hosted_weight_paths() -> list[str]:
url = f'https://huggingface.co/api/models/{HF_REPO}/tree/main?recursive=true'
with urllib.request.urlopen(url, timeout=60) as resp:
tree = json.load(resp)
prefixes = tuple(f'{f}/' for f in FAMILIES.values())
return sorted(
e['path'] for e in tree
if e['type'] == 'file' and e['path'].startswith(prefixes) and not e['path'].endswith('README.md')
)
def check_hosted() -> int:
"""Every hosted AGPL weight file is an output of a recipe, and vice versa."""
hosted = hosted_weight_paths()
outputs = [o for r in RECIPES for o in r.outputs()]
def covered(path: str) -> bool:
return any(path == o or (o.endswith('/') and path.startswith(o)) for o in outputs)
def present(output: str) -> bool:
return any(h == output or (output.endswith('/') and h.startswith(output)) for h in hosted)
orphans = [h for h in hosted if not covered(h)]
unhosted = [o for o in outputs if not present(o)]
print(f'{len(hosted)} hosted files under {", ".join(FAMILIES.values())}; '
f'{len(RECIPES)} recipes, {len(outputs)} outputs')
for h in orphans:
print(f' NOT REPRODUCED by any recipe: {h}')
for o in unhosted:
print(f' recipe output not hosted: {o}')
if orphans or unhosted:
return 1
print('OK: every hosted file is a recipe output, and every recipe output is hosted.')
return 0
def list_recipes() -> None:
for r in RECIPES:
ov = ','.join(f'{o.precision}{"→" + o.suffix if o.suffix else ""}' for o in r.openvino)
onnx = f'{"dynamic" if r.onnx and r.onnx.dynamic else "static"} opset{r.onnx.opset if r.onnx else ""}'
cm = (f'{"flex≤" + str(r.coreml.batch) if r.coreml.dynamic else "static"}'
f'{" fp16" if r.coreml.half else ""}') if r.coreml else '-'
int8 = f' int8 N={r.calib_samples}' if any(o.precision == 'int8' for o in r.openvino) else ''
print(f'{r.name:30s} {r.folder:28s} {r.imgsz:4d} onnx[{onnx}] coreml[{cm}] '
f'openvino[{ov}]{int8} '
f'ultralytics {r.toolchain.ultralytics} from {r.checkpoint}.pt')
def parse_args(argv: list[str]) -> argparse.Namespace:
p = argparse.ArgumentParser(description=__doc__.split('\n\n')[0],
formatter_class=argparse.RawDescriptionHelpFormatter)
sel = p.add_argument_group('what to build')
sel.add_argument('--all', action='store_true', help='every hosted recipe')
sel.add_argument('--family', action='append', choices=sorted(FAMILIES), help='one mirror folder')
sel.add_argument('--variant', action='append', help='a hosted file stem, e.g. camstack-yolov9m-320')
sel.add_argument('--list', action='store_true', help='print the recipes and exit')
sel.add_argument('--check-hosted', action='store_true',
help='compare the recipes with the live mirror and exit (non-zero on a gap)')
sel.add_argument('--dry-run', action='store_true', help='print the files that would be built')
p.add_argument('--formats', default=','.join(ALL_FORMATS), help='subset of onnx,coreml,openvino')
p.add_argument('--out', type=Path, default=Path('exported-models'))
p.add_argument('--calib-dir', type=Path, default=DEFAULT_CALIB_DIR)
p.add_argument('--verify-coreml', action='store_true',
help='predict N=1 and N=4 on each flexible-batch CoreML package (in a child process)')
p.add_argument('--calib-samples', type=int, help="override the recipe's INT8 calibration image count")
adhoc = p.add_argument_group('ad-hoc build (not a hosted file)')
adhoc.add_argument('--model', help='Ultralytics id, e.g. yolo26n, yolov9m, yolo26s-seg')
adhoc.add_argument('--imgsz', type=int, default=640)
adhoc.add_argument('--name', help='output stem (default camstack-<model>-<imgsz>)')
adhoc.add_argument('--folder', default='custom')
adhoc.add_argument('--dynamic', action='store_true', help='dynamic axes (ONNX) / flexible batch (CoreML)')
adhoc.add_argument('--opset', type=int, default=17)
adhoc.add_argument('--half', action='store_true', help='FP16 CoreML')
adhoc.add_argument('--openvino', help='IR precisions, comma-separated: fp32,fp16,int8')
return p.parse_args(argv)
def select(args: argparse.Namespace) -> list[Recipe]:
if args.model:
return [adhoc_recipe(args)]
by_name = {r.name: r for r in RECIPES}
unknown = [v for v in (args.variant or []) if v not in by_name]
if unknown:
sys.exit(f'unknown --variant {unknown}; see --list')
folders = {FAMILIES[f] for f in (args.family or [])}
return [r for r in RECIPES
if args.all or r.folder in folders or r.name in (args.variant or [])]
def main(argv: list[str]) -> int:
args = parse_args(argv)
if args.list:
list_recipes()
return 0
if args.check_hosted:
return check_hosted()
recipes = select(args)
if not recipes:
print('Nothing selected: pass --all, --family, --variant or --model (see --help).')
return 2
formats = {f for f in args.formats.split(',') if f}
if formats - set(ALL_FORMATS):
sys.exit(f'--formats accepts {",".join(ALL_FORMATS)}')
if args.dry_run:
for r in recipes:
for o in r.outputs():
print(args.out / o)
return 0
report = Report()
for r in recipes:
build(r, formats, args.out, args.calib_dir, args.calib_samples,
args.verify_coreml, report)
print('\nDone.' + (f' Failed: {report.failed}.' if report.failed else '')
+ (f' Skipped: {report.skipped}.' if report.skipped else ''))
return 1 if report.failed else 0
if __name__ == '__main__':
sys.exit(main(sys.argv[1:]))