Skip to content

Commit 9fc4e81

Browse files
solderzzcsolderzzc
authored
Merge pull request #166 from SharpAI/feature/onnx-coreml-inference
Feature/onnx coreml inference
2 parents 59cba25 + 10c4cbf commit 9fc4e81

2 files changed

Lines changed: 179 additions & 5 deletions

File tree

‎skills/detection/yolo-detection-2026/scripts/env_config.py‎

Lines changed: 100 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -661,6 +661,72 @@ def __init__(self, *args, **kwargs):
661661
_log("coremltools not available, loading without compute_units")
662662
return YOLO(model_path)
663663

664+
# ── ONNX model download from HuggingFace ──────────────────────────
665+
666+
# Maps model base name → onnx-community HuggingFace repo
667+
_ONNX_HF_REPOS = {
668+
"yolo26n": "onnx-community/yolo26n-ONNX",
669+
"yolo26s": "onnx-community/yolo26s-ONNX",
670+
"yolo26m": "onnx-community/yolo26m-ONNX",
671+
"yolo26l": "onnx-community/yolo26l-ONNX",
672+
}
673+
674+
def _download_onnx_from_hf(self, model_name: str, dest_path: Path) -> bool:
675+
"""Download pre-built ONNX model from onnx-community on HuggingFace.
676+
677+
Uses urllib (no extra dependencies). Downloads to dest_path.
678+
Returns True on success, False on failure.
679+
"""
680+
repo = self._ONNX_HF_REPOS.get(model_name)
681+
if not repo:
682+
_log(f"No HuggingFace repo for {model_name}")
683+
return False
684+
685+
url = f"https://huggingface.co/{repo}/resolve/main/onnx/model.onnx"
686+
names_url = None # class names not available on HF, use bundled nano names
687+
688+
_log(f"Downloading {model_name}.onnx from {repo}...")
689+
try:
690+
import urllib.request
691+
import shutil
692+
693+
# Download ONNX model
694+
tmp_path = str(dest_path) + ".download"
695+
with urllib.request.urlopen(url) as resp, open(tmp_path, 'wb') as f:
696+
shutil.copyfileobj(resp, f)
697+
698+
# Rename to final path
699+
Path(tmp_path).rename(dest_path)
700+
size_mb = dest_path.stat().st_size / 1e6
701+
_log(f"Downloaded {model_name}.onnx ({size_mb:.1f} MB)")
702+
703+
# Create class names JSON if missing (COCO 80 — same for all YOLO models)
704+
names_path = Path(str(dest_path).replace('.onnx', '_names.json'))
705+
if not names_path.exists():
706+
# Try copying from nano (which is shipped in the repo)
707+
nano_names = dest_path.parent / "yolo26n_names.json"
708+
if nano_names.exists():
709+
shutil.copy2(str(nano_names), str(names_path))
710+
_log(f"Copied class names from yolo26n_names.json")
711+
else:
712+
# Generate default COCO names
713+
import json
714+
coco_names = {str(i): f"class_{i}" for i in range(80)}
715+
with open(str(names_path), 'w') as f:
716+
json.dump(coco_names, f)
717+
_log("Generated default class names")
718+
719+
return True
720+
except Exception as e:
721+
_log(f"HuggingFace download failed: {e}")
722+
# Clean up partial download
723+
for p in [str(dest_path) + ".download", str(dest_path)]:
724+
try:
725+
Path(p).unlink(missing_ok=True)
726+
except Exception:
727+
pass
728+
return False
729+
664730
def _load_onnx_coreml(self, onnx_path: str):
665731
"""Load ONNX model with CoreMLExecutionProvider for fast GPU/ANE inference.
666732
@@ -674,11 +740,27 @@ def _load_onnx_coreml(self, onnx_path: str):
674740
active = session.get_providers()
675741
_log(f"ONNX+CoreML session: {active}")
676742

677-
# Get YOLO class names from the .pt model (needed for detection output)
678-
from ultralytics import YOLO
679-
pt_path = onnx_path.replace('.onnx', '.pt')
680-
pt_model = YOLO(pt_path)
681-
class_names = pt_model.names # {0: 'person', 1: 'bicycle', ...}
743+
# Load class names from companion JSON (avoids torch/ultralytics dep)
744+
import json
745+
names_path = onnx_path.replace('.onnx', '_names.json')
746+
try:
747+
with open(names_path) as f:
748+
raw = json.load(f)
749+
# JSON keys are strings; convert to int-keyed dict
750+
class_names = {int(k): v for k, v in raw.items()}
751+
_log(f"Loaded {len(class_names)} class names from {Path(names_path).name}")
752+
except FileNotFoundError:
753+
# Fallback: try loading from .pt if JSON doesn't exist
754+
try:
755+
from ultralytics import YOLO
756+
pt_path = onnx_path.replace('.onnx', '.pt')
757+
pt_model = YOLO(pt_path)
758+
class_names = pt_model.names
759+
_log(f"Loaded class names from {Path(pt_path).name} (fallback)")
760+
except Exception:
761+
# Last resort: use COCO 80-class defaults
762+
_log("WARNING: No class names found, using generic labels")
763+
class_names = {i: f"class_{i}" for i in range(80)}
682764

683765
return _OnnxCoreMLModel(session, class_names)
684766

@@ -709,6 +791,19 @@ def load_optimized(self, model_name: str, use_optimized: bool = True):
709791
except Exception as e:
710792
_log(f"Failed to load cached model: {e}")
711793

794+
# Try downloading pre-built ONNX from HuggingFace (no torch needed)
795+
if self.export_format == "onnx" and self._download_onnx_from_hf(model_name, optimized_path):
796+
try:
797+
if self.backend == "mps":
798+
model = self._load_onnx_coreml(str(optimized_path))
799+
else:
800+
model = YOLO(str(optimized_path))
801+
self.load_ms = (time.perf_counter() - t0) * 1000
802+
_log(f"Loaded HuggingFace ONNX model ({self.load_ms:.0f}ms)")
803+
return model, self.export_format
804+
except Exception as e:
805+
_log(f"Failed to load HF-downloaded model: {e}")
806+
712807
# Try exporting then loading
713808
pt_model = YOLO(f"{model_name}.pt")
714809
exported = self.export_model(pt_model, model_name)

‎skills/lib/env_config.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -661,6 +661,72 @@ def __init__(self, *args, **kwargs):
661661
_log("coremltools not available, loading without compute_units")
662662
return YOLO(model_path)
663663

664+
# ── ONNX model download from HuggingFace ──────────────────────────
665+
666+
# Maps model base name → onnx-community HuggingFace repo
667+
_ONNX_HF_REPOS = {
668+
"yolo26n": "onnx-community/yolo26n-ONNX",
669+
"yolo26s": "onnx-community/yolo26s-ONNX",
670+
"yolo26m": "onnx-community/yolo26m-ONNX",
671+
"yolo26l": "onnx-community/yolo26l-ONNX",
672+
}
673+
674+
def _download_onnx_from_hf(self, model_name: str, dest_path: Path) -> bool:
675+
"""Download pre-built ONNX model from onnx-community on HuggingFace.
676+
677+
Uses urllib (no extra dependencies). Downloads to dest_path.
678+
Returns True on success, False on failure.
679+
"""
680+
repo = self._ONNX_HF_REPOS.get(model_name)
681+
if not repo:
682+
_log(f"No HuggingFace repo for {model_name}")
683+
return False
684+
685+
url = f"https://huggingface.co/{repo}/resolve/main/onnx/model.onnx"
686+
names_url = None # class names not available on HF, use bundled nano names
687+
688+
_log(f"Downloading {model_name}.onnx from {repo}...")
689+
try:
690+
import urllib.request
691+
import shutil
692+
693+
# Download ONNX model
694+
tmp_path = str(dest_path) + ".download"
695+
with urllib.request.urlopen(url) as resp, open(tmp_path, 'wb') as f:
696+
shutil.copyfileobj(resp, f)
697+
698+
# Rename to final path
699+
Path(tmp_path).rename(dest_path)
700+
size_mb = dest_path.stat().st_size / 1e6
701+
_log(f"Downloaded {model_name}.onnx ({size_mb:.1f} MB)")
702+
703+
# Create class names JSON if missing (COCO 80 — same for all YOLO models)
704+
names_path = Path(str(dest_path).replace('.onnx', '_names.json'))
705+
if not names_path.exists():
706+
# Try copying from nano (which is shipped in the repo)
707+
nano_names = dest_path.parent / "yolo26n_names.json"
708+
if nano_names.exists():
709+
shutil.copy2(str(nano_names), str(names_path))
710+
_log(f"Copied class names from yolo26n_names.json")
711+
else:
712+
# Generate default COCO names
713+
import json
714+
coco_names = {str(i): f"class_{i}" for i in range(80)}
715+
with open(str(names_path), 'w') as f:
716+
json.dump(coco_names, f)
717+
_log("Generated default class names")
718+
719+
return True
720+
except Exception as e:
721+
_log(f"HuggingFace download failed: {e}")
722+
# Clean up partial download
723+
for p in [str(dest_path) + ".download", str(dest_path)]:
724+
try:
725+
Path(p).unlink(missing_ok=True)
726+
except Exception:
727+
pass
728+
return False
729+
664730
def _load_onnx_coreml(self, onnx_path: str):
665731
"""Load ONNX model with CoreMLExecutionProvider for fast GPU/ANE inference.
666732
@@ -725,6 +791,19 @@ def load_optimized(self, model_name: str, use_optimized: bool = True):
725791
except Exception as e:
726792
_log(f"Failed to load cached model: {e}")
727793

794+
# Try downloading pre-built ONNX from HuggingFace (no torch needed)
795+
if self.export_format == "onnx" and self._download_onnx_from_hf(model_name, optimized_path):
796+
try:
797+
if self.backend == "mps":
798+
model = self._load_onnx_coreml(str(optimized_path))
799+
else:
800+
model = YOLO(str(optimized_path))
801+
self.load_ms = (time.perf_counter() - t0) * 1000
802+
_log(f"Loaded HuggingFace ONNX model ({self.load_ms:.0f}ms)")
803+
return model, self.export_format
804+
except Exception as e:
805+
_log(f"Failed to load HF-downloaded model: {e}")
806+
728807
# Try exporting then loading
729808
pt_model = YOLO(f"{model_name}.pt")
730809
exported = self.export_model(pt_model, model_name)

0 commit comments

Comments
 (0)