File size: 1,675 Bytes
18b55d3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
from __future__ import annotations
from pathlib import Path
from huggingface_hub import hf_hub_download
from ultralytics import YOLO
from tqdm.auto import tqdm
REPOBINGSU = "Bingsu/adetailer"
REPOGUON = "guon/hand-eyes"
model_info = [
(REPOBINGSU, "deepfashion2_yolov8s-seg.pt"),
(REPOBINGSU, "face_yolov8m.pt"),
(REPOBINGSU, "face_yolov8n.pt"),
(REPOBINGSU, "face_yolov8n_v2.pt"),
(REPOBINGSU, "face_yolov8s.pt"),
(REPOBINGSU, "face_yolov9c.pt"),
(REPOBINGSU, "hand_yolov8n.pt"),
(REPOBINGSU, "hand_yolov8s.pt"),
(REPOBINGSU, "hand_yolov9c.pt"),
(REPOBINGSU, "person_yolov8m-seg.pt"),
(REPOBINGSU, "person_yolov8n-seg.pt"),
(REPOBINGSU, "person_yolov8s-seg.pt"),
(REPOGUON, "Eyes.pt"),
(REPOGUON, "PitEyeDetailer-v1-seg.pt"),
(REPOGUON, "PitHandDetailer-v1-seg.pt"),
(REPOGUON, "female_breast-v4.2.pt"),
(REPOGUON, "full_eyes_detect_v1.pt"),
(REPOGUON, "lips_v1.pt"),
]
def download(model: str, repo: str) -> str:
return hf_hub_download(
repo, model, local_dir="downloads", local_dir_use_symlinks=False
)
def export(path: str) -> str:
model = YOLO(path)
return model.export(format="onnx", dynamic=True, simplify=True)
def main():
save = Path("models")
save.mkdir(exist_ok=True)
for repo, model in tqdm(model_info):
path = download(model, repo)
onnx = export(path)
target_dir = save / Path(model).stem
(target_dir / "1").mkdir(exist_ok=True, parents=True)
(target_dir / "config.pbtxt").touch()
target = target_dir / "1" / "model.onnx"
Path(onnx).rename(target)
if __name__ == "__main__":
main()
|